use std::collections::VecDeque;
use std::sync::{Arc, Mutex};
use crate::buffer::PinnedBuffer;
use crate::error::Result;
use super::{PoolInner, PoolStats, PooledBuffer};
pub struct BufferPool {
inner: Arc<Mutex<PoolInner>>,
capacity: usize,
buffer_size: usize,
}
impl BufferPool {
pub fn new(capacity: usize, buffer_size: usize) -> Self {
assert!(capacity > 0, "Pool capacity must be greater than zero");
assert!(buffer_size > 0, "Buffer size must be greater than zero");
let mut available = VecDeque::with_capacity(capacity);
for _ in 0..capacity {
available.push_back(PinnedBuffer::with_capacity(buffer_size));
}
Self {
inner: Arc::new(Mutex::new(PoolInner {
available,
in_use: 0,
total_allocations: 0,
failed_allocations: 0,
})),
capacity,
buffer_size,
}
}
pub fn with_factory<F>(capacity: usize, mut buffer_factory: F) -> Self
where
F: FnMut() -> PinnedBuffer<[u8]>,
{
assert!(capacity > 0, "Pool capacity must be greater than zero");
let mut available = VecDeque::with_capacity(capacity);
let mut buffer_size = 0;
for i in 0..capacity {
let buffer = buffer_factory();
if i == 0 {
buffer_size = buffer.len();
} else {
assert_eq!(
buffer.len(),
buffer_size,
"All buffers must have the same size"
);
}
available.push_back(buffer);
}
Self {
inner: Arc::new(Mutex::new(PoolInner {
available,
in_use: 0,
total_allocations: 0,
failed_allocations: 0,
})),
capacity,
buffer_size,
}
}
pub fn try_get(&self) -> Result<Option<PooledBuffer>> {
let mut inner = PoolInner::lock(&self.inner)?;
if let Some(buffer) = inner.available.pop_front() {
inner.in_use += 1;
inner.total_allocations += 1;
Ok(Some(PooledBuffer::new(buffer, Arc::clone(&self.inner))))
} else {
inner.failed_allocations += 1;
Ok(None)
}
}
pub fn try_get_fast(&self) -> Result<Option<PooledBuffer>> {
{
let inner = PoolInner::lock(&self.inner)?;
if inner.available.is_empty() {
return Ok(None);
}
}
self.try_get()
}
pub fn get(&self) -> Option<PooledBuffer> {
self.try_get().unwrap_or(None)
}
pub fn get_blocking(&self) -> Result<PooledBuffer> {
let mut backoff_us = 1;
const MAX_BACKOFF_US: u64 = 1000;
loop {
if let Some(buffer) = self.try_get()? {
return Ok(buffer);
}
std::thread::sleep(std::time::Duration::from_micros(backoff_us));
backoff_us = (backoff_us * 2).min(MAX_BACKOFF_US);
}
}
pub fn capacity(&self) -> usize {
self.capacity
}
pub fn available(&self) -> Result<usize> {
let inner = PoolInner::lock(&self.inner)?;
Ok(inner.available.len())
}
pub fn in_use(&self) -> Result<usize> {
let inner = PoolInner::lock(&self.inner)?;
Ok(inner.in_use)
}
pub fn buffer_size(&self) -> usize {
self.buffer_size
}
pub fn stats(&self) -> PoolStats {
let inner = PoolInner::lock(&self.inner).unwrap();
let utilization = inner.in_use as f64 / self.capacity as f64;
let available = inner.available.len();
PoolStats {
capacity: self.capacity,
available,
in_use: inner.in_use,
buffer_size: self.buffer_size,
total_allocations: inner.total_allocations,
failed_allocations: inner.failed_allocations,
utilization,
total_buffers: self.capacity,
available_buffers: available,
in_use_buffers: inner.in_use,
}
}
pub fn is_empty(&self) -> Result<bool> {
let inner = PoolInner::lock(&self.inner)?;
Ok(inner.available.is_empty())
}
pub fn is_full(&self) -> Result<bool> {
let inner = PoolInner::lock(&self.inner)?;
Ok(inner.available.len() == self.capacity)
}
pub fn snapshot(&self) -> Result<(usize, usize, bool, bool)> {
let inner = PoolInner::lock(&self.inner)?;
let available = inner.available.len();
let in_use = inner.in_use;
let is_empty = inner.available.is_empty();
let is_full = available == self.capacity;
Ok((available, in_use, is_empty, is_full))
}
}
unsafe impl Send for BufferPool {}
unsafe impl Sync for BufferPool {}