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

nimbus/stateful/
enrollment.rs
