use std::cell::UnsafeCell;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use crate::builder::Backpressure;
use crate::record::{END_OF_BUFFER, MAX_RECORD_SIZE, VERSION};
pub(crate) const SLOT_SIZE: usize = {
#[cfg(all(target_arch = "aarch64", target_vendor = "apple"))]
{
128
}
#[cfg(not(all(target_arch = "aarch64", target_vendor = "apple")))]
{
64
}
};
pub(crate) const RING_SIZE: usize = 1_048_576;
#[inline(always)]
pub(crate) const fn align_up(n: u64, align: u64) -> u64 {
n.wrapping_add(align - 1) & !(align - 1)
}
#[repr(C, align(64))]
pub(crate) struct RingBuffer {
pub(crate) head: AtomicU64,
pub(crate) tail_cache: UnsafeCell<u64>,
_pad_p: [u8; 48],
pub(crate) tail: AtomicU64,
pub(crate) head_cache: UnsafeCell<u64>,
_pad_d: [u8; 48],
pub(crate) live: AtomicBool,
pub(crate) data: Box<[UnsafeCell<u8>]>,
}
unsafe impl Sync for RingBuffer {}
impl RingBuffer {
pub(crate) fn new() -> Self {
let bytes = vec![0u8; RING_SIZE].into_boxed_slice();
let data: Box<[UnsafeCell<u8>]> =
unsafe { Box::from_raw(Box::into_raw(bytes) as *mut [UnsafeCell<u8>]) };
Self {
head: AtomicU64::new(0),
tail_cache: UnsafeCell::new(0),
_pad_p: [0u8; 48],
tail: AtomicU64::new(0),
head_cache: UnsafeCell::new(0),
_pad_d: [0u8; 48],
live: AtomicBool::new(true),
data,
}
}
#[inline]
pub(crate) fn write_record(&self, record: &[u8], policy: Backpressure) -> bool {
debug_assert!(
record.len() <= MAX_RECORD_SIZE,
"invariant: record length must fit the u16 total_size field"
);
let total_size = record.len() as u64;
let aligned = align_up(total_size, SLOT_SIZE as u64);
let head = self.head.load(Ordering::Relaxed);
let offset = (head & (RING_SIZE as u64 - 1)) as usize;
let remaining_phys = (RING_SIZE - offset) as u64;
let wrap = aligned > remaining_phys;
let needed = if wrap {
aligned + remaining_phys
} else {
aligned
};
if !self.ensure_capacity(head, needed, policy) {
return false;
}
let base = self.data.as_ptr() as *mut u8;
let mut next = head;
if wrap {
unsafe { write_eob(base.add(offset), remaining_phys as u16) };
next = next.wrapping_add(remaining_phys);
}
let dst = (next & (RING_SIZE as u64 - 1)) as usize;
unsafe {
std::ptr::copy_nonoverlapping(record.as_ptr(), base.add(dst), record.len());
}
next = next.wrapping_add(aligned);
self.head.store(next, Ordering::Release);
true
}
fn ensure_capacity(&self, head: u64, needed: u64, policy: Backpressure) -> bool {
let cached = unsafe { *self.tail_cache.get() };
if head.wrapping_add(needed).wrapping_sub(cached) <= RING_SIZE as u64 {
return true;
}
loop {
let tail = self.tail.load(Ordering::Acquire);
unsafe { *self.tail_cache.get() = tail };
if head.wrapping_add(needed).wrapping_sub(tail) <= RING_SIZE as u64 {
return true;
}
match policy {
Backpressure::Drop => return false,
Backpressure::Block => std::hint::spin_loop(),
}
}
}
}
unsafe fn write_eob(ptr: *mut u8, span: u16) {
let size = span.to_le_bytes();
unsafe {
ptr.write(VERSION);
ptr.add(1).write(END_OF_BUFFER);
ptr.add(2).write(size[0]);
ptr.add(3).write(size[1]);
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::mem;
#[test]
fn constants_are_power_of_two() {
assert!(RING_SIZE.is_power_of_two());
assert!(SLOT_SIZE.is_power_of_two());
}
#[test]
fn ring_size_is_1mb() {
assert_eq!(RING_SIZE, 1024 * 1024);
}
#[test]
fn struct_alignment_is_64() {
assert_eq!(mem::align_of::<RingBuffer>(), 64);
}
#[test]
fn struct_size_is_multiple_of_alignment() {
let size = mem::size_of::<RingBuffer>();
assert_eq!(size % 64, 0);
}
#[test]
fn head_at_offset_zero() {
let rb = RingBuffer::new();
let base = &rb as *const RingBuffer as usize;
let head_addr = &rb.head as *const AtomicU64 as usize;
assert_eq!(head_addr - base, 0);
}
#[test]
fn tail_at_offset_64() {
let rb = RingBuffer::new();
let base = &rb as *const RingBuffer as usize;
let tail_addr = &rb.tail as *const AtomicU64 as usize;
assert_eq!(tail_addr - base, 64);
}
#[test]
fn head_and_tail_on_different_cache_lines() {
let rb = RingBuffer::new();
let head_addr = &rb.head as *const AtomicU64 as usize;
let tail_addr = &rb.tail as *const AtomicU64 as usize;
let diff = head_addr.abs_diff(tail_addr);
assert!(diff >= 64);
}
#[test]
fn new_initializes_head_to_zero() {
let rb = RingBuffer::new();
assert_eq!(rb.head.load(std::sync::atomic::Ordering::Relaxed), 0);
}
#[test]
fn new_initializes_tail_to_zero() {
let rb = RingBuffer::new();
assert_eq!(rb.tail.load(std::sync::atomic::Ordering::Relaxed), 0);
}
#[test]
fn new_initializes_tail_cache_to_zero() {
let rb = RingBuffer::new();
let val = unsafe { *rb.tail_cache.get() };
assert_eq!(val, 0);
}
#[test]
fn new_initializes_head_cache_to_zero() {
let rb = RingBuffer::new();
let val = unsafe { *rb.head_cache.get() };
assert_eq!(val, 0);
}
#[test]
fn new_sets_live_to_true() {
let rb = RingBuffer::new();
assert!(rb.live.load(std::sync::atomic::Ordering::Relaxed));
}
#[test]
fn new_zero_initializes_data_region() {
let rb = RingBuffer::new();
let data =
unsafe { std::slice::from_raw_parts(rb.data.as_ptr() as *const u8, RING_SIZE) };
assert_eq!(data.len(), RING_SIZE);
assert!(data.iter().all(|&b| b == 0));
}
#[test]
fn pad_fields_are_zeroed() {
let rb = RingBuffer::new();
assert!(rb._pad_p.iter().all(|&b| b == 0));
assert!(rb._pad_d.iter().all(|&b| b == 0));
}
#[test]
fn align_up_already_aligned() {
assert_eq!(align_up(64, 64), 64);
assert_eq!(align_up(128, 64), 128);
assert_eq!(align_up(0, 64), 0);
}
#[test]
fn align_up_not_aligned() {
assert_eq!(align_up(1, 64), 64);
assert_eq!(align_up(63, 64), 64);
assert_eq!(align_up(65, 64), 128);
}
#[test]
fn align_up_align_1_is_identity() {
assert_eq!(align_up(0, 1), 0);
assert_eq!(align_up(1, 1), 1);
assert_eq!(align_up(42, 1), 42);
assert_eq!(align_up(u64::MAX, 1), u64::MAX);
}
#[test]
fn align_up_power_of_two_aligns() {
assert_eq!(align_up(0, 2), 0);
assert_eq!(align_up(1, 2), 2);
assert_eq!(align_up(2, 2), 2);
assert_eq!(align_up(3, 2), 4);
assert_eq!(align_up(0, 8), 0);
assert_eq!(align_up(7, 8), 8);
assert_eq!(align_up(8, 8), 8);
assert_eq!(align_up(9, 8), 16);
}
#[test]
fn align_up_large_values() {
let n = u64::MAX - 100;
let aligned = align_up(n, 64);
assert_eq!(aligned % 64, 0);
assert!(aligned >= n);
}
#[test]
fn align_up_slot_size_boundary() {
let slot = SLOT_SIZE as u64;
assert_eq!(align_up(0, slot), 0);
assert_eq!(align_up(1, slot), slot);
assert_eq!(align_up(slot - 1, slot), slot);
assert_eq!(align_up(slot, slot), slot);
assert_eq!(align_up(slot + 1, slot), 2 * slot);
}
#[test]
fn bitmask_is_ring_size_minus_one() {
let mask = (RING_SIZE - 1) as u64;
assert_eq!(mask, 0xFFFFF);
assert_eq!(mask + 1, RING_SIZE as u64);
}
#[test]
fn head_wrap_with_bitmask() {
let head: u64 = RING_SIZE as u64 + 42;
let offset = head & (RING_SIZE as u64 - 1);
assert_eq!(offset, 42);
}
#[test]
fn head_at_exact_capacity_wraps_to_zero() {
let head: u64 = RING_SIZE as u64;
let offset = head & (RING_SIZE as u64 - 1);
assert_eq!(offset, 0);
}
#[test]
fn ring_buffer_is_send() {
fn assert_send<T: Send>() {}
assert_send::<RingBuffer>();
}
#[test]
fn ring_buffer_is_sync() {
fn assert_sync<T: Sync>() {}
assert_sync::<RingBuffer>();
}
#[test]
fn write_record_places_record_and_advances_head() {
let rb = RingBuffer::new();
let record = vec![0xABu8; 40];
assert!(rb.write_record(&record, Backpressure::Drop));
let aligned = align_up(40, SLOT_SIZE as u64);
assert_eq!(rb.head.load(Ordering::Relaxed), aligned);
let data =
unsafe { std::slice::from_raw_parts(rb.data.as_ptr() as *const u8, RING_SIZE) };
assert_eq!(&data[..40], &record[..]);
}
#[test]
fn write_record_wraps_with_eob_at_ring_end() {
let rb = RingBuffer::new();
let slot = SLOT_SIZE as u64;
let start = RING_SIZE as u64 - slot;
rb.head.store(start, Ordering::Relaxed);
rb.tail.store(start, Ordering::Relaxed);
unsafe { *rb.tail_cache.get() = start };
let record = vec![0xCDu8; slot as usize + 1];
let aligned = align_up(record.len() as u64, slot); assert!(rb.write_record(&record, Backpressure::Drop));
assert_eq!(rb.head.load(Ordering::Relaxed), start + slot + aligned);
let data =
unsafe { std::slice::from_raw_parts(rb.data.as_ptr() as *const u8, RING_SIZE) };
let eob = (start & (RING_SIZE as u64 - 1)) as usize;
assert_eq!(data[eob], VERSION);
assert_eq!(data[eob + 1], END_OF_BUFFER);
assert_eq!(
u16::from_le_bytes([data[eob + 2], data[eob + 3]]) as u64,
slot
);
assert_eq!(&data[..record.len()], &record[..]);
}
#[test]
fn write_record_drops_when_full_under_drop_policy() {
let rb = RingBuffer::new();
rb.head.store(RING_SIZE as u64, Ordering::Relaxed);
rb.tail.store(0, Ordering::Relaxed);
unsafe { *rb.tail_cache.get() = 0 };
let record = vec![0u8; 40];
assert!(!rb.write_record(&record, Backpressure::Drop));
assert_eq!(rb.head.load(Ordering::Relaxed), RING_SIZE as u64);
}
#[test]
fn write_record_block_unblocks_when_drain_frees_space() {
let rb = std::sync::Arc::new(RingBuffer::new());
rb.head.store(RING_SIZE as u64, Ordering::Relaxed);
rb.tail.store(0, Ordering::Relaxed);
unsafe { *rb.tail_cache.get() = 0 };
let drain = std::sync::Arc::clone(&rb);
let handle = std::thread::spawn(move || {
std::thread::sleep(std::time::Duration::from_millis(20));
drain.tail.store(RING_SIZE as u64, Ordering::Release);
});
let record = vec![0x5Au8; 40];
assert!(rb.write_record(&record, Backpressure::Block));
handle.join().unwrap();
let aligned = align_up(40, SLOT_SIZE as u64);
assert_eq!(rb.head.load(Ordering::Relaxed), RING_SIZE as u64 + aligned);
let data =
unsafe { std::slice::from_raw_parts(rb.data.as_ptr() as *const u8, RING_SIZE) };
assert_eq!(&data[..40], &record[..]);
}
}