Use cgroup quotas for calculating available_parallelism

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,80 @@ 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;
+
+    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"),

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