#![allow(clippy::missing_errors_doc)]
use std::fs::{File, OpenOptions};
use std::io;
use std::path::Path;
use std::sync::atomic::{fence, AtomicI64, AtomicU64, Ordering};
use memmap2::{MmapMut, MmapOptions};
use subetha_core::has_movdir64b;
use crate::shared_deque_khpd::LineItem;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PublishRadius {
Local,
Distant,
}
impl PublishRadius {
pub fn pick_auto() -> Self {
if has_movdir64b() {
Self::Distant
} else {
Self::Local
}
}
pub fn resolve(self) -> Self {
match self {
Self::Local => Self::Local,
Self::Distant => {
if has_movdir64b() {
Self::Distant
} else {
Self::Local
}
}
}
}
}
#[inline(always)]
fn prefetchw_slot(slot: *const KhlSlot) {
#[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 KHL_MAGIC: u64 = 0x574B_484C_0000_0002;
pub const KHL_SLOT_SIZE: usize = 64;
pub const KHL_ITEMS_PER_SLOT: usize = 3;
#[repr(C, align(64))]
pub struct KhlHeader {
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 KhlSlot {
pub packed_sequence: AtomicI64,
pub _reserved: u64,
pub items: [LineItem; KHL_ITEMS_PER_SLOT],
}
#[inline(always)]
pub const fn pack_seq(idx_value: i64, n_items: usize) -> i64 {
(idx_value << 2) | (n_items as i64 & 0x3)
}
#[inline(always)]
pub const fn unpack_idx(packed: i64) -> i64 {
packed >> 2
}
#[inline(always)]
pub const fn unpack_n_items(packed: i64) -> usize {
(packed & 0x3) as usize
}
pub const fn khl_file_size(capacity: usize) -> usize {
std::mem::size_of::<KhlHeader>() + capacity * KHL_SLOT_SIZE
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PushError {
Full,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Steal {
Success(StealResult),
Empty,
Retry,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct StealResult {
pub n_items: usize,
pub items: [LineItem; KHL_ITEMS_PER_SLOT],
}
pub struct SharedDequeKhl {
_file: File,
mmap: MmapMut,
capacity: usize,
capacity_mask: i64,
publish_radius: PublishRadius,
}
unsafe impl Send for SharedDequeKhl {}
unsafe impl Sync for SharedDequeKhl {}
impl SharedDequeKhl {
pub fn create<P: AsRef<Path>>(path: P, capacity: usize) -> io::Result<Self> {
let capacity = capacity.max(2).next_power_of_two();
let size = khl_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 KhlHeader;
unsafe {
(*header_ptr).magic = KHL_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::<KhlHeader>();
for i in 0..capacity {
let off = slots_start + i * KHL_SLOT_SIZE;
let slot_ptr = unsafe { mmap.as_mut_ptr().add(off) as *mut KhlSlot };
unsafe {
(*slot_ptr).packed_sequence =
AtomicI64::new(pack_seq(i as i64, 0));
(*slot_ptr)._reserved = 0;
(*slot_ptr).items = [LineItem::default(); KHL_ITEMS_PER_SLOT];
}
}
mmap.flush()?;
Ok(Self {
_file: file,
mmap,
capacity,
capacity_mask: (capacity as i64) - 1,
publish_radius: PublishRadius::pick_auto(),
})
}
pub fn open<P: AsRef<Path>>(path: P) -> 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::<KhlHeader>() {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"khl file too small to contain header",
));
}
let mmap = unsafe { MmapOptions::new().len(size).map_mut(&file)? };
let header_ptr = mmap.as_ptr() as *const KhlHeader;
let (magic, capacity) =
unsafe { ((*header_ptr).magic, (*header_ptr).capacity as usize) };
if magic != KHL_MAGIC {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
format!("khl magic mismatch: got {magic:#x}, want {KHL_MAGIC:#x}"),
));
}
if !capacity.is_power_of_two() || capacity < 2 {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
format!("khl capacity {capacity} is not pow2 >= 2"),
));
}
if size < khl_file_size(capacity) {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
format!(
"khl file size {size} below expected {}",
khl_file_size(capacity)
),
));
}
Ok(Self {
_file: file,
mmap,
capacity,
capacity_mask: (capacity as i64) - 1,
publish_radius: PublishRadius::pick_auto(),
})
}
pub fn capacity(&self) -> usize {
self.capacity
}
pub fn publish_radius(&self) -> PublishRadius {
self.publish_radius
}
pub fn with_publish_radius(mut self, radius: PublishRadius) -> Self {
self.publish_radius = radius.resolve();
self
}
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) -> &KhlHeader {
unsafe { &*(self.mmap.as_ptr() as *const KhlHeader) }
}
fn slot_ptr(&self, idx: i64) -> *mut KhlSlot {
let slot_idx = (idx & self.capacity_mask) as usize;
let off = std::mem::size_of::<KhlHeader>() + slot_idx * KHL_SLOT_SIZE;
unsafe { self.mmap.as_ptr().add(off) as *mut KhlSlot }
}
pub fn snapshot_size(&self) -> (i64, i64, i64) {
let h = self.header();
let head = h.head.load(Ordering::Acquire);
let tail = h.tail.load(Ordering::Acquire);
(head, tail, tail - head)
}
pub fn publish_batch(&self, items: &[LineItem]) -> Result<usize, PushError> {
if items.is_empty() {
return Ok(0);
}
let k = items.len();
let n_slots = k.div_ceil(KHL_ITEMS_PER_SLOT);
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_slots as i64) > self.capacity as i64 {
return Err(PushError::Full);
}
let base = tail_snapshot;
prefetchw_slot(self.slot_ptr(base));
let mut written = 0usize;
for slot_i in 0..n_slots {
let idx = base + slot_i as i64;
if slot_i + 1 < n_slots {
prefetchw_slot(self.slot_ptr(idx + 1));
}
let take = (k - written).min(KHL_ITEMS_PER_SLOT);
unsafe {
self.publish_slot_at(idx, &items[written..written + take]);
}
written += take;
}
h.tail.store(base + n_slots as i64, Ordering::Release);
Ok(k)
}
unsafe fn publish_slot_at(&self, idx: i64, items: &[LineItem]) {
let slot = self.slot_ptr(idx);
loop {
let packed = unsafe {
(*slot).packed_sequence.load(Ordering::Acquire)
};
let idx_value = unpack_idx(packed);
let diff = idx_value - idx;
if diff == 0 {
break;
}
if diff < 0 {
std::hint::spin_loop();
continue;
}
panic!(
"KHL producer protocol violation: slot[{}] idx_value={} ahead of idx={}",
idx & self.capacity_mask,
idx_value,
idx
);
}
match self.publish_radius {
PublishRadius::Local => {
unsafe {
let n = items.len();
for (i, item) in items.iter().enumerate() {
(*slot).items[i] = *item;
}
(*slot)
.packed_sequence
.store(pack_seq(idx + 1, n), Ordering::Release);
}
}
PublishRadius::Distant => {
unsafe {
self.publish_slot_movdir64b(slot, idx, items);
}
}
}
}
#[inline(always)]
unsafe fn publish_slot_movdir64b(
&self,
slot: *mut KhlSlot,
idx: i64,
items: &[LineItem],
) {
#[repr(C, align(64))]
struct SrcLine {
packed_sequence: i64,
_reserved: u64,
items: [LineItem; KHL_ITEMS_PER_SLOT],
}
let mut src = SrcLine {
packed_sequence: pack_seq(idx + 1, items.len()),
_reserved: 0,
items: [LineItem::default(); KHL_ITEMS_PER_SLOT],
};
for (i, item) in items.iter().enumerate() {
src.items[i] = *item;
}
let dst_ptr = slot as *mut u8;
let src_ptr = &src as *const SrcLine as *const u8;
#[cfg(target_arch = "x86_64")]
{
unsafe {
core::arch::asm!(
"movdir64b {dst}, [{src}]",
"sfence",
dst = in(reg) dst_ptr,
src = in(reg) src_ptr,
options(nostack, preserves_flags),
);
}
}
#[cfg(not(target_arch = "x86_64"))]
{
_ = dst_ptr;
_ = src_ptr;
_ = slot;
_ = idx;
unreachable!(
"publish_slot_movdir64b reached on non-x86_64 host; \
PublishRadius::resolve() returns Local there"
);
}
}
pub fn steal_slot(&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 packed = unsafe {
(*slot).packed_sequence.load(Ordering::Acquire)
};
let idx_value = unpack_idx(packed);
if idx_value != head + 1 {
return Steal::Retry;
}
let n = unpack_n_items(packed).min(KHL_ITEMS_PER_SLOT);
let won = h
.head
.compare_exchange(head, head + 1, Ordering::SeqCst, Ordering::Relaxed)
.is_ok();
if !won {
return Steal::Retry;
}
let result = unsafe {
StealResult {
n_items: n,
items: (*slot).items,
}
};
unsafe {
(*slot).packed_sequence.store(
pack_seq(head + self.capacity as i64, 0),
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_khl_{pid}_{nonce}_{name}.bin"));
p
}
fn u32_item(id: u32) -> LineItem {
LineItem::new(&id.to_le_bytes()).expect("item")
}
fn item_id(item: &LineItem) -> u32 {
u32::from_le_bytes(item.payload[..4].try_into().unwrap())
}
#[test]
fn publish_radius_matches_host() {
let path = temp_path("radius_auto");
let d = SharedDequeKhl::create(&path, 8).expect("create");
let r = d.publish_radius();
if subetha_core::has_movdir64b() {
assert_eq!(r, PublishRadius::Distant);
} else {
assert_eq!(r, PublishRadius::Local);
}
std::fs::remove_file(&path).ok();
}
#[test]
fn publish_radius_distant_resolves_to_local_without_movdir64b() {
let path = temp_path("radius_resolve");
let d = SharedDequeKhl::create(&path, 8)
.expect("create")
.with_publish_radius(PublishRadius::Distant);
if subetha_core::has_movdir64b() {
assert_eq!(d.publish_radius(), PublishRadius::Distant);
} else {
assert_eq!(d.publish_radius(), PublishRadius::Local);
}
std::fs::remove_file(&path).ok();
}
#[test]
fn publish_then_drain_works_under_both_radius_modes() {
let path = temp_path("radius_round_trip");
let d = SharedDequeKhl::create(&path, 8).expect("create");
let items: Vec<LineItem> = (1..=6u32).map(u32_item).collect();
d.publish_batch(&items).expect("publish");
let mut drained = Vec::new();
loop {
match d.steal_slot() {
Steal::Success(r) => {
for i in 0..r.n_items {
drained.push(item_id(&r.items[i]));
}
}
Steal::Empty => break,
Steal::Retry => continue,
}
}
assert_eq!(drained, vec![1, 2, 3, 4, 5, 6]);
std::fs::remove_file(&path).ok();
}
#[test]
fn create_then_open_round_trips_header() {
let path = temp_path("create_open");
let _d = SharedDequeKhl::create(&path, 8).expect("create");
let o = SharedDequeKhl::open(&path).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");
assert!(SharedDequeKhl::open(&path).is_err());
std::fs::remove_file(&path).ok();
}
#[test]
fn publish_batch_packs_three_items_per_slot() {
let path = temp_path("publish_batch_packs");
let d = SharedDequeKhl::create(&path, 64).expect("create");
let items: Vec<LineItem> = (1..=7u32).map(u32_item).collect();
let n = d.publish_batch(&items).expect("publish_batch");
assert_eq!(n, 7);
let (_, tail, sz) = d.snapshot_size();
assert_eq!(tail, 3);
assert_eq!(sz, 3);
std::fs::remove_file(&path).ok();
}
#[test]
fn publish_batch_empty_is_noop() {
let path = temp_path("publish_empty");
let d = SharedDequeKhl::create(&path, 4).expect("create");
assert_eq!(d.publish_batch(&[]).expect("noop"), 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_full");
let d = SharedDequeKhl::create(&path, 2).expect("create");
let first: Vec<LineItem> = (1..=6u32).map(u32_item).collect();
d.publish_batch(&first).expect("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_publication_order() {
let path = temp_path("steal_order");
let d = SharedDequeKhl::create(&path, 8).expect("create");
let items: Vec<LineItem> = (1..=7u32).map(u32_item).collect();
d.publish_batch(&items).expect("publish");
let mut drained = Vec::new();
loop {
match d.steal_slot() {
Steal::Success(r) => {
for i in 0..r.n_items {
drained.push(item_id(&r.items[i]));
}
}
Steal::Empty => break,
Steal::Retry => continue,
}
}
assert_eq!(drained, vec![1, 2, 3, 4, 5, 6, 7]);
std::fs::remove_file(&path).ok();
}
#[test]
fn close_owner_zeros_pid_and_advances_epoch() {
let path = temp_path("close");
let d = SharedDequeKhl::create(&path, 2).expect("create");
let before = d.header().epoch.load(O::Acquire);
d.close_owner();
assert_eq!(d.owner_pid(), 0);
assert_eq!(d.header().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(SharedDequeKhl::create(&path, 256).expect("create"));
let n: usize = 5_000;
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_slot() {
Steal::Success(r) => {
for i in 0..r.n_items {
consumed.fetch_add(1, O::Relaxed);
sum.fetch_add(
item_id(&r.items[i]) as usize,
O::Relaxed,
);
}
}
Steal::Empty | Steal::Retry => std::thread::yield_now(),
}
}
}));
}
let burst = 64usize;
let mut pushed = 0usize;
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(),
}
}
pushed += want;
}
for t in thieves {
t.join().expect("thief");
}
let expected: usize = (0..n).sum();
assert_eq!(sum.load(O::Relaxed), expected, "every item consumed once");
std::fs::remove_file(&path).ok();
}
}