1 | //! Definition of the `JoinAll` combinator, waiting for all of a list of futures |
2 | //! to finish. |
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, MaybeDone}; |
14 | |
15 | #[cfg_attr (target_os = "none" , cfg(target_has_atomic = "ptr" ))] |
16 | use crate::stream::{Collect, FuturesOrdered, StreamExt}; |
17 | |
18 | pub(crate) fn iter_pin_mut<T>(slice: Pin<&mut [T]>) -> impl Iterator<Item = Pin<&mut T>> { |
19 | // Safety: `std` _could_ make this unsound if it were to decide Pin's |
20 | // invariants aren't required to transmit through slices. Otherwise this has |
21 | // the same safety as a normal field pin projection. |
22 | unsafe { slice.get_unchecked_mut() }.iter_mut().map(|t: &mut T| unsafe { Pin::new_unchecked(pointer:t) }) |
23 | } |
24 | |
25 | #[must_use = "futures do nothing unless you `.await` or poll them" ] |
26 | /// Future for the [`join_all`] function. |
27 | pub struct JoinAll<F> |
28 | where |
29 | F: Future, |
30 | { |
31 | kind: JoinAllKind<F>, |
32 | } |
33 | |
34 | #[cfg_attr (target_os = "none" , cfg(target_has_atomic = "ptr" ))] |
35 | pub(crate) const SMALL: usize = 30; |
36 | |
37 | enum JoinAllKind<F> |
38 | where |
39 | F: Future, |
40 | { |
41 | Small { |
42 | elems: Pin<Box<[MaybeDone<F>]>>, |
43 | }, |
44 | #[cfg_attr (target_os = "none" , cfg(target_has_atomic = "ptr" ))] |
45 | Big { |
46 | fut: Collect<FuturesOrdered<F>, Vec<F::Output>>, |
47 | }, |
48 | } |
49 | |
50 | impl<F> fmt::Debug for JoinAll<F> |
51 | where |
52 | F: Future + fmt::Debug, |
53 | F::Output: fmt::Debug, |
54 | { |
55 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { |
56 | match self.kind { |
57 | JoinAllKind::Small { ref elems: &Pin]>> } => { |
58 | f.debug_struct("JoinAll" ).field(name:"elems" , value:elems).finish() |
59 | } |
60 | #[cfg_attr (target_os = "none" , cfg(target_has_atomic = "ptr" ))] |
61 | JoinAllKind::Big { ref fut: &Collect, …>, .. } => fmt::Debug::fmt(self:fut, f), |
62 | } |
63 | } |
64 | } |
65 | |
66 | /// Creates a future which represents a collection of the outputs of the futures |
67 | /// given. |
68 | /// |
69 | /// The returned future will drive execution for all of its underlying futures, |
70 | /// collecting the results into a destination `Vec<T>` in the same order as they |
71 | /// were provided. |
72 | /// |
73 | /// This function is only available when the `std` or `alloc` feature of this |
74 | /// library is activated, and it is activated by default. |
75 | /// |
76 | /// # See Also |
77 | /// |
78 | /// `join_all` will switch to the more powerful [`FuturesOrdered`] for performance |
79 | /// reasons if the number of futures is large. You may want to look into using it or |
80 | /// its counterpart [`FuturesUnordered`][crate::stream::FuturesUnordered] directly. |
81 | /// |
82 | /// Some examples for additional functionality provided by these are: |
83 | /// |
84 | /// * Adding new futures to the set even after it has been started. |
85 | /// |
86 | /// * Only polling the specific futures that have been woken. In cases where |
87 | /// you have a lot of futures this will result in much more efficient polling. |
88 | /// |
89 | /// # Examples |
90 | /// |
91 | /// ``` |
92 | /// # futures::executor::block_on(async { |
93 | /// use futures::future::join_all; |
94 | /// |
95 | /// async fn foo(i: u32) -> u32 { i } |
96 | /// |
97 | /// let futures = vec![foo(1), foo(2), foo(3)]; |
98 | /// |
99 | /// assert_eq!(join_all(futures).await, [1, 2, 3]); |
100 | /// # }); |
101 | /// ``` |
102 | pub fn join_all<I>(iter: I) -> JoinAll<I::Item> |
103 | where |
104 | I: IntoIterator, |
105 | I::Item: Future, |
106 | { |
107 | let iter: ::IntoIter = iter.into_iter(); |
108 | |
109 | #[cfg (target_os = "none" )] |
110 | #[cfg_attr (target_os = "none" , cfg(not(target_has_atomic = "ptr" )))] |
111 | { |
112 | let kind = |
113 | JoinAllKind::Small { elems: iter.map(MaybeDone::Future).collect::<Box<[_]>>().into() }; |
114 | |
115 | assert_future::<Vec<<I::Item as Future>::Output>, _>(JoinAll { kind }) |
116 | } |
117 | |
118 | #[cfg_attr (target_os = "none" , cfg(target_has_atomic = "ptr" ))] |
119 | { |
120 | let kind: JoinAllKind<::Item> = match iter.size_hint().1 { |
121 | Some(max: usize) if max <= SMALL => JoinAllKind::Small { |
122 | elems: iter.map(MaybeDone::Future).collect::<Box<[_]>>().into(), |
123 | }, |
124 | _ => JoinAllKind::Big { fut: iter.collect::<FuturesOrdered<_>>().collect() }, |
125 | }; |
126 | |
127 | assert_future::<Vec<<I::Item as Future>::Output>, _>(JoinAll { kind }) |
128 | } |
129 | } |
130 | |
131 | impl<F> Future for JoinAll<F> |
132 | where |
133 | F: Future, |
134 | { |
135 | type Output = Vec<F::Output>; |
136 | |
137 | fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { |
138 | match &mut self.kind { |
139 | JoinAllKind::Small { elems } => { |
140 | let mut all_done = true; |
141 | |
142 | for elem in iter_pin_mut(elems.as_mut()) { |
143 | if elem.poll(cx).is_pending() { |
144 | all_done = false; |
145 | } |
146 | } |
147 | |
148 | if all_done { |
149 | let mut elems = mem::replace(elems, Box::pin([])); |
150 | let result = |
151 | iter_pin_mut(elems.as_mut()).map(|e| e.take_output().unwrap()).collect(); |
152 | Poll::Ready(result) |
153 | } else { |
154 | Poll::Pending |
155 | } |
156 | } |
157 | #[cfg_attr (target_os = "none" , cfg(target_has_atomic = "ptr" ))] |
158 | JoinAllKind::Big { fut } => Pin::new(fut).poll(cx), |
159 | } |
160 | } |
161 | } |
162 | |
163 | impl<F: Future> FromIterator<F> for JoinAll<F> { |
164 | fn from_iter<T: IntoIterator<Item = F>>(iter: T) -> Self { |
165 | join_all(iter) |
166 | } |
167 | } |
168 | |