diff --git a/tokio/src/runtime/blocking/mod.rs b/tokio/src/runtime/blocking/mod.rs index ff400b33b..5c808335c 100644 --- a/tokio/src/runtime/blocking/mod.rs +++ b/tokio/src/runtime/blocking/mod.rs @@ -5,7 +5,7 @@ cfg_blocking_impl! { mod pool; - pub(crate) use pool::{spawn_blocking, BlockingPool, Spawner}; + pub(crate) use pool::{spawn_blocking, try_spawn_blocking, BlockingPool, Spawner}; mod schedule; mod shutdown; diff --git a/tokio/src/runtime/blocking/pool.rs b/tokio/src/runtime/blocking/pool.rs index 1784312de..0b9d2209c 100644 --- a/tokio/src/runtime/blocking/pool.rs +++ b/tokio/src/runtime/blocking/pool.rs @@ -65,10 +65,21 @@ where let rt = Handle::current(); let (task, handle) = task::joinable(BlockingTask::new(func)); - rt.blocking_spawner.spawn(task, &rt); + let _ = rt.blocking_spawner.spawn(task, &rt); handle } +#[allow(dead_code)] +pub(crate) fn try_spawn_blocking(func: F) -> Result<(), ()> +where + F: FnOnce() -> R + Send + 'static, +{ + let rt = Handle::current(); + + let (task, _handle) = task::joinable(BlockingTask::new(func)); + rt.blocking_spawner.spawn(task, &rt) +} + // ===== impl BlockingPool ===== impl BlockingPool { @@ -137,7 +148,7 @@ impl fmt::Debug for BlockingPool { // ===== impl Spawner ===== impl Spawner { - fn spawn(&self, task: Task, rt: &Handle) { + fn spawn(&self, task: Task, rt: &Handle) -> Result<(), ()> { let shutdown_tx = { let mut shared = self.inner.shared.lock().unwrap(); @@ -146,7 +157,7 @@ impl Spawner { task.shutdown(); // no need to even push this task; it would never get picked up - return; + return Err(()); } shared.queue.push_back(task); @@ -178,6 +189,8 @@ impl Spawner { if let Some(shutdown_tx) = shutdown_tx { self.spawn_thread(shutdown_tx, rt); } + + Ok(()) } fn spawn_thread(&self, shutdown_tx: shutdown::Sender, rt: &Handle) { @@ -217,9 +230,6 @@ impl Inner { run_task(task); shared = self.shared.lock().unwrap(); - if shared.shutdown { - break; // Need to increment idle before we exit - } } // IDLE diff --git a/tokio/src/runtime/mod.rs b/tokio/src/runtime/mod.rs index 3c56e1384..6922ef59b 100644 --- a/tokio/src/runtime/mod.rs +++ b/tokio/src/runtime/mod.rs @@ -198,7 +198,8 @@ mod blocking; use blocking::BlockingPool; cfg_blocking_impl! { - pub(crate) use blocking::spawn_blocking; + #[allow(unused_imports)] + pub(crate) use blocking::{spawn_blocking, try_spawn_blocking}; } mod builder; diff --git a/tokio/src/runtime/thread_pool/tests/loom_pool.rs b/tokio/src/runtime/thread_pool/tests/loom_pool.rs index 3151eaa62..98765ac02 100644 --- a/tokio/src/runtime/thread_pool/tests/loom_pool.rs +++ b/tokio/src/runtime/thread_pool/tests/loom_pool.rs @@ -8,6 +8,28 @@ use loom::sync::{Arc, Mutex}; use std::future::Future; use std::sync::atomic::Ordering::{Acquire, Relaxed, Release}; +#[test] +fn racy_shutdown() { + loom::model(|| { + let pool = mk_pool(1); + + // here's the case we want to exercise: + // + // a worker that still has tasks in its local queue gets sent to the blocking pool (due to + // block_in_place). the blocking pool is shut down, so drops the worker. the worker's + // shutdown method never gets run. + // + // we do this by spawning two tasks on one worker, the first of which does block_in_place, + // and then immediately drop the pool. + + pool.spawn(async { + crate::task::block_in_place(|| {}); + }); + pool.spawn(async {}); + drop(pool); + }); +} + #[test] fn pool_multi_spawn() { loom::model(|| { diff --git a/tokio/src/runtime/thread_pool/worker.rs b/tokio/src/runtime/thread_pool/worker.rs index 5e96a4422..a4e9d83be 100644 --- a/tokio/src/runtime/thread_pool/worker.rs +++ b/tokio/src/runtime/thread_pool/worker.rs @@ -212,12 +212,6 @@ impl Worker { return; } - // make sure no subsequent code thinks that it is on a worker - current::clear(); - - // Track that the worker is gone - self.gone.set(true); - // If this method is called, we need to move the entire worker onto a // separate (blocking) thread before returning. Once we return, the // caller is going to execute some blocking code which would otherwise @@ -259,7 +253,20 @@ impl Worker { }; // Give away the worker - runtime::spawn_blocking(move || worker.run()); + // + // Returns `Err` if the spawn failed due to the runtime shutting down + let res = runtime::try_spawn_blocking(move || worker.run()); + + // If the worker hand-off was successful, clear the local state. + // Otherwise, the runtime is in the process of shutting down, so we will + // just block on the worker. + if res.is_ok() { + // make sure no subsequent code thinks that it is on a worker + current::clear(); + + // Track that the worker is gone + self.gone.set(true); + } } }