hyperfoot 0.2.2

//! Linux cgroup v2 fast path for exact peak-memory and disk-IO accounting.
//!
//! Sampling can miss short-lived processes and always undercounts peak
//! memory between polls. cgroup v2 exposes exact, kernel-maintained
//! aggregates (`memory.peak`, `io.stat`) for every process that ever
//! belonged to the group, even after it exits — no polling gap. We only use
//! this when the calling process's own cgroup subtree is writable (typical
//! under a systemd user session); otherwise callers fall back to sampling.

use std::fs;
use std::path::{Path, PathBuf};
use std::sync::OnceLock;
use std::thread;
use std::time::{Duration, SystemTime, UNIX_EPOCH};

pub struct CGroup {
    dir: PathBuf,
}

#[derive(Debug, Default)]
pub struct CGroupStats {
    pub peak_memory_bytes: Option<u64>,
    pub disk_read_bytes: Option<u64>,
    pub disk_write_bytes: Option<u64>,
}

static RUN_PARENT: OnceLock<Option<PathBuf>> = OnceLock::new();

impl CGroup {
    /// Creates a fresh, uniquely-named leaf cgroup to run one benchmarked
    /// command in. Returns `None` if cgroup v2 isn't mounted or delegation
    /// isn't available to us.
    pub fn try_create() -> Option<CGroup> {
        let parent = RUN_PARENT.get_or_init(setup_run_parent).clone()?;
        let dir = parent.join(format!(
            "hyperfoot-run-{}-{}",
            std::process::id(),
            unique_suffix()
        ));
        fs::create_dir(&dir).ok()?;
        Some(CGroup { dir })
    }

    pub fn procs_path(&self) -> PathBuf {
        self.dir.join("cgroup.procs")
    }

    /// Waits (briefly) for all processes to leave the cgroup. Should already
    /// be empty by the time this is called, since the caller waits for the
    /// top-level spawned process to exit first, and it only exits once its
    /// own children have been reaped.
    pub fn wait_drain(&self) {
        let procs_path = self.dir.join("cgroup.procs");
        for _ in 0..50 {
            match fs::read_to_string(&procs_path) {
                Ok(contents) if contents.trim().is_empty() => return,
                Err(_) => return,
                _ => thread::sleep(Duration::from_millis(10)),
            }
        }
    }

    pub fn read_stats(&self) -> CGroupStats {
        let peak_memory_bytes = fs::read_to_string(self.dir.join("memory.peak"))
            .ok()
            .and_then(|s| s.trim().parse::<u64>().ok());

        let disk = fs::read_to_string(self.dir.join("io.stat"))
            .ok()
            .map(|contents| parse_io_stat(&contents));

        CGroupStats {
            peak_memory_bytes,
            disk_read_bytes: disk.map(|(r, _)| r),
            disk_write_bytes: disk.map(|(_, w)| w),
        }
    }

    pub fn cleanup(self) {
        let _ = fs::remove_dir(&self.dir);
    }
}

fn parse_io_stat(contents: &str) -> (u64, u64) {
    let mut read_bytes = 0u64;
    let mut write_bytes = 0u64;
    for line in contents.lines() {
        for field in line.split_whitespace().skip(1) {
            if let Some(v) = field.strip_prefix("rbytes=") {
                read_bytes += v.parse::<u64>().unwrap_or(0);
            } else if let Some(v) = field.strip_prefix("wbytes=") {
                write_bytes += v.parse::<u64>().unwrap_or(0);
            }
        }
    }
    (read_bytes, write_bytes)
}

/// One-time setup, memoized for the life of the process: moves this process
/// into its own leaf cgroup, then returns the (now member-free) original
/// cgroup as the place to create per-run leaves under.
///
/// This dance exists because of cgroup v2's "no internal process"
/// constraint: a cgroup can't have both direct member processes and
/// controller-enabled children at the same time. Our own process starts out
/// as a direct member of its cgroup, so `cpu`/`memory`/`io` can't be
/// enabled in `cgroup.subtree_control` until we vacate it — the same trick
/// `systemd-run --scope` performs for its target process.
fn setup_run_parent() -> Option<PathBuf> {
    let base = own_cgroup_dir()?;
    cleanup_stale_leaves(&base);

    let self_leaf = base.join(format!("hyperfoot-self-{}", std::process::id()));
    fs::create_dir(&self_leaf).ok()?;
    fs::write(
        self_leaf.join("cgroup.procs"),
        std::process::id().to_string(),
    )
    .ok()?;

    enable_controllers(&base);
    Some(base)
}

/// Removes any of our own leaf cgroups left behind by a previous run that
/// has since exited (cgroup dirs aren't auto-removed once empty).
fn cleanup_stale_leaves(base: &Path) {
    let Ok(entries) = fs::read_dir(base) else {
        return;
    };
    for entry in entries.flatten() {
        let name = entry.file_name();
        if !name.to_string_lossy().starts_with("hyperfoot-") {
            continue;
        }
        let path = entry.path();
        let is_empty = fs::read_to_string(path.join("cgroup.procs"))
            .map(|s| s.trim().is_empty())
            .unwrap_or(false);
        if is_empty {
            let _ = fs::remove_dir(&path);
        }
    }
}

/// Resolves the cgroup v2 directory the current process already belongs to,
/// which is the only place we're guaranteed to have write access to create
/// a child cgroup under (typically a systemd-delegated user session slice).
fn own_cgroup_dir() -> Option<PathBuf> {
    let contents = fs::read_to_string("/proc/self/cgroup").ok()?;
    // cgroup v2 (unified hierarchy) processes have exactly one line: "0::<path>".
    let rel = contents.lines().find_map(|l| l.strip_prefix("0::"))?;
    let path = Path::new("/sys/fs/cgroup").join(rel.trim_start_matches('/'));
    path.is_dir().then_some(path)
}

fn enable_controllers(base: &Path) {
    let Ok(controllers) = fs::read_to_string(base.join("cgroup.controllers")) else {
        return;
    };
    let wanted = ["cpu", "memory", "io"];
    let enable: Vec<String> = wanted
        .into_iter()
        .filter(|c| {
            controllers
                .split_whitespace()
                .any(|available| available == *c)
        })
        .map(|c| format!("+{c}"))
        .collect();
    if !enable.is_empty() {
        let _ = fs::write(base.join("cgroup.subtree_control"), enable.join(" "));
    }
}

fn unique_suffix() -> u128 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map(|d| d.as_nanos())
        .unwrap_or(0)
}