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 replaced = {
30            let mut registry = Self::get_write_lock()?;
31            registry.insert(id_hash, instance)
32        };
33
34        // Registered instances may run arbitrary shutdown logic when dropped.
35        // Keep that work outside the global registry lock in the unlikely event
36        // of an ID collision.
37        drop(replaced);
38
39        Some(id_hash)
40    }
41
42    pub fn register<T: Send + Sync + 'static>(instance: T) -> Option<u64> {
43        Self::register_arc(Arc::new(instance))
44    }
45
46    pub fn get_with_optional_id<T: Send + Sync + 'static>(id: Option<&u64>) -> Option<Arc<T>> {
47        id.and_then(|id_hash| Self::get::<T>(id_hash))
48    }
49
50    pub fn get<T: Send + Sync + 'static>(id: &u64) -> Option<Arc<T>> {
51        let instance = {
52            let registry = match REGISTRY.try_read_for(Duration::from_secs(5)) {
53                Some(guard) => guard,
54                None => {
55                    log_e!(TAG, "Failed to acquire read lock: Failed to lock REGISTRY");
56                    return None;
57                }
58            };
59
60            registry.get(id).cloned()
61        };
62
63        instance.and_then(|any_arc| match any_arc.downcast::<T>() {
64            Ok(t) => Some(t),
65            Err(_) => {
66                log_e!(
67                    TAG,
68                    "Failed to downcast instance with ref '{}' to generic type",
69                    id
70                );
71                None
72            }
73        })
74    }
75
76    pub fn get_raw(id: &u64) -> Option<Arc<dyn Any + Send + Sync>> {
77        let registry = match REGISTRY.try_read_for(Duration::from_secs(5)) {
78            Some(guard) => guard,
79            None => {
80                log_e!(TAG, "Failed to acquire read lock: Failed to lock REGISTRY");
81                return None;
82            }
83        };
84
85        registry.get(id).cloned()
86    }
87
88    pub fn remove(id: &u64) {
89        let removed = {
90            let mut registry = match Self::get_write_lock() {
91                Some(registry) => registry,
92                None => return,
93            };
94            registry.remove(id)
95        };
96
97        // Instance destructors can run shutdown hooks and callbacks.
98        drop(removed);
99    }
100
101    pub fn remove_all() {
102        let removed = {
103            let mut registry = match Self::get_write_lock() {
104                Some(registry) => registry,
105                None => return,
106            };
107            std::mem::take(&mut *registry)
108        };
109
110        // Drop all instances after the registry is available to other callers.
111        drop(removed);
112    }
113
114    fn get_write_lock() -> Option<RwLockWriteGuard<'static, AHashMap<u64, AnyInstance>>> {
115        match REGISTRY.try_write_for(Duration::from_secs(5)) {
116            Some(registry) => Some(registry),
117            None => {
118                log_e!(TAG, "Failed to acquire write lock: Failed to lock REGISTRY");
119                None
120            }
121        }
122    }
123}
124
125#[macro_export]
126macro_rules! get_instance_or_noop {
127    ($type:ty, $ref:expr) => {
128        match statsig_rust::InstanceRegistry::get::<$type>($ref) {
129            Some(instance) => instance,
130            None => {
131                $crate::log_w!(TAG, "{} Reference not found {}", stringify!($type), $ref);
132                return;
133            }
134        }
135    };
136}
137
138#[macro_export]
139macro_rules! get_instance_or_return {
140    ($type:ty, $ref:expr, $return_val:expr) => {
141        match statsig_rust::InstanceRegistry::get::<$type>($ref) {
142            Some(instance) => instance,
143            None => {
144                $crate::log_w!(TAG, "{} Reference not found {}", stringify!($type), $ref);
145                return $return_val;
146            }
147        }
148    };
149}
150
151#[macro_export]
152macro_rules! get_instance_or_else {
153    ($type:ty, $ref:expr, $else:expr) => {
154        match statsig_rust::InstanceRegistry::get::<$type>($ref) {
155            Some(instance) => instance,
156            None => $else,
157        }
158    };
159}
statsig_rust/instance_registry.rs

statsig_rust/
instance_registry.rs
