Skip to content

Commit

Permalink
Rollup merge of rust-lang#92697 - the8472:cgroups, r=joshtriplett
Browse files Browse the repository at this point in the history
Use cgroup quotas for calculating `available_parallelism`

Automated tests for this are possible but would require a bunch of assumptions. It requires root + a recent kernel, systemd and maybe docker. And even then it would need a helper binary since the test has to run in a separate process.

Limitations

* only supports cgroup v2 and assumes it's mounted under `/sys/fs/cgroup`
* procfs must be available
* the quota gets mixed into `sched_getaffinity`, so if the latter doesn't work then quota information gets ignored too

Manually tested via

```
// spawn a new cgroup scope for the current user
$ sudo systemd-run -p CPUQuota="300%" --uid=$(id -u) -tdS

// quota.rs
#![feature(available_parallelism)]
fn main() {
    println!("{:?}", std::thread::available_parallelism()); // prints Ok(3)
}
```

strace:

```
sched_getaffinity(3041643, 32, [0, 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]) = 32
openat(AT_FDCWD, "/proc/self/cgroup", O_RDONLY|O_CLOEXEC) = 3
statx(0, NULL, AT_STATX_SYNC_AS_STAT, STATX_ALL, NULL) = -1 EFAULT (Bad address)
statx(3, "", AT_STATX_SYNC_AS_STAT|AT_EMPTY_PATH, STATX_ALL, {stx_mask=STATX_BASIC_STATS|STATX_MNT_ID, stx_attributes=0, stx_mode=S_IFREG|0444, stx_size=0, ...}) = 0
lseek(3, 0, SEEK_CUR)                   = 0
read(3, "0::/system.slice/run-u31477.serv"..., 128) = 36
read(3, "", 92)                         = 0
close(3)                                = 0
statx(AT_FDCWD, "/sys/fs/cgroup/system.slice/run-u31477.service/cgroup.controllers", AT_STATX_SYNC_AS_STAT, STATX_ALL, {stx_mask=STATX_BASIC_STATS|STATX_MNT_ID, stx_attributes=0, stx_mode=S_IFREG|0444, stx_size=0, ...}) = 0
openat(AT_FDCWD, "/sys/fs/cgroup/system.slice/run-u31477.service/cpu.max", O_RDONLY|O_CLOEXEC) = 3
statx(3, "", AT_STATX_SYNC_AS_STAT|AT_EMPTY_PATH, STATX_ALL, {stx_mask=STATX_BASIC_STATS|STATX_MNT_ID, stx_attributes=0, stx_mode=S_IFREG|0644, stx_size=0, ...}) = 0
lseek(3, 0, SEEK_CUR)                   = 0
read(3, "300000 100000\n", 20)          = 14
read(3, "", 6)                          = 0
close(3)                                = 0
openat(AT_FDCWD, "/sys/fs/cgroup/system.slice/cpu.max", O_RDONLY|O_CLOEXEC) = 3
statx(3, "", AT_STATX_SYNC_AS_STAT|AT_EMPTY_PATH, STATX_ALL, {stx_mask=STATX_BASIC_STATS|STATX_MNT_ID, stx_attributes=0, stx_mode=S_IFREG|0644, stx_size=0, ...}) = 0
lseek(3, 0, SEEK_CUR)                   = 0
read(3, "max 100000\n", 20)             = 11
read(3, "", 9)                          = 0
close(3)                                = 0
openat(AT_FDCWD, "/sys/fs/cgroup/cpu.max", O_RDONLY|O_CLOEXEC) = -1 ENOENT (No such file or directory)
sched_getaffinity(0, 128, [0, 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]) = 40
```

r? `@joshtriplett`
cc `@yoshuawuyts`

Tracking issue and previous discussion: rust-lang#74479
  • Loading branch information
Dylan-DPC authored Mar 3, 2022
2 parents 06460fe + e18abbf commit 3edda1a
Show file tree
Hide file tree
Showing 2 changed files with 86 additions and 4 deletions.
85 changes: 82 additions & 3 deletions library/std/src/sys/unix/thread.rs
Original file line number Diff line number Diff line change
Expand Up @@ -279,10 +279,15 @@ pub fn available_parallelism() -> io::Result<NonZeroUsize> {
))] {
#[cfg(any(target_os = "android", target_os = "linux"))]
{
let quota = cgroup2_quota().max(1);
let mut set: libc::cpu_set_t = unsafe { mem::zeroed() };
if unsafe { libc::sched_getaffinity(0, mem::size_of::<libc::cpu_set_t>(), &mut set) } == 0 {
let count = unsafe { libc::CPU_COUNT(&set) };
return Ok(unsafe { NonZeroUsize::new_unchecked(count as usize) });
unsafe {
if libc::sched_getaffinity(0, mem::size_of::<libc::cpu_set_t>(), &mut set) == 0 {
let count = libc::CPU_COUNT(&set) as usize;
let count = count.min(quota);
// SAFETY: affinity mask can't be empty and the quota gets clamped to a minimum of 1
return Ok(NonZeroUsize::new_unchecked(count));
}
}
}
match unsafe { libc::sysconf(libc::_SC_NPROCESSORS_ONLN) } {
Expand Down Expand Up @@ -368,6 +373,80 @@ pub fn available_parallelism() -> io::Result<NonZeroUsize> {
}
}

/// Returns cgroup CPU quota in core-equivalents, rounded down, or usize::MAX if the quota cannot
/// be determined or is not set.
#[cfg(any(target_os = "android", target_os = "linux"))]
fn cgroup2_quota() -> usize {
use crate::ffi::OsString;
use crate::fs::{try_exists, File};
use crate::io::Read;
use crate::os::unix::ffi::OsStringExt;
use crate::path::PathBuf;

let mut quota = usize::MAX;

let _: Option<()> = try {
let mut buf = Vec::with_capacity(128);
// find our place in the cgroup hierarchy
File::open("/proc/self/cgroup").ok()?.read_to_end(&mut buf).ok()?;
let cgroup_path = buf
.split(|&c| c == b'\n')
.filter_map(|line| {
let mut fields = line.splitn(3, |&c| c == b':');
// expect cgroupv2 which has an empty 2nd field
if fields.nth(1) != Some(b"") {
return None;
}
let path = fields.last()?;
// skip leading slash
Some(path[1..].to_owned())
})
.next()?;
let cgroup_path = PathBuf::from(OsString::from_vec(cgroup_path));

let mut path = PathBuf::with_capacity(128);
let mut read_buf = String::with_capacity(20);

let cgroup_mount = "/sys/fs/cgroup";

path.push(cgroup_mount);
path.push(&cgroup_path);

path.push("cgroup.controllers");

// skip if we're not looking at cgroup2
if matches!(try_exists(&path), Err(_) | Ok(false)) {
return usize::MAX;
};

path.pop();

while path.starts_with(cgroup_mount) {
path.push("cpu.max");

read_buf.clear();

if File::open(&path).and_then(|mut f| f.read_to_string(&mut read_buf)).is_ok() {
let raw_quota = read_buf.lines().next()?;
let mut raw_quota = raw_quota.split(' ');
let limit = raw_quota.next()?;
let period = raw_quota.next()?;
match (limit.parse::<usize>(), period.parse::<usize>()) {
(Ok(limit), Ok(period)) => {
quota = quota.min(limit / period);
}
_ => {}
}
}

path.pop(); // pop filename
path.pop(); // pop dir
}
};

quota
}

#[cfg(all(
not(target_os = "linux"),
not(target_os = "freebsd"),
Expand Down
5 changes: 4 additions & 1 deletion library/std/src/thread/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1524,7 +1524,10 @@ fn _assert_sync_and_send() {
///
/// On Linux:
/// - It may overcount the amount of parallelism available when limited by a
/// process-wide affinity mask, or when affected by cgroup limits.
/// process-wide affinity mask or cgroup quotas and cgroup2 fs or `sched_getaffinity()` can't be
/// queried, e.g. due to sandboxing.
/// - It may undercount the amount of parallelism if the current thread's affinity mask
/// does not reflect the process' cpuset, e.g. due to pinned threads.
///
/// On all targets:
/// - It may overcount the amount of parallelism available when running in a VM
Expand Down

0 comments on commit 3edda1a

Please sign in to comment.