statsig_rust/
instance_registry.rs

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