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 pub fn new(labels: Vec<Vec<(&'static str, &'static str)>>) -> Self {
84 Self::with_le(labels, &BASIC_LE_CALCULATOR)
85 }
86 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 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 pub fn new(labels: Vec<(&'static str, &'static str)>) -> Self {
114 Self::with_le(labels, &BASIC_LE_CALCULATOR[..])
115 }
116 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 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 pub fn get_sample_count(&self) -> u64 {
133 self.count.load(Ordering::Relaxed)
134 }
135 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}