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

nimbus/stateful/
enrollment.rs
