1 | //! Definition of the `TryJoinAll` combinator, waiting for all of a list of |
2 | //! futures to finish with either success or error. |
3 | |
4 | use alloc::boxed::Box; |
5 | use alloc::vec::Vec; |
6 | use core::fmt; |
7 | use core::future::Future; |
8 | use core::iter::FromIterator; |
9 | use core::mem; |
10 | use core::pin::Pin; |
11 | use core::task::{Context, Poll}; |
12 | |
13 | use super::{assert_future, join_all, IntoFuture, TryFuture, TryMaybeDone}; |
14 | |
15 | #[cfg_attr (target_os = "none" , cfg(target_has_atomic = "ptr" ))] |
16 | use crate::stream::{FuturesOrdered, TryCollect, TryStreamExt}; |
17 | use crate::TryFutureExt; |
18 | |
19 | enum FinalState<E = ()> { |
20 | Pending, |
21 | AllDone, |
22 | Error(E), |
23 | } |
24 | |
25 | /// Future for the [`try_join_all`] function. |
26 | #[must_use = "futures do nothing unless you `.await` or poll them" ] |
27 | pub struct TryJoinAll<F> |
28 | where |
29 | F: TryFuture, |
30 | { |
31 | kind: TryJoinAllKind<F>, |
32 | } |
33 | |
34 | enum TryJoinAllKind<F> |
35 | where |
36 | F: TryFuture, |
37 | { |
38 | Small { |
39 | elems: Pin<Box<[TryMaybeDone<IntoFuture<F>>]>>, |
40 | }, |
41 | #[cfg_attr (target_os = "none" , cfg(target_has_atomic = "ptr" ))] |
42 | Big { |
43 | fut: TryCollect<FuturesOrdered<IntoFuture<F>>, Vec<F::Ok>>, |
44 | }, |
45 | } |
46 | |
47 | impl<F> fmt::Debug for TryJoinAll<F> |
48 | where |
49 | F: TryFuture + fmt::Debug, |
50 | F::Ok: fmt::Debug, |
51 | F::Error: fmt::Debug, |
52 | F::Output: fmt::Debug, |
53 | { |
54 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { |
55 | match self.kind { |
56 | TryJoinAllKind::Small { ref elems: &Pin>]>> } => { |
57 | f.debug_struct("TryJoinAll" ).field(name:"elems" , value:elems).finish() |
58 | } |
59 | #[cfg_attr (target_os = "none" , cfg(target_has_atomic = "ptr" ))] |
60 | TryJoinAllKind::Big { ref fut: &TryCollect, …>, .. } => fmt::Debug::fmt(self:fut, f), |
61 | } |
62 | } |
63 | } |
64 | |
65 | /// Creates a future which represents either a collection of the results of the |
66 | /// futures given or an error. |
67 | /// |
68 | /// The returned future will drive execution for all of its underlying futures, |
69 | /// collecting the results into a destination `Vec<T>` in the same order as they |
70 | /// were provided. |
71 | /// |
72 | /// If any future returns an error then all other futures will be canceled and |
73 | /// an error will be returned immediately. If all futures complete successfully, |
74 | /// however, then the returned future will succeed with a `Vec` of all the |
75 | /// successful results. |
76 | /// |
77 | /// This function is only available when the `std` or `alloc` feature of this |
78 | /// library is activated, and it is activated by default. |
79 | /// |
80 | /// # See Also |
81 | /// |
82 | /// `try_join_all` will switch to the more powerful [`FuturesOrdered`] for performance |
83 | /// reasons if the number of futures is large. You may want to look into using it or |
84 | /// it's counterpart [`FuturesUnordered`][crate::stream::FuturesUnordered] directly. |
85 | /// |
86 | /// Some examples for additional functionality provided by these are: |
87 | /// |
88 | /// * Adding new futures to the set even after it has been started. |
89 | /// |
90 | /// * Only polling the specific futures that have been woken. In cases where |
91 | /// you have a lot of futures this will result in much more efficient polling. |
92 | /// |
93 | /// |
94 | /// # Examples |
95 | /// |
96 | /// ``` |
97 | /// # futures::executor::block_on(async { |
98 | /// use futures::future::{self, try_join_all}; |
99 | /// |
100 | /// let futures = vec![ |
101 | /// future::ok::<u32, u32>(1), |
102 | /// future::ok::<u32, u32>(2), |
103 | /// future::ok::<u32, u32>(3), |
104 | /// ]; |
105 | /// |
106 | /// assert_eq!(try_join_all(futures).await, Ok(vec![1, 2, 3])); |
107 | /// |
108 | /// let futures = vec![ |
109 | /// future::ok::<u32, u32>(1), |
110 | /// future::err::<u32, u32>(2), |
111 | /// future::ok::<u32, u32>(3), |
112 | /// ]; |
113 | /// |
114 | /// assert_eq!(try_join_all(futures).await, Err(2)); |
115 | /// # }); |
116 | /// ``` |
117 | pub fn try_join_all<I>(iter: I) -> TryJoinAll<I::Item> |
118 | where |
119 | I: IntoIterator, |
120 | I::Item: TryFuture, |
121 | { |
122 | let iter = iter.into_iter().map(TryFutureExt::into_future); |
123 | |
124 | #[cfg (target_os = "none" )] |
125 | #[cfg_attr (target_os = "none" , cfg(not(target_has_atomic = "ptr" )))] |
126 | { |
127 | let kind = TryJoinAllKind::Small { |
128 | elems: iter.map(TryMaybeDone::Future).collect::<Box<[_]>>().into(), |
129 | }; |
130 | |
131 | assert_future::<Result<Vec<<I::Item as TryFuture>::Ok>, <I::Item as TryFuture>::Error>, _>( |
132 | TryJoinAll { kind }, |
133 | ) |
134 | } |
135 | |
136 | #[cfg_attr (target_os = "none" , cfg(target_has_atomic = "ptr" ))] |
137 | { |
138 | let kind = match iter.size_hint().1 { |
139 | Some(max) if max <= join_all::SMALL => TryJoinAllKind::Small { |
140 | elems: iter.map(TryMaybeDone::Future).collect::<Box<[_]>>().into(), |
141 | }, |
142 | _ => TryJoinAllKind::Big { fut: iter.collect::<FuturesOrdered<_>>().try_collect() }, |
143 | }; |
144 | |
145 | assert_future::<Result<Vec<<I::Item as TryFuture>::Ok>, <I::Item as TryFuture>::Error>, _>( |
146 | TryJoinAll { kind }, |
147 | ) |
148 | } |
149 | } |
150 | |
151 | impl<F> Future for TryJoinAll<F> |
152 | where |
153 | F: TryFuture, |
154 | { |
155 | type Output = Result<Vec<F::Ok>, F::Error>; |
156 | |
157 | fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { |
158 | match &mut self.kind { |
159 | TryJoinAllKind::Small { elems } => { |
160 | let mut state = FinalState::AllDone; |
161 | |
162 | for elem in join_all::iter_pin_mut(elems.as_mut()) { |
163 | match elem.try_poll(cx) { |
164 | Poll::Pending => state = FinalState::Pending, |
165 | Poll::Ready(Ok(())) => {} |
166 | Poll::Ready(Err(e)) => { |
167 | state = FinalState::Error(e); |
168 | break; |
169 | } |
170 | } |
171 | } |
172 | |
173 | match state { |
174 | FinalState::Pending => Poll::Pending, |
175 | FinalState::AllDone => { |
176 | let mut elems = mem::replace(elems, Box::pin([])); |
177 | let results = join_all::iter_pin_mut(elems.as_mut()) |
178 | .map(|e| e.take_output().unwrap()) |
179 | .collect(); |
180 | Poll::Ready(Ok(results)) |
181 | } |
182 | FinalState::Error(e) => { |
183 | let _ = mem::replace(elems, Box::pin([])); |
184 | Poll::Ready(Err(e)) |
185 | } |
186 | } |
187 | } |
188 | #[cfg_attr (target_os = "none" , cfg(target_has_atomic = "ptr" ))] |
189 | TryJoinAllKind::Big { fut } => Pin::new(fut).poll(cx), |
190 | } |
191 | } |
192 | } |
193 | |
194 | impl<F> FromIterator<F> for TryJoinAll<F> |
195 | where |
196 | F: TryFuture, |
197 | { |
198 | fn from_iter<T: IntoIterator<Item = F>>(iter: T) -> Self { |
199 | try_join_all(iter) |
200 | } |
201 | } |
202 | |