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/// Default cgroup headroom: use 85% of cgroup limit.
148///
149/// Rationale: Rust has no GC so no spike headroom needed (unlike JVM 75% / Go 80%).
150/// 15% headroom covers jemalloc fragmentation, kernel overhead, and page cache.
151const DEFAULT_CGROUP_HEADROOM: f64 = 0.85;
152
153/// Default pressure threshold: backpressure at 80% of effective limit.
154///
155/// With 85% headroom, backpressure activates at ~68% of actual cgroup limit.
156/// Matches OTel Collector's `limit_percentage: 80` philosophy.
157const DEFAULT_PRESSURE_THRESHOLD: f64 = 0.80;
158
159impl Default for MemoryGuardConfig {
160    fn default() -> Self {
161        Self {
162            limit_bytes: 0, // auto-detect
163            pressure_threshold: DEFAULT_PRESSURE_THRESHOLD,
164            cgroup_headroom: DEFAULT_CGROUP_HEADROOM,
165        }
166    }
167}
168
169impl MemoryGuardConfig {
170    /// Load from the config cascade, falling back to defaults.
171    ///
172    /// When the `config` feature is enabled and `config::setup()` has been
173    /// called, reads the `memory` key from the cascade. Otherwise returns
174    /// [`MemoryGuardConfig::default()`].
175    #[must_use]
176    pub fn from_cascade() -> Self {
177        #[cfg(feature = "config")]
178        {
179            if let Some(cfg) = crate::config::try_get()
180                && let Ok(memory) = cfg.unmarshal_key_registered::<Self>("memory")
181            {
182                return memory;
183            }
184        }
185        Self::default()
186    }
187
188    /// Create config from environment variables with a prefix.
189    ///
190    /// Reads standard env vars for memory configuration:
191    /// - `{PREFIX}_MEMORY_LIMIT_BYTES` -- explicit limit (0 or unset = auto-detect from cgroup)
192    /// - `{PREFIX}_MEMORY_PRESSURE_THRESHOLD` -- backpressure trigger (default 0.80)
193    /// - `{PREFIX}_MEMORY_CGROUP_HEADROOM` -- fraction of cgroup limit to use (default 0.85)
194    ///
195    /// # Example
196    ///
197    /// ```bash
198    /// DFE_MEMORY_LIMIT_BYTES=4294967296      # 4 GiB explicit
199    /// DFE_MEMORY_PRESSURE_THRESHOLD=0.75     # backpressure at 75%
200    /// DFE_MEMORY_CGROUP_HEADROOM=0.90        # use 90% of cgroup
201    /// ```
202    ///
203    /// ```rust,no_run
204    /// use hyperi_rustlib::memory::MemoryGuardConfig;
205    /// let config = MemoryGuardConfig::from_env("DFE");
206    /// ```
207    #[must_use]
208    #[cfg(feature = "config")]
209    pub fn from_env(prefix: &str) -> Self {
210        use crate::config::flat_env::flat_env_parsed;
211
212        let mut config = Self::default();
213
214        if let Some(v) = flat_env_parsed::<u64>(prefix, "MEMORY_LIMIT_BYTES") {
215            config.limit_bytes = v;
216        }
217        if let Some(v) = flat_env_parsed::<f64>(prefix, "MEMORY_PRESSURE_THRESHOLD") {
218            config.pressure_threshold = v;
219        }
220        if let Some(v) = flat_env_parsed::<f64>(prefix, "MEMORY_CGROUP_HEADROOM") {
221            config.cgroup_headroom = v;
222        }
223
224        config
225    }
226
227    /// Create config from environment variables without requiring `config` feature.
228    ///
229    /// Same as [`from_env`](Self::from_env) but uses `std::env` directly.
230    #[must_use]
231    pub fn from_env_raw(prefix: &str) -> Self {
232        let mut config = Self::default();
233
234        if let Some(v) = env_parsed::<u64>(prefix, "MEMORY_LIMIT_BYTES") {
235            config.limit_bytes = v;
236        }
237        if let Some(v) = env_parsed::<f64>(prefix, "MEMORY_PRESSURE_THRESHOLD") {
238            config.pressure_threshold = v;
239        }
240        if let Some(v) = env_parsed::<f64>(prefix, "MEMORY_CGROUP_HEADROOM") {
241            config.cgroup_headroom = v;
242        }
243
244        config
245    }
246
247    /// Validate the config, returning an error describing the first invalid
248    /// field. `pressure_threshold` and `cgroup_headroom` must each be a finite
249    /// fraction in `(0.0, 1.0]`. Call this at startup to fail fast on bad
250    /// config rather than relying on [`MemoryGuard::new`]'s defensive clamping.
251    ///
252    /// # Errors
253    ///
254    /// Returns `Err` with a human-readable message if a fraction field is
255    /// non-finite, `<= 0.0`, or `> 1.0`.
256    pub fn validate(&self) -> Result<(), String> {
257        check_fraction(self.pressure_threshold, "pressure_threshold")?;
258        check_fraction(self.cgroup_headroom, "cgroup_headroom")?;
259        Ok(())
260    }
261}
262
263/// Cgroup-aware memory tracking with backpressure signals.
264///
265/// Tracks application-level memory usage (not process RSS) and provides
266/// fast atomic checks for the hot path. Designed for data pipeline services
267/// where incoming data must be rejected (503) before hitting the container
268/// memory limit.
269///
270/// # Usage
271///
272/// ```rust,no_run
273/// use hyperi_rustlib::memory::{MemoryGuard, MemoryGuardConfig};
274///
275/// let guard = MemoryGuard::new(MemoryGuardConfig::default());
276///
277/// // On data arrival -- check before accepting
278/// let payload_len = 1024u64;
279/// if !guard.try_reserve(payload_len) {
280///     // return 503 -- backpressure
281/// }
282///
283/// // After data is flushed/sent
284/// guard.release(payload_len);
285///
286/// // Fast hot-path check
287/// if guard.under_pressure() {
288///     // return 503
289/// }
290/// ```
291pub struct MemoryGuard {
292    /// Current tracked bytes (application-level, not RSS).
293    current_bytes: AtomicU64,
294    /// Effective memory limit in bytes.
295    limit_bytes: u64,
296    /// Pressure threshold (0.0-1.0).
297    pressure_threshold: f64,
298    /// Fast boolean for hot-path pressure check.
299    under_pressure: AtomicBool,
300}
301
302impl MemoryGuard {
303    /// Create a new memory guard.
304    ///
305    /// If `config.limit_bytes` is 0, auto-detects from cgroup (K8s) or system memory,
306    /// then applies `cgroup_headroom` factor to leave room for process overhead.
307    #[must_use]
308    #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
309    pub fn new(config: MemoryGuardConfig) -> Self {
310        // Defensive: a non-finite / out-of-range threshold or headroom would
311        // produce a zero/NaN limit and a divide-by-zero pressure ratio. Clamp
312        // to the safe default and log loudly. Callers wanting hard rejection
313        // should call `config.validate()` at startup.
314        let pressure_threshold = sane_fraction(
315            config.pressure_threshold,
316            DEFAULT_PRESSURE_THRESHOLD,
317            "pressure_threshold",
318        );
319        let cgroup_headroom = sane_fraction(
320            config.cgroup_headroom,
321            DEFAULT_CGROUP_HEADROOM,
322            "cgroup_headroom",
323        );
324
325        let raw_limit = if config.limit_bytes > 0 {
326            config.limit_bytes
327        } else {
328            let detected = cgroup::detect_memory_limit();
329            // Apply headroom -- don't use 100% of cgroup limit
330            (detected as f64 * cgroup_headroom) as u64
331        };
332        // Never permit a zero effective limit: every pressure calculation
333        // divides by it.
334        let limit_bytes = raw_limit.max(1);
335
336        tracing::info!(limit_bytes, pressure_threshold, "memory guard initialised");
337
338        Self {
339            current_bytes: AtomicU64::new(0),
340            limit_bytes,
341            pressure_threshold,
342            under_pressure: AtomicBool::new(false),
343        }
344    }
345
346    /// Try to reserve bytes. Returns false if over the limit (backpressure).
347    ///
348    /// With a registered [`set_heap_source`], this is a projected-admission
349    /// check against the *true total heap* (`heap() + bytes <= limit`) and does
350    /// NOT mutate the reservation counter -- the allocator already accounts the
351    /// bytes once they are allocated, and frees them on drop, so no `release`
352    /// is needed. Without a source it is the classic atomic check-and-add on
353    /// the per-batch counter (rolled back if it would exceed the limit).
354    #[inline]
355    pub fn try_reserve(&self, bytes: u64) -> bool {
356        if let Some(heap) = heap_bytes() {
357            return heap + bytes <= self.limit_bytes;
358        }
359        let current = self.current_bytes.fetch_add(bytes, Ordering::Relaxed) + bytes;
360        if current > self.limit_bytes {
361            // Over limit -- roll back
362            self.current_bytes.fetch_sub(bytes, Ordering::Relaxed);
363            self.under_pressure.store(true, Ordering::Relaxed);
364            return false;
365        }
366        self.update_pressure(current);
367        true
368    }
369
370    /// Add bytes without checking the limit (for tracking only).
371    /// Use when data is already accepted and you just need to track it.
372    #[inline]
373    pub fn add_bytes(&self, bytes: u64) {
374        let new_total = self.current_bytes.fetch_add(bytes, Ordering::Relaxed) + bytes;
375        self.update_pressure(new_total);
376    }
377
378    /// Release bytes after data is flushed/sent/dropped.
379    ///
380    /// Uses saturating subtraction to prevent underflow wrapping.
381    #[inline]
382    pub fn release(&self, bytes: u64) {
383        let prev = self
384            .current_bytes
385            .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
386                Some(current.saturating_sub(bytes))
387            })
388            // Always succeeds (closure always returns Some).
389            .unwrap_or_else(|v| v);
390        self.update_pressure(prev.saturating_sub(bytes));
391    }
392
393    /// Fast hot-path pressure check.
394    ///
395    /// With a registered [`set_heap_source`], computes live from the true heap
396    /// (one atomic load + compare); otherwise reads the cached flag maintained
397    /// by `try_reserve`/`add_bytes`/`release`.
398    #[inline]
399    pub fn under_pressure(&self) -> bool {
400        if heap_bytes().is_some() {
401            return self.pressure_ratio() >= self.pressure_threshold;
402        }
403        self.under_pressure.load(Ordering::Relaxed)
404    }
405
406    /// Current pressure level.
407    #[inline]
408    pub fn pressure(&self) -> MemoryPressure {
409        let ratio = self.pressure_ratio();
410        if ratio >= self.pressure_threshold {
411            MemoryPressure::High
412        } else if ratio >= 0.5 {
413            MemoryPressure::Medium
414        } else {
415            MemoryPressure::Low
416        }
417    }
418
419    /// Current usage as fraction of limit (0.0 - 1.0+).
420    #[inline]
421    pub fn pressure_ratio(&self) -> f64 {
422        self.current_bytes() as f64 / self.limit_bytes as f64
423    }
424
425    /// Current memory usage in bytes.
426    ///
427    /// Returns the true total live heap when a [`set_heap_source`] is
428    /// registered, otherwise the sum of outstanding per-batch reservations.
429    #[inline]
430    pub fn current_bytes(&self) -> u64 {
431        heap_bytes().unwrap_or_else(|| self.current_bytes.load(Ordering::Relaxed))
432    }
433
434    /// Configured memory limit in bytes.
435    #[inline]
436    pub fn limit_bytes(&self) -> u64 {
437        self.limit_bytes
438    }
439
440    /// Update the pressure flag based on current usage.
441    #[inline]
442    fn update_pressure(&self, current: u64) {
443        let ratio = current as f64 / self.limit_bytes as f64;
444        self.under_pressure
445            .store(ratio >= self.pressure_threshold, Ordering::Relaxed);
446    }
447}
448
449#[cfg(test)]
450mod tests {
451    use super::*;
452
453    #[test]
454    fn test_memory_guard_default() {
455        let guard = MemoryGuard::new(MemoryGuardConfig {
456            limit_bytes: 1_000_000, // 1MB explicit
457            ..Default::default()
458        });
459        assert_eq!(guard.limit_bytes(), 1_000_000);
460        assert_eq!(guard.current_bytes(), 0);
461        assert!(!guard.under_pressure());
462        assert_eq!(guard.pressure(), MemoryPressure::Low);
463    }
464
465    #[test]
466    fn test_try_reserve_within_limit() {
467        let guard = MemoryGuard::new(MemoryGuardConfig {
468            limit_bytes: 1000,
469            ..Default::default()
470        });
471        assert!(guard.try_reserve(500));
472        assert_eq!(guard.current_bytes(), 500);
473    }
474
475    #[test]
476    fn test_try_reserve_over_limit() {
477        let guard = MemoryGuard::new(MemoryGuardConfig {
478            limit_bytes: 1000,
479            ..Default::default()
480        });
481        assert!(guard.try_reserve(500));
482        assert!(!guard.try_reserve(600)); // would exceed 1000
483        assert_eq!(guard.current_bytes(), 500); // rolled back
484        assert!(guard.under_pressure());
485    }
486
487    #[test]
488    fn test_release_reduces_pressure() {
489        let guard = MemoryGuard::new(MemoryGuardConfig {
490            limit_bytes: 1000,
491            pressure_threshold: 0.8,
492            ..Default::default()
493        });
494        guard.add_bytes(900); // 90% -- over threshold
495        assert!(guard.under_pressure());
496        assert_eq!(guard.pressure(), MemoryPressure::High);
497
498        guard.release(500); // down to 400 = 40%
499        assert!(!guard.under_pressure());
500        assert_eq!(guard.pressure(), MemoryPressure::Low);
501    }
502
503    #[test]
504    fn test_pressure_levels() {
505        let guard = MemoryGuard::new(MemoryGuardConfig {
506            limit_bytes: 1000,
507            pressure_threshold: 0.8,
508            ..Default::default()
509        });
510
511        // Low (< 50%)
512        guard.add_bytes(400);
513        assert_eq!(guard.pressure(), MemoryPressure::Low);
514
515        // Medium (50-80%)
516        guard.add_bytes(200); // 600 = 60%
517        assert_eq!(guard.pressure(), MemoryPressure::Medium);
518
519        // High (>= 80%)
520        guard.add_bytes(300); // 900 = 90%
521        assert_eq!(guard.pressure(), MemoryPressure::High);
522    }
523
524    #[test]
525    fn test_pressure_ratio() {
526        let guard = MemoryGuard::new(MemoryGuardConfig {
527            limit_bytes: 1000,
528            ..Default::default()
529        });
530        guard.add_bytes(250);
531        let ratio = guard.pressure_ratio();
532        assert!((ratio - 0.25).abs() < 0.001);
533    }
534
535    #[test]
536    fn test_release_saturating() {
537        let guard = MemoryGuard::new(MemoryGuardConfig {
538            limit_bytes: 1000,
539            ..Default::default()
540        });
541        guard.add_bytes(100);
542        guard.release(200); // release more than added -- saturates to 0
543        assert_eq!(
544            guard.current_bytes(),
545            0,
546            "over-release must saturate to 0, not wrap"
547        );
548        assert!(!guard.under_pressure());
549        assert_eq!(guard.pressure(), MemoryPressure::Low);
550
551        // Verify the guard is still functional after over-release
552        assert!(guard.try_reserve(500));
553        assert_eq!(guard.current_bytes(), 500);
554    }
555
556    #[test]
557    fn test_concurrent_reserve_release() {
558        use std::sync::Arc;
559        use std::thread;
560
561        let guard = Arc::new(MemoryGuard::new(MemoryGuardConfig {
562            limit_bytes: 100_000,
563            pressure_threshold: 0.8,
564            ..Default::default()
565        }));
566
567        let mut handles = vec![];
568        for _ in 0..10 {
569            let g = Arc::clone(&guard);
570            handles.push(thread::spawn(move || {
571                for _ in 0..100 {
572                    g.add_bytes(100);
573                    g.release(100);
574                }
575            }));
576        }
577        for h in handles {
578            h.join().unwrap();
579        }
580        // All bytes should be released -- may not be exactly 0 due to ordering
581        // but should be close (within one thread's batch)
582        assert!(
583            guard.current_bytes() < 1000,
584            "leaked bytes: {}",
585            guard.current_bytes()
586        );
587    }
588
589    #[test]
590    fn test_try_reserve_rollback_is_atomic() {
591        let guard = MemoryGuard::new(MemoryGuardConfig {
592            limit_bytes: 100,
593            ..Default::default()
594        });
595        assert!(guard.try_reserve(90));
596        assert!(!guard.try_reserve(20)); // over limit, rolled back
597        assert_eq!(guard.current_bytes(), 90); // not 110
598        assert!(guard.try_reserve(10)); // exactly at limit
599        assert_eq!(guard.current_bytes(), 100);
600    }
601
602    // Process-global heap source for the switch test. nextest isolates each
603    // test in its own process, so registering it here is contained to this
604    // test and does not leak into the per-batch-counter tests above. (This is
605    // the single test in this module that touches the global hook.)
606    static TEST_HEAP: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
607    fn test_heap_source() -> usize {
608        TEST_HEAP.load(Ordering::Relaxed)
609    }
610
611    #[test]
612    fn heap_source_overrides_read_path_and_admission() {
613        assert!(set_heap_source(test_heap_source), "first set wins");
614        assert!(
615            !set_heap_source(test_heap_source),
616            "second set is a no-op (first-wins)"
617        );
618
619        let guard = MemoryGuard::new(MemoryGuardConfig {
620            limit_bytes: 1_000,
621            pressure_threshold: 0.8,
622            ..Default::default()
623        });
624
625        // Reads come from the heap source, not the reservation counter.
626        TEST_HEAP.store(250, Ordering::Relaxed);
627        assert_eq!(guard.current_bytes(), 250);
628        assert!((guard.pressure_ratio() - 0.25).abs() < 0.001);
629        assert!(!guard.under_pressure());
630
631        // Pressure tracks the live heap -- including growth never reserved,
632        // which the per-batch counter would have been blind to.
633        TEST_HEAP.store(850, Ordering::Relaxed);
634        assert!(
635            guard.under_pressure(),
636            "85% live heap is over the 80% threshold"
637        );
638        assert_eq!(guard.pressure(), MemoryPressure::High);
639
640        // try_reserve is a projected-admission check against the true heap and
641        // does NOT mutate the reservation counter.
642        TEST_HEAP.store(900, Ordering::Relaxed);
643        assert!(guard.try_reserve(100), "900 + 100 == limit, admitted");
644        assert!(!guard.try_reserve(200), "900 + 200 > limit, rejected");
645        assert_eq!(
646            guard.current_bytes(),
647            900,
648            "counter untouched by try_reserve"
649        );
650    }
651
652    #[test]
653    fn test_config_defaults() {
654        let config = MemoryGuardConfig::default();
655        assert_eq!(config.limit_bytes, 0);
656        assert!((config.pressure_threshold - 0.80).abs() < 0.001);
657        assert!((config.cgroup_headroom - 0.85).abs() < 0.001);
658    }
659
660    #[test]
661    fn test_from_env_raw_defaults_when_unset() {
662        // With no env vars set, should return defaults
663        let config = MemoryGuardConfig::from_env_raw("TEST_MG_UNSET");
664        assert_eq!(config.limit_bytes, 0);
665        assert!((config.pressure_threshold - 0.80).abs() < 0.001);
666        assert!((config.cgroup_headroom - 0.85).abs() < 0.001);
667    }
668
669    #[test]
670    fn test_env_parsed_helper() {
671        // env_parsed returns None for unset vars
672        assert!(env_parsed::<u64>("NONEXISTENT_PREFIX_XYZ", "FOO").is_none());
673        assert!(env_parsed::<f64>("NONEXISTENT_PREFIX_XYZ", "BAR").is_none());
674    }
675
676    #[test]
677    fn test_guard_with_explicit_config_overrides() {
678        // Simulates what from_env would produce with overrides
679        let config = MemoryGuardConfig {
680            limit_bytes: 2_147_483_648,
681            pressure_threshold: 0.75,
682            cgroup_headroom: 0.90,
683        };
684        let guard = MemoryGuard::new(config);
685        assert_eq!(guard.limit_bytes(), 2_147_483_648);
686    }
687
688    #[test]
689    fn test_guard_with_custom_headroom() {
690        // 85% headroom on 1 GiB = 870 MiB effective
691        let config = MemoryGuardConfig {
692            limit_bytes: 0, // auto-detect
693            pressure_threshold: 0.80,
694            cgroup_headroom: 0.85,
695        };
696        let guard = MemoryGuard::new(config);
697        // Auto-detected, so limit should be 85% of system/cgroup memory
698        assert!(guard.limit_bytes() > 0);
699    }
700
701    #[test]
702    fn test_validate_accepts_defaults_and_rejects_bad_fractions() {
703        assert!(MemoryGuardConfig::default().validate().is_ok());
704
705        for bad in [0.0, -0.1, 1.5, f64::NAN, f64::INFINITY] {
706            let cfg = MemoryGuardConfig {
707                pressure_threshold: bad,
708                ..Default::default()
709            };
710            assert!(
711                cfg.validate().is_err(),
712                "pressure_threshold={bad} must be rejected"
713            );
714            let cfg = MemoryGuardConfig {
715                cgroup_headroom: bad,
716                ..Default::default()
717            };
718            assert!(
719                cfg.validate().is_err(),
720                "cgroup_headroom={bad} must be rejected"
721            );
722        }
723    }
724
725    #[test]
726    fn test_new_clamps_invalid_config_no_divide_by_zero() {
727        // A zero/NaN headroom with auto-detect could yield a zero limit ->
728        // divide-by-zero. A zero pressure_threshold would make every ratio
729        // "over". new() must clamp to safe defaults and keep ratios finite.
730        let guard = MemoryGuard::new(MemoryGuardConfig {
731            limit_bytes: 0,
732            pressure_threshold: 0.0,
733            cgroup_headroom: 0.0,
734        });
735        assert!(guard.limit_bytes() >= 1, "limit floored at >=1");
736        guard.add_bytes(10);
737        assert!(
738            guard.pressure_ratio().is_finite(),
739            "pressure ratio must be finite, not div-by-zero"
740        );
741    }
742
743    #[test]
744    fn test_new_with_nan_threshold_is_finite() {
745        let guard = MemoryGuard::new(MemoryGuardConfig {
746            limit_bytes: 1000,
747            pressure_threshold: f64::NAN,
748            cgroup_headroom: f64::NAN,
749        });
750        assert_eq!(guard.limit_bytes(), 1000);
751        guard.add_bytes(900);
752        // Clamped threshold (0.8 default) -> 90% is over -> under pressure.
753        assert!(guard.under_pressure());
754    }
755
756    #[test]
757    fn test_auto_detect_limit() {
758        // With limit_bytes = 0, should auto-detect from system
759        let guard = MemoryGuard::new(MemoryGuardConfig::default());
760        assert!(
761            guard.limit_bytes() > 0,
762            "auto-detected limit should be positive"
763        );
764        // Should be less than total system memory (headroom applied)
765    }
766}