iris_core/
memory_pool.rs

1//! Memory pool: reusable buffer allocation for GPU transfers and I/O.
2//!
3//! Provides `BufferPool` — a simple slab-style pool of `Vec<u8>` buffers
4//! that avoids repeated allocation churn. Used internally by iris-core and
5//! exposed for upper layers (iris-gpu, iris-js, etc.).
6
7use std::sync::Mutex;
8
9/// A pool of reusable byte buffers.
10///
11/// Each buffer starts at `chunk_size` bytes and grows as needed.
12/// When returned via `give_back`, the buffer is reset (length = 0)
13/// and made available for the next `take()` call.
14pub struct BufferPool {
15    /// Minimum size for each allocated buffer.
16    chunk_size: usize,
17    /// Maximum number of idle buffers to retain.
18    max_idle: usize,
19    /// Idle buffers ready for reuse.
20    idle: Mutex<Vec<Vec<u8>>>,
21}
22
23impl BufferPool {
24    /// Create a new buffer pool.
25    ///
26    /// * `chunk_size` — minimum capacity per buffer (default 4096).
27    /// * `max_idle`   — max number of buffers to keep in the pool (default 64).
28    pub fn new(chunk_size: usize, max_idle: usize) -> Self {
29        Self {
30            chunk_size,
31            max_idle,
32            idle: Mutex::new(Vec::with_capacity(max_idle)),
33        }
34    }
35
36    /// Take a buffer from the pool, or allocate a fresh one.
37    pub fn take(&self) -> Vec<u8> {
38        let mut idle = self.idle.lock().unwrap();
39        idle.pop().unwrap_or_else(|| Vec::with_capacity(self.chunk_size))
40    }
41
42    /// Take a buffer pre-sized to at least `min_capacity`.
43    pub fn take_sized(&self, min_capacity: usize) -> Vec<u8> {
44        let mut buf = self.take();
45        if buf.capacity() < min_capacity {
46            let additional = min_capacity - buf.len();
47            buf.reserve(additional);
48        }
49        buf
50    }
51
52    /// Return a buffer to the pool for reuse.
53    ///
54    /// The buffer is cleared (length set to 0) but its allocation is kept.
55    pub fn give_back(&self, mut buf: Vec<u8>) {
56        buf.clear();
57        let cap = buf.capacity();
58        // Only keep buffers that aren't excessively large (> 4× chunk_size).
59        if cap <= self.chunk_size * 4 {
60            let mut idle = self.idle.lock().unwrap();
61            if idle.len() < self.max_idle {
62                idle.push(buf);
63            }
64        }
65    }
66
67    /// Pre-allocate `count` buffers into the pool.
68    pub fn warm(&self, count: usize) {
69        let mut idle = self.idle.lock().unwrap();
70        let remaining = self.max_idle.saturating_sub(idle.len());
71        let to_add = count.min(remaining);
72        for _ in 0..to_add {
73            idle.push(Vec::with_capacity(self.chunk_size));
74        }
75    }
76
77    /// Number of buffers currently idle in the pool.
78    pub fn idle_count(&self) -> usize {
79        self.idle.lock().unwrap().len()
80    }
81}
82
83impl Default for BufferPool {
84    fn default() -> Self {
85        Self::new(4096, 64)
86    }
87}
88
89#[cfg(test)]
90mod tests {
91    use super::*;
92
93    #[test]
94    fn test_take_and_give_back() {
95        let pool = BufferPool::default();
96        let buf = pool.take();
97        assert!(buf.capacity() >= 4096);
98        assert!(buf.is_empty());
99        pool.give_back(buf);
100        assert_eq!(pool.idle_count(), 1);
101    }
102
103    #[test]
104    fn test_take_sized() {
105        let pool = BufferPool::default();
106        let buf = pool.take_sized(8192);
107        assert!(buf.capacity() >= 8192);
108    }
109
110    #[test]
111    fn test_warm() {
112        let pool = BufferPool::new(1024, 16);
113        pool.warm(8);
114        assert_eq!(pool.idle_count(), 8);
115    }
116
117    #[test]
118    fn test_max_idle_respected() {
119        let pool = BufferPool::new(64, 2);
120        let b1 = pool.take();
121        let b2 = pool.take();
122        let b3 = pool.take();
123        pool.give_back(b1);
124        pool.give_back(b2);
125        pool.give_back(b3);
126        assert_eq!(pool.idle_count(), 2);
127    }
128}
iris_core/memory_pool.rs

iris_core/
memory_pool.rs
