Provide owning Idle guard

jonhoo · jonhoo · commit b64160c4fd5f · 2019-06-26T16:06:27.000-04:00
diff --git a/Cargo.toml b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "hdrhistogram"
-version = "6.3.1"
+version = "6.3.2"
 edition = "2018"
 
 description = "A port of HdrHistogram to Rust"
diff --git a/src/sync/mod.rs b/src/sync/mod.rs
@@ -3,6 +3,8 @@
 use crate::errors::*;
 use crate::{Counter, Histogram};
 use std::borrow::Borrow;
+use std::borrow::BorrowMut;
+use std::marker::PhantomData;
 use std::ops::{AddAssign, Deref, DerefMut};
 use std::sync::{atomic, Arc, Mutex};
 use std::time;
@@ -88,15 +90,36 @@ struct Shared<C: Counter> {
     phase: atomic::AtomicUsize,
 }
 
+/// See [`IdleRecorder`]. This guard borrows the idle [`Recorder`].
+pub type IdleRecorderGuard<'a, C> = IdleRecorder<&'a mut Recorder<C>, C>;
+
 /// This guard denotes that a [`Recorder`] is currently idle, and should not be waited on by a
 /// [`SyncHistogram`] phase-shift.
-pub struct IdleRecorderGuard<'a, C: Counter>(&'a mut Recorder<C>);
+pub struct IdleRecorder<T, C: Counter>
+where
+    T: BorrowMut<Recorder<C>>,
+{
+    recorder: Option<T>,
+    c: PhantomData<C>,
+}
+
+impl<T, C: Counter> IdleRecorder<T, C>
+where
+    T: BorrowMut<Recorder<C>>,
+{
+    fn reactivate(&mut self) {
+        let recorder = if let Some(ref mut r) = self.recorder {
+            r
+        } else {
+            // already reactivated
+            return;
+        };
+
+        let recorder = recorder.borrow_mut();
 
-impl<'a, C: Counter> Drop for IdleRecorderGuard<'a, C> {
-    fn drop(&mut self) {
         // the Recorder is no longer idle, so the reader has to wait for us again
         // this basically means re-incrementing .recorders
-        let mut crit = self.0.shared.truth.lock().unwrap();
+        let mut crit = recorder.shared.truth.lock().unwrap();
         crit.recorders += 1;
 
         // we need to figure out what phase we're joining
@@ -107,13 +130,35 @@ impl<'a, C: Counter> Drop for IdleRecorderGuard<'a, C> {
         // to send), and bump the phase, all before we read it, which would lead us to believe that
         // we were already synchronized when in reality we were not, which would stall the reader
         // even if we issued more writes.
-        self.0.last_phase = self.0.shared.phase.load(atomic::Ordering::Acquire);
+        recorder.last_phase = recorder.shared.phase.load(atomic::Ordering::Acquire);
 
         // explicitly drop guard to ensure we don't accidentally drop it above
         drop(crit);
     }
 }
 
+impl<C: Counter> IdleRecorder<Recorder<C>, C> {
+    /// Mark the wrapped [`Recorder`] as active again and return it.
+    pub fn activate(mut self) -> Recorder<C> {
+        self.reactivate();
+        self.recorder.take().unwrap()
+    }
+
+    /// Clone the wrapped [`Recorder`].
+    pub fn recorder(&self) -> Recorder<C> {
+        self.recorder.as_ref().unwrap().clone()
+    }
+}
+
+impl<T, C: Counter> Drop for IdleRecorder<T, C>
+where
+    T: BorrowMut<Recorder<C>>,
+{
+    fn drop(&mut self) {
+        self.reactivate()
+    }
+}
+
 impl<C: Counter> Recorder<C> {
     fn with_hist<F, R>(&mut self, f: F) -> R
     where
@@ -139,11 +184,7 @@ impl<C: Counter> Recorder<C> {
         let _ = self.shared.sender.send(h).is_ok(); // if this is err, the reader went away
     }
 
-    /// Call this method if the Recorder will be idle for a while.
-    ///
-    /// Until the returned guard is dropped, the associated [`SyncHistogram`] will not wait for
-    /// this recorder on a phase shift.
-    pub fn idle(&mut self) -> IdleRecorderGuard<C> {
+    fn deactivate(&mut self) {
         let phase;
         {
             // we're leaving rotation, so we need to decrement .recorders
@@ -160,8 +201,31 @@ impl<C: Counter> Recorder<C> {
             }
         }
         self.last_phase = phase;
+    }
 
-        IdleRecorderGuard(self)
+    /// Call this method if the Recorder will be idle for a while.
+    ///
+    /// Until the returned guard is dropped, the associated [`SyncHistogram`] will not wait for
+    /// this recorder on a phase shift.
+    pub fn idle(&mut self) -> IdleRecorderGuard<C> {
+        self.deactivate();
+        IdleRecorder {
+            recorder: Some(self),
+            c: PhantomData,
+        }
+    }
+
+    /// Mark this `Recorder` as inactive.
+    ///
+    /// Until the returned guard is consumed, either by calling [`IdleRecorder::activate`] or by
+    /// dropping it, the associated [`SyncHistogram`] will not wait for this recorder on a phase
+    /// shift.
+    pub fn into_idle(mut self) -> IdleRecorder<Self, C> {
+        self.deactivate();
+        IdleRecorder {
+            recorder: Some(self),
+            c: PhantomData,
+        }
     }
 
     /// See [`Histogram::add`].
diff --git a/tests/sync.rs b/tests/sync.rs
@@ -1,7 +1,7 @@
 #[cfg(all(feature = "sync", test))]
 mod sync {
     use hdrhistogram::{sync::SyncHistogram, Histogram};
-    use std::sync::Arc;
+    use std::sync::{atomic, Arc};
     use std::{thread, time};
 
     const TRACKABLE_MAX: u64 = 3600 * 1000 * 1000;
@@ -214,6 +214,51 @@ mod sync {
         jh.join().unwrap();
     }
 
+    #[test]
+    fn clone_idle_recorder() {
+        let mut h: SyncHistogram<_> = Histogram::<u64>::new_with_max(TRACKABLE_MAX, SIGFIG)
+            .unwrap()
+            .into();
+
+        let done = Arc::new(atomic::AtomicBool::new(false));
+        let r = h.recorder().into_idle();
+        h.refresh(); // this should not block
+        h.refresh(); // nor should this
+        let mut r2 = r.recorder();
+        let d = Arc::clone(&done);
+        let jh = thread::spawn(move || {
+            let mut i = 0;
+            while !d.load(atomic::Ordering::SeqCst) {
+                r2 += TEST_VALUE_LEVEL;
+                i += 1;
+            }
+            i
+        });
+        h.refresh(); // this is released by the r2 += above
+        let mut r = r.activate();
+        // a call to refresh would block here now
+        let d = Arc::clone(&done);
+        let jh2 = thread::spawn(move || {
+            let mut i = 0;
+            while !d.load(atomic::Ordering::SeqCst) {
+                r += TEST_VALUE_LEVEL;
+                i += 1;
+            }
+            i
+        });
+
+        h.refresh(); // this is released by the second r2 _and_ the r += above
+
+        // tell recorders to exit
+        done.store(true, atomic::Ordering::SeqCst);
+        h.refresh(); // shouldn't block for long
+        let n = jh.join().unwrap() + jh2.join().unwrap();
+        h.refresh(); // no more recorders, so shouldn't block
+
+        assert_eq!(h.count_at(TEST_VALUE_LEVEL), n);
+        assert_eq!(h.len(), n);
+    }
+
     #[test]
     fn concurrent_writes() {
         let mut h: SyncHistogram<_> = Histogram::<u64>::new_with_max(TRACKABLE_MAX, SIGFIG)
