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
This commit is contained in:
The 8472 2022-03-03 00:35:47 +01:00
parent bac5523ea0
commit af6d2ed245

View File

@ -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 {
@ -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 mut quota = usize::MAX;
let cgroups_mount = PathBuf::from(OsString::from_vec(cgroups_mount));
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));
// 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 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);
}
_ => {}
}
}
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()
path.pop(); // pop filename
path.pop(); // pop dir
}
};
quota
}
#[cfg(all(