timvw_rust_runtime/
body.rs

1use std::fmt;
2use std::pin::Pin;
3use std::sync::Arc;
4use std::task::{Context, Poll};
5
6use bytes::Bytes;
7use futures_core::TryStream;
8use futures_util::TryStreamExt;
9use http_body::{Frame, SizeHint};
10use http_body_util::{BodyExt, StreamBody};
11use sync_wrapper::SyncStream;
12
13use crate::{BoxError, MaybeSend, MaybeSync};
14
15/// A body of any type: `Send` and `Sync` natively.
16#[cfg(not(target_family = "wasm"))]
17type DynBody = http_body_util::combinators::BoxBody<Bytes, BoxError>;
18#[cfg(target_family = "wasm")]
19type DynBody = Pin<Box<dyn http_body::Body<Data = Bytes, Error = BoxError>>>;
20
21/// Makes a fresh body for each attempt: `Send` and `Sync` natively.
22#[cfg(not(target_family = "wasm"))]
23pub(crate) type MakeBody = Arc<dyn Fn() -> Body + Send + Sync>;
24#[cfg(target_family = "wasm")]
25pub(crate) type MakeBody = Arc<dyn Fn() -> Body>;
26
27fn boxed<B>(body: B) -> DynBody
28where
29    B: http_body::Body<Data = Bytes, Error = BoxError> + MaybeSend + MaybeSync + 'static,
30{
31    #[cfg(not(target_family = "wasm"))]
32    {
33        body.boxed()
34    }
35    #[cfg(target_family = "wasm")]
36    {
37        Box::pin(body)
38    }
39}
40
41/// The body of a request or a response: bytes held in memory, a stream, or a stream that can be made again.
42///
43/// A body held in memory can be cloned ([`Body::try_clone`]), which is what lets a retry layer send a request again.
44/// So can a body made from a replayable [`Upload`](crate::Upload), such as a file, which is opened again for each
45/// attempt. Any other stream can be read once.
46pub struct Body {
47    kind: Kind,
48}
49
50enum Kind {
51    Full(Option<Bytes>),
52    Stream {
53        body: DynBody,
54        /// The length, when known and the body cannot tell (a stream of chunks).
55        len: Option<u64>,
56    },
57    Replay(Replay),
58}
59
60/// A body made again for each attempt: `make` starts a fresh stream, `current` is the one being read.
61struct Replay {
62    make: MakeBody,
63    len: Option<u64>,
64    current: Option<Box<Body>>,
65}
66
67impl Body {
68    /// An empty body.
69    pub fn empty() -> Self {
70        Self {
71            kind: Kind::Full(None),
72        }
73    }
74
75    /// A body that wraps another [`http_body::Body`], such as the body of a hyper or reqwest response. It must be
76    /// `Send` and `Sync`, except on WebAssembly ([`MaybeSend`], [`MaybeSync`]).
77    pub fn wrap<B>(body: B) -> Self
78    where
79        B: http_body::Body + MaybeSend + MaybeSync + 'static,
80        B::Data: Into<Bytes>,
81        B::Error: Into<BoxError>,
82    {
83        Self {
84            kind: Kind::Stream {
85                body: boxed(
86                    body.map_frame(|frame| frame.map_data(Into::into))
87                        .map_err(Into::into),
88                ),
89                len: None,
90            },
91        }
92    }
93
94    /// A body that reads the chunks of a stream, such as a file read with `tokio_util::io::ReaderStream`. It must be
95    /// `Send`, except on WebAssembly ([`MaybeSend`]).
96    pub fn wrap_stream<S>(stream: S) -> Self
97    where
98        S: TryStream + MaybeSend + 'static,
99        S::Ok: Into<Bytes>,
100        S::Error: Into<BoxError>,
101    {
102        Self::sized_stream(stream, None)
103    }
104
105    /// A stream of chunks whose total length, when known, is `len` (sent as `Content-Length`).
106    pub(crate) fn sized_stream<S>(stream: S, len: Option<u64>) -> Self
107    where
108        S: TryStream + MaybeSend + 'static,
109        S::Ok: Into<Bytes>,
110        S::Error: Into<BoxError>,
111    {
112        let frames = SyncStream::new(
113            stream
114                .into_stream()
115                .map_ok(|chunk| Frame::data(chunk.into()))
116                .map_err(Into::into),
117        );
118        Self {
119            kind: Kind::Stream {
120                body: boxed(StreamBody::new(frames)),
121                len,
122            },
123        }
124    }
125
126    /// A body that `make` starts afresh for each attempt, so it can be cloned; `len` is the length of every body
127    /// `make` returns, when known.
128    pub(crate) fn replay(make: MakeBody, len: Option<u64>) -> Self {
129        Self {
130            kind: Kind::Replay(Replay {
131                make,
132                len,
133                current: None,
134            }),
135        }
136    }
137
138    /// A copy of the body: a body held in memory, or a fresh start of one that can be made again (a file
139    /// [`Upload`](crate::Upload), say); `None` for a stream that can be read once.
140    pub fn try_clone(&self) -> Option<Self> {
141        match &self.kind {
142            Kind::Full(bytes) => Some(Self {
143                kind: Kind::Full(bytes.clone()),
144            }),
145            Kind::Stream { .. } => None,
146            Kind::Replay(replay) => Some(Self::replay(replay.make.clone(), replay.len)),
147        }
148    }
149
150    /// The bytes of a body held in memory, or the body itself for a stream.
151    pub(crate) fn into_bytes(self) -> Result<Bytes, Self> {
152        match self.kind {
153            Kind::Full(bytes) => Ok(bytes.unwrap_or_default()),
154            kind => Err(Self { kind }),
155        }
156    }
157
158    /// The bytes of a body held in memory; `None` for a stream.
159    pub fn as_bytes(&self) -> Option<&[u8]> {
160        match &self.kind {
161            Kind::Full(bytes) => Some(bytes.as_deref().unwrap_or_default()),
162            Kind::Stream { .. } | Kind::Replay(_) => None,
163        }
164    }
165
166    /// Reads the whole body into memory.
167    pub async fn collect_bytes(self) -> Result<Bytes, BoxError> {
168        match self.kind {
169            Kind::Full(bytes) => Ok(bytes.unwrap_or_default()),
170            Kind::Stream { body, .. } => Ok(body.collect().await?.to_bytes()),
171            kind @ Kind::Replay(_) => Ok(BodyExt::collect(Self { kind }).await?.to_bytes()),
172        }
173    }
174}
175
176impl Default for Body {
177    fn default() -> Self {
178        Self::empty()
179    }
180}
181
182impl fmt::Debug for Body {
183    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
184        match &self.kind {
185            Kind::Full(bytes) => f
186                .debug_struct("Body")
187                .field("len", &bytes.as_ref().map_or(0, Bytes::len))
188                .finish(),
189            Kind::Stream { len, .. } => f
190                .debug_struct("Body")
191                .field("stream", &true)
192                .field("len", len)
193                .finish(),
194            Kind::Replay(replay) => f
195                .debug_struct("Body")
196                .field("replayable", &true)
197                .field("len", &replay.len)
198                .finish(),
199        }
200    }
201}
202
203impl http_body::Body for Body {
204    type Data = Bytes;
205    type Error = BoxError;
206
207    fn poll_frame(
208        self: Pin<&mut Self>,
209        cx: &mut Context<'_>,
210    ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
211        match &mut self.get_mut().kind {
212            Kind::Full(bytes) => Poll::Ready(
213                bytes
214                    .take()
215                    .filter(|b| !b.is_empty())
216                    .map(|b| Ok(Frame::data(b))),
217            ),
218            Kind::Stream { body, .. } => Pin::new(body).poll_frame(cx),
219            Kind::Replay(replay) => {
220                let current = replay
221                    .current
222                    .get_or_insert_with(|| Box::new((replay.make)()));
223                Pin::new(&mut **current).poll_frame(cx)
224            }
225        }
226    }
227
228    fn is_end_stream(&self) -> bool {
229        match &self.kind {
230            Kind::Full(bytes) => bytes.as_ref().is_none_or(Bytes::is_empty),
231            Kind::Stream { body, .. } => body.is_end_stream(),
232            Kind::Replay(replay) => match &replay.current {
233                Some(current) => current.is_end_stream(),
234                None => replay.len == Some(0),
235            },
236        }
237    }
238
239    fn size_hint(&self) -> SizeHint {
240        match &self.kind {
241            Kind::Full(bytes) => SizeHint::with_exact(bytes.as_ref().map_or(0, |b| b.len() as u64)),
242            Kind::Stream { body, len } => {
243                len.map_or_else(|| body.size_hint(), SizeHint::with_exact)
244            }
245            Kind::Replay(replay) => replay
246                .len
247                .map_or_else(SizeHint::default, SizeHint::with_exact),
248        }
249    }
250}
251
252impl From<Bytes> for Body {
253    fn from(bytes: Bytes) -> Self {
254        Self {
255            kind: Kind::Full(Some(bytes)),
256        }
257    }
258}
259
260impl From<Vec<u8>> for Body {
261    fn from(bytes: Vec<u8>) -> Self {
262        Bytes::from(bytes).into()
263    }
264}
265
266impl From<String> for Body {
267    fn from(text: String) -> Self {
268        Bytes::from(text).into()
269    }
270}
271
272impl From<&'static str> for Body {
273    fn from(text: &'static str) -> Self {
274        Bytes::from_static(text.as_bytes()).into()
275    }
276}
277
278impl From<&'static [u8]> for Body {
279    fn from(bytes: &'static [u8]) -> Self {
280        Bytes::from_static(bytes).into()
281    }
282}
283
284#[cfg(test)]
285mod tests {
286    use super::*;
287
288    #[tokio::test]
289    async fn bytes_clone_and_read() {
290        let body = Body::from("hello");
291        let copy = body.try_clone().expect("bytes can be cloned");
292        assert_eq!(body.collect_bytes().await.unwrap(), "hello");
293        assert_eq!(copy.as_bytes(), Some(&b"hello"[..]));
294    }
295
296    #[tokio::test]
297    async fn streams_are_read_once() {
298        let chunks = futures_util::stream::iter(vec![
299            Ok::<_, std::io::Error>(Bytes::from("a")),
300            Ok(Bytes::from("b")),
301        ]);
302        let body = Body::wrap_stream(chunks);
303        assert!(body.try_clone().is_none());
304        assert_eq!(body.collect_bytes().await.unwrap(), "ab");
305    }
306
307    #[test]
308    fn empty_is_end_of_stream() {
309        assert!(http_body::Body::is_end_stream(&Body::empty()));
310        assert_eq!(
311            http_body::Body::size_hint(&Body::from("abc")).exact(),
312            Some(3)
313        );
314    }
315}
timvw_rust_runtime/body.rs

timvw_rust_runtime/
body.rs
