1 | use super::assert_future; |
2 | use crate::future::TryFutureExt; |
3 | use alloc::vec::Vec; |
4 | use core::iter::FromIterator; |
5 | use core::mem; |
6 | use core::pin::Pin; |
7 | use futures_core::future::{Future, TryFuture}; |
8 | use futures_core::task::{Context, Poll}; |
9 | |
10 | /// Future for the [`select_ok`] function. |
11 | #[derive(Debug)] |
12 | #[must_use = "futures do nothing unless you `.await` or poll them" ] |
13 | pub struct SelectOk<Fut> { |
14 | inner: Vec<Fut>, |
15 | } |
16 | |
17 | impl<Fut: Unpin> Unpin for SelectOk<Fut> {} |
18 | |
19 | /// Creates a new future which will select the first successful future over a list of futures. |
20 | /// |
21 | /// The returned future will wait for any future within `iter` to be ready and Ok. Unlike |
22 | /// `select_all`, this will only return the first successful completion, or the last |
23 | /// failure. This is useful in contexts where any success is desired and failures |
24 | /// are ignored, unless all the futures fail. |
25 | /// |
26 | /// This function is only available when the `std` or `alloc` feature of this |
27 | /// library is activated, and it is activated by default. |
28 | /// |
29 | /// # Panics |
30 | /// |
31 | /// This function will panic if the iterator specified contains no items. |
32 | pub fn select_ok<I>(iter: I) -> SelectOk<I::Item> |
33 | where |
34 | I: IntoIterator, |
35 | I::Item: TryFuture + Unpin, |
36 | { |
37 | let ret = SelectOk { inner: iter.into_iter().collect() }; |
38 | assert!(!ret.inner.is_empty(), "iterator provided to select_ok was empty" ); |
39 | assert_future::< |
40 | Result<(<I::Item as TryFuture>::Ok, Vec<I::Item>), <I::Item as TryFuture>::Error>, |
41 | _, |
42 | >(ret) |
43 | } |
44 | |
45 | impl<Fut: TryFuture + Unpin> Future for SelectOk<Fut> { |
46 | type Output = Result<(Fut::Ok, Vec<Fut>), Fut::Error>; |
47 | |
48 | fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { |
49 | // loop until we've either exhausted all errors, a success was hit, or nothing is ready |
50 | loop { |
51 | let item = |
52 | self.inner.iter_mut().enumerate().find_map(|(i, f)| match f.try_poll_unpin(cx) { |
53 | Poll::Pending => None, |
54 | Poll::Ready(e) => Some((i, e)), |
55 | }); |
56 | match item { |
57 | Some((idx, res)) => { |
58 | // always remove Ok or Err, if it's not the last Err continue looping |
59 | drop(self.inner.remove(idx)); |
60 | match res { |
61 | Ok(e) => { |
62 | let rest = mem::take(&mut self.inner); |
63 | return Poll::Ready(Ok((e, rest))); |
64 | } |
65 | Err(e) => { |
66 | if self.inner.is_empty() { |
67 | return Poll::Ready(Err(e)); |
68 | } |
69 | } |
70 | } |
71 | } |
72 | None => { |
73 | // based on the filter above, nothing is ready, return |
74 | return Poll::Pending; |
75 | } |
76 | } |
77 | } |
78 | } |
79 | } |
80 | |
81 | impl<Fut: TryFuture + Unpin> FromIterator<Fut> for SelectOk<Fut> { |
82 | fn from_iter<T: IntoIterator<Item = Fut>>(iter: T) -> Self { |
83 | select_ok(iter) |
84 | } |
85 | } |
86 | |