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
168
169
170
171
172
173
174
use std::{fs, os::unix::process::CommandExt as _};
use crate::v1::{Cgroup, UnifiedRepr};
/// Extension to the [`std::process::Command`] builder for attaching a command process to one or
/// more cgroups on start.
///
/// [`std::process::Command`]: https://doc.rust-lang.org/std/process/struct.Command.html
pub trait CommandExt {
/// Attaches a command process to a cgroup on start.
fn cgroup<C: Cgroup>(&mut self, cgroup: &mut C) -> &mut Self;
/// Attaches a command process to each subsystem supported by a [`UnifiedRepr`] on start.
///
/// [`UnifiedRepr`]: struct.UnifiedRepr.html
fn cgroups_unified_repr(&mut self, cgroups: &mut UnifiedRepr) -> &mut Self;
}
impl CommandExt for std::process::Command {
// NOTE: Keep the example below in sync with `README.md` and `lib.rs`
/// Attaches this command process to a cgroup on start.
///
/// The process will run within the cgroup from the beginning of its execution.
///
/// Multiple cgroups can be registered for the process attachment. The process will be attached
/// to the cgroups in order of their registration.
///
/// # Examples
///
/// ```no_run
/// # fn main() -> controlgroup::Result<()> {
/// use std::path::PathBuf;
/// use controlgroup::v1::{cpu, Cgroup, CgroupPath, SubsystemKind};
/// // Import extension trait
/// use controlgroup::v1::CommandExt as _;
///
/// let mut cgroup = cpu::Subsystem::new(
/// CgroupPath::new(SubsystemKind::Cpu, PathBuf::from("students/charlie")));
/// cgroup.create()?;
///
/// let mut child = std::process::Command::new("sleep")
/// .arg("1")
/// // Attach this command process to a cgroup on start
/// .cgroup(&mut cgroup)
/// // This process will run within the cgroup
/// .spawn()
/// .unwrap();
///
/// println!("{:?}", cgroup.stat()?);
///
/// child.wait().unwrap();
/// cgroup.delete()?;
///
/// # Ok(())
/// # }
/// ```
fn cgroup<C: Cgroup>(&mut self, cgroup: &mut C) -> &mut Self {
let path = cgroup.path().join("cgroup.procs");
unsafe { self.pre_exec(move || fs::write(&path, std::process::id().to_string())) }
// FIXME: is it safe to write to the same file in parallel?
}
/// Attaches this command process to each subsystem supported by a [`UnifiedRepr`] on start.
///
/// See [`cgroup`] for more information.
///
/// [`UnifiedRepr`]: struct.UnifiedRepr.html
/// [`cgroup`]: #method.cgroup
fn cgroups_unified_repr(&mut self, cgroups: &mut UnifiedRepr) -> &mut Self {
macro_rules! a {
( $($subsystem: ident),* $(, )? ) => { $(
if let Some(subsys) = cgroups.$subsystem() {
self.cgroup(subsys);
}
)* };
}
a! {
cpu_mut,
cpuset_mut,
cpuacct_mut,
memory_mut,
hugetlb_mut,
devices_mut,
blkio_mut,
rdma_mut,
net_prio_mut,
net_cls_mut,
pids_mut,
freezer_mut,
perf_event_mut,
}
self
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{
v1::{cpu, CgroupPath, SubsystemKind},
Pid, Result,
};
#[test]
fn test_command_ext_cgroup() -> Result<()> {
let mut cgroup =
cpu::Subsystem::new(CgroupPath::new(SubsystemKind::Cpu, gen_cgroup_name!()));
cgroup.create()?;
let mut child = std::process::Command::new("sleep")
.arg("1")
.cgroup(&mut cgroup)
.spawn()
.unwrap();
let pid = child.id();
assert_eq!(cgroup.procs().unwrap(), vec![Pid::from(pid)]);
child.wait()?;
cgroup.delete()
}
#[test]
fn test_command_ext_unified() -> Result<()> {
use crate::v1::cpuset;
use SubsystemKind::*;
let mut cgroups = UnifiedRepr::new(gen_cgroup_name!());
cgroups.skip_create(&[Cpuacct, NetCls]);
cgroups.create()?;
cgroups.cpuset_mut().unwrap().apply({
let id_set = [0].iter().copied().collect::<cpuset::IdSet>();
&cpuset::Resources {
cpus: Some(id_set.clone()),
mems: Some(id_set),
..cpuset::Resources::default()
}
.into()
})?;
let mut child = std::process::Command::new("sleep")
.arg("1")
.cgroups_unified_repr(&mut cgroups)
.spawn()
.unwrap();
let pid = Pid::from(child.id());
assert_eq!(
cgroups.procs().unwrap(),
hashmap! {
(BlkIo, vec![pid]),
(Cpu, vec![pid]),
(Cpuacct, vec![pid]),
(Cpuset, vec![pid]),
(Devices, vec![pid]),
(Freezer, vec![pid]),
(HugeTlb, vec![pid]),
(Memory, vec![pid]),
(NetCls, vec![pid]),
(NetPrio, vec![pid]),
(PerfEvent, vec![pid]),
(Pids, vec![pid]),
(Rdma, vec![pid]),
}
);
child.wait()?;
cgroups.delete()
}
}