use prometheus::{
core::{Collector, Desc, Describer, Number, Opts},
proto,
};
use std::{cmp::Ordering, marker::PhantomData};
pub type SourcedCounter<S> = SourcedMetric<Counter, S>;
pub type SourcedGauge<S> = SourcedMetric<Gauge, S>;
#[derive(Copy, Clone)]
pub enum Counter {}
#[derive(Copy, Clone)]
pub enum Gauge {}
#[derive(Debug, Clone)]
pub struct SourcedMetric<T, S> {
source: S,
desc: Desc,
_type: PhantomData<T>,
}
pub trait MetricSource: Sync + Send + Clone {
type N: Number;
fn collect(&self, set: impl FnMut(&[&str], Self::N));
}
impl<T: SourcedType, S: MetricSource> SourcedMetric<T, S> {
pub fn new(opts: &Opts, source: S) -> prometheus::Result<Self> {
let desc = opts.describe()?;
Ok(Self { source, desc, _type: PhantomData })
}
}
impl<T: SourcedType, S: MetricSource> Collector for SourcedMetric<T, S> {
fn desc(&self) -> Vec<&Desc> {
vec![&self.desc]
}
fn collect(&self) -> Vec<proto::MetricFamily> {
let mut counters = Vec::new();
self.source.collect(|label_values, value| {
let mut m = proto::Metric::default();
match T::proto() {
proto::MetricType::COUNTER => {
let mut c = proto::Counter::default();
c.set_value(value.into_f64());
m.set_counter(c);
},
proto::MetricType::GAUGE => {
let mut g = proto::Gauge::default();
g.set_value(value.into_f64());
m.set_gauge(g);
},
t => {
log::error!("Unsupported sourced metric type: {:?}", t);
},
}
debug_assert_eq!(self.desc.variable_labels.len(), label_values.len());
match self.desc.variable_labels.len().cmp(&label_values.len()) {
Ordering::Greater => {
log::warn!("Missing label values for sourced metric {}", self.desc.fq_name)
},
Ordering::Less => {
log::warn!("Too many label values for sourced metric {}", self.desc.fq_name)
},
Ordering::Equal => {},
}
m.set_label(
self.desc
.variable_labels
.iter()
.zip(label_values)
.map(|(l_name, l_value)| {
let mut l = proto::LabelPair::default();
l.set_name(l_name.to_string());
l.set_value(l_value.to_string());
l
})
.chain(self.desc.const_label_pairs.iter().cloned())
.collect::<Vec<_>>(),
);
counters.push(m);
});
let mut m = proto::MetricFamily::default();
m.set_name(self.desc.fq_name.clone());
m.set_help(self.desc.help.clone());
m.set_field_type(T::proto());
m.set_metric(counters);
vec![m]
}
}
pub trait SourcedType: private::Sealed + Sync + Send {
#[doc(hidden)]
fn proto() -> proto::MetricType;
}
impl SourcedType for Counter {
fn proto() -> proto::MetricType {
proto::MetricType::COUNTER
}
}
impl SourcedType for Gauge {
fn proto() -> proto::MetricType {
proto::MetricType::GAUGE
}
}
mod private {
pub trait Sealed {}
impl Sealed for super::Counter {}
impl Sealed for super::Gauge {}
}