use ParkThread in LocalPool

- split all state but park into separate struct
- various "poll" methods need to call park(Poll) as it might change the
  readiness of futures (expired timers, IO ready, ...) or spawn new
  futures (although ParkThread::park(Poll) won't)

Author: stbuehler

futures-executor/src/local_pool.rs

@@ -1,22 +1,85 @@
 use crate::enter;
+use crate::park::{Park, ParkDuration};
+use crate::park_thread::ParkThread;
 use futures_core::future::Future;
 use futures_core::stream::Stream;
 use futures_core::task::{Context, Poll};
-use futures_task::{waker_ref, ArcWake};
 use futures_task::{FutureObj, LocalFutureObj, LocalSpawn, Spawn, SpawnError};
 use futures_util::stream::FuturesUnordered;
 use futures_util::stream::StreamExt;
 use std::cell::RefCell;
 use std::ops::{Deref, DerefMut};
 use std::pin::pin;
 use std::rc::{Rc, Weak};
-use std::sync::{
-    atomic::{AtomicBool, Ordering},
-    Arc,
-};
-use std::thread::{self, Thread};
 use std::vec::Vec;
 
+// state outside park
+#[derive(Debug)]
+struct State {
+    pool: FuturesUnordered<LocalFutureObj<'static, ()>>,
+    incoming: Rc<Incoming>,
+}
+
+impl State {
+    /// Poll `self.pool`, re-filling it with any newly-spawned tasks.
+    /// Repeat until either the pool is empty, or it returns `Pending`.
+    ///
+    /// Returns `Ready` if the pool was empty, and `Pending` otherwise.
+    ///
+    /// NOTE: the pool may call `wake`, so `Pending` doesn't necessarily
+    /// mean that the pool can't make progress.
+    fn poll_pool(&mut self, cx: &mut Context<'_>) -> Poll<()> {
+        loop {
+            self.drain_incoming();
+
+            let pool_ret = self.pool.poll_next_unpin(cx);
+
+            // We queued up some new tasks; add them and poll again.
+            if !self.incoming.borrow().is_empty() {
+                continue;
+            }
+
+            match pool_ret {
+                Poll::Ready(Some(())) => continue,
+                Poll::Ready(None) => return Poll::Ready(()),
+                Poll::Pending => return Poll::Pending,
+            }
+        }
+    }
+
+    /// Empty the incoming queue of newly-spawned tasks.
+    fn drain_incoming(&mut self) {
+        let mut incoming = self.incoming.borrow_mut();
+        for task in incoming.drain(..) {
+            self.pool.push(task)
+        }
+    }
+
+    /// Runs all tasks and returns after completing one future or until no more progress
+    /// can be made. Returns `Ready` if one future was completed, `Pending` otherwise.
+    fn poll_pool_one(&mut self, cx: &mut Context<'_>) -> Poll<()> {
+        loop {
+            self.drain_incoming();
+
+            match self.pool.poll_next_unpin(cx) {
+                // Success!
+                Poll::Ready(Some(())) => return Poll::Ready(()),
+                // The pool was empty.
+                Poll::Ready(None) => return Poll::Pending,
+                Poll::Pending => (),
+            }
+
+            if !self.incoming.borrow().is_empty() {
+                // New tasks were spawned; try again.
+                continue;
+            } else {
+                // No more progress for now (but caller should check cx)
+                return Poll::Pending;
+            }
+        }
+    }
+}
+
 /// A single-threaded task pool for polling futures to completion.
 ///
 /// This executor allows you to multiplex any number of tasks onto a single
@@ -29,9 +92,9 @@ use std::vec::Vec;
 /// single-threaded, it supports a special form of task spawning for non-`Send`
 /// futures, via [`spawn_local_obj`](futures_task::LocalSpawn::spawn_local_obj).
 #[derive(Debug)]
-pub struct LocalPool {
-    pool: FuturesUnordered<LocalFutureObj<'static, ()>>,
-    incoming: Rc<Incoming>,
+pub struct LocalPool<P = ParkThread> {
+    park: P,
+    state: State,
 }
 
 /// A handle to a [`LocalPool`] that implements [`Spawn`](futures_task::Spawn).
@@ -42,80 +105,56 @@ pub struct LocalSpawner {
 
 type Incoming = RefCell<Vec<LocalFutureObj<'static, ()>>>;
 
-pub(crate) struct ThreadNotify {
-    /// The (single) executor thread.
-    thread: Thread,
-    /// A flag to ensure a wakeup (i.e. `unpark()`) is not "forgotten"
-    /// before the next `park()`, which may otherwise happen if the code
-    /// being executed as part of the future(s) being polled makes use of
-    /// park / unpark calls of its own, i.e. we cannot assume that no other
-    /// code uses park / unpark on the executing `thread`.
-    unparked: AtomicBool,
-}
-
-std::thread_local! {
-    static CURRENT_THREAD_NOTIFY: Arc<ThreadNotify> = Arc::new(ThreadNotify {
-        thread: thread::current(),
-        unparked: AtomicBool::new(false),
-    });
-}
-
-impl ArcWake for ThreadNotify {
-    fn wake_by_ref(arc_self: &Arc<Self>) {
-        // Make sure the wakeup is remembered until the next `park()`.
-        let unparked = arc_self.unparked.swap(true, Ordering::Release);
-        if !unparked {
-            // If the thread has not been unparked yet, it must be done
-            // now. If it was actually parked, it will run again,
-            // otherwise the token made available by `unpark`
-            // may be consumed before reaching `park()`, but `unparked`
-            // ensures it is not forgotten.
-            arc_self.thread.unpark();
-        }
-    }
-}
-
-// Set up and run a basic single-threaded spawner loop, invoking `f` on each
-// turn.
-fn run_executor<T, F: FnMut(&mut Context<'_>) -> Poll<T>>(mut f: F) -> T {
-    let _enter = enter().expect(
+fn run_executor<P: Park, T, F: FnMut(&mut Context<'_>) -> Poll<T>>(park: &mut P, mut f: F) -> T {
+    let mut enter = enter().expect(
         "cannot execute `LocalPool` executor from within \
-         another executor",
+                another executor",
     );
 
-    CURRENT_THREAD_NOTIFY.with(|thread_notify| {
-        let waker = waker_ref(thread_notify);
+    loop {
+        let waker = park.waker();
         let mut cx = Context::from_waker(&waker);
-        loop {
-            if let Poll::Ready(t) = f(&mut cx) {
-                return t;
-            }
 
-            // Wait for a wakeup.
-            while !thread_notify.unparked.swap(false, Ordering::Acquire) {
-                // No wakeup occurred. It may occur now, right before parking,
-                // but in that case the token made available by `unpark()`
-                // is guaranteed to still be available and `park()` is a no-op.
-                thread::park();
-            }
+        if let Poll::Ready(t) = f(&mut cx) {
+            return t;
         }
-    })
-}
-
-/// Check for a wakeup, but don't consume it.
-fn woken() -> bool {
-    CURRENT_THREAD_NOTIFY.with(|thread_notify| thread_notify.unparked.load(Ordering::Acquire))
+        park.park(&mut enter, ParkDuration::Block);
+    }
 }
 
 impl LocalPool {
     /// Create a new, empty pool of tasks.
+    // Not generic over `P: Default` - would break API (default type not working)
     pub fn new() -> Self {
-        Self { pool: FuturesUnordered::new(), incoming: Default::default() }
+        Self::new_with_park(ParkThread::default())
+    }
+}
+
+impl<P> LocalPool<P>
+where
+    P: Park,
+{
+    /// Create a new, empty pool of tasks.
+    pub fn new_with_park(park: P) -> Self {
+        LocalPool {
+            park,
+            state: State { pool: FuturesUnordered::new(), incoming: Default::default() },
+        }
+    }
+
+    /// Returns a reference to the underlying Park instance.
+    pub fn get_park(&self) -> &P {
+        &self.park
+    }
+
+    /// Returns a mutable reference to the underlying Park instance.
+    pub fn get_park_mut(&mut self) -> &mut P {
+        &mut self.park
     }
 
     /// Get a clonable handle to the pool as a [`Spawn`].
     pub fn spawner(&self) -> LocalSpawner {
-        LocalSpawner { incoming: Rc::downgrade(&self.incoming) }
+        LocalSpawner { incoming: Rc::downgrade(&self.state.incoming) }
     }
 
     /// Run all tasks in the pool to completion.
@@ -134,7 +173,7 @@ impl LocalPool {
     /// The function will block the calling thread until *all* tasks in the pool
     /// are complete, including any spawned while running existing tasks.
     pub fn run(&mut self) {
-        run_executor(|cx| self.poll_pool(cx))
+        self.run_executor(|state, cx| state.poll_pool(cx))
     }
 
     /// Runs all the tasks in the pool until the given future completes.
@@ -157,7 +196,7 @@ impl LocalPool {
     pub fn run_until<F: Future>(&mut self, future: F) -> F::Output {
         let mut future = pin!(future);
 
-        run_executor(|cx| {
+        self.run_executor(|state, cx| {
             {
                 // if our main task is done, so are we
                 let result = future.as_mut().poll(cx);
@@ -166,7 +205,7 @@ impl LocalPool {
                 }
             }
 
-            let _ = self.poll_pool(cx);
+            let _ = state.poll_pool(cx);
             Poll::Pending
         })
     }
@@ -200,29 +239,7 @@ impl LocalPool {
     /// further use of one of the pool's run or poll methods.
     /// Though only one task will be completed, progress may be made on multiple tasks.
     pub fn try_run_one(&mut self) -> bool {
-        run_executor(|cx| {
-            loop {
-                self.drain_incoming();
-
-                match self.pool.poll_next_unpin(cx) {
-                    // Success!
-                    Poll::Ready(Some(())) => return Poll::Ready(true),
-                    // The pool was empty.
-                    Poll::Ready(None) => return Poll::Ready(false),
-                    Poll::Pending => (),
-                }
-
-                if !self.incoming.borrow().is_empty() {
-                    // New tasks were spawned; try again.
-                    continue;
-                } else if woken() {
-                    // The pool yielded to us, but there's more progress to be made.
-                    return Poll::Pending;
-                } else {
-                    return Poll::Ready(false);
-                }
-            }
-        })
+        self.poll_executor(|state, cx| state.poll_pool_one(cx)).is_ready()
     }
 
     /// Runs all tasks in the pool and returns if no more progress can be made
@@ -251,58 +268,51 @@ impl LocalPool {
     /// of the pool's run or poll methods. While the function is running, all tasks
     /// in the pool will try to make progress.
     pub fn run_until_stalled(&mut self) {
-        run_executor(|cx| match self.poll_pool(cx) {
-            // The pool is empty.
-            Poll::Ready(()) => Poll::Ready(()),
-            Poll::Pending => {
-                if woken() {
-                    Poll::Pending
-                } else {
-                    // We're stalled for now.
-                    Poll::Ready(())
-                }
-            }
-        });
+        // Ready: all tasks finished, Pending: would block
+        let _ = self.poll_executor(|state, cx| state.poll_pool(cx));
     }
 
-    /// Poll `self.pool`, re-filling it with any newly-spawned tasks.
-    /// Repeat until either the pool is empty, or it returns `Pending`.
-    ///
-    /// Returns `Ready` if the pool was empty, and `Pending` otherwise.
-    ///
-    /// NOTE: the pool may call `wake`, so `Pending` doesn't necessarily
-    /// mean that the pool can't make progress.
-    fn poll_pool(&mut self, cx: &mut Context<'_>) -> Poll<()> {
-        loop {
-            self.drain_incoming();
+    // Set up and run a basic single-threaded spawner loop, polling `f` on each
+    // turn until it completes or [`P::park`] fails.
+    fn run_executor<T, F: FnMut(&mut State, &mut Context<'_>) -> Poll<T>>(
+        &mut self,
+        mut f: F,
+    ) -> T {
+        let park = &mut self.park;
+        let state = &mut self.state;
+        run_executor(park, |cx| f(state, cx))
+    }
 
-            let pool_ret = self.pool.poll_next_unpin(cx);
+    // Set up and run a basic single-threaded spawner loop, polling `f` on each
+    // turn until it completes or no progress can't be made anymore.
+    fn poll_executor<T, F: FnMut(&mut State, &mut Context<'_>) -> Poll<T>>(
+        &mut self,
+        mut f: F,
+    ) -> Poll<T> {
+        let _enter = enter().expect(
+            "cannot execute `LocalPool` executor from within \
+                    another executor",
+        );
 
-            // We queued up some new tasks; add them and poll again.
-            if !self.incoming.borrow().is_empty() {
-                continue;
-            }
+        loop {
+            let waker = self.park.waker();
+            let mut cx = Context::from_waker(&waker);
 
-            match pool_ret {
-                Poll::Ready(Some(())) => continue,
-                Poll::Ready(None) => return Poll::Ready(()),
-                Poll::Pending => return Poll::Pending,
+            if let Poll::Ready(t) = f(&mut self.state, &mut cx) {
+                return Poll::Ready(t);
+            }
+            if !self.park.poll() {
+                // can't make more progress
+                return Poll::Pending;
             }
-        }
-    }
-
-    /// Empty the incoming queue of newly-spawned tasks.
-    fn drain_incoming(&mut self) {
-        let mut incoming = self.incoming.borrow_mut();
-        for task in incoming.drain(..) {
-            self.pool.push(task)
         }
     }
 }
 
+// Not generic over `P: Default` - would break API (default type not working)
 impl Default for LocalPool {
     fn default() -> Self {
-        Self::new()
+        Self::new_with_park(ParkThread::default())
     }
 }
 
@@ -313,7 +323,7 @@ impl Default for LocalPool {
 /// Use a [`LocalPool`] if you need finer-grained control over spawned tasks.
 pub fn block_on<F: Future>(f: F) -> F::Output {
     let mut f = pin!(f);
-    run_executor(|cx| f.as_mut().poll(cx))
+    run_executor(&mut ParkThread::new(), |cx| f.as_mut().poll(cx))
 }
 
 /// Turn a stream into a blocking iterator.