1use crate::stream::{FuturesUnordered, StreamExt};
2use alloc::collections::binary_heap::{BinaryHeap, PeekMut};
3use core::cmp::Ordering;
4use core::fmt::{self, Debug};
5use core::iter::FromIterator;
6use core::pin::Pin;
7use futures_core::future::Future;
8use futures_core::ready;
9use futures_core::stream::Stream;
10use futures_core::{
11 task::{Context, Poll},
12 FusedStream,
13};
14use pin_project_lite::pin_project;
15
16pin_project! {
17 #[must_use = "futures do nothing unless you `.await` or poll them"]
18 #[derive(Debug)]
19 struct OrderWrapper<T> {
20 #[pin]
21 data: T, // A future or a future's output
22 index: isize,
23 }
24}
25
26impl<T> PartialEq for OrderWrapper<T> {
27 fn eq(&self, other: &Self) -> bool {
28 self.index == other.index
29 }
30}
31
32impl<T> Eq for OrderWrapper<T> {}
33
34impl<T> PartialOrd for OrderWrapper<T> {
35 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
36 Some(self.cmp(other))
37 }
38}
39
40impl<T> Ord for OrderWrapper<T> {
41 fn cmp(&self, other: &Self) -> Ordering {
42 // BinaryHeap is a max heap, so compare backwards here.
43 other.index.cmp(&self.index)
44 }
45}
46
47impl<T> Future for OrderWrapper<T>
48where
49 T: Future,
50{
51 type Output = OrderWrapper<T::Output>;
52
53 fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
54 let index = self.index;
55 self.project().data.poll(cx).map(|output| OrderWrapper { data: output, index })
56 }
57}
58
59/// An unbounded queue of futures.
60///
61/// This "combinator" is similar to [`FuturesUnordered`], but it imposes a FIFO
62/// order on top of the set of futures. While futures in the set will race to
63/// completion in parallel, results will only be returned in the order their
64/// originating futures were added to the queue.
65///
66/// Futures are pushed into this queue and their realized values are yielded in
67/// order. This structure is optimized to manage a large number of futures.
68/// Futures managed by [`FuturesOrdered`] will only be polled when they generate
69/// notifications. This reduces the required amount of work needed to coordinate
70/// large numbers of futures.
71///
72/// When a [`FuturesOrdered`] is first created, it does not contain any futures.
73/// Calling [`poll_next`](FuturesOrdered::poll_next) in this state will result
74/// in [`Poll::Ready(None)`](Poll::Ready) to be returned. Futures are submitted
75/// to the queue using [`push_back`](FuturesOrdered::push_back) (or
76/// [`push_front`](FuturesOrdered::push_front)); however, the future will
77/// **not** be polled at this point. [`FuturesOrdered`] will only poll managed
78/// futures when [`FuturesOrdered::poll_next`] is called. As such, it
79/// is important to call [`poll_next`](FuturesOrdered::poll_next) after pushing
80/// new futures.
81///
82/// If [`FuturesOrdered::poll_next`] returns [`Poll::Ready(None)`](Poll::Ready)
83/// this means that the queue is currently not managing any futures. A future
84/// may be submitted to the queue at a later time. At that point, a call to
85/// [`FuturesOrdered::poll_next`] will either return the future's resolved value
86/// **or** [`Poll::Pending`] if the future has not yet completed. When
87/// multiple futures are submitted to the queue, [`FuturesOrdered::poll_next`]
88/// will return [`Poll::Pending`] until the first future completes, even if
89/// some of the later futures have already completed.
90///
91/// Note that you can create a ready-made [`FuturesOrdered`] via the
92/// [`collect`](Iterator::collect) method, or you can start with an empty queue
93/// with the [`FuturesOrdered::new`] constructor.
94///
95/// This type is only available when the `std` or `alloc` feature of this
96/// library is activated, and it is activated by default.
97#[must_use = "streams do nothing unless polled"]
98pub struct FuturesOrdered<T: Future> {
99 in_progress_queue: FuturesUnordered<OrderWrapper<T>>,
100 queued_outputs: BinaryHeap<OrderWrapper<T::Output>>,
101 next_incoming_index: isize,
102 next_outgoing_index: isize,
103}
104
105impl<T: Future> Unpin for FuturesOrdered<T> {}
106
107impl<Fut: Future> FuturesOrdered<Fut> {
108 /// Constructs a new, empty `FuturesOrdered`
109 ///
110 /// The returned [`FuturesOrdered`] does not contain any futures and, in
111 /// this state, [`FuturesOrdered::poll_next`] will return
112 /// [`Poll::Ready(None)`](Poll::Ready).
113 pub fn new() -> Self {
114 Self {
115 in_progress_queue: FuturesUnordered::new(),
116 queued_outputs: BinaryHeap::new(),
117 next_incoming_index: 0,
118 next_outgoing_index: 0,
119 }
120 }
121
122 /// Returns the number of futures contained in the queue.
123 ///
124 /// This represents the total number of in-flight futures, both
125 /// those currently processing and those that have completed but
126 /// which are waiting for earlier futures to complete.
127 pub fn len(&self) -> usize {
128 self.in_progress_queue.len() + self.queued_outputs.len()
129 }
130
131 /// Returns `true` if the queue contains no futures
132 pub fn is_empty(&self) -> bool {
133 self.in_progress_queue.is_empty() && self.queued_outputs.is_empty()
134 }
135
136 /// Push a future into the queue.
137 ///
138 /// This function submits the given future to the internal set for managing.
139 /// This function will not call [`poll`](Future::poll) on the submitted
140 /// future. The caller must ensure that [`FuturesOrdered::poll_next`] is
141 /// called in order to receive task notifications.
142 #[deprecated(note = "use `push_back` instead")]
143 pub fn push(&mut self, future: Fut) {
144 self.push_back(future);
145 }
146
147 /// Pushes a future to the back of the queue.
148 ///
149 /// This function submits the given future to the internal set for managing.
150 /// This function will not call [`poll`](Future::poll) on the submitted
151 /// future. The caller must ensure that [`FuturesOrdered::poll_next`] is
152 /// called in order to receive task notifications.
153 pub fn push_back(&mut self, future: Fut) {
154 let wrapped = OrderWrapper { data: future, index: self.next_incoming_index };
155 self.next_incoming_index += 1;
156 self.in_progress_queue.push(wrapped);
157 }
158
159 /// Pushes a future to the front of the queue.
160 ///
161 /// This function submits the given future to the internal set for managing.
162 /// This function will not call [`poll`](Future::poll) on the submitted
163 /// future. The caller must ensure that [`FuturesOrdered::poll_next`] is
164 /// called in order to receive task notifications. This future will be
165 /// the next future to be returned complete.
166 pub fn push_front(&mut self, future: Fut) {
167 let wrapped = OrderWrapper { data: future, index: self.next_outgoing_index - 1 };
168 self.next_outgoing_index -= 1;
169 self.in_progress_queue.push(wrapped);
170 }
171}
172
173impl<Fut: Future> Default for FuturesOrdered<Fut> {
174 fn default() -> Self {
175 Self::new()
176 }
177}
178
179impl<Fut: Future> Stream for FuturesOrdered<Fut> {
180 type Item = Fut::Output;
181
182 fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
183 let this = &mut *self;
184
185 // Check to see if we've already received the next value
186 if let Some(next_output) = this.queued_outputs.peek_mut() {
187 if next_output.index == this.next_outgoing_index {
188 this.next_outgoing_index += 1;
189 return Poll::Ready(Some(PeekMut::pop(next_output).data));
190 }
191 }
192
193 loop {
194 match ready!(this.in_progress_queue.poll_next_unpin(cx)) {
195 Some(output) => {
196 if output.index == this.next_outgoing_index {
197 this.next_outgoing_index += 1;
198 return Poll::Ready(Some(output.data));
199 } else {
200 this.queued_outputs.push(output)
201 }
202 }
203 None => return Poll::Ready(None),
204 }
205 }
206 }
207
208 fn size_hint(&self) -> (usize, Option<usize>) {
209 let len = self.len();
210 (len, Some(len))
211 }
212}
213
214impl<Fut: Future> Debug for FuturesOrdered<Fut> {
215 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
216 write!(f, "FuturesOrdered {{ ... }}")
217 }
218}
219
220impl<Fut: Future> FromIterator<Fut> for FuturesOrdered<Fut> {
221 fn from_iter<T>(iter: T) -> Self
222 where
223 T: IntoIterator<Item = Fut>,
224 {
225 let acc = Self::new();
226 iter.into_iter().fold(acc, |mut acc, item| {
227 acc.push_back(item);
228 acc
229 })
230 }
231}
232
233impl<Fut: Future> FusedStream for FuturesOrdered<Fut> {
234 fn is_terminated(&self) -> bool {
235 self.in_progress_queue.is_terminated() && self.queued_outputs.is_empty()
236 }
237}
238
239impl<Fut: Future> Extend<Fut> for FuturesOrdered<Fut> {
240 fn extend<I>(&mut self, iter: I)
241 where
242 I: IntoIterator<Item = Fut>,
243 {
244 for item in iter {
245 self.push_back(item);
246 }
247 }
248}
249