lnc_network/
connection.rs1use 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 pub fn compact(&mut self) {
42 if self.buffer.capacity() > DEFAULT_BUFFER_SIZE * 4
43 && self.buffer.len() < DEFAULT_BUFFER_SIZE
44 {
45 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 unsafe { std::slice::from_raw_parts_mut(self.buffer.as_mut_ptr().add(len), cap - len) }
105 }
106
107 pub unsafe fn advance_write(&mut self, len: usize) {
113 let new_len = self.buffer.len() + len;
114 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}