hydra_sync/

channel.rs

//! A single-producer, single-consumer (SPSC) `lock-free` ring buffer for [`Bytes`](https://crates.io/crates/bytes) handles.
use anyhow::Result;
use bytes::Bytes;
use std::cell::UnsafeCell;
use std::mem::MaybeUninit;
use std::ops::{Deref, DerefMut};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use tokio::sync::Notify;

#[allow(unused)]
/// Cache line ~size on x86_64 / ARM.
const CACHE_LINE: usize = 64;

/// Forces a value onto its own cache line, preventing false sharing between threads.
#[repr(C, align(64))]
struct CachePadded<T>(T);

impl<T> CachePadded<T> {
    #[inline(always)]
    const fn pad(val: T) -> Self {
        Self(val)
    }
}

impl<T> Deref for CachePadded<T> {
    type Target = T;
    #[inline(always)]
    fn deref(&self) -> &T {
        &self.0
    }
}

impl<T> DerefMut for CachePadded<T> {
    #[inline(always)]
    fn deref_mut(&mut self) -> &mut T {
        &mut self.0
    }
}

/// The ring buffer is full; the push was rejected.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct FullError;
impl std::error::Error for FullError {}
impl std::fmt::Display for FullError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "ring buffer full")
    }
}

/// The ring buffer is empty; the pop was rejected.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct EmptyError;
impl std::error::Error for EmptyError {}
impl std::fmt::Display for EmptyError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "ring buffer empty")
    }
}

/// The peer handle was dropped; the channel is disconnected.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ClosedError;
impl std::error::Error for ClosedError {}
impl std::fmt::Display for ClosedError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "ring buffer peer dropped")
    }
}

/// Single ring buffer slot handle. (Wrapped in `UnsafeCell` because we can't trust the rust compiler)
type Slot = UnsafeCell<MaybeUninit<Bytes>>;

/// **Optimized** SPSC ring buffer.
///
/// Each slot stores immutable `Bytes` handle (24 bytes), `push/pop` move the handle —
/// the ring does not copy packet payload data. [Reference](https://rigtorp.se/ringbuffer/)
pub struct Ring {
    slots: Box<[Slot]>,
    capacity: usize,
    head: CachePadded<AtomicUsize>,
    tail: CachePadded<AtomicUsize>,
    /// Set when either handle is dropped; peers check this to detect disconnects.
    closed: CachePadded<AtomicBool>,
    /// Wakes a consumer blocked on an empty ring (`pop_async`).
    data_available: CachePadded<Arc<Notify>>,
    /// Wakes a producer blocked on a full ring (`push_async`).
    space_available: CachePadded<Arc<Notify>>,
}

unsafe impl Send for Ring {}
unsafe impl Sync for Ring {}

impl Ring {
    #[inline]
    /// Create a ring buffer with the given `capacity`,
    /// (must be a `power of two`, or it will choose `next power of two`) with minimum of capacity 2.
    pub fn create(mut capacity: usize) -> Result<Self> {
        if !capacity.is_power_of_two() {
            capacity = capacity.next_power_of_two();
        }
        capacity = capacity.max(2); // clamp at min 2

        let slots: Box<[Slot]> = (0..capacity)
            .map(|_| UnsafeCell::new(MaybeUninit::uninit()))
            .collect();

        Ok(Self {
            slots,
            capacity,
            head: CachePadded::pad(AtomicUsize::new(0)),
            tail: CachePadded::pad(AtomicUsize::new(0)),
            closed: CachePadded::pad(AtomicBool::new(false)),
            data_available: CachePadded::pad(Arc::new(Notify::new())),
            space_available: CachePadded::pad(Arc::new(Notify::new())),
        })
    }

    #[inline(always)]
    /// Cheap owned handle to await `space_available` without holding any ring reference.
    pub fn space_notify(&self) -> Arc<Notify> {
        Arc::clone(&self.space_available)
    }

    #[inline(always)]
    /// Cheap owned handle to await `data_available` without holding any ring reference.
    pub fn data_notify(&self) -> Arc<Notify> {
        Arc::clone(&self.data_available)
    }

    #[inline]
    /// Split into a [`Producer`] and [`Consumer`] handles.
    pub fn into_split(self) -> (Producer, Consumer) {
        let ring = Arc::new(self);
        (
            Producer {
                ring: Arc::clone(&ring),
                head_local: 0,
                tail_cached: 0,
            },
            Consumer {
                ring,
                tail_local: 0,
                head_cached: 0,
            },
        )
    }

    #[inline]
    /// Approximate number of occupied slots `(maybe stale)`.
    pub fn len(&self) -> usize {
        let h = self.head.load(Ordering::Relaxed);
        let t = self.tail.load(Ordering::Relaxed);
        (h.wrapping_sub(t)) & (self.capacity - 1)
    }

    #[inline]
    /// Returns `true` if the buffer appears full `(maybe stale)`.
    pub fn is_full(&self) -> bool {
        let h = self.head.load(Ordering::Relaxed);
        let t = self.tail.load(Ordering::Relaxed);
        ((h + 1) & (self.capacity - 1)) == t
    }

    #[inline]
    /// Returns `true` if the buffer appears empty `(maybe stale)`.
    pub fn is_empty(&self) -> bool {
        self.head.load(Ordering::Relaxed) == self.tail.load(Ordering::Relaxed)
    }

    #[inline]
    /// Returns `true` if either handle has been dropped `(peer disconnected)`.
    pub fn is_closed(&self) -> bool {
        self.closed.load(Ordering::Acquire)
    }

    #[inline]
    /// Approximate number of free slots `(maybe stale)`.
    pub fn free_slots(&self) -> usize {
        self.capacity - 1 - self.len()
    }

    #[inline]
    /// Total capacity of the ring buffer `(including the wasted sentinel slot)`.
    pub fn capacity(&self) -> usize {
        self.capacity
    }

    #[inline(always)]
    /// Raw mutable pointer to the slot array.
    ///
    /// SAFETY: `idx < capacity` always (masked by caller).
    /// `UnsafeCell` ensures that rust compiler won't do some UB stuffs,
    /// `(it gives a proper mutable pointer to the slot)`.
    fn slot_ptr(&self, idx: usize) -> *mut MaybeUninit<Bytes> {
        unsafe { self.slots.get_unchecked(idx).get() }
    }
}

impl Drop for Ring {
    fn drop(&mut self) {
        let h = *self.head.get_mut();
        let t = *self.tail.get_mut();
        let mut idx = t;
        while idx != h {
            unsafe {
                (*self.slots[idx].get()).assume_init_drop();
            }
            idx = (idx + 1) & (self.capacity - 1);
        }
    }
}

/// **Write half** of a ring buffer, only one `Producer` may exist per ring.
pub struct Producer {
    /// Shared ring buffer.
    ring: Arc<Ring>,
    /// Producer's exclusive local copy of the head index.
    /// (only the producer ever writes `head`, so this never goes stale)
    head_local: usize,
    /// Producer's cached copy of the consumer's tail index.
    tail_cached: usize,
}

impl Producer {
    /// Push `Bytes` by ownership. Returns `Err(FullError)` if the ring is full.
    /// The `Bytes` handle is moved into the ring slot (no payload data copy).
    #[inline(always)]
    pub fn push(&mut self, data: Bytes) -> Result<(), FullError> {
        let head = self.head_local;
        let next = (head + 1) & (self.ring.capacity - 1);

        // check if full by comparing the next head index with the cached tail index
        if next == self.tail_cached {
            self.tail_cached = self.ring.tail.load(Ordering::Acquire);
            if next == self.tail_cached {
                return Err(FullError);
            }
        }

        // write slots in place
        unsafe {
            self.ring.slot_ptr(head).write(MaybeUninit::new(data));
        }

        // update the head index with 'release ordering',
        // so that the data is visible to the consumer
        self.head_local = next;
        self.ring.head.store(next, Ordering::Release);
        // wake a consumer blocked in `pop_async` (permit stored if none waiting)
        self.ring.data_available.notify_one();
        Ok(())
    }

    #[inline(always)]
    /// Similar to [`push`](Self::push), but takes a slice of bytes, which `allocates and copies the data`.
    pub fn push_bytes(&mut self, data: &[u8]) -> Result<(), FullError> {
        self.push(Bytes::copy_from_slice(data)) // <- this allocates and copies the data
    }

    #[inline(always)]
    /// Push as many packets as possible, returns the `number of packets` pushed.
    pub fn push_batch(&mut self, packets: &[Bytes]) -> usize {
        if packets.is_empty() {
            return 0;
        }

        let head = self.head_local;
        let cap_mask = self.ring.capacity - 1;

        let mut used = (head.wrapping_sub(self.tail_cached)) & cap_mask;
        let mut free = (self.ring.capacity - 1) - used;

        // if the cached tail index is stale, refresh it and recalculate free slots
        if free < packets.len() {
            self.tail_cached = self.ring.tail.load(Ordering::Acquire);
            used = (head.wrapping_sub(self.tail_cached)) & cap_mask;
            free = (self.ring.capacity - 1) - used;
        }

        // push packets, up to the number of free slots
        let n = packets.len().min(free);
        if n == 0 {
            return 0;
        }

        // write all packets in place, wrapping around if necessary
        #[allow(clippy::needless_range_loop)]
        for i in 0..n {
            let idx = (head + i) & cap_mask;
            unsafe {
                self.ring
                    .slot_ptr(idx)
                    .write(MaybeUninit::new(packets.get_unchecked(i).clone())); // TODO: this clone can be copy or atomic ref count
            }
        }

        // update the head index
        self.head_local = (head + n) & cap_mask;
        self.ring.head.store(self.head_local, Ordering::Release);
        n
    }

    /// Async push: resolves once `data` has a slot. Waits on a [`Notify`] when
    /// full instead of spinning, returns `Err(ClosedError)` if the consumer was
    /// dropped. The lock-free fast path is identical to [`push`](Self::push).
    pub async fn push_async(&mut self, data: Bytes) -> Result<(), ClosedError> {
        loop {
            if self.try_push_shared(&data) {
                return Ok(());
            }
            if self.ring.closed.load(Ordering::Acquire) {
                return Err(ClosedError);
            }
            self.ring.space_available.notified().await;
        }
    }

    #[inline(always)]
    /// Like [`push`](Self::push), but clones `src` into the slot **only after** capacity is confirmed.
    fn try_push_shared(&mut self, src: &Bytes) -> bool {
        let head = self.head_local;
        let next = (head + 1) & (self.ring.capacity - 1);

        if next == self.tail_cached {
            self.tail_cached = self.ring.tail.load(Ordering::Acquire);
            if next == self.tail_cached {
                return false;
            }
        }

        unsafe {
            self.ring
                .slot_ptr(head)
                .write(MaybeUninit::new(src.clone()));
        }
        // both push paths MUST advance `head_local` together with the shared
        // atomic, or a later sync `push` would recompute a stale index and
        // overwrite a live slot
        self.head_local = next;
        self.ring.head.store(next, Ordering::Release);
        // wake a consumer blocked in `pop_async`
        self.ring.data_available.notify_one();
        true
    }

    #[inline(always)]
    /// Approximate number of occupied slots `(maybe stale)`.
    pub fn len(&self) -> usize {
        (self.head_local.wrapping_sub(self.tail_cached)) & (self.ring.capacity - 1)
    }

    #[inline(always)]
    /// Returns `true` if the ring appears full `(maybe stale)`.
    pub fn is_full(&self) -> bool {
        self.len() == self.ring.capacity - 1
    }

    #[inline(always)]
    /// Returns `true` if the ring appears empty `(maybe stale)`.
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    #[inline(always)]
    /// Returns `true` if the consumer handle has been dropped `(zombie ring)`.
    pub fn is_closed(&self) -> bool {
        self.ring.is_closed()
    }

    #[inline(always)]
    /// Cheap owned handle to await space-available **without** holding this `Producer`.
    pub fn space_notify(&self) -> Arc<Notify> {
        self.ring.space_notify()
    }

    /// Approximate number of free slots `(maybe stale)`.
    #[inline]
    pub fn free_slots(&self) -> usize {
        (self.ring.capacity - 1) - self.len()
    }

    /// Usable capacity `(total slots minus sentinel)`.
    #[inline]
    pub fn capacity(&self) -> usize {
        self.ring.capacity - 1
    }
}

impl Drop for Producer {
    fn drop(&mut self) {
        self.ring.closed.store(true, Ordering::Release);
        // wake any consumer blocked on pop so it can observe the disconnect;
        // late arrivals check the flag themselves before awaiting.
        self.ring.data_available.notify_waiters();
    }
}

/// **Read half** of a ring buffer, only one `Consumer` may exist per ring.
pub struct Consumer {
    /// Shared ring buffer.
    ring: Arc<Ring>,
    /// Consumer's exclusive local copy of the tail index.
    /// (only the consumer ever writes `tail`, so this never goes stale)
    tail_local: usize,
    /// Consumer's cached copy of the producer's head index.
    head_cached: usize,
}

impl Consumer {
    /// Pop a single `Bytes`. Returns `Err(EmptyError)` if the ring is empty.
    #[inline(always)]
    pub fn pop(&mut self) -> Result<Bytes, EmptyError> {
        let tail = self.tail_local;

        // check against the cached head index to see if the ring is empty
        if tail == self.head_cached {
            self.head_cached = self.ring.head.load(Ordering::Acquire);
            if tail == self.head_cached {
                return Err(EmptyError);
            }
        }

        // pull data
        let data = unsafe { self.ring.slot_ptr(tail).read().assume_init() };

        // update the tail index with 'release ordering',
        // so that the producer sees the slot as free
        let next = (tail + 1) & (self.ring.capacity - 1);
        self.tail_local = next;
        self.ring.tail.store(next, Ordering::Release);
        // wake a producer blocked in `push_async` (permit stored if none waiting)
        self.ring.space_available.notify_one();
        Ok(data)
    }

    /// Pop up to `buf.len()` packets into pre-allocated `Bytes` slots.
    /// Returns the `number of packets` popped.
    #[inline(always)]
    pub fn pop_batch(&mut self, buf: &mut [Bytes]) -> usize {
        if buf.is_empty() {
            return 0;
        }

        let tail = self.tail_local;
        let cap_mask = self.ring.capacity - 1;

        let mut available = (self.head_cached.wrapping_sub(tail)) & cap_mask;

        // check cached head index, if it's stale, refresh it and recalculate available slots
        if available < buf.len() {
            self.head_cached = self.ring.head.load(Ordering::Acquire);
            available = (self.head_cached.wrapping_sub(tail)) & cap_mask;
        }

        let n = buf.len().min(available);
        if n == 0 {
            return 0;
        }

        // pull data
        for i in 0..n {
            let idx = (tail + i) & cap_mask;
            unsafe {
                // safety; we have already checked that there are enough available slots
                *buf.get_unchecked_mut(i) = self.ring.slot_ptr(idx).read().assume_init();
            }
        }

        // update & return
        self.tail_local = (tail + n) & cap_mask;
        self.ring.tail.store(self.tail_local, Ordering::Release);
        n
    }

    /// Async pulls the next packet. Waits on a [`Notify`] when
    /// empty instead of spinning, returns `Err(ClosedError)` if the producer
    /// was dropped. The lock-free fast path is identical to [`pop`](Self::pop).
    pub async fn pop_async(&mut self) -> Result<Bytes, ClosedError> {
        loop {
            if let Ok(data) = self.pop() {
                return Ok(data);
            }
            if self.ring.closed.load(Ordering::Acquire) {
                return Err(ClosedError);
            }
            self.ring.data_available.notified().await;
        }
    }

    #[inline(always)]
    /// Approximate number of occupied slots (maybe stale).
    pub fn len(&self) -> usize {
        (self.head_cached.wrapping_sub(self.tail_local)) & (self.ring.capacity - 1)
    }

    #[inline]
    /// Returns `true` if the ring appears empty `(maybe stale)`.
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    #[inline(always)]
    /// Returns `true` if the producer handle has been dropped.
    pub fn is_closed(&self) -> bool {
        self.ring.is_closed()
    }

    #[inline]
    /// Approximate number of available items `(maybe stale)`.
    pub fn available(&self) -> usize {
        self.len()
    }

    #[inline]
    /// Usable capacity `(total slots minus sentinel)`.
    pub fn capacity(&self) -> usize {
        self.ring.capacity - 1
    }
}

impl Drop for Consumer {
    fn drop(&mut self) {
        self.ring.closed.store(true, Ordering::Release);
        // wake any producer blocked on push so it can observe the disconnect;
        // late arrivals check the flag themselves before awaiting.
        self.ring.space_available.notify_waiters();
    }
}