[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;