1use 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 fn get_prefs_with_state(&self) -> MapOfFeatureIdToPropertyNameToGeckoPrefState;
110
111 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 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 pub fn map_gecko_prefs_to_enrollment_slugs_and_update_store(
193 &self,
194 experiments: &[Experiment],
196 enrollments: &[ExperimentEnrollment],
198 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 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 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 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 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 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 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}