1 | use crate::future::{IntoFuture, TryFutureExt}; |
2 | use crate::stream::{Fuse, FuturesUnordered, IntoStream, StreamExt}; |
3 | use core::pin::Pin; |
4 | use futures_core::future::TryFuture; |
5 | use futures_core::stream::{Stream, TryStream}; |
6 | use futures_core::task::{Context, Poll}; |
7 | #[cfg (feature = "sink" )] |
8 | use futures_sink::Sink; |
9 | use pin_project_lite::pin_project; |
10 | |
11 | pin_project! { |
12 | /// Stream for the |
13 | /// [`try_buffer_unordered`](super::TryStreamExt::try_buffer_unordered) method. |
14 | #[derive(Debug)] |
15 | #[must_use = "streams do nothing unless polled" ] |
16 | pub struct TryBufferUnordered<St> |
17 | where St: TryStream |
18 | { |
19 | #[pin] |
20 | stream: Fuse<IntoStream<St>>, |
21 | in_progress_queue: FuturesUnordered<IntoFuture<St::Ok>>, |
22 | max: usize, |
23 | } |
24 | } |
25 | |
26 | impl<St> TryBufferUnordered<St> |
27 | where |
28 | St: TryStream, |
29 | St::Ok: TryFuture, |
30 | { |
31 | pub(super) fn new(stream: St, n: usize) -> Self { |
32 | Self { |
33 | stream: IntoStream::new(stream).fuse(), |
34 | in_progress_queue: FuturesUnordered::new(), |
35 | max: n, |
36 | } |
37 | } |
38 | |
39 | delegate_access_inner!(stream, St, (. .)); |
40 | } |
41 | |
42 | impl<St> Stream for TryBufferUnordered<St> |
43 | where |
44 | St: TryStream, |
45 | St::Ok: TryFuture<Error = St::Error>, |
46 | { |
47 | type Item = Result<<St::Ok as TryFuture>::Ok, St::Error>; |
48 | |
49 | fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { |
50 | let mut this = self.project(); |
51 | |
52 | // First up, try to spawn off as many futures as possible by filling up |
53 | // our queue of futures. Propagate errors from the stream immediately. |
54 | while this.in_progress_queue.len() < *this.max { |
55 | match this.stream.as_mut().poll_next(cx)? { |
56 | Poll::Ready(Some(fut)) => this.in_progress_queue.push(fut.into_future()), |
57 | Poll::Ready(None) | Poll::Pending => break, |
58 | } |
59 | } |
60 | |
61 | // Attempt to pull the next value from the in_progress_queue |
62 | match this.in_progress_queue.poll_next_unpin(cx) { |
63 | x @ Poll::Pending | x @ Poll::Ready(Some(_)) => return x, |
64 | Poll::Ready(None) => {} |
65 | } |
66 | |
67 | // If more values are still coming from the stream, we're not done yet |
68 | if this.stream.is_done() { |
69 | Poll::Ready(None) |
70 | } else { |
71 | Poll::Pending |
72 | } |
73 | } |
74 | } |
75 | |
76 | // Forwarding impl of Sink from the underlying stream |
77 | #[cfg (feature = "sink" )] |
78 | impl<S, Item, E> Sink<Item> for TryBufferUnordered<S> |
79 | where |
80 | S: TryStream + Sink<Item, Error = E>, |
81 | S::Ok: TryFuture<Error = E>, |
82 | { |
83 | type Error = E; |
84 | |
85 | delegate_sink!(stream, Item); |
86 | } |
87 | |