statsig_rust/interned_values/
interned_store.rs

1use std::{
2    borrow::Cow,
3    collections::hash_map::Entry,
4    sync::{Arc, OnceLock},
5    time::{Duration, Instant},
6};
7
8use ahash::AHashMap;
9use lazy_static::lazy_static;
10use parking_lot::Mutex;
11use serde_json::value::RawValue;
12
13use crate::{
14    evaluation::{
15        dynamic_returnable::DynamicReturnableValue,
16        evaluator_value::{EvaluatorValue, EvaluatorValueInner, MemoizedEvaluatorValue},
17    },
18    hashing,
19    interned_string::{InternedString, InternedStringValue},
20    log_d, log_e, log_w,
21    networking::ResponseData,
22    observability::ops_stats::OpsStatsForInstance,
23    specs_response::{
24        proto_specs::deserialize_protobuf,
25        spec_types::{Spec, SpecsResponseFull},
26        specs_hash_map::{SpecPointer, SpecsHashMap},
27    },
28    DynamicReturnable, StatsigErr,
29};
30
31const TAG: &str = "InternedStore";
32
33static IMMORTAL_DATA: OnceLock<ImmortalData> = OnceLock::new();
34
35lazy_static! {
36    static ref MUTABLE_DATA: Mutex<MutableData> = Mutex::new(MutableData::default());
37}
38
39/// Immortal vs Mutable Data
40/// ------------------------------------------------------------
41/// -`ImmortalData` is static and never changes. It will only exist if a successful call to `preload` is made. It is intentionally
42///  leaked so that it can be accessed across forks without incrementing the reference count.
43/// -`MutableData` is dynamic and changes over time as values are added and removed.
44/// ------------------------------------------------------------
45/// In all cases, we first check if there is a ImmortalData entry and then fallback to MutableData.
46#[derive(Default)]
47struct ImmortalData {
48    strings: AHashMap<u64, &'static str>,
49    returnables: AHashMap<u64, &'static RawValue>,
50    evaluator_values: AHashMap<u64, &'static MemoizedEvaluatorValue>,
51    feature_gates: AHashMap<u64, &'static Spec>,
52    dynamic_configs: AHashMap<u64, &'static Spec>,
53    layer_configs: AHashMap<u64, &'static Spec>,
54}
55
56#[derive(Default)]
57struct MutableData {
58    strings: AHashMap<u64, Arc<String>>,
59    returnables: AHashMap<u64, Arc<Box<RawValue>>>,
60    evaluator_values: AHashMap<u64, Arc<MemoizedEvaluatorValue>>,
61}
62
63pub trait Internable: Sized {
64    type Input<'a>;
65    fn intern(input: Self::Input<'_>) -> Self;
66}
67
68pub struct InternedStore;
69
70impl InternedStore {
71    pub fn preload(data: &[u8]) -> Result<(), StatsigErr> {
72        let start_time = Instant::now();
73
74        if IMMORTAL_DATA.get().is_some() {
75            log_e!(TAG, "Already preloaded");
76            return Err(StatsigErr::InvalidOperation(
77                "Already preloaded".to_string(),
78            ));
79        }
80
81        let specs_response = try_parse_as_json(data).or_else(|_| try_parse_as_proto(data))?;
82
83        let immortal = mutable_to_immortal(specs_response)?;
84
85        if IMMORTAL_DATA.set(immortal).is_err() {
86            return Err(StatsigErr::LockFailure(
87                "Failed to set IMMORTAL_DATA".to_string(),
88            ));
89        }
90
91        let end_time = Instant::now();
92        log_d!(
93            TAG,
94            "Preload took {}ms",
95            end_time.duration_since(start_time).as_millis()
96        );
97
98        Ok(())
99    }
100
101    pub fn get_or_intern_string<T: AsRef<str> + ToString>(value: T) -> InternedString {
102        let hash = hashing::hash_one(value.as_ref().as_bytes());
103
104        if let Some(string) = get_string_from_shared(hash) {
105            return InternedString::from_static(hash, string);
106        }
107
108        let ptr = get_string_from_local(hash, value);
109        InternedString::from_pointer(hash, ptr)
110    }
111
112    pub fn get_or_intern_returnable(value: Cow<'_, RawValue>) -> DynamicReturnable {
113        let raw_string = value.get();
114        match raw_string {
115            "true" => return DynamicReturnable::from_bool(true),
116            "false" => return DynamicReturnable::from_bool(false),
117            "null" => return DynamicReturnable::empty(),
118            _ => {}
119        }
120
121        let hash = hashing::hash_one(raw_string.as_bytes());
122
123        if let Some(returnable) = get_returnable_from_shared(hash) {
124            return DynamicReturnable::from_static(hash, returnable);
125        }
126
127        let ptr = get_returnable_from_local(hash, value);
128        DynamicReturnable::from_pointer(hash, ptr)
129    }
130
131    pub fn get_or_intern_evaluator_value(value: Cow<'_, RawValue>) -> EvaluatorValue {
132        let raw_string = value.get();
133        let hash = hashing::hash_one(raw_string.as_bytes());
134
135        if let Some(evaluator_value) = get_evaluator_value_from_shared(hash) {
136            return EvaluatorValue::from_static(hash, evaluator_value);
137        }
138
139        let ptr = get_evaluator_value_from_local(hash, value);
140        EvaluatorValue::from_pointer(hash, ptr)
141    }
142
143    pub fn replace_evaluator_value(hash: u64, evaluator_value: Arc<MemoizedEvaluatorValue>) {
144        let old = use_mutable_data("replace_evaluator_value", |data| {
145            data.evaluator_values.insert(hash, evaluator_value)
146        });
147        drop(old);
148    }
149
150    pub fn try_get_preloaded_evaluator_value(bytes: &[u8]) -> Option<EvaluatorValue> {
151        let hash = hashing::hash_one(bytes);
152        if let Some(evaluator_value) = get_evaluator_value_from_shared(hash) {
153            return Some(EvaluatorValue::from_static(hash, evaluator_value));
154        }
155
156        None
157    }
158
159    pub fn try_get_preloaded_returnable(bytes: &[u8]) -> Option<DynamicReturnable> {
160        match bytes {
161            b"true" => return Some(DynamicReturnable::from_bool(true)),
162            b"false" => return Some(DynamicReturnable::from_bool(false)),
163            b"null" => return Some(DynamicReturnable::empty()),
164            _ => {}
165        }
166
167        let hash = hashing::hash_one(bytes);
168        if let Some(returnable) = get_returnable_from_shared(hash) {
169            return Some(DynamicReturnable::from_static(hash, returnable));
170        }
171
172        None
173    }
174
175    pub fn try_get_preloaded_dynamic_config(name: &InternedString) -> Option<SpecPointer> {
176        match IMMORTAL_DATA.get() {
177            Some(shared) => shared
178                .dynamic_configs
179                .get(&name.hash)
180                .map(|s| SpecPointer::Static(s)),
181            None => None,
182        }
183    }
184
185    pub fn try_get_preloaded_layer_config(name: &InternedString) -> Option<SpecPointer> {
186        match IMMORTAL_DATA.get() {
187            Some(shared) => shared
188                .layer_configs
189                .get(&name.hash)
190                .map(|s| SpecPointer::Static(s)),
191            None => None,
192        }
193    }
194
195    pub fn try_get_preloaded_feature_gate(name: &InternedString) -> Option<SpecPointer> {
196        match IMMORTAL_DATA.get() {
197            Some(shared) => shared
198                .feature_gates
199                .get(&name.hash)
200                .map(|s| SpecPointer::Static(s)),
201            None => None,
202        }
203    }
204
205    pub fn release_returnable(hash: u64) {
206        let ptr = use_mutable_data("release_returnable", |data| {
207            try_release_entry(&mut data.returnables, hash)
208        });
209        drop(ptr);
210    }
211
212    pub fn release_string(hash: u64) {
213        let ptr = use_mutable_data("release_string", |data| {
214            try_release_entry(&mut data.strings, hash)
215        });
216        drop(ptr);
217    }
218
219    pub fn release_evaluator_value(hash: u64) {
220        let ptr = use_mutable_data("release_eval_value", |data| {
221            try_release_entry(&mut data.evaluator_values, hash)
222        });
223        drop(ptr);
224    }
225
226    #[cfg(test)]
227    pub fn get_memoized_len() -> (
228        /* strings */ usize,
229        /* returnables */ usize,
230        /* evaluator values */ usize,
231    ) {
232        match MUTABLE_DATA.try_lock() {
233            Some(memo) => (
234                memo.strings.len(),
235                memo.returnables.len(),
236                memo.evaluator_values.len(),
237            ),
238            None => (0, 0, 0),
239        }
240    }
241}
242
243// ------------------------------------------------------------------------------- [ Preloading ]
244
245fn try_parse_as_json(data: &[u8]) -> Result<SpecsResponseFull, StatsigErr> {
246    serde_json::from_slice(data)
247        .map_err(|e| StatsigErr::JsonParseError(TAG.to_string(), e.to_string()))
248}
249
250fn try_parse_as_proto(data: &[u8]) -> Result<SpecsResponseFull, StatsigErr> {
251    let current = SpecsResponseFull::default();
252    let mut next = SpecsResponseFull::default();
253
254    let mut response_data = ResponseData::from_bytes_with_headers(
255        data.to_vec(),
256        Some(std::collections::HashMap::from([(
257            "content-encoding".to_string(),
258            "statsig-br".to_string(),
259        )])),
260    );
261
262    let ops_stats = OpsStatsForInstance::new();
263
264    deserialize_protobuf(&ops_stats, &current, &mut next, &mut response_data)?;
265
266    Ok(next)
267}
268
269// ------------------------------------------------------------------------------- [ String ]
270
271fn get_string_from_shared(hash: u64) -> Option<&'static str> {
272    match IMMORTAL_DATA.get() {
273        Some(shared) => shared.strings.get(&hash).copied(),
274        None => None,
275    }
276}
277
278fn get_string_from_local<T: ToString>(hash: u64, value: T) -> Arc<String> {
279    let result = use_mutable_data("intern_string", |data| {
280        if let Some(string) = data.strings.get(&hash) {
281            return Some(string.clone());
282        }
283
284        let ptr = Arc::new(value.to_string());
285        data.strings.insert(hash, ptr.clone());
286        Some(ptr)
287    });
288
289    result.unwrap_or_else(|| {
290        log_w!(TAG, "Failed to get string from local");
291        Arc::new(value.to_string())
292    })
293}
294
295// ------------------------------------------------------------------------------- [ Returnable ]
296
297fn get_returnable_from_shared(hash: u64) -> Option<&'static RawValue> {
298    match IMMORTAL_DATA.get() {
299        Some(shared) => shared.returnables.get(&hash).copied(),
300        None => None,
301    }
302}
303
304fn get_returnable_from_local(hash: u64, value: Cow<RawValue>) -> Arc<Box<RawValue>> {
305    let result = use_mutable_data("intern_returnable", |data| {
306        if let Some(returnable) = data.returnables.get(&hash) {
307            return Some(returnable.clone());
308        }
309
310        None
311    });
312
313    if let Some(returnable) = result {
314        return returnable;
315    }
316
317    let owned = match value {
318        Cow::Borrowed(value) => value.to_owned(),
319        Cow::Owned(value) => value,
320    };
321
322    let ptr = Arc::new(owned);
323
324    use_mutable_data("intern_returnable", |data| {
325        data.returnables.insert(hash, ptr.clone());
326        Some(())
327    });
328
329    ptr
330}
331
332// ------------------------------------------------------------------------------- [ Evaluator Value ]
333
334fn get_evaluator_value_from_shared(hash: u64) -> Option<&'static MemoizedEvaluatorValue> {
335    match IMMORTAL_DATA.get() {
336        Some(shared) => shared.evaluator_values.get(&hash).copied(),
337        None => None,
338    }
339}
340
341fn get_evaluator_value_from_local(
342    hash: u64,
343    value: Cow<'_, RawValue>,
344) -> Arc<MemoizedEvaluatorValue> {
345    let result = use_mutable_data("eval_value_lookup", |data| {
346        if let Some(evaluator_value) = data.evaluator_values.get(&hash) {
347            return Some(evaluator_value.clone());
348        }
349
350        None
351    });
352
353    if let Some(evaluator_value) = result {
354        return evaluator_value;
355    }
356
357    // intentinonally done across two locks to avoid deadlock with InternedString creation
358    let ptr = Arc::new(MemoizedEvaluatorValue::from_raw_value(value));
359    let _ = use_mutable_data("intern_evaluator_value", |data| {
360        data.evaluator_values.insert(hash, ptr.clone());
361        Some(())
362    });
363
364    ptr
365}
366
367// ------------------------------------------------------------------------------- [ Helpers ]
368
369fn try_release_entry<T>(data: &mut AHashMap<u64, Arc<T>>, hash: u64) -> Option<Arc<T>> {
370    let found = match data.entry(hash) {
371        Entry::Occupied(entry) => entry,
372        Entry::Vacant(_) => return None,
373    };
374
375    let strong_count = Arc::strong_count(found.get());
376    if strong_count == 1 {
377        let value = found.remove();
378        // return the value so it isn't dropped while holding the lock
379        return Some(value);
380    }
381
382    None
383}
384
385fn use_mutable_data<T>(reason: &str, f: impl FnOnce(&mut MutableData) -> Option<T>) -> Option<T> {
386    let mut data = match MUTABLE_DATA.try_lock_for(Duration::from_secs(5)) {
387        Some(data) => data,
388        None => {
389            #[cfg(test)]
390            panic!("Failed to acquire lock for mutable data ({reason})");
391
392            #[cfg(not(test))]
393            {
394                log_e!(TAG, "Failed to acquire lock for mutable data ({reason})");
395                return None;
396            }
397        }
398    };
399
400    f(&mut data)
401}
402
403fn mutable_to_immortal(specs_res: SpecsResponseFull) -> Result<ImmortalData, StatsigErr> {
404    let mutable_data: MutableData = {
405        let mut mutable_data_lock = MUTABLE_DATA.lock();
406        std::mem::take(&mut *mutable_data_lock)
407    };
408    let mut immortal = ImmortalData::default();
409
410    for (hash, arc) in mutable_data.strings.into_iter() {
411        let raw = Arc::into_raw(arc);
412        let leaked: &'static str = unsafe { &*raw };
413        immortal.strings.insert(hash, leaked);
414    }
415
416    for (hash, returnable) in mutable_data.returnables.into_iter() {
417        let raw_returnable = Arc::into_raw(returnable);
418        let leaked = unsafe { &*raw_returnable };
419        immortal.returnables.insert(hash, leaked);
420    }
421
422    for (hash, evaluator_value) in mutable_data.evaluator_values.into_iter() {
423        let raw_evaluator_value = Arc::into_raw(evaluator_value);
424        let leaked = unsafe { &*raw_evaluator_value };
425        immortal.evaluator_values.insert(hash, leaked);
426    }
427
428    try_insert_specs(specs_res.feature_gates, &mut immortal.feature_gates);
429    try_insert_specs(specs_res.dynamic_configs, &mut immortal.dynamic_configs);
430    try_insert_specs(specs_res.layer_configs, &mut immortal.layer_configs);
431
432    Ok(immortal)
433}
434
435fn try_insert_specs(source: SpecsHashMap, destination: &mut AHashMap<u64, &'static Spec>) {
436    for (name, spec_ptr) in source.0.into_iter() {
437        let spec = match spec_ptr {
438            SpecPointer::Pointer(spec) => spec,
439            _ => continue,
440        };
441
442        if spec.checksum.is_none() {
443            // no point doint this if there is no checksum field to verify against later
444            continue;
445        }
446
447        let raw_spec = Arc::into_raw(spec);
448        let spec = unsafe { &*raw_spec };
449        destination.insert(name.hash, spec);
450    }
451}
452
453// ------------------------------------------------------------------------------- [ Helper Implementations ]
454
455impl EvaluatorValue {
456    fn from_static(hash: u64, evaluator_value: &'static MemoizedEvaluatorValue) -> Self {
457        Self {
458            hash,
459            inner: EvaluatorValueInner::Static(evaluator_value),
460        }
461    }
462
463    fn from_pointer(hash: u64, pointer: Arc<MemoizedEvaluatorValue>) -> Self {
464        Self {
465            hash,
466            inner: EvaluatorValueInner::Pointer(pointer),
467        }
468    }
469}
470
471impl DynamicReturnable {
472    fn from_static(hash: u64, returnable: &'static RawValue) -> Self {
473        Self {
474            hash,
475            value: DynamicReturnableValue::JsonStatic(returnable),
476        }
477    }
478
479    fn from_pointer(hash: u64, pointer: Arc<Box<RawValue>>) -> Self {
480        Self {
481            hash,
482            value: DynamicReturnableValue::JsonPointer(pointer),
483        }
484    }
485}
486
487impl InternedString {
488    fn from_static(hash: u64, string: &'static str) -> Self {
489        Self {
490            hash,
491            value: InternedStringValue::Static(string),
492        }
493    }
494
495    fn from_pointer(hash: u64, pointer: Arc<String>) -> Self {
496        Self {
497            hash,
498            value: InternedStringValue::Pointer(pointer),
499        }
500    }
501}
statsig_rust/interned_values/interned_store.rs

statsig_rust/interned_values/
interned_store.rs
