low_latency_data_structures/spmc/consumer.rs
1use std::marker::PhantomData;
2use std::sync::Arc;
3use std::sync::atomic::{Ordering, fence};
4
5use crate::mem::Allocation;
6use crate::spmc::queue::{Queue, Slot};
7
8#[repr(C, align(128))]
9pub(super) struct ConsumerState {
10 read_cursor: usize,
11 cached_write_cursor: usize,
12}
13
14/// Outcome of a [`Consumer::try_read`] call.
15///
16/// `try_read` is non-blocking and never allocates, so the return value
17/// distinguishes three cases the caller must handle: a successful read, an
18/// empty queue, or a detected lap. The enum is `#[must_use]` because
19/// silently dropping `Value` would lose the only copy of the data, and
20/// silently dropping `Lapped` would hide message loss.
21#[must_use = "the returned ReadResult indicates whether a value was read, the queue was empty, or the consumer was lapped"]
22#[derive(Debug, Clone, Copy, PartialEq, Eq)]
23pub enum ReadResult<T> {
24 /// A value was successfully read from the next slot.
25 Value(T),
26 /// The queue is currently empty.
27 Empty,
28 /// The producer lapped this consumer at least once. The read cursor has
29 /// been advanced to the most recently published slot; `skipped` is the
30 /// number of values that were never observed.
31 Lapped {
32 /// Number of values skipped over by the lap-recovery jump.
33 skipped: usize,
34 },
35}
36
37/// A reading handle of an SPMC broadcast queue.
38///
39/// Each consumer has its own private read cursor: every consumer observes
40/// every value the producer publishes, independently, unless it falls behind
41/// by more than `CAPACITY` slots.
42///
43/// `Consumer` is [`Send`] but not [`Sync`]: at most one thread may call
44/// [`try_read`](Self::try_read) on a given consumer at a time. To get more
45/// consumers, request more at construction via [`new`](crate::spmc::new).
46pub struct Consumer<T: bytemuck::AnyBitPattern, const CAPACITY: usize, AllocT: Allocation<Slot<T>>>
47{
48 state: ConsumerState,
49 inner: Arc<Queue<T, CAPACITY, AllocT>>,
50 _not_sync: PhantomData<*const ()>,
51}
52
53impl<T: bytemuck::AnyBitPattern, const CAPACITY: usize, AllocT: Allocation<Slot<T>>> std::fmt::Debug
54 for Consumer<T, CAPACITY, AllocT>
55{
56 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
57 f.debug_struct("Consumer")
58 .field("read_cursor", &self.state.read_cursor)
59 .field("cached_write_cursor", &self.state.cached_write_cursor)
60 .finish_non_exhaustive()
61 }
62}
63
64// SAFETY: It is Send on its own but we need to forbid the Sync.
65unsafe impl<T: bytemuck::AnyBitPattern, const CAPACITY: usize, AllocT: Allocation<Slot<T>>> Send
66 for Consumer<T, CAPACITY, AllocT>
67{
68}
69
70#[cfg(test)]
71mod tests {
72 use super::*;
73 use crate::mem::test_util::NeverAlloc;
74
75 static_assertions::assert_impl_all!(Consumer<u32, 2, NeverAlloc>: Send);
76 static_assertions::assert_not_impl_any!(Consumer<u32, 2, NeverAlloc>: Sync, Clone, Copy);
77}
78
79impl<T, const CAPACITY: usize, AllocT> Consumer<T, CAPACITY, AllocT>
80where
81 T: bytemuck::AnyBitPattern,
82 AllocT: Allocation<Slot<T>>,
83{
84 pub(super) fn new(queue: Arc<Queue<T, CAPACITY, AllocT>>) -> Self {
85 Self {
86 state: ConsumerState {
87 read_cursor: 0,
88 cached_write_cursor: 0,
89 },
90 inner: queue,
91 _not_sync: PhantomData,
92 }
93 }
94
95 /// Attempts to read the next value.
96 ///
97 /// Wait-free in the common case (value available, or queue empty);
98 /// lock-free when the producer is actively writing the target slot (the
99 /// consumer spin-loops until the producer publishes).
100 ///
101 /// # Protocol
102 ///
103 /// Each slot carries an even/odd seq number: even = stable, odd = mid
104 /// write. The consumer:
105 ///
106 /// 1. Loads the slot's seq. If odd, spin until it goes even.
107 /// 2. If the even seq is not the one expected for the current read
108 /// cursor, the producer has lapped us; reload the write cursor, jump
109 /// to the most recent slot, and return [`ReadResult::Lapped`].
110 /// 3. Otherwise read the slot's data, then re-load the seq. If the seq
111 /// is unchanged the read is committed and a [`ReadResult::Value`] is
112 /// returned; if the seq has moved, the read was torn and the loop
113 /// retries.
114 ///
115 /// # Examples
116 ///
117 /// ```
118 /// use low_latency_data_structures::spmc::{self, ReadResult, new};
119 /// use low_latency_data_structures::mem::global::GlobalAllocator;
120 ///
121 /// let (producer, [mut consumer]) = new::<u64, 16, 1, GlobalAllocator>(
122 /// spmc::Options::global_mlocked(),
123 /// );
124 /// assert_eq!(consumer.try_read(), ReadResult::Empty);
125 /// producer.publish(7);
126 /// assert_eq!(consumer.try_read(), ReadResult::Value(7));
127 /// ```
128 #[inline]
129 pub fn try_read(&mut self) -> ReadResult<T> {
130 if self.state.read_cursor == self.state.cached_write_cursor {
131 // It still may not be empty
132 if self.is_still_empty() {
133 return ReadResult::Empty;
134 }
135 }
136 let slot = self.slot(self.state.read_cursor);
137 // That's the expected seq_no after producer has written the item
138 let expected_seq = self.state.read_cursor.wrapping_mul(2).wrapping_add(2);
139 loop {
140 let seq1 = slot.seq.load(Ordering::Acquire);
141 // If producer currently writes it, we wait (by the way, it means that we were overlapped)
142 if seq1 & 1 == 1 {
143 std::hint::spin_loop();
144 continue;
145 }
146 if seq1 != expected_seq {
147 // We were overlapped. Reload cached_write_cursor before jumping:
148 // under sustained producer overflow the cached value is itself
149 // stale (cached_write_cursor - 1 may have been lapped past
150 // again), and without the reload the consumer gets pinned in
151 // a Lapped -> Lapped cycle with skipped=0 forever.
152 self.state.cached_write_cursor = self
153 .inner
154 .producer_state
155 .write_cursor
156 .load(Ordering::Acquire);
157 let new_r_cursor = self.state.cached_write_cursor.wrapping_sub(1);
158 let skipped = new_r_cursor.wrapping_sub(self.state.read_cursor);
159 self.state.read_cursor = new_r_cursor;
160 return ReadResult::Lapped { skipped };
161 }
162 // SAFETY: T: AnyBitPattern, so even a torn read materialises a valid T;
163 // the seq2 check below rejects the value if the read raced a producer write.
164 let attempted_read = unsafe { slot.data.get().read_volatile() };
165 // ASM: prevent the seq1 load from moving below the read
166 fence(Ordering::Acquire);
167 let seq2 = slot.seq.load(Ordering::Relaxed);
168 if seq1 == seq2 {
169 self.state.read_cursor = self.state.read_cursor.wrapping_add(1);
170 return ReadResult::Value(attempted_read);
171 }
172 // Torn-read, retry
173 }
174 }
175
176 #[cold]
177 fn is_still_empty(&mut self) -> bool {
178 self.state.cached_write_cursor = self
179 .inner
180 .producer_state
181 .write_cursor
182 .load(Ordering::Acquire);
183 self.state.read_cursor == self.state.cached_write_cursor
184 }
185
186 /// Wraps the given `i` around `CAPACITY - 1`.
187 ///
188 /// # Safety
189 ///
190 /// - The masked index is always in `[0, CAPACITY)`, so the derived
191 /// pointer stays inside the allocation.
192 /// - The allocation is aligned and dereferenceable per the `Allocator`
193 /// contract.
194 /// - Every slot is initialized in [`new`](super::new) before the
195 /// `Queue` is wrapped in `Arc`, so `assume_init_ref` is sound.
196 /// - Returning `&Slot<T>` while another thread mutates through
197 /// `slot.data` / `slot.seq` is legal because the mutable state sits
198 /// behind `UnsafeCell` and atomics.
199 #[inline(always)]
200 fn slot(&self, i: usize) -> &Slot<T> {
201 unsafe {
202 self.inner
203 .slots
204 .ptr()
205 .wrapping_add(i & (CAPACITY - 1))
206 .as_ref_unchecked()
207 .get()
208 .as_ref_unchecked()
209 .assume_init_ref()
210 }
211 }
212}