Skip to main content

agentd/supervisor/
cgroup.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2//! cgroup v2 memory awareness (read-only).
3//!
4//! Best-effort, **never required**: a cloud-native unit reports the memory
5//! budget its scheduler handed it so OOM risk is observable (logged at startup,
6//! and exposed as a `/metrics` gauge). Reads the unified cgroup v2 interface
7//! files directly under `/sys/fs/cgroup`; in a container with a cgroup
8//! namespace (the target shape) that path is the unit's *own* cgroup, so the
9//! direct read is correct.
10//! On a bare host it reflects the root cgroup (whole-host) — still informative.
11//! Any missing file / cgroup v1 / parse failure degrades to `None`.
12//!
13//! ## Active enforcement (best-effort, opt-in, never required)
14//!
15//! On top of the reads, when the operator opts in (`--cgroup auto|<path>` /
16//! `AGENTD_CGROUP`) and the cgroup-v2 tree is writable, each supervised run is
17//! placed in its own child cgroup so teardown can write **`cgroup.kill`** — the
18//! kernel then SIGKILLs the *entire* subtree atomically, catching processes that
19//! escaped the process group (`setsid`) which `killpg` + `PR_SET_PDEATHSIG`
20//! would miss — the worst leak agentd can suffer, since such a process outlives
21//! every other teardown mechanism. And [`under_memory_pressure`]
22//! lets the spawn-admission gates backpressure when the unit is at its
23//! `memory.high` soft limit. Every cgroup op is best-effort: if the tree isn't
24//! writable (no delegation, cgroup-v1, off-cgroup) the feature silently disables
25//! and the run falls back to the PDEATHSIG + kill-ladder path — agentd stays
26//! cgroup-*aware*, never cgroup-*requiring*.
27//!
28//! Note: hard resource *limits* on the child (`memory.max`/`pids.max`) need the
29//! parent to delegate controllers via `cgroup.subtree_control`, which fails
30//! (`EBUSY`) whenever the parent cgroup holds processes directly — common for a
31//! systemd unit without `Delegate=yes`. So limit-setting stays a deployment
32//! concern (size the pod's `resources.limits`); agentd adds the atomic teardown
33//! backstop and the soft-pressure backpressure, which need no delegation.
34
35use std::path::{Path, PathBuf};
36use std::sync::OnceLock;
37use std::sync::atomic::{AtomicU64, Ordering};
38use std::time::Duration;
39
40const CGROUP_ROOT: &str = "/sys/fs/cgroup";
41
42/// Backpressure new spawns once usage reaches this percentage of `memory.high`.
43const MEMORY_HIGH_BACKPRESSURE_PCT: u64 = 95;
44
45/// The resolved parent cgroup directory under which per-run child cgroups are
46/// created, set once at startup by [`configure`]. `Some(None)` = configured but
47/// the tree isn't writable (feature disabled); unset / `Some(None)` both mean
48/// [`CgroupGuard::for_run`] yields `None` and runs fall back to PDEATHSIG.
49static PARENT: OnceLock<Option<PathBuf>> = OnceLock::new();
50
51/// Per-run child-cgroup name counter (unique within this process).
52static RUN_SEQ: AtomicU64 = AtomicU64::new(0);
53
54/// Hard resource limits to write onto each per-run leaf cgroup, set once at
55/// startup by [`configure`]. Empty (the default) → leaves get no limits, only the
56/// `cgroup.kill` teardown backstop. Applying limits needs the parent to delegate
57/// the `memory`/`pids` controllers (see [`enable_controllers`]); where it can't
58/// (e.g. `--cgroup auto` under a busy unit cgroup → `EBUSY`) the writes no-op and
59/// the run still gets atomic teardown.
60static LIMITS: OnceLock<Limits> = OnceLock::new();
61
62/// Normalised hard limits for a per-run leaf cgroup. Each field, when set, is the
63/// exact string written to the corresponding cgroup-v2 interface file.
64#[derive(Debug, Clone, Default, PartialEq, Eq)]
65pub struct Limits {
66    /// Value for `memory.max` (`"max"` or a byte count).
67    pub memory_max: Option<String>,
68    /// Value for `pids.max` (`"max"` or a count).
69    pub pids_max: Option<String>,
70}
71
72impl Limits {
73    /// Build from raw `--cgroup-memory-max` / `--cgroup-pids-max` specs; an
74    /// unparseable spec is dropped (the limit just isn't applied).
75    pub fn from_specs(memory_max: Option<&str>, pids_max: Option<&str>) -> Limits {
76        Limits {
77            memory_max: memory_max.and_then(normalize_bytes),
78            pids_max: pids_max.and_then(normalize_count),
79        }
80    }
81
82    /// No limits requested → skip controller delegation + per-leaf writes entirely.
83    pub fn is_empty(&self) -> bool {
84        self.memory_max.is_none() && self.pids_max.is_none()
85    }
86}
87
88/// Normalise a memory-size spec to the bytes string `memory.max` expects: `"max"`
89/// (unlimited), a plain byte count, or a `K`/`M`/`G`-suffixed (1024-based) size.
90/// `None` for anything unparseable.
91fn normalize_bytes(s: &str) -> Option<String> {
92    let s = s.trim();
93    if s.eq_ignore_ascii_case("max") {
94        return Some("max".to_string());
95    }
96    let (digits, mult): (&str, u64) = match s.chars().last() {
97        Some(c) if c.is_ascii_digit() => (s, 1),
98        Some('K' | 'k') => (&s[..s.len() - 1], 1024),
99        Some('M' | 'm') => (&s[..s.len() - 1], 1024 * 1024),
100        Some('G' | 'g') => (&s[..s.len() - 1], 1024 * 1024 * 1024),
101        _ => return None,
102    };
103    let n: u64 = digits.trim().parse().ok()?;
104    n.checked_mul(mult).map(|b| b.to_string())
105}
106
107/// Normalise a pid-count spec for `pids.max`: `"max"` or a non-negative integer.
108fn normalize_count(s: &str) -> Option<String> {
109    let s = s.trim();
110    if s.eq_ignore_ascii_case("max") {
111        return Some("max".to_string());
112    }
113    s.parse::<u64>().ok().map(|n| n.to_string())
114}
115
116/// A point-in-time view of the unit's cgroup v2 memory interface.
117#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
118pub struct MemorySnapshot {
119    /// Hard limit (`memory.max`); `None` = unlimited (`"max"`) or unavailable.
120    pub max: Option<u64>,
121    /// Current charged usage (`memory.current`), in bytes.
122    pub current: Option<u64>,
123    /// Soft throttle threshold (`memory.high`); `None` = unset or unavailable.
124    pub high: Option<u64>,
125}
126
127impl MemorySnapshot {
128    /// Whether any cgroup v2 memory file was readable (i.e. we are cgroup-aware).
129    pub fn detected(&self) -> bool {
130        self.max.is_some() || self.current.is_some() || self.high.is_some()
131    }
132}
133
134/// Read the current cgroup v2 memory snapshot (best-effort; never fails).
135pub fn snapshot() -> MemorySnapshot {
136    MemorySnapshot {
137        max: memory_max(),
138        current: memory_current(),
139        high: memory_high(),
140    }
141}
142
143/// `memory.max` — the hard limit; `None` when unlimited (`"max"`) or unreadable.
144pub fn memory_max() -> Option<u64> {
145    read_mem(&Path::new(CGROUP_ROOT).join("memory.max"))
146}
147
148/// `memory.current` — current charged usage in bytes.
149pub fn memory_current() -> Option<u64> {
150    read_mem(&Path::new(CGROUP_ROOT).join("memory.current"))
151}
152
153/// `memory.high` — the soft (throttling) limit; `None` when unset (`"max"`).
154pub fn memory_high() -> Option<u64> {
155    read_mem(&Path::new(CGROUP_ROOT).join("memory.high"))
156}
157
158/// Read + parse one cgroup memory file. `None` on any I/O error or `"max"`.
159fn read_mem(path: &Path) -> Option<u64> {
160    std::fs::read_to_string(path)
161        .ok()
162        .and_then(|s| parse_mem(&s))
163}
164
165/// Parse a cgroup v2 memory value: a byte count, or `"max"` (unlimited → `None`).
166fn parse_mem(s: &str) -> Option<u64> {
167    match s.trim() {
168        "max" => None,
169        t => t.parse::<u64>().ok(),
170    }
171}
172
173// ---------------------------------------------------------------------------
174// Active enforcement: child-cgroup placement + `cgroup.kill` teardown backstop.
175// ---------------------------------------------------------------------------
176
177/// What [`configure`] settled on: the resolved parent dir (if armed) and whether
178/// requested hard limits will actually be enforced (controllers delegated).
179#[derive(Debug, Clone, PartialEq, Eq)]
180pub struct Configured {
181    /// The parent dir under which per-run leaves are created (feature is armed).
182    pub parent: PathBuf,
183    /// Limits that will be applied to each leaf (empty if none requested).
184    pub limits: Limits,
185    /// True when limits were requested but the controllers could **not** be
186    /// delegated (so the writes will no-op) — the caller should warn.
187    pub limits_unavailable: bool,
188}
189
190/// Resolve + probe the `--cgroup` spec ONCE at startup, arming per-run child
191/// cgroups. `spec` is `"auto"` (derive `<own-cgroup>/agentd` from
192/// `/proc/self/cgroup`) or an absolute path under `/sys/fs/cgroup`. Optional
193/// `memory_max`/`pids_max` specs request hard limits on each run's leaf; this
194/// best-effort delegates the controllers to the parent so the limits can take
195/// effect. Returns `None` when off / not writable (the feature stays dormant).
196/// Idempotent — the first call wins.
197pub fn configure(
198    spec: Option<&str>,
199    memory_max: Option<&str>,
200    pids_max: Option<&str>,
201) -> Option<Configured> {
202    let resolved = spec.and_then(resolve_parent).filter(|p| ensure_writable(p));
203    let limits = Limits::from_specs(memory_max, pids_max);
204    let mut limits_unavailable = false;
205    if let Some(p) = &resolved {
206        // Reclaim any `run-*` cgroups orphaned by prior crashed/abandoned runs (a
207        // wedged D-state task can outlive its guard's Drop), so a long-lived
208        // daemon can't slowly accumulate stale child cgroups across restarts.
209        sweep_stale(p);
210        // Delegate the controllers the requested limits need to the parent, so
211        // the per-run leaves get enforceable `memory.max`/`pids.max`. Fails
212        // (EBUSY) where the parent holds processes directly (e.g. `auto` under a
213        // busy unit cgroup) — limits then no-op, but teardown still works.
214        if !limits.is_empty() {
215            limits_unavailable = !enable_controllers(p, &limits);
216        }
217    }
218    // OnceLock::set fails only if already set (first call wins). Read back the
219    // stored values so the caller's log never disagrees with what governs runs.
220    let _ = PARENT.set(resolved);
221    let _ = LIMITS.set(limits);
222    let parent = PARENT.get().cloned().flatten()?;
223    Some(Configured {
224        parent,
225        limits: LIMITS.get().cloned().unwrap_or_default(),
226        limits_unavailable,
227    })
228}
229
230/// Delegate the `memory`/`pids` controllers (only those the requested limits
231/// need) to `parent` via its `cgroup.subtree_control`, so child leaves expose the
232/// matching interface files. Each controller is delegated with its OWN write: a
233/// `subtree_control` write is atomic, so a combined `+memory +pids` would fail
234/// wholesale when the parent can delegate only one (e.g. a `Delegate=pids` unit),
235/// needlessly dropping the achievable limit. Best-effort → returns whether
236/// **every requested** controller was delegated.
237fn enable_controllers(parent: &Path, limits: &Limits) -> bool {
238    let file = parent.join("cgroup.subtree_control");
239    let mut all_ok = true;
240    if limits.memory_max.is_some() {
241        all_ok &= write_cgroup(&file, "+memory");
242    }
243    if limits.pids_max.is_some() {
244        all_ok &= write_cgroup(&file, "+pids");
245    }
246    all_ok
247}
248
249/// Best-effort reclaim of stale per-run child cgroups under `parent`. Removes a
250/// `run-<pid>-*` cgroup only when its owning pid is **dead** (or is our own,
251/// freshly-reused pid) — so a concurrent sibling agentd sharing this parent is
252/// never torn down. `.probe-*` leftovers (from a crashed `ensure_writable`) are
253/// always reclaimed.
254fn sweep_stale(parent: &Path) {
255    let Ok(entries) = std::fs::read_dir(parent) else {
256        return;
257    };
258    let me = std::process::id();
259    for entry in entries.flatten() {
260        if stale_run(&entry.file_name().to_string_lossy(), me) {
261            let dir = entry.path();
262            let _ = std::fs::write(dir.join("cgroup.kill"), "1");
263            let _ = std::fs::remove_dir(&dir);
264        }
265    }
266}
267
268/// Whether a child-cgroup dir name denotes a reclaimable stale run: a
269/// `run-<pid>-*` whose pid is dead or our own (freshly-reused — we've made no run
270/// cgroups yet) pid, or a `.probe-*` leftover from a crashed `ensure_writable`. A
271/// live sibling's `run-*` (and any other name) is spared.
272fn stale_run(name: &str, me: u32) -> bool {
273    if let Some(rest) = name.strip_prefix("run-") {
274        match rest.split('-').next().and_then(|p| p.parse::<u32>().ok()) {
275            Some(pid) => pid == me || !pid_alive(pid),
276            None => false,
277        }
278    } else {
279        name.starts_with(".probe-")
280    }
281}
282
283/// Whether `pid` is a live process. `kill(pid, 0)` → `Ok`/`EPERM` = alive,
284/// `ESRCH` = dead. Conservative: an unrelated process that reused the pid reads
285/// as alive, so we skip (delay reclaim) rather than risk touching its cgroup.
286fn pid_alive(pid: u32) -> bool {
287    let rc = unsafe { libc::kill(pid as i32, 0) };
288    rc == 0 || std::io::Error::last_os_error().raw_os_error() == Some(libc::EPERM)
289}
290
291/// Resolve a `--cgroup` spec to an absolute parent directory (no I/O beyond
292/// reading `/proc/self/cgroup` for `auto`). `None` for an unusable spec.
293fn resolve_parent(spec: &str) -> Option<PathBuf> {
294    match spec.trim() {
295        "" => None,
296        "auto" => own_cgroup_dir().map(|d| d.join("agentd")),
297        // An explicit path must sit under the cgroup-v2 mount, by path COMPONENT
298        // (so `/sys/fs/cgroup-sibling` can't slip past a byte-prefix), with no
299        // `..`. A guard-rail, not a security boundary — `--cgroup` is operator-
300        // supplied and the operator already controls the process; a symlink
301        // component could still redirect, which the trust model accepts.
302        p if Path::new(p).is_absolute()
303            && Path::new(p).starts_with(CGROUP_ROOT)
304            && !p.contains("..") =>
305        {
306            Some(PathBuf::from(p))
307        }
308        _ => None,
309    }
310}
311
312/// The unit's own cgroup directory, from the `0::<path>` line of
313/// `/proc/self/cgroup` (cgroup-v2 unified hierarchy). `None` off cgroup-v2.
314fn own_cgroup_dir() -> Option<PathBuf> {
315    let content = std::fs::read_to_string("/proc/self/cgroup").ok()?;
316    let rel = content.lines().find_map(|l| l.strip_prefix("0::"))?.trim();
317    Some(Path::new(CGROUP_ROOT).join(rel.trim_start_matches('/')))
318}
319
320/// Probe that we can actually create + remove a child cgroup under `parent`
321/// (so the feature only arms where it works). Creates `parent` if needed.
322fn ensure_writable(parent: &Path) -> bool {
323    if std::fs::create_dir_all(parent).is_err() {
324        return false;
325    }
326    let probe = parent.join(format!(".probe-{}", std::process::id()));
327    match std::fs::create_dir(&probe) {
328        Ok(()) => {
329            let _ = std::fs::remove_dir(&probe);
330            true
331        }
332        Err(_) => false,
333    }
334}
335
336/// Whether the unit is at/over the backpressure fraction of its `memory.high`
337/// soft limit — a signal for the spawn-admission gates to refuse new subagents
338/// rather than push the cgroup into reclaim/OOM. Reads live each call; `false`
339/// when no cgroup / no `memory.high` set (can't tell → don't block).
340pub fn under_memory_pressure() -> bool {
341    // Read only the two files the predicate needs (not the full snapshot's three).
342    over_threshold(memory_current(), memory_high())
343}
344
345/// Pure predicate behind [`under_memory_pressure`] (testable without a cgroup).
346fn over_threshold(current: Option<u64>, high: Option<u64>) -> bool {
347    match (current, high) {
348        (Some(cur), Some(high)) if high > 0 => {
349            cur.saturating_mul(100) >= high.saturating_mul(MEMORY_HIGH_BACKPRESSURE_PCT)
350        }
351        _ => false,
352    }
353}
354
355/// A per-run child cgroup. Placing the root subagent here puts its whole subtree
356/// in the cgroup (membership inherits across `fork`), so [`kill_all`] tears the
357/// entire subtree down atomically. RAII: `Drop` kills + removes the cgroup.
358///
359/// [`kill_all`]: CgroupGuard::kill_all
360pub struct CgroupGuard {
361    dir: PathBuf,
362}
363
364impl CgroupGuard {
365    /// Create the per-run child cgroup under the configured parent, applying the
366    /// configured hard limits, or `None` when the feature is off / creation fails
367    /// (best-effort, never an error).
368    pub fn for_run() -> Option<CgroupGuard> {
369        let parent = PARENT.get().and_then(|o| o.clone())?;
370        let name = format!(
371            "run-{}-{}",
372            std::process::id(),
373            RUN_SEQ.fetch_add(1, Ordering::Relaxed)
374        );
375        let guard = Self::create(&parent, &name)?;
376        if let Some(limits) = LIMITS.get() {
377            guard.apply_limits(limits);
378        }
379        Some(guard)
380    }
381
382    /// Create a child cgroup `parent/name` (best-effort). Shared by `for_run`
383    /// and tests (which resolve a parent directly, bypassing the global).
384    fn create(parent: &Path, name: &str) -> Option<CgroupGuard> {
385        let dir = parent.join(name);
386        std::fs::create_dir_all(&dir).ok()?;
387        Some(CgroupGuard { dir })
388    }
389
390    /// Write the configured hard limits onto this leaf's `memory.max`/`pids.max`.
391    /// Best-effort: a write no-ops where the controller wasn't delegated (the
392    /// interface file is absent / read-only), so the run keeps `cgroup.kill`
393    /// teardown without the limit. Returns `(memory_ok, pids_ok)` for logging.
394    pub fn apply_limits(&self, limits: &Limits) -> (bool, bool) {
395        let memory_ok = match &limits.memory_max {
396            Some(v) => write_cgroup(&self.dir.join("memory.max"), v),
397            None => false,
398        };
399        let pids_ok = match &limits.pids_max {
400            Some(v) => write_cgroup(&self.dir.join("pids.max"), v),
401            None => false,
402        };
403        (memory_ok, pids_ok)
404    }
405
406    /// Move `pid` (and, by inheritance, its future descendants) into this
407    /// cgroup by writing its `cgroup.procs`. Best-effort → returns success.
408    pub fn place(&self, pid: i32) -> bool {
409        write_cgroup(&self.dir.join("cgroup.procs"), &pid.to_string())
410    }
411
412    /// Atomically SIGKILL every process in the subtree via `cgroup.kill` — the
413    /// backstop beyond `killpg`/PDEATHSIG. Best-effort → returns success.
414    pub fn kill_all(&self) -> bool {
415        write_cgroup(&self.dir.join("cgroup.kill"), "1")
416    }
417
418    /// The cgroup directory (for logging).
419    pub fn path(&self) -> &Path {
420        &self.dir
421    }
422
423    /// The `oom_kill` count from this leaf's `memory.events` — how many of its
424    /// processes the kernel OOM-killed (i.e. hit `memory.max`). `None` when the
425    /// memory controller isn't active (no limit / not delegated). Lets the
426    /// supervisor report a `memory.max` kill plainly instead of as a generic exit.
427    pub fn oom_kills(&self) -> Option<u64> {
428        parse_oom_kills(&std::fs::read_to_string(self.dir.join("memory.events")).ok()?)
429    }
430
431    /// Remove the cgroup dir; succeeds only once every member process has been
432    /// reaped (the kernel frees the cgroup then).
433    fn try_remove(&self) -> bool {
434        std::fs::remove_dir(&self.dir).is_ok()
435    }
436}
437
438impl Drop for CgroupGuard {
439    fn drop(&mut self) {
440        // Final backstop: SIGKILL anything left, then remove the (now-freeing)
441        // cgroup with a few short retries while the kernel reaps its members. A
442        // lingering empty dir is harmless, so this never blocks for long.
443        self.kill_all();
444        for _ in 0..5 {
445            if self.try_remove() {
446                return;
447            }
448            std::thread::sleep(Duration::from_millis(10));
449        }
450    }
451}
452
453/// Write a value to a cgroup control file. Best-effort: any error (no such
454/// controller, EACCES, EBUSY, off-cgroup) is swallowed → returns success.
455fn write_cgroup(path: &Path, value: &str) -> bool {
456    std::fs::write(path, value).is_ok()
457}
458
459/// Parse the `oom_kill` counter from a `memory.events` body (one `key value`
460/// pair per line). `None` when the key is absent.
461fn parse_oom_kills(events: &str) -> Option<u64> {
462    events
463        .lines()
464        .find_map(|l| l.strip_prefix("oom_kill "))?
465        .trim()
466        .parse()
467        .ok()
468}
469
470#[cfg(test)]
471mod tests {
472    use super::*;
473    use std::io::Write;
474
475    #[test]
476    fn parse_mem_handles_max_and_numbers() {
477        assert_eq!(parse_mem("max\n"), None); // unlimited
478        assert_eq!(parse_mem("max"), None);
479        assert_eq!(parse_mem("1073741824\n"), Some(1_073_741_824));
480        assert_eq!(parse_mem("0"), Some(0));
481        assert_eq!(parse_mem("garbage"), None);
482        assert_eq!(parse_mem(""), None);
483    }
484
485    #[test]
486    fn read_mem_reads_a_fixture_or_degrades_to_none() {
487        let mut f = tempfile::NamedTempFile::new().unwrap();
488        writeln!(f, "536870912").unwrap();
489        assert_eq!(read_mem(f.path()), Some(536_870_912));
490
491        let mut unlimited = tempfile::NamedTempFile::new().unwrap();
492        writeln!(unlimited, "max").unwrap();
493        assert_eq!(read_mem(unlimited.path()), None);
494
495        // missing file → None, never an error
496        assert_eq!(read_mem(Path::new("/nonexistent/agentd/memory.max")), None);
497    }
498
499    #[test]
500    fn snapshot_detected_reflects_any_readable_field() {
501        assert!(!MemorySnapshot::default().detected());
502        assert!(
503            MemorySnapshot {
504                max: Some(1),
505                ..Default::default()
506            }
507            .detected()
508        );
509    }
510
511    #[test]
512    fn over_threshold_backpressures_at_95_percent_of_high() {
513        // No high set, or no current → can't tell → never backpressure.
514        assert!(!over_threshold(Some(1_000), None));
515        assert!(!over_threshold(None, Some(1_000)));
516        assert!(!over_threshold(None, None));
517        assert!(!over_threshold(Some(1_000), Some(0))); // high==0 → ignore
518        // Below the fraction → allow; at/above → backpressure.
519        assert!(!over_threshold(Some(900), Some(1_000))); // 90% < 95%
520        assert!(over_threshold(Some(950), Some(1_000))); // exactly 95%
521        assert!(over_threshold(Some(1_000), Some(1_000))); // at high
522        assert!(over_threshold(Some(2_000), Some(1_000))); // over high
523    }
524
525    #[test]
526    fn normalize_bytes_handles_suffixes_max_and_garbage() {
527        assert_eq!(normalize_bytes("max").as_deref(), Some("max"));
528        assert_eq!(normalize_bytes("MAX").as_deref(), Some("max"));
529        assert_eq!(normalize_bytes("1048576").as_deref(), Some("1048576"));
530        assert_eq!(
531            normalize_bytes("512M").as_deref(),
532            Some((512 * 1024 * 1024).to_string().as_str())
533        );
534        assert_eq!(
535            normalize_bytes("2G").as_deref(),
536            Some((2u64 * 1024 * 1024 * 1024).to_string().as_str())
537        );
538        assert_eq!(
539            normalize_bytes("64k").as_deref(),
540            Some((64 * 1024).to_string().as_str())
541        );
542        assert_eq!(normalize_bytes(""), None);
543        assert_eq!(normalize_bytes("M"), None); // no digits
544        assert_eq!(normalize_bytes("12T"), None); // unsupported suffix
545        assert_eq!(normalize_bytes("abc"), None);
546    }
547
548    #[test]
549    fn normalize_count_handles_max_and_integers() {
550        assert_eq!(normalize_count("max").as_deref(), Some("max"));
551        assert_eq!(normalize_count("128").as_deref(), Some("128"));
552        assert_eq!(normalize_count("0").as_deref(), Some("0"));
553        assert_eq!(normalize_count(""), None);
554        assert_eq!(normalize_count("-1"), None);
555        assert_eq!(normalize_count("lots"), None);
556    }
557
558    #[test]
559    fn parse_oom_kills_reads_the_counter() {
560        let events = "low 0\nhigh 0\nmax 3\noom 1\noom_kill 2\noom_group_kill 0\n";
561        assert_eq!(parse_oom_kills(events), Some(2));
562        assert_eq!(parse_oom_kills("oom_kill 0\n"), Some(0));
563        assert_eq!(parse_oom_kills("low 0\nhigh 0\n"), None); // key absent
564        assert_eq!(parse_oom_kills(""), None);
565    }
566
567    #[test]
568    fn limits_from_specs_drops_unparseable() {
569        let l = Limits::from_specs(Some("256M"), Some("32"));
570        assert_eq!(
571            l.memory_max.as_deref(),
572            Some((256 * 1024 * 1024).to_string().as_str())
573        );
574        assert_eq!(l.pids_max.as_deref(), Some("32"));
575        assert!(!l.is_empty());
576        assert!(Limits::from_specs(None, None).is_empty());
577        assert!(Limits::from_specs(Some("nonsense"), None).is_empty()); // dropped
578    }
579
580    /// Live proof that hard limits are both applied and enforced: build a manager
581    /// cgroup that delegates the pids controller, create a leaf via the real
582    /// `CgroupGuard` path, set `pids.max=1`, and confirm a process inside the leaf
583    /// cannot `fork` (kernel `EAGAIN`). Skips cleanly where controller delegation
584    /// isn't available (limits are best-effort).
585    #[test]
586    fn limits_are_applied_and_pids_max_is_enforced() {
587        let mgr = Path::new(CGROUP_ROOT).join(format!("agentd-test-limits-{}", std::process::id()));
588        if std::fs::create_dir(&mgr).is_err() {
589            eprintln!("skip: cannot create a cgroup under {CGROUP_ROOT}");
590            return;
591        }
592        struct Cleanup(PathBuf);
593        impl Drop for Cleanup {
594            fn drop(&mut self) {
595                let _ = std::fs::write(self.0.join("cgroup.kill"), "1");
596                let _ = std::fs::remove_dir(&self.0);
597            }
598        }
599        let _mgr_cleanup = Cleanup(mgr.clone());
600
601        let limits = Limits::from_specs(Some("32M"), Some("1"));
602        if !enable_controllers(&mgr, &limits) {
603            eprintln!("skip: parent cannot delegate memory/pids controllers");
604            return;
605        }
606        let guard = CgroupGuard::create(&mgr, "leaf").expect("create leaf cgroup");
607        let (mem_ok, pids_ok) = guard.apply_limits(&limits);
608        assert!(pids_ok, "pids.max applied");
609        assert_eq!(
610            std::fs::read_to_string(guard.dir.join("pids.max"))
611                .unwrap()
612                .trim(),
613            "1"
614        );
615        if mem_ok {
616            assert_eq!(
617                std::fs::read_to_string(guard.dir.join("memory.max"))
618                    .unwrap()
619                    .trim(),
620                (32 * 1024 * 1024).to_string()
621            );
622        }
623
624        // Functional enforcement: a process migrated INTO the leaf is the 1 task
625        // pids.max=1 allows, so its own `fork` must be refused. Sync via a pipe so
626        // the probe only forks after it has been placed in the cgroup.
627        let mut fds = [0i32; 2];
628        assert_eq!(unsafe { libc::pipe(fds.as_mut_ptr()) }, 0, "pipe");
629        let (rfd, wfd) = (fds[0], fds[1]);
630        let pid = unsafe { libc::fork() };
631        assert!(pid >= 0, "fork probe");
632        if pid == 0 {
633            // Child: async-signal-safe calls only.
634            unsafe {
635                libc::close(wfd);
636                let mut b = [0u8; 1];
637                libc::read(rfd, b.as_mut_ptr() as *mut libc::c_void, 1); // wait until placed
638                let g = libc::fork();
639                if g == 0 {
640                    libc::_exit(0); // grandchild (only reached if enforcement failed)
641                }
642                if g < 0 {
643                    libc::_exit(0); // EXPECTED: fork refused with EAGAIN
644                }
645                let mut s = 0;
646                libc::waitpid(g, &mut s, 0);
647                libc::_exit(1); // fork unexpectedly succeeded
648            }
649        }
650        // SIGKILL + reap the probe even if an assertion below panics first, so a
651        // failing run never leaks the blocked-on-read child / a busy leaf cgroup.
652        struct ProbeGuard(Option<i32>);
653        impl Drop for ProbeGuard {
654            fn drop(&mut self) {
655                if let Some(pid) = self.0 {
656                    unsafe {
657                        libc::kill(pid, libc::SIGKILL);
658                        let mut s = 0;
659                        libc::waitpid(pid, &mut s, 0);
660                    }
661                }
662            }
663        }
664        let mut probe = ProbeGuard(Some(pid));
665
666        unsafe { libc::close(rfd) };
667        assert!(guard.place(pid), "migrate the probe into the leaf");
668        unsafe {
669            libc::write(wfd, b"x".as_ptr() as *const libc::c_void, 1);
670            libc::close(wfd);
671        }
672        let mut status = 0;
673        assert_eq!(
674            unsafe { libc::waitpid(pid, &mut status, 0) },
675            pid,
676            "reap probe"
677        );
678        probe.0 = None; // reaped — disarm the guard (avoid waitpid on a reused pid)
679        assert!(
680            libc::WIFEXITED(status) && libc::WEXITSTATUS(status) == 0,
681            "a fork inside the pids.max=1 cgroup must be refused (status={status})"
682        );
683    }
684
685    #[test]
686    fn resolve_parent_accepts_auto_and_in_mount_paths_only() {
687        assert_eq!(resolve_parent(""), None);
688        assert_eq!(resolve_parent("relative/path"), None);
689        assert_eq!(resolve_parent("/etc/passwd"), None); // outside the mount
690        assert_eq!(resolve_parent("/sys/fs/cgroup/../etc"), None); // no `..` escape
691        assert_eq!(resolve_parent("/sys/fs/cgroup-sibling/x"), None); // component check, not byte prefix
692        assert_eq!(
693            resolve_parent("/sys/fs/cgroup/foo/agentd"),
694            Some(PathBuf::from("/sys/fs/cgroup/foo/agentd"))
695        );
696        // `auto` resolves iff this host exposes cgroup-v2 (`0::` line); either a
697        // path under the mount or None — never a panic.
698        if let Some(p) = resolve_parent("auto") {
699            assert!(p.starts_with(CGROUP_ROOT));
700            assert!(p.ends_with("agentd"));
701        }
702    }
703
704    #[test]
705    fn stale_run_targets_dead_and_own_pid_only() {
706        let me = std::process::id();
707        // A pid above any possible pid_max → never assigned → always dead (ESRCH).
708        let dead = i32::MAX as u32;
709        assert!(
710            stale_run(&format!("run-{dead}-0"), me),
711            "dead pid → reclaim"
712        );
713        assert!(
714            stale_run(&format!("run-{me}-7"), me),
715            "our reused pid → reclaim"
716        );
717        assert!(stale_run(".probe-123", me), "probe leftover → reclaim");
718        assert!(!stale_run("run-1-0", me), "live sibling (pid 1) → spare");
719        assert!(!stale_run("unrelated", me), "non-run dir → spare");
720        assert!(!stale_run("run-notapid-0", me), "unparseable pid → spare");
721    }
722
723    /// Live backstop proof: a process that `setsid()`s out of our process group
724    /// (so `killpg` would MISS it) is still SIGKILLed by `cgroup.kill` once placed
725    /// in the cgroup. Skips cleanly where the cgroup-v2 tree isn't writable — the
726    /// feature is never required, so its absence must not fail the suite.
727    #[test]
728    fn cgroup_kill_reaps_a_process_that_left_the_process_group() {
729        let Some(parent) = resolve_parent("auto") else {
730            eprintln!("skip: no cgroup-v2 on this host");
731            return;
732        };
733        if !ensure_writable(&parent) {
734            eprintln!("skip: cgroup-v2 tree not writable (no delegation)");
735            return;
736        }
737        let cg = CgroupGuard::create(&parent, &format!("test-kill-{}", std::process::id()))
738            .expect("create child cgroup");
739
740        // Fork a child that leaves our process group then waits. Only
741        // async-signal-safe libc calls between fork and exit (no allocations).
742        let pid = unsafe { libc::fork() };
743        assert!(pid >= 0, "fork failed");
744        if pid == 0 {
745            unsafe {
746                libc::setsid(); // new session/pgroup → killpg(our pgid) can't reach it
747                libc::sleep(10); // cgroup.kill should end us well before this
748                libc::_exit(0); // if we get here, we were NOT killed → test fails below
749            }
750        }
751
752        assert!(cg.place(pid), "place the child pid into the cgroup");
753        assert!(cg.kill_all(), "write cgroup.kill");
754
755        let mut status = 0i32;
756        let reaped = unsafe { libc::waitpid(pid, &mut status, 0) };
757        assert_eq!(reaped, pid, "reaped the child");
758        assert!(
759            libc::WIFSIGNALED(status),
760            "child was SIGKILLed by cgroup.kill, not a clean exit (status={status})"
761        );
762        assert_eq!(libc::WTERMSIG(status), libc::SIGKILL, "killed by SIGKILL");
763        // `cg` Drop removes the now-empty cgroup dir.
764    }
765}