runtime: don't park in current_thread if before_park defers waker (#7835)

tokio/src/runtime/scheduler/current_thread/mod.rs

@@ -382,7 +382,7 @@ impl Context {
 
         // This check will fail if `before_park` spawns a task for us to run
         // instead of parking the thread
-        if core.tasks.is_empty() {
+        if !self.has_pending_work(&core) {
             // Park until the thread is signaled
             core.metrics.about_to_park();
             core.submit_metrics(handle);
@@ -414,6 +414,10 @@ impl Context {
         core
     }
 
+    fn has_pending_work(&self, core: &Core) -> bool {
+        !core.tasks.is_empty() || !self.defer.is_empty()
+    }
+
     fn park_internal(
         &self,
         core: Box<Core>,
@@ -775,7 +779,7 @@ impl CoreGuard<'_> {
                         None => {
                             core.metrics.end_processing_scheduled_tasks();
 
-                            core = if !context.defer.is_empty() {
+                            core = if context.has_pending_work(&core) {
                                 context.park_yield(core, handle)
                             } else {
                                 context.park(core, handle)

tokio/tests/rt_basic.rs

@@ -456,3 +456,44 @@ fn rt() -> Runtime {
         .build()
         .unwrap()
 }
+
+#[test]
+fn before_park_yields() {
+    use futures::task::ArcWake;
+    use std::sync::Arc;
+    use tokio::runtime::Builder;
+    use tokio::sync::Notify;
+
+    struct MyWaker(Notify);
+
+    impl ArcWake for MyWaker {
+        fn wake_by_ref(arc_self: &Arc<Self>) {
+            arc_self.0.notify_one();
+        }
+    }
+
+    let notify = Arc::new(MyWaker(Notify::new()));
+    let notify2 = notify.clone();
+    let waker = futures::task::waker(notify2);
+    let woken = Arc::new(AtomicBool::new(false));
+    let woken2 = woken.clone();
+
+    let rt = Builder::new_current_thread()
+        .enable_all()
+        .on_thread_park(move || {
+            if !woken2.swap(true, Ordering::SeqCst) {
+                let mut cx = Context::from_waker(&waker);
+                // `yield_now` pushes the waker to the defer slot.
+                let fut = std::pin::pin!(tokio::task::yield_now());
+                let _ = fut.poll(&mut cx);
+            }
+        })
+        .build()
+        .unwrap();
+
+    rt.block_on(async {
+        notify.0.notified().await;
+    });
+
+    assert!(woken.load(Ordering::SeqCst));
+}