llm_optimizer_processor/metrics/
registry.rs

1//! Global metrics registry for Prometheus metrics
2//!
3//! This module provides a singleton registry for all processor metrics,
4//! ensuring thread-safe access and metric registration.
5
6use prometheus_client::registry::Registry;
7use std::sync::{Arc, OnceLock};
8use parking_lot::RwLock;
9
10/// Global metrics registry instance
11pub static METRICS_REGISTRY: OnceLock<Arc<MetricsRegistry>> = OnceLock::new();
12
13/// Thread-safe registry for Prometheus metrics
14pub struct MetricsRegistry {
15    registry: Arc<RwLock<Registry>>,
16}
17
18impl MetricsRegistry {
19    /// Create a new metrics registry
20    pub fn new() -> Self {
21        Self {
22            registry: Arc::new(RwLock::new(Registry::default())),
23        }
24    }
25
26    /// Get the global metrics registry, initializing it if necessary
27    pub fn global() -> Arc<MetricsRegistry> {
28        METRICS_REGISTRY
29            .get_or_init(|| Arc::new(MetricsRegistry::new()))
30            .clone()
31    }
32
33    /// Get a reference to the internal registry
34    pub fn registry(&self) -> Arc<RwLock<Registry>> {
35        self.registry.clone()
36    }
37
38    /// Encode metrics in Prometheus text format
39    pub fn encode(&self) -> Result<String, crate::metrics::MetricsError> {
40        let registry = self.registry.read();
41        let mut buffer = String::new();
42
43        prometheus_client::encoding::text::encode(&mut buffer, &registry)
44            .map_err(|e| crate::metrics::MetricsError::EncodingError(e.to_string()))?;
45
46        Ok(buffer)
47    }
48}
49
50impl Default for MetricsRegistry {
51    fn default() -> Self {
52        Self::new()
53    }
54}
55
56impl Clone for MetricsRegistry {
57    fn clone(&self) -> Self {
58        Self {
59            registry: self.registry.clone(),
60        }
61    }
62}
63
64#[cfg(test)]
65mod tests {
66    use super::*;
67
68    #[test]
69    fn test_registry_creation() {
70        let registry = MetricsRegistry::new();
71        assert!(registry.encode().is_ok());
72    }
73
74    #[test]
75    fn test_global_registry() {
76        let registry1 = MetricsRegistry::global();
77        let registry2 = MetricsRegistry::global();
78
79        // Should return the same instance
80        assert!(Arc::ptr_eq(&registry1, &registry2));
81    }
82
83    #[test]
84    fn test_registry_clone() {
85        let registry1 = MetricsRegistry::new();
86        let registry2 = registry1.clone();
87
88        // Both should share the same underlying registry
89        assert!(Arc::ptr_eq(&registry1.registry(), &registry2.registry()));
90    }
91}
llm_optimizer_processor/metrics/registry.rs

llm_optimizer_processor/metrics/
registry.rs
