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