subms_mpsc_queue/features/
mpmc.rs

1//! Multi-consumer extension: bounded MPMC ring with tail-sequence CAS.
2//!
3//! Disruptor-style barrier: per-slot sequence numbers gate both
4//! producer claim (CAS the tail) and consumer claim (CAS the head).
5//! Multiple consumers race; the loser sees a stale head and retries
6//! with the new value. Optional [`MpmcQueue::cas_retries`] counts
7//! contention for callers wiring it through the `metrics` feature.
8//!
9//! Both [`MpmcQueue::try_enqueue`] and [`MpmcQueue::try_dequeue`] are wait-free in the
10//! uncontended case and bounded-retry under contention.
11
12use std::cell::UnsafeCell;
13use std::mem::MaybeUninit;
14use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
15
16/// Bounded MPMC ring queue. Capacity is rounded up to the next power
17/// of two (minimum 2).
18pub struct MpmcQueue<T> {
19    mask: usize,
20    slots: Box<[Slot<T>]>,
21    head: AtomicUsize,
22    tail: AtomicUsize,
23    cas_retries: AtomicU64,
24}
25
26struct Slot<T> {
27    seq: AtomicUsize,
28    value: UnsafeCell<MaybeUninit<T>>,
29}
30
31unsafe impl<T: Send> Sync for MpmcQueue<T> {}
32unsafe impl<T: Send> Send for MpmcQueue<T> {}
33
34impl<T> MpmcQueue<T> {
35    pub fn new(capacity: usize) -> Self {
36        let cap = capacity.next_power_of_two().max(2);
37        let mut slots = Vec::with_capacity(cap);
38        for i in 0..cap {
39            slots.push(Slot {
40                seq: AtomicUsize::new(i),
41                value: UnsafeCell::new(MaybeUninit::uninit()),
42            });
43        }
44        Self {
45            mask: cap - 1,
46            slots: slots.into_boxed_slice(),
47            head: AtomicUsize::new(0),
48            tail: AtomicUsize::new(0),
49            cas_retries: AtomicU64::new(0),
50        }
51    }
52
53    pub fn capacity(&self) -> usize {
54        self.mask + 1
55    }
56
57    /// Monotonic count of slots ever claimed by producers.
58    pub fn producer_index(&self) -> usize {
59        self.tail.load(Ordering::Acquire)
60    }
61
62    /// Monotonic count of slots ever claimed by consumers.
63    pub fn consumer_index(&self) -> usize {
64        self.head.load(Ordering::Acquire)
65    }
66
67    /// Total CAS retries (both producers losing tail-CAS and consumers
68    /// losing head-CAS). Useful for diagnosing contention; ignored by
69    /// the hot path otherwise.
70    pub fn cas_retries(&self) -> u64 {
71        self.cas_retries.load(Ordering::Relaxed)
72    }
73
74    /// Multi-producer enqueue. Returns `Err(value)` if the ring is
75    /// full.
76    pub fn try_enqueue(&self, value: T) -> Result<(), T> {
77        let mut tail = self.tail.load(Ordering::Relaxed);
78        loop {
79            let slot = &self.slots[tail & self.mask];
80            let seq = slot.seq.load(Ordering::Acquire);
81            let diff = seq.wrapping_sub(tail) as isize;
82            if diff == 0 {
83                match self.tail.compare_exchange_weak(
84                    tail,
85                    tail.wrapping_add(1),
86                    Ordering::Relaxed,
87                    Ordering::Relaxed,
88                ) {
89                    Ok(_) => {
90                        unsafe { (*slot.value.get()).write(value) };
91                        slot.seq.store(tail.wrapping_add(1), Ordering::Release);
92                        return Ok(());
93                    }
94                    Err(t) => {
95                        self.cas_retries.fetch_add(1, Ordering::Relaxed);
96                        tail = t;
97                    }
98                }
99            } else if diff < 0 {
100                return Err(value);
101            } else {
102                tail = self.tail.load(Ordering::Relaxed);
103            }
104        }
105    }
106
107    /// Multi-consumer dequeue. Returns `None` if the ring is empty.
108    pub fn try_dequeue(&self) -> Option<T> {
109        let mut head = self.head.load(Ordering::Relaxed);
110        loop {
111            let slot = &self.slots[head & self.mask];
112            let seq = slot.seq.load(Ordering::Acquire);
113            let diff = seq.wrapping_sub(head.wrapping_add(1)) as isize;
114            if diff == 0 {
115                match self.head.compare_exchange_weak(
116                    head,
117                    head.wrapping_add(1),
118                    Ordering::Relaxed,
119                    Ordering::Relaxed,
120                ) {
121                    Ok(_) => {
122                        let value = unsafe { (*slot.value.get()).assume_init_read() };
123                        slot.seq
124                            .store(head.wrapping_add(self.mask + 1), Ordering::Release);
125                        return Some(value);
126                    }
127                    Err(h) => {
128                        self.cas_retries.fetch_add(1, Ordering::Relaxed);
129                        head = h;
130                    }
131                }
132            } else if diff < 0 {
133                return None;
134            } else {
135                head = self.head.load(Ordering::Relaxed);
136            }
137        }
138    }
139
140    /// Drop everything currently readable and return the count. Any consumer
141    /// may call it, and other consumers keep draining alongside, so the count
142    /// is this caller's share rather than the queue's total.
143    pub fn clear(&self) -> usize {
144        let mut n = 0;
145        while self.try_dequeue().is_some() {
146            n += 1;
147        }
148        n
149    }
150
151    /// Approximate length.
152    pub fn len(&self) -> usize {
153        let h = self.head.load(Ordering::Acquire);
154        let t = self.tail.load(Ordering::Acquire);
155        t.wrapping_sub(h)
156    }
157
158    pub fn is_empty(&self) -> bool {
159        self.len() == 0
160    }
161
162    /// Best-effort fullness. Stale the instant any consumer drains a slot.
163    pub fn is_full(&self) -> bool {
164        self.len() >= self.capacity()
165    }
166}
167
168impl<T> Drop for MpmcQueue<T> {
169    fn drop(&mut self) {
170        self.clear();
171    }
172}
173
174#[cfg(test)]
175#[path = "mpmc_tests.rs"]
176mod tests;
subms_mpsc_queue/features/mpmc.rs

subms_mpsc_queue/features/
mpmc.rs
