1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
| | // 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},
};
// 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 enable_subtree_control(fd: &dyn AsFd) -> Result<(), String> {
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![];
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(())
}
fn cgroup_setup(args: ArgsOs) -> Result<(), String> {
let mut leaf = false;
let mut systemd_delegate = false;
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();
let Some(arg_) = arg.as_bytes().strip_prefix(b"--") else {
return Err("takes no short options".to_owned());
};
match arg_ {
b"" => break,
b"leaf" => leaf = true,
b"no-wait" => wait = false,
b"systemd-delegate" => systemd_delegate = true,
_ => 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 make child cgroup {child_cgroup_path:?}: {e}"
));
}
}
let child = openat2_simple(&cgroup, child_cgroup_path, OpenFlags::Directory)
.map_err(|e| format!("Cannot make 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}"))?;
}
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_delegate {
// 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)?;
}
}
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: &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);
}
}
}
|