use std::cell::UnsafeCell;
use std::mem::MaybeUninit;
use std::sync::atomic::{AtomicUsize, Ordering};
#[derive(Clone, Copy, Debug)]
pub struct RingConfig {
pub capacity: usize,
pub backpressure_threshold: usize,
}
impl Default for RingConfig {
fn default() -> Self {
Self {
capacity: 4095, backpressure_threshold: 2867, }
}
}
const RING_CAPACITY: usize = 4096;
const CACHE_LINE: usize = 128;
#[inline(always)]
#[allow(dead_code)]
fn prefetch_read<T>(ptr: *const T) {
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
unsafe {
std::arch::x86_64::_mm_prefetch(ptr as *const i8, std::arch::x86_64::_MM_HINT_T0);
}
#[cfg(not(any(target_arch = "x86", target_arch = "x86_64")))]
{
std::hint::black_box(ptr);
}
}
#[inline(always)]
#[allow(dead_code)]
fn prefetch_write<T>(ptr: *mut T) {
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
unsafe {
std::arch::x86_64::_mm_prefetch(ptr as *const i8, std::arch::x86_64::_MM_HINT_ET1);
}
#[cfg(not(any(target_arch = "x86", target_arch = "x86_64")))]
{
std::hint::black_box(ptr);
}
}
#[derive(Debug)]
pub struct SpscRingBuffer<T> {
buffer: UnsafeCell<[MaybeUninit<T>; RING_CAPACITY]>,
_pad1: [u8; CACHE_LINE],
head: AtomicUsize,
cached_tail: AtomicUsize,
_pad2: [u8; CACHE_LINE],
tail: AtomicUsize,
cached_head: AtomicUsize,
_pad3: [u8; CACHE_LINE],
overflow_count: AtomicUsize,
config: RingConfig,
}
impl<T> SpscRingBuffer<T> {
const CAPACITY: usize = RING_CAPACITY;
pub fn new(config: RingConfig) -> Self {
assert!(
config.capacity < Self::CAPACITY,
"capacity must be strictly less than RING_CAPACITY"
);
let buffer = UnsafeCell::new(unsafe {
MaybeUninit::<[MaybeUninit<T>; RING_CAPACITY]>::uninit().assume_init()
});
Self {
buffer,
_pad1: [0; CACHE_LINE],
head: AtomicUsize::new(0),
cached_tail: AtomicUsize::new(0),
_pad2: [0; CACHE_LINE],
tail: AtomicUsize::new(0),
cached_head: AtomicUsize::new(0),
_pad3: [0; CACHE_LINE],
overflow_count: AtomicUsize::new(0),
config,
}
}
pub fn push(&self, value: T) -> Result<(), T> {
let tail = self.tail.load(Ordering::Relaxed);
let mut head = self.cached_head.load(Ordering::Relaxed);
let mut current_len = tail.wrapping_sub(head) & (Self::CAPACITY - 1);
if current_len >= self.config.capacity {
head = self.head.load(Ordering::Acquire);
self.cached_head.store(head, Ordering::Relaxed);
current_len = tail.wrapping_sub(head) & (Self::CAPACITY - 1);
if current_len >= self.config.capacity {
self.overflow_count.fetch_add(1, Ordering::Relaxed);
return Err(value);
}
}
let next_tail = Self::increment(tail);
unsafe {
let buf = &mut *self.buffer.get();
buf[tail].write(value);
self.tail.store(next_tail, Ordering::Release);
}
Ok(())
}
pub fn pop(&self) -> Option<T> {
let head = self.head.load(Ordering::Relaxed);
let mut tail = self.cached_tail.load(Ordering::Relaxed);
if head == tail {
tail = self.tail.load(Ordering::Acquire);
self.cached_tail.store(tail, Ordering::Relaxed);
if head == tail {
return None;
}
}
let value = unsafe {
let buf = &mut *self.buffer.get();
let val = buf[head].assume_init_read();
self.head.store(Self::increment(head), Ordering::Release);
val
};
Some(value)
}
pub fn len(&self) -> usize {
let head = self.head.load(Ordering::Acquire);
let tail = self.tail.load(Ordering::Acquire);
tail.wrapping_sub(head) & (Self::CAPACITY - 1)
}
pub fn is_empty(&self) -> bool {
let head = self.head.load(Ordering::Relaxed);
let mut tail = self.cached_tail.load(Ordering::Relaxed);
if head == tail {
tail = self.tail.load(Ordering::Acquire);
self.cached_tail.store(tail, Ordering::Relaxed);
}
head == tail
}
pub fn capacity(&self) -> usize {
self.config.capacity
}
pub fn overflow_count(&self) -> usize {
self.overflow_count.load(Ordering::Relaxed)
}
pub fn empty_count(&self) -> usize {
self.config.capacity.saturating_sub(self.len())
}
pub fn should_backpressure(&self) -> bool {
let tail = self.tail.load(Ordering::Relaxed);
let mut head = self.cached_head.load(Ordering::Relaxed);
let mut len = tail.wrapping_sub(head) & (Self::CAPACITY - 1);
if len >= self.config.backpressure_threshold {
head = self.head.load(Ordering::Acquire);
self.cached_head.store(head, Ordering::Relaxed);
len = tail.wrapping_sub(head) & (Self::CAPACITY - 1);
}
len >= self.config.backpressure_threshold
}
pub fn write_available(&self) -> usize {
let tail = self.tail.load(Ordering::Relaxed);
let mut head = self.cached_head.load(Ordering::Relaxed);
let mut len = tail.wrapping_sub(head) & (Self::CAPACITY - 1);
if len >= self.config.capacity {
head = self.head.load(Ordering::Acquire);
self.cached_head.store(head, Ordering::Relaxed);
len = tail.wrapping_sub(head) & (Self::CAPACITY - 1);
}
self.config.capacity.saturating_sub(len)
}
pub fn read_available(&self) -> usize {
let head = self.head.load(Ordering::Relaxed);
let mut tail = self.cached_tail.load(Ordering::Relaxed);
if head == tail {
tail = self.tail.load(Ordering::Acquire);
self.cached_tail.store(tail, Ordering::Relaxed);
}
tail.wrapping_sub(head) & (Self::CAPACITY - 1)
}
pub fn wait_for_space(&self) {
let mut backoff = 1u32;
loop {
if !self.should_backpressure() {
return;
}
let limit = backoff.min(256);
for _ in 0..limit {
core::hint::spin_loop();
}
backoff = backoff.saturating_mul(2);
if backoff > 256 {
std::thread::yield_now();
backoff = 1;
}
}
}
pub fn wait_for_data(&self) {
let mut backoff = 1u32;
loop {
if !self.is_empty() {
return;
}
let limit = backoff.min(256);
for _ in 0..limit {
core::hint::spin_loop();
}
backoff = backoff.saturating_mul(2);
if backoff > 256 {
std::thread::yield_now();
backoff = 1;
}
}
}
pub fn push_batch(&self, slice: &[T]) -> usize
where
T: Copy,
{
let tail = self.tail.load(Ordering::Relaxed);
let mut head = self.cached_head.load(Ordering::Relaxed);
let mut current_len = tail.wrapping_sub(head) & (Self::CAPACITY - 1);
if current_len >= self.config.capacity {
head = self.head.load(Ordering::Acquire);
self.cached_head.store(head, Ordering::Relaxed);
current_len = tail.wrapping_sub(head) & (Self::CAPACITY - 1);
if current_len >= self.config.capacity {
self.overflow_count
.fetch_add(slice.len(), Ordering::Relaxed);
return 0;
}
}
let available = self.config.capacity - current_len;
let count = available.min(slice.len());
if count == 0 {
return 0;
}
if count < slice.len() {
self.overflow_count
.fetch_add(slice.len() - count, Ordering::Relaxed);
}
let buf = unsafe { &mut *self.buffer.get() };
let mask = Self::CAPACITY - 1;
let start = tail & mask;
let buf_ptr = buf.as_mut_ptr() as *mut T;
if start + count <= Self::CAPACITY {
unsafe {
std::ptr::copy_nonoverlapping(slice.as_ptr(), buf_ptr.add(start), count);
}
} else {
let first_chunk = Self::CAPACITY - start;
let second_chunk = count - first_chunk;
unsafe {
std::ptr::copy_nonoverlapping(slice.as_ptr(), buf_ptr.add(start), first_chunk);
std::ptr::copy_nonoverlapping(
slice.as_ptr().add(first_chunk),
buf_ptr,
second_chunk,
);
}
}
let new_tail = tail.wrapping_add(count) & mask;
self.tail.store(new_tail, Ordering::Release);
count
}
pub fn pop_batch(&self, out: &mut [std::mem::MaybeUninit<T>]) -> usize {
let head = self.head.load(Ordering::Relaxed);
let mut tail = self.cached_tail.load(Ordering::Relaxed);
if head == tail {
tail = self.tail.load(Ordering::Acquire);
self.cached_tail.store(tail, Ordering::Relaxed);
if head == tail {
return 0;
}
}
let available = tail.wrapping_sub(head) & (Self::CAPACITY - 1);
let count = available.min(out.len());
if count == 0 {
return 0;
}
let buf = unsafe { &mut *self.buffer.get() };
let mask = Self::CAPACITY - 1;
let start = head & mask;
let buf_ptr = buf.as_ptr() as *const T;
let out_ptr = out.as_mut_ptr() as *mut T;
if start + count <= Self::CAPACITY {
unsafe {
std::ptr::copy_nonoverlapping(buf_ptr.add(start), out_ptr, count);
}
} else {
let first_chunk = Self::CAPACITY - start;
let second_chunk = count - first_chunk;
unsafe {
std::ptr::copy_nonoverlapping(buf_ptr.add(start), out_ptr, first_chunk);
std::ptr::copy_nonoverlapping(buf_ptr, out_ptr.add(first_chunk), second_chunk);
}
}
let new_head = head.wrapping_add(count) & mask;
self.head.store(new_head, Ordering::Release);
count
}
fn increment(index: usize) -> usize {
(index + 1) & (Self::CAPACITY - 1)
}
}
unsafe impl<T: Send> Send for SpscRingBuffer<T> {}
unsafe impl<T: Send> Sync for SpscRingBuffer<T> {}
impl<T> Drop for SpscRingBuffer<T> {
fn drop(&mut self) {
while self.pop().is_some() {}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::Arc;
use std::thread;
#[test]
fn supports_fifo_push_and_pop() {
let ring = SpscRingBuffer::<i32>::new(RingConfig::default());
assert!(ring.push(1).is_ok());
assert!(ring.push(2).is_ok());
assert_eq!(ring.len(), 2);
assert_eq!(ring.pop(), Some(1));
assert_eq!(ring.pop(), Some(2));
assert!(ring.is_empty());
}
#[test]
fn returns_full_when_capacity_is_reached() {
let config = RingConfig {
capacity: 2,
backpressure_threshold: 1,
};
let ring = SpscRingBuffer::<i32>::new(config);
assert!(ring.push(1).is_ok());
assert!(ring.push(2).is_ok());
assert_eq!(ring.push(3), Err(3));
}
#[test]
fn supports_single_producer_single_consumer_threads() {
let ring = Arc::new(SpscRingBuffer::<i32>::new(RingConfig::default()));
let producer_ring = Arc::clone(&ring);
let consumer_ring = Arc::clone(&ring);
let producer = thread::spawn(move || {
for value in 0..1_000 {
loop {
if producer_ring.push(value).is_ok() {
break;
}
std::thread::yield_now();
}
}
});
let consumer = thread::spawn(move || {
let mut values = Vec::new();
while values.len() < 1_000 {
if let Some(value) = consumer_ring.pop() {
values.push(value);
} else {
std::thread::yield_now();
}
}
values
});
producer.join().unwrap();
let values = consumer.join().unwrap();
assert_eq!(values.len(), 1_000);
assert_eq!(values.first(), Some(&0));
assert_eq!(values.last(), Some(&999));
}
#[test]
fn handles_wrap_around_correctly() {
let ring = SpscRingBuffer::<i32>::new(RingConfig::default());
for i in 0..1020 {
ring.push(i).unwrap();
}
assert_eq!(ring.len(), 1020);
for i in 0..1020 {
assert_eq!(ring.pop(), Some(i));
}
assert!(ring.is_empty());
}
#[test]
fn backpressure_threshold_works() {
let config = RingConfig {
capacity: 100,
backpressure_threshold: 70,
};
let ring = SpscRingBuffer::<i32>::new(config);
for i in 0..69 {
ring.push(i).unwrap();
}
assert!(!ring.should_backpressure());
let result = ring.push(70);
assert!(result.is_ok()); assert!(ring.should_backpressure());
}
#[test]
fn wait_for_space_releases() {
let ring = Arc::new(SpscRingBuffer::<i32>::new(RingConfig::default()));
let consumer_ring = Arc::clone(&ring);
for i in 0..1020 {
ring.push(i).unwrap();
}
consumer_ring.pop();
let result = ring.push(1021);
assert!(result.is_ok());
}
#[test]
fn wait_for_data_releases() {
let ring = Arc::new(SpscRingBuffer::<i32>::new(RingConfig::default()));
let producer_ring = Arc::clone(&ring);
producer_ring.push(42).unwrap();
assert_eq!(ring.pop(), Some(42));
}
#[test]
fn push_batch_writes_multiple_elements() {
let ring = SpscRingBuffer::<u64>::new(RingConfig::default());
let data: Vec<u64> = (0..100).collect();
let n = ring.push_batch(&data);
assert_eq!(n, 100);
assert_eq!(ring.len(), 100);
}
#[test]
fn push_batch_respects_capacity() {
let config = RingConfig {
capacity: 10,
backpressure_threshold: 8,
};
let ring = SpscRingBuffer::<u64>::new(config);
let data: Vec<u64> = (0..100).collect();
let n = ring.push_batch(&data);
assert_eq!(n, 10);
assert_eq!(ring.len(), 10);
assert_eq!(ring.overflow_count(), 90);
}
#[test]
fn pop_batch_reads_multiple_elements() {
let ring = SpscRingBuffer::<u64>::new(RingConfig::default());
for i in 0..50u64 {
ring.push(i).unwrap();
}
let mut out = vec![std::mem::MaybeUninit::uninit(); 20];
let n = ring.pop_batch(&mut out);
assert_eq!(n, 20);
let drained: Vec<u64> = out[..n]
.iter()
.map(|x| unsafe { x.assume_init_read() })
.collect();
assert_eq!(drained, (0..20).collect::<Vec<u64>>());
assert_eq!(ring.len(), 30);
}
#[test]
fn pop_batch_returns_available_count() {
let ring = SpscRingBuffer::<u64>::new(RingConfig::default());
for i in 0..5u64 {
ring.push(i).unwrap();
}
let mut out = vec![std::mem::MaybeUninit::uninit(); 100];
let n = ring.pop_batch(&mut out);
assert_eq!(n, 5);
}
#[test]
fn push_pop_batch_round_trip_spsc() {
let ring = Arc::new(SpscRingBuffer::<u64>::new(RingConfig::default()));
let prod = Arc::clone(&ring);
let cons = Arc::clone(&ring);
let iters = 5_000u64;
let batch_size = 64;
let h_prod = std::thread::spawn(move || {
let mut sent = 0u64;
while sent < iters {
let chunk_size = batch_size.min((iters - sent) as usize);
let chunk: Vec<u64> = (sent..sent + chunk_size as u64).collect();
let n = prod.push_batch(&chunk);
sent += n as u64;
if n == 0 {
prod.wait_for_space();
}
}
});
let h_cons = std::thread::spawn(move || {
let mut received = 0u64;
let mut buf = vec![std::mem::MaybeUninit::uninit(); batch_size];
while received < iters {
let n = cons.pop_batch(&mut buf);
if n > 0 {
received += n as u64;
} else {
cons.wait_for_data();
}
}
});
h_prod.join().unwrap();
h_cons.join().unwrap();
assert!(ring.is_empty());
}
#[test]
fn write_read_available_tracking() {
let config = RingConfig {
capacity: 100,
backpressure_threshold: 70,
};
let ring = SpscRingBuffer::<u64>::new(config);
assert_eq!(ring.write_available(), 100);
assert_eq!(ring.read_available(), 0);
for i in 0..30u64 {
ring.push(i).unwrap();
}
assert_eq!(ring.write_available(), 70);
assert_eq!(ring.read_available(), 30);
ring.pop();
assert_eq!(ring.read_available(), 29);
let wa = ring.write_available();
assert!(
wa == 70 || wa == 71,
"write_available should be 70 (conservative) or 71 (true): got {wa}"
);
}
}