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}