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::{Stream, FusedStream};
9 use futures_core::task::{Context, Poll};
10
11 use super::assert_stream;
12 use crate::stream::{StreamExt, StreamFuture, FuturesUnordered};
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"]
27 pub struct SelectAll<St> {
28 inner: FuturesUnordered<StreamFuture<St>>,
29 }
30
31 impl<St: Debug> Debug for SelectAll<St> {
fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result32 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
33 write!(f, "SelectAll {{ ... }}")
34 }
35 }
36
37 impl<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)`.
new() -> Self42 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.
len(&self) -> usize49 pub fn len(&self) -> usize {
50 self.inner.len()
51 }
52
53 /// Returns `true` if the set contains no streams
is_empty(&self) -> bool54 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.
push(&mut self, stream: St)64 pub fn push(&mut self, stream: St) {
65 self.inner.push(stream.into_future());
66 }
67 }
68
69 impl<St: Stream + Unpin> Default for SelectAll<St> {
default() -> Self70 fn default() -> Self {
71 Self::new()
72 }
73 }
74
75 impl<St: Stream + Unpin> Stream for SelectAll<St> {
76 type Item = St::Item;
77
poll_next( mut self: Pin<&mut Self>, cx: &mut Context<'_>, ) -> Poll<Option<Self::Item>>78 fn poll_next(
79 mut self: Pin<&mut Self>,
80 cx: &mut Context<'_>,
81 ) -> Poll<Option<Self::Item>> {
82 loop {
83 match ready!(self.inner.poll_next_unpin(cx)) {
84 Some((Some(item), remaining)) => {
85 self.push(remaining);
86 return Poll::Ready(Some(item));
87 }
88 Some((None, _)) => {
89 // `FuturesUnordered` thinks it isn't terminated
90 // because it yielded a Some.
91 // We do not return, but poll `FuturesUnordered`
92 // in the next loop iteration.
93 }
94 None => return Poll::Ready(None),
95 }
96 }
97 }
98 }
99
100 impl<St: Stream + Unpin> FusedStream for SelectAll<St> {
is_terminated(&self) -> bool101 fn is_terminated(&self) -> bool {
102 self.inner.is_terminated()
103 }
104 }
105
106 /// Convert a list of streams into a `Stream` of results from the streams.
107 ///
108 /// This essentially takes a list of streams (e.g. a vector, an iterator, etc.)
109 /// and bundles them together into a single stream.
110 /// The stream will yield items as they become available on the underlying
111 /// streams internally, in the order they become available.
112 ///
113 /// Note that the returned set can also be used to dynamically push more
114 /// futures into the set as they become available.
115 ///
116 /// This function is only available when the `std` or `alloc` feature of this
117 /// library is activated, and it is activated by default.
select_all<I>(streams: I) -> SelectAll<I::Item> where I: IntoIterator, I::Item: Stream + Unpin118 pub fn select_all<I>(streams: I) -> SelectAll<I::Item>
119 where I: IntoIterator,
120 I::Item: Stream + Unpin
121 {
122 let mut set = SelectAll::new();
123
124 for stream in streams {
125 set.push(stream);
126 }
127
128 assert_stream::<<I::Item as Stream>::Item, _>(set)
129 }
130
131 impl<St: Stream + Unpin> FromIterator<St> for SelectAll<St> {
from_iter<T: IntoIterator<Item = St>>(iter: T) -> Self132 fn from_iter<T: IntoIterator<Item = St>>(iter: T) -> Self {
133 select_all(iter)
134 }
135 }
136
137 impl<St: Stream + Unpin> Extend<St> for SelectAll<St> {
extend<T: IntoIterator<Item = St>>(&mut self, iter: T)138 fn extend<T: IntoIterator<Item = St>>(&mut self, iter: T) {
139 for st in iter {
140 self.push(st)
141 }
142 }
143 }
144