Skip to main content

oxihuman_core/
buffer_pool.rs

1// Copyright (C) 2026 COOLJAPAN OU (Team KitaSan)
2// SPDX-License-Identifier: Apache-2.0
3#![allow(dead_code)]
4
5//! Pool of reusable byte buffers to reduce allocation pressure.
6
7/// A pool of byte buffers that can be checked out and returned.
8#[allow(dead_code)]
9#[derive(Debug, Clone)]
10pub struct BufferPool {
11    pool: Vec<Vec<u8>>,
12    buffer_size: usize,
13    max_pool_size: usize,
14    checkout_count: u64,
15}
16
17#[allow(dead_code)]
18impl BufferPool {
19    pub fn new(buffer_size: usize, max_pool_size: usize) -> Self {
20        Self {
21            pool: Vec::new(),
22            buffer_size,
23            max_pool_size,
24            checkout_count: 0,
25        }
26    }
27
28    pub fn checkout(&mut self) -> Vec<u8> {
29        self.checkout_count += 1;
30        self.pool
31            .pop()
32            .unwrap_or_else(|| vec![0u8; self.buffer_size])
33    }
34
35    pub fn return_buffer(&mut self, mut buf: Vec<u8>) {
36        if self.pool.len() < self.max_pool_size {
37            buf.iter_mut().for_each(|b| *b = 0);
38            self.pool.push(buf);
39        }
40    }
41
42    pub fn available(&self) -> usize {
43        self.pool.len()
44    }
45
46    pub fn buffer_size(&self) -> usize {
47        self.buffer_size
48    }
49
50    pub fn max_pool_size(&self) -> usize {
51        self.max_pool_size
52    }
53
54    pub fn total_checkouts(&self) -> u64 {
55        self.checkout_count
56    }
57
58    pub fn preallocate(&mut self, count: usize) {
59        let to_add = count.min(self.max_pool_size.saturating_sub(self.pool.len()));
60        for _ in 0..to_add {
61            self.pool.push(vec![0u8; self.buffer_size]);
62        }
63    }
64
65    pub fn shrink_to(&mut self, target: usize) {
66        while self.pool.len() > target {
67            self.pool.pop();
68        }
69    }
70
71    pub fn clear(&mut self) {
72        self.pool.clear();
73    }
74
75    pub fn total_memory(&self) -> usize {
76        self.pool.len() * self.buffer_size
77    }
78}
79
80#[cfg(test)]
81mod tests {
82    use super::*;
83
84    #[test]
85    fn new_pool_is_empty() {
86        let p = BufferPool::new(1024, 8);
87        assert_eq!(p.available(), 0);
88    }
89
90    #[test]
91    fn checkout_creates_buffer() {
92        let mut p = BufferPool::new(64, 4);
93        let buf = p.checkout();
94        assert_eq!(buf.len(), 64);
95    }
96
97    #[test]
98    fn return_and_reuse() {
99        let mut p = BufferPool::new(32, 4);
100        let buf = p.checkout();
101        p.return_buffer(buf);
102        assert_eq!(p.available(), 1);
103        let buf2 = p.checkout();
104        assert_eq!(buf2.len(), 32);
105        assert_eq!(p.available(), 0);
106    }
107
108    #[test]
109    fn return_beyond_max_drops() {
110        let mut p = BufferPool::new(16, 1);
111        let b1 = p.checkout();
112        let b2 = p.checkout();
113        p.return_buffer(b1);
114        p.return_buffer(b2);
115        assert_eq!(p.available(), 1);
116    }
117
118    #[test]
119    fn preallocate_fills_pool() {
120        let mut p = BufferPool::new(8, 5);
121        p.preallocate(3);
122        assert_eq!(p.available(), 3);
123    }
124
125    #[test]
126    fn shrink_to_reduces() {
127        let mut p = BufferPool::new(8, 10);
128        p.preallocate(5);
129        p.shrink_to(2);
130        assert_eq!(p.available(), 2);
131    }
132
133    #[test]
134    fn clear_empties_pool() {
135        let mut p = BufferPool::new(8, 4);
136        p.preallocate(4);
137        p.clear();
138        assert_eq!(p.available(), 0);
139    }
140
141    #[test]
142    fn total_memory_calculated() {
143        let mut p = BufferPool::new(100, 10);
144        p.preallocate(3);
145        assert_eq!(p.total_memory(), 300);
146    }
147
148    #[test]
149    fn checkout_count_increments() {
150        let mut p = BufferPool::new(8, 4);
151        let _ = p.checkout();
152        let _ = p.checkout();
153        assert_eq!(p.total_checkouts(), 2);
154    }
155
156    #[test]
157    fn buffer_size_returns_configured() {
158        let p = BufferPool::new(256, 4);
159        assert_eq!(p.buffer_size(), 256);
160    }
161}