about summary refs log tree commit diff
path: root/library/std/src/sys/unix/thread.rs
diff options
context:
space:
mode:
authorThe 8472 <git@infinite-source.de>2022-03-03 00:35:47 +0100
committerThe 8472 <git@infinite-source.de>2022-03-03 00:43:46 +0100
commitaf6d2ed24557694ff7d32bf2a29a6cd5aaade859 (patch)
tree9cf146a1f1d8ca537bb495c852bd045ee7f61943 /library/std/src/sys/unix/thread.rs
parentbac5523ea0025988361a092c1f0e7f4eb90f6ad7 (diff)
downloadrust-af6d2ed24557694ff7d32bf2a29a6cd5aaade859.tar.gz
rust-af6d2ed24557694ff7d32bf2a29a6cd5aaade859.zip
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
Diffstat (limited to 'library/std/src/sys/unix/thread.rs')
-rw-r--r--library/std/src/sys/unix/thread.rs120
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(