Skip to main content

gam_runtime/
resource.rs

1use crate::cgroup_memory::detect_cgroup_memory;
2pub use crate::cgroup_memory::{
3    CgroupMemoryAvailability, CgroupMemoryLimit, CgroupMemoryObservation, CgroupMemoryProbeFailure,
4    CgroupMemoryProbeFailureKind,
5};
6
7/// The library-default streamed row-chunk target (8 MiB), shared as a `const`
8/// so compile-time consumers (e.g. device tile geometry) stay in lockstep with
9/// [`ResourcePolicy::default_library`] without a runtime policy query.
10///
11/// **The 8 MiB value is unmeasured.** The paragraph above justifies there being
12/// ONE copy; it does not justify the number. No cache level, measured
13/// throughput, or consequence-for-being-wrong is recorded anywhere for it.
14/// Twelve transcribed copies across six crates were collapsed onto this
15/// definition (#2704) purely to remove duplication — that consolidation moved
16/// no value and measured nothing, and a single canonical copy of an arbitrary
17/// number must not be read as a decided one. Anyone retuning it is retuning an
18/// unsupported literal, not overriding a derived budget.
19///
20/// Consumers should size chunks with [`rows_for_target_bytes`], which performs
21/// the `target / (cols * 8)` division and floors the result at one row.
22///
23/// Re-declaring this value — as a module-local `const`, an inline literal, or
24/// any other spelling of 8 MiB — is caught by the workspace guard
25/// `tests/row_chunk_target_bytes_single_source_2704.rs`, which fails on any
26/// transcription of this value outside the exemption table it carries for the
27/// three sites that are a genuinely DIFFERENT quantity coinciding at 8 MiB.
28/// Import this constant instead, or derive from it under a name.
29pub const LIBRARY_ROW_CHUNK_TARGET_BYTES: usize = 8 * 1024 * 1024;
30
31#[derive(Clone, Debug)]
32pub struct ResourcePolicy {
33    pub max_single_materialization_bytes: usize,
34    pub max_operator_cache_bytes: usize,
35    pub max_spatial_distance_cache_bytes: usize,
36    pub max_owned_data_cache_bytes: usize,
37    pub row_chunk_target_bytes: usize,
38    pub derivative_storage_mode: DerivativeStorageMode,
39}
40
41pub const OWNED_DATA_CACHE_MAX_ENTRIES: usize = 2;
42
43// ─────────────────────────────────────────────────────────────────────────────
44// Process-wide memory governor
45// ─────────────────────────────────────────────────────────────────────────────
46
47/// Fraction of this process's memory **capacity** the governor is allowed to
48/// hand out as reservations: 3/4.
49///
50/// The ledger only accounts for the large, planned allocations that route
51/// through [`MemoryGovernor::try_reserve`] (dense design materializations,
52/// covariance blocks, sampler design assemblies). Everything else — allocator
53/// slack, thread stacks, code, and small per-iteration temporaries — lives in
54/// the remaining quarter.
55///
56/// The base quantity is *capacity* (`min(host total, binding cgroup hard
57/// limit)`), not free space, and that is the same separation #2684 established
58/// for the materialization ceiling — see
59/// `governor_budget_from_availability` for why the ledger cannot be an
60/// exception to it (#2702).
61const GOVERNOR_BUDGET_NUMERATOR: u128 = 3;
62const GOVERNOR_BUDGET_DENOMINATOR: u128 = 4;
63
64/// Which observation is the binding ceiling on memory available to this
65/// process.
66#[derive(Clone, Copy, Debug, PartialEq, Eq)]
67pub enum MemoryAvailabilitySource {
68    Host,
69    Cgroup,
70    HostAndCgroup,
71    CgroupProbeFailure,
72}
73
74impl std::fmt::Display for MemoryAvailabilitySource {
75    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
76        match self {
77            Self::Host => formatter.write_str("host"),
78            Self::Cgroup => formatter.write_str("cgroup"),
79            Self::HostAndCgroup => formatter.write_str("host and cgroup equally"),
80            Self::CgroupProbeFailure => formatter.write_str("cgroup probe failure"),
81        }
82    }
83}
84
85/// Provenance-preserving memory availability for the current process.
86///
87/// Two quantities live here and they answer different questions.
88///
89/// * **Available** (`available_bytes`) is *free space right now*: a finite
90///   cgroup-v1 or cgroup-v2 ceiling admits exactly `min(host available,
91///   reclaim-aware cgroup available)`. A v2 literal `memory.max = max` remains
92///   typed as unbounded and therefore defers to the host; v1's numeric
93///   unlimited sentinel participates exactly and likewise loses to a tighter
94///   host value. This quantity *moves*, and under pressure it goes to zero.
95/// * **Capacity** (`capacity_bytes`) is *how much memory this process could
96///   ever address*: `min(host total, the binding cgroup's hard limit)`. It is
97///   a property of the box and the cgroup configuration, not of the schedule,
98///   so it does not move when a sibling allocates.
99///
100/// The distinction is load-bearing (#2684). "Does this dense footprint fit
101/// here at all?" is a *capacity* question and its answer must be stationary —
102/// see [`process_memory_availability`] on why a route derived from a moving
103/// quantity is a defect. "Does this allocation fit *right now*?" is an
104/// availability question, and the joint ledger
105/// ([`MemoryGovernor::try_reserve`]) is the instrument for it. Reading the
106/// second where the first was meant is what let a cgroup a few pages from its
107/// limit refuse a 28,800-byte design on a host with 448 GB free.
108///
109/// If an active controller cannot be parsed exactly, admission fails closed
110/// with zero bytes — both available *and* capacity — while retaining the typed
111/// probe failure; it never silently inherits host capacity.
112#[derive(Clone, Debug, PartialEq, Eq)]
113pub struct MemoryAvailability {
114    host_available_bytes: u64,
115    host_total_bytes: u64,
116    cgroup: CgroupMemoryObservation,
117    available_bytes: u64,
118    capacity_bytes: u64,
119    limiting_source: MemoryAvailabilitySource,
120}
121
122impl MemoryAvailability {
123    pub(crate) fn from_observation(
124        host_available_bytes: u64,
125        host_total_bytes: u64,
126        cgroup: CgroupMemoryObservation,
127    ) -> Self {
128        use std::cmp::Ordering;
129
130        // Capacity never depends on what is currently resident: a finite
131        // cgroup ceiling clamps the host's total, an unbounded or absent
132        // controller defers to it, and a probe that failed closed admits
133        // nothing at all. A host whose reported available exceeds its reported
134        // total (a torn /proc/meminfo read) must not shrink capacity below
135        // what is demonstrably reachable, hence the max.
136        let capacity_bytes = match &cgroup {
137            CgroupMemoryObservation::NotPresent | CgroupMemoryObservation::V2Unbounded { .. } => {
138                host_total_bytes.max(host_available_bytes)
139            }
140            CgroupMemoryObservation::V2Limited(observation)
141            | CgroupMemoryObservation::V1Limited(observation) => host_total_bytes
142                .max(host_available_bytes)
143                .min(observation.limit_bytes()),
144            CgroupMemoryObservation::ProbeFailed(_) => 0,
145        };
146
147        let (available_bytes, limiting_source) = match &cgroup {
148            CgroupMemoryObservation::NotPresent | CgroupMemoryObservation::V2Unbounded { .. } => {
149                (host_available_bytes, MemoryAvailabilitySource::Host)
150            }
151            CgroupMemoryObservation::V2Limited(observation)
152            | CgroupMemoryObservation::V1Limited(observation) => {
153                match observation.available_bytes().cmp(&host_available_bytes) {
154                    Ordering::Less => (
155                        observation.available_bytes(),
156                        MemoryAvailabilitySource::Cgroup,
157                    ),
158                    Ordering::Equal => (
159                        host_available_bytes,
160                        MemoryAvailabilitySource::HostAndCgroup,
161                    ),
162                    Ordering::Greater => (host_available_bytes, MemoryAvailabilitySource::Host),
163                }
164            }
165            CgroupMemoryObservation::ProbeFailed(_) => {
166                (0, MemoryAvailabilitySource::CgroupProbeFailure)
167            }
168        };
169        Self {
170            host_available_bytes,
171            host_total_bytes,
172            cgroup,
173            available_bytes,
174            capacity_bytes,
175            limiting_source,
176        }
177    }
178
179    pub const fn host_available_bytes(&self) -> u64 {
180        self.host_available_bytes
181    }
182
183    pub const fn host_total_bytes(&self) -> u64 {
184        self.host_total_bytes
185    }
186
187    /// The stationary ceiling on memory this process could ever address:
188    /// `min(host total, binding cgroup hard limit)`, or zero when the cgroup
189    /// probe failed closed. Unlike [`Self::available_bytes`] this does not move
190    /// when other work allocates, which is what makes it usable as a *routing*
191    /// threshold (#2684).
192    pub const fn capacity_bytes(&self) -> u64 {
193        self.capacity_bytes
194    }
195
196    pub fn capacity_bytes_usize(&self) -> usize {
197        usize::try_from(self.capacity_bytes).unwrap_or(usize::MAX)
198    }
199
200    pub const fn cgroup(&self) -> &CgroupMemoryObservation {
201        &self.cgroup
202    }
203
204    pub const fn available_bytes(&self) -> u64 {
205        self.available_bytes
206    }
207
208    pub fn available_bytes_usize(&self) -> usize {
209        usize::try_from(self.available_bytes).unwrap_or(usize::MAX)
210    }
211
212    pub const fn limiting_source(&self) -> MemoryAvailabilitySource {
213        self.limiting_source
214    }
215}
216
217impl std::fmt::Display for MemoryAvailability {
218    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
219        match &self.cgroup {
220            CgroupMemoryObservation::ProbeFailed(failure) => write!(
221                formatter,
222                "0 bytes admitted because the active cgroup probe failed closed (host_available={}, failure={})",
223                self.host_available_bytes, failure,
224            ),
225            observation => write!(
226                formatter,
227                "{} bytes limited by {} (capacity={}, host_available={}, host_total={}, {})",
228                self.available_bytes,
229                self.limiting_source,
230                self.capacity_bytes,
231                self.host_available_bytes,
232                self.host_total_bytes,
233                observation,
234            ),
235        }
236    }
237}
238
239/// Number of completed OS/cgroup availability probes in this process.
240///
241/// Every [`resample_memory_availability`] call opens and parses `/proc/meminfo`
242/// plus the four to five cgroup files behind `detect_cgroup_memory`, so this
243/// counter is a direct census of that syscall traffic. It exists so a planner
244/// can assert *in a test* that its budget decisions cost a bounded number of
245/// probes rather than one per unit of work (#2560).
246pub fn memory_availability_probe_count() -> u64 {
247    MEMORY_AVAILABILITY_PROBES.load(std::sync::atomic::Ordering::Relaxed)
248}
249
250static MEMORY_AVAILABILITY_PROBES: std::sync::atomic::AtomicU64 =
251    std::sync::atomic::AtomicU64::new(0);
252
253/// Take a **fresh** reading of the OS and cgroup memory observations.
254///
255/// This is a syscall-bearing probe, not an accessor: it refreshes `sysinfo`
256/// (`/proc/meminfo`) and re-reads the cgroup limit/usage files on every call,
257/// and the value it returns *moves* — available memory is a live, shared
258/// quantity owned by the whole machine, not by this process.
259///
260/// Call this only where a live reading is the point (an OOM guard immediately
261/// before committing an allocation). Anything that *routes* — picks a
262/// strategy, sizes a plan, chooses an operator — must use
263/// [`process_memory_availability`] instead, because a route derived from a
264/// moving quantity makes the fitted answer depend on what else the box was
265/// doing (SPEC-20: a fit object must only ever come from a converged
266/// optimization, and reproducibly so).
267pub fn resample_memory_availability() -> MemoryAvailability {
268    static SYSTEM: OnceLock<Mutex<sysinfo::System>> = OnceLock::new();
269    let system = SYSTEM.get_or_init(|| Mutex::new(sysinfo::System::new()));
270    let mut system = system.lock().expect("sysinfo system mutex poisoned");
271    system.refresh_memory();
272    let cgroup = detect_cgroup_memory();
273    MEMORY_AVAILABILITY_PROBES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
274    MemoryAvailability::from_observation(
275        system.available_memory(),
276        system.total_memory(),
277        cgroup,
278    )
279}
280
281/// The process's memory availability: **one** observation, taken once, shared
282/// by every planner in the process.
283///
284/// This is the same reading the [`MemoryGovernor`] sized its budget from, so a
285/// planner that admits a plan against this figure and the governor that then
286/// accounts the allocation cannot disagree about how much memory the process
287/// believes it has. Two consequences follow, and both are the point:
288///
289/// * **Free.** It is an atomic-free borrow of an already-initialized static —
290///   no syscalls, no parsing, no allocation — so it can sit on the hottest
291///   inner loop without a caching layer of its own.
292/// * **Constant.** Two runs of the same fit in one process see byte-identical
293///   budgets, so any route chosen from it is reproducible. Planners used to
294///   defend against the movement themselves (a monotone high-water floor in
295///   BMS, a per-term captured field in the SAE); a stationary primitive is what
296///   makes those compensations unnecessary rather than load-bearing.
297pub fn process_memory_availability() -> &'static MemoryAvailability {
298    &MemoryGovernor::global().ledger.availability
299}
300
301/// The process's available-memory figure in bytes, saturating to `usize`.
302/// Convenience over [`process_memory_availability`] for planners that only
303/// need the scalar.
304pub fn process_available_memory_bytes() -> usize {
305    process_memory_availability().available_bytes_usize()
306}
307
308/// Convert one provenance-preserving availability observation to the process
309/// ledger budget: 3/4 of the observation's stationary **capacity**. A probe
310/// that failed closed reports zero capacity and therefore admits nothing, which
311/// is the authoritative exhausted-memory signal.
312///
313/// gam#2702. This used to read `available_bytes` — free space — and that made
314/// every reservation verdict in the process a function of what the rest of the
315/// box happened to be doing, in exactly the way #2684 removed one layer up:
316///
317/// * the observation behind it is taken **once**, in the [`MemoryGovernor`]'s
318///   `OnceLock`, so it never tracked exhaustion in the first place. It froze
319///   whatever free memory existed at the instant of first governed allocation
320///   and used that number for the rest of the process's life. A frozen sample
321///   of a moving quantity is neither a live guard nor a reproducible
322///   threshold;
323/// * two processes launched from one job cgroup therefore derived different
324///   budgets from the same configuration, because page cache and sibling
325///   processes charge that cgroup continuously. Measured consequence: three
326///   `gam` inference tests passed in a 71-test run and failed in a 6-test
327///   subset of the same binary at the same commit, with
328///   `resource policy refused exact coefficient-SE columns 0..1` — a refusal of
329///   a few kilobytes.
330///
331/// The ledger's job is to bound **this process's own** jointly-live governed
332/// allocations against the ceiling this process could ever address. Memory
333/// already committed elsewhere on the box is not memory this ledger can account
334/// for, release, or route around; reading it only makes a fit's verdict a
335/// function of its neighbours, which SPEC-20 forbids (see
336/// [`process_memory_availability`]). Under a job scheduler the cgroup hard
337/// limit *is* the per-job ceiling, so capacity is the bound that actually binds.
338fn governor_budget_from_availability(availability: &MemoryAvailability) -> usize {
339    stationary_headroom_of_capacity(availability)
340}
341
342/// Convert the same observation to the process's **stationary** materialization
343/// ceiling: the answer to "could a dense footprint this large ever live here?",
344/// which must not move when a sibling allocates (#2684). A probe that failed
345/// closed reports zero capacity and therefore admits nothing.
346fn governor_materialization_cap_from_availability(availability: &MemoryAvailability) -> usize {
347    stationary_headroom_of_capacity(availability)
348}
349
350/// The one arithmetic both process-wide ceilings are derived from: the headroom
351/// fraction applied to stationary capacity.
352///
353/// The ledger budget and the materialization cap are therefore equal by
354/// construction rather than by coincidence, and that equality is the point:
355/// they are two questions about the same ceiling ("does this fit alongside what
356/// this process already has live" and "could this ever fit here at all"), so a
357/// routing decision taken from the cap can never be contradicted by the ledger
358/// for a reason that is not another live reservation. Keeping the arithmetic in
359/// one place is what stops the two from drifting back apart (#2684, #2702).
360fn stationary_headroom_of_capacity(availability: &MemoryAvailability) -> usize {
361    let scaled = u128::from(availability.capacity_bytes()) * GOVERNOR_BUDGET_NUMERATOR
362        / GOVERNOR_BUDGET_DENOMINATOR;
363    usize::try_from(scaled).unwrap_or(usize::MAX)
364}
365
366/// Typed refusal from [`MemoryGovernor::try_reserve`].
367///
368/// Carries the full ledger evidence so callers can route to a chunked or
369/// matrix-free strategy (and so error messages explain *why* dense was
370/// refused). This is a routing signal, never an abort: the process still has
371/// its unreserved headroom, the requested allocation just does not fit the
372/// joint budget.
373#[derive(Debug, Clone, thiserror::Error, PartialEq, Eq)]
374pub enum MemoryReservationError {
375    #[error(
376        "{context}: cannot reserve {requested_bytes} bytes; {reserved_bytes} of {budget_bytes} bytes already reserved process-wide; detected availability: {availability}"
377    )]
378    BudgetExceeded {
379        context: Box<str>,
380        requested_bytes: usize,
381        reserved_bytes: usize,
382        budget_bytes: usize,
383        availability: MemoryAvailability,
384    },
385
386    #[error(
387        "{context}: dense allocation size overflow for {copies} copies of a {nrows}x{ncols} f64 matrix"
388    )]
389    SizeOverflow {
390        context: Box<str>,
391        nrows: usize,
392        ncols: usize,
393        copies: usize,
394    },
395}
396
397#[derive(Debug)]
398struct GovernorLedger {
399    budget_bytes: usize,
400    materialization_cap_bytes: usize,
401    availability: MemoryAvailability,
402    reserved_bytes: std::sync::atomic::AtomicUsize,
403}
404
405/// Process-wide byte-accounting governor for large allocations.
406///
407/// Every large planned allocation (dense design materialization, covariance
408/// block, sampler design assembly) reserves its byte footprint against one
409/// shared ledger via [`try_reserve`](Self::try_reserve) and holds the returned
410/// RAII [`MemoryReservation`] for as long as the allocation is live. Because
411/// the ledger is shared, allocations that are each individually acceptable can
412/// no longer *jointly* exceed memory: whichever request would tip the ledger
413/// past the budget gets a typed [`MemoryReservationError`] and routes to a
414/// chunked or matrix-free strategy instead. Strategy selection is thereby a
415/// continuous function of predicted live bytes vs remaining budget, not of
416/// row/column thresholds.
417///
418/// The global budget is sized once from this process's stationary capacity
419/// (host total memory, clamped by the binding cgroup's hard limit inside
420/// containers or under a job scheduler) — see `GOVERNOR_BUDGET_NUMERATOR` for
421/// the headroom rationale and `governor_budget_from_availability` for why it
422/// is not denominated in free memory (#2702).
423#[derive(Debug, Clone)]
424pub struct MemoryGovernor {
425    ledger: Arc<GovernorLedger>,
426}
427
428impl MemoryGovernor {
429    /// The process-wide governor. Budget detection runs once, on first use.
430    ///
431    /// This `OnceLock` is where the process's single availability observation
432    /// is taken; [`process_memory_availability`] hands the same observation to
433    /// every other planner rather than probing again.
434    pub fn global() -> &'static MemoryGovernor {
435        static GLOBAL: OnceLock<MemoryGovernor> = OnceLock::new();
436        GLOBAL.get_or_init(|| {
437            let availability = resample_memory_availability();
438            MemoryGovernor::with_detected_availability(availability)
439        })
440    }
441
442    fn with_detected_availability(availability: MemoryAvailability) -> Self {
443        let budget_bytes = governor_budget_from_availability(&availability);
444        let materialization_cap_bytes = governor_materialization_cap_from_availability(&availability);
445        Self {
446            ledger: Arc::new(GovernorLedger {
447                budget_bytes,
448                materialization_cap_bytes,
449                availability,
450                reserved_bytes: std::sync::atomic::AtomicUsize::new(0),
451            }),
452        }
453    }
454
455    /// Total bytes this process's ledger may ever have reserved at once: 3/4 of
456    /// its stationary capacity. Two processes launched the same way on the same
457    /// box derive the same number, whatever else is resident (#2702).
458    pub fn budget_bytes(&self) -> usize {
459        self.ledger.budget_bytes
460    }
461
462    pub fn availability(&self) -> MemoryAvailability {
463        self.ledger.availability.clone()
464    }
465
466    pub fn reserved_bytes(&self) -> usize {
467        self.ledger
468            .reserved_bytes
469            .load(std::sync::atomic::Ordering::Acquire)
470    }
471
472    /// Bytes still admissible: the budget less what is reserved *by this
473    /// process* right now. The only quantity here that moves is this process's
474    /// own live governed footprint — which is the one thing a caller can
475    /// actually route around.
476    pub fn remaining_bytes(&self) -> usize {
477        self.ledger
478            .budget_bytes
479            .saturating_sub(self.reserved_bytes())
480    }
481
482    /// Absolute ceiling for one governed operation: 3/4 of this process's
483    /// memory **capacity** (`min(host total, binding cgroup limit)`).
484    ///
485    /// This is a *routing* threshold — the question it answers is "could a
486    /// dense footprint this large ever live in this process?" — so it is
487    /// deliberately stationary. It is **not** the live budget: whether an
488    /// allocation fits *right now* is decided by [`Self::try_reserve`] against
489    /// the joint ledger, which returns a typed, routable refusal instead of an
490    /// abort. Since #2702 both are denominated in the same capacity, so the two
491    /// can only disagree because of a live reservation.
492    ///
493    /// Returning the live budget here was gam#2684: the budget was 3/4 of
494    /// *available* memory, so a cgroup sitting at its limit drove this ceiling
495    /// continuously to zero (measured at 53,248 bytes available with 448 GB
496    /// free on the host), and every caller comparing a request against it —
497    /// including `DenseDesignMatrix::to_dense`, which panics on refusal —
498    /// refused allocations as small as a 300x12 design. Threshold-shaped
499    /// routing decisions taken from a moving quantity also make the chosen
500    /// route depend on what else the box was doing, which SPEC-20 forbids
501    /// (see [`process_memory_availability`]); one consumer even hashes this
502    /// cap into a basis cache key.
503    pub fn single_materialization_cap_bytes(&self) -> usize {
504        self.ledger.materialization_cap_bytes
505    }
506
507    /// Reserve `bytes` against the joint ledger.
508    ///
509    /// On success the returned [`MemoryReservation`] must be held for as long
510    /// as the allocation it accounts for is live; dropping it releases the
511    /// bytes. On failure the caller receives the ledger evidence and is
512    /// expected to fall back to a chunked or matrix-free strategy.
513    pub fn try_reserve(
514        &self,
515        bytes: usize,
516        context: &str,
517    ) -> Result<MemoryReservation, MemoryReservationError> {
518        use std::sync::atomic::Ordering;
519        let mut current = self.ledger.reserved_bytes.load(Ordering::Relaxed);
520        loop {
521            let next = match current.checked_add(bytes) {
522                Some(next) if next <= self.ledger.budget_bytes => next,
523                _ => {
524                    return Err(MemoryReservationError::BudgetExceeded {
525                        context: context.into(),
526                        requested_bytes: bytes,
527                        reserved_bytes: current,
528                        budget_bytes: self.ledger.budget_bytes,
529                        availability: self.ledger.availability.clone(),
530                    });
531                }
532            };
533            match self.ledger.reserved_bytes.compare_exchange_weak(
534                current,
535                next,
536                Ordering::AcqRel,
537                Ordering::Relaxed,
538            ) {
539                Ok(_) => {
540                    return Ok(MemoryReservation {
541                        ledger: Arc::clone(&self.ledger),
542                        bytes,
543                    });
544                }
545                Err(observed) => current = observed,
546            }
547        }
548    }
549
550    /// Reserve the footprint of a dense `nrows × ncols` `f64` matrix.
551    /// Dimension-product overflow is reported as a budget refusal (an
552    /// allocation whose size cannot even be computed certainly does not fit).
553    pub fn try_reserve_dense_f64(
554        &self,
555        nrows: usize,
556        ncols: usize,
557        context: &str,
558    ) -> Result<MemoryReservation, MemoryReservationError> {
559        self.try_reserve_dense_f64_copies(nrows, ncols, 1, context)
560    }
561
562    /// Reserve the predicted live footprint of `copies` simultaneous dense
563    /// matrices with one atomic ledger charge.
564    pub fn try_reserve_dense_f64_copies(
565        &self,
566        nrows: usize,
567        ncols: usize,
568        copies: usize,
569        context: &str,
570    ) -> Result<MemoryReservation, MemoryReservationError> {
571        let bytes = dense_f64_bytes(nrows, ncols)
572            .and_then(|one| one.checked_mul(copies))
573            .ok_or_else(|| MemoryReservationError::SizeOverflow {
574                context: context.into(),
575                nrows,
576                ncols,
577                copies,
578            })?;
579        self.try_reserve(bytes, context)
580    }
581}
582
583/// Checked byte footprint of a dense `nrows × ncols` `f64` matrix.
584pub const fn dense_f64_bytes(nrows: usize, ncols: usize) -> Option<usize> {
585    match nrows.checked_mul(ncols) {
586        Some(cells) => cells.checked_mul(std::mem::size_of::<f64>()),
587        None => None,
588    }
589}
590
591/// RAII guard for bytes reserved on a [`MemoryGovernor`] ledger; dropping it
592/// releases the reservation. Hold it exactly as long as the accounted
593/// allocation is live.
594#[derive(Debug)]
595#[must_use = "dropping a memory reservation immediately releases its ledger charge"]
596pub struct MemoryReservation {
597    ledger: Arc<GovernorLedger>,
598    bytes: usize,
599}
600
601impl MemoryReservation {
602    pub fn bytes(&self) -> usize {
603        self.bytes
604    }
605
606    /// Couple this reservation to the value whose memory it accounts for.
607    pub fn bind<T>(self, value: T) -> Governed<T> {
608        Governed {
609            value,
610            reservation: self,
611        }
612    }
613}
614
615/// A value whose live memory is coupled to a process-wide reservation.
616///
617/// Large fallible materializations return this owner so the allocation cannot
618/// outlive its ledger charge. It dereferences to the wrapped value for normal
619/// ndarray and collection operations.
620#[derive(Debug)]
621#[must_use = "the governed value owns a live process-wide memory reservation"]
622pub struct Governed<T> {
623    value: T,
624    reservation: MemoryReservation,
625}
626
627impl<T> Governed<T> {
628    pub fn reserved_bytes(&self) -> usize {
629        self.reservation.bytes()
630    }
631}
632
633impl<T> std::ops::Deref for Governed<T> {
634    type Target = T;
635
636    fn deref(&self) -> &Self::Target {
637        &self.value
638    }
639}
640
641impl<T> std::ops::DerefMut for Governed<T> {
642    fn deref_mut(&mut self) -> &mut Self::Target {
643        &mut self.value
644    }
645}
646
647impl<T> AsRef<T> for Governed<T> {
648    fn as_ref(&self) -> &T {
649        &self.value
650    }
651}
652
653impl<T> AsMut<T> for Governed<T> {
654    fn as_mut(&mut self) -> &mut T {
655        &mut self.value
656    }
657}
658
659impl Drop for MemoryReservation {
660    fn drop(&mut self) {
661        self.ledger
662            .reserved_bytes
663            .fetch_sub(self.bytes, std::sync::atomic::Ordering::AcqRel);
664    }
665}
666
667/// Hints that flip strict mode on regardless of n/p — used when a code path
668/// is structurally operator-only and any dense fallback would be a bug.
669#[derive(Clone, Copy, Debug, Default)]
670pub struct ProblemHints {
671    pub marginal_slope_large_scale_active: bool,
672}
673
674#[derive(Clone, Copy, Debug, PartialEq, Eq)]
675pub enum DerivativeStorageMode {
676    /// Production exact-math: operator-backed, no dense fallback.
677    AnalyticOperatorRequired,
678    /// Allow dense materialization if under the single-materialization budget.
679    MaterializeIfSmall,
680    /// Dense materialization only permitted for diagnostic code paths.
681    DiagnosticsOnly,
682}
683
684#[derive(Clone, Debug)]
685pub struct MaterializationPolicy {
686    pub max_single_dense_bytes: usize,
687    pub max_cached_dense_bytes: usize,
688    pub row_chunk_target_bytes: usize,
689    pub allow_operator_materialization: bool,
690    pub allow_diagnostic_materialization: bool,
691}
692
693#[derive(Debug, thiserror::Error)]
694pub enum MatrixMaterializationError {
695    #[error(
696        "{context}: dense materialization of {nrows}x{ncols} requires {bytes} bytes (limit {limit_bytes})"
697    )]
698    TooLarge {
699        context: &'static str,
700        nrows: usize,
701        ncols: usize,
702        bytes: usize,
703        limit_bytes: usize,
704    },
705
706    #[error("{context}: operator does not implement chunked row access")]
707    MissingRowChunk { context: &'static str },
708
709    #[error("{context}: row materialization failed: {reason}")]
710    RowMaterializationFailed {
711        context: &'static str,
712        reason: String,
713    },
714
715    #[error("{context}: materialization forbidden by policy (mode={mode:?})")]
716    Forbidden {
717        context: &'static str,
718        mode: DerivativeStorageMode,
719    },
720
721    /// The process-wide [`MemoryGovernor`] could not reserve the requested live
722    /// footprint (or its checked byte size overflowed). Callers route to a
723    /// chunked or matrix-free strategy.
724    #[error(transparent)]
725    Reservation(#[from] MemoryReservationError),
726}
727
728pub trait ResidentBytes {
729    fn resident_bytes(&self) -> usize;
730}
731
732impl ResourcePolicy {
733    /// Conservative default suitable for general-purpose use.
734    ///
735    /// Uses `MaterializeIfSmall`: dense materialization is allowed only when
736    /// the matrix fits under `max_single_materialization_bytes`. This lets
737    /// small-data families that lack an implicit operator work out of the box,
738    /// while problems whose dense footprint does not fit real memory get a
739    /// typed refusal that forces the analytic-operator path. Set
740    /// `derivative_storage_mode = AnalyticOperatorRequired` explicitly to
741    /// reject all dense fallback.
742    ///
743    /// Scalar caps expose the governor's stationary capacity ceiling; they are
744    /// routing thresholds, not independent budgets, and they do not move with
745    /// load. Actual materializations and caches must reserve their complete
746    /// live footprint against the shared governor, so any combination of
747    /// categories is bounded by one ledger — that reservation, not these caps,
748    /// is what enforces "fits right now" (#2684).
749    pub fn default_library() -> Self {
750        Self::for_observed_memory(process_memory_availability())
751    }
752
753    /// The library-default policy a process would derive from `availability`.
754    ///
755    /// [`Self::default_library`] is exactly this function applied to the
756    /// process's own observation, so the routing thresholds a fit sees are a
757    /// pure function of the *observed environment* — and, because every scalar
758    /// below is taken from
759    /// [`MemoryAvailability::capacity_bytes`] rather than
760    /// [`MemoryAvailability::available_bytes`], a pure function of that
761    /// environment's stationary CAPACITY. Two observations of the same box that
762    /// differ only in how much memory happened to be free at the instant of the
763    /// probe therefore produce byte-identical policies, and any route selected
764    /// from one of these caps is reproducible across processes and across load
765    /// (#2684).
766    ///
767    /// Exposed rather than kept private so a test can hold capacity fixed and
768    /// vary free memory without racing the real machine; the live path uses the
769    /// same code, so the property under test is the shipped one.
770    pub fn for_observed_memory(availability: &MemoryAvailability) -> Self {
771        let single_cap = governor_materialization_cap_from_availability(availability);
772        Self {
773            max_single_materialization_bytes: single_cap,
774            max_operator_cache_bytes: single_cap,
775            max_spatial_distance_cache_bytes: single_cap,
776            max_owned_data_cache_bytes: single_cap,
777            row_chunk_target_bytes: LIBRARY_ROW_CHUNK_TARGET_BYTES,
778            derivative_storage_mode: DerivativeStorageMode::MaterializeIfSmall,
779        }
780    }
781
782    /// Strict mode that rejects every dense fallback. Use when you intend to
783    /// run only on operator-backed bases (large-scale Duchon/TPS, exact
784    /// GAMLSS marginal slope, CTN, etc.). The byte caps only govern the
785    /// residual diagnostic surfaces (materialization itself is forbidden by
786    /// the mode).
787    pub fn analytic_operator_required() -> Self {
788        let base = Self::default_library();
789        Self {
790            derivative_storage_mode: DerivativeStorageMode::AnalyticOperatorRequired,
791            ..base
792        }
793    }
794
795    /// Auto-derive the resource policy from the shape of the problem rather
796    /// than from an explicit CLI flag.
797    ///
798    /// Shape alone never flips a policy mode: doing so merely moves the old
799    /// row/column cliff to a byte threshold. Every non-structural path starts
800    /// permissive and makes its strategy decision from the operation's checked
801    /// predicted live bytes versus the governor's current remaining budget.
802    ///
803    /// `hints.marginal_slope_large_scale_active` forces strict mode regardless
804    /// of shape: that path is structurally operator-only and any dense
805    /// fallback would be a bug, not a memory question.
806    pub fn for_problem(hints: ProblemHints) -> Self {
807        if hints.marginal_slope_large_scale_active {
808            return Self::analytic_operator_required();
809        }
810        Self::default_library()
811    }
812
813    /// Permissive mode for small-data usage and tests. Admission still uses
814    /// the same process ledger; only the streaming chunk geometry differs.
815    pub fn permissive_small_data() -> Self {
816        let base = Self::default_library();
817        Self {
818            row_chunk_target_bytes: 64 * 1024 * 1024,
819            ..base
820        }
821    }
822
823    pub const fn material_policy(&self) -> MaterializationPolicy {
824        MaterializationPolicy {
825            max_single_dense_bytes: self.max_single_materialization_bytes,
826            max_cached_dense_bytes: self.max_operator_cache_bytes,
827            row_chunk_target_bytes: self.row_chunk_target_bytes,
828            allow_operator_materialization: matches!(
829                self.derivative_storage_mode,
830                DerivativeStorageMode::MaterializeIfSmall
831            ),
832            allow_diagnostic_materialization: !matches!(
833                self.derivative_storage_mode,
834                DerivativeStorageMode::AnalyticOperatorRequired
835            ),
836        }
837    }
838}
839
840/// Returns how many rows to stream per chunk so that each chunk uses approximately
841/// `target_bytes` given a row width of `cols` f64 entries.
842pub const fn rows_for_target_bytes(target_bytes: usize, cols: usize) -> usize {
843    let raw_bytes_per_row = cols.saturating_mul(std::mem::size_of::<f64>());
844    let bytes_per_row = if raw_bytes_per_row == 0 {
845        1
846    } else {
847        raw_bytes_per_row
848    };
849    let rows = target_bytes / bytes_per_row;
850    if rows == 0 { 1 } else { rows }
851}
852
853/// Select the row count for prediction-time covariance work.
854///
855/// A prediction row can keep roughly four `parameter_dim × local_dim`
856/// `f64` workspaces live while gradients and covariance solves are assembled.
857/// This central policy prevents predictors from drifting to different memory
858/// budgets and chunk bounds for the same operation.
859pub fn prediction_chunk_rows(parameter_dim: usize, local_dim: usize, total_rows: usize) -> usize {
860    const MIN_ROWS: usize = 16;
861    const MAX_ROWS: usize = 4096;
862
863    if total_rows == 0 {
864        return 1;
865    }
866    let live_f64_values_per_row = parameter_dim
867        .max(1)
868        .saturating_mul(local_dim.max(1))
869        .saturating_mul(4);
870    rows_for_target_bytes(
871        ResourcePolicy::default_library().row_chunk_target_bytes,
872        live_f64_values_per_row,
873    )
874    .clamp(MIN_ROWS, MAX_ROWS)
875    .min(total_rows)
876}
877
878use std::collections::{HashMap, VecDeque};
879use std::hash::{Hash, Hasher};
880use std::sync::{Arc, Mutex, OnceLock};
881
882/// Byte-limited LRU cache with an optional entry cap.
883///
884/// Unlike an entry-count-limited LRU, this cache tracks the resident byte cost
885/// of each value (via [`ResidentBytes`]) and evicts the least-recently-used
886/// entries until the total resident bytes fit under `max_bytes`. This is the
887/// correct policy for large-scale payloads where a single cache entry (e.g.
888/// an n*K distance matrix) can itself be multiple gigabytes and an entry-count
889/// cap would silently blow the memory budget. Small entry caps are still useful
890/// for payloads with known shape, such as owned PC data matrices shared across
891/// model blocks.
892pub struct ByteLruCache<K: Eq + Hash + Clone, V> {
893    /// One independent LRU partition per shard. A single shard (the default)
894    /// is byte-for-byte equivalent to the original single-`Mutex` cache; with
895    /// `shard_count > 1` the key hash selects the shard, so concurrent traffic
896    /// on distinct keys contends `1/shard_count` as often and each shard's
897    /// recency `VecDeque` is `1/shard_count` as long (the hit-path rescan is a
898    /// linear `position` lookup, so shrinking the per-shard order also cuts
899    /// per-access cost). Sharding is opt-in (`new_sharded`) precisely because
900    /// the byte budget is split across shards — that is correct for caches of
901    /// many small entries (e.g. cell-moment memos) but wrong for caches of a
902    /// few multi-GiB entries (distance matrices), which keep `shard_count == 1`.
903    shards: Box<[Mutex<ByteLruInner<K, V>>]>,
904    /// Per-shard byte budget. `shard_bytes * shards.len() >= max_bytes`.
905    shard_bytes: usize,
906    /// Per-shard entry budget, if any (`0` disables caching, as before).
907    shard_entries: Option<usize>,
908    max_bytes: usize,
909    governor: MemoryGovernor,
910}
911
912struct ByteLruInner<K, V> {
913    // The reservation is stored beside the value, so eviction and clear drop
914    // the process-wide charge at exactly the same time as cache ownership.
915    map: HashMap<K, (V, usize, MemoryReservation)>,
916    order: VecDeque<K>,
917    resident_bytes: usize,
918}
919
920impl<K: Eq + Hash + Clone, V: Clone + ResidentBytes> ByteLruCache<K, V> {
921    pub fn new(max_bytes: usize) -> Self {
922        Self::build(max_bytes, None, 1)
923    }
924
925    pub fn with_max_entries(max_bytes: usize, max_entries: usize) -> Self {
926        Self::build(max_bytes, Some(max_entries), 1)
927    }
928
929    /// Like [`new`](Self::new) but partitions the cache across `shard_count`
930    /// independently-locked LRU shards to cut lock contention under heavy
931    /// concurrent access. The byte budget is divided evenly across shards, so
932    /// this is only appropriate for caches holding many small entries.
933    pub fn new_sharded(max_bytes: usize, shard_count: usize) -> Self {
934        Self::build(max_bytes, None, shard_count)
935    }
936
937    /// Like [`with_max_entries`](Self::with_max_entries) but sharded; see
938    /// [`new_sharded`](Self::new_sharded).
939    pub fn with_max_entries_sharded(
940        max_bytes: usize,
941        max_entries: usize,
942        shard_count: usize,
943    ) -> Self {
944        Self::build(max_bytes, Some(max_entries), shard_count)
945    }
946
947    fn build(max_bytes: usize, max_entries: Option<usize>, shard_count: usize) -> Self {
948        Self::build_with_governor(
949            max_bytes,
950            max_entries,
951            shard_count,
952            MemoryGovernor::global().clone(),
953        )
954    }
955
956    fn build_with_governor(
957        max_bytes: usize,
958        max_entries: Option<usize>,
959        shard_count: usize,
960        governor: MemoryGovernor,
961    ) -> Self {
962        let shard_count = shard_count.max(1);
963        // Split the global budgets across shards, rounding up so the aggregate
964        // capacity never falls below the requested budget. With a single shard
965        // these equal the global budgets exactly (legacy behavior). A `0`
966        // entry budget still disables caching and must not be rounded up to 1.
967        let shard_bytes = max_bytes.div_ceil(shard_count);
968        let shard_entries = max_entries.map(|m| {
969            if m == 0 {
970                0
971            } else {
972                m.div_ceil(shard_count).max(1)
973            }
974        });
975        let shards = (0..shard_count)
976            .map(|_| {
977                Mutex::new(ByteLruInner {
978                    map: HashMap::new(),
979                    order: VecDeque::new(),
980                    resident_bytes: 0,
981                })
982            })
983            .collect::<Vec<_>>()
984            .into_boxed_slice();
985        Self {
986            shards,
987            shard_bytes,
988            shard_entries,
989            max_bytes,
990            governor,
991        }
992    }
993
994    #[inline]
995    fn shard(&self, key: &K) -> &Mutex<ByteLruInner<K, V>> {
996        if self.shards.len() == 1 {
997            return &self.shards[0];
998        }
999        let mut hasher = std::collections::hash_map::DefaultHasher::new();
1000        key.hash(&mut hasher);
1001        &self.shards[(hasher.finish() as usize) % self.shards.len()]
1002    }
1003
1004    pub fn get(&self, key: &K) -> Option<V> {
1005        // recover from poison
1006        let mut g = self.shard(key).lock().unwrap_or_else(|p| p.into_inner());
1007        let v = g.map.get(key)?.0.clone();
1008        // move to back (most-recently-used)
1009        if let Some(pos) = g.order.iter().position(|k| k == key) {
1010            let k = g
1011                .order
1012                .remove(pos)
1013                .expect("position() returned an in-bounds index into this same deque");
1014            g.order.push_back(k);
1015        }
1016        Some(v)
1017    }
1018
1019    pub fn insert(&self, key: K, value: V) {
1020        let charge = value.resident_bytes();
1021        let mut g = self.shard(&key).lock().unwrap_or_else(|p| p.into_inner());
1022
1023        // If already present, remove the old entry first so resident bytes stay
1024        // accurate and the LRU ordering reflects this insertion.
1025        if let Some((_old, old_charge, _reservation)) = g.map.remove(&key) {
1026            g.resident_bytes = g.resident_bytes.saturating_sub(old_charge);
1027            if let Some(pos) = g.order.iter().position(|k| k == &key) {
1028                g.order.remove(pos);
1029            }
1030        }
1031
1032        if charge > self.shard_bytes {
1033            // Too large to cache; skip insertion.
1034            return;
1035        }
1036
1037        if let Some(max_entries) = self.shard_entries {
1038            if max_entries == 0 {
1039                return;
1040            }
1041            while g.map.len() >= max_entries {
1042                if let Some(evict_key) = g.order.pop_front() {
1043                    if let Some((_v, c, _reservation)) = g.map.remove(&evict_key) {
1044                        g.resident_bytes = g.resident_bytes.saturating_sub(c);
1045                    }
1046                } else {
1047                    break;
1048                }
1049            }
1050        }
1051
1052        while g.resident_bytes + charge > self.shard_bytes {
1053            if let Some(evict_key) = g.order.pop_front() {
1054                if let Some((_v, c, _reservation)) = g.map.remove(&evict_key) {
1055                    g.resident_bytes = g.resident_bytes.saturating_sub(c);
1056                }
1057            } else {
1058                break;
1059            }
1060        }
1061
1062        let reservation = match self
1063            .governor
1064            .try_reserve(charge, "ByteLruCache resident entry")
1065        {
1066            Ok(reservation) => reservation,
1067            Err(_) => return,
1068        };
1069        g.map.insert(key.clone(), (value, charge, reservation));
1070        g.order.push_back(key);
1071        g.resident_bytes = g.resident_bytes.saturating_add(charge);
1072    }
1073
1074    pub fn resident_bytes(&self) -> usize {
1075        self.shards
1076            .iter()
1077            .map(|shard| {
1078                shard
1079                    .lock()
1080                    .unwrap_or_else(|p| p.into_inner())
1081                    .resident_bytes
1082            })
1083            .sum()
1084    }
1085
1086    pub const fn max_bytes(&self) -> usize {
1087        self.max_bytes
1088    }
1089
1090    pub fn len(&self) -> usize {
1091        self.shards
1092            .iter()
1093            .map(|shard| shard.lock().unwrap_or_else(|p| p.into_inner()).map.len())
1094            .sum()
1095    }
1096
1097    pub fn is_empty(&self) -> bool {
1098        self.len() == 0
1099    }
1100
1101    pub fn clear(&self) {
1102        for shard in self.shards.iter() {
1103            let mut g = shard.lock().unwrap_or_else(|p| p.into_inner());
1104            g.map.clear();
1105            g.order.clear();
1106            g.resident_bytes = 0;
1107        }
1108    }
1109}
1110
1111impl<K: Eq + Hash + Clone, V: Clone + ResidentBytes> std::fmt::Debug for ByteLruCache<K, V> {
1112    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1113        f.debug_struct("ByteLruCache")
1114            .field("resident_bytes", &self.resident_bytes())
1115            .field("max_bytes", &self.max_bytes)
1116            .field("shard_count", &self.shards.len())
1117            .field("shard_bytes", &self.shard_bytes)
1118            .field("shard_entries", &self.shard_entries)
1119            .finish()
1120    }
1121}
1122
1123/// Byte-accounting for `Arc<Array2<f64>>`.
1124///
1125/// Reports the full dense footprint of the owned array. Multiple `Arc`s
1126/// pointing to the same allocation will each report the full size; this is
1127/// the conservative accounting the caches want because a single residency in
1128/// the cache is what we are budgeting for.
1129impl ResidentBytes for Arc<ndarray::Array2<f64>> {
1130    fn resident_bytes(&self) -> usize {
1131        std::mem::size_of::<f64>()
1132            .saturating_mul(self.nrows())
1133            .saturating_mul(self.ncols())
1134    }
1135}
1136
1137/// Lazy-init cache safe to call from inside rayon par_iter.
1138///
1139/// `std::sync::OnceLock::get_or_init` parks racing threads on an OS
1140/// condition variable until the leader's init closure finishes. If the
1141/// leader's init closure itself dispatches a nested `into_par_iter`, the
1142/// parked threads are now unavailable as rayon workers, and the leader
1143/// blocks waiting for chunks that no one can service. Classic deadlock.
1144///
1145/// `RayonSafeOnce` removes the trap by computing the value *outside* any
1146/// lock. Concurrent racers may produce duplicate values; the first to
1147/// publish wins, the rest drop their result. No thread ever parks waiting
1148/// for another thread's init to finish, so nested rayon par_iter inside
1149/// the init closure is safe.
1150///
1151/// Use this in place of `OnceLock` whenever the init closure transitively
1152/// runs rayon work *and* the cache may be entered concurrently from
1153/// inside another rayon par_iter. The redundant-work cost on first race
1154/// is the price for never deadlocking; in practice the loser threads
1155/// throw away one round of work and steady-state is identical to
1156/// `OnceLock`.
1157pub struct RayonSafeOnce<T> {
1158    slot: std::sync::OnceLock<T>,
1159}
1160
1161impl<T> RayonSafeOnce<T> {
1162    pub const fn new() -> Self {
1163        Self {
1164            slot: std::sync::OnceLock::new(),
1165        }
1166    }
1167
1168    /// Returns the cached value if already populated.
1169    #[inline]
1170    pub fn get(&self) -> Option<&T> {
1171        self.slot.get()
1172    }
1173
1174    /// Returns the cached value, computing it if absent.
1175    ///
1176    /// The init closure runs WITHOUT holding any lock — calls from
1177    /// concurrent rayon workers may all run it, and all but the first
1178    /// to call `set` discard their result. This is the contract that
1179    /// keeps nested `into_par_iter` inside `init` from deadlocking on
1180    /// other workers parked on a `OnceLock`.
1181    ///
1182    /// Named `get_or_compute` (not `get_or_init`) so the codebase-level
1183    /// lint that bans `OnceLock::get_or_init` near rayon `par_iter` does
1184    /// not flag this safe-by-construction path.
1185    pub fn get_or_compute<F>(&self, init: F) -> &T
1186    where
1187        F: FnOnce() -> T,
1188    {
1189        if let Some(v) = self.slot.get() {
1190            return v;
1191        }
1192        let candidate = init();
1193        if self.slot.set(candidate).is_err() {
1194            log::trace!(
1195                "RayonSafeOnce: a concurrent initializer won the race; \
1196                 keeping its value and discarding this candidate"
1197            );
1198        }
1199        self.slot
1200            .get()
1201            .expect("RayonSafeOnce slot populated by set() above")
1202    }
1203}
1204
1205impl<T> Default for RayonSafeOnce<T> {
1206    fn default() -> Self {
1207        Self::new()
1208    }
1209}
1210
1211impl<T: Clone> Clone for RayonSafeOnce<T> {
1212    fn clone(&self) -> Self {
1213        let cloned = Self::new();
1214        if let Some(value) = self.slot.get() {
1215            // `get_or_init` rather than `set(..).expect(..)`: the slot of a
1216            // freshly constructed `Self` is empty, so the initializer always
1217            // runs — and this spelling has no `Result` to discard, which keeps
1218            // the impl's bound at `T: Clone` instead of forcing `T: Debug` on
1219            // every caller just to name a panic that cannot happen.
1220            cloned.slot.get_or_init(|| value.clone());
1221        }
1222        cloned
1223    }
1224}
1225
1226impl<T: std::fmt::Debug> std::fmt::Debug for RayonSafeOnce<T> {
1227    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1228        f.debug_struct("RayonSafeOnce")
1229            .field("slot", &self.slot.get())
1230            .finish()
1231    }
1232}
1233
1234#[cfg(test)]
1235mod byte_lru_tests {
1236    use super::*;
1237
1238    fn cache_test_governor(budget_bytes: usize) -> MemoryGovernor {
1239        let available_bytes = (budget_bytes as u128 * GOVERNOR_BUDGET_DENOMINATOR)
1240            .div_ceil(GOVERNOR_BUDGET_NUMERATOR);
1241        let available_bytes =
1242            u64::try_from(available_bytes).expect("test cache budget must fit in u64");
1243        MemoryGovernor::with_detected_availability(MemoryAvailability::from_observation(
1244            available_bytes,
1245            available_bytes,
1246            CgroupMemoryObservation::NotPresent,
1247        ))
1248    }
1249
1250    /// Fixed-charge value so byte-budget arithmetic in the tests is exact.
1251    #[derive(Clone, PartialEq, Debug)]
1252    struct Payload(u64);
1253    impl ResidentBytes for Payload {
1254        fn resident_bytes(&self) -> usize {
1255            8
1256        }
1257    }
1258
1259    #[test]
1260    fn single_shard_round_trips_and_evicts_by_bytes() {
1261        // 3 entries' worth of budget; a single shard preserves strict global LRU.
1262        let cache: ByteLruCache<u64, Payload> =
1263            ByteLruCache::build_with_governor(24, None, 1, cache_test_governor(24));
1264        for k in 0..3 {
1265            cache.insert(k, Payload(k));
1266        }
1267        assert_eq!(cache.len(), 3);
1268        assert_eq!(cache.resident_bytes(), 24);
1269        // Touch key 0 so it is most-recently-used, then overflow by one.
1270        assert_eq!(cache.get(&0), Some(Payload(0)));
1271        cache.insert(3, Payload(3));
1272        // Key 1 (now least-recently-used) is evicted; 0 survives the touch.
1273        assert_eq!(cache.len(), 3);
1274        assert_eq!(cache.get(&1), None);
1275        assert_eq!(cache.get(&0), Some(Payload(0)));
1276        assert_eq!(cache.get(&3), Some(Payload(3)));
1277    }
1278
1279    #[test]
1280    fn zero_entry_budget_disables_caching_in_every_shard() {
1281        let single: ByteLruCache<u64, Payload> = ByteLruCache::with_max_entries(1 << 20, 0);
1282        single.insert(7, Payload(7));
1283        assert_eq!(single.get(&7), None);
1284        let sharded: ByteLruCache<u64, Payload> =
1285            ByteLruCache::with_max_entries_sharded(1 << 20, 0, 16);
1286        sharded.insert(7, Payload(7));
1287        assert_eq!(sharded.get(&7), None);
1288    }
1289
1290    #[test]
1291    fn sharded_cache_retrieves_all_keys_and_respects_aggregate_budget() {
1292        // Generous budget split across 8 shards; every inserted key must be
1293        // retrievable and the aggregate residency must never exceed the global
1294        // budget (shard_bytes * shard_count, rounded up).
1295        let shard_count = 8usize;
1296        let max_bytes = 8 * 64; // 64 entries' worth, 8 per shard on average.
1297        let cache: ByteLruCache<u64, Payload> = ByteLruCache::build_with_governor(
1298            max_bytes,
1299            None,
1300            shard_count,
1301            cache_test_governor(max_bytes),
1302        );
1303        for k in 0..64u64 {
1304            cache.insert(k, Payload(k));
1305        }
1306        // Per-shard budgets sum to >= the requested global budget.
1307        assert!(cache.resident_bytes() <= max_bytes.div_ceil(shard_count) * shard_count);
1308        // Re-inserting then reading back a key returns the stored payload.
1309        cache.insert(123, Payload(123));
1310        assert_eq!(cache.get(&123), Some(Payload(123)));
1311        assert!(!cache.is_empty());
1312        cache.clear();
1313        assert_eq!(cache.len(), 0);
1314        assert_eq!(cache.resident_bytes(), 0);
1315    }
1316}
1317
1318#[cfg(test)]
1319mod resource_policy_tests {
1320    use super::*;
1321
1322    fn test_governor(budget_bytes: usize) -> MemoryGovernor {
1323        let available_bytes = (budget_bytes as u128 * GOVERNOR_BUDGET_DENOMINATOR)
1324            .div_ceil(GOVERNOR_BUDGET_NUMERATOR);
1325        let available_bytes =
1326            u64::try_from(available_bytes).expect("test budget has a representable observation");
1327        let availability = MemoryAvailability::from_observation(
1328            available_bytes,
1329            available_bytes,
1330            CgroupMemoryObservation::NotPresent,
1331        );
1332        let governor = MemoryGovernor::with_detected_availability(availability);
1333        assert_eq!(governor.budget_bytes(), budget_bytes);
1334        governor
1335    }
1336
1337    // ── rows_for_target_bytes ─────────────────────────────────────────────────
1338
1339    #[test]
1340    fn rows_for_target_bytes_exact_fit() {
1341        // 1 col × 8 bytes/f64; target 8 bytes → 1 row
1342        assert_eq!(rows_for_target_bytes(8, 1), 1);
1343    }
1344
1345    #[test]
1346    fn rows_for_target_bytes_multiple_rows() {
1347        // 1 col × 8 bytes/f64; target 80 bytes → 10 rows
1348        assert_eq!(rows_for_target_bytes(80, 1), 10);
1349    }
1350
1351    #[test]
1352    fn rows_for_target_bytes_multiple_cols() {
1353        // 4 cols × 8 = 32 bytes/row; target 128 → 4 rows
1354        assert_eq!(rows_for_target_bytes(128, 4), 4);
1355    }
1356
1357    #[test]
1358    fn rows_for_target_bytes_zero_target_returns_one() {
1359        // Zero target cannot give 0 rows — floor to 1
1360        assert_eq!(rows_for_target_bytes(0, 1), 1);
1361    }
1362
1363    #[test]
1364    fn rows_for_target_bytes_zero_cols_returns_non_zero() {
1365        // Zero cols → bytes_per_row falls back to 1 → rows = target
1366        assert_eq!(rows_for_target_bytes(100, 0), 100);
1367    }
1368
1369    #[test]
1370    fn rows_for_target_bytes_large_target() {
1371        // Canonical target, 1024 cols → 8192 bytes/row → rows = target/8192.
1372        // Taken by name rather than transcribed, so the case stays pinned to the
1373        // library target if that target ever moves (#2704).
1374        let target = LIBRARY_ROW_CHUNK_TARGET_BYTES;
1375        let cols = 1024_usize;
1376        let expected = target / (cols * std::mem::size_of::<f64>());
1377        assert_eq!(rows_for_target_bytes(target, cols), expected);
1378    }
1379
1380    #[test]
1381    fn prediction_chunks_share_the_runtime_byte_budget() {
1382        assert_eq!(prediction_chunk_rows(1024, 1, 100_000), 256);
1383        assert_eq!(prediction_chunk_rows(32, 2, 100_000), 4096);
1384    }
1385
1386    #[test]
1387    fn prediction_chunks_respect_dataset_bounds() {
1388        assert_eq!(prediction_chunk_rows(1, 1, 7), 7);
1389        assert_eq!(prediction_chunk_rows(1, 1, 0), 1);
1390    }
1391
1392    // ── ResourcePolicy::for_problem ──────────────────────────────────────────
1393
1394    #[test]
1395    fn for_problem_small_data_uses_materialize_if_small() {
1396        let p = ResourcePolicy::for_problem(ProblemHints::default());
1397        assert_eq!(
1398            p.derivative_storage_mode,
1399            DerivativeStorageMode::MaterializeIfSmall
1400        );
1401    }
1402
1403    #[test]
1404    fn for_problem_has_no_row_or_column_cliff() {
1405        let narrow = ResourcePolicy::for_problem(ProblemHints::default());
1406        let wide = ResourcePolicy::for_problem(ProblemHints::default());
1407        assert_eq!(
1408            narrow.derivative_storage_mode,
1409            DerivativeStorageMode::MaterializeIfSmall
1410        );
1411        assert_eq!(
1412            wide.derivative_storage_mode,
1413            DerivativeStorageMode::MaterializeIfSmall
1414        );
1415    }
1416
1417    #[test]
1418    fn for_problem_dimension_overflow_defers_to_typed_reservation() {
1419        let policy = ResourcePolicy::for_problem(ProblemHints::default());
1420        assert_eq!(
1421            policy.derivative_storage_mode,
1422            DerivativeStorageMode::MaterializeIfSmall
1423        );
1424    }
1425
1426    #[test]
1427    fn for_problem_marginal_slope_hint_is_strict() {
1428        let p = ResourcePolicy::for_problem(ProblemHints {
1429            marginal_slope_large_scale_active: true,
1430        });
1431        assert_eq!(
1432            p.derivative_storage_mode,
1433            DerivativeStorageMode::AnalyticOperatorRequired
1434        );
1435    }
1436
1437    // ── ResourcePolicy::material_policy ─────────────────────────────────────
1438
1439    #[test]
1440    fn material_policy_default_library_allows_operator_and_diagnostics() {
1441        let mp = ResourcePolicy::default_library().material_policy();
1442        assert!(mp.allow_operator_materialization);
1443        assert!(mp.allow_diagnostic_materialization);
1444    }
1445
1446    #[test]
1447    fn material_policy_analytic_operator_required_blocks_both() {
1448        let mp = ResourcePolicy::analytic_operator_required().material_policy();
1449        assert!(!mp.allow_operator_materialization);
1450        assert!(!mp.allow_diagnostic_materialization);
1451    }
1452
1453    #[test]
1454    fn material_policy_propagates_byte_limits() {
1455        let policy = ResourcePolicy::default_library();
1456        let mp = policy.material_policy();
1457        assert_eq!(
1458            mp.max_single_dense_bytes,
1459            policy.max_single_materialization_bytes
1460        );
1461        assert_eq!(mp.max_cached_dense_bytes, policy.max_operator_cache_bytes);
1462        assert_eq!(mp.row_chunk_target_bytes, policy.row_chunk_target_bytes);
1463    }
1464
1465    // ── MemoryGovernor ledger ────────────────────────────────────────────────
1466
1467    #[test]
1468    fn reservations_account_and_release_on_drop() {
1469        let governor = test_governor(1_000);
1470        assert_eq!(governor.remaining_bytes(), 1_000);
1471        let first = governor.try_reserve(600, "test-first").expect("fits");
1472        assert_eq!(governor.reserved_bytes(), 600);
1473        assert_eq!(governor.remaining_bytes(), 400);
1474        assert_eq!(first.bytes(), 600);
1475        drop(first);
1476        assert_eq!(governor.reserved_bytes(), 0);
1477        assert_eq!(governor.remaining_bytes(), 1_000);
1478    }
1479
1480    #[test]
1481    fn jointly_excessive_reservations_are_refused_with_evidence() {
1482        // Two allocations that each fit alone must not be jointly grantable —
1483        // this is exactly the independent-budgets failure the ledger exists
1484        // to prevent.
1485        let governor = test_governor(1_000);
1486        let availability = governor.availability();
1487        let held = governor.try_reserve(600, "test-held").expect("fits alone");
1488        let refusal = governor
1489            .try_reserve(600, "test-joint")
1490            .expect_err("600 + 600 exceeds the 1000-byte budget");
1491        assert_eq!(
1492            refusal,
1493            MemoryReservationError::BudgetExceeded {
1494                context: "test-joint".into(),
1495                requested_bytes: 600,
1496                reserved_bytes: 600,
1497                budget_bytes: 1_000,
1498                availability,
1499            }
1500        );
1501        // After releasing the holder, the same request succeeds: refusal is a
1502        // routing signal, not a terminal state.
1503        drop(held);
1504        let refreshed = governor
1505            .try_reserve(600, "test-joint")
1506            .expect("fits after release");
1507        assert_eq!(refreshed.bytes(), 600);
1508    }
1509
1510    #[test]
1511    fn dense_reservation_uses_checked_footprint() {
1512        let governor = test_governor(1 << 20);
1513        let ok = governor
1514            .try_reserve_dense_f64(1024, 64, "test-dense")
1515            .expect("512 KiB fits in 1 MiB");
1516        assert_eq!(ok.bytes(), 1024 * 64 * 8);
1517        drop(ok);
1518        // Dimension-product overflow must refuse, never wrap into a tiny
1519        // spurious reservation.
1520        governor
1521            .try_reserve_dense_f64(usize::MAX, 2, "test-overflow")
1522            .expect_err("overflowing footprint cannot be reserved");
1523        assert!(matches!(
1524            governor.try_reserve_dense_f64(usize::MAX, 2, "test-overflow"),
1525            Err(MemoryReservationError::SizeOverflow { .. })
1526        ));
1527    }
1528
1529    #[test]
1530    fn concurrent_reservations_never_oversubscribe() {
1531        let governor = std::sync::Arc::new(test_governor(1_000));
1532        let granted = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
1533        let barrier = std::sync::Arc::new(std::sync::Barrier::new(9));
1534        std::thread::scope(|scope| {
1535            for _ in 0..8 {
1536                let governor = std::sync::Arc::clone(&governor);
1537                let granted = std::sync::Arc::clone(&granted);
1538                let barrier = std::sync::Arc::clone(&barrier);
1539                scope.spawn(move || {
1540                    let held = governor.try_reserve(200, "test-race");
1541                    if held.is_ok() {
1542                        granted.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
1543                    }
1544                    barrier.wait();
1545                    assert!(governor.reserved_bytes() <= governor.budget_bytes());
1546                    barrier.wait();
1547                    drop(held);
1548                });
1549            }
1550            barrier.wait();
1551            assert_eq!(granted.load(std::sync::atomic::Ordering::SeqCst), 5);
1552            assert_eq!(governor.reserved_bytes(), 1_000);
1553            barrier.wait();
1554        });
1555        assert_eq!(governor.reserved_bytes(), 0);
1556    }
1557
1558    #[test]
1559    fn global_policy_caps_are_one_shared_admission_ceiling() {
1560        let governor = MemoryGovernor::global();
1561        assert_eq!(
1562            governor.single_materialization_cap_bytes(),
1563            governor_materialization_cap_from_availability(&governor.availability())
1564        );
1565        // Both ceilings are denominated in the same stationary capacity since
1566        // #2702, so they are equal by construction; a divergence here would mean
1567        // one of them had picked up a second base quantity again.
1568        assert_eq!(
1569            governor.single_materialization_cap_bytes(),
1570            governor.budget_bytes(),
1571            "the routing cap and the ledger budget are one ceiling asked two questions"
1572        );
1573        let policy = ResourcePolicy::default_library();
1574        assert_eq!(
1575            policy.max_single_materialization_bytes,
1576            governor.single_materialization_cap_bytes()
1577        );
1578        let strict = ResourcePolicy::analytic_operator_required();
1579        assert_eq!(
1580            strict.max_single_materialization_bytes,
1581            policy.max_single_materialization_bytes
1582        );
1583    }
1584
1585    #[test]
1586    fn governed_value_holds_and_releases_its_charge() {
1587        let governor = test_governor(64);
1588        let governed = governor
1589            .try_reserve(32, "governed-value")
1590            .expect("reservation fits")
1591            .bind(vec![0_u8; 32]);
1592        assert_eq!(governed.len(), 32);
1593        assert_eq!(governed.reserved_bytes(), 32);
1594        assert_eq!(governor.reserved_bytes(), 32);
1595        drop(governed);
1596        assert_eq!(governor.reserved_bytes(), 0);
1597    }
1598
1599    /// The `available_bytes` assertions here are the availability derivation.
1600    /// The budget assertions alongside them are deliberately denominated in
1601    /// CAPACITY (#2702): free space decides whether an allocation fits *now*,
1602    /// which is the ledger's live `reserved_bytes` question, not the size of the
1603    /// ceiling the ledger measures against.
1604    #[test]
1605    fn memory_availability_distinguishes_host_cgroup_and_exhaustion() {
1606        let host_only =
1607            MemoryAvailability::from_observation(1_000, 4_000, CgroupMemoryObservation::NotPresent);
1608        assert_eq!(host_only.available_bytes(), 1_000);
1609        assert_eq!(host_only.limiting_source(), MemoryAvailabilitySource::Host);
1610        assert_eq!(host_only.capacity_bytes(), 4_000);
1611        assert_eq!(governor_budget_from_availability(&host_only), 3_000);
1612
1613        // Free space at zero with the box's capacity unchanged: the budget is
1614        // the same 3/4 of 4,000, because a host that is momentarily full has not
1615        // become a smaller host.
1616        let exhausted_host =
1617            MemoryAvailability::from_observation(0, 4_000, CgroupMemoryObservation::NotPresent);
1618        assert_eq!(exhausted_host.available_bytes(), 0);
1619        assert_eq!(governor_budget_from_availability(&exhausted_host), 3_000);
1620
1621        let finite_cgroup = MemoryAvailability::from_observation(
1622            1_000,
1623            4_000_000_000,
1624            CgroupMemoryObservation::V2Limited(CgroupMemoryAvailability::fixture(
1625                "/fixture/leaf",
1626                600,
1627                200,
1628                0,
1629                1,
1630            )),
1631        );
1632        assert_eq!(finite_cgroup.available_bytes(), 400);
1633        assert_eq!(
1634            finite_cgroup.limiting_source(),
1635            MemoryAvailabilitySource::Cgroup
1636        );
1637        // Capacity is the cgroup's 600-byte hard limit, not its 400 free bytes.
1638        assert_eq!(finite_cgroup.capacity_bytes(), 600);
1639        assert_eq!(governor_budget_from_availability(&finite_cgroup), 450);
1640
1641        let exhausted_cgroup = MemoryAvailability::from_observation(
1642            1_000,
1643            4_000_000_000,
1644            CgroupMemoryObservation::V2Limited(CgroupMemoryAvailability::fixture(
1645                "/fixture/leaf",
1646                600,
1647                600,
1648                0,
1649                1,
1650            )),
1651        );
1652        assert_eq!(exhausted_cgroup.available_bytes(), 0);
1653        assert_eq!(
1654            exhausted_cgroup.limiting_source(),
1655            MemoryAvailabilitySource::Cgroup
1656        );
1657        // Same cgroup, charged to its limit: same ceiling, so same budget as the
1658        // headroom-having reading above. That equality IS the #2702 property.
1659        assert_eq!(exhausted_cgroup.capacity_bytes(), 600);
1660        assert_eq!(governor_budget_from_availability(&exhausted_cgroup), 450);
1661
1662        // A finite cgroup with more headroom than the host remains visible as
1663        // provenance, but the host is the binding observation.
1664        let host_is_tighter = MemoryAvailability::from_observation(
1665            1_000,
1666            4_000_000_000,
1667            CgroupMemoryObservation::V2Limited(CgroupMemoryAvailability::fixture(
1668                "/fixture/leaf",
1669                8_000,
1670                2_000,
1671                0,
1672                1,
1673            )),
1674        );
1675        assert_eq!(host_is_tighter.available_bytes(), 1_000);
1676        assert_eq!(
1677            host_is_tighter.limiting_source(),
1678            MemoryAvailabilitySource::Host
1679        );
1680
1681        let equal_cgroup_ceiling = MemoryAvailability::from_observation(
1682            1_000,
1683            4_000_000_000,
1684            CgroupMemoryObservation::V2Limited(CgroupMemoryAvailability::fixture(
1685                "/fixture/leaf",
1686                1_200,
1687                200,
1688                0,
1689                1,
1690            )),
1691        );
1692        assert_eq!(equal_cgroup_ceiling.available_bytes(), 1_000);
1693        assert_eq!(
1694            equal_cgroup_ceiling.limiting_source(),
1695            MemoryAvailabilitySource::HostAndCgroup
1696        );
1697
1698        let both_exhausted = MemoryAvailability::from_observation(
1699            0,
1700            4_000_000_000,
1701            CgroupMemoryObservation::V2Limited(CgroupMemoryAvailability::fixture(
1702                "/fixture/leaf",
1703                600,
1704                600,
1705                0,
1706                1,
1707            )),
1708        );
1709        assert_eq!(
1710            both_exhausted.limiting_source(),
1711            MemoryAvailabilitySource::HostAndCgroup
1712        );
1713
1714        // A numeric cgroup-v2 `memory.max = 0` is distinct from the literal
1715        // `max` token and remains an authoritative hard-zero ceiling.
1716        let zero_ceiling = MemoryAvailability::from_observation(
1717            8_000,
1718            4_000_000_000,
1719            CgroupMemoryObservation::V2Limited(CgroupMemoryAvailability::fixture(
1720                "/fixture/leaf",
1721                0,
1722                0,
1723                0,
1724                1,
1725            )),
1726        );
1727        assert_eq!(zero_ceiling.available_bytes(), 0);
1728        assert_eq!(
1729            zero_ceiling.limiting_source(),
1730            MemoryAvailabilitySource::Cgroup
1731        );
1732        assert_eq!(governor_budget_from_availability(&zero_ceiling), 0);
1733    }
1734
1735    #[test]
1736    fn literal_unlimited_cgroup_defers_to_host_available_memory_2317() {
1737        let unlimited = MemoryAvailability::from_observation(
1738            2_430_926_848,
1739            4_000_000_000,
1740            CgroupMemoryObservation::V2Unbounded {
1741                cgroup_path: "/fixture/leaf".into(),
1742                inspected_levels: 3,
1743            },
1744        );
1745        assert_eq!(unlimited.available_bytes(), 2_430_926_848);
1746        assert_eq!(unlimited.limiting_source(), MemoryAvailabilitySource::Host);
1747        // An unbounded controller defers to the host for capacity too, so the
1748        // budget is 3/4 of the host's 4 GB total (#2702).
1749        assert_eq!(unlimited.capacity_bytes(), 4_000_000_000);
1750        assert_eq!(governor_budget_from_availability(&unlimited), 3_000_000_000);
1751        assert!(format!("{unlimited}").contains("unbounded cgroup-v2"));
1752    }
1753
1754    #[test]
1755    fn finite_cgroup_v1_headroom_participates_in_the_same_exact_minimum() {
1756        let availability = MemoryAvailability::from_observation(
1757            8_000,
1758            4_000_000_000,
1759            CgroupMemoryObservation::V1Limited(CgroupMemoryAvailability::fixture(
1760                "/sys/fs/cgroup/memory/slurm/job",
1761                4_000,
1762                1_500,
1763                500,
1764                4,
1765            )),
1766        );
1767        assert_eq!(availability.available_bytes(), 3_000);
1768        assert_eq!(
1769            availability.limiting_source(),
1770            MemoryAvailabilitySource::Cgroup
1771        );
1772        // The v1 controller's 4,000-byte hard limit is the capacity, and the
1773        // budget is 3/4 of it regardless of the 1,500 bytes charged to it.
1774        assert_eq!(availability.capacity_bytes(), 4_000);
1775        assert_eq!(governor_budget_from_availability(&availability), 3_000);
1776        let evidence = format!("{availability}");
1777        assert!(evidence.contains("cgroup-v1"));
1778        assert!(evidence.contains("available=3000"));
1779    }
1780
1781    /// gam#2684. A cgroup pinned at its hard limit reports almost no *available*
1782    /// memory while its *capacity* is unchanged. The numbers below are the ones
1783    /// measured on an MSI compute node (`--mem=6g`, host with 448 GB free):
1784    /// the shipped probe reported `available=53_248` while the job's ceiling was
1785    /// still 6 GiB. The routing cap must come from the ceiling — otherwise a
1786    /// 300x12 f64 design (28,800 bytes) is refused.
1787    ///
1788    /// #2702 extended that to the ledger budget, which this test used to assert
1789    /// shrank to 39,936 bytes here: the observation is taken once per process
1790    /// and never refreshed, so a free-denominated budget did not track
1791    /// exhaustion — it froze one process's view of the schedule. A large
1792    /// reservation still routes instead of allocating, but because the ledger
1793    /// accounts THIS process's live footprint, not because the ceiling moved.
1794    #[test]
1795    fn a_cgroup_at_its_limit_moves_neither_the_budget_nor_the_materialization_cap_2684_2702() {
1796        const DESIGN_300X12_BYTES: usize = 300 * 12 * 8;
1797        let limit_bytes = 6 * 1024 * 1024 * 1024_u64;
1798        let at_the_limit = MemoryAvailability::from_observation(
1799            448_648_040_448,
1800            527_799_400 * 1024,
1801            CgroupMemoryObservation::V1Limited(CgroupMemoryAvailability::fixture(
1802                "/sys/fs/cgroup/memory/slurm/uid_81060/job_14615476",
1803                limit_bytes,
1804                limit_bytes - 53_248,
1805                0,
1806                6,
1807            )),
1808        );
1809        assert_eq!(at_the_limit.available_bytes(), 53_248);
1810        assert_eq!(at_the_limit.capacity_bytes(), limit_bytes);
1811        let governor = MemoryGovernor::with_detected_availability(at_the_limit);
1812        // The budget is the job's ceiling, not the 53,248 bytes that happened to
1813        // be free when the probe ran.
1814        assert_eq!(governor.budget_bytes(), (limit_bytes as usize) / 4 * 3);
1815        // It still refuses what cannot fit: 8 GiB does not fit a 6 GiB job.
1816        assert!(governor.try_reserve(8 << 30, "larger-than-the-job").is_err());
1817        // And the few-kilobyte reservation #2702 was filed for is admitted at
1818        // exactly the reading that refused it — the SE path asks for two dense
1819        // f64 copies of one column of a p=64 transformed Hessian.
1820        let se_chunk = governor
1821            .try_reserve_dense_f64_copies(64, 1, 2, "coefficient-SE solve chunk")
1822            .expect("a 1,024-byte SE chunk must be admissible in a 6 GiB job");
1823        assert_eq!(se_chunk.bytes(), 1_024);
1824        drop(se_chunk);
1825        // The routing ceiling does not move either.
1826        assert_eq!(
1827            governor.single_materialization_cap_bytes(),
1828            (limit_bytes as usize) / 4 * 3
1829        );
1830        assert!(governor.single_materialization_cap_bytes() > DESIGN_300X12_BYTES);
1831
1832        // Falsification guard: the cap must still be able to say no. A cgroup
1833        // whose whole ceiling is smaller than the design refuses it, so the
1834        // assertion above cannot be satisfied by a cap that always admits.
1835        let tiny_ceiling = MemoryAvailability::from_observation(
1836            448_648_040_448,
1837            527_799_400 * 1024,
1838            CgroupMemoryObservation::V1Limited(CgroupMemoryAvailability::fixture(
1839                "/fixture/tiny",
1840                1_024,
1841                0,
1842                0,
1843                1,
1844            )),
1845        );
1846        assert_eq!(tiny_ceiling.capacity_bytes(), 1_024);
1847        assert_eq!(
1848            MemoryGovernor::with_detected_availability(tiny_ceiling)
1849                .single_materialization_cap_bytes(),
1850            768
1851        );
1852
1853        // Second falsification guard: an idle cgroup with the same ceiling must
1854        // report the same cap, or the cap is still tracking the schedule.
1855        let idle = MemoryAvailability::from_observation(
1856            448_648_040_448,
1857            527_799_400 * 1024,
1858            CgroupMemoryObservation::V1Limited(CgroupMemoryAvailability::fixture(
1859                "/sys/fs/cgroup/memory/slurm/uid_81060/job_14615476",
1860                limit_bytes,
1861                92_827_648,
1862                28_672,
1863                6,
1864            )),
1865        );
1866        assert!(idle.available_bytes() > 6_000_000_000);
1867        assert_eq!(
1868            MemoryGovernor::with_detected_availability(idle).single_materialization_cap_bytes(),
1869            (limit_bytes as usize) / 4 * 3
1870        );
1871    }
1872
1873    #[test]
1874    fn malformed_active_cgroup_fails_closed_with_typed_evidence() {
1875        let availability = MemoryAvailability::from_observation(
1876            8_000,
1877            4_000_000_000,
1878            CgroupMemoryObservation::ProbeFailed(CgroupMemoryProbeFailure::fixture(
1879                CgroupMemoryProbeFailureKind::InvalidCounter,
1880                "/fixture/leaf/memory.current",
1881                "expected an unsigned byte count",
1882            )),
1883        );
1884        assert_eq!(availability.available_bytes(), 0);
1885        // Fail-closed covers capacity too: an unreadable controller must not
1886        // hand out a routing ceiling either (#2684 keeps this policy intact).
1887        assert_eq!(availability.capacity_bytes(), 0);
1888        assert_eq!(
1889            availability.limiting_source(),
1890            MemoryAvailabilitySource::CgroupProbeFailure
1891        );
1892        assert_eq!(governor_budget_from_availability(&availability), 0);
1893        assert_eq!(
1894            governor_materialization_cap_from_availability(&availability),
1895            0
1896        );
1897        let evidence = format!("{availability}");
1898        assert!(evidence.contains("failed closed"));
1899        assert!(evidence.contains("invalid-counter"));
1900    }
1901
1902    #[test]
1903    fn compressed_macos_observation_keeps_xnu_available_memory_positive() {
1904        // #2316's healthy 8 GiB macOS host had more compressed than
1905        // free+inactive pages. sysinfo 0.33 subtracted compressor pages and
1906        // saturated to zero; 0.38 follows XNU and reports
1907        // (active + inactive + free) * page_size.
1908        let xnu_available = (75_514_u64 + 69_056 + 3_802) * 16_384;
1909        assert_eq!(xnu_available, 2_430_926_848);
1910        let availability = MemoryAvailability::from_observation(
1911            xnu_available,
1912            8 * 1024 * 1024 * 1024,
1913            CgroupMemoryObservation::NotPresent,
1914        );
1915        assert_eq!(availability.available_bytes(), xnu_available);
1916        // Capacity is the host's 8 GiB, so the budget is 6 GiB (#2702); the
1917        // 2.43 GB free figure decides nothing but the availability provenance.
1918        assert_eq!(availability.capacity_bytes(), 8 * 1024 * 1024 * 1024);
1919        assert_eq!(
1920            governor_budget_from_availability(&availability),
1921            6 * 1024 * 1024 * 1024
1922        );
1923    }
1924}
1925
1926/// gam#2702: a reservation verdict is a function of the request and this
1927/// process's own live footprint — never of what else was on the box, and never
1928/// of what this process happened to do earlier.
1929///
1930/// The filed incident: three `gam` inference tests passed in a 71-test run of
1931/// the `inference` binary and failed in a 6-test subset of that same binary at
1932/// that same commit, all three with
1933/// `resource policy refused exact coefficient-SE columns 0..1`. The refused
1934/// allocation was two dense f64 copies of one column — kilobytes. The ledger
1935/// budget was `3/4 x FREE memory` sampled once, at whichever moment this
1936/// process first touched the governor, so processes launched from one job
1937/// cgroup derived different budgets according to how much page cache and how
1938/// many sibling test processes were charged to that cgroup at their own instant
1939/// of first touch.
1940///
1941/// #2684 had already removed exactly this mechanism from the materialization
1942/// ceiling. What follows asserts the ledger is not an exception, on the shipped
1943/// derivation, with the pre-fix arithmetic spelled out as the falsification
1944/// control.
1945#[cfg(test)]
1946mod governor_budget_is_capacity_determined_2702_tests {
1947    use super::*;
1948
1949    /// The MSI compute node the incident was measured on: a `--mem=8g` job on a
1950    /// box with hundreds of GB free.
1951    const HOST_AVAILABLE_BYTES: u64 = 448_648_040_448;
1952    const HOST_TOTAL_BYTES: u64 = 527_799_400 * 1024;
1953    const JOB_LIMIT_BYTES: u64 = 8 * 1024 * 1024 * 1024;
1954
1955    /// The refused allocation, restated from the failing path
1956    /// (`gam-solve` `optimizer.rs`): the exact factorized coefficient-SE solve
1957    /// reserves two dense `p_t x chunk` f64 workspaces at once. `p_t = 512` and
1958    /// `chunk = 1` is the smallest request that path can ever make on a model of
1959    /// this width — one column at a time.
1960    const SE_TRANSFORMED_ROWS: usize = 512;
1961    /// One column, two copies, eight bytes per f64.
1962    const SE_CHUNK_BYTES: usize = SE_TRANSFORMED_ROWS * 8 * 2;
1963
1964    /// One cgroup read at load `charged`, everything else held fixed.
1965    fn one_job_cgroup_at_load(charged: u64) -> MemoryAvailability {
1966        crate::test_support::simulated_cgroup_memory_environment(
1967            HOST_AVAILABLE_BYTES,
1968            HOST_TOTAL_BYTES,
1969            JOB_LIMIT_BYTES,
1970            charged,
1971        )
1972    }
1973
1974    /// The derivation this fix replaced, kept here as the control: 3/4 of FREE
1975    /// memory. A test that never evaluates it cannot show that the assertions
1976    /// below had a way to fail.
1977    fn pre_2702_free_denominated_budget(availability: &MemoryAvailability) -> usize {
1978        let scaled = u128::from(availability.available_bytes()) * GOVERNOR_BUDGET_NUMERATOR
1979            / GOVERNOR_BUDGET_DENOMINATOR;
1980        usize::try_from(scaled).unwrap_or(usize::MAX)
1981    }
1982
1983    #[test]
1984    fn one_job_observed_at_two_load_levels_yields_one_budget_and_one_verdict() {
1985        // The two readings the incident's two runs took: a cgroup with room, and
1986        // the same cgroup four kilobytes from its limit — the state any cgroup
1987        // settles into once its job has read a few gigabytes of build artifacts.
1988        let roomy = one_job_cgroup_at_load(92_827_648);
1989        let pinned = one_job_cgroup_at_load(JOB_LIMIT_BYTES - 4_096);
1990
1991        // The arms are genuinely different observations, or nothing below is
1992        // being tested.
1993        assert!(roomy.available_bytes() > 7_000_000_000);
1994        assert_eq!(pinned.available_bytes(), 4_096);
1995        // ... and they are genuinely the same job.
1996        assert_eq!(roomy.capacity_bytes(), JOB_LIMIT_BYTES);
1997        assert_eq!(pinned.capacity_bytes(), JOB_LIMIT_BYTES);
1998
1999        // FALSIFICATION CONTROL. Under the pre-#2702 derivation the pinned arm's
2000        // budget was 3,072 bytes: below the SE chunk, so that arm refused and the
2001        // roomy arm admitted. The assertions that follow therefore had a way to
2002        // fail, and this is the exact inversion the issue reported.
2003        assert_eq!(pre_2702_free_denominated_budget(&pinned), 3_072);
2004        assert!(pre_2702_free_denominated_budget(&pinned) < SE_CHUNK_BYTES);
2005        assert!(pre_2702_free_denominated_budget(&roomy) > SE_CHUNK_BYTES);
2006
2007        let roomy_governor = MemoryGovernor::with_detected_availability(roomy);
2008        let pinned_governor = MemoryGovernor::with_detected_availability(pinned);
2009
2010        assert_eq!(roomy_governor.budget_bytes(), pinned_governor.budget_bytes());
2011        assert_eq!(
2012            pinned_governor.budget_bytes(),
2013            (JOB_LIMIT_BYTES as usize) / 4 * 3
2014        );
2015        assert_eq!(
2016            roomy_governor.remaining_bytes(),
2017            pinned_governor.remaining_bytes()
2018        );
2019
2020        // The verdict itself, taken through the shipped reservation call on both
2021        // arms: the same request, the same answer.
2022        for governor in [&roomy_governor, &pinned_governor] {
2023            let reservation = governor
2024                .try_reserve_dense_f64_copies(
2025                    SE_TRANSFORMED_ROWS,
2026                    1,
2027                    2,
2028                    "factorized coefficient-SE solve chunk",
2029                )
2030                .expect("an 8 KiB SE chunk is admissible in an 8 GiB job at any load");
2031            assert_eq!(reservation.bytes(), SE_CHUNK_BYTES);
2032        }
2033    }
2034
2035    #[test]
2036    fn a_request_larger_than_the_job_is_still_refused_at_every_load() {
2037        // The ceiling must keep saying no, or the test above is satisfied by a
2038        // governor that admits everything.
2039        for charged in [0, JOB_LIMIT_BYTES / 2, JOB_LIMIT_BYTES - 4_096] {
2040            let governor = MemoryGovernor::with_detected_availability(one_job_cgroup_at_load(charged));
2041            let refusal = governor
2042                .try_reserve(16 * 1024 * 1024 * 1024, "twice the job's ceiling")
2043                .expect_err("16 GiB cannot be admitted in an 8 GiB job");
2044            match refusal {
2045                MemoryReservationError::BudgetExceeded { budget_bytes, .. } => {
2046                    assert_eq!(budget_bytes, (JOB_LIMIT_BYTES as usize) / 4 * 3);
2047                }
2048                other => panic!("expected a budget refusal naming the ceiling, got {other:?}"),
2049            }
2050        }
2051    }
2052
2053    #[test]
2054    fn the_verdict_does_not_depend_on_what_this_process_did_earlier() {
2055        // The property the issue asks for, stated as history-independence: a
2056        // request's verdict must be the same whether or not the process has
2057        // already built and dropped something large. Reservations are the only
2058        // process state the ledger has, and a released one must leave no trace.
2059        let governor = MemoryGovernor::with_detected_availability(one_job_cgroup_at_load(0));
2060        let request = || {
2061            governor
2062                .try_reserve_dense_f64_copies(
2063                    SE_TRANSFORMED_ROWS,
2064                    1,
2065                    2,
2066                    "factorized coefficient-SE solve chunk",
2067                )
2068                .map(|reservation| reservation.bytes())
2069        };
2070
2071        let before = request().expect("admissible on a fresh ledger");
2072        {
2073            // Something big enough that a free-denominated ledger would have
2074            // been left visibly poorer by it: all but one SE chunk of the budget.
2075            let bulk = governor
2076                .try_reserve(
2077                    governor.remaining_bytes() - SE_CHUNK_BYTES,
2078                    "prior work in this process",
2079                )
2080                .expect("the bulk reservation is exactly the remaining budget");
2081            assert_eq!(governor.remaining_bytes(), SE_CHUNK_BYTES);
2082            // While it is live the ledger honestly reports the pressure: this is
2083            // the one thing that MAY change a verdict, and the next request still
2084            // fits because the bulk left exactly one chunk.
2085            assert!(request().is_ok());
2086            drop(bulk);
2087        }
2088        assert_eq!(governor.reserved_bytes(), 0);
2089        assert_eq!(request().expect("admissible again"), before);
2090    }
2091}