Skip to content

Commit af6d2ed

Browse files
committed
hardcode /sys/fs/cgroup instead of doing a lookup via mountinfo
this avoids parsing mountinfo which can be huge on some systems and something might be emulating cgroup fs for sandboxing reasons which means it wouldn't show up as mountpoint additionally the new implementation operates on a single pathbuffer, reducing allocations
1 parent bac5523 commit af6d2ed

File tree

1 file changed

+67
-53
lines changed

1 file changed

+67
-53
lines changed

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

+67-53
Original file line numberDiff line numberDiff line change
@@ -279,7 +279,7 @@ 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);
282+
let quota = cgroup2_quota().max(1);
283283
let mut set: libc::cpu_set_t = unsafe { mem::zeroed() };
284284
unsafe {
285285
if libc::sched_getaffinity(0, mem::size_of::<libc::cpu_set_t>(), &mut set) == 0 {
@@ -373,64 +373,78 @@ pub fn available_parallelism() -> io::Result<NonZeroUsize> {
373373
}
374374
}
375375

376+
/// Returns cgroup CPU quota in core-equivalents, rounded down, or usize::MAX if the quota cannot
377+
/// be determined or is not set.
376378
#[cfg(any(target_os = "android", target_os = "linux"))]
377-
fn cgroup2_quota() -> Option<usize> {
379+
fn cgroup2_quota() -> usize {
378380
use crate::ffi::OsString;
379-
use crate::fs::{read, read_to_string, File};
380-
use crate::io::{BufRead, BufReader};
381+
use crate::fs::{try_exists, File};
382+
use crate::io::Read;
381383
use crate::os::unix::ffi::OsStringExt;
382384
use crate::path::PathBuf;
383385

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,
386+
let mut quota = usize::MAX;
387+
388+
let _: Option<()> = try {
389+
let mut buf = Vec::with_capacity(128);
390+
// find our place in the cgroup hierarchy
391+
File::open("/proc/self/cgroup").ok()?.read_to_end(&mut buf).ok()?;
392+
let cgroup_path = buf
393+
.split(|&c| c == b'\n')
394+
.filter_map(|line| {
395+
let mut fields = line.splitn(3, |&c| c == b':');
396+
// expect cgroupv2 which has an empty 2nd field
397+
if fields.nth(1) != Some(b"") {
398+
return None;
399+
}
400+
let path = fields.last()?;
401+
// skip leading slash
402+
Some(path[1..].to_owned())
403+
})
404+
.next()?;
405+
let cgroup_path = PathBuf::from(OsString::from_vec(cgroup_path));
406+
407+
let mut path = PathBuf::with_capacity(128);
408+
let mut read_buf = String::with_capacity(20);
409+
410+
let cgroup_mount = "/sys/fs/cgroup";
411+
412+
path.push(cgroup_mount);
413+
path.push(&cgroup_path);
414+
415+
path.push("cgroup.controllers");
416+
417+
// skip if we're not looking at cgroup2
418+
if matches!(try_exists(&path), Err(_) | Ok(false)) {
419+
return usize::MAX;
420+
};
421+
422+
path.pop();
423+
424+
while path.starts_with(cgroup_mount) {
425+
path.push("cpu.max");
426+
427+
read_buf.clear();
428+
429+
if File::open(&path).and_then(|mut f| f.read_to_string(&mut read_buf)).is_ok() {
430+
let raw_quota = read_buf.lines().next()?;
431+
let mut raw_quota = raw_quota.split(' ');
432+
let limit = raw_quota.next()?;
433+
let period = raw_quota.next()?;
434+
match (limit.parse::<usize>(), period.parse::<usize>()) {
435+
(Ok(limit), Ok(period)) => {
436+
quota = quota.min(limit / period);
437+
}
438+
_ => {}
439+
}
431440
}
432-
})
433-
.min()
441+
442+
path.pop(); // pop filename
443+
path.pop(); // pop dir
444+
}
445+
};
446+
447+
quota
434448
}
435449

436450
#[cfg(all(

0 commit comments

Comments
 (0)