use crossbeam_channel::{Receiver, Sender, bounded};
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
pub const DEFAULT_SLOT_SIZE: usize = 200 * 1024 * 1024;
pub const DEFAULT_NUM_SLOTS: usize = 8;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct PoolPlan {
pub num_slots: usize,
pub slot_size: usize,
pub slice_size: usize,
}
impl PoolPlan {
pub fn bytes(&self) -> u64 {
self.num_slots as u64 * self.slot_size as u64
}
pub fn default_for(num_workers: usize) -> Self {
PoolPlan {
num_slots: DEFAULT_NUM_SLOTS,
slot_size: DEFAULT_SLOT_SIZE,
slice_size: Self::slice_size_for(num_workers),
}
}
pub fn slice_size_for(num_workers: usize) -> usize {
(DEFAULT_SLOT_SIZE / num_workers.max(1)).max(1)
}
pub fn plan(input_bytes: u64, num_workers: usize, budget_bytes: u64) -> Self {
let slice_size = Self::slice_size_for(num_workers);
let slice = slice_size as u64;
let default_slices = (DEFAULT_NUM_SLOTS as u64) * (DEFAULT_SLOT_SIZE as u64) / slice;
let needed = input_bytes.div_ceil(slice).max(1);
let affordable = (budget_bytes / slice).max(1);
let slices = needed.min(affordable).min(default_slices);
if slices >= default_slices {
return Self::default_for(num_workers);
}
let num_slots = (slices as usize).min(DEFAULT_NUM_SLOTS).max(1);
let per_slot = ((slices as usize) / num_slots).max(1);
let slot_size = (per_slot * slice_size).min(DEFAULT_SLOT_SIZE);
PoolPlan { num_slots, slot_size, slice_size }
}
}
#[derive(Copy, Clone)]
struct SlotPtr(*mut u8);
unsafe impl Send for SlotPtr {}
unsafe impl Sync for SlotPtr {}
pub struct Round {
pub slot_id: u32,
ptr: *const u8,
pub len: usize,
pub skip: bool,
pub file_index: u64,
pub fdata_offset: u64,
pub chunk_seq: u32,
}
unsafe impl Send for Round {}
impl Round {
pub unsafe fn as_slice<'a>(&self) -> &'a [u8] {
unsafe { std::slice::from_raw_parts(self.ptr, self.len) }
}
}
#[derive(Clone)]
pub struct Ejector {
inner: Arc<EjectorInner>,
}
struct EjectorInner {
free_tx: Sender<u32>,
outstanding: Vec<AtomicUsize>,
}
impl Ejector {
pub fn release_one(&self, slot_id: u32) {
let prev = self.inner.outstanding[slot_id as usize].fetch_sub(1, Ordering::AcqRel);
debug_assert!(prev >= 1, "release_one underflow on slot {slot_id}");
if prev == 1 {
self.inner.free_tx.send(slot_id).ok();
}
}
}
pub struct Magazine {
slot_size: usize,
slice_size: usize,
base: Vec<SlotPtr>,
_mem: Vec<Box<[u8]>>, free_rx: Receiver<u32>,
ret: Ejector,
}
impl Magazine {
pub fn new(num_slots: usize, slot_size: usize, num_workers: usize) -> Self {
let slice_size = (slot_size / num_workers.max(1)).max(1);
Self::with_slice_size(num_slots, slot_size, slice_size)
}
pub fn from_plan(plan: PoolPlan) -> Self {
Self::with_slice_size(plan.num_slots, plan.slot_size, plan.slice_size)
}
pub fn with_slice_size(num_slots: usize, slot_size: usize, slice_size: usize) -> Self {
assert!(num_slots > 0 && slot_size > 0 && slice_size > 0);
assert!(
slice_size <= slot_size,
"slice_size {slice_size} exceeds slot_size {slot_size}: a slice must fit \
contiguously in one slot"
);
let mut mem: Vec<Box<[u8]>> = (0..num_slots)
.map(|_| vec![0u8; slot_size].into_boxed_slice())
.collect();
let base: Vec<SlotPtr> = mem.iter_mut().map(|b| SlotPtr(b.as_mut_ptr())).collect();
let (free_tx, free_rx) = bounded(num_slots);
for id in 0..num_slots as u32 {
free_tx.send(id).expect("free channel send during init");
}
let outstanding = (0..num_slots).map(|_| AtomicUsize::new(0)).collect();
Magazine {
slot_size,
slice_size,
base,
_mem: mem,
free_rx,
ret: Ejector { inner: Arc::new(EjectorInner { free_tx, outstanding }) },
}
}
pub fn slot_size(&self) -> usize {
self.slot_size
}
pub fn slice_size(&self) -> usize {
self.slice_size
}
pub fn num_slots(&self) -> usize {
self.base.len()
}
pub fn returner(&self) -> Ejector {
self.ret.clone()
}
pub fn claim(&self) -> Option<Clip<'_>> {
let slot_id = self.free_rx.recv().ok()?;
Some(Clip { pool: self, slot_id, cursor: 0, slices: Vec::new() })
}
}
pub struct Clip<'a> {
pool: &'a Magazine,
slot_id: u32,
cursor: usize,
slices: Vec<Round>,
}
impl<'a> Clip<'a> {
pub fn slot_id(&self) -> u32 {
self.slot_id
}
pub fn remaining(&self) -> usize {
self.pool.slot_size - self.cursor
}
pub fn writable(&mut self, max: usize) -> &mut [u8] {
let n = max.min(self.remaining());
let base = self.pool.base[self.slot_id as usize].0;
unsafe { std::slice::from_raw_parts_mut(base.add(self.cursor), n) }
}
pub fn commit_slice(
&mut self,
len: usize,
skip: bool,
file_index: u64,
fdata_offset: u64,
chunk_seq: u32,
) {
debug_assert!(len <= self.remaining());
let base = self.pool.base[self.slot_id as usize].0 as *const u8;
let ptr = unsafe { base.add(self.cursor) };
self.slices.push(Round {
slot_id: self.slot_id,
ptr,
len,
skip,
file_index,
fdata_offset,
chunk_seq,
});
self.cursor += len;
}
#[must_use]
pub fn publish(self) -> Vec<Round> {
let n = self.slices.len();
if n == 0 {
self.pool.ret.inner.free_tx.send(self.slot_id).ok();
return Vec::new();
}
self.pool.ret.inner.outstanding[self.slot_id as usize].store(n, Ordering::Release);
self.slices
}
}
#[cfg(test)]
mod tests {
use super::*;
macro_rules! assert_emit {
($check:expr, $ok:expr, $($detail:tt)+) => {{
let __ok: bool = $ok;
let __detail = format!($($detail)+);
#[cfg(feature = "testmatrix")]
crate::functional_status("znippy-common/magazine", $check, __ok, &__detail);
assert!(__ok, "znippy-common/magazine::{} — {}", $check, __detail);
}};
}
#[test]
fn slot_returns_only_after_last_slice() {
let pool = Magazine::new(2, 64, 4);
assert_eq!(pool.slice_size(), 16);
let ret = pool.returner();
let mut a = pool.claim().unwrap();
let _b = pool.claim().unwrap();
let n = { a.writable(10).len().min(10) };
a.commit_slice(n, false, 0, 0, 0);
let n2 = { a.writable(20).len().min(20) };
a.commit_slice(n2, false, 1, 0, 0);
let slices = a.publish();
assert_eq!(slices.len(), 2);
let slot_a = slices[0].slot_id;
ret.release_one(slot_a);
let held_after_first = pool.claim_now().is_none();
assert_emit!(
"slot_held_until_last_slice",
held_after_first,
"slot={slot_a} not reclaimable while 1/2 slices still outstanding"
);
ret.release_one(slot_a);
let freed_after_last = pool.claim_now().is_some();
assert_emit!(
"slot_freed_after_last_slice",
freed_after_last,
"slot={slot_a} reclaimed once outstanding counter hit zero"
);
}
#[test]
fn pool_plan_is_bounded_by_input_and_budget() {
const WORKERS: usize = 16;
let slice = PoolPlan::slice_size_for(WORKERS);
let huge = u64::MAX / 4;
let default_bytes = (DEFAULT_NUM_SLOTS * DEFAULT_SLOT_SIZE) as u64;
let fat = PoolPlan::plan(huge, WORKERS, huge);
assert_eq!(
fat,
PoolPlan::default_for(WORKERS),
"an unconstrained plan must be the historical 8 × 200 MiB geometry"
);
assert_eq!(fat.bytes(), default_bytes, "unconstrained reservation is 1.6 GiB");
let five_mib = 5 * 1024 * 1024;
let tiny = PoolPlan::plan(five_mib, WORKERS, huge);
assert!(
tiny.bytes() < default_bytes / 100,
"5 MiB of input reserved {} bytes — must be a sliver of the old 1.6 GiB",
tiny.bytes()
);
assert!(
tiny.bytes() >= five_mib,
"the pool must still be able to hold the input in flight ({} < {five_mib})",
tiny.bytes()
);
let budget = 64 * 1024 * 1024;
let pod = PoolPlan::plan(huge, WORKERS, budget);
assert!(
pod.bytes() <= budget,
"planned {} bytes over a {budget}-byte budget",
pod.bytes()
);
let starved = PoolPlan::plan(huge, WORKERS, 1);
assert_eq!(starved.num_slots, 1);
assert_eq!(starved.slot_size, slice, "the floor is exactly one slice");
for p in [fat, tiny, pod, starved] {
assert_eq!(
p.slice_size, slice,
"slice_size moved with the memory plan — that changes the big/small \
partition and every big-file chunk boundary, i.e. the archive bytes"
);
assert!(p.slice_size <= p.slot_size, "a slice must fit in a slot");
assert!(p.num_slots >= 1 && p.slot_size >= 1);
}
for workers in [1usize, 2, 3, 7, 16, 29, 32, 64] {
let ss = PoolPlan::slice_size_for(workers) as u64;
for budget in [
1u64, 7, ss - 1, ss, ss + 1, ss * 3, ss * 8, ss * 9, ss * 17,
64 << 20, 100 << 20, 256 << 20, 1 << 30,
] {
let p = PoolPlan::plan(huge, workers, budget);
assert_eq!(p.slice_size as u64, ss, "slice_size moved (workers={workers})");
assert!(p.slice_size <= p.slot_size, "slice must fit in a slot");
let allowed = budget.max(ss).min(default_bytes);
assert!(
p.bytes() <= allowed,
"workers={workers} budget={budget}: planned {} bytes, allowed {allowed}",
p.bytes()
);
}
for input in [0u64, 1, ss / 2, ss, ss * 5, ss * 300, 1 << 30] {
let p = PoolPlan::plan(input, workers, huge);
let allowed = input.max(ss).next_multiple_of(ss).min(default_bytes) + ss;
assert!(
p.bytes() <= allowed,
"workers={workers} input={input}: planned {} bytes for {input} of input",
p.bytes()
);
}
}
}
#[test]
fn a_starved_plan_still_cycles_slots() {
let plan = PoolPlan::plan(u64::MAX / 4, 8, 1);
let pool = Magazine::from_plan(plan);
assert_eq!(pool.num_slots(), 1, "starved plan is a single slot");
assert_eq!(pool.slice_size(), plan.slice_size);
let ret = pool.returner();
let mut clip = pool.claim().unwrap();
let n = clip.writable(plan.slice_size).len();
assert_eq!(n, plan.slice_size, "the one slot holds exactly one full slice");
clip.commit_slice(n, false, 0, 0, 0);
let slices = clip.publish();
assert_eq!(slices.len(), 1);
assert!(pool.claim_now().is_none(), "the only slot is still outstanding");
ret.release_one(slices[0].slot_id);
assert!(pool.claim_now().is_some(), "released slot returns to the free list");
}
#[test]
fn writable_clamps_to_remaining() {
let pool = Magazine::new(1, 32, 4);
let mut f = pool.claim().unwrap();
let full = f.writable(1000).len() == 32;
f.commit_slice(30, false, 0, 0, 0);
let clamped = f.remaining() == 2 && f.writable(1000).len() == 2;
assert_emit!(
"writable_clamps_to_remaining",
full && clamped,
"writable() never exceeds the slot's remaining bytes (no OOB into the next slot)"
);
}
impl Magazine {
fn claim_now(&self) -> Option<u32> {
self.free_rx.try_recv().ok()
}
}
}