nordisk/znippy-zoomies

gatling/src/window_plan.rs

//! **The read-window plan** — one shared rule for how much compressed input a
//! parallel decoder reads at a time, from the input size `S` (a file's or an
//! OCI layer's size; unknown on a pipe) and the cores `N`.
//!
//! ```text
//! segment     = clamp(S / (N × k), min_seg, max_seg)   k pieces per core
//! read_window = N × segment                            one round of work for all cores
//! ramp-up     = the first window is N × min_seg, then each one doubles up to read_window
//! small input = S ≤ 2 × N × min_seg: read it all at once, no ramp
//! unknown S   = ramp from N × min_seg, doubling up to N × max_seg
//! ```
//!
//! The ramp gets the first output out fast; the steady window keeps every core
//! fed while the next window is read (two in flight). `min_seg` and `max_seg`
//! are per format: a format with small natural blocks (bzip2, ≤ 900 KB) takes
//! a lower `min_seg`.

/// Per-format sizes.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Plan {
    /// Smallest segment worth a core, bytes.
    pub min_seg: u64,
    /// Largest segment, bytes.
    pub max_seg: u64,
    /// Pieces per core over the whole input (load balance).
    pub k: u64,
}

impl Plan {
    /// The steady segment size for input size `s` on `n` cores.
    pub fn segment(&self, s: Option<u64>, n: usize) -> u64 {
        let n = n.max(1) as u64;
        match s {
            Some(s) => (s / (n * self.k.max(1))).clamp(self.min_seg, self.max_seg),
            None => self.max_seg,
        }
    }

    /// The steady read window for input size `s` on `n` cores.
    pub fn read_window(&self, s: Option<u64>, n: usize) -> u64 {
        n.max(1) as u64 * self.segment(s, n)
    }
}

/// The read windows, in order: their sizes add up to `s` when it is known.
pub fn windows(s: Option<u64>, n: usize, plan: Plan) -> Windows {
    let n64 = n.max(1) as u64;
    let small = s.is_some_and(|s| s <= 2 * n64 * plan.min_seg);
    Windows {
        left: s,
        next: if small {
            s.unwrap_or(0)
        } else {
            n64 * plan.min_seg
        },
        steady: plan.read_window(s, n),
        small,
    }
}

/// Iterator of read-window sizes (see [`windows`]).
#[derive(Debug, Clone)]
pub struct Windows {
    left: Option<u64>,
    next: u64,
    steady: u64,
    small: bool,
}

impl Iterator for Windows {
    type Item = u64;
    fn next(&mut self) -> Option<u64> {
        let mut w = self.next.max(1);
        if let Some(left) = self.left {
            if left == 0 {
                return None;
            }
            w = w.min(left);
            self.left = Some(left - w);
        }
        if !self.small {
            self.next = (self.next * 2)
                .min(self.steady)
                .max(self.next.min(self.steady));
        }
        Some(w)
    }
}