diff --git a/tokio/src/runtime/handle.rs b/tokio/src/runtime/handle.rs index 7dff91448f1..4d365281129 100644 --- a/tokio/src/runtime/handle.rs +++ b/tokio/src/runtime/handle.rs @@ -301,7 +301,7 @@ impl Handle { .expect("failed to park thread") } - pub(crate) fn shutdown(mut self) { + pub(crate) fn shutdown(&mut self) { self.spawner.shutdown(); } } diff --git a/tokio/src/runtime/mod.rs b/tokio/src/runtime/mod.rs index 52532ec6f12..48b1d0e7e6d 100644 --- a/tokio/src/runtime/mod.rs +++ b/tokio/src/runtime/mod.rs @@ -275,6 +275,9 @@ cfg_rt! { /// The runtime executor is either a thread-pool or a current-thread executor. #[derive(Debug)] enum Kind { + /// Runtime is shutting down + Gone, + /// Execute all tasks on the current-thread. CurrentThread(BasicScheduler), @@ -447,6 +450,7 @@ cfg_rt! { let _enter = self.enter(); match &self.kind { + Kind::Gone => panic!("Runtime is shutting down"), Kind::CurrentThread(exec) => exec.block_on(future), #[cfg(feature = "rt-multi-thread")] Kind::ThreadPool(exec) => exec.block_on(future), @@ -560,4 +564,12 @@ cfg_rt! { self.shutdown_timeout(Duration::from_nanos(0)) } } + + impl Drop for Runtime { + fn drop(&mut self) { + let _guard = self.enter(); + // Shutdown the runtime whilst it is still active + self.kind = Kind::Gone; + } + } } diff --git a/tokio/tests/task_abort.rs b/tokio/tests/task_abort.rs index 8f621683faa..f9e6651fa2e 100644 --- a/tokio/tests/task_abort.rs +++ b/tokio/tests/task_abort.rs @@ -1,7 +1,7 @@ #![warn(rust_2018_idioms)] #![cfg(feature = "full")] -use std::sync::Arc; +use std::sync::{mpsc, Arc}; use std::thread::sleep; use std::time::Duration; @@ -233,3 +233,44 @@ fn test_abort_task_that_panics_on_drop_returned() { assert!(handle.await.unwrap_err().is_panic()); }); } + +struct SpawnOnDrop(mpsc::Sender); + +impl Drop for SpawnOnDrop { + fn drop(&mut self) { + let res = std::panic::catch_unwind(|| { + tokio::spawn(async move { + println!("did something"); + }); + }); + self.0.send(res.is_ok()).unwrap(); + } +} + +/// Checks that aborting a task whose destructor panics has the expected result. +#[test] +fn test_task_that_spawns_task_on_drop() { + let (tx, rx) = mpsc::channel(); + + { + let rt = tokio::runtime::Builder::new_current_thread() + .enable_time() + .build() + .unwrap(); + + rt.block_on(async move { + let _handle = tokio::spawn(async move { + // Make sure the Arc is moved into the task + let _spawn_dropped = SpawnOnDrop(tx); + println!("task started"); + tokio::time::sleep(std::time::Duration::new(1, 0)).await + }); + + // wait for task to sleep. + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + }); + } + + // Check that spawning the task did not panic + assert!(rx.recv().unwrap()); +}