diff --git a/library/std/src/sys/unix/thread.rs b/library/std/src/sys/unix/thread.rs index 933210e1ff0..ff01ce27333 100644 --- a/library/std/src/sys/unix/thread.rs +++ b/library/std/src/sys/unix/thread.rs @@ -279,7 +279,7 @@ pub fn available_parallelism() -> io::Result { ))] { #[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::(), &mut set) == 0 { @@ -373,64 +373,78 @@ pub fn available_parallelism() -> io::Result { } } +/// 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 { +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::(), period.parse::()) { + (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::(), period.parse::()) { - (Ok(limit), Ok(period)) => Some(limit / period), - _ => None, - } - }) - .min() + path.pop(); // pop filename + path.pop(); // pop dir + } + }; + + quota } #[cfg(all(