use std::borrow::Borrow;
use std::collections::BTreeMap;
use std::marker::PhantomData;
use std::ops::Deref;
use prometheus::core::{
Atomic, GenericCounter, GenericCounterVec, GenericGauge, GenericGaugeVec, MetricVec,
MetricVecBuilder,
};
use prometheus::{Histogram, HistogramVec};
pub trait PromLabelsExt<'a> {
fn get_from_metric_vec<P: MetricVecBuilder>(
&self,
vec: &MetricVec<P>,
) -> <P as MetricVecBuilder>::M;
fn remove_from_metric_vec<P: MetricVecBuilder>(
&self,
vec: &MetricVec<P>,
) -> Result<(), prometheus::Error>;
}
impl<'a> PromLabelsExt<'a> for &'a [&'a str] {
fn get_from_metric_vec<P: MetricVecBuilder>(
&self,
vec: &MetricVec<P>,
) -> <P as MetricVecBuilder>::M {
vec.get_metric_with_label_values(self)
.expect("retrieving a metric by label values")
}
fn remove_from_metric_vec<P: MetricVecBuilder>(
&self,
vec: &MetricVec<P>,
) -> Result<(), prometheus::Error> {
vec.remove_label_values(self)
}
}
impl PromLabelsExt<'static> for Vec<String> {
fn get_from_metric_vec<P: MetricVecBuilder>(
&self,
vec: &MetricVec<P>,
) -> <P as MetricVecBuilder>::M {
let labels: Vec<&str> = self.iter().map(String::as_str).collect();
vec.get_metric_with_label_values(labels.as_slice())
.expect("retrieving a metric by label values")
}
fn remove_from_metric_vec<P: MetricVecBuilder>(
&self,
vec: &MetricVec<P>,
) -> Result<(), prometheus::Error> {
let labels: Vec<&str> = self.iter().map(String::as_str).collect();
vec.remove_label_values(labels.as_slice())
}
}
impl<'a> PromLabelsExt<'a> for Vec<&'a str> {
fn get_from_metric_vec<P: MetricVecBuilder>(
&self,
vec: &MetricVec<P>,
) -> <P as MetricVecBuilder>::M {
vec.get_metric_with_label_values(self.as_slice())
.expect("retrieving a metric by label values")
}
fn remove_from_metric_vec<P: MetricVecBuilder>(
&self,
vec: &MetricVec<P>,
) -> Result<(), prometheus::Error> {
vec.remove_label_values(self.as_slice())
}
}
impl PromLabelsExt<'static> for BTreeMap<String, String> {
fn get_from_metric_vec<P: MetricVecBuilder>(
&self,
vec: &MetricVec<P>,
) -> <P as MetricVecBuilder>::M {
let labels = self.iter().map(|(k, v)| (k.as_str(), v.as_str())).collect();
vec.get_metric_with(&labels)
.expect("retrieving a metric by label values")
}
fn remove_from_metric_vec<P: MetricVecBuilder>(
&self,
vec: &MetricVec<P>,
) -> Result<(), prometheus::Error> {
let labels = self.iter().map(|(k, v)| (k.as_str(), v.as_str())).collect();
vec.remove(&labels)
}
}
impl<'a> PromLabelsExt<'a> for BTreeMap<&'a str, &'a str> {
fn get_from_metric_vec<P: MetricVecBuilder>(
&self,
vec: &MetricVec<P>,
) -> <P as MetricVecBuilder>::M {
let labels = self.iter().map(|(k, v)| (*k, *v)).collect();
vec.get_metric_with(&labels)
.expect("retrieving a metric by label values")
}
fn remove_from_metric_vec<P: MetricVecBuilder>(
&self,
vec: &MetricVec<P>,
) -> Result<(), prometheus::Error> {
let labels = self.iter().map(|(k, v)| (*k, *v)).collect();
vec.remove(&labels)
}
}
#[derive(Debug)]
pub struct DeleteOnDropHistogram<'a, L>
where
L: PromLabelsExt<'a>,
{
inner: Histogram,
labels: L,
vec: HistogramVec,
_phantom: &'a PhantomData<()>,
}
impl<'a, L> DeleteOnDropHistogram<'a, L>
where
L: PromLabelsExt<'a>,
{
fn from_metric_vector(vec: HistogramVec, labels: L) -> Self {
let inner = labels.get_from_metric_vec(&vec);
Self {
inner,
labels,
vec,
_phantom: &PhantomData,
}
}
}
impl<'a, L> Deref for DeleteOnDropHistogram<'a, L>
where
L: PromLabelsExt<'a>,
{
type Target = Histogram;
fn deref(&self) -> &Self::Target {
&self.inner
}
}
impl<'a, L> Borrow<Histogram> for DeleteOnDropHistogram<'a, L>
where
L: PromLabelsExt<'a>,
{
fn borrow(&self) -> &Histogram {
&self.inner
}
}
impl<'a, L> Drop for DeleteOnDropHistogram<'a, L>
where
L: PromLabelsExt<'a>,
{
fn drop(&mut self) {
if self.labels.remove_from_metric_vec(&self.vec).is_err() {
}
}
}
#[derive(Debug)]
pub struct DeleteOnDropCounter<'a, P, L>
where
P: Atomic,
L: PromLabelsExt<'a>,
{
inner: GenericCounter<P>,
labels: L,
vec: GenericCounterVec<P>,
_phantom: &'a PhantomData<()>,
}
impl<'a, P, L> DeleteOnDropCounter<'a, P, L>
where
P: Atomic,
L: PromLabelsExt<'a>,
{
fn from_metric_vector(vec: GenericCounterVec<P>, labels: L) -> Self {
let inner = labels.get_from_metric_vec(&vec);
Self {
inner,
labels,
vec,
_phantom: &PhantomData,
}
}
}
impl<'a, P, L> Deref for DeleteOnDropCounter<'a, P, L>
where
P: Atomic,
L: PromLabelsExt<'a>,
{
type Target = GenericCounter<P>;
fn deref(&self) -> &GenericCounter<P> {
&self.inner
}
}
impl<'a, P, L> Borrow<GenericCounter<P>> for DeleteOnDropCounter<'a, P, L>
where
P: Atomic,
L: PromLabelsExt<'a>,
{
fn borrow(&self) -> &GenericCounter<P> {
&self.inner
}
}
impl<'a, P, L> Drop for DeleteOnDropCounter<'a, P, L>
where
P: Atomic,
L: PromLabelsExt<'a>,
{
fn drop(&mut self) {
if self.labels.remove_from_metric_vec(&self.vec).is_err() {
}
}
}
pub trait CounterVecExt {
type CounterType: Atomic;
fn get_delete_on_drop_counter<'a, L: PromLabelsExt<'a>>(
&self,
labels: L,
) -> DeleteOnDropCounter<'a, Self::CounterType, L>;
}
impl<P: Atomic> CounterVecExt for GenericCounterVec<P> {
type CounterType = P;
fn get_delete_on_drop_counter<'a, L: PromLabelsExt<'a>>(
&self,
labels: L,
) -> DeleteOnDropCounter<'a, Self::CounterType, L> {
DeleteOnDropCounter::from_metric_vector(self.clone(), labels)
}
}
pub trait HistogramVecExt {
fn get_delete_on_drop_histogram<'a, L: PromLabelsExt<'a>>(
&self,
labels: L,
) -> DeleteOnDropHistogram<'a, L>;
}
impl HistogramVecExt for HistogramVec {
fn get_delete_on_drop_histogram<'a, L: PromLabelsExt<'a>>(
&self,
labels: L,
) -> DeleteOnDropHistogram<'a, L> {
DeleteOnDropHistogram::from_metric_vector(self.clone(), labels)
}
}
#[derive(Debug)]
pub struct DeleteOnDropGauge<'a, P, L>
where
P: Atomic,
L: PromLabelsExt<'a>,
{
inner: GenericGauge<P>,
labels: L,
vec: GenericGaugeVec<P>,
_phantom: &'a PhantomData<()>,
}
impl<'a, P, L> DeleteOnDropGauge<'a, P, L>
where
P: Atomic,
L: PromLabelsExt<'a>,
{
fn from_metric_vector(vec: GenericGaugeVec<P>, labels: L) -> Self {
let inner = labels.get_from_metric_vec(&vec);
Self {
inner,
labels,
vec,
_phantom: &PhantomData,
}
}
}
impl<'a, P, L> Deref for DeleteOnDropGauge<'a, P, L>
where
P: Atomic,
L: PromLabelsExt<'a>,
{
type Target = GenericGauge<P>;
fn deref(&self) -> &GenericGauge<P> {
&self.inner
}
}
impl<'a, P, L> Borrow<GenericGauge<P>> for DeleteOnDropGauge<'a, P, L>
where
P: Atomic,
L: PromLabelsExt<'a>,
{
fn borrow(&self) -> &GenericGauge<P> {
&self.inner
}
}
impl<'a, P, L> Drop for DeleteOnDropGauge<'a, P, L>
where
P: Atomic,
L: PromLabelsExt<'a>,
{
fn drop(&mut self) {
if self.labels.remove_from_metric_vec(&self.vec).is_err() {
}
}
}
pub trait GaugeVecExt {
type GaugeType: Atomic;
fn get_delete_on_drop_gauge<'a, L: PromLabelsExt<'a>>(
&self,
labels: L,
) -> DeleteOnDropGauge<'a, Self::GaugeType, L>;
}
impl<P: Atomic> GaugeVecExt for GenericGaugeVec<P> {
type GaugeType = P;
fn get_delete_on_drop_gauge<'a, L: PromLabelsExt<'a>>(
&self,
labels: L,
) -> DeleteOnDropGauge<'a, Self::GaugeType, L> {
DeleteOnDropGauge::from_metric_vector(self.clone(), labels)
}
}
#[cfg(test)]
mod test {
use prometheus::core::{AtomicI64, AtomicU64};
use prometheus::IntGaugeVec;
use crate::metric;
use crate::metrics::{IntCounterVec, MetricsRegistry};
use super::*;
#[crate::test]
fn dropping_counters() {
let reg = MetricsRegistry::new();
let vec: IntCounterVec = reg.register(metric!(
name: "test_metric",
help: "a test metric",
var_labels: ["dimension"]));
let dims: &[&str] = &["one"];
let metric_1 = vec.get_delete_on_drop_counter(dims);
metric_1.inc();
let metrics = reg.gather();
assert_eq!(metrics.len(), 1);
let reported_vec = &metrics[0];
assert_eq!(reported_vec.get_name(), "test_metric");
let dims = reported_vec.get_metric();
assert_eq!(dims.len(), 1);
assert_eq!(dims[0].get_label()[0].get_value(), "one");
drop(metric_1);
let metrics = reg.gather();
assert_eq!(metrics.len(), 0);
let string_labels: Vec<String> = ["owned"].iter().map(ToString::to_string).collect();
struct Ownership {
counter: DeleteOnDropCounter<'static, AtomicU64, Vec<String>>,
}
let metric_owned = Ownership {
counter: vec.get_delete_on_drop_counter(string_labels),
};
metric_owned.counter.inc();
let metrics = reg.gather();
assert_eq!(metrics.len(), 1);
let reported_vec = &metrics[0];
assert_eq!(reported_vec.get_name(), "test_metric");
let dims = reported_vec.get_metric();
assert_eq!(dims.len(), 1);
assert_eq!(dims[0].get_label()[0].get_value(), "owned");
drop(metric_owned);
let metrics = reg.gather();
assert_eq!(metrics.len(), 0);
}
#[crate::test]
fn dropping_gauges() {
let reg = MetricsRegistry::new();
let vec: IntGaugeVec = reg.register(metric!(
name: "test_metric",
help: "a test metric",
var_labels: ["dimension"]));
let dims: &[&str] = &["one"];
let metric_1 = vec.get_delete_on_drop_gauge(dims);
metric_1.set(666);
let metrics = reg.gather();
assert_eq!(metrics.len(), 1);
let reported_vec = &metrics[0];
assert_eq!(reported_vec.get_name(), "test_metric");
let dims = reported_vec.get_metric();
assert_eq!(dims.len(), 1);
assert_eq!(dims[0].get_label()[0].get_value(), "one");
drop(metric_1);
let metrics = reg.gather();
assert_eq!(metrics.len(), 0);
let string_labels: Vec<String> = ["owned"].iter().map(ToString::to_string).collect();
struct Ownership {
gauge: DeleteOnDropGauge<'static, AtomicI64, Vec<String>>,
}
let metric_owned = Ownership {
gauge: vec.get_delete_on_drop_gauge(string_labels),
};
metric_owned.gauge.set(666);
let metrics = reg.gather();
assert_eq!(metrics.len(), 1);
let reported_vec = &metrics[0];
assert_eq!(reported_vec.get_name(), "test_metric");
let dims = reported_vec.get_metric();
assert_eq!(dims.len(), 1);
assert_eq!(dims[0].get_label()[0].get_value(), "owned");
drop(metric_owned);
let metrics = reg.gather();
assert_eq!(metrics.len(), 0);
}
}