//! One cgroup per run, under a daemon unit that systemd delegated with
//! `Delegate=memory pids` and `DelegateSubgroup=daemon`
//! (`nix/agent-modules/mcp.nix`). The daemon process lives in the unit's
//! `daemon/` leaf; each run gets `<unit>/<kind>/<name>/` with
//! `memory.oom.group=1`, so the OOM killer takes a whole run, never part of
//! one and never the daemon.
//!
//! systemd enables no controller below a service's own cgroup, so
//! [`RunCgroups::init`] writes `+memory +pids` into the unit's and the
//! `<kind>/` node's `cgroup.subtree_control` itself. Without that the run
//! cgroups would have no `memory.oom.group` file.
//!
//! Two ways into a run's cgroup, both before the run's program executes, so
//! everything it later forks is born inside:
//! - [`enter_in_pre_exec`], from a `Command::pre_exec` hook, for a spawn
//!   this crate's caller builds itself;
//! - [`entry`], for a spawn built by a library with no such hook:
//!   it runs this daemon's own binary, whose `main` calls [`exec_entry`]
//!   first thing, which enters the cgroup and then execs the real program.
//!   It also sets the run's `oom_score_adj`, and does that alone when the
//!   run has no cgroup.

use std::ffi::OsStr;
use std::fs::File;
use std::io;
use std::os::fd::RawFd;
use std::os::unix::process::CommandExt;
use std::path::{Path, PathBuf};
use std::time::Duration;

/// The leaf systemd puts the daemon process in. Must equal the unit's
/// `DelegateSubgroup=` in `nix/agent-modules/mcp.nix`.
const DAEMON_SUBGROUP: &str = "daemon";

const CGROUP_ROOT: &str = "/sys/fs/cgroup";
const CONTROLLERS: &str = "+memory +pids";

/// The program [`exec_entry`] execs. Its presence is what makes a process
/// an entry wrapper rather than the daemon.
const ENTRY_PROGRAM_ENV: &str = "HIVE_RUN_ENTRY_PROGRAM";
const ENTRY_CGROUP_ENV: &str = "HIVE_RUN_ENTRY_CGROUP";
const ENTRY_OOM_SCORE_ADJ_ENV: &str = "HIVE_RUN_ENTRY_OOM_SCORE_ADJ";

/// How long a run's cgroup may stay populated after the run ended before
/// removal is left to the next daemon start.
const RELEASE_GRACE: Duration = Duration::from_secs(30);
const RELEASE_POLL: Duration = Duration::from_millis(250);

/// The `<unit>/<kind>/` node run cgroups are created under.
#[derive(Debug)]
pub struct RunCgroups {
    parent: PathBuf,
}

impl RunCgroups {
    /// Set up `<unit>/<kind>/` for this process's unit. `None`, with one
    /// warning, when this process is not in a delegated `daemon/` leaf on
    /// cgroup v2 (a test, a dev shell, a unit without `Delegate=`): runs then
    /// stay in the daemon's own cgroup.
    #[must_use]
    pub fn init(kind: &str) -> Option<Self> {
        let result = std::fs::read_to_string("/proc/self/cgroup")
            .and_then(|own| Self::init_at(Path::new(CGROUP_ROOT), &own, kind));
        match result {
            Ok(cgroups) => {
                tracing::info!(parent = %cgroups.parent.display(), "run cgroups ready");
                Some(cgroups)
            }
            Err(e) => {
                tracing::warn!(
                    error = %e,
                    "run cgroups unavailable; runs share the daemon's cgroup"
                );
                None
            }
        }
    }

    /// [`RunCgroups::init`] against cgroup2 mounted at `root`, with
    /// `proc_self_cgroup` the content of `/proc/self/cgroup`.
    fn init_at(root: &Path, proc_self_cgroup: &str, kind: &str) -> io::Result<Self> {
        let own = proc_self_cgroup
            .lines()
            .find_map(|line| line.strip_prefix("0::"))
            .ok_or_else(|| io::Error::other("no cgroup v2 entry in /proc/self/cgroup"))?;
        let own = root.join(own.trim().trim_start_matches('/'));
        if own.file_name().and_then(|n| n.to_str()) != Some(DAEMON_SUBGROUP) {
            return Err(io::Error::other(format!(
                "{} is not a delegated `{DAEMON_SUBGROUP}` leaf",
                own.display()
            )));
        }
        let unit = own
            .parent()
            .ok_or_else(|| io::Error::other("daemon leaf has no parent"))?;
        std::fs::write(unit.join("cgroup.subtree_control"), CONTROLLERS)?;
        let parent = unit.join(kind);
        match std::fs::create_dir(&parent) {
            Ok(()) => {}
            Err(e) if e.kind() == io::ErrorKind::AlreadyExists => {}
            Err(e) => return Err(e),
        }
        std::fs::write(parent.join("cgroup.subtree_control"), CONTROLLERS)?;
        // Runs that outlived the previous daemon's grace leave empty dirs;
        // a populated one belongs to a process still alive and stays.
        for entry in std::fs::read_dir(&parent)?.flatten() {
            if entry.file_type().is_ok_and(|t| t.is_dir()) {
                let _ = std::fs::remove_dir(entry.path());
            }
        }
        Ok(Self { parent })
    }

    /// Create the run cgroup `name`, with `memory.oom.group` set.
    ///
    /// # Errors
    ///
    /// `AlreadyExists` when a populated cgroup of that name is still there,
    /// or the underlying failure creating it or writing `memory.oom.group`.
    pub fn create(&self, name: &str) -> io::Result<RunCgroup> {
        let dir = self.parent.join(name);
        match std::fs::create_dir(&dir) {
            Err(e) if e.kind() == io::ErrorKind::AlreadyExists => {
                std::fs::remove_dir(&dir).map_err(|_| e)?;
                std::fs::create_dir(&dir)?;
            }
            other => other?,
        }
        let cgroup = RunCgroup { dir };
        std::fs::write(cgroup.dir.join("memory.oom.group"), "1")?;
        Ok(cgroup)
    }
}

/// One run's cgroup. Dropping it removes the directory once the run's last
/// process has left, waiting up to 30 seconds on a background thread.
#[derive(Debug)]
pub struct RunCgroup {
    dir: PathBuf,
}

impl RunCgroup {
    /// The run's `cgroup.procs`, open for [`enter_in_pre_exec`]. Keep it
    /// open until the spawn has returned.
    ///
    /// # Errors
    ///
    /// The failure opening the file.
    pub fn procs(&self) -> io::Result<File> {
        File::options()
            .write(true)
            .open(self.dir.join("cgroup.procs"))
    }
}

/// The program and extra environment that run `program` at `oom_score_adj`,
/// inside `cgroup` when there is one: this daemon's own binary, whose `main`
/// starts with [`exec_entry`].
#[must_use]
pub fn entry(
    program: &Path,
    oom_score_adj: i32,
    cgroup: Option<&RunCgroup>,
) -> (PathBuf, Vec<(String, String)>) {
    let mut env = vec![
        (
            ENTRY_PROGRAM_ENV.to_owned(),
            program.to_string_lossy().into_owned(),
        ),
        (
            ENTRY_OOM_SCORE_ADJ_ENV.to_owned(),
            oom_score_adj.to_string(),
        ),
    ];
    if let Some(cgroup) = cgroup {
        env.push((
            ENTRY_CGROUP_ENV.to_owned(),
            cgroup.dir.to_string_lossy().into_owned(),
        ));
    }
    (PathBuf::from("/proc/self/exe"), env)
}

impl Drop for RunCgroup {
    fn drop(&mut self) {
        match std::fs::remove_dir(&self.dir) {
            Err(e) if e.raw_os_error() == Some(libc::EBUSY) => {
                let dir = std::mem::take(&mut self.dir);
                std::thread::spawn(move || release_when_empty(&dir));
            }
            Err(e) if e.kind() != io::ErrorKind::NotFound => {
                tracing::warn!(dir = %self.dir.display(), error = %e, "removing run cgroup failed");
            }
            _ => {}
        }
    }
}

/// Retry removing a still-populated run cgroup until it empties or
/// [`RELEASE_GRACE`] runs out; past that the next daemon start removes it.
fn release_when_empty(dir: &Path) {
    let deadline = std::time::Instant::now() + RELEASE_GRACE;
    while std::time::Instant::now() < deadline {
        std::thread::sleep(RELEASE_POLL);
        match std::fs::remove_dir(dir) {
            Err(e) if e.raw_os_error() == Some(libc::EBUSY) => {}
            _ => return,
        }
    }
    tracing::info!(dir = %dir.display(), "run cgroup still populated after the run ended; left in place");
}

/// Move the calling process into the cgroup whose `cgroup.procs` is `procs`.
/// Only `write(2)`, so it is safe between fork and exec.
///
/// # Errors
///
/// The `write(2)` failure.
pub fn enter_in_pre_exec(procs: RawFd) -> io::Result<()> {
    // SAFETY: a plain write of a static buffer to a caller-owned fd.
    if unsafe { libc::write(procs, b"0".as_ptr().cast(), 1) } < 0 {
        return Err(io::Error::last_os_error());
    }
    Ok(())
}

/// When this process was started by [`entry`]: enter the run's
/// cgroup if it has one, set its `oom_score_adj`, and exec the run's program with this
/// process's arguments; on failure print the reason and exit 127. Otherwise
/// return, doing nothing. Call it first in `main`.
pub fn exec_entry() {
    let Some(program) = std::env::var_os(ENTRY_PROGRAM_ENV) else {
        return;
    };
    let err = enter_and_exec(&program);
    eprintln!("run cgroup entry failed: {err}");
    std::process::exit(127);
}

fn enter_and_exec(program: &OsStr) -> io::Error {
    if let Some(cgroup) = std::env::var_os(ENTRY_CGROUP_ENV)
        && let Err(e) = std::fs::write(Path::new(&cgroup).join("cgroup.procs"), "0")
    {
        return e;
    }
    if let Some(adj) = std::env::var_os(ENTRY_OOM_SCORE_ADJ_ENV)
        && let Err(e) = std::fs::write("/proc/self/oom_score_adj", adj.as_encoded_bytes())
    {
        return e;
    }
    std::process::Command::new(program)
        .args(std::env::args_os().skip(1))
        .env_remove(ENTRY_PROGRAM_ENV)
        .env_remove(ENTRY_CGROUP_ENV)
        .env_remove(ENTRY_OOM_SCORE_ADJ_ENV)
        .exec()
}

#[cfg(test)]
mod tests {
    use super::*;

    const KIND: &str = "runs";

    fn unit_tree() -> (tempfile::TempDir, PathBuf) {
        let root = tempfile::tempdir().expect("tempdir");
        let unit = root.path().join("system.slice/x.service");
        std::fs::create_dir_all(unit.join(DAEMON_SUBGROUP)).expect("daemon leaf");
        (root, unit)
    }

    const OWN: &str = "0::/system.slice/x.service/daemon\n";

    #[test]
    fn init_enables_controllers_on_the_unit_and_the_kind_node() {
        let (root, unit) = unit_tree();
        let cgroups = RunCgroups::init_at(root.path(), OWN, KIND).expect("init");
        assert_eq!(cgroups.parent, unit.join(KIND));
        for node in [&unit, &unit.join(KIND)] {
            let control =
                std::fs::read_to_string(node.join("cgroup.subtree_control")).expect("written");
            assert_eq!(control, CONTROLLERS);
        }
    }

    #[test]
    fn init_refuses_a_process_outside_a_daemon_leaf() {
        let (root, unit) = unit_tree();
        let err = RunCgroups::init_at(root.path(), "0::/system.slice/x.service\n", KIND)
            .expect_err("not a daemon leaf");
        assert!(err.to_string().contains(DAEMON_SUBGROUP), "{err}");
        assert!(!unit.join("cgroup.subtree_control").exists());
    }

    #[test]
    fn init_refuses_cgroup_v1() {
        let (root, _unit) = unit_tree();
        assert!(RunCgroups::init_at(root.path(), "1:memory:/x\n", KIND).is_err());
    }

    #[test]
    fn init_removes_empty_leftover_runs_and_keeps_populated_ones() {
        let (root, unit) = unit_tree();
        let empty = unit.join(KIND).join("old-0");
        let populated = unit.join(KIND).join("old-1");
        std::fs::create_dir_all(&empty).expect("empty");
        std::fs::create_dir_all(&populated).expect("populated");
        std::fs::write(populated.join("cgroup.procs"), "42").expect("procs");
        RunCgroups::init_at(root.path(), OWN, KIND).expect("init");
        assert!(!empty.exists());
        assert!(populated.exists());
    }

    #[test]
    fn create_sets_oom_group_and_drop_removes_the_dir() {
        let (root, unit) = unit_tree();
        let cgroups = RunCgroups::init_at(root.path(), OWN, KIND).expect("init");
        let run = cgroups.create("alpha-0").expect("create");
        let dir = unit.join(KIND).join("alpha-0");
        assert_eq!(run.dir, dir);
        assert_eq!(
            std::fs::read_to_string(dir.join("memory.oom.group")).expect("written"),
            "1"
        );
        // cgroupfs drops a cgroup's interface files with the directory; a
        // plain tempdir needs them gone first.
        std::fs::remove_file(dir.join("memory.oom.group")).expect("rm");
        drop(run);
        assert!(!dir.exists());
    }

    #[test]
    fn create_replaces_an_empty_leftover_of_the_same_name() {
        let (root, unit) = unit_tree();
        let cgroups = RunCgroups::init_at(root.path(), OWN, KIND).expect("init");
        std::fs::create_dir(unit.join(KIND).join("t1")).expect("leftover");
        cgroups.create("t1").expect("create over an empty leftover");
    }

    #[test]
    fn entry_names_this_binary_and_the_run() {
        let (root, _unit) = unit_tree();
        let cgroups = RunCgroups::init_at(root.path(), OWN, KIND).expect("init");
        let run = cgroups.create("alpha-0").expect("create");
        let (program, env) = entry(Path::new("claude"), 1000, Some(&run));
        assert_eq!(program, Path::new("/proc/self/exe"));
        let get = |k: &str| env.iter().find(|(n, _)| n == k).map(|(_, v)| v.as_str());
        assert_eq!(get(ENTRY_PROGRAM_ENV), Some("claude"));
        assert_eq!(get(ENTRY_CGROUP_ENV), run.dir.to_str());
        assert_eq!(get(ENTRY_OOM_SCORE_ADJ_ENV), Some("1000"));
    }

    #[test]
    fn entry_without_a_cgroup_still_sets_the_score() {
        let (program, env) = entry(Path::new("claude"), 1000, None);
        assert_eq!(program, Path::new("/proc/self/exe"));
        let get = |k: &str| env.iter().find(|(n, _)| n == k).map(|(_, v)| v.as_str());
        assert_eq!(get(ENTRY_PROGRAM_ENV), Some("claude"));
        assert_eq!(get(ENTRY_OOM_SCORE_ADJ_ENV), Some("1000"));
        assert_eq!(get(ENTRY_CGROUP_ENV), None);
    }

    #[test]
    fn enter_in_pre_exec_writes_zero_to_the_fd() {
        let dir = tempfile::tempdir().expect("tempdir");
        let path = dir.path().join("cgroup.procs");
        std::fs::write(&path, "").expect("create");
        let procs = File::options().write(true).open(&path).expect("open");
        enter_in_pre_exec(std::os::fd::AsRawFd::as_raw_fd(&procs)).expect("write");
        assert_eq!(std::fs::read_to_string(&path).expect("read"), "0");
    }
}