From 90f683650e46dc1c1ea46b554888639fb289708f Mon Sep 17 00:00:00 2001
From: Zak Cutner <zak@cloudflare.com>
Date: Tue, 28 May 2024 17:15:41 +0100
Subject: [PATCH] Provide a method to reset a counter

Signed-off-by: Zak Cutner <zak@cloudflare.com>
---
 src/metrics/counter.rs | 28 ++++++++++++++++++++++++++++
 1 file changed, 28 insertions(+)

diff --git a/src/metrics/counter.rs b/src/metrics/counter.rs
index c1c5e511..51f8e774 100644
--- a/src/metrics/counter.rs
+++ b/src/metrics/counter.rs
@@ -84,6 +84,11 @@ impl<N, A: Atomic<N>> Counter<N, A> {
         self.value.inc_by(v)
     }
 
+    /// Resets the [`Counter`] to `0`, returning the previous value.
+    pub fn reset(&self) -> N {
+        self.value.reset()
+    }
+
     /// Get the current value of the [`Counter`].
     pub fn get(&self) -> N {
         self.value.get()
@@ -110,6 +115,9 @@ pub trait Atomic<N> {
     /// Increase the value.
     fn inc_by(&self, v: N) -> N;
 
+    /// Reset the value to `0`.
+    fn reset(&self) -> N;
+
     /// Get the the value.
     fn get(&self) -> N;
 }
@@ -124,6 +132,10 @@ impl Atomic<u64> for AtomicU64 {
         self.fetch_add(v, Ordering::Relaxed)
     }
 
+    fn reset(&self) -> u64 {
+        self.swap(Default::default(), Ordering::Relaxed)
+    }
+
     fn get(&self) -> u64 {
         self.load(Ordering::Relaxed)
     }
@@ -138,6 +150,10 @@ impl Atomic<u32> for AtomicU32 {
         self.fetch_add(v, Ordering::Relaxed)
     }
 
+    fn reset(&self) -> u32 {
+        self.swap(Default::default(), Ordering::Relaxed)
+    }
+
     fn get(&self) -> u32 {
         self.load(Ordering::Relaxed)
     }
@@ -164,6 +180,10 @@ impl Atomic<f64> for AtomicU64 {
         old_f64
     }
 
+    fn reset(&self) -> f64 {
+        f64::from_bits(self.swap(Default::default(), Ordering::Relaxed))
+    }
+
     fn get(&self) -> f64 {
         f64::from_bits(self.load(Ordering::Relaxed))
     }
@@ -231,6 +251,14 @@ mod tests {
         assert_eq!(1, counter.get());
     }
 
+    #[test]
+    fn inc_reset_and_get() {
+        let counter: Counter = Counter::default();
+        assert_eq!(0, counter.inc());
+        assert_eq!(1, counter.reset());
+        assert_eq!(0, counter.get());
+    }
+
     #[cfg(not(any(target_arch = "mips", target_arch = "powerpc")))]
     #[test]
     fn f64_stored_in_atomic_u64() {