torrust_metrics/
sample_collection.rs

1use std::collections::HashMap;
2use std::collections::hash_map::Iter;
3use std::fmt::Write as _;
4
5use serde::{Deserialize, Deserializer, Serialize, Serializer};
6use torrust_clock::DurationSinceUnixEpoch;
7
8use super::counter::Counter;
9use super::gauge::Gauge;
10use super::label::LabelSet;
11use super::prometheus::PrometheusSerializable;
12use super::sample::Sample;
13use crate::sample::Measurement;
14
15#[derive(Debug, Clone, Default, PartialEq)]
16pub struct SampleCollection<T> {
17    samples: HashMap<LabelSet, Measurement<T>>,
18}
19
20impl<T> SampleCollection<T> {
21    /// Creates a new `MetricKindCollection` from a vector of metrics
22    ///
23    /// # Errors
24    ///
25    /// Returns an error if there are duplicate `LabelSets` in the provided
26    /// samples.
27    pub fn new(samples: Vec<Sample<T>>) -> Result<Self, Error> {
28        let mut map: HashMap<LabelSet, Measurement<T>> = HashMap::with_capacity(samples.len());
29
30        for sample in samples {
31            let (label_set, sample_data): (LabelSet, Measurement<T>) = sample.into();
32
33            let label_set_clone = label_set.clone();
34
35            if let Some(_old_measurement) = map.insert(label_set, sample_data) {
36                return Err(Error::DuplicateLabelSetInList {
37                    label_set: label_set_clone,
38                });
39            }
40        }
41
42        Ok(Self { samples: map })
43    }
44
45    #[must_use]
46    pub fn get(&self, label: &LabelSet) -> Option<&Measurement<T>> {
47        self.samples.get(label)
48    }
49
50    #[must_use]
51    pub fn len(&self) -> usize {
52        self.samples.len()
53    }
54
55    #[must_use]
56    pub fn is_empty(&self) -> bool {
57        self.samples.is_empty()
58    }
59
60    #[must_use]
61    #[allow(clippy::iter_without_into_iter)]
62    pub fn iter(&self) -> Iter<'_, LabelSet, Measurement<T>> {
63        self.samples.iter()
64    }
65}
66
67#[derive(thiserror::Error, Debug, Clone)]
68pub enum Error {
69    #[error("Found duplicate label set in list. Label set must be unique in a SampleCollection.")]
70    DuplicateLabelSetInList { label_set: LabelSet },
71}
72
73impl SampleCollection<Counter> {
74    pub fn increment(&mut self, label_set: &LabelSet, time: DurationSinceUnixEpoch) {
75        let sample = self
76            .samples
77            .entry(label_set.clone())
78            .or_insert_with(|| Measurement::new(Counter::default(), time));
79
80        sample.increment(time);
81    }
82
83    pub fn absolute(&mut self, label_set: &LabelSet, value: u64, time: DurationSinceUnixEpoch) {
84        let sample = self
85            .samples
86            .entry(label_set.clone())
87            .or_insert_with(|| Measurement::new(Counter::default(), time));
88
89        sample.absolute(value, time);
90    }
91}
92
93impl SampleCollection<Gauge> {
94    pub fn set(&mut self, label_set: &LabelSet, value: f64, time: DurationSinceUnixEpoch) {
95        let sample = self
96            .samples
97            .entry(label_set.clone())
98            .or_insert_with(|| Measurement::new(Gauge::default(), time));
99
100        sample.set(value, time);
101    }
102
103    pub fn increment(&mut self, label_set: &LabelSet, time: DurationSinceUnixEpoch) {
104        let sample = self
105            .samples
106            .entry(label_set.clone())
107            .or_insert_with(|| Measurement::new(Gauge::default(), time));
108
109        sample.increment(time);
110    }
111
112    pub fn decrement(&mut self, label_set: &LabelSet, time: DurationSinceUnixEpoch) {
113        let sample = self
114            .samples
115            .entry(label_set.clone())
116            .or_insert_with(|| Measurement::new(Gauge::default(), time));
117
118        sample.decrement(time);
119    }
120}
121
122impl<T: Serialize> Serialize for SampleCollection<T> {
123    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
124    where
125        S: Serializer,
126    {
127        let mut samples: Vec<Sample<&T>> = vec![];
128
129        for (label_set, sample_data) in &self.samples {
130            samples.push(Sample::new(sample_data.value(), sample_data.recorded_at(), label_set.clone()));
131        }
132
133        samples.serialize(serializer)
134    }
135}
136
137impl<'de, T> Deserialize<'de> for SampleCollection<T>
138where
139    T: Deserialize<'de>,
140{
141    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
142    where
143        D: Deserializer<'de>,
144    {
145        let samples = Vec::<Sample<T>>::deserialize(deserializer)?;
146
147        let sample_collection = SampleCollection::new(samples).map_err(serde::de::Error::custom)?;
148
149        Ok(sample_collection)
150    }
151}
152
153impl<T: PrometheusSerializable> PrometheusSerializable for SampleCollection<T> {
154    fn to_prometheus(&self) -> String {
155        let mut output = String::new();
156
157        for (label_set, sample_data) in &self.samples {
158            if label_set.is_empty() {
159                let _ = write!(output, "{}", sample_data.to_prometheus());
160            } else {
161                let _ = write!(output, "{} {}", label_set.to_prometheus(), sample_data.to_prometheus());
162            }
163        }
164
165        output
166    }
167}
168
169#[cfg(test)]
170mod tests {
171    use torrust_clock::DurationSinceUnixEpoch;
172
173    use crate::counter::Counter;
174    use crate::label::LabelSet;
175    use crate::sample::Sample;
176    use crate::sample_collection::SampleCollection;
177
178    fn sample_update_time() -> DurationSinceUnixEpoch {
179        DurationSinceUnixEpoch::from_secs(1_743_552_000)
180    }
181
182    #[test]
183    fn it_should_fail_trying_to_create_a_sample_collection_with_duplicate_label_sets() {
184        let samples = vec![
185            Sample::new(Counter::default(), sample_update_time(), LabelSet::default()),
186            Sample::new(Counter::default(), sample_update_time(), LabelSet::default()),
187        ];
188
189        let result = SampleCollection::new(samples);
190
191        assert!(result.is_err());
192    }
193
194    #[test]
195    fn it_should_return_a_sample_searching_by_label_set_with_one_empty_label_set() {
196        let label_set = LabelSet::default();
197
198        let sample = Sample::new(Counter::default(), sample_update_time(), label_set.clone());
199
200        let collection = SampleCollection::new(vec![sample.clone()]).unwrap();
201
202        let retrieved = collection.get(&label_set);
203
204        assert_eq!(retrieved.unwrap(), sample.measurement());
205    }
206
207    #[test]
208    fn it_should_return_a_sample_searching_by_label_set_with_two_label_sets() {
209        let label_set_1 = LabelSet::from(vec![("label_name_1", "label value 1")]);
210        let label_set_2 = LabelSet::from(vec![("label_name_2", "label value 2")]);
211
212        let sample_1 = Sample::new(Counter::new(1), sample_update_time(), label_set_1.clone());
213        let sample_2 = Sample::new(Counter::new(2), sample_update_time(), label_set_2.clone());
214
215        let collection = SampleCollection::new(vec![sample_1.clone(), sample_2.clone()]).unwrap();
216
217        let retrieved = collection.get(&label_set_1);
218        assert_eq!(retrieved.unwrap(), sample_1.measurement());
219
220        let retrieved = collection.get(&label_set_2);
221        assert_eq!(retrieved.unwrap(), sample_2.measurement());
222    }
223
224    #[test]
225    fn it_should_return_the_number_of_samples_in_the_collection() {
226        let samples = vec![Sample::new(Counter::default(), sample_update_time(), LabelSet::default())];
227        let collection = SampleCollection::new(samples).unwrap();
228        assert_eq!(collection.len(), 1);
229    }
230
231    #[test]
232    fn it_should_return_zero_number_of_samples_when_empty() {
233        let empty = SampleCollection::<Counter>::default();
234        assert_eq!(empty.len(), 0);
235    }
236
237    #[test]
238    fn it_should_indicate_is_it_is_empty() {
239        let empty = SampleCollection::<Counter>::default();
240        assert!(empty.is_empty());
241
242        let samples = vec![Sample::new(Counter::default(), sample_update_time(), LabelSet::default())];
243        let collection = SampleCollection::new(samples).unwrap();
244        assert!(!collection.is_empty());
245    }
246
247    #[test]
248    fn it_should_allow_iterating_samples() {
249        let label_set = LabelSet::from(vec![("key", "val")]);
250        let sample = Sample::new(Counter::new(5), sample_update_time(), label_set.clone());
251        let collection = SampleCollection::new(vec![sample]).unwrap();
252
253        let mut count = 0;
254        for (ls, meas) in collection.iter() {
255            assert_eq!(ls, &label_set);
256            assert_eq!(meas.value(), &Counter::new(5));
257            count += 1;
258        }
259        assert_eq!(count, 1);
260    }
261
262    mod json_serialization {
263        use crate::counter::Counter;
264        use crate::label::LabelSet;
265        use crate::sample::Sample;
266        use crate::sample_collection::SampleCollection;
267        use crate::sample_collection::tests::sample_update_time;
268
269        #[test]
270        fn it_should_be_serializable_and_deserializable_for_json_format() {
271            let sample = Sample::new(Counter::default(), sample_update_time(), LabelSet::default());
272            let collection = SampleCollection::new(vec![sample]).unwrap();
273
274            let serialized = serde_json::to_string(&collection).unwrap();
275            let deserialized: SampleCollection<Counter> = serde_json::from_str(&serialized).unwrap();
276
277            assert_eq!(deserialized, collection);
278        }
279
280        #[test]
281        fn it_should_fail_deserializing_from_json_with_duplicate_label_sets() {
282            let samples = vec![
283                Sample::new(Counter::default(), sample_update_time(), LabelSet::default()),
284                Sample::new(Counter::default(), sample_update_time(), LabelSet::default()),
285            ];
286
287            let serialized = serde_json::to_string(&samples).unwrap();
288
289            let result: Result<SampleCollection<Counter>, _> = serde_json::from_str(&serialized);
290
291            assert!(result.is_err());
292        }
293    }
294
295    mod prometheus_serialization {
296        use crate::counter::Counter;
297        use crate::label::LabelSet;
298        use crate::prometheus::PrometheusSerializable;
299        use crate::sample::Sample;
300        use crate::sample_collection::SampleCollection;
301        use crate::sample_collection::tests::sample_update_time;
302        use crate::tests::format_prometheus_output;
303
304        #[test]
305        fn it_should_be_exportable_to_prometheus_format_when_empty() {
306            let sample = Sample::new(Counter::default(), sample_update_time(), LabelSet::default());
307            let collection = SampleCollection::new(vec![sample]).unwrap();
308
309            let prometheus_output = collection.to_prometheus();
310
311            assert!(!prometheus_output.is_empty());
312        }
313
314        #[test]
315        fn it_should_be_exportable_to_prometheus_format() {
316            let sample = Sample::new(
317                Counter::new(1),
318                sample_update_time(),
319                LabelSet::from(vec![("labe_name_1", "label value value 1")]),
320            );
321
322            let collection = SampleCollection::new(vec![sample]).unwrap();
323
324            let prometheus_output = collection.to_prometheus();
325
326            let expected_prometheus_output = format_prometheus_output("{labe_name_1=\"label value value 1\"} 1");
327
328            assert_eq!(prometheus_output, expected_prometheus_output);
329        }
330    }
331
332    #[cfg(test)]
333    mod for_counters {
334
335        use std::ops::Add;
336
337        use super::super::LabelSet;
338        use super::*;
339
340        #[test]
341        fn it_should_increment_the_counter_for_a_preexisting_label_set() {
342            let label_set = LabelSet::default();
343            let mut collection = SampleCollection::<Counter>::default();
344
345            // Initialize the sample
346            collection.increment(&label_set, sample_update_time());
347
348            // Verify initial state
349            let sample = collection.get(&label_set).unwrap();
350            assert_eq!(sample.value(), &Counter::new(1));
351
352            // Increment again
353            collection.increment(&label_set, sample_update_time());
354            let sample = collection.get(&label_set).unwrap();
355            assert_eq!(*sample.value(), Counter::new(2));
356        }
357
358        #[test]
359        fn it_should_allow_increment_the_counter_for_a_non_existent_label_set() {
360            let label_set = LabelSet::default();
361            let mut collection = SampleCollection::<Counter>::default();
362
363            // Increment a non-existent label
364            collection.increment(&label_set, sample_update_time());
365
366            // Verify the label exists
367            assert!(collection.get(&label_set).is_some());
368            let sample = collection.get(&label_set).unwrap();
369            assert_eq!(*sample.value(), Counter::new(1));
370        }
371
372        #[test]
373        fn it_should_update_the_latest_update_time_when_incremented() {
374            let label_set = LabelSet::default();
375            let initial_time = sample_update_time();
376
377            let mut collection = SampleCollection::<Counter>::default();
378            collection.increment(&label_set, initial_time);
379
380            // Increment with a new time
381            let new_time = initial_time.add(DurationSinceUnixEpoch::from_secs(1));
382            collection.increment(&label_set, new_time);
383
384            let sample = collection.get(&label_set).unwrap();
385            assert_eq!(sample.recorded_at(), new_time);
386            assert_eq!(*sample.value(), Counter::new(2));
387        }
388
389        #[test]
390        fn it_should_increment_the_counter_for_multiple_labels() {
391            let label1 = LabelSet::from([("name", "value1")]);
392            let label2 = LabelSet::from([("name", "value2")]);
393            let now = sample_update_time();
394
395            let mut collection = SampleCollection::<Counter>::default();
396
397            collection.increment(&label1, now);
398            collection.increment(&label2, now);
399
400            assert_eq!(collection.get(&label1).unwrap().value(), &Counter::new(1));
401            assert_eq!(collection.get(&label2).unwrap().value(), &Counter::new(1));
402            assert_eq!(collection.len(), 2);
403        }
404
405        #[test]
406        fn it_should_allow_setting_absolute_value_for_a_counter() {
407            let label_set = LabelSet::default();
408            let mut collection = SampleCollection::<Counter>::default();
409
410            // Set absolute value for a non-existent label
411            collection.absolute(&label_set, 42, sample_update_time());
412
413            // Verify the label exists and has the absolute value
414            assert!(collection.get(&label_set).is_some());
415            let sample = collection.get(&label_set).unwrap();
416            assert_eq!(*sample.value(), Counter::new(42));
417        }
418
419        #[test]
420        fn it_should_allow_setting_absolute_value_for_existing_counter() {
421            let label_set = LabelSet::default();
422            let mut collection = SampleCollection::<Counter>::default();
423
424            // Initialize the sample with increment
425            collection.increment(&label_set, sample_update_time());
426
427            // Verify initial state
428            let sample = collection.get(&label_set).unwrap();
429            assert_eq!(sample.value(), &Counter::new(1));
430
431            // Set absolute value
432            collection.absolute(&label_set, 100, sample_update_time());
433            let sample = collection.get(&label_set).unwrap();
434            assert_eq!(*sample.value(), Counter::new(100));
435        }
436
437        #[test]
438        fn it_should_update_time_when_setting_absolute_value() {
439            let label_set = LabelSet::default();
440            let initial_time = sample_update_time();
441            let mut collection = SampleCollection::<Counter>::default();
442
443            // Set absolute value with initial time
444            collection.absolute(&label_set, 50, initial_time);
445
446            // Set absolute value with a new time
447            let new_time = initial_time.add(DurationSinceUnixEpoch::from_secs(1));
448            collection.absolute(&label_set, 75, new_time);
449
450            let sample = collection.get(&label_set).unwrap();
451            assert_eq!(sample.recorded_at(), new_time);
452            assert_eq!(*sample.value(), Counter::new(75));
453        }
454    }
455
456    #[cfg(test)]
457    mod for_gauges {
458
459        use std::ops::Add;
460
461        use super::super::LabelSet;
462        use super::*;
463        use crate::gauge::Gauge;
464
465        #[test]
466        fn it_should_allow_setting_the_gauge_for_a_preexisting_label_set() {
467            let label_set = LabelSet::default();
468            let mut collection = SampleCollection::<Gauge>::default();
469
470            // Initialize the sample
471            collection.set(&label_set, 1.0, sample_update_time());
472
473            // Verify initial state
474            let sample = collection.get(&label_set).unwrap();
475            assert_eq!(sample.value(), &Gauge::new(1.0));
476
477            // Set again
478            collection.set(&label_set, 2.0, sample_update_time());
479            let sample = collection.get(&label_set).unwrap();
480            assert_eq!(*sample.value(), Gauge::new(2.0));
481        }
482
483        #[test]
484        fn it_should_allow_setting_the_gauge_for_a_non_existent_label_set() {
485            let label_set = LabelSet::default();
486            let mut collection = SampleCollection::<Gauge>::default();
487
488            // Set a non-existent label
489            collection.set(&label_set, 1.0, sample_update_time());
490
491            // Verify the label exists
492            assert!(collection.get(&label_set).is_some());
493            let sample = collection.get(&label_set).unwrap();
494            assert_eq!(*sample.value(), Gauge::new(1.0));
495        }
496
497        #[test]
498        fn it_should_update_the_latest_update_time_when_setting() {
499            let label_set = LabelSet::default();
500            let initial_time = sample_update_time();
501
502            let mut collection = SampleCollection::<Gauge>::default();
503            collection.set(&label_set, 1.0, initial_time);
504
505            // Set with a new time
506            let new_time = initial_time.add(DurationSinceUnixEpoch::from_secs(1));
507            collection.set(&label_set, 2.0, new_time);
508
509            let sample = collection.get(&label_set).unwrap();
510            assert_eq!(sample.recorded_at(), new_time);
511            assert_eq!(*sample.value(), Gauge::new(2.0));
512        }
513
514        #[test]
515        fn it_should_allow_setting_the_gauge_for_multiple_labels() {
516            let label1 = LabelSet::from([("name", "value1")]);
517            let label2 = LabelSet::from([("name", "value2")]);
518            let now = sample_update_time();
519
520            let mut collection = SampleCollection::<Gauge>::default();
521
522            collection.set(&label1, 1.0, now);
523            collection.set(&label2, 2.0, now);
524
525            assert_eq!(collection.get(&label1).unwrap().value(), &Gauge::new(1.0));
526            assert_eq!(collection.get(&label2).unwrap().value(), &Gauge::new(2.0));
527            assert_eq!(collection.len(), 2);
528        }
529
530        #[test]
531        fn it_should_allow_incrementing_the_gauge() {
532            let label_set = LabelSet::default();
533            let mut collection = SampleCollection::<Gauge>::default();
534
535            // Initialize the sample
536            collection.set(&label_set, 1.0, sample_update_time());
537
538            // Increment
539            collection.increment(&label_set, sample_update_time());
540            let sample = collection.get(&label_set).unwrap();
541            assert_eq!(*sample.value(), Gauge::new(2.0));
542        }
543
544        #[test]
545        fn it_should_allow_decrementing_the_gauge() {
546            let label_set = LabelSet::default();
547            let mut collection = SampleCollection::<Gauge>::default();
548
549            // Initialize the sample
550            collection.set(&label_set, 1.0, sample_update_time());
551
552            // Increment
553            collection.decrement(&label_set, sample_update_time());
554            let sample = collection.get(&label_set).unwrap();
555            assert_eq!(*sample.value(), Gauge::new(0.0));
556        }
557
558        #[test]
559        fn it_should_create_a_default_gauge_when_decrementing_a_nonexistent_label_set() {
560            let label_set = LabelSet::default();
561            let mut collection = SampleCollection::<Gauge>::default();
562
563            // Decrement without prior set or increment — triggers the or_insert_with path
564            collection.decrement(&label_set, sample_update_time());
565            let sample = collection.get(&label_set).unwrap();
566            assert_eq!(*sample.value(), Gauge::new(-1.0));
567        }
568    }
569}
torrust_metrics/sample_collection.rs

torrust_metrics/
sample_collection.rs
