statsig_rust/override_adapter/
statsig_local_override_adapter.rs

1use crate::event_logging::exposable_string::ExposableString;
2use crate::{log_d, read_lock_or_return, write_lock_or_noop, OverrideAdapter, StatsigUser};
3use std::{collections::HashMap, sync::RwLock};
4
5use crate::evaluation::evaluator_result::EvaluatorResult;
6use crate::specs_response::spec_types::Spec;
7use serde_json::Value;
8
9enum ExperimentOverrides {
10    Value(HashMap<String, Value>),
11    GroupName(String),
12}
13
14#[derive(Default)]
15struct OverrideStore {
16    pub gate: HashMap<String, HashMap<String, bool>>,
17    pub config: HashMap<String, HashMap<String, HashMap<String, Value>>>,
18    pub experiment: HashMap<String, HashMap<String, ExperimentOverrides>>,
19    pub layer: HashMap<String, HashMap<String, HashMap<String, Value>>>,
20}
21
22const TAG: &str = stringify!(StatsigLocalOverrideAdapter);
23const LOCAL_OVERRIDE_REASON: &str = "LocalOverride";
24const NO_ID_OVERRIDE: &str = "__STATSIG_NO_ID__";
25
26lazy_static::lazy_static! {
27    static ref OVERRIDE_RULE_ID: ExposableString = ExposableString::new("override".to_string());
28}
29
30#[derive(Default)]
31pub struct StatsigLocalOverrideAdapter {
32    store: RwLock<OverrideStore>,
33}
34
35impl StatsigLocalOverrideAdapter {
36    #[must_use]
37    pub fn new() -> Self {
38        Self::default()
39    }
40}
41
42impl OverrideAdapter for StatsigLocalOverrideAdapter {
43    fn get_gate_override(
44        &self,
45        user: &StatsigUser,
46        gate_name: &str,
47        result: &mut EvaluatorResult<'_>,
48    ) -> bool {
49        let store = read_lock_or_return!(TAG, self.store, false);
50
51        let gate_overrides = match store.gate.get(gate_name) {
52            Some(overrides) => overrides,
53            None => return false,
54        };
55
56        log_d!(TAG, "gate_overrides found for {}", gate_name);
57
58        find_override_for_user(
59            user,
60            gate_overrides,
61            |value, res| {
62                res.bool_value = *value;
63            },
64            result,
65        )
66    }
67
68    fn get_dynamic_config_override(
69        &self,
70        user: &StatsigUser,
71        dynamic_config_name: &str,
72        result: &mut EvaluatorResult<'_>,
73    ) -> bool {
74        let store = read_lock_or_return!(TAG, self.store, false);
75
76        let config_overrides = match store.config.get(dynamic_config_name) {
77            Some(overrides) => overrides,
78            None => return false,
79        };
80
81        find_override_for_user(
82            user,
83            config_overrides,
84            |value, res| {
85                res.json_value = Some(value.clone());
86            },
87            result,
88        )
89    }
90
91    fn get_experiment_override(
92        &self,
93        user: &StatsigUser,
94        experiment_name: &str,
95        result: &mut EvaluatorResult<'_>,
96        opt_spec: Option<&Spec>,
97    ) -> bool {
98        let store = read_lock_or_return!(TAG, self.store, false);
99
100        let experiment_overrides = match store.experiment.get(experiment_name) {
101            Some(overrides) => overrides,
102            None => return false,
103        };
104
105        find_override_for_user(
106            user,
107            experiment_overrides,
108            |override_value, res| match override_value {
109                ExperimentOverrides::Value(map) => {
110                    res.json_value = Some(map.clone());
111                }
112                ExperimentOverrides::GroupName(group_name) => {
113                    res.json_value = get_experiment_with_group_name(opt_spec, group_name);
114                }
115            },
116            result,
117        )
118    }
119
120    fn get_layer_override(
121        &self,
122        user: &StatsigUser,
123        layer_name: &str,
124        result: &mut EvaluatorResult<'_>,
125    ) -> bool {
126        let store = read_lock_or_return!(TAG, self.store, false);
127
128        let layer_overrides = match store.layer.get(layer_name) {
129            Some(overrides) => overrides,
130            None => return false,
131        };
132
133        find_override_for_user(
134            user,
135            layer_overrides,
136            |value, res| {
137                res.json_value = Some(value.clone());
138            },
139            result,
140        )
141    }
142
143    fn override_gate(&self, key: &str, value: bool, id: Option<&str>) {
144        let mut store = write_lock_or_noop!(TAG, self.store);
145        let id_str = id.unwrap_or(NO_ID_OVERRIDE);
146        store
147            .gate
148            .entry(key.to_string())
149            .or_default()
150            .insert(id_str.to_string(), value);
151    }
152
153    fn override_dynamic_config(&self, key: &str, value: HashMap<String, Value>, id: Option<&str>) {
154        let mut store = write_lock_or_noop!(TAG, self.store);
155        let id_str = id.unwrap_or(NO_ID_OVERRIDE);
156        store
157            .config
158            .entry(key.to_string())
159            .or_default()
160            .insert(id_str.to_string(), value);
161    }
162
163    fn override_experiment(&self, key: &str, value: HashMap<String, Value>, id: Option<&str>) {
164        let mut store = write_lock_or_noop!(TAG, self.store);
165        let id_str = id.unwrap_or(NO_ID_OVERRIDE);
166        store
167            .experiment
168            .entry(key.to_string())
169            .or_default()
170            .insert(id_str.to_string(), ExperimentOverrides::Value(value));
171    }
172
173    fn override_experiment_by_group_name(&self, key: &str, group_name: &str, id: Option<&str>) {
174        let mut store = write_lock_or_noop!(TAG, self.store);
175        let id_str = id.unwrap_or(NO_ID_OVERRIDE);
176        store.experiment.entry(key.to_string()).or_default().insert(
177            id_str.to_string(),
178            ExperimentOverrides::GroupName(group_name.to_string()),
179        );
180    }
181
182    fn override_layer(&self, key: &str, value: HashMap<String, Value>, id: Option<&str>) {
183        let mut store = write_lock_or_noop!(TAG, self.store);
184        let id_str = id.unwrap_or(NO_ID_OVERRIDE);
185        store
186            .layer
187            .entry(key.to_string())
188            .or_default()
189            .insert(id_str.to_string(), value);
190    }
191
192    fn remove_gate_override(&self, key: &str, id: Option<&str>) {
193        let mut store = write_lock_or_noop!(TAG, self.store);
194        match id {
195            None => {
196                store.gate.remove(key);
197            }
198            Some(id_str) => {
199                if let Some(overrides) = store.gate.get_mut(key) {
200                    overrides.remove(&id_str.to_string());
201                }
202            }
203        }
204    }
205
206    fn remove_dynamic_config_override(&self, key: &str, id: Option<&str>) {
207        let mut store = write_lock_or_noop!(TAG, self.store);
208        match id {
209            None => {
210                store.config.remove(key);
211            }
212            Some(id_str) => {
213                if let Some(overrides) = store.config.get_mut(key) {
214                    overrides.remove(&id_str.to_string());
215                }
216            }
217        }
218    }
219
220    fn remove_experiment_override(&self, key: &str, id: Option<&str>) {
221        let mut store = write_lock_or_noop!(TAG, self.store);
222        match id {
223            None => {
224                store.experiment.remove(key);
225            }
226            Some(id_str) => {
227                if let Some(overrides) = store.experiment.get_mut(key) {
228                    overrides.remove(&id_str.to_string());
229                }
230            }
231        }
232    }
233
234    fn remove_layer_override(&self, key: &str, id: Option<&str>) {
235        let mut store = write_lock_or_noop!(TAG, self.store);
236        match id {
237            None => {
238                store.layer.remove(key);
239            }
240            Some(id_str) => {
241                if let Some(overrides) = store.layer.get_mut(key) {
242                    overrides.remove(&id_str.to_string());
243                }
244            }
245        }
246    }
247
248    fn remove_all_overrides(&self) {
249        let mut store = write_lock_or_noop!(TAG, self.store);
250        store.gate.clear();
251        store.config.clear();
252        store.experiment.clear();
253        store.layer.clear();
254    }
255}
256
257fn find_override_for_user<T, F>(
258    user: &StatsigUser,
259    overrides: &HashMap<String, T>,
260    apply_override: F,
261    result: &mut EvaluatorResult<'_>,
262) -> bool
263where
264    F: Fn(&T, &mut EvaluatorResult<'_>),
265{
266    if check_user_id_override(user, overrides, &apply_override, result) {
267        return true;
268    }
269
270    if check_custom_ids_override(user, overrides, &apply_override, result) {
271        return true;
272    }
273
274    check_default_override(overrides, apply_override, result)
275}
276
277fn mark_result_as_override(result: &mut EvaluatorResult<'_>) {
278    result.override_reason = Some(LOCAL_OVERRIDE_REASON);
279    result.rule_id = Some(&OVERRIDE_RULE_ID);
280}
281
282fn check_default_override<T, F>(
283    overrides: &HashMap<String, T>,
284    apply_override: F,
285    result: &mut EvaluatorResult<'_>,
286) -> bool
287where
288    F: Fn(&T, &mut EvaluatorResult<'_>),
289{
290    if let Some(override_value) = overrides.get(NO_ID_OVERRIDE) {
291        log_d!(TAG, "default override found");
292        apply_override(override_value, result);
293        mark_result_as_override(result);
294        return true;
295    }
296    false
297}
298
299fn check_user_id_override<T, F>(
300    user: &StatsigUser,
301    overrides: &HashMap<String, T>,
302    apply_override: F,
303    result: &mut EvaluatorResult<'_>,
304) -> bool
305where
306    F: Fn(&T, &mut EvaluatorResult<'_>),
307{
308    let user_id = match &user.data.user_id {
309        Some(id) => id,
310        None => return false,
311    };
312
313    let id_string = match &user_id.string_value {
314        Some(s) => &s.value,
315        None => return false,
316    };
317
318    let override_value = match overrides.get(id_string) {
319        Some(v) => v,
320        None => return false,
321    };
322
323    log_d!(TAG, "override found for user ID {}", id_string);
324    apply_override(override_value, result);
325    mark_result_as_override(result);
326    true
327}
328
329fn check_custom_ids_override<T, F>(
330    user: &StatsigUser,
331    overrides: &HashMap<String, T>,
332    apply_override: F,
333    result: &mut EvaluatorResult<'_>,
334) -> bool
335where
336    F: Fn(&T, &mut EvaluatorResult<'_>),
337{
338    let custom_ids = match &user.data.custom_ids {
339        Some(ids) => ids,
340        None => return false,
341    };
342
343    for custom_id_value in custom_ids.values() {
344        let id_string = match &custom_id_value.string_value {
345            Some(s) => &s.value,
346            None => continue,
347        };
348
349        let override_value = match overrides.get(id_string) {
350            Some(v) => v,
351            None => continue,
352        };
353
354        log_d!(TAG, "override found for custom ID {}", id_string);
355        apply_override(override_value, result);
356        mark_result_as_override(result);
357        return true;
358    }
359
360    false
361}
362
363fn get_experiment_with_group_name(
364    opt_spec: Option<&Spec>,
365    group_name: &str,
366) -> Option<HashMap<String, Value>> {
367    let spec = opt_spec?;
368
369    for rule in &spec.rules {
370        let rule_group_name = match &rule.group_name {
371            Some(rule_group_name) => rule_group_name,
372            None => continue,
373        };
374
375        if rule_group_name == group_name {
376            return rule.return_value.get_json();
377        }
378    }
379
380    None
381}