multipart_write/stream/
collect_complete.rs

1use crate::{FusedMultipartWrite, MultipartWrite};
2
3use futures_core::future::{FusedFuture, Future};
4use futures_core::ready;
5use futures_core::stream::Stream;
6use std::fmt::{self, Debug, Formatter};
7use std::pin::Pin;
8use std::task::{Context, Poll};
9
10pin_project_lite::pin_project! {
11    /// Future for [`collect_complete`].
12    ///
13    /// [`collect_complete`]: super::MultipartStreamExt::collect_complete
14    #[must_use = "futures do nothing unless polled"]
15    pub struct CollectComplete<St: Stream, Wr> {
16        #[pin]
17        writer: Wr,
18        #[pin]
19        stream: Option<St>,
20        buffered: Option<St::Item>,
21        is_terminated: bool,
22    }
23}
24
25impl<St: Stream, Wr> CollectComplete<St, Wr> {
26    pub(super) fn new(stream: St, writer: Wr) -> Self {
27        Self {
28            writer,
29            stream: Some(stream),
30            buffered: None,
31            is_terminated: false,
32        }
33    }
34}
35
36impl<St, Wr> FusedFuture for CollectComplete<St, Wr>
37where
38    St: Stream,
39    Wr: FusedMultipartWrite<St::Item>,
40{
41    fn is_terminated(&self) -> bool {
42        self.writer.is_terminated() || self.is_terminated
43    }
44}
45
46impl<St, Wr> Future for CollectComplete<St, Wr>
47where
48    St: Stream,
49    Wr: MultipartWrite<St::Item>,
50{
51    type Output = Result<Wr::Output, Wr::Error>;
52
53    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
54        let mut this = self.project();
55
56        loop {
57            if this.buffered.is_some() {
58                loop {
59                    match this.writer.as_mut().poll_ready(cx) {
60                        Poll::Ready(Ok(())) => break,
61                        Poll::Ready(Err(e)) => return Poll::Ready(Err(e)),
62                        Poll::Pending => match this.writer.as_mut().poll_flush(cx)? {
63                            Poll::Pending => return Poll::Pending,
64                            Poll::Ready(()) => {}
65                        },
66                    }
67                }
68
69                let _ = this
70                    .writer
71                    .as_mut()
72                    .start_send(this.buffered.take().unwrap())?;
73            }
74
75            let Some(st) = this.stream.as_mut().as_pin_mut() else {
76                ready!(this.writer.as_mut().poll_flush(cx))?;
77                let output = ready!(this.writer.as_mut().poll_complete(cx));
78                *this.is_terminated = true;
79                return Poll::Ready(output);
80            };
81
82            match ready!(st.poll_next(cx)) {
83                Some(it) => *this.buffered = Some(it),
84                None => this.stream.set(None),
85            }
86        }
87    }
88}
89
90impl<St, Wr> Debug for CollectComplete<St, Wr>
91where
92    St: Stream + Debug,
93    St::Item: Debug,
94    Wr: Debug,
95{
96    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
97        f.debug_struct("CollectComplete")
98            .field("writer", &self.writer)
99            .field("stream", &self.stream)
100            .field("buffered", &self.buffered)
101            .field("is_terminated", &self.is_terminated)
102            .finish()
103    }
104}