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