use std::path::{Path, PathBuf};
pub(crate) const DEFAULT_MEMORY_FRACTION: f64 = 0.7;
pub(crate) const DYNAMIC_CAP_FLOOR_BYTES: u64 = 16 * 1024 * 1024;
pub(crate) const MEMORY_TICK_SECS: u64 = 10;
#[must_use]
pub(crate) fn read_memory_budget() -> u64 {
cgroup_v2_memory_max().unwrap_or_else(|| mem_available().unwrap_or(256 * 1024 * 1024))
}
#[must_use]
pub(crate) fn cap_for_fraction(fraction: f64, budget: u64) -> u64 {
let f = fraction.clamp(0.0, 1.0);
((budget as f64) * f) as u64
}
fn cgroup_v2_memory_max() -> Option<u64> {
let contents = std::fs::read_to_string("/proc/self/cgroup").ok()?;
let path = contents.lines().find_map(|l| l.strip_prefix("0::"))?;
let mut dir = if path.trim().is_empty() {
PathBuf::from("/sys/fs/cgroup")
} else {
PathBuf::from(format!("/sys/fs/cgroup{}", path.trim()))
};
loop {
if let Ok(s) = std::fs::read_to_string(dir.join("memory.max")) {
let t = s.trim();
if t != "max" {
if let Ok(v) = t.parse::<u64>() {
if v > 0 {
return Some(v);
}
}
}
}
if dir == Path::new("/sys/fs/cgroup") || !dir.pop() {
return None;
}
}
}
fn mem_available() -> Option<u64> {
let s = std::fs::read_to_string("/proc/meminfo").ok()?;
for line in s.lines() {
if let Some(rest) = line.strip_prefix("MemAvailable:") {
let kb = rest
.trim()
.trim_end_matches("kB")
.trim()
.parse::<u64>()
.ok()?;
return Some(kb.saturating_mul(1024));
}
}
None
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn cap_is_fraction_of_budget() {
let budget = 1_000_000_000;
let cap = cap_for_fraction(DEFAULT_MEMORY_FRACTION, budget);
assert_eq!(cap, 700_000_000);
}
#[test]
fn fraction_is_clamped() {
assert_eq!(cap_for_fraction(2.0, 100), 100);
assert_eq!(cap_for_fraction(-1.0, 100), 0);
}
#[test]
fn cgroup_root_unlimited_falls_back() {
assert!(cgroup_v2_memory_max().is_none() || cgroup_v2_memory_max().is_some_and(|v| v > 0));
}
#[test]
fn budget_is_nonzero() {
assert!(read_memory_budget() > 0);
}
#[test]
fn cap_never_exceeds_budget() {
assert!(cap_for_fraction(0.7, 100) <= 100);
}
}