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:
parent
bac5523ea0
commit
af6d2ed245
@ -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(
|
||||
|
Loading…
Reference in New Issue
Block a user