buffer_pool/
buffer.rs

1// Copyright (C) 2025, Cloudflare, Inc.
2// All rights reserved.
3//
4// Redistribution and use in source and binary forms, with or without
5// modification, are permitted provided that the following conditions are
6// met:
7//
8//     * Redistributions of source code must retain the above copyright notice,
9//       this list of conditions and the following disclaimer.
10//
11//     * Redistributions in binary form must reproduce the above copyright
12//       notice, this list of conditions and the following disclaimer in the
13//       documentation and/or other materials provided with the distribution.
14//
15// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
16// IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
17// THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
18// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
19// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
20// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
21// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
22// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
23// LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
24// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
25// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
26
27use std::ops::Deref;
28use std::ops::DerefMut;
29
30use crate::buffer_pool::consume_buffer_total_bytes;
31use crate::Reuse;
32
33/// A convinience wrapper around Vec that allows to "consume" data from the
34/// front *without* shifting.
35///
36/// This is not unlike `VecDeque` but more ergonomic
37/// for the operations we require. Conceptually `VecDeque` is two slices, and
38/// this is one slice. Also there is no `set_len` for `VecDeque`, so it has to
39/// be converted to `Vec` and then back again.
40#[derive(Default, Debug)]
41pub struct ConsumeBuffer {
42    inner: Vec<u8>,
43    head: usize,
44}
45
46impl Deref for ConsumeBuffer {
47    type Target = [u8];
48
49    fn deref(&self) -> &Self::Target {
50        &self.inner[self.head..]
51    }
52}
53
54impl DerefMut for ConsumeBuffer {
55    fn deref_mut(&mut self) -> &mut Self::Target {
56        &mut self.inner[self.head..]
57    }
58}
59
60impl Reuse for ConsumeBuffer {
61    fn reuse(&mut self, val: usize) -> bool {
62        let old_capacity = self.inner.capacity();
63        self.inner.clear();
64        self.inner.shrink_to(val);
65        self.update_metrics_after_resize(old_capacity);
66        self.head = 0;
67        self.inner.capacity() > 0
68    }
69
70    fn capacity(&self) -> usize {
71        self.inner.capacity()
72    }
73}
74
75impl Drop for ConsumeBuffer {
76    fn drop(&mut self) {
77        consume_buffer_total_bytes().dec_by(self.inner.capacity() as u64);
78    }
79}
80
81impl ConsumeBuffer {
82    pub fn from_vec(inner: Vec<u8>) -> Self {
83        consume_buffer_total_bytes().inc_by(inner.capacity() as u64);
84        ConsumeBuffer { inner, head: 0 }
85    }
86
87    pub fn into_vec(mut self) -> Vec<u8> {
88        // `ConsumeBuffer` implements Drop, so use `take` instead of moving the
89        // vector out directly. Update metrics manually because this object no
90        // longer owns the buffer.
91        let mut inner = std::mem::take(&mut self.inner);
92        consume_buffer_total_bytes().dec_by(inner.capacity() as u64);
93        inner.drain(0..self.head);
94        inner
95    }
96
97    pub fn pop_front(&mut self, count: usize) {
98        assert!(self.head + count <= self.inner.len());
99        self.head += count;
100    }
101
102    pub fn expand(&mut self, count: usize) {
103        let old_capacity = self.inner.capacity();
104        self.inner.reserve_exact(count);
105        self.update_metrics_after_resize(old_capacity);
106        // SAFETY: u8 is always initialized and we reserved the capacity.
107        unsafe { self.inner.set_len(count) };
108    }
109
110    pub fn truncate(&mut self, count: usize) {
111        self.inner.truncate(self.head + count);
112    }
113
114    pub fn add_prefix(&mut self, prefix: &[u8]) -> bool {
115        if self.head < prefix.len() {
116            return false;
117        }
118
119        self.head -= prefix.len();
120        self.inner[self.head..self.head + prefix.len()].copy_from_slice(prefix);
121
122        true
123    }
124
125    fn update_metrics_after_resize(&mut self, old_capacity: usize) {
126        let new_capacity = self.inner.capacity();
127        if new_capacity < old_capacity {
128            consume_buffer_total_bytes()
129                .dec_by((old_capacity - new_capacity) as u64);
130        } else if new_capacity > old_capacity {
131            consume_buffer_total_bytes()
132                .inc_by((new_capacity - old_capacity) as u64);
133        }
134    }
135}
136
137impl<'a> Extend<&'a u8> for ConsumeBuffer {
138    fn extend<T: IntoIterator<Item = &'a u8>>(&mut self, iter: T) {
139        let old_capacity = self.inner.capacity();
140        self.inner.extend(iter);
141        self.update_metrics_after_resize(old_capacity);
142    }
143}
144
145#[cfg(test)]
146mod tests {
147    use super::*;
148
149    #[test]
150    fn test_metrics() {
151        {
152            let mut buf = ConsumeBuffer::default();
153            assert_eq!(buf.inner.capacity(), 0);
154            assert_eq!(consume_buffer_total_bytes().get(), 0);
155
156            buf.extend(&[0, 1, 2, 3, 4]);
157            assert_eq!(
158                consume_buffer_total_bytes().get(),
159                buf.inner.capacity() as u64
160            );
161        }
162        // drop buf
163        assert_eq!(consume_buffer_total_bytes().get(), 0);
164
165        let mut buf_a = ConsumeBuffer::from_vec(vec![0, 1, 2, 3, 4]);
166        let buf_b = ConsumeBuffer::from_vec(vec![5, 6, 7]);
167        assert_eq!(
168            consume_buffer_total_bytes().get(),
169            (buf_a.inner.capacity() + buf_b.inner.capacity()) as u64
170        );
171
172        buf_a.expand(100000);
173        assert_eq!(
174            consume_buffer_total_bytes().get(),
175            (buf_a.inner.capacity() + buf_b.inner.capacity()) as u64
176        );
177
178        buf_b.into_vec();
179        assert_eq!(
180            consume_buffer_total_bytes().get(),
181            buf_a.inner.capacity() as u64
182        );
183
184        drop(buf_a);
185        assert_eq!(consume_buffer_total_bytes().get(), 0);
186    }
187}
buffer_pool/buffer.rs

buffer_pool/
buffer.rs
