From af6d2ed24557694ff7d32bf2a29a6cd5aaade859 Mon Sep 17 00:00:00 2001
From: The 8472 <git@infinite-source.de>
Date: Thu, 3 Mar 2022 00:35:47 +0100
Subject: [PATCH] hardcode /sys/fs/cgroup instead of doing a lookup via
 mountinfo

this avoids parsing mountinfo which can be huge on some systems and
something might be emulating cgroup fs for sandboxing reasons which means
it wouldn't show up as mountpoint

additionally the new implementation operates on a single pathbuffer, reducing allocations
---
 library/std/src/sys/unix/thread.rs | 114 ++++++++++++++++-------------
 1 file changed, 64 insertions(+), 50 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 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::<usize>(), period.parse::<usize>()) {
+                    (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::<usize>(), period.parse::<usize>()) {
-                (Ok(limit), Ok(period)) => Some(limit / period),
-                _ => None,
-            }
-        })
-        .min()
+            path.pop(); // pop filename
+            path.pop(); // pop dir
+        }
+    };
+
+    quota
 }
 
 #[cfg(all(