diff --git a/tokio/tests/rt_common.rs b/tokio/tests/rt_common.rs index 5d959365b94..1f7a378549b 100644 --- a/tokio/tests/rt_common.rs +++ b/tokio/tests/rt_common.rs @@ -748,7 +748,7 @@ rt_test! { #[test] fn wake_while_rt_is_dropping() { - use tokio::task; + use tokio::sync::Barrier; use core::sync::atomic::{AtomicBool, Ordering}; let drop_triggered = Arc::new(AtomicBool::new(false)); @@ -766,11 +766,16 @@ rt_test! { let (tx2, rx2) = oneshot::channel(); let (tx3, rx3) = oneshot::channel(); + let barrier = Arc::new(Barrier::new(4)); + let barrier1 = barrier.clone(); + let barrier2 = barrier.clone(); + let barrier3 = barrier.clone(); + let rt = rt(); rt.spawn(async move { // Ensure a waker gets stored in oneshot 1. - let _ = rx1.await; + let _ = tokio::join!(rx1, barrier1.wait()); tx3.send(()).unwrap(); }); @@ -791,20 +796,18 @@ rt_test! { // Just a sanity check that this entire thing actually happened set_drop_triggered.store(true, Ordering::Relaxed); }); - let _ = rx2.await; + let _ = tokio::join!(rx2, barrier2.wait()); }); rt.spawn(async move { - let _ = rx3.await; + let _ = tokio::join!(rx3, barrier3.wait()); // 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; - }); + // Wait until every oneshot channel has been polled. + rt.block_on(barrier.wait()); // Drop the rt drop(rt);