Skip to main content

metrics_lib/
registry.rs

1//! Thread-safe registry for storing and retrieving metrics by name.
2//!
3//! The `Registry` uses one `RwLock<HashMap>` per enabled metric type.
4//! Read-heavy workloads (metric *lookups*) proceed with shared read locks and
5//! minimal contention. Write locks are only acquired when a new metric name is
6//! registered, which is expected to be an infrequent, one-time-per-name event.
7
8#[cfg(any(
9    feature = "count",
10    feature = "gauge",
11    feature = "timer",
12    feature = "meter"
13))]
14use std::collections::HashMap;
15#[cfg(any(
16    feature = "count",
17    feature = "gauge",
18    feature = "timer",
19    feature = "meter"
20))]
21use std::sync::{Arc, RwLock};
22
23#[cfg(feature = "count")]
24use crate::Counter;
25#[cfg(feature = "gauge")]
26use crate::Gauge;
27#[cfg(feature = "meter")]
28use crate::RateMeter;
29#[cfg(feature = "timer")]
30use crate::Timer;
31
32/// A thread-safe registry for storing metrics by name.
33///
34/// Stores one `Arc`-wrapped instance per unique metric name for each enabled
35/// metric type. Repeated lookups for the same name return the same `Arc`
36/// (pointer equality holds). Only metric types enabled via Cargo features are
37/// compiled in; attempting to call a method for a disabled type will result in
38/// a compile-time error.
39///
40/// Cache-line aligned to prevent false sharing in concurrent environments.
41#[repr(align(64))]
42pub struct Registry {
43    #[cfg(feature = "count")]
44    counters: RwLock<HashMap<String, Arc<Counter>>>,
45    #[cfg(feature = "gauge")]
46    gauges: RwLock<HashMap<String, Arc<Gauge>>>,
47    #[cfg(feature = "timer")]
48    timers: RwLock<HashMap<String, Arc<Timer>>>,
49    #[cfg(feature = "meter")]
50    rate_meters: RwLock<HashMap<String, Arc<RateMeter>>>,
51}
52
53impl Registry {
54    /// Create a new empty registry.
55    pub fn new() -> Self {
56        Self {
57            #[cfg(feature = "count")]
58            counters: RwLock::new(HashMap::new()),
59            #[cfg(feature = "gauge")]
60            gauges: RwLock::new(HashMap::new()),
61            #[cfg(feature = "timer")]
62            timers: RwLock::new(HashMap::new()),
63            #[cfg(feature = "meter")]
64            rate_meters: RwLock::new(HashMap::new()),
65        }
66    }
67
68    /// Get or create a counter by name.
69    ///
70    /// Requires the `count` feature.
71    #[cfg(feature = "count")]
72    pub fn get_or_create_counter(&self, name: &str) -> Arc<Counter> {
73        // Fast path: try read lock first
74        if let Ok(counters) = self.counters.read() {
75            if let Some(counter) = counters.get(name) {
76                return counter.clone();
77            }
78        }
79
80        // Slow path: write lock to create new counter
81        let mut counters = self.counters.write().unwrap_or_else(|e| e.into_inner());
82        counters
83            .entry(name.to_string())
84            .or_insert_with(|| Arc::new(Counter::new()))
85            .clone()
86    }
87
88    /// Get or create a gauge by name.
89    ///
90    /// Requires the `gauge` feature.
91    #[cfg(feature = "gauge")]
92    pub fn get_or_create_gauge(&self, name: &str) -> Arc<Gauge> {
93        // Fast path: try read lock first
94        if let Ok(gauges) = self.gauges.read() {
95            if let Some(gauge) = gauges.get(name) {
96                return gauge.clone();
97            }
98        }
99
100        // Slow path: write lock to create new gauge
101        let mut gauges = self.gauges.write().unwrap_or_else(|e| e.into_inner());
102        gauges
103            .entry(name.to_string())
104            .or_insert_with(|| Arc::new(Gauge::new()))
105            .clone()
106    }
107
108    /// Get or create a timer by name.
109    ///
110    /// Requires the `timer` feature.
111    #[cfg(feature = "timer")]
112    pub fn get_or_create_timer(&self, name: &str) -> Arc<Timer> {
113        // Fast path: try read lock first
114        if let Ok(timers) = self.timers.read() {
115            if let Some(timer) = timers.get(name) {
116                return timer.clone();
117            }
118        }
119
120        // Slow path: write lock to create new timer
121        let mut timers = self.timers.write().unwrap_or_else(|e| e.into_inner());
122        timers
123            .entry(name.to_string())
124            .or_insert_with(|| Arc::new(Timer::new()))
125            .clone()
126    }
127
128    /// Get or create a rate meter by name.
129    ///
130    /// Requires the `meter` feature.
131    #[cfg(feature = "meter")]
132    pub fn get_or_create_rate_meter(&self, name: &str) -> Arc<RateMeter> {
133        // Fast path: try read lock first
134        if let Ok(rate_meters) = self.rate_meters.read() {
135            if let Some(rate_meter) = rate_meters.get(name) {
136                return rate_meter.clone();
137            }
138        }
139
140        // Slow path: write lock to create new rate meter
141        let mut rate_meters = self.rate_meters.write().unwrap_or_else(|e| e.into_inner());
142        rate_meters
143            .entry(name.to_string())
144            .or_insert_with(|| Arc::new(RateMeter::new()))
145            .clone()
146    }
147
148    /// Get all registered counter names. Requires the `count` feature.
149    #[cfg(feature = "count")]
150    pub fn counter_names(&self) -> Vec<String> {
151        self.counters
152            .read()
153            .unwrap_or_else(|e| e.into_inner())
154            .keys()
155            .cloned()
156            .collect()
157    }
158
159    /// Get all registered gauge names. Requires the `gauge` feature.
160    #[cfg(feature = "gauge")]
161    pub fn gauge_names(&self) -> Vec<String> {
162        self.gauges
163            .read()
164            .unwrap_or_else(|e| e.into_inner())
165            .keys()
166            .cloned()
167            .collect()
168    }
169
170    /// Get all registered timer names. Requires the `timer` feature.
171    #[cfg(feature = "timer")]
172    pub fn timer_names(&self) -> Vec<String> {
173        self.timers
174            .read()
175            .unwrap_or_else(|e| e.into_inner())
176            .keys()
177            .cloned()
178            .collect()
179    }
180
181    /// Get all registered rate meter names. Requires the `meter` feature.
182    #[cfg(feature = "meter")]
183    pub fn rate_meter_names(&self) -> Vec<String> {
184        self.rate_meters
185            .read()
186            .unwrap_or_else(|e| e.into_inner())
187            .keys()
188            .cloned()
189            .collect()
190    }
191
192    /// Get total number of registered metrics across all enabled metric types.
193    pub fn metric_count(&self) -> usize {
194        #[allow(unused_mut)]
195        let mut total = 0;
196        #[cfg(feature = "count")]
197        {
198            total += self
199                .counters
200                .read()
201                .unwrap_or_else(|e| e.into_inner())
202                .len();
203        }
204        #[cfg(feature = "gauge")]
205        {
206            total += self.gauges.read().unwrap_or_else(|e| e.into_inner()).len();
207        }
208        #[cfg(feature = "timer")]
209        {
210            total += self.timers.read().unwrap_or_else(|e| e.into_inner()).len();
211        }
212        #[cfg(feature = "meter")]
213        {
214            total += self
215                .rate_meters
216                .read()
217                .unwrap_or_else(|e| e.into_inner())
218                .len();
219        }
220        total
221    }
222
223    /// Clear all metrics from the registry.
224    pub fn clear(&self) {
225        #[cfg(feature = "count")]
226        self.counters
227            .write()
228            .unwrap_or_else(|e| e.into_inner())
229            .clear();
230        #[cfg(feature = "gauge")]
231        self.gauges
232            .write()
233            .unwrap_or_else(|e| e.into_inner())
234            .clear();
235        #[cfg(feature = "timer")]
236        self.timers
237            .write()
238            .unwrap_or_else(|e| e.into_inner())
239            .clear();
240        #[cfg(feature = "meter")]
241        self.rate_meters
242            .write()
243            .unwrap_or_else(|e| e.into_inner())
244            .clear();
245    }
246}
247
248impl Default for Registry {
249    fn default() -> Self {
250        Self::new()
251    }
252}
253
254// `Registry` is Send + Sync automatically because `RwLock<HashMap<…>>` is
255// Send + Sync when its contents are Send + Sync. No unsafe impls required.
256
257#[cfg(test)]
258// Registry integration tests require the default metric features.
259#[cfg(all(feature = "count", feature = "gauge", feature = "timer"))]
260mod tests {
261    use super::*;
262    use std::sync::Arc;
263    use std::thread;
264    #[test]
265    fn test_counter_registration() {
266        let registry = Registry::new();
267
268        let counter1 = registry.get_or_create_counter("requests");
269        let counter2 = registry.get_or_create_counter("requests");
270
271        // Should return the same instance
272        assert!(Arc::ptr_eq(&counter1, &counter2));
273        assert_eq!(registry.metric_count(), 1);
274    }
275
276    #[test]
277    fn test_gauge_registration() {
278        let registry = Registry::new();
279
280        let gauge1 = registry.get_or_create_gauge("cpu_usage");
281        let gauge2 = registry.get_or_create_gauge("cpu_usage");
282
283        // Should return the same instance
284        assert!(Arc::ptr_eq(&gauge1, &gauge2));
285        assert_eq!(registry.metric_count(), 1);
286    }
287
288    #[test]
289    fn test_timer_registration() {
290        let registry = Registry::new();
291
292        let timer1 = registry.get_or_create_timer("db_query");
293        let timer2 = registry.get_or_create_timer("db_query");
294
295        // Should return the same instance
296        assert!(Arc::ptr_eq(&timer1, &timer2));
297        assert_eq!(registry.metric_count(), 1);
298    }
299
300    #[test]
301    #[cfg(feature = "meter")]
302    fn test_rate_meter_registration() {
303        let registry = Registry::new();
304
305        let meter1 = registry.get_or_create_rate_meter("api_calls");
306        let meter2 = registry.get_or_create_rate_meter("api_calls");
307
308        // Should return the same instance
309        assert!(Arc::ptr_eq(&meter1, &meter2));
310        assert_eq!(registry.metric_count(), 1);
311    }
312
313    #[test]
314    #[cfg(feature = "meter")]
315    fn test_mixed_metrics() {
316        let registry = Registry::new();
317
318        let _counter = registry.get_or_create_counter("requests");
319        let _gauge = registry.get_or_create_gauge("cpu_usage");
320        let _timer = registry.get_or_create_timer("db_query");
321        let _meter = registry.get_or_create_rate_meter("api_calls");
322
323        assert_eq!(registry.metric_count(), 4);
324        assert_eq!(registry.counter_names().len(), 1);
325        assert_eq!(registry.gauge_names().len(), 1);
326        assert_eq!(registry.timer_names().len(), 1);
327        assert_eq!(registry.rate_meter_names().len(), 1);
328    }
329
330    #[test]
331    fn test_concurrent_access() {
332        let registry = Arc::new(Registry::new());
333        let mut handles = vec![];
334
335        // Spawn multiple threads accessing the same counter
336        for _i in 0..10 {
337            let registry = registry.clone();
338            let handle = thread::spawn(move || {
339                let counter = registry.get_or_create_counter("concurrent_test");
340                counter.inc();
341                counter.get()
342            });
343            handles.push(handle);
344        }
345
346        // Wait for all threads to complete
347        for handle in handles {
348            handle.join().unwrap();
349        }
350
351        // Should have exactly one counter registered
352        assert_eq!(registry.metric_count(), 1);
353        let counter = registry.get_or_create_counter("concurrent_test");
354        assert_eq!(counter.get(), 10);
355    }
356
357    #[test]
358    fn test_clear() {
359        let registry = Registry::new();
360
361        let _counter = registry.get_or_create_counter("requests");
362        let _gauge = registry.get_or_create_gauge("cpu_usage");
363
364        assert_eq!(registry.metric_count(), 2);
365
366        registry.clear();
367        assert_eq!(registry.metric_count(), 0);
368    }
369
370    #[test]
371    fn test_metric_names() {
372        let registry = Registry::new();
373
374        let _counter1 = registry.get_or_create_counter("requests");
375        let _counter2 = registry.get_or_create_counter("errors");
376        let _gauge1 = registry.get_or_create_gauge("cpu_usage");
377
378        let counter_names = registry.counter_names();
379        let gauge_names = registry.gauge_names();
380
381        assert_eq!(counter_names.len(), 2);
382        assert_eq!(gauge_names.len(), 1);
383        assert!(counter_names.contains(&"requests".to_string()));
384        assert!(counter_names.contains(&"errors".to_string()));
385        assert!(gauge_names.contains(&"cpu_usage".to_string()));
386    }
387
388    #[test]
389    #[cfg(feature = "meter")]
390    fn test_duplicate_names_across_types_are_independent() {
391        let registry = Registry::new();
392
393        let c = registry.get_or_create_counter("same_name");
394        let g = registry.get_or_create_gauge("same_name");
395        let t = registry.get_or_create_timer("same_name");
396        let r = registry.get_or_create_rate_meter("same_name");
397
398        // All should be registered independently (4 metrics total)
399        assert_eq!(registry.metric_count(), 4);
400
401        // And they should be distinct types; at least ensure their addresses differ pairwise
402        let c_addr = Arc::as_ptr(&c) as usize;
403        let g_addr = Arc::as_ptr(&g) as usize;
404        let t_addr = Arc::as_ptr(&t) as usize;
405        let r_addr = Arc::as_ptr(&r) as usize;
406
407        assert_ne!(c_addr, g_addr);
408        assert_ne!(c_addr, t_addr);
409        assert_ne!(c_addr, r_addr);
410        assert_ne!(g_addr, t_addr);
411        assert_ne!(g_addr, r_addr);
412        assert_ne!(t_addr, r_addr);
413    }
414
415    #[test]
416    fn test_clear_then_recreate_returns_new_instances() {
417        let registry = Registry::new();
418
419        let counter_before = registry.get_or_create_counter("requests");
420        let gauge_before = registry.get_or_create_gauge("cpu");
421        assert_eq!(registry.metric_count(), 2);
422
423        // Clear the registry; previously returned Arcs still exist but registry is empty
424        registry.clear();
425        assert_eq!(registry.metric_count(), 0);
426
427        let counter_after = registry.get_or_create_counter("requests");
428        let gauge_after = registry.get_or_create_gauge("cpu");
429
430        // New instances should NOT be ptr-equal to the old ones
431        assert!(!Arc::ptr_eq(&counter_before, &counter_after));
432        assert!(!Arc::ptr_eq(&gauge_before, &gauge_after));
433    }
434
435    #[test]
436    fn test_concurrent_duplicate_registration_singleton_per_name() {
437        let registry = Arc::new(Registry::new());
438        let mut handles = vec![];
439
440        for _ in 0..16 {
441            let r = registry.clone();
442            handles.push(thread::spawn(move || r.get_or_create_timer("dup")));
443        }
444
445        let first = registry.get_or_create_timer("dup");
446        for h in handles {
447            let timer = h.join().unwrap();
448            assert!(Arc::ptr_eq(&first, &timer));
449        }
450        assert_eq!(registry.metric_count(), 1);
451    }
452}