mikronova/openzl-sys-rs

src/budget.rs

//! **A per-thread allocation budget for OpenZL.**
//!
//! OpenZL sizes its internal stream buffers from the frame's graph, not from the
//! decompressed size in the frame header. A 36-byte frame can make it allocate
//! and fill gigabytes before it fails (the heimdall fuzz find, lane zlwasm), and
//! OpenZL 0.3.0 has no allocator parameter on its public API.
//!
//! [`with_alloc_budget`] runs a closure with a budget on THIS thread: every
//! allocation OpenZL makes through its allocation layer (`ZL_malloc`,
//! `ZL_calloc`, `ZL_realloc` and the arenas' raw allocator, all in
//! `src/openzl/common/allocation.c`) is charged against it, and one that does not
//! fit returns NULL, which OpenZL reports as an allocation error. A thread with no
//! budget pays one thread-local read per allocation and is otherwise unchanged.
//!
//! The charge is cumulative: a free gives nothing back and a `realloc` is charged
//! its full new size. That over-counts an honest decode (a buffer that grows by
//! doubling is charged about twice its final size), which is the safe direction;
//! size the limit from measurements.
//!
//! Not charged: allocations OpenZL makes outside its allocation layer (Huffman
//! tables, error strings, the zstd decode context), which the format bounds.

/// What a budgeted run spent.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct AllocBudget {
    /// The limit the closure ran under, in bytes.
    pub limit: usize,
    /// Bytes charged against it (cumulative, see the module docs).
    pub used: usize,
    /// True when at least one allocation was refused for not fitting.
    pub exceeded: bool,
}

#[repr(C)]
#[derive(Clone, Copy)]
struct RawBudget {
    active: core::ffi::c_int,
    exceeded: core::ffi::c_int,
    remaining: usize,
}

extern "C" {
    fn zlsys_budget_get(out: *mut RawBudget);
    fn zlsys_budget_set(new: *const RawBudget);
}

fn get() -> RawBudget {
    let mut b = RawBudget { active: 0, exceeded: 0, remaining: 0 };
    // SAFETY: writes one plain struct this thread owns.
    unsafe { zlsys_budget_get(&mut b) };
    b
}

fn set(b: &RawBudget) {
    // SAFETY: reads one plain struct; the C side copies it into a thread-local.
    unsafe { zlsys_budget_set(b) };
}

/// Puts the outer budget back when the closure returns or unwinds, and reads
/// what the inner one spent.
struct Restore {
    outer: RawBudget,
    limit: usize,
    report: Option<AllocBudget>,
}

impl Restore {
    fn finish(&mut self) -> AllocBudget {
        if let Some(r) = self.report {
            return r;
        }
        let inner = get();
        set(&self.outer);
        let r = AllocBudget {
            limit: self.limit,
            used: self.limit - inner.remaining,
            exceeded: inner.exceeded != 0,
        };
        self.report = Some(r);
        r
    }
}

impl Drop for Restore {
    fn drop(&mut self) {
        self.finish();
    }
}

/// Whether OpenZL's allocation layer is routed through the budget in this
/// build. False only when the crate links pre-built libraries
/// (`OPENZL_SYS_LIB_DIR`), whose `allocation.c` is upstream's.
pub fn alloc_budget_enforced() -> bool {
    !cfg!(zlsys_budget_not_enforced)
}

/// Run `f` with an OpenZL allocation budget of `limit` bytes on this thread.
///
/// Budgets nest: an inner call runs under its own limit and the outer budget is
/// restored afterwards (the inner run is not charged to it). The budget is
/// restored on unwind as well.
pub fn with_alloc_budget<R>(limit: usize, f: impl FnOnce() -> R) -> (R, AllocBudget) {
    let mut restore = Restore { outer: get(), limit, report: None };
    set(&RawBudget { active: 1, exceeded: 0, remaining: limit });
    let r = f();
    let spent = restore.finish();
    (r, spent)
}