Re: [PATCH v5 02/19] tools: Add control group manager

[Alyssa Ross <hi@alyssa.is>]

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

> The cgroup-setup Rust program can create and purge cgroups.  It can also
> wait for one to become empty, spawn a program in a cgroup, and more.  In
> the future, it will also support cgroup-based resource control.  Locking
> is used to ensure that concurrent invocations are safe.
>
> Signed-off-by: Demi Marie Obenour <demiobenour@gmail.com>
> ---
>  .codespellrc                          |   2 +-
>  host/rootfs/default.nix               |   6 +-
>  pkgs/default.nix                      |   1 +
>  tools/cgroup-setup/Cargo.lock         |  67 ++++++++
>  tools/cgroup-setup/Cargo.lock.license |   2 +
>  tools/cgroup-setup/Cargo.toml         |  10 ++
>  tools/cgroup-setup/default.nix        |  22 +++
>  tools/cgroup-setup/src/cgroup.rs      | 308 ++++++++++++++++++++++++++++++++++
>  tools/cgroup-setup/src/main.rs        | 186 ++++++++++++++++++++
>  9 files changed, 600 insertions(+), 4 deletions(-)

Looking much better, thank you!

> +fn push_child_fds(fds: &mut Vec<(Rc<RefCell<Dir>>, PathBuf)>, fd: OwnedFd) {

Would it not make more sense to take Dir than OwnedFd?

> +    // The rustix source code shows that Dir::new() never fails.
> +    let child_fd = Rc::new(RefCell::new(Dir::new(fd).unwrap()));
> +    while let Some(element) = child_fd.borrow_mut().next() {
> +        let element = element.expect("Iterating through a cgroup directory failed?");
> +        if element.file_type() != rustix::fs::FileType::Directory {
> +            continue;
> +        }
> +        match element.file_name().to_bytes() {
> +            b"." | b".." => {}
> +            other => {
> +                let other = Path::new(OsStr::from_bytes(other)).to_owned();
> +                assert_single_component(&other);
> +                fds.push((child_fd.clone(), other));

The data structures used here are still very confusing.  Why are we
storing a reference to the same file descriptor in every entry in the
Vec?

> +            }
> +        }
> +    }
> +}
> +
> +// Remove all subdirectories of the given directory recursively,
> +// but not the directory itself.  The directory file descriptor
> +// is closed.
> +//
> +// 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.
> +fn remove_child_directories(dirfd: OwnedFd) -> Result<(), Errno> {
> +    let mut fds = Vec::new();
> +    // Push the children of this directory onto the stack.
> +    push_child_fds(&mut fds, dirfd);
> +    while let Some((d, path)) = fds.pop() {

Couldn't we call push_child_fds() once here, rather than twice as is
currently done?  (And then consider inlining it, depending on how
complex it's looking at the time.)

> +        assert_single_component(&path);
> +        // Try to delete the directory.  If that fails because there are child
> +        // directories, push the child directories onto the stack, then push
> +        // this directory again.
> +        match rustix::fs::unlinkat(
> +            // The rustix source code shows that Dir::fd() never fails.
> +            d.borrow().fd().unwrap(),
> +            &path,
> +            AtFlags::REMOVEDIR,
> +        ) {
> +            Err(Errno::NOTEMPTY) => {}
> +            Ok(()) => continue,
> +            Err(bad) => return Err(bad),
> +        }
> +        let fd = openat2_simple(
> +            // The rustix source code shows that Dir::fd() never fails.
> +            &d.borrow().fd().unwrap(),
> +            &path,
> +            OFlags::DIRECTORY | OFlags::RDONLY,
> +        )?;
> +        // Process child directories first, then attempt to delete the
> +        // directory again.
> +        fds.push((d, path));
> +        push_child_fds(&mut fds, fd);
> +    }
> +    Ok(())
> +}

> +impl Cgroup {
> +    pub fn open_beneath(&self, path: &Path, flags: OFlags) -> Result<OwnedFd, Errno> {
> +        openat2_simple(self, path, flags)
> +    }

This method looks pretty redundant now.

> +    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}"))?;

Can we not use openat2_simple here?  It's missing e.g. NOCTTY.

> +        let mut cgroup = Self {
> +            fd: vec![(cgroup_root)],
> +        };
> +
> +        let path = prepend_current_cgroup_if_needed(path);
> +        for component in path.components() {
> +            let component = match component {
> +                Component::Normal(component) => component,
> +                _ => unreachable!(),
> +            };

I think it would be slightly more idiomatic to do:

let Component::Normal(component) = component else {
  unreachable!()
};

> +            let sub_fd = cgroup
> +                .open_beneath(Path::new(component), OFlags::RDONLY | OFlags::DIRECTORY)
> +                .map_err(|e| format!("Cannot open sub-cgroup {component:?}: {e}"))?;
> +            // Take a shared lock on the *previous* file descriptor.
> +            rustix::fs::flock(&cgroup, FlockOperation::LockShared)
> +                .map_err(|e| format!("Cannot lock sub-cgroup {component:?}: {e}"))?;
> +            cgroup.fd.push(sub_fd);
> +        }
> +        // Take an exclusive lock on the final file descriptor.
> +        rustix::fs::flock(&cgroup, FlockOperation::LockExclusive)
> +            .map_err(|e| format!("Cannot lock {path:?}: {e}"))?;
> +        Ok(cgroup)
> +    }
> +
> +    pub fn wait_for_empty(fd: &dyn AsFd) -> std::io::Result<()> {
> +        let wait_file = openat2_simple(fd, Path::new("cgroup.events"), OFlags::RDONLY)?;
> +        let poll_fd = wait_file.as_raw_fd();
> +        let mut wait_fd = File::from(wait_file);
> +        let mut fds = libc::pollfd {
> +            fd: poll_fd,

I would inline poll_fd here.  RawFd is easy to misuse, so I like to
avoid having them hang around.

> +            events: libc::POLLPRI | libc::POLLERR,
> +            revents: 0,
> +        };
> +        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;
> +            }
> +            // 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::new(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 self.open_beneath(path, OFlags::RDONLY | OFlags::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.
> +        rustix::fs::flock(&sub_fd, FlockOperation::LockExclusive)
> +            .map_err(|e| format!("Cannot lock sub-cgroup: {e}"))?;
> +
> +        // Drop the exclusive lock on the original cgroup,
> +        // This avoids blocking concurrent operations on other
> +        // child cgroups while the cgroup is being purged,
> +        // or while waiting for programs to exit.
> +        rustix::fs::flock(&self, FlockOperation::LockShared)
> +            .map_err(|e| format!("Cannot relock: {e}"))?;

Could you add some extra explanation here of why it's okay for the
exclusive lock to be temporarily dropped here?

I'm wondering whether taking a lock, then dropping it temporarily is a
sign that we're taking the lock too early in the first place, and should
scope it better to where it's actually needed.

> +
> +        // 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();
> +            })?;
> +
> +        // Remove the child cgroup and its contents recursively.
> +        remove_child_directories(sub_fd).map_err(|e| format!("Cannot remove: {e}"))?;
> +        // Re-take an exclusive lock on the parent of the cgroup being purged.
> +        // Otherwise, a concurrent instance of cgroup-setup might create a cgroup
> +        // only for this one to delete it.  The other instance could then try to
> +        // create a sub-cgroup of a deleted cgroup, which would fail.  Waiting
> +        // until nobody is using the parent cgroup ensures these problems can't
> +        // happen.
> +        //
> +        // This must happen *after* the lock on the cgroup being purged is released.
> +        // Another instance of the program might have a shared lock on the parent
> +        // and be waiting for an exclusive lock on the child.  Trying to take an
> +        // exclusive lock on the parent while a lock is held on the child would
> +        // result in an ABBA deadlock.
> +        rustix::fs::flock(&self, FlockOperation::LockExclusive)
> +            .map_err(|e| format!("Cannot re-lock exclusively: {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..e8d9e9d7c2111857cc25b36433d8f33ddb28c27b
> --- /dev/null
> +++ b/tools/cgroup-setup/src/main.rs
> @@ -0,0 +1,186 @@
> +// SPDX-License-Identifier: EUPL-1.2+
> +// SPDX-FileCopyrightText: 2026 Demi Marie Obenour <demiobenour@gmail.com>
> +
> +mod cgroup;
> +
> +use cgroup::{Cgroup, openat2_simple, write_value};
> +use rustix::{
> +    fs::{FlockOperation, Mode, OFlags, XattrFlags},
> +    io::Errno,
> +};
> +use std::{
> +    env::ArgsOs,
> +    ffi::OsStr,
> +    fs::File,
> +    io::Read as _,
> +    os::unix::prelude::*,
> +    path::{Path, PathBuf},
> +};
> +
> +// 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)
> +        .map(|()| (path.parent().unwrap(), Path::new(path.file_name().unwrap())))

Doing this with map rather than ? is a little strange.

> +}
> +
> +fn read_control_file(fd: &dyn AsFd, p: &Path) -> Result<Vec<u8>, String> {
> +    let mut buf = Vec::new();
> +    File::from(
> +        openat2_simple(&fd, Path::new(p), OFlags::RDONLY)
> +            .map_err(|e| format!("Cannot open {p:?}: {e}"))?,
> +    )
> +    .read_to_end(&mut buf)
> +    .map_err(|e| format!("Cannot read {p:?}: {e}"))?;
> +    Ok(buf)
> +}

Nothing control-file-specific about this method.  It just reads a file.
And a bit odd for write_value to be in cgroup.rs while this is here.

> +
> +fn enable_subtree_control(fd: &dyn AsFd) -> Result<(), String> {

Would it not make sense for this to be an instance method on Cgroup,
since it's a Cgroup-specific operation?

> +    let p = Path::new("cgroup.controllers");
> +    let buf = read_control_file(fd, p)?;

p is only used here, so can just be inlined.  If read_control_file took
AsRef<Path> like the standard library functions do, you wouldn't even
need to construct the path here.

> +    let mut subtree = vec![];
> +    for controller in buf.split(|&b| b == b' ').filter(|e| !e.is_empty()) {

Are there ever likely to be empty works in this file?

> +        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(())
> +}
> +
> +fn cgroup_setup(mut args: ArgsOs) -> Result<(), String> {
> +    let mut leaf = false;
> +    let mut cgroup_path;
> +    let mut systemd_compat = false;
> +    let mut wait = true;
> +    loop {
> +        cgroup_path = args.next();
> +        let Some(ref arg_) = cgroup_path else {
> +            break;
> +        };
> +        let arg_ = arg_.as_bytes();
> +        if arg_ == b"--" {
> +            cgroup_path = args.next();

This cgroup_path thing is a bit complicated.  I feel like this could
probably be cleaned up with a peekable iterator and a while loop.

> +            break;
> +        }
> +        if !arg_.starts_with(b"-") {
> +            break;
> +        }
> +
> +        if !arg_.starts_with(b"--") {
> +            return Err("takes no short options".to_owned());
> +        }
> +
> +        match &arg_[2..] {
> +            b"leaf" => leaf = true,
> +            b"wait" => wait = true,
> +            b"no-wait" => wait = false,
> +            b"systemd-compat" => systemd_compat = true,

Why do we have --wait and --no-wait, but no --no-leaf or --no-systemd-compat?

> +            arg => return Err(format!("unknown long option {:?}", OsStr::from_bytes(arg))),
> +        }
> +    }
> +    let Some(cgroup_path) = cgroup_path.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 make child cgroup: {e}")),
> +    }
> +    let child = cgroup
> +        .open_beneath(child_cgroup_path, OFlags::RDONLY | OFlags::DIRECTORY)
> +        .map_err(|e| format!("Cannot make child cgroup: {e}"))?;
> +    if wait {
> +        // While waiting, only hold an exclusive lock on the child, not the parent.
> +        rustix::fs::flock(&child, FlockOperation::LockExclusive)
> +            .map_err(|e| format!("Cannot take an exclusive lock on child cgroup: {e}"))?;
> +        rustix::fs::flock(&cgroup, FlockOperation::LockShared)
> +            .map_err(|e| format!("Cannot downgrade lock on cgroup to a shared lock: {e}"))?;
> +        Cgroup::wait_for_empty(&child)
> +            .map_err(|e| format!("Cannot wait for {parent_cgroup_path:?} to be empty: {e}"))?;
> +    }
> +    let pid = std::process::id().to_string();
> +    if leaf {
> +        if args.len() != 0 {
> +            // If we aren't delegating any cgroups, don't create a sub-cgroup.
> +            write_value(&child, Path::new("cgroup.procs"), pid.as_bytes())
> +                .map_err(|e| format!("Cannot move process to child cgroup: {e}"))?;
> +        }
> +    } else {
> +        // 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 make child cgroup: {e}")),
> +        }
> +        if args.len() != 0 {
> +            let child_proc_path = Path::new(cgroup::DEFAULT_LEAF).join(Path::new("cgroup.procs"));
> +            write_value(&child, &child_proc_path, pid.as_bytes())
> +                .map_err(|e| format!("Cannot move process to child cgroup: {e}"))?;
> +        }
> +        if systemd_compat {
> +            // systemd-aware programs expect to have user.delegate=1
> +            // and to set cgroup.subtree_control themselves
> +            rustix::fs::fsetxattr(&child, c"user.delegate", b"1", XattrFlags::empty()).map_err(
> +                |e| format!("Cannot enable cgroup delegation in {parent_cgroup_path:?}: {e}"),
> +            )?
> +        } else {
> +            // Spectrum's programs do not check for user.delegate=1
> +            // and expect the caller to set cgroup.subtree_control.
> +            enable_subtree_control(&child)?;
> +        }
> +    }

So looking at this I still see several different modes and am wondering
whether we could simplify this further.

 • Why do we need a separate leaf mode?  Why not just still use a
   $inner.service in that case?

 • What would the consequences be if we took the systemd_compat branch
   for a non-cgroup-aware Spectrum program?

> +    let Some(program_name) = args.next() else {
> +        return Ok(());
> +    };
> +    let e = std::process::Command::new(&program_name).args(args).exec();
> +    Err(format!("Cannot spawn child {:?}: {}", program_name, e))
> +}
> +
> +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: &OsStr, args: ArgsOs) -> Result<(), String> {
> +    match prog_name
> +        .as_bytes()
> +        .split(|&b| b == b'/')
> +        .next_back()
> +        .unwrap()

prog_name.file_name(), where prog_name is &Path?

> +    {
> +        b"cgroup-setup" => cgroup_setup(args),
> +        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(&prog_name, args) {
> +        Ok(()) => {}
> +        Err(e) => {
> +            eprintln!("{prog_name:?}: {}", e);
> +            std::process::exit(1);
> +        }
> +    }
> +}
>
> --
> 2.55.0