refactor

Author: josecelano

packages/primitives/src/metrics/metric.rs

@@ -1,10 +1,12 @@
 use derive_more::Display;
 use serde::{Deserialize, Serialize};
 
+use super::counter::Counter;
 use super::label_set::LabelSet;
 use super::prometheus::PrometheusSerializable;
 use super::sample::Sample;
 use super::sample_collection::SampleCollection;
+use crate::DurationSinceUnixEpoch;
 
 #[derive(Debug, Clone, Default, PartialEq)]
 pub struct Metric<T> {
@@ -32,6 +34,12 @@ impl<T> Metric<T> {
     }
 }
 
+impl Metric<Counter> {
+    pub fn increment(&mut self, labels: &LabelSet, time: DurationSinceUnixEpoch) {
+        self.sample_collection.increment(labels, time);
+    }
+}
+
 impl<T: PrometheusSerializable> PrometheusSerializable for Metric<T> {
     fn to_prometheus(&self) -> String {
         let mut output = String::new();

packages/primitives/src/metrics/metric_collection.rs

@@ -158,31 +158,24 @@ impl<T> MetricKindCollection<T> {
 }
 
 impl MetricKindCollection<Counter> {
+    /// Increments the counter for the given metric name and labels.
+    ///
+    /// If the metric name does not exist, it will be created.
+    ///
     /// # Panics
     ///
     /// Panics if the metric name already exists in the collection.
     pub fn increment(&mut self, name: &MetricName, labels: &LabelSet, time: DurationSinceUnixEpoch) {
-        // Create the metric if it doesn't exist
         if !self.metrics.contains_key(name) {
-            let new_metric = Metric::new(name.clone(), SampleCollection { samples: HashMap::new() });
-            self.metrics.insert(name.clone(), new_metric);
+            self.metrics.insert(
+                name.clone(),
+                Metric::new(name.clone(), SampleCollection { samples: HashMap::new() }),
+            );
         }
 
         let metric = self.metrics.get_mut(name).unwrap();
 
-        // Use entry API to handle existing or new sample
-        let sample = metric
-            .sample_collection
-            .samples
-            .entry(labels.clone())
-            .or_insert_with(|| Sample {
-                value: Counter::new(0),
-                update_at: time,
-                labels: labels.clone(),
-            });
-
-        sample.value.increment(1);
-        sample.update_at = time;
+        metric.increment(labels, time);
     }
 
     #[must_use]
@@ -242,17 +235,33 @@ mod tests {
     use crate::metrics::sample::Sample;
 
     #[test]
-    fn test_increase_counter_in_metric_collection() {
+    fn it_should_increase_a_preexistent_counter() {
         let time = DurationSinceUnixEpoch::from_secs(1_743_552_000);
         let label_set: LabelSet = (LabelName::new("label_name"), LabelValue::new("value")).into();
 
-        let mut metric_collection = MetricCollection {
-            counters: MetricKindCollection::new(vec![Metric::new(
+        let mut metric_collection = MetricCollection::new(
+            MetricKindCollection::new(vec![Metric::new(
                 MetricName::new("test_counter"),
                 SampleCollection::new(vec![Sample::new(Counter::new(0), time, label_set.clone())]),
             )]),
-            gauges: MetricKindCollection::new(vec![]),
-        };
+            MetricKindCollection::new(vec![]),
+        );
+
+        metric_collection.increase_counter(&MetricName::new("test_counter"), &label_set, time);
+        metric_collection.increase_counter(&MetricName::new("test_counter"), &label_set, time);
+
+        assert_eq!(
+            metric_collection.get_counter_value(&MetricName::new("test_counter"), &label_set),
+            2
+        );
+    }
+
+    #[test]
+    fn it_should_automatically_create_a_counter_when_increasing_if_it_does_not_exist() {
+        let time = DurationSinceUnixEpoch::from_secs(1_743_552_000);
+        let label_set: LabelSet = (LabelName::new("label_name"), LabelValue::new("value")).into();
+
+        let mut metric_collection = MetricCollection::new(MetricKindCollection::new(vec![]), MetricKindCollection::new(vec![]));
 
         metric_collection.increase_counter(&MetricName::new("test_counter"), &label_set, time);
         metric_collection.increase_counter(&MetricName::new("test_counter"), &label_set, time);

packages/primitives/src/metrics/sample.rs

@@ -1,6 +1,8 @@
 use chrono::DateTime;
 use serde::{Serialize, Serializer};
 
+use super::counter::Counter;
+use super::gauge::Gauge;
 use super::label_set::LabelSet;
 use super::prometheus::PrometheusSerializable;
 use crate::DurationSinceUnixEpoch;
@@ -44,7 +46,7 @@ impl<T> Sample<T> {
         self.update_at
     }
 
-    pub fn set_update_at(&mut self, time: DurationSinceUnixEpoch) {
+    fn set_update_at(&mut self, time: DurationSinceUnixEpoch) {
         self.update_at = time;
     }
 }
@@ -55,6 +57,20 @@ impl<T: PrometheusSerializable> PrometheusSerializable for Sample<T> {
     }
 }
 
+impl Sample<Counter> {
+    pub fn increment(&mut self, time: DurationSinceUnixEpoch) {
+        self.value.increment(1);
+        self.set_update_at(time);
+    }
+}
+
+impl Sample<Gauge> {
+    pub fn set(&mut self, value: f64, time: DurationSinceUnixEpoch) {
+        self.value.set(value);
+        self.set_update_at(time);
+    }
+}
+
 /// Serializes the `update_at` field as a string in ISO 8601 format (RFC 3339).
 ///
 /// # Errors

packages/primitives/src/metrics/sample_collection.rs

@@ -2,9 +2,12 @@ use std::collections::HashMap;
 
 use serde::{Serialize, Serializer};
 
+use super::counter::Counter;
+use super::gauge::Gauge;
 use super::label_set::LabelSet;
 use super::prometheus::PrometheusSerializable;
 use super::sample::Sample;
+use crate::DurationSinceUnixEpoch;
 
 #[derive(Debug, Clone, Default, PartialEq)]
 pub struct SampleCollection<T> {
@@ -38,6 +41,28 @@ impl<T> SampleCollection<T> {
     }
 }
 
+impl SampleCollection<Counter> {
+    pub fn increment(&mut self, labels: &LabelSet, time: DurationSinceUnixEpoch) {
+        let sample = self
+            .samples
+            .entry(labels.clone())
+            .or_insert_with(|| Sample::new(Counter::new(0), time, labels.clone()));
+
+        sample.increment(time);
+    }
+}
+
+impl SampleCollection<Gauge> {
+    pub fn set(&mut self, labels: &LabelSet, value: f64, time: DurationSinceUnixEpoch) {
+        let sample = self
+            .samples
+            .entry(labels.clone())
+            .or_insert_with(|| Sample::new(Gauge::new(0.0), time, labels.clone()));
+
+        sample.set(value, time);
+    }
+}
+
 impl<T: Serialize> Serialize for SampleCollection<T> {
     fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
     where