From e7fc2bdde9cd8e407234bc1a2fa0e867b246a632 Mon Sep 17 00:00:00 2001 From: Bryan Donlan Date: Mon, 6 Jul 2020 23:52:43 +0000 Subject: [PATCH] provide a way to drop a runtime in an async context 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_now()` function which can be used to shutdown a runtime without blocking at all, as an out for cases where this really is necessary. --- tokio/src/runtime/blocking/shutdown.rs | 24 +++++++--- tokio/src/runtime/mod.rs | 30 +++++++++++++ tokio/tests/task_blocking.rs | 61 ++++++++++++++++++++++++++ 3 files changed, 108 insertions(+), 7 deletions(-) diff --git a/tokio/src/runtime/blocking/shutdown.rs b/tokio/src/runtime/blocking/shutdown.rs index f3c60ee301d..e76a7013552 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) { - 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 fc016e96ae8..00453d6f105 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_now()`, 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_now(); + /// }); + /// } + /// ``` + pub fn shutdown_now(self) { + self.shutdown_timeout(Duration::from_nanos(0)) + } } diff --git a/tokio/tests/task_blocking.rs b/tokio/tests/task_blocking.rs index 72fed01e961..58c577ff53b 100644 --- a/tokio/tests/task_blocking.rs +++ b/tokio/tests/task_blocking.rs @@ -96,3 +96,64 @@ fn no_block_in_basic_block_on() { task::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_now(); + }); +}