Skip to main content

low_latency_data_structures/spsc/
producer.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 pushing handle of an SPSC FIFO queue.
11///
12/// Created together with its paired [`Consumer`](crate::spsc::Consumer) by
13/// [`new`](crate::spsc::new).
14/// `Producer` is [`Send`] but not [`Sync`]: at most one thread may push at a
15/// time.
16pub struct Producer<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 Producer<T, CAPACITY, A>
28where
29    A: Allocator,
30{
31    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
32        f.debug_struct("Producer")
33            .field("capacity", &CAPACITY)
34            .finish_non_exhaustive()
35    }
36}
37
38// SAFETY: Producer 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 Producer<T, CAPACITY, A>
41where
42    T: Send,
43    A: Allocator,
44    A::Allocation<T>: Send,
45{
46}
47
48impl<T, const CAPACITY: usize, A> Producer<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    /// Pushes `item` onto the queue.
64    ///
65    /// Wait-free. Never blocks, never allocates. Returns `None` on success.
66    /// If the queue is full the item is returned unchanged as `Some(item)`,
67    /// so the caller can retry or back off without losing data.
68    ///
69    /// # Examples
70    ///
71    /// ```
72    /// use low_latency_data_structures::spsc::{self, new};
73    /// use low_latency_data_structures::mem::global::GlobalAllocator;
74    ///
75    /// let (mut producer, mut consumer) = new::<u64, 2, GlobalAllocator>(
76    ///     spsc::Options::global_mlocked(),
77    /// );
78    /// assert_eq!(producer.push(1), None);
79    /// assert_eq!(producer.push(2), None);
80    /// // Queue is full; item is handed back so we can retry later.
81    /// assert_eq!(producer.push(3), Some(3));
82    /// # let _ = consumer;
83    /// ```
84    #[inline]
85    #[must_use = "if the queue is full, the returned item must be handled (e.g. retried) or it is silently dropped"]
86    pub fn push(&mut self, item: T) -> Option<T> {
87        if self.cached_tail.wrapping_sub(self.cached_head) == CAPACITY {
88            if self.still_full() {
89                return Some(item);
90            }
91        }
92        let slot_ptr = self
93            .slots_ptr
94            .wrapping_add(self.cached_tail & (CAPACITY - 1));
95        // SAFETY: slot_ptr can't point to something after the slots buffer because of `% capacity`
96        // above. And it can be converted to a reference to T because T is self-contained bitwise
97        // (&T is 'static during the with_mut closure).
98        unsafe {
99            slot_ptr
100                .as_ref_unchecked()
101                .with_mut(|ptr| ptr.cast::<T>().write(item))
102        };
103        self.cached_tail += 1;
104        self.inner
105            .tail
106            .store(self.cached_tail, atomic::Ordering::Release);
107        None
108    }
109
110    #[cold]
111    fn still_full(&mut self) -> bool {
112        // consumer may have moved the head
113        self.cached_head = self.inner.head.load(atomic::Ordering::Acquire);
114        if self.cached_tail.wrapping_sub(self.cached_head) == CAPACITY {
115            return true;
116        }
117        false
118    }
119}
120
121#[cfg(test)]
122mod tests {
123    use super::*;
124    use crate::mem::test_util::NeverAllocator;
125
126    // Shouldn't be possible to construct Arc<Producer<T>> and then use it from different threads as it
127    // will break the requirement of *Single* producer *Single* consumer queue.
128    static_assertions::assert_not_impl_any!(Producer<u32, 0, NeverAllocator>: Sync);
129}