use crate::compat::{lock, vec, Any, Box, Mutex, String, Vec};
use alloc::sync::Arc;
pub trait Poolable: Default + Clone {
fn reset(&mut self);
}
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct PoolConfig {
pub initial_size: usize,
pub max_size: usize,
pub growth_factor: f32,
}
impl Default for PoolConfig {
fn default() -> Self {
Self { initial_size: 16, max_size: 1024, growth_factor: 1.5 }
}
}
pub struct ObjectPool<T: Poolable> {
pool: Vec<T>,
config: PoolConfig,
allocated: usize,
}
impl<T: Poolable> ObjectPool<T> {
pub fn new(config: PoolConfig) -> Self {
let initial = config.initial_size.min(config.max_size);
let mut pool = Vec::with_capacity(initial);
for _ in 0..initial {
pool.push(T::default());
}
Self { pool, config, allocated: 0 }
}
pub fn acquire(&mut self) -> T {
if let Some(mut obj) = self.pool.pop() {
obj.reset();
self.allocated += 1;
obj
} else {
self.allocated += 1;
T::default()
}
}
pub fn release(&mut self, mut obj: T) {
obj.reset();
if self.pool.len() < self.config.max_size {
self.pool.push(obj);
}
self.allocated = self.allocated.saturating_sub(1);
}
pub fn available(&self) -> usize {
self.pool.len()
}
pub fn allocated(&self) -> usize {
self.allocated
}
pub fn capacity(&self) -> usize {
self.pool.capacity()
}
pub fn clear(&mut self) {
self.pool.clear();
self.allocated = 0;
}
pub fn shrink_to_fit(&mut self) {
self.pool.shrink_to_fit();
}
}
impl<T: Poolable> Default for ObjectPool<T> {
fn default() -> Self {
Self::new(PoolConfig::default())
}
}
pub struct SharedPool<T: Poolable + Send> {
pool: Arc<Mutex<ObjectPool<T>>>,
}
impl<T: Poolable + Send> SharedPool<T> {
pub fn new(config: PoolConfig) -> Self {
Self { pool: Arc::new(Mutex::new(ObjectPool::new(config))) }
}
pub fn acquire(&self) -> T {
lock(&self.pool).acquire()
}
pub fn release(&self, obj: T) {
lock(&self.pool).release(obj);
}
pub fn stats(&self) -> PoolStats {
let pool = lock(&self.pool);
PoolStats {
available: pool.available(),
allocated: pool.allocated(),
capacity: pool.capacity(),
}
}
}
impl<T: Poolable + Send> Clone for SharedPool<T> {
fn clone(&self) -> Self {
Self { pool: Arc::clone(&self.pool) }
}
}
impl<T: Poolable + Send> Default for SharedPool<T> {
fn default() -> Self {
Self::new(PoolConfig::default())
}
}
#[derive(Debug, Clone, Copy)]
pub struct PoolStats {
pub available: usize,
pub allocated: usize,
pub capacity: usize,
}
pub struct PoolManager {
pools: Vec<Box<dyn Any + Send>>,
}
impl PoolManager {
pub fn new() -> Self {
Self { pools: Vec::new() }
}
pub fn register<T: Poolable + Send + 'static>(&mut self, config: PoolConfig) -> SharedPool<T> {
let pool = SharedPool::new(config);
self.pools.push(Box::new(pool.clone()));
pool
}
pub fn clear_all(&mut self) {
self.pools.clear();
}
}
crate::impl_default_via_new!(PoolManager);
pub struct BufferPool {
buffers: Vec<Vec<u8>>,
buffer_size: usize,
max_buffers: usize,
}
impl BufferPool {
pub fn new(buffer_size: usize, initial_count: usize, max_buffers: usize) -> Self {
let initial_count = initial_count.min(max_buffers);
let mut buffers = Vec::with_capacity(initial_count);
for _ in 0..initial_count {
buffers.push(vec![0u8; buffer_size]);
}
Self { buffers, buffer_size, max_buffers }
}
pub fn acquire(&mut self) -> Vec<u8> {
match self.buffers.pop() {
Some(mut buffer) => {
buffer.resize(self.buffer_size, 0);
buffer
}
None => vec![0u8; self.buffer_size],
}
}
pub fn acquire_sized(&mut self, size: usize) -> Vec<u8> {
if size <= self.buffer_size {
let mut buf = self.acquire();
buf.clear();
buf.resize(size, 0);
buf
} else {
vec![0u8; size]
}
}
pub fn release(&mut self, mut buffer: Vec<u8>) {
if buffer.capacity() == self.buffer_size && self.buffers.len() < self.max_buffers {
buffer.clear();
self.buffers.push(buffer);
}
}
pub fn available(&self) -> usize {
self.buffers.len()
}
pub fn buffer_size(&self) -> usize {
self.buffer_size
}
pub fn clear(&mut self) {
self.buffers.clear();
}
}
impl Default for BufferPool {
fn default() -> Self {
Self::new(4096, 4, 64)
}
}
pub struct StringPool {
strings: Vec<String>,
default_capacity: usize,
max_strings: usize,
}
impl StringPool {
pub fn new(default_capacity: usize, initial_count: usize, max_strings: usize) -> Self {
let initial_count = initial_count.min(max_strings);
let mut strings = Vec::with_capacity(initial_count);
for _ in 0..initial_count {
strings.push(String::with_capacity(default_capacity));
}
Self { strings, default_capacity, max_strings }
}
pub fn acquire(&mut self) -> String {
self.strings.pop().unwrap_or_else(|| String::with_capacity(self.default_capacity))
}
pub fn release(&mut self, mut s: String) {
s.clear();
if s.capacity() >= self.default_capacity && self.strings.len() < self.max_strings {
self.strings.push(s);
}
}
pub fn available(&self) -> usize {
self.strings.len()
}
pub fn clear(&mut self) {
self.strings.clear();
}
}
impl Default for StringPool {
fn default() -> Self {
Self::new(64, 8, 128)
}
}
pub struct VecPool<T> {
vecs: Vec<Vec<T>>,
default_capacity: usize,
max_vecs: usize,
}
impl<T> VecPool<T> {
pub fn new(default_capacity: usize, initial_count: usize, max_vecs: usize) -> Self {
let initial_count = initial_count.min(max_vecs);
let mut vecs = Vec::with_capacity(initial_count);
for _ in 0..initial_count {
vecs.push(Vec::with_capacity(default_capacity));
}
Self { vecs, default_capacity, max_vecs }
}
pub fn acquire(&mut self) -> Vec<T> {
self.vecs.pop().unwrap_or_else(|| Vec::with_capacity(self.default_capacity))
}
pub fn release(&mut self, mut v: Vec<T>) {
v.clear();
if v.capacity() >= self.default_capacity && self.vecs.len() < self.max_vecs {
self.vecs.push(v);
}
}
pub fn available(&self) -> usize {
self.vecs.len()
}
pub fn clear(&mut self) {
self.vecs.clear();
}
}
impl<T> Default for VecPool<T> {
fn default() -> Self {
Self::new(16, 4, 64)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[derive(Default, Clone)]
struct TestObject {
value: i32,
data: Vec<u8>,
}
impl Poolable for TestObject {
fn reset(&mut self) {
self.value = 0;
self.data.clear();
}
}
#[test]
fn test_object_pool() {
let mut pool: ObjectPool<TestObject> =
ObjectPool::new(PoolConfig { initial_size: 4, max_size: 8, growth_factor: 1.5 });
assert_eq!(pool.available(), 4);
let obj1 = pool.acquire();
assert_eq!(pool.allocated(), 1);
pool.release(obj1);
assert_eq!(pool.available(), 4);
}
#[test]
fn test_buffer_pool() {
let mut pool = BufferPool::new(1024, 2, 4);
let buf1 = pool.acquire();
assert_eq!(buf1.len(), 1024);
pool.release(buf1);
assert_eq!(pool.available(), 2);
}
#[test]
fn buffer_pool_re_acquire_restores_the_documented_length() {
let mut pool = BufferPool::new(16, 1, 4);
let first = pool.acquire();
assert_eq!(first.len(), 16, "first acquire is a full buffer");
pool.release(first);
let second = pool.acquire();
assert_eq!(second.len(), 16, "re-acquire must restore the length, not return len=0");
assert!(second.capacity() >= 16);
}
#[test]
fn buffer_pool_acquire_sized_then_acquire_restores_the_length() {
let mut pool = BufferPool::new(16, 1, 4);
let small = pool.acquire_sized(8);
assert_eq!(small.len(), 8, "acquire_sized returns exactly the requested size");
pool.release(small);
let full = pool.acquire();
assert_eq!(full.len(), 16, "a normal acquire is always a full buffer");
}
#[test]
fn pool_constructors_clamp_initial_fill_to_the_maximum() {
let object_pool = ObjectPool::<TestObject>::new(PoolConfig {
initial_size: 3,
max_size: 1,
growth_factor: 1.5,
});
assert_eq!(object_pool.available(), 1, "object pool is bounded by max_size");
let buffer_pool = BufferPool::new(8, 3, 1);
assert_eq!(buffer_pool.available(), 1, "buffer pool is bounded by max_buffers");
let string_pool = StringPool::new(4, 3, 1);
assert_eq!(string_pool.available(), 1, "string pool is bounded by max_strings");
let vec_pool: VecPool<u8> = VecPool::new(4, 3, 1);
assert_eq!(vec_pool.available(), 1, "vec pool is bounded by max_vecs");
}
#[test]
fn pool_constructors_preserve_zero_max_and_normal_cases() {
assert_eq!(
ObjectPool::<TestObject>::new(PoolConfig {
initial_size: 3,
max_size: 0,
growth_factor: 1.5,
})
.available(),
0
);
assert_eq!(BufferPool::new(8, 3, 0).available(), 0);
assert_eq!(StringPool::new(4, 3, 0).available(), 0);
assert_eq!(VecPool::<u8>::new(4, 3, 0).available(), 0);
assert_eq!(
ObjectPool::<TestObject>::new(PoolConfig {
initial_size: 2,
max_size: 8,
growth_factor: 1.5,
})
.available(),
2
);
assert_eq!(BufferPool::new(8, 2, 8).available(), 2);
assert_eq!(StringPool::new(4, 2, 8).available(), 2);
assert_eq!(VecPool::<u8>::new(4, 2, 8).available(), 2);
}
}