cros_async: Add a function to run one future to completion

Adding the ability to run one future to completion will allow for
driving the top level future through the executer when only one future
is needed. For example, if a future-rs combinator is used to collect
several sub futures.

TEST=update doc test
cargo test

Change-Id: Idd1121310a3043bb4110853e5e72eef3bd06950e
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/crosvm/+/2173970
Reviewed-by: Dylan Reid <dgreid@chromium.org>
Tested-by: Dylan Reid <dgreid@chromium.org>
Tested-by: kokoro <noreply+kokoro@google.com>
Commit-Queue: Dylan Reid <dgreid@chromium.org>

2 files changed, 60 insertions, 1 deletions

diff --git a/cros_async/src/executor.rs b/cros_async/src/executor.rs
index ac788f1..6641ec7 100644
--- a/cros_async/src/executor.rs
+++ b/cros_async/src/executor.rs
@@ -10,6 +10,8 @@ use std::rc::Rc;
 use std::task::Waker;
 use std::task::{Context, Poll};
 
+use futures::future::FutureExt;
+
 use crate::waker::create_waker;
 
 /// Represents a future executor that can be run. Implementers of the trait will take a list of
@@ -149,6 +151,49 @@ impl FutureList for UnitFutures {
     }
 }
 
+// Execute one future until it completes.
+pub(crate) struct RunOne<F: Future + Unpin> {
+    fut: F,
+    fut_state: FutureState,
+    added_futures: UnitFutures,
+}
+
+impl<F: Future + Unpin> RunOne<F> {
+    pub fn new(f: F) -> RunOne<F> {
+        RunOne {
+            fut: f,
+            fut_state: FutureState::new(),
+            added_futures: UnitFutures::new(),
+        }
+    }
+}
+
+impl<F: Future + Unpin> FutureList for RunOne<F> {
+    type Output = F::Output;
+
+    fn futures_mut(&mut self) -> &mut UnitFutures {
+        &mut self.added_futures
+    }
+
+    fn poll_results(&mut self) -> Option<Self::Output> {
+        let _ = self.added_futures.poll_results();
+
+        if self.fut_state.needs_poll.replace(false) {
+            let mut ctx = Context::from_waker(&self.fut_state.waker);
+            // The future impls `Unpin`, use `poll_unpin` to avoid wrapping it in
+            // `Pin` to call `poll`.
+            if let Poll::Ready(o) = self.fut.poll_unpin(&mut ctx) {
+                return Some(o);
+            }
+        };
+        None
+    }
+
+    fn any_ready(&self) -> bool {
+        self.added_futures.any_ready() || self.fut_state.needs_poll.get()
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
diff --git a/cros_async/src/lib.rs b/cros_async/src/lib.rs
index c3450bc..00109a8 100644
--- a/cros_async/src/lib.rs
+++ b/cros_async/src/lib.rs
@@ -64,7 +64,7 @@ mod waker;
 pub use executor::Executor;
 pub use select::SelectResult;
 
-use executor::UnitFutures;
+use executor::{RunOne, UnitFutures};
 use fd_executor::{FdExecutor, Result};
 use std::future::Future;
 
@@ -88,6 +88,20 @@ pub fn empty_executor() -> Result<impl Executor> {
     FdExecutor::new(UnitFutures::new())
 }
 
+/// Creates a FdExecutor that runs one future to completion.
+///
+///  # Example
+///
+///    ```
+///    use cros_async::run_one;
+///
+///    let fut = async { 55 };
+///    assert_eq!(Ok(55),run_one(Box::pin(fut)));
+///    ```
+pub fn run_one<F: Future + Unpin>(fut: F) -> Result<F::Output> {
+    FdExecutor::new(RunOne::new(fut)).and_then(|mut ex| ex.run())
+}
+
 // Select helpers to run until any future completes.
 
 /// Creates an executor that runs the two given futures until one completes, returning a tuple