use std::collections::VecDeque;
use std::ptr;
use std::sync::atomic::{AtomicPtr, AtomicU32, AtomicU64, AtomicUsize, Ordering};
use std::task::{Context, Poll, Waker};
#[repr(align(64))]
pub struct TreiberStack {
head: AtomicU64,
next: Box<[AtomicU32]>,
}
impl TreiberStack {
#[must_use]
pub fn new(size: usize) -> Self {
let mut next = Vec::with_capacity(size);
for _ in 0..size {
next.push(AtomicU32::new(0));
}
if size > 0 {
next[size - 1].store(u32::MAX, Ordering::Relaxed);
}
Self {
head: AtomicU64::new(u64::from(u32::MAX)),
next: next.into_boxed_slice(),
}
}
#[must_use]
pub(crate) fn new_empty(size: usize) -> Self {
let mut next = Vec::with_capacity(size);
for _ in 0..size {
next.push(AtomicU32::new(0));
}
Self {
head: AtomicU64::new(u64::from(u32::MAX)),
next: next.into_boxed_slice(),
}
}
#[inline]
pub fn push(&self, idx: u32) {
let mut head = self.head.load(Ordering::Acquire);
loop {
let head_idx = (head & 0xFFFF_FFFF) as u32;
let tag = (head >> 32) as u32;
self.next[idx as usize].store(head_idx, Ordering::Release);
let new_head = (u64::from(tag.wrapping_add(1)) << 32) | u64::from(idx);
match self.head.compare_exchange_weak(
head,
new_head,
Ordering::Release,
Ordering::Acquire,
) {
Ok(_) => break,
Err(actual) => head = actual,
}
}
}
#[inline]
pub fn pop(&self) -> Option<u32> {
let mut head = self.head.load(Ordering::Acquire);
loop {
let head_idx = (head & 0xFFFF_FFFF) as u32;
if head_idx == u32::MAX {
return None;
}
let tag = (head >> 32) as u32;
let next = self.next[head_idx as usize].load(Ordering::Acquire);
let new_head = (u64::from(tag.wrapping_add(1)) << 32) | u64::from(next);
match self.head.compare_exchange_weak(
head,
new_head,
Ordering::Release,
Ordering::Acquire,
) {
Ok(_) => return Some(head_idx),
Err(actual) => head = actual,
}
}
}
#[inline]
pub(crate) fn write_next_slot(&self, idx: usize, next_idx: u32) {
self.next[idx].store(next_idx, Ordering::Relaxed);
}
#[inline]
pub(crate) fn set_head(&self, head_idx: u32) {
let initial_head = u64::from(head_idx);
self.head.store(initial_head, Ordering::Relaxed);
}
}
pub struct BufferPool {
arena_ptr: *mut u8,
layout: std::alloc::Layout,
chunk_size: usize,
free: TreiberStack,
}
unsafe impl Send for BufferPool {}
unsafe impl Sync for BufferPool {}
impl BufferPool {
#[must_use]
pub fn new(total_chunks: usize, chunk_size: usize) -> Self {
let total_chunks_bounded = total_chunks.max(1);
let chunk_size_bounded = chunk_size.max(1);
let layout =
std::alloc::Layout::from_size_align(total_chunks_bounded * chunk_size_bounded, 4096)
.expect("Invalid layout alignment for BufferPool");
let arena_ptr = unsafe { std::alloc::alloc(layout) };
if arena_ptr.is_null() {
std::alloc::handle_alloc_error(layout);
}
let free = TreiberStack::new_empty(total_chunks_bounded);
for i in 0..total_chunks_bounded {
let next_idx = if i == total_chunks_bounded - 1 {
u32::MAX
} else {
(i + 1) as u32
};
free.write_next_slot(i, next_idx);
}
free.set_head(0);
Self {
arena_ptr,
layout,
chunk_size: chunk_size_bounded,
free,
}
}
#[inline]
pub const fn get_ptr(&self, idx: u32) -> *mut u8 {
unsafe { self.arena_ptr.add(idx as usize * self.chunk_size) }
}
#[inline]
pub const fn chunk_size(&self) -> usize {
self.chunk_size
}
#[inline]
pub fn acquire(&self) -> Option<u32> {
self.free.pop()
}
#[inline]
pub fn release(&self, idx: u32) {
self.free.push(idx);
}
}
impl Drop for BufferPool {
fn drop(&mut self) {
unsafe {
std::alloc::dealloc(self.arena_ptr, self.layout);
}
}
}
const WAITING: usize = 0;
const REGISTERING: usize = 0b01;
const WAKING: usize = 0b10;
#[repr(align(64))]
pub struct AtomicWakerSlot {
state: AtomicUsize,
waker: std::cell::UnsafeCell<Option<Waker>>,
}
impl Default for AtomicWakerSlot {
fn default() -> Self {
Self::new()
}
}
impl AtomicWakerSlot {
#[must_use]
pub const fn new() -> Self {
Self {
state: AtomicUsize::new(WAITING),
waker: std::cell::UnsafeCell::new(None),
}
}
#[inline]
pub fn register(&self, waker: &Waker) {
match self.state.compare_exchange(
WAITING,
REGISTERING,
Ordering::Acquire,
Ordering::Acquire,
) {
Ok(_) => {
let slot = unsafe { &mut *self.waker.get() };
let already_registered = slot.as_ref().is_some_and(|w| w.will_wake(waker));
if !already_registered {
*slot = Some(waker.clone());
}
if self
.state
.compare_exchange(REGISTERING, WAITING, Ordering::AcqRel, Ordering::Acquire)
.is_err()
{
let pending = unsafe { &mut *self.waker.get() }.take();
self.state.store(WAITING, Ordering::Release);
if let Some(w) = pending {
w.wake();
}
}
}
Err(_) => {
waker.wake_by_ref();
}
}
}
#[inline]
pub fn take_and_wake(&self) {
match self
.state
.compare_exchange(WAITING, WAKING, Ordering::Acquire, Ordering::Acquire)
{
Ok(_) => {
let waker = unsafe { &mut *self.waker.get() }.take();
self.state.store(WAITING, Ordering::Release);
if let Some(w) = waker {
w.wake();
}
}
Err(_) => {
self.state.fetch_or(WAKING, Ordering::AcqRel);
}
}
}
}
impl Drop for AtomicWakerSlot {
fn drop(&mut self) {
drop(self.waker.get_mut().take());
}
}
unsafe impl Send for AtomicWakerSlot {}
unsafe impl Sync for AtomicWakerSlot {}
#[repr(align(64))]
pub struct MpmcStack<T> {
head: AtomicPtr<Node<T>>,
len: AtomicUsize,
}
#[repr(align(64))]
struct Node<T> {
value: T,
next: *mut Self,
}
impl<T> Default for MpmcStack<T> {
fn default() -> Self {
Self::new()
}
}
impl<T> MpmcStack<T> {
#[must_use]
pub const fn new() -> Self {
Self {
head: AtomicPtr::new(ptr::null_mut()),
len: AtomicUsize::new(0),
}
}
#[inline]
pub fn push(&self, value: T) {
let node = Box::into_raw(Box::new(Node {
value,
next: ptr::null_mut(),
}));
let mut head = self.head.load(Ordering::Relaxed);
loop {
unsafe { (*node).next = head };
match self
.head
.compare_exchange_weak(head, node, Ordering::Release, Ordering::Relaxed)
{
Ok(_) => break,
Err(actual) => head = actual,
}
}
self.len.fetch_add(1, Ordering::Relaxed);
}
#[inline]
pub fn pop(&self) -> Option<T> {
let mut head = self.head.load(Ordering::Acquire);
loop {
if head.is_null() {
return None;
}
let next = unsafe { (*head).next };
match self
.head
.compare_exchange_weak(head, next, Ordering::Acquire, Ordering::Acquire)
{
Ok(_) => {
self.len.fetch_sub(1, Ordering::Relaxed);
let boxed = unsafe { Box::from_raw(head) };
return Some(boxed.value);
}
Err(actual) => head = actual,
}
}
}
#[inline]
pub fn drain_all(&self) -> Vec<T> {
let mut head = self.head.swap(ptr::null_mut(), Ordering::AcqRel);
let mut out = Vec::new();
while !head.is_null() {
let boxed = unsafe { Box::from_raw(head) };
head = boxed.next;
out.push(boxed.value);
}
self.len.store(0, Ordering::Relaxed);
out.reverse();
out
}
#[inline(always)]
pub fn is_empty(&self) -> bool {
self.head.load(Ordering::Relaxed).is_null()
}
#[inline(always)]
pub fn len(&self) -> usize {
self.len.load(Ordering::Relaxed)
}
#[allow(non_snake_case)]
#[inline]
pub fn drain_into_vec_deque(&self, target: &mut VecDeque<T>) {
let mut current = self.head.swap(std::ptr::null_mut(), Ordering::Acquire);
let mut prev = std::ptr::null_mut();
while !current.is_null() {
unsafe {
let next = (*current).next;
(*current).next = prev;
prev = current;
current = next;
}
}
let mut node = prev;
while !node.is_null() {
unsafe {
let BoxedNode = Box::from_raw(node);
target.push_back(BoxedNode.value);
node = BoxedNode.next;
}
}
}
}
unsafe impl<T: Send> Send for MpmcStack<T> {}
unsafe impl<T: Send> Sync for MpmcStack<T> {}
impl<T> Drop for MpmcStack<T> {
fn drop(&mut self) {
while self.pop().is_some() {}
}
}
pub struct SpscQueue<T> {
head: CacheAlignedUsize,
tail: CacheAlignedUsize,
buffer: Box<[std::mem::MaybeUninit<T>]>,
capacity: usize,
}
#[repr(align(64))]
struct CacheAlignedUsize {
value: AtomicUsize,
}
unsafe impl<T: Send> Send for SpscQueue<T> {}
unsafe impl<T: Send> Sync for SpscQueue<T> {}
impl<T> SpscQueue<T> {
#[must_use]
pub fn new(capacity: usize) -> Self {
assert!(capacity.is_power_of_two());
let mut buffer = Vec::with_capacity(capacity);
for _ in 0..capacity {
buffer.push(std::mem::MaybeUninit::uninit());
}
Self {
head: CacheAlignedUsize {
value: AtomicUsize::new(0),
},
tail: CacheAlignedUsize {
value: AtomicUsize::new(0),
},
buffer: buffer.into_boxed_slice(),
capacity,
}
}
#[inline]
pub fn push(&self, value: T) -> Result<(), T> {
let tail = self.tail.value.load(Ordering::Relaxed);
let head = self.head.value.load(Ordering::Acquire);
if tail.wrapping_sub(head) == self.capacity {
return Err(value);
}
let mask = self.capacity - 1;
let idx = tail & mask;
unsafe {
let ptr = self.buffer[idx].as_ptr().cast_mut();
ptr.write(value);
}
self.tail
.value
.store(tail.wrapping_add(1), Ordering::Release);
Ok(())
}
#[inline]
pub fn pop(&self) -> Option<T> {
let head = self.head.value.load(Ordering::Relaxed);
let tail = self.tail.value.load(Ordering::Acquire);
if head == tail {
return None;
}
let mask = self.capacity - 1;
let idx = head & mask;
let value = unsafe {
let ptr = self.buffer[idx].as_ptr();
ptr.read()
};
self.head
.value
.store(head.wrapping_add(1), Ordering::Release);
Some(value)
}
#[inline(always)]
pub fn is_empty(&self) -> bool {
let head = self.head.value.load(Ordering::Relaxed);
let tail = self.tail.value.load(Ordering::Acquire);
head == tail
}
#[inline(always)]
pub fn is_full(&self) -> bool {
let tail = self.tail.value.load(Ordering::Relaxed);
let head = self.head.value.load(Ordering::Acquire);
tail.wrapping_sub(head) == self.capacity
}
#[inline(always)]
pub fn len(&self) -> usize {
let head = self.head.value.load(Ordering::Relaxed);
let tail = self.tail.value.load(Ordering::Acquire);
tail.wrapping_sub(head)
}
#[inline(always)]
pub const fn capacity(&self) -> usize {
self.capacity
}
}
impl<T: Copy> SpscQueue<T> {
#[inline]
pub fn push_slice(&self, src: &[T]) -> usize {
let tail = self.tail.value.load(Ordering::Relaxed);
let head = self.head.value.load(Ordering::Acquire);
let free = self.capacity - tail.wrapping_sub(head);
let to_push = free.min(src.len());
if to_push == 0 {
return 0;
}
let mask = self.capacity - 1;
let start = tail & mask;
let first = to_push.min(self.capacity - start);
unsafe {
let base = self.buffer.as_ptr().cast_mut().cast::<T>();
ptr::copy_nonoverlapping(src.as_ptr(), base.add(start), first);
if to_push > first {
ptr::copy_nonoverlapping(src.as_ptr().add(first), base, to_push - first);
}
}
self.tail
.value
.store(tail.wrapping_add(to_push), Ordering::Release);
to_push
}
#[inline]
pub fn pop_slice(&self, dst: &mut [T]) -> usize {
let head = self.head.value.load(Ordering::Relaxed);
let tail = self.tail.value.load(Ordering::Acquire);
let avail = tail.wrapping_sub(head);
let to_pop = avail.min(dst.len());
if to_pop == 0 {
return 0;
}
let mask = self.capacity - 1;
let start = head & mask;
let first = to_pop.min(self.capacity - start);
unsafe {
let base = self.buffer.as_ptr().cast::<T>();
ptr::copy_nonoverlapping(base.add(start), dst.as_mut_ptr(), first);
if to_pop > first {
ptr::copy_nonoverlapping(base, dst.as_mut_ptr().add(first), to_pop - first);
}
}
self.head
.value
.store(head.wrapping_add(to_pop), Ordering::Release);
to_pop
}
}
impl<T> Drop for SpscQueue<T> {
fn drop(&mut self) {
while self.pop().is_some() {}
}
}
#[repr(align(64))]
pub struct OnceSlot<T> {
ptr: AtomicPtr<T>,
waker: AtomicWakerSlot,
}
impl<T> Default for OnceSlot<T> {
fn default() -> Self {
Self::new()
}
}
impl<T> OnceSlot<T> {
#[must_use]
pub const fn new() -> Self {
Self {
ptr: AtomicPtr::new(ptr::null_mut()),
waker: AtomicWakerSlot::new(),
}
}
#[inline]
pub fn set(&self, value: T) {
let boxed = Box::into_raw(Box::new(value));
let prev = self.ptr.swap(boxed, Ordering::AcqRel);
debug_assert!(
prev.is_null(),
"OnceSlot::set called more than once — second value leaked"
);
self.waker.take_and_wake();
}
#[inline]
pub fn poll(&self, cx: &Context<'_>) -> Poll<T> {
let mut p = self.ptr.load(Ordering::Acquire);
if !p.is_null() {
p = self.ptr.swap(ptr::null_mut(), Ordering::AcqRel);
if !p.is_null() {
return Poll::Ready(*unsafe { Box::from_raw(p) });
}
}
self.waker.register(cx.waker());
p = self.ptr.swap(ptr::null_mut(), Ordering::AcqRel);
if !p.is_null() {
return Poll::Ready(*unsafe { Box::from_raw(p) });
}
Poll::Pending
}
}
impl<T> Drop for OnceSlot<T> {
fn drop(&mut self) {
let p = *self.ptr.get_mut();
if !p.is_null() {
drop(unsafe { Box::from_raw(p) });
}
}
}
unsafe impl<T: Send> Send for OnceSlot<T> {}
unsafe impl<T: Send> Sync for OnceSlot<T> {}
#[repr(C)]
pub struct BoundedMpmcQueue<T> {
buffer: Box<[Cell<T>]>,
capacity: usize,
_pad1: [u64; 8],
enqueue_pos: AtomicUsize,
_pad2: [u64; 8],
dequeue_pos: AtomicUsize,
_pad3: [u64; 8],
}
#[repr(align(64))]
struct Cell<T> {
sequence: AtomicUsize,
data: std::cell::UnsafeCell<std::mem::MaybeUninit<T>>,
}
impl<T> BoundedMpmcQueue<T> {
#[must_use]
pub fn new(capacity: usize) -> Self {
let capacity = capacity.max(1);
let mut buffer = Vec::with_capacity(capacity);
for i in 0..capacity {
buffer.push(Cell {
sequence: AtomicUsize::new(i),
data: std::cell::UnsafeCell::new(std::mem::MaybeUninit::uninit()),
});
}
Self {
buffer: buffer.into_boxed_slice(),
capacity,
_pad1: [0; 8],
enqueue_pos: AtomicUsize::new(0),
_pad2: [0; 8],
dequeue_pos: AtomicUsize::new(0),
_pad3: [0; 8],
}
}
#[must_use]
pub const fn capacity(&self) -> usize {
self.capacity
}
#[inline]
pub fn try_push(&self, value: T) -> Result<(), T> {
let mut pos = self.enqueue_pos.load(Ordering::Relaxed);
loop {
let cell = &self.buffer[pos % self.capacity];
let seq = cell.sequence.load(Ordering::Acquire);
#[allow(clippy::cast_possible_wrap)]
let diff = seq as isize - pos as isize;
match diff.cmp(&0) {
std::cmp::Ordering::Equal => {
match self.enqueue_pos.compare_exchange_weak(
pos,
pos.wrapping_add(1),
Ordering::Relaxed,
Ordering::Relaxed,
) {
Ok(_) => {
unsafe { (*cell.data.get()).write(value) };
cell.sequence.store(pos.wrapping_add(1), Ordering::Release);
return Ok(());
}
Err(actual) => pos = actual,
}
}
std::cmp::Ordering::Less => return Err(value),
std::cmp::Ordering::Greater => pos = self.enqueue_pos.load(Ordering::Relaxed),
}
}
}
#[inline]
pub fn try_pop(&self) -> Option<T> {
let mut pos = self.dequeue_pos.load(Ordering::Relaxed);
loop {
let cell = &self.buffer[pos % self.capacity];
let seq = cell.sequence.load(Ordering::Acquire);
#[allow(clippy::cast_possible_wrap)]
let diff = seq as isize - pos.wrapping_add(1) as isize;
match diff.cmp(&0) {
std::cmp::Ordering::Equal => {
match self.dequeue_pos.compare_exchange_weak(
pos,
pos.wrapping_add(1),
Ordering::Relaxed,
Ordering::Relaxed,
) {
Ok(_) => {
let value = unsafe { (*cell.data.get()).assume_init_read() };
cell.sequence
.store(pos.wrapping_add(self.capacity), Ordering::Release);
return Some(value);
}
Err(actual) => pos = actual,
}
}
std::cmp::Ordering::Less => return None,
std::cmp::Ordering::Greater => pos = self.dequeue_pos.load(Ordering::Relaxed),
}
}
}
#[must_use]
#[inline(always)]
pub fn enqueue_snapshot(&self) -> usize {
self.enqueue_pos.load(Ordering::Acquire)
}
#[inline(always)]
pub fn drain_until(&self, target: usize, mut on_pop: impl FnMut(T)) {
#[allow(clippy::cast_possible_wrap)]
while (target.wrapping_sub(self.dequeue_pos.load(Ordering::Relaxed)) as isize) > 0 {
match self.try_pop() {
Some(value) => on_pop(value),
None => std::hint::spin_loop(),
}
}
}
#[must_use]
#[inline(always)]
pub fn len(&self) -> usize {
let enq = self.enqueue_pos.load(Ordering::Relaxed);
let deq = self.dequeue_pos.load(Ordering::Relaxed);
enq.saturating_sub(deq)
}
#[must_use]
#[inline(always)]
pub fn is_empty(&self) -> bool {
self.len() == 0
}
}
impl<T> Drop for BoundedMpmcQueue<T> {
fn drop(&mut self) {
while self.try_pop().is_some() {}
}
}
unsafe impl<T: Send> Send for BoundedMpmcQueue<T> {}
unsafe impl<T: Send> Sync for BoundedMpmcQueue<T> {}