mcp_postgres/
buffers.rs

1use std::sync::Arc;
2use parking_lot::Mutex;
3
4/// Per-connection result buffer pool for reusing allocated buffers.
5/// Reduces allocation overhead during query result processing.
6pub struct ResultBuffer {
7    data: Vec<u8>,
8}
9
10impl ResultBuffer {
11    const DEFAULT_CAPACITY: usize = 4096;
12
13    pub fn new() -> Self {
14        Self {
15            data: Vec::with_capacity(Self::DEFAULT_CAPACITY),
16        }
17    }
18
19    pub fn as_mut(&mut self) -> &mut Vec<u8> {
20        &mut self.data
21    }
22
23    pub fn as_slice(&self) -> &[u8] {
24        &self.data
25    }
26
27    pub fn clear(&mut self) {
28        self.data.clear();
29    }
30
31    pub fn len(&self) -> usize {
32        self.data.len()
33    }
34
35    pub fn is_empty(&self) -> bool {
36        self.data.is_empty()
37    }
38}
39
40/// Thread-local result buffer for reducing allocations
41/// Each tokio task gets its own buffer to avoid contention
42pub struct BufferPool {
43    buffers: Arc<Mutex<Vec<ResultBuffer>>>,
44    max_cached: usize,
45}
46
47impl BufferPool {
48    pub fn new(max_cached: usize) -> Self {
49        Self {
50            buffers: Arc::new(Mutex::new(Vec::with_capacity(max_cached))),
51            max_cached,
52        }
53    }
54
55    pub fn acquire(&self) -> ResultBuffer {
56        let mut buffers = self.buffers.lock();
57        buffers.pop().unwrap_or_else(ResultBuffer::new)
58    }
59
60    pub fn release(&self, mut buffer: ResultBuffer) {
61        buffer.clear();
62        let mut buffers = self.buffers.lock();
63        if buffers.len() < self.max_cached {
64            buffers.push(buffer);
65        }
66    }
67
68    pub fn size(&self) -> usize {
69        self.buffers.lock().len()
70    }
71}
72
73#[cfg(test)]
74mod tests {
75    use super::*;
76
77    #[test]
78    fn test_result_buffer_creation() {
79        let buf = ResultBuffer::new();
80        assert!(buf.is_empty());
81        assert!(buf.len() == 0);
82    }
83
84    #[test]
85    fn test_buffer_pool_acquire_release() {
86        let pool = BufferPool::new(2);
87        let buf = pool.acquire();
88        assert_eq!(pool.size(), 0);
89        pool.release(buf);
90        assert_eq!(pool.size(), 1);
91    }
92
93    #[test]
94    fn test_buffer_pool_respects_max_cached() {
95        let pool = BufferPool::new(1);
96        let buf1 = pool.acquire();
97        let buf2 = pool.acquire();
98        pool.release(buf1);
99        pool.release(buf2);
100        // Only 1 should be cached
101        assert_eq!(pool.size(), 1);
102    }
103}
mcp_postgres/buffers.rs

mcp_postgres/
buffers.rs
