rama_http/layer/retry/
body.rs

1use bytes::Bytes;
2use rama_http_types::dep::http_body;
3
4#[derive(Debug, Clone)]
5/// A body that can be clone and used for requests that have to be rertried.
6pub struct RetryBody {
7    bytes: Option<Bytes>,
8}
9
10impl RetryBody {
11    pub(crate) fn new(bytes: Bytes) -> Self {
12        RetryBody { bytes: Some(bytes) }
13    }
14
15    #[cfg(test)]
16    pub(crate) fn empty() -> Self {
17        RetryBody { bytes: None }
18    }
19
20    /// Turn this body into bytes.
21    pub fn into_bytes(self) -> Option<Bytes> {
22        self.bytes
23    }
24}
25
26impl crate::dep::http_body::Body for RetryBody {
27    type Data = Bytes;
28    type Error = rama_core::error::BoxError;
29
30    fn poll_frame(
31        mut self: std::pin::Pin<&mut Self>,
32        _cx: &mut std::task::Context<'_>,
33    ) -> std::task::Poll<Option<Result<http_body::Frame<Self::Data>, Self::Error>>> {
34        std::task::Poll::Ready(
35            self.bytes
36                .take()
37                .map(|bytes| Ok(http_body::Frame::data(bytes))),
38        )
39    }
40
41    fn is_end_stream(&self) -> bool {
42        self.bytes.is_none()
43    }
44
45    fn size_hint(&self) -> http_body::SizeHint {
46        http_body::SizeHint::with_exact(
47            self.bytes
48                .as_ref()
49                .map(|b| b.len() as u64)
50                .unwrap_or_default(),
51        )
52    }
53}
54
55impl From<RetryBody> for crate::Body {
56    fn from(body: RetryBody) -> Self {
57        match body.bytes {
58            Some(bytes) => bytes.into(),
59            None => crate::Body::empty(),
60        }
61    }
62}
63
64#[cfg(test)]
65mod tests {
66    use super::*;
67    use crate::BodyExtractExt;
68
69    #[tokio::test]
70    async fn consume_retry_body() {
71        let body = RetryBody::new(Bytes::from("hello"));
72        let s = body.try_into_string().await.unwrap();
73        assert_eq!(s, "hello");
74    }
75}
rama_http/layer/retry/body.rs

rama_http/layer/retry/
body.rs
