dtact_util/lockfree.rs
1//! Shared lock-free building blocks used by every module in this crate's
2//! native backends.
3//!
4//! Covers `timer`, `fs`, and, going forward, `process`/`signal`/`stream`.
5//! Nothing here takes an `std::sync::Mutex`/`Condvar` on any hot path —
6//! completion state is plain atomics, waker storage is a single wait-free
7//! `AtomicPtr<Waker>` swap, and cross-thread handoff queues are lock-free
8//! Treiber stacks, not `Mutex<Vec<_>>`.
9
10use std::collections::VecDeque;
11use std::ptr;
12use std::sync::atomic::{AtomicPtr, AtomicU32, AtomicU64, AtomicUsize, Ordering};
13use std::task::{Context, Poll, Waker};
14
15// =============================================================================
16// TreiberStack — lock-free, ABA-safe, index-based free-list
17// =============================================================================
18// Moved here from `io::native` (previously a private copy inside that
19// module) so every native backend that wants a preallocated slot pool —
20// `io`, and now `fs` — shares one implementation instead of each hand-
21// rolling its own. The tagged 64-bit head (32-bit index + 32-bit
22// generation tag packed together) makes push/pop immune to the classic
23// ABA problem on a lock-free stack: a popped-then-repushed index can't be
24// mistaken for "unchanged" because the tag always advances.
25/// Lock-free, ABA-safe, index-based free-list of `size` slots (indices
26/// `0..size`), all initially free.
27#[repr(align(64))]
28pub struct TreiberStack {
29 head: AtomicU64,
30 next: Box<[AtomicU32]>,
31}
32
33impl TreiberStack {
34 /// Build a stack holding indices `0..size`, all free.
35 #[must_use]
36 pub fn new(size: usize) -> Self {
37 let mut next = Vec::with_capacity(size);
38 for _ in 0..size {
39 next.push(AtomicU32::new(0));
40 }
41 if size > 0 {
42 next[size - 1].store(u32::MAX, Ordering::Relaxed);
43 }
44 Self {
45 head: AtomicU64::new(u64::from(u32::MAX)),
46 next: next.into_boxed_slice(),
47 }
48 }
49
50 /// Build a stack holding indices `0..size`, all free.
51 #[must_use]
52 pub(crate) fn new_empty(size: usize) -> Self {
53 let mut next = Vec::with_capacity(size);
54 for _ in 0..size {
55 next.push(AtomicU32::new(0));
56 }
57 Self {
58 head: AtomicU64::new(u64::from(u32::MAX)),
59 next: next.into_boxed_slice(),
60 }
61 }
62
63 /// Return `idx` to the free-list. `idx` must have previously come from
64 /// [`Self::pop`] (or from initial construction) and not currently be
65 /// held by anything else — pushing an index that's already free is a
66 /// caller bug (double-free of the index) and will corrupt the
67 /// free-list for subsequent `pop`s.
68 #[inline]
69 pub fn push(&self, idx: u32) {
70 let mut head = self.head.load(Ordering::Acquire);
71 loop {
72 let head_idx = (head & 0xFFFF_FFFF) as u32;
73 let tag = (head >> 32) as u32;
74
75 self.next[idx as usize].store(head_idx, Ordering::Release);
76 let new_head = (u64::from(tag.wrapping_add(1)) << 32) | u64::from(idx);
77
78 match self.head.compare_exchange_weak(
79 head,
80 new_head,
81 Ordering::Release,
82 Ordering::Acquire,
83 ) {
84 Ok(_) => break,
85 Err(actual) => head = actual,
86 }
87 }
88 }
89
90 /// Take an index off the free-list, or `None` if it's exhausted.
91 #[inline]
92 pub fn pop(&self) -> Option<u32> {
93 let mut head = self.head.load(Ordering::Acquire);
94 loop {
95 let head_idx = (head & 0xFFFF_FFFF) as u32;
96 if head_idx == u32::MAX {
97 return None;
98 }
99 let tag = (head >> 32) as u32;
100 let next = self.next[head_idx as usize].load(Ordering::Acquire);
101 let new_head = (u64::from(tag.wrapping_add(1)) << 32) | u64::from(next);
102
103 match self.head.compare_exchange_weak(
104 head,
105 new_head,
106 Ordering::Release,
107 Ordering::Acquire,
108 ) {
109 Ok(_) => return Some(head_idx),
110 Err(actual) => head = actual,
111 }
112 }
113 }
114
115 /// Linear initialization helper for single-threaded setup paths
116 #[inline]
117 pub(crate) fn write_next_slot(&self, idx: usize, next_idx: u32) {
118 self.next[idx].store(next_idx, Ordering::Relaxed);
119 }
120
121 /// Single-threaded initialization helper to set the initial entry head
122 #[inline]
123 pub(crate) fn set_head(&self, head_idx: u32) {
124 let initial_head = u64::from(head_idx);
125 self.head.store(initial_head, Ordering::Relaxed);
126 }
127}
128
129// =============================================================================
130// BufferPool — page-aligned arena carved into fixed-size chunks, handed
131// out/reclaimed via a TreiberStack free-list
132// =============================================================================
133// Also moved here from `io::native` (previously private, and duplicated
134// in spirit by `fs`'s earlier per-op `Vec<u8>`/`Box<OpState>` allocations
135// before this pass). One arena `alloc()` up front, then `acquire()`/
136// `release()` are index-stack push/pop — no allocator call, no lock, on
137// the per-operation hot path.
138/// Page-aligned arena carved into `total_chunks` fixed-size chunks, handed
139/// out and reclaimed via a [`TreiberStack`] free-list.
140pub struct BufferPool {
141 arena_ptr: *mut u8,
142 layout: std::alloc::Layout,
143 chunk_size: usize,
144 free: TreiberStack,
145}
146
147// SAFETY: `arena_ptr` points at a heap allocation this `BufferPool`
148// exclusively owns (freed only in `Drop`); handing out `*mut u8` chunk
149// pointers via `get_ptr` doesn't alias Rust-level state, and the
150// free-list itself (`TreiberStack`) is already lock-free/thread-safe.
151unsafe impl Send for BufferPool {}
152// SAFETY: same reasoning as `Send` above — concurrent `acquire`/`release`
153// only ever goes through the already-thread-safe `TreiberStack`.
154unsafe impl Sync for BufferPool {}
155
156impl BufferPool {
157 /// Allocate a page-aligned arena of `total_chunks` chunks of
158 /// `chunk_size` bytes each, with all chunks initially free.
159 ///
160 /// # Panics
161 ///
162 /// Panics if `total_chunks * chunk_size` (each floored to at least 1)
163 /// overflows or produces a layout with invalid size/alignment for the
164 /// global allocator (`Layout::from_size_align` failure), or if the
165 /// allocator itself fails to satisfy the allocation
166 /// (`handle_alloc_error`).
167 #[must_use]
168 pub fn new(total_chunks: usize, chunk_size: usize) -> Self {
169 let total_chunks_bounded = total_chunks.max(1);
170 let chunk_size_bounded = chunk_size.max(1);
171
172 let layout =
173 std::alloc::Layout::from_size_align(total_chunks_bounded * chunk_size_bounded, 4096)
174 .expect("Invalid layout alignment for BufferPool");
175
176 let arena_ptr = unsafe { std::alloc::alloc(layout) };
177 if arena_ptr.is_null() {
178 std::alloc::handle_alloc_error(layout);
179 }
180
181 // Initialize empty shell with zero contention
182 let free = TreiberStack::new_empty(total_chunks_bounded);
183
184 // Link chunk i directly to chunk i + 1 using Relaxed writes
185 for i in 0..total_chunks_bounded {
186 let next_idx = if i == total_chunks_bounded - 1 {
187 u32::MAX
188 } else {
189 (i + 1) as u32
190 };
191 free.write_next_slot(i, next_idx);
192 }
193
194 // Point head directly to chunk 0
195 free.set_head(0);
196
197 Self {
198 arena_ptr,
199 layout,
200 chunk_size: chunk_size_bounded,
201 free,
202 }
203 }
204
205 /// Raw pointer to the start of chunk `idx` within the arena. `idx`
206 /// must be `< total_chunks` as passed to [`Self::new`] — out-of-range
207 /// indices produce a pointer outside the arena allocation, which is
208 /// unsound to dereference (this function itself performs no
209 /// dereference, only pointer arithmetic).
210 #[inline]
211 pub const fn get_ptr(&self, idx: u32) -> *mut u8 {
212 unsafe { self.arena_ptr.add(idx as usize * self.chunk_size) }
213 }
214
215 /// The fixed chunk size (in bytes) this pool was constructed with.
216 #[inline]
217 pub const fn chunk_size(&self) -> usize {
218 self.chunk_size
219 }
220
221 /// Borrow a chunk index from the pool; `None` if exhausted.
222 #[inline]
223 pub fn acquire(&self) -> Option<u32> {
224 self.free.pop()
225 }
226
227 /// Return a chunk index to the pool.
228 #[inline]
229 pub fn release(&self, idx: u32) {
230 self.free.push(idx);
231 }
232}
233
234impl Drop for BufferPool {
235 fn drop(&mut self) {
236 unsafe {
237 std::alloc::dealloc(self.arena_ptr, self.layout);
238 }
239 }
240}
241
242// =============================================================================
243// AtomicWakerSlot — single-slot waker storage, no per-register allocation
244// =============================================================================
245// Single-`AtomicUsize`-state-machine design (the same algorithm as tokio's
246// production `AtomicWaker`, which this is a from-scratch reimplementation
247// of against the same well-reviewed shape — not a copy of tokio's source).
248// The waker itself lives inline in an `UnsafeCell<Option<Waker>>`; the
249// `AtomicUsize` state is the *sole* arbiter of who's allowed to touch that
250// cell, so there is exactly one atomic to reason about per critical
251// section — no torn reads across two independently-updated atomics.
252//
253// This replaces an earlier `Box::into_raw`/`Box::from_raw`-per-register
254// version: correct (a swap-based ownership handoff is easy to reason
255// about), but it paid a heap allocation *and* a deallocation on every
256// single `register()` call, including the extremely common case where the
257// exact same task polls the exact same listener repeatedly without the
258// registered waker ever changing. That version itself replaced an even
259// earlier two-separate-atomics-plus-spinlock design that had a real,
260// hard-to-pin-down soundness bug (intermittent heap corruption under
261// concurrent register/wake traffic) — the bug there was two independently
262// racy atomics (`data`, `vtable`) with no single point of truth for
263// "who owns the slot right now". The state machine below avoids that
264// specific failure mode structurally: there is only one atomic, and every
265// participant CASes it before touching the `UnsafeCell`, so "who owns the
266// slot" is always unambiguous.
267const WAITING: usize = 0;
268const REGISTERING: usize = 0b01;
269const WAKING: usize = 0b10;
270
271/// Wait-free-on-the-fast-path single-slot waker storage.
272///
273/// `register` skips re-storing when the incoming waker already wakes the
274/// same task as what's registered ([`Waker::will_wake`]), and even when it
275/// does need to store, does so into an inline cell — no heap allocation on
276/// the `register`/`take_and_wake` path at all, unlike an `AtomicPtr<Waker>`
277/// swap-based design.
278#[repr(align(64))]
279pub struct AtomicWakerSlot {
280 state: AtomicUsize,
281 waker: std::cell::UnsafeCell<Option<Waker>>,
282}
283
284impl Default for AtomicWakerSlot {
285 fn default() -> Self {
286 Self::new()
287 }
288}
289
290impl AtomicWakerSlot {
291 /// An empty slot with no waker registered.
292 #[must_use]
293 pub const fn new() -> Self {
294 Self {
295 state: AtomicUsize::new(WAITING),
296 waker: std::cell::UnsafeCell::new(None),
297 }
298 }
299
300 /// Store `waker`, replacing whatever was previously registered (unless
301 /// it already [`Waker::will_wake`] the same task, in which case this is
302 /// a no-op beyond the CAS). If a [`Self::take_and_wake`] races with
303 /// this call and observes the slot mid-registration, this call takes
304 /// over waking `waker` itself before returning, so a delivery racing a
305 /// registration is never lost.
306 #[inline]
307 pub fn register(&self, waker: &Waker) {
308 match self.state.compare_exchange(
309 WAITING,
310 REGISTERING,
311 Ordering::Acquire,
312 Ordering::Acquire,
313 ) {
314 Ok(_) => {
315 // SAFETY: the CAS above is the sole gate on touching
316 // `waker` — we now hold the only `REGISTERING` token in
317 // existence for this slot (the CAS is exclusive), and
318 // `take_and_wake` never touches the cell when it observes
319 // anything other than `WAITING`. No other code path reads
320 // or writes the cell while we hold this state.
321 let slot = unsafe { &mut *self.waker.get() };
322 let already_registered = slot.as_ref().is_some_and(|w| w.will_wake(waker));
323 if !already_registered {
324 *slot = Some(waker.clone());
325 }
326 // Release back to WAITING — unless a concurrent
327 // `take_and_wake` set the `WAKING` bit while we were
328 // registering, in which case it deferred the actual wake
329 // to us (it saw we were mid-registration and couldn't
330 // safely touch the cell itself).
331 if self
332 .state
333 .compare_exchange(REGISTERING, WAITING, Ordering::AcqRel, Ordering::Acquire)
334 .is_err()
335 {
336 // SAFETY: same as above — we still hold exclusive
337 // access; the only other party that could have
338 // touched `state` here is `take_and_wake`, and it's
339 // structurally forbidden from touching the cell once
340 // it sees we're `REGISTERING`.
341 let pending = unsafe { &mut *self.waker.get() }.take();
342 self.state.store(WAITING, Ordering::Release);
343 if let Some(w) = pending {
344 w.wake();
345 }
346 }
347 }
348 Err(_) => {
349 // Someone else currently holds the slot (either
350 // `take_and_wake` is mid-delivery, or — for a listener
351 // type with a single-poller contract that a concurrent
352 // `register` would itself violate — an overlapping
353 // register). Either way, the safe fallback that can never
354 // lose a wakeup is to wake `waker` directly rather than
355 // storing it: worst case this causes one spurious extra
356 // poll, never a missed one.
357 waker.wake_by_ref();
358 }
359 }
360 }
361
362 /// Take whatever waker is registered (if any) and wake it.
363 #[inline]
364 pub fn take_and_wake(&self) {
365 match self
366 .state
367 .compare_exchange(WAITING, WAKING, Ordering::Acquire, Ordering::Acquire)
368 {
369 Ok(_) => {
370 // SAFETY: we hold the only `WAKING` token for this slot,
371 // and `register` never touches the cell once it observes
372 // anything other than `WAITING`, so this access is
373 // exclusive.
374 let waker = unsafe { &mut *self.waker.get() }.take();
375 self.state.store(WAITING, Ordering::Release);
376 if let Some(w) = waker {
377 w.wake();
378 }
379 }
380 Err(_) => {
381 // A `register` is currently in flight (state is
382 // `REGISTERING`, possibly with `WAKING` already OR'd in by
383 // an even earlier racing `take_and_wake` — either way, at
384 // most one delivery needs to be recorded, and ORing the
385 // bit in is enough to tell `register` "wake whatever you
386 // end up storing before you return", which it does. We
387 // don't touch the cell ourselves here — only whoever's
388 // currently `REGISTERING` is allowed to.
389 self.state.fetch_or(WAKING, Ordering::AcqRel);
390 }
391 }
392 }
393}
394
395impl Drop for AtomicWakerSlot {
396 fn drop(&mut self) {
397 // SAFETY: `&mut self` — no concurrent access is possible.
398 drop(self.waker.get_mut().take());
399 }
400}
401
402// SAFETY: `waker`'s `UnsafeCell` is only ever touched by whichever thread
403// currently holds the `REGISTERING` or `WAKING` token in `state` — the CAS
404// on `state` is the single, unambiguous point of truth for exclusive
405// access, so no two threads ever read or write the cell concurrently.
406unsafe impl Send for AtomicWakerSlot {}
407// SAFETY: same reasoning as `Send`.
408unsafe impl Sync for AtomicWakerSlot {}
409
410// =============================================================================
411// MpmcStack<T> — lock-free multi-producer multi-consumer Treiber stack
412// =============================================================================
413// Used as the cross-thread handoff for "many task threads submit ops, one
414// worker thread drains and issues them" (fs::uring_linux's SQE queue,
415// timer's per-bucket entry lists). Ordering within a bucket/batch is
416// irrelevant for both use sites, so a stack (LIFO) is as good as any other
417// MPMC structure and is the simplest one that's genuinely lock-free.
418/// Lock-free multi-producer multi-consumer stack (LIFO), used as a
419/// cross-thread handoff queue where ordering within a batch doesn't
420/// matter. See the module-level comment above for the intended use
421/// sites.
422#[repr(align(64))]
423pub struct MpmcStack<T> {
424 head: AtomicPtr<Node<T>>,
425 len: AtomicUsize,
426}
427
428#[repr(align(64))]
429struct Node<T> {
430 value: T,
431 next: *mut Self,
432}
433
434impl<T> Default for MpmcStack<T> {
435 fn default() -> Self {
436 Self::new()
437 }
438}
439
440impl<T> MpmcStack<T> {
441 /// An empty stack.
442 #[must_use]
443 pub const fn new() -> Self {
444 Self {
445 head: AtomicPtr::new(ptr::null_mut()),
446 len: AtomicUsize::new(0),
447 }
448 }
449
450 /// Push `value` onto the stack. Never blocks, never fails (heap
451 /// allocation aside).
452 #[inline]
453 pub fn push(&self, value: T) {
454 let node = Box::into_raw(Box::new(Node {
455 value,
456 next: ptr::null_mut(),
457 }));
458 let mut head = self.head.load(Ordering::Relaxed);
459 loop {
460 unsafe { (*node).next = head };
461 match self
462 .head
463 .compare_exchange_weak(head, node, Ordering::Release, Ordering::Relaxed)
464 {
465 Ok(_) => break,
466 Err(actual) => head = actual,
467 }
468 }
469 self.len.fetch_add(1, Ordering::Relaxed);
470 }
471
472 /// Pop the most-recently-pushed value, or `None` if empty.
473 #[inline]
474 pub fn pop(&self) -> Option<T> {
475 let mut head = self.head.load(Ordering::Acquire);
476 loop {
477 if head.is_null() {
478 return None;
479 }
480 let next = unsafe { (*head).next };
481 match self
482 .head
483 .compare_exchange_weak(head, next, Ordering::Acquire, Ordering::Acquire)
484 {
485 Ok(_) => {
486 self.len.fetch_sub(1, Ordering::Relaxed);
487 let boxed = unsafe { Box::from_raw(head) };
488 return Some(boxed.value);
489 }
490 Err(actual) => head = actual,
491 }
492 }
493 }
494
495 /// Atomically take the entire stack's contents as a `Vec`, leaving the
496 /// stack empty. O(1) swap of the head pointer plus an O(n) linked-list
497 /// walk to materialize the `Vec` — no CAS retries beyond the single
498 /// head swap regardless of `n`.
499 #[inline]
500 pub fn drain_all(&self) -> Vec<T> {
501 let mut head = self.head.swap(ptr::null_mut(), Ordering::AcqRel);
502 let mut out = Vec::new();
503 while !head.is_null() {
504 let boxed = unsafe { Box::from_raw(head) };
505 head = boxed.next;
506 out.push(boxed.value);
507 }
508 self.len.store(0, Ordering::Relaxed);
509 // LIFO push order means `drain_all` naturally yields
510 // most-recently-pushed-first; reverse so batches submit in
511 // roughly FIFO order (cosmetic — correctness doesn't depend on it).
512 out.reverse();
513 out
514 }
515
516 /// Whether the stack currently has no elements. Racy under concurrent
517 /// push/pop from other threads — a `true` result can be stale by the
518 /// time the caller acts on it, same as any lock-free length check.
519 #[inline(always)]
520 pub fn is_empty(&self) -> bool {
521 self.head.load(Ordering::Relaxed).is_null()
522 }
523
524 /// Approximate current length. Racy under concurrent push/pop for the
525 /// same reason as [`Self::is_empty`].
526 #[inline(always)]
527 pub fn len(&self) -> usize {
528 self.len.load(Ordering::Relaxed)
529 }
530
531 /// Zero-allocation drainage directly into the existing consumer buffer
532 #[allow(non_snake_case)]
533 #[inline]
534 pub fn drain_into_vec_deque(&self, target: &mut VecDeque<T>) {
535 // Collect from Treiber stack via atomic swap
536 let mut current = self.head.swap(std::ptr::null_mut(), Ordering::Acquire);
537
538 // Since a Treiber stack is LIFO, reversing elements as we walk the pointer chain
539 // yields correct FIFO ordering.
540 let mut prev = std::ptr::null_mut();
541 while !current.is_null() {
542 unsafe {
543 let next = (*current).next;
544 (*current).next = prev;
545 prev = current;
546 current = next;
547 }
548 }
549
550 // Push the correctly ordered items into the reused buffer capacity
551 let mut node = prev;
552 while !node.is_null() {
553 unsafe {
554 let BoxedNode = Box::from_raw(node);
555 target.push_back(BoxedNode.value);
556 node = BoxedNode.next;
557 }
558 }
559 }
560}
561
562// SAFETY: nodes are heap-boxed and moved between threads only via the
563// atomic `head` pointer swap (`push`/`pop`/`drain_all`); whichever thread
564// wins the CAS/swap has sole ownership of the node it took, so `T: Send`
565// is sufficient for the stack itself to be `Send`.
566unsafe impl<T: Send> Send for MpmcStack<T> {}
567// SAFETY: all mutation goes through the atomic `head`/`len` fields; no
568// two threads ever get concurrent mutable access to the same `Node<T>`.
569unsafe impl<T: Send> Sync for MpmcStack<T> {}
570
571impl<T> Drop for MpmcStack<T> {
572 fn drop(&mut self) {
573 while self.pop().is_some() {}
574 }
575}
576
577// =============================================================================
578// SpscQueue<T> — cache-aligned, lock-free single-producer/single-consumer
579// ring buffer
580// =============================================================================
581// Moved here from `io::native` (previously private) for the same reason as
582// `TreiberStack`/`BufferPool`: `stream`'s native duplex-pipe backend needs
583// exactly this shape — one writer, one reader, fixed capacity, no lock —
584// for each direction of a pipe, so it reuses this implementation rather
585// than hand-rolling its own ring buffer.
586//
587// No outer `repr(align(64))` — `head`/`tail` already each own a cache line
588// via `CacheAlignedUsize`, which is what actually matters for avoiding
589// false sharing between producer and consumer; aligning the container
590// itself only pads the start of `buffer` for no benefit.
591/// Cache-aligned, lock-free single-producer/single-consumer ring buffer.
592///
593/// Fixed power-of-two `capacity`. See the module-level comment above for
594/// why `head`/`tail` are each cache-line-aligned separately instead of
595/// aligning the whole struct.
596pub struct SpscQueue<T> {
597 head: CacheAlignedUsize,
598 tail: CacheAlignedUsize,
599 buffer: Box<[std::mem::MaybeUninit<T>]>,
600 capacity: usize,
601}
602
603#[repr(align(64))]
604struct CacheAlignedUsize {
605 value: AtomicUsize,
606}
607
608// SAFETY: exactly one producer thread ever calls `push` and exactly one
609// consumer thread ever calls `pop` (the SPSC contract callers must
610// uphold); `head`/`tail` atomics establish the happens-before edges that
611// make handing a `T` from the producer's `push` to the consumer's `pop`
612// sound, so `T: Send` is sufficient for the queue itself to be `Send`.
613unsafe impl<T: Send> Send for SpscQueue<T> {}
614// SAFETY: same SPSC contract as `Send` — no two threads ever access the
615// same slot concurrently, so sharing `&SpscQueue<T>` across threads
616// (one producer, one consumer) is sound whenever `T: Send`.
617unsafe impl<T: Send> Sync for SpscQueue<T> {}
618
619impl<T> SpscQueue<T> {
620 /// Build an empty queue. `capacity` must be a power of two.
621 ///
622 /// # Panics
623 ///
624 /// Panics if `capacity` is not a power of two.
625 #[must_use]
626 pub fn new(capacity: usize) -> Self {
627 assert!(capacity.is_power_of_two());
628 let mut buffer = Vec::with_capacity(capacity);
629 for _ in 0..capacity {
630 buffer.push(std::mem::MaybeUninit::uninit());
631 }
632 Self {
633 head: CacheAlignedUsize {
634 value: AtomicUsize::new(0),
635 },
636 tail: CacheAlignedUsize {
637 value: AtomicUsize::new(0),
638 },
639 buffer: buffer.into_boxed_slice(),
640 capacity,
641 }
642 }
643
644 /// Single-producer push. Returns `Err(value)` (handing the value back)
645 /// if the queue is full — never blocks.
646 ///
647 /// # Errors
648 ///
649 /// Returns `Err(value)`, handing the value straight back to the
650 /// caller, if the ring buffer is currently full (`tail - head` has
651 /// reached `capacity`). This is not a fault — it is the normal
652 /// backpressure signal for a bounded SPSC queue; the caller decides
653 /// whether to retry, drop the value, or apply its own backoff.
654 #[inline]
655 pub fn push(&self, value: T) -> Result<(), T> {
656 let tail = self.tail.value.load(Ordering::Relaxed);
657 let head = self.head.value.load(Ordering::Acquire);
658 if tail.wrapping_sub(head) == self.capacity {
659 return Err(value);
660 }
661 let mask = self.capacity - 1;
662 let idx = tail & mask;
663 unsafe {
664 let ptr = self.buffer[idx].as_ptr().cast_mut();
665 ptr.write(value);
666 }
667 self.tail
668 .value
669 .store(tail.wrapping_add(1), Ordering::Release);
670 Ok(())
671 }
672
673 /// Single-consumer pop. Returns `None` if the queue is empty.
674 #[inline]
675 pub fn pop(&self) -> Option<T> {
676 let head = self.head.value.load(Ordering::Relaxed);
677 let tail = self.tail.value.load(Ordering::Acquire);
678 if head == tail {
679 return None;
680 }
681 let mask = self.capacity - 1;
682 let idx = head & mask;
683 let value = unsafe {
684 let ptr = self.buffer[idx].as_ptr();
685 ptr.read()
686 };
687 self.head
688 .value
689 .store(head.wrapping_add(1), Ordering::Release);
690 Some(value)
691 }
692
693 /// Whether the queue currently has no elements.
694 #[inline(always)]
695 pub fn is_empty(&self) -> bool {
696 let head = self.head.value.load(Ordering::Relaxed);
697 let tail = self.tail.value.load(Ordering::Acquire);
698 head == tail
699 }
700
701 /// Whether the queue is currently at capacity (the next [`Self::push`]
702 /// would be rejected).
703 #[inline(always)]
704 pub fn is_full(&self) -> bool {
705 let tail = self.tail.value.load(Ordering::Relaxed);
706 let head = self.head.value.load(Ordering::Acquire);
707 tail.wrapping_sub(head) == self.capacity
708 }
709
710 /// Number of elements currently queued.
711 #[inline(always)]
712 pub fn len(&self) -> usize {
713 let head = self.head.value.load(Ordering::Relaxed);
714 let tail = self.tail.value.load(Ordering::Acquire);
715 tail.wrapping_sub(head)
716 }
717
718 /// The fixed capacity this queue was constructed with.
719 #[inline(always)]
720 pub const fn capacity(&self) -> usize {
721 self.capacity
722 }
723}
724
725impl<T: Copy> SpscQueue<T> {
726 /// Single-producer bulk push: copy as many elements from `src` as fit
727 /// into the ring, returning the count actually pushed (`0` if full).
728 ///
729 /// This is the bulk analogue of [`push`](Self::push) and exists purely
730 /// as a hot-path optimization for `T: Copy` element types (e.g. the
731 /// byte pipe in `stream::native`): pushing a run of N bytes one
732 /// [`push`](Self::push) call at a time pays N separate atomic
733 /// load/store pairs and N bounds-checked writes, whereas this issues at
734 /// most two `copy_nonoverlapping` runs (the ring may wrap once) and a
735 /// single `tail` store. Behaviour is otherwise identical — same
736 /// capacity accounting, same `Release` publication of the new tail.
737 #[inline]
738 pub fn push_slice(&self, src: &[T]) -> usize {
739 let tail = self.tail.value.load(Ordering::Relaxed);
740 let head = self.head.value.load(Ordering::Acquire);
741 let free = self.capacity - tail.wrapping_sub(head);
742 let to_push = free.min(src.len());
743 if to_push == 0 {
744 return 0;
745 }
746 let mask = self.capacity - 1;
747 let start = tail & mask;
748 // The ring may wrap: `first` covers the run from `start` to the end
749 // of the backing buffer, the remainder wraps to the front.
750 let first = to_push.min(self.capacity - start);
751 // SAFETY: `MaybeUninit<T>` has the same layout as `T`; `to_push`
752 // never exceeds the free capacity, so both runs stay in bounds and
753 // never overlap the consumer's occupied region.
754 unsafe {
755 let base = self.buffer.as_ptr().cast_mut().cast::<T>();
756 ptr::copy_nonoverlapping(src.as_ptr(), base.add(start), first);
757 if to_push > first {
758 ptr::copy_nonoverlapping(src.as_ptr().add(first), base, to_push - first);
759 }
760 }
761 self.tail
762 .value
763 .store(tail.wrapping_add(to_push), Ordering::Release);
764 to_push
765 }
766
767 /// Single-consumer bulk pop: copy up to `dst.len()` queued elements into
768 /// `dst`, returning the count actually popped (`0` if empty).
769 ///
770 /// Bulk analogue of [`pop`](Self::pop); see [`push_slice`](Self::push_slice)
771 /// for the rationale.
772 #[inline]
773 pub fn pop_slice(&self, dst: &mut [T]) -> usize {
774 let head = self.head.value.load(Ordering::Relaxed);
775 let tail = self.tail.value.load(Ordering::Acquire);
776 let avail = tail.wrapping_sub(head);
777 let to_pop = avail.min(dst.len());
778 if to_pop == 0 {
779 return 0;
780 }
781 let mask = self.capacity - 1;
782 let start = head & mask;
783 let first = to_pop.min(self.capacity - start);
784 // SAFETY: `to_pop` never exceeds the number of occupied slots, so
785 // every element read here was published by a prior `push`/`push_slice`
786 // (observed under the `Acquire` load of `tail` above) and is
787 // initialized; the two runs stay in bounds and never overlap the
788 // producer's free region.
789 unsafe {
790 let base = self.buffer.as_ptr().cast::<T>();
791 ptr::copy_nonoverlapping(base.add(start), dst.as_mut_ptr(), first);
792 if to_pop > first {
793 ptr::copy_nonoverlapping(base, dst.as_mut_ptr().add(first), to_pop - first);
794 }
795 }
796 self.head
797 .value
798 .store(head.wrapping_add(to_pop), Ordering::Release);
799 to_pop
800 }
801}
802
803impl<T> Drop for SpscQueue<T> {
804 fn drop(&mut self) {
805 while self.pop().is_some() {}
806 }
807}
808
809// =============================================================================
810// OnceSlot<T> — a single-fire, wait-free async result cell
811// =============================================================================
812// The generalization of the `PENDING`-sentinel-`AtomicI64` pattern
813// `fs::iocp_windows`/`fs::uring_linux` use, for ops whose result isn't a
814// plain integer (a `std::process::ExitStatus`, a `(usize, Vec<u8>)` read
815// result, etc). One `AtomicPtr<T>` starts null; `set` heap-boxes the value
816// and swaps it in; `poll` follows the same double-check-around-
817// waker-registration shape as every other completion primitive in this
818// module. Exactly one `set` call is ever expected per `OnceSlot` — calling
819// it twice is a caller bug, not something this type tries to paper over
820// (debug-asserted, not defended against in release).
821/// A single-fire, wait-free async result cell.
822///
823/// The generalization of the `PENDING`-sentinel-`AtomicI64` pattern used
824/// elsewhere in this crate, for results that aren't a plain integer. See
825/// the module doc above for the exactly-once `set` contract.
826#[repr(align(64))]
827pub struct OnceSlot<T> {
828 ptr: AtomicPtr<T>,
829 waker: AtomicWakerSlot,
830}
831
832impl<T> Default for OnceSlot<T> {
833 fn default() -> Self {
834 Self::new()
835 }
836}
837
838impl<T> OnceSlot<T> {
839 /// An empty, not-yet-completed slot.
840 #[must_use]
841 pub const fn new() -> Self {
842 Self {
843 ptr: AtomicPtr::new(ptr::null_mut()),
844 waker: AtomicWakerSlot::new(),
845 }
846 }
847
848 /// Complete this slot with `value`, waking whatever's polling it.
849 /// Must be called at most once per `OnceSlot`.
850 #[inline]
851 pub fn set(&self, value: T) {
852 let boxed = Box::into_raw(Box::new(value));
853 let prev = self.ptr.swap(boxed, Ordering::AcqRel);
854 debug_assert!(
855 prev.is_null(),
856 "OnceSlot::set called more than once — second value leaked"
857 );
858 self.waker.take_and_wake();
859 }
860
861 /// Poll for completion, registering `cx`'s waker if not yet complete.
862 ///
863 /// Both checks `swap` the pointer out (not just `load`) before
864 /// reconstructing the `Box` — polling again after an already-observed
865 /// `Ready` is not something callers are expected to do (standard
866 /// `Future` contract), but doing it anyway must not double-free, and
867 /// a `load`-then-`Box::from_raw` on the fast path would leave a
868 /// dangling non-null pointer behind for exactly that case.
869 #[inline]
870 pub fn poll(&self, cx: &Context<'_>) -> Poll<T> {
871 let mut p = self.ptr.load(Ordering::Acquire);
872
873 if !p.is_null() {
874 // Data is present! Safely isolate it by swapping it out
875 p = self.ptr.swap(ptr::null_mut(), Ordering::AcqRel);
876 if !p.is_null() {
877 return Poll::Ready(*unsafe { Box::from_raw(p) });
878 }
879 }
880
881 // Fallback path: register the waker
882 self.waker.register(cx.waker());
883
884 // Double-check after registration using a swap to prevent race conditions
885 p = self.ptr.swap(ptr::null_mut(), Ordering::AcqRel);
886 if !p.is_null() {
887 return Poll::Ready(*unsafe { Box::from_raw(p) });
888 }
889
890 Poll::Pending
891 }
892}
893
894impl<T> Drop for OnceSlot<T> {
895 fn drop(&mut self) {
896 let p = *self.ptr.get_mut();
897 if !p.is_null() {
898 drop(unsafe { Box::from_raw(p) });
899 }
900 }
901}
902
903// SAFETY: `set` heap-boxes the value and atomically swaps it into `ptr`;
904// `poll` atomically swaps it back out. Whichever side observes the
905// non-null pointer has sole ownership, so `T: Send` suffices for `Send`.
906unsafe impl<T: Send> Send for OnceSlot<T> {}
907// SAFETY: same reasoning as `Send` — every access to the boxed value goes
908// through an atomic swap, never a shared read of `&T` from two threads.
909unsafe impl<T: Send> Sync for OnceSlot<T> {}
910
911// =============================================================================
912// BoundedMpmcQueue<T> — lock-free bounded multi-producer/multi-consumer
913// FIFO queue (Dmitry Vyukov's classic bounded MPMC ring-buffer algorithm)
914// =============================================================================
915// Used by `sync::mpsc`'s bounded channel in place of a
916// `std::sync::Mutex<VecDeque<T>>`: that mutex, plus the register/wake
917// traffic it causes under backpressure, measured multiple times slower
918// than `tokio::sync::mpsc` under contention (see
919// `benches/sync_performance.rs`'s `mpsc_multi_producer_throughput`) — a
920// real, user-facing performance defect, not a micro-optimization. This
921// queue removes the mutex entirely: each slot carries its own `sequence`
922// counter, and producers/consumers claim a slot via a single CAS on a
923// shared position counter, retrying only on a lost race — no thread ever
924// blocks another out for the duration of a push/pop.
925//
926// This is the same well-known algorithm behind (e.g.) folly's
927// `MPMCQueue` and boost::lockfree::queue's bounded mode — not a novel
928// design — reimplemented here from the published algorithm rather than
929// vendored, to keep this crate dependency-free.
930/// Lock-free bounded multi-producer/multi-consumer FIFO queue of fixed
931/// `capacity`.
932///
933/// `try_push`/`try_pop` never block; they return `Err`/`None`
934/// immediately if the queue is full/empty rather than waiting — callers
935/// needing to wait layer their own backoff/parking on top (this is what
936/// `sync::mpsc` does with its `WaitQueue`-based registration).
937#[repr(C)]
938pub struct BoundedMpmcQueue<T> {
939 buffer: Box<[Cell<T>]>,
940 capacity: usize,
941 _pad1: [u64; 8],
942 enqueue_pos: AtomicUsize,
943 _pad2: [u64; 8],
944 dequeue_pos: AtomicUsize,
945 _pad3: [u64; 8],
946}
947
948#[repr(align(64))]
949struct Cell<T> {
950 /// See the algorithm's invariant in `try_push`/`try_pop`: a cell at
951 /// buffer index `i` cycles through `sequence` values `i`, `i+1`,
952 /// `i+1+capacity`, `i+1+2*capacity`, ... — "ready to be written for
953 /// enqueue position `i`", "ready to be read", "ready to be written
954 /// for enqueue position `i+capacity`", etc.
955 sequence: AtomicUsize,
956 data: std::cell::UnsafeCell<std::mem::MaybeUninit<T>>,
957}
958
959impl<T> BoundedMpmcQueue<T> {
960 /// Build an empty queue holding up to `capacity` elements (rounded up
961 /// to at least 1).
962 #[must_use]
963 pub fn new(capacity: usize) -> Self {
964 let capacity = capacity.max(1);
965 let mut buffer = Vec::with_capacity(capacity);
966 for i in 0..capacity {
967 buffer.push(Cell {
968 sequence: AtomicUsize::new(i),
969 data: std::cell::UnsafeCell::new(std::mem::MaybeUninit::uninit()),
970 });
971 }
972 Self {
973 buffer: buffer.into_boxed_slice(),
974 capacity,
975 _pad1: [0; 8],
976 enqueue_pos: AtomicUsize::new(0),
977 _pad2: [0; 8],
978 dequeue_pos: AtomicUsize::new(0),
979 _pad3: [0; 8],
980 }
981 }
982
983 /// The fixed capacity this queue was constructed with.
984 #[must_use]
985 pub const fn capacity(&self) -> usize {
986 self.capacity
987 }
988
989 /// Push `value`, returning it back in `Err` if the queue is currently
990 /// full. Never blocks.
991 ///
992 /// # Errors
993 /// Returns `value` back, unmodified, if the queue is at `capacity`.
994 #[inline]
995 pub fn try_push(&self, value: T) -> Result<(), T> {
996 let mut pos = self.enqueue_pos.load(Ordering::Relaxed);
997 loop {
998 let cell = &self.buffer[pos % self.capacity];
999 let seq = cell.sequence.load(Ordering::Acquire);
1000 #[allow(clippy::cast_possible_wrap)]
1001 let diff = seq as isize - pos as isize;
1002 match diff.cmp(&0) {
1003 std::cmp::Ordering::Equal => {
1004 // This cell is free for our position. Race every
1005 // other producer to claim it by advancing the shared
1006 // counter; on success we — and only we — own the
1007 // cell until we publish it below.
1008 match self.enqueue_pos.compare_exchange_weak(
1009 pos,
1010 pos.wrapping_add(1),
1011 Ordering::Relaxed,
1012 Ordering::Relaxed,
1013 ) {
1014 Ok(_) => {
1015 // SAFETY: winning the CAS above is exclusive
1016 // ownership of this cell until the `Release`
1017 // store below publishes it to a consumer.
1018 unsafe { (*cell.data.get()).write(value) };
1019 cell.sequence.store(pos.wrapping_add(1), Ordering::Release);
1020 return Ok(());
1021 }
1022 Err(actual) => pos = actual,
1023 }
1024 }
1025 std::cmp::Ordering::Less => return Err(value),
1026 std::cmp::Ordering::Greater => pos = self.enqueue_pos.load(Ordering::Relaxed),
1027 }
1028 }
1029 }
1030
1031 /// Pop the oldest pushed value, or `None` if the queue is currently
1032 /// empty. Never blocks.
1033 #[inline]
1034 pub fn try_pop(&self) -> Option<T> {
1035 let mut pos = self.dequeue_pos.load(Ordering::Relaxed);
1036 loop {
1037 let cell = &self.buffer[pos % self.capacity];
1038 let seq = cell.sequence.load(Ordering::Acquire);
1039 #[allow(clippy::cast_possible_wrap)]
1040 let diff = seq as isize - pos.wrapping_add(1) as isize;
1041 match diff.cmp(&0) {
1042 std::cmp::Ordering::Equal => {
1043 match self.dequeue_pos.compare_exchange_weak(
1044 pos,
1045 pos.wrapping_add(1),
1046 Ordering::Relaxed,
1047 Ordering::Relaxed,
1048 ) {
1049 Ok(_) => {
1050 // SAFETY: winning the CAS above is exclusive
1051 // ownership of this cell's current value —
1052 // the `Acquire` load of `sequence` paired
1053 // with the producer's `Release` store made
1054 // its write visible.
1055 let value = unsafe { (*cell.data.get()).assume_init_read() };
1056 cell.sequence
1057 .store(pos.wrapping_add(self.capacity), Ordering::Release);
1058 return Some(value);
1059 }
1060 Err(actual) => pos = actual,
1061 }
1062 }
1063 std::cmp::Ordering::Less => return None,
1064 std::cmp::Ordering::Greater => pos = self.dequeue_pos.load(Ordering::Relaxed),
1065 }
1066 }
1067 }
1068
1069 /// A snapshot of how many `try_push` calls have *claimed* a slot so
1070 /// far (via winning the `enqueue_pos` CAS), including any still
1071 /// in-flight (claimed but not yet published). Paired with
1072 /// [`Self::try_pop_until`] so a drain-everything caller can wait out
1073 /// an in-flight push rather than mistaking it for "queue empty".
1074 #[must_use]
1075 #[inline(always)]
1076 pub fn enqueue_snapshot(&self) -> usize {
1077 self.enqueue_pos.load(Ordering::Acquire)
1078 }
1079
1080 /// Pop everything pushed up to (not including) `target` (an
1081 /// [`Self::enqueue_snapshot`] taken earlier), spin-waiting out any
1082 /// push that's claimed a slot but not yet published rather than
1083 /// treating a transient empty read as "done".
1084 ///
1085 /// Ordinary [`Self::try_pop`] alone can't safely implement "drain
1086 /// every currently-registered entry": it returns `None` the instant
1087 /// dequeue catches up to a slot that's been *claimed* (CAS won) but
1088 /// not yet *published* (its value written and `sequence` released),
1089 /// which looks identical to a genuinely empty queue. A caller that
1090 /// stops draining on the first `None` can permanently miss that
1091 /// about-to-be-published entry if this was the last drain anyone
1092 /// will ever perform (see `sync::broadcast`'s module doc for the
1093 /// real, reproducible hang this caused via `WaitQueue::wake_all`).
1094 /// Looping until `target` is reached closes that window: any slot
1095 /// below `target` was claimed *before* this call started, so it can
1096 /// only be transiently unpublished, never permanently absent.
1097 #[inline(always)]
1098 pub fn drain_until(&self, target: usize, mut on_pop: impl FnMut(T)) {
1099 #[allow(clippy::cast_possible_wrap)]
1100 while (target.wrapping_sub(self.dequeue_pos.load(Ordering::Relaxed)) as isize) > 0 {
1101 match self.try_pop() {
1102 Some(value) => on_pop(value),
1103 None => std::hint::spin_loop(),
1104 }
1105 }
1106 }
1107
1108 /// Approximate current length — racy under concurrent push/pop, same
1109 /// caveat as every other lock-free length check in this module.
1110 #[must_use]
1111 #[inline(always)]
1112 pub fn len(&self) -> usize {
1113 let enq = self.enqueue_pos.load(Ordering::Relaxed);
1114 let deq = self.dequeue_pos.load(Ordering::Relaxed);
1115 enq.saturating_sub(deq)
1116 }
1117
1118 /// Whether the queue currently has no elements. Racy, same caveat as
1119 /// [`Self::len`].
1120 #[must_use]
1121 #[inline(always)]
1122 pub fn is_empty(&self) -> bool {
1123 self.len() == 0
1124 }
1125}
1126
1127impl<T> Drop for BoundedMpmcQueue<T> {
1128 fn drop(&mut self) {
1129 while self.try_pop().is_some() {}
1130 }
1131}
1132
1133// SAFETY: every `data` access is gated by first winning the corresponding
1134// `enqueue_pos`/`dequeue_pos` CAS for that specific cell — the `sequence`
1135// Acquire/Release pairing establishes happens-before between a
1136// producer's write and the consumer's read, and the CAS itself ensures
1137// no two producers (or two consumers) ever claim the same cell at the
1138// same generation. `T: Send` suffices since ownership transfers cleanly
1139// producer-thread -> consumer-thread.
1140unsafe impl<T: Send> Send for BoundedMpmcQueue<T> {}
1141// SAFETY: same reasoning as `Send`.
1142unsafe impl<T: Send> Sync for BoundedMpmcQueue<T> {}