wireshift_core/buffer/
pool.rs

1use crossbeam_queue::ArrayQueue;
2use std::sync::Arc;
3
4use super::{allocate_overflow_buffer, allocate_reusable_buffer, Buffer, Owned, PooledBuffer};
5use crate::error::{Error, Result};
6
7#[derive(Debug)]
8pub(crate) struct BufferPoolInner {
9    pub(crate) capacity: usize,
10    pub(crate) max_buffers: usize,
11    pub(crate) overflow_on_exhaustion: bool,
12    pub(crate) free: ArrayQueue<PooledBuffer>,
13}
14
15impl Drop for BufferPoolInner {
16    fn drop(&mut self) {
17        while let Some(buffer) = self.free.pop() {
18            buffer.deallocate();
19        }
20    }
21}
22
23/// Fixed-capacity buffer pool with explicit lifecycle transitions.
24///
25/// ```rust
26/// use wireshift::BufferPool;
27///
28/// let pool = BufferPool::new(1024, 4)?;
29/// let buf = pool.acquire()?;
30/// assert_eq!(buf.capacity(), 1024);
31/// # Ok::<(), wireshift::Error>(())
32/// ```
33#[derive(Clone, Debug)]
34pub struct BufferPool {
35    inner: Arc<BufferPoolInner>,
36}
37
38impl BufferPool {
39    /// Creates a new pool.
40    pub fn new(capacity: usize, max_buffers: usize) -> Result<Self> {
41        Self::with_overflow_strategy(capacity, max_buffers, false)
42    }
43
44    /// Creates the registered-buffer pool profile used by wireshift's hot read paths.
45    ///
46    /// This pool eagerly allocates 32 reusable 256 KiB buffers and falls back
47    /// to standalone heap buffers when all pooled entries are in flight.
48    pub fn registered() -> Result<Self> {
49        Self::with_overflow_strategy(
50            super::REGISTERED_BUFFER_CAPACITY,
51            super::REGISTERED_BUFFER_COUNT,
52            true,
53        )
54    }
55
56    fn with_overflow_strategy(
57        capacity: usize,
58        max_buffers: usize,
59        overflow_on_exhaustion: bool,
60    ) -> Result<Self> {
61        if capacity == 0 {
62            return Err(Error::validation(
63                "buffer capacity must be greater than zero",
64                "construct the pool with a non-zero buffer capacity",
65            ));
66        }
67        if max_buffers == 0 {
68            return Err(Error::validation(
69                "max_buffers must be greater than zero",
70                "construct the pool with at least one reusable buffer",
71            ));
72        }
73        let free = ArrayQueue::new(max_buffers);
74        for _ in 0..max_buffers {
75            let buffer = allocate_reusable_buffer(capacity)?;
76            if let Err(buffer) = free.push(buffer) {
77                buffer.deallocate();
78                return Err(Error::resource_exhausted(
79                    "failed to initialize buffer pool free queue",
80                    "increase the max_buffers capacity to fit all allocations",
81                ));
82            }
83        }
84        Ok(Self {
85            inner: Arc::new(BufferPoolInner {
86                capacity,
87                max_buffers,
88                overflow_on_exhaustion,
89                free,
90            }),
91        })
92    }
93
94    /// Acquires a buffer for new I/O operations.
95    pub fn acquire(&self) -> Result<Buffer<Owned>> {
96        if let Some(buffer) = self.inner.free.pop() {
97            return Ok(Buffer::new(self.inner.clone(), buffer, 0));
98        }
99        if !self.inner.overflow_on_exhaustion {
100            return Err(Error::resource_exhausted(
101                format!(
102                    "buffer pool exhausted (capacity: {}, max_buffers: {})",
103                    self.inner.capacity, self.inner.max_buffers
104                ),
105                "release completed buffers back to the pool or increase max_buffers",
106            ));
107        }
108        Ok(Buffer::new(
109            self.inner.clone(),
110            allocate_overflow_buffer(self.inner.capacity),
111            0,
112        ))
113    }
114
115    /// The configured per-buffer capacity.
116    #[must_use]
117    pub fn capacity(&self) -> usize {
118        self.inner.capacity
119    }
120
121    /// Returns the number of reusable buffers kept in the lock-free queue.
122    #[must_use]
123    pub fn max_buffers(&self) -> usize {
124        self.inner.max_buffers
125    }
126}
wireshift_core/buffer/pool.rs

wireshift_core/buffer/
pool.rs
