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>

1 files changed, 17 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