• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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