statsig_rust/
instance_registry.rs

1use ahash::AHashMap;
2use lazy_static::lazy_static;
3use parking_lot::{RwLock, RwLockWriteGuard};
4use std::any::Any;
5use std::sync::Arc;
6use std::time::Duration;
7use uuid::Uuid;
8
9use crate::hashing::hash_one;
10use crate::log_e;
11
12type AnyInstance = Arc<dyn Any + Send + Sync>;
13
14lazy_static! {
15    static ref REGISTRY: RwLock<AHashMap<u64, AnyInstance>> = RwLock::new(AHashMap::default());
16}
17
18const TAG: &str = "InstanceRegistry";
19
20pub struct InstanceRegistry;
21
22impl InstanceRegistry {
23    pub fn register_arc<T: Send + Sync + 'static>(instance: Arc<T>) -> Option<u64> {
24        let full_type_name = std::any::type_name::<T>();
25        let short_type_name = full_type_name.split("::").last().unwrap_or(full_type_name);
26        let id_tuple = (short_type_name, Uuid::new_v4());
27        let id_hash = hash_one(id_tuple);
28
29        let mut registry = Self::get_write_lock()?;
30        registry.insert(id_hash, instance);
31
32        Some(id_hash)
33    }
34
35    pub fn register<T: Send + Sync + 'static>(instance: T) -> Option<u64> {
36        Self::register_arc(Arc::new(instance))
37    }
38
39    pub fn get_with_optional_id<T: Send + Sync + 'static>(id: Option<&u64>) -> Option<Arc<T>> {
40        id.and_then(|id_hash| Self::get::<T>(id_hash))
41    }
42
43    pub fn get<T: Send + Sync + 'static>(id: &u64) -> Option<Arc<T>> {
44        let registry = match REGISTRY.try_read_for(Duration::from_secs(5)) {
45            Some(guard) => guard,
46            None => {
47                log_e!(TAG, "Failed to acquire read lock: Failed to lock REGISTRY");
48                return None;
49            }
50        };
51
52        registry
53            .get(id)
54            .and_then(|any_arc| match any_arc.clone().downcast::<T>() {
55                Ok(t) => Some(t),
56                Err(_) => {
57                    log_e!(
58                        TAG,
59                        "Failed to downcast instance with ref '{}' to generic type",
60                        id
61                    );
62                    None
63                }
64            })
65    }
66
67    pub fn get_raw(id: &u64) -> Option<Arc<dyn Any + Send + Sync>> {
68        let registry = match REGISTRY.try_read_for(Duration::from_secs(5)) {
69            Some(guard) => guard,
70            None => {
71                log_e!(TAG, "Failed to acquire read lock: Failed to lock REGISTRY");
72                return None;
73            }
74        };
75
76        registry.get(id).cloned()
77    }
78
79    pub fn remove(id: &u64) {
80        let mut registry = match Self::get_write_lock() {
81            Some(registry) => registry,
82            None => return,
83        };
84        registry.remove(id);
85    }
86
87    pub fn remove_all() {
88        let mut registry = match Self::get_write_lock() {
89            Some(registry) => registry,
90            None => return,
91        };
92        registry.clear();
93    }
94
95    fn get_write_lock() -> Option<RwLockWriteGuard<'static, AHashMap<u64, AnyInstance>>> {
96        match REGISTRY.try_write_for(Duration::from_secs(5)) {
97            Some(registry) => Some(registry),
98            None => {
99                log_e!(TAG, "Failed to acquire write lock: Failed to lock REGISTRY");
100                None
101            }
102        }
103    }
104}
105
106#[macro_export]
107macro_rules! get_instance_or_noop {
108    ($type:ty, $ref:expr) => {
109        match statsig_rust::InstanceRegistry::get::<$type>($ref) {
110            Some(instance) => instance,
111            None => {
112                $crate::log_w!(TAG, "{} Reference not found {}", stringify!($type), $ref);
113                return;
114            }
115        }
116    };
117}
118
119#[macro_export]
120macro_rules! get_instance_or_return {
121    ($type:ty, $ref:expr, $return_val:expr) => {
122        match statsig_rust::InstanceRegistry::get::<$type>($ref) {
123            Some(instance) => instance,
124            None => {
125                $crate::log_w!(TAG, "{} Reference not found {}", stringify!($type), $ref);
126                return $return_val;
127            }
128        }
129    };
130}
131
132#[macro_export]
133macro_rules! get_instance_or_else {
134    ($type:ty, $ref:expr, $else:expr) => {
135        match statsig_rust::InstanceRegistry::get::<$type>($ref) {
136            Some(instance) => instance,
137            None => $else,
138        }
139    };
140}
statsig_rust/instance_registry.rs

statsig_rust/
instance_registry.rs
