Skip to main content

low_latency_data_structures/spsc/
consumer.rs

1use std::marker::PhantomData;
2use std::mem::MaybeUninit;
3use std::sync::atomic;
4
5use crate::mem::{Allocation, Allocator};
6use crate::shim::cell::UnsafeCell;
7use crate::shim::sync::Arc;
8use crate::spsc::queue::Queue;
9
10/// The popping handle of an SPSC FIFO queue.
11///
12/// Created together with its paired [`Producer`](crate::spsc::Producer) by
13/// [`new`](crate::spsc::new).
14/// `Consumer` is [`Send`] but not [`Sync`]: at most one thread may pop at a
15/// time.
16pub struct Consumer<T, const CAPACITY: usize, A>
17where
18    A: Allocator,
19{
20    inner: Arc<Queue<T, CAPACITY, A>>,
21    slots_ptr: *mut UnsafeCell<MaybeUninit<T>>,
22    cached_head: usize,
23    cached_tail: usize,
24    _not_sync: PhantomData<*const ()>,
25}
26
27impl<T, const CAPACITY: usize, A> std::fmt::Debug for Consumer<T, CAPACITY, A>
28where
29    A: Allocator,
30{
31    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
32        f.debug_struct("Consumer")
33            .field("capacity", &CAPACITY)
34            .finish_non_exhaustive()
35    }
36}
37
38// SAFETY: Consumer is Send because the underlying Queue is Send when both T
39// and the allocation are Send; PhantomData<*const ()> blocks Sync.
40unsafe impl<T, const CAPACITY: usize, A> Send for Consumer<T, CAPACITY, A>
41where
42    T: Send,
43    A: Allocator,
44    A::Allocation<T>: Send,
45{
46}
47
48impl<T, const CAPACITY: usize, A> Consumer<T, CAPACITY, A>
49where
50    A: Allocator,
51{
52    pub(super) fn new(queue: Arc<Queue<T, CAPACITY, A>>) -> Self {
53        let slots_ptr = queue.slots_allocation.ptr();
54        Self {
55            inner: queue,
56            slots_ptr,
57            cached_head: 0,
58            cached_tail: 0,
59            _not_sync: PhantomData,
60        }
61    }
62
63    /// Pops the next item from the queue.
64    ///
65    /// Wait-free. Returns `Some(item)` on success, `None` if the queue is
66    /// currently empty.
67    ///
68    /// # Examples
69    ///
70    /// ```
71    /// use low_latency_data_structures::spsc::{self, new};
72    /// use low_latency_data_structures::mem::global::GlobalAllocator;
73    ///
74    /// let (mut producer, mut consumer) = new::<u64, 4, GlobalAllocator>(
75    ///     spsc::Options::global_mlocked(),
76    /// );
77    /// assert_eq!(consumer.pop(), None);
78    /// let _ = producer.push(42);
79    /// assert_eq!(consumer.pop(), Some(42));
80    /// ```
81    #[inline]
82    #[must_use = "ignoring the popped item silently drops it"]
83    pub fn pop(&mut self) -> Option<T> {
84        if self.cached_head == self.cached_tail {
85            if self.still_empty() {
86                return None;
87            }
88        }
89        let slot_ptr = self
90            .slots_ptr
91            .wrapping_add(self.cached_head & (CAPACITY - 1));
92        // SAFETY: we read the cached_tail value that was released some time ago, it means we are
93        // guaranteed to see written value here. And it's not copied more than once because we
94        // increment head on the next line.
95        let item = unsafe {
96            slot_ptr
97                .as_ref_unchecked()
98                .with_mut(|ptr| ptr.cast::<T>().read())
99        };
100        self.cached_head += 1;
101        self.inner
102            .head
103            .store(self.cached_head, atomic::Ordering::Release);
104        Some(item)
105    }
106
107    #[cold]
108    fn still_empty(&mut self) -> bool {
109        // producer may have written something
110        self.cached_tail = self.inner.tail.load(atomic::Ordering::Acquire);
111        self.cached_head == self.cached_tail
112    }
113}
114
115#[cfg(test)]
116mod tests {
117    use super::*;
118    use crate::mem::test_util::NeverAllocator;
119
120    // Shouldn't be possible to construct Arc<Consumer<T>> and then use it from different threads as it
121    // will break the requirement of *Single* producer *Single* consumer queue.
122    static_assertions::assert_not_impl_any!(Consumer<u32, 0, NeverAllocator>: Sync);
123}