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