diff --git a/tokio-executor/src/global.rs b/tokio-executor/src/global.rs index 801341fcf..0e5c47740 100644 --- a/tokio-executor/src/global.rs +++ b/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(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(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); - - 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}; diff --git a/tokio-executor/src/lib.rs b/tokio-executor/src/lib.rs index 8bbf92b2e..9abdf91c9 100644 --- a/tokio-executor/src/lib.rs +++ b/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; diff --git a/tokio-reactor/src/lib.rs b/tokio-reactor/src/lib.rs index 727ae536c..c28a7ae6b 100644 --- a/tokio-reactor/src/lib.rs +++ b/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(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; diff --git a/tokio-timer/src/clock/clock.rs b/tokio-timer/src/clock/clock.rs index fbba6981d..e65d668e4 100644 --- a/tokio-timer/src/clock/clock.rs +++ b/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>, } +/// 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> = Cell::new(None) @@ -114,26 +121,35 @@ pub fn with_default(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>); - - 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); + }); + } +} diff --git a/tokio-timer/src/clock/mod.rs b/tokio-timer/src/clock/mod.rs index 1791bc758..8acb83179 100644 --- a/tokio-timer/src/clock/mod.rs +++ b/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; diff --git a/tokio-timer/src/timer/handle.rs b/tokio-timer/src/timer/handle.rs index c9598b9ae..3129ec9bf 100644 --- a/tokio-timer/src/timer/handle.rs +++ b/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, } +/// 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> = RefCell::new(None) @@ -64,42 +71,34 @@ pub fn with_default(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; + }); + } +} diff --git a/tokio-timer/src/timer/mod.rs b/tokio-timer/src/timer/mod.rs index c79a82a26..31eb0afbb 100644 --- a/tokio-timer/src/timer/mod.rs +++ b/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;