[0.1.x] add set_default to 0.1 executor, timer, and reactor (#1725)

This commit adds `set_default` drop guard style APIs for setting the
default reactor, executor, and timer. These are similar to the APIs used
in `tokio` 0.2

In addition to having potentially better ergonomics than the
`with_default` closure APIs, the drop-guard based APIs will be helpful
in rewriting the `tokio-compat` crate to wrap the existing tokio 0.2
runtime, rather than constructing its own runtime.

Because the runtime does not expose an `around_worker` API, it cannot
currently be used with the 0.1 `with_default` method of setting the
reactor, timer, and executor. This means that tokio-compat must
duplicate a lot of existing code from `tokio` to construct the runtime,
which is unfortunate (and has the potential to introduce errors). On the
other hand, we can use the drop guard APIs with `before_start` and
`after_stop`, by storing the drop guards in a thread-local. This will
allow `tokio-compat` to wrap the 0.2 runtime, reducing code duplication.
Also, this will allow the blocking pool to be used on the compat
runtime, which is currently impossible (as the blocking APIs are private
to `tokio`).

Signed-off-by: Eliza Weisman <[email protected]>

tokio-timer/src/clock/clock.rs

@@ -5,6 +5,7 @@ use tokio_executor::Enter;
 
 use std::cell::Cell;
 use std::fmt;
+use std::marker::PhantomData;
 use std::sync::Arc;
 use std::time::Instant;
 
@@ -20,6 +21,12 @@ pub struct Clock {
     now: Option<Arc<dyn Now>>,
 }
 
+/// A guard that resets the current `Clock` to `None` when dropped.
+#[derive(Debug)]
+pub struct DefaultGuard<'a> {
+    _lifetime: PhantomData<&'a ()>,
+}
+
 thread_local! {
     /// Thread-local tracking the current clock
     static CLOCK: Cell<Option<*const Clock>> = Cell::new(None)
@@ -114,26 +121,35 @@ pub fn with_default<F, R>(clock: &Clock, enter: &mut Enter, f: F) -> R
 where
     F: FnOnce(&mut Enter) -> R,
 {
+    let _guard = set_default(clock);
+
+    f(enter)
+}
+
+/// Sets `clock` as the default clock, returning a guard that unsets it on drop.
+///
+/// # Panics
+///
+/// This function panics if there already is a default clock set.
+pub fn set_default(clock: &Clock) -> DefaultGuard<'_> {
     CLOCK.with(|cell| {
         assert!(
             cell.get().is_none(),
             "default clock already set for execution context"
         );
 
-        // Ensure that the clock is removed from the thread-local context
-        // when leaving the scope. This handles cases that involve panicking.
-        struct Reset<'a>(&'a Cell<Option<*const Clock>>);
-
-        impl<'a> Drop for Reset<'a> {
-            fn drop(&mut self) {
-                self.0.set(None);
-            }
-        }
-
-        let _reset = Reset(cell);
-
         cell.set(Some(clock as *const Clock));
 
-        f(enter)
+        DefaultGuard {
+            _lifetime: PhantomData,
+        }
     })
 }
+
+impl<'a> Drop for DefaultGuard<'a> {
+    fn drop(&mut self) {
+        let _ = CLOCK.try_with(|cell| {
+            cell.set(None);
+        });
+    }
+}

tokio-timer/src/clock/mod.rs

@@ -19,5 +19,5 @@
 mod clock;
 mod now;
 
-pub use self::clock::{now, with_default, Clock};
+pub use self::clock::{now, set_default, with_default, Clock, DefaultGuard};
 pub use self::now::Now;

tokio-timer/src/timer/handle.rs

@@ -5,6 +5,7 @@ use tokio_executor::Enter;
 
 use std::cell::RefCell;
 use std::fmt;
+use std::marker::PhantomData;
 use std::sync::{Arc, Weak};
 use std::time::{Duration, Instant};
 
@@ -44,6 +45,12 @@ pub(crate) struct HandlePriv {
     inner: Weak<Inner>,
 }
 
+/// A guard that resets the current timer to `None` when dropped.
+#[derive(Debug)]
+pub struct DefaultGuard<'a> {
+    _lifetime: PhantomData<&'a ()>,
+}
+
 thread_local! {
     /// Tracks the timer for the current execution context.
     static CURRENT_TIMER: RefCell<Option<HandlePriv>> = RefCell::new(None)
@@ -64,42 +71,34 @@ pub fn with_default<F, R>(handle: &Handle, enter: &mut Enter, f: F) -> R
 where
     F: FnOnce(&mut Enter) -> R,
 {
-    // Ensure that the timer is removed from the thread-local context
-    // when leaving the scope. This handles cases that involve panicking.
-    struct Reset;
-
-    impl Drop for Reset {
-        fn drop(&mut self) {
-            CURRENT_TIMER.with(|current| {
-                let mut current = current.borrow_mut();
-                *current = None;
-            });
-        }
-    }
-
-    // This ensures the value for the current timer gets reset even if there is
-    // a panic.
-    let _r = Reset;
+    let _guard = set_default(handle);
+    f(enter)
+}
 
+/// Sets `handle` as the default timer, returning a guard that unsets it on drop.
+///
+/// # Panics
+///
+/// This function panics if there already is a default timer set.
+pub fn set_default(handle: &Handle) -> DefaultGuard<'_> {
     CURRENT_TIMER.with(|current| {
-        {
-            let mut current = current.borrow_mut();
+        let mut current = current.borrow_mut();
 
-            assert!(
-                current.is_none(),
-                "default Tokio timer already set \
-                 for execution context"
-            );
+        assert!(
+            current.is_none(),
+            "default Tokio timer already set \
+             for execution context"
+        );
 
-            let handle = handle
-                .as_priv()
-                .unwrap_or_else(|| panic!("`handle` does not reference a timer"));
+        let handle = handle
+            .as_priv()
+            .unwrap_or_else(|| panic!("`handle` does not reference a timer"));
 
-            *current = Some(handle.clone());
-        }
-
-        f(enter)
-    })
+        *current = Some(handle.clone());
+    });
+    DefaultGuard {
+        _lifetime: PhantomData,
+    }
 }
 
 impl Handle {
@@ -194,3 +193,12 @@ impl fmt::Debug for HandlePriv {
         write!(f, "HandlePriv")
     }
 }
+
+impl<'a> Drop for DefaultGuard<'a> {
+    fn drop(&mut self) {
+        let _ = CURRENT_TIMER.try_with(|current| {
+            let mut current = current.borrow_mut();
+            *current = None;
+        });
+    }
+}

tokio-timer/src/timer/mod.rs

@@ -45,7 +45,7 @@ use self::entry::Entry;
 use self::stack::Stack;
 
 pub(crate) use self::handle::HandlePriv;
-pub use self::handle::{with_default, Handle};
+pub use self::handle::{set_default, with_default, DefaultGuard, Handle};
 pub use self::now::{Now, SystemNow};
 pub(crate) use self::registration::Registration;