rt: yield_now defers task until after driver poll

tokio/src/runtime/context.rs

@@ -9,7 +9,6 @@ use crate::util::rand::{FastRand, RngSeed};
 cfg_rt! {
     use crate::runtime::{scheduler, task::Id, Defer};
 
-
     use std::cell::RefCell;
     use std::marker::PhantomData;
     use std::time::Duration;
@@ -33,6 +32,7 @@ struct Context {
 
     /// Yielded task wakers are stored here and notified after resource drivers
     /// are polled.
+    #[cfg(feature = "rt")]
     defer: RefCell<Option<Defer>>,
 
     #[cfg(any(feature = "rt", feature = "macros"))]

tokio/src/runtime/defer.rs

@@ -1,8 +1,7 @@
-use std::collections::VecDeque;
 use std::task::Waker;
 
 pub(crate) struct Defer {
-    deferred: VecDeque<Waker>,
+    deferred: Vec<Waker>,
 }
 
 impl Defer {
@@ -13,7 +12,7 @@ impl Defer {
     }
 
     pub(crate) fn defer(&mut self, waker: Waker) {
-        self.deferred.push_back(waker);
+        self.deferred.push(waker);
     }
 
     pub(crate) fn is_empty(&self) -> bool {

tokio/tests/rt_common.rs

@@ -9,6 +9,8 @@ macro_rules! rt_test {
         mod current_thread_scheduler {
             $($t)*
 
+            const NUM_WORKERS: usize = 1;
+
             fn rt() -> Arc<Runtime> {
                 tokio::runtime::Builder::new_current_thread()
                     .enable_all()
@@ -22,6 +24,8 @@ macro_rules! rt_test {
         mod threaded_scheduler_4_threads {
             $($t)*
 
+            const NUM_WORKERS: usize = 4;
+
             fn rt() -> Arc<Runtime> {
                 tokio::runtime::Builder::new_multi_thread()
                     .worker_threads(4)
@@ -36,6 +40,8 @@ macro_rules! rt_test {
         mod threaded_scheduler_1_thread {
             $($t)*
 
+            const NUM_WORKERS: usize = 1;
+
             fn rt() -> Arc<Runtime> {
                 tokio::runtime::Builder::new_multi_thread()
                     .worker_threads(1)
@@ -652,7 +658,12 @@ rt_test! {
         for _ in 0..100 {
             rt.spawn(async {
                 loop {
-                    tokio::task::yield_now().await;
+                    // Don't use Tokio's `yield_now()` to avoid special defer
+                    // logic.
+                    let _: () = futures::future::poll_fn(|cx| {
+                        cx.waker().wake_by_ref();
+                        std::task::Poll::Pending
+                    }).await;
                 }
             });
         }
@@ -680,6 +691,64 @@ rt_test! {
         });
     }
 
+    /// Tests that yielded tasks are not scheduled until **after** resource
+    /// drivers are polled.
+    #[test]
+    #[cfg(not(target_os="wasi"))]
+    fn yield_defers_until_park() {
+        use std::sync::atomic::{AtomicBool, Ordering::SeqCst};
+        use std::sync::Barrier;
+
+        let rt = rt();
+
+        let flag = Arc::new(AtomicBool::new(false));
+        let barrier = Arc::new(Barrier::new(NUM_WORKERS));
+
+        rt.block_on(async {
+            // Make sure other workers cannot steal tasks
+            for _ in 0..(NUM_WORKERS-1) {
+                let flag = flag.clone();
+                let barrier = barrier.clone();
+
+                tokio::spawn(async move {
+                    barrier.wait();
+
+                    while !flag.load(SeqCst) {
+                        std::thread::sleep(std::time::Duration::from_millis(1));
+                    }
+                });
+            }
+
+            barrier.wait();
+
+            tokio::spawn(async move {
+                // Create a TCP litener
+                let listener = TcpListener::bind("0.0.0.0:0").await.unwrap();
+                let addr = listener.local_addr().unwrap();
+
+                tokio::join!(
+                    async {
+                        // Done blocking intentionally
+                        let _socket = std::net::TcpStream::connect(addr).unwrap();
+
+                        // Yield until connected
+                        let mut cnt = 0;
+                        while !flag.load(SeqCst){
+                            tokio::task::yield_now().await;
+                            cnt += 1;
+                        }
+
+                        assert!(cnt < 3, "actual={}", cnt);
+                    },
+                    async {
+                        let _ = listener.accept().await.unwrap();
+                        flag.store(true, SeqCst);
+                    }
+                );
+            }).await.unwrap();
+        });
+    }
+
     #[cfg(not(target_os="wasi"))] // Wasi does not support threads
     #[test]
     fn client_server_block_on() {