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}