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::fmt;
30use std::future::Future;
31use std::marker::PhantomData;
32use std::pin::Pin;
33use std::sync::atomic::{AtomicBool, Ordering};
34use std::sync::Arc;
35use std::task::{Context, Poll};
36
37use crossbeam_queue::ArrayQueue;
38use futures_util::task::AtomicWaker;
39
40struct Inner<T> {
41    queue: ArrayQueue<T>,
42    producer_closed: AtomicBool,
43    waker: AtomicWaker,
44}
45
46/// Constructor namespace for a bounded single-producer single-consumer ring.
47///
48/// `N` is the maximum supported capacity. Use [`Self::new`] to allocate a ring
49/// with exactly `N` slots, or [`Self::with_capacity`] to choose a smaller
50/// runtime capacity while keeping the type-level upper bound.
51#[derive(Debug, Default)]
52pub struct SpscRing<T, const N: usize>(PhantomData<T>);
53
54/// Producer half of an [`SpscRing`].
55pub struct SpscProducer<T, const N: usize> {
56    inner: Arc<Inner<T>>,
57}
58
59/// Consumer half of an [`SpscRing`].
60pub struct SpscConsumer<T, const N: usize> {
61    inner: Arc<Inner<T>>,
62}
63
64/// Future returned by [`SpscConsumer::pop_async`].
65#[must_use = "futures do nothing unless awaited or polled"]
66pub struct PopFuture<'a, T, const N: usize> {
67    consumer: &'a SpscConsumer<T, N>,
68}
69
70/// Feature-gated [`futures_core::Stream`] wrapper around an [`SpscConsumer`].
71#[cfg(feature = "futures-stream")]
72#[cfg_attr(docsrs, doc(cfg(feature = "futures-stream")))]
73#[must_use = "streams do nothing unless polled"]
74pub struct SpscConsumerStream<'a, T, const N: usize> {
75    consumer: &'a SpscConsumer<T, N>,
76}
77
78#[allow(clippy::new_ret_no_self)]
79impl<T, const N: usize> SpscRing<T, N> {
80    /// Creates a ring with capacity `N`.
81    ///
82    /// # Panics
83    ///
84    /// Panics if `N` is 0.
85    #[must_use]
86    pub fn new() -> (SpscProducer<T, N>, SpscConsumer<T, N>) {
87        Self::with_capacity(N)
88    }
89
90    /// Creates a ring with a runtime capacity up to the type-level maximum `N`.
91    ///
92    /// # Panics
93    ///
94    /// Panics if `capacity` is 0 or larger than `N`.
95    #[must_use]
96    pub fn with_capacity(capacity: usize) -> (SpscProducer<T, N>, SpscConsumer<T, N>) {
97        assert!(N > 0, "SpscRing capacity must be > 0");
98        assert!(capacity > 0, "SpscRing capacity must be > 0");
99        assert!(
100            capacity <= N,
101            "SpscRing capacity {capacity} exceeds type maximum {N}"
102        );
103
104        let inner = Arc::new(Inner {
105            queue: ArrayQueue::new(capacity),
106            producer_closed: AtomicBool::new(false),
107            waker: AtomicWaker::new(),
108        });
109
110        (
111            SpscProducer {
112                inner: Arc::clone(&inner),
113            },
114            SpscConsumer { inner },
115        )
116    }
117}
118
119impl<T, const N: usize> fmt::Debug for SpscProducer<T, N> {
120    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
121        f.debug_struct("SpscProducer")
122            .field("buffered", &self.buffered_count())
123            .field("capacity", &self.capacity())
124            .finish_non_exhaustive()
125    }
126}
127
128impl<T, const N: usize> fmt::Debug for SpscConsumer<T, N> {
129    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
130        f.debug_struct("SpscConsumer")
131            .field("buffered", &self.buffered_count())
132            .field("capacity", &self.capacity())
133            .field("is_closed", &self.is_closed())
134            .finish_non_exhaustive()
135    }
136}
137
138impl<T, const N: usize> fmt::Debug for PopFuture<'_, T, N> {
139    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
140        f.debug_struct("PopFuture").finish_non_exhaustive()
141    }
142}
143
144#[cfg(feature = "futures-stream")]
145impl<T, const N: usize> fmt::Debug for SpscConsumerStream<'_, T, N> {
146    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
147        f.debug_struct("SpscConsumerStream").finish_non_exhaustive()
148    }
149}
150
151impl<T, const N: usize> SpscProducer<T, N> {
152    /// Attempts to push an item into the ring without blocking.
153    ///
154    /// # Errors
155    ///
156    /// Returns `Err(item)` if the ring is currently full.
157    pub fn push(&self, item: T) -> Result<(), T> {
158        match self.inner.queue.push(item) {
159            Ok(()) => {
160                self.inner.waker.wake();
161                Ok(())
162            }
163            Err(item) => Err(item),
164        }
165    }
166
167    /// Pushes an item into the ring, overwriting the oldest buffered entry if
168    /// necessary.
169    ///
170    /// Returns the displaced oldest item when an overwrite happens.
171    pub fn push_overwrite(&self, item: T) -> Option<T> {
172        let dropped = self.inner.queue.force_push(item);
173        self.inner.waker.wake();
174        dropped
175    }
176
177    /// Returns the current buffered item count.
178    #[must_use]
179    pub fn buffered_count(&self) -> usize {
180        self.inner.queue.len()
181    }
182
183    /// Returns the runtime capacity of the ring.
184    #[must_use]
185    pub fn capacity(&self) -> usize {
186        self.inner.queue.capacity()
187    }
188}
189
190impl<T, const N: usize> Drop for SpscProducer<T, N> {
191    fn drop(&mut self) {
192        self.inner.producer_closed.store(true, Ordering::Release);
193        self.inner.waker.wake();
194    }
195}
196
197impl<T, const N: usize> SpscConsumer<T, N> {
198    /// Attempts to pop the next buffered item without blocking.
199    #[must_use]
200    pub fn pop(&self) -> Option<T> {
201        self.inner.queue.pop()
202    }
203
204    /// Returns a future that resolves to the next buffered item, or `None` once
205    /// the producer has been dropped and the ring is empty.
206    pub const fn pop_async(&self) -> PopFuture<'_, T, N> {
207        PopFuture { consumer: self }
208    }
209
210    /// Returns the current buffered item count.
211    #[must_use]
212    pub fn buffered_count(&self) -> usize {
213        self.inner.queue.len()
214    }
215
216    /// Returns the runtime capacity of the ring.
217    #[must_use]
218    pub fn capacity(&self) -> usize {
219        self.inner.queue.capacity()
220    }
221
222    /// Returns `true` if the producer has been dropped.
223    #[must_use]
224    pub fn is_closed(&self) -> bool {
225        self.inner.producer_closed.load(Ordering::Acquire)
226    }
227
228    #[cfg(feature = "futures-stream")]
229    #[cfg_attr(docsrs, doc(cfg(feature = "futures-stream")))]
230    pub const fn stream(&self) -> SpscConsumerStream<'_, T, N> {
231        SpscConsumerStream { consumer: self }
232    }
233
234    fn poll_pop(&self, cx: &Context<'_>) -> Poll<Option<T>> {
235        if let Some(item) = self.pop() {
236            return Poll::Ready(Some(item));
237        }
238
239        if self.is_closed() {
240            return Poll::Ready(self.pop());
241        }
242
243        self.inner.waker.register(cx.waker());
244
245        if let Some(item) = self.pop() {
246            return Poll::Ready(Some(item));
247        }
248
249        if self.is_closed() {
250            return Poll::Ready(self.pop());
251        }
252
253        Poll::Pending
254    }
255}
256
257impl<T, const N: usize> Future for PopFuture<'_, T, N> {
258    type Output = Option<T>;
259
260    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
261        self.consumer.poll_pop(cx)
262    }
263}
264
265#[cfg(feature = "futures-stream")]
266impl<T, const N: usize> futures_core::Stream for SpscConsumerStream<'_, T, N> {
267    type Item = T;
268
269    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
270        self.consumer.poll_pop(cx)
271    }
272}
273
274#[cfg(test)]
275mod tests {
276    #[cfg(feature = "futures-stream")]
277    use std::future::poll_fn;
278    #[cfg(feature = "futures-stream")]
279    use std::pin::Pin;
280    use std::sync::{mpsc, Arc};
281    use std::task::{Context, Poll, Wake, Waker};
282    use std::thread;
283    use std::time::{Duration, Instant};
284
285    use super::{SpscConsumer, SpscProducer, SpscRing};
286
287    #[test]
288    fn preserves_sequence_in_single_thread() {
289        let (producer, consumer) = SpscRing::<u32, 4>::new();
290
291        assert_eq!(producer.push(1), Ok(()));
292        assert_eq!(producer.push(2), Ok(()));
293        assert_eq!(producer.push(3), Ok(()));
294
295        assert_eq!(consumer.pop(), Some(1));
296        assert_eq!(consumer.pop(), Some(2));
297        assert_eq!(consumer.pop(), Some(3));
298        assert_eq!(consumer.pop(), None);
299    }
300
301    #[test]
302    fn overwrite_drops_oldest_item() {
303        let (producer, consumer) = SpscRing::<u32, 2>::new();
304
305        assert_eq!(producer.push_overwrite(10), None);
306        assert_eq!(producer.push_overwrite(20), None);
307        assert_eq!(producer.push_overwrite(30), Some(10));
308
309        assert_eq!(consumer.pop(), Some(20));
310        assert_eq!(consumer.pop(), Some(30));
311        assert_eq!(consumer.pop(), None);
312    }
313
314    #[test]
315    fn producer_calls_return_immediately_when_full() {
316        let (producer, _consumer) = SpscRing::<u64, 1>::new();
317        assert_eq!(producer.push(7), Ok(()));
318
319        let start = Instant::now();
320        let mut expected_drop = Some(7);
321        for value in 0..100_000 {
322            assert_eq!(producer.push(value), Err(value));
323            assert_eq!(producer.push_overwrite(value), expected_drop);
324            expected_drop = Some(value);
325        }
326
327        assert!(
328            start.elapsed() < Duration::from_secs(2),
329            "producer operations took too long while the ring stayed full"
330        );
331    }
332
333    #[test]
334    fn pop_async_drains_then_closes() {
335        let (producer, consumer) = SpscRing::<u32, 8>::new();
336        producer.push(1).unwrap();
337        producer.push(2).unwrap();
338        drop(producer);
339
340        assert_eq!(pollster::block_on(consumer.pop_async()), Some(1));
341        assert_eq!(pollster::block_on(consumer.pop_async()), Some(2));
342        assert_eq!(pollster::block_on(consumer.pop_async()), None);
343    }
344
345    #[test]
346    fn concurrent_producer_consumer_preserve_order() {
347        let (producer, consumer) = SpscRing::<u64, 1024>::new();
348        let producer_thread = thread::spawn(move || {
349            for expected in 0..50_000_u64 {
350                let mut item = expected;
351                loop {
352                    match producer.push(item) {
353                        Ok(()) => break,
354                        Err(returned) => {
355                            item = returned;
356                            std::hint::spin_loop();
357                        }
358                    }
359                }
360            }
361        });
362
363        for expected in 0..50_000_u64 {
364            let actual = pollster::block_on(consumer.pop_async());
365            assert_eq!(actual, Some(expected));
366        }
367        assert_eq!(pollster::block_on(consumer.pop_async()), None);
368
369        producer_thread.join().unwrap();
370    }
371
372    struct NoopWake;
373
374    impl Wake for NoopWake {
375        fn wake(self: Arc<Self>) {}
376    }
377
378    fn spin_pop(consumer: &SpscConsumer<u32, 1>, cx: &Context<'_>) -> Option<u32> {
379        loop {
380            if let Poll::Ready(item) = consumer.poll_pop(cx) {
381                return item;
382            }
383            std::hint::spin_loop();
384        }
385    }
386
387    #[test]
388    fn item_pushed_just_before_close_is_delivered() {
389        let (producers_tx, producers_rx) = mpsc::channel::<SpscProducer<u32, 1>>();
390        let worker = thread::spawn(move || {
391            for producer in producers_rx {
392                producer.push(7).unwrap();
393                drop(producer);
394            }
395        });
396        let waker = Waker::from(Arc::new(NoopWake));
397        let cx = Context::from_waker(&waker);
398
399        for _ in 0..20_000 {
400            let (producer, consumer) = SpscRing::<u32, 1>::new();
401            producers_tx.send(producer).unwrap();
402            assert_eq!(spin_pop(&consumer, &cx), Some(7));
403            assert_eq!(spin_pop(&consumer, &cx), None);
404        }
405
406        drop(producers_tx);
407        worker.join().unwrap();
408    }
409
410    #[cfg(feature = "futures-stream")]
411    #[test]
412    fn stream_wrapper_yields_items() {
413        use futures_core::Stream;
414
415        let (producer, consumer) = SpscRing::<u32, 4>::new();
416        let mut stream = consumer.stream();
417
418        producer.push(11).unwrap();
419        drop(producer);
420
421        let first = pollster::block_on(poll_fn(|cx| Pin::new(&mut stream).poll_next(cx)));
422        let second = pollster::block_on(poll_fn(|cx| Pin::new(&mut stream).poll_next(cx)));
423
424        assert_eq!(first, Some(11));
425        assert_eq!(second, None);
426    }
427}
doom_fish_utils/spsc.rs

doom_fish_utils/
spsc.rs
