Added functions to partially progress a local pool.

fixes #1369

Author: SrTobi

futures-executor/src/local_pool.rs

@@ -74,6 +74,18 @@ fn run_executor<T, F: FnMut(&mut Context<'_>) -> Poll<T>>(mut f: F) -> T {
     })
 }
 
+fn poll_executor<T, F: FnMut(&mut Context<'_>) -> T>(mut f: F) -> T {
+    let _enter = enter()
+        .expect("cannot execute `LocalPool` executor from within \
+                 another executor");
+
+    CURRENT_THREAD_NOTIFY.with(|thread_notify| {
+        let waker = waker_ref(thread_notify);
+        let mut cx = Context::from_waker(&waker);
+        f(&mut cx)
+    })
+}
+
 impl LocalPool {
     /// Create a new, empty pool of tasks.
     pub fn new() -> LocalPool {
@@ -135,8 +147,8 @@ impl LocalPool {
     /// The function will block the calling thread *only* until the future `f`
     /// completes; there may still be incomplete tasks in the pool, which will
     /// be inert after the call completes, but can continue with further use of
-    /// `run` or `run_until`. While the function is running, however, all tasks
-    /// in the pool will try to make progress.
+    /// one of the pool's run or poll mothods. While the function is running,
+    /// however, all tasks in the pool will try to make progress.
     pub fn run_until<F: Future>(&mut self, future: F) -> F::Output {
         pin_mut!(future);
 
@@ -154,20 +166,94 @@ impl LocalPool {
         })
     }
 
+    /// Runs all tasks and returns after completing one future or until no more progress
+    /// can be made. Returns `true` if one future was completed, `false` otherwise.
+    ///
+    /// ```
+    /// #![feature(futures_api)]
+    /// use futures::executor::LocalPool;
+    /// use futures::task::LocalSpawnExt;
+    /// use futures::future::{ready, empty};
+    ///
+    /// let mut pool = LocalPool::new();
+    /// let mut spawner = pool.spawner();
+    ///
+    /// spawner.spawn_local(ready(()));
+    /// spawner.spawn_local(ready(()));
+    /// spawner.spawn_local(empty());
+    ///
+    /// // Run the two ready tasks and return true for them.
+    /// pool.poll_one(); // returns true after completing one of the ready futures
+    /// pool.poll_one(); // returns true after completing the other ready future
+    ///
+    /// // the remaining task can not be completed
+    /// pool.poll_one(); // returns false
+    /// ```
+    ///
+    /// This function will not block the calling thread and will return the moment
+    /// that there are no tasks left for which progress can be made or after exactly one
+    /// task was completed; Remaining incomplete tasks in the pool can continue with
+    /// 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 poll_one(&mut self) -> bool {
+        poll_executor(|ctx| {
+            let ret = self.poll_pool_once(ctx);
+
+            // return if we really have executed a future
+            match ret {
+                Poll::Ready(Some(_)) => true,
+                _ => false
+            }
+        })
+    }
+
+    /// Runs all tasks in the pool and returns if no more progress can be made
+    /// on any task.
+    ///
+    /// ```
+    /// #![feature(futures_api)]
+    /// use futures::executor::LocalPool;
+    /// use futures::task::LocalSpawnExt;
+    /// use futures::future::{ready, empty};
+    ///
+    /// let mut pool = LocalPool::new();
+    /// let mut spawner = pool.spawner();
+    ///
+    /// spawner.spawn_local(ready(()));
+    /// spawner.spawn_local(ready(()));
+    /// spawner.spawn_local(empty());
+    ///
+    /// // Runs the two ready task and returns.
+    /// // The empty task remains in the pool.
+    /// pool.poll();
+    /// ```
+    ///
+    /// This function will not block the calling thread and will return the moment
+    /// that there are no tasks left for which progress can be made;
+    /// remaining incomplete tasks in the pool can continue with further use of one
+    /// 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 poll(&mut self) {
+        poll_executor(|ctx| {
+            loop {
+                let result = self.poll_pool_once(ctx);
+
+                // if there are no more ready futures exit
+                match result {
+                    Poll::Pending | Poll::Ready(None) => return,
+                    _ => continue
+                }
+            }
+        })
+    }
+
     // Make maximal progress on the entire pool of spawned task, returning `Ready`
     // if the pool is empty and `Pending` if no further progress can be made.
     fn poll_pool(&mut self, cx: &mut Context<'_>) -> Poll<()> {
         // state for the FuturesUnordered, which will never be used
         loop {
-            // empty the incoming queue of newly-spawned tasks
-            {
-                let mut incoming = self.incoming.borrow_mut();
-                for task in incoming.drain(..) {
-                    self.pool.push(task)
-                }
-            }
+            let ret = self.poll_pool_once(cx);
 
-            let 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;
@@ -181,6 +267,20 @@ impl LocalPool {
             }
         }
     }
+
+    // Try make minimal progress on the pool of spawned tasks
+    fn poll_pool_once(&mut self, cx: &mut Context<'_>) -> Poll<Option<()>> {
+        // empty the incoming queue of newly-spawned tasks
+        {
+            let mut incoming = self.incoming.borrow_mut();
+            for task in incoming.drain(..) {
+                self.pool.push(task)
+            }
+        }
+
+        // try to execute the next ready future
+        self.pool.poll_next_unpin(cx)
+    }
 }
 
 impl Default for LocalPool {

futures-executor/tests/local_pool.rs

@@ -2,8 +2,8 @@
 
 use futures::channel::oneshot;
 use futures::executor::LocalPool;
-use futures::future::{Future, lazy};
-use futures::task::{Context, Poll, Spawn, LocalSpawn};
+use futures::future::{Future, lazy, poll_fn};
+use futures::task::{Context, Poll, Spawn, LocalSpawn, Waker};
 use std::cell::{Cell, RefCell};
 use std::pin::Pin;
 use std::rc::Rc;
@@ -58,6 +58,13 @@ fn run_until_executes_spawned() {
     pool.run_until(rx).unwrap();
 }
 
+#[test]
+fn run_returns_if_empty() {
+    let mut pool = LocalPool::new();
+    pool.run();
+    pool.run();
+}
+
 #[test]
 fn run_executes_spawned() {
     let cnt = Rc::new(Cell::new(0));
@@ -103,6 +110,130 @@ fn run_spawn_many() {
     assert_eq!(cnt.get(), ITER);
 }
 
+#[test]
+fn poll_one_returns_if_empty() {
+    let mut pool = LocalPool::new();
+    assert!(pool.poll_one() == false);
+}
+
+#[test]
+fn poll_one_executes_one_ready() {
+    const ITER: usize = 200;
+
+    let cnt = Rc::new(Cell::new(0));
+
+    let mut pool = LocalPool::new();
+    let mut spawn = pool.spawner();
+
+    for _ in 0..ITER {
+        spawn.spawn_local_obj(Box::pin(pending()).into()).unwrap();
+
+        let cnt = cnt.clone();
+        spawn.spawn_local_obj(Box::pin(lazy(move |_| {
+            cnt.set(cnt.get() + 1);
+            ()
+        })).into()).unwrap();
+
+        spawn.spawn_local_obj(Box::pin(pending()).into()).unwrap();
+    }
+
+    for i in 0..ITER {
+        assert_eq!(cnt.get(), i);
+        assert!(pool.poll_one());
+        assert_eq!(cnt.get(), i + 1);
+    }
+    assert!(pool.poll_one() == false);
+}
+
+#[test]
+fn poll_one_returns_on_no_progress() {
+    const ITER: usize = 10;
+
+    let cnt = Rc::new(Cell::new(0));
+
+    let mut pool = LocalPool::new();
+    let mut spawn = pool.spawner();
+
+    let waker: Rc<Cell<Option<Waker>>> = Rc::new(Cell::new(None));
+    {
+        let cnt = cnt.clone();
+        let waker = waker.clone();
+        spawn.spawn_local_obj(Box::pin(poll_fn(move |ctx| {
+            cnt.set(cnt.get() + 1);
+            waker.set(Some(ctx.waker().clone()));
+            if cnt.get() == ITER {
+                Poll::Ready(())
+            } else {
+                Poll::Pending
+            }
+        })).into()).unwrap();
+    }
+
+    for i in 0..ITER - 1 {
+        assert_eq!(cnt.get(), i);
+        assert!(!pool.poll_one());
+        assert_eq!(cnt.get(), i + 1);
+        let w = waker.take();
+        assert!(w.is_some());
+        w.unwrap().wake();
+    }
+    assert!(pool.poll_one());
+    assert_eq!(cnt.get(), ITER);
+}
+
+#[test]
+fn poll_returns_if_empty() {
+    let mut pool = LocalPool::new();
+    pool.poll();
+    pool.poll();
+}
+
+#[test]
+fn poll_returns_multiple_times() {
+    let mut pool = LocalPool::new();
+    let mut spawn = pool.spawner();
+    let cnt = Rc::new(Cell::new(0));
+
+    let cnt1 = cnt.clone();
+    spawn.spawn_local_obj(Box::pin(lazy(move |_|{ cnt1.set(cnt1.get() + 1) })).into()).unwrap();
+    pool.poll();
+    assert_eq!(cnt.get(), 1);
+
+    let cnt2 = cnt.clone();
+    spawn.spawn_local_obj(Box::pin(lazy(move |_|{ cnt2.set(cnt2.get() + 1) })).into()).unwrap();
+    pool.poll();
+    assert_eq!(cnt.get(), 2);
+}
+
+#[test]
+fn poll_executes_all_ready() {
+    const ITER: usize = 200;
+    const PER_ITER: usize = 3;
+
+    let cnt = Rc::new(Cell::new(0));
+
+    let mut pool = LocalPool::new();
+    let mut spawn = pool.spawner();
+
+    for i in 0..ITER {
+        for _ in 0..PER_ITER {
+            spawn.spawn_local_obj(Box::pin(pending()).into()).unwrap();
+
+            let cnt = cnt.clone();
+            spawn.spawn_local_obj(Box::pin(lazy(move |_| {
+                cnt.set(cnt.get() + 1);
+                ()
+            })).into()).unwrap();
+
+            // also add some pending tasks to test if they are ignored
+            spawn.spawn_local_obj(Box::pin(pending()).into()).unwrap();
+        }
+        assert_eq!(cnt.get(), i * PER_ITER);
+        pool.poll();
+        assert_eq!(cnt.get(), (i + 1) * PER_ITER);
+    }
+}
+
 #[test]
 #[should_panic]
 fn nesting_run() {
@@ -113,6 +244,21 @@ fn nesting_run() {
         let mut pool = LocalPool::new();
         pool.run();
     })).into()).unwrap();
+
+    pool.run();
+}
+
+#[test]
+#[should_panic]
+fn nesting_run_poll() {
+    let mut pool = LocalPool::new();
+    let mut spawn = pool.spawner();
+
+    spawn.spawn_obj(Box::pin(lazy(|_| {
+        let mut pool = LocalPool::new();
+        pool.poll();
+    })).into()).unwrap();
+
     pool.run();
 }