diff --git a/tokio/src/runtime/basic_scheduler.rs b/tokio/src/runtime/basic_scheduler.rs index 7634e9c60..872d0d5b8 100644 --- a/tokio/src/runtime/basic_scheduler.rs +++ b/tokio/src/runtime/basic_scheduler.rs @@ -2,6 +2,7 @@ use crate::future::poll_fn; use crate::loom::sync::atomic::AtomicBool; use crate::loom::sync::Mutex; use crate::park::{Park, Unpark}; +use crate::runtime::context::EnterGuard; use crate::runtime::stats::{RuntimeStats, WorkerStatsBatcher}; use crate::runtime::task::{self, JoinHandle, OwnedTasks, Schedule, Task}; use crate::runtime::Callback; @@ -29,6 +30,12 @@ pub(crate) struct BasicScheduler { /// Sendable task spawner spawner: Spawner, + + /// This is usually None, but right before dropping the BasicScheduler, it + /// is changed to `Some` with the context being the runtime's own context. + /// This ensures that any tasks dropped in the `BasicScheduler`s destructor + /// run in that runtime's context. + context_guard: Option, } /// The inner scheduler that owns the task queue and the main parker P. @@ -160,6 +167,7 @@ impl BasicScheduler

{ inner, notify: Notify::new(), spawner, + context_guard: None, } } @@ -210,6 +218,10 @@ impl BasicScheduler

{ basic_scheduler: self, }) } + + pub(super) fn set_context_guard(&mut self, guard: EnterGuard) { + self.context_guard = Some(guard); + } } impl Inner

{ diff --git a/tokio/src/runtime/blocking/pool.rs b/tokio/src/runtime/blocking/pool.rs index 8e25ffa5a..77ab49586 100644 --- a/tokio/src/runtime/blocking/pool.rs +++ b/tokio/src/runtime/blocking/pool.rs @@ -8,7 +8,6 @@ use crate::runtime::builder::ThreadNameFn; use crate::runtime::context; use crate::runtime::task::{self, JoinHandle}; use crate::runtime::{Builder, Callback, Handle}; -use crate::util::error::CONTEXT_MISSING_ERROR; use std::collections::{HashMap, VecDeque}; use std::fmt; @@ -81,7 +80,7 @@ where F: FnOnce() -> R + Send + 'static, R: Send + 'static, { - let rt = context::current().expect(CONTEXT_MISSING_ERROR); + let rt = context::current(); rt.spawn_blocking(func) } diff --git a/tokio/src/runtime/context.rs b/tokio/src/runtime/context.rs index 7ed652066..1f44a5340 100644 --- a/tokio/src/runtime/context.rs +++ b/tokio/src/runtime/context.rs @@ -1,5 +1,5 @@ //! Thread local runtime context -use crate::runtime::Handle; +use crate::runtime::{Handle, TryCurrentError}; use std::cell::RefCell; @@ -7,47 +7,73 @@ thread_local! { static CONTEXT: RefCell> = RefCell::new(None) } -pub(crate) fn current() -> Option { - CONTEXT.with(|ctx| ctx.borrow().clone()) +pub(crate) fn try_current() -> Result { + match CONTEXT.try_with(|ctx| ctx.borrow().clone()) { + Ok(Some(handle)) => Ok(handle), + Ok(None) => Err(TryCurrentError::new_no_context()), + Err(_access_error) => Err(TryCurrentError::new_thread_local_destroyed()), + } +} + +pub(crate) fn current() -> Handle { + match try_current() { + Ok(handle) => handle, + Err(e) => panic!("{}", e), + } } cfg_io_driver! { pub(crate) fn io_handle() -> crate::runtime::driver::IoHandle { - CONTEXT.with(|ctx| { + match CONTEXT.try_with(|ctx| { let ctx = ctx.borrow(); ctx.as_ref().expect(crate::util::error::CONTEXT_MISSING_ERROR).io_handle.clone() - }) + }) { + Ok(io_handle) => io_handle, + Err(_) => panic!("{}", crate::util::error::THREAD_LOCAL_DESTROYED_ERROR), + } } } cfg_signal_internal! { #[cfg(unix)] pub(crate) fn signal_handle() -> crate::runtime::driver::SignalHandle { - CONTEXT.with(|ctx| { + match CONTEXT.try_with(|ctx| { let ctx = ctx.borrow(); ctx.as_ref().expect(crate::util::error::CONTEXT_MISSING_ERROR).signal_handle.clone() - }) + }) { + Ok(signal_handle) => signal_handle, + Err(_) => panic!("{}", crate::util::error::THREAD_LOCAL_DESTROYED_ERROR), + } } } cfg_time! { pub(crate) fn time_handle() -> crate::runtime::driver::TimeHandle { - CONTEXT.with(|ctx| { + match CONTEXT.try_with(|ctx| { let ctx = ctx.borrow(); ctx.as_ref().expect(crate::util::error::CONTEXT_MISSING_ERROR).time_handle.clone() - }) + }) { + Ok(time_handle) => time_handle, + Err(_) => panic!("{}", crate::util::error::THREAD_LOCAL_DESTROYED_ERROR), + } } cfg_test_util! { pub(crate) fn clock() -> Option { - CONTEXT.with(|ctx| (*ctx.borrow()).as_ref().map(|ctx| ctx.clock.clone())) + match CONTEXT.try_with(|ctx| (*ctx.borrow()).as_ref().map(|ctx| ctx.clock.clone())) { + Ok(clock) => clock, + Err(_) => panic!("{}", crate::util::error::THREAD_LOCAL_DESTROYED_ERROR), + } } } } cfg_rt! { pub(crate) fn spawn_handle() -> Option { - CONTEXT.with(|ctx| (*ctx.borrow()).as_ref().map(|ctx| ctx.spawner.clone())) + match CONTEXT.try_with(|ctx| (*ctx.borrow()).as_ref().map(|ctx| ctx.spawner.clone())) { + Ok(spawner) => spawner, + Err(_) => panic!("{}", crate::util::error::THREAD_LOCAL_DESTROYED_ERROR), + } } } @@ -55,10 +81,22 @@ cfg_rt! { /// /// [`Handle`]: Handle pub(crate) fn enter(new: Handle) -> EnterGuard { - CONTEXT.with(|ctx| { - let old = ctx.borrow_mut().replace(new); - EnterGuard(old) - }) + match try_enter(new) { + Some(guard) => guard, + None => panic!("{}", crate::util::error::THREAD_LOCAL_DESTROYED_ERROR), + } +} + +/// Sets this [`Handle`] as the current active [`Handle`]. +/// +/// [`Handle`]: Handle +pub(crate) fn try_enter(new: Handle) -> Option { + CONTEXT + .try_with(|ctx| { + let old = ctx.borrow_mut().replace(new); + EnterGuard(old) + }) + .ok() } #[derive(Debug)] diff --git a/tokio/src/runtime/handle.rs b/tokio/src/runtime/handle.rs index dbd7f47eb..cd1cb760a 100644 --- a/tokio/src/runtime/handle.rs +++ b/tokio/src/runtime/handle.rs @@ -1,7 +1,7 @@ use crate::runtime::blocking::{BlockingTask, NoopSchedule}; use crate::runtime::task::{self, JoinHandle}; use crate::runtime::{blocking, context, driver, Spawner}; -use crate::util::error::CONTEXT_MISSING_ERROR; +use crate::util::error::{CONTEXT_MISSING_ERROR, THREAD_LOCAL_DESTROYED_ERROR}; use std::future::Future; use std::marker::PhantomData; @@ -110,7 +110,7 @@ impl Handle { /// # } /// ``` pub fn current() -> Self { - context::current().expect(CONTEXT_MISSING_ERROR) + context::current() } /// Returns a Handle view over the currently running Runtime @@ -119,7 +119,7 @@ impl Handle { /// /// Contrary to `current`, this never panics pub fn try_current() -> Result { - context::current().ok_or(TryCurrentError(())) + context::try_current() } cfg_stats! { @@ -334,17 +334,60 @@ impl Handle { } /// Error returned by `try_current` when no Runtime has been started -pub struct TryCurrentError(()); +#[derive(Debug)] +pub struct TryCurrentError { + kind: TryCurrentErrorKind, +} -impl fmt::Debug for TryCurrentError { +impl TryCurrentError { + pub(crate) fn new_no_context() -> Self { + Self { + kind: TryCurrentErrorKind::NoContext, + } + } + + pub(crate) fn new_thread_local_destroyed() -> Self { + Self { + kind: TryCurrentErrorKind::ThreadLocalDestroyed, + } + } + + /// Returns true if the call failed because there is currently no runtime in + /// the Tokio context. + pub fn is_missing_context(&self) -> bool { + matches!(self.kind, TryCurrentErrorKind::NoContext) + } + + /// Returns true if the call failed because the Tokio context thread-local + /// had been destroyed. This can usually only happen if in the destructor of + /// other thread-locals. + pub fn is_thread_local_destroyed(&self) -> bool { + matches!(self.kind, TryCurrentErrorKind::ThreadLocalDestroyed) + } +} + +enum TryCurrentErrorKind { + NoContext, + ThreadLocalDestroyed, +} + +impl fmt::Debug for TryCurrentErrorKind { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("TryCurrentError").finish() + use TryCurrentErrorKind::*; + match self { + NoContext => f.write_str("NoContext"), + ThreadLocalDestroyed => f.write_str("ThreadLocalDestroyed"), + } } } impl fmt::Display for TryCurrentError { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.write_str(CONTEXT_MISSING_ERROR) + use TryCurrentErrorKind::*; + match self.kind { + NoContext => f.write_str(CONTEXT_MISSING_ERROR), + ThreadLocalDestroyed => f.write_str(THREAD_LOCAL_DESTROYED_ERROR), + } } } diff --git a/tokio/src/runtime/mod.rs b/tokio/src/runtime/mod.rs index a07c16571..96bb47c1d 100644 --- a/tokio/src/runtime/mod.rs +++ b/tokio/src/runtime/mod.rs @@ -205,7 +205,7 @@ cfg_rt! { use self::enter::enter; mod handle; - pub use handle::{EnterGuard, Handle}; + pub use handle::{EnterGuard, Handle, TryCurrentError}; mod spawner; use self::spawner::Spawner; @@ -537,7 +537,7 @@ cfg_rt! { /// ``` pub fn shutdown_timeout(mut self, duration: Duration) { // Wakeup and shutdown all the worker threads - self.handle.shutdown(); + self.handle.clone().shutdown(); self.blocking_pool.shutdown(Some(duration)); } @@ -571,4 +571,30 @@ cfg_rt! { self.shutdown_timeout(Duration::from_nanos(0)) } } + + #[allow(clippy::single_match)] // there are comments in the error branch, so we don't want if-let + impl Drop for Runtime { + fn drop(&mut self) { + match &mut self.kind { + Kind::CurrentThread(basic) => { + // This ensures that tasks spawned on the basic runtime are dropped inside the + // runtime's context. + match self::context::try_enter(self.handle.clone()) { + Some(guard) => basic.set_context_guard(guard), + None => { + // The context thread-local has alread been destroyed. + // + // We don't set the guard in this case. Calls to tokio::spawn in task + // destructors would fail regardless if this happens. + }, + } + }, + #[cfg(feature = "rt-multi-thread")] + Kind::ThreadPool(_) => { + // The threaded scheduler drops its tasks on its worker threads, which is + // already in the runtime's context. + }, + } + } + } } diff --git a/tokio/src/task/builder.rs b/tokio/src/task/builder.rs index e46bdefe9..f991fc65e 100644 --- a/tokio/src/task/builder.rs +++ b/tokio/src/task/builder.rs @@ -1,5 +1,4 @@ #![allow(unreachable_pub)] -use crate::util::error::CONTEXT_MISSING_ERROR; use crate::{runtime::context, task::JoinHandle}; use std::future::Future; @@ -98,8 +97,6 @@ impl<'a> Builder<'a> { Function: FnOnce() -> Output + Send + 'static, Output: Send + 'static, { - context::current() - .expect(CONTEXT_MISSING_ERROR) - .spawn_blocking_inner(function, self.name) + context::current().spawn_blocking_inner(function, self.name) } } diff --git a/tokio/src/util/error.rs b/tokio/src/util/error.rs index 0e52364a7..8f252c0c9 100644 --- a/tokio/src/util/error.rs +++ b/tokio/src/util/error.rs @@ -7,3 +7,11 @@ pub(crate) const CONTEXT_MISSING_ERROR: &str = /// Error string explaining that the Tokio context is shutting down and cannot drive timers. pub(crate) const RUNTIME_SHUTTING_DOWN_ERROR: &str = "A Tokio 1.x context was found, but it is being shutdown."; + +// some combinations of features might not use this +#[allow(dead_code)] +/// Error string explaining that the Tokio context is not available because the +/// thread-local storing it has been destroyed. This usually only happens during +/// destructors of other thread-locals. +pub(crate) const THREAD_LOCAL_DESTROYED_ERROR: &str = + "The Tokio context thread-local variable has been destroyed."; diff --git a/tokio/tests/rt_basic.rs b/tokio/tests/rt_basic.rs index 4b1bdadc1..70056b16f 100644 --- a/tokio/tests/rt_basic.rs +++ b/tokio/tests/rt_basic.rs @@ -3,10 +3,14 @@ use tokio::runtime::Runtime; use tokio::sync::oneshot; +use tokio::time::{timeout, Duration}; use tokio_test::{assert_err, assert_ok}; +use std::future::Future; +use std::pin::Pin; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::task::{Context, Poll}; use std::thread; -use tokio::time::{timeout, Duration}; mod support { pub(crate) mod mpsc_stream; @@ -135,6 +139,35 @@ fn acquire_mutex_in_drop() { drop(rt); } +#[test] +fn drop_tasks_in_context() { + static SUCCESS: AtomicBool = AtomicBool::new(false); + + struct ContextOnDrop; + + impl Future for ContextOnDrop { + type Output = (); + + fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<()> { + Poll::Pending + } + } + + impl Drop for ContextOnDrop { + fn drop(&mut self) { + if tokio::runtime::Handle::try_current().is_ok() { + SUCCESS.store(true, Ordering::SeqCst); + } + } + } + + let rt = rt(); + rt.spawn(ContextOnDrop); + drop(rt); + + assert!(SUCCESS.load(Ordering::SeqCst)); +} + #[test] #[should_panic( expected = "A Tokio 1.x context was found, but timers are disabled. Call `enable_time` on the runtime builder to enable timers."