nimbus/stateful/
nimbus_client.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 std::collections::HashSet;
6use std::fmt::Debug;
7use std::path::PathBuf;
8use std::sync::{Arc, Mutex, MutexGuard};
9
10use chrono::{DateTime, NaiveDateTime, Utc};
11use once_cell::sync::OnceCell;
12use remote_settings::RemoteSettingsService;
13use serde_json::Value;
14use uuid::Uuid;
15
16use crate::defaults::Defaults;
17use crate::enrollment::{
18    EnrolledFeature, EnrollmentChangeEvent, EnrollmentChangeEventType, EnrollmentsEvolver,
19    ExperimentEnrollment, PreviousGeckoPrefState,
20};
21use crate::error::{BehaviorError, info, warn};
22use crate::evaluator::{
23    CalculatedAttributes, ExperimentAvailable, TargetingAttributes, get_calculated_attributes,
24    is_experiment_available,
25};
26use crate::json::{JsonObject, PrefValue};
27use crate::metrics::{
28    DatabaseLoadExtraDef, DatabaseMigrationExtraDef, EnrollmentStatusExtraDef,
29    FeatureExposureExtraDef, MalformedFeatureConfigExtraDef, MetricsHandler,
30};
31use crate::schema::parse_experiments;
32use crate::stateful::behavior::EventStore;
33use crate::stateful::client::{NimbusServerSettings, SettingsClient, create_client};
34use crate::stateful::dbcache::DatabaseCache;
35use crate::stateful::enrollment::{
36    enroll_in_firefox_lab, get_experiment_participation, get_rollout_participation,
37    opt_in_with_branch, opt_out, reset_telemetry_identifiers, set_experiment_participation,
38    set_rollout_participation, unenroll_for_pref, unenroll_from_all_firefox_labs,
39    unenroll_from_firefox_lab,
40};
41use crate::stateful::firefox_labs::{
42    FirefoxLabsEnrollResult, FirefoxLabsEnrollStatus, FirefoxLabsMetadata,
43    FirefoxLabsUnenrollResult, FirefoxLabsUnenrollStatus,
44};
45use crate::stateful::gecko_prefs::{
46    GeckoPref, GeckoPrefHandler, GeckoPrefState, GeckoPrefStore, OriginalGeckoPref, PrefBranch,
47    PrefEnrollmentData, PrefUnenrollReason,
48};
49use crate::stateful::matcher::AppContext;
50use crate::stateful::persistence::{Database, StoreId, Writer};
51use crate::stateful::targeting::{RecordedContext, execute_event_queries, validate_event_queries};
52use crate::stateful::updating::{read_and_remove_pending_experiments, write_pending_experiments};
53use crate::strings::fmt_with_map;
54use crate::{
55    AvailableExperiment, AvailableRandomizationUnits, EnrolledExperiment, EnrollmentStatus,
56};
57use crate::{Experiment, ExperimentBranch, NimbusError, NimbusTargetingHelper, Result};
58
59const DB_KEY_NIMBUS_ID: &str = "nimbus-id";
60pub const DB_KEY_INSTALLATION_DATE: &str = "installation-date";
61pub const DB_KEY_UPDATE_DATE: &str = "update-date";
62pub const DB_KEY_APP_VERSION: &str = "app-version";
63pub const DB_KEY_FETCH_ENABLED: &str = "fetch-enabled";
64
65// The main `NimbusClient` struct must not expose any methods that make an `&mut self`,
66// in order to be compatible with the uniffi's requirements on objects. This is a helper
67// struct to contain the bits that do actually need to be mutable, so they can be
68// protected by a Mutex.
69#[derive(Default)]
70pub struct InternalMutableState {
71    pub(crate) available_randomization_units: AvailableRandomizationUnits,
72    pub(crate) install_date: Option<DateTime<Utc>>,
73    pub(crate) update_date: Option<DateTime<Utc>>,
74    // Application level targeting attributes
75    pub(crate) targeting_attributes: TargetingAttributes,
76}
77
78impl InternalMutableState {
79    pub(crate) fn update_time_to_now(&mut self, now: DateTime<Utc>) {
80        self.targeting_attributes
81            .update_time_to_now(now, &self.install_date, &self.update_date);
82    }
83}
84
85/// Nimbus is the main struct representing the experiments state
86/// It should hold all the information needed to communicate a specific user's
87/// experimentation status
88pub struct NimbusClient {
89    settings_client: Mutex<Box<dyn SettingsClient + Send>>,
90    pub(crate) mutable_state: Mutex<InternalMutableState>,
91    app_context: AppContext,
92    pub(crate) db: OnceCell<Database>,
93    // Manages an in-memory cache so that we can answer certain requests
94    // without doing (or waiting for) IO.
95    database_cache: DatabaseCache,
96    db_path: PathBuf,
97    coenrolling_feature_ids: Vec<String>,
98    event_store: Arc<Mutex<EventStore>>,
99    recorded_context: Option<Arc<dyn RecordedContext>>,
100    pub(crate) gecko_prefs: Option<Arc<GeckoPrefStore>>,
101    metrics_handler: Arc<dyn MetricsHandler>,
102}
103
104impl NimbusClient {
105    // This constructor *must* not do any kind of I/O since it might be called on the main
106    // thread in the gecko Javascript stack, hence the use of OnceCell for the db.
107    #[allow(clippy::too_many_arguments)]
108    pub fn new<P: Into<PathBuf>>(
109        app_context: AppContext,
110        recorded_context: Option<Arc<dyn RecordedContext>>,
111        coenrolling_feature_ids: Vec<String>,
112        db_path: P,
113        metrics_handler: Arc<dyn MetricsHandler>,
114        gecko_pref_handler: Option<Arc<dyn GeckoPrefHandler>>,
115        remote_settings_info: Option<NimbusServerSettings>,
116    ) -> Result<Self> {
117        let settings_client = Mutex::new(create_client(remote_settings_info)?);
118
119        let targeting_attributes: TargetingAttributes = app_context.clone().into();
120        let mutable_state = Mutex::new(InternalMutableState {
121            available_randomization_units: Default::default(),
122            targeting_attributes,
123            install_date: Default::default(),
124            update_date: Default::default(),
125        });
126
127        let mut prefs = None;
128        if let Some(handler) = gecko_pref_handler {
129            prefs = Some(Arc::new(GeckoPrefStore::new(handler)));
130        }
131
132        info!(
133            "Initialized NimbusClient with: app_context = {:?}; recorded_context = {:?}",
134            app_context,
135            recorded_context
136                .as_ref()
137                .map(|rc| serde_json::Value::Object(rc.to_json()))
138                .unwrap_or(serde_json::Value::Null)
139        );
140
141        Ok(Self {
142            settings_client,
143            mutable_state,
144            app_context,
145            database_cache: Default::default(),
146            db_path: db_path.into(),
147            coenrolling_feature_ids,
148            db: OnceCell::default(),
149            event_store: Arc::default(),
150            recorded_context,
151            gecko_prefs: prefs,
152            metrics_handler,
153        })
154    }
155
156    pub fn with_targeting_attributes(&mut self, targeting_attributes: TargetingAttributes) {
157        let mut state = self.mutable_state.lock().unwrap();
158        state.targeting_attributes = targeting_attributes;
159    }
160
161    pub fn get_targeting_attributes(&self) -> TargetingAttributes {
162        let mut state = self.mutable_state.lock().unwrap();
163        state.update_time_to_now(Utc::now());
164        state.targeting_attributes.clone()
165    }
166
167    pub fn initialize(&self) -> Result<()> {
168        let db = self.db()?;
169        // We're not actually going to write, we just want to exclude concurrent writers.
170        let mut writer = db.write()?;
171
172        let mut state = self.mutable_state.lock().unwrap();
173        self.begin_initialize(db, &mut writer, &mut state)?;
174        self.end_initialize(db, writer, &mut state)?;
175
176        Ok(())
177    }
178
179    // These are tasks which should be in the initialize and apply_pending_experiments
180    // but should happen before the enrollment calculations are done.
181    fn begin_initialize(
182        &self,
183        db: &Database,
184        writer: &mut Writer,
185        state: &mut MutexGuard<InternalMutableState>,
186    ) -> Result<()> {
187        self.read_or_create_nimbus_id(db, writer, state)?;
188        self.update_ta_install_dates(db, writer, state)?;
189        self.event_store
190            .lock()
191            .expect("unable to lock event_store mutex")
192            .read_from_db(db)?;
193
194        if let Some(recorded_context) = &self.recorded_context {
195            let targeting_helper = self.create_targeting_helper_with_context(match serde_json::to_value(
196                &state.targeting_attributes,
197            ) {
198                Ok(v) => v,
199                Err(e) => return Err(NimbusError::JSONError("targeting_helper = nimbus::stateful::nimbus_client::NimbusClient::begin_initialize::serde_json::to_value".into(), e.to_string()))
200            });
201            execute_event_queries(&**recorded_context, targeting_helper.as_ref())?;
202            state
203                .targeting_attributes
204                .set_recorded_context(recorded_context.to_json());
205        }
206
207        if let Some(gecko_prefs) = &self.gecko_prefs {
208            gecko_prefs.initialize()?;
209        }
210
211        Ok(())
212    }
213
214    // These are tasks which should be in the initialize and apply_pending_experiments
215    // but should happen after the enrollment calculations are done.
216    fn end_initialize(
217        &self,
218        db: &Database,
219        writer: Writer,
220        state: &mut MutexGuard<InternalMutableState>,
221    ) -> Result<()> {
222        self.update_ta_active_experiments(db, &writer, state)?;
223        let coenrolling_ids = self
224            .coenrolling_feature_ids
225            .iter()
226            .map(|s| s.as_str())
227            .collect();
228        self.database_cache.commit_and_update(
229            db,
230            writer,
231            &coenrolling_ids,
232            self.gecko_prefs.clone(),
233            true,
234        )?;
235        Ok(())
236    }
237
238    pub fn get_enrollment_by_feature(&self, feature_id: String) -> Result<Option<EnrolledFeature>> {
239        self.database_cache.get_enrollment_by_feature(&feature_id)
240    }
241
242    // Note: the contract for this function is that it never blocks on IO.
243    pub fn get_experiment_branch(&self, slug: String) -> Result<Option<String>> {
244        self.database_cache.get_experiment_branch(&slug)
245    }
246
247    pub fn get_feature_config_variables(&self, feature_id: String) -> Result<Option<String>> {
248        Ok(
249            if let Some(s) = self
250                .database_cache
251                .get_feature_config_variables(&feature_id)?
252            {
253                self.record_feature_activation_if_needed(&feature_id);
254                Some(s)
255            } else {
256                None
257            },
258        )
259    }
260
261    pub fn get_experiment_branches(&self, slug: String) -> Result<Vec<ExperimentBranch>> {
262        self.get_all_experiments()?
263            .into_iter()
264            .find(|e| e.slug == slug)
265            .map(|e| e.branches.into_iter().map(|b| b.into()).collect())
266            .ok_or(NimbusError::NoSuchExperiment(slug))
267    }
268
269    pub fn get_experiment_participation(&self) -> Result<bool> {
270        let db = self.db()?;
271        let reader = db.read()?;
272        get_experiment_participation(db, &reader)
273    }
274
275    pub fn get_rollout_participation(&self) -> Result<bool> {
276        let db = self.db()?;
277        let reader = db.read()?;
278        get_rollout_participation(db, &reader)
279    }
280
281    pub fn set_experiment_participation(
282        &self,
283        user_participating: bool,
284    ) -> Result<Vec<EnrollmentChangeEvent>> {
285        let db = self.db()?;
286        let mut writer = db.write()?;
287        let mut state = self.mutable_state.lock().unwrap();
288        set_experiment_participation(db, &mut writer, user_participating)?;
289
290        let existing_experiments: Vec<Experiment> =
291            db.get_store(StoreId::Experiments).collect_all(&writer)?;
292        let events = self.evolve_experiments(db, &mut writer, &mut state, &existing_experiments)?;
293        let res = self.end_initialize(db, writer, &mut state);
294        self.record_enrollment_status_telemetry(&mut state)?;
295        res?;
296        Ok(events)
297    }
298
299    pub fn set_rollout_participation(
300        &self,
301        user_participating: bool,
302    ) -> Result<Vec<EnrollmentChangeEvent>> {
303        let db = self.db()?;
304        let mut writer = db.write()?;
305        let mut state = self.mutable_state.lock().unwrap();
306        set_rollout_participation(db, &mut writer, user_participating)?;
307
308        let existing_experiments: Vec<Experiment> =
309            db.get_store(StoreId::Experiments).collect_all(&writer)?;
310        let events = self.evolve_experiments(db, &mut writer, &mut state, &existing_experiments)?;
311        let res = self.end_initialize(db, writer, &mut state);
312        self.record_enrollment_status_telemetry(&mut state)?;
313        res?;
314        Ok(events)
315    }
316
317    pub fn get_active_experiments(&self) -> Result<Vec<EnrolledExperiment>> {
318        self.database_cache.get_active_experiments()
319    }
320
321    pub fn get_all_experiments(&self) -> Result<Vec<Experiment>> {
322        let db = self.db()?;
323        let reader = db.read()?;
324        db.get_store(StoreId::Experiments)
325            .collect_all::<Experiment, _>(&reader)
326    }
327
328    pub fn get_available_experiments(&self) -> Result<Vec<AvailableExperiment>> {
329        let th = self.create_targeting_helper(None)?;
330        Ok(self
331            .get_all_experiments()?
332            .into_iter()
333            .filter(|exp| {
334                is_experiment_available(&th, exp, false) == ExperimentAvailable::Available
335            })
336            .map(|exp| exp.into())
337            .collect())
338    }
339
340    pub fn opt_in_with_branch(
341        &self,
342        experiment_slug: String,
343        branch: String,
344    ) -> Result<Vec<EnrollmentChangeEvent>> {
345        let db = self.db()?;
346        let mut writer = db.write()?;
347        let result = opt_in_with_branch(db, &mut writer, &experiment_slug, &branch)?;
348        let mut state = self.mutable_state.lock().unwrap();
349        self.end_initialize(db, writer, &mut state)?;
350        Ok(result)
351    }
352
353    pub fn opt_out(&self, experiment_slug: String) -> Result<Vec<EnrollmentChangeEvent>> {
354        let db = self.db()?;
355        let mut writer = db.write()?;
356        let result = opt_out(
357            db,
358            &mut writer,
359            &experiment_slug,
360            self.gecko_prefs.as_deref(),
361        )?;
362        let mut state = self.mutable_state.lock().unwrap();
363        self.end_initialize(db, writer, &mut state)?;
364        Ok(result)
365    }
366
367    pub fn fetch_experiments(&self) -> Result<()> {
368        if !self.is_fetch_enabled()? {
369            return Ok(());
370        }
371        info!("fetching experiments");
372        let settings_client = self.settings_client.lock().unwrap();
373        let new_experiments = settings_client.fetch_experiments()?;
374        let db = self.db()?;
375        let mut writer = db.write()?;
376        write_pending_experiments(db, &mut writer, new_experiments)?;
377        writer.commit()?;
378        Ok(())
379    }
380
381    pub fn set_fetch_enabled(&self, allow: bool) -> Result<()> {
382        let db = self.db()?;
383        let mut writer = db.write()?;
384        db.get_store(StoreId::Meta)
385            .put(&mut writer, DB_KEY_FETCH_ENABLED, &allow)?;
386        writer.commit()?;
387        Ok(())
388    }
389
390    pub(crate) fn is_fetch_enabled(&self) -> Result<bool> {
391        let db = self.db()?;
392        let reader = db.read()?;
393        let enabled = db
394            .get_store(StoreId::Meta)
395            .get(&reader, DB_KEY_FETCH_ENABLED)?
396            .unwrap_or(true);
397        Ok(enabled)
398    }
399
400    /**
401     * Calculate the days since install and days since update on the targeting_attributes.
402     */
403    fn update_ta_install_dates(
404        &self,
405        db: &Database,
406        writer: &mut Writer,
407        state: &mut MutexGuard<InternalMutableState>,
408    ) -> Result<()> {
409        // Only set install_date and update_date with this method if it hasn't been set already.
410        // This cuts down on deriving the dates at runtime, but also allows us to use
411        // the test methods set_install_date() and set_update_date() to set up
412        // scenarios for test.
413        if state.install_date.is_none() {
414            let installation_date = self.get_installation_date(db, writer)?;
415            state.install_date = Some(installation_date);
416        }
417        if state.update_date.is_none() {
418            let update_date = self.get_update_date(db, writer)?;
419            state.update_date = Some(update_date);
420        }
421        state.update_time_to_now(Utc::now());
422
423        Ok(())
424    }
425
426    /**
427     * Calculates the active_experiments based on current enrollments for the targeting attributes.
428     */
429    fn update_ta_active_experiments(
430        &self,
431        db: &Database,
432        writer: &Writer,
433        state: &mut MutexGuard<InternalMutableState>,
434    ) -> Result<()> {
435        let enrollments_store = db.get_store(StoreId::Enrollments);
436        let prev_enrollments: Vec<ExperimentEnrollment> = enrollments_store.collect_all(writer)?;
437
438        state
439            .targeting_attributes
440            .update_enrollments(&prev_enrollments);
441
442        Ok(())
443    }
444
445    fn evolve_experiments(
446        &self,
447        db: &Database,
448        writer: &mut Writer,
449        state: &mut InternalMutableState,
450        experiments: &[Experiment],
451    ) -> Result<Vec<EnrollmentChangeEvent>> {
452        let mut targeting_helper = NimbusTargetingHelper::with_targeting_attributes(
453            &state.targeting_attributes,
454            self.event_store.clone(),
455            self.gecko_prefs.clone(),
456        );
457        if let Some(ref recorded_context) = self.recorded_context {
458            recorded_context.record();
459        }
460        let coenrolling_feature_ids = self
461            .coenrolling_feature_ids
462            .iter()
463            .map(|s| s.as_str())
464            .collect();
465        let mut evolver = EnrollmentsEvolver::new(
466            &state.available_randomization_units,
467            &mut targeting_helper,
468            &coenrolling_feature_ids,
469        );
470        evolver.evolve_enrollments_in_db(db, writer, experiments, self.gecko_prefs.as_deref())
471    }
472
473    pub fn apply_pending_experiments(&self) -> Result<Vec<EnrollmentChangeEvent>> {
474        info!("updating experiment list");
475        let db = self.db()?;
476        let mut writer = db.write()?;
477
478        // We'll get the pending experiments which were stored for us, either by fetch_experiments
479        // or by set_experiments_locally.
480        let pending_updates = read_and_remove_pending_experiments(db, &mut writer)?;
481        let mut state = self.mutable_state.lock().unwrap();
482        self.begin_initialize(db, &mut writer, &mut state)?;
483
484        let should_record_enrollment_status = pending_updates.is_some();
485        let res = match pending_updates {
486            Some(new_experiments) => {
487                self.update_ta_active_experiments(db, &writer, &mut state)?;
488                // Perform the enrollment calculations if there are pending experiments.
489                self.evolve_experiments(db, &mut writer, &mut state, &new_experiments)?
490            }
491            None => vec![],
492        };
493
494        // Finish up any cleanup, e.g. copying from database in to memory.
495        let end_init_res = self.end_initialize(db, writer, &mut state);
496        if should_record_enrollment_status {
497            self.record_enrollment_status_telemetry(&mut state)?;
498        }
499        end_init_res?;
500        Ok(res)
501    }
502
503    #[allow(deprecated)] // Bug 1960256 - use of deprecated chrono functions.
504    fn get_installation_date(&self, db: &Database, writer: &mut Writer) -> Result<DateTime<Utc>> {
505        // we first check our context
506        if let Some(context_installation_date) = self.app_context.installation_date {
507            let res = DateTime::<Utc>::from_naive_utc_and_offset(
508                NaiveDateTime::from_timestamp_opt(context_installation_date / 1_000, 0).unwrap(),
509                Utc,
510            );
511            info!("[Nimbus] Retrieved date from Context: {}", res);
512            return Ok(res);
513        }
514        let store = db.get_store(StoreId::Meta);
515        let persisted_installation_date: Option<DateTime<Utc>> =
516            store.get(writer, DB_KEY_INSTALLATION_DATE)?;
517        Ok(
518            if let Some(installation_date) = persisted_installation_date {
519                installation_date
520            } else {
521                Utc::now()
522            },
523        )
524    }
525
526    fn get_update_date(&self, db: &Database, writer: &mut Writer) -> Result<DateTime<Utc>> {
527        let store = db.get_store(StoreId::Meta);
528
529        let persisted_app_version: Option<String> = store.get(writer, DB_KEY_APP_VERSION)?;
530        let update_date: Option<DateTime<Utc>> = store.get(writer, DB_KEY_UPDATE_DATE)?;
531        Ok(
532            match (
533                persisted_app_version,
534                &self.app_context.app_version,
535                update_date,
536            ) {
537                // The app been run before, but has not just been updated.
538                (Some(persisted), Some(current), Some(date)) if persisted == *current => date,
539                // The app has been run before, and just been updated.
540                (Some(persisted), Some(current), _) if persisted != *current => {
541                    let now = Utc::now();
542                    store.put(writer, DB_KEY_APP_VERSION, current)?;
543                    store.put(writer, DB_KEY_UPDATE_DATE, &now)?;
544                    now
545                }
546                // The app has just been installed
547                (None, Some(current), _) => {
548                    let now = Utc::now();
549                    store.put(writer, DB_KEY_APP_VERSION, current)?;
550                    store.put(writer, DB_KEY_UPDATE_DATE, &now)?;
551                    now
552                }
553                // The current version is not available, or the persisted date is not available.
554                (_, _, Some(date)) => date,
555                // Either way, this doesn't appear to be a good production environment.
556                _ => Utc::now(),
557            },
558        )
559    }
560
561    pub fn set_experiments_locally(&self, experiments_json: String) -> Result<()> {
562        let new_experiments = parse_experiments(&experiments_json)?;
563        let db = self.db()?;
564        let mut writer = db.write()?;
565        write_pending_experiments(db, &mut writer, new_experiments)?;
566        writer.commit()?;
567        Ok(())
568    }
569
570    /// Reset all enrollments and experiments in the database.
571    ///
572    /// This should only be used in testing.
573    pub fn reset_enrollments(&self) -> Result<()> {
574        let db = self.db()?;
575        let mut writer = db.write()?;
576        let mut state = self.mutable_state.lock().unwrap();
577        db.clear_experiments_and_enrollments(&mut writer)?;
578        self.end_initialize(db, writer, &mut state)?;
579        Ok(())
580    }
581
582    /// Reset internal state in response to application-level telemetry reset.
583    ///
584    /// When the user resets their telemetry state in the consuming application, we need learn
585    /// the new values of any external randomization units, and we need to reset any unique
586    /// identifiers used internally by the SDK. If we don't then we risk accidentally tracking
587    /// across the telemetry reset, since we could use Nimbus metrics to link their pings from
588    /// before and after the reset.
589    ///
590    pub fn reset_telemetry_identifiers(&self) -> Result<Vec<EnrollmentChangeEvent>> {
591        let mut events = vec![];
592        let db = self.db()?;
593        let mut writer = db.write()?;
594        let mut state = self.mutable_state.lock().unwrap();
595        // If we have no `nimbus_id` when we can safely assume that there's
596        // no other experiment state that needs to be reset.
597        let store = db.get_store(StoreId::Meta);
598        if store.get::<String, _>(&writer, DB_KEY_NIMBUS_ID)?.is_some() {
599            // Each enrollment state now opts out because we don't want to leak information between resets.
600            events = reset_telemetry_identifiers(db, &mut writer)?;
601
602            // Remove any stored event counts
603            db.clear_event_count_data(&mut writer)?;
604
605            // The `nimbus_id` itself is a unique identifier.
606            // N.B. we do this last, as a signal that all data has been reset.
607            store.delete(&mut writer, DB_KEY_NIMBUS_ID)?;
608            self.end_initialize(db, writer, &mut state)?;
609        }
610
611        // (No need to commit `writer` if the above check was false, since we didn't change anything)
612        state.available_randomization_units = Default::default();
613        state.targeting_attributes.nimbus_id = None;
614
615        Ok(events)
616    }
617
618    pub fn nimbus_id(&self) -> Result<Uuid> {
619        let db = self.db()?;
620        let mut writer = db.write()?;
621        let mut state = self.mutable_state.lock().unwrap();
622        let uuid = self.read_or_create_nimbus_id(db, &mut writer, &mut state)?;
623
624        // We don't know whether we needed to generate and save the uuid, so
625        // we commit just in case - this is hopefully close to a noop in that
626        // case!
627        writer.commit()?;
628        Ok(uuid)
629    }
630
631    /// Return the nimbus ID from the database, or create a new one and write it
632    /// to the database.
633    ///
634    /// The internal state will be updated with the nimbus ID.
635    fn read_or_create_nimbus_id(
636        &self,
637        db: &Database,
638        writer: &mut Writer,
639        state: &mut MutexGuard<'_, InternalMutableState>,
640    ) -> Result<Uuid> {
641        let store = db.get_store(StoreId::Meta);
642        let nimbus_id = match store.get(writer, DB_KEY_NIMBUS_ID)? {
643            Some(nimbus_id) => nimbus_id,
644            None => {
645                let nimbus_id = Uuid::new_v4();
646                store.put(writer, DB_KEY_NIMBUS_ID, &nimbus_id)?;
647                nimbus_id
648            }
649        };
650
651        state.available_randomization_units.nimbus_id = Some(nimbus_id.to_string());
652        state.targeting_attributes.nimbus_id = Some(nimbus_id.to_string());
653
654        Ok(nimbus_id)
655    }
656
657    // Sets the nimbus ID - TEST ONLY - should not be exposed to real clients.
658    // (Useful for testing so you can have some control over what experiments
659    // are enrolled)
660    pub fn set_nimbus_id(&self, uuid: &Uuid) -> Result<()> {
661        let db = self.db()?;
662        let mut writer = db.write()?;
663        db.get_store(StoreId::Meta)
664            .put(&mut writer, DB_KEY_NIMBUS_ID, uuid)?;
665        writer.commit()?;
666        Ok(())
667    }
668
669    pub(crate) fn db(&self) -> Result<&Database> {
670        self.db
671            .get_or_try_init(|| Database::new(&self.db_path, self.metrics_handler.clone()))
672    }
673
674    fn merge_additional_context(&self, context: Option<JsonObject>) -> Result<Value> {
675        let context = context.map(Value::Object);
676        let targeting = match serde_json::to_value(self.get_targeting_attributes()) {
677            Ok(v) => v,
678            Err(e) => return Err(NimbusError::JSONError("targeting = nimbus::stateful::nimbus_client::NimbusClient::merge_additional_context::serde_json::to_value".into(), e.to_string()))
679        };
680        let context = match context {
681            Some(v) => v.defaults(&targeting)?,
682            None => targeting,
683        };
684
685        Ok(context)
686    }
687
688    pub fn create_targeting_helper(
689        &self,
690        additional_context: Option<JsonObject>,
691    ) -> Result<Arc<NimbusTargetingHelper>> {
692        let context = self.merge_additional_context(additional_context)?;
693        let helper =
694            NimbusTargetingHelper::new(context, self.event_store.clone(), self.gecko_prefs.clone());
695        Ok(Arc::new(helper))
696    }
697
698    pub fn create_targeting_helper_with_context(
699        &self,
700        context: Value,
701    ) -> Arc<NimbusTargetingHelper> {
702        Arc::new(NimbusTargetingHelper::new(
703            context,
704            self.event_store.clone(),
705            self.gecko_prefs.clone(),
706        ))
707    }
708
709    pub fn create_string_helper(
710        &self,
711        additional_context: Option<JsonObject>,
712    ) -> Result<Arc<NimbusStringHelper>> {
713        let context = self.merge_additional_context(additional_context)?;
714        let helper = NimbusStringHelper::new(context.as_object().unwrap().to_owned());
715        Ok(Arc::new(helper))
716    }
717
718    /// Records an event for the purposes of behavioral targeting.
719    ///
720    /// This function is used to record and persist data used for the behavioral
721    /// targeting such as "core-active" user targeting.
722    pub fn record_event(&self, event_id: String, count: i64) -> Result<()> {
723        let mut event_store = self.event_store.lock().unwrap();
724        event_store.record_event(count as u64, &event_id, None)?;
725        event_store.persist_data(self.db()?)?;
726        Ok(())
727    }
728
729    /// Records an event for the purposes of behavioral targeting.
730    ///
731    /// This differs from the `record_event` method in that the event is recorded as if it were
732    /// recorded `seconds_ago` in the past. This makes it very useful for testing.
733    pub fn record_past_event(&self, event_id: String, seconds_ago: i64, count: i64) -> Result<()> {
734        if seconds_ago < 0 {
735            return Err(NimbusError::BehaviorError(BehaviorError::InvalidDuration(
736                "Time duration in the past must be positive".to_string(),
737            )));
738        }
739        let mut event_store = self.event_store.lock().unwrap();
740        event_store.record_past_event(
741            count as u64,
742            &event_id,
743            None,
744            chrono::Duration::seconds(seconds_ago),
745        )?;
746        event_store.persist_data(self.db()?)?;
747        Ok(())
748    }
749
750    /// Advances the event store's concept of `now` artificially.
751    ///
752    /// This works alongside `record_event` and `record_past_event` for testing purposes.
753    pub fn advance_event_time(&self, by_seconds: i64) -> Result<()> {
754        if by_seconds < 0 {
755            return Err(NimbusError::BehaviorError(BehaviorError::InvalidDuration(
756                "Time duration in the future must be positive".to_string(),
757            )));
758        }
759        let mut event_store = self.event_store.lock().unwrap();
760        event_store.advance_datum(chrono::Duration::seconds(by_seconds));
761        Ok(())
762    }
763
764    /// Clear all events in the Nimbus event store.
765    ///
766    /// This should only be used in testing or cases where the previous event store is no longer viable.
767    pub fn clear_events(&self) -> Result<()> {
768        let mut event_store = self.event_store.lock().unwrap();
769        event_store.clear(self.db()?)?;
770        Ok(())
771    }
772
773    pub fn event_store(&self) -> Arc<Mutex<EventStore>> {
774        self.event_store.clone()
775    }
776
777    pub fn dump_state_to_log(&self) -> Result<()> {
778        let experiments = self.get_active_experiments()?;
779        info!("{0: <65}| {1: <30}| {2}", "Slug", "Features", "Branch");
780        for exp in &experiments {
781            info!(
782                "{0: <65}| {1: <30}| {2}",
783                &exp.slug,
784                &exp.feature_ids.join(", "),
785                &exp.branch_slug
786            );
787        }
788        Ok(())
789    }
790
791    /// Given a Gecko pref state and a pref unenroll reason, unenroll from an experiment
792    pub fn unenroll_for_gecko_pref(
793        &self,
794        pref_state: GeckoPrefState,
795        pref_unenroll_reason: PrefUnenrollReason,
796    ) -> Result<Vec<EnrollmentChangeEvent>> {
797        let mut events = Vec::new();
798        if let Some(prefs) = self.gecko_prefs.clone() {
799            {
800                let mut pref_store_state = prefs.get_mutable_pref_state();
801                pref_store_state.update_pref_state(&pref_state);
802            }
803            let enrollments = self
804                .database_cache
805                .get_enrollments_for_pref(&pref_state.gecko_pref.pref)?;
806
807            let db = self.db()?;
808            let mut writer = db.write()?;
809
810            if let Some(enrollments) = enrollments {
811                for experiment_slug in enrollments {
812                    unenroll_for_pref(
813                        db,
814                        &mut writer,
815                        &experiment_slug,
816                        pref_unenroll_reason,
817                        &pref_state.gecko_pref.pref,
818                        self.gecko_prefs.as_deref(),
819                        &mut events,
820                    )?;
821                }
822            } else {
823                warn!(
824                    "Could not find enrollment. Could unenrollment already occurred through another preference?"
825                )
826            }
827
828            let mut state = self.mutable_state.lock().unwrap();
829            self.end_initialize(db, writer, &mut state)?;
830        }
831        Ok(events)
832    }
833
834    pub fn register_previous_gecko_pref_states(
835        &self,
836        gecko_pref_states: &[GeckoPrefState],
837    ) -> Result<()> {
838        let all_prev_gecko_pref_states =
839            super::gecko_prefs::build_prev_gecko_pref_states(gecko_pref_states);
840
841        let db = self.db()?;
842        let mut writer = db.write()?;
843
844        for (experiment_slug, prev_gecko_pref_states) in all_prev_gecko_pref_states {
845            Self::add_prev_gecko_pref_state_for_experiment(
846                db,
847                &mut writer,
848                &experiment_slug,
849                prev_gecko_pref_states,
850            )?;
851        }
852
853        let coenrolling_ids = self
854            .coenrolling_feature_ids
855            .iter()
856            .map(|s| s.as_str())
857            .collect();
858
859        // Registering the Gecko original values does not require a Gecko update,
860        // but the cache does need to be refreshed.
861        self.database_cache.commit_and_update(
862            db,
863            writer,
864            &coenrolling_ids,
865            self.gecko_prefs.clone(),
866            false,
867        )?;
868
869        Ok(())
870    }
871
872    pub(crate) fn add_prev_gecko_pref_state_for_experiment(
873        db: &Database,
874        writer: &mut Writer,
875        experiment_slug: &str,
876        prev_gecko_pref_states: Vec<PreviousGeckoPrefState>,
877    ) -> Result<()> {
878        let enrollments = db.get_store(StoreId::Enrollments);
879
880        if let Ok(Some(existing_enrollment)) =
881            enrollments.get::<ExperimentEnrollment, Writer>(writer, experiment_slug)
882        {
883            // Previous states are only valid on Enrolled experiments
884            let updated_states =
885                existing_enrollment.on_add_gecko_pref_states(prev_gecko_pref_states);
886            enrollments.put(writer, experiment_slug, &updated_states)?;
887        }
888        Ok(())
889    }
890
891    pub fn get_previous_gecko_pref_states(
892        &self,
893        experiment_slug: String,
894    ) -> Result<Option<Vec<PreviousGeckoPrefState>>> {
895        let db = self.db()?;
896        let reader = db.read()?;
897
898        Ok(db
899            .get_store(StoreId::Enrollments)
900            .get::<ExperimentEnrollment, _>(&reader, &experiment_slug)?
901            .and_then(|enrollment| {
902                if let EnrollmentStatus::Enrolled {
903                    prev_gecko_pref_states: prev_gecko_pref_state,
904                    ..
905                } = enrollment.status
906                {
907                    prev_gecko_pref_state
908                } else {
909                    None
910                }
911            }))
912    }
913
914    pub fn set_install_time(&mut self, then: DateTime<Utc>) {
915        let mut state = self.mutable_state.lock().unwrap();
916        state.install_date = Some(then);
917        state.update_time_to_now(Utc::now());
918    }
919
920    pub fn set_update_time(&mut self, then: DateTime<Utc>) {
921        let mut state = self.mutable_state.lock().unwrap();
922        state.update_date = Some(then);
923        state.update_time_to_now(Utc::now());
924    }
925
926    /// This is only called from `get_feature_config_variables` which is itself is cached with
927    /// thread safety in the FeatureHolder.kt and FeatureHolder.swift
928    fn record_feature_activation_if_needed(&self, feature_id: &str) {
929        if let Ok(Some(f)) = self.database_cache.get_enrollment_by_feature(feature_id)
930            && f.branch.is_some()
931            && !self.coenrolling_feature_ids.contains(&f.feature_id)
932        {
933            self.metrics_handler.record_feature_activation(f.into());
934        }
935    }
936
937    pub fn record_feature_exposure(&self, feature_id: String, slug: Option<String>) {
938        let event = if let Some(slug) = slug {
939            if let Ok(Some(branch)) = self.database_cache.get_experiment_branch(&slug) {
940                Some(FeatureExposureExtraDef {
941                    feature_id,
942                    branch: Some(branch),
943                    slug,
944                })
945            } else {
946                None
947            }
948        } else if let Ok(Some(f)) = self.database_cache.get_enrollment_by_feature(&feature_id) {
949            if f.branch.is_some() {
950                Some(f.into())
951            } else {
952                None
953            }
954        } else {
955            None
956        };
957
958        if let Some(event) = event {
959            self.metrics_handler.record_feature_exposure(event);
960        }
961    }
962
963    pub fn record_malformed_feature_config(&self, feature_id: String, part_id: String) {
964        let event = if let Ok(Some(f)) = self.database_cache.get_enrollment_by_feature(&feature_id)
965        {
966            MalformedFeatureConfigExtraDef::from_feature_and_part(f, part_id)
967        } else {
968            MalformedFeatureConfigExtraDef::new(feature_id, part_id)
969        };
970        self.metrics_handler.record_malformed_feature_config(event);
971    }
972
973    fn record_enrollment_status_telemetry(
974        &self,
975        state: &mut MutexGuard<InternalMutableState>,
976    ) -> Result<()> {
977        let targeting_helper = NimbusTargetingHelper::new(
978            state.targeting_attributes.clone(),
979            self.event_store.clone(),
980            self.gecko_prefs.clone(),
981        );
982        let experiments = self.database_cache.get_experiments()?;
983        let experiments = experiments
984            .iter()
985            .filter(|exp| {
986                is_experiment_available(&targeting_helper, exp, true)
987                    == ExperimentAvailable::Available
988            })
989            .map(|exp| &*exp.slug)
990            .collect::<HashSet<&str>>();
991        self.metrics_handler.record_enrollment_statuses(
992            self.database_cache
993                .get_enrollments()?
994                .into_iter()
995                .filter_map(|e| match experiments.contains(&*e.slug) {
996                    true => Some(e.into()),
997                    false => None,
998                })
999                .collect(),
1000        );
1001        self.metrics_handler.submit_targeting_context();
1002        Ok(())
1003    }
1004
1005    pub fn get_available_firefox_labs(&self) -> Result<Vec<FirefoxLabsMetadata>> {
1006        let targeting_attributes = self.get_targeting_attributes();
1007        let targeting_helper = NimbusTargetingHelper::with_targeting_attributes(
1008            &targeting_attributes,
1009            self.event_store.clone(),
1010            self.gecko_prefs.clone(),
1011        );
1012        self.database_cache
1013            .get_available_firefox_labs_metadata(&targeting_helper, &self.coenrolling_feature_ids)
1014    }
1015
1016    pub fn enroll_in_firefox_lab(&self, slug: &str) -> Result<FirefoxLabsEnrollResult> {
1017        let feature_conflict = self
1018            .database_cache
1019            .check_for_feature_conflict(slug, &self.coenrolling_feature_ids)?;
1020
1021        let db = self.db()?;
1022        let mut writer = db.write()?;
1023        let result = enroll_in_firefox_lab(db, &mut writer, slug, feature_conflict);
1024        let mut state = self.mutable_state.lock().unwrap();
1025        self.end_initialize(db, writer, &mut state)?;
1026        result
1027    }
1028
1029    pub fn unenroll_from_firefox_lab(&self, slug: &str) -> Result<FirefoxLabsUnenrollResult> {
1030        let db = self.db()?;
1031        let mut writer = db.write()?;
1032        let result = unenroll_from_firefox_lab(db, &mut writer, slug, self.gecko_prefs.as_deref());
1033        let mut state = self.mutable_state.lock().unwrap();
1034        self.end_initialize(db, writer, &mut state)?;
1035        result
1036    }
1037
1038    pub fn unenroll_from_all_firefox_labs(&self) -> Result<Vec<EnrollmentChangeEvent>> {
1039        let db = self.db()?;
1040        let mut writer = db.write()?;
1041        let result = unenroll_from_all_firefox_labs(db, &mut writer, self.gecko_prefs.as_deref());
1042        let mut state = self.mutable_state.lock().unwrap();
1043        self.end_initialize(db, writer, &mut state)?;
1044        result
1045    }
1046
1047    #[cfg(test)]
1048    pub fn get_experiment_enrollment(&self, slug: &str) -> Result<Option<ExperimentEnrollment>> {
1049        self.database_cache.get_experiment_enrollment(slug)
1050    }
1051}
1052
1053pub struct NimbusStringHelper {
1054    context: JsonObject,
1055}
1056
1057impl NimbusStringHelper {
1058    fn new(context: JsonObject) -> Self {
1059        Self { context }
1060    }
1061
1062    pub fn get_uuid(&self, template: String) -> Option<String> {
1063        if template.contains("{uuid}") {
1064            let uuid = Uuid::new_v4();
1065            Some(uuid.to_string())
1066        } else {
1067            None
1068        }
1069    }
1070
1071    pub fn string_format(&self, template: String, uuid: Option<String>) -> String {
1072        match uuid {
1073            Some(uuid) => {
1074                let mut map = self.context.clone();
1075                map.insert("uuid".to_string(), Value::String(uuid));
1076                fmt_with_map(&template, &map)
1077            }
1078            _ => fmt_with_map(&template, &self.context),
1079        }
1080    }
1081}
1082
1083#[cfg(feature = "stateful-uniffi-bindings")]
1084uniffi::custom_type!(JsonObject, String, {
1085    remote,
1086    try_lift: |val| {
1087        let json: Value = serde_json::from_str(&val)?;
1088
1089        match json.as_object() {
1090            Some(obj) => Ok(obj.clone()),
1091            _ => Err(uniffi::deps::anyhow::anyhow!(
1092                "Unexpected JSON-non-object in the bagging area"
1093            )),
1094        }
1095    },
1096    lower: |obj| serde_json::Value::Object(obj).to_string(),
1097});
1098
1099#[cfg(feature = "stateful-uniffi-bindings")]
1100uniffi::custom_type!(PrefValue, String, {
1101    remote,
1102    try_lift: |val| {
1103        // Raw strings that are not valid JSON (e.g. pref values read directly from Gecko)
1104        // should be treated as JSON string values.
1105        let json: Value = match serde_json::from_str(&val) {
1106            Ok(json) => json,
1107            Err(_) => Value::String(val),
1108        };
1109        let is_valid_pref_type = json.is_string() || json.is_boolean()
1110            || (json.is_number() && !json.is_f64()) || json.is_null();
1111        if is_valid_pref_type {
1112            Ok(json)
1113        } else {
1114            Err(anyhow::anyhow!(format!("Value {} is not a string, boolean, number, or null, or is a float", json)))
1115        }
1116    },
1117    lower: |val| {
1118        val.to_string()
1119    }
1120});
1121
1122#[cfg(feature = "stateful-uniffi-bindings")]
1123uniffi::include_scaffolding!("nimbus");