1use 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 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 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 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 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 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
61pub 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
196pub 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}