use std::cell::UnsafeCell;
use std::mem::MaybeUninit;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
#[repr(align(128))]
pub(crate) struct Padded(pub(crate) AtomicUsize);
pub(crate) struct Inner<T> {
pub(crate) head: Padded,
pub(crate) tail: Padded,
pub(crate) capacity: usize,
pub(crate) mask: usize,
pub(crate) buf: Box<[UnsafeCell<MaybeUninit<T>>]>,
}
unsafe impl<T: Send> Sync for Inner<T> {}
unsafe impl<T: Send> Send for Inner<T> {}
impl<T> Drop for Inner<T> {
fn drop(&mut self) {
let head = self.head.0.load(Ordering::Relaxed);
let tail = self.tail.0.load(Ordering::Relaxed);
for i in head..tail {
unsafe {
(*self.buf[i & self.mask].get()).assume_init_drop();
}
}
}
}
pub struct Producer<T> {
pub(crate) inner: Arc<Inner<T>>,
pub(crate) cached_head: usize,
}
unsafe impl<T: Send> Send for Producer<T> {}
pub struct Consumer<T> {
pub(crate) inner: Arc<Inner<T>>,
pub(crate) cached_tail: usize,
}
unsafe impl<T: Send> Send for Consumer<T> {}
pub struct SpscRingBuffer;
impl SpscRingBuffer {
pub fn with_capacity<T>(requested_capacity: usize) -> (Producer<T>, Consumer<T>) {
let cap = requested_capacity.max(2).next_power_of_two();
let mut buf = Vec::with_capacity(cap);
for _ in 0..cap {
buf.push(UnsafeCell::new(MaybeUninit::<T>::uninit()));
}
let inner = Arc::new(Inner {
head: Padded(AtomicUsize::new(0)),
tail: Padded(AtomicUsize::new(0)),
capacity: cap,
mask: cap - 1,
buf: buf.into_boxed_slice(),
});
(
Producer {
inner: inner.clone(),
cached_head: 0,
},
Consumer {
inner,
cached_tail: 0,
},
)
}
}
impl<T> Producer<T> {
pub fn try_push(&mut self, value: T) -> Result<(), T> {
let tail = self.inner.tail.0.load(Ordering::Relaxed);
if tail.wrapping_sub(self.cached_head) == self.inner.capacity {
self.cached_head = self.inner.head.0.load(Ordering::Acquire);
if tail.wrapping_sub(self.cached_head) == self.inner.capacity {
return Err(value);
}
}
unsafe {
(*self.inner.buf[tail & self.inner.mask].get()).write(value);
}
self.inner
.tail
.0
.store(tail.wrapping_add(1), Ordering::Release);
Ok(())
}
pub fn capacity(&self) -> usize {
self.inner.capacity
}
}
impl<T> Consumer<T> {
pub fn try_pop(&mut self) -> Option<T> {
let head = self.inner.head.0.load(Ordering::Relaxed);
if head == self.cached_tail {
self.cached_tail = self.inner.tail.0.load(Ordering::Acquire);
if head == self.cached_tail {
return None;
}
}
let value = unsafe { (*self.inner.buf[head & self.inner.mask].get()).assume_init_read() };
self.inner
.head
.0
.store(head.wrapping_add(1), Ordering::Release);
Some(value)
}
pub fn capacity(&self) -> usize {
self.inner.capacity
}
}
#[cfg(feature = "harness")]
pub mod recipe;
#[cfg(any(
feature = "bulk",
feature = "wait-strategies",
feature = "mpsc-fan-in",
feature = "mpmc-disruptor",
feature = "metrics",
))]
pub mod features;
#[cfg(feature = "metrics")]
pub use features::metrics::{InstrumentedSpsc, RingMetrics, RingMetricsSnapshot};
#[cfg(feature = "mpmc-disruptor")]
pub use features::mpmc_disruptor::{DisruptorConsumer, DisruptorProducer, MpmcDisruptor};
#[cfg(feature = "mpsc-fan-in")]
pub use features::mpsc_fan_in::{MpscFanIn, MpscFanInConsumer, MpscFanInProducer};
#[cfg(feature = "wait-strategies")]
pub use features::wait_strategies::{
BlockingSpscConsumer, BlockingSpscProducer, BusySpin, ParkStrategy, WaitStrategy, YieldStrategy,
};