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
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167

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: Option<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) {
        if let Some(value) = self.value {
            self.metric.add(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: Some(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
    }

    /// Reference to the labels that will be used when incrementing the metric.
    pub fn labels(&self) -> &L {
        &self.labels
    }

    /// Mutable reference to the labels that will be used when incrementing the metric.
    pub fn labels_mut(&mut self) -> &mut L {
        &mut self.labels
    }

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

    /// Cancel the increment, consuming the guard.
    pub fn cancel(&mut self) {
        self.value = None;
    }
}

/// 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: Option<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)
    }

    /// Cancel the increment, consuming the guard.
    pub fn cancel(&mut self) {
        self.value = None;
    }
}

/// When dropped, a [`DeferredAdd`] guard will increment its counter.
impl<'a, P: Atomic> Drop for DeferredAdd<'a, P> {
    fn drop(&mut self) {
        if let Some(value) = self.value {
            self.metric.inc_by(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: Some(v),
            metric: self,
        }
    }
}