futures_util/stream/
select_all.rs

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

futures_util/stream/
select_all.rs
