use std::path::{Path, PathBuf};
use std::sync::OnceLock;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;
const CGROUP_ROOT: &str = "/sys/fs/cgroup";
const MEMORY_HIGH_BACKPRESSURE_PCT: u64 = 95;
static PARENT: OnceLock<Option<PathBuf>> = OnceLock::new();
static RUN_SEQ: AtomicU64 = AtomicU64::new(0);
static LIMITS: OnceLock<Limits> = OnceLock::new();
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct Limits {
pub memory_max: Option<String>,
pub pids_max: Option<String>,
}
impl Limits {
pub fn from_specs(memory_max: Option<&str>, pids_max: Option<&str>) -> Limits {
Limits {
memory_max: memory_max.and_then(normalize_bytes),
pids_max: pids_max.and_then(normalize_count),
}
}
pub fn is_empty(&self) -> bool {
self.memory_max.is_none() && self.pids_max.is_none()
}
}
fn normalize_bytes(s: &str) -> Option<String> {
let s = s.trim();
if s.eq_ignore_ascii_case("max") {
return Some("max".to_string());
}
let (digits, mult): (&str, u64) = match s.chars().last() {
Some(c) if c.is_ascii_digit() => (s, 1),
Some('K' | 'k') => (&s[..s.len() - 1], 1024),
Some('M' | 'm') => (&s[..s.len() - 1], 1024 * 1024),
Some('G' | 'g') => (&s[..s.len() - 1], 1024 * 1024 * 1024),
_ => return None,
};
let n: u64 = digits.trim().parse().ok()?;
n.checked_mul(mult).map(|b| b.to_string())
}
fn normalize_count(s: &str) -> Option<String> {
let s = s.trim();
if s.eq_ignore_ascii_case("max") {
return Some("max".to_string());
}
s.parse::<u64>().ok().map(|n| n.to_string())
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct MemorySnapshot {
pub max: Option<u64>,
pub current: Option<u64>,
pub high: Option<u64>,
}
impl MemorySnapshot {
pub fn detected(&self) -> bool {
self.max.is_some() || self.current.is_some() || self.high.is_some()
}
}
pub fn snapshot() -> MemorySnapshot {
MemorySnapshot {
max: memory_max(),
current: memory_current(),
high: memory_high(),
}
}
pub fn memory_max() -> Option<u64> {
read_mem(&Path::new(CGROUP_ROOT).join("memory.max"))
}
pub fn memory_current() -> Option<u64> {
read_mem(&Path::new(CGROUP_ROOT).join("memory.current"))
}
pub fn memory_high() -> Option<u64> {
read_mem(&Path::new(CGROUP_ROOT).join("memory.high"))
}
fn read_mem(path: &Path) -> Option<u64> {
std::fs::read_to_string(path)
.ok()
.and_then(|s| parse_mem(&s))
}
fn parse_mem(s: &str) -> Option<u64> {
match s.trim() {
"max" => None,
t => t.parse::<u64>().ok(),
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Configured {
pub parent: PathBuf,
pub limits: Limits,
pub limits_unavailable: bool,
}
pub fn configure(
spec: Option<&str>,
memory_max: Option<&str>,
pids_max: Option<&str>,
) -> Option<Configured> {
let resolved = spec.and_then(resolve_parent).filter(|p| ensure_writable(p));
let limits = Limits::from_specs(memory_max, pids_max);
let mut limits_unavailable = false;
if let Some(p) = &resolved {
sweep_stale(p);
if !limits.is_empty() {
limits_unavailable = !enable_controllers(p, &limits);
}
}
let _ = PARENT.set(resolved);
let _ = LIMITS.set(limits);
let parent = PARENT.get().cloned().flatten()?;
Some(Configured {
parent,
limits: LIMITS.get().cloned().unwrap_or_default(),
limits_unavailable,
})
}
fn enable_controllers(parent: &Path, limits: &Limits) -> bool {
let file = parent.join("cgroup.subtree_control");
let mut all_ok = true;
if limits.memory_max.is_some() {
all_ok &= write_cgroup(&file, "+memory");
}
if limits.pids_max.is_some() {
all_ok &= write_cgroup(&file, "+pids");
}
all_ok
}
fn sweep_stale(parent: &Path) {
let Ok(entries) = std::fs::read_dir(parent) else {
return;
};
let me = std::process::id();
for entry in entries.flatten() {
if stale_run(&entry.file_name().to_string_lossy(), me) {
let dir = entry.path();
let _ = std::fs::write(dir.join("cgroup.kill"), "1");
let _ = std::fs::remove_dir(&dir);
}
}
}
fn stale_run(name: &str, me: u32) -> bool {
if let Some(rest) = name.strip_prefix("run-") {
match rest.split('-').next().and_then(|p| p.parse::<u32>().ok()) {
Some(pid) => pid == me || !pid_alive(pid),
None => false,
}
} else {
name.starts_with(".probe-")
}
}
fn pid_alive(pid: u32) -> bool {
let rc = unsafe { libc::kill(pid as i32, 0) };
rc == 0 || std::io::Error::last_os_error().raw_os_error() == Some(libc::EPERM)
}
fn resolve_parent(spec: &str) -> Option<PathBuf> {
match spec.trim() {
"" => None,
"auto" => own_cgroup_dir().map(|d| d.join("agentd")),
p if Path::new(p).is_absolute()
&& Path::new(p).starts_with(CGROUP_ROOT)
&& !p.contains("..") =>
{
Some(PathBuf::from(p))
}
_ => None,
}
}
fn own_cgroup_dir() -> Option<PathBuf> {
let content = std::fs::read_to_string("/proc/self/cgroup").ok()?;
let rel = content.lines().find_map(|l| l.strip_prefix("0::"))?.trim();
Some(Path::new(CGROUP_ROOT).join(rel.trim_start_matches('/')))
}
fn ensure_writable(parent: &Path) -> bool {
if std::fs::create_dir_all(parent).is_err() {
return false;
}
let probe = parent.join(format!(".probe-{}", std::process::id()));
match std::fs::create_dir(&probe) {
Ok(()) => {
let _ = std::fs::remove_dir(&probe);
true
}
Err(_) => false,
}
}
pub fn under_memory_pressure() -> bool {
over_threshold(memory_current(), memory_high())
}
fn over_threshold(current: Option<u64>, high: Option<u64>) -> bool {
match (current, high) {
(Some(cur), Some(high)) if high > 0 => {
cur.saturating_mul(100) >= high.saturating_mul(MEMORY_HIGH_BACKPRESSURE_PCT)
}
_ => false,
}
}
pub struct CgroupGuard {
dir: PathBuf,
}
impl CgroupGuard {
pub fn for_run() -> Option<CgroupGuard> {
let parent = PARENT.get().and_then(|o| o.clone())?;
let name = format!(
"run-{}-{}",
std::process::id(),
RUN_SEQ.fetch_add(1, Ordering::Relaxed)
);
let guard = Self::create(&parent, &name)?;
if let Some(limits) = LIMITS.get() {
guard.apply_limits(limits);
}
Some(guard)
}
fn create(parent: &Path, name: &str) -> Option<CgroupGuard> {
let dir = parent.join(name);
std::fs::create_dir_all(&dir).ok()?;
Some(CgroupGuard { dir })
}
pub fn apply_limits(&self, limits: &Limits) -> (bool, bool) {
let memory_ok = match &limits.memory_max {
Some(v) => write_cgroup(&self.dir.join("memory.max"), v),
None => false,
};
let pids_ok = match &limits.pids_max {
Some(v) => write_cgroup(&self.dir.join("pids.max"), v),
None => false,
};
(memory_ok, pids_ok)
}
pub fn place(&self, pid: i32) -> bool {
write_cgroup(&self.dir.join("cgroup.procs"), &pid.to_string())
}
pub fn kill_all(&self) -> bool {
write_cgroup(&self.dir.join("cgroup.kill"), "1")
}
pub fn path(&self) -> &Path {
&self.dir
}
pub fn oom_kills(&self) -> Option<u64> {
parse_oom_kills(&std::fs::read_to_string(self.dir.join("memory.events")).ok()?)
}
fn try_remove(&self) -> bool {
std::fs::remove_dir(&self.dir).is_ok()
}
}
impl Drop for CgroupGuard {
fn drop(&mut self) {
self.kill_all();
for _ in 0..5 {
if self.try_remove() {
return;
}
std::thread::sleep(Duration::from_millis(10));
}
}
}
fn write_cgroup(path: &Path, value: &str) -> bool {
std::fs::write(path, value).is_ok()
}
fn parse_oom_kills(events: &str) -> Option<u64> {
events
.lines()
.find_map(|l| l.strip_prefix("oom_kill "))?
.trim()
.parse()
.ok()
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::Write;
#[test]
fn parse_mem_handles_max_and_numbers() {
assert_eq!(parse_mem("max\n"), None); assert_eq!(parse_mem("max"), None);
assert_eq!(parse_mem("1073741824\n"), Some(1_073_741_824));
assert_eq!(parse_mem("0"), Some(0));
assert_eq!(parse_mem("garbage"), None);
assert_eq!(parse_mem(""), None);
}
#[test]
fn read_mem_reads_a_fixture_or_degrades_to_none() {
let mut f = tempfile::NamedTempFile::new().unwrap();
writeln!(f, "536870912").unwrap();
assert_eq!(read_mem(f.path()), Some(536_870_912));
let mut unlimited = tempfile::NamedTempFile::new().unwrap();
writeln!(unlimited, "max").unwrap();
assert_eq!(read_mem(unlimited.path()), None);
assert_eq!(read_mem(Path::new("/nonexistent/agentd/memory.max")), None);
}
#[test]
fn snapshot_detected_reflects_any_readable_field() {
assert!(!MemorySnapshot::default().detected());
assert!(
MemorySnapshot {
max: Some(1),
..Default::default()
}
.detected()
);
}
#[test]
fn over_threshold_backpressures_at_95_percent_of_high() {
assert!(!over_threshold(Some(1_000), None));
assert!(!over_threshold(None, Some(1_000)));
assert!(!over_threshold(None, None));
assert!(!over_threshold(Some(1_000), Some(0))); assert!(!over_threshold(Some(900), Some(1_000))); assert!(over_threshold(Some(950), Some(1_000))); assert!(over_threshold(Some(1_000), Some(1_000))); assert!(over_threshold(Some(2_000), Some(1_000))); }
#[test]
fn normalize_bytes_handles_suffixes_max_and_garbage() {
assert_eq!(normalize_bytes("max").as_deref(), Some("max"));
assert_eq!(normalize_bytes("MAX").as_deref(), Some("max"));
assert_eq!(normalize_bytes("1048576").as_deref(), Some("1048576"));
assert_eq!(
normalize_bytes("512M").as_deref(),
Some((512 * 1024 * 1024).to_string().as_str())
);
assert_eq!(
normalize_bytes("2G").as_deref(),
Some((2u64 * 1024 * 1024 * 1024).to_string().as_str())
);
assert_eq!(
normalize_bytes("64k").as_deref(),
Some((64 * 1024).to_string().as_str())
);
assert_eq!(normalize_bytes(""), None);
assert_eq!(normalize_bytes("M"), None); assert_eq!(normalize_bytes("12T"), None); assert_eq!(normalize_bytes("abc"), None);
}
#[test]
fn normalize_count_handles_max_and_integers() {
assert_eq!(normalize_count("max").as_deref(), Some("max"));
assert_eq!(normalize_count("128").as_deref(), Some("128"));
assert_eq!(normalize_count("0").as_deref(), Some("0"));
assert_eq!(normalize_count(""), None);
assert_eq!(normalize_count("-1"), None);
assert_eq!(normalize_count("lots"), None);
}
#[test]
fn parse_oom_kills_reads_the_counter() {
let events = "low 0\nhigh 0\nmax 3\noom 1\noom_kill 2\noom_group_kill 0\n";
assert_eq!(parse_oom_kills(events), Some(2));
assert_eq!(parse_oom_kills("oom_kill 0\n"), Some(0));
assert_eq!(parse_oom_kills("low 0\nhigh 0\n"), None); assert_eq!(parse_oom_kills(""), None);
}
#[test]
fn limits_from_specs_drops_unparseable() {
let l = Limits::from_specs(Some("256M"), Some("32"));
assert_eq!(
l.memory_max.as_deref(),
Some((256 * 1024 * 1024).to_string().as_str())
);
assert_eq!(l.pids_max.as_deref(), Some("32"));
assert!(!l.is_empty());
assert!(Limits::from_specs(None, None).is_empty());
assert!(Limits::from_specs(Some("nonsense"), None).is_empty()); }
#[test]
fn limits_are_applied_and_pids_max_is_enforced() {
let mgr = Path::new(CGROUP_ROOT).join(format!("agentd-test-limits-{}", std::process::id()));
if std::fs::create_dir(&mgr).is_err() {
eprintln!("skip: cannot create a cgroup under {CGROUP_ROOT}");
return;
}
struct Cleanup(PathBuf);
impl Drop for Cleanup {
fn drop(&mut self) {
let _ = std::fs::write(self.0.join("cgroup.kill"), "1");
let _ = std::fs::remove_dir(&self.0);
}
}
let _mgr_cleanup = Cleanup(mgr.clone());
let limits = Limits::from_specs(Some("32M"), Some("1"));
if !enable_controllers(&mgr, &limits) {
eprintln!("skip: parent cannot delegate memory/pids controllers");
return;
}
let guard = CgroupGuard::create(&mgr, "leaf").expect("create leaf cgroup");
let (mem_ok, pids_ok) = guard.apply_limits(&limits);
assert!(pids_ok, "pids.max applied");
assert_eq!(
std::fs::read_to_string(guard.dir.join("pids.max"))
.unwrap()
.trim(),
"1"
);
if mem_ok {
assert_eq!(
std::fs::read_to_string(guard.dir.join("memory.max"))
.unwrap()
.trim(),
(32 * 1024 * 1024).to_string()
);
}
let mut fds = [0i32; 2];
assert_eq!(unsafe { libc::pipe(fds.as_mut_ptr()) }, 0, "pipe");
let (rfd, wfd) = (fds[0], fds[1]);
let pid = unsafe { libc::fork() };
assert!(pid >= 0, "fork probe");
if pid == 0 {
unsafe {
libc::close(wfd);
let mut b = [0u8; 1];
libc::read(rfd, b.as_mut_ptr() as *mut libc::c_void, 1); let g = libc::fork();
if g == 0 {
libc::_exit(0); }
if g < 0 {
libc::_exit(0); }
let mut s = 0;
libc::waitpid(g, &mut s, 0);
libc::_exit(1); }
}
struct ProbeGuard(Option<i32>);
impl Drop for ProbeGuard {
fn drop(&mut self) {
if let Some(pid) = self.0 {
unsafe {
libc::kill(pid, libc::SIGKILL);
let mut s = 0;
libc::waitpid(pid, &mut s, 0);
}
}
}
}
let mut probe = ProbeGuard(Some(pid));
unsafe { libc::close(rfd) };
assert!(guard.place(pid), "migrate the probe into the leaf");
unsafe {
libc::write(wfd, b"x".as_ptr() as *const libc::c_void, 1);
libc::close(wfd);
}
let mut status = 0;
assert_eq!(
unsafe { libc::waitpid(pid, &mut status, 0) },
pid,
"reap probe"
);
probe.0 = None; assert!(
libc::WIFEXITED(status) && libc::WEXITSTATUS(status) == 0,
"a fork inside the pids.max=1 cgroup must be refused (status={status})"
);
}
#[test]
fn resolve_parent_accepts_auto_and_in_mount_paths_only() {
assert_eq!(resolve_parent(""), None);
assert_eq!(resolve_parent("relative/path"), None);
assert_eq!(resolve_parent("/etc/passwd"), None); assert_eq!(resolve_parent("/sys/fs/cgroup/../etc"), None); assert_eq!(resolve_parent("/sys/fs/cgroup-sibling/x"), None); assert_eq!(
resolve_parent("/sys/fs/cgroup/foo/agentd"),
Some(PathBuf::from("/sys/fs/cgroup/foo/agentd"))
);
if let Some(p) = resolve_parent("auto") {
assert!(p.starts_with(CGROUP_ROOT));
assert!(p.ends_with("agentd"));
}
}
#[test]
fn stale_run_targets_dead_and_own_pid_only() {
let me = std::process::id();
let dead = i32::MAX as u32;
assert!(
stale_run(&format!("run-{dead}-0"), me),
"dead pid → reclaim"
);
assert!(
stale_run(&format!("run-{me}-7"), me),
"our reused pid → reclaim"
);
assert!(stale_run(".probe-123", me), "probe leftover → reclaim");
assert!(!stale_run("run-1-0", me), "live sibling (pid 1) → spare");
assert!(!stale_run("unrelated", me), "non-run dir → spare");
assert!(!stale_run("run-notapid-0", me), "unparseable pid → spare");
}
#[test]
fn cgroup_kill_reaps_a_process_that_left_the_process_group() {
let Some(parent) = resolve_parent("auto") else {
eprintln!("skip: no cgroup-v2 on this host");
return;
};
if !ensure_writable(&parent) {
eprintln!("skip: cgroup-v2 tree not writable (no delegation)");
return;
}
let cg = CgroupGuard::create(&parent, &format!("test-kill-{}", std::process::id()))
.expect("create child cgroup");
let pid = unsafe { libc::fork() };
assert!(pid >= 0, "fork failed");
if pid == 0 {
unsafe {
libc::setsid(); libc::sleep(10); libc::_exit(0); }
}
assert!(cg.place(pid), "place the child pid into the cgroup");
assert!(cg.kill_all(), "write cgroup.kill");
let mut status = 0i32;
let reaped = unsafe { libc::waitpid(pid, &mut status, 0) };
assert_eq!(reaped, pid, "reaped the child");
assert!(
libc::WIFSIGNALED(status),
"child was SIGKILLed by cgroup.kill, not a clean exit (status={status})"
);
assert_eq!(libc::WTERMSIG(status), libc::SIGKILL, "killed by SIGKILL");
}
}