leyline/core/
body_stream.rs

1use std::future::Future;
2use std::pin::Pin;
3use std::task::{Context, Poll, ready};
4use std::time::Duration;
5
6use bytes::Bytes;
7use futures_util::Stream;
8use tokio::sync::mpsc;
9use tokio::time::Sleep;
10use tokio_util::sync::{CancellationToken, WaitForCancellationFutureOwned};
11
12use crate::core::config::HostPass;
13
14use crate::core::session::decompress::Decoder;
15use crate::trace::BodyWatch;
16
17pub struct BodyStream {
18    rx: mpsc::Receiver<std::io::Result<Bytes>>,
19    read_timeout: Option<Duration>,
20    idle: Option<Pin<Box<Sleep>>>,
21    body_limit: Option<Duration>,
22    body_deadline: Option<Pin<Box<Sleep>>>,
23    decoded: Option<Box<Decoded>>,
24    watch: Option<BodyWatch>,
25    shutdown: Option<Pin<Box<WaitForCancellationFutureOwned>>>,
26    ended: bool,
27    pass: Option<HostPass>,
28}
29
30struct Decoded {
31    raw: BodyStream,
32    decoder: Decoder,
33    done: bool,
34}
35
36impl Decoded {
37    fn poll(&mut self, cx: &mut Context<'_>) -> Poll<Option<std::io::Result<Bytes>>> {
38        loop {
39            if self.done {
40                return Poll::Ready(None);
41            }
42            let mut out = Vec::new();
43            let step = match Pin::new(&mut self.raw).poll_next(cx) {
44                Poll::Pending => return Poll::Pending,
45                Poll::Ready(Some(Err(e))) => {
46                    self.done = true;
47                    return Poll::Ready(Some(Err(e)));
48                }
49                Poll::Ready(Some(Ok(chunk))) => self.decoder.feed(&chunk, &mut out),
50                Poll::Ready(None) => {
51                    self.done = true;
52                    self.decoder.finish(&mut out)
53                }
54            };
55            if let Err(e) = step {
56                self.done = true;
57                return Poll::Ready(Some(Err(e.into_io())));
58            }
59            if !out.is_empty() {
60                return Poll::Ready(Some(Ok(Bytes::from(out))));
61            }
62        }
63    }
64}
65
66impl BodyStream {
67    pub(crate) fn new(rx: mpsc::Receiver<std::io::Result<Bytes>>) -> Self {
68        Self {
69            rx,
70            read_timeout: None,
71            idle: None,
72            body_limit: None,
73            body_deadline: None,
74            decoded: None,
75            watch: None,
76            shutdown: None,
77            ended: false,
78            pass: None,
79        }
80    }
81
82    pub(crate) fn decoded(mut raw: BodyStream, decoder: Decoder) -> Self {
83        let (_, rx) = mpsc::channel(1);
84        let watch = raw.watch.take();
85        Self {
86            watch,
87            decoded: Some(Box::new(Decoded {
88                raw,
89                decoder,
90                done: false,
91            })),
92            ..Self::new(rx)
93        }
94    }
95
96    pub(crate) fn from_bytes(buf: Bytes) -> Self {
97        let (tx, rx) = mpsc::channel(1);
98        if !buf.is_empty() {
99            tx.try_send(Ok(buf))
100                .expect("fresh channel with capacity 1 has room");
101        }
102        drop(tx);
103        Self::new(rx)
104    }
105
106    pub(crate) fn set_read_timeout(&mut self, timeout: Option<Duration>) {
107        self.read_timeout = timeout;
108    }
109
110    pub(crate) fn set_body_timeout(&mut self, timeout: Option<Duration>) {
111        self.body_limit = timeout;
112        self.body_deadline = None;
113    }
114
115    fn poll_body_deadline(&mut self, cx: &mut Context<'_>) -> Poll<std::io::Error> {
116        if let Some(limit) = self.body_limit.take() {
117            self.body_deadline = Some(Box::pin(tokio::time::sleep(limit)));
118        }
119        let Some(sleep) = self.body_deadline.as_mut() else {
120            return Poll::Pending;
121        };
122        ready!(sleep.as_mut().poll(cx));
123        Poll::Ready(timed_out(
124            "body timeout: response body not complete within the configured body timeout",
125        ))
126    }
127
128    pub(crate) fn stop_on(&mut self, token: &CancellationToken) {
129        self.shutdown = Some(Box::pin(token.clone().cancelled_owned()));
130    }
131
132    pub(crate) fn hold(&mut self, pass: Option<HostPass>) {
133        self.pass = pass;
134    }
135
136    fn poll_shutdown(&mut self, cx: &mut Context<'_>) -> bool {
137        self.shutdown
138            .as_mut()
139            .is_some_and(|wait| wait.as_mut().poll(cx).is_ready())
140    }
141
142    fn finish(&mut self, item: Option<&std::io::Result<Bytes>>) {
143        if self.ended {
144            return;
145        }
146        self.ended = true;
147        self.rx.close();
148        let (_, closed) = mpsc::channel(1);
149        drop(std::mem::replace(&mut self.rx, closed));
150        self.decoded = None;
151        self.idle = None;
152        self.body_limit = None;
153        self.body_deadline = None;
154        self.shutdown = None;
155        self.pass = None;
156        if let Some(mut watch) = self.watch.take() {
157            watch.observe(item);
158        }
159    }
160
161    fn end_with(&mut self, err: std::io::Error) -> Poll<Option<std::io::Result<Bytes>>> {
162        let item = Some(Err(err));
163        self.finish(item.as_ref());
164        Poll::Ready(item)
165    }
166
167    pub(crate) fn watch_end(&mut self) {
168        self.watch = BodyWatch::begin();
169    }
170
171    fn poll_raw(&mut self, cx: &mut Context<'_>) -> Poll<Option<std::io::Result<Bytes>>> {
172        match self.rx.poll_recv(cx) {
173            Poll::Ready(item) => {
174                self.idle = None;
175                Poll::Ready(item)
176            }
177            Poll::Pending => {
178                let Some(timeout) = self.read_timeout else {
179                    return Poll::Pending;
180                };
181                let sleep = self
182                    .idle
183                    .get_or_insert_with(|| Box::pin(tokio::time::sleep(timeout)));
184                match sleep.as_mut().poll(cx) {
185                    Poll::Ready(()) => {
186                        self.idle = None;
187                        Poll::Ready(Some(Err(timed_out(
188                            "read timeout: no response body chunk within the configured read_timeout",
189                        ))))
190                    }
191                    Poll::Pending => Poll::Pending,
192                }
193            }
194        }
195    }
196}
197
198fn timed_out(message: &'static str) -> std::io::Error {
199    std::io::Error::new(std::io::ErrorKind::TimedOut, message)
200}
201
202impl Stream for BodyStream {
203    type Item = std::io::Result<Bytes>;
204
205    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
206        let this = self.get_mut();
207        if this.ended {
208            return Poll::Ready(None);
209        }
210        if this.poll_shutdown(cx) {
211            return this.end_with(std::io::Error::other(crate::core::error::Error::shut_down()));
212        }
213        if let Poll::Ready(err) = this.poll_body_deadline(cx) {
214            return this.end_with(err);
215        }
216        let polled = match this.decoded.as_mut() {
217            Some(decoded) => decoded.poll(cx),
218            None => this.poll_raw(cx),
219        };
220        match &polled {
221            Poll::Ready(item @ Some(Ok(_))) => {
222                if let Some(watch) = this.watch.as_mut() {
223                    watch.observe(item.as_ref());
224                }
225            }
226            Poll::Ready(item) => this.finish(item.as_ref()),
227            Poll::Pending => {}
228        }
229        polled
230    }
231}
232
233impl std::fmt::Debug for BodyStream {
234    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
235        f.debug_struct("BodyStream").finish_non_exhaustive()
236    }
237}
leyline/core/body_stream.rs

leyline/core/
body_stream.rs
