Skip to main content

subms_spsc_ring_buffer/features/
wait_strategies.rs

1//! Blocking wrappers around the base wait-free SPSC ring.
2//!
3//! The base `Producer` / `Consumer` return immediately on full / empty. These
4//! wrappers take a `WaitStrategy` and block when the ring isn't ready. Three
5//! strategies are provided:
6//!
7//! - `BusySpin` - tight `spin_loop` hint; lowest wakeup latency, highest CPU.
8//! - `YieldStrategy` - calls `thread::yield_now` between retries; lets other
9//!   threads run, decent default for over-subscribed cores.
10//! - `ParkStrategy` - `thread::park` + producer-side `unpark`; lowest CPU,
11//!   adds a syscall and a few microseconds of wakeup latency.
12//!
13//! All wrappers are wait-free in the base case (slot available immediately).
14//! They become blocking ONLY when the ring is full / empty.
15
16use std::sync::Arc;
17use std::sync::atomic::{AtomicBool, Ordering};
18use std::thread::{self, Thread};
19
20use crate::{Consumer, Producer};
21
22/// Strategy for backing off when the ring is full (producer) or empty
23/// (consumer). Implementations should be cheap to construct and `Send`able
24/// to the owning thread.
25pub trait WaitStrategy: Send {
26    /// Wait briefly. Called between retries when the ring isn't ready.
27    fn wait(&mut self);
28    /// Signal the waiter. Default is a no-op; only `ParkStrategy` overrides.
29    fn signal(&self) {}
30}
31
32/// Tight CPU spin. Use when the producer and consumer are on dedicated cores
33/// and latency matters more than CPU usage.
34pub struct BusySpin;
35impl WaitStrategy for BusySpin {
36    fn wait(&mut self) {
37        std::hint::spin_loop();
38    }
39}
40
41/// `thread::yield_now` between retries. Good default when the threads share
42/// cores with other work.
43pub struct YieldStrategy;
44impl WaitStrategy for YieldStrategy {
45    fn wait(&mut self) {
46        thread::yield_now();
47    }
48}
49
50/// Park the calling thread; the opposite-end wrapper `unpark`s it on progress.
51/// Lowest CPU; adds wakeup latency.
52///
53/// Both ends share a `Parker` instance via `Arc` so the producer can wake the
54/// consumer and vice versa.
55pub struct ParkStrategy {
56    parker: Arc<Parker>,
57    is_producer: bool,
58}
59
60impl ParkStrategy {
61    /// Build a matched `(producer_strategy, consumer_strategy)` pair that
62    /// share a parking state.
63    pub fn pair() -> (Self, Self) {
64        let p = Arc::new(Parker::new());
65        (
66            Self {
67                parker: p.clone(),
68                is_producer: true,
69            },
70            Self {
71                parker: p,
72                is_producer: false,
73            },
74        )
75    }
76}
77
78impl WaitStrategy for ParkStrategy {
79    fn wait(&mut self) {
80        // Park the calling thread, recording it as the side that's blocked.
81        // The opposite end's `signal()` will unpark it.
82        if self.is_producer {
83            self.parker.park_producer();
84        } else {
85            self.parker.park_consumer();
86        }
87    }
88
89    fn signal(&self) {
90        // Wake the OPPOSITE side (the producer wakes the consumer and vice versa).
91        if self.is_producer {
92            self.parker.unpark_consumer();
93        } else {
94            self.parker.unpark_producer();
95        }
96    }
97}
98
99/// Internal: cross-end park state. Stores the parked Thread handle for each
100/// side and an `unparked` flag to defeat lost-wakeup races.
101struct Parker {
102    producer: parking_lot::Mutex<Option<Thread>>,
103    consumer: parking_lot::Mutex<Option<Thread>>,
104    producer_unparked: AtomicBool,
105    consumer_unparked: AtomicBool,
106}
107
108impl Parker {
109    fn new() -> Self {
110        Self {
111            producer: parking_lot::Mutex::new(None),
112            consumer: parking_lot::Mutex::new(None),
113            producer_unparked: AtomicBool::new(false),
114            consumer_unparked: AtomicBool::new(false),
115        }
116    }
117
118    fn park_producer(&self) {
119        // Fast path: prior unpark already pending - consume it without sleeping.
120        if self.producer_unparked.swap(false, Ordering::Acquire) {
121            return;
122        }
123        {
124            let mut slot = self.producer.lock();
125            *slot = Some(thread::current());
126        }
127        // Re-check after registering: the unpark might have raced with our
128        // registration; if it did, swap returns true and we skip the park.
129        if self.producer_unparked.swap(false, Ordering::Acquire) {
130            return;
131        }
132        thread::park();
133        // Clear any residual flag set during the park.
134        self.producer_unparked.store(false, Ordering::Release);
135    }
136
137    fn park_consumer(&self) {
138        if self.consumer_unparked.swap(false, Ordering::Acquire) {
139            return;
140        }
141        {
142            let mut slot = self.consumer.lock();
143            *slot = Some(thread::current());
144        }
145        if self.consumer_unparked.swap(false, Ordering::Acquire) {
146            return;
147        }
148        thread::park();
149        self.consumer_unparked.store(false, Ordering::Release);
150    }
151
152    fn unpark_producer(&self) {
153        self.producer_unparked.store(true, Ordering::Release);
154        if let Some(t) = self.producer.lock().take() {
155            t.unpark();
156        }
157    }
158
159    fn unpark_consumer(&self) {
160        self.consumer_unparked.store(true, Ordering::Release);
161        if let Some(t) = self.consumer.lock().take() {
162            t.unpark();
163        }
164    }
165}
166
167// Mini lock impl - no external dep, mirrors a basic mutex. We need it for
168// safe handoff of Thread handles between producer and consumer.
169mod parking_lot {
170    use std::cell::UnsafeCell;
171    use std::sync::atomic::{AtomicBool, Ordering};
172
173    pub struct Mutex<T> {
174        locked: AtomicBool,
175        inner: UnsafeCell<T>,
176    }
177
178    unsafe impl<T: Send> Sync for Mutex<T> {}
179    unsafe impl<T: Send> Send for Mutex<T> {}
180
181    pub struct Guard<'a, T> {
182        m: &'a Mutex<T>,
183    }
184
185    impl<T> Mutex<T> {
186        pub fn new(value: T) -> Self {
187            Self {
188                locked: AtomicBool::new(false),
189                inner: UnsafeCell::new(value),
190            }
191        }
192
193        pub fn lock(&self) -> Guard<'_, T> {
194            while self
195                .locked
196                .compare_exchange_weak(false, true, Ordering::Acquire, Ordering::Relaxed)
197                .is_err()
198            {
199                std::hint::spin_loop();
200            }
201            Guard { m: self }
202        }
203    }
204
205    impl<T> std::ops::Deref for Guard<'_, T> {
206        type Target = T;
207        fn deref(&self) -> &T {
208            unsafe { &*self.m.inner.get() }
209        }
210    }
211
212    impl<T> std::ops::DerefMut for Guard<'_, T> {
213        fn deref_mut(&mut self) -> &mut T {
214            unsafe { &mut *self.m.inner.get() }
215        }
216    }
217
218    impl<T> Drop for Guard<'_, T> {
219        fn drop(&mut self) {
220            self.m.locked.store(false, Ordering::Release);
221        }
222    }
223}
224
225/// Blocking producer: `push(value)` waits for a free slot using the strategy.
226pub struct BlockingSpscProducer<T, S: WaitStrategy> {
227    inner: Producer<T>,
228    strategy: S,
229}
230
231impl<T, S: WaitStrategy> BlockingSpscProducer<T, S> {
232    pub fn new(producer: Producer<T>, strategy: S) -> Self {
233        Self {
234            inner: producer,
235            strategy,
236        }
237    }
238
239    /// Block until the value can be pushed.
240    pub fn push(&mut self, mut value: T) {
241        loop {
242            match self.inner.try_push(value) {
243                Ok(()) => {
244                    // Wake the consumer if it was parked.
245                    self.strategy.signal();
246                    return;
247                }
248                Err(returned) => {
249                    value = returned;
250                    self.strategy.wait();
251                }
252            }
253        }
254    }
255
256    /// Non-blocking try_push; pass-through to the underlying ring.
257    pub fn try_push(&mut self, value: T) -> Result<(), T> {
258        let r = self.inner.try_push(value);
259        if r.is_ok() {
260            self.strategy.signal();
261        }
262        r
263    }
264
265    pub fn capacity(&self) -> usize {
266        self.inner.capacity()
267    }
268}
269
270/// Blocking consumer: `pop()` waits for an item using the strategy.
271pub struct BlockingSpscConsumer<T, S: WaitStrategy> {
272    inner: Consumer<T>,
273    strategy: S,
274}
275
276impl<T, S: WaitStrategy> BlockingSpscConsumer<T, S> {
277    pub fn new(consumer: Consumer<T>, strategy: S) -> Self {
278        Self {
279            inner: consumer,
280            strategy,
281        }
282    }
283
284    /// Block until an item is available.
285    pub fn pop(&mut self) -> T {
286        loop {
287            if let Some(v) = self.inner.try_pop() {
288                // Wake the producer if it was parked on a full ring.
289                self.strategy.signal();
290                return v;
291            }
292            self.strategy.wait();
293        }
294    }
295
296    pub fn try_pop(&mut self) -> Option<T> {
297        let v = self.inner.try_pop();
298        if v.is_some() {
299            self.strategy.signal();
300        }
301        v
302    }
303
304    pub fn capacity(&self) -> usize {
305        self.inner.capacity()
306    }
307}
308
309#[cfg(test)]
310#[path = "wait_strategies_tests.rs"]
311mod tests;