1 files changed, 67 insertions, 53 deletions

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<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 cgroups_mount = PathBuf::from(OsString::from_vec(cgroups_mount));
-
-    // 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 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,
+    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);
+                    }
+                    _ => {}
+                }
             }
-        })
-        .min()
+
+            path.pop(); // pop filename
+            path.pop(); // pop dir
+        }
+    };
+
+    quota
 }
 
 #[cfg(all(