mz_ore/metrics/
delete_on_drop.rs

Help
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;
33
34use prometheus::core::{
35    Atomic, GenericCounter, GenericCounterVec, GenericGauge, GenericGaugeVec, Metric, MetricVec,
36    MetricVecBuilder,
37};
38use prometheus::{Histogram, HistogramVec};
39
40/// The `prometheus` API uses the `HashMap` type to pass metrics labels, so we have to allow its
41/// usage when calling that API.
42#[allow(clippy::disallowed_types)]
43type PromLabelMap<'a> = std::collections::HashMap<&'a str, &'a str>;
44
45/// A trait that allows being generic over [`MetricVec`]s.
46pub trait MetricVec_: Sized {
47    /// The associated Metric collected.
48    type M: Metric;
49
50    /// See [`MetricVec::get_metric_with_label_values`].
51    fn get_metric_with_label_values(&self, vals: &[&str]) -> Result<Self::M, prometheus::Error>;
52
53    /// See [`MetricVec::get_metric_with`].
54    fn get_metric_with(&self, labels: &PromLabelMap) -> Result<Self::M, prometheus::Error>;
55
56    /// See [`MetricVec::remove_label_values`].
57    fn remove_label_values(&self, vals: &[&str]) -> Result<(), prometheus::Error>;
58
59    /// See [`MetricVec::remove`].
60    fn remove(&self, labels: &PromLabelMap) -> Result<(), prometheus::Error>;
61}
62
63impl<P: MetricVecBuilder> MetricVec_ for MetricVec<P> {
64    type M = P::M;
65
66    fn get_metric_with_label_values(&self, vals: &[&str]) -> prometheus::Result<Self::M> {
67        self.get_metric_with_label_values(vals)
68    }
69
70    fn get_metric_with(&self, labels: &PromLabelMap) -> Result<Self::M, prometheus::Error> {
71        self.get_metric_with(labels)
72    }
73
74    fn remove_label_values(&self, vals: &[&str]) -> Result<(), prometheus::Error> {
75        self.remove_label_values(vals)
76    }
77
78    fn remove(&self, labels: &PromLabelMap) -> Result<(), prometheus::Error> {
79        self.remove(labels)
80    }
81}
82
83/// Extension trait for metrics vectors.
84///
85/// It adds a method to create a concrete metric from the vector that gets removed from the vector
86/// when the concrete metric is dropped.
87pub trait MetricVecExt: MetricVec_ {
88    /// Returns a metric that deletes its labels from this metrics vector when dropped.
89    fn get_delete_on_drop_metric<L: PromLabelsExt>(&self, labels: L)
90    -> DeleteOnDropMetric<Self, L>;
91}
92
93impl<V: MetricVec_ + Clone> MetricVecExt for V {
94    fn get_delete_on_drop_metric<L>(&self, labels: L) -> DeleteOnDropMetric<Self, L>
95    where
96        L: PromLabelsExt,
97    {
98        DeleteOnDropMetric::from_metric_vector(self.clone(), labels)
99    }
100}
101
102/// An extension trait for types that are valid (or convertible into) prometheus labels:
103/// slices/vectors of strings, and [`BTreeMap`]s.
104pub trait PromLabelsExt {
105    /// Returns or creates a metric with the given metric label values.
106    /// Panics if retrieving the metric returns an error.
107    fn get_from_metric_vec<V: MetricVec_>(&self, vec: &V) -> V::M;
108
109    /// Removes a metric with these labels from a metrics vector.
110    fn remove_from_metric_vec<V: MetricVec_>(&self, vec: &V) -> Result<(), prometheus::Error>;
111}
112
113impl PromLabelsExt for &[&str] {
114    fn get_from_metric_vec<V: MetricVec_>(&self, vec: &V) -> V::M {
115        vec.get_metric_with_label_values(self)
116            .expect("retrieving a metric by label values")
117    }
118
119    fn remove_from_metric_vec<V: MetricVec_>(&self, vec: &V) -> Result<(), prometheus::Error> {
120        vec.remove_label_values(self)
121    }
122}
123
124impl PromLabelsExt for Vec<String> {
125    fn get_from_metric_vec<V: MetricVec_>(&self, vec: &V) -> V::M {
126        let labels: Vec<&str> = self.iter().map(String::as_str).collect();
127        vec.get_metric_with_label_values(labels.as_slice())
128            .expect("retrieving a metric by label values")
129    }
130
131    fn remove_from_metric_vec<V: MetricVec_>(&self, vec: &V) -> Result<(), prometheus::Error> {
132        let labels: Vec<&str> = self.iter().map(String::as_str).collect();
133        vec.remove_label_values(labels.as_slice())
134    }
135}
136
137impl PromLabelsExt for Vec<&str> {
138    fn get_from_metric_vec<V: MetricVec_>(&self, vec: &V) -> V::M {
139        vec.get_metric_with_label_values(self.as_slice())
140            .expect("retrieving a metric by label values")
141    }
142
143    fn remove_from_metric_vec<V: MetricVec_>(&self, vec: &V) -> Result<(), prometheus::Error> {
144        vec.remove_label_values(self.as_slice())
145    }
146}
147
148impl PromLabelsExt for BTreeMap<String, String> {
149    fn get_from_metric_vec<V: MetricVec_>(&self, vec: &V) -> V::M {
150        let labels = self.iter().map(|(k, v)| (k.as_str(), v.as_str())).collect();
151        vec.get_metric_with(&labels)
152            .expect("retrieving a metric by label values")
153    }
154
155    fn remove_from_metric_vec<V: MetricVec_>(&self, vec: &V) -> Result<(), prometheus::Error> {
156        let labels = self.iter().map(|(k, v)| (k.as_str(), v.as_str())).collect();
157        vec.remove(&labels)
158    }
159}
160
161impl PromLabelsExt for BTreeMap<&str, &str> {
162    fn get_from_metric_vec<V: MetricVec_>(&self, vec: &V) -> V::M {
163        let labels = self.iter().map(|(k, v)| (*k, *v)).collect();
164        vec.get_metric_with(&labels)
165            .expect("retrieving a metric by label values")
166    }
167
168    fn remove_from_metric_vec<V: MetricVec_>(&self, vec: &V) -> Result<(), prometheus::Error> {
169        let labels = self.iter().map(|(k, v)| (*k, *v)).collect();
170        vec.remove(&labels)
171    }
172}
173
174/// A [`Metric`] wrapper that deletes its labels from the vec when it is dropped.
175///
176/// It adds a method to create a concrete metric from the vector that gets removed from the vector
177/// when the concrete metric is dropped.
178///
179/// NOTE: This type implements [`Borrow`], which imposes some constraints on implementers. To
180/// ensure these constraints, do *not* implement any of the `Eq`, `Ord`, or `Hash` traits on this.
181/// type.
182#[derive(Debug, Clone)]
183pub struct DeleteOnDropMetric<V, L>
184where
185    V: MetricVec_,
186    L: PromLabelsExt,
187{
188    inner: V::M,
189    labels: L,
190    vec: V,
191}
192
193impl<V, L> DeleteOnDropMetric<V, L>
194where
195    V: MetricVec_,
196    L: PromLabelsExt,
197{
198    fn from_metric_vector(vec: V, labels: L) -> Self {
199        let inner = labels.get_from_metric_vec(&vec);
200        Self { inner, labels, vec }
201    }
202}
203
204impl<V, L> Deref for DeleteOnDropMetric<V, L>
205where
206    V: MetricVec_,
207    L: PromLabelsExt,
208{
209    type Target = V::M;
210
211    fn deref(&self) -> &Self::Target {
212        &self.inner
213    }
214}
215
216impl<V, L> Drop for DeleteOnDropMetric<V, L>
217where
218    V: MetricVec_,
219    L: PromLabelsExt,
220{
221    fn drop(&mut self) {
222        if self.labels.remove_from_metric_vec(&self.vec).is_err() {
223            // ignore.
224        }
225    }
226}
227
228/// A [`GenericCounter`] wrapper that deletes its labels from the vec when it is dropped.
229pub type DeleteOnDropCounter<P, L> = DeleteOnDropMetric<GenericCounterVec<P>, L>;
230
231impl<P, L> Borrow<GenericCounter<P>> for DeleteOnDropCounter<P, L>
232where
233    P: Atomic,
234    L: PromLabelsExt,
235{
236    fn borrow(&self) -> &GenericCounter<P> {
237        &self.inner
238    }
239}
240
241/// A [`GenericGauge`] wrapper that deletes its labels from the vec when it is dropped.
242pub type DeleteOnDropGauge<P, L> = DeleteOnDropMetric<GenericGaugeVec<P>, L>;
243
244impl<P, L> Borrow<GenericGauge<P>> for DeleteOnDropGauge<P, L>
245where
246    P: Atomic,
247    L: PromLabelsExt,
248{
249    fn borrow(&self) -> &GenericGauge<P> {
250        &self.inner
251    }
252}
253
254/// A [`Histogram`] wrapper that deletes its labels from the vec when it is dropped.
255pub type DeleteOnDropHistogram<L> = DeleteOnDropMetric<HistogramVec, L>;
256
257impl<L> Borrow<Histogram> for DeleteOnDropHistogram<L>
258where
259    L: PromLabelsExt,
260{
261    fn borrow(&self) -> &Histogram {
262        &self.inner
263    }
264}
265
266#[cfg(test)]
267mod test {
268    use prometheus::IntGaugeVec;
269    use prometheus::core::{AtomicI64, AtomicU64};
270
271    use crate::metric;
272    use crate::metrics::{IntCounterVec, MetricsRegistry};
273
274    use super::*;
275
276    #[crate::test]
277    fn dropping_counters() {
278        let reg = MetricsRegistry::new();
279        let vec: IntCounterVec = reg.register(metric!(
280            name: "test_metric",
281            help: "a test metric",
282            var_labels: ["dimension"]));
283
284        let dims: &[&str] = &["one"];
285        let metric_1 = vec.get_delete_on_drop_metric(dims);
286        metric_1.inc();
287
288        let metrics = reg.gather();
289        assert_eq!(metrics.len(), 1);
290        let reported_vec = &metrics[0];
291        assert_eq!(reported_vec.get_name(), "test_metric");
292        let dims = reported_vec.get_metric();
293        assert_eq!(dims.len(), 1);
294        assert_eq!(dims[0].get_label()[0].get_value(), "one");
295
296        drop(metric_1);
297        let metrics = reg.gather();
298        assert_eq!(metrics.len(), 0);
299
300        let string_labels: Vec<String> = ["owned"].iter().map(ToString::to_string).collect();
301        struct Ownership {
302            counter: DeleteOnDropCounter<AtomicU64, Vec<String>>,
303        }
304        let metric_owned = Ownership {
305            counter: vec.get_delete_on_drop_metric(string_labels),
306        };
307        metric_owned.counter.inc();
308
309        let metrics = reg.gather();
310        assert_eq!(metrics.len(), 1);
311        let reported_vec = &metrics[0];
312        assert_eq!(reported_vec.get_name(), "test_metric");
313        let dims = reported_vec.get_metric();
314        assert_eq!(dims.len(), 1);
315        assert_eq!(dims[0].get_label()[0].get_value(), "owned");
316
317        drop(metric_owned);
318        let metrics = reg.gather();
319        assert_eq!(metrics.len(), 0);
320    }
321
322    #[crate::test]
323    fn dropping_gauges() {
324        let reg = MetricsRegistry::new();
325        let vec: IntGaugeVec = reg.register(metric!(
326            name: "test_metric",
327            help: "a test metric",
328            var_labels: ["dimension"]));
329
330        let dims: &[&str] = &["one"];
331        let metric_1 = vec.get_delete_on_drop_metric(dims);
332        metric_1.set(666);
333
334        let metrics = reg.gather();
335        assert_eq!(metrics.len(), 1);
336        let reported_vec = &metrics[0];
337        assert_eq!(reported_vec.get_name(), "test_metric");
338        let dims = reported_vec.get_metric();
339        assert_eq!(dims.len(), 1);
340        assert_eq!(dims[0].get_label()[0].get_value(), "one");
341
342        drop(metric_1);
343        let metrics = reg.gather();
344        assert_eq!(metrics.len(), 0);
345
346        let string_labels: Vec<String> = ["owned"].iter().map(ToString::to_string).collect();
347        struct Ownership {
348            gauge: DeleteOnDropGauge<AtomicI64, Vec<String>>,
349        }
350        let metric_owned = Ownership {
351            gauge: vec.get_delete_on_drop_metric(string_labels),
352        };
353        metric_owned.gauge.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.get_name(), "test_metric");
359        let dims = reported_vec.get_metric();
360        assert_eq!(dims.len(), 1);
361        assert_eq!(dims[0].get_label()[0].get_value(), "owned");
362
363        drop(metric_owned);
364        let metrics = reg.gather();
365        assert_eq!(metrics.len(), 0);
366    }
367}
mz_ore/metrics/delete_on_drop.rs

mz_ore/metrics/
delete_on_drop.rs
