wireshift-core 0.1.1

Core typed operations, buffers, and backend traits for wireshift
Documentation
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();
        }
    }
}

/// Fixed-capacity buffer pool with explicit lifecycle transitions.
///
/// ```rust
/// use wireshift::BufferPool;
///
/// let pool = BufferPool::new(1024, 4)?;
/// let buf = pool.acquire()?;
/// assert_eq!(buf.capacity(), 1024);
/// # Ok::<(), wireshift::Error>(())
/// ```
#[derive(Clone, Debug)]
pub struct BufferPool {
    inner: Arc<BufferPoolInner>,
}

impl BufferPool {
    /// Creates a new pool.
    pub fn new(capacity: usize, max_buffers: usize) -> Result<Self> {
        Self::with_overflow_strategy(capacity, max_buffers, false)
    }

    /// Creates the registered-buffer pool profile used by wireshift's hot read paths.
    ///
    /// This pool eagerly allocates 32 reusable 256 KiB buffers and falls back
    /// to standalone heap buffers when all pooled entries are in flight.
    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,
            }),
        })
    }

    /// Acquires a buffer for new I/O operations.
    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,
        ))
    }

    /// The configured per-buffer capacity.
    #[must_use]
    pub fn capacity(&self) -> usize {
        self.inner.capacity
    }

    /// Returns the number of reusable buffers kept in the lock-free queue.
    #[must_use]
    pub fn max_buffers(&self) -> usize {
        self.inner.max_buffers
    }
}