statsig_rust/
instance_registry.rs

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

statsig_rust/
instance_registry.rs
