runtime: avoid redundant unpark in current_thread scheduler (#7834)

Author: Darksonn

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

@@ -15,7 +15,7 @@ use crate::util::{waker_ref, RngSeedGenerator, Wake, WakerRef};
 use std::cell::RefCell;
 use std::collections::VecDeque;
 use std::future::{poll_fn, Future};
-use std::sync::atomic::Ordering::{AcqRel, Release};
+use std::sync::atomic::Ordering::{AcqRel, Acquire, Release};
 use std::task::Poll::{Pending, Ready};
 use std::task::Waker;
 use std::thread::ThreadId;
@@ -380,8 +380,8 @@ impl Context {
             core = c;
         }
 
-        // This check will fail if `before_park` spawns a task for us to run
-        // instead of parking the thread
+        // If `before_park` spawns a task (or otherwise schedules work for us), then we should not
+        // park the thread.
         if !self.has_pending_work(&core) {
             // Park until the thread is signaled
             core.metrics.about_to_park();
@@ -415,7 +415,7 @@ impl Context {
     }
 
     fn has_pending_work(&self, core: &Core) -> bool {
-        !core.tasks.is_empty() || !self.defer.is_empty()
+        !core.tasks.is_empty() || !self.defer.is_empty() || self.handle.shared.woken.load(Acquire)
     }
 
     fn park_internal(
@@ -724,8 +724,20 @@ impl Wake for Handle {
 
     /// Wake by reference
     fn wake_by_ref(arc_self: &Arc<Self>) {
-        arc_self.shared.woken.store(true, Release);
-        arc_self.driver.unpark();
+        let already_woken = arc_self.shared.woken.swap(true, Release);
+
+        if !already_woken {
+            use scheduler::Context::CurrentThread;
+
+            // If we are already running on the runtime, then it's not required to wake up the
+            // runtime.
+            context::with_scheduler(|maybe_cx| match maybe_cx {
+                Some(CurrentThread(cx)) if Arc::ptr_eq(arc_self, &cx.handle) => {}
+                _ => {
+                    arc_self.driver.unpark();
+                }
+            });
+        }
     }
 }
 

tokio/tests/rt_common_before_park.rs

@@ -0,0 +1,92 @@
+#![warn(rust_2018_idioms)]
+#![cfg(feature = "full")]
+#![cfg(not(target_os = "wasi"))] // Wasi doesn't support threads
+
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::Arc;
+use tokio::runtime::Builder;
+use tokio::sync::Notify;
+
+#[test]
+fn before_park_wakes_block_on_task() {
+    let notify = Arc::new(Notify::new());
+    let notify2 = notify.clone();
+    let woken = Arc::new(AtomicBool::new(false));
+    let woken2 = woken.clone();
+
+    let rt = Builder::new_current_thread()
+        .enable_all()
+        .on_thread_park(move || {
+            // Only wake once to avoid busy loop if something goes wrong,
+            // though in this test we expect it to unpark immediately.
+            if !woken2.swap(true, Ordering::SeqCst) {
+                notify2.notify_one();
+            }
+        })
+        .build()
+        .unwrap();
+
+    rt.block_on(async {
+        // This will block until `notify` is notified.
+        // `before_park` should run when the runtime is about to park.
+        // It will notify `notify`, which should wake this task.
+        // The runtime should then see the task is woken and NOT park.
+        notify.notified().await;
+    });
+
+    assert!(woken.load(Ordering::SeqCst));
+}
+
+#[test]
+fn before_park_spawns_task() {
+    let notify = Arc::new(Notify::new());
+    let notify2 = notify.clone();
+    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 notify = notify2.clone();
+                tokio::spawn(async move {
+                    notify.notify_one();
+                });
+            }
+        })
+        .build()
+        .unwrap();
+
+    rt.block_on(async {
+        // This will block until `notify` is notified.
+        // `before_park` should run when the runtime is about to park.
+        // It will spawn a task that notifies `notify`.
+        // The runtime should see the new task and NOT park.
+        // If it parks, it will deadlock.
+        notify.notified().await;
+    });
+
+    assert!(woken.load(Ordering::SeqCst));
+}
+
+#[test]
+fn wake_from_other_thread_block_on() {
+    let rt = Builder::new_current_thread().enable_all().build().unwrap();
+    let handle = rt.handle().clone();
+    let notify = Arc::new(Notify::new());
+    let notify2 = notify.clone();
+
+    let th = std::thread::spawn(move || {
+        // Give the main thread time to park
+        std::thread::sleep(std::time::Duration::from_millis(5));
+        handle.block_on(async move {
+            notify2.notify_one();
+        });
+    });
+
+    rt.block_on(async {
+        notify.notified().await;
+    });
+
+    th.join().unwrap();
+}

tokio/tests/rt_metrics.rs

@@ -152,7 +152,7 @@ fn worker_park_unpark_count() {
     let metrics = rt.metrics();
     rt.block_on(rt.spawn(async {})).unwrap();
     drop(rt);
-    assert!(2 <= metrics.worker_park_unpark_count(0));
+    assert_eq!(0, metrics.worker_park_unpark_count(0) % 2);
 
     let rt = threaded();
     let metrics = rt.metrics();