1use 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
17pub 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
27pub 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
33pub 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
49pub 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
60pub 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
66pub 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
72pub 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
88pub fn exponential_bucket_latencies(max_latency: f64) -> Option<Vec<f64>> {
90 exponential_bucket_interval(0.001_f64, max_latency)
91}
92
93pub 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
103enum MeasurementUnit {
105 Milliseconds,
107 Microseconds,
109}
110
111pub 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 pub fn finish(mut self) -> f64 {
130 self.finish_by_ref()
131 }
132
133 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 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
166pub trait MeasureLatency: Sized {
168 fn measure_latency(&self) -> ActiveMeasurementGuard<'_, Self>;
171
172 fn measure_latency_us(&self) -> ActiveMeasurementGuard<'_, Self>;
175
176 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 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 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 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 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 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}