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

statsig_rust/interned_values/
interned_store.rs
