use std::sync::{Arc, RwLock};
#[derive(Debug)]
pub struct BufferPool {
buffers: RwLock<Vec<Arc<RwLock<Vec<u8>>>>>,
buffer_size: usize,
max_buffers: usize,
}
impl BufferPool {
#[must_use]
pub fn new(count: usize, buffer_size: usize) -> Self {
let buffers: Vec<_> = (0..count)
.map(|_| Arc::new(RwLock::new(vec![0u8; buffer_size])))
.collect();
Self {
buffers: RwLock::new(buffers),
buffer_size,
max_buffers: count,
}
}
#[must_use]
pub fn with_capacity(max_buffers: usize, buffer_size: usize) -> Self {
Self {
buffers: RwLock::new(Vec::with_capacity(max_buffers)),
buffer_size,
max_buffers,
}
}
#[must_use]
pub fn acquire(&self) -> Option<Arc<RwLock<Vec<u8>>>> {
self.buffers.write().ok()?.pop()
}
#[must_use]
pub fn acquire_or_alloc(&self) -> Arc<RwLock<Vec<u8>>> {
self.acquire()
.unwrap_or_else(|| Arc::new(RwLock::new(vec![0u8; self.buffer_size])))
}
pub fn release(&self, buffer: Arc<RwLock<Vec<u8>>>) {
if let Ok(mut buffers) = self.buffers.write() {
if buffers.len() < self.max_buffers {
if let Ok(mut guard) = buffer.write() {
guard.fill(0);
}
buffers.push(buffer);
}
}
}
#[must_use]
pub fn available(&self) -> usize {
self.buffers.read().map(|b| b.len()).unwrap_or(0)
}
#[must_use]
pub fn buffer_size(&self) -> usize {
self.buffer_size
}
#[must_use]
pub fn max_buffers(&self) -> usize {
self.max_buffers
}
}
impl Default for BufferPool {
fn default() -> Self {
Self::new(4, 4096)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_new() {
let pool = BufferPool::new(4, 1024);
assert_eq!(pool.available(), 4);
assert_eq!(pool.buffer_size(), 1024);
assert_eq!(pool.max_buffers(), 4);
}
#[test]
fn test_with_capacity() {
let pool = BufferPool::with_capacity(8, 2048);
assert_eq!(pool.available(), 0);
assert_eq!(pool.buffer_size(), 2048);
assert_eq!(pool.max_buffers(), 8);
}
#[test]
fn test_acquire_release() {
let pool = BufferPool::new(2, 1024);
assert_eq!(pool.available(), 2);
let buf1 = pool.acquire().expect("acquire should succeed");
assert_eq!(pool.available(), 1);
let buf2 = pool.acquire().expect("acquire should succeed");
assert_eq!(pool.available(), 0);
assert!(pool.acquire().is_none());
pool.release(buf1);
assert_eq!(pool.available(), 1);
pool.release(buf2);
assert_eq!(pool.available(), 2);
}
#[test]
fn test_acquire_or_alloc() {
let pool = BufferPool::new(0, 1024);
assert_eq!(pool.available(), 0);
let buffer = pool.acquire_or_alloc();
assert_eq!(buffer.read().expect("read lock should succeed").len(), 1024);
}
#[test]
fn test_buffer_contents() {
let pool = BufferPool::new(1, 64);
let buffer = pool.acquire().expect("acquire should succeed");
{
let mut guard = buffer.write().expect("write lock should succeed");
guard[0] = 42;
guard[63] = 255;
}
{
let guard = buffer.read().expect("read lock should succeed");
assert_eq!(guard[0], 42);
assert_eq!(guard[63], 255);
}
pool.release(buffer);
let buffer = pool.acquire().expect("acquire should succeed");
{
let guard = buffer.read().expect("read lock should succeed");
assert_eq!(guard[0], 0);
assert_eq!(guard[63], 0);
}
}
#[test]
fn test_default() {
let pool = BufferPool::default();
assert_eq!(pool.available(), 4);
assert_eq!(pool.buffer_size(), 4096);
}
#[test]
fn test_release_at_capacity() {
let pool = BufferPool::new(2, 1024);
let extra_buffer = Arc::new(RwLock::new(vec![0u8; 1024]));
pool.release(extra_buffer);
assert_eq!(pool.available(), 2); }
}