Skip to content

Commit d823462

Browse files
committed
add cgroupv1 support to available_parallelism
1 parent 09d52bc commit d823462

File tree

3 files changed

+128
-48
lines changed

3 files changed

+128
-48
lines changed

library/std/src/lib.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -274,6 +274,7 @@
274274
#![feature(hasher_prefixfree_extras)]
275275
#![feature(hashmap_internals)]
276276
#![feature(int_error_internals)]
277+
#![feature(is_some_with)]
277278
#![feature(maybe_uninit_slice)]
278279
#![feature(maybe_uninit_write_slice)]
279280
#![feature(mixed_integer_ops)]

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

Lines changed: 126 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -285,7 +285,7 @@ pub fn available_parallelism() -> io::Result<NonZeroUsize> {
285285
))] {
286286
#[cfg(any(target_os = "android", target_os = "linux"))]
287287
{
288-
let quota = cgroup2_quota().max(1);
288+
let quota = cgroups::quota().max(1);
289289
let mut set: libc::cpu_set_t = unsafe { mem::zeroed() };
290290
unsafe {
291291
if libc::sched_getaffinity(0, mem::size_of::<libc::cpu_set_t>(), &mut set) == 0 {
@@ -379,49 +379,77 @@ pub fn available_parallelism() -> io::Result<NonZeroUsize> {
379379
}
380380
}
381381

382-
/// Returns cgroup CPU quota in core-equivalents, rounded down, or usize::MAX if the quota cannot
383-
/// be determined or is not set.
384382
#[cfg(any(target_os = "android", target_os = "linux"))]
385-
fn cgroup2_quota() -> usize {
383+
mod cgroups {
386384
use crate::ffi::OsString;
387385
use crate::fs::{try_exists, File};
388386
use crate::io::Read;
389387
use crate::os::unix::ffi::OsStringExt;
390388
use crate::path::PathBuf;
389+
use crate::str::from_utf8;
391390

392-
let mut quota = usize::MAX;
393-
if cfg!(miri) {
394-
// Attempting to open a file fails under default flags due to isolation.
395-
// And Miri does not have parallelism anyway.
396-
return quota;
397-
}
398-
399-
let _: Option<()> = try {
400-
let mut buf = Vec::with_capacity(128);
401-
// find our place in the cgroup hierarchy
402-
File::open("/proc/self/cgroup").ok()?.read_to_end(&mut buf).ok()?;
403-
let cgroup_path = buf
404-
.split(|&c| c == b'\n')
405-
.filter_map(|line| {
406-
let mut fields = line.splitn(3, |&c| c == b':');
407-
// expect cgroupv2 which has an empty 2nd field
408-
if fields.nth(1) != Some(b"") {
409-
return None;
410-
}
411-
let path = fields.last()?;
412-
// skip leading slash
413-
Some(path[1..].to_owned())
414-
})
415-
.next()?;
416-
let cgroup_path = PathBuf::from(OsString::from_vec(cgroup_path));
391+
enum Cgroup {
392+
V1,
393+
V2,
394+
}
395+
396+
/// Returns cgroup CPU quota in core-equivalents, rounded down, or usize::MAX if the quota cannot
397+
/// be determined or is not set.
398+
pub(super) fn quota() -> usize {
399+
let mut quota = usize::MAX;
400+
if cfg!(miri) {
401+
// Attempting to open a file fails under default flags due to isolation.
402+
// And Miri does not have parallelism anyway.
403+
return quota;
404+
}
405+
406+
let _: Option<()> = try {
407+
let mut buf = Vec::with_capacity(128);
408+
// find our place in the cgroup hierarchy
409+
File::open("/proc/self/cgroup").ok()?.read_to_end(&mut buf).ok()?;
410+
let (cgroup_path, version) = buf
411+
.split(|&c| c == b'\n')
412+
.filter_map(|line| {
413+
let mut fields = line.splitn(3, |&c| c == b':');
414+
// 2nd field is a list of controllers for v1 or empty for v2
415+
let version = match fields.nth(1) {
416+
Some(b"") => Some(Cgroup::V2),
417+
Some(controllers)
418+
if from_utf8(controllers)
419+
.is_ok_and(|c| c.split(",").any(|c| c == "cpu")) =>
420+
{
421+
Some(Cgroup::V1)
422+
}
423+
_ => None,
424+
}?;
425+
426+
let path = fields.last()?;
427+
// skip leading slash
428+
Some((path[1..].to_owned(), version))
429+
})
430+
.next()?;
431+
let cgroup_path = PathBuf::from(OsString::from_vec(cgroup_path));
432+
433+
quota = match version {
434+
Cgroup::V1 => quota_v1(cgroup_path),
435+
Cgroup::V2 => quota_v2(cgroup_path),
436+
};
437+
};
438+
439+
quota
440+
}
441+
442+
fn quota_v2(group_path: PathBuf) -> usize {
443+
let mut quota = usize::MAX;
417444

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

448+
// standard mount location defined in file-hierarchy(7) manpage
421449
let cgroup_mount = "/sys/fs/cgroup";
422450

423451
path.push(cgroup_mount);
424-
path.push(&cgroup_path);
452+
path.push(&group_path);
425453

426454
path.push("cgroup.controllers");
427455

@@ -432,30 +460,81 @@ fn cgroup2_quota() -> usize {
432460

433461
path.pop();
434462

435-
while path.starts_with(cgroup_mount) {
436-
path.push("cpu.max");
463+
let _: Option<()> = try {
464+
while path.starts_with(cgroup_mount) {
465+
path.push("cpu.max");
466+
467+
read_buf.clear();
468+
469+
if File::open(&path).and_then(|mut f| f.read_to_string(&mut read_buf)).is_ok() {
470+
let raw_quota = read_buf.lines().next()?;
471+
let mut raw_quota = raw_quota.split(' ');
472+
let limit = raw_quota.next()?;
473+
let period = raw_quota.next()?;
474+
match (limit.parse::<usize>(), period.parse::<usize>()) {
475+
(Ok(limit), Ok(period)) => {
476+
quota = quota.min(limit / period);
477+
}
478+
_ => {}
479+
}
480+
}
437481

438-
read_buf.clear();
482+
path.pop(); // pop filename
483+
path.pop(); // pop dir
484+
}
485+
};
439486

440-
if File::open(&path).and_then(|mut f| f.read_to_string(&mut read_buf)).is_ok() {
441-
let raw_quota = read_buf.lines().next()?;
442-
let mut raw_quota = raw_quota.split(' ');
443-
let limit = raw_quota.next()?;
444-
let period = raw_quota.next()?;
445-
match (limit.parse::<usize>(), period.parse::<usize>()) {
446-
(Ok(limit), Ok(period)) => {
447-
quota = quota.min(limit / period);
448-
}
487+
quota
488+
}
489+
490+
fn quota_v1(group_path: PathBuf) -> usize {
491+
let mut quota = usize::MAX;
492+
let mut path = PathBuf::with_capacity(128);
493+
let mut read_buf = String::with_capacity(20);
494+
495+
// Hardcode commonly used locations mentioned in the cgroups(7) manpage
496+
// since scanning mountinfo can be expensive on some systems.
497+
// This isn't exactly standardized since cgroupv1 was meant to allow flexibly
498+
// mixing and matching controller hierarchies.
499+
let mounts = ["/sys/fs/cgroup/cpu", "/sys/fs/cgroup/cpu,cpuacct"];
500+
501+
for mount in mounts {
502+
path.clear();
503+
path.push(mount);
504+
path.push(&group_path);
505+
506+
// skip if we guessed the mount incorrectly
507+
if matches!(try_exists(&path), Err(_) | Ok(false)) {
508+
continue;
509+
}
510+
511+
while path.starts_with(mount) {
512+
let mut parse_file = |name| {
513+
path.push(name);
514+
read_buf.clear();
515+
516+
let mut f = File::open(&path).ok()?;
517+
f.read_to_string(&mut read_buf).ok()?;
518+
let parsed = read_buf.trim().parse::<usize>().ok()?;
519+
520+
path.pop();
521+
Some(parsed)
522+
};
523+
524+
let limit = parse_file("cpu.cfs_quota_us");
525+
let period = parse_file("cpu.cfs_period_us");
526+
527+
match (limit, period) {
528+
(Some(limit), Some(period)) => quota = quota.min(limit / period),
449529
_ => {}
450530
}
451-
}
452531

453-
path.pop(); // pop filename
454-
path.pop(); // pop dir
532+
path.pop();
533+
}
455534
}
456-
};
457535

458-
quota
536+
quota
537+
}
459538
}
460539

461540
#[cfg(all(

library/std/src/thread/mod.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1571,7 +1571,7 @@ fn _assert_sync_and_send() {
15711571
///
15721572
/// On Linux:
15731573
/// - It may overcount the amount of parallelism available when limited by a
1574-
/// process-wide affinity mask or cgroup quotas and cgroup2 fs or `sched_getaffinity()` can't be
1574+
/// process-wide affinity mask or cgroup quotas and `sched_getaffinity()` or cgroup fs can't be
15751575
/// queried, e.g. due to sandboxing.
15761576
/// - It may undercount the amount of parallelism if the current thread's affinity mask
15771577
/// does not reflect the process' cpuset, e.g. due to pinned threads.

0 commit comments

Comments
 (0)