[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-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!(
-                current.is_none(),
-                "default Tokio reactor already set \
-                 for execution context"
-            );
+        assert!(
+            current.is_none(),
+            "default Tokio reactor already set \
+             for execution context"
+        );
 
-            let handle = match handle.as_priv() {
-                Some(handle) => handle,
-                None => {
-                    panic!("`handle` does not reference a reactor");
-                }
-            };
+        let handle = match handle.as_priv() {
+            Some(handle) => handle,
+            None => {
+                panic!("`handle` does not reference a reactor");
+            }
+        };
 
-            *current = Some(handle.clone());
-        }
-
-        f(enter)
-    })
+        *current = Some(handle.clone());
+    });
+    DefaultGuard {
+        _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;