rt: fix flaky wake_while_rt_is_dropping test (#4698)

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);