trait-net 0.7.1

use crate::metrics::{AsStatusLabel, Observer};
use prometheus::{
    core::{Collector, Desc},
    proto::MetricFamily,
    IntGaugeVec, Opts,
};
use std::{iter::once_with, ops::Deref};

#[derive(Clone)]
pub struct StatusCount(IntGaugeVec);

impl StatusCount {
    pub fn new(opts: Opts, label_names: &[&str]) -> prometheus::Result<Self> {
        Self::new_with(opts, label_names, "status")
    }

    pub fn new_with(
        opts: Opts,
        label_names: &[&str],
        status_label_name: &str,
    ) -> prometheus::Result<Self> {
        let labels: Vec<_> = label_names
            .into_iter()
            .map(|s| *s)
            .chain(once_with(|| status_label_name))
            .collect();
        Ok(Self(IntGaugeVec::new(opts, &labels)?))
    }

    pub fn make_observer(&self, labels: &[&str]) -> StatusCountObserver {
        StatusCountObserver {
            gauge: self.0.clone(),
            labels: labels.into_iter().map(|s| s.to_string()).collect(),
        }
    }
}

impl Deref for StatusCount {
    type Target = IntGaugeVec;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl Collector for StatusCount {
    fn desc(&self) -> Vec<&Desc> {
        self.0.desc()
    }

    fn collect(&self) -> Vec<MetricFamily> {
        self.0.collect()
    }
}

pub struct StatusCountObserver {
    gauge: IntGaugeVec,
    labels: Vec<String>,
}

impl<Out: ?Sized + AsStatusLabel> Observer<Out> for StatusCountObserver {
    fn on_poll_ready(&mut self, output: &Out) {
        let status_label_name = output.as_status_label();
        let labels: Vec<_> = self
            .labels
            .iter()
            .map(|s| s.as_str())
            .chain(once_with(|| status_label_name.as_str()))
            .collect();
        self.gauge.with_label_values(&labels).inc();
    }
}