Skip to content

Commit

Permalink
hardcode /sys/fs/cgroup instead of doing a lookup via mountinfo
Browse files Browse the repository at this point in the history
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
  • Loading branch information
the8472 committed Mar 2, 2022
1 parent bac5523 commit af6d2ed
Showing 1 changed file with 67 additions and 53 deletions.
120 changes: 67 additions & 53 deletions library/std/src/sys/unix/thread.rs
Expand Up @@ -279,7 +279,7 @@ 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 quota = cgroup2_quota().max(1);
let mut set: libc::cpu_set_t = unsafe { mem::zeroed() };
unsafe {
if libc::sched_getaffinity(0, mem::size_of::<libc::cpu_set_t>(), &mut set) == 0 {
Expand Down Expand Up @@ -373,64 +373,78 @@ 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() -> Option<usize> {
fn cgroup2_quota() -> usize {
use crate::ffi::OsString;
use crate::fs::{read, read_to_string, File};
use crate::io::{BufRead, BufReader};
use crate::fs::{try_exists, File};
use crate::io::Read;
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,
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);
}
_ => {}
}
}
})
.min()

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

quota
}

#[cfg(all(
Expand Down

0 comments on commit af6d2ed

Please sign in to comment.