diff --git a/library/std/src/sys/unix/thread.rs b/library/std/src/sys/unix/thread.rs index cf8cf5ad49f73..ff01ce2733329 100644 --- a/library/std/src/sys/unix/thread.rs +++ b/library/std/src/sys/unix/thread.rs @@ -279,10 +279,15 @@ pub fn available_parallelism() -> io::Result { ))] { #[cfg(any(target_os = "android", target_os = "linux"))] { + let quota = cgroup2_quota().max(1); let mut set: libc::cpu_set_t = unsafe { mem::zeroed() }; - if unsafe { libc::sched_getaffinity(0, mem::size_of::(), &mut set) } == 0 { - let count = unsafe { libc::CPU_COUNT(&set) }; - return Ok(unsafe { NonZeroUsize::new_unchecked(count as usize) }); + unsafe { + if libc::sched_getaffinity(0, mem::size_of::(), &mut set) == 0 { + let count = libc::CPU_COUNT(&set) as usize; + let count = count.min(quota); + // SAFETY: affinity mask can't be empty and the quota gets clamped to a minimum of 1 + return Ok(NonZeroUsize::new_unchecked(count)); + } } } match unsafe { libc::sysconf(libc::_SC_NPROCESSORS_ONLN) } { @@ -368,6 +373,80 @@ 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() -> usize { + use crate::ffi::OsString; + use crate::fs::{try_exists, File}; + use crate::io::Read; + use crate::os::unix::ffi::OsStringExt; + use crate::path::PathBuf; + + 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::(), period.parse::()) { + (Ok(limit), Ok(period)) => { + quota = quota.min(limit / period); + } + _ => {} + } + } + + path.pop(); // pop filename + path.pop(); // pop dir + } + }; + + quota +} + #[cfg(all( not(target_os = "linux"), not(target_os = "freebsd"), diff --git a/library/std/src/thread/mod.rs b/library/std/src/thread/mod.rs index beb606099341e..09d1e714ab6dd 100644 --- a/library/std/src/thread/mod.rs +++ b/library/std/src/thread/mod.rs @@ -1524,7 +1524,10 @@ fn _assert_sync_and_send() { /// /// On Linux: /// - It may overcount the amount of parallelism available when limited by a -/// process-wide affinity mask, or when affected by cgroup limits. +/// process-wide affinity mask or cgroup quotas and cgroup2 fs or `sched_getaffinity()` can't be +/// queried, e.g. due to sandboxing. +/// - It may undercount the amount of parallelism if the current thread's affinity mask +/// does not reflect the process' cpuset, e.g. due to pinned threads. /// /// On all targets: /// - It may overcount the amount of parallelism available when running in a VM