runtime: drop basic scheduler tasks inside context (#4213)

Author: Darksonn

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<P: Park> {
 
     /// 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<EnterGuard>,
 }
 
 /// The inner scheduler that owns the task queue and the main parker P.
@@ -160,6 +167,7 @@ impl<P: Park> BasicScheduler<P> {
             inner,
             notify: Notify::new(),
             spawner,
+            context_guard: None,
         }
     }
 
@@ -210,6 +218,10 @@ impl<P: Park> BasicScheduler<P> {
             basic_scheduler: self,
         })
     }
+
+    pub(super) fn set_context_guard(&mut self, guard: EnterGuard) {
+        self.context_guard = Some(guard);
+    }
 }
 
 impl<P: Park> Inner<P> {

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)
 }
 

tokio/src/runtime/context.rs

@@ -1,64 +1,102 @@
 //! Thread local runtime context
-use crate::runtime::Handle;
+use crate::runtime::{Handle, TryCurrentError};
 
 use std::cell::RefCell;
 
 thread_local! {
     static CONTEXT: RefCell<Option<Handle>> = RefCell::new(None)
 }
 
-pub(crate) fn current() -> Option<Handle> {
-    CONTEXT.with(|ctx| ctx.borrow().clone())
+pub(crate) fn try_current() -> Result<Handle, crate::runtime::TryCurrentError> {
+    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<crate::runtime::driver::Clock> {
-            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<crate::runtime::Spawner> {
-        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),
+        }
     }
 }
 
 /// Sets this [`Handle`] as the current active [`Handle`].
 ///
 /// [`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<EnterGuard> {
+    CONTEXT
+        .try_with(|ctx| {
+            let old = ctx.borrow_mut().replace(new);
+            EnterGuard(old)
+        })
+        .ok()
 }
 
 #[derive(Debug)]

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<Self, TryCurrentError> {
-        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 TryCurrentError {
+    pub(crate) fn new_no_context() -> Self {
+        Self {
+            kind: TryCurrentErrorKind::NoContext,
+        }
+    }
 
-impl fmt::Debug for TryCurrentError {
+    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),
+        }
     }
 }
 

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.
+                },
+            }
+        }
+    }
 }

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)
     }
 }

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.";

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."