Bug 2038991 - Update Cache when calling register_previous_gecko_pref_states

components/nimbus/src/stateful/dbcache.rs

@@ -43,18 +43,21 @@ impl DatabaseCache {
     //
     // This function must be passed a `&Database` and a `Writer`, which it
     // will commit before updating the in-memory cache. This is a slightly weird
-    // API but it helps encorce two important properties:
+    // API but it helps enforce two important properties:
     //
     //  * By requiring a `Writer`, we ensure mutual exclusion of other db writers
     //    and thus prevent the possibility of caching stale data.
     //  * By taking ownership of the `Writer`, we ensure that the calling code
     //    updates the cache after all of its writes have been performed.
+    //  * `update_gecko_prefs` - Pass true for regular enrollment changes. Pass false
+    //     when the Gecko prefs do not need to be synced with Gecko.
     pub fn commit_and_update(
         &self,
         db: &Database,
         writer: Writer,
         coenrolling_ids: &HashSet<&str>,
         gecko_pref_store: Option<Arc<GeckoPrefStore>>,
+        update_gecko_prefs: bool,
     ) -> Result<()> {
         // By passing in the active `writer` we read the state of enrollments
         // as written by the calling code, before it's committed to the db.
@@ -80,6 +83,7 @@ impl DatabaseCache {
                 &experiments,
                 &enrollments,
                 &experiments_by_slug,
+                update_gecko_prefs,
             )
         });
 

components/nimbus/src/stateful/gecko_prefs.rs

@@ -221,6 +221,8 @@ impl GeckoPrefStore {
         enrollments: &[ExperimentEnrollment],
         // contains slug of enrolled branch
         experiments_by_slug: &HashMap<String, EnrolledExperiment>,
+        // when true, sets Gecko prefs using the handler. when false, only updates internal store state
+        update_gecko_prefs: bool,
     ) -> HashMap<String, HashSet<String>> {
         struct RecipeData<'a> {
             experiment: &'a Experiment,
@@ -325,17 +327,20 @@ impl GeckoPrefStore {
             }
         }
 
-        // obtain a list of all Gecko pref states for which there is an enrollment value
-        let mut set_state_list = Vec::new();
-        state.gecko_prefs_with_state.iter().for_each(|(_, props)| {
-            props.iter().for_each(|(_, pref_state)| {
-                if pref_state.enrollment_value.is_some() {
-                    set_state_list.push(pref_state.clone());
-                }
+        // setting with Gecko is used for true updates, sometimes the state just needs to be recalculated (e.g., registering the original value on the enrollment)
+        if update_gecko_prefs {
+            // obtain a list of all Gecko pref states for which there is an enrollment value
+            let mut set_state_list = Vec::new();
+            state.gecko_prefs_with_state.iter().for_each(|(_, props)| {
+                props.iter().for_each(|(_, pref_state)| {
+                    if pref_state.enrollment_value.is_some() {
+                        set_state_list.push(pref_state.clone());
+                    }
+                });
             });
-        });
-        // tell the handler to set the aforementioned Gecko prefs
-        self.handler.set_gecko_prefs_state(set_state_list);
+            // tell the handler to set the aforementioned Gecko prefs
+            self.handler.set_gecko_prefs_state(set_state_list);
+        }
 
         results
     }

components/nimbus/src/stateful/nimbus_client.rs

@@ -227,6 +227,7 @@ impl NimbusClient {
             writer,
             &coenrolling_ids,
             self.gecko_prefs.clone(),
+            true,
         )?;
         Ok(())
     }
@@ -846,7 +847,22 @@ impl NimbusClient {
             )?;
         }
 
-        writer.commit()?;
+        let coenrolling_ids = self
+            .coenrolling_feature_ids
+            .iter()
+            .map(|s| s.as_str())
+            .collect();
+
+        // Registering the Gecko original values does not require a Gecko update,
+        // but the cache does need to be refreshed.
+        self.database_cache.commit_and_update(
+            db,
+            writer,
+            &coenrolling_ids,
+            self.gecko_prefs.clone(),
+            false,
+        )?;
+
         Ok(())
     }
 

components/nimbus/src/tests/helpers.rs

@@ -177,6 +177,7 @@ impl TestMetrics {
 pub struct TestGeckoPrefHandlerState {
     pub prefs_set: Option<Vec<GeckoPrefState>>,
     pub original_prefs_state: Option<Vec<OriginalGeckoPref>>,
+    pub set_gecko_prefs_state_call_count: u32,
 }
 
 #[cfg(feature = "stateful")]
@@ -193,6 +194,7 @@ impl TestGeckoPrefHandler {
             state: Mutex::new(TestGeckoPrefHandlerState {
                 prefs_set: None,
                 original_prefs_state: None,
+                set_gecko_prefs_state_call_count: 0,
             }),
         }
     }
@@ -205,10 +207,12 @@ impl GeckoPrefHandler for TestGeckoPrefHandler {
     }
 
     fn set_gecko_prefs_state(&self, new_prefs_state: Vec<GeckoPrefState>) {
-        self.state
+        let mut state = self
+            .state
             .lock()
-            .expect("Unable to lock TestGeckoPrefHandler state")
-            .prefs_set = Some(new_prefs_state);
+            .expect("Unable to lock TestGeckoPrefHandler state");
+        state.prefs_set = Some(new_prefs_state);
+        state.set_gecko_prefs_state_call_count += 1;
     }
 
     fn set_gecko_prefs_original_values(&self, original_prefs_state: Vec<OriginalGeckoPref>) {

components/nimbus/src/tests/stateful/test_gecko_prefs.rs

@@ -56,6 +56,7 @@ fn test_gecko_pref_store_map_gecko_prefs_to_enrollment_slugs_and_update_store()
         &experiments,
         &experiment_enrollments,
         &experiments_by_slug,
+        true,
     );
 
     let handler = unsafe {
@@ -147,6 +148,7 @@ fn test_gecko_pref_store_map_gecko_prefs_to_enrollment_slugs_and_update_store_ex
         &experiments,
         &experiment_enrollments,
         &experiments_by_slug,
+        true,
     );
 
     let handler = unsafe {

components/nimbus/src/tests/stateful/test_nimbus.rs

@@ -2174,9 +2174,28 @@ fn register_previous_gecko_pref_states() -> Result<()> {
         gecko_pref_state_2.clone(),
         gecko_pref_state_3.clone(),
     ];
+
+    let call_count_before = client
+        .get_gecko_pref_store()
+        .state
+        .lock()
+        .unwrap()
+        .set_gecko_prefs_state_call_count;
+
     let registration = client.register_previous_gecko_pref_states(&gecko_pref_states);
     assert!(registration.is_ok());
 
+    // Registration must not send pref values to Gecko.
+    assert_eq!(
+        call_count_before,
+        client
+            .get_gecko_pref_store()
+            .state
+            .lock()
+            .unwrap()
+            .set_gecko_prefs_state_call_count
+    );
+
     let db = client.db()?;
     let reader = db.read()?;
     let mut enrollments: Vec<ExperimentEnrollment> =