use prometheus_client::{
encoding::{EncodeLabelSet, EncodeLabelValue},
metrics::{counter::Counter as DefaultCounter, family::Family, gauge::Gauge},
};
use std::sync::atomic::Ordering;
#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, EncodeLabelSet)]
pub struct Label {
status: Status,
}
#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, EncodeLabelValue)]
pub enum Status {
Success,
Failure,
Invalid,
Dropped,
Timeout,
}
pub type Counter = Family<Label, DefaultCounter>;
pub trait CounterExt {
fn guard(&self, status: Status) -> CounterGuard;
fn inc(&self, status: Status);
fn inc_by(&self, status: Status, n: u64);
}
impl CounterExt for Counter {
fn guard(&self, status: Status) -> CounterGuard {
CounterGuard {
metric: self.clone(),
status,
}
}
fn inc(&self, status: Status) {
self.get_or_create(&Label { status }).inc();
}
fn inc_by(&self, status: Status, n: u64) {
self.get_or_create(&Label { status }).inc_by(n);
}
}
pub struct CounterGuard {
metric: Counter,
status: Status,
}
impl CounterGuard {
pub const fn set(&mut self, status: Status) {
self.status = status;
}
}
impl Drop for CounterGuard {
fn drop(&mut self) {
self.metric.inc(self.status);
}
}
pub trait GaugeExt {
fn try_set<T: TryInto<i64>>(&self, val: T) -> Result<i64, T::Error>;
fn try_set_max<T: TryInto<i64> + Copy>(&self, val: T) -> Result<i64, T::Error>;
}
impl GaugeExt for Gauge {
fn try_set<T: TryInto<i64>>(&self, val: T) -> Result<i64, T::Error> {
let val = val.try_into()?;
let out = self.set(val);
Ok(out)
}
fn try_set_max<T: TryInto<i64> + Copy>(&self, val: T) -> Result<i64, T::Error> {
let val = val.try_into()?;
Ok(self.inner().fetch_max(val, Ordering::Relaxed))
}
}