Skip to main content

lnc_network/
connection.rs

1use bytes::{Bytes, BytesMut};
2
3const DEFAULT_BUFFER_SIZE: usize = 64 * 1024;
4
5pub struct ConnectionBuffer {
6    buffer: BytesMut,
7    read_position: usize,
8}
9
10impl ConnectionBuffer {
11    pub fn new() -> Self {
12        Self::with_capacity(DEFAULT_BUFFER_SIZE)
13    }
14
15    pub fn with_capacity(capacity: usize) -> Self {
16        Self {
17            buffer: BytesMut::with_capacity(capacity),
18            read_position: 0,
19        }
20    }
21
22    pub fn extend(&mut self, data: &[u8]) {
23        self.buffer.extend_from_slice(data);
24    }
25
26    pub fn slice(&self, start: usize, len: usize) -> Bytes {
27        self.buffer.clone().freeze().slice(start..start + len)
28    }
29
30    pub fn consume(&mut self, len: usize) {
31        if len >= self.buffer.len() {
32            self.buffer.clear();
33            self.read_position = 0;
34        } else {
35            let _ = self.buffer.split_to(len);
36        }
37    }
38
39    /// Compact the buffer to reduce memory usage when oversized.
40    /// Uses split_off pattern to avoid to_vec() allocation per standards ยง18.1.
41    pub fn compact(&mut self) {
42        if self.buffer.capacity() > DEFAULT_BUFFER_SIZE * 4
43            && self.buffer.len() < DEFAULT_BUFFER_SIZE
44        {
45            // Zero-copy approach: create new buffer and copy existing data
46            // This is unavoidable for compaction but uses extend_from_slice
47            // which is more efficient than to_vec() + extend_from_slice
48            let len = self.buffer.len();
49            let mut new_buffer = BytesMut::with_capacity(DEFAULT_BUFFER_SIZE);
50            new_buffer.extend_from_slice(&self.buffer[..len]);
51            self.buffer = new_buffer;
52        }
53    }
54
55    #[inline]
56    #[must_use]
57    pub fn as_slice(&self) -> &[u8] {
58        &self.buffer[..]
59    }
60
61    #[inline]
62    #[must_use]
63    pub fn len(&self) -> usize {
64        self.buffer.len()
65    }
66
67    #[inline]
68    #[must_use]
69    pub fn is_empty(&self) -> bool {
70        self.buffer.is_empty()
71    }
72
73    #[inline]
74    #[must_use]
75    pub fn capacity(&self) -> usize {
76        self.buffer.capacity()
77    }
78
79    #[inline]
80    #[must_use]
81    pub fn remaining(&self) -> usize {
82        self.buffer.len().saturating_sub(self.read_position)
83    }
84
85    pub fn reserve(&mut self, additional: usize) {
86        self.buffer.reserve(additional);
87    }
88
89    pub fn clear(&mut self) {
90        self.buffer.clear();
91        self.read_position = 0;
92    }
93
94    pub fn get_write_buffer(&mut self, min_size: usize) -> &mut [u8] {
95        if self.buffer.spare_capacity_mut().len() < min_size {
96            self.buffer.reserve(min_size);
97        }
98
99        let len = self.buffer.len();
100        let cap = self.buffer.capacity();
101
102        // SAFETY: We're getting a mutable slice to the spare capacity
103        // which will be initialized by the caller before the length is increased
104        unsafe { std::slice::from_raw_parts_mut(self.buffer.as_mut_ptr().add(len), cap - len) }
105    }
106
107    /// Advance the write position of the buffer.
108    ///
109    /// # Safety
110    /// The caller must ensure that `len` bytes have been written to the
111    /// spare capacity returned by `spare_capacity_mut()`.
112    pub unsafe fn advance_write(&mut self, len: usize) {
113        let new_len = self.buffer.len() + len;
114        // SAFETY: Caller guarantees that len bytes have been written to the spare capacity
115        unsafe {
116            self.buffer.set_len(new_len);
117        }
118    }
119}
120
121impl Default for ConnectionBuffer {
122    fn default() -> Self {
123        Self::new()
124    }
125}
126
127#[cfg(test)]
128#[allow(clippy::unwrap_used)]
129mod tests {
130    use super::*;
131
132    #[test]
133    fn test_connection_buffer_extend() {
134        let mut buf = ConnectionBuffer::new();
135        buf.extend(b"hello");
136        buf.extend(b" world");
137
138        assert_eq!(buf.as_slice(), b"hello world");
139        assert_eq!(buf.len(), 11);
140    }
141
142    #[test]
143    fn test_connection_buffer_slice() {
144        let mut buf = ConnectionBuffer::new();
145        buf.extend(b"hello world");
146
147        let slice = buf.slice(0, 5);
148        assert_eq!(&slice[..], b"hello");
149
150        let slice2 = buf.slice(6, 5);
151        assert_eq!(&slice2[..], b"world");
152    }
153
154    #[test]
155    fn test_connection_buffer_consume() {
156        let mut buf = ConnectionBuffer::new();
157        buf.extend(b"hello world");
158
159        buf.consume(6);
160        assert_eq!(buf.as_slice(), b"world");
161        assert_eq!(buf.len(), 5);
162    }
163}