use std::cell::UnsafeCell;
use std::mem::MaybeUninit;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::Arc;
#[repr(align(64))]
struct CachePadded<T>(T);
struct Ring<T> {
slots: Box<[UnsafeCell<MaybeUninit<T>>]>,
cap: usize,
head: CachePadded<AtomicUsize>,
tail: CachePadded<AtomicUsize>,
closed: AtomicBool,
}
unsafe impl<T: Send> Send for Ring<T> {}
unsafe impl<T: Send> Sync for Ring<T> {}
impl<T> Drop for Ring<T> {
fn drop(&mut self) {
let head = *self.head.0.get_mut();
let tail = *self.tail.0.get_mut();
for pos in head..tail {
let idx = pos % self.cap;
unsafe { (*self.slots[idx].get()).assume_init_drop() };
}
}
}
pub struct Producer<T> {
ring: Arc<Ring<T>>,
}
pub struct Consumer<T> {
ring: Arc<Ring<T>>,
}
pub use crate::TrySendError;
pub fn channel<T>(capacity: usize) -> (Producer<T>, Consumer<T>) {
assert!(capacity >= 1, "SPSC capacity must be >= 1");
let mut slots = Vec::with_capacity(capacity);
for _ in 0..capacity {
slots.push(UnsafeCell::new(MaybeUninit::uninit()));
}
let ring = Arc::new(Ring {
slots: slots.into_boxed_slice(),
cap: capacity,
head: CachePadded(AtomicUsize::new(0)),
tail: CachePadded(AtomicUsize::new(0)),
closed: AtomicBool::new(false),
});
(
Producer { ring: ring.clone() },
Consumer { ring },
)
}
impl<T> Producer<T> {
pub fn try_push(&self, item: T) -> Result<(), TrySendError<T>> {
let ring = &self.ring;
let tail = ring.tail.0.load(Ordering::Relaxed);
let head = ring.head.0.load(Ordering::Acquire);
if tail - head == ring.cap {
return Err(TrySendError::Full(item));
}
let idx = tail % ring.cap;
unsafe { (*ring.slots[idx].get()).write(item) };
ring.tail.0.store(tail + 1, Ordering::Release);
Ok(())
}
pub fn push_blocking(&self, mut item: T) {
loop {
match self.try_push(item) {
Ok(()) => return,
Err(TrySendError::Full(returned)) | Err(TrySendError::Closed(returned)) => {
item = returned;
std::hint::spin_loop();
}
}
}
}
}
impl<T> Drop for Producer<T> {
fn drop(&mut self) {
self.ring.closed.store(true, Ordering::Release);
}
}
impl<T> Consumer<T> {
pub fn try_pop(&self) -> Option<T> {
let ring = &self.ring;
let head = ring.head.0.load(Ordering::Relaxed);
let tail = ring.tail.0.load(Ordering::Acquire);
if head == tail {
return None;
}
let idx = head % ring.cap;
let item = unsafe { (*ring.slots[idx].get()).assume_init_read() };
ring.head.0.store(head + 1, Ordering::Release);
Some(item)
}
pub fn is_closed(&self) -> bool {
self.ring.closed.load(Ordering::Acquire)
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::atomic::AtomicUsize;
#[test]
fn push_pop_fifo_order() {
let (tx, rx) = channel::<u32>(4);
tx.try_push(1).ok();
tx.try_push(2).ok();
tx.try_push(3).ok();
assert_eq!(rx.try_pop(), Some(1));
assert_eq!(rx.try_pop(), Some(2));
assert_eq!(rx.try_pop(), Some(3));
assert_eq!(rx.try_pop(), None);
}
#[test]
fn full_returns_item() {
let (tx, _rx) = channel::<u32>(2);
assert!(tx.try_push(1).is_ok());
assert!(tx.try_push(2).is_ok());
match tx.try_push(3) {
Err(TrySendError::Full(v)) => assert_eq!(v, 3),
other => panic!("expected Full, got {other:?}"),
}
}
#[test]
fn wraps_around_many_times() {
let (tx, rx) = channel::<usize>(3);
for i in 0..1000 {
tx.try_push(i).expect("space (drained each iter)");
assert_eq!(rx.try_pop(), Some(i));
}
assert_eq!(rx.try_pop(), None);
}
#[test]
fn drops_unconsumed_items_on_teardown() {
static DROPS: AtomicUsize = AtomicUsize::new(0);
struct Tracked;
impl Drop for Tracked {
fn drop(&mut self) {
DROPS.fetch_add(1, Ordering::SeqCst);
}
}
{
let (tx, _rx) = channel::<Tracked>(4);
tx.try_push(Tracked).ok();
tx.try_push(Tracked).ok();
}
assert_eq!(DROPS.load(Ordering::SeqCst), 2);
}
#[test]
fn closed_after_producer_drops() {
let (tx, rx) = channel::<u8>(2);
assert!(!rx.is_closed());
drop(tx);
assert!(rx.is_closed());
}
#[test]
fn concurrent_producer_consumer() {
let (tx, rx) = channel::<u64>(64);
let producer = std::thread::spawn(move || {
for i in 0..100_000u64 {
tx.push_blocking(i);
}
});
let mut next = 0u64;
while next < 100_000 {
if let Some(v) = rx.try_pop() {
assert_eq!(v, next); next += 1;
} else {
std::hint::spin_loop();
}
}
producer.join().unwrap();
}
}