gitdataai/lib/track/metrics.rs

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<MetricsRegistryInner>,
}

struct MetricsRegistryInner {
    registry: Registry,
    counters: Mutex<HashMap<String, prometheus::Counter>>,
    histograms: Mutex<HashMap<String, prometheus::Histogram>>,
    counter_vecs: Mutex<HashMap<String, CounterVec>>,
    histogram_vecs: Mutex<HashMap<String, HistogramVec>>,
    gauges: Mutex<HashMap<String, Gauge>>,
}

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<prometheus::Counter> {
        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<f64>,
    ) -> prometheus::Result<prometheus::Histogram> {
        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<CounterVec> {
        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<f64>,
    ) -> prometheus::Result<HistogramVec> {
        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<prometheus::Gauge> {
        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<String, String> {
        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<prometheus::proto::MetricFamily> {
        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<KeyValue> = 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<u64>>,
        >,
    > = 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"));
    }
}