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;