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/. */
4use crate::{
5    enrollment::{
6        map_enrollments, EnrollmentChangeEvent, EnrollmentChangeEventType, EnrollmentsEvolver,
7        ExperimentEnrollment,
8    },
9    error::{debug, warn, Result},
10    stateful::{
11        gecko_prefs::PrefUnenrollReason,
12        persistence::{Database, Readable, StoreId, Writer},
13    },
14    EnrolledExperiment, EnrollmentStatus, Experiment,
15};
16
17const DB_KEY_GLOBAL_USER_PARTICIPATION: &str = "user-opt-in";
18const DEFAULT_GLOBAL_USER_PARTICIPATION: bool = true;
19
20impl EnrollmentsEvolver<'_> {
21    /// Convenient wrapper around `evolve_enrollments` that fetches the current state of experiments,
22    /// enrollments and user participation from the database.
23    pub(crate) fn evolve_enrollments_in_db(
24        &mut self,
25        db: &Database,
26        writer: &mut Writer,
27        next_experiments: &[Experiment],
28    ) -> Result<Vec<EnrollmentChangeEvent>> {
29        // Get the state from the db.
30        let is_user_participating = get_global_user_participation(db, writer)?;
31        let experiments_store = db.get_store(StoreId::Experiments);
32        let enrollments_store = db.get_store(StoreId::Enrollments);
33        let prev_experiments: Vec<Experiment> = experiments_store.collect_all(writer)?;
34        let prev_enrollments: Vec<ExperimentEnrollment> = enrollments_store.collect_all(writer)?;
35        // Calculate the changes.
36        let (next_enrollments, enrollments_change_events) = self.evolve_enrollments(
37            is_user_participating,
38            &prev_experiments,
39            next_experiments,
40            &prev_enrollments,
41        )?;
42        let next_enrollments = map_enrollments(&next_enrollments);
43        // Write the changes to the Database.
44        enrollments_store.clear(writer)?;
45        for enrollment in next_enrollments.values() {
46            enrollments_store.put(writer, &enrollment.slug, *enrollment)?;
47        }
48        experiments_store.clear(writer)?;
49        for experiment in next_experiments {
50            // Sanity check.
51            if !next_enrollments.contains_key(&experiment.slug) {
52                error_support::report_error!("nimbus-evolve-enrollments", "evolve_enrollments_in_db: experiment '{}' has no enrollment, dropping to keep database consistent", &experiment.slug);
53                continue;
54            }
55            experiments_store.put(writer, &experiment.slug, experiment)?;
56        }
57        Ok(enrollments_change_events)
58    }
59}
60
61/// Return information about all enrolled experiments.
62/// Note this does not include rollouts
63pub fn get_enrollments<'r>(
64    db: &Database,
65    reader: &'r impl Readable<'r>,
66) -> Result<Vec<EnrolledExperiment>> {
67    let enrollments: Vec<ExperimentEnrollment> =
68        db.get_store(StoreId::Enrollments).collect_all(reader)?;
69    let mut result = Vec::with_capacity(enrollments.len());
70    for enrollment in enrollments {
71        debug!("Have enrollment: {:?}", enrollment);
72        if let EnrollmentStatus::Enrolled { branch, .. } = &enrollment.status {
73            match db
74                .get_store(StoreId::Experiments)
75                .get::<Experiment, _>(reader, &enrollment.slug)?
76            {
77                Some(experiment) => {
78                    result.push(EnrolledExperiment {
79                        feature_ids: experiment.get_feature_ids(),
80                        slug: experiment.slug,
81                        user_facing_name: experiment.user_facing_name,
82                        user_facing_description: experiment.user_facing_description,
83                        branch_slug: branch.to_string(),
84                    });
85                }
86                _ => {
87                    warn!(
88                        "Have enrollment {:?} but no matching experiment!",
89                        enrollment
90                    );
91                }
92            };
93        }
94    }
95    Ok(result)
96}
97
98pub fn opt_in_with_branch(
99    db: &Database,
100    writer: &mut Writer,
101    experiment_slug: &str,
102    branch: &str,
103) -> Result<Vec<EnrollmentChangeEvent>> {
104    let mut events = vec![];
105    if let Ok(Some(exp)) = db
106        .get_store(StoreId::Experiments)
107        .get::<Experiment, Writer>(writer, experiment_slug)
108    {
109        let enrollment = ExperimentEnrollment::from_explicit_opt_in(&exp, branch, &mut events);
110        db.get_store(StoreId::Enrollments)
111            .put(writer, experiment_slug, &enrollment.unwrap())?;
112    } else {
113        events.push(EnrollmentChangeEvent {
114            experiment_slug: experiment_slug.to_string(),
115            branch_slug: branch.to_string(),
116            reason: Some("does-not-exist".to_string()),
117            change: EnrollmentChangeEventType::EnrollFailed,
118        });
119    }
120
121    Ok(events)
122}
123
124pub fn opt_out(
125    db: &Database,
126    writer: &mut Writer,
127    experiment_slug: &str,
128) -> Result<Vec<EnrollmentChangeEvent>> {
129    let mut events = vec![];
130    let enr_store = db.get_store(StoreId::Enrollments);
131    if let Ok(Some(existing_enrollment)) =
132        enr_store.get::<ExperimentEnrollment, Writer>(writer, experiment_slug)
133    {
134        let updated_enrollment = &existing_enrollment.on_explicit_opt_out(&mut events);
135        enr_store.put(writer, experiment_slug, updated_enrollment)?;
136    } else {
137        events.push(EnrollmentChangeEvent {
138            experiment_slug: experiment_slug.to_string(),
139            branch_slug: "N/A".to_string(),
140            reason: Some("does-not-exist".to_string()),
141            change: EnrollmentChangeEventType::UnenrollFailed,
142        });
143    }
144
145    Ok(events)
146}
147
148pub fn unenroll_for_pref(
149    db: &Database,
150    writer: &mut Writer,
151    experiment_slug: &str,
152    unenroll_reason: PrefUnenrollReason,
153) -> Result<Vec<EnrollmentChangeEvent>> {
154    let mut events = vec![];
155    let enr_store = db.get_store(StoreId::Enrollments);
156    if let Ok(Some(existing_enrollment)) =
157        enr_store.get::<ExperimentEnrollment, Writer>(writer, experiment_slug)
158    {
159        let updated_enrollment =
160            &existing_enrollment.on_pref_unenroll(unenroll_reason, &mut events);
161        enr_store.put(writer, experiment_slug, updated_enrollment)?;
162    } else {
163        events.push(EnrollmentChangeEvent {
164            experiment_slug: experiment_slug.to_string(),
165            branch_slug: "N/A".to_string(),
166            reason: Some("does-not-exist".to_string()),
167            change: EnrollmentChangeEventType::UnenrollFailed,
168        });
169    }
170
171    Ok(events)
172}
173
174pub fn get_global_user_participation<'r>(
175    db: &Database,
176    reader: &'r impl Readable<'r>,
177) -> Result<bool> {
178    let store = db.get_store(StoreId::Meta);
179    let opted_in = store.get::<bool, _>(reader, DB_KEY_GLOBAL_USER_PARTICIPATION)?;
180    if let Some(opted_in) = opted_in {
181        Ok(opted_in)
182    } else {
183        Ok(DEFAULT_GLOBAL_USER_PARTICIPATION)
184    }
185}
186
187pub fn set_global_user_participation(
188    db: &Database,
189    writer: &mut Writer,
190    opt_in: bool,
191) -> Result<()> {
192    let store = db.get_store(StoreId::Meta);
193    store.put(writer, DB_KEY_GLOBAL_USER_PARTICIPATION, &opt_in)
194}
195
196/// Reset unique identifiers in response to application-level telemetry reset.
197///
198pub fn reset_telemetry_identifiers(
199    db: &Database,
200    writer: &mut Writer,
201) -> Result<Vec<EnrollmentChangeEvent>> {
202    let mut events = vec![];
203    let store = db.get_store(StoreId::Enrollments);
204    let enrollments: Vec<ExperimentEnrollment> = store.collect_all(writer)?;
205    let updated_enrollments = enrollments
206        .iter()
207        .map(|enrollment| enrollment.reset_telemetry_identifiers(&mut events));
208    store.clear(writer)?;
209    for enrollment in updated_enrollments {
210        store.put(writer, &enrollment.slug, &enrollment)?;
211    }
212    Ok(events)
213}