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>

3 files changed, 108 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))
+    }
 }
diff --git a/tokio/tests/task_blocking.rs b/tokio/tests/task_blocking.rs
index b4019868..50c070a3 100644
--- a/tokio/tests/task_blocking.rs
+++ b/tokio/tests/task_blocking.rs
@@ -115,3 +115,64 @@ fn can_enter_basic_rt_from_within_block_in_place() {
         })
     });
 }
+
+#[test]
+fn useful_panic_message_when_dropping_rt_in_rt() {
+    use std::panic::{catch_unwind, AssertUnwindSafe};
+
+    let mut outer = tokio::runtime::Builder::new()
+        .threaded_scheduler()
+        .build()
+        .unwrap();
+
+    let result = catch_unwind(AssertUnwindSafe(|| {
+        outer.block_on(async {
+            let _ = tokio::runtime::Builder::new()
+                .basic_scheduler()
+                .build()
+                .unwrap();
+        });
+    }));
+
+    assert!(result.is_err());
+    let err = result.unwrap_err();
+    let err: &'static str = err.downcast_ref::<&'static str>().unwrap();
+
+    assert!(
+        err.find("Cannot drop a runtime").is_some(),
+        "Wrong panic message: {:?}",
+        err
+    );
+}
+
+#[test]
+fn can_shutdown_with_zero_timeout_in_runtime() {
+    let mut outer = tokio::runtime::Builder::new()
+        .threaded_scheduler()
+        .build()
+        .unwrap();
+
+    outer.block_on(async {
+        let rt = tokio::runtime::Builder::new()
+            .basic_scheduler()
+            .build()
+            .unwrap();
+        rt.shutdown_timeout(Duration::from_nanos(0));
+    });
+}
+
+#[test]
+fn can_shutdown_now_in_runtime() {
+    let mut outer = tokio::runtime::Builder::new()
+        .threaded_scheduler()
+        .build()
+        .unwrap();
+
+    outer.block_on(async {
+        let rt = tokio::runtime::Builder::new()
+            .basic_scheduler()
+            .build()
+            .unwrap();
+        rt.shutdown_background();
+    });
+}