fluke_buffet/
piece.rs

1//! Types for performing vectored I/O.
2
3use http::header::HeaderName;
4use std::{collections::VecDeque, fmt, ops::Deref, rc::Rc, str::Utf8Error};
5
6use crate::{Roll, RollStr};
7
8/// A piece of data (arbitrary bytes) with a stable address, suitable for
9/// passing to the kernel (io_uring writes).
10#[derive(Clone)]
11pub enum Piece {
12    Full {
13        core: PieceCore,
14    },
15    Slice {
16        core: PieceCore,
17        start: usize,
18        len: usize,
19    },
20}
21
22impl Piece {
23    /// Returns an empty piece
24    pub fn empty() -> Self {
25        Self::Full {
26            core: PieceCore::Static(&[]),
27        }
28    }
29}
30
31#[derive(Clone)]
32pub enum PieceCore {
33    Static(&'static [u8]),
34    Vec(Rc<Vec<u8>>),
35    Roll(Roll),
36    HeaderName(HeaderName),
37}
38
39impl<T> From<T> for Piece
40where
41    T: Into<PieceCore>,
42{
43    #[inline(always)]
44    fn from(t: T) -> Self {
45        Piece::Full { core: t.into() }
46    }
47}
48
49impl From<&'static [u8]> for PieceCore {
50    #[inline(always)]
51    fn from(slice: &'static [u8]) -> Self {
52        PieceCore::Static(slice)
53    }
54}
55
56impl From<&'static str> for PieceCore {
57    #[inline(always)]
58    fn from(slice: &'static str) -> Self {
59        PieceCore::Static(slice.as_bytes())
60    }
61}
62
63impl From<Vec<u8>> for PieceCore {
64    #[inline(always)]
65    fn from(vec: Vec<u8>) -> Self {
66        PieceCore::Vec(Rc::new(vec))
67    }
68}
69
70impl From<Roll> for PieceCore {
71    #[inline(always)]
72    fn from(roll: Roll) -> Self {
73        PieceCore::Roll(roll)
74    }
75}
76
77impl From<()> for PieceCore {
78    #[inline(always)]
79    fn from(_empty: ()) -> Self {
80        PieceCore::Static(&[])
81    }
82}
83
84impl From<PieceStr> for Piece {
85    fn from(s: PieceStr) -> Self {
86        s.piece
87    }
88}
89
90impl From<HeaderName> for PieceCore {
91    #[inline(always)]
92    fn from(name: HeaderName) -> Self {
93        PieceCore::HeaderName(name)
94    }
95}
96
97impl Deref for PieceCore {
98    type Target = [u8];
99
100    fn deref(&self) -> &Self::Target {
101        self.as_ref()
102    }
103}
104
105impl Deref for Piece {
106    type Target = [u8];
107
108    fn deref(&self) -> &Self::Target {
109        self.as_ref()
110    }
111}
112
113impl AsRef<[u8]> for PieceCore {
114    fn as_ref(&self) -> &[u8] {
115        match self {
116            PieceCore::Static(slice) => slice,
117            PieceCore::Vec(vec) => vec.as_ref(),
118            PieceCore::Roll(roll) => roll.as_ref(),
119            PieceCore::HeaderName(name) => name.as_str().as_bytes(),
120        }
121    }
122}
123
124impl Piece {
125    fn core(&self) -> &PieceCore {
126        match self {
127            Piece::Full { core } => core,
128            Piece::Slice { core, .. } => core,
129        }
130    }
131
132    /// Split the piece into two at the given index.
133    /// The original piece will be consumed.
134    /// Returns a tuple of the two pieces.
135    pub fn split_at(self, middle: usize) -> (Self, Self) {
136        let len = self.len();
137        assert!(middle <= len);
138
139        match self {
140            Piece::Full { core } => (
141                Self::Slice {
142                    core: core.clone(),
143                    start: 0,
144                    len: middle,
145                },
146                Self::Slice {
147                    core,
148                    start: middle,
149                    len: len - middle,
150                },
151            ),
152            Piece::Slice { core, start, len } => (
153                Self::Slice {
154                    core: core.clone(),
155                    start,
156                    len: middle,
157                },
158                Self::Slice {
159                    core,
160                    start: start + middle,
161                    len: len - middle,
162                },
163            ),
164        }
165    }
166}
167
168impl AsRef<[u8]> for Piece {
169    fn as_ref(&self) -> &[u8] {
170        let ptr = self.core().as_ref();
171        if let Piece::Slice { start, len, .. } = self {
172            &ptr[*start..][..*len]
173        } else {
174            ptr
175        }
176    }
177}
178
179impl Piece {
180    // Decode as utf-8 (owned)
181    pub fn to_str(self) -> Result<PieceStr, Utf8Error> {
182        _ = std::str::from_utf8(&self[..])?;
183        Ok(PieceStr { piece: self })
184    }
185
186    /// Convert to [PieceStr].
187    ///
188    /// # Safety
189    /// UB if not utf-8. Typically only used in parsers.
190    pub unsafe fn to_string_unchecked(self) -> PieceStr {
191        PieceStr { piece: self }
192    }
193}
194
195impl Piece {
196    #[inline(always)]
197    pub fn len(&self) -> usize {
198        self.as_ref().len()
199    }
200
201    #[inline(always)]
202    pub fn is_empty(&self) -> bool {
203        self.len() == 0
204    }
205}
206
207/// A list of [Piece], suitable for issuing vectored writes via io_uring.
208#[derive(Default)]
209pub struct PieceList {
210    // note: we can't use smallvec here, because the address of
211    // the piece list must be stable for the kernel to take
212    // ownership of it.
213    //
214    // we could however do our own memory pooling.
215    pub(crate) pieces: VecDeque<Piece>,
216}
217
218impl PieceList {
219    /// Create a new piece list with a single chunk
220    pub fn single(piece: impl Into<Piece>) -> Self {
221        Self {
222            pieces: [piece.into()].into(),
223        }
224    }
225
226    /// Add a single chunk to the back of the list
227    pub fn push_back(&mut self, chunk: impl Into<Piece>) {
228        let chunk = chunk.into();
229        if !chunk.is_empty() {
230            self.pieces.push_back(chunk);
231        }
232    }
233
234    /// Add a single chunk to the front of the list
235    pub fn push_front(&mut self, chunk: impl Into<Piece>) {
236        let chunk = chunk.into();
237        if !chunk.is_empty() {
238            self.pieces.push_front(chunk);
239        }
240    }
241
242    /// Add a single chunk to the back list and return self
243    pub fn followed_by(mut self, chunk: impl Into<Piece>) -> Self {
244        self.push_back(chunk);
245        self
246    }
247
248    /// Add a single chunk to the front of the list and return self
249    pub fn preceded_by(mut self, chunk: impl Into<Piece>) -> Self {
250        self.push_front(chunk);
251        self
252    }
253
254    /// Returns total length
255    pub fn len(&self) -> usize {
256        self.pieces.iter().map(|c| c.len()).sum()
257    }
258
259    pub fn num_pieces(&self) -> usize {
260        self.pieces.len()
261    }
262
263    pub fn is_empty(&self) -> bool {
264        self.pieces.is_empty() || self.len() == 0
265    }
266
267    pub fn clear(&mut self) {
268        self.pieces.clear();
269    }
270
271    pub fn into_vec_deque(self) -> VecDeque<Piece> {
272        self.pieces
273    }
274}
275
276impl From<VecDeque<Piece>> for PieceList {
277    fn from(pieces: VecDeque<Piece>) -> Self {
278        Self { pieces }
279    }
280}
281impl From<PieceList> for VecDeque<Piece> {
282    fn from(list: PieceList) -> Self {
283        list.pieces
284    }
285}
286
287/// A piece of data with a stable address that's _also_
288/// valid utf-8.
289#[derive(Clone)]
290pub struct PieceStr {
291    piece: Piece,
292}
293
294impl fmt::Debug for PieceStr {
295    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> std::fmt::Result {
296        fmt::Debug::fmt(&self[..], f)
297    }
298}
299
300impl fmt::Display for PieceStr {
301    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> std::fmt::Result {
302        f.pad(self)
303    }
304}
305
306impl Deref for PieceStr {
307    type Target = str;
308
309    fn deref(&self) -> &Self::Target {
310        unsafe { std::str::from_utf8_unchecked(&self.piece) }
311    }
312}
313
314impl AsRef<str> for PieceStr {
315    fn as_ref(&self) -> &str {
316        self
317    }
318}
319
320impl PieceStr {
321    /// Returns the underlying bytes (borrowed)
322    pub fn as_bytes(&self) -> &[u8] {
323        self.piece.as_ref()
324    }
325
326    /// Returns the underlying bytes (owned)
327    pub fn into_inner(self) -> Piece {
328        self.piece
329    }
330}
331
332impl From<&'static str> for PieceStr {
333    fn from(s: &'static str) -> Self {
334        PieceStr {
335            piece: PieceCore::Static(s.as_bytes()).into(),
336        }
337    }
338}
339
340impl From<String> for PieceStr {
341    fn from(s: String) -> Self {
342        PieceStr {
343            piece: PieceCore::Vec(Rc::new(s.into_bytes())).into(),
344        }
345    }
346}
347
348impl From<RollStr> for PieceStr {
349    fn from(s: RollStr) -> Self {
350        PieceStr {
351            piece: PieceCore::Roll(s.into_inner()).into(),
352        }
353    }
354}
355
356#[cfg(test)]
357mod tests {
358    use crate::{Piece, PieceCore};
359
360    #[test]
361    fn test_slice() {
362        // test that slicing works correctly for a
363        // piece made from a &'static u8
364        let piece: Piece = PieceCore::Static("französisch".as_bytes()).into();
365        // split so that "l" is "franz"
366        let (first_name, last_name) = piece.split_at(5);
367        assert_eq!(&first_name[..], "franz".as_bytes());
368        assert_eq!(&last_name[..], "ösisch".as_bytes());
369
370        // test edge cases, zero-length left
371        let piece: Piece = PieceCore::Static("französisch".as_bytes()).into();
372        let (first_name, last_name) = piece.split_at(0);
373        assert_eq!(&first_name[..], "".as_bytes());
374        assert_eq!(&last_name[..], "französisch".as_bytes());
375
376        // test edge cases, zero-length right
377        let piece: Piece = PieceCore::Static("französisch".as_bytes()).into();
378        let (first_name, last_name) = piece.split_at(12);
379        assert_eq!(&first_name[..], "französisch".as_bytes());
380        assert_eq!(&last_name[..], "".as_bytes());
381
382        // edge case: empty piece being split into two
383        let piece: Piece = PieceCore::Static(b"").into();
384        let (first_name, last_name) = piece.split_at(0);
385        assert_eq!(&first_name[..], "".as_bytes());
386        assert_eq!(&last_name[..], "".as_bytes());
387    }
388}
fluke_buffet/piece.rs

fluke_buffet/
piece.rs
