nimbus/stateful/
enrollment.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::iter;
6
7use crate::enrollment::Participation;
8use crate::enrollment::{
9    EnrollmentChangeEvent, EnrollmentChangeEventType, EnrollmentsEvolver, ExperimentEnrollment,
10    map_enrollments,
11};
12use crate::error::{Result, debug, warn};
13use crate::stateful::gecko_prefs::GeckoPrefStore;
14use crate::stateful::gecko_prefs::PrefUnenrollReason;
15use crate::stateful::persistence::{
16    DB_KEY_EXPERIMENT_PARTICIPATION, DB_KEY_ROLLOUT_PARTICIPATION,
17    DEFAULT_EXPERIMENT_PARTICIPATION, DEFAULT_ROLLOUT_PARTICIPATION,
18};
19use crate::stateful::persistence::{Database, Readable, StoreId, Writer};
20use crate::{EnrolledExperiment, EnrollmentStatus, Experiment};
21
22impl EnrollmentsEvolver<'_> {
23    /// Convenient wrapper around `evolve_enrollments` that fetches the current state of experiments,
24    /// enrollments and user participation from the database.
25    pub(crate) fn evolve_enrollments_in_db(
26        &mut self,
27        db: &Database,
28        writer: &mut Writer,
29        next_experiments: &[Experiment],
30        gecko_pref_store: Option<&GeckoPrefStore>,
31    ) -> Result<Vec<EnrollmentChangeEvent>> {
32        // Get separate participation states from the db
33        let is_participating_in_experiments = get_experiment_participation(db, writer)?;
34        let is_participating_in_rollouts = get_rollout_participation(db, writer)?;
35
36        let participation = Participation {
37            in_experiments: is_participating_in_experiments,
38            in_rollouts: is_participating_in_rollouts,
39        };
40
41        let experiments_store = db.get_store(StoreId::Experiments);
42        let enrollments_store = db.get_store(StoreId::Enrollments);
43        let prev_experiments: Vec<Experiment> = experiments_store.collect_all(writer)?;
44        let prev_enrollments: Vec<ExperimentEnrollment> = enrollments_store.collect_all(writer)?;
45        // Calculate the changes.
46        let (next_enrollments, enrollments_change_events) = self.evolve_enrollments(
47            participation,
48            &prev_experiments,
49            next_experiments,
50            &prev_enrollments,
51            gecko_pref_store,
52        )?;
53        let next_enrollments = map_enrollments(&next_enrollments);
54        // Write the changes to the Database.
55        enrollments_store.clear(writer)?;
56        for enrollment in next_enrollments.values() {
57            enrollments_store.put(writer, &enrollment.slug, *enrollment)?;
58        }
59        experiments_store.clear(writer)?;
60        for experiment in next_experiments {
61            // Sanity check.
62            if !next_enrollments.contains_key(&experiment.slug) {
63                error_support::report_error!(
64                    "nimbus-evolve-enrollments",
65                    "evolve_enrollments_in_db: experiment '{}' has no enrollment, dropping to keep database consistent",
66                    &experiment.slug
67                );
68                continue;
69            }
70            experiments_store.put(writer, &experiment.slug, experiment)?;
71        }
72        Ok(enrollments_change_events)
73    }
74}
75
76/// Return information about all enrolled experiments.
77/// Note this does not include rollouts
78pub fn get_enrollments<'r>(
79    db: &Database,
80    reader: &'r impl Readable<'r>,
81) -> Result<Vec<EnrolledExperiment>> {
82    let enrollments: Vec<ExperimentEnrollment> =
83        db.get_store(StoreId::Enrollments).collect_all(reader)?;
84    let mut result = Vec::with_capacity(enrollments.len());
85    for enrollment in enrollments {
86        debug!("Have enrollment: {:?}", enrollment);
87        if let EnrollmentStatus::Enrolled { branch, .. } = &enrollment.status {
88            match db
89                .get_store(StoreId::Experiments)
90                .get::<Experiment, _>(reader, &enrollment.slug)?
91            {
92                Some(experiment) => {
93                    result.push(EnrolledExperiment {
94                        feature_ids: experiment.get_feature_ids(),
95                        slug: experiment.slug,
96                        user_facing_name: experiment.user_facing_name,
97                        user_facing_description: experiment.user_facing_description,
98                        branch_slug: branch.to_string(),
99                    });
100                }
101                _ => {
102                    warn!(
103                        "Have enrollment {:?} but no matching experiment!",
104                        enrollment
105                    );
106                }
107            };
108        }
109    }
110    Ok(result)
111}
112
113pub fn opt_in_with_branch(
114    db: &Database,
115    writer: &mut Writer,
116    experiment_slug: &str,
117    branch: &str,
118) -> Result<Vec<EnrollmentChangeEvent>> {
119    let mut events = vec![];
120    if let Ok(Some(exp)) = db
121        .get_store(StoreId::Experiments)
122        .get::<Experiment, Writer>(writer, experiment_slug)
123    {
124        let enrollment = ExperimentEnrollment::from_explicit_opt_in(&exp, branch, &mut events);
125        db.get_store(StoreId::Enrollments)
126            .put(writer, experiment_slug, &enrollment.unwrap())?;
127    } else {
128        events.push(EnrollmentChangeEvent {
129            experiment_slug: experiment_slug.to_string(),
130            branch_slug: branch.to_string(),
131            reason: Some("does-not-exist".to_string()),
132            change: EnrollmentChangeEventType::EnrollFailed,
133            feature_ids: vec![],
134        });
135    }
136
137    Ok(events)
138}
139
140fn get_enrollment_and_experiment(
141    db: &Database,
142    writer: &mut Writer,
143    experiment_slug: &str,
144) -> Result<(Option<ExperimentEnrollment>, Option<Experiment>)> {
145    // TODO(bug 2038055): Compute this using the database cache.
146    let maybe_enrollment: Option<ExperimentEnrollment> = db
147        .get_store(StoreId::Enrollments)
148        .get(writer, experiment_slug)?;
149    let maybe_experiment: Option<Experiment> = db
150        .get_store(StoreId::Experiments)
151        .get(writer, experiment_slug)?;
152
153    // We are technically guaranteed at this time that if an active enrollment
154    // exists in the enrollments store that the corresponding experiment must
155    // also exist in the experiment store.
156    //
157    // This is only not true during apply_pending_experiments.
158    Ok((maybe_enrollment, maybe_experiment))
159}
160
161pub fn opt_out(
162    db: &Database,
163    writer: &mut Writer,
164    experiment_slug: &str,
165    gecko_prefs: Option<&GeckoPrefStore>,
166) -> Result<Vec<EnrollmentChangeEvent>> {
167    let mut events = vec![];
168
169    match get_enrollment_and_experiment(db, writer, experiment_slug) {
170        Ok((Some(existing_enrollment), maybe_experiment)) => {
171            let updated_enrollment = &existing_enrollment.on_explicit_opt_out(
172                maybe_experiment.as_ref(),
173                &mut events,
174                gecko_prefs,
175            );
176
177            db.get_store(StoreId::Enrollments)
178                .put(writer, experiment_slug, updated_enrollment)?;
179        }
180
181        _ => {
182            events.push(EnrollmentChangeEvent {
183                experiment_slug: experiment_slug.to_string(),
184                branch_slug: "N/A".to_string(),
185                reason: Some("does-not-exist".to_string()),
186                change: EnrollmentChangeEventType::UnenrollFailed,
187                feature_ids: vec![],
188            });
189        }
190    }
191
192    Ok(events)
193}
194
195#[cfg(feature = "stateful")]
196pub fn unenroll_for_pref(
197    db: &Database,
198    writer: &mut Writer,
199    experiment_slug: &str,
200    unenroll_reason: PrefUnenrollReason,
201    triggering_pref_name: &str,
202    gecko_pref_store: Option<&GeckoPrefStore>,
203    events: &mut Vec<EnrollmentChangeEvent>,
204) -> Result<()> {
205    match get_enrollment_and_experiment(db, writer, experiment_slug) {
206        Ok((Some(existing_enrollment), maybe_experiment)) => {
207            existing_enrollment
208                .maybe_revert_unchanged_gecko_pref_states(triggering_pref_name, gecko_pref_store);
209
210            let updated_enrollment = &existing_enrollment.on_pref_unenroll(
211                unenroll_reason,
212                maybe_experiment.as_ref(),
213                events,
214            );
215            db.get_store(StoreId::Enrollments)
216                .put(writer, experiment_slug, updated_enrollment)?;
217        }
218
219        _ => {
220            events.push(EnrollmentChangeEvent {
221                experiment_slug: experiment_slug.to_string(),
222                branch_slug: "N/A".to_string(),
223                reason: Some("does-not-exist".to_string()),
224                change: EnrollmentChangeEventType::UnenrollFailed,
225                feature_ids: vec![],
226            });
227        }
228    }
229
230    Ok(())
231}
232
233pub fn get_experiment_participation<'r>(
234    db: &Database,
235    reader: &'r impl Readable<'r>,
236) -> Result<bool> {
237    let store = db.get_store(StoreId::Meta);
238    let opted_in = store.get::<bool, _>(reader, DB_KEY_EXPERIMENT_PARTICIPATION)?;
239    if let Some(opted_in) = opted_in {
240        Ok(opted_in)
241    } else {
242        Ok(DEFAULT_EXPERIMENT_PARTICIPATION)
243    }
244}
245
246pub fn get_rollout_participation<'r>(db: &Database, reader: &'r impl Readable<'r>) -> Result<bool> {
247    let store = db.get_store(StoreId::Meta);
248    let opted_in = store.get::<bool, _>(reader, DB_KEY_ROLLOUT_PARTICIPATION)?;
249    if let Some(opted_in) = opted_in {
250        Ok(opted_in)
251    } else {
252        Ok(DEFAULT_ROLLOUT_PARTICIPATION)
253    }
254}
255
256pub fn set_experiment_participation(
257    db: &Database,
258    writer: &mut Writer,
259    opt_in: bool,
260) -> Result<()> {
261    let store = db.get_store(StoreId::Meta);
262    store.put(writer, DB_KEY_EXPERIMENT_PARTICIPATION, &opt_in)
263}
264
265pub fn set_rollout_participation(db: &Database, writer: &mut Writer, opt_in: bool) -> Result<()> {
266    let store = db.get_store(StoreId::Meta);
267    store.put(writer, DB_KEY_ROLLOUT_PARTICIPATION, &opt_in)
268}
269
270/// Reset unique identifiers in response to application-level telemetry reset.
271///
272pub fn reset_telemetry_identifiers(
273    db: &Database,
274    writer: &mut Writer,
275) -> Result<Vec<EnrollmentChangeEvent>> {
276    let mut events = vec![];
277    let store = db.get_store(StoreId::Enrollments);
278    let enrollments: Vec<ExperimentEnrollment> = store.collect_all(writer)?;
279    // TODO(bug 2038055): Compute this using the database cache.
280    let experiments: Vec<Option<Experiment>> = enrollments
281        .iter()
282        .map(|enrollment| {
283            db.get_store(StoreId::Experiments)
284                .get::<Experiment, _>(writer, &enrollment.slug)
285        })
286        .collect::<Result<_>>()?;
287
288    let updated_enrollments =
289        iter::zip(enrollments, experiments).map(|(enrollment, experiment)| {
290            enrollment.reset_telemetry_identifiers(experiment.as_ref(), &mut events)
291        });
292    store.clear(writer)?;
293    for enrollment in updated_enrollments {
294        store.put(writer, &enrollment.slug, &enrollment)?;
295    }
296    Ok(events)
297}
nimbus/stateful/enrollment.rs

nimbus/stateful/
enrollment.rs
