nimbus/stateful/
gecko_prefs.rs

1/* This Source Code Form is subject to the terms of the Mozilla Public
2 * License, v. 2.0. If a copy of the MPL was not distributed with this
3 * file, You can obtain one at https://mozilla.org/MPL/2.0/. */
4use std::cmp::Ordering;
5use std::collections::{HashMap, HashSet};
6use std::fmt::{Display, Formatter};
7use std::sync::{Arc, Mutex, MutexGuard};
8
9use serde_derive::{Deserialize, Serialize};
10use serde_json::Value;
11
12use crate::enrollment::{EnrollmentStatus, ExperimentEnrollment, PreviousGeckoPrefState};
13use crate::error::Result;
14use crate::json::PrefValue;
15use crate::{EnrolledExperiment, Experiment, NimbusError};
16
17#[derive(Debug, Clone, Serialize, Deserialize, Hash, PartialEq, Eq, Copy)]
18#[serde(rename_all = "lowercase")]
19pub enum PrefBranch {
20    Default,
21    User,
22}
23
24impl Display for PrefBranch {
25    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
26        match self {
27            PrefBranch::Default => f.write_str("default"),
28            PrefBranch::User => f.write_str("user"),
29        }
30    }
31}
32
33#[derive(Debug, Clone, Serialize, Deserialize)]
34pub struct GeckoPref {
35    pub pref: String,
36    pub branch: PrefBranch,
37}
38
39#[derive(Debug, Clone, Serialize, Deserialize)]
40pub struct PrefEnrollmentData {
41    pub experiment_slug: String,
42    pub pref_value: PrefValue,
43    pub feature_id: String,
44    pub variable: String,
45}
46
47#[derive(Debug, Clone, Serialize, Deserialize)]
48pub struct GeckoPrefState {
49    pub gecko_pref: GeckoPref,
50    pub gecko_value: Option<PrefValue>,
51    pub enrollment_value: Option<PrefEnrollmentData>,
52    pub is_user_set: bool,
53}
54
55impl GeckoPrefState {
56    pub fn new(pref: &str, branch: Option<PrefBranch>) -> Self {
57        Self {
58            gecko_pref: GeckoPref {
59                pref: pref.into(),
60                branch: branch.unwrap_or(PrefBranch::Default),
61            },
62            gecko_value: None,
63            enrollment_value: None,
64            is_user_set: false,
65        }
66    }
67
68    pub fn with_gecko_value(mut self, value: PrefValue) -> Self {
69        self.gecko_value = Some(value);
70        self
71    }
72
73    pub fn with_enrollment_value(mut self, pref_enrollment_data: PrefEnrollmentData) -> Self {
74        self.enrollment_value = Some(pref_enrollment_data);
75        self
76    }
77
78    pub fn set_by_user(mut self) -> Self {
79        self.is_user_set = true;
80        self
81    }
82}
83
84#[derive(Debug, Clone, Serialize, Deserialize, Hash, Eq, PartialEq, Copy)]
85pub enum PrefUnenrollReason {
86    Changed,
87    FailedToSet,
88}
89
90// The pre-experiment original state of a Gecko pref. Values may be used to set on Gecko to restore the pref to the original state.
91// ⚠️ Attention : Changes to this type should be accompanied by a new test  ⚠️
92// ⚠️ in `src/stateful/tests/test_enrollment_bw_compat.rs` below, and may require a DB migration. ⚠️
93#[derive(Deserialize, Serialize, Debug, Clone, Hash, Eq, PartialEq)]
94pub struct OriginalGeckoPref {
95    pub pref: String,
96    pub branch: PrefBranch,
97    pub value: Option<PrefValue>,
98}
99
100impl<'a> From<&'a GeckoPrefState> for OriginalGeckoPref {
101    fn from(state: &'a GeckoPrefState) -> Self {
102        Self {
103            pref: state.gecko_pref.pref.clone(),
104            branch: state.gecko_pref.branch,
105            value: state.gecko_value.clone(),
106        }
107    }
108}
109
110pub type MapOfFeatureIdToPropertyNameToGeckoPrefState =
111    HashMap<String, HashMap<String, GeckoPrefState>>;
112
113pub fn create_feature_prop_pref_map(
114    list: Vec<(&str, &str, GeckoPrefState)>,
115) -> MapOfFeatureIdToPropertyNameToGeckoPrefState {
116    list.iter().fold(
117        HashMap::new(),
118        |mut feature_map, (feature_id, prop_name, pref_state)| {
119            feature_map
120                .entry(feature_id.to_string())
121                .or_default()
122                .insert(prop_name.to_string(), pref_state.clone());
123            feature_map
124        },
125    )
126}
127
128#[uniffi::trait_interface]
129pub trait GeckoPrefHandler: Send + Sync {
130    /// Used to obtain the prefs values from Gecko
131    fn get_prefs_with_state(&self) -> MapOfFeatureIdToPropertyNameToGeckoPrefState;
132
133    /// Used to set the state for each pref based on enrollments
134    fn set_gecko_prefs_state(&self, new_prefs_state: Vec<GeckoPrefState>);
135
136    /// Used to set back to the original state for each pref based on the original Gecko value
137    fn set_gecko_prefs_original_values(&self, original_gecko_prefs: Vec<OriginalGeckoPref>);
138}
139
140#[derive(Default)]
141pub struct GeckoPrefStoreState {
142    pub gecko_prefs_with_state: MapOfFeatureIdToPropertyNameToGeckoPrefState,
143}
144
145impl GeckoPrefStoreState {
146    pub fn update_pref_state(&mut self, new_pref_state: &GeckoPrefState) -> bool {
147        self.gecko_prefs_with_state
148            .iter_mut()
149            .find_map(|(_, props)| {
150                props.iter_mut().find_map(|(_, pref_state)| {
151                    if pref_state.gecko_pref.pref == new_pref_state.gecko_pref.pref {
152                        *pref_state = new_pref_state.clone();
153                        Some(true)
154                    } else {
155                        None
156                    }
157                })
158            })
159            .is_some()
160    }
161}
162
163pub struct GeckoPrefStore {
164    // This is Arc<Box<_>> because of FFI
165    pub handler: Arc<dyn GeckoPrefHandler>,
166    pub state: Mutex<GeckoPrefStoreState>,
167}
168
169impl GeckoPrefStore {
170    pub fn new(handler: Arc<dyn GeckoPrefHandler>) -> Self {
171        Self {
172            handler,
173            state: Mutex::new(GeckoPrefStoreState::default()),
174        }
175    }
176
177    pub fn initialize(&self) -> Result<()> {
178        let prefs = self.handler.get_prefs_with_state();
179        let mut state = self
180            .state
181            .lock()
182            .expect("Unable to lock GeckoPrefStore state");
183        state.gecko_prefs_with_state = prefs;
184
185        Ok(())
186    }
187
188    pub fn get_mutable_pref_state(&self) -> MutexGuard<'_, GeckoPrefStoreState> {
189        self.state
190            .lock()
191            .expect("Unable to lock GeckoPrefStore state")
192    }
193
194    pub fn pref_is_user_set(&self, pref: &str) -> bool {
195        let state = self.get_mutable_pref_state();
196        state
197            .gecko_prefs_with_state
198            .iter()
199            .find_map(|(_, props)| {
200                props.iter().find_map(|(_, gecko_pref_state)| {
201                    if gecko_pref_state.gecko_pref.pref == pref {
202                        Some(gecko_pref_state.is_user_set)
203                    } else {
204                        None
205                    }
206                })
207            })
208            .unwrap_or(false)
209    }
210
211    /// This method accomplishes a number of tasks important to the Gecko pref enrollment workflow.
212    /// 1. It returns a map of pref string to a vector of enrolled recipes in which the value for
213    ///    the enrolled branch's feature values includes the property of that feature that sets the
214    ///    aforementioned pref.
215    /// 2. It updates the GeckoPrefStore state, such that the appropriate GeckoPrefState's
216    ///    `enrollment_value` reflects the appropriate value.
217    pub fn map_gecko_prefs_to_enrollment_slugs_and_update_store(
218        &self,
219        // contains full experiment metadata
220        experiments: &[Experiment],
221        // contains enrollment status for a given experiment
222        enrollments: &[ExperimentEnrollment],
223        // contains slug of enrolled branch
224        experiments_by_slug: &HashMap<String, EnrolledExperiment>,
225        // when true, sets Gecko prefs using the handler. when false, only updates internal store state
226        update_gecko_prefs: bool,
227    ) -> HashMap<String, HashSet<String>> {
228        struct RecipeData<'a> {
229            experiment: &'a Experiment,
230            experiment_enrollment: &'a ExperimentEnrollment,
231            branch_slug: &'a str,
232        }
233
234        let mut state = self.get_mutable_pref_state();
235
236        /* List of tuples that contain recipe slug, rollout bool, list of feature ids, and
237         * branch, in that order.
238         */
239        let mut recipe_data: Vec<RecipeData> = vec![];
240
241        for experiment_enrollment in enrollments {
242            let experiment = match experiments
243                .iter()
244                .find(|experiment| experiment.slug == experiment_enrollment.slug)
245            {
246                Some(exp) => exp,
247                None => continue,
248            };
249            recipe_data.push(RecipeData {
250                experiment,
251                experiment_enrollment,
252                branch_slug: match experiments_by_slug.get(&experiment.slug) {
253                    Some(ee) => &ee.branch_slug,
254                    None => continue,
255                },
256            });
257        }
258        // sort `recipe_data` such that rollouts are applied before experiments
259        recipe_data.sort_by(
260            |a, b| match (a.experiment.is_rollout, b.experiment.is_rollout) {
261                (true, false) => Ordering::Less,
262                (false, true) => Ordering::Greater,
263                _ => Ordering::Equal,
264            },
265        );
266
267        /* This map will ultimately be returned from the function, as a map of pref strings to
268         * relevant enrolled recipe slugs, 'relevant' meaning experiments whose enrolled branch
269         * values apply a value to a prop for which there is a Gecko pref.
270         *
271         * We start by iterating mutably over the map of features to props to gecko prefs.
272         */
273        let mut results: HashMap<String, HashSet<String>> = HashMap::new();
274
275        for (feature_name, props) in state.gecko_prefs_with_state.iter_mut() {
276            let mut has_matching_recipes = false;
277            for RecipeData {
278                experiment:
279                    Experiment {
280                        slug,
281                        feature_ids,
282                        branches,
283                        ..
284                    },
285                experiment_enrollment,
286                branch_slug,
287            } in &recipe_data
288            {
289                if feature_ids.contains(feature_name)
290                    && matches!(
291                        experiment_enrollment.status,
292                        EnrollmentStatus::Enrolled { .. }
293                    )
294                {
295                    let branch = match branches.iter().find(|branch| &branch.slug == branch_slug) {
296                        Some(b) => b,
297                        None => continue,
298                    };
299                    has_matching_recipes = true;
300                    for (feature, prop_name, prop_value) in branch.get_feature_props_and_values() {
301                        if feature == *feature_name && props.contains_key(&prop_name) {
302                            // set the enrollment_value for this gecko pref.
303                            // rollouts and experiments on the same feature will
304                            // both set the value here, but rollouts will happen
305                            // first, and will therefore be overridden by
306                            // experiments.
307                            props.entry(prop_name.clone()).and_modify(|pref_state| {
308                                pref_state.enrollment_value = Some(PrefEnrollmentData {
309                                    experiment_slug: slug.clone(),
310                                    pref_value: prop_value.clone(),
311                                    feature_id: feature,
312                                    variable: prop_name,
313                                });
314                                results
315                                    .entry(pref_state.gecko_pref.pref.clone())
316                                    .or_default()
317                                    .insert(slug.clone());
318                            });
319                        }
320                    }
321                }
322            }
323
324            if !has_matching_recipes {
325                for (_, pref_state) in props.iter_mut() {
326                    pref_state.enrollment_value = None;
327                }
328            }
329        }
330
331        // 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)
332        if update_gecko_prefs {
333            // obtain a list of all Gecko pref states for which there is an enrollment value
334            let mut set_state_list = Vec::new();
335            state.gecko_prefs_with_state.iter().for_each(|(_, props)| {
336                props.iter().for_each(|(_, pref_state)| {
337                    if pref_state.enrollment_value.is_some() {
338                        set_state_list.push(pref_state.clone());
339                    }
340                });
341            });
342            // tell the handler to set the aforementioned Gecko prefs
343            self.handler.set_gecko_prefs_state(set_state_list);
344        }
345
346        results
347    }
348}
349
350pub fn query_gecko_pref_store(
351    gecko_pref_store: Option<Arc<GeckoPrefStore>>,
352    args: &[Value],
353) -> Result<Value> {
354    if args.len() != 1 {
355        return Err(NimbusError::TransformParameterError(
356            "gecko_pref transform preferenceIsUserSet requires exactly 1 parameter".into(),
357        ));
358    }
359
360    let gecko_pref = match serde_json::from_value::<String>(args.first().unwrap().clone()) {
361        Ok(v) => v,
362        Err(e) => return Err(NimbusError::JSONError("gecko_pref = nimbus::stateful::gecko_prefs::query_gecko_prefs_store::serde_json::from_value".into(), e.to_string()))
363    };
364
365    Ok(gecko_pref_store
366        .map(|store| Value::Bool(store.pref_is_user_set(&gecko_pref)))
367        .unwrap_or(Value::Bool(false)))
368}
369
370pub(crate) type MapOfExperimentSlugToPreviousState = HashMap<String, Vec<PreviousGeckoPrefState>>;
371pub(crate) fn build_prev_gecko_pref_states(
372    states: &[GeckoPrefState],
373) -> MapOfExperimentSlugToPreviousState {
374    let mut original_gecko_states = MapOfExperimentSlugToPreviousState::new();
375    for state in states {
376        let Some(enrollment_value) = &state.enrollment_value else {
377            continue;
378        };
379        original_gecko_states
380            .entry(enrollment_value.experiment_slug.clone())
381            .or_default()
382            .push(PreviousGeckoPrefState {
383                original_value: state.into(),
384                feature_id: enrollment_value.feature_id.clone(),
385                variable: enrollment_value.variable.clone(),
386            });
387    }
388    original_gecko_states
389}
nimbus/stateful/gecko_prefs.rs

nimbus/stateful/
gecko_prefs.rs
