statsig_rust/
instance_store.rs

1use lazy_static::lazy_static;
2use std::collections::HashMap;
3use std::sync::{Arc, RwLock};
4use uuid::Uuid;
5
6use crate::{
7    log_d, log_e, log_w, Statsig, StatsigHttpEventLoggingAdapter, StatsigHttpSpecsAdapter,
8    StatsigLocalFileEventLoggingAdapter, StatsigLocalFileSpecsAdapter, StatsigOptions, StatsigUser,
9};
10
11const TAG: &str = stringify!(InstanceStore);
12
13#[macro_export]
14macro_rules! get_instance_or_noop {
15    ($type:ty, $ref:expr) => {
16        match INST_STORE.get::<$type>($ref) {
17            Some(instance) => instance,
18            None => {
19                $crate::log_w!(TAG, "{} Reference not found {}", stringify!($type), $ref);
20                return;
21            }
22        }
23    };
24}
25
26#[macro_export]
27macro_rules! get_instance_or_return {
28    ($type:ty, $ref:expr, $return_val:expr) => {
29        match INST_STORE.get::<$type>($ref) {
30            Some(instance) => instance,
31            None => {
32                $crate::log_w!(TAG, "{} Reference not found {}", stringify!($type), $ref);
33                return $return_val;
34            }
35        }
36    };
37}
38
39#[macro_export]
40macro_rules! get_instance_or_else {
41    ($type:ty, $ref:expr, $else:expr) => {
42        match INST_STORE.get::<$type>($ref) {
43            Some(instance) => instance,
44            None => $else,
45        }
46    };
47}
48
49macro_rules! impl_boxable_instance {
50    ($type:ty, $variant:ident, $prefix:expr) => {
51        impl BoxableInstance for $type {
52            fn from_box(boxed: &BoxedInstance) -> Option<Arc<Self>> {
53                if let BoxedInstance::$variant(inner) = boxed {
54                    Some(inner.clone())
55                } else {
56                    log_e!(TAG, "Invalid box type");
57                    None
58                }
59            }
60
61            fn into_box(self) -> BoxedInstance {
62                BoxedInstance::$variant(Arc::new(self))
63            }
64
65            fn get_display_value_static() -> String {
66                stringify!($type).to_string()
67            }
68
69            fn get_display_value(&self) -> String {
70                stringify!($type).to_string()
71            }
72
73            fn get_prefix_value(&self) -> String {
74                $prefix.to_string()
75            }
76        }
77    };
78}
79
80macro_rules! impl_all_instances {
81    ($(($type:ty, $variant:ident) => $prefix:expr),* $(,)?) => {
82        $(
83            impl_boxable_instance!($type, $variant, $prefix);
84        )*
85    }
86}
87
88lazy_static! {
89    pub static ref INST_STORE: InstanceStore = InstanceStore::new();
90}
91
92const MAX_STORED_INSTANCES: usize = 400_000;
93
94pub enum BoxedInstance {
95    Statsig(Arc<Statsig>),
96    StatsigOptions(Arc<StatsigOptions>),
97    StatsigUser(Arc<StatsigUser>),
98    StatsigHttpSpecsAdapter(Arc<StatsigHttpSpecsAdapter>),
99    StatsigLocalFileSpecsAdapter(Arc<StatsigLocalFileSpecsAdapter>),
100    StatsigHttpEventLoggingAdapter(Arc<StatsigHttpEventLoggingAdapter>),
101    StatsigLocalFileEventLoggingAdapter(Arc<StatsigLocalFileEventLoggingAdapter>),
102}
103
104pub trait BoxableInstance {
105    fn from_box(boxed: &BoxedInstance) -> Option<Arc<Self>>;
106    fn into_box(self) -> BoxedInstance;
107    fn get_display_value_static() -> String;
108    fn get_display_value(&self) -> String;
109    fn get_prefix_value(&self) -> String;
110}
111
112impl_all_instances! {
113    (Statsig, Statsig) => "stsg",
114    (StatsigOptions, StatsigOptions) => "opts",
115    (StatsigUser, StatsigUser) => "usr",
116    (StatsigHttpSpecsAdapter, StatsigHttpSpecsAdapter) => "spc_http",
117    (StatsigLocalFileSpecsAdapter, StatsigLocalFileSpecsAdapter) => "spc_file",
118    (StatsigHttpEventLoggingAdapter, StatsigHttpEventLoggingAdapter) => "evt_http",
119    (StatsigLocalFileEventLoggingAdapter, StatsigLocalFileEventLoggingAdapter) => "evt_file",
120}
121
122pub struct InstanceStore {
123    instances: RwLock<HashMap<String, BoxedInstance>>,
124}
125
126impl Default for InstanceStore {
127    fn default() -> Self {
128        Self::new()
129    }
130}
131
132impl InstanceStore {
133    #[must_use]
134    pub fn new() -> Self {
135        Self {
136            instances: RwLock::new(HashMap::new()),
137        }
138    }
139
140    pub fn add(&self, boxable: impl BoxableInstance) -> Option<String> {
141        let mut instances = self.instances.write().ok()?;
142
143        if instances.len() >= MAX_STORED_INSTANCES {
144            log_e!(
145                TAG,
146                "Too many {} references created. Max ID limit reached.",
147                boxable.get_display_value()
148            );
149            return None;
150        }
151
152        let id_prefix = boxable.get_prefix_value();
153        let mut retries = 0;
154
155        loop {
156            let id = format!("{}_{}", id_prefix, Uuid::new_v4());
157
158            // Check for collisions
159            if !instances.contains_key(&id) {
160                log_d!(TAG, "Added {} {}", boxable.get_display_value(), &id);
161                instances.insert(id.clone(), boxable.into_box());
162                return Some(id);
163            }
164
165            retries += 1;
166            if retries > 10 {
167                let err_msg = format!(
168                    "Failed to generate a unique ID for {} after multiple attempts.",
169                    boxable.get_display_value()
170                );
171                log_e!(TAG, "{}", err_msg);
172                return None;
173            }
174
175            log_w!(TAG, "Collision, retrying...");
176        }
177    }
178
179    pub fn get<T: BoxableInstance>(&self, id: &str) -> Option<Arc<T>> {
180        let instances = if let Ok(instances) = self.instances.read() {
181            instances
182        } else {
183            log_e!(TAG, "Instance store is poisoned");
184            return None;
185        };
186
187        let found = if let Some(inst) = instances.get(id) {
188            inst
189        } else {
190            log_d!(
191                TAG,
192                "{} instance not found for ID {}",
193                T::get_display_value_static(),
194                id
195            );
196            return None;
197        };
198
199        if let Some(inst) = T::from_box(found) {
200            Some(inst)
201        } else {
202            log_e!(
203                TAG,
204                "Invalid box type for {}",
205                T::get_display_value_static()
206            );
207            None
208        }
209    }
210
211    pub fn get_with_optional_id<T: BoxableInstance>(&self, id: Option<&String>) -> Option<Arc<T>> {
212        match id {
213            Some(id) => self.get(id),
214            None => None,
215        }
216    }
217
218    pub fn remove(&self, id: &str) -> Option<BoxedInstance> {
219        self.instances.write().ok()?.remove(id)
220    }
221
222    pub fn remove_all(&self) {
223        if let Ok(mut instances) = self.instances.write() {
224            instances.clear();
225        }
226    }
227}
statsig_rust/instance_store.rs

statsig_rust/
instance_store.rs
