1use crate::{Counter, Gauge, RateMeter, Timer};
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
45 .entry(name.to_string())
46 .or_insert_with(|| Arc::new(Counter::new()))
47 .clone()
48 }
49
50 pub fn get_or_create_gauge(&self, name: &str) -> Arc<Gauge> {
52 if let Ok(gauges) = self.gauges.read() {
54 if let Some(gauge) = gauges.get(name) {
55 return gauge.clone();
56 }
57 }
58
59 let mut gauges = self.gauges.write().unwrap();
61 gauges
62 .entry(name.to_string())
63 .or_insert_with(|| Arc::new(Gauge::new()))
64 .clone()
65 }
66
67 pub fn get_or_create_timer(&self, name: &str) -> Arc<Timer> {
69 if let Ok(timers) = self.timers.read() {
71 if let Some(timer) = timers.get(name) {
72 return timer.clone();
73 }
74 }
75
76 let mut timers = self.timers.write().unwrap();
78 timers
79 .entry(name.to_string())
80 .or_insert_with(|| Arc::new(Timer::new()))
81 .clone()
82 }
83
84 pub fn get_or_create_rate_meter(&self, name: &str) -> Arc<RateMeter> {
86 if let Ok(rate_meters) = self.rate_meters.read() {
88 if let Some(rate_meter) = rate_meters.get(name) {
89 return rate_meter.clone();
90 }
91 }
92
93 let mut rate_meters = self.rate_meters.write().unwrap();
95 rate_meters
96 .entry(name.to_string())
97 .or_insert_with(|| Arc::new(RateMeter::new()))
98 .clone()
99 }
100
101 pub fn counter_names(&self) -> Vec<String> {
103 self.counters.read().unwrap().keys().cloned().collect()
104 }
105
106 pub fn gauge_names(&self) -> Vec<String> {
108 self.gauges.read().unwrap().keys().cloned().collect()
109 }
110
111 pub fn timer_names(&self) -> Vec<String> {
113 self.timers.read().unwrap().keys().cloned().collect()
114 }
115
116 pub fn rate_meter_names(&self) -> Vec<String> {
118 self.rate_meters.read().unwrap().keys().cloned().collect()
119 }
120
121 pub fn metric_count(&self) -> usize {
123 self.counters.read().unwrap().len()
124 + self.gauges.read().unwrap().len()
125 + self.timers.read().unwrap().len()
126 + self.rate_meters.read().unwrap().len()
127 }
128
129 pub fn clear(&self) {
131 self.counters.write().unwrap().clear();
132 self.gauges.write().unwrap().clear();
133 self.timers.write().unwrap().clear();
134 self.rate_meters.write().unwrap().clear();
135 }
136}
137
138impl Default for Registry {
139 fn default() -> Self {
140 Self::new()
141 }
142}
143
144unsafe impl Send for Registry {}
145unsafe impl Sync for Registry {}
146
147#[cfg(test)]
148mod tests {
149 use super::*;
150 use std::sync::Arc;
151 use std::thread;
152
153 #[test]
154 fn test_counter_registration() {
155 let registry = Registry::new();
156
157 let counter1 = registry.get_or_create_counter("requests");
158 let counter2 = registry.get_or_create_counter("requests");
159
160 assert!(Arc::ptr_eq(&counter1, &counter2));
162 assert_eq!(registry.metric_count(), 1);
163 }
164
165 #[test]
166 fn test_gauge_registration() {
167 let registry = Registry::new();
168
169 let gauge1 = registry.get_or_create_gauge("cpu_usage");
170 let gauge2 = registry.get_or_create_gauge("cpu_usage");
171
172 assert!(Arc::ptr_eq(&gauge1, &gauge2));
174 assert_eq!(registry.metric_count(), 1);
175 }
176
177 #[test]
178 fn test_timer_registration() {
179 let registry = Registry::new();
180
181 let timer1 = registry.get_or_create_timer("db_query");
182 let timer2 = registry.get_or_create_timer("db_query");
183
184 assert!(Arc::ptr_eq(&timer1, &timer2));
186 assert_eq!(registry.metric_count(), 1);
187 }
188
189 #[test]
190 fn test_rate_meter_registration() {
191 let registry = Registry::new();
192
193 let meter1 = registry.get_or_create_rate_meter("api_calls");
194 let meter2 = registry.get_or_create_rate_meter("api_calls");
195
196 assert!(Arc::ptr_eq(&meter1, &meter2));
198 assert_eq!(registry.metric_count(), 1);
199 }
200
201 #[test]
202 fn test_mixed_metrics() {
203 let registry = Registry::new();
204
205 let _counter = registry.get_or_create_counter("requests");
206 let _gauge = registry.get_or_create_gauge("cpu_usage");
207 let _timer = registry.get_or_create_timer("db_query");
208 let _meter = registry.get_or_create_rate_meter("api_calls");
209
210 assert_eq!(registry.metric_count(), 4);
211 assert_eq!(registry.counter_names().len(), 1);
212 assert_eq!(registry.gauge_names().len(), 1);
213 assert_eq!(registry.timer_names().len(), 1);
214 assert_eq!(registry.rate_meter_names().len(), 1);
215 }
216
217 #[test]
218 fn test_concurrent_access() {
219 let registry = Arc::new(Registry::new());
220 let mut handles = vec![];
221
222 for _i in 0..10 {
224 let registry = registry.clone();
225 let handle = thread::spawn(move || {
226 let counter = registry.get_or_create_counter("concurrent_test");
227 counter.inc();
228 counter.get()
229 });
230 handles.push(handle);
231 }
232
233 for handle in handles {
235 handle.join().unwrap();
236 }
237
238 assert_eq!(registry.metric_count(), 1);
240 let counter = registry.get_or_create_counter("concurrent_test");
241 assert_eq!(counter.get(), 10);
242 }
243
244 #[test]
245 fn test_clear() {
246 let registry = Registry::new();
247
248 let _counter = registry.get_or_create_counter("requests");
249 let _gauge = registry.get_or_create_gauge("cpu_usage");
250
251 assert_eq!(registry.metric_count(), 2);
252
253 registry.clear();
254 assert_eq!(registry.metric_count(), 0);
255 }
256
257 #[test]
258 fn test_metric_names() {
259 let registry = Registry::new();
260
261 let _counter1 = registry.get_or_create_counter("requests");
262 let _counter2 = registry.get_or_create_counter("errors");
263 let _gauge1 = registry.get_or_create_gauge("cpu_usage");
264
265 let counter_names = registry.counter_names();
266 let gauge_names = registry.gauge_names();
267
268 assert_eq!(counter_names.len(), 2);
269 assert_eq!(gauge_names.len(), 1);
270 assert!(counter_names.contains(&"requests".to_string()));
271 assert!(counter_names.contains(&"errors".to_string()));
272 assert!(gauge_names.contains(&"cpu_usage".to_string()));
273 }
274}