znippy_common/

slotpool.rs

//! Magazine — shared large-slot buffer pool for the no-barrier slice pipeline.
//!
//! Design: TODO_NOW.md "THE FINAL SOLUTION 2026". Replaces the per-thread
//! ChunkRevolver rings. One reader thread fills shared slots; N workers pull
//! slices off a single global queue (built by the caller) and never wait for a
//! slot to complete — a worker that finishes a slice grabs the next one from
//! ANY slot.
//!
//! Lifecycle of a slot:
//!   FREE → (reader `claim`s) FILLING → (reader `publish`es) DRAINING
//!        → (workers `release_one` each slice; last one frees it) FREE
//!
//! Packing (done by the reader via `Clip`):
//!   - small file (≤ slice_size): read into the slot at the cursor, committed
//!     as one variable-width slice; many small files coalesce into one slot.
//!   - big file (> slice_size): cut into slice_size pieces, each committed as a
//!     slice; the file spills across as many slots as needed.
//!   Invariant: a slice ⊆ one file AND ⊆ one slot (contiguous in its slot).
//!
//! Safety: while a slot is FILLING only the reader touches it (exclusive, via
//! `writable`). After `publish` the reader drops the `Clip` and only workers
//! read the slot (shared, via `Round::as_slice`). The slot is not handed back to
//! the reader (`claim`) until its outstanding counter hits zero, so the
//! &mut/&  windows never overlap.

use crossbeam_channel::{Receiver, Sender, bounded};
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};

/// Slot size the pool uses when memory is plentiful and the input is large —
/// the historical hardcoded value, kept exactly so the fat path is unchanged.
pub const DEFAULT_SLOT_SIZE: usize = 200 * 1024 * 1024;

/// Slot count the pool uses when memory is plentiful and the input is large.
/// `DEFAULT_NUM_SLOTS * DEFAULT_SLOT_SIZE` = 1.6 GiB, the old unconditional
/// reservation.
pub const DEFAULT_NUM_SLOTS: usize = 8;

/// The geometry of one pass's slot pool: how many slots, how big each is, and —
/// the load-bearing part — the `slice_size` at which the reader cuts files.
///
/// ## Why `slice_size` is planned separately from `slot_size`
/// `slice_size` is an **output-affecting** quantity: it is the big/small
/// partition threshold in `compress_dir` AND the chunk length a big file is cut
/// into, so it lands in `chunk_seq` / `fdata_offset` / every chunk checksum.
/// Change it and the archive changes.
///
/// `slot_size` and `num_slots` are **not** output-affecting. They decide only how
/// many slices coalesce into one buffer and how many buffers are in flight — the
/// reader publishes and claims a fresh slot whenever the next slice does not fit
/// ([`Clip::remaining`]), and the compress sink writes in producer order
/// regardless. So the pool's memory footprint can be shrunk to fit the input or a
/// memory budget while the bytes on disk stay **identical**.
///
/// That is exactly what [`PoolPlan::plan`] does: `slice_size` is pinned to
/// `DEFAULT_SLOT_SIZE / num_workers` — the value the old hardcoded pool implied —
/// and only the reservation moves.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct PoolPlan {
    pub num_slots: usize,
    pub slot_size: usize,
    pub slice_size: usize,
}

impl PoolPlan {
    /// Bytes this plan will allocate up front.
    pub fn bytes(&self) -> u64 {
        self.num_slots as u64 * self.slot_size as u64
    }

    /// The unshrunk geometry: `DEFAULT_NUM_SLOTS × DEFAULT_SLOT_SIZE` (1.6 GiB).
    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),
        }
    }

    /// The output-affecting cut length. Depends ONLY on `num_workers`, never on
    /// how much of the pool we could afford — that invariant is what keeps a
    /// memory-shrunk run byte-identical to a fat-box run on the same host.
    pub fn slice_size_for(num_workers: usize) -> usize {
        (DEFAULT_SLOT_SIZE / num_workers.max(1)).max(1)
    }

    /// Plan a pool that is no larger than it needs to be.
    ///
    /// The reservation is the smallest of three bounds:
    ///  * **the old default** — `DEFAULT_NUM_SLOTS × DEFAULT_SLOT_SIZE`; this
    ///    function never reserves *more* than the code did before,
    ///  * **the input** — a pass that will read `input_bytes` can never keep more
    ///    than `ceil(input_bytes / slice_size)` slices in flight, so slots beyond
    ///    that are pure waste (a 5 MiB staging tree no longer reserves 1.6 GiB),
    ///  * **the budget** — `budget_bytes` from the caller (see
    ///    `common_config::slot_pool_budget_bytes`, which reads the cgroup limit),
    ///    so a constrained pod shrinks instead of being OOM-killed.
    ///
    /// The floor is one slot of one slice: the pipeline must be able to hold the
    /// largest cut it can produce, so it always makes progress no matter how
    /// small the budget claims to be.
    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;

        // Slices the *default* pool holds. Reaching this bound means nothing is
        // constrained — return the historical geometry EXACTLY (not a rounded
        // reconstruction of it), so the fat/bench path is bit-for-bit the old code.
        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);
        }

        // Spread the affordable slices over as many slots as we can (more slots =
        // more reader/worker overlap). Round the per-slot count DOWN: rounding up
        // would multiply the shortfall by `num_slots` and blow the budget the
        // whole plan exists to respect.
        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 }
    }
}

/// Raw base pointer into a slot buffer. Send+Sync by the lifecycle discipline
/// documented above (same contract as `chunkrevolver::SendPtr`).
#[derive(Copy, Clone)]
struct SlotPtr(*mut u8);
unsafe impl Send for SlotPtr {}
unsafe impl Sync for SlotPtr {}

/// One unit of work. Borrows bytes inside slot `slot_id`; valid until the slot
/// is released. Carries everything needed to build a `ChunkMeta` later.
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,
}

// Safety: the pointer addresses an immutable region of a published slot. No
// thread writes the slot between publish and release, so sharing the read-only
// view across worker threads is sound.
unsafe impl Send for Round {}

impl Round {
    /// Borrow the slice bytes.
    ///
    /// Safety: the caller must hold this `Round` only while its slot is
    /// unreleased (i.e. call `Ejector::release_one` for this slot only after
    /// the returned borrow is dropped / the bytes are written).
    pub unsafe fn as_slice<'a>(&self) -> &'a [u8] {
        unsafe { std::slice::from_raw_parts(self.ptr, self.len) }
    }
}

/// Shared handle workers use to release slots back to the reader. Cloneable
/// (it is an `Arc` internally) — move a clone into each worker thread.
#[derive(Clone)]
pub struct Ejector {
    inner: Arc<EjectorInner>,
}

struct EjectorInner {
    free_tx: Sender<u32>,
    outstanding: Vec<AtomicUsize>,
}

impl Ejector {
    /// One slice of `slot_id` has been fully consumed (compressed, or pwritten
    /// for the skip path). When the slot's last slice is done, the slot returns
    /// to the free list so the reader may claim it again.
    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 {
            // Never blocks: we only ever return ids we previously took out, so
            // the bounded free channel can hold them all.
            self.inner.free_tx.send(slot_id).ok();
        }
    }
}

/// Pool of `num_slots` reusable slot buffers. Owned by the reader thread.
pub struct Magazine {
    slot_size: usize,
    slice_size: usize,
    base: Vec<SlotPtr>,
    _mem: Vec<Box<[u8]>>, // backing storage, kept alive for the pool's lifetime
    free_rx: Receiver<u32>,
    ret: Ejector,
}

impl Magazine {
    /// Allocate the pool. `slice_size = slot_size / num_workers` (≥1), the
    /// granularity at which the reader cuts big files. All slots start FREE.
    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)
    }

    /// Allocate the pool from a [`PoolPlan`] — the memory-bounded entry point.
    pub fn from_plan(plan: PoolPlan) -> Self {
        Self::with_slice_size(plan.num_slots, plan.slot_size, plan.slice_size)
    }

    /// Allocate the pool with `slice_size` given **explicitly** rather than
    /// derived from `slot_size`.
    ///
    /// This is the seam that lets the reservation shrink without moving a single
    /// output byte: `slice_size` is the cut length (output-affecting — it lands in
    /// chunk boundaries and checksums), `slot_size × num_slots` is only how much
    /// buffer is held in flight. See [`PoolPlan`].
    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();
        // Take base pointers AFTER the buffers are in their final Vec slots; the
        // heap data addresses are stable from here on (we never reallocate mem).
        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()
    }

    /// Handle workers use to release slots. Clone it into each worker thread.
    pub fn returner(&self) -> Ejector {
        self.ret.clone()
    }

    /// Reader: claim a free slot to fill. Blocks until one is free — this block
    /// is the ONLY backpressure in the pipeline. Returns `None` once the pool is
    /// shut down and drained (all return senders gone).
    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() })
    }
}

/// Reader-side handle for filling one claimed slot. Pack files into it via
/// `writable` + `commit_slice`, then `publish` to hand the slices to the queue.
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
    }

    /// Free space left in the slot.
    pub fn remaining(&self) -> usize {
        self.pool.slot_size - self.cursor
    }

    /// Writable view at the cursor, up to `max` bytes (clamped to remaining).
    /// The reader reads file bytes directly into this — no intermediate buffer.
    /// The borrow must be dropped before calling `commit_slice`.
    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;
        // Safety: exclusive access — this slot is FILLING and only the reader
        // (holding &mut self) touches it; no slices are published yet.
        unsafe { std::slice::from_raw_parts_mut(base.add(self.cursor), n) }
    }

    /// Commit the `len` bytes just written at the cursor as one slice and advance.
    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;
        // Safety: cursor stays within the slot (asserted above).
        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;
    }

    /// Publish the slot: arm its outstanding counter and return the slices for
    /// the caller to push onto the global queue. An empty slot is immediately
    /// returned to the free list (and an empty Vec is returned).
    #[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::*;

    /// `assert!`-with-emit for the Magazine slot-lifecycle invariants. The
    /// `assert!` is the gate (a plain `cargo test` still fails on a broken
    /// invariant); the emit turns the REAL verdict into a nornir matrix row under
    /// `--features testmatrix`. HONEST detail: the Magazine is the forked slot
    /// pool that replaced the per-thread `ChunkRevolver` rings (TODO_NOW "THE
    /// FINAL SOLUTION 2026"); its `Send`/`Sync` is unsafe-by-discipline, so these
    /// lifecycle tests are the standing proof the discipline holds.
    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() {
        // 2 slots, 64 bytes each, 4 workers → slice_size 16.
        let pool = Magazine::new(2, 64, 4);
        assert_eq!(pool.slice_size(), 16);
        let ret = pool.returner();

        // Claim both slots; pool is now empty.
        let mut a = pool.claim().unwrap();
        let _b = pool.claim().unwrap();

        // Pack two small "files" into slot a.
        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;

        // Releasing the first slice must NOT free the slot yet.
        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"
        );

        // Releasing the last slice frees it.
        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"
        );
    }

    /// MEMORY BOUND: the pool reservation follows the input and the budget, and
    /// the cut length — the only output-affecting quantity — never moves.
    ///
    /// The old code reserved `8 × 200 MiB = 1.6 GiB` before opening a file,
    /// whatever the input size and whatever the memory limit; that is what
    /// OOM-killed a constrained backup pod. Inject real byte counts, assert real
    /// reservations.
    #[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;

        // 1. UNCONSTRAINED — big input, plentiful RAM — is EXACTLY the old
        //    geometry. The fat/bench path must not move.
        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");

        // 2. TINY INPUT, plentiful RAM: a 5 MiB staging tree no longer reserves
        //    1.6 GiB — the reservation is bounded by the bytes there are to read.
        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()
        );

        // 3. CONSTRAINED POD: big input, 64 MiB budget → the reservation obeys the
        //    budget instead of OOM-killing the pod.
        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()
        );

        // 4. ABSURD BUDGET: still makes progress — one slot holding one slice.
        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");

        // 5. THE INVARIANT: slice_size is identical in every one of them, so all
        //    four write byte-identical archives — only the RAM footprint differs.
        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);
        }

        // 6. SWEEP — the budget is honoured at EVERY size, not just the round
        //    ones. (Rounding the per-slot slice count UP instead of DOWN broke
        //    exactly here: a 64 MiB budget planned 110 MiB, because a 1-slice
        //    shortfall got multiplied by 8 slots.)
        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");
                // The floor is one slice: below that the pipeline cannot run at
                // all, so a budget under one slice buys exactly one slice.
                let allowed = budget.max(ss).min(default_bytes);
                assert!(
                    p.bytes() <= allowed,
                    "workers={workers} budget={budget}: planned {} bytes, allowed {allowed}",
                    p.bytes()
                );
            }
            // And the input bound holds independently of the budget.
            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()
                );
            }
        }
    }

    /// A shrunk pool is a WORKING pool: pack and drain slices through the
    /// smallest plan `plan` can return and prove the lifecycle still holds.
    #[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 {
        /// Non-blocking claim, for tests only.
        fn claim_now(&self) -> Option<u32> {
            self.free_rx.try_recv().ok()
        }
    }
}