Use cgroup quotas for calculating available_parallelism
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:🧵:available_parallelism()); // prints Ok(3)
}
```
Caveats
* cgroup v1 is ignored
* funky mountpoints (containing spaces, newlines or control chars) for cgroupfs will not be handled correctly since that would require unescaping /proc/self/mountinfo
The escaping behavior of procfs seems to be undocumented. systemd and docker default to `/sys/fs/cgroup` so it should be fine for most systems.
* quota will be ignored when `sched_getaffinity` doesn't work
* assumes procfs is mounted under `/proc` and cgroupfs mounted and readable somewhere in the directory tree
This commit is contained in:
parent
8769f4ef2f
commit
bac5523ea0
1 changed files with 68 additions and 3 deletions
|
@ -279,10 +279,15 @@ pub fn available_parallelism() -> io::Result<NonZeroUsize> {
|
|||
))] {
|
||||
#[cfg(any(target_os = "android", target_os = "linux"))]
|
||||
{
|
||||
let quota = cgroup2_quota().unwrap_or(usize::MAX).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) } {
|
||||
|
@ -368,6 +373,66 @@ pub fn available_parallelism() -> io::Result<NonZeroUsize> {
|
|||
}
|
||||
}
|
||||
|
||||
#[cfg(any(target_os = "android", target_os = "linux"))]
|
||||
fn cgroup2_quota() -> Option<usize> {
|
||||
use crate::ffi::OsString;
|
||||
use crate::fs::{read, read_to_string, File};
|
||||
use crate::io::{BufRead, BufReader};
|
||||
use crate::os::unix::ffi::OsStringExt;
|
||||
use crate::path::PathBuf;
|
||||
|
||||
// find cgroup2 fs
|
||||
let cgroups_mount = BufReader::new(File::open("/proc/self/mountinfo").ok()?)
|
||||
.split(b'\n')
|
||||
.map_while(Result::ok)
|
||||
.filter_map(|line| {
|
||||
let fields: Vec<_> = line.split(|&c| c == b' ').collect();
|
||||
let suffix_at = fields.iter().position(|f| f == b"-")?;
|
||||
let fs_type = fields[suffix_at + 1];
|
||||
if fs_type == b"cgroup2" { Some(fields[4].to_owned()) } else { None }
|
||||
})
|
||||
.next()?;
|
||||
|
||||
let cgroups_mount = PathBuf::from(OsString::from_vec(cgroups_mount));
|
||||
|
||||
// find our place in the hierarchy
|
||||
let cgroup_path = read("/proc/self/cgroup")
|
||||
.ok()?
|
||||
.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));
|
||||
|
||||
// walk hierarchy and take the minimum quota
|
||||
cgroup_path
|
||||
.ancestors()
|
||||
.filter_map(|level| {
|
||||
let cgroup_path = cgroups_mount.join(level);
|
||||
let quota = match read_to_string(cgroup_path.join("cpu.max")) {
|
||||
Ok(quota) => quota,
|
||||
_ => return None,
|
||||
};
|
||||
let quota = quota.lines().next()?;
|
||||
let mut quota = quota.split(' ');
|
||||
let limit = quota.next()?;
|
||||
let period = quota.next()?;
|
||||
match (limit.parse::<usize>(), period.parse::<usize>()) {
|
||||
(Ok(limit), Ok(period)) => Some(limit / period),
|
||||
_ => None,
|
||||
}
|
||||
})
|
||||
.min()
|
||||
}
|
||||
|
||||
#[cfg(all(
|
||||
not(target_os = "linux"),
|
||||
not(target_os = "freebsd"),
|
||||
|
|
Loading…
Add table
Reference in a new issue