try_select

Signed-off-by: Yoshua Wuyts <[email protected]>

Author: yoshuawuyts

src/future/future/mod.rs

@@ -1,6 +1,8 @@
 mod select;
+mod try_select;
 
 use select::Select;
+use try_select::TrySelect;
 
 extension_trait! {
     use std::pin::Pin;
@@ -104,7 +106,7 @@ extension_trait! {
 
     pub trait FutureExt: std::future::Future {
         #[doc = r#"
-            Waits for either one of several similarly-typed futures to complete.
+            Waits for one of two similarly-typed futures to complete.
 
             Awaits multiple futures simultaneously, returning all results once complete.
 
@@ -120,7 +122,7 @@ extension_trait! {
             # Examples
 
             ```
-            # futures::executor::block_on(async {
+            # async_std::task::block_on(async {
             use async_std::prelude::*;
             use async_std::future;
 
@@ -133,12 +135,51 @@ extension_trait! {
             # });
             ```
         "#]
-        fn select<F>(self, other: F) -> Select<Self, F>
+        fn select<F>(self, other: F) ->
+            impl Future<Output = Self::Output> [Select<Self, F>]
+            where
+                Self: Sized + Future,
+                F: std::future::Future<Output = <Self as Future>::Output>,
+        {
+            Select::new(self, other)
+        }
+
+        #[doc = r#"
+            Waits for one of two similarly-typed fallible futures to complete.
+
+            Awaits multiple futures simultaneously, returning all results once complete.
+
+            `try_select` is similar to [`select`], but keeps going if a future
+            resolved to an error until all futures have been resolved. In which case
+            the error of the `other` future will be returned.
+
+            # Examples
+
+            ```
+            # fn main() -> std::io::Result<()> { async_std::task::block_on(async {
+            #
+            use async_std::prelude::*;
+            use async_std::future;
+            use std::io::{Error, ErrorKind};
+
+            let a = future::pending::<Result<_, Error>>();
+            let b = future::ready(Err(Error::from(ErrorKind::Other)));
+            let c = future::ready(Ok(1u8));
+
+            let f = a.try_select(b).try_select(c);
+            assert_eq!(f.await?, 1u8);
+            #
+            # Ok(()) }) }
+            ```
+        "#]
+        fn try_select<F, T, E>(self, other: F) ->
+            impl Future<Output = Self::Output> [TrySelect<Self, F>]
             where
                 Self: Sized,
+                Self: Future<Output = Result<T, E>>,
                 F: Future<Output = Self::Output>,
         {
-            Select::new(self, other)
+            TrySelect::new(self, other)
         }
     }
 

src/future/future/select.rs

@@ -3,8 +3,7 @@ use std::pin::Pin;
 use pin_project_lite::pin_project;
 use async_macros::MaybeDone;
 
-use crate::future::Future;
-// use crate::future::MaybeDone;
+use std::future::Future;
 use crate::task::{Context, Poll};
 
 pin_project! {

src/future/future/try_select.rs

@@ -0,0 +1,63 @@
+use std::pin::Pin;
+
+use async_macros::MaybeDone;
+use pin_project_lite::pin_project;
+
+use std::future::Future;
+use crate::task::{Context, Poll};
+
+pin_project! {
+    #[allow(missing_docs)]
+    #[allow(missing_debug_implementations)]
+    pub struct TrySelect<L, R> where L: Future, R: Future<Output = L::Output> {
+        #[pin] left: MaybeDone<L>,
+        #[pin] right: MaybeDone<R>,
+    }
+}
+
+impl<L, R> TrySelect<L, R>
+where
+    L: Future,
+    R: Future<Output = L::Output>,
+{
+    pub(crate) fn new(left: L, right: R) -> Self {
+        Self {
+            left: MaybeDone::new(left),
+            right: MaybeDone::new(right),
+        }
+    }
+}
+
+impl<L, R, T, E> Future for TrySelect<L, R>
+where
+    L: Future<Output = Result<T, E>>,
+    R: Future<Output = L::Output>,
+{
+    type Output = L::Output;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let this = self.project();
+        let mut left_errored = false;
+
+        // Check if the left future is ready & successful. Continue if not.
+        let mut left = this.left;
+        if Future::poll(Pin::new(&mut left), cx).is_ready() {
+            if left.as_ref().output().unwrap().is_ok() {
+                return Poll::Ready(left.take().unwrap());
+            } else {
+                left_errored = true;
+            }
+        }
+
+        // Check if the right future is ready & successful. Return err if left
+        // future also resolved to err. Continue if not.
+        let mut right = this.right;
+        if Future::poll(Pin::new(&mut right), cx).is_ready() {
+            if right.as_ref().output().unwrap().is_ok() || left_errored {
+                return Poll::Ready(right.take().unwrap());
+            }
+        }
+
+        Poll::Pending
+    }
+}