base_http/body/
buffered.rs

1// Released under the MIT License.
2// Copyright, 2019-2026, by Samuel Williams.
3// Copyright, 2020, by Bryan Powell.
4// Copyright, 2025, by William T. Nelson.
5
6use super::{BodyFrame, ReadableBody};
7use base_scheduler::{Cancellation, Cancelled};
8
9/// An in-memory, replayable sequence of chunks.
10///
11/// Reads clone each chunk, following Socketry's buffered body. The application
12/// chooses its owner: a `Vec<u8>` yields independent data, whereas an immutable
13/// shared owner can clone cheaply. Closing preserves the chunks for rewinding.
14#[derive(Clone, Debug)]
15pub struct BufferedBody<D = Vec<u8>> {
16    chunks: Chunks<D>,
17    index: usize,
18    length: Option<u64>,
19}
20
21// Keep the common static response allocation-free without changing the public
22// slice view or the replay cursor. Multiple chunks retain their original order.
23#[derive(Clone, Debug)]
24enum Chunks<D> {
25    Empty,
26    One(D),
27    Many(Vec<D>),
28}
29
30impl<D> Chunks<D> {
31    fn as_slice(&self) -> &[D] {
32        match self {
33            Self::Empty => &[],
34            Self::One(chunk) => std::slice::from_ref(chunk),
35            Self::Many(chunks) => chunks,
36        }
37    }
38
39    fn push(&mut self, chunk: D) {
40        match self {
41            Self::Empty => *self = Self::One(chunk),
42            Self::Many(chunks) => chunks.push(chunk),
43            Self::One(_) => {
44                let Self::One(first) = std::mem::replace(self, Self::Empty) else {
45                    unreachable!();
46                };
47                *self = Self::Many(vec![first, chunk]);
48            }
49        }
50    }
51}
52
53impl<D: AsRef<[u8]> + Clone + Send> BufferedBody<D> {
54    pub fn new(chunks: impl IntoIterator<Item = D>) -> Self {
55        let mut iter = chunks.into_iter();
56        let chunks = match iter.next() {
57            None => Chunks::Empty,
58            Some(first) => match iter.next() {
59                None => Chunks::One(first),
60                Some(second) => {
61                    let mut chunks = vec![first, second];
62                    chunks.extend(iter);
63                    Chunks::Many(chunks)
64                }
65            },
66        };
67        let length = chunks.as_slice().iter().try_fold(0u64, |sum, chunk| {
68            u64::try_from(chunk.as_ref().len())
69                .ok()
70                .and_then(|length| sum.checked_add(length))
71        });
72        Self {
73            chunks,
74            index: 0,
75            length,
76        }
77    }
78
79    /// All original chunks, including data already read.
80    pub fn chunks(&self) -> &[D] {
81        self.chunks.as_slice()
82    }
83
84    /// Read immediately available memory without scheduling an operation.
85    pub fn read_chunk(&mut self) -> Option<D> {
86        let chunk = self.chunks.as_slice().get(self.index)?.clone();
87        self.index += 1;
88        Some(chunk)
89    }
90
91    /// A buffered representation with an independent cursor at the beginning.
92    pub fn buffered(&self) -> Self {
93        Self {
94            chunks: self.chunks.clone(),
95            index: 0,
96            length: self.length,
97        }
98    }
99
100    /// Append a chunk, including after earlier chunks have been read or closed.
101    pub fn write(&mut self, chunk: D) {
102        self.length = self
103            .length
104            .zip(u64::try_from(chunk.as_ref().len()).ok())
105            .and_then(|(length, added)| length.checked_add(added));
106        self.chunks.push(chunk);
107    }
108
109    pub fn clear(&mut self) {
110        self.chunks = Chunks::Empty;
111        self.index = 0;
112        self.length = Some(0);
113    }
114}
115
116impl<D> Default for BufferedBody<D> {
117    fn default() -> Self {
118        Self {
119            chunks: Chunks::Empty,
120            index: 0,
121            length: Some(0),
122        }
123    }
124}
125
126impl<D: AsRef<[u8]> + Clone + Send> From<D> for BufferedBody<D> {
127    fn from(chunk: D) -> Self {
128        Self::new([chunk])
129    }
130}
131
132impl<D: AsRef<[u8]> + Clone + Send> ReadableBody for BufferedBody<D> {
133    type Data = D;
134    type Error = Cancelled;
135
136    fn may_have_trailers(&self) -> bool {
137        false
138    }
139
140    fn is_ready(&self) -> bool {
141        true
142    }
143
144    fn is_empty(&self) -> bool {
145        self.index >= self.chunks.as_slice().len()
146    }
147
148    fn length(&self) -> Option<u64> {
149        self.length
150    }
151
152    fn is_rewindable(&self) -> bool {
153        true
154    }
155
156    fn rewind(&mut self) -> bool {
157        self.index = 0;
158        true
159    }
160
161    fn close(&mut self, _error: Option<Self::Error>) {
162        self.index = self.chunks.as_slice().len();
163    }
164
165    async fn read(
166        &mut self,
167        cancellation: &Cancellation,
168    ) -> Result<BodyFrame<Self::Data>, Self::Error> {
169        cancellation.check()?;
170        Ok(match self.read_chunk() {
171            Some(data) => BodyFrame::Data {
172                data,
173                end: self.is_empty(),
174            },
175            None => BodyFrame::End,
176        })
177    }
178}
base_http/body/buffered.rs

base_http/body/
buffered.rs
