diff options
Diffstat (limited to 'library/std/src/sys/unix/thread.rs')
| -rw-r--r-- | library/std/src/sys/unix/thread.rs | 90 |
1 files changed, 87 insertions, 3 deletions
diff --git a/library/std/src/sys/unix/thread.rs b/library/std/src/sys/unix/thread.rs index cf8cf5ad49f..2d5d306ed62 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<NonZeroUsize> { ))] { #[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::<libc::cpu_set_t>(), &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::<libc::cpu_set_t>(), &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,85 @@ 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() -> 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; + if cfg!(miri) { + // Attempting to open a file fails under default flags due to isolation. + // And Miri does not have parallelism anyway. + return quota; + } + + 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); + } + _ => {} + } + } + + path.pop(); // pop filename + path.pop(); // pop dir + } + }; + + quota +} + #[cfg(all( not(target_os = "linux"), not(target_os = "freebsd"), |