Skip to main content

hyperi_rustlib/memory/
guard.rs

1// Project:   hyperi-rustlib
2// File:      src/memory/guard.rs
3// Purpose:   Memory guard with backpressure signals
4// Language:  Rust
5//
6// License:   BUSL-1.1
7// Copyright: (c) 2026 HYPERI PTY LIMITED
8
9//! Memory guard with backpressure signals.
10
11use std::sync::OnceLock;
12use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
13
14use super::cgroup;
15
16/// Process-wide total-heap byte source, set once at startup.
17///
18/// See [`set_heap_source`] for the rationale. Allocator-agnostic: any
19/// `fn() -> usize` returning live heap bytes (e.g. `cap::Cap::allocated`,
20/// jemalloc `stats.allocated`).
21static HEAP_SOURCE: OnceLock<fn() -> usize> = OnceLock::new();
22
23/// Register a process-wide source of total live-heap bytes.
24///
25/// When set, every [`MemoryGuard`] switches its read path
26/// ([`current_bytes`](MemoryGuard::current_bytes), pressure checks, and
27/// [`try_reserve`](MemoryGuard::try_reserve) admission) from the per-batch
28/// reservation counter to this source -- a cheap, accurate, *total-process*
29/// heap figure that also catches growth the per-batch reservations never see
30/// (e.g. a transform ballooning a `Vec`).
31///
32/// **Why a global hook and not a dependency:** a tracking allocator must be
33/// the binary's single `#[global_allocator]`, which is the *application's*
34/// choice, not a library's -- and rustlib is `#![forbid(unsafe_code)]`, so it
35/// cannot implement one anyway. The application installs its allocator and
36/// wires it here in a few lines. This keeps rustlib allocator-agnostic with no
37/// allocator dependency in its graph.
38///
39/// The first call wins and returns `true`; later calls are a no-op and return
40/// `false` (the existing source is kept). Call once at startup, before
41/// constructing guards.
42///
43/// The application picks a tracking allocator -- prefer an actively-maintained
44/// one such as `tikv-jemalloc-ctl` (`stats.allocated`); the `cap` crate also
45/// works but is effectively unmaintained (last release 2023).
46///
47/// ```ignore
48/// // In the application binary, using jemalloc:
49/// #[global_allocator]
50/// static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
51///
52/// fn main() {
53///     hyperi_rustlib::memory::set_heap_source(|| {
54///         tikv_jemalloc_ctl::epoch::advance().ok();
55///         tikv_jemalloc_ctl::stats::allocated::read().unwrap_or(0)
56///     });
57///     // ... build ServiceRuntime / MemoryGuard ...
58/// }
59/// ```
60#[must_use]
61pub fn set_heap_source(source: fn() -> usize) -> bool {
62    HEAP_SOURCE.set(source).is_ok()
63}
64
65/// Read the registered total-heap source, if any.
66#[inline]
67fn heap_bytes() -> Option<u64> {
68    HEAP_SOURCE.get().map(|f| f() as u64)
69}
70
71/// Read an env var `{PREFIX}_{SUFFIX}` and parse it.
72fn env_parsed<T: std::str::FromStr>(prefix: &str, suffix: &str) -> Option<T> {
73    std::env::var(format!("{prefix}_{suffix}"))
74        .ok()
75        .and_then(|v| v.parse().ok())
76}
77
78/// Memory pressure levels.
79#[derive(Debug, Clone, Copy, PartialEq, Eq)]
80pub enum MemoryPressure {
81    /// Usage below 50% of limit.
82    Low,
83    /// Usage between 50% and pressure_threshold.
84    Medium,
85    /// Usage above pressure_threshold -- apply backpressure.
86    High,
87}
88
89/// Configuration for `MemoryGuard`.
90///
91/// When the `config` feature is enabled, this can be loaded from the config
92/// cascade under the `memory` key:
93///
94/// ```yaml
95/// memory:
96///   limit_bytes: 0           # 0 = auto-detect from cgroup/system
97///   pressure_threshold: 0.80 # backpressure at 80% of effective limit
98///   cgroup_headroom: 0.85    # use 85% of cgroup limit
99/// ```
100#[derive(Debug, Clone, serde::Deserialize, serde::Serialize)]
101pub struct MemoryGuardConfig {
102    /// Explicit memory limit in bytes. 0 = auto-detect from cgroup/system.
103    #[serde(default)]
104    pub limit_bytes: u64,
105    /// Fraction of limit at which backpressure activates (default 0.8).
106    #[serde(default = "default_pressure_threshold")]
107    pub pressure_threshold: f64,
108    /// Fraction of cgroup limit to use as the effective limit (default 0.85).
109    /// Leaves headroom for the process itself (stack, code, etc.).
110    #[serde(default = "default_cgroup_headroom")]
111    pub cgroup_headroom: f64,
112}
113
114fn default_pressure_threshold() -> f64 {
115    DEFAULT_PRESSURE_THRESHOLD
116}
117
118fn default_cgroup_headroom() -> f64 {
119    DEFAULT_CGROUP_HEADROOM
120}
121
122/// A fraction is valid iff it is finite and within `(0.0, 1.0]`.
123fn check_fraction(v: f64, name: &str) -> Result<(), String> {
124    if !v.is_finite() || v <= 0.0 || v > 1.0 {
125        return Err(format!(
126            "memory.{name} must be a finite fraction in (0.0, 1.0], got {v}"
127        ));
128    }
129    Ok(())
130}
131
132/// Return `v` if it is a valid fraction, else log an error and substitute
133/// `default`. Defensive guard so a bad config cannot produce a zero/`NaN`
134/// limit and a divide-by-zero pressure ratio.
135fn sane_fraction(v: f64, default: f64, name: &str) -> f64 {
136    if check_fraction(v, name).is_err() {
137        tracing::error!(
138            value = v,
139            "invalid memory.{name} (need finite fraction in (0,1]); using default {default}"
140        );
141        default
142    } else {
143        v
144    }
145}
146
147/// Effective auto-detected limit: cgroup limit * headroom, capped at the soft
148/// throttle (`memory.high`) when that is set lower.
149///
150/// The kernel reclaims hard and throttles allocations at `memory.high` (before
151/// the `memory.max` OOM-kill), so admitting past it courts a latency cliff.
152/// Pure, so the cap logic is unit-testable without touching the real cgroup
153/// files.
154#[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
155fn effective_auto_limit(detected: u64, headroom: f64, high: Option<u64>) -> u64 {
156    let headroom_limit = (detected as f64 * headroom) as u64;
157    match high {
158        Some(h) => headroom_limit.min(h),
159        None => headroom_limit,
160    }
161}
162
163/// Default cgroup headroom: use 85% of cgroup limit.
164///
165/// Rationale: Rust has no GC so no spike headroom needed (unlike JVM 75% / Go 80%).
166/// 15% headroom covers jemalloc fragmentation, kernel overhead, and page cache.
167const DEFAULT_CGROUP_HEADROOM: f64 = 0.85;
168
169/// Default pressure threshold: backpressure at 80% of effective limit.
170///
171/// With 85% headroom, backpressure activates at ~68% of actual cgroup limit.
172/// Matches OTel Collector's `limit_percentage: 80` philosophy.
173const DEFAULT_PRESSURE_THRESHOLD: f64 = 0.80;
174
175impl Default for MemoryGuardConfig {
176    fn default() -> Self {
177        Self {
178            limit_bytes: 0, // auto-detect
179            pressure_threshold: DEFAULT_PRESSURE_THRESHOLD,
180            cgroup_headroom: DEFAULT_CGROUP_HEADROOM,
181        }
182    }
183}
184
185impl MemoryGuardConfig {
186    /// Load from the config cascade, falling back to defaults.
187    ///
188    /// When the `config` feature is enabled and `config::setup()` has been
189    /// called, reads the `memory` key from the cascade. Otherwise returns
190    /// [`MemoryGuardConfig::default()`].
191    #[must_use]
192    pub fn from_cascade() -> Self {
193        #[cfg(feature = "config")]
194        {
195            if let Some(cfg) = crate::config::try_get()
196                && let Ok(memory) = cfg.unmarshal_key_registered::<Self>("memory")
197            {
198                return memory;
199            }
200        }
201        Self::default()
202    }
203
204    /// Create config from environment variables with a prefix.
205    ///
206    /// Reads standard env vars for memory configuration:
207    /// - `{PREFIX}_MEMORY_LIMIT_BYTES` -- explicit limit (0 or unset = auto-detect from cgroup)
208    /// - `{PREFIX}_MEMORY_PRESSURE_THRESHOLD` -- backpressure trigger (default 0.80)
209    /// - `{PREFIX}_MEMORY_CGROUP_HEADROOM` -- fraction of cgroup limit to use (default 0.85)
210    ///
211    /// # Example
212    ///
213    /// ```bash
214    /// DFE_MEMORY_LIMIT_BYTES=4294967296      # 4 GiB explicit
215    /// DFE_MEMORY_PRESSURE_THRESHOLD=0.75     # backpressure at 75%
216    /// DFE_MEMORY_CGROUP_HEADROOM=0.90        # use 90% of cgroup
217    /// ```
218    ///
219    /// ```rust,no_run
220    /// use hyperi_rustlib::memory::MemoryGuardConfig;
221    /// let config = MemoryGuardConfig::from_env("DFE");
222    /// ```
223    #[must_use]
224    #[cfg(feature = "config")]
225    pub fn from_env(prefix: &str) -> Self {
226        use crate::config::flat_env::flat_env_parsed;
227
228        let mut config = Self::default();
229
230        if let Some(v) = flat_env_parsed::<u64>(prefix, "MEMORY_LIMIT_BYTES") {
231            config.limit_bytes = v;
232        }
233        if let Some(v) = flat_env_parsed::<f64>(prefix, "MEMORY_PRESSURE_THRESHOLD") {
234            config.pressure_threshold = v;
235        }
236        if let Some(v) = flat_env_parsed::<f64>(prefix, "MEMORY_CGROUP_HEADROOM") {
237            config.cgroup_headroom = v;
238        }
239
240        config
241    }
242
243    /// Create config from environment variables without requiring `config` feature.
244    ///
245    /// Same as [`from_env`](Self::from_env) but uses `std::env` directly.
246    #[must_use]
247    pub fn from_env_raw(prefix: &str) -> Self {
248        let mut config = Self::default();
249
250        if let Some(v) = env_parsed::<u64>(prefix, "MEMORY_LIMIT_BYTES") {
251            config.limit_bytes = v;
252        }
253        if let Some(v) = env_parsed::<f64>(prefix, "MEMORY_PRESSURE_THRESHOLD") {
254            config.pressure_threshold = v;
255        }
256        if let Some(v) = env_parsed::<f64>(prefix, "MEMORY_CGROUP_HEADROOM") {
257            config.cgroup_headroom = v;
258        }
259
260        config
261    }
262
263    /// Validate the config, returning an error describing the first invalid
264    /// field. `pressure_threshold` and `cgroup_headroom` must each be a finite
265    /// fraction in `(0.0, 1.0]`. Call this at startup to fail fast on bad
266    /// config rather than relying on [`MemoryGuard::new`]'s defensive clamping.
267    ///
268    /// # Errors
269    ///
270    /// Returns `Err` with a human-readable message if a fraction field is
271    /// non-finite, `<= 0.0`, or `> 1.0`.
272    pub fn validate(&self) -> Result<(), String> {
273        check_fraction(self.pressure_threshold, "pressure_threshold")?;
274        check_fraction(self.cgroup_headroom, "cgroup_headroom")?;
275        Ok(())
276    }
277}
278
279/// Cgroup-aware memory tracking with backpressure signals.
280///
281/// Tracks application-level memory usage (not process RSS) and provides
282/// fast atomic checks for the hot path. Designed for data pipeline services
283/// where incoming data must be rejected (503) before hitting the container
284/// memory limit.
285///
286/// # Usage
287///
288/// ```rust,no_run
289/// use hyperi_rustlib::memory::{MemoryGuard, MemoryGuardConfig};
290///
291/// let guard = MemoryGuard::new(MemoryGuardConfig::default());
292///
293/// // On data arrival -- check before accepting
294/// let payload_len = 1024u64;
295/// if !guard.try_reserve(payload_len) {
296///     // return 503 -- backpressure
297/// }
298///
299/// // After data is flushed/sent
300/// guard.release(payload_len);
301///
302/// // Fast hot-path check
303/// if guard.under_pressure() {
304///     // return 503
305/// }
306/// ```
307pub struct MemoryGuard {
308    /// Current tracked bytes (application-level, not RSS).
309    current_bytes: AtomicU64,
310    /// Effective memory limit in bytes.
311    limit_bytes: u64,
312    /// Pressure threshold (0.0-1.0).
313    pressure_threshold: f64,
314    /// Fast boolean for hot-path pressure check.
315    under_pressure: AtomicBool,
316}
317
318impl MemoryGuard {
319    /// Create a new memory guard.
320    ///
321    /// If `config.limit_bytes` is 0, auto-detects from cgroup (K8s) or system memory,
322    /// then applies `cgroup_headroom` factor to leave room for process overhead.
323    #[must_use]
324    #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
325    pub fn new(config: MemoryGuardConfig) -> Self {
326        // Defensive: a non-finite / out-of-range threshold or headroom would
327        // produce a zero/NaN limit and a divide-by-zero pressure ratio. Clamp
328        // to the safe default and log loudly. Callers wanting hard rejection
329        // should call `config.validate()` at startup.
330        let pressure_threshold = sane_fraction(
331            config.pressure_threshold,
332            DEFAULT_PRESSURE_THRESHOLD,
333            "pressure_threshold",
334        );
335        let cgroup_headroom = sane_fraction(
336            config.cgroup_headroom,
337            DEFAULT_CGROUP_HEADROOM,
338            "cgroup_headroom",
339        );
340
341        let raw_limit = if config.limit_bytes > 0 {
342            config.limit_bytes
343        } else {
344            effective_auto_limit(
345                cgroup::detect_memory_limit(),
346                cgroup_headroom,
347                cgroup::detect_memory_high(),
348            )
349        };
350        // Never permit a zero effective limit: every pressure calculation
351        // divides by it.
352        let limit_bytes = raw_limit.max(1);
353
354        tracing::info!(limit_bytes, pressure_threshold, "memory guard initialised");
355
356        Self {
357            current_bytes: AtomicU64::new(0),
358            limit_bytes,
359            pressure_threshold,
360            under_pressure: AtomicBool::new(false),
361        }
362    }
363
364    /// Try to reserve bytes. Returns false if over the limit (backpressure).
365    ///
366    /// With a registered [`set_heap_source`], this is a projected-admission
367    /// check against the *true total heap* (`heap() + bytes <= limit`) and does
368    /// NOT mutate the reservation counter -- the allocator already accounts the
369    /// bytes once they are allocated, and frees them on drop, so no `release`
370    /// is needed. Without a source it is the classic atomic check-and-add on
371    /// the per-batch counter (rolled back if it would exceed the limit).
372    #[inline]
373    pub fn try_reserve(&self, bytes: u64) -> bool {
374        if let Some(heap) = heap_bytes() {
375            return heap + bytes <= self.limit_bytes;
376        }
377        let current = self.current_bytes.fetch_add(bytes, Ordering::Relaxed) + bytes;
378        if current > self.limit_bytes {
379            // Over limit -- roll back
380            self.current_bytes.fetch_sub(bytes, Ordering::Relaxed);
381            self.under_pressure.store(true, Ordering::Relaxed);
382            return false;
383        }
384        self.update_pressure(current);
385        true
386    }
387
388    /// Add bytes without checking the limit (for tracking only).
389    /// Use when data is already accepted and you just need to track it.
390    #[inline]
391    pub fn add_bytes(&self, bytes: u64) {
392        let new_total = self.current_bytes.fetch_add(bytes, Ordering::Relaxed) + bytes;
393        self.update_pressure(new_total);
394    }
395
396    /// Release bytes after data is flushed/sent/dropped.
397    ///
398    /// Uses saturating subtraction to prevent underflow wrapping.
399    #[inline]
400    pub fn release(&self, bytes: u64) {
401        let prev = self
402            .current_bytes
403            .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
404                Some(current.saturating_sub(bytes))
405            })
406            // Always succeeds (closure always returns Some).
407            .unwrap_or_else(|v| v);
408        self.update_pressure(prev.saturating_sub(bytes));
409    }
410
411    /// Fast hot-path pressure check.
412    ///
413    /// With a registered [`set_heap_source`], computes live from the true heap
414    /// (one atomic load + compare); otherwise reads the cached flag maintained
415    /// by `try_reserve`/`add_bytes`/`release`.
416    #[inline]
417    pub fn under_pressure(&self) -> bool {
418        if heap_bytes().is_some() {
419            return self.pressure_ratio() >= self.pressure_threshold;
420        }
421        self.under_pressure.load(Ordering::Relaxed)
422    }
423
424    /// Current pressure level.
425    #[inline]
426    pub fn pressure(&self) -> MemoryPressure {
427        let ratio = self.pressure_ratio();
428        if ratio >= self.pressure_threshold {
429            MemoryPressure::High
430        } else if ratio >= 0.5 {
431            MemoryPressure::Medium
432        } else {
433            MemoryPressure::Low
434        }
435    }
436
437    /// Current usage as fraction of limit (0.0 - 1.0+).
438    #[inline]
439    pub fn pressure_ratio(&self) -> f64 {
440        self.current_bytes() as f64 / self.limit_bytes as f64
441    }
442
443    /// Current memory usage in bytes.
444    ///
445    /// Returns the true total live heap when a [`set_heap_source`] is
446    /// registered, otherwise the sum of outstanding per-batch reservations.
447    #[inline]
448    pub fn current_bytes(&self) -> u64 {
449        heap_bytes().unwrap_or_else(|| self.current_bytes.load(Ordering::Relaxed))
450    }
451
452    /// Configured memory limit in bytes.
453    #[inline]
454    pub fn limit_bytes(&self) -> u64 {
455        self.limit_bytes
456    }
457
458    /// Update the pressure flag based on current usage.
459    #[inline]
460    fn update_pressure(&self, current: u64) {
461        let ratio = current as f64 / self.limit_bytes as f64;
462        self.under_pressure
463            .store(ratio >= self.pressure_threshold, Ordering::Relaxed);
464    }
465}
466
467#[cfg(test)]
468mod tests {
469    use super::*;
470
471    #[test]
472    fn test_memory_guard_default() {
473        let guard = MemoryGuard::new(MemoryGuardConfig {
474            limit_bytes: 1_000_000, // 1MB explicit
475            ..Default::default()
476        });
477        assert_eq!(guard.limit_bytes(), 1_000_000);
478        assert_eq!(guard.current_bytes(), 0);
479        assert!(!guard.under_pressure());
480        assert_eq!(guard.pressure(), MemoryPressure::Low);
481    }
482
483    #[test]
484    fn test_try_reserve_within_limit() {
485        let guard = MemoryGuard::new(MemoryGuardConfig {
486            limit_bytes: 1000,
487            ..Default::default()
488        });
489        assert!(guard.try_reserve(500));
490        assert_eq!(guard.current_bytes(), 500);
491    }
492
493    #[test]
494    fn test_try_reserve_over_limit() {
495        let guard = MemoryGuard::new(MemoryGuardConfig {
496            limit_bytes: 1000,
497            ..Default::default()
498        });
499        assert!(guard.try_reserve(500));
500        assert!(!guard.try_reserve(600)); // would exceed 1000
501        assert_eq!(guard.current_bytes(), 500); // rolled back
502        assert!(guard.under_pressure());
503    }
504
505    #[test]
506    fn test_release_reduces_pressure() {
507        let guard = MemoryGuard::new(MemoryGuardConfig {
508            limit_bytes: 1000,
509            pressure_threshold: 0.8,
510            ..Default::default()
511        });
512        guard.add_bytes(900); // 90% -- over threshold
513        assert!(guard.under_pressure());
514        assert_eq!(guard.pressure(), MemoryPressure::High);
515
516        guard.release(500); // down to 400 = 40%
517        assert!(!guard.under_pressure());
518        assert_eq!(guard.pressure(), MemoryPressure::Low);
519    }
520
521    #[test]
522    fn test_pressure_levels() {
523        let guard = MemoryGuard::new(MemoryGuardConfig {
524            limit_bytes: 1000,
525            pressure_threshold: 0.8,
526            ..Default::default()
527        });
528
529        // Low (< 50%)
530        guard.add_bytes(400);
531        assert_eq!(guard.pressure(), MemoryPressure::Low);
532
533        // Medium (50-80%)
534        guard.add_bytes(200); // 600 = 60%
535        assert_eq!(guard.pressure(), MemoryPressure::Medium);
536
537        // High (>= 80%)
538        guard.add_bytes(300); // 900 = 90%
539        assert_eq!(guard.pressure(), MemoryPressure::High);
540    }
541
542    #[test]
543    fn test_pressure_ratio() {
544        let guard = MemoryGuard::new(MemoryGuardConfig {
545            limit_bytes: 1000,
546            ..Default::default()
547        });
548        guard.add_bytes(250);
549        let ratio = guard.pressure_ratio();
550        assert!((ratio - 0.25).abs() < 0.001);
551    }
552
553    #[test]
554    fn test_release_saturating() {
555        let guard = MemoryGuard::new(MemoryGuardConfig {
556            limit_bytes: 1000,
557            ..Default::default()
558        });
559        guard.add_bytes(100);
560        guard.release(200); // release more than added -- saturates to 0
561        assert_eq!(
562            guard.current_bytes(),
563            0,
564            "over-release must saturate to 0, not wrap"
565        );
566        assert!(!guard.under_pressure());
567        assert_eq!(guard.pressure(), MemoryPressure::Low);
568
569        // Verify the guard is still functional after over-release
570        assert!(guard.try_reserve(500));
571        assert_eq!(guard.current_bytes(), 500);
572    }
573
574    #[test]
575    fn test_concurrent_reserve_release() {
576        use std::sync::Arc;
577        use std::thread;
578
579        let guard = Arc::new(MemoryGuard::new(MemoryGuardConfig {
580            limit_bytes: 100_000,
581            pressure_threshold: 0.8,
582            ..Default::default()
583        }));
584
585        let mut handles = vec![];
586        for _ in 0..10 {
587            let g = Arc::clone(&guard);
588            handles.push(thread::spawn(move || {
589                for _ in 0..100 {
590                    g.add_bytes(100);
591                    g.release(100);
592                }
593            }));
594        }
595        for h in handles {
596            h.join().unwrap();
597        }
598        // All bytes should be released -- may not be exactly 0 due to ordering
599        // but should be close (within one thread's batch)
600        assert!(
601            guard.current_bytes() < 1000,
602            "leaked bytes: {}",
603            guard.current_bytes()
604        );
605    }
606
607    #[test]
608    fn test_try_reserve_rollback_is_atomic() {
609        let guard = MemoryGuard::new(MemoryGuardConfig {
610            limit_bytes: 100,
611            ..Default::default()
612        });
613        assert!(guard.try_reserve(90));
614        assert!(!guard.try_reserve(20)); // over limit, rolled back
615        assert_eq!(guard.current_bytes(), 90); // not 110
616        assert!(guard.try_reserve(10)); // exactly at limit
617        assert_eq!(guard.current_bytes(), 100);
618    }
619
620    // Process-global heap source for the switch test. nextest isolates each
621    // test in its own process, so registering it here is contained to this
622    // test and does not leak into the per-batch-counter tests above. (This is
623    // the single test in this module that touches the global hook.)
624    static TEST_HEAP: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
625    fn test_heap_source() -> usize {
626        TEST_HEAP.load(Ordering::Relaxed)
627    }
628
629    #[test]
630    fn heap_source_overrides_read_path_and_admission() {
631        assert!(set_heap_source(test_heap_source), "first set wins");
632        assert!(
633            !set_heap_source(test_heap_source),
634            "second set is a no-op (first-wins)"
635        );
636
637        let guard = MemoryGuard::new(MemoryGuardConfig {
638            limit_bytes: 1_000,
639            pressure_threshold: 0.8,
640            ..Default::default()
641        });
642
643        // Reads come from the heap source, not the reservation counter.
644        TEST_HEAP.store(250, Ordering::Relaxed);
645        assert_eq!(guard.current_bytes(), 250);
646        assert!((guard.pressure_ratio() - 0.25).abs() < 0.001);
647        assert!(!guard.under_pressure());
648
649        // Pressure tracks the live heap -- including growth never reserved,
650        // which the per-batch counter would have been blind to.
651        TEST_HEAP.store(850, Ordering::Relaxed);
652        assert!(
653            guard.under_pressure(),
654            "85% live heap is over the 80% threshold"
655        );
656        assert_eq!(guard.pressure(), MemoryPressure::High);
657
658        // try_reserve is a projected-admission check against the true heap and
659        // does NOT mutate the reservation counter.
660        TEST_HEAP.store(900, Ordering::Relaxed);
661        assert!(guard.try_reserve(100), "900 + 100 == limit, admitted");
662        assert!(!guard.try_reserve(200), "900 + 200 > limit, rejected");
663        assert_eq!(
664            guard.current_bytes(),
665            900,
666            "counter untouched by try_reserve"
667        );
668    }
669
670    #[test]
671    fn effective_auto_limit_caps_at_memory_high() {
672        // headroom 0.85 of 1000 = 850; high 600 is lower -> cap at the soft
673        // throttle so we shed before the kernel does.
674        assert_eq!(effective_auto_limit(1000, 0.85, Some(600)), 600);
675        // high above the headroom limit -> headroom wins (no change).
676        assert_eq!(effective_auto_limit(1000, 0.85, Some(900)), 850);
677        // no memory.high in force -> headroom limit.
678        assert_eq!(effective_auto_limit(1000, 0.85, None), 850);
679    }
680
681    #[test]
682    fn test_config_defaults() {
683        let config = MemoryGuardConfig::default();
684        assert_eq!(config.limit_bytes, 0);
685        assert!((config.pressure_threshold - 0.80).abs() < 0.001);
686        assert!((config.cgroup_headroom - 0.85).abs() < 0.001);
687    }
688
689    #[test]
690    fn test_from_env_raw_defaults_when_unset() {
691        // With no env vars set, should return defaults
692        let config = MemoryGuardConfig::from_env_raw("TEST_MG_UNSET");
693        assert_eq!(config.limit_bytes, 0);
694        assert!((config.pressure_threshold - 0.80).abs() < 0.001);
695        assert!((config.cgroup_headroom - 0.85).abs() < 0.001);
696    }
697
698    #[test]
699    fn test_env_parsed_helper() {
700        // env_parsed returns None for unset vars
701        assert!(env_parsed::<u64>("NONEXISTENT_PREFIX_XYZ", "FOO").is_none());
702        assert!(env_parsed::<f64>("NONEXISTENT_PREFIX_XYZ", "BAR").is_none());
703    }
704
705    #[test]
706    fn test_guard_with_explicit_config_overrides() {
707        // Simulates what from_env would produce with overrides
708        let config = MemoryGuardConfig {
709            limit_bytes: 2_147_483_648,
710            pressure_threshold: 0.75,
711            cgroup_headroom: 0.90,
712        };
713        let guard = MemoryGuard::new(config);
714        assert_eq!(guard.limit_bytes(), 2_147_483_648);
715    }
716
717    #[test]
718    fn test_guard_with_custom_headroom() {
719        // 85% headroom on 1 GiB = 870 MiB effective
720        let config = MemoryGuardConfig {
721            limit_bytes: 0, // auto-detect
722            pressure_threshold: 0.80,
723            cgroup_headroom: 0.85,
724        };
725        let guard = MemoryGuard::new(config);
726        // Auto-detected, so limit should be 85% of system/cgroup memory
727        assert!(guard.limit_bytes() > 0);
728    }
729
730    #[test]
731    fn test_validate_accepts_defaults_and_rejects_bad_fractions() {
732        assert!(MemoryGuardConfig::default().validate().is_ok());
733
734        for bad in [0.0, -0.1, 1.5, f64::NAN, f64::INFINITY] {
735            let cfg = MemoryGuardConfig {
736                pressure_threshold: bad,
737                ..Default::default()
738            };
739            assert!(
740                cfg.validate().is_err(),
741                "pressure_threshold={bad} must be rejected"
742            );
743            let cfg = MemoryGuardConfig {
744                cgroup_headroom: bad,
745                ..Default::default()
746            };
747            assert!(
748                cfg.validate().is_err(),
749                "cgroup_headroom={bad} must be rejected"
750            );
751        }
752    }
753
754    #[test]
755    fn test_new_clamps_invalid_config_no_divide_by_zero() {
756        // A zero/NaN headroom with auto-detect could yield a zero limit ->
757        // divide-by-zero. A zero pressure_threshold would make every ratio
758        // "over". new() must clamp to safe defaults and keep ratios finite.
759        let guard = MemoryGuard::new(MemoryGuardConfig {
760            limit_bytes: 0,
761            pressure_threshold: 0.0,
762            cgroup_headroom: 0.0,
763        });
764        assert!(guard.limit_bytes() >= 1, "limit floored at >=1");
765        guard.add_bytes(10);
766        assert!(
767            guard.pressure_ratio().is_finite(),
768            "pressure ratio must be finite, not div-by-zero"
769        );
770    }
771
772    #[test]
773    fn test_new_with_nan_threshold_is_finite() {
774        let guard = MemoryGuard::new(MemoryGuardConfig {
775            limit_bytes: 1000,
776            pressure_threshold: f64::NAN,
777            cgroup_headroom: f64::NAN,
778        });
779        assert_eq!(guard.limit_bytes(), 1000);
780        guard.add_bytes(900);
781        // Clamped threshold (0.8 default) -> 90% is over -> under pressure.
782        assert!(guard.under_pressure());
783    }
784
785    #[test]
786    fn test_auto_detect_limit() {
787        // With limit_bytes = 0, should auto-detect from system
788        let guard = MemoryGuard::new(MemoryGuardConfig::default());
789        assert!(
790            guard.limit_bytes() > 0,
791            "auto-detected limit should be positive"
792        );
793        // Should be less than total system memory (headroom applied)
794    }
795}