base_scheduler/
heap_buffer_pool.rs

1// Released under the MIT License.
2// Copyright, 2026, by Samuel Williams.
3
4use crate::{BufferLease, Cancellation, Cancelled};
5use event_listener::Event;
6#[cfg(not(test))]
7use std::process::abort;
8use std::{
9    cell::UnsafeCell,
10    future::{Future, poll_fn},
11    io,
12    pin::pin,
13    ptr::NonNull,
14    sync::{
15        Arc, Mutex,
16        atomic::{AtomicUsize, Ordering},
17    },
18    task::Poll,
19};
20#[cfg(test)]
21use tests::abort;
22
23/// A bounded, shareable collection of stable initialized slots in one contiguous slab. Construction reserves all payload storage; acquiring, freezing, slicing, and returning leases allocate no payload or reference-count wrappers.
24///
25/// Pools may be shared across connections and file pipelines. Exhaustion waits with cancellation instead of growing the pool. Retained views count against capacity. This portable implementation of `BufferPool` uses a heap allocation. Backends may supply another pool implementation with the same ownership contract.
26#[derive(Clone)]
27pub struct HeapBufferPool(pub(crate) Arc<Pool>);
28
29pub(crate) struct Slot {
30    pub(crate) readers: AtomicUsize,
31}
32
33pub(crate) struct Pool {
34    slab: Arc<Slab>,
35    waiters: Mutex<usize>,
36}
37
38// Availability and lifetime transitions share one atomic word, so closure cannot miss a concurrent return and exhaustion cannot miss a wakeup.
39const CLOSED: usize = 1 << (usize::BITS - 1);
40const WAITING: usize = CLOSED >> 1;
41const AVAILABLE: usize = WAITING - 1;
42const WIDTH: usize = usize::BITS as usize - 2;
43
44// Unit tests reach the ownership limit through real acquisitions and clones. Production keeps the same conservative ceiling as Arc's reference count.
45#[cfg(not(test))]
46const COUNT_LIMIT: usize = isize::MAX as usize;
47#[cfg(test)]
48const COUNT_LIMIT: usize = 128;
49
50// Each physical slot has one persistent strong ownership claim. Its availability bit, mutable lease, or complete group of frozen views owns that claim. Pool handles retain a separate owner while acquisition is possible. Closing the pool retires idle claims; outstanding leases retire theirs when returned.
51pub(crate) struct Slab {
52    pub(crate) slots: Box<[Slot]>,
53    bytes: Box<[UnsafeCell<u8>]>,
54    offset: usize,
55    stride: usize,
56    // One bit per stable slot; the final word never advertises unused indices.
57    free: Box<[AtomicUsize]>,
58    returned: Event,
59    size: usize,
60    alignment: usize,
61}
62
63// SAFETY: A set availability bit, one mutable lease, or immutable views counted in readers exclusively owns each slot. Only the mutable lease can mutate bytes. Clearing an availability bit acquires the final owner's release before accessing that slot. Slots and the byte slab never move or resize.
64unsafe impl Sync for Slab {}
65
66impl HeapBufferPool {
67    /// Reserve `buffers` slots of `buffer_size` bytes each, with an alignment of 1. This is [`with_options`](Self::with_options) applied to [`BufferPoolOptions::new`](crate::BufferPoolOptions::new).
68    pub fn new(buffers: usize, buffer_size: usize) -> io::Result<Self> {
69        Self::with_options(crate::BufferPoolOptions::new(buffers, buffer_size))
70    }
71
72    /// Reserve one contiguous initialized allocation. Slot starts satisfy the requested alignment; padding is outside the usable payload budget.
73    ///
74    /// Returns [`io::ErrorKind::InvalidInput`] if `buffers` or `buffer_size` is zero, the alignment is not a power of two, or the total allocation is not representable.
75    pub fn with_options(options: crate::BufferPoolOptions) -> io::Result<Self> {
76        options.validate()?;
77        let crate::BufferPoolOptions {
78            buffers,
79            buffer_size,
80            alignment,
81        } = options;
82        let stride = (buffer_size + alignment - 1) & !(alignment - 1);
83        let allocation_size = (buffers - 1) * stride + buffer_size + alignment - 1;
84        let bytes = vec![0u8; allocation_size].into_boxed_slice();
85        let offset = (bytes.as_ptr() as usize).wrapping_neg() & (alignment - 1);
86        // SAFETY: UnsafeCell<u8> has exactly u8's layout and drop behavior. The allocation remains owned, initialized, and fixed in size; references are created only for individually owned slots.
87        let bytes = unsafe { Box::from_raw(Box::into_raw(bytes) as *mut [UnsafeCell<u8>]) };
88        let returned = Event::new();
89        returned.notify(0);
90        let free = (0..buffers.div_ceil(WIDTH))
91            .map(|word| {
92                let bits = (buffers - word * WIDTH).min(WIDTH);
93                AtomicUsize::new((1 << bits) - 1)
94            })
95            .collect();
96        let slab = Arc::new(Slab {
97            slots: (0..buffers)
98                .map(|_| Slot {
99                    readers: AtomicUsize::new(0),
100                })
101                .collect(),
102            bytes,
103            offset,
104            stride,
105            free,
106            returned,
107            size: buffer_size,
108            alignment,
109        });
110        // Mint one ownership claim per physical slot. Checkout, freeze, and return transfer it without updating the slab reference count.
111        for _ in 0..buffers {
112            let _ = Arc::into_raw(Arc::clone(&slab));
113        }
114        Ok(Self(Arc::new(Pool {
115            slab,
116            waiters: Mutex::new(0),
117        })))
118    }
119
120    /// The total number of slots, including leased and frozen ones.
121    pub fn capacity(&self) -> usize {
122        self.0.slab.slots.len()
123    }
124    /// The usable length, in bytes, of each slot.
125    pub fn buffer_size(&self) -> usize {
126        self.0.slab.size
127    }
128
129    /// The requested minimum alignment, in bytes, of each slot's first byte.
130    pub fn alignment(&self) -> usize {
131        self.0.slab.alignment
132    }
133
134    /// Acquire exclusive writable storage immediately, or return `None` without allocating.
135    pub fn try_acquire(&self) -> Option<BufferLease> {
136        for (word_index, word) in self.0.slab.free.iter().enumerate() {
137            let mut available = word.load(Ordering::Relaxed);
138            while available & AVAILABLE != 0 {
139                let bit = 1 << (available & AVAILABLE).trailing_zeros();
140                match word.compare_exchange_weak(
141                    available,
142                    available & !bit,
143                    Ordering::Acquire,
144                    Ordering::Relaxed,
145                ) {
146                    Ok(_) => {
147                        let index = word_index * WIDTH + bit.trailing_zeros() as usize;
148                        // SAFETY: Clearing the bit transfers its persistent slot claim to this lease.
149                        let slab =
150                            unsafe { NonNull::new_unchecked(Arc::as_ptr(&self.0.slab).cast_mut()) };
151                        return Some(BufferLease::new(slab, index, self.0.slab.size));
152                    }
153                    Err(current) => available = current,
154                }
155            }
156        }
157        None
158    }
159
160    /// Wait for one lease. Pre-cancellation leaves capacity unchanged; a successfully acquired lease wins a later cancellation race.
161    pub async fn acquire(&self, cancellation: &Cancellation) -> Result<BufferLease, Cancelled> {
162        cancellation.check()?;
163        if let Some(buffer) = self.try_acquire() {
164            return Ok(buffer);
165        }
166        let mut cancelled = pin!(cancellation.cancelled());
167        let _waiting = PoolWaiter::new(&self.0);
168        loop {
169            event_listener::listener!(self.0.slab.returned => listener);
170            if let Some(buffer) = self.try_acquire() {
171                return Ok(buffer);
172            }
173            let mut listener = pin!(listener);
174            poll_fn(|context| {
175                if cancelled.as_mut().poll(context).is_ready() {
176                    return Poll::Ready(Err(Cancelled));
177                }
178                listener.as_mut().poll(context).map(|_| Ok(()))
179            })
180            .await?;
181        }
182    }
183}
184
185impl Drop for Pool {
186    fn drop(&mut self) {
187        // Closing and publishing a return share the same atomic word. Every captured availability bit transfers its claim to this drain; subsequent returns retire their own claim. The pool retains its separate slab owner throughout the drain.
188        for word in &self.slab.free {
189            let mut available = word.swap(CLOSED, Ordering::AcqRel) & AVAILABLE;
190            while available != 0 {
191                available &= available - 1;
192                // SAFETY: Each captured bit owns exactly one claim minted at setup.
193                unsafe { Arc::decrement_strong_count(Arc::as_ptr(&self.slab)) };
194            }
195        }
196    }
197}
198
199struct PoolWaiter<'a>(&'a Pool);
200
201impl<'a> PoolWaiter<'a> {
202    fn new(pool: &'a Pool) -> Self {
203        let mut waiters = pool.waiters.lock().unwrap();
204        if *waiters == 0 {
205            for word in &pool.slab.free {
206                word.fetch_or(WAITING, Ordering::Release);
207            }
208        }
209        if *waiters >= COUNT_LIMIT {
210            abort();
211        }
212        *waiters += 1;
213        Self(pool)
214    }
215}
216
217impl Drop for PoolWaiter<'_> {
218    fn drop(&mut self) {
219        let mut waiters = self.0.waiters.lock().unwrap();
220        *waiters -= 1;
221        if *waiters == 0 {
222            for word in &self.0.slab.free {
223                word.fetch_and(!WAITING, Ordering::Release);
224            }
225        }
226    }
227}
228
229impl Slab {
230    pub(crate) fn pointer(&self, index: usize) -> *mut u8 {
231        assert!(index < self.slots.len());
232        // SAFETY: Every slot's complete payload fits in this stable allocation. UnsafeCell permits mutation of the individually leased range without forming a mutable reference covering other slots.
233        unsafe {
234            self.bytes
235                .as_ptr()
236                .cast::<u8>()
237                .cast_mut()
238                .add(self.offset + index * self.stride)
239        }
240    }
241
242    /// Consume a slot claim. After publication, the pool may close and reclaim it immediately.
243    ///
244    /// SAFETY: `slab` must represent the unique mutable lease or final frozen reader's persistent slot claim.
245    pub(crate) unsafe fn release(slab: NonNull<Self>, index: usize) {
246        // SAFETY: The caller owns this slot's claim until it is published or retired below.
247        let word = unsafe { &slab.as_ref().free[index / WIDTH] };
248        let bit = 1 << (index % WIDTH);
249        let mut state = word.load(Ordering::Relaxed);
250        let mut notification = None;
251        loop {
252            if state & CLOSED != 0 {
253                // SAFETY: A return after closure retires its claim instead of publishing it.
254                unsafe { Arc::decrement_strong_count(slab.as_ptr()) };
255                return;
256            }
257            if state & WAITING != 0 && notification.is_none() {
258                // SAFETY: The slot claim is still ours. A temporary owner protects notification after publication, including concurrent pool closure. This reference-count update is confined to exhausted pools.
259                unsafe {
260                    Arc::increment_strong_count(slab.as_ptr());
261                    notification = Some(Arc::from_raw(slab.as_ptr()));
262                }
263            }
264            match word.compare_exchange_weak(
265                state,
266                state | bit,
267                Ordering::Release,
268                Ordering::Relaxed,
269            ) {
270                Ok(_) => break,
271                Err(current) => state = current,
272            }
273        }
274        // Do not access `slab` or `word` after publication without a separate owner.
275        if let Some(slab) = notification {
276            slab.returned.notify_additional(1);
277        }
278    }
279
280    pub(crate) fn retain(&self, index: usize) {
281        // Frozen views share their slot's one slab claim, so guard their own count.
282        if self.slots[index].readers.fetch_add(1, Ordering::Relaxed) >= COUNT_LIMIT {
283            abort();
284        }
285    }
286
287    pub(crate) unsafe fn release_reader(slab: NonNull<Self>, index: usize) {
288        // SAFETY: A frozen reader retains its slot claim through the final decrement.
289        if unsafe {
290            slab.as_ref().slots[index]
291                .readers
292                .fetch_sub(1, Ordering::AcqRel)
293        } == 1
294        {
295            // SAFETY: The final frozen reader transfers the unique slot claim.
296            unsafe { Self::release(slab, index) };
297        }
298    }
299}
300
301impl crate::BufferPool for HeapBufferPool {
302    type Buffer = BufferLease;
303    fn capacity(&self) -> usize {
304        HeapBufferPool::capacity(self)
305    }
306    fn buffer_size(&self) -> usize {
307        HeapBufferPool::buffer_size(self)
308    }
309    fn alignment(&self) -> usize {
310        HeapBufferPool::alignment(self)
311    }
312    fn try_acquire(&self) -> Option<Self::Buffer> {
313        HeapBufferPool::try_acquire(self)
314    }
315    async fn acquire(&self, cancellation: &Cancellation) -> Result<Self::Buffer, Cancelled> {
316        HeapBufferPool::acquire(self, cancellation).await
317    }
318}
319
320#[cfg(test)]
321mod tests;
base_scheduler/heap_buffer_pool.rs

base_scheduler/
heap_buffer_pool.rs
