znippy_common/

poolplan.rs

//! `PoolPlan` — znippy's **cut policy** for a compress pass.
//!
//! ## Why this is znippy's and not the pool's
//!
//! The slot pool itself is `znippy_zoomies::chunk_revolver::ChunkRevolver` in
//! coalescing mode — one canonical implementation, upstream, shared with every
//! other consumer. What lives here is the part that is **not** pool mechanics:
//! how big a slice znippy cuts a file into, and how much buffer it is willing to
//! reserve to do it.
//!
//! `slice_size` is **output-affecting**. It is the big/small partition threshold
//! in `compress_dir` AND the length a big file is cut into, so it lands in
//! `chunk_seq`, in `fdata_offset`, and in 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, and the sink writes in producer order either
//! way.
//!
//! That split is the whole point of planning them separately: the pool's memory
//! footprint can be shrunk to fit a cgroup limit while the bytes on disk stay
//! **identical**. A pool crate cannot own that decision, because it does not
//! know what an archive is.
//!
//! Until 2026-08-25 this module also carried `Magazine`/`Clip`/`Round`/`Ejector`
//! — a fork of pool mechanics that now exist canonically upstream as
//! `ChunkRevolver::with_coalescing` / `Filler` / `PackedSlice` / `Returner`.
//! The fork was deleted rather than kept: measured over 133 650 real files
//! (985.7 MiB), the canonical implementation beat the replica in 5 of 5
//! interleaved runs by 5-15 %, with 7.9x fewer allocations.

/// 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 }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    /// 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()
                );
            }
        }
    }
}