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_vec,
8    register_int_counter_vec, HistogramVec, IntCounterVec, Opts,
9};
10
11use crate::time::Instant;
12
13const LINERA_NAMESPACE: &str = "linera";
14
15/// Wrapper around Prometheus register_int_counter_vec! macro which also sets the `linera` namespace
16pub fn register_int_counter_vec(
17    name: &str,
18    description: &str,
19    label_names: &[&str],
20) -> IntCounterVec {
21    let counter_opts = Opts::new(name, description).namespace(LINERA_NAMESPACE);
22    register_int_counter_vec!(counter_opts, label_names).expect("IntCounter can be created")
23}
24
25/// Wrapper around Prometheus `register_histogram_vec!` macro which also sets the `linera` namespace
26pub fn register_histogram_vec(
27    name: &str,
28    description: &str,
29    label_names: &[&str],
30    buckets: Option<Vec<f64>>,
31) -> HistogramVec {
32    let histogram_opts = if let Some(buckets) = buckets {
33        histogram_opts!(name, description, buckets).namespace(LINERA_NAMESPACE)
34    } else {
35        histogram_opts!(name, description).namespace(LINERA_NAMESPACE)
36    };
37
38    register_histogram_vec!(histogram_opts, label_names).expect("Histogram can be created")
39}
40
41/// Construct the bucket interval exponentially starting from a value and an ending value.
42pub fn exponential_bucket_interval(start_value: f64, end_value: f64) -> Option<Vec<f64>> {
43    let quot = end_value / start_value;
44    let factor = 3.0_f64;
45    let count_approx = quot.ln() / factor.ln();
46    let count = count_approx.round() as usize;
47    let mut buckets = exponential_buckets(start_value, factor, count)
48        .expect("Exponential buckets creation should not fail!");
49    if let Some(last) = buckets.last() {
50        if *last < end_value {
51            buckets.push(end_value);
52        }
53    }
54    Some(buckets)
55}
56
57/// Construct the latencies exponentially starting from 0.001 and ending at the maximum latency
58pub fn exponential_bucket_latencies(max_latency: f64) -> Option<Vec<f64>> {
59    exponential_bucket_interval(0.001_f64, max_latency)
60}
61
62/// Construct the bucket interval linearly starting from a value and an ending value.
63pub fn linear_bucket_interval(start_value: f64, width: f64, end_value: f64) -> Option<Vec<f64>> {
64    let count = (end_value - start_value) / width;
65    let count = count.round() as usize;
66    let mut buckets = linear_buckets(start_value, width, count)
67        .expect("Linear buckets creation should not fail!");
68    buckets.push(end_value);
69    Some(buckets)
70}
71
72/// A guard for an active latency measurement.
73///
74/// Finishes the measurement when dropped, and then updates the `Metric`.
75pub struct ActiveMeasurementGuard<'metric, Metric>
76where
77    Metric: MeasureLatency,
78{
79    start: Instant,
80    metric: Option<&'metric Metric>,
81}
82
83impl<Metric> ActiveMeasurementGuard<'_, Metric>
84where
85    Metric: MeasureLatency,
86{
87    /// Finishes the measurement, updates the `Metric` and returns the measured latency in
88    /// milliseconds.
89    pub fn finish(mut self) -> f64 {
90        self.finish_by_ref()
91    }
92
93    /// Finishes the measurement without taking ownership of this [`ActiveMeasurementGuard`],
94    /// updates the `Metric` and returns the measured latency in milliseconds.
95    fn finish_by_ref(&mut self) -> f64 {
96        match self.metric.take() {
97            Some(metric) => {
98                let latency = self.start.elapsed().as_secs_f64() * 1000.0;
99                metric.finish_measurement(latency);
100                latency
101            }
102            None => {
103                // This is getting called from `Drop` after `finish` has already been
104                // executed
105                f64::NAN
106            }
107        }
108    }
109}
110
111impl<Metric> Drop for ActiveMeasurementGuard<'_, Metric>
112where
113    Metric: MeasureLatency,
114{
115    fn drop(&mut self) {
116        self.finish_by_ref();
117    }
118}
119
120/// An extension trait for metrics that can be used to measure latencies.
121pub trait MeasureLatency: Sized {
122    /// Starts measuring the latency, finishing when the returned
123    /// [`ActiveMeasurementGuard`] is dropped.
124    fn measure_latency(&self) -> ActiveMeasurementGuard<'_, Self>;
125
126    /// Updates the metric with measured latency in `milliseconds`.
127    fn finish_measurement(&self, milliseconds: f64);
128}
129
130impl MeasureLatency for HistogramVec {
131    fn measure_latency(&self) -> ActiveMeasurementGuard<'_, Self> {
132        ActiveMeasurementGuard {
133            start: Instant::now(),
134            metric: Some(self),
135        }
136    }
137
138    fn finish_measurement(&self, milliseconds: f64) {
139        self.with_label_values(&[]).observe(milliseconds);
140    }
141}
142
143#[cfg(test)]
144mod tests {
145    use super::*;
146
147    // Helper function for approximate floating point comparison
148    fn assert_float_vec_eq(left: &[f64], right: &[f64]) {
149        const EPSILON: f64 = 1e-10;
150
151        assert_eq!(left.len(), right.len(), "Vectors have different lengths");
152        for (i, (l, r)) in left.iter().zip(right.iter()).enumerate() {
153            assert!(
154                (l - r).abs() < EPSILON,
155                "Vectors differ at index {}: {} != {}",
156                i,
157                l,
158                r
159            );
160        }
161    }
162
163    #[test]
164    fn test_linear_bucket_interval() {
165        // Case 1: Width divides range evenly - small values
166        let buckets = linear_bucket_interval(0.05, 0.01, 0.1).unwrap();
167        assert_float_vec_eq(&buckets, &[0.05, 0.06, 0.07, 0.08, 0.09, 0.1]);
168
169        // Case 2: Width divides range evenly - large values
170        let buckets = linear_bucket_interval(100.0, 50.0, 500.0).unwrap();
171        assert_float_vec_eq(
172            &buckets,
173            &[
174                100.0, 150.0, 200.0, 250.0, 300.0, 350.0, 400.0, 450.0, 500.0,
175            ],
176        );
177
178        // Case 3: Width doesn't divide range evenly - small values
179        let buckets = linear_bucket_interval(0.05, 0.12, 0.5).unwrap();
180        assert_float_vec_eq(&buckets, &[0.05, 0.17, 0.29, 0.41, 0.5]);
181
182        // Case 4: Width doesn't divide range evenly - large values
183        let buckets = linear_bucket_interval(100.0, 150.0, 500.0).unwrap();
184        assert_float_vec_eq(&buckets, &[100.0, 250.0, 400.0, 500.0]);
185    }
186}