doom_fish_utils/
spsc.rs

1//! Single-producer single-consumer lock-free bounded ring buffer.
2//!
3//! Designed for the `CoreAudio` render-thread → async-consumer producer-consumer
4//! pattern. The ring itself never takes a mutex and never allocates after it
5//! has been constructed. Pushing an item or dropping the producer wakes the
6//! consumer's registered waker on the producer thread, though, and that runs
7//! executor code that may lock, allocate or make a system call. If the
8//! producer must stay lock-free, have the consumer poll with
9//! [`SpscConsumer::pop`] rather than await [`SpscConsumer::pop_async`], so that
10//! no waker is registered.
11//!
12//! Internally this wrapper uses a pre-allocated bounded queue plus an
13//! [`AtomicWaker`] so the consumer can await
14//! the next item without forcing the producer to block.
15//!
16//! # Example
17//!
18//! ```no_run
19//! use doom_fish_utils::spsc::SpscRing;
20//!
21//! # async fn run() {
22//! let (producer, consumer) = SpscRing::<u32, 256>::new();
23//!
24//! producer.push(1).unwrap();
25//! assert_eq!(consumer.pop_async().await, Some(1));
26//! # }
27//! ```
28
29use std::cell::UnsafeCell;
30use std::fmt;
31use std::future::Future;
32use std::marker::PhantomData;
33use std::mem::MaybeUninit;
34use std::pin::Pin;
35use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
36use std::sync::Arc;
37use std::task::{Context, Poll};
38
39use futures_util::task::AtomicWaker;
40
41struct Slot<T> {
42    stamp: AtomicUsize,
43    value: UnsafeCell<MaybeUninit<T>>,
44}
45
46#[repr(align(128))]
47struct Cursor(AtomicUsize);
48
49struct Ring<T> {
50    slots: Box<[Slot<T>]>,
51    lap: usize,
52    read: Cursor,
53    write: Cursor,
54}
55
56unsafe impl<T: Send> Sync for Ring<T> {}
57
58impl<T> Ring<T> {
59    fn new(capacity: usize) -> Self {
60        let slots = (0..capacity)
61            .map(|index| Slot {
62                stamp: AtomicUsize::new(index * 2),
63                value: UnsafeCell::new(MaybeUninit::uninit()),
64            })
65            .collect();
66
67        Self {
68            slots,
69            lap: capacity.next_power_of_two() * 2,
70            read: Cursor(AtomicUsize::new(0)),
71            write: Cursor(AtomicUsize::new(0)),
72        }
73    }
74
75    fn next(&self, position: usize) -> usize {
76        if ((position & (self.lap - 1)) >> 1) + 1 == self.slots.len() {
77            (position & !(self.lap - 1)).wrapping_add(self.lap)
78        } else {
79            position.wrapping_add(2)
80        }
81    }
82
83    fn push(&self, item: T, overwrite: bool) -> Result<Option<T>, T> {
84        let mut position = self.write.0.load(Ordering::SeqCst);
85
86        loop {
87            let slot = &self.slots[(position & (self.lap - 1)) >> 1];
88            let stamp = slot.stamp.load(Ordering::Acquire);
89            let next = self.next(position);
90
91            if stamp == position {
92                if self
93                    .write
94                    .0
95                    .compare_exchange_weak(position, next, Ordering::SeqCst, Ordering::SeqCst)
96                    .is_ok()
97                {
98                    unsafe { (*slot.value.get()).write(item) };
99                    slot.stamp
100                        .store(position.wrapping_add(1), Ordering::Release);
101                    return Ok(None);
102                }
103            } else {
104                let oldest = position.wrapping_sub(self.lap);
105                if stamp == oldest.wrapping_add(1) && self.read.0.load(Ordering::SeqCst) == oldest {
106                    if !overwrite {
107                        return Err(item);
108                    }
109                    if self
110                        .read
111                        .0
112                        .compare_exchange(
113                            oldest,
114                            self.next(oldest),
115                            Ordering::SeqCst,
116                            Ordering::SeqCst,
117                        )
118                        .is_ok()
119                    {
120                        self.write.0.store(next, Ordering::SeqCst);
121                        let displaced = unsafe { (*slot.value.get()).as_mut_ptr().replace(item) };
122                        slot.stamp
123                            .store(position.wrapping_add(1), Ordering::Release);
124                        return Ok(Some(displaced));
125                    }
126                }
127            }
128
129            std::hint::spin_loop();
130            position = self.write.0.load(Ordering::SeqCst);
131        }
132    }
133
134    fn pop(&self) -> Option<T> {
135        let mut position = self.read.0.load(Ordering::SeqCst);
136
137        loop {
138            let slot = &self.slots[(position & (self.lap - 1)) >> 1];
139            let stamp = slot.stamp.load(Ordering::Acquire);
140            if stamp == position.wrapping_add(1) {
141                if self
142                    .read
143                    .0
144                    .compare_exchange_weak(
145                        position,
146                        self.next(position),
147                        Ordering::SeqCst,
148                        Ordering::SeqCst,
149                    )
150                    .is_ok()
151                {
152                    let item = unsafe { (*slot.value.get()).assume_init_read() };
153                    slot.stamp
154                        .store(position.wrapping_add(self.lap), Ordering::Release);
155                    return Some(item);
156                }
157            } else if stamp == position && self.write.0.load(Ordering::SeqCst) == position {
158                return None;
159            }
160
161            std::hint::spin_loop();
162            position = self.read.0.load(Ordering::SeqCst);
163        }
164    }
165
166    fn len(&self) -> usize {
167        loop {
168            let read = self.read.0.load(Ordering::SeqCst);
169            let write = self.write.0.load(Ordering::SeqCst);
170            if read == self.read.0.load(Ordering::SeqCst) {
171                let read_index = (read & (self.lap - 1)) >> 1;
172                let write_index = (write & (self.lap - 1)) >> 1;
173                return if read == write {
174                    0
175                } else if write_index > read_index {
176                    write_index - read_index
177                } else {
178                    self.slots.len() - read_index + write_index
179                };
180            }
181            std::hint::spin_loop();
182        }
183    }
184}
185
186impl<T> Drop for Ring<T> {
187    fn drop(&mut self) {
188        while let Some(item) = self.pop() {
189            drop(item);
190        }
191    }
192}
193
194struct Inner<T> {
195    queue: Ring<T>,
196    producer_closed: AtomicBool,
197    waker: AtomicWaker,
198}
199
200/// Constructor namespace for a bounded single-producer single-consumer ring.
201///
202/// `N` is the maximum supported capacity. Use [`Self::new`] to allocate a ring
203/// with exactly `N` slots, or [`Self::with_capacity`] to choose a smaller
204/// runtime capacity while keeping the type-level upper bound.
205#[derive(Debug, Default)]
206pub struct SpscRing<T, const N: usize>(PhantomData<T>);
207
208/// Producer half of an [`SpscRing`].
209pub struct SpscProducer<T, const N: usize> {
210    inner: Arc<Inner<T>>,
211}
212
213/// Consumer half of an [`SpscRing`].
214pub struct SpscConsumer<T, const N: usize> {
215    inner: Arc<Inner<T>>,
216}
217
218/// Future returned by [`SpscConsumer::pop_async`].
219#[must_use = "futures do nothing unless awaited or polled"]
220pub struct PopFuture<'a, T, const N: usize> {
221    consumer: &'a SpscConsumer<T, N>,
222}
223
224/// Feature-gated [`futures_core::Stream`] wrapper around an [`SpscConsumer`].
225#[cfg(feature = "futures-stream")]
226#[cfg_attr(docsrs, doc(cfg(feature = "futures-stream")))]
227#[must_use = "streams do nothing unless polled"]
228pub struct SpscConsumerStream<'a, T, const N: usize> {
229    consumer: &'a SpscConsumer<T, N>,
230}
231
232#[allow(clippy::new_ret_no_self)]
233impl<T, const N: usize> SpscRing<T, N> {
234    /// Creates a ring with capacity `N`.
235    ///
236    /// # Panics
237    ///
238    /// Panics if `N` is 0.
239    #[must_use]
240    pub fn new() -> (SpscProducer<T, N>, SpscConsumer<T, N>) {
241        Self::with_capacity(N)
242    }
243
244    /// Creates a ring with a runtime capacity up to the type-level maximum `N`.
245    ///
246    /// # Panics
247    ///
248    /// Panics if `capacity` is 0 or larger than `N`.
249    #[must_use]
250    pub fn with_capacity(capacity: usize) -> (SpscProducer<T, N>, SpscConsumer<T, N>) {
251        assert!(N > 0, "SpscRing capacity must be > 0");
252        assert!(capacity > 0, "SpscRing capacity must be > 0");
253        assert!(
254            capacity <= N,
255            "SpscRing capacity {capacity} exceeds type maximum {N}"
256        );
257
258        let inner = Arc::new(Inner {
259            queue: Ring::new(capacity),
260            producer_closed: AtomicBool::new(false),
261            waker: AtomicWaker::new(),
262        });
263
264        (
265            SpscProducer {
266                inner: Arc::clone(&inner),
267            },
268            SpscConsumer { inner },
269        )
270    }
271}
272
273impl<T, const N: usize> fmt::Debug for SpscProducer<T, N> {
274    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
275        f.debug_struct("SpscProducer")
276            .field("buffered", &self.buffered_count())
277            .field("capacity", &self.capacity())
278            .finish_non_exhaustive()
279    }
280}
281
282impl<T, const N: usize> fmt::Debug for SpscConsumer<T, N> {
283    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
284        f.debug_struct("SpscConsumer")
285            .field("buffered", &self.buffered_count())
286            .field("capacity", &self.capacity())
287            .field("is_closed", &self.is_closed())
288            .finish_non_exhaustive()
289    }
290}
291
292impl<T, const N: usize> fmt::Debug for PopFuture<'_, T, N> {
293    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
294        f.debug_struct("PopFuture").finish_non_exhaustive()
295    }
296}
297
298#[cfg(feature = "futures-stream")]
299impl<T, const N: usize> fmt::Debug for SpscConsumerStream<'_, T, N> {
300    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
301        f.debug_struct("SpscConsumerStream").finish_non_exhaustive()
302    }
303}
304
305impl<T, const N: usize> SpscProducer<T, N> {
306    /// Attempts to push an item into the ring without blocking.
307    ///
308    /// # Errors
309    ///
310    /// Returns `Err(item)` if the ring is currently full.
311    pub fn push(&self, item: T) -> Result<(), T> {
312        match self.inner.queue.push(item, false) {
313            Ok(_) => {
314                self.inner.waker.wake();
315                Ok(())
316            }
317            Err(item) => Err(item),
318        }
319    }
320
321    /// Pushes an item into the ring, overwriting the oldest buffered entry if
322    /// necessary.
323    ///
324    /// Returns the displaced oldest item when an overwrite happens.
325    pub fn push_overwrite(&self, item: T) -> Option<T> {
326        let dropped = self
327            .inner
328            .queue
329            .push(item, true)
330            .unwrap_or_else(|_| unreachable!());
331        self.inner.waker.wake();
332        dropped
333    }
334
335    /// Returns the current buffered item count.
336    #[must_use]
337    pub fn buffered_count(&self) -> usize {
338        self.inner.queue.len()
339    }
340
341    /// Returns the runtime capacity of the ring.
342    #[must_use]
343    pub fn capacity(&self) -> usize {
344        self.inner.queue.slots.len()
345    }
346}
347
348impl<T, const N: usize> Drop for SpscProducer<T, N> {
349    fn drop(&mut self) {
350        self.inner.producer_closed.store(true, Ordering::Release);
351        self.inner.waker.wake();
352    }
353}
354
355impl<T, const N: usize> SpscConsumer<T, N> {
356    /// Attempts to pop the next buffered item without blocking.
357    #[must_use]
358    pub fn pop(&self) -> Option<T> {
359        self.inner.queue.pop()
360    }
361
362    /// Returns a future that resolves to the next buffered item, or `None` once
363    /// the producer has been dropped and the ring is empty.
364    pub const fn pop_async(&self) -> PopFuture<'_, T, N> {
365        PopFuture { consumer: self }
366    }
367
368    /// Returns the current buffered item count.
369    #[must_use]
370    pub fn buffered_count(&self) -> usize {
371        self.inner.queue.len()
372    }
373
374    /// Returns the runtime capacity of the ring.
375    #[must_use]
376    pub fn capacity(&self) -> usize {
377        self.inner.queue.slots.len()
378    }
379
380    /// Returns `true` if the producer has been dropped.
381    #[must_use]
382    pub fn is_closed(&self) -> bool {
383        self.inner.producer_closed.load(Ordering::Acquire)
384    }
385
386    #[cfg(feature = "futures-stream")]
387    #[cfg_attr(docsrs, doc(cfg(feature = "futures-stream")))]
388    pub const fn stream(&self) -> SpscConsumerStream<'_, T, N> {
389        SpscConsumerStream { consumer: self }
390    }
391
392    fn poll_pop(&self, cx: &Context<'_>) -> Poll<Option<T>> {
393        if let Some(item) = self.pop() {
394            return Poll::Ready(Some(item));
395        }
396
397        if self.is_closed() {
398            return Poll::Ready(self.pop());
399        }
400
401        self.inner.waker.register(cx.waker());
402
403        if let Some(item) = self.pop() {
404            return Poll::Ready(Some(item));
405        }
406
407        if self.is_closed() {
408            return Poll::Ready(self.pop());
409        }
410
411        Poll::Pending
412    }
413}
414
415impl<T, const N: usize> Future for PopFuture<'_, T, N> {
416    type Output = Option<T>;
417
418    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
419        self.consumer.poll_pop(cx)
420    }
421}
422
423#[cfg(feature = "futures-stream")]
424impl<T, const N: usize> futures_core::Stream for SpscConsumerStream<'_, T, N> {
425    type Item = T;
426
427    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
428        self.consumer.poll_pop(cx)
429    }
430}
431
432#[cfg(test)]
433mod tests {
434    #[cfg(feature = "futures-stream")]
435    use std::future::poll_fn;
436    #[cfg(feature = "futures-stream")]
437    use std::pin::Pin;
438    use std::sync::{mpsc, Arc};
439    use std::task::{Context, Poll, Wake, Waker};
440    use std::thread;
441    use std::time::{Duration, Instant};
442
443    use super::{SpscConsumer, SpscProducer, SpscRing};
444
445    #[test]
446    fn preserves_sequence_in_single_thread() {
447        let (producer, consumer) = SpscRing::<u32, 4>::new();
448
449        assert_eq!(producer.push(1), Ok(()));
450        assert_eq!(producer.push(2), Ok(()));
451        assert_eq!(producer.push(3), Ok(()));
452
453        assert_eq!(consumer.pop(), Some(1));
454        assert_eq!(consumer.pop(), Some(2));
455        assert_eq!(consumer.pop(), Some(3));
456        assert_eq!(consumer.pop(), None);
457    }
458
459    #[test]
460    fn overwrite_drops_oldest_item() {
461        let (producer, consumer) = SpscRing::<u32, 2>::new();
462
463        assert_eq!(producer.push_overwrite(10), None);
464        assert_eq!(producer.push_overwrite(20), None);
465        assert_eq!(producer.push_overwrite(30), Some(10));
466
467        assert_eq!(consumer.pop(), Some(20));
468        assert_eq!(consumer.pop(), Some(30));
469        assert_eq!(consumer.pop(), None);
470    }
471
472    #[test]
473    fn producer_calls_return_immediately_when_full() {
474        let (producer, _consumer) = SpscRing::<u64, 1>::new();
475        assert_eq!(producer.push(7), Ok(()));
476
477        let start = Instant::now();
478        let mut expected_drop = Some(7);
479        for value in 0..100_000 {
480            assert_eq!(producer.push(value), Err(value));
481            assert_eq!(producer.push_overwrite(value), expected_drop);
482            expected_drop = Some(value);
483        }
484
485        assert!(
486            start.elapsed() < Duration::from_secs(2),
487            "producer operations took too long while the ring stayed full"
488        );
489    }
490
491    #[test]
492    fn pop_async_drains_then_closes() {
493        let (producer, consumer) = SpscRing::<u32, 8>::new();
494        producer.push(1).unwrap();
495        producer.push(2).unwrap();
496        drop(producer);
497
498        assert_eq!(pollster::block_on(consumer.pop_async()), Some(1));
499        assert_eq!(pollster::block_on(consumer.pop_async()), Some(2));
500        assert_eq!(pollster::block_on(consumer.pop_async()), None);
501    }
502
503    #[test]
504    fn concurrent_producer_consumer_preserve_order() {
505        let (producer, consumer) = SpscRing::<u64, 1024>::new();
506        let producer_thread = thread::spawn(move || {
507            for expected in 0..50_000_u64 {
508                let mut item = expected;
509                loop {
510                    match producer.push(item) {
511                        Ok(()) => break,
512                        Err(returned) => {
513                            item = returned;
514                            std::hint::spin_loop();
515                        }
516                    }
517                }
518            }
519        });
520
521        for expected in 0..50_000_u64 {
522            let actual = pollster::block_on(consumer.pop_async());
523            assert_eq!(actual, Some(expected));
524        }
525        assert_eq!(pollster::block_on(consumer.pop_async()), None);
526
527        producer_thread.join().unwrap();
528    }
529
530    struct NoopWake;
531
532    impl Wake for NoopWake {
533        fn wake(self: Arc<Self>) {}
534    }
535
536    fn spin_pop(consumer: &SpscConsumer<u32, 1>, cx: &Context<'_>) -> Option<u32> {
537        loop {
538            if let Poll::Ready(item) = consumer.poll_pop(cx) {
539                return item;
540            }
541            std::hint::spin_loop();
542        }
543    }
544
545    #[test]
546    fn item_pushed_just_before_close_is_delivered() {
547        let (producers_tx, producers_rx) = mpsc::channel::<SpscProducer<u32, 1>>();
548        let worker = thread::spawn(move || {
549            for producer in producers_rx {
550                producer.push(7).unwrap();
551                drop(producer);
552            }
553        });
554        let waker = Waker::from(Arc::new(NoopWake));
555        let cx = Context::from_waker(&waker);
556
557        for _ in 0..20_000 {
558            let (producer, consumer) = SpscRing::<u32, 1>::new();
559            producers_tx.send(producer).unwrap();
560            assert_eq!(spin_pop(&consumer, &cx), Some(7));
561            assert_eq!(spin_pop(&consumer, &cx), None);
562        }
563
564        drop(producers_tx);
565        worker.join().unwrap();
566    }
567
568    #[cfg(feature = "futures-stream")]
569    #[test]
570    fn stream_wrapper_yields_items() {
571        use futures_core::Stream;
572
573        let (producer, consumer) = SpscRing::<u32, 4>::new();
574        let mut stream = consumer.stream();
575
576        producer.push(11).unwrap();
577        drop(producer);
578
579        let first = pollster::block_on(poll_fn(|cx| Pin::new(&mut stream).poll_next(cx)));
580        let second = pollster::block_on(poll_fn(|cx| Pin::new(&mut stream).poll_next(cx)));
581
582        assert_eq!(first, Some(11));
583        assert_eq!(second, None);
584    }
585}
doom_fish_utils/spsc.rs

doom_fish_utils/
spsc.rs
