fix: insert metric automatically if it does not exist

Author: josecelano

packages/primitives/src/metrics/metric_collection.rs

@@ -7,6 +7,7 @@ use super::counter::Counter;
 use super::gauge::Gauge;
 use super::label_set::LabelSet;
 use super::metric::Metric;
+use super::sample::Sample;
 use super::MetricName;
 use crate::metrics::sample_collection::SampleCollection;
 use crate::DurationSinceUnixEpoch;
@@ -42,7 +43,15 @@ impl MetricCollection {
         self.counters.get_value(name, labels)
     }
 
+    /// # Panics
+    ///
+    /// Panics if a gauge with the same name already exists.
     pub fn increase_counter(&mut self, name: &MetricName, labels: &LabelSet, time: DurationSinceUnixEpoch) {
+        assert!(
+            !self.gauges.metrics.contains_key(name),
+            "Cannot create counter with name '{name}': a gauge with this name already exists",
+        );
+
         self.counters.increment(name, labels, time);
     }
 
@@ -53,7 +62,15 @@ impl MetricCollection {
         self.gauges.get_value(name, labels)
     }
 
+    /// # Panics
+    ///
+    /// Panics if a counter with the same name already exists.
     pub fn set_gauge(&mut self, name: &MetricName, labels: &LabelSet, value: f64, time: DurationSinceUnixEpoch) {
+        assert!(
+            !self.counters.metrics.contains_key(name),
+            "Cannot create gauge with name '{name}': a counter with this name already exists"
+        );
+
         self.gauges.set(name, labels, value, time);
     }
 
@@ -182,13 +199,30 @@ impl<T> MetricKindCollection<T> {
 }
 
 impl MetricKindCollection<Counter> {
+    /// # Panics
+    ///
+    /// Panics if the metric name already exists in the collection.
     pub fn increment(&mut self, name: &MetricName, labels: &LabelSet, time: DurationSinceUnixEpoch) {
-        if let Some(metric) = self.metrics.get_mut(name) {
-            if let Some(sample) = metric.samples.get_mut(labels) {
-                sample.value.increment(1);
-                sample.update_at = time;
-            }
+        // Create the metric if it doesn't exist
+        if !self.metrics.contains_key(name) {
+            let new_metric = Metric {
+                name: name.clone(),
+                samples: SampleCollection { samples: HashMap::new() },
+            };
+            self.metrics.insert(name.clone(), new_metric);
         }
+
+        let metric = self.metrics.get_mut(name).unwrap();
+
+        // Use entry API to handle existing or new sample
+        let sample = metric.samples.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;
     }
 
     #[must_use]
@@ -201,13 +235,35 @@ impl MetricKindCollection<Counter> {
 }
 
 impl MetricKindCollection<Gauge> {
+    /// # Panics
+    ///
+    /// Panics if the metric name already exists in the collection.
     pub fn set(&mut self, name: &MetricName, labels: &LabelSet, value: f64, time: DurationSinceUnixEpoch) {
-        if let Some(metric) = self.metrics.get_mut(name) {
-            // Direct hashmap lookup
-            if let Some(sample) = metric.samples.get_mut(labels) {
+        // Create the metric if it doesn't exist
+        if !self.metrics.contains_key(name) {
+            let new_metric = Metric {
+                name: name.clone(),
+                samples: SampleCollection { samples: HashMap::new() },
+            };
+            self.metrics.insert(name.clone(), new_metric);
+        }
+
+        let metric = self.metrics.get_mut(name).unwrap();
+
+        // Use entry API to handle existing or new sample
+        match metric.samples.samples.entry(labels.clone()) {
+            std::collections::hash_map::Entry::Occupied(mut entry) => {
+                let sample = entry.get_mut();
                 sample.value.set(value);
                 sample.update_at = time;
             }
+            std::collections::hash_map::Entry::Vacant(entry) => {
+                entry.insert(Sample {
+                    value: Gauge::new(value),
+                    update_at: time,
+                    labels: labels.clone(),
+                });
+            }
         }
     }