| 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: impl Future = 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 = 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 | |