mz_ore/metrics/
delete_on_drop.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License in the LICENSE file at the
6// root of this repository, or online at
7//
8//     http://www.apache.org/licenses/LICENSE-2.0
9//
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16//! Support for metrics that get removed from their corresponding metrics vector when dropped.
17//!
18//! # Ownership & life times
19//!
20//! This kind of data type is realized by a struct that retains ownership of the _labels_ used to
21//! create the spin-off metric. The created metric follows these rules:
22//! * When passing references, the metric must not outlive the references to the labels used to create
23//!   it: A `'static` slice of static strings means the metric is allowed to live for the `'static`
24//!   lifetime as well.
25//! * Metrics created from references to dynamically constructed labels can only live as long as those
26//!   labels do.
27//! * When using owned data (an extension over what Prometheus allows, which only lets you use
28//!   references to refer to labels), the created metric is also allowed to live for `'static`.
29
30use std::borrow::Borrow;
31use std::collections::BTreeMap;
32use std::ops::Deref;
33use std::sync::Arc;
34
35use prometheus::core::{
36    Atomic, GenericCounter, GenericCounterVec, GenericGauge, GenericGaugeVec, Metric, MetricVec,
37    MetricVecBuilder,
38};
39use prometheus::{Histogram, HistogramVec};
40
41/// The `prometheus` API uses the `HashMap` type to pass metrics labels, so we have to allow its
42/// usage when calling that API.
43#[allow(clippy::disallowed_types)]
44type PromLabelMap<'a> = std::collections::HashMap<&'a str, &'a str>;
45
46/// A trait that allows being generic over [`MetricVec`]s.
47pub trait MetricVec_: Sized {
48    /// The associated Metric collected.
49    type M: Metric;
50
51    /// See [`MetricVec::get_metric_with_label_values`].
52    fn get_metric_with_label_values(&self, vals: &[&str]) -> Result<Self::M, prometheus::Error>;
53
54    /// See [`MetricVec::get_metric_with`].
55    fn get_metric_with(&self, labels: &PromLabelMap) -> Result<Self::M, prometheus::Error>;
56
57    /// See [`MetricVec::remove_label_values`].
58    fn remove_label_values(&self, vals: &[&str]) -> Result<(), prometheus::Error>;
59
60    /// See [`MetricVec::remove`].
61    fn remove(&self, labels: &PromLabelMap) -> Result<(), prometheus::Error>;
62}
63
64impl<P: MetricVecBuilder> MetricVec_ for MetricVec<P> {
65    type M = P::M;
66
67    fn get_metric_with_label_values(&self, vals: &[&str]) -> prometheus::Result<Self::M> {
68        self.get_metric_with_label_values(vals)
69    }
70
71    fn get_metric_with(&self, labels: &PromLabelMap) -> Result<Self::M, prometheus::Error> {
72        self.get_metric_with(labels)
73    }
74
75    fn remove_label_values(&self, vals: &[&str]) -> Result<(), prometheus::Error> {
76        self.remove_label_values(vals)
77    }
78
79    fn remove(&self, labels: &PromLabelMap) -> Result<(), prometheus::Error> {
80        self.remove(labels)
81    }
82}
83
84/// Extension trait for metrics vectors.
85///
86/// It adds a method to create a concrete metric from the vector that gets removed from the vector
87/// when the concrete metric is dropped.
88pub trait MetricVecExt: MetricVec_ {
89    /// Returns a metric that deletes its labels from this metrics vector when dropped.
90    fn get_delete_on_drop_metric<L: PromLabelsExt>(&self, labels: L)
91    -> DeleteOnDropMetric<Self, L>;
92}
93
94impl<V: MetricVec_ + Clone> MetricVecExt for V {
95    fn get_delete_on_drop_metric<L>(&self, labels: L) -> DeleteOnDropMetric<Self, L>
96    where
97        L: PromLabelsExt,
98    {
99        DeleteOnDropMetric::from_metric_vector(self.clone(), labels)
100    }
101}
102
103/// An extension trait for types that are valid (or convertible into) prometheus labels:
104/// slices/vectors of strings, and [`BTreeMap`]s.
105pub trait PromLabelsExt {
106    /// Returns or creates a metric with the given metric label values.
107    /// Panics if retrieving the metric returns an error.
108    fn get_from_metric_vec<V: MetricVec_>(&self, vec: &V) -> V::M;
109
110    /// Removes a metric with these labels from a metrics vector.
111    fn remove_from_metric_vec<V: MetricVec_>(&self, vec: &V) -> Result<(), prometheus::Error>;
112}
113
114impl PromLabelsExt for &[&str] {
115    fn get_from_metric_vec<V: MetricVec_>(&self, vec: &V) -> V::M {
116        vec.get_metric_with_label_values(self)
117            .expect("retrieving a metric by label values")
118    }
119
120    fn remove_from_metric_vec<V: MetricVec_>(&self, vec: &V) -> Result<(), prometheus::Error> {
121        vec.remove_label_values(self)
122    }
123}
124
125impl PromLabelsExt for Vec<String> {
126    fn get_from_metric_vec<V: MetricVec_>(&self, vec: &V) -> V::M {
127        let labels: Vec<&str> = self.iter().map(String::as_str).collect();
128        vec.get_metric_with_label_values(labels.as_slice())
129            .expect("retrieving a metric by label values")
130    }
131
132    fn remove_from_metric_vec<V: MetricVec_>(&self, vec: &V) -> Result<(), prometheus::Error> {
133        let labels: Vec<&str> = self.iter().map(String::as_str).collect();
134        vec.remove_label_values(labels.as_slice())
135    }
136}
137
138impl PromLabelsExt for Vec<&str> {
139    fn get_from_metric_vec<V: MetricVec_>(&self, vec: &V) -> V::M {
140        vec.get_metric_with_label_values(self.as_slice())
141            .expect("retrieving a metric by label values")
142    }
143
144    fn remove_from_metric_vec<V: MetricVec_>(&self, vec: &V) -> Result<(), prometheus::Error> {
145        vec.remove_label_values(self.as_slice())
146    }
147}
148
149impl PromLabelsExt for BTreeMap<String, String> {
150    fn get_from_metric_vec<V: MetricVec_>(&self, vec: &V) -> V::M {
151        let labels = self.iter().map(|(k, v)| (k.as_str(), v.as_str())).collect();
152        vec.get_metric_with(&labels)
153            .expect("retrieving a metric by label values")
154    }
155
156    fn remove_from_metric_vec<V: MetricVec_>(&self, vec: &V) -> Result<(), prometheus::Error> {
157        let labels = self.iter().map(|(k, v)| (k.as_str(), v.as_str())).collect();
158        vec.remove(&labels)
159    }
160}
161
162impl PromLabelsExt for BTreeMap<&str, &str> {
163    fn get_from_metric_vec<V: MetricVec_>(&self, vec: &V) -> V::M {
164        let labels = self.iter().map(|(k, v)| (*k, *v)).collect();
165        vec.get_metric_with(&labels)
166            .expect("retrieving a metric by label values")
167    }
168
169    fn remove_from_metric_vec<V: MetricVec_>(&self, vec: &V) -> Result<(), prometheus::Error> {
170        let labels = self.iter().map(|(k, v)| (*k, *v)).collect();
171        vec.remove(&labels)
172    }
173}
174
175/// A [`Metric`] wrapper that deletes its labels from the vec when the last clone is dropped.
176///
177/// Cloning is ref-counted: the labeled metric is unregistered from its parent vector exactly
178/// once, when the final clone drops. This means cloning is always safe — earlier versions of
179/// this type derived `Clone` directly, which made every clone's `Drop` unregister the metric
180/// and silently disabled the metric for surviving clones.
181///
182/// NOTE: This type implements [`Borrow`], which imposes some constraints on implementers. To
183/// ensure these constraints, do *not* implement any of the `Eq`, `Ord`, or `Hash` traits on this.
184/// type.
185#[derive(Debug, Clone)]
186pub struct DeleteOnDropMetric<V, L>
187where
188    V: MetricVec_,
189    L: PromLabelsExt,
190{
191    inner: V::M,
192    /// Shared cleanup handle. The label is removed from `vec` only when the last clone drops.
193    /// Held purely for its `Drop` side effect — never read directly.
194    #[allow(dead_code)]
195    cleanup: Arc<DeleteOnDropCleanup<V, L>>,
196}
197
198/// Owns the data needed to unregister a labeled metric from its parent vector. Held inside an
199/// [`Arc`] by [`DeleteOnDropMetric`] so that registration is reversed exactly once, when the
200/// last clone of the metric drops.
201#[derive(Debug)]
202struct DeleteOnDropCleanup<V, L>
203where
204    V: MetricVec_,
205    L: PromLabelsExt,
206{
207    labels: L,
208    vec: V,
209}
210
211impl<V, L> DeleteOnDropMetric<V, L>
212where
213    V: MetricVec_,
214    L: PromLabelsExt,
215{
216    fn from_metric_vector(vec: V, labels: L) -> Self {
217        let inner = labels.get_from_metric_vec(&vec);
218        Self {
219            inner,
220            cleanup: Arc::new(DeleteOnDropCleanup { labels, vec }),
221        }
222    }
223}
224
225impl<V, L> Deref for DeleteOnDropMetric<V, L>
226where
227    V: MetricVec_,
228    L: PromLabelsExt,
229{
230    type Target = V::M;
231
232    fn deref(&self) -> &Self::Target {
233        &self.inner
234    }
235}
236
237impl<V, L> Drop for DeleteOnDropCleanup<V, L>
238where
239    V: MetricVec_,
240    L: PromLabelsExt,
241{
242    fn drop(&mut self) {
243        if self.labels.remove_from_metric_vec(&self.vec).is_err() {
244            // ignore.
245        }
246    }
247}
248
249/// A [`GenericCounter`] wrapper that deletes its labels from the vec when it is dropped.
250pub type DeleteOnDropCounter<P, L> = DeleteOnDropMetric<GenericCounterVec<P>, L>;
251
252impl<P, L> Borrow<GenericCounter<P>> for DeleteOnDropCounter<P, L>
253where
254    P: Atomic,
255    L: PromLabelsExt,
256{
257    fn borrow(&self) -> &GenericCounter<P> {
258        &self.inner
259    }
260}
261
262/// A [`GenericGauge`] wrapper that deletes its labels from the vec when it is dropped.
263pub type DeleteOnDropGauge<P, L> = DeleteOnDropMetric<GenericGaugeVec<P>, L>;
264
265impl<P, L> Borrow<GenericGauge<P>> for DeleteOnDropGauge<P, L>
266where
267    P: Atomic,
268    L: PromLabelsExt,
269{
270    fn borrow(&self) -> &GenericGauge<P> {
271        &self.inner
272    }
273}
274
275/// A [`Histogram`] wrapper that deletes its labels from the vec when it is dropped.
276pub type DeleteOnDropHistogram<L> = DeleteOnDropMetric<HistogramVec, L>;
277
278impl<L> Borrow<Histogram> for DeleteOnDropHistogram<L>
279where
280    L: PromLabelsExt,
281{
282    fn borrow(&self) -> &Histogram {
283        &self.inner
284    }
285}
286
287#[cfg(test)]
288mod test {
289    use prometheus::IntGaugeVec;
290    use prometheus::core::{AtomicI64, AtomicU64};
291
292    use crate::metric;
293    use crate::metrics::{IntCounterVec, MetricsRegistry};
294
295    use super::*;
296
297    #[crate::test]
298    fn dropping_counters() {
299        let reg = MetricsRegistry::new();
300        let vec: IntCounterVec = reg.register(metric!(
301            name: "test_metric",
302            help: "a test metric",
303            var_labels: ["dimension"]));
304
305        let dims: &[&str] = &["one"];
306        let metric_1 = vec.get_delete_on_drop_metric(dims);
307        metric_1.inc();
308
309        let metrics = reg.gather();
310        assert_eq!(metrics.len(), 1);
311        let reported_vec = &metrics[0];
312        assert_eq!(reported_vec.name(), "test_metric");
313        let dims = reported_vec.get_metric();
314        assert_eq!(dims.len(), 1);
315        assert_eq!(dims[0].get_label()[0].value(), "one");
316
317        drop(metric_1);
318        let metrics = reg.gather();
319        assert_eq!(metrics.len(), 0);
320
321        let string_labels: Vec<String> = ["owned"].iter().map(ToString::to_string).collect();
322        struct Ownership {
323            counter: DeleteOnDropCounter<AtomicU64, Vec<String>>,
324        }
325        let metric_owned = Ownership {
326            counter: vec.get_delete_on_drop_metric(string_labels),
327        };
328        metric_owned.counter.inc();
329
330        let metrics = reg.gather();
331        assert_eq!(metrics.len(), 1);
332        let reported_vec = &metrics[0];
333        assert_eq!(reported_vec.name(), "test_metric");
334        let dims = reported_vec.get_metric();
335        assert_eq!(dims.len(), 1);
336        assert_eq!(dims[0].get_label()[0].value(), "owned");
337
338        drop(metric_owned);
339        let metrics = reg.gather();
340        assert_eq!(metrics.len(), 0);
341    }
342
343    #[crate::test]
344    fn dropping_gauges() {
345        let reg = MetricsRegistry::new();
346        let vec: IntGaugeVec = reg.register(metric!(
347            name: "test_metric",
348            help: "a test metric",
349            var_labels: ["dimension"]));
350
351        let dims: &[&str] = &["one"];
352        let metric_1 = vec.get_delete_on_drop_metric(dims);
353        metric_1.set(666);
354
355        let metrics = reg.gather();
356        assert_eq!(metrics.len(), 1);
357        let reported_vec = &metrics[0];
358        assert_eq!(reported_vec.name(), "test_metric");
359        let dims = reported_vec.get_metric();
360        assert_eq!(dims.len(), 1);
361        assert_eq!(dims[0].get_label()[0].value(), "one");
362
363        drop(metric_1);
364        let metrics = reg.gather();
365        assert_eq!(metrics.len(), 0);
366
367        let string_labels: Vec<String> = ["owned"].iter().map(ToString::to_string).collect();
368        struct Ownership {
369            gauge: DeleteOnDropGauge<AtomicI64, Vec<String>>,
370        }
371        let metric_owned = Ownership {
372            gauge: vec.get_delete_on_drop_metric(string_labels),
373        };
374        metric_owned.gauge.set(666);
375
376        let metrics = reg.gather();
377        assert_eq!(metrics.len(), 1);
378        let reported_vec = &metrics[0];
379        assert_eq!(reported_vec.name(), "test_metric");
380        let dims = reported_vec.get_metric();
381        assert_eq!(dims.len(), 1);
382        assert_eq!(dims[0].get_label()[0].value(), "owned");
383
384        drop(metric_owned);
385        let metrics = reg.gather();
386        assert_eq!(metrics.len(), 0);
387    }
388
389    /// Cloning a `DeleteOnDropMetric` must not unregister the labeled metric until the last
390    /// clone is dropped. Regression test for CLU-63.
391    #[crate::test]
392    fn clones_share_registration() {
393        let reg = MetricsRegistry::new();
394        let vec: IntCounterVec = reg.register(metric!(
395            name: "test_metric",
396            help: "a test metric",
397            var_labels: ["dimension"]));
398
399        let metric_1 = vec.get_delete_on_drop_metric(&["one"][..]);
400        let metric_2 = metric_1.clone();
401        metric_1.inc();
402        metric_2.inc();
403
404        // Dropping one clone must not unregister the metric.
405        drop(metric_1);
406        let gathered = reg.gather();
407        assert_eq!(gathered.len(), 1);
408        assert_eq!(gathered[0].get_metric()[0].get_counter().get_value(), 2.0);
409
410        // Increment via the surviving clone is still observable.
411        metric_2.inc();
412        let gathered = reg.gather();
413        assert_eq!(gathered[0].get_metric()[0].get_counter().get_value(), 3.0);
414
415        // Dropping the last clone unregisters the metric.
416        drop(metric_2);
417        assert_eq!(reg.gather().len(), 0);
418    }
419}
mz_ore/metrics/delete_on_drop.rs

mz_ore/metrics/
delete_on_drop.rs
