prometheus-utils 0.6.3

use crate::guards::DeferredAddWithLabels;
use prometheus::{
    register_histogram_vec, register_int_counter_vec, register_int_gauge_vec, HistogramTimer,
    HistogramVec, IntCounterVec, IntGaugeVec,
};
use std::marker::PhantomData;

/// A sequence of values for Prometheus labels
pub type LabelValues<'a> = Vec<&'a str>;

/// The `Labels` trait applies to values intended to generate Prometheus labels for a metric.
///
/// A metric in Prometheus can include any number of labels. Each label has a fixed name,
/// and when events are emitted for the metric, those events must include values for each
/// of the labels. Using labels makes it possible to easily see a metric in aggregate (i.e.,
/// to see totals regardless of label values), or to query specific kinds of events by
/// filtering label values.
///
/// Thus, for example, rather than having a separate metric for each kind of error that
/// arises, we can produce one metric with an "err" label, whose value will reflect which
/// error has occurred. That simplifies the top-level list of metrics, and makes it easier
/// to build queries to aggregate specific kinds of error events.
///
/// This trait adds some extra guard rails on top of the prometheus-rs crate, so that when
/// we emit a labeled metric we can use a custom type to represent the labels, rather than
/// working directly with slices of string slices. When defining a labeled metric, you should
/// also define a new type representing its labels that implements the `Labels` trait.
/// Then, when emitting events for the metric, labels are passed in using this custom type,
/// which rules out several kinds of bugs (like missing or incorrectly ordered label values).
///
/// You can define labeled metrics using types like [`IntCounterWithLabels`], which are
/// parameterized by a type that implements `Labels`.
///
/// [`IntCounterWithLabels`]: struct.IntCounterWithLabels.html
pub trait Labels {
    /// The names of the labels that will be defined for the corresponding metric.
    fn label_names() -> Vec<&'static str>;

    /// Labels values to seed the metric with initially.
    ///
    /// Since Prometheus doesn't know the possible values a label will take on, when we set
    /// up a labeled metric by default no values will appear until events are emitted. But
    /// for discoverability, it's helpful to initialize a labeled metric with some possible
    /// label values (at count 0) even if no events for the metric have occurred.
    ///
    /// The label values provided by this function are used to pre-populate the metric at
    /// count 0. **The values do _not_ need to be exhaustive**; it's fine for events to emit
    /// label values that are not included here.
    fn possible_label_values() -> Vec<LabelValues<'static>>;

    /// The actual label values to provide when emitting an event to Prometheus.
    ///
    /// The sequence of values should correspond to the names provided in `label_names`,
    /// in order.
    fn label_values(&self) -> LabelValues;
}

/// A Prometheus integer counter metric, with labels described by the type `L`.
///
/// The type `L` must implement the [`Labels`] trait; see the documentation for that trait
/// for an overview of Prometheus metric labels.
///
/// [`Labels`]: trait.Labels.html
pub struct IntCounterWithLabels<L: Labels> {
    metric: IntCounterVec,
    _labels: PhantomData<L>,
}

impl<L: Labels> IntCounterWithLabels<L> {
    /// Construct and immediately register a new `IntCounterWithLabels` instance.
    pub fn register_new(name: &str, help: &str) -> IntCounterWithLabels<L> {
        let metric = register_int_counter_vec!(name, help, &L::label_names()).unwrap();

        for vals in L::possible_label_values() {
            metric.with_label_values(&vals).inc_by(0);
        }

        Self {
            metric,
            _labels: PhantomData,
        }
    }

    /// Get the value of the counter with the provided `labels`.
    pub fn get(&self, labels: &L) -> u64 {
        self.metric.with_label_values(&labels.label_values()).get()
    }

    /// Increment the metric by `1`, using the provided `labels` for the event.
    pub fn inc(&self, labels: &L) {
        self.metric.with_label_values(&labels.label_values()).inc();
    }

    /// Increment the metric by `v`, using the provided `labels` for the event.
    pub fn add(&self, v: u64, labels: &L) {
        self.metric
            .with_label_values(&labels.label_values())
            .inc_by(v);
    }

    /// Creates a guard value that will increment the metric by `1`, using the provided `labels`,
    /// once dropped.
    ///
    /// Prior to dropping, the labels can be altered using [`DeferredAddWithLabels::with_labels`].
    #[must_use]
    pub fn deferred_inc<'a>(&'a self, labels: L) -> DeferredAddWithLabels<'a, L> {
        DeferredAddWithLabels::new(self, 1, labels)
    }

    /// Creates a guard value that will increment the metric by `v`, using the provided `labels`,
    /// once dropped.
    ///
    /// Prior to dropping, the labels can be altered using [`DeferredAddWithLabels::with_labels`].
    #[must_use]
    pub fn deferred_add<'a>(&'a self, v: u64, labels: L) -> DeferredAddWithLabels<'a, L> {
        DeferredAddWithLabels::new(self, v, labels)
    }
}

/// A Prometheus integer gauge metric, with labels described by the type `L`.
///
/// The type `L` must implement the [`Labels`] trait; see the documentation for that trait
/// for an overview of Prometheus metric labels.
///
/// [`Labels`]: trait.Labels.html
pub struct IntGaugeWithLabels<L: Labels> {
    metric: IntGaugeVec,
    _labels: PhantomData<L>,
}

impl<L: Labels> IntGaugeWithLabels<L> {
    /// Construct and immediately register a new `IntGaugeWithLabels` instance.
    pub fn register_new(name: &str, help: &str) -> IntGaugeWithLabels<L> {
        let metric = register_int_gauge_vec!(name, help, &L::label_names()).unwrap();

        // Note: for gauges, unlike counters, we don't need to -- and should not! -- prepopulate
        // the metric with the possible labels. Unlike counters, which are only updated when an
        // event occurs, gauges _always_ have a value. Moreover, we cannot make assumptions about
        // an initial gauge value (unlike an initial 0 value for counters).

        Self {
            metric,
            _labels: PhantomData,
        }
    }

    /// Get the value of the gauge with the provided `labels`.
    pub fn get(&self, labels: &L) -> i64 {
        self.metric.with_label_values(&labels.label_values()).get()
    }

    /// Set the value of the gauge with the provided `labels`.
    pub fn set(&self, labels: &L, value: i64) {
        self.metric
            .with_label_values(&labels.label_values())
            .set(value);
    }

    /// Add `value` to the gauge with the provided `labels`.
    pub fn add(&self, labels: &L, value: i64) {
        self.metric
            .with_label_values(&labels.label_values())
            .add(value);
    }

    /// Subtract `value` from the gauge with the provided `labels`.
    pub fn sub(&self, labels: &L, value: i64) {
        self.metric
            .with_label_values(&labels.label_values())
            .sub(value);
    }

    /// Increment the gauge by `1`, using the provided `labels` for the event.
    pub fn inc(&self, labels: &L) {
        self.metric.with_label_values(&labels.label_values()).inc();
    }

    /// Decrement the gauge by `1`, using the provided `labels` for the event.
    pub fn dec(&self, labels: &L) {
        self.metric.with_label_values(&labels.label_values()).dec();
    }
}

/// A Prometheus histogram metric, with labels described by the type `L`.
///
/// The type `L` must implement the [`Labels`] trait; see the documentation for that trait
/// for an overview of Prometheus metric labels.
///
/// [`Labels`]: trait.Labels.html
pub struct HistogramWithLabels<L: Labels> {
    metric: HistogramVec,
    _labels: PhantomData<L>,
}

impl<L: Labels> HistogramWithLabels<L> {
    /// Construct and immediately register a new `HistogramWithLabels` instance.
    pub fn register_new(name: &str, help: &str) -> Self {
        let metric = register_histogram_vec!(name, help, &L::label_names()).unwrap();

        // Note: for histograms, like gauges, we don't need to -- and should not! -- prepopulate
        // the metric with the possible labels.

        Self {
            metric,
            _labels: PhantomData,
        }
    }

    /// Construct and immediately register a new `HistogramWithLabels` instance.
    ///
    /// This will use the provided `buckets` when registering the underlying [`HistogramVec`].
    pub fn register_new_with_buckets(name: &str, help: &str, buckets: Vec<f64>) -> Self {
        let metric = register_histogram_vec!(name, help, &L::label_names(), buckets).unwrap();

        // Note: for histograms, like gauges, we don't need to -- and should not! -- prepopulate
        // the metric with the possible labels.

        Self {
            metric,
            _labels: PhantomData,
        }
    }

    /// Add a single observation to the histogram with the provided `labels`.
    pub fn observe(&self, labels: &L, value: f64) {
        self.metric
            .with_label_values(&labels.label_values())
            .observe(value);
    }

    /// Return a [`HistogramTimer`] to track a duration, using the provided `labels`.
    pub fn start_timer(&self, labels: &L) -> HistogramTimer {
        self.metric
            .with_label_values(&labels.label_values())
            .start_timer()
    }

    /// Observe execution time of a closure, in seconds.
    pub fn observe_closure_duration<F, T>(&self, labels: &L, f: F) -> T
    where
        F: FnOnce() -> T,
    {
        self.metric
            .with_label_values(&labels.label_values())
            .observe_closure_duration(f)
    }

    /// Return accumulated sum of all samples, using the provided `labels`.
    pub fn get_sample_sum(&self, labels: &L) -> f64 {
        self.metric
            .with_label_values(&labels.label_values())
            .get_sample_sum()
    }

    /// Return count of all samples, using the provided `labels`.
    pub fn get_sample_count(&self, labels: &L) -> u64 {
        self.metric
            .with_label_values(&labels.label_values())
            .get_sample_count()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use lazy_static::lazy_static;
    use test_labels::*;

    #[allow(dead_code)]
    mod test_labels {
        use crate::label_enum;
        label_enum! {
            pub enum Animal {
                Bird,
                Cat,
                Dog,
            }
        }
        label_enum! {
            pub enum Size {
                Small,
                Large,
            }
        }
    }

    struct TestLabels {
        animal: Animal,
        size: Size,
    }

    impl Labels for TestLabels {
        fn label_names() -> Vec<&'static str> {
            vec!["animal", "size"]
        }
        fn possible_label_values() -> Vec<LabelValues<'static>> {
            Default::default() // elide this for brevity.
        }
        fn label_values(&self) -> LabelValues {
            let Self { animal, size } = self;
            vec![animal.as_str(), size.as_str()]
        }
    }

    /// Show that an [`IntCounterWithLabels`] can have its value accessed.
    #[test]
    fn int_counter_with_labels_get_works() {
        lazy_static! {
            static ref COUNTER: IntCounterWithLabels<TestLabels> =
                IntCounterWithLabels::register_new(
                    "test_label_counter",
                    "a labeled counter for tests"
                );
        }

        // Increment the counter for some various combinations of label values.
        COUNTER.inc(&TestLabels {
            animal: Animal::Bird,
            size: Size::Large,
        });
        COUNTER.inc(&TestLabels {
            animal: Animal::Bird,
            size: Size::Small,
        });
        COUNTER.inc(&TestLabels {
            animal: Animal::Cat,
            size: Size::Large,
        });

        assert_eq!(
            COUNTER.get(&TestLabels {
                animal: Animal::Bird,
                size: Size::Large,
            }),
            1
        );
        assert_eq!(
            COUNTER.get(&TestLabels {
                animal: Animal::Cat,
                size: Size::Small,
            }),
            0
        );
    }

    /// Show that an [`IntGaugeWithLabels`] can have its value accessed.
    #[test]
    fn int_gauge_with_labels_get_works() {
        lazy_static! {
            static ref GAUGE: IntGaugeWithLabels<TestLabels> =
                IntGaugeWithLabels::register_new("test_label_gauge", "a labeled gauge for tests");
        }

        // Increment the gauge for some various combinations of label values.
        GAUGE.inc(&TestLabels {
            animal: Animal::Bird,
            size: Size::Large,
        });
        GAUGE.inc(&TestLabels {
            animal: Animal::Bird,
            size: Size::Small,
        });
        GAUGE.inc(&TestLabels {
            animal: Animal::Cat,
            size: Size::Large,
        });

        assert_eq!(
            GAUGE.get(&TestLabels {
                animal: Animal::Bird,
                size: Size::Large,
            }),
            1
        );
        assert_eq!(
            GAUGE.get(&TestLabels {
                animal: Animal::Cat,
                size: Size::Small,
            }),
            0
        );
    }
}