From 6406328176cdecf15cad69b327597a4d4d0b8e20 Mon Sep 17 00:00:00 2001 From: Carl Lerche Date: Wed, 8 Jan 2020 21:23:10 -0800 Subject: [PATCH] rt: fix threaded scheduler shutdown deadlock (#2074) Previously, if an IO event was received during the runtime shutdown process, it was possible to enter a deadlock. This was due to the scheduler shutdown logic not expecting tasks to get scheduled once the worker was in the shutdown process. This patch fixes the deadlock by checking the queues for new tasks after each call to park. If a new task is received, it is forcefully shutdown. Fixes #2061 --- tokio/src/runtime/thread_pool/worker.rs | 40 +++++++---- tokio/tests/rt_basic.rs | 59 ---------------- tokio/tests/rt_common.rs | 94 ++++++++++++++++++++++++- 3 files changed, 120 insertions(+), 73 deletions(-) diff --git a/tokio/src/runtime/thread_pool/worker.rs b/tokio/src/runtime/thread_pool/worker.rs index 44df8b74d..298ef06de 100644 --- a/tokio/src/runtime/thread_pool/worker.rs +++ b/tokio/src/runtime/thread_pool/worker.rs @@ -533,27 +533,41 @@ impl GenerationGuard<'_> { // Transition all tasks owned by the worker to canceled. self.owned().owned_tasks.shutdown(); - // First, drain all tasks from both the local & global queue. - while let Some(task) = self.owned().work_queue.pop_local_first() { - task.shutdown(); - } - - // Notify all workers in case they have pending tasks to drop - // - // Not super efficient, but we are also shutting down. - self.worker.slices.notify_all(); + // Always notify the first time around. This flushes any released tasks + // that happened before the call to `Worker::shutdown` + let mut notify = true; // The worker can only shutdown once there are no further owned tasks. - while !self.owned().owned_tasks.is_empty() { + loop { + // First, drain all tasks from both the local & global queue. + while let Some(task) = self.owned().work_queue.pop_local_first() { + notify = true; + task.shutdown(); + } + + if notify { + // If any tasks are shutdown, they may be pushed on another + // worker's `pending_drop` stack. However, we don't know which + // workers need to be notified, so we just notify all of them. + // Since this is a shutdown process, excessive notification is + // not a huge deal. + self.worker.slices.notify_all(); + notify = false; + } + + // Try draining more tasks + self.drain_tasks_pending_drop(); + + if self.owned().owned_tasks.is_empty() { + break; + } + // Wait until task that this worker owns are released. // // `transition_to_parked` is not called as we are not working // anymore. When a task is released, the owning worker is unparked // directly. self.park_mut().park().expect("park failed"); - - // Try draining more tasks - self.drain_tasks_pending_drop(); } } diff --git a/tokio/tests/rt_basic.rs b/tokio/tests/rt_basic.rs index 39250c4cb..38a726928 100644 --- a/tokio/tests/rt_basic.rs +++ b/tokio/tests/rt_basic.rs @@ -63,65 +63,6 @@ fn acquire_mutex_in_drop() { drop(rt); } -#[test] -fn wake_while_rt_is_dropping() { - use tokio::task; - - struct OnDrop(F); - - impl Drop for OnDrop { - fn drop(&mut self) { - (self.0)() - } - } - - let (tx1, rx1) = oneshot::channel(); - let (tx2, rx2) = oneshot::channel(); - let (tx3, rx3) = oneshot::channel(); - - let mut rt = rt(); - - let h1 = rt.handle().clone(); - - rt.handle().spawn(async move { - // Ensure a waker gets stored in oneshot 1. - let _ = rx1.await; - tx3.send(()).unwrap(); - }); - - rt.handle().spawn(async move { - // When this task is dropped, we'll be "closing remotes". - // We spawn a new task that owns the `tx1`, to move its Drop - // out of here. - // - // Importantly, the oneshot 1 has a waker already stored, so - // the eventual drop here will try to re-schedule again. - let mut opt_tx1 = Some(tx1); - let _d = OnDrop(move || { - let tx1 = opt_tx1.take().unwrap(); - h1.spawn(async move { - tx1.send(()).unwrap(); - }); - }); - let _ = rx2.await; - }); - - rt.handle().spawn(async move { - let _ = rx3.await; - // We'll never get here, but once task 3 drops, this will - // force task 2 to re-schedule since it's waiting on oneshot 2. - tx2.send(()).unwrap(); - }); - - // Tick the loop - rt.block_on(async { - task::yield_now().await; - }); - - // Drop the rt - drop(rt); -} - fn rt() -> Runtime { tokio::runtime::Builder::new() .basic_scheduler() diff --git a/tokio/tests/rt_common.rs b/tokio/tests/rt_common.rs index 15f5de6ce..31edd10a2 100644 --- a/tokio/tests/rt_common.rs +++ b/tokio/tests/rt_common.rs @@ -41,7 +41,7 @@ fn send_sync_bound() { } rt_test! { - use tokio::net::{TcpListener, TcpStream}; + use tokio::net::{TcpListener, TcpStream, UdpSocket}; use tokio::prelude::*; use tokio::runtime::Runtime; use tokio::sync::oneshot; @@ -617,6 +617,98 @@ rt_test! { assert_ok!(drop_rx.recv()); } + #[test] + fn wake_while_rt_is_dropping() { + use tokio::task; + + struct OnDrop(F); + + impl Drop for OnDrop { + fn drop(&mut self) { + (self.0)() + } + } + + let (tx1, rx1) = oneshot::channel(); + let (tx2, rx2) = oneshot::channel(); + let (tx3, rx3) = oneshot::channel(); + + let mut rt = rt(); + + let h1 = rt.handle().clone(); + + rt.handle().spawn(async move { + // Ensure a waker gets stored in oneshot 1. + let _ = rx1.await; + tx3.send(()).unwrap(); + }); + + rt.handle().spawn(async move { + // When this task is dropped, we'll be "closing remotes". + // We spawn a new task that owns the `tx1`, to move its Drop + // out of here. + // + // Importantly, the oneshot 1 has a waker already stored, so + // the eventual drop here will try to re-schedule again. + let mut opt_tx1 = Some(tx1); + let _d = OnDrop(move || { + let tx1 = opt_tx1.take().unwrap(); + h1.spawn(async move { + tx1.send(()).unwrap(); + }); + }); + let _ = rx2.await; + }); + + rt.handle().spawn(async move { + let _ = rx3.await; + // We'll never get here, but once task 3 drops, this will + // force task 2 to re-schedule since it's waiting on oneshot 2. + tx2.send(()).unwrap(); + }); + + // Tick the loop + rt.block_on(async { + task::yield_now().await; + }); + + // Drop the rt + drop(rt); + } + + #[test] + fn io_notify_while_shutting_down() { + use std::net::Ipv6Addr; + + for _ in 1..100 { + let mut runtime = rt(); + + runtime.block_on(async { + let socket = UdpSocket::bind((Ipv6Addr::LOCALHOST, 0)).await.unwrap(); + let addr = socket.local_addr().unwrap(); + let (mut recv_half, mut send_half) = socket.split(); + + tokio::spawn(async move { + let mut buf = [0]; + loop { + recv_half.recv_from(&mut buf).await.unwrap(); + std::thread::sleep(Duration::from_millis(2)); + } + }); + + tokio::spawn(async move { + let buf = [0]; + loop { + send_half.send_to(&buf, &addr).await.unwrap(); + tokio::time::delay_for(Duration::from_millis(1)).await; + } + }); + + tokio::time::delay_for(Duration::from_millis(5)).await; + }); + } + } + #[test] fn runtime_in_thread_local() { use std::cell::RefCell;