mz_ore/metrics/
aggregation.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//! Scrape-time aggregation of per-object values into metric families whose
17//! cardinality does not depend on the number of objects.
18//!
19//! A subsystem that tracks one value per object (per shard, per replica, ...)
20//! would naively export one series per object. Instead, its collector pushes
21//! each object's current values into an [`AggregationSnapshot`] once per
22//! scrape and emits what [`AggregationSnapshot::finish`] returns. Per
23//! [`AggregatedFamily`] that is:
24//!
25//! * `<name>_percentile`: one gauge per entry of [`PERCENTILES`], labeled
26//!   `percentile`.
27//! * `<name>_sum`: one gauge, the sum of all values, if [`AggregatedFamily::sum`]
28//!   is set.
29//! * `<name>_topk`: one gauge for each of the `k` objects with the largest
30//!   non-zero value, labeled by [`AggregationKey`]. An object that leaves the
31//!   top `k` is simply not emitted on the next scrape.
32//!
33//! All are exact over the snapshot. Percentiles of different processes cannot
34//! be combined into a percentile over all their objects, so these describe one
35//! process's objects.
36//!
37//! NOTE: Built at runtime, not with [`metric!`](crate::metric), so
38//! `bin/gen-metrics-catalog` misses them and `doc/user/data/metrics.yml` has
39//! no entry for them.
40
41use std::marker::PhantomData;
42
43use itertools::Itertools;
44use prometheus::Opts;
45use prometheus::core::{Collector, Desc};
46use prometheus::proto::MetricFamily;
47
48use crate::cast::CastLossy;
49use crate::metrics::UIntGauge;
50use crate::metrics::raw::UIntGaugeVec;
51
52/// The `percentile` label values of a `_percentile` family and the
53/// percentile each reports, as a fraction of the objects.
54///
55/// Each is the nearest-rank value: of `n` values sorted ascending, the one at
56/// 1-based rank `ceil(q * n)`, and at least rank 1. So `min` and `max` are the
57/// smallest and largest values, and every reported value is one some object
58/// actually has.
59pub const PERCENTILES: &[(&str, f64)] = &[
60    ("min", 0.0),
61    ("p25", 0.25),
62    ("p50", 0.5),
63    ("p75", 0.75),
64    ("p90", 0.9),
65    ("p99", 0.99),
66    ("p999", 0.999),
67    ("p9999", 0.9999),
68    ("max", 1.0),
69];
70
71/// Identifies one object in the `_topk` families.
72pub trait AggregationKey {
73    /// The label names of a `_topk` series.
74    const LABEL_NAMES: &'static [&'static str];
75
76    /// The label values of this object, one per [`Self::LABEL_NAMES`] entry.
77    /// Only called for objects that make a top-K.
78    fn label_values(&self) -> Vec<String>;
79}
80
81/// The definition of one aggregated per-object value.
82#[derive(Debug, Clone)]
83pub struct AggregatedFamily {
84    /// The name the per-object metric would have.
85    pub name: String,
86    /// The help text of the per-object metric.
87    pub help: String,
88    /// Whether to emit `<name>_sum`. Unset for values whose sum means nothing,
89    /// such as a per-object maximum.
90    pub sum: bool,
91    /// If set, also emit the `k` largest objects as `<name>_topk`.
92    pub top_k: Option<usize>,
93}
94
95impl AggregatedFamily {
96    fn percentile_vec(&self) -> UIntGaugeVec {
97        let opts = Opts::new(
98            format!("{}_percentile", self.name),
99            format!("{} (percentiles over this process's objects)", self.help),
100        );
101        UIntGaugeVec::new(opts, &["percentile"]).expect("valid percentile family")
102    }
103
104    fn sum_gauge(&self) -> UIntGauge {
105        let opts = Opts::new(
106            format!("{}_sum", self.name),
107            format!(
108                "{} (sum over this process's objects, which other processes may also count)",
109                self.help
110            ),
111        );
112        UIntGauge::with_opts(opts).expect("valid sum family")
113    }
114
115    fn top_k_vec<K: AggregationKey>(&self, k: usize) -> UIntGaugeVec {
116        let opts = Opts::new(
117            format!("{}_topk", self.name),
118            format!("{} (the {k} largest objects)", self.help),
119        );
120        UIntGaugeVec::new(opts, K::LABEL_NAMES).expect("valid top-k family")
121    }
122}
123
124/// A fixed set of [`AggregatedFamily`]s over objects keyed by `K`.
125#[derive(Debug)]
126pub struct AggregatedFamilies<K> {
127    families: Vec<AggregatedFamily>,
128    descs: Vec<Desc>,
129    _key: PhantomData<fn(&K)>,
130}
131
132impl<K: AggregationKey> AggregatedFamilies<K> {
133    /// Panics on an invalid name, like registering an invalid
134    /// [`metric!`](crate::metric) does.
135    pub fn new(families: Vec<AggregatedFamily>) -> Self {
136        let mut descs = Vec::new();
137        for f in &families {
138            descs.extend(f.percentile_vec().desc().into_iter().cloned());
139            if f.sum {
140                descs.extend(f.sum_gauge().desc().into_iter().cloned());
141            }
142            if let Some(k) = f.top_k {
143                descs.extend(f.top_k_vec::<K>(k).desc().into_iter().cloned());
144            }
145        }
146        AggregatedFamilies {
147            families,
148            descs,
149            _key: PhantomData,
150        }
151    }
152
153    /// The descriptors of every emitted family, for `Collector::desc`.
154    pub fn descs(&self) -> Vec<&Desc> {
155        self.descs.iter().collect()
156    }
157
158    /// Starts an empty snapshot.
159    pub fn snapshot(&self) -> AggregationSnapshot<'_, K> {
160        AggregationSnapshot {
161            families: self,
162            keys: Vec::new(),
163            columns: vec![Vec::new(); self.families.len()],
164        }
165    }
166}
167
168/// The current values of a set of objects.
169///
170/// Pushing only copies, so a caller can push under its own lock and call
171/// [`Self::finish`], which does the sorting, after releasing it.
172#[derive(Debug)]
173pub struct AggregationSnapshot<'a, K> {
174    families: &'a AggregatedFamilies<K>,
175    keys: Vec<K>,
176    /// Per family, `(value, index into keys)`.
177    columns: Vec<Vec<(u64, usize)>>,
178}
179
180impl<K: AggregationKey> AggregationSnapshot<'_, K> {
181    /// Records one object, with one value per family in definition order.
182    pub fn push(&mut self, key: K, values: &[u64]) {
183        for (column, value) in self.columns.iter_mut().zip_eq(values) {
184            column.push((*value, self.keys.len()));
185        }
186        self.keys.push(key);
187    }
188
189    /// Folds the snapshot into the aggregated metric families.
190    pub fn finish(self) -> Vec<MetricFamily> {
191        let mut out = Vec::new();
192        for (f, mut column) in self.families.families.iter().zip_eq(self.columns) {
193            // Ascending by value. The index breaks ties, so the top-K is
194            // deterministic.
195            column.sort_unstable();
196
197            // Built fresh every scrape, so no series outlives the snapshot.
198            let percentiles = f.percentile_vec();
199            if !column.is_empty() {
200                let n = f64::cast_lossy(column.len());
201                for (label, q) in PERCENTILES {
202                    let rank = usize::cast_lossy((q * n).ceil());
203                    let (value, _) = column[rank.max(1) - 1];
204                    percentiles.with_label_values(&[*label]).set(value);
205                }
206            }
207            out.extend(percentiles.collect());
208
209            if f.sum {
210                let sum = f.sum_gauge();
211                sum.set(column.iter().map(|(v, _)| v).sum());
212                out.extend(sum.collect());
213            }
214
215            if let Some(k) = f.top_k {
216                let top_k = f.top_k_vec::<K>(k);
217                // An object with nothing to report is not a heavy hitter.
218                for &(value, idx) in column.iter().rev().take_while(|(v, _)| *v > 0).take(k) {
219                    top_k
220                        .with_label_values(&self.keys[idx].label_values())
221                        .set(value);
222                }
223                out.extend(top_k.collect());
224            }
225        }
226        out
227    }
228}
229
230#[cfg(test)]
231mod tests {
232    use std::collections::BTreeMap;
233
234    use super::*;
235
236    struct Obj(u64);
237
238    impl AggregationKey for Obj {
239        const LABEL_NAMES: &'static [&'static str] = &["obj"];
240        fn label_values(&self) -> Vec<String> {
241            vec![self.0.to_string()]
242        }
243    }
244
245    /// Aggregates one family `m` with a sum and top-10 over objects with the
246    /// given values, keyed by the value itself.
247    fn aggregate(values: impl IntoIterator<Item = u64>) -> Vec<MetricFamily> {
248        aggregate_with_sum(values, true)
249    }
250
251    fn aggregate_with_sum(values: impl IntoIterator<Item = u64>, sum: bool) -> Vec<MetricFamily> {
252        let families = AggregatedFamilies::<Obj>::new(vec![AggregatedFamily {
253            name: "m".into(),
254            help: "help".into(),
255            sum,
256            top_k: Some(10),
257        }]);
258        let mut snapshot = families.snapshot();
259        for v in values {
260            snapshot.push(Obj(v), &[v]);
261        }
262        snapshot.finish()
263    }
264
265    /// The series of family `name`, as label value to value.
266    fn series(families: &[MetricFamily], name: &str) -> BTreeMap<String, u64> {
267        let family = families.iter().find(|f| f.name() == name).expect("family");
268        family
269            .get_metric()
270            .iter()
271            .map(|m| {
272                (
273                    m.get_label()[0].value().to_string(),
274                    u64::cast_lossy(m.get_gauge().value()),
275                )
276            })
277            .collect()
278    }
279
280    fn percentiles(values: [u64; 9]) -> BTreeMap<String, u64> {
281        PERCENTILES
282            .iter()
283            .map(|(l, _)| l.to_string())
284            .zip_eq(values)
285            .collect()
286    }
287
288    fn objs(values: impl IntoIterator<Item = u64>) -> BTreeMap<String, u64> {
289        values.into_iter().map(|v| (v.to_string(), v)).collect()
290    }
291
292    #[crate::test]
293    fn percentiles_are_exact() {
294        // Pushed in reverse to check that the values are sorted.
295        assert_eq!(
296            series(&aggregate((1..=10_000).rev()), "m_percentile"),
297            percentiles([1, 2500, 5000, 7500, 9000, 9900, 9990, 9999, 10_000])
298        );
299        // Nearest rank rounds up: with 100 values, p99.9 and p99.99 are the max.
300        assert_eq!(
301            series(&aggregate(1..=100), "m_percentile"),
302            percentiles([1, 25, 50, 75, 90, 99, 100, 100, 100])
303        );
304        assert_eq!(series(&aggregate([7]), "m_percentile"), percentiles([7; 9]));
305        assert!(series(&aggregate([]), "m_percentile").is_empty());
306    }
307
308    #[crate::test]
309    fn sum_is_exact() {
310        let sum = |values: Vec<u64>| {
311            let families = aggregate(values);
312            let family = families.iter().find(|f| f.name() == "m_sum").expect("sum");
313            u64::cast_lossy(family.get_metric()[0].get_gauge().value())
314        };
315        assert_eq!(sum((1..=100).collect()), 5050);
316        assert_eq!(sum(vec![]), 0);
317    }
318
319    #[crate::test]
320    fn sum_is_omitted_when_unset() {
321        let families = aggregate_with_sum(1..=10, false);
322        assert!(families.iter().all(|f| f.name() != "m_sum"));
323        assert!(families.iter().any(|f| f.name() == "m_percentile"));
324    }
325
326    #[crate::test]
327    fn topk_is_the_largest_non_zero_values() {
328        assert_eq!(series(&aggregate(1..=100), "m_topk"), objs(91..=100));
329        assert_eq!(series(&aggregate([0, 7, 0]), "m_topk"), objs([7]));
330        // Membership is recomputed each time, so a departed object is gone.
331        assert_eq!(series(&aggregate(1..=5), "m_topk"), objs(1..=5));
332    }
333
334    #[crate::test]
335    fn cardinality_is_flat_in_object_count() {
336        let series_count =
337            |n| -> usize { aggregate(1..=n).iter().map(|f| f.get_metric().len()).sum() };
338        assert_eq!(series_count(10), PERCENTILES.len() + 1 + 10);
339        assert_eq!(series_count(10), series_count(10_000));
340    }
341}
mz_ore/metrics/aggregation.rs

mz_ore/metrics/
aggregation.rs
