multipart_write/stream/
write_complete.rs

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