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