use crossbeam_queue::ArrayQueue;
use std::sync::Arc;
use super::{allocate_overflow_buffer, allocate_reusable_buffer, Buffer, Owned, PooledBuffer};
use crate::error::{Error, Result};
#[derive(Debug)]
pub(crate) struct BufferPoolInner {
pub(crate) capacity: usize,
pub(crate) max_buffers: usize,
pub(crate) overflow_on_exhaustion: bool,
pub(crate) free: ArrayQueue<PooledBuffer>,
}
impl Drop for BufferPoolInner {
fn drop(&mut self) {
while let Some(buffer) = self.free.pop() {
buffer.deallocate();
}
}
}
#[derive(Clone, Debug)]
pub struct BufferPool {
inner: Arc<BufferPoolInner>,
}
impl BufferPool {
pub fn new(capacity: usize, max_buffers: usize) -> Result<Self> {
Self::with_overflow_strategy(capacity, max_buffers, false)
}
pub fn registered() -> Result<Self> {
Self::with_overflow_strategy(
super::REGISTERED_BUFFER_CAPACITY,
super::REGISTERED_BUFFER_COUNT,
true,
)
}
fn with_overflow_strategy(
capacity: usize,
max_buffers: usize,
overflow_on_exhaustion: bool,
) -> Result<Self> {
if capacity == 0 {
return Err(Error::validation(
"buffer capacity must be greater than zero",
"construct the pool with a non-zero buffer capacity",
));
}
if max_buffers == 0 {
return Err(Error::validation(
"max_buffers must be greater than zero",
"construct the pool with at least one reusable buffer",
));
}
let free = ArrayQueue::new(max_buffers);
for _ in 0..max_buffers {
let buffer = allocate_reusable_buffer(capacity)?;
if let Err(buffer) = free.push(buffer) {
buffer.deallocate();
return Err(Error::resource_exhausted(
"failed to initialize buffer pool free queue",
"increase the max_buffers capacity to fit all allocations",
));
}
}
Ok(Self {
inner: Arc::new(BufferPoolInner {
capacity,
max_buffers,
overflow_on_exhaustion,
free,
}),
})
}
pub fn acquire(&self) -> Result<Buffer<Owned>> {
if let Some(buffer) = self.inner.free.pop() {
return Ok(Buffer::new(self.inner.clone(), buffer, 0));
}
if !self.inner.overflow_on_exhaustion {
return Err(Error::resource_exhausted(
format!(
"buffer pool exhausted (capacity: {}, max_buffers: {})",
self.inner.capacity, self.inner.max_buffers
),
"release completed buffers back to the pool or increase max_buffers",
));
}
Ok(Buffer::new(
self.inner.clone(),
allocate_overflow_buffer(self.inner.capacity),
0,
))
}
#[must_use]
pub fn capacity(&self) -> usize {
self.inner.capacity
}
#[must_use]
pub fn max_buffers(&self) -> usize {
self.inner.max_buffers
}
}