From 93c63dbcdecda2aaca331a1df9746f9774483800 Mon Sep 17 00:00:00 2001
From: Wenbo Zhang <ethercflow@gmail.com>
Date: Tue, 3 Dec 2024 16:49:36 +0800
Subject: [PATCH] feat(metrics/family): add `get` method to `Family` (#234)

Signed-off-by: Wenbo Zhang <wenbo.zhang@iomesh.com>
---
 CHANGELOG.md          |  4 +++
 src/metrics/family.rs | 64 +++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 65 insertions(+), 3 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 945d128f..685c8b66 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -34,9 +34,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Supoort `Arc<String>` for `EncodeLabelValue`.
   See [PR 217].
 
+- Added `get` method to `Family`.
+  See [PR 234].
+
 [PR 173]: https://github.com/prometheus/client_rust/pull/173
 [PR 216]: https://github.com/prometheus/client_rust/pull/216
 [PR 217]: https://github.com/prometheus/client_rust/pull/217
+[PR 234]: https://github.com/prometheus/client_rust/pull/234
 
 ### Fixed
 
diff --git a/src/metrics/family.rs b/src/metrics/family.rs
index 2f23b198..1a76cf8f 100644
--- a/src/metrics/family.rs
+++ b/src/metrics/family.rs
@@ -226,9 +226,7 @@ impl<S: Clone + std::hash::Hash + Eq, M, C: MetricConstructor<M>> Family<S, M, C
     /// family.get_or_create(&vec![("method".to_owned(), "GET".to_owned())]).inc();
     /// ```
     pub fn get_or_create(&self, label_set: &S) -> MappedRwLockReadGuard<M> {
-        if let Ok(metric) =
-            RwLockReadGuard::try_map(self.metrics.read(), |metrics| metrics.get(label_set))
-        {
+        if let Some(metric) = self.get(label_set) {
             return metric;
         }
 
@@ -247,6 +245,23 @@ impl<S: Clone + std::hash::Hash + Eq, M, C: MetricConstructor<M>> Family<S, M, C
         })
     }
 
+    /// Access a metric with the given label set, returning None if one
+    /// does not yet exist.
+    ///
+    /// ```
+    /// # use prometheus_client::metrics::counter::{Atomic, Counter};
+    /// # use prometheus_client::metrics::family::Family;
+    /// #
+    /// let family = Family::<Vec<(String, String)>, Counter>::default();
+    ///
+    /// if let Some(metric) = family.get(&vec![("method".to_owned(), "GET".to_owned())]) {
+    ///     metric.inc();
+    /// };
+    /// ```
+    pub fn get(&self, label_set: &S) -> Option<MappedRwLockReadGuard<M>> {
+        RwLockReadGuard::try_map(self.metrics.read(), |metrics| metrics.get(label_set)).ok()
+    }
+
     /// Remove a label set from the metric family.
     ///
     /// Returns a bool indicating if a label set was removed or not.
@@ -452,4 +467,47 @@ mod tests {
                 .get()
         );
     }
+
+    #[test]
+    fn test_get() {
+        let family = Family::<Vec<(String, String)>, Counter>::default();
+
+        // Test getting a non-existent metric.
+        let non_existent = family.get(&vec![("method".to_string(), "GET".to_string())]);
+        assert!(non_existent.is_none());
+
+        // Create a metric.
+        family
+            .get_or_create(&vec![("method".to_string(), "GET".to_string())])
+            .inc();
+
+        // Test getting an existing metric.
+        let existing = family.get(&vec![("method".to_string(), "GET".to_string())]);
+        assert!(existing.is_some());
+        assert_eq!(existing.unwrap().get(), 1);
+
+        // Test getting a different non-existent metric.
+        let another_non_existent = family.get(&vec![("method".to_string(), "POST".to_string())]);
+        assert!(another_non_existent.is_none());
+
+        // Test modifying the metric through the returned reference.
+        if let Some(metric) = family.get(&vec![("method".to_string(), "GET".to_string())]) {
+            metric.inc();
+        }
+
+        // Verify the modification.
+        let modified = family.get(&vec![("method".to_string(), "GET".to_string())]);
+        assert_eq!(modified.unwrap().get(), 2);
+
+        // Test with a different label set type.
+        let string_family = Family::<String, Counter>::default();
+        string_family.get_or_create(&"test".to_string()).inc();
+
+        let string_metric = string_family.get(&"test".to_string());
+        assert!(string_metric.is_some());
+        assert_eq!(string_metric.unwrap().get(), 1);
+
+        let non_existent_string = string_family.get(&"non_existent".to_string());
+        assert!(non_existent_string.is_none());
+    }
 }