1use crate::{Counter, Gauge, Timer, RateMeter};
7use std::collections::HashMap;
8use std::sync::{Arc, RwLock};
9
10#[repr(align(64))]
15pub struct Registry {
16 counters: RwLock<HashMap<String, Arc<Counter>>>,
17 gauges: RwLock<HashMap<String, Arc<Gauge>>>,
18 timers: RwLock<HashMap<String, Arc<Timer>>>,
19 rate_meters: RwLock<HashMap<String, Arc<RateMeter>>>,
20}
21
22impl Registry {
23 pub fn new() -> Self {
25 Self {
26 counters: RwLock::new(HashMap::new()),
27 gauges: RwLock::new(HashMap::new()),
28 timers: RwLock::new(HashMap::new()),
29 rate_meters: RwLock::new(HashMap::new()),
30 }
31 }
32
33 pub fn get_or_create_counter(&self, name: &str) -> Arc<Counter> {
35 if let Ok(counters) = self.counters.read() {
37 if let Some(counter) = counters.get(name) {
38 return counter.clone();
39 }
40 }
41
42 let mut counters = self.counters.write().unwrap();
44 counters.entry(name.to_string())
45 .or_insert_with(|| Arc::new(Counter::new()))
46 .clone()
47 }
48
49 pub fn get_or_create_gauge(&self, name: &str) -> Arc<Gauge> {
51 if let Ok(gauges) = self.gauges.read() {
53 if let Some(gauge) = gauges.get(name) {
54 return gauge.clone();
55 }
56 }
57
58 let mut gauges = self.gauges.write().unwrap();
60 gauges.entry(name.to_string())
61 .or_insert_with(|| Arc::new(Gauge::new()))
62 .clone()
63 }
64
65 pub fn get_or_create_timer(&self, name: &str) -> Arc<Timer> {
67 if let Ok(timers) = self.timers.read() {
69 if let Some(timer) = timers.get(name) {
70 return timer.clone();
71 }
72 }
73
74 let mut timers = self.timers.write().unwrap();
76 timers.entry(name.to_string())
77 .or_insert_with(|| Arc::new(Timer::new()))
78 .clone()
79 }
80
81 pub fn get_or_create_rate_meter(&self, name: &str) -> Arc<RateMeter> {
83 if let Ok(rate_meters) = self.rate_meters.read() {
85 if let Some(rate_meter) = rate_meters.get(name) {
86 return rate_meter.clone();
87 }
88 }
89
90 let mut rate_meters = self.rate_meters.write().unwrap();
92 rate_meters.entry(name.to_string())
93 .or_insert_with(|| Arc::new(RateMeter::new()))
94 .clone()
95 }
96
97 pub fn counter_names(&self) -> Vec<String> {
99 self.counters.read().unwrap().keys().cloned().collect()
100 }
101
102 pub fn gauge_names(&self) -> Vec<String> {
104 self.gauges.read().unwrap().keys().cloned().collect()
105 }
106
107 pub fn timer_names(&self) -> Vec<String> {
109 self.timers.read().unwrap().keys().cloned().collect()
110 }
111
112 pub fn rate_meter_names(&self) -> Vec<String> {
114 self.rate_meters.read().unwrap().keys().cloned().collect()
115 }
116
117 pub fn metric_count(&self) -> usize {
119 self.counters.read().unwrap().len() +
120 self.gauges.read().unwrap().len() +
121 self.timers.read().unwrap().len() +
122 self.rate_meters.read().unwrap().len()
123 }
124
125 pub fn clear(&self) {
127 self.counters.write().unwrap().clear();
128 self.gauges.write().unwrap().clear();
129 self.timers.write().unwrap().clear();
130 self.rate_meters.write().unwrap().clear();
131 }
132}
133
134impl Default for Registry {
135 fn default() -> Self {
136 Self::new()
137 }
138}
139
140unsafe impl Send for Registry {}
141unsafe impl Sync for Registry {}
142
143#[cfg(test)]
144mod tests {
145 use super::*;
146 use std::sync::Arc;
147 use std::thread;
148
149 #[test]
150 fn test_counter_registration() {
151 let registry = Registry::new();
152
153 let counter1 = registry.get_or_create_counter("requests");
154 let counter2 = registry.get_or_create_counter("requests");
155
156 assert!(Arc::ptr_eq(&counter1, &counter2));
158 assert_eq!(registry.metric_count(), 1);
159 }
160
161 #[test]
162 fn test_gauge_registration() {
163 let registry = Registry::new();
164
165 let gauge1 = registry.get_or_create_gauge("cpu_usage");
166 let gauge2 = registry.get_or_create_gauge("cpu_usage");
167
168 assert!(Arc::ptr_eq(&gauge1, &gauge2));
170 assert_eq!(registry.metric_count(), 1);
171 }
172
173 #[test]
174 fn test_timer_registration() {
175 let registry = Registry::new();
176
177 let timer1 = registry.get_or_create_timer("db_query");
178 let timer2 = registry.get_or_create_timer("db_query");
179
180 assert!(Arc::ptr_eq(&timer1, &timer2));
182 assert_eq!(registry.metric_count(), 1);
183 }
184
185 #[test]
186 fn test_rate_meter_registration() {
187 let registry = Registry::new();
188
189 let meter1 = registry.get_or_create_rate_meter("api_calls");
190 let meter2 = registry.get_or_create_rate_meter("api_calls");
191
192 assert!(Arc::ptr_eq(&meter1, &meter2));
194 assert_eq!(registry.metric_count(), 1);
195 }
196
197 #[test]
198 fn test_mixed_metrics() {
199 let registry = Registry::new();
200
201 let _counter = registry.get_or_create_counter("requests");
202 let _gauge = registry.get_or_create_gauge("cpu_usage");
203 let _timer = registry.get_or_create_timer("db_query");
204 let _meter = registry.get_or_create_rate_meter("api_calls");
205
206 assert_eq!(registry.metric_count(), 4);
207 assert_eq!(registry.counter_names().len(), 1);
208 assert_eq!(registry.gauge_names().len(), 1);
209 assert_eq!(registry.timer_names().len(), 1);
210 assert_eq!(registry.rate_meter_names().len(), 1);
211 }
212
213 #[test]
214 fn test_concurrent_access() {
215 let registry = Arc::new(Registry::new());
216 let mut handles = vec![];
217
218 for _i in 0..10 {
220 let registry = registry.clone();
221 let handle = thread::spawn(move || {
222 let counter = registry.get_or_create_counter("concurrent_test");
223 counter.inc();
224 counter.get()
225 });
226 handles.push(handle);
227 }
228
229 for handle in handles {
231 handle.join().unwrap();
232 }
233
234 assert_eq!(registry.metric_count(), 1);
236 let counter = registry.get_or_create_counter("concurrent_test");
237 assert_eq!(counter.get(), 10);
238 }
239
240 #[test]
241 fn test_clear() {
242 let registry = Registry::new();
243
244 let _counter = registry.get_or_create_counter("requests");
245 let _gauge = registry.get_or_create_gauge("cpu_usage");
246
247 assert_eq!(registry.metric_count(), 2);
248
249 registry.clear();
250 assert_eq!(registry.metric_count(), 0);
251 }
252
253 #[test]
254 fn test_metric_names() {
255 let registry = Registry::new();
256
257 let _counter1 = registry.get_or_create_counter("requests");
258 let _counter2 = registry.get_or_create_counter("errors");
259 let _gauge1 = registry.get_or_create_gauge("cpu_usage");
260
261 let counter_names = registry.counter_names();
262 let gauge_names = registry.gauge_names();
263
264 assert_eq!(counter_names.len(), 2);
265 assert_eq!(gauge_names.len(), 1);
266 assert!(counter_names.contains(&"requests".to_string()));
267 assert!(counter_names.contains(&"errors".to_string()));
268 assert!(gauge_names.contains(&"cpu_usage".to_string()));
269 }
270}