Demi Marie Obenour <demiobenour@gmail.com> writes:

> This program has two modes:
>
> 1. Create a control group if it doesn't exist, optionally wait for other
>    programs in it to exit, and exec another program in it.
>
> 2. Purge a control group: kill all programs in it, then delete it.
>
> Locking is used to ensure that concurrent invocations are safe.
>
> Signed-off-by: Demi Marie Obenour <demiobenour@gmail.com>

The structure of the program is looking really good now.  All remaining
comments are minor, except for making sure we're doing the right thing
with enabling controllers.  Pay attention to naming — good names are
really important for making it clear to readers what a program does.

> diff --git a/tools/cgroup-setup/src/cgroup.rs b/tools/cgroup-setup/src/cgroup.rs
> new file mode 100644
> index 0000000000000000000000000000000000000000..463d3fc87ee34ccfbf399a337e2f49d9031728ee
> --- /dev/null
> +++ b/tools/cgroup-setup/src/cgroup.rs
> @@ -0,0 +1,289 @@
> +// SPDX-FileCopyrightText: 2026 Demi Marie Obenour <demiobenour@gmail.com>
> +// SPDX-License-Identifier: EUPL-1.2+
> +
> +use std::ffi::OsStr;
> +use std::fs::File;
> +use std::io::{Read as _, Seek as _, Write as _};
> +use std::os::unix::prelude::*;
> +
> +use std::path::{Component, Path, PathBuf};
> +
> +use rustix::fs::{AtFlags, CWD, Dir, FlockOperation};
> +use rustix::path;
> +use rustix::{
> +    fs::{Mode, OFlags, ResolveFlags},
> +    io::Errno,
> +};
> +
> +pub enum OpenFlags {
> +    Read,
> +    Write,
> +    Directory,
> +}
> +
> +#[derive(Debug)]
> +pub(crate) struct Cgroup {
> +    fd: Vec<OwnedFd>,
> +}
> +
> +impl AsFd for Cgroup {
> +    fn as_fd(&self) -> BorrowedFd<'_> {
> +        self.fd.last().unwrap().as_fd()
> +    }
> +}
> +
> +fn assert_single_component(component: &Path) {
> +    match component.as_os_str().as_bytes() {
> +        b"" | b"." | b".." => panic!("bad component"),
> +        c if c.contains(&b'\0') => panic!("NUL in component"),
> +        c if c.contains(&b'/') => panic!("/ in component"),
> +        _ => {}
> +    }
> +}
> +
> +// Wrapper around openat2() with better defaults.
> +pub fn openat2_simple(
> +    fd: impl AsFd,
> +    path: impl path::Arg,
> +    flags: OpenFlags,
> +) -> Result<OwnedFd, Errno> {
> +    rustix::fs::openat2(
> +        fd.as_fd(),
> +        path,
> +        OFlags::CLOEXEC
> +            | match flags {
> +                OpenFlags::Read => OFlags::RDONLY | OFlags::NOCTTY,
> +                OpenFlags::Write => OFlags::WRONLY | OFlags::NOCTTY,
> +                OpenFlags::Directory => OFlags::RDONLY | OFlags::DIRECTORY,
> +            },
> +        Mode::empty(),
> +        ResolveFlags::NO_SYMLINKS | ResolveFlags::NO_MAGICLINKS | ResolveFlags::NO_XDEV,
> +    )
> +}
> +
> +pub const DEFAULT_LEAF: &str = "$inner.service";
> +
> +pub fn check_path(path: &Path) -> Result<(), String> {
> +    let bytes = path.as_os_str().as_bytes();
> +    // Path::components() skips ., so use string manipulation instead.
> +    for component in bytes[path.is_absolute() as usize..].split(|&b| b == b'/') {
> +        if matches!(component, b"" | b"." | b"..") {
> +            return Err(format!("cgroup path {path:?} isn't canonical"));
> +        }
> +    }
> +    Ok(())
> +}
> +
> +// Remove all subdirectories of the given directory recursively, but not the
> +// directory itself.  The directory file descriptor is closed.

[nit] That's clear from the type signature, so probably doesn't need to
be explicitly documented.

> +//
> +// This isn't the most efficient possible algorithm, but simplicity is more
> +// important than performance in this case.  Also, it keeps open more file
> +// descriptors than strictly necessary, but Spectrum runs with a very high limit
> +// for the number of open file descriptors, and it uses shallow control group
> +// hierarchies.
> +//
> +// This uses a recursive algorithm, but so does std::fs::remove_dir_all().  Trying
> +// to be more robust than the standard library is not worthwhile.  In particular,
> +// the standard library function must be safe on systems where untrusted users (or
> +// even network endpoints!) can create deeply nested directory trees, whereas in
> +// Spectrum cgroups are only writeable by root.
> +fn remove_recursively(mut dirfd: Dir, remaining_depth: usize) -> Result<(), Errno> {
> +    if remaining_depth < 1 {
> +        panic!("control groups too deeply nested");
> +    }
> +    while let Some(element) = dirfd.next() {

"element" is a bit of an odd name for a directory entry, no?

> +        let parent_fd = dirfd.fd().unwrap();

Could be lifted out of the loop, right?

> +        let element = element.expect("Iterating through a cgroup directory failed?");
> +        let path = element.file_name();

"name" would probably be clearer than "path", since we know it's a
single component (and it's consistent with the file_name method).

> +        if element.file_type() != rustix::fs::FileType::Directory || path == c"." || path == c".." {
> +            continue;
> +        }
> +        let fd = openat2_simple(parent_fd, path, OpenFlags::Directory)?;
> +        remove_recursively(Dir::new(fd).unwrap(), remaining_depth - 1)?;
> +        rustix::fs::unlinkat(parent_fd, path, AtFlags::REMOVEDIR)?;
> +    }
> +    Ok(())
> +}
> +
> +fn assert_simple_path(current_cgroup: &Path) {
> +    let current_cgroup = current_cgroup.as_os_str().as_bytes();
> +    if !matches!(current_cgroup, b"" | b".") {
> +        for component in current_cgroup.split(|&b| b == b'/') {
> +            assert_single_component(Path::new(OsStr::from_bytes(component)));
> +        }
> +    }
> +}

Very non-obvious what this does — of course you'll get a lot of "single
component"s if you split a path on /.  What you actually want to do is
just check for no null bytes or .. components, right?  Given
assert_single_component is doing a more specific check than just single
components, it should be renamed accordingly.  (Although I'd struggle to
think of a name, because I still find what it's checking, and where we
check it, to be a bit arbitrary, especially when it's a path that's come
from the kernel…)

With check_path as well we have a confusing collection of subtly
different, inconsistently named path checking functions.  These should
be unified if possible, named systematically if not, and in either case
it should be clear from the name what the function is for.

> +// Convert the cgroup path to one relative to /sys/fs/cgroup.
> +//
> +// If the path starts with /, the leading / is removed and the result is returned
> +// without further processing.  Otherwise, the current cgroup is read from
> +// /proc/thread-self/cgroup.  If its last component is $inner.service, that is
> +// removed.  Finally, the current cgroup is prepended to the provided cgroup path,
> +// with a single / as separator.  The result of this operation is returned.
> +fn prepend_current_cgroup_if_needed(path: &Path) -> Result<PathBuf, String> {
> +    if let Ok(suffix) = path.strip_prefix("/") {
> +        return Ok(suffix.to_owned());
> +    }
> +    assert_simple_path(path);
> +    // /proc/thread-self is the same as /proc/self, except for the current thread
> +    // instead of the initial thread.  In this case, the two are identical, but
> +    // using /proc/thread-self is better practice as it is correct in more cases.
> +    // Reading /proc/thread-self/cgroup should never fail unless the system is
> +    // seriously broken.
> +    let current_cgroup =
> +        std::fs::read("/proc/thread-self/cgroup").expect("cannot read /proc/thread-self/cgroup");
> +    // Using this on a system without cgroups v2 mounted is user error
> +    // and not supported.
> +    let current_cgroup = current_cgroup
> +        .strip_prefix(b"0::/")
> +        .and_then(|e| e.strip_suffix(b"\n"))
> +        .ok_or_else(|| {
> +            "/proc/thread-self/cgroup doesn't start with 0::/ or doesn't end with a newline.\n\
> +            Either cgroups aren't in use at all, or you are using cgroups v1."
> +                .to_owned()
> +        })?;
> +    let mut current_cgroup = PathBuf::from(OsStr::from_bytes(current_cgroup));
> +    assert_simple_path(&current_cgroup);
> +    // Strip the implied $inner.service suffix.
> +    // This is used to satisfy the "no internal processes" rule.
> +    if current_cgroup.ends_with(Path::new(DEFAULT_LEAF)) {
> +        assert!(current_cgroup.pop());
> +    }
> +    // "." refers to the current cgroup.
> +    if path != Path::new(".") {
> +        current_cgroup.push(path);
> +    }
> +    Ok(current_cgroup)
> +}
> +
> +pub(crate) fn write_value(fd: &dyn AsFd, name: &Path, value: &[u8]) -> Result<(), String> {
> +    let fd = openat2_simple(fd, name, OpenFlags::Write)
> +        .map_err(|e| format!("Cannot open {name:?}: {e}"))?;
> +    File::from(fd).write_all(value).map_err(|e| {
> +        format!(
> +            "Cannot write {:?} to {name:?}: {e}",
> +            OsStr::from_bytes(value)
> +        )
> +    })
> +}
> +
> +impl Cgroup {
> +    pub fn new(path: &Path) -> Result<Self, String> {
> +        let cgroup_root = rustix::fs::openat2(
> +            CWD,
> +            Path::new("/sys/fs/cgroup"),
> +            OFlags::CLOEXEC | OFlags::DIRECTORY | OFlags::RDONLY,
> +            Mode::empty(),
> +            ResolveFlags::NO_SYMLINKS | ResolveFlags::NO_MAGICLINKS,
> +        )
> +        .map_err(|e| format!("Cannot open /sys/fs/cgroup: {e}"))?;
> +        // It's simpler to always have the root cgroup at the bottom of the stack,
> +        // even though no lock needs to be taken on it.  Otherwise, one would need
> +        // to special-case the cgroup root.  One could remove the first element if
> +        // there is more than one element in the vector, but that's not worth it.
> +        // cgroup-setup doesn't operate in an environment where FDs are a limited
> +        // resource.
> +        let mut cgroup = Self {
> +            fd: vec![cgroup_root],
> +        };
> +
> +        let path = prepend_current_cgroup_if_needed(path)?;
> +        for component in path.components() {
> +            let Component::Normal(component) = component else {
> +                unreachable!()
> +            };
> +            let sub_fd = openat2_simple(&cgroup, component, OpenFlags::Directory)
> +                .map_err(|e| format!("Cannot open sub-cgroup {component:?}: {e}"))?;
> +            // Take a shared lock on the cgroup.
> +            rustix::fs::flock(&sub_fd, FlockOperation::LockShared)
> +                .map_err(|e| format!("Cannot lock sub-cgroup {component:?}: {e}"))?;
> +            cgroup.fd.push(sub_fd);
> +        }
> +        Ok(cgroup)
> +    }
> +
> +    pub fn wait_for_empty(fd: &dyn AsFd) -> std::io::Result<()> {
> +        let wait_file = openat2_simple(fd, c"cgroup.events", OpenFlags::Read)?;
> +        let mut wait_fd = File::from(wait_file);
> +        let mut v = vec![];
> +        loop {
> +            v.clear();
> +            wait_fd
> +                .seek(std::io::SeekFrom::Start(0))
> +                .expect("Seek on control group file should succeed");
> +            wait_fd
> +                .read_to_end(&mut v)
> +                .expect("reading from control group should work");
> +            // Check that the cgroup isn't already empty.  If it was,
> +            // the kernel would not send an event and poll() would wait
> +            // forever.
> +            if v.split(|&c| c == b'\n').any(|line| line == b"populated 0") {
> +                break;
> +            }
> +            let mut fds = libc::pollfd {
> +                fd: wait_fd.as_raw_fd(),
> +                events: libc::POLLPRI | libc::POLLERR,
> +                revents: 0,
> +            };
> +            // SAFETY: FFI call, valid arguments, fds contains 1 element
> +            if unsafe { libc::poll(&raw mut fds, 1, -1) } != 1 {
> +                panic!("poll failed");
> +            }
> +        }
> +        Ok(())
> +    }
> +
> +    pub fn purge_child(&mut self, path: &Path) -> Result<(), String> {
> +        assert_single_component(path);
> +        // See if we can just delete the child directly.
> +        match rustix::fs::unlinkat(&self, path, AtFlags::REMOVEDIR) {
> +            // If the cgroup was successfully deleted, or if it
> +            // has already been deleted, we are done.
> +            Ok(()) | Err(Errno::NOENT) => return Ok(()),
> +            // If this cgroup is in use, keep going.
> +            Err(Errno::BUSY) => {}
> +            Err(e) => return Err(format!("Cannot purge {path:?}: {e}")),
> +        }
> +
> +        let sub_fd = match openat2_simple(&self, path, OpenFlags::Directory) {
> +            Ok(sub_fd) => sub_fd,
> +            Err(Errno::NOENT) => return Ok(()),
> +            Err(e) => {
> +                return Err(format!("Cannot open sub-cgroup {path:?}: {e}",));
> +            }
> +        };
> +
> +        // Take an exclusive lock on the cgroup that is about to be removed.  This
> +        // avoids concurrent executions of this program operating on deleted
> +        // sub-cgroups.  Dir::new() doesn't expose a reference to its internal FD
> +        // so it must be delayed until later.

Yes it does?  It's Dir::fd.  You used it elsewhere already.  It's fine
to delay Dir::new but this comment is not correct.

> +        rustix::fs::flock(&sub_fd, FlockOperation::LockExclusive)
> +            .map_err(|e| format!("Cannot lock sub-cgroup: {e}"))?;
> +
> +        // Kill all processes in the child cgroup.
> +        write_value(&sub_fd, Path::new("cgroup.kill"), b"1")?;
> +
> +        // Wait for the child cgroup to become empty.
> +        Self::wait_for_empty(&sub_fd)
> +            .map_err(|e| format!("Cannot wait for cgroup to become empty: {e}"))
> +            .inspect_err(|_| {
> +                self.fd.pop().unwrap();

Why do we need to do this?  What's the matching push?  Why should
failing to wait for a child cgroup to be empty mean we unlock its
parent?

> +            })?;
> +
> +        // Remove the child cgroup and its contents recursively.
> +        remove_recursively(Dir::new(sub_fd).unwrap(), 1000)
> +            .map_err(|e| format!("Cannot remove: {e}"))?;
> +
> +        // Delete the cgroup.  If it's been re-created in the meantime and is
> +        // currently in use, this is not an error.  Another process deleting the
> +        // cgroup is also not an error.  Both of these can happen because of the
> +        // time period between remove_child_directories() closing the file
> +        // descriptor (releasing its lock) and the above call to flock().
> +        match rustix::fs::unlinkat(&self, path, AtFlags::REMOVEDIR) {
> +            Ok(()) | Err(Errno::BUSY) | Err(Errno::NOENT) => Ok(()),
> +            Err(e) => Err(format!("Cannot delete: {e}")),
> +        }
> +    }
> +}
> diff --git a/tools/cgroup-setup/src/main.rs b/tools/cgroup-setup/src/main.rs
> new file mode 100644
> index 0000000000000000000000000000000000000000..c757a4ad37c812ef5ce5249dc1bc3104ec246eef
> --- /dev/null
> +++ b/tools/cgroup-setup/src/main.rs
> @@ -0,0 +1,168 @@
> +// SPDX-FileCopyrightText: 2026 Demi Marie Obenour <demiobenour@gmail.com>
> +// SPDX-License-Identifier: EUPL-1.2+
> +
> +mod cgroup;
> +
> +use cgroup::{Cgroup, OpenFlags, openat2_simple, write_value};
> +use rustix::{
> +    fs::{FlockOperation, Mode, XattrFlags},
> +    io::Errno,
> +};
> +use std::{
> +    env::ArgsOs,
> +    fs::File,
> +    io::Read as _,
> +    os::unix::prelude::*,
> +    path::{Path, PathBuf},
> +};
> +
> +fn enable_subtree_control(fd: &dyn AsFd) -> Result<(), String> {
> +    rustix::fs::fsetxattr(fd, c"user.delegate", b"1", XattrFlags::empty())
> +        .map_err(|e| format!("Cannot enable cgroup delegation: {e}"))?;
> +    let mut buf = Vec::new();
> +    File::from(
> +        openat2_simple(fd, c"cgroup.controllers", OpenFlags::Read)
> +            .map_err(|e| format!("Cannot open cgroup.controllers: {e}"))?,
> +    )
> +    .read_to_end(&mut buf)
> +    .map_err(|e| format!("Cannot read cgroup.controllers: {e}"))?;
> +    let mut subtree = vec![];
> +    if buf.is_empty() {
> +        return Ok(());
> +    }
> +    for controller in buf.split(|&b| b == b' ') {
> +        if !subtree.is_empty() {
> +            subtree.push(b' ');
> +        }
> +        subtree.push(b'+');
> +        subtree.extend_from_slice(controller);
> +    }
> +    if !subtree.is_empty() {
> +        write_value(&fd, Path::new("cgroup.subtree_control"), &subtree)?;
> +    }
> +    Ok(())
> +}

My memory of our conversation on a call last week is that we found it
undesirable to enable every controller, since that causes behaviour
surprising action-at-a-distance behaviour changes.  Rather specific
requested controllers should be enabled when necessary, right?

> +
> +fn spawn_in_cgroup(
> +    mut args: std::iter::Peekable<ArgsOs>,
> +    cgroup: Option<&dyn AsFd>,
> +) -> Result<(), String> {

We're not spawning anything if all we're doing is an exec.  It should be
called exec_in_cgroup.

> +    let Some(program_name) = args.next() else {
> +        return Ok(());
> +    };
> +    if let Some(cgroup) = cgroup {
> +        let pid = std::process::id().to_string();
> +        write_value(
> +            cgroup,
> +            Path::new("$inner.service/cgroup.procs"),
> +            pid.as_bytes(),
> +        )
> +        .map_err(|e| format!("Cannot move process to child cgroup: {e}"))?;
> +    }
> +    let e = std::process::Command::new(&program_name).args(args).exec();
> +    Err(format!("Cannot spawn child {program_name:?}: {e}"))

Cannot *exec*.

> +}
> +
> +// Check that the path is canonical,
> +// then split it into basename and filename.
> +fn split_path(path: &Path) -> Result<(&Path, &Path), String> {
> +    cgroup::check_path(path)?;
> +    Ok((path.parent().unwrap(), Path::new(path.file_name().unwrap())))
> +}
> +
> +fn cgroup_setup(args: ArgsOs) -> Result<(), String> {
> +    let mut wait = true;
> +    let mut args = args.peekable();
> +    while let Some(arg) = args.peek() {
> +        if !arg.as_bytes().starts_with(b"-") {
> +            break;
> +        }
> +        let arg = args.next().unwrap();

I'd find let _ = args.next() slightly clearer, because then it's clear
we're not interested in the value, since we already have it.

> +        let Some(option) = arg.as_bytes().strip_prefix(b"--") else {
> +            return Err("takes no short options".to_owned());
> +        };
> +        match option {
> +            b"" => break,
> +            b"no-wait" => wait = false,
> +            _ => return Err(format!("unknown long option {arg:?}")),
> +        }
> +    }
> +    let Some(cgroup_path) = args.next().map(PathBuf::from) else {
> +        return Err("have no positional arguments, expected at least 1".to_owned());
> +    };
> +
> +    let (parent_cgroup_path, child_cgroup_path) = split_path(&cgroup_path)?;
> +    let cgroup = Cgroup::new(parent_cgroup_path)?;
> +    match rustix::fs::mkdirat(&cgroup, child_cgroup_path, Mode::from_raw_mode(0o755)) {
> +        Ok(()) | Err(Errno::EXIST) => {}
> +        Err(e) => {
> +            return Err(format!(
> +                "Cannot create child cgroup {child_cgroup_path:?}: {e}"
> +            ));
> +        }
> +    }
> +
> +    let child = openat2_simple(&cgroup, child_cgroup_path, OpenFlags::Directory)
> +        .map_err(|e| format!("Cannot open child cgroup: {e}"))?;
> +
> +    // While waiting, hold an exclusive lock on the child.
> +    // This avoids two processes both waiting for the same cgroup to become
> +    // empty, then spawning processes in the same cgroup.
> +    rustix::fs::flock(&child, FlockOperation::LockExclusive)
> +        .map_err(|e| format!("Cannot take an exclusive lock on child cgroup: {e}"))?;
> +    if wait {
> +        Cgroup::wait_for_empty(&child)
> +            .map_err(|e| format!("Cannot wait for {parent_cgroup_path:?} to be empty: {e}"))?;
> +    }
> +
> +    // Spectrum's programs (such as this one) expect cgroup.subtree_control
> +    // to be set by the program that created the cgroup.  systemd-aware
> +    // programs, like systemd-udevd, expect user.delegate=1 to be set.
> +    enable_subtree_control(&child)?;
> +
> +    // If the child process will need to manage cgroups itself, it will need
> +    // to set up a sub-cgroup due to the "no internal processes" rule.  It's
> +    // simplest to just do it automatically.  If the cgroup already exists,
> +    // that isn't an error.
> +    match rustix::fs::mkdirat(&child, cgroup::DEFAULT_LEAF, Mode::from_raw_mode(0o755)) {
> +        Ok(()) | Err(Errno::EXIST) => {}
> +        Err(e) => return Err(format!("Cannot create $inner.service cgroup: {e}")),
> +    }
> +
> +    spawn_in_cgroup(args, Some(&child))
> +}
> +
> +fn cgroup_purge(mut args: ArgsOs) -> Result<(), String> {
> +    if args.len() != 1 {
> +        return Err("usage: cgroup-purge CGROUP_TO_PURGE".to_owned());
> +    }
> +    let arg = args.next().unwrap();
> +    let (parent, child) = split_path(Path::new(&arg))?;
> +    Cgroup::new(parent)?.purge_child(child)
> +}
> +
> +fn run(prog_name: &Path, args: ArgsOs) -> Result<(), String> {
> +    match prog_name.file_name().map(|f| f.as_bytes()) {
> +        Some(b"cgroup-setup") => cgroup_setup(args),
> +        Some(b"cgroup-purge") => cgroup_purge(args),
> +        _ => Err(format!(
> +            "must be invoked as \"cgroup-setup\" or \
> +                 \"cgroup-purge\", got {prog_name:?}",
> +        )),
> +    }
> +}
> +
> +fn main() {
> +    let mut args = std::env::args_os();
> +    let Some(prog_name) = args.next() else {
> +        eprintln!("No command line arguments (argv[0] is NULL)");
> +        std::process::exit(1);
> +    };
> +    match run(Path::new(&prog_name), args) {
> +        Ok(()) => {}
> +        Err(e) => {
> +            eprintln!("{prog_name:?}: {}", e);
> +            std::process::exit(1);
> +        }
> +    }
> +}