Bug 2003370 - Use and react to Gecko preferences stored values

This patch adds functionality to use the new PreviousGeckoPrefState.

It adds:

* `set_gecko_prefs_original_values` for external handling of setting prefs back to original values
* Mechanisms to return to a previous states when:
  * `on_experiment_updated`
    * Certain situations and as determined in will_pref_experiment_change
  * `on_experiment_ended`
  * `on_opt_out`

Author: ohall-m

components/nimbus/android/src/test/java/org/mozilla/experiments/nimbus/NimbusTests.kt

@@ -43,6 +43,7 @@ import org.mozilla.experiments.nimbus.internal.GeckoPrefHandler
 import org.mozilla.experiments.nimbus.internal.GeckoPrefState
 import org.mozilla.experiments.nimbus.internal.JsonObject
 import org.mozilla.experiments.nimbus.internal.NimbusException
+import org.mozilla.experiments.nimbus.internal.OriginalGeckoPref
 import org.mozilla.experiments.nimbus.internal.PrefBranch
 import org.mozilla.experiments.nimbus.internal.PrefEnrollmentData
 import org.mozilla.experiments.nimbus.internal.PrefUnenrollReason
@@ -860,6 +861,7 @@ class NimbusTests {
             ),
         ),
         var setValues: List<GeckoPrefState>? = null,
+        var originalGeckoPrefValues: List<OriginalGeckoPref>? = null,
     ) : GeckoPrefHandler {
         override fun getPrefsWithState(): Map<String, Map<String, GeckoPrefState>> {
             return internalMap
@@ -868,6 +870,10 @@ class NimbusTests {
         override fun setGeckoPrefsState(newPrefsState: List<GeckoPrefState>) {
             setValues = newPrefsState
         }
+
+        override fun setGeckoPrefsOriginalValues(originalGeckoPrefs: List<OriginalGeckoPref>) {
+            originalGeckoPrefValues = originalGeckoPrefs
+        }
     }
 
     @Test
@@ -889,6 +895,21 @@ class NimbusTests {
         assertEquals("42", handler.setValues?.get(0)?.enrollmentValue?.prefValue)
     }
 
+    @Test
+    fun `GeckoPrefHandler setGeckoPrefsOriginalValues function`() {
+      val handler = TestGeckoPrefHandler()
+      val originalValues = listOf(
+          OriginalGeckoPref(
+              pref = "pref.number",
+              branch = PrefBranch.DEFAULT,
+              value = "1",
+          ),
+      )
+      handler.setGeckoPrefsOriginalValues(originalValues)
+      assertEquals(1, handler.originalGeckoPrefValues?.size)
+      assertEquals("pref.number", handler.originalGeckoPrefValues?.get(0)?.pref)
+    }
+
     @Test
     fun `unenroll for gecko pref functions`() {
         val handler = TestGeckoPrefHandler()

components/nimbus/src/enrollment.rs

@@ -2,7 +2,7 @@
 // License, v. 2.0. If a copy of the MPL was not distributed with this
 // file, You can obtain one at https://mozilla.org/MPL/2.0/.
 #[cfg(feature = "stateful")]
-use crate::stateful::gecko_prefs::{OriginalGeckoPref, PrefUnenrollReason};
+use crate::stateful::gecko_prefs::{GeckoPrefStore, OriginalGeckoPref, PrefUnenrollReason};
 use crate::{
     defaults::Defaults,
     error::{debug, warn, NimbusError, Result},
@@ -11,6 +11,8 @@ use crate::{
     SLUG_REPLACEMENT_PATTERN,
 };
 use serde_derive::*;
+#[cfg(feature = "stateful")]
+use std::sync::Arc;
 use std::{
     collections::{HashMap, HashSet},
     fmt::{Display, Formatter, Result as FmtResult},
@@ -145,6 +147,24 @@ pub struct PreviousGeckoPrefState {
     pub variable: String,
 }
 
+#[cfg(feature = "stateful")]
+impl PreviousGeckoPrefState {
+    pub(crate) fn on_revert_to_prev_gecko_pref_states(
+        prev_gecko_pref_states: &[Self],
+        gecko_pref_store: &Option<Arc<GeckoPrefStore>>,
+    ) {
+        if let Some(store) = gecko_pref_store {
+            let original_values: Vec<_> = prev_gecko_pref_states
+                .iter()
+                .map(|state| state.original_value.clone())
+                .collect();
+            store
+                .handler
+                .set_gecko_prefs_original_values(original_values);
+        }
+    }
+}
+
 // Every experiment has an ExperimentEnrollment, even when we aren't enrolled.
 // ⚠️ Attention : Changes to this type should be accompanied by a new test  ⚠️
 // ⚠️ in `src/stateful/tests/test_enrollment_bw_compat.rs` below, and may require a DB migration. ⚠️
@@ -232,6 +252,7 @@ impl ExperimentEnrollment {
         available_randomization_units: &AvailableRandomizationUnits,
         updated_experiment: &Experiment,
         targeting_helper: &NimbusTargetingHelper,
+        #[cfg(feature = "stateful")] gecko_pref_store: &Option<Arc<GeckoPrefStore>>,
         out_enrollment_events: &mut Vec<EnrollmentChangeEvent>,
     ) -> Result<Self> {
         Ok(match &self.status {
@@ -265,12 +286,34 @@ impl ExperimentEnrollment {
                         "Existing experiment enrollment '{}' is now disqualified (global opt-out)",
                         &self.slug
                     );
+                    #[cfg(feature = "stateful")]
+                    if let EnrollmentStatus::Enrolled {
+                        prev_gecko_pref_states: Some(prev_gecko_pref_states),
+                        ..
+                    } = &self.status
+                    {
+                        PreviousGeckoPrefState::on_revert_to_prev_gecko_pref_states(
+                            prev_gecko_pref_states,
+                            gecko_pref_store,
+                        );
+                    }
                     let updated_enrollment =
                         self.disqualify_from_enrolled(DisqualifiedReason::OptOut);
                     out_enrollment_events.push(updated_enrollment.get_change_event());
                     updated_enrollment
                 } else if !updated_experiment.has_branch(branch) {
                     // The branch we were in disappeared!
+                    #[cfg(feature = "stateful")]
+                    if let EnrollmentStatus::Enrolled {
+                        prev_gecko_pref_states: Some(prev_gecko_pref_states),
+                        ..
+                    } = &self.status
+                    {
+                        PreviousGeckoPrefState::on_revert_to_prev_gecko_pref_states(
+                            prev_gecko_pref_states,
+                            gecko_pref_store,
+                        );
+                    }
                     let updated_enrollment =
                         self.disqualify_from_enrolled(DisqualifiedReason::Error);
                     out_enrollment_events.push(updated_enrollment.get_change_event());
@@ -285,6 +328,20 @@ impl ExperimentEnrollment {
                         updated_experiment,
                         targeting_helper,
                     )?;
+
+                    #[cfg(feature = "stateful")]
+                    if self.will_pref_experiment_change(updated_experiment, &evaluated_enrollment) {
+                        if let EnrollmentStatus::Enrolled {
+                            prev_gecko_pref_states: Some(prev_gecko_pref_states),
+                            ..
+                        } = &self.status
+                        {
+                            PreviousGeckoPrefState::on_revert_to_prev_gecko_pref_states(
+                                prev_gecko_pref_states,
+                                gecko_pref_store,
+                            );
+                        }
+                    }
                     match evaluated_enrollment.status {
                         EnrollmentStatus::Error { .. } => {
                             let updated_enrollment =
@@ -369,6 +426,7 @@ impl ExperimentEnrollment {
     /// from the database after `PREVIOUS_ENROLLMENTS_GC_TIME`.
     fn on_experiment_ended(
         &self,
+        #[cfg(feature = "stateful")] gecko_pref_store: &Option<Arc<GeckoPrefStore>>,
         out_enrollment_events: &mut Vec<EnrollmentChangeEvent>,
     ) -> Option<Self> {
         debug!(
@@ -382,6 +440,17 @@ impl ExperimentEnrollment {
             | EnrollmentStatus::WasEnrolled { .. }
             | EnrollmentStatus::Error { .. } => return None, // We were never enrolled anyway, simply delete the enrollment record from the DB.
         };
+        #[cfg(feature = "stateful")]
+        if let EnrollmentStatus::Enrolled {
+            prev_gecko_pref_states: Some(prev_gecko_pref_states),
+            ..
+        } = &self.status
+        {
+            PreviousGeckoPrefState::on_revert_to_prev_gecko_pref_states(
+                prev_gecko_pref_states,
+                gecko_pref_store,
+            );
+        }
         let enrollment = Self {
             slug: self.slug.clone(),
             status: EnrollmentStatus::WasEnrolled {
@@ -399,9 +468,22 @@ impl ExperimentEnrollment {
     pub(crate) fn on_explicit_opt_out(
         &self,
         out_enrollment_events: &mut Vec<EnrollmentChangeEvent>,
+        #[cfg(feature = "stateful")] gecko_pref_store: &Option<Arc<GeckoPrefStore>>,
     ) -> ExperimentEnrollment {
         match self.status {
             EnrollmentStatus::Enrolled { .. } => {
+                #[cfg(feature = "stateful")]
+                if let EnrollmentStatus::Enrolled {
+                    prev_gecko_pref_states: Some(prev_gecko_pref_states),
+                    ..
+                } = &self.status
+                {
+                    PreviousGeckoPrefState::on_revert_to_prev_gecko_pref_states(
+                        prev_gecko_pref_states,
+                        gecko_pref_store,
+                    );
+                }
+
                 let enrollment = self.disqualify_from_enrolled(DisqualifiedReason::OptOut);
                 out_enrollment_events.push(enrollment.get_change_event());
                 enrollment
@@ -555,6 +637,71 @@ impl ExperimentEnrollment {
             | EnrollmentStatus::Error { .. } => self.clone(),
         }
     }
+
+    #[cfg(feature = "stateful")]
+    pub(crate) fn will_pref_experiment_change(
+        &self,
+        updated_experiment: &Experiment,
+        updated_enrollment: &ExperimentEnrollment,
+    ) -> bool {
+        let (original_prev_gecko_pref_states, original_branch_slug) = match &self.status {
+            EnrollmentStatus::Enrolled {
+                prev_gecko_pref_states: Some(prev_gecko_pref_states),
+                branch,
+                ..
+            } => (prev_gecko_pref_states, branch),
+            // Can't change if it isn't a pref experiment
+            _ => {
+                return false;
+            }
+        };
+
+        let updated_branch_slug = match &updated_enrollment.status {
+            EnrollmentStatus::Enrolled { branch, .. } => branch,
+            // If we are no longer going to be enrolled, then a change happened
+            _ => {
+                return true;
+            }
+        };
+
+        // Branch changed
+        if updated_branch_slug != original_branch_slug {
+            return true;
+        }
+
+        // Couldn't get a branch, something changed
+        let Some(updated_branch) = updated_experiment.get_branch(updated_branch_slug) else {
+            return true;
+        };
+
+        let updated_features = updated_branch.get_feature_configs();
+        let original_feature_ids: HashSet<&String> = original_prev_gecko_pref_states
+            .iter()
+            .map(|state| &state.feature_id)
+            .collect();
+
+        // Amount of features should be the same
+        if updated_features.len() != original_feature_ids.len() {
+            return true;
+        }
+
+        for original_state in original_prev_gecko_pref_states {
+            let matching_feature = updated_features
+                .iter()
+                .find(|config| config.feature_id == original_state.feature_id);
+
+            // If original feature isn't present, then something changed
+            let Some(updated_feature) = matching_feature else {
+                return true;
+            };
+
+            // Property key should still exist in the feature's value map
+            if !updated_feature.value.contains_key(&original_state.variable) {
+                return true;
+            }
+        }
+        false
+    }
 }
 
 // ⚠️ Attention : Changes to this type should be accompanied by a new test  ⚠️
@@ -650,6 +797,7 @@ impl<'a> EnrollmentsEvolver<'a> {
         prev_experiments: &[E],
         next_experiments: &[Experiment],
         prev_enrollments: &[ExperimentEnrollment],
+        #[cfg(feature = "stateful")] gecko_pref_store: &Option<Arc<GeckoPrefStore>>,
     ) -> Result<(Vec<ExperimentEnrollment>, Vec<EnrollmentChangeEvent>)>
     where
         E: ExperimentMetadata + Clone,
@@ -671,6 +819,8 @@ impl<'a> EnrollmentsEvolver<'a> {
             &prev_rollouts,
             &next_rollouts,
             &ro_enrollments,
+            #[cfg(feature = "stateful")]
+            gecko_pref_store,
         )?;
 
         enrollments.extend(next_ro_enrollments);
@@ -695,6 +845,8 @@ impl<'a> EnrollmentsEvolver<'a> {
             &prev_experiments,
             &next_experiments,
             &prev_enrollments,
+            #[cfg(feature = "stateful")]
+            gecko_pref_store,
         )?;
 
         enrollments.extend(next_exp_enrollments);
@@ -711,6 +863,7 @@ impl<'a> EnrollmentsEvolver<'a> {
         prev_experiments: &[E],
         next_experiments: &[Experiment],
         prev_enrollments: &[ExperimentEnrollment],
+        #[cfg(feature = "stateful")] gecko_pref_store: &Option<Arc<GeckoPrefStore>>,
     ) -> Result<(Vec<ExperimentEnrollment>, Vec<EnrollmentChangeEvent>)>
     where
         E: ExperimentMetadata + Clone,
@@ -750,6 +903,8 @@ impl<'a> EnrollmentsEvolver<'a> {
                 next_experiments_map.get(slug).copied(),
                 Some(prev_enrollment),
                 &mut enrollment_events,
+                #[cfg(feature = "stateful")]
+                gecko_pref_store,
             ) {
                 Ok(enrollment) => enrollment,
                 Err(e) => {
@@ -851,6 +1006,8 @@ impl<'a> EnrollmentsEvolver<'a> {
                     Some(next_experiment),
                     prev_enrollment,
                     &mut enrollment_events,
+                    #[cfg(feature = "stateful")]
+                    gecko_pref_store,
                 ) {
                     Ok(enrollment) => enrollment,
                     Err(e) => {
@@ -944,6 +1101,7 @@ impl<'a> EnrollmentsEvolver<'a> {
         next_experiment: Option<&Experiment>,
         prev_enrollment: Option<&ExperimentEnrollment>,
         out_enrollment_events: &mut Vec<EnrollmentChangeEvent>, // out param containing the events we'd like to emit to glean.
+        #[cfg(feature = "stateful")] gecko_pref_store: &Option<Arc<GeckoPrefStore>>,
     ) -> Result<Option<ExperimentEnrollment>>
     where
         E: ExperimentMetadata + Clone,
@@ -972,16 +1130,20 @@ impl<'a> EnrollmentsEvolver<'a> {
                 out_enrollment_events,
             )?),
             // Experiment deleted remotely.
-            (Some(_), None, Some(enrollment)) => {
-                enrollment.on_experiment_ended(out_enrollment_events)
-            }
+            (Some(_), None, Some(enrollment)) => enrollment.on_experiment_ended(
+                #[cfg(feature = "stateful")]
+                gecko_pref_store,
+                out_enrollment_events,
+            ),
             // Known experiment.
             (Some(_), Some(experiment), Some(enrollment)) => {
                 Some(enrollment.on_experiment_updated(
                     is_user_participating,
                     self.available_randomization_units,
                     experiment,
                     &targeting_helper,
+                    #[cfg(feature = "stateful")]
+                    gecko_pref_store,
                     out_enrollment_events,
                 )?)
             }

components/nimbus/src/nimbus.udl

@@ -128,6 +128,9 @@ callback interface GeckoPrefHandler {
     record<string, record<string, GeckoPrefState>> get_prefs_with_state();
 
     void set_gecko_prefs_state(sequence<GeckoPrefState> new_prefs_state);
+
+    void set_gecko_prefs_original_values(sequence<OriginalGeckoPref> original_gecko_prefs);
+
 };
 
 dictionary GeckoPref {