1 //! An unbounded set of streams
2
3 use core::fmt::{self, Debug};
4 use core::iter::FromIterator;
5 use core::pin::Pin;
6
7 use futures_core::ready;
8 use futures_core::stream::{FusedStream, Stream};
9 use futures_core::task::{Context, Poll};
10
11 use pin_project_lite::pin_project;
12
13 use super::assert_stream;
14 use crate::stream::{futures_unordered, FuturesUnordered, StreamExt, StreamFuture};
15
16 pin_project! {
17 /// An unbounded set of streams
18 ///
19 /// This "combinator" provides the ability to maintain a set of streams
20 /// and drive them all to completion.
21 ///
22 /// Streams are pushed into this set and their realized values are
23 /// yielded as they become ready. Streams will only be polled when they
24 /// generate notifications. This allows to coordinate a large number of streams.
25 ///
26 /// Note that you can create a ready-made `SelectAll` via the
27 /// `select_all` function in the `stream` module, or you can start with an
28 /// empty set with the `SelectAll::new` constructor.
29 #[must_use = "streams do nothing unless polled"]
30 pub struct SelectAll<St> {
31 #[pin]
32 inner: FuturesUnordered<StreamFuture<St>>,
33 }
34 }
35
36 impl<St: Debug> Debug for SelectAll<St> {
fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result37 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
38 write!(f, "SelectAll {{ ... }}")
39 }
40 }
41
42 impl<St: Stream + Unpin> SelectAll<St> {
43 /// Constructs a new, empty `SelectAll`
44 ///
45 /// The returned `SelectAll` does not contain any streams and, in this
46 /// state, `SelectAll::poll` will return `Poll::Ready(None)`.
new() -> Self47 pub fn new() -> Self {
48 Self { inner: FuturesUnordered::new() }
49 }
50
51 /// Returns the number of streams contained in the set.
52 ///
53 /// This represents the total number of in-flight streams.
len(&self) -> usize54 pub fn len(&self) -> usize {
55 self.inner.len()
56 }
57
58 /// Returns `true` if the set contains no streams
is_empty(&self) -> bool59 pub fn is_empty(&self) -> bool {
60 self.inner.is_empty()
61 }
62
63 /// Push a stream into the set.
64 ///
65 /// This function submits the given stream to the set for managing. This
66 /// function will not call `poll` on the submitted stream. The caller must
67 /// ensure that `SelectAll::poll` is called in order to receive task
68 /// notifications.
push(&mut self, stream: St)69 pub fn push(&mut self, stream: St) {
70 self.inner.push(stream.into_future());
71 }
72
73 /// Returns an iterator that allows inspecting each stream in the set.
iter(&self) -> Iter<'_, St>74 pub fn iter(&self) -> Iter<'_, St> {
75 Iter(self.inner.iter())
76 }
77
78 /// Returns an iterator that allows modifying each stream in the set.
iter_mut(&mut self) -> IterMut<'_, St>79 pub fn iter_mut(&mut self) -> IterMut<'_, St> {
80 IterMut(self.inner.iter_mut())
81 }
82
83 /// Clears the set, removing all streams.
clear(&mut self)84 pub fn clear(&mut self) {
85 self.inner.clear()
86 }
87 }
88
89 impl<St: Stream + Unpin> Default for SelectAll<St> {
default() -> Self90 fn default() -> Self {
91 Self::new()
92 }
93 }
94
95 impl<St: Stream + Unpin> Stream for SelectAll<St> {
96 type Item = St::Item;
97
poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>>98 fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
99 loop {
100 match ready!(self.inner.poll_next_unpin(cx)) {
101 Some((Some(item), remaining)) => {
102 self.push(remaining);
103 return Poll::Ready(Some(item));
104 }
105 Some((None, _)) => {
106 // `FuturesUnordered` thinks it isn't terminated
107 // because it yielded a Some.
108 // We do not return, but poll `FuturesUnordered`
109 // in the next loop iteration.
110 }
111 None => return Poll::Ready(None),
112 }
113 }
114 }
115 }
116
117 impl<St: Stream + Unpin> FusedStream for SelectAll<St> {
is_terminated(&self) -> bool118 fn is_terminated(&self) -> bool {
119 self.inner.is_terminated()
120 }
121 }
122
123 /// Convert a list of streams into a `Stream` of results from the streams.
124 ///
125 /// This essentially takes a list of streams (e.g. a vector, an iterator, etc.)
126 /// and bundles them together into a single stream.
127 /// The stream will yield items as they become available on the underlying
128 /// streams internally, in the order they become available.
129 ///
130 /// Note that the returned set can also be used to dynamically push more
131 /// streams into the set as they become available.
132 ///
133 /// This function is only available when the `std` or `alloc` feature of this
134 /// library is activated, and it is activated by default.
select_all<I>(streams: I) -> SelectAll<I::Item> where I: IntoIterator, I::Item: Stream + Unpin,135 pub fn select_all<I>(streams: I) -> SelectAll<I::Item>
136 where
137 I: IntoIterator,
138 I::Item: Stream + Unpin,
139 {
140 let mut set = SelectAll::new();
141
142 for stream in streams {
143 set.push(stream);
144 }
145
146 assert_stream::<<I::Item as Stream>::Item, _>(set)
147 }
148
149 impl<St: Stream + Unpin> FromIterator<St> for SelectAll<St> {
from_iter<T: IntoIterator<Item = St>>(iter: T) -> Self150 fn from_iter<T: IntoIterator<Item = St>>(iter: T) -> Self {
151 select_all(iter)
152 }
153 }
154
155 impl<St: Stream + Unpin> Extend<St> for SelectAll<St> {
extend<T: IntoIterator<Item = St>>(&mut self, iter: T)156 fn extend<T: IntoIterator<Item = St>>(&mut self, iter: T) {
157 for st in iter {
158 self.push(st)
159 }
160 }
161 }
162
163 impl<St: Stream + Unpin> IntoIterator for SelectAll<St> {
164 type Item = St;
165 type IntoIter = IntoIter<St>;
166
into_iter(self) -> Self::IntoIter167 fn into_iter(self) -> Self::IntoIter {
168 IntoIter(self.inner.into_iter())
169 }
170 }
171
172 impl<'a, St: Stream + Unpin> IntoIterator for &'a SelectAll<St> {
173 type Item = &'a St;
174 type IntoIter = Iter<'a, St>;
175
into_iter(self) -> Self::IntoIter176 fn into_iter(self) -> Self::IntoIter {
177 self.iter()
178 }
179 }
180
181 impl<'a, St: Stream + Unpin> IntoIterator for &'a mut SelectAll<St> {
182 type Item = &'a mut St;
183 type IntoIter = IterMut<'a, St>;
184
into_iter(self) -> Self::IntoIter185 fn into_iter(self) -> Self::IntoIter {
186 self.iter_mut()
187 }
188 }
189
190 /// Immutable iterator over all streams in the unordered set.
191 #[derive(Debug)]
192 pub struct Iter<'a, St: Unpin>(futures_unordered::Iter<'a, StreamFuture<St>>);
193
194 /// Mutable iterator over all streams in the unordered set.
195 #[derive(Debug)]
196 pub struct IterMut<'a, St: Unpin>(futures_unordered::IterMut<'a, StreamFuture<St>>);
197
198 /// Owned iterator over all streams in the unordered set.
199 #[derive(Debug)]
200 pub struct IntoIter<St: Unpin>(futures_unordered::IntoIter<StreamFuture<St>>);
201
202 impl<'a, St: Stream + Unpin> Iterator for Iter<'a, St> {
203 type Item = &'a St;
204
next(&mut self) -> Option<Self::Item>205 fn next(&mut self) -> Option<Self::Item> {
206 let st = self.0.next()?;
207 let next = st.get_ref();
208 // This should always be true because FuturesUnordered removes completed futures.
209 debug_assert!(next.is_some());
210 next
211 }
212
size_hint(&self) -> (usize, Option<usize>)213 fn size_hint(&self) -> (usize, Option<usize>) {
214 self.0.size_hint()
215 }
216 }
217
218 impl<St: Stream + Unpin> ExactSizeIterator for Iter<'_, St> {}
219
220 impl<'a, St: Stream + Unpin> Iterator for IterMut<'a, St> {
221 type Item = &'a mut St;
222
next(&mut self) -> Option<Self::Item>223 fn next(&mut self) -> Option<Self::Item> {
224 let st = self.0.next()?;
225 let next = st.get_mut();
226 // This should always be true because FuturesUnordered removes completed futures.
227 debug_assert!(next.is_some());
228 next
229 }
230
size_hint(&self) -> (usize, Option<usize>)231 fn size_hint(&self) -> (usize, Option<usize>) {
232 self.0.size_hint()
233 }
234 }
235
236 impl<St: Stream + Unpin> ExactSizeIterator for IterMut<'_, St> {}
237
238 impl<St: Stream + Unpin> Iterator for IntoIter<St> {
239 type Item = St;
240
next(&mut self) -> Option<Self::Item>241 fn next(&mut self) -> Option<Self::Item> {
242 let st = self.0.next()?;
243 let next = st.into_inner();
244 // This should always be true because FuturesUnordered removes completed futures.
245 debug_assert!(next.is_some());
246 next
247 }
248
size_hint(&self) -> (usize, Option<usize>)249 fn size_hint(&self) -> (usize, Option<usize>) {
250 self.0.size_hint()
251 }
252 }
253
254 impl<St: Stream + Unpin> ExactSizeIterator for IntoIter<St> {}
255