[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-executor/src/global.rs

@@ -3,6 +3,7 @@ use super::{Enter, Executor, SpawnError};
 use futures::{future, Future};
 
 use std::cell::Cell;
+use std::marker::PhantomData;
 
 /// Executes futures on the default executor for the current execution context.
 ///
@@ -19,6 +20,13 @@ pub struct DefaultExecutor {
     _dummy: (),
 }
 
+/// Ensures that the executor is removed from the thread-local context
+/// when leaving the scope. This handles cases that involve panicking.
+#[derive(Debug)]
+pub struct DefaultGuard<'a> {
+    _lifetime: PhantomData<&'a ()>,
+}
+
 impl DefaultExecutor {
     /// Returns a handle to the default executor for the current context.
     ///
@@ -174,6 +182,20 @@ pub fn with_default<T, F, R>(executor: &mut T, enter: &mut Enter, f: F) -> R
 where
     T: Executor,
     F: FnOnce(&mut Enter) -> R,
+{
+    let _guard = set_default(executor);
+    f(enter)
+}
+
+/// Sets `executor` as the default executor, returning a guard that unsets it when
+/// dropped.
+///
+/// # Panics
+///
+/// This function panics if there already is a default executor set.
+pub fn set_default<T>(executor: &mut T) -> DefaultGuard<'_>
+where
+    T: Executor,
 {
     EXECUTOR.with(|cell| {
         match cell.get() {
@@ -183,18 +205,6 @@ where
             _ => {}
         }
 
-        // Ensure that the executor is removed from the thread-local context
-        // when leaving the scope. This handles cases that involve panicking.
-        struct Reset<'a>(&'a Cell<State>);
-
-        impl<'a> Drop for Reset<'a> {
-            fn drop(&mut self) {
-                self.0.set(State::Empty);
-            }
-        }
-
-        let _reset = Reset(cell);
-
         // While scary, this is safe. The function takes a
         // `&mut Executor`, which guarantees that the reference lives for the
         // duration of `with_default`.
@@ -205,9 +215,11 @@ where
         let executor = unsafe { hide_lt(executor as &mut _ as *mut _) };
 
         cell.set(State::Ready(executor));
+    });
 
-        f(enter)
+    DefaultGuard {
-    })
+        _lifetime: PhantomData,
+    }
 }
 
 unsafe fn hide_lt<'a>(p: *mut (dyn Executor + 'a)) -> *mut (dyn Executor + 'static) {
@@ -215,6 +227,14 @@ unsafe fn hide_lt<'a>(p: *mut (dyn Executor + 'a)) -> *mut (dyn Executor + 'stat
     mem::transmute(p)
 }
 
+impl<'a> Drop for DefaultGuard<'a> {
+    fn drop(&mut self) {
+        let _ = EXECUTOR.try_with(|cell| {
+            cell.set(State::Empty);
+        });
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::{with_default, DefaultExecutor, Executor};

tokio-executor/src/lib.rs

@@ -64,5 +64,5 @@ mod typed;
 pub use enter::{enter, exit, Enter, EnterError};
 pub use error::SpawnError;
 pub use executor::Executor;
-pub use global::{spawn, with_default, DefaultExecutor};
+pub use global::{set_default, spawn, with_default, DefaultExecutor, DefaultGuard};
 pub use typed::TypedExecutor;

tokio-reactor/src/lib.rs

@@ -68,6 +68,7 @@ use tokio_sync::task::AtomicTask;
 use std::cell::RefCell;
 use std::error::Error;
 use std::io;
+use std::marker::PhantomData;
 use std::mem;
 #[cfg(all(unix, not(target_os = "fuchsia")))]
 use std::os::unix::io::{AsRawFd, RawFd};
@@ -133,6 +134,13 @@ pub struct SetFallbackError(());
 #[doc(hidden)]
 pub type SetDefaultError = SetFallbackError;
 
+/// Ensure that the default reactor is removed from the thread-local context
+/// when leaving the scope. This handles cases that involve panicking.
+#[derive(Debug)]
+pub struct DefaultGuard<'a> {
+    _lifetime: PhantomData<&'a ()>,
+}
+
 #[test]
 fn test_handle_size() {
     use std::mem;
@@ -197,45 +205,40 @@ pub fn with_default<F, R>(handle: &Handle, enter: &mut Enter, f: F) -> R
 where
     F: FnOnce(&mut Enter) -> R,
 {
-    // Ensure that the executor 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_REACTOR.with(|current| {
-                let mut current = current.borrow_mut();
-                *current = None;
-            });
-        }
-    }
-
     // This ensures the value for the current reactor gets reset even if there
     // is a panic.
-    let _r = Reset;
+    let _guard = set_default(handle);
+    f(enter)
+}
 
+/// Sets `handle` as the default reactor, returning a guard that unsets it when
+/// dropped.
+///
+/// # Panics
+///
+/// This function panics if there already is a default reactor set.
+pub fn set_default(handle: &Handle) -> DefaultGuard<'_> {
     CURRENT_REACTOR.with(|current| {
-        {
+        let mut current = current.borrow_mut();
-            let mut current = current.borrow_mut();
 
-            assert!(
+        assert!(
-                current.is_none(),
+            current.is_none(),
-                "default Tokio reactor already set \
+            "default Tokio reactor already set \
-                 for execution context"
+             for execution context"
-            );
+        );
 
-            let handle = match handle.as_priv() {
+        let handle = match handle.as_priv() {
-                Some(handle) => handle,
+            Some(handle) => handle,
-                None => {
+            None => {
-                    panic!("`handle` does not reference a reactor");
+                panic!("`handle` does not reference a reactor");
-                }
+            }
-            };
+        };
 
-            *current = Some(handle.clone());
+        *current = Some(handle.clone());
-        }
+    });
-
+    DefaultGuard {
-        f(enter)
+        _lifetime: PhantomData,
-    })
+    }
 }
 
 impl Reactor {
@@ -743,6 +746,15 @@ impl Direction {
     }
 }
 
+impl<'a> Drop for DefaultGuard<'a> {
+    fn drop(&mut self) {
+        let _ = CURRENT_REACTOR.try_with(|current| {
+            let mut current = current.borrow_mut();
+            *current = None;
+        });
+    }
+}
+
 #[cfg(unix)]
 mod platform {
     use mio::unix::UnixReady;

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
+    let _guard = set_default(handle);
-    // when leaving the scope. This handles cases that involve panicking.
+    f(enter)
-    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;
 
+/// 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!(
+        assert!(
-                current.is_none(),
+            current.is_none(),
-                "default Tokio timer already set \
+            "default Tokio timer already set \
-                 for execution context"
+             for execution context"
-            );
+        );
 
-            let handle = handle
+        let handle = handle
-                .as_priv()
+            .as_priv()
-                .unwrap_or_else(|| panic!("`handle` does not reference a timer"));
+            .unwrap_or_else(|| panic!("`handle` does not reference a timer"));
 
-            *current = Some(handle.clone());
+        *current = Some(handle.clone());
-        }
+    });
-
+    DefaultGuard {
-        f(enter)
+        _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;