aster/config/
experiments.rs

1use super::base::Config;
2use anyhow::Result;
3use std::collections::HashMap;
4
5/// It is the ground truth for init experiments. The experiment names in users' experiment list but not
6/// in the list will be remove from user list; The experiment names in the ground-truth list but not
7/// in users' experiment list will be added to user list with default value false;
8/// TODO: keep this up to date with the experimental-features.md documentation page
9const ALL_EXPERIMENTS: &[(&str, bool)] = &[];
10
11/// Experiment configuration management
12pub struct ExperimentManager;
13
14impl ExperimentManager {
15    /// Get all experiments and their configurations
16    ///
17    /// - Ensures the user's experiment list is synchronized with `ALL_EXPERIMENTS`.
18    /// - Adds missing experiments from `ALL_EXPERIMENTS` with the default value.
19    /// - Removes experiments not in `ALL_EXPERIMENTS`.
20    pub fn get_all() -> Result<Vec<(String, bool)>> {
21        let config = Config::global();
22        let mut experiments: HashMap<String, bool> =
23            config.get_param("experiments").unwrap_or_default();
24        Self::refresh_experiments(&mut experiments);
25
26        Ok(experiments.into_iter().collect())
27    }
28
29    /// Enable or disable an experiment
30    pub fn set_enabled(name: &str, enabled: bool) -> Result<()> {
31        let config = Config::global();
32        let mut experiments: HashMap<String, bool> = config
33            .get_param("experiments")
34            .unwrap_or_else(|_| HashMap::new());
35        Self::refresh_experiments(&mut experiments);
36        experiments.insert(name.to_string(), enabled);
37
38        config.set_param("experiments", experiments)?;
39        Ok(())
40    }
41
42    /// Check if an experiment is enabled
43    pub fn is_enabled(name: &str) -> Result<bool> {
44        let experiments = Self::get_all()?;
45        let experiments_map: HashMap<String, bool> = experiments.into_iter().collect();
46        Ok(*experiments_map.get(name).unwrap_or(&false))
47    }
48
49    fn refresh_experiments(experiments: &mut HashMap<String, bool>) {
50        // Add missing experiments from `ALL_EXPERIMENTS`
51        for &(key, default_value) in ALL_EXPERIMENTS {
52            experiments.entry(key.to_string()).or_insert(default_value);
53        }
54
55        // Remove experiments not present in `ALL_EXPERIMENTS`
56        experiments.retain(|key, _| ALL_EXPERIMENTS.iter().any(|(k, _)| k == key));
57    }
58}
aster/config/experiments.rs

aster/config/
experiments.rs
