oxirs_chat/memory_optimization/
pooling.rs

1//! Memory pooling for efficient buffer reuse
2
3use anyhow::{anyhow, Result};
4use std::collections::VecDeque;
5
6/// Memory pool for buffer reuse
7pub struct MemoryPool {
8    max_size: usize,
9    current_usage: usize,
10    buffers: VecDeque<Vec<u8>>,
11}
12
13impl MemoryPool {
14    pub fn new(max_size: usize) -> Self {
15        Self {
16            max_size,
17            current_usage: 0,
18            buffers: VecDeque::new(),
19        }
20    }
21
22    /// Allocate a buffer from the pool
23    pub fn allocate(&mut self, size: usize) -> Result<PooledBuffer> {
24        // Try to find a suitable buffer in the pool
25        if let Some(buffer) = self.find_suitable_buffer(size) {
26            return Ok(PooledBuffer::Pooled { data: buffer });
27        }
28
29        // Check if we have room to allocate
30        if self.current_usage + size > self.max_size {
31            return Err(anyhow!("Pool exhausted"));
32        }
33
34        // Allocate new buffer
35        let buffer = vec![0u8; size];
36        self.current_usage += size;
37
38        Ok(PooledBuffer::Pooled { data: buffer })
39    }
40
41    /// Return a buffer to the pool
42    pub fn deallocate(&mut self, buffer: Vec<u8>) {
43        let size = buffer.len();
44
45        // Only keep buffers if we have room
46        if self.buffers.len() < 100 {
47            // Keep up to 100 buffers
48            self.buffers.push_back(buffer);
49        } else {
50            self.current_usage = self.current_usage.saturating_sub(size);
51        }
52    }
53
54    /// Find a suitable buffer from the pool
55    fn find_suitable_buffer(&mut self, size: usize) -> Option<Vec<u8>> {
56        // Find buffer that's at least `size` bytes
57        for i in 0..self.buffers.len() {
58            if self.buffers[i].len() >= size {
59                return self.buffers.remove(i);
60            }
61        }
62        None
63    }
64
65    /// Get current pool usage
66    pub fn current_usage(&self) -> usize {
67        self.current_usage
68    }
69
70    /// Get available space
71    pub fn available(&self) -> usize {
72        self.max_size.saturating_sub(self.current_usage)
73    }
74}
75
76/// Pooled buffer with automatic return to pool
77pub enum PooledBuffer {
78    Pooled { data: Vec<u8> },
79    Heap { data: Vec<u8> },
80}
81
82impl PooledBuffer {
83    pub fn new_heap(size: usize) -> Result<Self> {
84        Ok(Self::Heap {
85            data: vec![0u8; size],
86        })
87    }
88
89    pub fn len(&self) -> usize {
90        match self {
91            Self::Pooled { data } => data.len(),
92            Self::Heap { data } => data.len(),
93        }
94    }
95
96    pub fn is_empty(&self) -> bool {
97        self.len() == 0
98    }
99
100    pub fn as_slice(&self) -> &[u8] {
101        match self {
102            Self::Pooled { data } => data,
103            Self::Heap { data } => data,
104        }
105    }
106
107    pub fn as_mut_slice(&mut self) -> &mut [u8] {
108        match self {
109            Self::Pooled { data } => data,
110            Self::Heap { data } => data,
111        }
112    }
113}
114
115#[cfg(test)]
116mod tests {
117    use super::*;
118
119    #[test]
120    fn test_pool_allocation() {
121        let mut pool = MemoryPool::new(10240);
122
123        let buffer = pool.allocate(1024).expect("should succeed");
124        assert_eq!(buffer.len(), 1024);
125        assert_eq!(pool.current_usage(), 1024);
126    }
127
128    #[test]
129    fn test_pool_reuse() {
130        let mut pool = MemoryPool::new(10240);
131
132        let buffer1 = pool.allocate(1024).expect("should succeed");
133        let data = match buffer1 {
134            PooledBuffer::Pooled { data } => data,
135            PooledBuffer::Heap { data } => data,
136        };
137
138        pool.deallocate(data);
139
140        let buffer2 = pool.allocate(1024).expect("should succeed");
141        assert_eq!(buffer2.len(), 1024);
142    }
143
144    #[test]
145    fn test_pool_exhaustion() {
146        let mut pool = MemoryPool::new(1024);
147
148        let _b1 = pool.allocate(512).expect("should succeed");
149        let _b2 = pool.allocate(512).expect("should succeed");
150
151        // Pool should be exhausted
152        let result = pool.allocate(1);
153        assert!(result.is_err());
154    }
155
156    #[test]
157    fn test_pool_available() {
158        let mut pool = MemoryPool::new(10240);
159
160        assert_eq!(pool.available(), 10240);
161
162        let _buffer = pool.allocate(1024).expect("should succeed");
163        assert_eq!(pool.available(), 9216);
164    }
165}
oxirs_chat/memory_optimization/pooling.rs

oxirs_chat/memory_optimization/
pooling.rs
