about summary refs log tree commit diff
path: root/library/std/src/sys
diff options
context:
space:
mode:
authorThe 8472 <git@infinite-source.de>2022-01-09 19:39:02 +0100
committerThe 8472 <git@infinite-source.de>2022-03-03 00:43:45 +0100
commitbac5523ea0025988361a092c1f0e7f4eb90f6ad7 (patch)
tree825750b995d9f2c9b1c2bab49094f7196acb5673 /library/std/src/sys
parent8769f4ef2fe1efddd1f072485f97f568e7328f79 (diff)
downloadrust-bac5523ea0025988361a092c1f0e7f4eb90f6ad7.tar.gz
rust-bac5523ea0025988361a092c1f0e7f4eb90f6ad7.zip
Use cgroup quotas for calculating `available_parallelism`
Manually tested via


```
// spawn a new cgroup scope for the current user
$ sudo systemd-run -p CPUQuota="300%" --uid=$(id -u) -tdS


// quota.rs
#![feature(available_parallelism)]
fn main() {
    println!("{:?}", std::thread::available_parallelism()); // prints Ok(3)
}
```


Caveats

* cgroup v1 is ignored
* funky mountpoints (containing spaces, newlines or control chars) for cgroupfs will not be handled correctly since that would require unescaping /proc/self/mountinfo
  The escaping behavior of procfs seems to be undocumented. systemd and docker default to `/sys/fs/cgroup` so it should be fine for most systems.
* quota will be ignored when `sched_getaffinity` doesn't work
* assumes procfs is mounted under `/proc` and cgroupfs mounted and readable somewhere in the directory tree
Diffstat (limited to 'library/std/src/sys')
-rw-r--r--library/std/src/sys/unix/thread.rs71
1 files changed, 68 insertions, 3 deletions
diff --git a/library/std/src/sys/unix/thread.rs b/library/std/src/sys/unix/thread.rs
index cf8cf5ad49f..933210e1ff0 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().unwrap_or(usize::MAX).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,66 @@ pub fn available_parallelism() -> io::Result<NonZeroUsize> {
     }
 }
 
+#[cfg(any(target_os = "android", target_os = "linux"))]
+fn cgroup2_quota() -> Option<usize> {
+    use crate::ffi::OsString;
+    use crate::fs::{read, read_to_string, File};
+    use crate::io::{BufRead, BufReader};
+    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,
+            }
+        })
+        .min()
+}
+
 #[cfg(all(
     not(target_os = "linux"),
     not(target_os = "freebsd"),