nimbus/stateful/
enrollment.rs
1use 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 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 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 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 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 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
58pub 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
167pub 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}