waitfree_sync/
spsc.rs

1//! A wait-free single-producer single-consumer (SPSC) queue to send data to another thread.
2//! It is based on the improved FastForward queue.
3//!
4//! This is similar to [`std::sync::mpsc`], but restricted to a single producer and a single
5//! consumer and backed by a fixed-size ring buffer instead of an unbounded linked list. Because
6//! of this, [`Sender::try_send`] and [`Receiver::try_recv`] never block and never allocate: they
7//! return immediately instead of parking the calling thread the way `std`'s blocking `send`/`recv`
8//! do.
9//!
10//! # Example
11//! ```rust
12//! use waitfree_sync::spsc;
13//!
14//! //                            Type ──╮   ╭─ Capacity
15//! let (mut tx, mut rx) = spsc::spsc::<u64>(8);
16//! tx.try_send(234);
17//! assert_eq!(rx.try_recv(),Ok(234u64));
18//! ```
19//!
20//! # Behavior for full and empty queue.
21//! If the queue is full, [`Sender::try_send`] returns [`SendError::NoSpaceLeft`].
22//! If the queue is empty, [`Receiver::try_recv`] returns [`TryRecvError::Empty`].
23//!
24use crate::import::{Arc, AtomicBool, Ordering, UnsafeCell};
25use core::error::Error;
26use crossbeam_utils::CachePadded;
27use std::{fmt::Debug, sync::atomic::AtomicUsize};
28
29/// Create a new wait-free SPSC queue. The `capacity` must be a power of two, which is validate during runtime.
30/// # Panic
31/// Panics if the `capacity` is not a power of two.
32/// # Example
33/// ```rust
34/// use waitfree_sync::spsc;
35///
36/// //               Data type ──╮   ╭─ Capacity
37/// let (tx, rx) = spsc::spsc::<u64>(8);
38/// ```
39pub fn spsc<T>(capacity: usize) -> (Sender<T>, Receiver<T>) {
40    if !is_power_of_two(capacity) {
41        panic!("The SIZE must be a power of 2")
42    }
43
44    let chan = Arc::new(Spsc::new(capacity));
45
46    let r = Receiver::new(chan.clone());
47    let w = Sender::new(chan);
48
49    (w, r)
50}
51
52const fn is_power_of_two(x: usize) -> bool {
53    let c = x.wrapping_sub(1);
54    (x != 0) && (x != 1) && ((x & c) == 0)
55}
56
57/// An error returned from the [`Sender::try_send`] function on a [`Sender`].
58///
59/// The error contains the data being sent as a payload so it can be recovered.
60#[derive(Clone, Debug, PartialEq)]
61pub enum SendError<T> {
62    /// The queue is full. The receiving side of the queue must collect items.
63    NoSpaceLeft(T),
64    /// The receiving end of a channel is disconnected, implying that the data could never be received.
65    ReceiverSideDropped(T),
66}
67impl<T: Debug> Error for SendError<T> {}
68impl<T> core::fmt::Display for SendError<T> {
69    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
70        match self {
71            SendError::NoSpaceLeft(_) => write!(f, "No space left in the SPSC queue."),
72            SendError::ReceiverSideDropped(_) => {
73                write!(f, "Receiver side of the SPSC queue dropped.")
74            }
75        }
76    }
77}
78impl<T> SendError<T> {
79    /// Returns the value which was tried to be sent to the queue.
80    pub fn into_value(self) -> T {
81        match self {
82            SendError::NoSpaceLeft(val) => val,
83            SendError::ReceiverSideDropped(val) => val,
84        }
85    }
86}
87
88/// This enumeration is the list of the possible reasons that [`Receiver::try_recv`] could not return data when called.
89#[derive(Clone, Debug, PartialEq)]
90pub enum TryRecvError {
91    /// This queue is currently empty, but the Sender have not yet disconnected, so data may yet become available.
92    Empty,
93    /// The queues sending half has become disconnected, and there will never be any more data received on it.
94    Disconnected,
95}
96impl Error for TryRecvError {}
97impl core::fmt::Display for TryRecvError {
98    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
99        match self {
100            TryRecvError::Empty => write!(f, "No data available in the SPSC queue."),
101            TryRecvError::Disconnected => {
102                write!(f, "Sender side of the SPSC queue dropped.")
103            }
104        }
105    }
106}
107
108#[derive(Debug)]
109struct Slot<T> {
110    value: UnsafeCell<Option<T>>,
111    occupied: CachePadded<AtomicBool>,
112}
113impl<T> Slot<T> {
114    fn new() -> Self {
115        Self {
116            value: UnsafeCell::new(None),
117            occupied: CachePadded::new(false.into()),
118        }
119    }
120}
121
122#[derive(Debug)]
123struct Spsc<T> {
124    mem: Box<[Slot<T>]>,
125    // The mask is written when this structure is created and is then only read.
126    // Therefore, we do not need Atomic here.
127    mask: usize,
128    read: CachePadded<AtomicUsize>,
129    write: CachePadded<AtomicUsize>,
130}
131
132impl<T> Spsc<T> {
133    fn new(size: usize) -> Self {
134        let mut buffer = Vec::with_capacity(size);
135        for _ in 0..size {
136            buffer.push(Slot::new());
137        }
138        let buffer: Box<[Slot<T>]> = buffer.into_boxed_slice();
139        Spsc {
140            mem: buffer,
141            mask: size - 1,
142            read: CachePadded::new(0.into()),
143            write: CachePadded::new(0.into()),
144        }
145    }
146
147    #[inline]
148    fn capacity(&self) -> usize {
149        self.mask + 1
150    }
151
152    #[inline]
153    fn len(&self) -> usize {
154        self.write
155            .load(Ordering::Relaxed)
156            .saturating_sub(self.read.load(Ordering::Relaxed))
157    }
158}
159
160/// The receiving side of the [spsc] queue.
161#[derive(Debug)]
162pub struct Receiver<T> {
163    spsc: Arc<Spsc<T>>,
164}
165unsafe impl<T: Send> Send for Receiver<T> {}
166unsafe impl<T: Send> Sync for Receiver<T> {}
167
168impl<T> Receiver<T> {
169    fn new(spsc: Arc<Spsc<T>>) -> Self {
170        Receiver { spsc }
171    }
172}
173
174impl<T> Receiver<T> {
175    /// Retrieve the next available element from the queue without blocking.
176    ///
177    /// Returns [`TryRecvError::Empty`] if the queue is currently empty,
178    /// or [`TryRecvError::Disconnected`] if the [Sender] has been dropped and
179    /// no further items can arrive.
180    pub fn try_recv(&mut self) -> Result<T, TryRecvError> {
181        let read = self.spsc.read.load(Ordering::Relaxed);
182        let rpos = read & self.spsc.mask;
183        let slot = unsafe { self.spsc.mem.get_unchecked(rpos) };
184        if !slot.occupied.load(Ordering::Acquire) {
185            if Arc::strong_count(&self.spsc) < 2 {
186                Err(TryRecvError::Disconnected)
187            } else {
188                Err(TryRecvError::Empty)
189            }
190        } else {
191            #[cfg(not(loom))]
192            let val = unsafe { slot.value.get().replace(None) };
193            #[cfg(loom)]
194            let val = unsafe { slot.value.get_mut().with(|ptr| ptr.replace(None)) };
195
196            slot.occupied.store(false, Ordering::Release);
197            // self.read = self.read.wrapping_add(1);
198            self.spsc
199                .read
200                .store(read.wrapping_add(1), Ordering::Relaxed);
201            Ok(val.ok_or(TryRecvError::Empty)?)
202        }
203    }
204    /// Peeks the next element in the queue without removing it.
205    #[cfg(not(loom))] // We can't return a reference to an UnsafeCell of loom.
206    pub fn peek(&self) -> Option<&T> {
207        let rpos = self.spsc.read.load(Ordering::Relaxed) & self.spsc.mask;
208        let slot = unsafe { self.spsc.mem.get_unchecked(rpos) };
209        if !slot.occupied.load(Ordering::Acquire) {
210            None
211        } else {
212            let val = unsafe { &*slot.value.get() };
213            val.as_ref()
214        }
215    }
216
217    /// Returns the total number of items that the queue can hold at most.
218    #[inline]
219    pub fn capacity(&self) -> usize {
220        // SAFETY: This is safe because we only read size which is never written.
221        self.spsc.capacity()
222    }
223
224    /// Returns the number of items in the queue.
225    /// # WARNING
226    /// This length is only a best-effort estimate.
227    /// It is computed from relaxed atomic and is NOT a linearizable value.
228    /// It may be temporarily incorrect (including over/under-counting) due to
229    /// reordering and visibility delays across threads.
230    #[inline]
231    pub fn len(&self) -> usize {
232        self.spsc.len()
233    }
234
235    /// Returns true if the queue is empty.
236    /// # WARNING
237    /// This length is only a best-effort estimate.
238    /// It is computed from relaxed atomic and is NOT a linearizable value.
239    /// It may be temporarily incorrect (including over/under-counting) due to
240    /// reordering and visibility delays across threads.
241    #[inline]
242    pub fn is_empty(&self) -> bool {
243        self.spsc.len() == 0
244    }
245}
246
247/// The sending side of the [spsc] queue.
248#[derive(Debug)]
249pub struct Sender<T> {
250    spsc: Arc<Spsc<T>>,
251}
252unsafe impl<T: Send> Send for Sender<T> {}
253unsafe impl<T: Send> Sync for Sender<T> {}
254impl<T> Sender<T> {
255    fn new(spsc: Arc<Spsc<T>>) -> Self {
256        Sender { spsc }
257    }
258}
259
260impl<T> Sender<T> {
261    /// Attempts to send a value to the queue without blocking.
262    ///
263    /// Because this queue has a fixed capacity, sending returns [`SendError::NoSpaceLeft`]
264    /// instead of blocking or growing the buffer when the queue is full.
265    /// Returns [`SendError::ReceiverSideDropped`] if the [Receiver] has been dropped.
266    pub fn try_send(&mut self, data: T) -> Result<(), SendError<T>> {
267        let write = self.spsc.write.load(Ordering::Relaxed);
268        let wpos = write & self.spsc.mask;
269
270        if Arc::strong_count(&self.spsc) < 2 {
271            return Err(SendError::ReceiverSideDropped(data));
272        }
273
274        let slot = unsafe { self.spsc.mem.get_unchecked(wpos) };
275        if slot.occupied.load(Ordering::Acquire) {
276            Err(SendError::NoSpaceLeft(data))
277        } else {
278            #[cfg(not(loom))]
279            unsafe {
280                slot.value.get().write(Some(data))
281            };
282            #[cfg(loom)]
283            unsafe {
284                slot.value.get_mut().with(|ptr| ptr.write(Some(data)))
285            };
286            slot.occupied.store(true, Ordering::Release);
287            self.spsc
288                .write
289                .store(write.wrapping_add(1), Ordering::Relaxed);
290            Ok(())
291        }
292    }
293
294    /// Returns the total number of items that the queue can hold at most.
295    #[inline]
296    pub fn capacity(&self) -> usize {
297        // SAFETY: This is safe because we only read size which is never written.
298        self.spsc.capacity()
299    }
300
301    /// Returns the number of items in the queue.
302    /// # WARNING
303    /// This length is only a best-effort estimate.
304    /// It is computed from relaxed atomic and is NOT a linearizable value.
305    /// It may be temporarily incorrect (including over/under-counting) due to
306    /// reordering and visibility delays across threads.
307    #[inline]
308    pub fn len(&self) -> usize {
309        self.spsc.len()
310    }
311
312    /// Returns true if the queue is empty.
313    /// # WARNING
314    /// This length is only a best-effort estimate.
315    /// It is computed from relaxed atomic and is NOT a linearizable value.
316    /// It may be temporarily incorrect (including over/under-counting) due to
317    /// reordering and visibility delays across threads.
318    #[inline]
319    pub fn is_empty(&self) -> bool {
320        self.spsc.len() == 0
321    }
322}
323
324#[cfg(not(loom))]
325#[cfg(test)]
326mod test {
327    #[cfg(loom)]
328    use loom::thread;
329    #[cfg(not(loom))]
330    use std::thread;
331
332    use super::*;
333
334    #[test]
335    fn smoke() {
336        let (mut w, mut r) = spsc(4);
337        w.try_send(vec![0; 15]).unwrap();
338        w.try_send(vec![0; 16]).unwrap();
339        w.try_send(vec![0; 17]).unwrap();
340        w.try_send(vec![0; 18]).unwrap();
341
342        assert_eq!(r.try_recv(), Ok(vec![0; 15]));
343        assert_eq!(r.try_recv(), Ok(vec![0; 16]));
344        assert_eq!(r.try_recv(), Ok(vec![0; 17]));
345        assert_eq!(r.try_recv(), Ok(vec![0; 18]));
346    }
347
348    #[test]
349    fn test_is_power_of_two() {
350        assert!(!is_power_of_two(0));
351        assert!(!is_power_of_two(1));
352        assert!(is_power_of_two(2));
353        assert!(!is_power_of_two(3));
354        assert!(is_power_of_two(4));
355        assert!(!is_power_of_two(5));
356        assert!(!is_power_of_two(6));
357        assert!(!is_power_of_two(7));
358        assert!(is_power_of_two(8));
359        assert!(!is_power_of_two(9));
360
361        assert!(!is_power_of_two(15));
362        assert!(is_power_of_two(16));
363        assert!(!is_power_of_two(17));
364
365        assert!(!is_power_of_two(31));
366        assert!(is_power_of_two(32));
367        assert!(!is_power_of_two(33));
368    }
369
370    #[test]
371    fn test_drop_read_side() {
372        let (mut write, read) = spsc::<i32>(4);
373
374        assert_eq!(write.try_send(1), Ok(()));
375        assert_eq!(write.len(), 1);
376        assert_eq!(write.try_send(2), Ok(()));
377        assert_eq!(write.len(), 2);
378        drop(read);
379        assert_eq!(write.try_send(3), Err(SendError::ReceiverSideDropped(3)));
380        assert_eq!(write.len(), 2);
381        assert_eq!(write.try_send(4), Err(SendError::ReceiverSideDropped(4)));
382        assert_eq!(write.len(), 2);
383        assert_eq!(write.try_send(5), Err(SendError::ReceiverSideDropped(5)));
384        assert_eq!(write.len(), 2);
385    }
386
387    #[test]
388    fn test_drop_write_side() {
389        let (mut write, mut read) = spsc::<i32>(4);
390
391        write.try_send(0).unwrap();
392        write.try_send(1).unwrap();
393        assert_eq!(read.try_recv(), Ok(0));
394        drop(write);
395        assert_eq!(read.try_recv(), Ok(1));
396    }
397
398    #[test]
399    fn test_full_empty() {
400        let (mut write, mut read) = spsc::<i32>(4);
401        assert_eq!(write.try_send(1), Ok(()));
402        assert_eq!(write.len(), 1);
403        assert_eq!(write.try_send(2), Ok(()));
404        assert_eq!(write.len(), 2);
405        assert_eq!(write.try_send(3), Ok(()));
406        assert_eq!(write.len(), 3);
407        assert_eq!(write.try_send(4), Ok(()));
408        assert_eq!(write.len(), 4);
409        assert_eq!(write.try_send(5), Err(SendError::NoSpaceLeft(5)));
410        assert_eq!(write.len(), 4);
411
412        assert_eq!(read.try_recv(), Ok(1));
413        assert_eq!(write.len(), 3);
414        assert_eq!(write.try_send(6), Ok(()));
415        assert_eq!(write.len(), 4);
416        assert_eq!(read.try_recv(), Ok(2));
417        assert_eq!(write.len(), 3);
418        assert_eq!(read.try_recv(), Ok(3));
419        assert_eq!(write.len(), 2);
420        assert_eq!(read.try_recv(), Ok(4));
421        assert_eq!(write.len(), 1);
422        assert_eq!(read.try_recv(), Ok(6));
423        assert_eq!(read.try_recv(), Err(TryRecvError::Empty));
424    }
425
426    #[test]
427    fn test_drop_one_side() {
428        let (mut write, read) = spsc::<i32>(4);
429        assert_eq!(write.try_send(1), Ok(()));
430        assert_eq!(write.len(), 1);
431        assert_eq!(write.try_send(2), Ok(()));
432        assert_eq!(write.len(), 2);
433        drop(read);
434        assert_eq!(write.try_send(3), Err(SendError::ReceiverSideDropped(3)));
435        assert_eq!(write.len(), 2);
436        assert_eq!(write.try_send(4), Err(SendError::ReceiverSideDropped(4)));
437        assert_eq!(write.len(), 2);
438        assert_eq!(write.try_send(5), Err(SendError::ReceiverSideDropped(5)));
439        assert_eq!(write.len(), 2);
440    }
441
442    #[test]
443    fn test_peek() {
444        let (mut w, mut r) = spsc(4);
445        w.try_send(vec![0; 15]).unwrap();
446        w.try_send(vec![0; 16]).unwrap();
447        w.try_send(vec![0; 17]).unwrap();
448        w.try_send(vec![0; 18]).unwrap();
449
450        assert_eq!(r.peek(), Some(&vec![0; 15]));
451        assert_eq!(r.try_recv(), Ok(vec![0; 15]));
452        assert_eq!(r.peek(), Some(&vec![0; 16]));
453        assert_eq!(r.try_recv(), Ok(vec![0; 16]));
454        assert_eq!(r.peek(), Some(&vec![0; 17]));
455        assert_eq!(r.try_recv(), Ok(vec![0; 17]));
456        assert_eq!(r.peek(), Some(&vec![0; 18]));
457        assert_eq!(r.peek(), Some(&vec![0; 18]));
458        assert_eq!(r.peek(), Some(&vec![0; 18]));
459        assert_eq!(r.try_recv(), Ok(vec![0; 18]));
460        assert_eq!(r.peek(), None);
461    }
462
463    #[test]
464    fn test_peek_threaded() {
465        let (mut sender, mut receiver) = spsc(4);
466
467        let writer_thread = thread::spawn(move || {
468            thread::park();
469            for i in 0..4 {
470                assert_eq!(sender.try_send([i; 50]), Ok(()));
471            }
472        });
473        let reader_thread = thread::spawn(move || {
474            thread::park();
475            let mut i = 0;
476            while i < 4 {
477                if let Some(val) = receiver.peek() {
478                    let first_entry = val[0];
479                    for entry in val {
480                        assert_eq!(*entry, first_entry);
481                    }
482                    let val = receiver.try_recv().unwrap();
483                    let first_entry = val[0];
484                    for entry in val {
485                        assert_eq!(entry, first_entry);
486                    }
487                    i += 1;
488                }
489            }
490        });
491        writer_thread.thread().unpark();
492        reader_thread.thread().unpark();
493        assert!(writer_thread.join().is_ok());
494        assert!(reader_thread.join().is_ok());
495    }
496}
waitfree_sync/spsc.rs

waitfree_sync/
spsc.rs
