usiem/components/metrics/
histogram.rs

1use std::sync::{
2    atomic::{AtomicI64, AtomicU64, Ordering},
3    Arc,
4};
5
6use serde::{
7    ser::{SerializeSeq, SerializeStruct},
8    Serialize, Serializer,
9};
10
11use super::prometheus::Encoder;
12
13pub const BASIC_LE_CALCULATOR: [f64; 10] =
14    [0.001, 0.005, 0.01, 0.05, 0.1, 0.5, 1.0, 2.0, 5.0, 10.0];
15
16pub const HISTOGRAM_MODIFIER: i64 = 100_000;
17
18#[derive(Debug, Clone)]
19pub struct HistogramVec {
20    pub metrics: Vec<Histogram>,
21}
22
23#[derive(Debug, Clone)]
24pub struct Histogram {
25    pub sum: Arc<AtomicI64>,
26    pub count: Arc<AtomicU64>,
27    pub counters: HistogramBucket,
28    pub labels: Vec<(&'static str, &'static str)>,
29}
30
31#[derive(Debug, Clone)]
32pub struct HistogramBucket {
33    pub observations: Vec<Observation>,
34}
35
36#[derive(Debug, Clone)]
37pub struct Observation {
38    pub le: f64,
39    pub value: Arc<AtomicU64>,
40}
41impl Observation {
42    pub fn new(le: f64) -> Self {
43        Self {
44            value: Arc::new(AtomicU64::new(0)),
45            le,
46        }
47    }
48}
49
50impl HistogramBucket {
51    pub fn new(le: &[f64]) -> Self {
52        Self {
53            observations: le.iter().map(|v| Observation::new(*v)).collect(),
54        }
55    }
56    pub fn observe(&self, value: f64) {
57        for observation in &self.observations {
58            if value > observation.le {
59                continue;
60            }
61            observation.value.fetch_add(1, Ordering::SeqCst);
62        }
63    }
64    pub fn get_count(&self, le: f64) -> Option<u64> {
65        for observation in &self.observations {
66            if le >= observation.le {
67                return Some(observation.value.load(Ordering::Relaxed));
68            }
69        }
70        None
71    }
72}
73
74fn normalize_f64(numb: f64) -> i64 {
75    (numb * (HISTOGRAM_MODIFIER as f64)) as i64
76}
77fn normalize_i64(numb: i64) -> f64 {
78    numb as f64 / HISTOGRAM_MODIFIER as f64
79}
80
81impl HistogramVec {
82    /// Creates a new histogram with static labels. The metric won't accept a new label once created.
83    pub fn new(labels: Vec<Vec<(&'static str, &'static str)>>) -> Self {
84        Self::with_le(labels, &BASIC_LE_CALCULATOR)
85    }
86    /// Creates a new histogram with static labels and custom values for the LE (lesser than or equals) parameter. The metric won't accept a new label once created.
87    pub fn with_le(labels: Vec<Vec<(&'static str, &'static str)>>, le: &[f64]) -> Self {
88        let mut metrics = Vec::with_capacity(labels.len());
89        for label in labels {
90            metrics.push(Histogram::with_le(label, le));
91        }
92        Self { metrics }
93    }
94    /// Obtains a histogram with the specified labels
95    pub fn with_labels(&self, labels: &[(&str, &str)]) -> Option<&Histogram> {
96        'cnt: for hist in &self.metrics {
97            if hist.labels.len() != labels.len() {
98                continue;
99            }
100            for ((name1, value1), (name2, value2)) in hist.labels.iter().zip(labels.iter()) {
101                if name1 != name2 || value1 != value2 {
102                    continue 'cnt;
103                }
104            }
105            return Some(hist);
106        }
107        None
108    }
109}
110
111impl Histogram {
112    /// Creates a new Histogram with static labels and the default le params: [0.001, 0.005, 0.01, 0.05, 0.1, 0.5, 1.0, 2.0, 5.0, 10.0]
113    pub fn new(labels: Vec<(&'static str, &'static str)>) -> Self {
114        Self::with_le(labels, &BASIC_LE_CALCULATOR[..])
115    }
116    /// Creates a new Histogram with static labels and a custom LE
117    pub fn with_le(labels: Vec<(&'static str, &'static str)>, le: &[f64]) -> Self {
118        Self {
119            sum: Arc::new(AtomicI64::new(0)),
120            count: Arc::new(AtomicU64::new(0)),
121            counters: HistogramBucket::new(le),
122            labels,
123        }
124    }
125    /// Observe a value
126    pub fn observe(&self, value: f64) {
127        self.count.fetch_add(1, Ordering::SeqCst);
128        self.sum.fetch_add(normalize_f64(value), Ordering::SeqCst);
129        self.counters.observe(value);
130    }
131    /// Get the number of samples
132    pub fn get_sample_count(&self) -> u64 {
133        self.count.load(Ordering::Relaxed)
134    }
135    /// Get the acumulated sample value
136    pub fn get_sample_sum(&self) -> f64 {
137        normalize_i64(self.sum.load(Ordering::Relaxed))
138    }
139}
140
141impl Encoder for HistogramVec {
142    fn encode<W: std::fmt::Write>(
143        &self,
144        f: &mut W,
145        name: &str,
146        description: &str,
147        help: bool,
148    ) -> Result<(), std::fmt::Error> {
149        if self.metrics.is_empty() {
150            return Ok(());
151        }
152        if help {
153            f.write_str("# HELP ")?;
154            f.write_str(name)?;
155            f.write_str(" ")?;
156            f.write_str(description)?;
157            f.write_str("\n")?;
158        }
159        f.write_str("# TYPE ")?;
160        f.write_str(name)?;
161        f.write_str(" histogram \n")?;
162        for counter in &self.metrics {
163            counter.encode(f, name, description, help)?;
164        }
165        Ok(())
166    }
167}
168
169impl Encoder for Histogram {
170    fn encode<W: std::fmt::Write>(
171        &self,
172        f: &mut W,
173        name: &str,
174        _description: &str,
175        _help: bool,
176    ) -> Result<(), std::fmt::Error> {
177        let count = self.get_sample_count();
178        if count == 0 {
179            return Ok(());
180        }
181        for observation in &self.counters.observations {
182            f.write_str(name)?;
183            f.write_str("_bucket{")?;
184            for (name, value) in &self.labels {
185                if value.is_empty() {
186                    continue;
187                }
188                f.write_str(name)?;
189                f.write_str("=\"")?;
190                f.write_str(value)?;
191                f.write_str("\",")?;
192            }
193            f.write_str("le=")?;
194            f.write_fmt(format_args!("{}", observation.le))?;
195            f.write_str("} ")?;
196            f.write_fmt(format_args!(
197                "{}",
198                observation.value.load(Ordering::Relaxed)
199            ))?;
200            f.write_str("\n")?;
201        }
202        f.write_str(name)?;
203        f.write_str("_sum{")?;
204        let mut i = 0;
205        for (name, value) in &self.labels {
206            i += 1;
207            if value.is_empty() {
208                continue;
209            }
210            f.write_str(name)?;
211            f.write_str("=\"")?;
212            f.write_str(value)?;
213            f.write_str("\"")?;
214            if i != self.labels.len() {
215                f.write_str(",")?;
216            }
217        }
218        f.write_str("} ")?;
219        f.write_fmt(format_args!("{}", self.get_sample_sum()))?;
220        f.write_str("\n")?;
221        f.write_str(name)?;
222        f.write_str("_count{")?;
223        let mut i = 0;
224        for (name, value) in &self.labels {
225            i += 1;
226            if value.is_empty() {
227                continue;
228            }
229            f.write_str(name)?;
230            f.write_str("=\"")?;
231            f.write_str(value)?;
232            f.write_str("\"")?;
233            if i != self.labels.len() {
234                f.write_str(",")?;
235            }
236        }
237        f.write_str("} ")?;
238        f.write_fmt(format_args!("{}", count))?;
239        f.write_str("\n")?;
240        Ok(())
241    }
242}
243
244impl Serialize for HistogramVec {
245    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
246    where
247        S: Serializer,
248    {
249        let mut state = serializer.serialize_seq(Some(self.metrics.len()))?;
250        for metric in &self.metrics {
251            state.serialize_element(metric)?;
252        }
253        state.end()
254    }
255}
256
257impl Serialize for Histogram {
258    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
259    where
260        S: Serializer,
261    {
262        let mut state = serializer.serialize_struct("Histogram", 4)?;
263        state.serialize_field("labels", &self.labels)?;
264        state.serialize_field("count", &self.get_sample_count())?;
265        state.serialize_field("sum", &self.get_sample_sum())?;
266        state.serialize_field("buckets", &self.counters)?;
267        state.end()
268    }
269}
270
271impl Serialize for HistogramBucket {
272    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
273    where
274        S: Serializer,
275    {
276        let mut state = serializer.serialize_seq(Some(self.observations.len()))?;
277        for observation in &self.observations {
278            state.serialize_element(observation)?;
279        }
280        state.end()
281    }
282}
283
284impl Serialize for Observation {
285    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
286    where
287        S: Serializer,
288    {
289        let mut state = serializer.serialize_struct("Observation", 2)?;
290        state.serialize_field("v", &self.value.load(Ordering::Relaxed))?;
291        state.serialize_field("le", &self.le)?;
292        state.end()
293    }
294}
295#[cfg(test)]
296mod tst {
297    use crate::components::metrics::{
298        histogram::HistogramVec, label_combinations, prometheus::Encoder,
299    };
300    #[test]
301    fn should_create_complex_static_metric() {
302        let name_values = vec!["a", "b", "c"];
303        let v1_values = vec!["d", "e", "f"];
304        let v2_values = vec!["g", "h", "i"];
305
306        let labels = vec![
307            ("name", &name_values[..]),
308            ("v1", &v1_values[..]),
309            ("v2", &v2_values[..]),
310        ];
311        let labels = label_combinations(&labels[..]);
312        let metric = HistogramVec::new(labels);
313        let hist = metric
314            .with_labels(&[("name", "a"), ("v1", "d"), ("v2", "g")])
315            .unwrap();
316        hist.observe(0.001);
317        assert_eq!(1, hist.get_sample_count());
318        assert_eq!(0.001, hist.get_sample_sum());
319
320        assert_eq!(Some(1), hist.counters.get_count(0.001));
321    }
322
323    #[test]
324    fn gauge_should_be_encoded_in_prometheus() {
325        let name_values = vec!["a", "b", "c"];
326        let v1_values = vec!["d", "e", "f"];
327        let v2_values = vec!["g", "h", "i"];
328
329        let labels = vec![
330            ("name", &name_values[..]),
331            ("v1", &v1_values[..]),
332            ("v2", &v2_values[..]),
333        ];
334        let labels = label_combinations(&labels[..]);
335        let metric = HistogramVec::new(labels);
336        let hist = metric
337            .with_labels(&[("name", "a"), ("v1", "d"), ("v2", "g")])
338            .unwrap();
339        hist.observe(0.001);
340        let mut st = String::with_capacity(1_000_000);
341        metric
342            .encode(&mut st, "simple_gauge", "Simple Gauge metric", true)
343            .unwrap();
344        assert_eq!(
345            r#"# HELP simple_gauge Simple Gauge metric
346# TYPE simple_gauge histogram 
347simple_gauge_bucket{name="a",v1="d",v2="g",le=0.001} 1
348simple_gauge_bucket{name="a",v1="d",v2="g",le=0.005} 1
349simple_gauge_bucket{name="a",v1="d",v2="g",le=0.01} 1
350simple_gauge_bucket{name="a",v1="d",v2="g",le=0.05} 1
351simple_gauge_bucket{name="a",v1="d",v2="g",le=0.1} 1
352simple_gauge_bucket{name="a",v1="d",v2="g",le=0.5} 1
353simple_gauge_bucket{name="a",v1="d",v2="g",le=1} 1
354simple_gauge_bucket{name="a",v1="d",v2="g",le=2} 1
355simple_gauge_bucket{name="a",v1="d",v2="g",le=5} 1
356simple_gauge_bucket{name="a",v1="d",v2="g",le=10} 1
357simple_gauge_sum{name="a",v1="d",v2="g"} 0.001
358simple_gauge_count{name="a",v1="d",v2="g"} 1
359"#,
360            st
361        );
362    }
363
364    #[test]
365    fn gauge_should_be_encoded_in_prometheus_without_label_name() {
366        let name_values = vec!["", "b", "c"];
367        let v1_values = vec!["d", "e", "f"];
368        let v2_values = vec!["g", "h", "i"];
369
370        let labels = vec![
371            ("name", &name_values[..]),
372            ("v1", &v1_values[..]),
373            ("v2", &v2_values[..]),
374        ];
375        let labels = label_combinations(&labels[..]);
376        let metric = HistogramVec::new(labels);
377        let hist = metric
378            .with_labels(&[("name", ""), ("v1", "d"), ("v2", "g")])
379            .unwrap();
380        hist.observe(0.001);
381        let mut st = String::with_capacity(1_000_000);
382        metric
383            .encode(&mut st, "simple_gauge", "Simple Gauge metric", true)
384            .unwrap();
385        assert_eq!(
386            r#"# HELP simple_gauge Simple Gauge metric
387# TYPE simple_gauge histogram 
388simple_gauge_bucket{v1="d",v2="g",le=0.001} 1
389simple_gauge_bucket{v1="d",v2="g",le=0.005} 1
390simple_gauge_bucket{v1="d",v2="g",le=0.01} 1
391simple_gauge_bucket{v1="d",v2="g",le=0.05} 1
392simple_gauge_bucket{v1="d",v2="g",le=0.1} 1
393simple_gauge_bucket{v1="d",v2="g",le=0.5} 1
394simple_gauge_bucket{v1="d",v2="g",le=1} 1
395simple_gauge_bucket{v1="d",v2="g",le=2} 1
396simple_gauge_bucket{v1="d",v2="g",le=5} 1
397simple_gauge_bucket{v1="d",v2="g",le=10} 1
398simple_gauge_sum{v1="d",v2="g"} 0.001
399simple_gauge_count{v1="d",v2="g"} 1
400"#,
401            st
402        );
403    }
404}