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    pub fn new(buffers: usize, buffer_size: usize) -> io::Result<Self> {
68        Self::with_options(crate::BufferPoolOptions::new(buffers, buffer_size))
69    }
70
71    /// Reserve one contiguous initialized allocation. Slot starts satisfy the requested alignment; padding is outside the usable payload budget.
72    pub fn with_options(options: crate::BufferPoolOptions) -> io::Result<Self> {
73        options.validate()?;
74        let crate::BufferPoolOptions {
75            buffers,
76            buffer_size,
77            alignment,
78        } = options;
79        let stride = (buffer_size + alignment - 1) & !(alignment - 1);
80        let allocation_size = (buffers - 1) * stride + buffer_size + alignment - 1;
81        let bytes = vec![0u8; allocation_size].into_boxed_slice();
82        let offset = (bytes.as_ptr() as usize).wrapping_neg() & (alignment - 1);
83        // 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.
84        let bytes = unsafe { Box::from_raw(Box::into_raw(bytes) as *mut [UnsafeCell<u8>]) };
85        let returned = Event::new();
86        returned.notify(0);
87        let free = (0..buffers.div_ceil(WIDTH))
88            .map(|word| {
89                let bits = (buffers - word * WIDTH).min(WIDTH);
90                AtomicUsize::new((1 << bits) - 1)
91            })
92            .collect();
93        let slab = Arc::new(Slab {
94            slots: (0..buffers)
95                .map(|_| Slot {
96                    readers: AtomicUsize::new(0),
97                })
98                .collect(),
99            bytes,
100            offset,
101            stride,
102            free,
103            returned,
104            size: buffer_size,
105            alignment,
106        });
107        // Mint one ownership claim per physical slot. Checkout, freeze, and return transfer it without updating the slab reference count.
108        for _ in 0..buffers {
109            let _ = Arc::into_raw(Arc::clone(&slab));
110        }
111        Ok(Self(Arc::new(Pool {
112            slab,
113            waiters: Mutex::new(0),
114        })))
115    }
116
117    pub fn capacity(&self) -> usize {
118        self.0.slab.slots.len()
119    }
120    pub fn buffer_size(&self) -> usize {
121        self.0.slab.size
122    }
123
124    pub fn alignment(&self) -> usize {
125        self.0.slab.alignment
126    }
127
128    /// Acquire exclusive writable storage immediately, or return `None` without allocating.
129    pub fn try_acquire(&self) -> Option<BufferLease> {
130        for (word_index, word) in self.0.slab.free.iter().enumerate() {
131            let mut available = word.load(Ordering::Relaxed);
132            while available & AVAILABLE != 0 {
133                let bit = 1 << (available & AVAILABLE).trailing_zeros();
134                match word.compare_exchange_weak(
135                    available,
136                    available & !bit,
137                    Ordering::Acquire,
138                    Ordering::Relaxed,
139                ) {
140                    Ok(_) => {
141                        let index = word_index * WIDTH + bit.trailing_zeros() as usize;
142                        // SAFETY: Clearing the bit transfers its persistent slot claim to this lease.
143                        let slab =
144                            unsafe { NonNull::new_unchecked(Arc::as_ptr(&self.0.slab).cast_mut()) };
145                        return Some(BufferLease::new(slab, index, self.0.slab.size));
146                    }
147                    Err(current) => available = current,
148                }
149            }
150        }
151        None
152    }
153
154    /// Wait for one lease. Pre-cancellation leaves capacity unchanged; a successfully acquired lease wins a later cancellation race.
155    pub async fn acquire(&self, cancellation: &Cancellation) -> Result<BufferLease, Cancelled> {
156        cancellation.check()?;
157        if let Some(buffer) = self.try_acquire() {
158            return Ok(buffer);
159        }
160        let mut cancelled = pin!(cancellation.cancelled());
161        let _waiting = PoolWaiter::new(&self.0);
162        loop {
163            event_listener::listener!(self.0.slab.returned => listener);
164            if let Some(buffer) = self.try_acquire() {
165                return Ok(buffer);
166            }
167            let mut listener = pin!(listener);
168            poll_fn(|context| {
169                if cancelled.as_mut().poll(context).is_ready() {
170                    return Poll::Ready(Err(Cancelled));
171                }
172                listener.as_mut().poll(context).map(|_| Ok(()))
173            })
174            .await?;
175        }
176    }
177}
178
179impl Drop for Pool {
180    fn drop(&mut self) {
181        // 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.
182        for word in &self.slab.free {
183            let mut available = word.swap(CLOSED, Ordering::AcqRel) & AVAILABLE;
184            while available != 0 {
185                available &= available - 1;
186                // SAFETY: Each captured bit owns exactly one claim minted at setup.
187                unsafe { Arc::decrement_strong_count(Arc::as_ptr(&self.slab)) };
188            }
189        }
190    }
191}
192
193struct PoolWaiter<'a>(&'a Pool);
194
195impl<'a> PoolWaiter<'a> {
196    fn new(pool: &'a Pool) -> Self {
197        let mut waiters = pool.waiters.lock().unwrap();
198        if *waiters == 0 {
199            for word in &pool.slab.free {
200                word.fetch_or(WAITING, Ordering::Release);
201            }
202        }
203        if *waiters >= COUNT_LIMIT {
204            abort();
205        }
206        *waiters += 1;
207        Self(pool)
208    }
209}
210
211impl Drop for PoolWaiter<'_> {
212    fn drop(&mut self) {
213        let mut waiters = self.0.waiters.lock().unwrap();
214        *waiters -= 1;
215        if *waiters == 0 {
216            for word in &self.0.slab.free {
217                word.fetch_and(!WAITING, Ordering::Release);
218            }
219        }
220    }
221}
222
223impl Slab {
224    pub(crate) fn pointer(&self, index: usize) -> *mut u8 {
225        assert!(index < self.slots.len());
226        // 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.
227        unsafe {
228            self.bytes
229                .as_ptr()
230                .cast::<u8>()
231                .cast_mut()
232                .add(self.offset + index * self.stride)
233        }
234    }
235
236    /// Consume a slot claim. After publication, the pool may close and reclaim it immediately.
237    ///
238    /// SAFETY: `slab` must represent the unique mutable lease or final frozen reader's persistent slot claim.
239    pub(crate) unsafe fn release(slab: NonNull<Self>, index: usize) {
240        // SAFETY: The caller owns this slot's claim until it is published or retired below.
241        let word = unsafe { &slab.as_ref().free[index / WIDTH] };
242        let bit = 1 << (index % WIDTH);
243        let mut state = word.load(Ordering::Relaxed);
244        let mut notification = None;
245        loop {
246            if state & CLOSED != 0 {
247                // SAFETY: A return after closure retires its claim instead of publishing it.
248                unsafe { Arc::decrement_strong_count(slab.as_ptr()) };
249                return;
250            }
251            if state & WAITING != 0 && notification.is_none() {
252                // 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.
253                unsafe {
254                    Arc::increment_strong_count(slab.as_ptr());
255                    notification = Some(Arc::from_raw(slab.as_ptr()));
256                }
257            }
258            match word.compare_exchange_weak(
259                state,
260                state | bit,
261                Ordering::Release,
262                Ordering::Relaxed,
263            ) {
264                Ok(_) => break,
265                Err(current) => state = current,
266            }
267        }
268        // Do not access `slab` or `word` after publication without a separate owner.
269        if let Some(slab) = notification {
270            slab.returned.notify_additional(1);
271        }
272    }
273
274    pub(crate) fn retain(&self, index: usize) {
275        // Frozen views share their slot's one slab claim, so guard their own count.
276        if self.slots[index].readers.fetch_add(1, Ordering::Relaxed) >= COUNT_LIMIT {
277            abort();
278        }
279    }
280
281    pub(crate) unsafe fn release_reader(slab: NonNull<Self>, index: usize) {
282        // SAFETY: A frozen reader retains its slot claim through the final decrement.
283        if unsafe {
284            slab.as_ref().slots[index]
285                .readers
286                .fetch_sub(1, Ordering::AcqRel)
287        } == 1
288        {
289            // SAFETY: The final frozen reader transfers the unique slot claim.
290            unsafe { Self::release(slab, index) };
291        }
292    }
293}
294
295impl crate::BufferPool for HeapBufferPool {
296    type Buffer = BufferLease;
297    fn capacity(&self) -> usize {
298        HeapBufferPool::capacity(self)
299    }
300    fn buffer_size(&self) -> usize {
301        HeapBufferPool::buffer_size(self)
302    }
303    fn alignment(&self) -> usize {
304        HeapBufferPool::alignment(self)
305    }
306    fn try_acquire(&self) -> Option<Self::Buffer> {
307        HeapBufferPool::try_acquire(self)
308    }
309    async fn acquire(&self, cancellation: &Cancellation) -> Result<Self::Buffer, Cancelled> {
310        HeapBufferPool::acquire(self, cancellation).await
311    }
312}
313
314#[cfg(test)]
315mod tests;
base_scheduler/heap_buffer_pool.rs

base_scheduler/
heap_buffer_pool.rs
