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: 1023, backpressure_threshold: 716, }
}
}
const RING_CAPACITY: usize = 1024;
#[derive(Debug)]
pub struct SpscRingBuffer<T> {
buffer: UnsafeCell<[MaybeUninit<T>; RING_CAPACITY]>,
head: AtomicUsize,
tail: AtomicUsize,
overflow_count: AtomicUsize,
empty_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,
head: AtomicUsize::new(0),
tail: AtomicUsize::new(0),
overflow_count: AtomicUsize::new(0),
empty_count: AtomicUsize::new(Self::CAPACITY - config.capacity),
config,
}
}
pub fn push(&self, value: T) -> Result<(), T> {
let tail = self.tail.load(Ordering::Acquire);
let head = self.head.load(Ordering::Acquire);
let current_len = if tail >= head {
tail - head
} else {
Self::CAPACITY - head + tail
};
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::Acquire);
let tail = self.tail.load(Ordering::Acquire);
let is_empty = if head >= tail { head == tail } else { false };
if is_empty {
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 {
self.head.load(Ordering::Acquire) == self.tail.load(Ordering::Acquire)
}
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.empty_count.load(Ordering::Relaxed)
}
pub fn should_backpressure(&self) -> bool {
let len = self.len();
len >= self.config.backpressure_threshold
}
pub fn wait_for_space(&self) {
while self.should_backpressure() {
std::thread::yield_now();
}
}
pub fn wait_for_data(&self) {
while self.is_empty() {
std::thread::yield_now();
}
}
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));
}
}