leyline/core/
body_stream.rs

1use std::future::Future;
2use std::pin::Pin;
3use std::task::{Context, Poll};
4use std::time::Duration;
5
6use bytes::Bytes;
7use futures_util::Stream;
8use tokio::sync::mpsc;
9use tokio::time::Sleep;
10
11pub struct BodyStream {
12    rx: mpsc::Receiver<std::io::Result<Bytes>>,
13    read_timeout: Option<Duration>,
14    idle: Option<Pin<Box<Sleep>>>,
15}
16
17impl BodyStream {
18    pub(crate) fn new(rx: mpsc::Receiver<std::io::Result<Bytes>>) -> Self {
19        Self {
20            rx,
21            read_timeout: None,
22            idle: None,
23        }
24    }
25
26    pub(crate) fn from_bytes(buf: Bytes) -> Self {
27        let (tx, rx) = mpsc::channel(1);
28        if !buf.is_empty() {
29            tx.try_send(Ok(buf))
30                .expect("fresh channel with capacity 1 has room");
31        }
32        drop(tx);
33        Self::new(rx)
34    }
35
36    pub(crate) fn set_read_timeout(&mut self, timeout: Option<Duration>) {
37        self.read_timeout = timeout;
38    }
39}
40
41impl Stream for BodyStream {
42    type Item = std::io::Result<Bytes>;
43
44    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
45        let this = self.get_mut();
46        match this.rx.poll_recv(cx) {
47            Poll::Ready(item) => {
48                this.idle = None;
49                Poll::Ready(item)
50            }
51            Poll::Pending => {
52                let Some(timeout) = this.read_timeout else {
53                    return Poll::Pending;
54                };
55                let sleep = this
56                    .idle
57                    .get_or_insert_with(|| Box::pin(tokio::time::sleep(timeout)));
58                match sleep.as_mut().poll(cx) {
59                    Poll::Ready(()) => {
60                        this.idle = None;
61                        Poll::Ready(Some(Err(std::io::Error::new(
62                            std::io::ErrorKind::TimedOut,
63                            "read timeout: no response body chunk within the configured read_timeout",
64                        ))))
65                    }
66                    Poll::Pending => Poll::Pending,
67                }
68            }
69        }
70    }
71}
72
73impl std::fmt::Debug for BodyStream {
74    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
75        f.debug_struct("BodyStream").finish_non_exhaustive()
76    }
77}
leyline/core/body_stream.rs

leyline/core/
body_stream.rs
