risingwave_common_metrics/
guarded_metrics.rs

1// Copyright 2024 RisingWave Labs
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 at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::any::type_name;
16use std::collections::{HashMap, HashSet};
17use std::fmt::{Debug, Formatter};
18use std::ops::Deref;
19use std::sync::{Arc, LazyLock};
20
21use parking_lot::Mutex;
22use prometheus::core::{
23    Atomic, AtomicF64, AtomicI64, AtomicU64, Collector, Desc, GenericCounter, GenericLocalCounter,
24    MetricVec, MetricVecBuilder,
25};
26use prometheus::local::{LocalHistogram, LocalIntCounter};
27use prometheus::proto::MetricFamily;
28use prometheus::{Gauge, Histogram, IntCounter, IntGauge};
29use thiserror_ext::AsReport;
30use tracing::warn;
31
32#[macro_export]
33macro_rules! register_guarded_histogram_vec_with_registry {
34    ($NAME:expr, $HELP:expr, $LABELS_NAMES:expr, $REGISTRY:expr $(,)?) => {{
35        $crate::register_guarded_histogram_vec_with_registry! {
36            {prometheus::histogram_opts!($NAME, $HELP)},
37            $LABELS_NAMES,
38            $REGISTRY
39        }
40    }};
41    ($NAME:expr, $HELP:expr, $LABELS_NAMES:expr, $BUCKETS:expr, $REGISTRY:expr $(,)?) => {{
42        $crate::register_guarded_histogram_vec_with_registry! {
43            {prometheus::histogram_opts!($NAME, $HELP, $BUCKETS)},
44            $LABELS_NAMES,
45            $REGISTRY
46        }
47    }};
48    ($HOPTS:expr, $LABELS_NAMES:expr, $REGISTRY:expr $(,)?) => {{
49        let inner = prometheus::HistogramVec::new($HOPTS, $LABELS_NAMES);
50        inner.and_then(|inner| {
51            let inner = $crate::__extract_histogram_builder(inner);
52            let label_guarded = $crate::LabelGuardedHistogramVec::new(inner, { $LABELS_NAMES });
53            let result = ($REGISTRY).register(Box::new(label_guarded.clone()));
54            result.map(move |()| label_guarded)
55        })
56    }};
57}
58
59#[macro_export]
60macro_rules! register_guarded_gauge_vec_with_registry {
61    ($NAME:expr, $HELP:expr, $LABELS_NAMES:expr, $REGISTRY:expr $(,)?) => {{
62        let inner = prometheus::GaugeVec::new(prometheus::opts!($NAME, $HELP), $LABELS_NAMES);
63        inner.and_then(|inner| {
64            let inner = $crate::__extract_gauge_builder(inner);
65            let label_guarded = $crate::LabelGuardedGaugeVec::new(inner, { $LABELS_NAMES });
66            let result = ($REGISTRY).register(Box::new(label_guarded.clone()));
67            result.map(move |()| label_guarded)
68        })
69    }};
70}
71
72#[macro_export]
73macro_rules! register_guarded_int_gauge_vec_with_registry {
74    ($NAME:expr, $HELP:expr, $LABELS_NAMES:expr, $REGISTRY:expr $(,)?) => {{
75        let inner = prometheus::IntGaugeVec::new(prometheus::opts!($NAME, $HELP), $LABELS_NAMES);
76        inner.and_then(|inner| {
77            let inner = $crate::__extract_gauge_builder(inner);
78            let label_guarded = $crate::LabelGuardedIntGaugeVec::new(inner, { $LABELS_NAMES });
79            let result = ($REGISTRY).register(Box::new(label_guarded.clone()));
80            result.map(move |()| label_guarded)
81        })
82    }};
83}
84
85#[macro_export]
86macro_rules! register_guarded_uint_gauge_vec_with_registry {
87    ($NAME:expr, $HELP:expr, $LABELS_NAMES:expr, $REGISTRY:expr $(,)?) => {{
88        let inner = prometheus::core::GenericGaugeVec::<prometheus::core::AtomicU64>::new(
89            prometheus::opts!($NAME, $HELP),
90            $LABELS_NAMES,
91        );
92        inner.and_then(|inner| {
93            let inner = $crate::__extract_gauge_builder(inner);
94            let label_guarded = $crate::LabelGuardedUintGaugeVec::new(inner, { $LABELS_NAMES });
95            let result = ($REGISTRY).register(Box::new(label_guarded.clone()));
96            result.map(move |()| label_guarded)
97        })
98    }};
99}
100
101#[macro_export]
102macro_rules! register_guarded_int_counter_vec_with_registry {
103    ($NAME:expr, $HELP:expr, $LABELS_NAMES:expr, $REGISTRY:expr $(,)?) => {{
104        let inner = prometheus::IntCounterVec::new(prometheus::opts!($NAME, $HELP), $LABELS_NAMES);
105        inner.and_then(|inner| {
106            let inner = $crate::__extract_counter_builder(inner);
107            let label_guarded = $crate::LabelGuardedIntCounterVec::new(inner, { $LABELS_NAMES });
108            let result = ($REGISTRY).register(Box::new(label_guarded.clone()));
109            result.map(move |()| label_guarded)
110        })
111    }};
112}
113
114// put TAITs in a separate module to avoid "non-defining opaque type use in defining scope"
115mod tait {
116    use prometheus::core::{
117        Atomic, GenericCounter, GenericCounterVec, GenericGauge, GenericGaugeVec, MetricVec,
118        MetricVecBuilder,
119    };
120    use prometheus::{Histogram, HistogramVec};
121
122    pub type VecBuilderOfCounter<P: Atomic> = impl MetricVecBuilder<M = GenericCounter<P>>;
123    pub type VecBuilderOfGauge<P: Atomic> = impl MetricVecBuilder<M = GenericGauge<P>>;
124    pub type VecBuilderOfHistogram = impl MetricVecBuilder<M = Histogram>;
125
126    #[define_opaque(VecBuilderOfCounter)]
127    pub fn __extract_counter_builder<P: Atomic>(
128        vec: GenericCounterVec<P>,
129    ) -> MetricVec<VecBuilderOfCounter<P>> {
130        vec
131    }
132
133    #[define_opaque(VecBuilderOfGauge)]
134    pub fn __extract_gauge_builder<P: Atomic>(
135        vec: GenericGaugeVec<P>,
136    ) -> MetricVec<VecBuilderOfGauge<P>> {
137        vec
138    }
139
140    #[define_opaque(VecBuilderOfHistogram)]
141    pub fn __extract_histogram_builder(vec: HistogramVec) -> MetricVec<VecBuilderOfHistogram> {
142        vec
143    }
144}
145pub use tait::*;
146
147use crate::UintGauge;
148
149pub type LabelGuardedHistogramVec = LabelGuardedMetricVec<VecBuilderOfHistogram>;
150pub type LabelGuardedIntCounterVec = LabelGuardedMetricVec<VecBuilderOfCounter<AtomicU64>>;
151pub type LabelGuardedIntGaugeVec = LabelGuardedMetricVec<VecBuilderOfGauge<AtomicI64>>;
152pub type LabelGuardedUintGaugeVec = LabelGuardedMetricVec<VecBuilderOfGauge<AtomicU64>>;
153pub type LabelGuardedGaugeVec = LabelGuardedMetricVec<VecBuilderOfGauge<AtomicF64>>;
154
155pub type LabelGuardedHistogram = LabelGuardedMetric<Histogram>;
156pub type LabelGuardedIntCounter = LabelGuardedMetric<IntCounter>;
157pub type LabelGuardedIntGauge = LabelGuardedMetric<IntGauge>;
158pub type LabelGuardedUintGauge = LabelGuardedMetric<UintGauge>;
159pub type LabelGuardedGauge = LabelGuardedMetric<Gauge>;
160
161pub type LabelGuardedLocalHistogram = LabelGuardedMetric<LocalHistogram>;
162pub type LabelGuardedLocalIntCounter = LabelGuardedMetric<LocalIntCounter>;
163
164fn gen_test_label<const N: usize>() -> [&'static str; N] {
165    const TEST_LABELS: [&str; 5] = ["test1", "test2", "test3", "test4", "test5"];
166    (0..N)
167        .map(|i| TEST_LABELS[i])
168        .collect::<Vec<_>>()
169        .try_into()
170        .unwrap()
171}
172
173#[derive(Default)]
174struct LabelGuardedMetricsInfo {
175    labeled_metrics_count: HashMap<Box<[String]>, usize>,
176    uncollected_removed_labels: HashSet<Box<[String]>>,
177}
178
179impl LabelGuardedMetricsInfo {
180    fn register_new_label<V: AsRef<str>>(mutex: &Arc<Mutex<Self>>, labels: &[V]) -> LabelGuard {
181        let mut guard = mutex.lock();
182        let label_string = labels
183            .iter()
184            .map(|label| label.as_ref().to_owned())
185            .collect::<Vec<_>>()
186            .into_boxed_slice();
187        guard.uncollected_removed_labels.remove(&label_string);
188        *guard
189            .labeled_metrics_count
190            .entry(label_string.clone())
191            .or_insert(0) += 1;
192        LabelGuard {
193            labels: label_string,
194            info: mutex.clone(),
195        }
196    }
197}
198
199/// An RAII metrics vec with labels.
200///
201/// `LabelGuardedMetricVec` enhances the [`MetricVec`] to ensure the set of labels to be
202/// correctly removed from the Prometheus client once being dropped. This is useful for metrics
203/// that are associated with an object that can be dropped, such as streaming jobs, fragments,
204/// actors, batch tasks, etc.
205///
206/// When a set labels is dropped, it will record it in the `uncollected_removed_labels` set.
207/// Once the metrics has been collected, it will finally remove the metrics of the labels.
208///
209/// See also [`LabelGuardedMetricsInfo`] and [`LabelGuard::drop`].
210///
211/// # Arguments
212///
213/// * `T` - The type of the raw metrics vec.
214/// * `N` - The number of labels.
215#[derive(Clone)]
216pub struct LabelGuardedMetricVec<T: MetricVecBuilder> {
217    inner: MetricVec<T>,
218    info: Arc<Mutex<LabelGuardedMetricsInfo>>,
219    labels: Box<[&'static str]>,
220}
221
222impl<T: MetricVecBuilder> Debug for LabelGuardedMetricVec<T> {
223    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
224        f.debug_struct(format!("LabelGuardedMetricVec<{}>", type_name::<T>()).as_str())
225            .field("label", &self.labels)
226            .finish()
227    }
228}
229
230impl<T: MetricVecBuilder> Collector for LabelGuardedMetricVec<T> {
231    fn desc(&self) -> Vec<&Desc> {
232        self.inner.desc()
233    }
234
235    fn collect(&self) -> Vec<MetricFamily> {
236        let mut guard = self.info.lock();
237        let ret = self.inner.collect();
238        for labels in guard.uncollected_removed_labels.drain() {
239            if let Err(e) = self.inner.remove_label_values(&labels) {
240                warn!(
241                    error = %e.as_report(),
242                    "err when delete metrics of {:?} of labels {:?}",
243                    self.inner.desc().first().expect("should have desc").fq_name,
244                    self.labels,
245                );
246            }
247        }
248        ret
249    }
250}
251
252impl<T: MetricVecBuilder> LabelGuardedMetricVec<T> {
253    pub fn new(inner: MetricVec<T>, labels: &[&'static str]) -> Self {
254        Self {
255            inner,
256            info: Default::default(),
257            labels: labels.to_vec().into_boxed_slice(),
258        }
259    }
260
261    /// This is similar to the `with_label_values` of the raw metrics vec.
262    /// We need to pay special attention that, unless for some special purpose,
263    /// we should not drop the returned `LabelGuardedMetric` immediately after
264    /// using it, such as `metrics.with_guarded_label_values(...).inc();`,
265    /// because after dropped the label will be regarded as not used any more,
266    /// and the internal raw metrics will be removed and reset.
267    ///
268    /// Instead, we should store the returned `LabelGuardedMetric` in a scope with longer
269    /// lifetime so that the labels can be regarded as being used in its whole life scope.
270    /// This is also the recommended way to use the raw metrics vec.
271    ///
272    /// A short-lived guard is safe for a gauge that is always updated with `set`: collection
273    /// observes the value before removing a dropped label, and the next `set` replaces the value
274    /// completely. Counters and histograms must retain the guard for their intended lifetime so
275    /// that values are not reset between collections.
276    pub fn with_guarded_label_values<V: AsRef<str> + std::fmt::Debug>(
277        &self,
278        labels: &[V],
279    ) -> LabelGuardedMetric<T::M> {
280        let guard = LabelGuardedMetricsInfo::register_new_label(&self.info, labels);
281        let inner = self.inner.with_label_values(labels);
282        LabelGuardedMetric {
283            inner,
284            _guard: Arc::new(guard),
285        }
286    }
287
288    pub fn with_test_label<const N: usize>(&self) -> LabelGuardedMetric<T::M> {
289        let labels = gen_test_label::<N>();
290        self.with_guarded_label_values(&labels)
291    }
292}
293
294impl LabelGuardedIntCounterVec {
295    pub fn test_int_counter_vec<const N: usize>() -> Self {
296        let registry = prometheus::Registry::new();
297        let labels = gen_test_label::<N>();
298        register_guarded_int_counter_vec_with_registry!("test", "test", &labels, &registry).unwrap()
299    }
300}
301
302impl LabelGuardedIntGaugeVec {
303    pub fn test_int_gauge_vec<const N: usize>() -> Self {
304        let registry = prometheus::Registry::new();
305        let labels = gen_test_label::<N>();
306        register_guarded_int_gauge_vec_with_registry!("test", "test", &labels, &registry).unwrap()
307    }
308}
309
310impl LabelGuardedGaugeVec {
311    pub fn test_gauge_vec<const N: usize>() -> Self {
312        let registry = prometheus::Registry::new();
313        let labels = gen_test_label::<N>();
314        register_guarded_gauge_vec_with_registry!("test", "test", &labels, &registry).unwrap()
315    }
316}
317
318impl LabelGuardedHistogramVec {
319    pub fn test_histogram_vec<const N: usize>() -> Self {
320        let registry = prometheus::Registry::new();
321        let labels = gen_test_label::<N>();
322        register_guarded_histogram_vec_with_registry!("test", "test", &labels, &registry).unwrap()
323    }
324}
325
326#[derive(Clone)]
327struct LabelGuard {
328    labels: Box<[String]>,
329    info: Arc<Mutex<LabelGuardedMetricsInfo>>,
330}
331
332impl Drop for LabelGuard {
333    fn drop(&mut self) {
334        let mut guard = self.info.lock();
335        let count = guard.labeled_metrics_count.get_mut(&self.labels).expect(
336            "should exist because the current existing dropping one means the count is not zero",
337        );
338        *count -= 1;
339        if *count == 0 {
340            guard
341                .labeled_metrics_count
342                .remove(&self.labels)
343                .expect("should exist");
344            guard.uncollected_removed_labels.insert(self.labels.clone());
345        }
346    }
347}
348
349#[derive(Clone)]
350pub struct LabelGuardedMetric<T> {
351    inner: T,
352    _guard: Arc<LabelGuard>,
353}
354
355impl<T> Debug for LabelGuardedMetric<T> {
356    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
357        f.debug_struct("LabelGuardedMetric").finish()
358    }
359}
360
361impl<T> Deref for LabelGuardedMetric<T> {
362    type Target = T;
363
364    fn deref(&self) -> &Self::Target {
365        &self.inner
366    }
367}
368
369impl LabelGuardedHistogram {
370    pub fn test_histogram<const N: usize>() -> Self {
371        LabelGuardedHistogramVec::test_histogram_vec::<N>().with_test_label::<N>()
372    }
373}
374
375impl LabelGuardedIntCounter {
376    pub fn test_int_counter<const N: usize>() -> Self {
377        LabelGuardedIntCounterVec::test_int_counter_vec::<N>().with_test_label::<N>()
378    }
379}
380
381impl LabelGuardedIntGauge {
382    pub fn test_int_gauge<const N: usize>() -> Self {
383        LabelGuardedIntGaugeVec::test_int_gauge_vec::<N>().with_test_label::<N>()
384    }
385}
386
387impl LabelGuardedGauge {
388    pub fn test_gauge<const N: usize>() -> Self {
389        LabelGuardedGaugeVec::test_gauge_vec::<N>().with_test_label::<N>()
390    }
391}
392
393pub type LazyLabelGuardedMetrics<T: MetricVecBuilder> =
394    LazyLock<LabelGuardedMetric<T::M>, impl FnOnce() -> LabelGuardedMetric<T::M>>;
395
396impl<T: MetricVecBuilder> LabelGuardedMetricVec<T> {
397    #[define_opaque(LazyLabelGuardedMetrics)]
398    pub fn lazy_guarded_metrics(self, labels: Vec<String>) -> LazyLabelGuardedMetrics<T> {
399        LazyLock::new(move || self.with_guarded_label_values(&labels))
400    }
401}
402
403pub type LazyLabelGuardedHistogram = LazyLabelGuardedMetrics<VecBuilderOfHistogram>;
404pub type LazyLabelGuardedIntCounter = LazyLabelGuardedMetrics<VecBuilderOfCounter<AtomicU64>>;
405pub type LazyLabelGuardedIntGauge = LazyLabelGuardedMetrics<VecBuilderOfGauge<AtomicI64>>;
406pub type LazyLabelGuardedUintGauge = LazyLabelGuardedMetrics<VecBuilderOfGauge<AtomicU64>>;
407pub type LazyLabelGuardedGauge = LazyLabelGuardedMetrics<VecBuilderOfGauge<AtomicF64>>;
408
409pub trait MetricWithLocal {
410    type Local;
411    fn local(&self) -> Self::Local;
412}
413
414impl MetricWithLocal for Histogram {
415    type Local = LocalHistogram;
416
417    fn local(&self) -> Self::Local {
418        self.local()
419    }
420}
421
422impl<P: Atomic> MetricWithLocal for GenericCounter<P> {
423    type Local = GenericLocalCounter<P>;
424
425    fn local(&self) -> Self::Local {
426        self.local()
427    }
428}
429
430impl<T: MetricWithLocal> LabelGuardedMetric<T> {
431    pub fn local(&self) -> LabelGuardedMetric<T::Local> {
432        LabelGuardedMetric {
433            inner: self.inner.local(),
434            _guard: self._guard.clone(),
435        }
436    }
437}
438
439#[cfg(test)]
440mod tests {
441    use prometheus::core::Collector;
442
443    use crate::{LabelGuardedIntCounterVec, LabelGuardedIntGaugeVec};
444
445    #[test]
446    fn test_label_guarded_metrics_drop() {
447        let vec = LabelGuardedIntCounterVec::test_int_counter_vec::<3>();
448        let m1_1 = vec.with_guarded_label_values(&["1", "2", "3"]);
449        assert_eq!(1, vec.collect().pop().unwrap().get_metric().len());
450        let m1_2 = vec.with_guarded_label_values(&["1", "2", "3"]);
451        let m1_3 = m1_2.clone();
452        assert_eq!(1, vec.collect().pop().unwrap().get_metric().len());
453        let m2 = vec.with_guarded_label_values(&["2", "2", "3"]);
454        assert_eq!(2, vec.collect().pop().unwrap().get_metric().len());
455        drop(m1_3);
456        assert_eq!(2, vec.collect().pop().unwrap().get_metric().len());
457        assert_eq!(2, vec.collect().pop().unwrap().get_metric().len());
458        drop(m2);
459        assert_eq!(2, vec.collect().pop().unwrap().get_metric().len());
460        assert_eq!(1, vec.collect().pop().unwrap().get_metric().len());
461        drop(m1_1);
462        assert_eq!(1, vec.collect().pop().unwrap().get_metric().len());
463        assert_eq!(1, vec.collect().pop().unwrap().get_metric().len());
464        drop(m1_2);
465        assert_eq!(1, vec.collect().pop().unwrap().get_metric().len());
466        assert_eq!(0, vec.collect().pop().unwrap().get_metric().len());
467    }
468
469    #[test]
470    fn test_short_lived_guarded_gauge() {
471        let vec = LabelGuardedIntGaugeVec::test_int_gauge_vec::<1>();
472
473        vec.with_guarded_label_values(&["1"]).set(7);
474        vec.with_guarded_label_values(&["1"]).set(11);
475        let collected = vec.collect().pop().unwrap();
476        assert_eq!(
477            11.0,
478            collected.get_metric()[0]
479                .get_gauge()
480                .as_ref()
481                .unwrap()
482                .value()
483        );
484        assert!(vec.collect().pop().unwrap().get_metric().is_empty());
485
486        vec.with_guarded_label_values(&["1"]).set(13);
487        let collected = vec.collect().pop().unwrap();
488        assert_eq!(
489            13.0,
490            collected.get_metric()[0]
491                .get_gauge()
492                .as_ref()
493                .unwrap()
494                .value()
495        );
496    }
497}
risingwave_common_metrics/guarded_metrics.rs

risingwave_common_metrics/
guarded_metrics.rs
