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..300a14657bf 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 b4019868506..50c070a355a 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(); + }); +}