multipart_write/stream/
collect_writer.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_writer`].
12    ///
13    /// [`collect_writer`]: super::MultipartStreamExt::collect_writer
14    #[must_use = "futures do nothing unless polled"]
15    pub struct CollectWriter<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> CollectWriter<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 CollectWriter<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 CollectWriter<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                let output = ready!(this.writer.as_mut().poll_complete(cx));
77                *this.is_terminated = true;
78                return Poll::Ready(output);
79            };
80
81            match ready!(st.poll_next(cx)) {
82                Some(it) => *this.buffered = Some(it),
83                None => this.stream.set(None),
84            }
85        }
86    }
87}
88
89impl<St, Wr> Debug for CollectWriter<St, Wr>
90where
91    St: Stream + Debug,
92    St::Item: Debug,
93    Wr: Debug,
94{
95    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
96        f.debug_struct("CollectWriter")
97            .field("writer", &self.writer)
98            .field("stream", &self.stream)
99            .field("buffered", &self.buffered)
100            .field("is_terminated", &self.is_terminated)
101            .finish()
102    }
103}
multipart_write/stream/collect_writer.rs

multipart_write/stream/
collect_writer.rs
