rama_http/layer/retry/
body.rs

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

rama_http/layer/retry/
body.rs
