subms_spsc_ring_buffer/features/
wait_strategies.rs1use std::sync::Arc;
17use std::sync::atomic::{AtomicBool, Ordering};
18use std::thread::{self, Thread};
19
20use crate::{Consumer, Producer};
21
22pub trait WaitStrategy: Send {
26 fn wait(&mut self);
28 fn signal(&self) {}
30}
31
32pub struct BusySpin;
35impl WaitStrategy for BusySpin {
36 fn wait(&mut self) {
37 std::hint::spin_loop();
38 }
39}
40
41pub struct YieldStrategy;
44impl WaitStrategy for YieldStrategy {
45 fn wait(&mut self) {
46 thread::yield_now();
47 }
48}
49
50pub struct ParkStrategy {
56 parker: Arc<Parker>,
57 is_producer: bool,
58}
59
60impl ParkStrategy {
61 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 if self.is_producer {
83 self.parker.park_producer();
84 } else {
85 self.parker.park_consumer();
86 }
87 }
88
89 fn signal(&self) {
90 if self.is_producer {
92 self.parker.unpark_consumer();
93 } else {
94 self.parker.unpark_producer();
95 }
96 }
97}
98
99struct 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 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 if self.producer_unparked.swap(false, Ordering::Acquire) {
130 return;
131 }
132 thread::park();
133 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
167mod 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
225pub 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 pub fn push(&mut self, mut value: T) {
241 loop {
242 match self.inner.try_push(value) {
243 Ok(()) => {
244 self.strategy.signal();
246 return;
247 }
248 Err(returned) => {
249 value = returned;
250 self.strategy.wait();
251 }
252 }
253 }
254 }
255
256 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
270pub 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 pub fn pop(&mut self) -> T {
286 loop {
287 if let Some(v) = self.inner.try_pop() {
288 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;