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

[Demi Marie Obenour <demiobenour@gmail.com>]

[-- Attachment #1.1: Type: text/plain, Size: 23782 bytes --]

On 8/3/26 08:47, Alyssa Ross wrote:
> 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?

Easy fix.

>> +    // 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?

Consider the recursive implementation (in pseudo-Rust):

fn recursive_remove(fd) {
    for entry in get_entries(&fd) {
        if entry.is_dir_and_not_dot_or_dotdot() {
            let directory = open_dir(&fd, &entry.path())?;
            recursive_remove(directory)?;
            remove_dir(&fd, entry.path())?;
        }
    }
}

The compiler knows that fd will stay open through recursive calls,
so this doesn't need any unsafe code.  Using an explicit stack takes
away this information from the compiler, so unsafe code is required.
Using Rc<RefCell<Dir>> avoids the need for unsafe code at a cost
in performance.

For what it is worth, the standard library implementation of
fs::remove_dir_all() is recursive.  Standard library security hole?


>> +            }
>> +        }
>> +    }
>> +}
>> +
>> +// 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.)

Can you provide an example?  I don't see how to make this change
while preserving semantics.  Only directories meant for deletion
can appear on the stack, and the root of the traversal must not
be deleted (yet).

>> +        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.

Will remove.

>> +    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.

NOCTTY is redundant with DIRECTORY, as directories can't be terminals.
openat2_simple() uses NO_XDEV, which won't work here because
/sys/fs/cgroup is not the same filesystem as /.

>> +        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!()
> };

Will fix in v6.

>> +            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.

I will move `fds` into the inner loop.

>> +            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.

Indeed so.  Programs that are modifying a cgroup need an exclusive
lock on it.  Adding or removing to the cgroup does *not* count as
modification: both operations are idempotent, removing an in-use
cgroup fails with -EBUSY, and operating on a deleted cgroup fails
with -ENODEV or -ENOENT depending on what one is doing.  Operations on
control files *do* require an exclusive lock.

>> +
>> +        // 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.

Will change in v6.

>> +}
>> +
>> +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.

It's only used once, so I'll inline it.

>> +
>> +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?

We don't create a Cgroup struct for the child cgroup
FD on which this function is called.

>> +    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.

I will just inline the whole read_control_file function.
It's only used once.

>> +    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?

No, there will not be unless there is a kernel bug.

>> +        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.

Easy to fix.

>> +            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?

Spectrum never needs --wait, so I'll remove it.

>> +            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?

$inner.service is just wasteful and makes it harder to inspect the cgroup
tree by hand.

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

Non-cgroup-aware programs would be fine, but nested calls to cgroup-setup
would break because they need the enable_subtree_control() call.  However,
in the future, I would like to check the user.delegate xattr to determine
if one can safely write to the control files of the cgroup or if the cgroup
is owned by another 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?

Works for me :)

>> +    {
>> +        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


-- 
Sincerely,
Demi Marie Obenour (she/her/hers)


[-- Attachment #2: OpenPGP digital signature --]
[-- Type: application/pgp-signature, Size: 833 bytes --]