statsig_rust/
interned_value_store.rs

1use std::{
2    borrow::Cow,
3    sync::{Arc, Weak},
4    time::Duration,
5};
6
7use ahash::HashMap;
8use parking_lot::Mutex;
9use serde_json::value::RawValue;
10
11use crate::log_e;
12
13pub trait FromRawValue {
14    fn from_raw_value(raw_value: Cow<'_, RawValue>) -> Self;
15}
16
17pub struct InternedValueStore<T: FromRawValue> {
18    tag: &'static str,
19    pub(crate) values: Mutex<HashMap<u64, Weak<T>>>,
20}
21
22impl<T: FromRawValue> InternedValueStore<T> {
23    pub fn new(tag: &'static str) -> Self {
24        Self {
25            tag,
26            values: Mutex::new(HashMap::default()),
27        }
28    }
29
30    pub fn get_or_create_interned_value(
31        &self,
32        hash: u64,
33        raw_value: Cow<'_, RawValue>,
34    ) -> (u64, Arc<T>) {
35        if let Some(value) = self.try_get_interned_value(hash) {
36            return (hash, value);
37        }
38
39        let value = T::from_raw_value(raw_value);
40        let value_arc = Arc::new(value);
41        self.set_interned_value(hash, &value_arc);
42
43        (hash, value_arc)
44    }
45
46    pub fn try_remove_interned_value(&self, hash: u64) {
47        let mut memoized_values = match self.values.try_lock_for(Duration::from_secs(5)) {
48            Some(values) => values,
49            None => return,
50        };
51
52        let found = match memoized_values.get(&hash) {
53            Some(value) => value,
54            None => return,
55        };
56
57        if found.strong_count() == 1 {
58            memoized_values.remove(&hash);
59        }
60    }
61
62    pub fn set_interned_value(&self, hash: u64, value: &Arc<T>) {
63        let mut memoized_values = match self.values.try_lock_for(Duration::from_secs(5)) {
64            Some(values) => values,
65            None => return,
66        };
67        memoized_values.insert(hash, Arc::downgrade(value));
68    }
69
70    pub fn try_get_interned_value(&self, hash: u64) -> Option<Arc<T>> {
71        let memoized_values = match self.values.try_lock_for(Duration::from_secs(5)) {
72            Some(values) => values,
73            None => {
74                log_e!(self.tag, "Failed to lock interned values map");
75                return None;
76            }
77        };
78
79        memoized_values.get(&hash).and_then(|value| value.upgrade())
80    }
81}
82
83#[macro_export]
84macro_rules! impl_interned_value {
85    ($struct_name:ident, $memoized_type:ident) => {
86        lazy_static::lazy_static! {
87            pub(crate) static ref INTERNED_STORE: $crate::interned_value_store::InternedValueStore<$memoized_type> = $crate::interned_value_store::InternedValueStore::new(stringify!($struct_name));
88        }
89
90        impl Drop for $struct_name {
91            fn drop(&mut self) {
92                INTERNED_STORE.try_remove_interned_value(self.hash);
93            }
94        }
95
96        impl $struct_name {
97            #[allow(dead_code)]
98            fn get_or_create_memoized(raw_value: Cow<'_, RawValue>) -> (u64, std::sync::Arc<$memoized_type>) {
99                let hash = $crate::hashing::hash_one(raw_value.get());
100                INTERNED_STORE.get_or_create_interned_value(hash, raw_value)
101            }
102        }
103    };
104}
statsig_rust/interned_value_store.rs

statsig_rust/
interned_value_store.rs
