linera_base/
prometheus_util.rs

1// Copyright (c) Zefchain Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4//! This module defines utility functions for interacting with Prometheus (logging metrics, etc)
5
6use prometheus::{
7    exponential_buckets, histogram_opts, linear_buckets, register_histogram,
8    register_histogram_vec, register_int_counter, register_int_counter_vec, register_int_gauge,
9    register_int_gauge_vec, Histogram, HistogramVec, IntCounter, IntCounterVec, IntGauge,
10    IntGaugeVec, Opts,
11};
12
13use crate::time::Instant;
14
15const LINERA_NAMESPACE: &str = "linera";
16
17/// Wrapper around Prometheus `register_int_counter_vec!` macro which also sets the `linera` namespace
18pub fn register_int_counter_vec(
19    name: &str,
20    description: &str,
21    label_names: &[&str],
22) -> IntCounterVec {
23    let counter_opts = Opts::new(name, description).namespace(LINERA_NAMESPACE);
24    register_int_counter_vec!(counter_opts, label_names).expect("IntCounter can be created")
25}
26
27/// Wrapper around Prometheus `register_int_counter!` macro which also sets the `linera` namespace
28pub fn register_int_counter(name: &str, description: &str) -> IntCounter {
29    let counter_opts = Opts::new(name, description).namespace(LINERA_NAMESPACE);
30    register_int_counter!(counter_opts).expect("IntCounter can be created")
31}
32
33/// Wrapper around Prometheus `register_histogram_vec!` macro which also sets the `linera` namespace
34pub fn register_histogram_vec(
35    name: &str,
36    description: &str,
37    label_names: &[&str],
38    buckets: Option<Vec<f64>>,
39) -> HistogramVec {
40    let histogram_opts = if let Some(buckets) = buckets {
41        histogram_opts!(name, description, buckets).namespace(LINERA_NAMESPACE)
42    } else {
43        histogram_opts!(name, description).namespace(LINERA_NAMESPACE)
44    };
45
46    register_histogram_vec!(histogram_opts, label_names).expect("Histogram can be created")
47}
48
49/// Wrapper around Prometheus `register_histogram!` macro which also sets the `linera` namespace
50pub fn register_histogram(name: &str, description: &str, buckets: Option<Vec<f64>>) -> Histogram {
51    let histogram_opts = if let Some(buckets) = buckets {
52        histogram_opts!(name, description, buckets).namespace(LINERA_NAMESPACE)
53    } else {
54        histogram_opts!(name, description).namespace(LINERA_NAMESPACE)
55    };
56
57    register_histogram!(histogram_opts).expect("Histogram can be created")
58}
59
60/// Wrapper around Prometheus `register_int_gauge!` macro which also sets the `linera` namespace
61pub fn register_int_gauge(name: &str, description: &str) -> IntGauge {
62    let gauge_opts = Opts::new(name, description).namespace(LINERA_NAMESPACE);
63    register_int_gauge!(gauge_opts).expect("IntGauge can be created")
64}
65
66/// Wrapper around Prometheus `register_int_gauge_vec!` macro which also sets the `linera` namespace
67pub fn register_int_gauge_vec(name: &str, description: &str, label_names: &[&str]) -> IntGaugeVec {
68    let gauge_opts = Opts::new(name, description).namespace(LINERA_NAMESPACE);
69    register_int_gauge_vec!(gauge_opts, label_names).expect("IntGauge can be created")
70}
71
72/// Construct the bucket interval exponentially starting from a value and an ending value.
73pub fn exponential_bucket_interval(start_value: f64, end_value: f64) -> Option<Vec<f64>> {
74    let quot = end_value / start_value;
75    let factor = 3.0_f64;
76    let count_approx = quot.ln() / factor.ln();
77    let count = count_approx.round() as usize;
78    let mut buckets = exponential_buckets(start_value, factor, count)
79        .expect("Exponential buckets creation should not fail!");
80    if let Some(last) = buckets.last() {
81        if *last < end_value {
82            buckets.push(end_value);
83        }
84    }
85    Some(buckets)
86}
87
88/// Construct the latencies exponentially starting from 0.001 and ending at the maximum latency
89pub fn exponential_bucket_latencies(max_latency: f64) -> Option<Vec<f64>> {
90    exponential_bucket_interval(0.001_f64, max_latency)
91}
92
93/// Construct the bucket interval linearly starting from a value and an ending value.
94pub fn linear_bucket_interval(start_value: f64, width: f64, end_value: f64) -> Option<Vec<f64>> {
95    let count = (end_value - start_value) / width;
96    let count = count.round() as usize;
97    let mut buckets = linear_buckets(start_value, width, count)
98        .expect("Linear buckets creation should not fail!");
99    buckets.push(end_value);
100    Some(buckets)
101}
102
103/// The unit of measurement for latency metrics.
104enum MeasurementUnit {
105    /// Measure latency in milliseconds.
106    Milliseconds,
107    /// Measure latency in microseconds.
108    Microseconds,
109}
110
111/// A guard for an active latency measurement.
112///
113/// Finishes the measurement when dropped, and then updates the `Metric`.
114pub struct ActiveMeasurementGuard<'metric, Metric>
115where
116    Metric: MeasureLatency,
117{
118    start: Instant,
119    metric: Option<&'metric Metric>,
120    unit: MeasurementUnit,
121}
122
123impl<Metric> ActiveMeasurementGuard<'_, Metric>
124where
125    Metric: MeasureLatency,
126{
127    /// Finishes the measurement, updates the `Metric` and returns the measured latency in
128    /// the unit specified when the measurement was started.
129    pub fn finish(mut self) -> f64 {
130        self.finish_by_ref()
131    }
132
133    /// Finishes the measurement without taking ownership of this [`ActiveMeasurementGuard`],
134    /// updates the `Metric` and returns the measured latency in the unit specified when
135    /// the measurement was started.
136    fn finish_by_ref(&mut self) -> f64 {
137        match self.metric.take() {
138            Some(metric) => {
139                let latency = match self.unit {
140                    MeasurementUnit::Milliseconds => self.start.elapsed().as_secs_f64() * 1000.0,
141                    MeasurementUnit::Microseconds => {
142                        self.start.elapsed().as_secs_f64() * 1_000_000.0
143                    }
144                };
145                metric.finish_measurement(latency);
146                latency
147            }
148            None => {
149                // This is getting called from `Drop` after `finish` has already been
150                // executed
151                f64::NAN
152            }
153        }
154    }
155}
156
157impl<Metric> Drop for ActiveMeasurementGuard<'_, Metric>
158where
159    Metric: MeasureLatency,
160{
161    fn drop(&mut self) {
162        self.finish_by_ref();
163    }
164}
165
166/// An extension trait for metrics that can be used to measure latencies.
167pub trait MeasureLatency: Sized {
168    /// Starts measuring the latency in milliseconds, finishing when the returned
169    /// [`ActiveMeasurementGuard`] is dropped.
170    fn measure_latency(&self) -> ActiveMeasurementGuard<'_, Self>;
171
172    /// Starts measuring the latency in microseconds, finishing when the returned
173    /// [`ActiveMeasurementGuard`] is dropped.
174    fn measure_latency_us(&self) -> ActiveMeasurementGuard<'_, Self>;
175
176    /// Updates the metric with measured latency in `milliseconds`.
177    fn finish_measurement(&self, milliseconds: f64);
178}
179
180impl MeasureLatency for HistogramVec {
181    fn measure_latency(&self) -> ActiveMeasurementGuard<'_, Self> {
182        ActiveMeasurementGuard {
183            start: Instant::now(),
184            metric: Some(self),
185            unit: MeasurementUnit::Milliseconds,
186        }
187    }
188
189    fn measure_latency_us(&self) -> ActiveMeasurementGuard<'_, Self> {
190        ActiveMeasurementGuard {
191            start: Instant::now(),
192            metric: Some(self),
193            unit: MeasurementUnit::Microseconds,
194        }
195    }
196
197    fn finish_measurement(&self, milliseconds: f64) {
198        self.with_label_values(&[]).observe(milliseconds);
199    }
200}
201
202impl MeasureLatency for Histogram {
203    fn measure_latency(&self) -> ActiveMeasurementGuard<'_, Self> {
204        ActiveMeasurementGuard {
205            start: Instant::now(),
206            metric: Some(self),
207            unit: MeasurementUnit::Milliseconds,
208        }
209    }
210
211    fn measure_latency_us(&self) -> ActiveMeasurementGuard<'_, Self> {
212        ActiveMeasurementGuard {
213            start: Instant::now(),
214            metric: Some(self),
215            unit: MeasurementUnit::Microseconds,
216        }
217    }
218
219    fn finish_measurement(&self, milliseconds: f64) {
220        self.observe(milliseconds);
221    }
222}
223
224#[cfg(test)]
225mod tests {
226    use super::*;
227
228    // Helper function for approximate floating point comparison
229    fn assert_float_vec_eq(left: &[f64], right: &[f64]) {
230        const EPSILON: f64 = 1e-10;
231
232        assert_eq!(left.len(), right.len(), "Vectors have different lengths");
233        for (i, (l, r)) in left.iter().zip(right.iter()).enumerate() {
234            assert!(
235                (l - r).abs() < EPSILON,
236                "Vectors differ at index {}: {} != {}",
237                i,
238                l,
239                r
240            );
241        }
242    }
243
244    #[test]
245    fn test_linear_bucket_interval() {
246        // Case 1: Width divides range evenly - small values
247        let buckets = linear_bucket_interval(0.05, 0.01, 0.1).unwrap();
248        assert_float_vec_eq(&buckets, &[0.05, 0.06, 0.07, 0.08, 0.09, 0.1]);
249
250        // Case 2: Width divides range evenly - large values
251        let buckets = linear_bucket_interval(100.0, 50.0, 500.0).unwrap();
252        assert_float_vec_eq(
253            &buckets,
254            &[
255                100.0, 150.0, 200.0, 250.0, 300.0, 350.0, 400.0, 450.0, 500.0,
256            ],
257        );
258
259        // Case 3: Width doesn't divide range evenly - small values
260        let buckets = linear_bucket_interval(0.05, 0.12, 0.5).unwrap();
261        assert_float_vec_eq(&buckets, &[0.05, 0.17, 0.29, 0.41, 0.5]);
262
263        // Case 4: Width doesn't divide range evenly - large values
264        let buckets = linear_bucket_interval(100.0, 150.0, 500.0).unwrap();
265        assert_float_vec_eq(&buckets, &[100.0, 250.0, 400.0, 500.0]);
266    }
267}