mz_ore/
metrics.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//! Metrics for materialize systems.
17//!
18//! The idea here is that each subsystem keeps its metrics in a scoped-to-it struct, which gets
19//! registered (once) to the server's (or a test's) prometheus registry.
20//!
21//! Instead of using prometheus's (very verbose) metrics definitions, we rely on type inference to
22//! reduce the verbosity a little bit. A typical subsystem will look like the following:
23//!
24//! ```rust
25//! # use mz_ore::metrics::{MetricsRegistry, IntCounter};
26//! # use mz_ore::metric;
27//! #[derive(Debug, Clone)] // Note that prometheus metrics can safely be cloned
28//! struct Metrics {
29//!     pub bytes_sent: IntCounter,
30//! }
31//!
32//! impl Metrics {
33//!     pub fn register_into(registry: &MetricsRegistry) -> Metrics {
34//!         Metrics {
35//!             bytes_sent: registry.register(metric!(
36//!                 name: "mz_pg_sent_bytes",
37//!                 help: "total number of bytes sent here",
38//!             )),
39//!         }
40//!     }
41//! }
42//! ```
43
44use std::any::Any;
45use std::fmt;
46use std::fmt::{Debug, Formatter};
47use std::future::Future;
48use std::pin::Pin;
49use std::sync::{Arc, Mutex};
50use std::task::{Context, Poll};
51use std::time::{Duration, Instant};
52
53use derivative::Derivative;
54use pin_project::pin_project;
55use prometheus::core::{
56    Atomic, AtomicF64, AtomicI64, AtomicU64, Collector, Desc, GenericCounter, GenericCounterVec,
57    GenericGauge, GenericGaugeVec,
58};
59use prometheus::proto::MetricFamily;
60use prometheus::{HistogramOpts, Registry};
61
62mod delete_on_drop;
63
64pub use delete_on_drop::*;
65pub use prometheus::Opts as PrometheusOpts;
66
67/// Define a metric for use in materialize.
68#[macro_export]
69macro_rules! metric {
70    (
71        name: $name:expr,
72        help: $help:expr
73        $(, subsystem: $subsystem_name:expr)?
74        $(, const_labels: { $($cl_key:expr => $cl_value:expr ),* })?
75        $(, var_labels: [ $($vl_name:expr),* ])?
76        $(, buckets: $bk_name:expr)?
77        $(, visibility: $visibility:expr)?
78        $(, tags: [ $($tag:expr),* $(,)? ])?
79        $(,)?
80    ) => {{
81        let const_labels = (&[
82            $($(
83                ($cl_key.to_string(), $cl_value.to_string()),
84            )*)?
85        ]).into_iter().cloned().collect();
86        let var_labels = vec![
87            $(
88                $($vl_name.into(),)*
89            )?];
90        #[allow(unused_mut)]
91        let mut mk_opts = $crate::metrics::MakeCollectorOpts {
92            opts: $crate::metrics::PrometheusOpts::new($name, $help)
93                $(.subsystem( $subsystem_name ))?
94                .const_labels(const_labels)
95                .variable_labels(var_labels),
96            buckets: None,
97        };
98        // Set buckets if passed
99        $(mk_opts.buckets = Some($bk_name);)*
100        // `visibility` is documentation metadata for the metrics catalog
101        // (`bin/gen-metrics-catalog`).
102        // It has no runtime effect; we only type-check it here so a bad value is
103        // a compile error rather than silently ignored.
104        $(let _: $crate::metrics::MetricVisibility = $visibility;)?
105        // `tags` is documentation metadata for the metrics catalog.
106        // It has no runtime effect.
107        $($(let _: $crate::metrics::MetricTag = $tag;)*)?
108        mk_opts
109    }}
110}
111
112/// Options for MakeCollector. This struct should be instantiated using the metric macro.
113#[derive(Debug, Clone)]
114pub struct MakeCollectorOpts {
115    /// Common Prometheus options
116    pub opts: PrometheusOpts,
117    /// Buckets to be used with Histogram and HistogramVec. Must be set to create Histogram types
118    /// and must not be set for other types.
119    pub buckets: Option<Vec<f64>>,
120}
121
122/// This is documentation metadata: it is set via the optional `visibility:`
123/// field of [`metric!`] and consumed by the metrics catalog
124/// (`bin/gen-metrics-catalog`), which reads it from the source tree to produce
125/// the user-facing metrics reference. It has no effect on the metric at
126/// runtime.
127#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, serde::Serialize)]
128#[serde(rename_all = "snake_case")]
129pub enum MetricVisibility {
130    /// A metric we use for internal development.
131    #[default]
132    Internal,
133    /// A metric we want customers to build dashboards and
134    /// alerts on. We do not guarantee stability for this group
135    /// of metrics.
136    Public,
137}
138
139/// A tag categorizing a metric in the user-facing metrics catalog.
140///
141/// It is set via the optional `tags:` field of [`metric!`] and
142/// consumed by the metrics catalog (`bin/gen-metrics-catalog`).
143/// It has no effect on the metric at runtime.
144///
145/// A metric may carry many tags, or none (the default). A tag names the
146/// grouping the metric is presented under in user-facing documentation.
147#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize)]
148#[serde(rename_all = "kebab-case")]
149pub enum MetricTag {
150    /// SQL Control plane metrics (client connections, availability, catalog).
151    Environment,
152    /// Metrics for compute objects (indexes, materialized views).
153    Compute,
154    /// Metrics for sources.
155    Source,
156    /// Metrics for sinks.
157    Sink,
158}
159
160/// The materialize metrics registry.
161#[derive(Clone, Derivative)]
162#[derivative(Debug)]
163pub struct MetricsRegistry {
164    inner: Registry,
165    #[derivative(Debug = "ignore")]
166    postprocessors: Arc<Mutex<Vec<Box<dyn FnMut(&mut Vec<MetricFamily>) + Send + Sync>>>>,
167}
168
169/// A wrapper for metrics to require delete on drop semantics
170///
171/// The wrapper behaves like regular metrics but only provides functions to create delete-on-drop
172/// variants. This way, no metrics of this type can be leaked.
173///
174/// In situations where the delete-on-drop behavior is not desired or in legacy code, use the raw
175/// variants of the metrics, as defined in [self::raw].
176#[derive(Clone)]
177pub struct DeleteOnDropWrapper<M> {
178    inner: M,
179}
180
181impl<M: MakeCollector + Debug> Debug for DeleteOnDropWrapper<M> {
182    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
183        self.inner.fmt(f)
184    }
185}
186
187impl<M: Collector> Collector for DeleteOnDropWrapper<M> {
188    fn desc(&self) -> Vec<&Desc> {
189        self.inner.desc()
190    }
191
192    fn collect(&self) -> Vec<MetricFamily> {
193        self.inner.collect()
194    }
195}
196
197impl<M: MakeCollector> MakeCollector for DeleteOnDropWrapper<M> {
198    fn make_collector(opts: MakeCollectorOpts) -> Self {
199        DeleteOnDropWrapper {
200            inner: M::make_collector(opts),
201        }
202    }
203}
204
205impl<M: MetricVecExt> DeleteOnDropWrapper<M> {
206    /// Returns a metric that deletes its labels from this metrics vector when dropped.
207    pub fn get_delete_on_drop_metric<L: PromLabelsExt>(
208        &self,
209        labels: L,
210    ) -> DeleteOnDropMetric<M, L> {
211        self.inner.get_delete_on_drop_metric(labels)
212    }
213}
214
215/// The unsigned integer version of [`Gauge`]. Provides better performance if
216/// metric values are all unsigned integers.
217pub type UIntGauge = GenericGauge<AtomicU64>;
218
219/// Delete-on-drop shadow of Prometheus [prometheus::CounterVec].
220pub type CounterVec = DeleteOnDropWrapper<prometheus::CounterVec>;
221/// Delete-on-drop shadow of Prometheus [prometheus::Gauge].
222pub type Gauge = DeleteOnDropWrapper<prometheus::Gauge>;
223/// Delete-on-drop shadow of Prometheus [prometheus::GaugeVec].
224pub type GaugeVec = DeleteOnDropWrapper<prometheus::GaugeVec>;
225/// Delete-on-drop shadow of Prometheus [prometheus::HistogramVec].
226pub type HistogramVec = DeleteOnDropWrapper<prometheus::HistogramVec>;
227/// Delete-on-drop shadow of Prometheus [prometheus::IntCounterVec].
228pub type IntCounterVec = DeleteOnDropWrapper<prometheus::IntCounterVec>;
229/// Delete-on-drop shadow of Prometheus [prometheus::IntGaugeVec].
230pub type IntGaugeVec = DeleteOnDropWrapper<prometheus::IntGaugeVec>;
231/// Delete-on-drop shadow of Prometheus [raw::UIntGaugeVec].
232pub type UIntGaugeVec = DeleteOnDropWrapper<raw::UIntGaugeVec>;
233
234use crate::assert_none;
235
236pub use prometheus::{Counter, Histogram, IntCounter, IntGauge};
237
238/// Access to non-delete-on-drop vector types
239pub mod raw {
240    use prometheus::core::{AtomicU64, GenericGaugeVec};
241
242    /// The unsigned integer version of [`GaugeVec`].
243    /// Provides better performance if metric values are all unsigned integers.
244    pub type UIntGaugeVec = GenericGaugeVec<AtomicU64>;
245
246    pub use prometheus::{CounterVec, Gauge, GaugeVec, HistogramVec, IntCounterVec, IntGaugeVec};
247}
248
249impl MetricsRegistry {
250    /// Creates a new metrics registry.
251    pub fn new() -> Self {
252        MetricsRegistry {
253            inner: Registry::new(),
254            postprocessors: Arc::new(Mutex::new(vec![])),
255        }
256    }
257
258    /// Register a metric defined with the [`metric`] macro.
259    pub fn register<M>(&self, opts: MakeCollectorOpts) -> M
260    where
261        M: MakeCollector,
262    {
263        let collector = M::make_collector(opts);
264        self.inner.register(Box::new(collector.clone())).unwrap();
265        collector
266    }
267
268    /// Registers a gauge whose value is computed when observed.
269    pub fn register_computed_gauge<P>(
270        &self,
271        opts: MakeCollectorOpts,
272        f: impl Fn() -> P::T + Send + Sync + 'static,
273    ) -> ComputedGenericGauge<P>
274    where
275        P: Atomic + 'static,
276    {
277        let gauge = ComputedGenericGauge {
278            gauge: GenericGauge::make_collector(opts),
279            f: Arc::new(f),
280        };
281        self.inner.register(Box::new(gauge.clone())).unwrap();
282        gauge
283    }
284
285    /// Register a pre-defined prometheus collector.
286    pub fn register_collector<C: 'static + prometheus::core::Collector>(&self, collector: C) {
287        self.inner
288            .register(Box::new(collector))
289            .expect("registering pre-defined metrics collector");
290    }
291
292    /// Register a pre-defined collector and return a handle that unregisters it on drop.
293    ///
294    /// Use this for collectors with a bounded lifetime (e.g. a metric sink that is torn down with
295    /// its dataflow), so the collector's series stop being scraped once the owner drops the handle.
296    /// The returned value is opaque: the caller only needs to hold it for as long as the collector
297    /// should stay registered, then drop it.
298    ///
299    /// `prometheus::Registry::unregister` matches a collector by the id of its `Desc`s, not by
300    /// object identity, so the guard keeps a clone of the collector and hands it back to
301    /// `unregister` on drop.
302    ///
303    /// If a collector with the same descriptor id is already registered this soft-panics and
304    /// returns a handle that owns no registration. A duplicate registration is a logic error, but
305    /// panicking here would run on a worker thread and could crash the process, so it degrades to
306    /// missing series rather than taking down the whole scrape. The contract is that a caller
307    /// replacing a collector must drop the old handle before registering the new one: doing so
308    /// unregisters the old descriptor id first, so the re-registration succeeds cleanly.
309    pub fn register_collector_with_dropper<C>(&self, collector: C) -> Box<dyn Any + Send + Sync>
310    where
311        C: 'static + prometheus::core::Collector + Clone + Send + Sync,
312    {
313        match self.inner.register(Box::new(collector.clone())) {
314            Ok(()) => {
315                // `prometheus::Registry` is `Arc`-backed, so this clone is cheap and shares the
316                // same underlying registry the collector was registered into.
317                let registry = self.inner.clone();
318                Box::new(scopeguard::guard(collector, move |c| {
319                    let _ = registry.unregister(Box::new(c));
320                }))
321            }
322            Err(e) => {
323                crate::soft_panic_or_log!("collector already registered: {e}");
324                // Nothing was registered, so the handle must not unregister anything on drop.
325                Box::new(())
326            }
327        }
328    }
329
330    /// Registers a metric postprocessor.
331    ///
332    /// Postprocessors are invoked on every call to [`MetricsRegistry::gather`]
333    /// in the order that they are registered.
334    pub fn register_postprocessor<F>(&self, f: F)
335    where
336        F: FnMut(&mut Vec<MetricFamily>) + Send + Sync + 'static,
337    {
338        let mut postprocessors = self.postprocessors.lock().expect("lock poisoned");
339        postprocessors.push(Box::new(f));
340    }
341
342    /// Gather all the metrics from the metrics registry for reporting.
343    ///
344    /// This function invokes the postprocessors on all gathered metrics (see
345    /// [`MetricsRegistry::register_postprocessor`]) in the order the
346    /// postprocessors were registered.
347    ///
348    /// See also [`prometheus::Registry::gather`].
349    pub fn gather(&self) -> Vec<MetricFamily> {
350        let mut metrics = self.inner.gather();
351        let mut postprocessors = self.postprocessors.lock().expect("lock poisoned");
352        for postprocessor in &mut *postprocessors {
353            postprocessor(&mut metrics);
354        }
355        metrics
356    }
357}
358
359/// A wrapper for creating prometheus metrics more conveniently.
360///
361/// Together with the [`metric`] macro, this trait is mainly used by [`MetricsRegistry`] and should
362/// not normally be used outside the metric registration flow.
363pub trait MakeCollector: Collector + Clone + 'static {
364    /// Creates a new collector.
365    fn make_collector(opts: MakeCollectorOpts) -> Self;
366}
367
368impl<T> MakeCollector for GenericCounter<T>
369where
370    T: Atomic + 'static,
371{
372    fn make_collector(mk_opts: MakeCollectorOpts) -> Self {
373        assert_none!(mk_opts.buckets);
374        Self::with_opts(mk_opts.opts).expect("defining a counter")
375    }
376}
377
378impl<T> MakeCollector for GenericCounterVec<T>
379where
380    T: Atomic + 'static,
381{
382    fn make_collector(mk_opts: MakeCollectorOpts) -> Self {
383        assert_none!(mk_opts.buckets);
384        let labels: Vec<String> = mk_opts.opts.variable_labels.clone();
385        let label_refs: Vec<&str> = labels.iter().map(String::as_str).collect();
386        Self::new(mk_opts.opts, label_refs.as_slice()).expect("defining a counter vec")
387    }
388}
389
390impl<T> MakeCollector for GenericGauge<T>
391where
392    T: Atomic + 'static,
393{
394    fn make_collector(mk_opts: MakeCollectorOpts) -> Self {
395        assert_none!(mk_opts.buckets);
396        Self::with_opts(mk_opts.opts).expect("defining a gauge")
397    }
398}
399
400impl<T> MakeCollector for GenericGaugeVec<T>
401where
402    T: Atomic + 'static,
403{
404    fn make_collector(mk_opts: MakeCollectorOpts) -> Self {
405        assert_none!(mk_opts.buckets);
406        let labels = mk_opts.opts.variable_labels.clone();
407        let labels = &labels.iter().map(|x| x.as_str()).collect::<Vec<_>>();
408        Self::new(mk_opts.opts, labels).expect("defining a gauge vec")
409    }
410}
411
412impl MakeCollector for Histogram {
413    fn make_collector(mk_opts: MakeCollectorOpts) -> Self {
414        assert!(mk_opts.buckets.is_some());
415        Self::with_opts(HistogramOpts {
416            common_opts: mk_opts.opts,
417            buckets: mk_opts.buckets.unwrap(),
418        })
419        .expect("defining a histogram")
420    }
421}
422
423impl MakeCollector for raw::HistogramVec {
424    fn make_collector(mk_opts: MakeCollectorOpts) -> Self {
425        assert!(mk_opts.buckets.is_some());
426        let labels = mk_opts.opts.variable_labels.clone();
427        let labels = &labels.iter().map(|x| x.as_str()).collect::<Vec<_>>();
428        Self::new(
429            HistogramOpts {
430                common_opts: mk_opts.opts,
431                buckets: mk_opts.buckets.unwrap(),
432            },
433            labels,
434        )
435        .expect("defining a histogram vec")
436    }
437}
438
439/// A [`Gauge`] whose value is computed whenever it is observed.
440pub struct ComputedGenericGauge<P>
441where
442    P: Atomic,
443{
444    gauge: GenericGauge<P>,
445    f: Arc<dyn Fn() -> P::T + Send + Sync>,
446}
447
448impl<P> fmt::Debug for ComputedGenericGauge<P>
449where
450    P: Atomic + fmt::Debug,
451{
452    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
453        f.debug_struct("ComputedGenericGauge")
454            .field("gauge", &self.gauge)
455            .finish_non_exhaustive()
456    }
457}
458
459impl<P> Clone for ComputedGenericGauge<P>
460where
461    P: Atomic,
462{
463    fn clone(&self) -> ComputedGenericGauge<P> {
464        ComputedGenericGauge {
465            gauge: self.gauge.clone(),
466            f: Arc::clone(&self.f),
467        }
468    }
469}
470
471impl<T> Collector for ComputedGenericGauge<T>
472where
473    T: Atomic,
474{
475    fn desc(&self) -> Vec<&prometheus::core::Desc> {
476        self.gauge.desc()
477    }
478
479    fn collect(&self) -> Vec<MetricFamily> {
480        self.gauge.set((self.f)());
481        self.gauge.collect()
482    }
483}
484
485impl<P> ComputedGenericGauge<P>
486where
487    P: Atomic,
488{
489    /// Computes the current value of the gauge.
490    pub fn get(&self) -> P::T {
491        (self.f)()
492    }
493}
494
495/// A [`ComputedGenericGauge`] for 64-bit floating point numbers.
496pub type ComputedGauge = ComputedGenericGauge<AtomicF64>;
497
498/// A [`ComputedGenericGauge`] for 64-bit signed integers.
499pub type ComputedIntGauge = ComputedGenericGauge<AtomicI64>;
500
501/// A [`ComputedGenericGauge`] for 64-bit unsigned integers.
502pub type ComputedUIntGauge = ComputedGenericGauge<AtomicU64>;
503
504/// Exposes combinators that report metrics related to the execution of a [`Future`] to prometheus.
505pub trait MetricsFutureExt<F> {
506    /// Records the number of seconds it takes a [`Future`] to complete according to "the clock on
507    /// the wall".
508    ///
509    /// More specifically, it records the instant at which the `Future` was first polled, and the
510    /// instant at which the `Future` completes. Then reports the duration between those two
511    /// instances to the provided metric.
512    ///
513    /// # Wall Time vs Execution Time
514    ///
515    /// There is also [`MetricsFutureExt::exec_time`], which measures how long a [`Future`] spent
516    /// executing, instead of how long it took to complete. For example, a network request may have
517    /// a wall time of 1 second, meanwhile it's execution time may have only been 50ms. The 950ms
518    /// delta would be how long the [`Future`] waited for a response from the network.
519    ///
520    /// # Uses
521    ///
522    /// Recording the wall time can be useful for monitoring latency, for example the latency of a
523    /// SQL request.
524    ///
525    /// Note: You must call either [`observe`] to record the execution time to a [`Histogram`] or
526    /// [`inc_by`] to record to a [`Counter`]. The following will not compile:
527    ///
528    /// ```compile_fail
529    /// use mz_ore::metrics::MetricsFutureExt;
530    ///
531    /// # let _ = async {
532    /// async { Ok(()) }
533    ///     .wall_time()
534    ///     .await;
535    /// # };
536    /// ```
537    ///
538    /// [`observe`]: WallTimeFuture::observe
539    /// [`inc_by`]: WallTimeFuture::inc_by
540    fn wall_time(self) -> WallTimeFuture<F, UnspecifiedMetric>;
541
542    /// Records the total number of seconds for which a [`Future`] was executing.
543    ///
544    /// More specifically, every time the `Future` is polled it records how long that individual
545    /// call took, and maintains a running sum until the `Future` completes. Then we report that
546    /// duration to the provided metric.
547    ///
548    /// # Wall Time vs Execution Time
549    ///
550    /// There is also [`MetricsFutureExt::wall_time`], which measures how long a [`Future`] took to
551    /// complete, instead of how long it spent executing. For example, a network request may have
552    /// a wall time of 1 second, meanwhile it's execution time may have only been 50ms. The 950ms
553    /// delta would be how long the [`Future`] waited for a response from the network.
554    ///
555    /// # Uses
556    ///
557    /// Recording execution time can be useful if you want to monitor [`Future`]s that could be
558    /// sensitive to CPU usage. For example, if you have a single logical control thread you'll
559    /// want to make sure that thread never spends too long running a single `Future`. Reporting
560    /// the execution time of `Future`s running on this thread can help ensure there is no
561    /// unexpected blocking.
562    ///
563    /// Note: You must call either [`observe`] to record the execution time to a [`Histogram`] or
564    /// [`inc_by`] to record to a [`Counter`]. The following will not compile:
565    ///
566    /// ```compile_fail
567    /// use mz_ore::metrics::MetricsFutureExt;
568    ///
569    /// # let _ = async {
570    /// async { Ok(()) }
571    ///     .exec_time()
572    ///     .await;
573    /// # };
574    /// ```
575    ///
576    /// [`observe`]: ExecTimeFuture::observe
577    /// [`inc_by`]: ExecTimeFuture::inc_by
578    fn exec_time(self) -> ExecTimeFuture<F, UnspecifiedMetric>;
579}
580
581impl<F: Future> MetricsFutureExt<F> for F {
582    fn wall_time(self) -> WallTimeFuture<F, UnspecifiedMetric> {
583        WallTimeFuture {
584            fut: self,
585            metric: UnspecifiedMetric(()),
586            start: None,
587            filter: None,
588        }
589    }
590
591    fn exec_time(self) -> ExecTimeFuture<F, UnspecifiedMetric> {
592        ExecTimeFuture {
593            fut: self,
594            metric: UnspecifiedMetric(()),
595            running_duration: Duration::from_millis(0),
596            filter: None,
597        }
598    }
599}
600
601/// Future returned by [`MetricsFutureExt::wall_time`].
602#[must_use = "futures do nothing unless you `.await` or poll them"]
603#[pin_project]
604pub struct WallTimeFuture<F, Metric> {
605    /// The inner [`Future`] that we're recording the wall time for.
606    #[pin]
607    fut: F,
608    /// Prometheus metric that we'll report to.
609    metric: Metric,
610    /// [`Instant`] at which the [`Future`] was first polled.
611    start: Option<Instant>,
612    /// Optional filter that determines if we observe the wall time of this [`Future`].
613    filter: Option<Box<dyn FnMut(Duration) -> bool + Send + Sync>>,
614}
615
616impl<F: Debug, M: Debug> fmt::Debug for WallTimeFuture<F, M> {
617    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
618        f.debug_struct("WallTimeFuture")
619            .field("fut", &self.fut)
620            .field("metric", &self.metric)
621            .field("start", &self.start)
622            .field("filter", &self.filter.is_some())
623            .finish()
624    }
625}
626
627impl<F> WallTimeFuture<F, UnspecifiedMetric> {
628    /// Sets the recored metric to be a [`prometheus::Histogram`].
629    ///
630    /// ```text
631    /// my_future
632    ///     .wall_time()
633    ///     .observe(metrics.slow_queries_hist.with_label_values(&["select"]))
634    /// ```
635    pub fn observe(
636        self,
637        histogram: prometheus::Histogram,
638    ) -> WallTimeFuture<F, prometheus::Histogram> {
639        WallTimeFuture {
640            fut: self.fut,
641            metric: histogram,
642            start: self.start,
643            filter: self.filter,
644        }
645    }
646
647    /// Sets the recored metric to be a [`prometheus::Counter`].
648    ///
649    /// ```text
650    /// my_future
651    ///     .wall_time()
652    ///     .inc_by(metrics.slow_queries.with_label_values(&["select"]))
653    /// ```
654    pub fn inc_by(self, counter: prometheus::Counter) -> WallTimeFuture<F, prometheus::Counter> {
655        WallTimeFuture {
656            fut: self.fut,
657            metric: counter,
658            start: self.start,
659            filter: self.filter,
660        }
661    }
662
663    /// Sets the recorded duration in a specific f64.
664    pub fn set_at(self, place: &mut f64) -> WallTimeFuture<F, &mut f64> {
665        WallTimeFuture {
666            fut: self.fut,
667            metric: place,
668            start: self.start,
669            filter: self.filter,
670        }
671    }
672}
673
674impl<F, M> WallTimeFuture<F, M> {
675    /// Specifies a filter which much return `true` for the wall time to be recorded.
676    ///
677    /// This can be particularly useful if you have a high volume `Future` and you only want to
678    /// record ones that take a long time to complete.
679    pub fn with_filter(
680        mut self,
681        filter: impl FnMut(Duration) -> bool + Send + Sync + 'static,
682    ) -> Self {
683        self.filter = Some(Box::new(filter));
684        self
685    }
686}
687
688impl<F: Future, M: DurationMetric> Future for WallTimeFuture<F, M> {
689    type Output = F::Output;
690
691    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
692        let this = self.project();
693
694        if this.start.is_none() {
695            *this.start = Some(Instant::now());
696        }
697
698        let result = match this.fut.poll(cx) {
699            Poll::Ready(r) => r,
700            Poll::Pending => return Poll::Pending,
701        };
702        let duration = Instant::now().duration_since(this.start.expect("timer to be started"));
703
704        let pass = this
705            .filter
706            .as_mut()
707            .map(|filter| filter(duration))
708            .unwrap_or(true);
709        if pass {
710            this.metric.record(duration.as_secs_f64())
711        }
712
713        Poll::Ready(result)
714    }
715}
716
717/// Future returned by [`MetricsFutureExt::exec_time`].
718#[must_use = "futures do nothing unless you `.await` or poll them"]
719#[pin_project]
720pub struct ExecTimeFuture<F, Metric> {
721    /// The inner [`Future`] that we're recording the wall time for.
722    #[pin]
723    fut: F,
724    /// Prometheus metric that we'll report to.
725    metric: Metric,
726    /// Total [`Duration`] for which this [`Future`] has been executing.
727    running_duration: Duration,
728    /// Optional filter that determines if we observe the execution time of this [`Future`].
729    filter: Option<Box<dyn FnMut(Duration) -> bool + Send + Sync>>,
730}
731
732impl<F: Debug, M: Debug> fmt::Debug for ExecTimeFuture<F, M> {
733    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
734        f.debug_struct("ExecTimeFuture")
735            .field("fut", &self.fut)
736            .field("metric", &self.metric)
737            .field("running_duration", &self.running_duration)
738            .field("filter", &self.filter.is_some())
739            .finish()
740    }
741}
742
743impl<F> ExecTimeFuture<F, UnspecifiedMetric> {
744    /// Sets the recored metric to be a [`prometheus::Histogram`].
745    ///
746    /// ```text
747    /// my_future
748    ///     .exec_time()
749    ///     .observe(metrics.slow_queries_hist.with_label_values(&["select"]))
750    /// ```
751    pub fn observe(
752        self,
753        histogram: prometheus::Histogram,
754    ) -> ExecTimeFuture<F, prometheus::Histogram> {
755        ExecTimeFuture {
756            fut: self.fut,
757            metric: histogram,
758            running_duration: self.running_duration,
759            filter: self.filter,
760        }
761    }
762
763    /// Sets the recored metric to be a [`prometheus::Counter`].
764    ///
765    /// ```text
766    /// my_future
767    ///     .exec_time()
768    ///     .inc_by(metrics.slow_queries.with_label_values(&["select"]))
769    /// ```
770    pub fn inc_by(self, counter: prometheus::Counter) -> ExecTimeFuture<F, prometheus::Counter> {
771        ExecTimeFuture {
772            fut: self.fut,
773            metric: counter,
774            running_duration: self.running_duration,
775            filter: self.filter,
776        }
777    }
778}
779
780impl<F, M> ExecTimeFuture<F, M> {
781    /// Specifies a filter which much return `true` for the execution time to be recorded.
782    pub fn with_filter(
783        mut self,
784        filter: impl FnMut(Duration) -> bool + Send + Sync + 'static,
785    ) -> Self {
786        self.filter = Some(Box::new(filter));
787        self
788    }
789}
790
791impl<F: Future, M: DurationMetric> Future for ExecTimeFuture<F, M> {
792    type Output = F::Output;
793
794    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
795        let this = self.project();
796
797        let start = Instant::now();
798        let result = this.fut.poll(cx);
799        let duration = Instant::now().duration_since(start);
800
801        *this.running_duration = this.running_duration.saturating_add(duration);
802
803        let result = match result {
804            Poll::Ready(result) => result,
805            Poll::Pending => return Poll::Pending,
806        };
807
808        let duration = *this.running_duration;
809        let pass = this
810            .filter
811            .as_mut()
812            .map(|filter| filter(duration))
813            .unwrap_or(true);
814        if pass {
815            this.metric.record(duration.as_secs_f64());
816        }
817
818        Poll::Ready(result)
819    }
820}
821
822/// A type level flag used to ensure callers specify the kind of metric to record for
823/// [`MetricsFutureExt`].
824///
825/// For example, `WallTimeFuture<F, M>` only implements [`Future`] for `M` that implements
826/// `DurationMetric` which [`UnspecifiedMetric`] does not. This forces users at build time to
827/// call [`WallTimeFuture::observe`] or [`WallTimeFuture::inc_by`].
828#[derive(Debug)]
829pub struct UnspecifiedMetric(());
830
831/// A trait makes recording a duration generic over different prometheus metrics. This allows us to
832/// de-dupe the implemenation of [`Future`] for our wrapper Futures like [`WallTimeFuture`] and
833/// [`ExecTimeFuture`] over different kinds of prometheus metrics.
834trait DurationMetric {
835    fn record(&mut self, seconds: f64);
836}
837
838impl DurationMetric for prometheus::Histogram {
839    fn record(&mut self, seconds: f64) {
840        self.observe(seconds)
841    }
842}
843
844impl DurationMetric for prometheus::Counter {
845    fn record(&mut self, seconds: f64) {
846        self.inc_by(seconds)
847    }
848}
849
850// An implementation of `DurationMetric` that lets the user take the recorded
851// value and use it elsewhere.
852impl DurationMetric for &'_ mut f64 {
853    fn record(&mut self, seconds: f64) {
854        **self = seconds;
855    }
856}
857
858/// Register the Tokio runtime's metrics in our metrics registry.
859#[cfg(feature = "async")]
860pub fn register_runtime_metrics(
861    name: &'static str,
862    runtime_metrics: tokio::runtime::RuntimeMetrics,
863    registry: &MetricsRegistry,
864) {
865    macro_rules! register {
866        ($method:ident, $doc:literal) => {
867            let metrics = runtime_metrics.clone();
868            registry.register_computed_gauge::<prometheus::core::AtomicU64>(
869                crate::metric!(
870                    name: concat!("mz_tokio_", stringify!($method)),
871                    help: $doc,
872                    const_labels: {"runtime" => name},
873                ),
874                move || <u64 as crate::cast::CastFrom<_>>::cast_from(metrics.$method()),
875            );
876        };
877    }
878
879    macro_rules! register_per_worker {
880        ($method:ident, $doc:literal) => {
881            let metrics = runtime_metrics.clone();
882            registry.register_computed_gauge::<prometheus::core::AtomicU64>(
883                crate::metric!(
884                    name: concat!("mz_tokio_", stringify!($method)),
885                    help: $doc,
886                    const_labels: {"runtime" => name},
887                ),
888                move || {
889                    (0..metrics.num_workers())
890                        .map(|i| <u64 as crate::cast::CastFrom<_>>::cast_from(metrics.$method(i)))
891                        .sum::<u64>()
892                },
893            );
894        };
895    }
896
897    macro_rules! register_per_worker_duration_secs {
898        ($method:ident, $doc:literal) => {
899            let metrics = runtime_metrics.clone();
900            registry.register_computed_gauge::<prometheus::core::AtomicF64>(
901                crate::metric!(
902                    name: concat!("mz_tokio_", stringify!($method)),
903                    help: $doc,
904                    const_labels: {"runtime" => name},
905                ),
906                move || {
907                    (0..metrics.num_workers())
908                        .map(|i| metrics.$method(i).as_secs_f64())
909                        .sum::<f64>()
910                },
911            );
912        };
913    }
914
915    register!(
916        num_workers,
917        "The number of worker threads used by the runtime."
918    );
919    register!(
920        num_alive_tasks,
921        "The current number of alive tasks in the runtime."
922    );
923    register!(
924        global_queue_depth,
925        "The number of tasks currently scheduled in the runtime's global queue."
926    );
927    register_per_worker_duration_secs!(
928        worker_total_busy_duration,
929        "The amount of time the worker threads have been busy, in seconds."
930    );
931    register_per_worker!(
932        worker_park_count,
933        "The total number of times the worker threads have parked."
934    );
935    register_per_worker!(
936        worker_park_unpark_count,
937        "The total number of times the worker threads have parked and unparked."
938    );
939
940    #[cfg(tokio_unstable)]
941    {
942        register!(
943            num_blocking_threads,
944            "The number of additional threads spawned by the runtime."
945        );
946        register!(
947            num_idle_blocking_threads,
948            "The number of idle threads which have spawned by the runtime for spawn_blocking calls."
949        );
950        register_per_worker!(
951            worker_local_queue_depth,
952            "The number of tasks currently scheduled in the workers' local queues."
953        );
954        register!(
955            blocking_queue_depth,
956            "The number of tasks currently scheduled in the blocking thread pool, spawned using spawn_blocking."
957        );
958        register!(
959            spawned_tasks_count,
960            "The number of tasks spawned in this runtime since it was created."
961        );
962        register!(
963            remote_schedule_count,
964            "The number of tasks scheduled from outside of the runtime."
965        );
966        register!(
967            budget_forced_yield_count,
968            "The number of times that tasks have been forced to yield back to the scheduler after exhausting their task budgets."
969        );
970        register_per_worker!(
971            worker_noop_count,
972            "The number of times the given worker thread unparked but performed no work before parking again."
973        );
974        register_per_worker!(
975            worker_steal_count,
976            "The number of tasks the given worker thread stole from another worker thread."
977        );
978        register_per_worker!(
979            worker_steal_operations,
980            "The number of times the given worker thread stole tasks from another worker thread."
981        );
982        register_per_worker!(
983            worker_poll_count,
984            "The number of tasks the given worker thread has polled."
985        );
986        register_per_worker!(
987            worker_local_schedule_count,
988            "The number of tasks scheduled from within the runtime on the given worker's local queue."
989        );
990        register_per_worker!(
991            worker_overflow_count,
992            "The number of times the given worker thread saturated its local queue."
993        );
994        register_per_worker_duration_secs!(
995            worker_mean_poll_time,
996            "The mean duration of task polls in seconds."
997        );
998    }
999}
1000
1001/// Returns the `(name, help, labels, source)` of every Tokio runtime metric
1002/// registered by [`register_runtime_metrics`].
1003#[cfg(feature = "async")]
1004pub fn describe_runtime_metrics() -> Vec<(String, String, Vec<String>, &'static str)> {
1005    // A current-thread runtime is enough to enumerate the metrics; we only read
1006    // their names, help text, and labels, never their values.
1007    let runtime = tokio::runtime::Builder::new_current_thread()
1008        .build()
1009        .expect("building a current-thread runtime");
1010    let registry = MetricsRegistry::new();
1011    register_runtime_metrics("describe", runtime.handle().metrics(), &registry);
1012    registry
1013        .gather()
1014        .into_iter()
1015        .map(|mf| {
1016            // Every series in a family shares the same label keys, so the first
1017            // metric's labels are representative.
1018            let mut labels: Vec<String> = mf
1019                .get_metric()
1020                .first()
1021                .map(|m| m.get_label().iter().map(|l| l.name().to_owned()).collect())
1022                .unwrap_or_default();
1023            labels.sort();
1024            labels.dedup();
1025            (mf.name().to_owned(), mf.help().to_owned(), labels, file!())
1026        })
1027        .collect()
1028}
1029
1030#[cfg(test)]
1031mod tests {
1032    use std::time::Duration;
1033
1034    use prometheus::{CounterVec, HistogramVec};
1035
1036    use crate::stats::histogram_seconds_buckets;
1037
1038    use super::{MetricsFutureExt, MetricsRegistry};
1039
1040    struct Metrics {
1041        pub wall_time_hist: HistogramVec,
1042        pub wall_time_cnt: CounterVec,
1043        pub exec_time_hist: HistogramVec,
1044        pub exec_time_cnt: CounterVec,
1045    }
1046
1047    impl Metrics {
1048        pub fn register_into(registry: &MetricsRegistry) -> Self {
1049            Self {
1050                wall_time_hist: registry.register(metric!(
1051                    name: "wall_time_hist",
1052                    help: "help",
1053                    var_labels: ["action"],
1054                    buckets: histogram_seconds_buckets(0.000_128, 8.0),
1055                )),
1056                wall_time_cnt: registry.register(metric!(
1057                    name: "wall_time_cnt",
1058                    help: "help",
1059                    var_labels: ["action"],
1060                )),
1061                exec_time_hist: registry.register(metric!(
1062                    name: "exec_time_hist",
1063                    help: "help",
1064                    var_labels: ["action"],
1065                    buckets: histogram_seconds_buckets(0.000_128, 8.0),
1066                )),
1067                exec_time_cnt: registry.register(metric!(
1068                    name: "exec_time_cnt",
1069                    help: "help",
1070                    var_labels: ["action"],
1071                )),
1072            }
1073        }
1074    }
1075
1076    #[crate::test]
1077    #[cfg_attr(miri, ignore)] // unsupported operation: integer-to-pointer casts and `ptr::from_exposed_addr` are not supported with `-Zmiri-strict-provenance`
1078    fn smoke_test_metrics_future_ext() {
1079        let runtime = tokio::runtime::Builder::new_current_thread()
1080            .enable_time()
1081            .build()
1082            .expect("failed to start runtime");
1083        let registry = MetricsRegistry::new();
1084        let metrics = Metrics::register_into(&registry);
1085
1086        // Record the walltime and execution time of an async sleep.
1087        let async_sleep_future = async {
1088            tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
1089        };
1090        runtime.block_on(
1091            async_sleep_future
1092                .wall_time()
1093                .observe(metrics.wall_time_hist.with_label_values(&["async_sleep_w"]))
1094                .exec_time()
1095                .observe(metrics.exec_time_hist.with_label_values(&["async_sleep_e"])),
1096        );
1097
1098        let reports = registry.gather();
1099
1100        let exec_family = reports
1101            .iter()
1102            .find(|m| m.name() == "exec_time_hist")
1103            .expect("metric not found");
1104        let exec_metric = exec_family.get_metric();
1105        assert_eq!(exec_metric.len(), 1);
1106        assert_eq!(exec_metric[0].get_label()[0].value(), "async_sleep_e");
1107
1108        let exec_histogram = exec_metric[0].get_histogram();
1109        assert_eq!(exec_histogram.get_sample_count(), 1);
1110        // This future will normally complete very quickly, but it's hard to guarantee any particular
1111        // timing in an arbitrary test environment, so we don't assert on it here.
1112
1113        let wall_family = reports
1114            .iter()
1115            .find(|m| m.name() == "wall_time_hist")
1116            .expect("metric not found");
1117        let wall_metric = wall_family.get_metric();
1118        assert_eq!(wall_metric.len(), 1);
1119        assert_eq!(wall_metric[0].get_label()[0].value(), "async_sleep_w");
1120
1121        let wall_histogram = wall_metric[0].get_histogram();
1122        assert_eq!(wall_histogram.get_sample_count(), 1);
1123        // The 13th bucket is 512ms, which the wall time should be longer than, but is also much
1124        // faster than the actual execution time of the async sleep.
1125        assert_eq!(wall_histogram.get_bucket()[12].cumulative_count(), 0);
1126
1127        // Reset the registery to make collecting metrics easier.
1128        let registry = MetricsRegistry::new();
1129        let metrics = Metrics::register_into(&registry);
1130
1131        // Record the walltime and execution time of a thread sleep.
1132        let thread_sleep_future = async {
1133            std::thread::sleep(std::time::Duration::from_secs(1));
1134        };
1135        runtime.block_on(
1136            thread_sleep_future
1137                .wall_time()
1138                .with_filter(|duration| duration < Duration::from_millis(10))
1139                .inc_by(metrics.wall_time_cnt.with_label_values(&["thread_sleep_w"]))
1140                .exec_time()
1141                .inc_by(metrics.exec_time_cnt.with_label_values(&["thread_sleep_e"])),
1142        );
1143
1144        let reports = registry.gather();
1145
1146        let exec_family = reports
1147            .iter()
1148            .find(|m| m.name() == "exec_time_cnt")
1149            .expect("metric not found");
1150        let exec_metric = exec_family.get_metric();
1151        assert_eq!(exec_metric.len(), 1);
1152        assert_eq!(exec_metric[0].get_label()[0].value(), "thread_sleep_e");
1153
1154        let exec_counter = exec_metric[0].get_counter();
1155        // Since we're synchronously sleeping the execution time will be long.
1156        assert!(exec_counter.value() >= 1.0);
1157
1158        let wall_family = reports
1159            .iter()
1160            .find(|m| m.name() == "wall_time_cnt")
1161            .expect("metric not found");
1162        let wall_metric = wall_family.get_metric();
1163        assert_eq!(wall_metric.len(), 1);
1164
1165        let wall_counter = wall_metric[0].get_counter();
1166        // We filtered wall time to < 10ms, so our wall time metric should be filtered out.
1167        assert_eq!(wall_counter.value(), 0.0);
1168    }
1169
1170    #[crate::test]
1171    fn collector_drop_handle_unregisters() {
1172        use prometheus::IntGauge;
1173
1174        let registry = MetricsRegistry::new();
1175        let gauge = IntGauge::new("mz_test_guarded", "help").unwrap();
1176        gauge.set(7);
1177        let before = registry.gather().len();
1178
1179        let handle = registry.register_collector_with_dropper(gauge.clone());
1180        assert_eq!(registry.gather().len(), before + 1);
1181
1182        // Dropping the handle unregisters the collector, so its series stops being scraped.
1183        drop(handle);
1184        assert_eq!(registry.gather().len(), before);
1185    }
1186
1187    #[crate::test]
1188    fn register_drop_then_reregister() {
1189        use prometheus::IntGauge;
1190
1191        // Two collectors sharing a descriptor id: re-registering the second while the first is
1192        // still registered would collide. This locks in the contract that dropping the old handle
1193        // first clears the id so the re-registration succeeds cleanly, the ordering a caller
1194        // replacing a collector (e.g. a metric sink re-rendered on reconciliation) must uphold.
1195        let registry = MetricsRegistry::new();
1196        let old = IntGauge::new("mz_test_reregister", "help").unwrap();
1197        let new = IntGauge::new("mz_test_reregister", "help").unwrap();
1198        let before = registry.gather().len();
1199
1200        let handle = registry.register_collector_with_dropper(old);
1201        assert_eq!(registry.gather().len(), before + 1);
1202
1203        // Drop old before registering new, matching the required ordering.
1204        drop(handle);
1205        let handle = registry.register_collector_with_dropper(new);
1206        assert_eq!(registry.gather().len(), before + 1);
1207
1208        drop(handle);
1209        assert_eq!(registry.gather().len(), before);
1210    }
1211}
mz_ore/metrics.rs

mz_ore/
metrics.rs
