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    ///
38    /// The returned buffer has length 0 and capacity >= `chunk_size`.
39    pub fn take(&self) -> Vec<u8> {
40        let mut idle = self.idle.lock().unwrap();
41        idle.pop().unwrap_or_else(|| Vec::with_capacity(self.chunk_size))
42    }
43
44    /// Take a buffer pre-sized to at least `min_capacity`.
45    pub fn take_sized(&self, min_capacity: usize) -> Vec<u8> {
46        let mut buf = self.take();
47        if buf.capacity() < min_capacity {
48            let additional = min_capacity - buf.len();
49            buf.reserve(additional);
50        }
51        buf
52    }
53
54    /// Return a buffer to the pool for reuse.
55    ///
56    /// The buffer is cleared (length set to 0) but its allocation is kept.
57    pub fn give_back(&self, mut buf: Vec<u8>) {
58        buf.clear();
59        let cap = buf.capacity();
60        // Only keep buffers that aren't excessively large (> 4× chunk_size).
61        if cap <= self.chunk_size * 4 {
62            let mut idle = self.idle.lock().unwrap();
63            if idle.len() < self.max_idle {
64                idle.push(buf);
65            }
66        }
67    }
68
69    /// Pre-allocate `count` buffers into the pool.
70    pub fn warm(&self, count: usize) {
71        let mut idle = self.idle.lock().unwrap();
72        let remaining = self.max_idle.saturating_sub(idle.len());
73        let to_add = count.min(remaining);
74        for _ in 0..to_add {
75            idle.push(Vec::with_capacity(self.chunk_size));
76        }
77    }
78
79    /// Number of buffers currently idle in the pool.
80    pub fn idle_count(&self) -> usize {
81        self.idle.lock().unwrap().len()
82    }
83}
84
85impl Default for BufferPool {
86    fn default() -> Self {
87        Self::new(4096, 64)
88    }
89}
90
91#[cfg(test)]
92mod tests {
93    use super::*;
94
95    #[test]
96    fn test_take_and_give_back() {
97        let pool = BufferPool::default();
98        let buf = pool.take();
99        assert!(buf.capacity() >= 4096);
100        assert!(buf.is_empty());
101        pool.give_back(buf);
102        assert_eq!(pool.idle_count(), 1);
103    }
104
105    #[test]
106    fn test_take_sized() {
107        let pool = BufferPool::default();
108        let buf = pool.take_sized(8192);
109        assert!(buf.capacity() >= 8192);
110    }
111
112    #[test]
113    fn test_warm() {
114        let pool = BufferPool::new(1024, 16);
115        pool.warm(8);
116        assert_eq!(pool.idle_count(), 8);
117    }
118
119    #[test]
120    fn test_max_idle_respected() {
121        let pool = BufferPool::new(64, 2);
122        let b1 = pool.take();
123        let b2 = pool.take();
124        let b3 = pool.take();
125        pool.give_back(b1);
126        pool.give_back(b2);
127        pool.give_back(b3);
128        assert_eq!(pool.idle_count(), 2);
129    }
130}
iris_core/memory_pool.rs

iris_core/
memory_pool.rs
