prometheus_extensions/

lib.rs

//! Prometheus extensions for richer metric collection.
//!
//! This crate provides three utilities that complement the
//! [`prometheus`](https://docs.rs/prometheus) crate:
//!
//! | Type | Purpose |
//! |------|---------|
//! | [`AggregateCounter`] | A `CounterVec` wrapper that automatically emits an extra **unlabeled total** alongside every per-label counter. |
//! | [`ScientificEncoder`] | A lightweight Prometheus text-format encoder that renders counter values in scientific notation (`1.23E4`) with a trailing comma after the last label — matching the format Kafka JMX exporters produce. |
//! | [`Sensor`] | A lock-free exponential moving average (EMA) gauge for tracking rates or throughput. |
//!
//! # Quick start
//!
//! ```rust
//! use prometheus::Opts;
//! use prometheus::core::Collector;
//! use prometheus_extensions::{AggregateCounter, ScientificEncoder};
//!
//! // 1. Create an aggregate counter
//! let counter = AggregateCounter::new(
//!     Opts::new("http_requests_total", "Total HTTP requests"),
//!     &["method"],
//! ).unwrap();
//!
//! counter.with_label_values(&["GET"]).inc_by(100.0);
//! counter.with_label_values(&["POST"]).inc_by(42.0);
//!
//! // Collecting yields 3 metrics: the unlabeled total (142) plus the two labeled ones.
//! let families = counter.collect();
//! assert_eq!(families[0].get_metric().len(), 3);
//!
//! // 2. Encode in Kafka-compatible scientific notation
//! let encoder = ScientificEncoder::new();
//! let mut buf = Vec::new();
//! encoder.encode(&families, &mut buf).unwrap();
//! let output = String::from_utf8(buf).unwrap();
//! assert!(output.contains("http_requests_total 1.4200000000E2"));
//! ```
//!
//! # Sensor (EMA)
//!
//! ```rust
//! use prometheus_extensions::Sensor;
//!
//! let sensor = Sensor::new(0.05); // alpha = 0.05
//! sensor.measure(100.0);
//! sensor.measure(200.0);
//! // EMA tracks the smoothed value
//! let value = sensor.get();
//! assert!(value > 0.0);
//! ```

#![deny(missing_docs)]
#![deny(warnings)]

use prometheus::core::{Atomic, AtomicF64, Collector, Desc};
use prometheus::proto::MetricFamily;
use prometheus::{Counter, CounterVec, Opts};

use std::io::Write;

// ---------------------------------------------------------------------------
// AggregateCounter
// ---------------------------------------------------------------------------

/// A Prometheus [`CounterVec`] wrapper that automatically produces an extra
/// **unlabeled aggregate** metric (the sum of all label combinations) in
/// addition to the per-label counters.
///
/// This is useful when you want a single metric name to expose both a total
/// and a per-dimension breakdown without maintaining a separate `Counter`.
///
/// # Example
///
/// ```rust
/// use prometheus::Opts;
/// use prometheus::core::Collector;
/// use prometheus_extensions::AggregateCounter;
///
/// let counter = AggregateCounter::new(
///     Opts::new("rpc_calls_total", "Total RPC calls"),
///     &["service"],
/// ).unwrap();
///
/// counter.with_label_values(&["auth"]).inc_by(10.0);
/// counter.with_label_values(&["billing"]).inc_by(20.0);
///
/// let families = counter.collect();
/// let metrics = families[0].get_metric();
///
/// // First metric is the aggregate (no labels, value = 30)
/// assert_eq!(metrics[0].get_label().len(), 0);
/// assert_eq!(metrics[0].get_counter().value(), 30.0);
///
/// // Remaining metrics are the per-label counters
/// assert_eq!(metrics.len(), 3); // 1 aggregate + 2 labeled
/// ```
#[derive(Clone)]
pub struct AggregateCounter {
    labeled_counter: CounterVec,
    aggregate_desc: Desc,
}

impl AggregateCounter {
    /// Create a new `AggregateCounter`.
    ///
    /// `opts` defines the metric name and help string.
    /// `label_names` are the label dimensions for the per-label counters.
    pub fn new(opts: Opts, label_names: &[&str]) -> Result<Self, prometheus::Error> {
        let labeled_counter = CounterVec::new(opts.clone(), label_names)?;

        let aggregate_desc = Desc::new(
            opts.name.clone(),
            opts.help.clone(),
            vec![],
            opts.const_labels.clone(),
        )?;

        Ok(AggregateCounter {
            labeled_counter,
            aggregate_desc,
        })
    }

    /// Return a [`Counter`] for the given label values, creating it if it
    /// does not already exist.
    pub fn with_label_values(&self, vals: &[&str]) -> Counter {
        self.labeled_counter.with_label_values(vals)
    }
}

impl Collector for AggregateCounter {
    fn desc(&self) -> Vec<&Desc> {
        let mut descs = vec![&self.aggregate_desc];
        descs.extend(self.labeled_counter.desc());
        descs
    }

    fn collect(&self) -> Vec<MetricFamily> {
        let labeled_metrics = self.labeled_counter.collect();

        if labeled_metrics.is_empty() {
            return vec![];
        }

        let mut metric_family = labeled_metrics[0].clone();

        // Sum all labeled counter values into a single aggregate.
        let total: f64 = metric_family
            .get_metric()
            .iter()
            .map(|m| m.get_counter().value())
            .sum();

        let mut aggregate_metric = prometheus::proto::Metric::new();
        let mut counter = prometheus::proto::Counter::new();
        counter.set_value(total);
        aggregate_metric.set_counter(counter);

        // Prepend the aggregate metric before the labeled ones.
        let mut all_metrics = vec![aggregate_metric];
        all_metrics.extend(metric_family.take_metric().into_iter());
        metric_family.set_metric(all_metrics.into());

        vec![metric_family]
    }
}

// ---------------------------------------------------------------------------
// ScientificEncoder
// ---------------------------------------------------------------------------

/// A Prometheus text-format encoder that renders counter values in scientific
/// notation (`{:.10E}`) and appends a trailing comma after the last label.
///
/// This matches the output format produced by Kafka's JMX-to-Prometheus
/// exporters, making it useful for dashboards that expect that convention.
///
/// # Example
///
/// ```rust
/// use prometheus::Opts;
/// use prometheus::core::Collector;
/// use prometheus_extensions::{AggregateCounter, ScientificEncoder};
///
/// let counter = AggregateCounter::new(
///     Opts::new("bytes_in", "Bytes received"),
///     &["topic"],
/// ).unwrap();
/// counter.with_label_values(&["events"]).inc_by(4.6384186519e10);
///
/// let families = counter.collect();
/// let encoder = ScientificEncoder::new();
/// let mut buf = Vec::new();
/// encoder.encode(&families, &mut buf).unwrap();
///
/// let output = String::from_utf8(buf).unwrap();
/// assert!(output.contains("bytes_in{topic=\"events\",}"));
/// assert!(output.contains("E10"));
/// ```
#[derive(Default)]
pub struct ScientificEncoder;

impl ScientificEncoder {
    /// Create a new encoder.
    pub fn new() -> Self {
        Self
    }

    /// Encode `metric_families` into the Prometheus text exposition format
    /// with scientific-notation counter values.
    pub fn encode<W: Write>(
        &self,
        metric_families: &[MetricFamily],
        writer: &mut W,
    ) -> Result<(), std::io::Error> {
        for mf in metric_families {
            writeln!(writer, "# HELP {} {}", mf.name(), mf.help())?;
            writeln!(writer, "# TYPE {} counter", mf.name())?;

            for metric in mf.get_metric() {
                let value = metric.get_counter().value();

                if metric.get_label().is_empty() {
                    // Aggregate metric (no labels)
                    writeln!(writer, "{} {:.10E}", mf.name(), value)?;
                } else {
                    // Labeled metric — format with trailing comma
                    let mut label_str = String::new();
                    for label in metric.get_label() {
                        if !label_str.is_empty() {
                            label_str.push(',');
                        }
                        label_str.push_str(&format!(
                            "{}=\"{}\"",
                            label.name(),
                            label.value()
                        ));
                    }
                    label_str.push(',');

                    writeln!(writer, "{}{{{}}} {:.10E}", mf.name(), label_str, value)?;
                }
            }
        }
        Ok(())
    }
}

// ---------------------------------------------------------------------------
// Sensor (EMA)
// ---------------------------------------------------------------------------

/// A lock-free exponential moving average (EMA) gauge.
///
/// Useful for tracking smoothed rates or throughput where you want to dampen
/// spikes. The smoothing factor `alpha` controls responsiveness:
///
/// - `alpha` close to 1.0 → responds quickly to new values (noisy)
/// - `alpha` close to 0.0 → responds slowly (smooth)
///
/// A typical value for per-second throughput is `0.05`.
///
/// # Example
///
/// ```rust
/// use prometheus_extensions::Sensor;
///
/// let sensor = Sensor::new(0.05);
/// for _ in 0..100 {
///     sensor.measure(1000.0);
/// }
/// // After many identical measurements, EMA converges to the input value.
/// assert!((sensor.get() - 1000.0).abs() < 10.0);
/// ```
pub struct Sensor {
    ema: AtomicF64,
    alpha: f64,
}

impl Sensor {
    /// Create a new sensor with the given smoothing factor `alpha`.
    ///
    /// # Panics
    ///
    /// Panics if `alpha` is not in the range `(0.0, 1.0]`.
    pub fn new(alpha: f64) -> Self {
        assert!(
            alpha > 0.0 && alpha <= 1.0,
            "alpha must be in (0.0, 1.0], got {alpha}"
        );
        Self {
            ema: AtomicF64::new(0.0),
            alpha,
        }
    }

    /// Feed a new measurement into the EMA.
    pub fn measure(&self, value: f64) {
        let current = self.ema.get();
        let new_value = self.alpha * value + (1.0 - self.alpha) * current;
        self.ema.set(new_value);
    }

    /// Read the current smoothed value.
    pub fn get(&self) -> f64 {
        self.ema.get()
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;
    use prometheus::core::Collector;

    #[test]
    fn aggregate_counter_produces_both_metrics() {
        let opts = Opts::new("test_counter", "Test counter for aggregation");
        let counter = AggregateCounter::new(opts, &["label"]).unwrap();

        counter.with_label_values(&["value1"]).inc_by(10.0);
        counter.with_label_values(&["value2"]).inc_by(20.0);
        counter.with_label_values(&["value3"]).inc_by(30.0);

        let metric_families = counter.collect();
        assert_eq!(metric_families.len(), 1);

        let mf = &metric_families[0];
        assert_eq!(mf.name(), "test_counter");
        assert_eq!(mf.get_metric().len(), 4); // 1 aggregate + 3 labeled

        // First metric is the aggregate (no labels)
        let aggregate = &mf.get_metric()[0];
        assert_eq!(aggregate.get_label().len(), 0);
        assert_eq!(aggregate.get_counter().value(), 60.0);

        // Remaining are labeled
        for i in 1..4 {
            let metric = &mf.get_metric()[i];
            assert_eq!(metric.get_label().len(), 1);
            let label_value = metric.get_label()[0].value();
            let counter_value = metric.get_counter().value();
            match label_value {
                "value1" => assert_eq!(counter_value, 10.0),
                "value2" => assert_eq!(counter_value, 20.0),
                "value3" => assert_eq!(counter_value, 30.0),
                _ => panic!("Unexpected label value: {label_value}"),
            }
        }
    }

    #[test]
    fn aggregate_counter_empty_collect() {
        let opts = Opts::new("empty_counter", "No observations yet");
        let counter = AggregateCounter::new(opts, &["x"]).unwrap();
        let families = counter.collect();
        // CounterVec with no observations still returns a family with 0 metrics
        // or an empty vec depending on prometheus version — either is fine.
        if !families.is_empty() {
            let total: f64 = families[0]
                .get_metric()
                .iter()
                .map(|m| m.get_counter().value())
                .sum();
            assert_eq!(total, 0.0);
        }
    }

    #[test]
    fn scientific_encoder_format() {
        let opts = Opts::new("test_kafka_metrics", "Test counter for Kafka format");
        let counter = AggregateCounter::new(opts, &["topic"]).unwrap();

        counter
            .with_label_values(&["hits"])
            .inc_by(4.6384186519e10);
        counter.with_label_values(&["logs"]).inc_by(1.2345e9);

        let families = counter.collect();
        let encoder = ScientificEncoder::new();
        let mut buf = Vec::new();
        encoder.encode(&families, &mut buf).unwrap();
        let output = String::from_utf8(buf).unwrap();

        assert!(output.contains("# HELP test_kafka_metrics Test counter for Kafka format"));
        assert!(output.contains("# TYPE test_kafka_metrics counter"));
        assert!(output.contains("test_kafka_metrics 4.7618686519E10"));
        assert!(output.contains("test_kafka_metrics{topic=\"hits\",} 4.6384186519E10"));
        assert!(output.contains("test_kafka_metrics{topic=\"logs\",} 1.2345000000E9"));
    }

    #[test]
    fn sensor_converges() {
        let sensor = Sensor::new(0.05);
        for _ in 0..1000 {
            sensor.measure(500.0);
        }
        assert!(
            (sensor.get() - 500.0).abs() < 1.0,
            "EMA should converge to 500.0, got {}",
            sensor.get()
        );
    }

    #[test]
    fn sensor_starts_at_zero() {
        let sensor = Sensor::new(0.1);
        assert_eq!(sensor.get(), 0.0);
    }

    #[test]
    #[should_panic(expected = "alpha must be in (0.0, 1.0]")]
    fn sensor_rejects_zero_alpha() {
        Sensor::new(0.0);
    }

    #[test]
    #[should_panic(expected = "alpha must be in (0.0, 1.0]")]
    fn sensor_rejects_negative_alpha() {
        Sensor::new(-0.5);
    }

    #[test]
    fn sensor_alpha_one_tracks_instantly() {
        let sensor = Sensor::new(1.0);
        sensor.measure(42.0);
        assert_eq!(sensor.get(), 42.0);
        sensor.measure(99.0);
        assert_eq!(sensor.get(), 99.0);
    }
}