santh-tracing 0.2.0

//! Prometheus-backed metrics implementation.

use std::collections::HashMap;
use std::sync::{Mutex, OnceLock};

use prometheus::{CounterVec, GaugeVec, HistogramOpts, HistogramVec, Opts};

use crate::get_tool_name;

enum MetricVec {
    Counter(CounterVec),
    Histogram(HistogramVec),
    Gauge(GaugeVec),
}

static METRICS: OnceLock<Mutex<HashMap<String, MetricVec>>> = OnceLock::new();

fn metric_key(name: &str, label_names: &[&str]) -> String {
    let mut key = name.to_owned();
    for ln in label_names {
        key.push('|');
        key.push_str(ln);
    }
    key
}

fn get_or_create_counter(name: &str, labels: &[(&str, &str)]) -> Option<CounterVec> {
    let map = METRICS.get_or_init(|| Mutex::new(HashMap::new()));
    let label_names: Vec<&str> = labels.iter().map(|(k, _)| *k).collect();
    let key = metric_key(name, &label_names);

    {
        let m = map.lock().ok()?;
        if let Some(MetricVec::Counter(vec)) = m.get(&key) {
            return Some(vec.clone());
        }
    }

    let opts = Opts::new(name, format!("santh counter {name}"));
    let vec = CounterVec::new(opts, &label_names).ok()?;
    let _ = prometheus::default_registry().register(Box::new(vec.clone()));
    map.lock()
        .ok()?
        .insert(key, MetricVec::Counter(vec.clone()));
    Some(vec)
}

fn get_or_create_histogram(name: &str, labels: &[(&str, &str)]) -> Option<HistogramVec> {
    let map = METRICS.get_or_init(|| Mutex::new(HashMap::new()));
    let label_names: Vec<&str> = labels.iter().map(|(k, _)| *k).collect();
    let key = metric_key(name, &label_names);

    {
        let m = map.lock().ok()?;
        if let Some(MetricVec::Histogram(vec)) = m.get(&key) {
            return Some(vec.clone());
        }
    }

    let opts = HistogramOpts::new(name, format!("santh histogram {name}"));
    let vec = HistogramVec::new(opts, &label_names).ok()?;
    let _ = prometheus::default_registry().register(Box::new(vec.clone()));
    map.lock()
        .ok()?
        .insert(key, MetricVec::Histogram(vec.clone()));
    Some(vec)
}

fn get_or_create_gauge(name: &str, labels: &[(&str, &str)]) -> Option<GaugeVec> {
    let map = METRICS.get_or_init(|| Mutex::new(HashMap::new()));
    let label_names: Vec<&str> = labels.iter().map(|(k, _)| *k).collect();
    let key = metric_key(name, &label_names);

    {
        let m = map.lock().ok()?;
        if let Some(MetricVec::Gauge(vec)) = m.get(&key) {
            return Some(vec.clone());
        }
    }

    let opts = Opts::new(name, format!("santh gauge {name}"));
    let vec = GaugeVec::new(opts, &label_names).ok()?;
    let _ = prometheus::default_registry().register(Box::new(vec.clone()));
    map.lock().ok()?.insert(key, MetricVec::Gauge(vec.clone()));
    Some(vec)
}

/// Increment a counter metric.
///
/// Metrics are automatically namespaced under `santh_<tool_name>_<name>`.
pub fn counter(name: &str, labels: &[(&str, &str)]) {
    let tool = get_tool_name();
    let full_name = format!("santh_{tool}_{name}");
    if let Some(vec) = get_or_create_counter(&full_name, labels) {
        let values: Vec<&str> = labels.iter().map(|(_, v)| *v).collect();
        vec.with_label_values(&values).inc();
    }
}

/// Record a histogram observation.
pub fn histogram(name: &str, value: f64, labels: &[(&str, &str)]) {
    let tool = get_tool_name();
    let full_name = format!("santh_{tool}_{name}");
    if let Some(vec) = get_or_create_histogram(&full_name, labels) {
        let values: Vec<&str> = labels.iter().map(|(_, v)| *v).collect();
        vec.with_label_values(&values).observe(value);
    }
}

/// Set a gauge value.
pub fn gauge(name: &str, value: f64, labels: &[(&str, &str)]) {
    let tool = get_tool_name();
    let full_name = format!("santh_{tool}_{name}");
    if let Some(vec) = get_or_create_gauge(&full_name, labels) {
        let values: Vec<&str> = labels.iter().map(|(_, v)| *v).collect();
        vec.with_label_values(&values).set(value);
    }
}