Skip to content

Commit bac5523

Browse files
committed
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::thread::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
1 parent 8769f4e commit bac5523

File tree

1 file changed

+68
-3
lines changed

1 file changed

+68
-3
lines changed

library/std/src/sys/unix/thread.rs

+68-3
Original file line numberDiff line numberDiff line change
@@ -279,10 +279,15 @@ pub fn available_parallelism() -> io::Result<NonZeroUsize> {
279279
))] {
280280
#[cfg(any(target_os = "android", target_os = "linux"))]
281281
{
282+
let quota = cgroup2_quota().unwrap_or(usize::MAX).max(1);
282283
let mut set: libc::cpu_set_t = unsafe { mem::zeroed() };
283-
if unsafe { libc::sched_getaffinity(0, mem::size_of::<libc::cpu_set_t>(), &mut set) } == 0 {
284-
let count = unsafe { libc::CPU_COUNT(&set) };
285-
return Ok(unsafe { NonZeroUsize::new_unchecked(count as usize) });
284+
unsafe {
285+
if libc::sched_getaffinity(0, mem::size_of::<libc::cpu_set_t>(), &mut set) == 0 {
286+
let count = libc::CPU_COUNT(&set) as usize;
287+
let count = count.min(quota);
288+
// SAFETY: affinity mask can't be empty and the quota gets clamped to a minimum of 1
289+
return Ok(NonZeroUsize::new_unchecked(count));
290+
}
286291
}
287292
}
288293
match unsafe { libc::sysconf(libc::_SC_NPROCESSORS_ONLN) } {
@@ -368,6 +373,66 @@ pub fn available_parallelism() -> io::Result<NonZeroUsize> {
368373
}
369374
}
370375

376+
#[cfg(any(target_os = "android", target_os = "linux"))]
377+
fn cgroup2_quota() -> Option<usize> {
378+
use crate::ffi::OsString;
379+
use crate::fs::{read, read_to_string, File};
380+
use crate::io::{BufRead, BufReader};
381+
use crate::os::unix::ffi::OsStringExt;
382+
use crate::path::PathBuf;
383+
384+
// find cgroup2 fs
385+
let cgroups_mount = BufReader::new(File::open("/proc/self/mountinfo").ok()?)
386+
.split(b'\n')
387+
.map_while(Result::ok)
388+
.filter_map(|line| {
389+
let fields: Vec<_> = line.split(|&c| c == b' ').collect();
390+
let suffix_at = fields.iter().position(|f| f == b"-")?;
391+
let fs_type = fields[suffix_at + 1];
392+
if fs_type == b"cgroup2" { Some(fields[4].to_owned()) } else { None }
393+
})
394+
.next()?;
395+
396+
let cgroups_mount = PathBuf::from(OsString::from_vec(cgroups_mount));
397+
398+
// find our place in the hierarchy
399+
let cgroup_path = read("/proc/self/cgroup")
400+
.ok()?
401+
.split(|&c| c == b'\n')
402+
.filter_map(|line| {
403+
let mut fields = line.splitn(3, |&c| c == b':');
404+
// expect cgroupv2 which has an empty 2nd field
405+
if fields.nth(1) != Some(b"") {
406+
return None;
407+
}
408+
let path = fields.last()?;
409+
// skip leading slash
410+
Some(path[1..].to_owned())
411+
})
412+
.next()?;
413+
let cgroup_path = PathBuf::from(OsString::from_vec(cgroup_path));
414+
415+
// walk hierarchy and take the minimum quota
416+
cgroup_path
417+
.ancestors()
418+
.filter_map(|level| {
419+
let cgroup_path = cgroups_mount.join(level);
420+
let quota = match read_to_string(cgroup_path.join("cpu.max")) {
421+
Ok(quota) => quota,
422+
_ => return None,
423+
};
424+
let quota = quota.lines().next()?;
425+
let mut quota = quota.split(' ');
426+
let limit = quota.next()?;
427+
let period = quota.next()?;
428+
match (limit.parse::<usize>(), period.parse::<usize>()) {
429+
(Ok(limit), Ok(period)) => Some(limit / period),
430+
_ => None,
431+
}
432+
})
433+
.min()
434+
}
435+
371436
#[cfg(all(
372437
not(target_os = "linux"),
373438
not(target_os = "freebsd"),

0 commit comments

Comments
 (0)