1#[cfg(loom)]
12use loom::sync::atomic::{AtomicUsize, Ordering};
13use std::cell::UnsafeCell;
14use std::mem::MaybeUninit;
15#[cfg(not(loom))]
16use std::sync::atomic::{AtomicUsize, Ordering};
17
18#[derive(Clone, Copy, Debug)]
20pub struct RingConfig {
21 pub capacity: usize,
23 pub backpressure_threshold: usize,
25}
26
27impl Default for RingConfig {
28 fn default() -> Self {
29 Self {
30 capacity: 4095, backpressure_threshold: 2867, }
33 }
34}
35
36#[cfg(loom)]
38const RING_CAPACITY: usize = 16;
39#[cfg(not(loom))]
40const RING_CAPACITY: usize = 4096;
41const CACHE_LINE: usize = 128; #[inline(always)]
45#[allow(dead_code)]
46fn prefetch_read<T>(ptr: *const T) {
47 #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
48 unsafe {
49 std::arch::x86_64::_mm_prefetch(ptr as *const i8, std::arch::x86_64::_MM_HINT_T0);
50 }
51 #[cfg(not(any(target_arch = "x86", target_arch = "x86_64")))]
52 {
53 std::hint::black_box(ptr);
55 }
56}
57
58#[inline(always)]
60#[allow(dead_code)]
61fn prefetch_write<T>(ptr: *mut T) {
62 #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
63 unsafe {
64 std::arch::x86_64::_mm_prefetch(ptr as *const i8, std::arch::x86_64::_MM_HINT_ET1);
65 }
66 #[cfg(not(any(target_arch = "x86", target_arch = "x86_64")))]
67 {
68 std::hint::black_box(ptr);
69 }
70}
71
72#[derive(Debug)]
80pub struct SpscRingBuffer<T> {
81 buffer: UnsafeCell<[MaybeUninit<T>; RING_CAPACITY]>,
83
84 _pad1: [u8; CACHE_LINE],
86 head: AtomicUsize,
88 cached_tail: AtomicUsize,
90
91 _pad2: [u8; CACHE_LINE],
93 tail: AtomicUsize,
95 cached_head: AtomicUsize,
97
98 _pad3: [u8; CACHE_LINE],
100 overflow_count: AtomicUsize,
102 config: RingConfig,
104}
105
106impl<T> SpscRingBuffer<T> {
107 const CAPACITY: usize = RING_CAPACITY;
108
109 pub fn new(config: RingConfig) -> Self {
111 assert!(
112 config.capacity < Self::CAPACITY,
113 "capacity must be strictly less than RING_CAPACITY"
114 );
115
116 let buffer = UnsafeCell::new(unsafe {
118 MaybeUninit::<[MaybeUninit<T>; RING_CAPACITY]>::uninit().assume_init()
119 });
120
121 Self {
122 buffer,
123 _pad1: [0; CACHE_LINE],
124 head: AtomicUsize::new(0),
125 cached_tail: AtomicUsize::new(0),
126 _pad2: [0; CACHE_LINE],
127 tail: AtomicUsize::new(0),
128 cached_head: AtomicUsize::new(0),
129 _pad3: [0; CACHE_LINE],
130 overflow_count: AtomicUsize::new(0),
131 config,
132 }
133 }
134
135 pub fn push(&self, value: T) -> Result<(), T> {
140 let tail = self.tail.load(Ordering::Relaxed) & (Self::CAPACITY - 1);
144
145 let mut head = self.cached_head.load(Ordering::Relaxed);
147 let mut current_len = tail.wrapping_sub(head) & (Self::CAPACITY - 1);
148
149 if current_len >= self.config.capacity {
150 head = self.head.load(Ordering::Acquire);
152 self.cached_head.store(head, Ordering::Relaxed);
153 current_len = tail.wrapping_sub(head) & (Self::CAPACITY - 1);
154
155 if current_len >= self.config.capacity {
156 self.overflow_count.fetch_add(1, Ordering::Relaxed);
157 return Err(value);
158 }
159 }
160
161 let next_tail = (tail + 1) & (Self::CAPACITY - 1);
162
163 unsafe {
165 let buf_ptr = self.buffer.get() as *mut std::mem::MaybeUninit<T>;
166 (*buf_ptr.add(tail)).write(value);
167
168 self.tail.store(next_tail, Ordering::Release);
170 }
171
172 Ok(())
173 }
174
175 pub fn pop(&self) -> Option<T> {
179 let head = self.head.load(Ordering::Relaxed) & (Self::CAPACITY - 1);
183
184 let mut tail = self.cached_tail.load(Ordering::Relaxed);
186
187 if head == tail {
188 tail = self.tail.load(Ordering::Acquire);
190 self.cached_tail.store(tail, Ordering::Relaxed);
191
192 if head == tail {
193 return None;
194 }
195 }
196
197 let value = unsafe {
199 let buf_ptr = self.buffer.get() as *const std::mem::MaybeUninit<T>;
200 let val = (*buf_ptr.add(head)).assume_init_read();
201
202 self.head
204 .store((head + 1) & (Self::CAPACITY - 1), Ordering::Release);
205
206 val
207 };
208
209 Some(value)
210 }
211
212 pub fn len(&self) -> usize {
214 let head = self.head.load(Ordering::Acquire);
215 let tail = self.tail.load(Ordering::Acquire);
216
217 tail.wrapping_sub(head) & (Self::CAPACITY - 1)
218 }
219
220 pub fn is_empty(&self) -> bool {
222 let head = self.head.load(Ordering::Relaxed);
223 let mut tail = self.cached_tail.load(Ordering::Relaxed);
224
225 if head == tail {
226 tail = self.tail.load(Ordering::Acquire);
227 self.cached_tail.store(tail, Ordering::Relaxed);
228 }
229
230 head == tail
231 }
232
233 pub fn capacity(&self) -> usize {
235 self.config.capacity
236 }
237
238 pub fn overflow_count(&self) -> usize {
240 self.overflow_count.load(Ordering::Relaxed)
241 }
242
243 pub fn empty_count(&self) -> usize {
245 self.config.capacity.saturating_sub(self.len())
246 }
247
248 pub fn should_backpressure(&self) -> bool {
250 let tail = self.tail.load(Ordering::Relaxed);
251 let mut head = self.cached_head.load(Ordering::Relaxed);
252 let mut len = tail.wrapping_sub(head) & (Self::CAPACITY - 1);
253
254 if len >= self.config.backpressure_threshold {
255 head = self.head.load(Ordering::Acquire);
256 self.cached_head.store(head, Ordering::Relaxed);
257 len = tail.wrapping_sub(head) & (Self::CAPACITY - 1);
258 }
259
260 len >= self.config.backpressure_threshold
261 }
262
263 pub fn write_available(&self) -> usize {
265 let tail = self.tail.load(Ordering::Relaxed);
266 let mut head = self.cached_head.load(Ordering::Relaxed);
267 let mut len = tail.wrapping_sub(head) & (Self::CAPACITY - 1);
268 if len >= self.config.capacity {
269 head = self.head.load(Ordering::Acquire);
270 self.cached_head.store(head, Ordering::Relaxed);
271 len = tail.wrapping_sub(head) & (Self::CAPACITY - 1);
272 }
273 self.config.capacity.saturating_sub(len)
274 }
275
276 pub fn read_available(&self) -> usize {
278 let head = self.head.load(Ordering::Relaxed);
279 let mut tail = self.cached_tail.load(Ordering::Relaxed);
280 if head == tail {
281 tail = self.tail.load(Ordering::Acquire);
282 self.cached_tail.store(tail, Ordering::Relaxed);
283 }
284 tail.wrapping_sub(head) & (Self::CAPACITY - 1)
285 }
286
287 pub fn wait_for_space(&self) {
292 let mut backoff = 1u32;
293 loop {
294 if !self.should_backpressure() {
295 return;
296 }
297 let limit = backoff.min(256);
299 for _ in 0..limit {
300 core::hint::spin_loop();
301 }
302 backoff = backoff.saturating_mul(2);
303 if backoff > 256 {
304 std::thread::yield_now();
305 backoff = 1;
306 }
307 }
308 }
309
310 pub fn wait_for_data(&self) {
315 let mut backoff = 1u32;
316 loop {
317 if !self.is_empty() {
318 return;
319 }
320 let limit = backoff.min(256);
321 for _ in 0..limit {
322 core::hint::spin_loop();
323 }
324 backoff = backoff.saturating_mul(2);
325 if backoff > 256 {
326 std::thread::yield_now();
327 backoff = 1;
328 }
329 }
330 }
331
332 pub fn push_batch(&self, slice: &[T]) -> usize
340 where
341 T: Copy,
342 {
343 let tail = self.tail.load(Ordering::Relaxed);
344 let mut head = self.cached_head.load(Ordering::Relaxed);
345 let mut current_len = tail.wrapping_sub(head) & (Self::CAPACITY - 1);
346
347 if current_len >= self.config.capacity {
348 head = self.head.load(Ordering::Acquire);
349 self.cached_head.store(head, Ordering::Relaxed);
350 current_len = tail.wrapping_sub(head) & (Self::CAPACITY - 1);
351 if current_len >= self.config.capacity {
352 self.overflow_count
354 .fetch_add(slice.len(), Ordering::Relaxed);
355 return 0;
356 }
357 }
358
359 let available = self.config.capacity - current_len;
360 let count = available.min(slice.len());
361 if count == 0 {
362 return 0;
363 }
364 if count < slice.len() {
365 self.overflow_count
368 .fetch_add(slice.len() - count, Ordering::Relaxed);
369 }
370
371 let buf_ptr = self.buffer.get() as *mut T;
372 let mask = Self::CAPACITY - 1;
373 let start = tail & mask;
374 if start + count <= Self::CAPACITY {
375 unsafe {
376 std::ptr::copy_nonoverlapping(slice.as_ptr(), buf_ptr.add(start), count);
377 }
378 } else {
379 let first_chunk = Self::CAPACITY - start;
380 let second_chunk = count - first_chunk;
381 unsafe {
382 std::ptr::copy_nonoverlapping(slice.as_ptr(), buf_ptr.add(start), first_chunk);
383 std::ptr::copy_nonoverlapping(
384 slice.as_ptr().add(first_chunk),
385 buf_ptr,
386 second_chunk,
387 );
388 }
389 }
390 let new_tail = tail.wrapping_add(count) & mask;
391 self.tail.store(new_tail, Ordering::Release);
392 count
393 }
394
395 pub fn pop_batch(&self, out: &mut [std::mem::MaybeUninit<T>]) -> usize {
401 let head = self.head.load(Ordering::Relaxed);
402 let mut tail = self.cached_tail.load(Ordering::Relaxed);
403
404 if head == tail {
405 tail = self.tail.load(Ordering::Acquire);
406 self.cached_tail.store(tail, Ordering::Relaxed);
407 if head == tail {
408 return 0;
409 }
410 }
411
412 let available = tail.wrapping_sub(head) & (Self::CAPACITY - 1);
413 let count = available.min(out.len());
414 if count == 0 {
415 return 0;
416 }
417
418 let buf_ptr = self.buffer.get() as *const T;
419 let mask = Self::CAPACITY - 1;
420 let start = head & mask;
421 let out_ptr = out.as_mut_ptr() as *mut T;
422 if start + count <= Self::CAPACITY {
423 unsafe {
424 std::ptr::copy_nonoverlapping(buf_ptr.add(start), out_ptr, count);
425 }
426 } else {
427 let first_chunk = Self::CAPACITY - start;
428 let second_chunk = count - first_chunk;
429 unsafe {
430 std::ptr::copy_nonoverlapping(buf_ptr.add(start), out_ptr, first_chunk);
431 std::ptr::copy_nonoverlapping(buf_ptr, out_ptr.add(first_chunk), second_chunk);
432 }
433 }
434 let new_head = head.wrapping_add(count) & mask;
435 self.head.store(new_head, Ordering::Release);
436 count
437 }
438}
439
440unsafe impl<T: Send> Send for SpscRingBuffer<T> {}
444unsafe impl<T: Send> Sync for SpscRingBuffer<T> {}
445
446impl<T> Drop for SpscRingBuffer<T> {
447 fn drop(&mut self) {
448 while self.pop().is_some() {}
450 }
451}
452
453#[cfg(test)]
454mod tests {
455 use super::*;
456 use std::sync::Arc;
457 use std::thread;
458
459 #[test]
460 fn supports_fifo_push_and_pop() {
461 let ring = SpscRingBuffer::<i32>::new(RingConfig::default());
462 assert!(ring.push(1).is_ok());
463 assert!(ring.push(2).is_ok());
464 assert_eq!(ring.len(), 2);
465 assert_eq!(ring.pop(), Some(1));
466 assert_eq!(ring.pop(), Some(2));
467 assert!(ring.is_empty());
468 }
469
470 #[test]
471 fn returns_full_when_capacity_is_reached() {
472 let config = RingConfig {
473 capacity: 2,
474 backpressure_threshold: 1,
475 };
476 let ring = SpscRingBuffer::<i32>::new(config);
477 assert!(ring.push(1).is_ok());
478 assert!(ring.push(2).is_ok());
479 assert_eq!(ring.push(3), Err(3));
480 }
481
482 #[test]
483 fn supports_single_producer_single_consumer_threads() {
484 let ring = Arc::new(SpscRingBuffer::<i32>::new(RingConfig::default()));
485 let producer_ring = Arc::clone(&ring);
486 let consumer_ring = Arc::clone(&ring);
487
488 let producer = thread::spawn(move || {
489 for value in 0..1_000 {
490 loop {
491 if producer_ring.push(value).is_ok() {
492 break;
493 }
494 std::thread::yield_now();
495 }
496 }
497 });
498
499 let consumer = thread::spawn(move || {
500 let mut values = Vec::new();
501 while values.len() < 1_000 {
502 if let Some(value) = consumer_ring.pop() {
503 values.push(value);
504 } else {
505 std::thread::yield_now();
506 }
507 }
508 values
509 });
510
511 producer.join().unwrap();
512 let values = consumer.join().unwrap();
513
514 assert_eq!(values.len(), 1_000);
515 assert_eq!(values.first(), Some(&0));
516 assert_eq!(values.last(), Some(&999));
517 }
518
519 #[test]
520 fn handles_wrap_around_correctly() {
521 let ring = SpscRingBuffer::<i32>::new(RingConfig::default());
522
523 for i in 0..1020 {
525 ring.push(i).unwrap();
526 }
527
528 assert_eq!(ring.len(), 1020);
529
530 for i in 0..1020 {
532 assert_eq!(ring.pop(), Some(i));
533 }
534
535 assert!(ring.is_empty());
536 }
537
538 #[test]
539 fn backpressure_threshold_works() {
540 let config = RingConfig {
541 capacity: 100,
542 backpressure_threshold: 70,
543 };
544 let ring = SpscRingBuffer::<i32>::new(config);
545
546 for i in 0..69 {
548 ring.push(i).unwrap();
549 }
550 assert!(!ring.should_backpressure());
551
552 let result = ring.push(70);
554 assert!(result.is_ok()); assert!(ring.should_backpressure());
556 }
557
558 #[test]
559 fn wait_for_space_releases() {
560 let ring = Arc::new(SpscRingBuffer::<i32>::new(RingConfig::default()));
561 let consumer_ring = Arc::clone(&ring);
562
563 for i in 0..1020 {
565 ring.push(i).unwrap();
566 }
567
568 consumer_ring.pop();
570
571 let result = ring.push(1021);
573 assert!(result.is_ok());
574 }
575
576 #[test]
577 fn wait_for_data_releases() {
578 let ring = Arc::new(SpscRingBuffer::<i32>::new(RingConfig::default()));
579 let producer_ring = Arc::clone(&ring);
580
581 producer_ring.push(42).unwrap();
585
586 assert_eq!(ring.pop(), Some(42));
588 }
589
590 #[test]
591 fn push_batch_writes_multiple_elements() {
592 let ring = SpscRingBuffer::<u64>::new(RingConfig::default());
593 let data: Vec<u64> = (0..100).collect();
594 let n = ring.push_batch(&data);
595 assert_eq!(n, 100);
596 assert_eq!(ring.len(), 100);
597 }
598
599 #[test]
600 fn push_batch_respects_capacity() {
601 let config = RingConfig {
602 capacity: 10,
603 backpressure_threshold: 8,
604 };
605 let ring = SpscRingBuffer::<u64>::new(config);
606 let data: Vec<u64> = (0..100).collect();
607 let n = ring.push_batch(&data);
608 assert_eq!(n, 10);
609 assert_eq!(ring.len(), 10);
610 assert_eq!(ring.overflow_count(), 90);
613 }
614
615 #[test]
616 fn pop_batch_reads_multiple_elements() {
617 let ring = SpscRingBuffer::<u64>::new(RingConfig::default());
618 for i in 0..50u64 {
619 ring.push(i).unwrap();
620 }
621 let mut out = vec![std::mem::MaybeUninit::uninit(); 20];
622 let n = ring.pop_batch(&mut out);
623 assert_eq!(n, 20);
624 let drained: Vec<u64> = out[..n]
625 .iter()
626 .map(|x| unsafe { x.assume_init_read() })
627 .collect();
628 assert_eq!(drained, (0..20).collect::<Vec<u64>>());
629 assert_eq!(ring.len(), 30);
630 }
631
632 #[test]
633 fn pop_batch_returns_available_count() {
634 let ring = SpscRingBuffer::<u64>::new(RingConfig::default());
635 for i in 0..5u64 {
636 ring.push(i).unwrap();
637 }
638 let mut out = vec![std::mem::MaybeUninit::uninit(); 100];
639 let n = ring.pop_batch(&mut out);
640 assert_eq!(n, 5);
641 }
642
643 #[test]
644 fn push_pop_batch_round_trip_spsc() {
645 let ring = Arc::new(SpscRingBuffer::<u64>::new(RingConfig::default()));
646 let prod = Arc::clone(&ring);
647 let cons = Arc::clone(&ring);
648 let iters = 5_000u64;
649 let batch_size = 64;
650
651 let h_prod = std::thread::spawn(move || {
652 let mut sent = 0u64;
653 while sent < iters {
654 let chunk_size = batch_size.min((iters - sent) as usize);
655 let chunk: Vec<u64> = (sent..sent + chunk_size as u64).collect();
656 let n = prod.push_batch(&chunk);
657 sent += n as u64;
658 if n == 0 {
659 prod.wait_for_space();
660 }
661 }
662 });
663
664 let h_cons = std::thread::spawn(move || {
665 let mut received = 0u64;
666 let mut buf = vec![std::mem::MaybeUninit::uninit(); batch_size];
667 while received < iters {
668 let n = cons.pop_batch(&mut buf);
669 if n > 0 {
670 received += n as u64;
671 } else {
672 cons.wait_for_data();
673 }
674 }
675 });
676
677 h_prod.join().unwrap();
678 h_cons.join().unwrap();
679 assert!(ring.is_empty());
680 }
681
682 #[test]
683 fn write_read_available_tracking() {
684 let config = RingConfig {
685 capacity: 100,
686 backpressure_threshold: 70,
687 };
688 let ring = SpscRingBuffer::<u64>::new(config);
689 assert_eq!(ring.write_available(), 100);
690 assert_eq!(ring.read_available(), 0);
691
692 for i in 0..30u64 {
693 ring.push(i).unwrap();
694 }
695 assert_eq!(ring.write_available(), 70);
696 assert_eq!(ring.read_available(), 30);
697
698 ring.pop();
702 assert_eq!(ring.read_available(), 29);
703 let wa = ring.write_available();
704 assert!(
705 wa == 70 || wa == 71,
706 "write_available should be 70 (conservative) or 71 (true): got {wa}"
707 );
708 }
709}