use std::{ collections::HashMap, sync::{Arc, Mutex}, }; use opentelemetry::KeyValue; use prometheus::{ CounterVec, Encoder, Gauge, HistogramOpts, HistogramVec, Opts, Registry, TextEncoder, }; #[derive(Clone)] pub struct MetricsRegistry { inner: Arc, } struct MetricsRegistryInner { registry: Registry, counters: Mutex>, histograms: Mutex>, counter_vecs: Mutex>, histogram_vecs: Mutex>, gauges: Mutex>, } impl MetricsRegistry { pub fn new() -> Self { Self { inner: Arc::new(MetricsRegistryInner { registry: Registry::new(), counters: Mutex::new(HashMap::new()), histograms: Mutex::new(HashMap::new()), counter_vecs: Mutex::new(HashMap::new()), histogram_vecs: Mutex::new(HashMap::new()), gauges: Mutex::new(HashMap::new()), }), } } pub fn registry(&self) -> &Registry { &self.inner.registry } pub fn register_counter( &self, name: &str, help: &str, ) -> prometheus::Result { let mut counters = self.inner.counters.lock().expect("metrics mutex poisoned"); if let Some(counter) = counters.get(name) { return Ok(counter.clone()); } let counter = prometheus::Counter::new(name, help)?; self.inner.registry.register(Box::new(counter.clone()))?; counters.insert(name.to_string(), counter.clone()); Ok(counter) } pub fn register_histogram( &self, name: &str, help: &str, buckets: Vec, ) -> prometheus::Result { let mut histograms = self .inner .histograms .lock() .expect("metrics mutex poisoned"); if let Some(histogram) = histograms.get(name) { return Ok(histogram.clone()); } let opts = prometheus::HistogramOpts::new(name, help).buckets(buckets); let histogram = prometheus::Histogram::with_opts(opts)?; self.inner.registry.register(Box::new(histogram.clone()))?; histograms.insert(name.to_string(), histogram.clone()); Ok(histogram) } pub fn register_counter_vec( &self, name: &str, help: &str, labels: &[&str], ) -> prometheus::Result { let mut counter_vecs = self .inner .counter_vecs .lock() .expect("metrics mutex poisoned"); if let Some(cv) = counter_vecs.get(name) { return Ok(cv.clone()); } let opts = Opts::new(name, help); let cv = CounterVec::new(opts, labels)?; self.inner.registry.register(Box::new(cv.clone()))?; counter_vecs.insert(name.to_string(), cv.clone()); Ok(cv) } pub fn register_histogram_vec( &self, name: &str, help: &str, labels: &[&str], buckets: Vec, ) -> prometheus::Result { let mut histogram_vecs = self .inner .histogram_vecs .lock() .expect("metrics mutex poisoned"); if let Some(hv) = histogram_vecs.get(name) { return Ok(hv.clone()); } let opts = HistogramOpts::new(name, help).buckets(buckets); let hv = HistogramVec::new(opts, labels)?; self.inner.registry.register(Box::new(hv.clone()))?; histogram_vecs.insert(name.to_string(), hv.clone()); Ok(hv) } pub fn register_gauge( &self, name: &str, help: &str, ) -> prometheus::Result { let mut gauges = self.inner.gauges.lock().expect("metrics mutex poisoned"); if let Some(gauge) = gauges.get(name) { return Ok(gauge.clone()); } let gauge = prometheus::Gauge::new(name, help)?; self.inner.registry.register(Box::new(gauge.clone()))?; gauges.insert(name.to_string(), gauge.clone()); Ok(gauge) } pub fn encode(&self) -> Result { let mut buffer = Vec::new(); let encoder = TextEncoder::new(); encoder .encode(&self.inner.registry.gather(), &mut buffer) .map_err(|e| format!("failed to encode metrics: {e}"))?; String::from_utf8(buffer).map_err(|e| format!("invalid utf8: {e}")) } pub fn gather(&self) -> Vec { self.inner.registry.gather() } } impl Default for MetricsRegistry { fn default() -> Self { Self::new() } } pub fn record_otel_counter( name: &'static str, value: u64, labels: &[(&'static str, String)], ) { if value == 0 { return; } let attrs: Vec = labels .iter() .map(|(key, value)| KeyValue::new(*key, value.clone())) .collect(); static COUNTERS: std::sync::OnceLock< std::sync::Mutex< HashMap<&'static str, opentelemetry::metrics::Counter>, >, > = std::sync::OnceLock::new(); let counters = COUNTERS.get_or_init(|| std::sync::Mutex::new(HashMap::new())); let mut map = counters.lock().expect("otel counter mutex poisoned"); let counter = map .entry(name) .or_insert_with(|| { opentelemetry::global::meter("gitdataai") .u64_counter(name) .build() }) .clone(); drop(map); counter.add(value, &attrs); } #[cfg(test)] mod tests { use super::MetricsRegistry; #[test] fn repeated_counter_vec_registration_reuses_metric() { let registry = MetricsRegistry::new(); let first = registry .register_counter_vec("test_events_total", "Test events", &["kind"]) .expect("first registration should succeed"); let second = registry .register_counter_vec("test_events_total", "Test events", &["kind"]) .expect("second registration should reuse existing metric"); first.with_label_values(&["a"]).inc(); second.with_label_values(&["a"]).inc(); let encoded = registry.encode().expect("metrics should encode"); assert!(encoded.contains("test_events_total{kind=\"a\"} 2")); } #[test] fn repeated_histogram_vec_registration_reuses_metric() { let registry = MetricsRegistry::new(); let first = registry .register_histogram_vec( "test_duration_seconds", "Test duration", &["kind"], vec![0.1, 1.0], ) .expect("first registration should succeed"); let second = registry .register_histogram_vec( "test_duration_seconds", "Test duration", &["kind"], vec![0.1, 1.0], ) .expect("second registration should reuse existing metric"); first.with_label_values(&["a"]).observe(0.2); second.with_label_values(&["a"]).observe(0.3); let encoded = registry.encode().expect("metrics should encode"); assert!(encoded.contains("test_duration_seconds_count{kind=\"a\"} 2")); } }