From e2c37f4253349333387be659e8d44aa127e1c10d Mon Sep 17 00:00:00 2001
From: Dario Saccavino <kathoum@gmail.com>
Date: Sat, 9 Mar 2024 16:58:35 +0100
Subject: [PATCH] Add get functions to Registry (#457)

---
 metrics-util/src/registry/mod.rs | 26 ++++++++++++++++++++++++++
 1 file changed, 26 insertions(+)

diff --git a/metrics-util/src/registry/mod.rs b/metrics-util/src/registry/mod.rs
index 68a0e5f1..4ac6580c 100644
--- a/metrics-util/src/registry/mod.rs
+++ b/metrics-util/src/registry/mod.rs
@@ -389,6 +389,27 @@ where
         });
         histograms
     }
+
+    /// Gets a copy of an existing counter.
+    pub fn get_counter(&self, key: &K) -> Option<S::Counter> {
+        let (hash, shard) = self.get_hash_and_shard_for_counter(key);
+        let shard_read = shard.read().unwrap_or_else(PoisonError::into_inner);
+        shard_read.raw_entry().from_key_hashed_nocheck(hash, key).map(|(_, v)| v.clone())
+    }
+
+    /// Gets a copy of an existing gauge.
+    pub fn get_gauge(&self, key: &K) -> Option<S::Gauge> {
+        let (hash, shard) = self.get_hash_and_shard_for_gauge(key);
+        let shard_read = shard.read().unwrap_or_else(PoisonError::into_inner);
+        shard_read.raw_entry().from_key_hashed_nocheck(hash, key).map(|(_, v)| v.clone())
+    }
+
+    /// Gets a copy of an existing histogram.
+    pub fn get_histogram(&self, key: &K) -> Option<S::Histogram> {
+        let (hash, shard) = self.get_hash_and_shard_for_histogram(key);
+        let shard_read = shard.read().unwrap_or_else(PoisonError::into_inner);
+        shard_read.raw_entry().from_key_hashed_nocheck(hash, key).map(|(_, v)| v.clone())
+    }
 }
 
 #[cfg(test)]
@@ -406,6 +427,8 @@ mod tests {
         let entries = registry.get_counter_handles();
         assert_eq!(entries.len(), 0);
 
+        assert!(registry.get_counter(&key).is_none());
+
         registry.get_or_create_counter(&key, |c: &Arc<AtomicU64>| c.increment(1));
 
         let initial_entries = registry.get_counter_handles();
@@ -430,6 +453,9 @@ mod tests {
         assert_eq!(ukey, key);
         assert_eq!(uvalue.load(Ordering::SeqCst), 2);
 
+        let value = registry.get_counter(&key).expect("failed to get entry");
+        assert!(Arc::ptr_eq(&value, &uvalue));
+
         assert!(registry.delete_counter(&key));
 
         let entries = registry.get_counter_handles();