Skip to main content

hexz_core/cache/
buffer_pool.rs

1//! Reusable buffer pool for decompression output buffers.
2//!
3//! Eliminates per-block `Vec<u8>` allocations during decompression by pooling
4//! and reusing buffers of common sizes. This is especially important for the
5//! parallel decompression path where N threads hitting the global allocator
6//! concurrently causes contention.
7//!
8//! # Design
9//!
10//! The pool is a simple `Mutex<Vec<Vec<u8>>>` stack. Buffers are checked out
11//! for decompression and returned after the decompressed data has been copied
12//! or converted to `Bytes`.
13//!
14//! When a buffer is needed for cache insertion (converting to `Bytes`), the
15//! buffer is consumed and not returned to the pool. The pool naturally reaches
16//! steady state as the block cache fills up and stops triggering new
17//! decompressions.
18
19use std::sync::Mutex;
20
21/// A pool of reusable `Vec<u8>` buffers for decompression.
22///
23/// Thread-safe via internal `Mutex`. The pool stores buffers sorted by capacity
24/// to enable efficient size-matched checkout.
25pub struct BufferPool {
26    /// Stack of available buffers, largest capacity last for O(1) pop.
27    buffers: Mutex<Vec<Vec<u8>>>,
28    /// Maximum number of buffers to retain in the pool.
29    max_buffers: usize,
30}
31
32impl BufferPool {
33    /// Creates a new buffer pool.
34    ///
35    /// # Parameters
36    ///
37    /// - `max_buffers`: Maximum number of idle buffers to retain.
38    ///   Excess buffers returned via `checkin` are dropped.
39    pub fn new(max_buffers: usize) -> Self {
40        Self {
41            buffers: Mutex::new(Vec::with_capacity(max_buffers)),
42            max_buffers,
43        }
44    }
45
46    /// Checks out a buffer with at least `capacity` bytes.
47    ///
48    /// Returns a pooled buffer if one of sufficient size is available,
49    /// otherwise allocates a new one. The returned buffer has length 0
50    /// but capacity >= `capacity`.
51    pub fn checkout(&self, capacity: usize) -> Vec<u8> {
52        if let Ok(mut pool) = self.buffers.lock() {
53            // Find the first buffer with sufficient capacity (linear scan is
54            // fine since max_buffers is small, typically 8-32).
55            if let Some(idx) = pool.iter().position(|b| b.capacity() >= capacity) {
56                let mut buf = pool.swap_remove(idx);
57                buf.clear();
58                return buf;
59            }
60        }
61        Vec::with_capacity(capacity)
62    }
63
64    /// Returns a buffer to the pool for future reuse.
65    ///
66    /// If the pool is full, the buffer is dropped. Buffers are only worth
67    /// pooling if they have meaningful capacity (the caller should not return
68    /// tiny buffers).
69    pub fn checkin(&self, buf: Vec<u8>) {
70        if buf.capacity() == 0 {
71            return;
72        }
73        if let Ok(mut pool) = self.buffers.lock() {
74            if pool.len() < self.max_buffers {
75                pool.push(buf);
76            }
77            // else: drop buf (pool is full)
78        }
79        // else: lock poisoned, just drop the buffer
80    }
81}
82
83#[cfg(test)]
84mod tests {
85    use super::*;
86
87    #[test]
88    fn checkout_returns_sufficient_capacity() {
89        let pool = BufferPool::new(4);
90        let buf = pool.checkout(1024);
91        assert!(buf.capacity() >= 1024);
92        assert_eq!(buf.len(), 0);
93    }
94
95    #[test]
96    fn checkin_and_reuse() {
97        let pool = BufferPool::new(4);
98
99        // Allocate and return a buffer
100        let mut buf = pool.checkout(1024);
101        buf.extend_from_slice(&[42u8; 512]);
102        let cap = buf.capacity();
103        pool.checkin(buf);
104
105        // Should get the same buffer back (same capacity, cleared)
106        let buf2 = pool.checkout(1024);
107        assert_eq!(buf2.capacity(), cap);
108        assert_eq!(buf2.len(), 0);
109    }
110
111    #[test]
112    fn respects_max_buffers() {
113        let pool = BufferPool::new(2);
114
115        pool.checkin(Vec::with_capacity(1024));
116        pool.checkin(Vec::with_capacity(1024));
117        pool.checkin(Vec::with_capacity(1024)); // should be dropped
118
119        let guard = pool.buffers.lock().unwrap();
120        assert_eq!(guard.len(), 2);
121    }
122
123    #[test]
124    fn checkout_allocates_when_empty() {
125        let pool = BufferPool::new(4);
126        let buf = pool.checkout(65536);
127        assert!(buf.capacity() >= 65536);
128    }
129
130    #[test]
131    fn checkout_skips_too_small_buffers() {
132        let pool = BufferPool::new(4);
133        pool.checkin(Vec::with_capacity(512));
134
135        // Request larger than available
136        let buf = pool.checkout(1024);
137        assert!(buf.capacity() >= 1024);
138
139        // The small buffer should still be in the pool
140        let guard = pool.buffers.lock().unwrap();
141        assert_eq!(guard.len(), 1);
142    }
143}