pulling_gauge.rs

use std::{collections::HashMap, fmt, sync::Arc};

use crate::{
    core::Collector,
    proto::{Gauge, Metric, MetricFamily, MetricType},
};
use protobuf::RepeatedField;

/// A prometheus gauge that reads the value from a provided function on every
/// collect run. Used to expose metrics that are better to be polled as opposed
/// to continuously set.
#[derive(Clone)]
pub struct PullingGauge {
    desc: crate::core::Desc,
    value: Arc<Box<dyn Fn() -> f64 + Send + Sync>>,
}

impl fmt::Debug for PullingGauge {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("PullingGauge")
            .field("desc", &self.desc)
            .field("value", &"<opaque>")
            .finish()
    }
}

impl PullingGauge {
    /// Create a new PullingGauge.
    pub fn new<S1: Into<String>, S2: Into<String>>(
        fq_name: S1,
        help: S2,
        value: Box<dyn Fn() -> f64 + Send + Sync>,
    ) -> crate::Result<Self> {
        Ok(PullingGauge {
            value: Arc::new(value),
            desc: crate::core::Desc::new(fq_name.into(), help.into(), Vec::new(), HashMap::new())?,
        })
    }

    /// Create a new PullingGauge and register it with the global registry.
    pub fn create_and_register<S1: Into<String>, S2: Into<String>>(
        name: S1,
        help: S2,
        value: Box<dyn Fn() -> f64 + Send + Sync>,
    ) -> crate::Result<Self> {
        let metric = PullingGauge::new(name, help, value)?;
        crate::register(Box::new(metric.clone())).map(|_| metric)
    }

    fn metric(&self) -> Metric {
        let mut gauge = Gauge::default();
        let getter = &self.value;
        gauge.set_value(getter());

        let mut metric = Metric::default();
        metric.set_gauge(gauge);

        metric
    }
}

impl Collector for PullingGauge {
    fn desc(&self) -> Vec<&crate::core::Desc> {
        vec![&self.desc]
    }

    fn collect(&self) -> Vec<crate::proto::MetricFamily> {
        let mut m = MetricFamily::default();
        m.set_name(self.desc.fq_name.clone());
        m.set_help(self.desc.help.clone());
        m.set_field_type(MetricType::GAUGE);
        m.set_metric(RepeatedField::from_vec(vec![self.metric()]));
        vec![m]
    }
}