1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143

use crate::{IntCounterWithLabels, Labels};
use prometheus::core::{Atomic, AtomicF64, AtomicI64, GenericCounter, GenericGauge, Number};

/// An RAII-style guard for an [`AtomicI64`] gauge.
///
/// Created by the methods on the [`GuardedGauge`] extension trait.
pub type IntGaugeGuard = GenericGaugeGuard<AtomicI64>;

/// An RAII-style guard for an [`AtomicF64`] gauge.
///
/// Created by the methods on the [`GuardedGauge`] extension trait.
pub type GaugeGuard = GenericGaugeGuard<AtomicF64>;

/// An RAII-style guard for situations where we want to increment a gauge and then ensure that there
/// is always a corresponding decrement.
///
/// Created by the methods on the [`GuardedGauge`] extension trait.
pub struct GenericGaugeGuard<P: Atomic + 'static> {
    value: P::T,
    gauge: &'static GenericGauge<P>,
}

/// When a gauge guard is dropped, it will perform the corresponding decrement.
impl<P: Atomic + 'static> Drop for GenericGaugeGuard<P> {
    fn drop(&mut self) {
        self.gauge.sub(self.value);
    }
}

/// An extension trait for [`GenericGauge`] to provide methods for temporarily modifying a gauge.
pub trait GuardedGauge<P: Atomic + 'static> {
    /// Increase the gauge by 1 while the guard exists.
    #[must_use]
    fn guarded_inc(&'static self) -> GenericGaugeGuard<P>;

    /// Increase the gauge by the given increment while the guard exists.
    #[must_use]
    fn guarded_add(&'static self, v: P::T) -> GenericGaugeGuard<P>;
}

impl<P: Atomic + 'static> GuardedGauge<P> for GenericGauge<P> {
    fn guarded_inc(&'static self) -> GenericGaugeGuard<P> {
        self.inc();
        GenericGaugeGuard {
            value: <P::T as Number>::from_i64(1),
            gauge: self,
        }
    }

    fn guarded_add(&'static self, v: P::T) -> GenericGaugeGuard<P> {
        self.add(v);
        GenericGaugeGuard {
            value: v,
            gauge: self,
        }
    }
}

/// A guard that will automatically increment a labeled metric when dropped.
///
/// Created by calling [`IntCounterWithLabels::deferred_inc`].
pub struct DeferredAddWithLabels<'a, L: Labels> {
    value: u64,
    metric: &'a IntCounterWithLabels<L>,
    labels: L,
}

/// When dropped, a [`DeferredAddWithLabels`] guard will increment its counter.
impl<'a, L: Labels> Drop for DeferredAddWithLabels<'a, L> {
    fn drop(&mut self) {
        self.metric.add(self.value, &self.labels)
    }
}

impl<'a, L: Labels> DeferredAddWithLabels<'a, L> {
    /// Create a new deferred increment guard.
    //
    // This is not exposed in the public interface, these should only be acquired through
    // `deferred_inc`.
    pub(crate) fn new(metric: &'a IntCounterWithLabels<L>, value: u64, labels: L) -> Self {
        Self {
            value,
            metric,
            labels,
        }
    }

    /// Update the labels to use when incrementing the metric.
    pub fn with_labels(mut self, new_labels: L) -> DeferredAddWithLabels<'a, L> {
        self.labels = new_labels;
        self
    }

    /// Eagerly perform the increment, consuming the guard.
    pub fn complete_add(self) {
        drop(self)
    }
}

/// A guard that will automatically increment a [`GenericCounter`] when dropped.
///
/// Created by the methods on the [`DeferredCounter`] extension trait.
pub struct DeferredAdd<'a, P: Atomic> {
    value: P::T,
    metric: &'a GenericCounter<P>,
}

impl<'a, P: Atomic> DeferredAdd<'a, P> {
    /// Eagerly perform the increment, consuming the guard.
    pub fn complete_add(self) {
        drop(self)
    }
}

/// When dropped, a [`DeferredAdd`] guard will increment its counter.
impl<'a, P: Atomic> Drop for DeferredAdd<'a, P> {
    fn drop(&mut self) {
        self.metric.inc_by(self.value);
    }
}

/// An extension trait for [`GenericCounter`] to provide methods for incrementing a counter once
/// an RAII-style guard has been dropped.
pub trait DeferredCounter<P: Atomic + 'static> {
    /// Increase the counter by `1` when the guard is dropped.
    #[must_use]
    fn deferred_inc(&'static self) -> DeferredAdd<P> {
        self.deferred_add(<P::T as Number>::from_i64(1))
    }

    /// Increase the counter by `v` when the guard is dropped.
    #[must_use]
    fn deferred_add(&'static self, v: P::T) -> DeferredAdd<P>;
}

impl<P: Atomic + 'static> DeferredCounter<P> for GenericCounter<P> {
    fn deferred_add(&'static self, v: P::T) -> DeferredAdd<P> {
        DeferredAdd {
            value: v,
            metric: self,
        }
    }
}