feox_server/io/
buffer_pool.rs

1use bytes::{Bytes, BytesMut};
2use parking_lot::Mutex;
3use std::collections::VecDeque;
4use std::sync::Arc;
5
6/// Thread-safe buffer pool
7pub struct BufferPool {
8    /// Pre-allocated buffers
9    buffers: Vec<BufferSlot>,
10
11    /// Free buffer indices
12    free_list: Arc<Mutex<VecDeque<usize>>>,
13
14    /// Buffer size
15    buffer_size: usize,
16}
17
18struct BufferSlot {
19    data: Vec<u8>,
20    in_use: Arc<Mutex<bool>>,
21}
22
23impl BufferPool {
24    pub fn new(num_buffers: usize, buffer_size: usize) -> Self {
25        let mut buffers = Vec::with_capacity(num_buffers);
26        let mut free_list = VecDeque::with_capacity(num_buffers);
27
28        for i in 0..num_buffers {
29            // Pre-allocate buffer
30            let data = vec![0u8; buffer_size];
31
32            buffers.push(BufferSlot {
33                data,
34                in_use: Arc::new(Mutex::new(false)),
35            });
36
37            free_list.push_back(i);
38        }
39
40        Self {
41            buffers,
42            free_list: Arc::new(Mutex::new(free_list)),
43            buffer_size,
44        }
45    }
46
47    /// Get a buffer for reading
48    pub fn get_buffer(&self) -> Option<BufferHandle> {
49        let mut free_list = self.free_list.lock();
50
51        if let Some(index) = free_list.pop_front() {
52            let slot = &self.buffers[index];
53            *slot.in_use.lock() = true;
54
55            Some(BufferHandle {
56                index,
57                data: slot.data.as_ptr(),
58                len: self.buffer_size,
59                pool: self as *const BufferPool,
60            })
61        } else {
62            None
63        }
64    }
65
66    /// Return a buffer to the pool
67    fn return_buffer(&self, index: usize) {
68        if index < self.buffers.len() {
69            let slot = &self.buffers[index];
70            *slot.in_use.lock() = false;
71            self.free_list.lock().push_back(index);
72        }
73    }
74
75    /// Get a BytesMut buffer for writing
76    pub fn get_buffer_mut(&self) -> Option<BufferHandleMut> {
77        let mut free_list = self.free_list.lock();
78
79        if let Some(index) = free_list.pop_front() {
80            let slot = &self.buffers[index];
81            *slot.in_use.lock() = true;
82
83            Some(BufferHandleMut {
84                index,
85                buffer: BytesMut::from(&slot.data[..]),
86                pool: self as *const BufferPool,
87            })
88        } else {
89            None
90        }
91    }
92
93    /// Create a Bytes view into a buffer
94    pub fn buffer_to_bytes(&self, index: usize, len: usize) -> Option<Bytes> {
95        if index >= self.buffers.len() || len > self.buffer_size {
96            return None;
97        }
98
99        let data = &self.buffers[index].data[..len];
100        Some(Bytes::copy_from_slice(data))
101    }
102}
103
104/// Handle to a buffer from the pool
105pub struct BufferHandle {
106    index: usize,
107    data: *const u8,
108    len: usize,
109    pool: *const BufferPool,
110}
111
112impl BufferHandle {
113    pub fn index(&self) -> usize {
114        self.index
115    }
116
117    pub fn as_slice(&self) -> &[u8] {
118        unsafe { std::slice::from_raw_parts(self.data, self.len) }
119    }
120
121    pub fn as_mut_slice(&mut self) -> &mut [u8] {
122        unsafe { std::slice::from_raw_parts_mut(self.data as *mut u8, self.len) }
123    }
124}
125
126impl Drop for BufferHandle {
127    fn drop(&mut self) {
128        unsafe {
129            let pool = &*self.pool;
130            pool.return_buffer(self.index);
131        }
132    }
133}
134
135/// Mutable handle to a buffer
136pub struct BufferHandleMut {
137    index: usize,
138    buffer: BytesMut,
139    pool: *const BufferPool,
140}
141
142impl BufferHandleMut {
143    pub fn buffer_mut(&mut self) -> &mut BytesMut {
144        &mut self.buffer
145    }
146
147    pub fn freeze(mut self) -> Bytes {
148        self.buffer.split().freeze()
149    }
150}
151
152impl Drop for BufferHandleMut {
153    fn drop(&mut self) {
154        unsafe {
155            let pool = &*self.pool;
156            pool.return_buffer(self.index);
157        }
158    }
159}
160
161// Make handles Send + Sync
162unsafe impl Send for BufferHandle {}
163unsafe impl Sync for BufferHandle {}
164unsafe impl Send for BufferHandleMut {}
165unsafe impl Sync for BufferHandleMut {}
feox_server/io/buffer_pool.rs

feox_server/io/
buffer_pool.rs
