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;
@@ -182,13 +183,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 +219,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(),
+                });
+            }
         }
     }