#![allow(clippy::missing_errors_doc)]
use std::fs::{File, OpenOptions};
use std::io;
use std::path::Path;
use std::sync::atomic::{AtomicI64, AtomicU64, Ordering, fence};
use memmap2::{MmapMut, MmapOptions};
use parking_lot::Mutex;
use crate::shared_deque_khpd::LineItem;
#[inline(always)]
fn prefetch_slot(slot: *const LcrqJobSlot) {
#[cfg(target_arch = "x86_64")]
{
unsafe {
core::arch::asm!(
"prefetchw [{ptr}]",
ptr = in(reg) slot,
options(nostack, preserves_flags),
);
}
}
#[cfg(not(target_arch = "x86_64"))]
{
_ = slot;
}
}
pub const LOH_MAGIC: u64 = 0x574C_4F48_0000_0001;
pub const LOH_SLOT_SIZE: usize = 64;
pub const DEFAULT_LIFO_CAP: usize = 256;
#[repr(C, align(64))]
pub struct LohHeader {
pub magic: u64,
pub capacity: u64,
pub owner_pid: AtomicU64,
pub epoch: AtomicU64,
pub _pad_meta: [u8; 24],
pub tail: AtomicI64,
pub _pad_tail: [u8; 56],
pub head: AtomicI64,
pub _pad_head: [u8; 56],
}
#[repr(C, align(64))]
pub struct LcrqJobSlot {
pub sequence: AtomicI64,
pub item: LineItem,
pub _pad: [u8; 40],
}
pub const fn loh_file_size(capacity: usize) -> usize {
std::mem::size_of::<LohHeader>() + capacity * LOH_SLOT_SIZE
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PushError {
Full,
LifoFull,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Steal {
Success(StealResult),
Empty,
Retry,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct StealResult {
pub item: LineItem,
}
pub struct SharedDequeLoh {
_file: File,
mmap: MmapMut,
capacity: usize,
capacity_mask: i64,
flush_threshold: usize,
lifo_cap: usize,
local_lifo: Mutex<Vec<LineItem>>,
}
unsafe impl Send for SharedDequeLoh {}
unsafe impl Sync for SharedDequeLoh {}
impl SharedDequeLoh {
pub fn create<P: AsRef<Path>>(
path: P,
capacity: usize,
flush_threshold: usize,
) -> io::Result<Self> {
let capacity = capacity.max(2).next_power_of_two();
let size = loh_file_size(capacity);
let file = OpenOptions::new()
.read(true)
.write(true)
.create(true)
.truncate(true)
.open(path.as_ref())?;
file.set_len(size as u64)?;
let mut mmap = unsafe { MmapOptions::new().len(size).map_mut(&file)? };
let header_ptr = mmap.as_mut_ptr() as *mut LohHeader;
unsafe {
(*header_ptr).magic = LOH_MAGIC;
(*header_ptr).capacity = capacity as u64;
(*header_ptr).owner_pid = AtomicU64::new(std::process::id() as u64);
(*header_ptr).epoch = AtomicU64::new(0);
std::ptr::write_bytes((*header_ptr)._pad_meta.as_mut_ptr(), 0, 24);
(*header_ptr).tail = AtomicI64::new(0);
std::ptr::write_bytes((*header_ptr)._pad_tail.as_mut_ptr(), 0, 56);
(*header_ptr).head = AtomicI64::new(0);
std::ptr::write_bytes((*header_ptr)._pad_head.as_mut_ptr(), 0, 56);
}
let slots_start = std::mem::size_of::<LohHeader>();
for i in 0..capacity {
let off = slots_start + i * LOH_SLOT_SIZE;
let slot_ptr = unsafe { mmap.as_mut_ptr().add(off) as *mut LcrqJobSlot };
unsafe {
(*slot_ptr).sequence = AtomicI64::new(i as i64);
(*slot_ptr).item = LineItem::default();
std::ptr::write_bytes((*slot_ptr)._pad.as_mut_ptr(), 0, 40);
}
}
mmap.flush()?;
let flush_threshold = flush_threshold.max(1);
Ok(Self {
_file: file,
mmap,
capacity,
capacity_mask: (capacity as i64) - 1,
flush_threshold,
lifo_cap: DEFAULT_LIFO_CAP,
local_lifo: Mutex::new(Vec::with_capacity(DEFAULT_LIFO_CAP)),
})
}
pub fn open<P: AsRef<Path>>(path: P, flush_threshold: usize) -> io::Result<Self> {
let file = OpenOptions::new()
.read(true)
.write(true)
.open(path.as_ref())?;
let size = file.metadata()?.len() as usize;
if size < std::mem::size_of::<LohHeader>() {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"loh file too small to contain header",
));
}
let mmap = unsafe { MmapOptions::new().len(size).map_mut(&file)? };
let header_ptr = mmap.as_ptr() as *const LohHeader;
let (magic, capacity) =
unsafe { ((*header_ptr).magic, (*header_ptr).capacity as usize) };
if magic != LOH_MAGIC {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
format!("loh magic mismatch: got {magic:#x}, want {LOH_MAGIC:#x}"),
));
}
if !capacity.is_power_of_two() || capacity < 2 {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
format!("loh capacity {capacity} is not a power of two >= 2"),
));
}
if size < loh_file_size(capacity) {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
format!(
"loh file size {size} below expected {}",
loh_file_size(capacity)
),
));
}
let flush_threshold = flush_threshold.max(1);
Ok(Self {
_file: file,
mmap,
capacity,
capacity_mask: (capacity as i64) - 1,
flush_threshold,
lifo_cap: DEFAULT_LIFO_CAP,
local_lifo: Mutex::new(Vec::with_capacity(DEFAULT_LIFO_CAP)),
})
}
pub fn capacity(&self) -> usize {
self.capacity
}
pub fn flush_threshold(&self) -> usize {
self.flush_threshold
}
pub fn owner_pid(&self) -> u64 {
self.header().owner_pid.load(Ordering::Acquire)
}
pub fn close_owner(&self) {
self.header().owner_pid.store(0, Ordering::Release);
self.header().epoch.fetch_add(1, Ordering::Release);
}
fn header(&self) -> &LohHeader {
unsafe { &*(self.mmap.as_ptr() as *const LohHeader) }
}
fn slot_ptr(&self, idx: i64) -> *mut LcrqJobSlot {
let slot_idx = (idx & self.capacity_mask) as usize;
let off = std::mem::size_of::<LohHeader>() + slot_idx * LOH_SLOT_SIZE;
unsafe { self.mmap.as_ptr().add(off) as *mut LcrqJobSlot }
}
pub fn snapshot_size(&self) -> (i64, i64, i64, usize) {
let h = self.header();
let head = h.head.load(Ordering::Acquire);
let tail = h.tail.load(Ordering::Acquire);
let lifo_len = self.local_lifo.try_lock().map(|g| g.len()).unwrap_or(0);
(head, tail, tail - head, lifo_len)
}
pub fn push(&self, item: LineItem) -> Result<(), PushError> {
let mut lifo = self.local_lifo.lock();
if lifo.len() >= self.lifo_cap {
return Err(PushError::LifoFull);
}
lifo.push(item);
if lifo.len() >= self.flush_threshold {
if let Err(e) = self.flush_locked(&mut lifo) {
lifo.pop();
return Err(e);
}
}
Ok(())
}
pub fn flush(&self) -> Result<usize, PushError> {
let mut lifo = self.local_lifo.lock();
self.flush_locked(&mut lifo)
}
pub fn publish_batch(&self, items: &[LineItem]) -> Result<usize, PushError> {
if items.is_empty() {
return Ok(0);
}
let n = items.len();
let h = self.header();
let head_snapshot = h.head.load(Ordering::Acquire);
let tail_snapshot = h.tail.load(Ordering::Relaxed);
if (tail_snapshot - head_snapshot + n as i64) > self.capacity as i64 {
return Err(PushError::Full);
}
let base = h.tail.fetch_add(n as i64, Ordering::AcqRel);
prefetch_slot(self.slot_ptr(base));
for (i, item) in items.iter().enumerate() {
let idx = base + i as i64;
if i + 1 < n {
prefetch_slot(self.slot_ptr(idx + 1));
}
unsafe {
self.publish_at(idx, *item);
}
}
Ok(n)
}
fn flush_locked(&self, lifo: &mut Vec<LineItem>) -> Result<usize, PushError> {
let n = lifo.len();
if n == 0 {
return Ok(0);
}
let h = self.header();
let head_snapshot = h.head.load(Ordering::Acquire);
let tail_snapshot = h.tail.load(Ordering::Relaxed);
if (tail_snapshot - head_snapshot + n as i64) > self.capacity as i64 {
return Err(PushError::Full);
}
let base = h.tail.fetch_add(n as i64, Ordering::AcqRel);
for (i, item) in lifo.drain(..).enumerate() {
let idx = base + i as i64;
unsafe {
self.publish_at(idx, item);
}
}
Ok(n)
}
unsafe fn publish_at(&self, idx: i64, item: LineItem) {
let slot = self.slot_ptr(idx);
loop {
let seq = unsafe { (*slot).sequence.load(Ordering::Acquire) };
let diff = seq - idx;
if diff == 0 {
break;
}
if diff < 0 {
std::hint::spin_loop();
continue;
}
panic!(
"LOH producer protocol violation: slot[{}] seq={} ahead of idx={}",
idx & self.capacity_mask,
seq,
idx
);
}
unsafe {
(*slot).item = item;
(*slot).sequence.store(idx + 1, Ordering::Release);
}
}
pub fn pop_local(&self) -> Option<LineItem> {
let mut lifo = self.local_lifo.lock();
lifo.pop()
}
pub fn steal(&self) -> Steal {
let h = self.header();
let head = h.head.load(Ordering::Acquire);
fence(Ordering::SeqCst);
let tail = h.tail.load(Ordering::Acquire);
if head >= tail {
return Steal::Empty;
}
let slot = self.slot_ptr(head);
let seq = unsafe { (*slot).sequence.load(Ordering::Acquire) };
if seq != head + 1 {
return Steal::Retry;
}
let won = h
.head
.compare_exchange(head, head + 1, Ordering::SeqCst, Ordering::Relaxed)
.is_ok();
if !won {
return Steal::Retry;
}
let result = unsafe {
StealResult {
item: (*slot).item,
}
};
unsafe {
(*slot)
.sequence
.store(head + self.capacity as i64, Ordering::Release);
}
Steal::Success(result)
}
pub fn flush_to_disk(&self) -> io::Result<()> {
self.mmap.flush()
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering as O};
use std::thread;
fn temp_path(name: &str) -> std::path::PathBuf {
let mut p = std::env::temp_dir();
let pid = std::process::id();
let nonce = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_nanos())
.unwrap_or(0);
p.push(format!("subetha_loh_{pid}_{nonce}_{name}.bin"));
p
}
fn u32_item(id: u32) -> LineItem {
LineItem::new(&id.to_le_bytes()).expect("build item")
}
fn item_id(item: &LineItem) -> u32 {
u32::from_le_bytes(item.payload[..4].try_into().unwrap())
}
#[test]
fn create_then_open_round_trips_header() {
let path = temp_path("create_open");
let _d = SharedDequeLoh::create(&path, 8, 4).expect("create");
let o = SharedDequeLoh::open(&path, 4).expect("open");
assert_eq!(o.capacity(), 8);
assert_eq!(o.owner_pid(), std::process::id() as u64);
std::fs::remove_file(&path).ok();
}
#[test]
fn open_rejects_bad_magic() {
let path = temp_path("bad_magic");
std::fs::write(&path, vec![0xCDu8; 8192]).expect("seed");
let r = SharedDequeLoh::open(&path, 4);
assert!(r.is_err());
std::fs::remove_file(&path).ok();
}
#[test]
fn push_and_explicit_flush_migrates() {
let path = temp_path("flush");
let d = SharedDequeLoh::create(&path, 8, usize::MAX).expect("create");
for i in 0..3u32 {
d.push(u32_item(i)).expect("push");
}
let (head, tail, sz, lifo_len) = d.snapshot_size();
assert_eq!(head, 0);
assert_eq!(tail, 0);
assert_eq!(sz, 0);
assert_eq!(lifo_len, 3);
let n = d.flush().expect("flush");
assert_eq!(n, 3);
let (_, tail, sz, lifo_len) = d.snapshot_size();
assert_eq!(tail, 3);
assert_eq!(sz, 3);
assert_eq!(lifo_len, 0);
std::fs::remove_file(&path).ok();
}
#[test]
fn push_auto_flushes_at_threshold() {
let path = temp_path("autoflush");
let d = SharedDequeLoh::create(&path, 8, 4).expect("create");
for i in 0..4u32 {
d.push(u32_item(i)).expect("push");
}
let (_, tail, sz, lifo_len) = d.snapshot_size();
assert_eq!(tail, 4);
assert_eq!(sz, 4);
assert_eq!(lifo_len, 0);
std::fs::remove_file(&path).ok();
}
#[test]
fn publish_batch_migrates_in_fifo_order() {
let path = temp_path("publish_batch");
let d = SharedDequeLoh::create(&path, 64, usize::MAX).expect("create");
let items: Vec<LineItem> = (1..=5u32).map(u32_item).collect();
let n = d.publish_batch(&items).expect("publish_batch");
assert_eq!(n, 5);
let (_, tail, sz, lifo_len) = d.snapshot_size();
assert_eq!(tail, 5);
assert_eq!(sz, 5);
assert_eq!(lifo_len, 0);
for expected in 1..=5u32 {
loop {
match d.steal() {
Steal::Success(r) => {
assert_eq!(item_id(&r.item), expected);
break;
}
Steal::Empty | Steal::Retry => std::thread::yield_now(),
}
}
}
assert!(matches!(d.steal(), Steal::Empty));
std::fs::remove_file(&path).ok();
}
#[test]
fn publish_batch_empty_is_noop() {
let path = temp_path("publish_batch_empty");
let d = SharedDequeLoh::create(&path, 4, usize::MAX).expect("create");
let n = d.publish_batch(&[]).expect("publish_batch empty");
assert_eq!(n, 0);
let (_, tail, sz, _) = d.snapshot_size();
assert_eq!(tail, 0);
assert_eq!(sz, 0);
std::fs::remove_file(&path).ok();
}
#[test]
fn publish_batch_full_returns_full() {
let path = temp_path("publish_batch_full");
let d = SharedDequeLoh::create(&path, 4, usize::MAX).expect("create");
let items: Vec<LineItem> = (1..=4u32).map(u32_item).collect();
d.publish_batch(&items).expect("publish first batch");
let err = d
.publish_batch(&[u32_item(99)])
.expect_err("publish past capacity");
assert_eq!(err, PushError::Full);
std::fs::remove_file(&path).ok();
}
#[test]
fn steal_drains_in_fifo_order_after_flush() {
let path = temp_path("fifo");
let d = SharedDequeLoh::create(&path, 8, usize::MAX).expect("create");
for i in 1..=3u32 {
d.push(u32_item(i)).expect("push");
}
d.flush().expect("flush");
for expected in 1..=3u32 {
loop {
match d.steal() {
Steal::Success(slot) => {
assert_eq!(item_id(&slot.item), expected);
break;
}
Steal::Empty | Steal::Retry => std::thread::yield_now(),
}
}
}
assert!(matches!(d.steal(), Steal::Empty));
std::fs::remove_file(&path).ok();
}
#[test]
fn pop_local_drains_lifo_in_lifo_order() {
let path = temp_path("pop_local_lifo");
let d = SharedDequeLoh::create(&path, 4, usize::MAX).expect("create");
for i in 1..=3u32 {
d.push(u32_item(i)).expect("push");
}
for expected in (1..=3u32).rev() {
let e = d.pop_local().expect("pop_local");
assert_eq!(item_id(&e), expected);
}
assert!(d.pop_local().is_none());
std::fs::remove_file(&path).ok();
}
#[test]
fn ring_full_at_capacity() {
let path = temp_path("full");
let d = SharedDequeLoh::create(&path, 2, usize::MAX).expect("create");
d.push(u32_item(1)).expect("push");
d.push(u32_item(2)).expect("push");
let n = d.flush().expect("flush");
assert_eq!(n, 2);
d.push(u32_item(3)).expect("push to lifo");
let err = d.flush().expect_err("flush past capacity");
assert_eq!(err, PushError::Full);
std::fs::remove_file(&path).ok();
}
#[test]
fn close_owner_zeros_pid_and_advances_epoch() {
let path = temp_path("close");
let d = SharedDequeLoh::create(&path, 2, 1).expect("create");
assert_eq!(d.owner_pid(), std::process::id() as u64);
let h = d.header();
let before = h.epoch.load(O::Acquire);
d.close_owner();
assert_eq!(d.owner_pid(), 0);
assert_eq!(h.epoch.load(O::Acquire), before + 1);
std::fs::remove_file(&path).ok();
}
#[test]
fn concurrent_thieves_no_double_take() {
let path = temp_path("stress");
let d = Arc::new(SharedDequeLoh::create(&path, 128, 8).expect("create"));
let n = 5_000usize;
let consumed = Arc::new(AtomicUsize::new(0));
let sum = Arc::new(AtomicUsize::new(0));
let mut thieves = Vec::new();
for _ in 0..2 {
let d = Arc::clone(&d);
let consumed = Arc::clone(&consumed);
let sum = Arc::clone(&sum);
thieves.push(thread::spawn(move || {
while consumed.load(O::Relaxed) < n {
match d.steal() {
Steal::Success(slot) => {
consumed.fetch_add(1, O::Relaxed);
sum.fetch_add(item_id(&slot.item) as usize, O::Relaxed);
}
Steal::Empty | Steal::Retry => std::thread::yield_now(),
}
}
}));
}
for i in 0..n {
loop {
match d.push(u32_item(i as u32)) {
Ok(()) => break,
Err(PushError::LifoFull) | Err(PushError::Full) => {
std::thread::yield_now();
d.flush().ok();
}
}
}
}
loop {
match d.flush() {
Ok(_) => break,
Err(PushError::Full) => std::thread::yield_now(),
Err(e) => panic!("terminal flush: {e:?}"),
}
}
for h in thieves {
h.join().expect("thief");
}
let expected: usize = (0..n).sum();
assert_eq!(
sum.load(O::Relaxed),
expected,
"every slot consumed once"
);
std::fs::remove_file(&path).ok();
}
#[test]
fn publish_batch_stress_two_thieves() {
let path = temp_path("publish_batch_stress");
let d = Arc::new(
SharedDequeLoh::create(&path, 256, usize::MAX).expect("create"),
);
let n = 5_000usize;
let consumed = Arc::new(AtomicUsize::new(0));
let sum = Arc::new(AtomicUsize::new(0));
let mut thieves = Vec::new();
for _ in 0..2 {
let d = Arc::clone(&d);
let consumed = Arc::clone(&consumed);
let sum = Arc::clone(&sum);
thieves.push(thread::spawn(move || {
while consumed.load(O::Relaxed) < n {
match d.steal() {
Steal::Success(slot) => {
consumed.fetch_add(1, O::Relaxed);
sum.fetch_add(item_id(&slot.item) as usize, O::Relaxed);
}
Steal::Empty | Steal::Retry => std::thread::yield_now(),
}
}
}));
}
let mut pushed = 0usize;
let burst = 64usize;
while pushed < n {
let want = burst.min(n - pushed);
let batch: Vec<LineItem> = (0..want)
.map(|j| u32_item((pushed + j) as u32))
.collect();
loop {
match d.publish_batch(&batch) {
Ok(_) => break,
Err(PushError::Full) => std::thread::yield_now(),
Err(other) => panic!("publish_batch: {other:?}"),
}
}
pushed += want;
}
for t in thieves {
t.join().expect("thief");
}
let expected: usize = (0..n).sum();
assert_eq!(
sum.load(O::Relaxed),
expected,
"publish_batch stress: every item consumed once"
);
std::fs::remove_file(&path).ok();
}
}