risingwave_common_metrics/
relabeled_metric.rs

1// Copyright 2025 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 prometheus::core::{Collector, MetricVec, MetricVecBuilder};
16use prometheus::{HistogramVec, IntCounterVec};
17
18use crate::{
19    LabelGuardedHistogramVec, LabelGuardedIntCounterVec, LabelGuardedIntGaugeVec,
20    LabelGuardedMetric, LabelGuardedMetricVec, MetricLevel,
21};
22
23/// For all `Relabeled*Vec` below,
24/// - when `metric_level` <= `relabel_threshold`, they behave exactly the same as their inner
25///   metric.
26/// - when `metric_level` > `relabel_threshold`, the first `relabel_num` labels are rewrite to "" when
27///   calling `with_label_values`. That's means the metric vec is aggregated into a single metric.
28///
29/// These wrapper classes add a `metric_level` field to corresponding metric.
30/// We could have use one single struct to represent all `MetricVec<T: MetricVecBuilder>`, rather
31/// than specializing them one by one. However, that's undoable because prometheus crate doesn't
32/// export `MetricVecBuilder` implementation like `HistogramVecBuilder`.
33///
34/// ## Note
35///
36/// CAUTION! Relabelling might cause expected result!
37///
38/// For counters (including histogram because it uses counters internally), it's usually natural
39/// to sum up the count from multiple labels.
40///
41/// For the rest (such as Gauge), the semantics becomes "any/last of the recorded value". Please be
42/// cautious.
43#[derive(Clone, Debug)]
44pub struct RelabeledMetricVec<M> {
45    relabel_threshold: MetricLevel,
46    metric_level: MetricLevel,
47    metric: M,
48
49    /// The first `relabel_num` labels will be relabeled to empty string
50    ///
51    /// For example, if `relabel_num` is 1, and the input labels are `["actor_id",
52    /// "fragment_id", "table_id"]`, when threshold is reached, the label values will be
53    /// `["", "<original_fragment_id>", "<original_table_id>"]`.
54    relabel_num: usize,
55}
56
57impl<M> RelabeledMetricVec<M> {
58    pub fn with_metric_level(
59        metric_level: MetricLevel,
60        metric: M,
61        relabel_threshold: MetricLevel,
62    ) -> Self {
63        Self {
64            relabel_threshold,
65            metric_level,
66            metric,
67            relabel_num: usize::MAX,
68        }
69    }
70
71    pub fn with_metric_level_relabel_n(
72        metric_level: MetricLevel,
73        metric: M,
74        relabel_threshold: MetricLevel,
75        relabel_num: usize,
76    ) -> Self {
77        Self {
78            relabel_threshold,
79            metric_level,
80            metric,
81            relabel_num,
82        }
83    }
84}
85
86#[easy_ext::ext(MetricVecRelabelExt)]
87impl<M> M
88where
89    M: Sized,
90{
91    /// Equivalent to [`RelabeledMetricVec::with_metric_level`].
92    pub fn relabel(
93        self,
94        metric_level: MetricLevel,
95        relabel_threshold: MetricLevel,
96    ) -> RelabeledMetricVec<M> {
97        RelabeledMetricVec::with_metric_level(metric_level, self, relabel_threshold)
98    }
99
100    /// Equivalent to [`RelabeledMetricVec::with_metric_level_relabel_n`].
101    pub fn relabel_n(
102        self,
103        metric_level: MetricLevel,
104        relabel_threshold: MetricLevel,
105        relabel_num: usize,
106    ) -> RelabeledMetricVec<M> {
107        RelabeledMetricVec::with_metric_level_relabel_n(
108            metric_level,
109            self,
110            relabel_threshold,
111            relabel_num,
112        )
113    }
114
115    /// Equivalent to [`RelabeledMetricVec::with_metric_level_relabel_n`] with `metric_level` set to
116    /// `MetricLevel::Debug` and `relabel_num` set to 1.
117    pub fn relabel_debug_1(self, relabel_threshold: MetricLevel) -> RelabeledMetricVec<M> {
118        RelabeledMetricVec::with_metric_level_relabel_n(
119            MetricLevel::Debug,
120            self,
121            relabel_threshold,
122            1,
123        )
124    }
125}
126
127impl<T: MetricVecBuilder> RelabeledMetricVec<MetricVec<T>> {
128    pub fn with_label_values<V: AsRef<str> + std::fmt::Debug>(&self, vals: &[V]) -> T::M {
129        if self.metric_level > self.relabel_threshold {
130            // relabel first n labels to empty string
131            let mut relabeled_vals = vals
132                .iter()
133                .map(|v| v.as_ref().to_owned())
134                .collect::<Vec<_>>();
135            for label in relabeled_vals.iter_mut().take(self.relabel_num) {
136                *label = String::new();
137            }
138            return self.metric.with_label_values(&relabeled_vals);
139        }
140        self.metric.with_label_values(vals)
141    }
142}
143
144impl<T: MetricVecBuilder> RelabeledMetricVec<LabelGuardedMetricVec<T>> {
145    pub fn with_guarded_label_values<V: AsRef<str> + std::fmt::Debug>(
146        &self,
147        vals: &[V],
148    ) -> LabelGuardedMetric<T::M> {
149        if self.metric_level > self.relabel_threshold {
150            // relabel first n labels to empty string
151            let mut relabeled_vals = vals
152                .iter()
153                .map(|v| v.as_ref().to_owned())
154                .collect::<Vec<_>>();
155            for label in relabeled_vals.iter_mut().take(self.relabel_num) {
156                *label = String::new();
157            }
158            return self.metric.with_guarded_label_values(&relabeled_vals);
159        }
160        self.metric.with_guarded_label_values(vals)
161    }
162}
163
164impl<T: Collector> Collector for RelabeledMetricVec<T> {
165    fn desc(&self) -> Vec<&prometheus::core::Desc> {
166        self.metric.desc()
167    }
168
169    fn collect(&self) -> Vec<prometheus::proto::MetricFamily> {
170        self.metric.collect()
171    }
172}
173
174pub type RelabeledCounterVec = RelabeledMetricVec<IntCounterVec>;
175pub type RelabeledHistogramVec = RelabeledMetricVec<HistogramVec>;
176
177pub type RelabeledGuardedHistogramVec = RelabeledMetricVec<LabelGuardedHistogramVec>;
178pub type RelabeledGuardedIntCounterVec = RelabeledMetricVec<LabelGuardedIntCounterVec>;
179
180/// CAUTION! Relabelling a Gauge might cause expected result!
181///
182/// See [`RelabeledMetricVec`] for details.
183pub type RelabeledGuardedIntGaugeVec = RelabeledMetricVec<LabelGuardedIntGaugeVec>;
risingwave_common_metrics/relabeled_metric.rs

risingwave_common_metrics/
relabeled_metric.rs
