nimbus/stateful/
dbcache.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/. */
4
5use std::collections::{HashMap, HashSet};
6use std::sync::{Arc, RwLock};
7
8use crate::enrollment::{
9    EnrolledFeature, EnrolledFeatureConfig, ExperimentEnrollment, map_features_by_feature_id,
10};
11use crate::error::{NimbusError, Result, warn};
12use crate::evaluator::{ExperimentAvailable, is_experiment_available};
13use crate::stateful::enrollment::get_enrollments;
14use crate::stateful::firefox_labs::FirefoxLabsMetadata;
15use crate::stateful::gecko_prefs::GeckoPrefStore;
16use crate::stateful::persistence::{Database, StoreId, Writer};
17use crate::targeting::NimbusTargetingHelper;
18use crate::{EnrolledExperiment, Experiment};
19
20// This module manages an in-memory cache of the database, so that some
21// functions exposed by nimbus can return results without blocking on any
22// IO. Consumers are expected to call our public `update()` function whenever
23// the database might have changed.
24
25// This struct is the cached data. This is never mutated, but instead
26// recreated every time the cache is updated.
27struct CachedData {
28    pub experiments: Vec<Experiment>,
29    pub enrollments: Vec<ExperimentEnrollment>,
30    pub experiments_by_slug: HashMap<String, EnrolledExperiment>,
31    pub features_by_feature_id: HashMap<String, EnrolledFeatureConfig>,
32    pub gecko_pref_to_enrollment_slugs: Option<HashMap<String, HashSet<String>>>,
33}
34
35// This is the public cache API. Each NimbusClient can create one of these and
36// it lives as long as the client - it encapsulates the synchronization needed
37// to allow the cache to work correctly.
38#[derive(Default)]
39pub struct DatabaseCache {
40    data: RwLock<Option<CachedData>>,
41}
42
43impl DatabaseCache {
44    // Call this function whenever it's possible that anything cached by this
45    // struct (eg, our enrollments) might have changed.
46    //
47    // This function must be passed a `&Database` and a `Writer`, which it
48    // will commit before updating the in-memory cache. This is a slightly weird
49    // API but it helps enforce two important properties:
50    //
51    //  * By requiring a `Writer`, we ensure mutual exclusion of other db writers
52    //    and thus prevent the possibility of caching stale data.
53    //  * By taking ownership of the `Writer`, we ensure that the calling code
54    //    updates the cache after all of its writes have been performed.
55    //  * `update_gecko_prefs` - Pass true for regular enrollment changes. Pass false
56    //     when the Gecko prefs do not need to be synced with Gecko.
57    pub fn commit_and_update(
58        &self,
59        db: &Database,
60        writer: Writer,
61        coenrolling_ids: &HashSet<&str>,
62        gecko_pref_store: Option<Arc<GeckoPrefStore>>,
63        update_gecko_prefs: bool,
64    ) -> Result<()> {
65        // By passing in the active `writer` we read the state of enrollments
66        // as written by the calling code, before it's committed to the db.
67        let enrollments = get_enrollments(db, &writer)?;
68
69        // Build a lookup table for experiments by experiment slug.
70        // This will be used for get_experiment_branch() and get_active_experiments()
71        let mut experiments_by_slug = HashMap::with_capacity(enrollments.len());
72        for e in enrollments {
73            experiments_by_slug.insert(e.slug.clone(), e);
74        }
75
76        let enrollments: Vec<ExperimentEnrollment> =
77            db.get_store(StoreId::Enrollments).collect_all(&writer)?;
78        let experiments: Vec<Experiment> =
79            db.get_store(StoreId::Experiments).collect_all(&writer)?;
80
81        let features_by_feature_id =
82            map_features_by_feature_id(&enrollments, &experiments, coenrolling_ids);
83
84        let gecko_pref_to_enrollment_slugs = gecko_pref_store.map(|store| {
85            store.map_gecko_prefs_to_enrollment_slugs_and_update_store(
86                &experiments,
87                &enrollments,
88                &experiments_by_slug,
89                update_gecko_prefs,
90            )
91        });
92
93        // This is where testing tools would override i.e. replace experimental feature configurations.
94        // i.e. testing tools would cause custom feature configs to be stored in a Store.
95        // Here, we get those overrides out of the store, and merge it with this map.
96
97        // This is where rollouts (promoted experiments on a given feature) will be merged in to the feature variables.
98
99        let data = CachedData {
100            experiments,
101            enrollments,
102            experiments_by_slug,
103            features_by_feature_id,
104            gecko_pref_to_enrollment_slugs,
105        };
106
107        // Try to commit the change to disk and update the cache as close
108        // together in time as possible. This leaves a small window where another
109        // thread could read new data from disk but see old data in the cache,
110        // but that seems benign in practice given the way we use the cache.
111        // The alternative would be to lock the cache while we commit to disk,
112        // and we don't want to risk blocking the main thread.
113        writer.commit()?;
114        let mut cached = self.data.write().unwrap();
115        cached.replace(data);
116        Ok(())
117    }
118
119    // Abstracts safely referencing our cached data.
120    //
121    // WARNING: because this manages locking, the callers of this need to be
122    // careful regarding deadlocks - if the callback takes other own locks then
123    // there's a risk of locks being taken in an inconsistent order. However,
124    // there's nothing this code specifically can do about that.
125    fn get_data<T, F>(&self, func: F) -> Result<T>
126    where
127        F: FnOnce(&CachedData) -> T,
128    {
129        match *self.data.read().unwrap() {
130            None => {
131                warn!("DatabaseCache attempting to read data before initialization is completed");
132                Err(NimbusError::DatabaseNotReady)
133            }
134            Some(ref data) => Ok(func(data)),
135        }
136    }
137
138    pub fn get_experiment_branch(&self, id: &str) -> Result<Option<String>> {
139        self.get_data(|data| -> Option<String> {
140            data.experiments_by_slug
141                .get(id)
142                .map(|experiment| experiment.branch_slug.clone())
143        })
144    }
145
146    // This gives access to the feature JSON. We pass it as a string because uniffi doesn't
147    // support JSON yet.
148    pub fn get_feature_config_variables(&self, feature_id: &str) -> Result<Option<String>> {
149        self.get_data(|data| {
150            let enrolled_feature = data.features_by_feature_id.get(feature_id)?;
151            let string = serde_json::to_string(&enrolled_feature.feature.value).unwrap();
152            Some(string)
153        })
154    }
155
156    pub fn get_enrollment_by_feature(&self, feature_id: &str) -> Result<Option<EnrolledFeature>> {
157        self.get_data(|data| {
158            data.features_by_feature_id
159                .get(feature_id)
160                .map(|feature| feature.into())
161        })
162    }
163
164    pub fn get_active_experiments(&self) -> Result<Vec<EnrolledExperiment>> {
165        self.get_data(|data| {
166            data.experiments_by_slug
167                .values()
168                .map(|e| e.to_owned())
169                .collect::<Vec<EnrolledExperiment>>()
170        })
171    }
172
173    pub fn get_experiments(&self) -> Result<Vec<Experiment>> {
174        self.get_data(|data| data.experiments.to_vec())
175    }
176
177    pub fn get_enrollments(&self) -> Result<Vec<ExperimentEnrollment>> {
178        self.get_data(|data| data.enrollments.to_owned())
179    }
180
181    pub fn get_enrollments_for_pref(&self, pref: &str) -> Result<Option<HashSet<String>>> {
182        self.get_data(|data| {
183            if let Some(a) = &data.gecko_pref_to_enrollment_slugs {
184                Ok(a.get(pref).cloned())
185            } else {
186                Ok(None)
187            }
188        })?
189    }
190
191    pub fn check_for_feature_conflict(
192        &self,
193        slug: &str,
194        coenrolling_feature_ids: &[String],
195    ) -> Result<Option<bool>> {
196        self.get_data(|data| {
197            if data.experiments_by_slug.contains_key(slug) {
198                // Cannot conflict with itself.
199                return Some(false);
200            }
201
202            if let Some(experiment) = data.experiments.iter().find(|e| e.slug == slug) {
203                let coenrolling_feature_ids: HashSet<&str> =
204                    coenrolling_feature_ids.iter().map(|s| s.as_ref()).collect();
205
206                let enrolled_feature_ids =
207                    compute_enrolled_feature_ids(&data.experiments_by_slug, true);
208
209                Some(!features_available(
210                    experiment,
211                    &enrolled_feature_ids,
212                    &coenrolling_feature_ids,
213                ))
214            } else {
215                None
216            }
217        })
218    }
219
220    pub fn get_available_firefox_labs_metadata(
221        &self,
222        targeting_helper: &NimbusTargetingHelper,
223        coenrolling_feature_ids: &[String],
224    ) -> Result<Vec<FirefoxLabsMetadata>> {
225        let mut all_labs: Vec<_> = self.get_data(|data| {
226            let enrolled_feature_ids =
227                compute_enrolled_feature_ids(&data.experiments_by_slug, true);
228
229            let coenrolling_feature_ids: HashSet<&str> =
230                coenrolling_feature_ids.iter().map(|s| s.as_ref()).collect();
231
232            data.experiments
233                .iter()
234                .filter_map(|experiment| {
235                    if !experiment.is_firefox_labs_opt_in {
236                        return None;
237                    }
238
239                    let enrolled = data.experiments_by_slug.contains_key(&experiment.slug);
240
241                    // We call is_experiment_available with is_release=true
242                    // because being able to enroll in experiments for different channels is a bug.
243                    //
244                    // See-also https://bugzilla.mozilla.org/show_bug.cgi?id=1909348
245                    if is_experiment_available(targeting_helper, experiment, true)
246                        == ExperimentAvailable::Available
247                        && (enrolled
248                            || (features_available(
249                                experiment,
250                                &enrolled_feature_ids,
251                                &coenrolling_feature_ids,
252                            ) && !experiment.is_enrollment_paused))
253                    {
254                        experiment.get_firefox_labs_metadata(enrolled)
255                    } else {
256                        None
257                    }
258                })
259                .collect()
260        })?;
261
262        // XXX: This is maybe only useful for tests, but at least we get a
263        // stable order.
264        all_labs.sort_by(|e1, e2| Ord::cmp(&e1.slug, &e2.slug));
265
266        Ok(all_labs)
267    }
268
269    #[cfg(test)]
270    pub fn get_experiment_enrollment(&self, slug: &str) -> Result<Option<ExperimentEnrollment>> {
271        self.get_data(|data| data.enrollments.iter().find(|e| e.slug == slug).cloned())
272    }
273}
274
275fn compute_enrolled_feature_ids(
276    experiments_by_slug: &HashMap<String, EnrolledExperiment>,
277    is_rollout: bool,
278) -> HashSet<&str> {
279    experiments_by_slug
280        .values()
281        .filter(|e| e.is_rollout == is_rollout)
282        .flat_map(|e| e.feature_ids.iter())
283        .map(|f| f.as_ref())
284        .collect()
285}
286
287fn features_available(
288    experiment: &Experiment,
289    enrolled_feature_ids: &HashSet<&str>,
290    coenrolling_feature_ids: &HashSet<&str>,
291) -> bool {
292    for feature_id in &experiment.feature_ids {
293        if enrolled_feature_ids.contains(&**feature_id)
294            && !coenrolling_feature_ids.contains(&**feature_id)
295        {
296            return false;
297        }
298    }
299
300    true
301}
nimbus/stateful/dbcache.rs

nimbus/stateful/
dbcache.rs
