mz_ore/metrics/
deregister_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//! Labeled metrics that stay registered while any handle for their labels is alive.
17//!
18//! Each [`DeleteOnDropMetric`](super::DeleteOnDropMetric) obtained from a vector removes its
19//! labels when it drops, even while another handle for the same labels, obtained separately, is
20//! still recording. The survivor then records into a series that is no longer exported. This
21//! happens when every timely worker creates a metric but only one keeps using it, and when a
22//! restarted dataflow creates its metrics before the previous incarnation drops its own.
23//!
24//! [`DeregisterOnDropVec`] counts handles per label set across all calls, so the series is
25//! removed only when the last handle for its labels drops.
26
27use std::collections::BTreeMap;
28use std::fmt::{self, Debug, Formatter};
29use std::ops::Deref;
30use std::sync::{Arc, Mutex, Weak};
31
32use prometheus::HistogramVec;
33use prometheus::core::{Collector, Desc, GenericCounterVec, GenericGaugeVec};
34use prometheus::proto::MetricFamily;
35
36use crate::metrics::{MakeCollector, MakeCollectorOpts, MetricVec_};
37
38/// The live registration of each label set, keyed by label values.
39type Registrations<V> = Arc<Mutex<BTreeMap<Vec<String>, Weak<Registration<V>>>>>;
40
41/// A metric vector whose handles share one registration per label set.
42///
43/// Construct it with [`MetricsRegistry::register`](super::MetricsRegistry::register).
44#[derive(Clone)]
45pub struct DeregisterOnDropVec<V: MetricVec_> {
46    inner: V,
47    registrations: Registrations<V>,
48}
49
50impl<V: MetricVec_ + Debug> Debug for DeregisterOnDropVec<V> {
51    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
52        self.inner.fmt(f)
53    }
54}
55
56impl<V: MetricVec_ + Collector> Collector for DeregisterOnDropVec<V> {
57    fn desc(&self) -> Vec<&Desc> {
58        self.inner.desc()
59    }
60
61    fn collect(&self) -> Vec<MetricFamily> {
62        self.inner.collect()
63    }
64}
65
66impl<V: MetricVec_ + MakeCollector> MakeCollector for DeregisterOnDropVec<V> {
67    fn make_collector(opts: MakeCollectorOpts) -> Self {
68        DeregisterOnDropVec {
69            inner: V::make_collector(opts),
70            registrations: Default::default(),
71        }
72    }
73}
74
75impl<V: MetricVec_ + Clone> DeregisterOnDropVec<V> {
76    /// Returns the metric for `labels`, which stays registered until every handle returned for
77    /// the same `labels` has dropped.
78    ///
79    /// Panics if the number of labels does not match the vector's label names.
80    pub fn get_deregister_on_drop_metric(&self, labels: Vec<String>) -> DeregisterOnDropMetric<V> {
81        let mut registrations = self.registrations.lock().expect("lock poisoned");
82        // Fetching while holding the lock keeps a concurrent `Registration::drop` from removing
83        // the series between this fetch and the registration below.
84        let inner = {
85            let values: Vec<&str> = labels.iter().map(String::as_str).collect();
86            self.inner
87                .get_metric_with_label_values(&values)
88                .expect("retrieving a metric by label values")
89        };
90        let registration = match registrations.get(&labels).and_then(Weak::upgrade) {
91            Some(registration) => registration,
92            None => {
93                let registration = Arc::new(Registration {
94                    labels: labels.clone(),
95                    vec: self.clone(),
96                });
97                registrations.insert(labels, Arc::downgrade(&registration));
98                registration
99            }
100        };
101        DeregisterOnDropMetric {
102            inner,
103            _registration: registration,
104        }
105    }
106}
107
108/// Removes a label set from its vector when the last handle for it drops.
109struct Registration<V: MetricVec_> {
110    labels: Vec<String>,
111    vec: DeregisterOnDropVec<V>,
112}
113
114impl<V: MetricVec_> Debug for Registration<V> {
115    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
116        f.debug_struct("Registration")
117            .field("labels", &self.labels)
118            .finish_non_exhaustive()
119    }
120}
121
122impl<V: MetricVec_> Drop for Registration<V> {
123    fn drop(&mut self) {
124        let mut registrations = self.vec.registrations.lock().expect("lock poisoned");
125        // A handle obtained after the last one for these labels dropped, but before this ran,
126        // re-registered the labels and still records into the series.
127        if registrations
128            .get(&self.labels)
129            .is_some_and(|live| live.strong_count() > 0)
130        {
131            return;
132        }
133        registrations.remove(&self.labels);
134        let values: Vec<&str> = self.labels.iter().map(String::as_str).collect();
135        // Ignore the error for a series an earlier registration of these labels already removed.
136        let _ = self.vec.inner.remove_label_values(&values);
137    }
138}
139
140/// A labeled metric that stays registered while any handle for its labels is alive.
141#[derive(Debug, Clone)]
142pub struct DeregisterOnDropMetric<V: MetricVec_> {
143    inner: V::M,
144    _registration: Arc<Registration<V>>,
145}
146
147impl<V: MetricVec_> Deref for DeregisterOnDropMetric<V> {
148    type Target = V::M;
149
150    fn deref(&self) -> &Self::Target {
151        &self.inner
152    }
153}
154
155/// A [`GenericGauge`](prometheus::core::GenericGauge) that stays registered while any handle for
156/// its labels is alive.
157pub type DeregisterOnDropGauge<P> = DeregisterOnDropMetric<GenericGaugeVec<P>>;
158
159/// A [`GenericCounter`](prometheus::core::GenericCounter) that stays registered while any handle
160/// for its labels is alive.
161pub type DeregisterOnDropCounter<P> = DeregisterOnDropMetric<GenericCounterVec<P>>;
162
163/// A [`Histogram`](prometheus::Histogram) that stays registered while any handle for its labels
164/// is alive.
165pub type DeregisterOnDropHistogram = DeregisterOnDropMetric<HistogramVec>;
166
167#[cfg(test)]
168mod test {
169    use prometheus::IntCounterVec;
170
171    use crate::metric;
172    use crate::metrics::MetricsRegistry;
173
174    use super::*;
175
176    fn labels(value: &str) -> Vec<String> {
177        vec![value.to_string()]
178    }
179
180    #[crate::test]
181    fn independent_handles_share_registration() {
182        let reg = MetricsRegistry::new();
183        let vec: DeregisterOnDropVec<IntCounterVec> = reg.register(metric!(
184            name: "test_metric",
185            help: "a test metric",
186            var_labels: ["dimension"]));
187
188        let first = vec.get_deregister_on_drop_metric(labels("one"));
189        let second = vec.get_deregister_on_drop_metric(labels("one"));
190        first.inc();
191        second.inc();
192
193        drop(first);
194        let gathered = reg.gather();
195        assert_eq!(
196            gathered.len(),
197            1,
198            "series must survive while a handle is alive"
199        );
200        assert_eq!(gathered[0].get_metric()[0].get_counter().value(), 2.0);
201
202        second.inc();
203        let gathered = reg.gather();
204        assert_eq!(gathered[0].get_metric()[0].get_counter().value(), 3.0);
205
206        drop(second);
207        assert_eq!(
208            reg.gather().len(),
209            0,
210            "last handle must deregister the series"
211        );
212        assert!(vec.registrations.lock().expect("lock poisoned").is_empty());
213    }
214
215    #[crate::test]
216    fn label_sets_are_independent() {
217        let reg = MetricsRegistry::new();
218        let vec: DeregisterOnDropVec<IntCounterVec> = reg.register(metric!(
219            name: "test_metric",
220            help: "a test metric",
221            var_labels: ["dimension"]));
222
223        let one = vec.get_deregister_on_drop_metric(labels("one"));
224        let two = vec.get_deregister_on_drop_metric(labels("two"));
225        one.inc();
226        two.inc();
227
228        drop(one);
229        let gathered = reg.gather();
230        let series = gathered[0].get_metric();
231        assert_eq!(series.len(), 1);
232        assert_eq!(series[0].get_label()[0].value(), "two");
233
234        drop(two);
235        assert_eq!(reg.gather().len(), 0);
236    }
237
238    #[crate::test]
239    fn reregistration_after_last_drop_starts_fresh() {
240        let reg = MetricsRegistry::new();
241        let vec: DeregisterOnDropVec<IntCounterVec> = reg.register(metric!(
242            name: "test_metric",
243            help: "a test metric",
244            var_labels: ["dimension"]));
245
246        let first = vec.get_deregister_on_drop_metric(labels("one"));
247        first.inc_by(5);
248        drop(first);
249
250        let second = vec.get_deregister_on_drop_metric(labels("one"));
251        second.inc();
252        let gathered = reg.gather();
253        assert_eq!(gathered[0].get_metric()[0].get_counter().value(), 1.0);
254    }
255}
mz_ore/metrics/deregister_on_drop.rs

mz_ore/metrics/
deregister_on_drop.rs
