track root task's awake status, and skip blocking on pollables if awake

Author: pchickey

src/runtime/block_on.rs

@@ -1,11 +1,10 @@
 use super::{Reactor, REACTOR};
 
-use core::future::Future;
-use core::pin::pin;
-use core::task::Waker;
-use core::task::{Context, Poll};
+use std::future::Future;
+use std::pin::pin;
+use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::Arc;
-use std::task::Wake;
+use std::task::{Context, Poll, Wake, Waker};
 
 /// Start the event loop
 pub fn block_on<Fut>(fut: Fut) -> Fut::Output
@@ -24,30 +23,50 @@ where
     let mut fut = pin!(fut);
 
     // Create a new context to be passed to the future.
-    let waker = noop_waker();
+    let root = Arc::new(RootWaker::new());
+    let waker = Waker::from(root.clone());
     let mut cx = Context::from_waker(&waker);
 
     // Either the future completes and we return, or some IO is happening
     // and we wait.
     let res = loop {
         match fut.as_mut().poll(&mut cx) {
             Poll::Ready(res) => break res,
-            Poll::Pending => reactor.block_until(),
+            Poll::Pending => {
+                // If some non-pollable based future has marked the root task
+                // as awake, reset and poll again. otherwise, block until a
+                // pollable wakes a future.
+                if root.is_awake() {
+                    root.reset()
+                } else {
+                    reactor.block_on_pollables()
+                }
+            }
         }
     };
     // Clear the singleton
     REACTOR.replace(None);
     res
 }
 
-/// Construct a new no-op waker
-// NOTE: we can remove this once <https://github.com/rust-lang/rust/issues/98286> lands
-fn noop_waker() -> Waker {
-    struct NoopWaker;
-
-    impl Wake for NoopWaker {
-        fn wake(self: Arc<Self>) {}
+struct RootWaker {
+    wake: AtomicBool,
+}
+impl RootWaker {
+    fn new() -> Self {
+        Self {
+            wake: AtomicBool::new(false),
+        }
+    }
+    fn is_awake(&self) -> bool {
+        self.wake.load(Ordering::Relaxed)
+    }
+    fn reset(&self) {
+        self.wake.store(false, Ordering::Relaxed);
+    }
+}
+impl Wake for RootWaker {
+    fn wake(self: Arc<Self>) {
+        self.wake.store(true, Ordering::Relaxed);
     }
-
-    Waker::from(Arc::new(NoopWaker))
 }

src/runtime/reactor.rs

@@ -126,19 +126,8 @@ impl Reactor {
         }
     }
 
-    /// Block until new events are ready. Calls the respective wakers once done.
-    ///
-    /// # On Wakers and single-threaded runtimes
-    ///
-    /// At first glance it might seem silly that this goes through the motions
-    /// of calling the wakers. The main waker we create here is a `noop` waker:
-    /// it does nothing. However, it is common and encouraged to use wakers to
-    /// distinguish between events. Concurrency primitives may construct their
-    /// own wakers to keep track of identity and wake more precisely. We do not
-    /// control the wakers construted by other libraries, and it is for this
-    /// reason that we have to call all the wakers - even if by default they
-    /// will do nothing.
-    pub(crate) fn block_until(&self) {
+    /// Block until at least one pending pollable is ready, waking a pending future.
+    pub(crate) fn block_on_pollables(&self) {
         let reactor = self.inner.borrow();
 
         // We're about to wait for a number of pollables. When they wake we get