risingwave_common_metrics/
relabeled_metric.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::collections::HashMap;
16use std::sync::atomic::{AtomicU64, Ordering};
17use std::sync::{Arc, Weak};
18
19use parking_lot::Mutex;
20use prometheus::core::{Collector, MetricVec, MetricVecBuilder};
21use prometheus::{HistogramVec, IntCounterVec};
22
23use crate::{
24    LabelGuardedHistogramVec, LabelGuardedIntCounterVec, LabelGuardedIntGaugeVec,
25    LabelGuardedMetric, LabelGuardedMetricVec, LazyLabelGuardedMetrics, MetricLevel,
26};
27
28/// For all `Relabeled*Vec` below,
29/// - when `metric_level` <= `relabel_threshold`, they behave exactly the same as their inner
30///   metric.
31/// - when `metric_level` > `relabel_threshold`, the first `relabel_num` labels are rewrite to "" when
32///   calling `with_label_values`. That's means the metric vec is aggregated into a single metric.
33///
34/// These wrapper classes add a `metric_level` field to corresponding metric.
35/// We could have use one single struct to represent all `MetricVec<T: MetricVecBuilder>`, rather
36/// than specializing them one by one. However, that's undoable because prometheus crate doesn't
37/// export `MetricVecBuilder` implementation like `HistogramVecBuilder`.
38///
39/// ## Note
40///
41/// CAUTION! Relabelling might cause expected result!
42///
43/// For counters (including histogram because it uses counters internally), it's usually natural
44/// to sum up the count from multiple labels.
45///
46/// For the rest (such as Gauge), the semantics becomes "any/last of the recorded value". Please be
47/// cautious.
48#[derive(Clone, Debug)]
49pub struct RelabeledMetricVec<M> {
50    relabel_threshold: MetricLevel,
51    metric_level: MetricLevel,
52    metric: M,
53
54    /// The first `relabel_num` labels will be relabeled to empty string
55    ///
56    /// For example, if `relabel_num` is 1, and the input labels are `["actor_id",
57    /// "fragment_id", "table_id"]`, when threshold is reached, the label values will be
58    /// `["", "<original_fragment_id>", "<original_table_id>"]`.
59    relabel_num: usize,
60}
61
62impl<M> RelabeledMetricVec<M> {
63    pub fn with_metric_level(
64        metric_level: MetricLevel,
65        metric: M,
66        relabel_threshold: MetricLevel,
67    ) -> Self {
68        Self {
69            relabel_threshold,
70            metric_level,
71            metric,
72            relabel_num: usize::MAX,
73        }
74    }
75
76    pub fn with_metric_level_relabel_n(
77        metric_level: MetricLevel,
78        metric: M,
79        relabel_threshold: MetricLevel,
80        relabel_num: usize,
81    ) -> Self {
82        Self {
83            relabel_threshold,
84            metric_level,
85            metric,
86            relabel_num,
87        }
88    }
89
90    fn relabel_impl<V: AsRef<str> + std::fmt::Debug>(&self, vals: &[V]) -> Option<Vec<String>> {
91        if self.metric_level > self.relabel_threshold {
92            // relabel first n labels to empty string
93            let mut relabeled_vals = vals
94                .iter()
95                .map(|v| v.as_ref().to_owned())
96                .collect::<Vec<_>>();
97            for label in relabeled_vals.iter_mut().take(self.relabel_num) {
98                *label = String::new();
99            }
100            Some(relabeled_vals)
101        } else {
102            None
103        }
104    }
105}
106
107#[easy_ext::ext(MetricVecRelabelExt)]
108impl<M> M
109where
110    M: Sized,
111{
112    /// Equivalent to [`RelabeledMetricVec::with_metric_level`].
113    pub fn relabel(
114        self,
115        metric_level: MetricLevel,
116        relabel_threshold: MetricLevel,
117    ) -> RelabeledMetricVec<M> {
118        RelabeledMetricVec::with_metric_level(metric_level, self, relabel_threshold)
119    }
120
121    /// Equivalent to [`RelabeledMetricVec::with_metric_level_relabel_n`].
122    pub fn relabel_n(
123        self,
124        metric_level: MetricLevel,
125        relabel_threshold: MetricLevel,
126        relabel_num: usize,
127    ) -> RelabeledMetricVec<M> {
128        RelabeledMetricVec::with_metric_level_relabel_n(
129            metric_level,
130            self,
131            relabel_threshold,
132            relabel_num,
133        )
134    }
135
136    /// Equivalent to [`RelabeledMetricVec::with_metric_level_relabel_n`] with `metric_level` set to
137    /// `MetricLevel::Debug` and `relabel_num` set to 1.
138    pub fn relabel_debug_1(self, relabel_threshold: MetricLevel) -> RelabeledMetricVec<M> {
139        RelabeledMetricVec::with_metric_level_relabel_n(
140            MetricLevel::Debug,
141            self,
142            relabel_threshold,
143            1,
144        )
145    }
146}
147
148impl<T: MetricVecBuilder> RelabeledMetricVec<MetricVec<T>> {
149    pub fn with_label_values<V: AsRef<str> + std::fmt::Debug>(&self, vals: &[V]) -> T::M {
150        if let Some(relabeled_vals) = self.relabel_impl(vals) {
151            return self.metric.with_label_values(&relabeled_vals);
152        }
153        self.metric.with_label_values(vals)
154    }
155
156    pub fn reset(&self) {
157        self.metric.reset();
158    }
159}
160
161impl<T: MetricVecBuilder> RelabeledMetricVec<LabelGuardedMetricVec<T>> {
162    pub fn with_guarded_label_values<V: AsRef<str> + std::fmt::Debug>(
163        &self,
164        vals: &[V],
165    ) -> LabelGuardedMetric<T::M> {
166        if let Some(relabeled_vals) = self.relabel_impl(vals) {
167            return self.metric.with_guarded_label_values(&relabeled_vals);
168        }
169        self.metric.with_guarded_label_values(vals)
170    }
171
172    pub fn lazy_guarded_metrics(&self, labels: Vec<String>) -> LazyLabelGuardedMetrics<T> {
173        if let Some(relabeled_vals) = self.relabel_impl(labels.as_slice()) {
174            return self.metric.clone().lazy_guarded_metrics(relabeled_vals);
175        }
176        self.metric.clone().lazy_guarded_metrics(labels)
177    }
178}
179
180impl<T: Collector> Collector for RelabeledMetricVec<T> {
181    fn desc(&self) -> Vec<&prometheus::core::Desc> {
182        self.metric.desc()
183    }
184
185    fn collect(&self) -> Vec<prometheus::proto::MetricFamily> {
186        self.metric.collect()
187    }
188}
189
190pub type RelabeledCounterVec = RelabeledMetricVec<IntCounterVec>;
191pub type RelabeledHistogramVec = RelabeledMetricVec<HistogramVec>;
192
193pub type RelabeledGuardedHistogramVec = RelabeledMetricVec<LabelGuardedHistogramVec>;
194pub type RelabeledGuardedIntCounterVec = RelabeledMetricVec<LabelGuardedIntCounterVec>;
195
196/// CAUTION! Relabelling a Gauge might cause expected result!
197///
198/// See [`RelabeledMetricVec`] for details.
199pub type RelabeledGuardedIntGaugeVec = RelabeledMetricVec<LabelGuardedIntGaugeVec>;
200
201/// Defines how live gauge values are combined after their labels are relabeled.
202#[derive(Clone, Copy, Debug, PartialEq, Eq)]
203pub enum GaugeAggregation {
204    Sum,
205    Min,
206    Max,
207}
208
209/// Extension methods for relabeling guarded integer gauges with explicit aggregation semantics.
210pub trait IntGaugeVecRelabelExt {
211    /// Relabels the first label according to the configured metric level and aggregates live
212    /// contributions whenever that effective label value is empty.
213    fn relabel_debug_1_with_aggregation(
214        self,
215        metric_level: MetricLevel,
216        aggregation: GaugeAggregation,
217    ) -> RelabeledAggregatedIntGaugeVec;
218}
219
220impl IntGaugeVecRelabelExt for LabelGuardedIntGaugeVec {
221    fn relabel_debug_1_with_aggregation(
222        self,
223        metric_level: MetricLevel,
224        aggregation: GaugeAggregation,
225    ) -> RelabeledAggregatedIntGaugeVec {
226        RelabeledAggregatedIntGaugeVec {
227            metric: RelabeledMetricVec::with_metric_level_relabel_n(
228                MetricLevel::Debug,
229                self,
230                metric_level,
231                1,
232            ),
233            aggregation,
234            groups: Arc::new(AggregatedGaugeGroups::default()),
235        }
236    }
237}
238
239/// A guarded integer gauge vec that preserves gauge semantics after actor labels are relabeled.
240#[derive(Clone, Debug)]
241pub struct RelabeledAggregatedIntGaugeVec {
242    metric: RelabeledGuardedIntGaugeVec,
243    aggregation: GaugeAggregation,
244    groups: Arc<AggregatedGaugeGroups>,
245}
246
247impl RelabeledAggregatedIntGaugeVec {
248    pub fn with_guarded_label_values<V: AsRef<str> + std::fmt::Debug>(
249        &self,
250        vals: &[V],
251    ) -> RelabeledAggregatedIntGauge {
252        let relabeled_vals = self.metric.relabel_impl(vals);
253        let first_label_is_empty = match &relabeled_vals {
254            Some(labels) => labels.first().is_some_and(String::is_empty),
255            None => vals.first().is_some_and(|label| label.as_ref().is_empty()),
256        };
257
258        if !first_label_is_empty {
259            let gauge = match relabeled_vals {
260                Some(labels) => self.metric.metric.with_guarded_label_values(&labels),
261                None => self.metric.metric.with_guarded_label_values(vals),
262            };
263            return RelabeledAggregatedIntGauge {
264                inner: RelabeledAggregatedIntGaugeInner::Passthrough(gauge),
265            };
266        }
267
268        let labels = relabeled_vals
269            .unwrap_or_else(|| vals.iter().map(|label| label.as_ref().to_owned()).collect())
270            .into_boxed_slice();
271        let group = self
272            .groups
273            .get_or_create(labels, self.aggregation, |labels| {
274                self.metric.metric.with_guarded_label_values(labels)
275            });
276        let id = self.groups.next_id.fetch_add(1, Ordering::Relaxed);
277        RelabeledAggregatedIntGauge {
278            inner: RelabeledAggregatedIntGaugeInner::Aggregated(Arc::new(
279                AggregatedGaugeContribution { id, group },
280            )),
281        }
282    }
283}
284
285impl Collector for RelabeledAggregatedIntGaugeVec {
286    fn desc(&self) -> Vec<&prometheus::core::Desc> {
287        self.metric.desc()
288    }
289
290    fn collect(&self) -> Vec<prometheus::proto::MetricFamily> {
291        self.metric.collect()
292    }
293}
294
295#[derive(Debug, Default)]
296struct AggregatedGaugeGroups {
297    inner: Mutex<HashMap<Box<[String]>, Weak<AggregatedGaugeGroup>>>,
298    next_id: AtomicU64,
299}
300
301impl AggregatedGaugeGroups {
302    fn get_or_create(
303        self: &Arc<Self>,
304        labels: Box<[String]>,
305        aggregation: GaugeAggregation,
306        create_gauge: impl FnOnce(&[String]) -> LabelGuardedMetric<prometheus::IntGauge>,
307    ) -> Arc<AggregatedGaugeGroup> {
308        let mut groups = self.inner.lock();
309        if let Some(group) = groups.get(&labels).and_then(Weak::upgrade) {
310            return group;
311        }
312
313        let group = Arc::new(AggregatedGaugeGroup {
314            gauge: create_gauge(&labels),
315            aggregation,
316            state: Mutex::new(AggregatedGaugeState::default()),
317            labels: labels.clone(),
318            groups: Arc::downgrade(self),
319        });
320        groups.insert(labels, Arc::downgrade(&group));
321        group
322    }
323}
324
325#[derive(Debug, Default)]
326struct AggregatedGaugeState {
327    values: HashMap<u64, i64>,
328    aggregate: Option<i64>,
329}
330
331#[derive(Debug)]
332struct AggregatedGaugeGroup {
333    gauge: LabelGuardedMetric<prometheus::IntGauge>,
334    aggregation: GaugeAggregation,
335    state: Mutex<AggregatedGaugeState>,
336    labels: Box<[String]>,
337    groups: Weak<AggregatedGaugeGroups>,
338}
339
340impl AggregatedGaugeGroup {
341    fn set(&self, id: u64, value: i64) {
342        let mut state = self.state.lock();
343        self.set_locked(&mut state, id, value);
344    }
345
346    fn set_locked(&self, state: &mut AggregatedGaugeState, id: u64, value: i64) {
347        let previous = state.values.insert(id, value);
348        let aggregate = match self.aggregation {
349            GaugeAggregation::Sum => state
350                .aggregate
351                .unwrap_or(0)
352                .wrapping_sub(previous.unwrap_or(0))
353                .wrapping_add(value),
354            GaugeAggregation::Min => {
355                if previous == state.aggregate && previous.is_some_and(|old| value > old) {
356                    *state
357                        .values
358                        .values()
359                        .min()
360                        .expect("the updated value exists")
361                } else {
362                    state.aggregate.map_or(value, |current| current.min(value))
363                }
364            }
365            GaugeAggregation::Max => {
366                if previous == state.aggregate && previous.is_some_and(|old| value < old) {
367                    *state
368                        .values
369                        .values()
370                        .max()
371                        .expect("the updated value exists")
372                } else {
373                    state.aggregate.map_or(value, |current| current.max(value))
374                }
375            }
376        };
377        state.aggregate = Some(aggregate);
378        self.gauge.set(aggregate);
379    }
380
381    fn add(&self, id: u64, delta: i64) {
382        let mut state = self.state.lock();
383        let current = state.values.get(&id).copied().unwrap_or(0);
384        self.set_locked(&mut state, id, current.wrapping_add(delta));
385    }
386
387    fn get(&self) -> i64 {
388        self.state.lock().aggregate.unwrap_or(0)
389    }
390
391    fn remove(&self, id: u64) {
392        let mut state = self.state.lock();
393        let Some(previous) = state.values.remove(&id) else {
394            return;
395        };
396        let aggregate = match self.aggregation {
397            GaugeAggregation::Sum => state.aggregate.unwrap_or(0).wrapping_sub(previous),
398            GaugeAggregation::Min => state.values.values().copied().min().unwrap_or(0),
399            GaugeAggregation::Max => state.values.values().copied().max().unwrap_or(0),
400        };
401        state.aggregate = (!state.values.is_empty()).then_some(aggregate);
402        self.gauge.set(aggregate);
403    }
404}
405
406impl Drop for AggregatedGaugeGroup {
407    fn drop(&mut self) {
408        let Some(groups) = self.groups.upgrade() else {
409            return;
410        };
411        let mut entries = groups.inner.lock();
412        if entries
413            .get(&self.labels)
414            .is_some_and(|entry| std::ptr::eq(entry.as_ptr(), self))
415        {
416            entries.remove(&self.labels);
417        }
418    }
419}
420
421#[derive(Clone, Debug)]
422pub struct RelabeledAggregatedIntGauge {
423    inner: RelabeledAggregatedIntGaugeInner,
424}
425
426#[derive(Clone, Debug)]
427enum RelabeledAggregatedIntGaugeInner {
428    Passthrough(LabelGuardedMetric<prometheus::IntGauge>),
429    Aggregated(Arc<AggregatedGaugeContribution>),
430}
431
432#[derive(Debug)]
433struct AggregatedGaugeContribution {
434    id: u64,
435    group: Arc<AggregatedGaugeGroup>,
436}
437
438impl Drop for AggregatedGaugeContribution {
439    fn drop(&mut self) {
440        self.group.remove(self.id);
441    }
442}
443
444impl RelabeledAggregatedIntGauge {
445    /// Creates a single-contribution gauge for tests and placeholder metrics.
446    pub fn test_int_gauge<const N: usize>(_aggregation: GaugeAggregation) -> Self {
447        Self {
448            inner: RelabeledAggregatedIntGaugeInner::Passthrough(
449                crate::LabelGuardedIntGauge::test_int_gauge::<N>(),
450            ),
451        }
452    }
453
454    pub fn set(&self, value: i64) {
455        match &self.inner {
456            RelabeledAggregatedIntGaugeInner::Passthrough(gauge) => gauge.set(value),
457            RelabeledAggregatedIntGaugeInner::Aggregated(contribution) => {
458                contribution.group.set(contribution.id, value);
459            }
460        }
461    }
462
463    /// Updates this contribution when a value exists, or excludes it from aggregation otherwise.
464    /// A passthrough gauge keeps the conventional zero value when no value exists.
465    pub fn set_optional(&self, value: Option<i64>) {
466        match &self.inner {
467            RelabeledAggregatedIntGaugeInner::Passthrough(gauge) => {
468                gauge.set(value.unwrap_or_default());
469            }
470            RelabeledAggregatedIntGaugeInner::Aggregated(contribution) => match value {
471                Some(value) => contribution.group.set(contribution.id, value),
472                None => contribution.group.remove(contribution.id),
473            },
474        }
475    }
476
477    pub fn add(&self, value: i64) {
478        match &self.inner {
479            RelabeledAggregatedIntGaugeInner::Passthrough(gauge) => gauge.add(value),
480            RelabeledAggregatedIntGaugeInner::Aggregated(contribution) => {
481                contribution.group.add(contribution.id, value);
482            }
483        }
484    }
485
486    pub fn sub(&self, value: i64) {
487        self.add(value.wrapping_neg());
488    }
489
490    pub fn inc(&self) {
491        self.add(1);
492    }
493
494    pub fn dec(&self) {
495        self.sub(1);
496    }
497
498    pub fn get(&self) -> i64 {
499        match &self.inner {
500            RelabeledAggregatedIntGaugeInner::Passthrough(gauge) => gauge.get(),
501            RelabeledAggregatedIntGaugeInner::Aggregated(contribution) => contribution.group.get(),
502        }
503    }
504}
505
506#[cfg(test)]
507mod tests {
508    use prometheus::Registry;
509
510    use super::*;
511
512    fn test_vec(
513        registry: &Registry,
514        level: MetricLevel,
515        aggregation: GaugeAggregation,
516    ) -> RelabeledAggregatedIntGaugeVec {
517        crate::register_guarded_int_gauge_vec_with_registry!(
518            "test_relabel_aggregated_gauge",
519            "test",
520            &["actor_id", "key"],
521            registry,
522        )
523        .unwrap()
524        .relabel_debug_1_with_aggregation(level, aggregation)
525    }
526
527    fn values(registry: &Registry) -> Vec<i64> {
528        let Some(family) = registry
529            .gather()
530            .into_iter()
531            .find(|family| family.name() == "test_relabel_aggregated_gauge")
532        else {
533            return vec![];
534        };
535        family
536            .get_metric()
537            .iter()
538            .map(|metric| metric.get_gauge().value() as i64)
539            .collect()
540    }
541
542    #[test]
543    fn test_sum_aggregation_update_clone_and_drop() {
544        for level in [
545            MetricLevel::Disabled,
546            MetricLevel::Critical,
547            MetricLevel::Info,
548        ] {
549            let registry = Registry::new();
550            let vec = test_vec(&registry, level, GaugeAggregation::Sum);
551            let actor_1 = vec.with_guarded_label_values(&["1", "shared"]);
552            let actor_2 = vec.with_guarded_label_values(&["2", "shared"]);
553
554            actor_1.set(5);
555            actor_2.set(7);
556            assert_eq!(values(&registry), [12]);
557
558            actor_1.set(9);
559            assert_eq!(values(&registry), [16]);
560
561            let actor_1_clone = actor_1.clone();
562            drop(actor_1);
563            assert_eq!(values(&registry), [16]);
564            drop(actor_1_clone);
565            assert_eq!(values(&registry), [7]);
566
567            drop(actor_2);
568            assert_eq!(values(&registry), [0]);
569            assert!(values(&registry).is_empty());
570        }
571    }
572
573    #[test]
574    fn test_min_aggregation_ignores_uninitialized_handles() {
575        for level in [
576            MetricLevel::Disabled,
577            MetricLevel::Critical,
578            MetricLevel::Info,
579        ] {
580            let registry = Registry::new();
581            let vec = test_vec(&registry, level, GaugeAggregation::Min);
582            let actor_1 = vec.with_guarded_label_values(&["1", "shared"]);
583            let actor_2 = vec.with_guarded_label_values(&["2", "shared"]);
584
585            actor_1.set(100);
586            assert_eq!(values(&registry), [100]);
587            actor_2.set(200);
588            assert_eq!(values(&registry), [100]);
589            actor_1.set(300);
590            assert_eq!(values(&registry), [200]);
591            drop(actor_2);
592            assert_eq!(values(&registry), [300]);
593        }
594    }
595
596    #[test]
597    fn test_min_aggregation_ignores_absent_contributions() {
598        for level in [MetricLevel::Critical, MetricLevel::Info] {
599            let registry = Registry::new();
600            let vec = test_vec(&registry, level, GaugeAggregation::Min);
601            let empty_actor = vec.with_guarded_label_values(&["1", "shared"]);
602            let pending_actor = vec.with_guarded_label_values(&["2", "shared"]);
603
604            empty_actor.set_optional(None);
605            pending_actor.set_optional(Some(100));
606            assert_eq!(values(&registry), [100]);
607
608            empty_actor.set_optional(Some(200));
609            assert_eq!(values(&registry), [100]);
610
611            pending_actor.set_optional(None);
612            assert_eq!(values(&registry), [200]);
613
614            empty_actor.set_optional(None);
615            assert_eq!(values(&registry), [0]);
616        }
617    }
618
619    #[test]
620    fn test_max_aggregation_update_clone_and_drop() {
621        for level in [
622            MetricLevel::Disabled,
623            MetricLevel::Critical,
624            MetricLevel::Info,
625        ] {
626            let registry = Registry::new();
627            let vec = test_vec(&registry, level, GaugeAggregation::Max);
628            let actor_1 = vec.with_guarded_label_values(&["1", "shared"]);
629            let actor_2 = vec.with_guarded_label_values(&["2", "shared"]);
630
631            actor_1.set(100);
632            assert_eq!(values(&registry), [100]);
633            actor_2.set(200);
634            assert_eq!(values(&registry), [200]);
635            actor_1.set(300);
636            assert_eq!(values(&registry), [300]);
637            actor_2.set(300);
638            actor_1.set(100);
639            assert_eq!(values(&registry), [300]);
640            actor_2.set(50);
641            assert_eq!(values(&registry), [100]);
642
643            let actor_1_clone = actor_1.clone();
644            drop(actor_1);
645            assert_eq!(values(&registry), [100]);
646            drop(actor_1_clone);
647            assert_eq!(values(&registry), [50]);
648            drop(actor_2);
649            assert_eq!(values(&registry), [0]);
650            assert!(values(&registry).is_empty());
651        }
652    }
653
654    #[test]
655    fn test_max_aggregation_ignores_absent_contributions() {
656        for level in [MetricLevel::Critical, MetricLevel::Info] {
657            let registry = Registry::new();
658            let vec = test_vec(&registry, level, GaugeAggregation::Max);
659            let actor_1 = vec.with_guarded_label_values(&["1", "shared"]);
660            let actor_2 = vec.with_guarded_label_values(&["2", "shared"]);
661
662            // Uninitialized and absent handles must not contribute a zero maximum.
663            actor_1.set(-100);
664            assert_eq!(values(&registry), [-100]);
665            actor_2.set_optional(None);
666            assert_eq!(values(&registry), [-100]);
667            actor_2.set_optional(Some(-50));
668            assert_eq!(values(&registry), [-50]);
669            actor_2.set_optional(None);
670            assert_eq!(values(&registry), [-100]);
671            actor_1.set_optional(None);
672            assert_eq!(values(&registry), [0]);
673            actor_2.set(-200);
674            assert_eq!(values(&registry), [-200]);
675        }
676    }
677
678    #[test]
679    fn test_debug_passthrough_and_distinct_groups() {
680        let registry = Registry::new();
681        let vec = test_vec(&registry, MetricLevel::Debug, GaugeAggregation::Sum);
682        let actor_1 = vec.with_guarded_label_values(&["1", "shared"]);
683        let actor_2 = vec.with_guarded_label_values(&["2", "shared"]);
684        actor_1.set(5);
685        actor_2.set(7);
686        let mut collected = values(&registry);
687        collected.sort_unstable();
688        assert_eq!(collected, [5, 7]);
689        let mut actor_ids = registry
690            .gather()
691            .into_iter()
692            .find(|family| family.name() == "test_relabel_aggregated_gauge")
693            .unwrap()
694            .get_metric()
695            .iter()
696            .map(|metric| crate::get_label::<String>(metric, "actor_id").unwrap())
697            .collect::<Vec<_>>();
698        actor_ids.sort_unstable();
699        assert_eq!(actor_ids, ["1", "2"]);
700
701        let registry = Registry::new();
702        let vec = test_vec(&registry, MetricLevel::Info, GaugeAggregation::Sum);
703        let first = vec.with_guarded_label_values(&["1", "first"]);
704        let second = vec.with_guarded_label_values(&["2", "second"]);
705        first.set(5);
706        second.set(7);
707        let mut collected = values(&registry);
708        collected.sort_unstable();
709        assert_eq!(collected, [5, 7]);
710    }
711
712    #[test]
713    fn test_aggregation_follows_the_effective_first_label() {
714        let registry = Registry::new();
715        let vec = test_vec(&registry, MetricLevel::Debug, GaugeAggregation::Sum);
716        let first = vec.with_guarded_label_values(&["", "shared"]);
717        let second = vec.with_guarded_label_values(&["", "shared"]);
718
719        first.set(5);
720        second.set(7);
721
722        assert_eq!(values(&registry), [12]);
723    }
724
725    #[test]
726    fn test_concurrent_adds_and_sets_keep_exported_value_consistent() {
727        let registry = Registry::new();
728        let vec = test_vec(&registry, MetricLevel::Info, GaugeAggregation::Sum);
729        let gauge = vec.with_guarded_label_values(&["1", "shared"]);
730        let threads = (0..4)
731            .map(|thread_index| {
732                let gauge = gauge.clone();
733                std::thread::spawn(move || {
734                    for value in 0..10_000 {
735                        if thread_index % 2 == 0 {
736                            gauge.add(1);
737                        } else {
738                            gauge.set(value);
739                        }
740                    }
741                })
742            })
743            .collect::<Vec<_>>();
744        for thread in threads {
745            thread.join().unwrap();
746        }
747        assert_eq!(values(&registry), [gauge.get()]);
748    }
749}
risingwave_common_metrics/relabeled_metric.rs

risingwave_common_metrics/
relabeled_metric.rs
