Skip to main content

vane_core/
buffer.rs

1//! Fixed buffer pool (`IORING_REGISTER_BUFFERS` semantics, `IO-02`).
2//!
3//! Every worker pre-allocates `CAP` slots of `BUF_SIZE` bytes up front —
4//! the request hot path never calls `malloc`. Slots have *stable addresses*,
5//! which is what allows the io_uring backend to register them with the ring
6//! (`IORING_REGISTER_BUFFERS`) and read with `buf_index = slot`. The mio
7//! backend dereferences the same slots at event time.
8
9/// Default per-slot buffer size (single read batch).
10pub const DEFAULT_BUF_SIZE: usize = 4096;
11
12/// Default slots per worker (power of two, io_uring registration-friendly).
13pub const DEFAULT_POOL_SIZE: usize = 1024;
14
15/// Pre-allocated, fixed-size buffer pool owned by one worker.
16pub struct BufferPool {
17    /// Stable-address slot storage: `slots[i]` is `BUF_SIZE` bytes.
18    slots: Box<[Box<[u8; DEFAULT_BUF_SIZE]>]>,
19    /// LIFO free list of slot indices (worker-local, single thread).
20    free: Vec<u32>,
21    /// Slot size in bytes.
22    buf_size: usize,
23}
24
25impl BufferPool {
26    /// Pre-allocates `capacity` slots of `buf_size` bytes.
27    ///
28    /// Returns `None` when `buf_size != DEFAULT_BUF_SIZE` and the io_uring
29    /// feature would require matching registration granularity — callers
30    /// should keep the default size.
31    #[must_use]
32    pub fn new(capacity: usize, buf_size: usize) -> Option<Self> {
33        if buf_size != DEFAULT_BUF_SIZE {
34            return None; // registration granularity is fixed
35        }
36        let mut slots = Vec::with_capacity(capacity);
37        let mut free = Vec::with_capacity(capacity);
38        for i in 0..capacity {
39            slots.push(vec![0u8; buf_size].into_boxed_slice().try_into().ok()?);
40            free.push(i as u32);
41        }
42        Some(Self {
43            slots: slots.into_boxed_slice(),
44            free,
45            buf_size,
46        })
47    }
48
49    /// Slot size.
50    #[must_use]
51    #[inline]
52    pub fn buf_size(&self) -> usize {
53        self.buf_size
54    }
55
56    /// Number of slots.
57    #[must_use]
58    #[inline]
59    pub fn capacity(&self) -> usize {
60        self.slots.len()
61    }
62
63    /// Free slot count.
64    #[must_use]
65    #[inline]
66    pub fn free_slots(&self) -> usize {
67        self.free.len()
68    }
69
70    /// Takes a slot, or `None` when the pool is exhausted (callers stop
71    /// reading — TCP backpressure does the rest).
72    #[inline]
73    pub fn take(&mut self) -> Option<u32> {
74        self.free.pop()
75    }
76
77    /// Returns a slot to the pool.
78    #[inline]
79    pub fn release(&mut self, slot: u32) {
80        debug_assert!((slot as usize) < self.slots.len(), "slot out of range");
81        self.free.push(slot);
82    }
83
84    /// Reads the slot's bytes (stable address, safe to hand to the kernel).
85    #[must_use]
86    #[inline]
87    pub fn slot(&self, slot: u32) -> &[u8; DEFAULT_BUF_SIZE] {
88        &self.slots[slot as usize]
89    }
90
91    /// Mutable access to a slot's bytes.
92    #[inline]
93    pub fn slot_mut(&mut self, slot: u32) -> &mut [u8; DEFAULT_BUF_SIZE] {
94        &mut self.slots[slot as usize]
95    }
96}
97
98#[cfg(test)]
99mod tests {
100    use super::*;
101
102    #[test]
103    fn take_release() {
104        let mut pool = BufferPool::new(4, DEFAULT_BUF_SIZE).expect("default size");
105        assert_eq!(pool.free_slots(), 4);
106        let a = pool.take().expect("slot");
107        let b = pool.take().expect("slot");
108        assert_eq!(pool.free_slots(), 2);
109        pool.release(a);
110        assert_eq!(pool.free_slots(), 3);
111        let _c = pool.take();
112        let d = pool.take();
113        assert!(d.is_some());
114        pool.release(b); // b was taken by value; return it directly
115    }
116
117    #[test]
118    fn exhaustion() {
119        let mut pool = BufferPool::new(2, DEFAULT_BUF_SIZE).expect("default size");
120        let _ = pool.take();
121        let _ = pool.take();
122        assert!(pool.take().is_none(), "pool exhausted");
123    }
124}