prometheus_utils/
guards.rs

1use crate::{IntCounterWithLabels, Labels};
2use prometheus::core::{Atomic, AtomicF64, AtomicI64, GenericCounter, GenericGauge, Number};
3
4/// An RAII-style guard for an [`AtomicI64`] gauge.
5///
6/// Created by the methods on the [`GuardedGauge`] extension trait.
7pub type IntGaugeGuard = GenericGaugeGuard<AtomicI64>;
8
9/// An RAII-style guard for an [`AtomicF64`] gauge.
10///
11/// Created by the methods on the [`GuardedGauge`] extension trait.
12pub type GaugeGuard = GenericGaugeGuard<AtomicF64>;
13
14/// An RAII-style guard for situations where we want to increment a gauge and then ensure that there
15/// is always a corresponding decrement.
16///
17/// Created by the methods on the [`GuardedGauge`] extension trait.
18pub struct GenericGaugeGuard<P: Atomic + 'static> {
19    value: P::T,
20    gauge: &'static GenericGauge<P>,
21}
22
23/// When a gauge guard is dropped, it will perform the corresponding decrement.
24impl<P: Atomic + 'static> Drop for GenericGaugeGuard<P> {
25    fn drop(&mut self) {
26        self.gauge.sub(self.value);
27    }
28}
29
30/// An extension trait for [`GenericGauge`] to provide methods for temporarily modifying a gauge.
31pub trait GuardedGauge<P: Atomic + 'static> {
32    /// Increase the gauge by 1 while the guard exists.
33    #[must_use]
34    fn guarded_inc(&'static self) -> GenericGaugeGuard<P>;
35
36    /// Increase the gauge by the given increment while the guard exists.
37    #[must_use]
38    fn guarded_add(&'static self, v: P::T) -> GenericGaugeGuard<P>;
39}
40
41impl<P: Atomic + 'static> GuardedGauge<P> for GenericGauge<P> {
42    fn guarded_inc(&'static self) -> GenericGaugeGuard<P> {
43        self.inc();
44        GenericGaugeGuard {
45            value: <P::T as Number>::from_i64(1),
46            gauge: self,
47        }
48    }
49
50    fn guarded_add(&'static self, v: P::T) -> GenericGaugeGuard<P> {
51        self.add(v);
52        GenericGaugeGuard {
53            value: v,
54            gauge: self,
55        }
56    }
57}
58
59/// A guard that will automatically increment a labeled metric when dropped.
60///
61/// Created by calling [`IntCounterWithLabels::deferred_inc`].
62pub struct DeferredAddWithLabels<'a, L: Labels> {
63    value: Option<u64>,
64    metric: &'a IntCounterWithLabels<L>,
65    labels: L,
66}
67
68/// When dropped, a [`DeferredAddWithLabels`] guard will increment its counter.
69impl<'a, L: Labels> Drop for DeferredAddWithLabels<'a, L> {
70    fn drop(&mut self) {
71        if let Some(value) = self.value {
72            self.metric.add(value, &self.labels)
73        }
74    }
75}
76
77impl<'a, L: Labels> DeferredAddWithLabels<'a, L> {
78    /// Create a new deferred increment guard.
79    //
80    // This is not exposed in the public interface, these should only be acquired through
81    // `deferred_inc`.
82    pub(crate) fn new(metric: &'a IntCounterWithLabels<L>, value: u64, labels: L) -> Self {
83        Self {
84            value: Some(value),
85            metric,
86            labels,
87        }
88    }
89
90    /// Update the labels to use when incrementing the metric.
91    pub fn with_labels(mut self, new_labels: L) -> DeferredAddWithLabels<'a, L> {
92        self.labels = new_labels;
93        self
94    }
95
96    /// Reference to the labels that will be used when incrementing the metric.
97    pub fn labels(&self) -> &L {
98        &self.labels
99    }
100
101    /// Mutable reference to the labels that will be used when incrementing the metric.
102    pub fn labels_mut(&mut self) -> &mut L {
103        &mut self.labels
104    }
105
106    /// Eagerly perform the increment, consuming the guard.
107    pub fn complete_add(self) {
108        drop(self)
109    }
110
111    /// Cancel the increment, consuming the guard.
112    pub fn cancel(&mut self) {
113        self.value = None;
114    }
115}
116
117/// A guard that will automatically increment a [`GenericCounter`] when dropped.
118///
119/// Created by the methods on the [`DeferredCounter`] extension trait.
120pub struct DeferredAdd<'a, P: Atomic> {
121    value: Option<P::T>,
122    metric: &'a GenericCounter<P>,
123}
124
125impl<'a, P: Atomic> DeferredAdd<'a, P> {
126    /// Eagerly perform the increment, consuming the guard.
127    pub fn complete_add(self) {
128        drop(self)
129    }
130
131    /// Cancel the increment, consuming the guard.
132    pub fn cancel(&mut self) {
133        self.value = None;
134    }
135}
136
137/// When dropped, a [`DeferredAdd`] guard will increment its counter.
138impl<'a, P: Atomic> Drop for DeferredAdd<'a, P> {
139    fn drop(&mut self) {
140        if let Some(value) = self.value {
141            self.metric.inc_by(value);
142        }
143    }
144}
145
146/// An extension trait for [`GenericCounter`] to provide methods for incrementing a counter once
147/// an RAII-style guard has been dropped.
148pub trait DeferredCounter<P: Atomic + 'static> {
149    /// Increase the counter by `1` when the guard is dropped.
150    #[must_use]
151    fn deferred_inc(&'static self) -> DeferredAdd<P> {
152        self.deferred_add(<P::T as Number>::from_i64(1))
153    }
154
155    /// Increase the counter by `v` when the guard is dropped.
156    #[must_use]
157    fn deferred_add(&'static self, v: P::T) -> DeferredAdd<P>;
158}
159
160impl<P: Atomic + 'static> DeferredCounter<P> for GenericCounter<P> {
161    fn deferred_add(&'static self, v: P::T) -> DeferredAdd<P> {
162        DeferredAdd {
163            value: Some(v),
164            metric: self,
165        }
166    }
167}
prometheus_utils/guards.rs

prometheus_utils/
guards.rs
