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/// Handles obtained by separate calls for the same labels each unregister them on drop. Use
183/// [`DeregisterOnDropVec`](super::DeregisterOnDropVec) when several owners hold the same labels.
184///
185/// NOTE: This type implements [`Borrow`], which imposes some constraints on implementers. To
186/// ensure these constraints, do *not* implement any of the `Eq`, `Ord`, or `Hash` traits on this.
187/// type.
188#[derive(Debug, Clone)]
189pub struct DeleteOnDropMetric<V, L>
190where
191    V: MetricVec_,
192    L: PromLabelsExt,
193{
194    inner: V::M,
195    /// Shared cleanup handle. The label is removed from `vec` only when the last clone drops.
196    /// Held purely for its `Drop` side effect — never read directly.
197    #[allow(dead_code)]
198    cleanup: Arc<DeleteOnDropCleanup<V, L>>,
199}
200
201/// Owns the data needed to unregister a labeled metric from its parent vector. Held inside an
202/// [`Arc`] by [`DeleteOnDropMetric`] so that registration is reversed exactly once, when the
203/// last clone of the metric drops.
204#[derive(Debug)]
205struct DeleteOnDropCleanup<V, L>
206where
207    V: MetricVec_,
208    L: PromLabelsExt,
209{
210    labels: L,
211    vec: V,
212}
213
214impl<V, L> DeleteOnDropMetric<V, L>
215where
216    V: MetricVec_,
217    L: PromLabelsExt,
218{
219    fn from_metric_vector(vec: V, labels: L) -> Self {
220        let inner = labels.get_from_metric_vec(&vec);
221        Self {
222            inner,
223            cleanup: Arc::new(DeleteOnDropCleanup { labels, vec }),
224        }
225    }
226}
227
228impl<V, L> Deref for DeleteOnDropMetric<V, L>
229where
230    V: MetricVec_,
231    L: PromLabelsExt,
232{
233    type Target = V::M;
234
235    fn deref(&self) -> &Self::Target {
236        &self.inner
237    }
238}
239
240impl<V, L> Drop for DeleteOnDropCleanup<V, L>
241where
242    V: MetricVec_,
243    L: PromLabelsExt,
244{
245    fn drop(&mut self) {
246        if self.labels.remove_from_metric_vec(&self.vec).is_err() {
247            // ignore.
248        }
249    }
250}
251
252/// A [`GenericCounter`] wrapper that deletes its labels from the vec when it is dropped.
253pub type DeleteOnDropCounter<P, L> = DeleteOnDropMetric<GenericCounterVec<P>, L>;
254
255impl<P, L> Borrow<GenericCounter<P>> for DeleteOnDropCounter<P, L>
256where
257    P: Atomic,
258    L: PromLabelsExt,
259{
260    fn borrow(&self) -> &GenericCounter<P> {
261        &self.inner
262    }
263}
264
265/// A [`GenericGauge`] wrapper that deletes its labels from the vec when it is dropped.
266pub type DeleteOnDropGauge<P, L> = DeleteOnDropMetric<GenericGaugeVec<P>, L>;
267
268impl<P, L> Borrow<GenericGauge<P>> for DeleteOnDropGauge<P, L>
269where
270    P: Atomic,
271    L: PromLabelsExt,
272{
273    fn borrow(&self) -> &GenericGauge<P> {
274        &self.inner
275    }
276}
277
278/// A [`Histogram`] wrapper that deletes its labels from the vec when it is dropped.
279pub type DeleteOnDropHistogram<L> = DeleteOnDropMetric<HistogramVec, L>;
280
281impl<L> Borrow<Histogram> for DeleteOnDropHistogram<L>
282where
283    L: PromLabelsExt,
284{
285    fn borrow(&self) -> &Histogram {
286        &self.inner
287    }
288}
289
290#[cfg(test)]
291mod test {
292    use prometheus::IntGaugeVec;
293    use prometheus::core::{AtomicI64, AtomicU64};
294
295    use crate::metric;
296    use crate::metrics::{IntCounterVec, MetricsRegistry};
297
298    use super::*;
299
300    #[crate::test]
301    fn dropping_counters() {
302        let reg = MetricsRegistry::new();
303        let vec: IntCounterVec = reg.register(metric!(
304            name: "test_metric",
305            help: "a test metric",
306            var_labels: ["dimension"]));
307
308        let dims: &[&str] = &["one"];
309        let metric_1 = vec.get_delete_on_drop_metric(dims);
310        metric_1.inc();
311
312        let metrics = reg.gather();
313        assert_eq!(metrics.len(), 1);
314        let reported_vec = &metrics[0];
315        assert_eq!(reported_vec.name(), "test_metric");
316        let dims = reported_vec.get_metric();
317        assert_eq!(dims.len(), 1);
318        assert_eq!(dims[0].get_label()[0].value(), "one");
319
320        drop(metric_1);
321        let metrics = reg.gather();
322        assert_eq!(metrics.len(), 0);
323
324        let string_labels: Vec<String> = ["owned"].iter().map(ToString::to_string).collect();
325        struct Ownership {
326            counter: DeleteOnDropCounter<AtomicU64, Vec<String>>,
327        }
328        let metric_owned = Ownership {
329            counter: vec.get_delete_on_drop_metric(string_labels),
330        };
331        metric_owned.counter.inc();
332
333        let metrics = reg.gather();
334        assert_eq!(metrics.len(), 1);
335        let reported_vec = &metrics[0];
336        assert_eq!(reported_vec.name(), "test_metric");
337        let dims = reported_vec.get_metric();
338        assert_eq!(dims.len(), 1);
339        assert_eq!(dims[0].get_label()[0].value(), "owned");
340
341        drop(metric_owned);
342        let metrics = reg.gather();
343        assert_eq!(metrics.len(), 0);
344    }
345
346    #[crate::test]
347    fn dropping_gauges() {
348        let reg = MetricsRegistry::new();
349        let vec: IntGaugeVec = reg.register(metric!(
350            name: "test_metric",
351            help: "a test metric",
352            var_labels: ["dimension"]));
353
354        let dims: &[&str] = &["one"];
355        let metric_1 = vec.get_delete_on_drop_metric(dims);
356        metric_1.set(666);
357
358        let metrics = reg.gather();
359        assert_eq!(metrics.len(), 1);
360        let reported_vec = &metrics[0];
361        assert_eq!(reported_vec.name(), "test_metric");
362        let dims = reported_vec.get_metric();
363        assert_eq!(dims.len(), 1);
364        assert_eq!(dims[0].get_label()[0].value(), "one");
365
366        drop(metric_1);
367        let metrics = reg.gather();
368        assert_eq!(metrics.len(), 0);
369
370        let string_labels: Vec<String> = ["owned"].iter().map(ToString::to_string).collect();
371        struct Ownership {
372            gauge: DeleteOnDropGauge<AtomicI64, Vec<String>>,
373        }
374        let metric_owned = Ownership {
375            gauge: vec.get_delete_on_drop_metric(string_labels),
376        };
377        metric_owned.gauge.set(666);
378
379        let metrics = reg.gather();
380        assert_eq!(metrics.len(), 1);
381        let reported_vec = &metrics[0];
382        assert_eq!(reported_vec.name(), "test_metric");
383        let dims = reported_vec.get_metric();
384        assert_eq!(dims.len(), 1);
385        assert_eq!(dims[0].get_label()[0].value(), "owned");
386
387        drop(metric_owned);
388        let metrics = reg.gather();
389        assert_eq!(metrics.len(), 0);
390    }
391
392    /// Cloning a `DeleteOnDropMetric` must not unregister the labeled metric until the last
393    /// clone is dropped. Regression test for CLU-63.
394    #[crate::test]
395    fn clones_share_registration() {
396        let reg = MetricsRegistry::new();
397        let vec: IntCounterVec = reg.register(metric!(
398            name: "test_metric",
399            help: "a test metric",
400            var_labels: ["dimension"]));
401
402        let metric_1 = vec.get_delete_on_drop_metric(&["one"][..]);
403        let metric_2 = metric_1.clone();
404        metric_1.inc();
405        metric_2.inc();
406
407        // Dropping one clone must not unregister the metric.
408        drop(metric_1);
409        let gathered = reg.gather();
410        assert_eq!(gathered.len(), 1);
411        assert_eq!(gathered[0].get_metric()[0].get_counter().value(), 2.0);
412
413        // Increment via the surviving clone is still observable.
414        metric_2.inc();
415        let gathered = reg.gather();
416        assert_eq!(gathered[0].get_metric()[0].get_counter().value(), 3.0);
417
418        // Dropping the last clone unregisters the metric.
419        drop(metric_2);
420        assert_eq!(reg.gather().len(), 0);
421    }
422}
mz_ore/metrics/delete_on_drop.rs

mz_ore/metrics/
delete_on_drop.rs
