scalo 2.12.2

// Project:   scalo
// File:      src/memory/guard.rs
// Purpose:   Memory guard with backpressure signals
// Language:  Rust
//
// License:   Apache-2.0
// Copyright: (c) 2026 HYPERI PTY LIMITED

//! Memory guard with backpressure signals.

use std::sync::OnceLock;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};

use super::cgroup;
use super::usage::{UsageReader, UsageSource};

/// Process-wide total-heap byte source, set once at startup.
///
/// See [`set_heap_source`] for the rationale. Allocator-agnostic: any
/// `fn() -> usize` returning live heap bytes (e.g. `cap::Cap::allocated`,
/// jemalloc `stats.allocated`).
static HEAP_SOURCE: OnceLock<fn() -> usize> = OnceLock::new();

/// Register a process-wide source of total live-heap bytes, overriding the
/// [`UsageSource`] the guard would otherwise detect.
///
/// This is rarely needed. The default source is the cgroup's `memory.current`,
/// which is both allocator-independent and the number the OOM killer acts on;
/// an allocator statistic is narrower, since it cannot see thread stacks,
/// mmap'd buffers, or pages the allocator retains after a free. Register one
/// only where the allocator figure is the one you want to gate on.
///
/// When set, every [`MemoryGuard`] switches its read path
/// ([`current_bytes`](MemoryGuard::current_bytes), pressure checks, and
/// [`try_reserve`](MemoryGuard::try_reserve) admission) to this source.
///
/// **Why a global hook and not a dependency:** a tracking allocator must be
/// the binary's single `#[global_allocator]`, which is the *application's*
/// choice, not a library's -- and scalo is `#![forbid(unsafe_code)]`, so it
/// cannot implement one anyway. The application installs its allocator and
/// wires it here in a few lines. This keeps scalo allocator-agnostic with no
/// allocator dependency in its graph.
///
/// The first call wins and returns `true`; later calls are a no-op and return
/// `false` (the existing source is kept). Call once at startup, before
/// constructing guards.
///
/// The application picks a tracking allocator -- prefer an actively-maintained
/// one such as `tikv-jemalloc-ctl` (`stats.allocated`); the `cap` crate also
/// works but is effectively unmaintained (last release 2023).
///
/// ```ignore
/// // In the application binary, using jemalloc:
/// #[global_allocator]
/// static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
///
/// fn main() {
///     scalo::memory::set_heap_source(|| {
///         tikv_jemalloc_ctl::epoch::advance().ok();
///         tikv_jemalloc_ctl::stats::allocated::read().unwrap_or(0)
///     });
///     // ... build ServiceRuntime / MemoryGuard ...
/// }
/// ```
#[must_use]
pub fn set_heap_source(source: fn() -> usize) -> bool {
    HEAP_SOURCE.set(source).is_ok()
}

/// Read the registered total-heap source, if any.
#[inline]
fn heap_bytes() -> Option<u64> {
    HEAP_SOURCE.get().map(|f| f() as u64)
}

/// Read an env var `{PREFIX}_{SUFFIX}` and parse it.
fn env_parsed<T: std::str::FromStr>(prefix: &str, suffix: &str) -> Option<T> {
    std::env::var(format!("{prefix}_{suffix}"))
        .ok()
        .and_then(|v| v.parse().ok())
}

/// Memory pressure levels.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MemoryPressure {
    /// Usage below 50% of limit.
    Low,
    /// Usage between 50% and pressure_threshold.
    Medium,
    /// Usage above pressure_threshold -- apply backpressure.
    High,
}

/// Configuration for `MemoryGuard`.
///
/// When the `config` feature is enabled, this can be loaded from the config
/// cascade under the `memory` key:
///
/// ```yaml
/// memory:
///   limit_bytes: 0           # 0 = auto-detect from cgroup/system
///   pressure_threshold: 0.80 # backpressure at 80% of effective limit
///   cgroup_headroom: 0.85    # use 85% of cgroup limit
/// ```
#[derive(Debug, Clone, serde::Deserialize, serde::Serialize)]
#[cfg_attr(feature = "config-schema", derive(schemars::JsonSchema))]
pub struct MemoryGuardConfig {
    /// Explicit memory limit in bytes. 0 = auto-detect from cgroup/system.
    #[serde(default)]
    pub limit_bytes: u64,
    /// Fraction of limit at which backpressure activates (default 0.8).
    #[serde(default = "default_pressure_threshold")]
    pub pressure_threshold: f64,
    /// Fraction of cgroup limit to use as the effective limit (default 0.85).
    /// Leaves headroom for the process itself (stack, code, etc.).
    #[serde(default = "default_cgroup_headroom")]
    pub cgroup_headroom: f64,
}

fn default_pressure_threshold() -> f64 {
    DEFAULT_PRESSURE_THRESHOLD
}

fn default_cgroup_headroom() -> f64 {
    DEFAULT_CGROUP_HEADROOM
}

/// A fraction is valid iff it is finite and within `(0.0, 1.0]`.
fn check_fraction(v: f64, name: &str) -> Result<(), String> {
    if !v.is_finite() || v <= 0.0 || v > 1.0 {
        return Err(format!(
            "memory.{name} must be a finite fraction in (0.0, 1.0], got {v}"
        ));
    }
    Ok(())
}

/// Return `v` if it is a valid fraction, else log an error and substitute
/// `default`. Defensive guard so a bad config cannot produce a zero/`NaN`
/// limit and a divide-by-zero pressure ratio.
fn sane_fraction(v: f64, default: f64, name: &str) -> f64 {
    if check_fraction(v, name).is_err() {
        tracing::error!(
            value = v,
            "invalid memory.{name} (need finite fraction in (0,1]); using default {default}"
        );
        default
    } else {
        v
    }
}

/// Effective auto-detected limit: cgroup limit * headroom, capped at the soft
/// throttle (`memory.high`) when that is set lower.
///
/// The kernel reclaims hard and throttles allocations at `memory.high` (before
/// the `memory.max` OOM-kill), so admitting past it courts a latency cliff.
/// Pure, so the cap logic is unit-testable without touching the real cgroup
/// files.
#[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
fn effective_auto_limit(detected: u64, headroom: f64, high: Option<u64>) -> u64 {
    let headroom_limit = (detected as f64 * headroom) as u64;
    match high {
        Some(h) => headroom_limit.min(h),
        None => headroom_limit,
    }
}

/// Default cgroup headroom: use 85% of cgroup limit.
///
/// Rationale: Rust has no GC so no spike headroom needed (unlike JVM 75% / Go 80%).
/// 15% headroom covers jemalloc fragmentation, kernel overhead, and page cache.
const DEFAULT_CGROUP_HEADROOM: f64 = 0.85;

/// Default pressure threshold: backpressure at 80% of effective limit.
///
/// With 85% headroom, backpressure activates at ~68% of actual cgroup limit.
/// Matches OTel Collector's `limit_percentage: 80` philosophy.
const DEFAULT_PRESSURE_THRESHOLD: f64 = 0.80;

impl Default for MemoryGuardConfig {
    fn default() -> Self {
        Self {
            limit_bytes: 0, // auto-detect
            pressure_threshold: DEFAULT_PRESSURE_THRESHOLD,
            cgroup_headroom: DEFAULT_CGROUP_HEADROOM,
        }
    }
}

impl MemoryGuardConfig {
    /// Load from the config cascade, falling back to defaults.
    ///
    /// When the `config` feature is enabled and `config::setup()` has been
    /// called, reads the `memory` key from the cascade. Otherwise returns
    /// [`MemoryGuardConfig::default()`].
    #[must_use]
    pub fn from_cascade() -> Self {
        #[cfg(feature = "config")]
        {
            if let Some(cfg) = crate::config::try_get()
                && let Ok(memory) = cfg.unmarshal_key_registered::<Self>("memory")
            {
                return memory;
            }
        }
        Self::default()
    }

    /// Create config from environment variables with a prefix.
    ///
    /// Reads standard env vars for memory configuration:
    /// - `{PREFIX}_MEMORY_LIMIT_BYTES` -- explicit limit (0 or unset = auto-detect from cgroup)
    /// - `{PREFIX}_MEMORY_PRESSURE_THRESHOLD` -- backpressure trigger (default 0.80)
    /// - `{PREFIX}_MEMORY_CGROUP_HEADROOM` -- fraction of cgroup limit to use (default 0.85)
    ///
    /// # Example
    ///
    /// ```bash
    /// DFE_MEMORY_LIMIT_BYTES=4294967296      # 4 GiB explicit
    /// DFE_MEMORY_PRESSURE_THRESHOLD=0.75     # backpressure at 75%
    /// DFE_MEMORY_CGROUP_HEADROOM=0.90        # use 90% of cgroup
    /// ```
    ///
    /// ```rust,no_run
    /// use scalo::memory::MemoryGuardConfig;
    /// let config = MemoryGuardConfig::from_env("DFE");
    /// ```
    #[must_use]
    #[cfg(feature = "config")]
    pub fn from_env(prefix: &str) -> Self {
        use crate::config::flat_env::flat_env_parsed;

        let mut config = Self::default();

        if let Some(v) = flat_env_parsed::<u64>(prefix, "MEMORY_LIMIT_BYTES") {
            config.limit_bytes = v;
        }
        if let Some(v) = flat_env_parsed::<f64>(prefix, "MEMORY_PRESSURE_THRESHOLD") {
            config.pressure_threshold = v;
        }
        if let Some(v) = flat_env_parsed::<f64>(prefix, "MEMORY_CGROUP_HEADROOM") {
            config.cgroup_headroom = v;
        }

        config
    }

    /// Create config from environment variables without requiring `config` feature.
    ///
    /// Same as [`from_env`](Self::from_env) but uses `std::env` directly.
    #[must_use]
    pub fn from_env_raw(prefix: &str) -> Self {
        let mut config = Self::default();

        if let Some(v) = env_parsed::<u64>(prefix, "MEMORY_LIMIT_BYTES") {
            config.limit_bytes = v;
        }
        if let Some(v) = env_parsed::<f64>(prefix, "MEMORY_PRESSURE_THRESHOLD") {
            config.pressure_threshold = v;
        }
        if let Some(v) = env_parsed::<f64>(prefix, "MEMORY_CGROUP_HEADROOM") {
            config.cgroup_headroom = v;
        }

        config
    }

    /// Validate the config, returning an error describing the first invalid
    /// field. `pressure_threshold` and `cgroup_headroom` must each be a finite
    /// fraction in `(0.0, 1.0]`. Call this at startup to fail fast on bad
    /// config rather than relying on [`MemoryGuard::new`]'s defensive clamping.
    ///
    /// # Errors
    ///
    /// Returns `Err` with a human-readable message if a fraction field is
    /// non-finite, `<= 0.0`, or `> 1.0`.
    pub fn validate(&self) -> Result<(), String> {
        check_fraction(self.pressure_threshold, "pressure_threshold")?;
        check_fraction(self.cgroup_headroom, "cgroup_headroom")?;
        Ok(())
    }
}

/// Cgroup-aware memory tracking with backpressure signals.
///
/// Reads what the kernel charges this process -- cgroup v2 `memory.current` by
/// default, see [`UsageSource`] -- and compares it against the cgroup limit,
/// so the signal is the one the OOM killer acts on and does not depend on
/// which allocator the binary installed. Designed for data pipeline services
/// where incoming data must be rejected (503) before hitting the container
/// memory limit.
///
/// # Usage
///
/// ```rust,no_run
/// use scalo::memory::{MemoryGuard, MemoryGuardConfig};
///
/// let guard = MemoryGuard::new(MemoryGuardConfig::default());
///
/// // On data arrival -- check before accepting
/// let payload_len = 1024u64;
/// if !guard.try_reserve(payload_len) {
///     // return 503 -- backpressure
/// }
///
/// // After data is flushed/sent
/// guard.release(payload_len);
///
/// // Fast hot-path check
/// if guard.under_pressure() {
///     // return 503
/// }
/// ```
pub struct MemoryGuard {
    /// Outstanding per-batch reservations, from `try_reserve`/`add_bytes` less
    /// `release`. Lease accounting, not the process's memory usage.
    reserved_bytes: AtomicU64,
    /// Where the process's real usage is read from.
    usage: UsageReader,
    /// Effective memory limit in bytes.
    limit_bytes: u64,
    /// Pressure threshold (0.0-1.0).
    pressure_threshold: f64,
    /// Fast hot-path pressure flag. Only consulted when no usage source is
    /// readable, where the reservation counter is all there is.
    under_pressure: AtomicBool,
}

impl MemoryGuard {
    /// Create a new memory guard, detecting where to read usage from.
    ///
    /// If `config.limit_bytes` is 0, auto-detects from cgroup (K8s) or system memory,
    /// then applies `cgroup_headroom` factor to leave room for process overhead.
    #[must_use]
    pub fn new(config: MemoryGuardConfig) -> Self {
        Self::with_usage_source(config, UsageSource::detect())
    }

    /// Create a memory guard reading usage from `source` instead of the
    /// detected one. For a test fixture directory, or a host where detection
    /// picks the wrong cgroup.
    #[must_use]
    #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
    pub fn with_usage_source(config: MemoryGuardConfig, source: UsageSource) -> Self {
        // Defensive: a non-finite / out-of-range threshold or headroom would
        // produce a zero/NaN limit and a divide-by-zero pressure ratio. Clamp
        // to the safe default and log loudly. Callers wanting hard rejection
        // should call `config.validate()` at startup.
        let pressure_threshold = sane_fraction(
            config.pressure_threshold,
            DEFAULT_PRESSURE_THRESHOLD,
            "pressure_threshold",
        );
        let cgroup_headroom = sane_fraction(
            config.cgroup_headroom,
            DEFAULT_CGROUP_HEADROOM,
            "cgroup_headroom",
        );

        let raw_limit = if config.limit_bytes > 0 {
            config.limit_bytes
        } else {
            effective_auto_limit(
                cgroup::detect_memory_limit(),
                cgroup_headroom,
                cgroup::detect_memory_high(),
            )
        };
        // Never permit a zero effective limit: every pressure calculation
        // divides by it.
        let limit_bytes = raw_limit.max(1);

        let usage = UsageReader::new(source);
        let usage_source = if HEAP_SOURCE.get().is_some() {
            "explicit"
        } else {
            usage.source().name()
        };
        let usage_bytes = heap_bytes().or_else(|| usage.read());
        tracing::info!(
            limit_bytes,
            pressure_threshold,
            usage_source,
            ?usage_bytes,
            "memory guard initialised"
        );
        if usage_bytes.is_none() {
            tracing::warn!(
                "no cgroup or procfs memory accounting readable: the guard sees \
                 only bytes callers reserve and release, not the process's usage"
            );
        }

        Self {
            reserved_bytes: AtomicU64::new(0),
            usage,
            limit_bytes,
            pressure_threshold,
            under_pressure: AtomicBool::new(false),
        }
    }

    /// Try to reserve bytes. Returns false if over the limit (backpressure).
    ///
    /// This is a projected-admission check against the process's real usage
    /// (`usage() + bytes <= limit`) and does NOT mutate the reservation
    /// counter -- the kernel charges the bytes once they are allocated and
    /// uncharges them on drop, so no `release` is needed to keep the check
    /// honest. An admission is charged to the usage reader's ledger so the
    /// callers behind it in the same cache window see it. Only where no usage
    /// source is readable does it fall back to the classic atomic
    /// check-and-add on the per-batch counter (rolled back if it would exceed
    /// the limit).
    #[inline]
    pub fn try_reserve(&self, bytes: u64) -> bool {
        if let Some(heap) = heap_bytes() {
            // A registered heap source is read live, so it needs no ledger.
            return heap.saturating_add(bytes) <= self.limit_bytes;
        }
        if let Some(estimate) = self.usage.estimate() {
            // Saturating, so an absurd `bytes` is refused rather than wrapping
            // past the limit and being admitted.
            if estimate.saturating_add(bytes) > self.limit_bytes {
                return false;
            }
            self.usage.admit(bytes);
            return true;
        }
        let current = self.reserved_bytes.fetch_add(bytes, Ordering::Relaxed) + bytes;
        if current > self.limit_bytes {
            // Over limit -- roll back
            self.reserved_bytes.fetch_sub(bytes, Ordering::Relaxed);
            self.under_pressure.store(true, Ordering::Relaxed);
            return false;
        }
        self.update_pressure(current);
        true
    }

    /// Add bytes to the reservation counter without checking the limit.
    /// Use when data is already accepted and you just need to track the lease.
    #[inline]
    pub fn add_bytes(&self, bytes: u64) {
        // Charged to the ledger as well, so a lease paired with `release` nets
        // to zero there instead of eroding another caller's admissions.
        self.usage.admit(bytes);
        let new_total = self.reserved_bytes.fetch_add(bytes, Ordering::Relaxed) + bytes;
        self.update_pressure(new_total);
    }

    /// Release bytes after data is flushed/sent/dropped.
    ///
    /// Uses saturating subtraction to prevent underflow wrapping.
    #[inline]
    pub fn release(&self, bytes: u64) {
        self.usage.forget(bytes);
        let prev = self
            .reserved_bytes
            .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
                Some(current.saturating_sub(bytes))
            })
            // Always succeeds (closure always returns Some).
            .unwrap_or_else(|v| v);
        self.update_pressure(prev.saturating_sub(bytes));
    }

    /// Fast hot-path pressure check.
    ///
    /// Computes live from the usage source, whose file read is cached for a
    /// short interval and topped up with the bytes admitted since. Only where
    /// no source is readable does it read the cached flag maintained by
    /// `try_reserve`/`add_bytes`/`release`.
    #[inline]
    pub fn under_pressure(&self) -> bool {
        if self.usage_bytes().is_some() {
            return self.pressure_ratio() >= self.pressure_threshold;
        }
        self.under_pressure.load(Ordering::Relaxed)
    }

    /// Current pressure level.
    #[inline]
    pub fn pressure(&self) -> MemoryPressure {
        let ratio = self.pressure_ratio();
        if ratio >= self.pressure_threshold {
            MemoryPressure::High
        } else if ratio >= 0.5 {
            MemoryPressure::Medium
        } else {
            MemoryPressure::Low
        }
    }

    /// Current usage as fraction of limit (0.0 - 1.0+).
    #[inline]
    pub fn pressure_ratio(&self) -> f64 {
        self.current_bytes() as f64 / self.limit_bytes as f64
    }

    /// Current memory usage in bytes: what the kernel charges this process,
    /// read from the guard's [`UsageSource`], plus the bytes admitted since
    /// that reading was taken.
    ///
    /// A registered [`set_heap_source`] overrides it. Where neither is
    /// readable this falls back to [`reserved_bytes`](Self::reserved_bytes),
    /// which sees only what callers reserved by hand.
    #[inline]
    pub fn current_bytes(&self) -> u64 {
        self.usage_bytes()
            .unwrap_or_else(|| self.reserved_bytes.load(Ordering::Relaxed))
    }

    /// Outstanding per-batch reservations: `try_reserve`/`add_bytes` less
    /// `release`. This is lease accounting -- how many bytes callers say they
    /// are holding -- not the process's memory usage, which is
    /// [`current_bytes`](Self::current_bytes).
    #[inline]
    pub fn reserved_bytes(&self) -> u64 {
        self.reserved_bytes.load(Ordering::Relaxed)
    }

    /// The name of the usage source in force, for logs and diagnostics.
    #[must_use]
    pub fn usage_source(&self) -> &'static str {
        if HEAP_SOURCE.get().is_some() {
            "explicit"
        } else {
            self.usage.source().name()
        }
    }

    /// Process usage from the explicit heap source, else the usage source plus
    /// its ledger of bytes admitted since the sample. `None` only where
    /// neither is readable.
    #[inline]
    fn usage_bytes(&self) -> Option<u64> {
        heap_bytes().or_else(|| self.usage.estimate())
    }

    /// Configured memory limit in bytes.
    #[inline]
    pub fn limit_bytes(&self) -> u64 {
        self.limit_bytes
    }

    /// Update the pressure flag based on current usage.
    #[inline]
    fn update_pressure(&self, current: u64) {
        let ratio = current as f64 / self.limit_bytes as f64;
        self.under_pressure
            .store(ratio >= self.pressure_threshold, Ordering::Relaxed);
    }
}

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

    /// A guard pinned to the reservation counter: the byte-lease ladder these
    /// tests exercise, with no host cgroup reading underneath it.
    fn reservation_guard(config: MemoryGuardConfig) -> MemoryGuard {
        MemoryGuard::with_usage_source(config, UsageSource::Reservations)
    }

    /// Write `name`->`contents` files into a fresh temp dir and return it.
    fn cgroup_fixture(files: &[(&str, &str)]) -> tempfile::TempDir {
        let dir = tempfile::tempdir().expect("tempdir");
        for (name, contents) in files {
            std::fs::write(dir.path().join(name), contents).expect("write fixture");
        }
        dir
    }

    #[test]
    fn test_memory_guard_default() {
        let guard = reservation_guard(MemoryGuardConfig {
            limit_bytes: 1_000_000, // 1MB explicit
            ..Default::default()
        });
        assert_eq!(guard.limit_bytes(), 1_000_000);
        assert_eq!(guard.current_bytes(), 0);
        assert!(!guard.under_pressure());
        assert_eq!(guard.pressure(), MemoryPressure::Low);
    }

    #[test]
    fn test_try_reserve_within_limit() {
        let guard = reservation_guard(MemoryGuardConfig {
            limit_bytes: 1000,
            ..Default::default()
        });
        assert!(guard.try_reserve(500));
        assert_eq!(guard.reserved_bytes(), 500);
    }

    #[test]
    fn test_try_reserve_over_limit() {
        let guard = reservation_guard(MemoryGuardConfig {
            limit_bytes: 1000,
            ..Default::default()
        });
        assert!(guard.try_reserve(500));
        assert!(!guard.try_reserve(600)); // would exceed 1000
        assert_eq!(guard.reserved_bytes(), 500); // rolled back
        assert!(guard.under_pressure());
    }

    #[test]
    fn test_release_reduces_pressure() {
        let guard = reservation_guard(MemoryGuardConfig {
            limit_bytes: 1000,
            pressure_threshold: 0.8,
            ..Default::default()
        });
        guard.add_bytes(900); // 90% -- over threshold
        assert!(guard.under_pressure());
        assert_eq!(guard.pressure(), MemoryPressure::High);

        guard.release(500); // down to 400 = 40%
        assert!(!guard.under_pressure());
        assert_eq!(guard.pressure(), MemoryPressure::Low);
    }

    #[test]
    fn test_pressure_levels() {
        let guard = reservation_guard(MemoryGuardConfig {
            limit_bytes: 1000,
            pressure_threshold: 0.8,
            ..Default::default()
        });

        // Low (< 50%)
        guard.add_bytes(400);
        assert_eq!(guard.pressure(), MemoryPressure::Low);

        // Medium (50-80%)
        guard.add_bytes(200); // 600 = 60%
        assert_eq!(guard.pressure(), MemoryPressure::Medium);

        // High (>= 80%)
        guard.add_bytes(300); // 900 = 90%
        assert_eq!(guard.pressure(), MemoryPressure::High);
    }

    #[test]
    fn test_pressure_ratio() {
        let guard = reservation_guard(MemoryGuardConfig {
            limit_bytes: 1000,
            ..Default::default()
        });
        guard.add_bytes(250);
        let ratio = guard.pressure_ratio();
        assert!((ratio - 0.25).abs() < 0.001);
    }

    #[test]
    fn test_release_saturating() {
        let guard = reservation_guard(MemoryGuardConfig {
            limit_bytes: 1000,
            ..Default::default()
        });
        guard.add_bytes(100);
        guard.release(200); // release more than added -- saturates to 0
        assert_eq!(
            guard.reserved_bytes(),
            0,
            "over-release must saturate to 0, not wrap"
        );
        assert!(!guard.under_pressure());
        assert_eq!(guard.pressure(), MemoryPressure::Low);

        // Verify the guard is still functional after over-release
        assert!(guard.try_reserve(500));
        assert_eq!(guard.reserved_bytes(), 500);
    }

    #[test]
    fn test_concurrent_reserve_release() {
        use std::sync::Arc;
        use std::thread;

        let guard = Arc::new(reservation_guard(MemoryGuardConfig {
            limit_bytes: 100_000,
            pressure_threshold: 0.8,
            ..Default::default()
        }));

        let mut handles = vec![];
        for _ in 0..10 {
            let g = Arc::clone(&guard);
            handles.push(thread::spawn(move || {
                for _ in 0..100 {
                    g.add_bytes(100);
                    g.release(100);
                }
            }));
        }
        for h in handles {
            h.join().unwrap();
        }
        // All bytes should be released -- may not be exactly 0 due to ordering
        // but should be close (within one thread's batch)
        assert!(
            guard.reserved_bytes() < 1000,
            "leaked bytes: {}",
            guard.reserved_bytes()
        );
    }

    #[test]
    fn test_try_reserve_rollback_is_atomic() {
        let guard = reservation_guard(MemoryGuardConfig {
            limit_bytes: 100,
            ..Default::default()
        });
        assert!(guard.try_reserve(90));
        assert!(!guard.try_reserve(20)); // over limit, rolled back
        assert_eq!(guard.reserved_bytes(), 90); // not 110
        assert!(guard.try_reserve(10)); // exactly at limit
        assert_eq!(guard.reserved_bytes(), 100);
    }

    #[test]
    fn pressure_ratio_follows_the_cgroup_current_file() {
        // The bug this guards: with jemalloc installed and no heap source
        // registered, the guard read a reservation counter that never moved
        // while the process held 189 MiB. It now reads what the kernel charges.
        let limit = 536_870_912u64; // matches memory.max in the fixture
        let dir = cgroup_fixture(&[("memory.current", "0\n"), ("memory.max", "536870912\n")]);
        let guard = MemoryGuard::with_usage_source(
            MemoryGuardConfig {
                limit_bytes: limit,
                pressure_threshold: 0.8,
                ..Default::default()
            },
            UsageSource::CgroupV2(dir.path().to_path_buf()),
        );

        // 198299648 bytes = the 189 MiB the transform actually held.
        std::fs::write(dir.path().join("memory.current"), "198299648\n").expect("write");
        std::thread::sleep(std::time::Duration::from_millis(60)); // past the read cache
        let ratio = guard.pressure_ratio();
        assert!(
            (ratio - 0.369).abs() < 0.01,
            "189 MiB of 512 MiB is 0.369, got {ratio}"
        );
        assert_eq!(guard.current_bytes(), 198_299_648);
        assert_eq!(guard.pressure(), MemoryPressure::Low);
        assert!(!guard.under_pressure());

        // Push past the threshold: 450 MiB of 512 MiB is 0.879.
        std::fs::write(dir.path().join("memory.current"), "471859200\n").expect("write");
        std::thread::sleep(std::time::Duration::from_millis(60));
        assert!(
            guard.under_pressure(),
            "88% of the cgroup limit is over the 80% threshold"
        );
        assert_eq!(guard.pressure(), MemoryPressure::High);
        assert!(
            !guard.try_reserve(64 * 1024 * 1024),
            "admission is projected against the cgroup, so 450 + 64 MiB is refused"
        );
    }

    /// Ten MiB limit against a cgroup file that stays at zero, so admission can
    /// only be refused by the ledger of bytes taken since that reading.
    fn ledger_guard(dir: &tempfile::TempDir) -> MemoryGuard {
        MemoryGuard::with_usage_source(
            MemoryGuardConfig {
                limit_bytes: 10 * 1024 * 1024,
                ..Default::default()
            },
            UsageSource::CgroupV2(dir.path().to_path_buf()),
        )
    }

    #[test]
    fn a_burst_inside_one_cache_window_is_refused() {
        // The overshoot this guards: the kernel figure is up to CACHE_INTERVAL
        // stale, so admitting against it alone lets a burst through unbounded.
        let dir = cgroup_fixture(&[("memory.current", "0\n")]);
        let guard = ledger_guard(&dir);
        let six_mib = 6 * 1024 * 1024;

        assert!(guard.try_reserve(six_mib), "6 MiB of a 10 MiB limit fits");
        assert!(
            !guard.try_reserve(six_mib),
            "a second 6 MiB is 12 MiB of a 10 MiB limit, against an unchanged file"
        );
    }

    #[test]
    fn release_discharges_the_ledger() {
        let dir = cgroup_fixture(&[("memory.current", "0\n")]);
        let guard = ledger_guard(&dir);
        let six_mib = 6 * 1024 * 1024;

        assert!(guard.try_reserve(six_mib));
        guard.release(six_mib);
        assert!(
            guard.try_reserve(six_mib),
            "the released bytes no longer count against the next admission"
        );
    }

    #[test]
    fn a_fresh_sample_clears_the_ledger() {
        let dir = cgroup_fixture(&[("memory.current", "0\n")]);
        let guard = ledger_guard(&dir);
        let six_mib = 6 * 1024 * 1024;

        assert!(guard.try_reserve(six_mib));
        std::thread::sleep(std::time::Duration::from_millis(60)); // past the read cache
        assert!(
            guard.try_reserve(six_mib),
            "the new reading accounts for those bytes, so the ledger restarts"
        );
    }

    #[test]
    fn the_ledger_adds_to_what_the_kernel_already_charges() {
        let dir = cgroup_fixture(&[("memory.current", "4194304\n")]); // 4 MiB
        let guard = ledger_guard(&dir);

        assert_eq!(guard.current_bytes(), 4 * 1024 * 1024);
        assert!(guard.try_reserve(2 * 1024 * 1024));
        assert_eq!(
            guard.current_bytes(),
            6 * 1024 * 1024,
            "4 MiB charged plus 2 MiB admitted against that reading"
        );
    }

    #[test]
    fn reserved_bytes_and_current_bytes_are_different_numbers() {
        let dir = cgroup_fixture(&[("memory.current", "104857600\n")]);
        let guard = MemoryGuard::with_usage_source(
            MemoryGuardConfig {
                limit_bytes: 536_870_912,
                ..Default::default()
            },
            UsageSource::CgroupV2(dir.path().to_path_buf()),
        );

        guard.add_bytes(4096);
        assert_eq!(
            guard.reserved_bytes(),
            4096,
            "the caller's outstanding lease"
        );
        assert_eq!(
            guard.current_bytes(),
            104_861_696,
            "what the kernel charges, plus the lease taken against that reading"
        );
    }

    #[test]
    fn usage_source_is_named_for_the_init_log() {
        let dir = cgroup_fixture(&[("memory.current", "1024\n")]);
        let guard = MemoryGuard::with_usage_source(
            MemoryGuardConfig {
                limit_bytes: 1000,
                ..Default::default()
            },
            UsageSource::CgroupV2(dir.path().to_path_buf()),
        );
        assert_eq!(guard.usage_source(), "cgroup-v2");

        let guard = reservation_guard(MemoryGuardConfig {
            limit_bytes: 1000,
            ..Default::default()
        });
        assert_eq!(guard.usage_source(), "reservations");
    }

    // Process-global heap source for the switch test. nextest isolates each
    // test in its own process, so registering it here is contained to this
    // test and does not leak into the per-batch-counter tests above. (This is
    // the single test in this module that touches the global hook.)
    static TEST_HEAP: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
    fn test_heap_source() -> usize {
        TEST_HEAP.load(Ordering::Relaxed)
    }

    #[test]
    fn heap_source_overrides_read_path_and_admission() {
        assert!(set_heap_source(test_heap_source), "first set wins");
        assert!(
            !set_heap_source(test_heap_source),
            "second set is a no-op (first-wins)"
        );

        let guard = MemoryGuard::new(MemoryGuardConfig {
            limit_bytes: 1_000,
            pressure_threshold: 0.8,
            ..Default::default()
        });

        // Reads come from the heap source, not the reservation counter.
        TEST_HEAP.store(250, Ordering::Relaxed);
        assert_eq!(guard.current_bytes(), 250);
        assert!((guard.pressure_ratio() - 0.25).abs() < 0.001);
        assert!(!guard.under_pressure());

        // Pressure tracks the live heap -- including growth never reserved,
        // which the per-batch counter would have been blind to.
        TEST_HEAP.store(850, Ordering::Relaxed);
        assert!(
            guard.under_pressure(),
            "85% live heap is over the 80% threshold"
        );
        assert_eq!(guard.pressure(), MemoryPressure::High);

        // try_reserve is a projected-admission check against the true heap and
        // does NOT mutate the reservation counter.
        TEST_HEAP.store(900, Ordering::Relaxed);
        assert!(guard.try_reserve(100), "900 + 100 == limit, admitted");
        assert!(!guard.try_reserve(200), "900 + 200 > limit, rejected");
        assert_eq!(guard.current_bytes(), 900, "still the heap source");
        assert_eq!(
            guard.reserved_bytes(),
            0,
            "counter untouched by try_reserve"
        );
        assert_eq!(guard.usage_source(), "explicit");
    }

    #[test]
    fn effective_auto_limit_caps_at_memory_high() {
        // headroom 0.85 of 1000 = 850; high 600 is lower -> cap at the soft
        // throttle so we shed before the kernel does.
        assert_eq!(effective_auto_limit(1000, 0.85, Some(600)), 600);
        // high above the headroom limit -> headroom wins (no change).
        assert_eq!(effective_auto_limit(1000, 0.85, Some(900)), 850);
        // no memory.high in force -> headroom limit.
        assert_eq!(effective_auto_limit(1000, 0.85, None), 850);
    }

    #[test]
    fn test_config_defaults() {
        let config = MemoryGuardConfig::default();
        assert_eq!(config.limit_bytes, 0);
        assert!((config.pressure_threshold - 0.80).abs() < 0.001);
        assert!((config.cgroup_headroom - 0.85).abs() < 0.001);
    }

    #[test]
    fn test_from_env_raw_defaults_when_unset() {
        // With no env vars set, should return defaults
        let config = MemoryGuardConfig::from_env_raw("TEST_MG_UNSET");
        assert_eq!(config.limit_bytes, 0);
        assert!((config.pressure_threshold - 0.80).abs() < 0.001);
        assert!((config.cgroup_headroom - 0.85).abs() < 0.001);
    }

    #[test]
    fn test_env_parsed_helper() {
        // env_parsed returns None for unset vars
        assert!(env_parsed::<u64>("NONEXISTENT_PREFIX_XYZ", "FOO").is_none());
        assert!(env_parsed::<f64>("NONEXISTENT_PREFIX_XYZ", "BAR").is_none());
    }

    #[test]
    fn test_guard_with_explicit_config_overrides() {
        // Simulates what from_env would produce with overrides
        let config = MemoryGuardConfig {
            limit_bytes: 2_147_483_648,
            pressure_threshold: 0.75,
            cgroup_headroom: 0.90,
        };
        let guard = MemoryGuard::new(config);
        assert_eq!(guard.limit_bytes(), 2_147_483_648);
    }

    #[test]
    fn test_guard_with_custom_headroom() {
        // 85% headroom on 1 GiB = 870 MiB effective
        let config = MemoryGuardConfig {
            limit_bytes: 0, // auto-detect
            pressure_threshold: 0.80,
            cgroup_headroom: 0.85,
        };
        let guard = MemoryGuard::new(config);
        // Auto-detected, so limit should be 85% of system/cgroup memory
        assert!(guard.limit_bytes() > 0);
    }

    #[test]
    fn test_validate_accepts_defaults_and_rejects_bad_fractions() {
        assert!(MemoryGuardConfig::default().validate().is_ok());

        for bad in [0.0, -0.1, 1.5, f64::NAN, f64::INFINITY] {
            let cfg = MemoryGuardConfig {
                pressure_threshold: bad,
                ..Default::default()
            };
            assert!(
                cfg.validate().is_err(),
                "pressure_threshold={bad} must be rejected"
            );
            let cfg = MemoryGuardConfig {
                cgroup_headroom: bad,
                ..Default::default()
            };
            assert!(
                cfg.validate().is_err(),
                "cgroup_headroom={bad} must be rejected"
            );
        }
    }

    #[test]
    fn test_new_clamps_invalid_config_no_divide_by_zero() {
        // A zero/NaN headroom with auto-detect could yield a zero limit ->
        // divide-by-zero. A zero pressure_threshold would make every ratio
        // "over". new() must clamp to safe defaults and keep ratios finite.
        let guard = reservation_guard(MemoryGuardConfig {
            limit_bytes: 0,
            pressure_threshold: 0.0,
            cgroup_headroom: 0.0,
        });
        assert!(guard.limit_bytes() >= 1, "limit floored at >=1");
        guard.add_bytes(10);
        assert!(
            guard.pressure_ratio().is_finite(),
            "pressure ratio must be finite, not div-by-zero"
        );
    }

    #[test]
    fn test_new_with_nan_threshold_is_finite() {
        let guard = reservation_guard(MemoryGuardConfig {
            limit_bytes: 1000,
            pressure_threshold: f64::NAN,
            cgroup_headroom: f64::NAN,
        });
        assert_eq!(guard.limit_bytes(), 1000);
        guard.add_bytes(900);
        // Clamped threshold (0.8 default) -> 90% is over -> under pressure.
        assert!(guard.under_pressure());
    }

    #[test]
    fn test_auto_detect_limit() {
        // With limit_bytes = 0, should auto-detect from system
        let guard = MemoryGuard::new(MemoryGuardConfig::default());
        assert!(
            guard.limit_bytes() > 0,
            "auto-detected limit should be positive"
        );
        // Should be less than total system memory (headroom applied)
    }
}