provide a way to drop a runtime in an async context (#2646)

Dropping a runtime normally involves waiting for any outstanding blocking tasks
to complete. When this drop happens in an asynchronous context, we previously
would issue a cryptic panic due to trying to block in an asynchronous context.

This change improves the panic message, and adds a `shutdown_blocking()` function
which can be used to shutdown a runtime without blocking at all, as an out for
cases where this really is necessary.

Co-authored-by: Bryan Donlan <bdonlan@amazon.com>
Co-authored-by: Alice Ryhl <alice@ryhl.io>

2 files changed, 47 insertions, 7 deletions

diff --git a/tokio/src/runtime/blocking/shutdown.rs b/tokio/src/runtime/blocking/shutdown.rs
index f3c60ee3..e76a7013 100644
--- a/tokio/src/runtime/blocking/shutdown.rs
+++ b/tokio/src/runtime/blocking/shutdown.rs
@@ -33,15 +33,25 @@ impl Receiver {
     /// duration. If `timeout` is `None`, then the thread is blocked until the
     /// shutdown signal is received.
     pub(crate) fn wait(&mut self, timeout: Option<Duration>) {
-        use crate::runtime::enter::{enter, try_enter};
+        use crate::runtime::enter::try_enter;
 
-        let mut e = if std::thread::panicking() {
-            match try_enter(false) {
-                Some(enter) => enter,
-                _ => return,
+        if timeout == Some(Duration::from_nanos(0)) {
+            return;
+        }
+
+        let mut e = match try_enter(false) {
+            Some(enter) => enter,
+            _ => {
+                if std::thread::panicking() {
+                    // Don't panic in a panic
+                    return;
+                } else {
+                    panic!(
+                        "Cannot drop a runtime in a context where blocking is not allowed. \
+                        This happens when a runtime is dropped from within an asynchronous context."
+                    );
+                }
             }
-        } else {
-            enter(false)
         };
 
         // The oneshot completes with an Err
diff --git a/tokio/src/runtime/mod.rs b/tokio/src/runtime/mod.rs
index fc016e96..300a1465 100644
--- a/tokio/src/runtime/mod.rs
+++ b/tokio/src/runtime/mod.rs
@@ -548,4 +548,34 @@ impl Runtime {
         } = self;
         blocking_pool.shutdown(Some(duration));
     }
+
+    /// Shutdown the runtime, without waiting for any spawned tasks to shutdown.
+    ///
+    /// This can be useful if you want to drop a runtime from within another runtime.
+    /// Normally, dropping a runtime will block indefinitely for spawned blocking tasks
+    /// to complete, which would normally not be permitted within an asynchronous context.
+    /// By calling `shutdown_background()`, you can drop the runtime from such a context.
+    ///
+    /// Note however, that because we do not wait for any blocking tasks to complete, this
+    /// may result in a resource leak (in that any blocking tasks are still running until they
+    /// return.
+    ///
+    /// This function is equivalent to calling `shutdown_timeout(Duration::of_nanos(0))`.
+    ///
+    /// ```
+    /// use tokio::runtime::Runtime;
+    ///
+    /// fn main() {
+    ///    let mut runtime = Runtime::new().unwrap();
+    ///
+    ///    runtime.block_on(async move {
+    ///        let inner_runtime = Runtime::new().unwrap();
+    ///        // ...
+    ///        inner_runtime.shutdown_background();
+    ///    });
+    /// }
+    /// ```
+    pub fn shutdown_background(self) {
+        self.shutdown_timeout(Duration::from_nanos(0))
+    }
 }