diff --git a/library/std/src/sys/unix/thread.rs b/library/std/src/sys/unix/thread.rs index c5b905685700b..af281a255cc26 100644 --- a/library/std/src/sys/unix/thread.rs +++ b/library/std/src/sys/unix/thread.rs @@ -381,19 +381,27 @@ pub fn available_parallelism() -> io::Result { #[cfg(any(target_os = "android", target_os = "linux"))] mod cgroups { + //! Currently not covered + //! * cgroup v2 in non-standard mountpoints + //! * paths containing control characters or spaces, since those would be escaped in procfs + //! output and we don't unescape + use crate::borrow::Cow; use crate::ffi::OsString; use crate::fs::{try_exists, File}; use crate::io::Read; + use crate::io::{BufRead, BufReader}; use crate::os::unix::ffi::OsStringExt; + use crate::path::Path; use crate::path::PathBuf; use crate::str::from_utf8; + #[derive(PartialEq)] enum Cgroup { V1, V2, } - /// Returns cgroup CPU quota in core-equivalents, rounded down, or usize::MAX if the quota cannot + /// Returns cgroup CPU quota in core-equivalents, rounded down or usize::MAX if the quota cannot /// be determined or is not set. pub(super) fn quota() -> usize { let mut quota = usize::MAX; @@ -407,27 +415,30 @@ mod cgroups { 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, version) = buf - .split(|&c| c == b'\n') - .filter_map(|line| { + let (cgroup_path, version) = + buf.split(|&c| c == b'\n').fold(None, |previous, line| { let mut fields = line.splitn(3, |&c| c == b':'); // 2nd field is a list of controllers for v1 or empty for v2 let version = match fields.nth(1) { - Some(b"") => Some(Cgroup::V2), + Some(b"") => Cgroup::V2, Some(controllers) if from_utf8(controllers) .is_ok_and(|c| c.split(",").any(|c| c == "cpu")) => { - Some(Cgroup::V1) + Cgroup::V1 } - _ => None, - }?; + _ => return previous, + }; + + // already-found v1 trumps v2 since it explicitly specifies its controllers + if previous.is_some() && version == Cgroup::V2 { + return previous; + } let path = fields.last()?; // skip leading slash Some((path[1..].to_owned(), version)) - }) - .next()?; + })?; let cgroup_path = PathBuf::from(OsString::from_vec(cgroup_path)); quota = match version { @@ -493,14 +504,21 @@ mod cgroups { let mut read_buf = String::with_capacity(20); // Hardcode commonly used locations mentioned in the cgroups(7) manpage - // since scanning mountinfo can be expensive on some systems. - // This isn't exactly standardized since cgroupv1 was meant to allow flexibly - // mixing and matching controller hierarchies. - let mounts = ["/sys/fs/cgroup/cpu", "/sys/fs/cgroup/cpu,cpuacct"]; + // if that doesn't work scan mountinfo and adjust `group_path` for bind-mounts + let mounts: &[fn(&Path) -> Option<(_, &Path)>] = &[ + |p| Some((Cow::Borrowed("/sys/fs/cgroup/cpu"), p)), + |p| Some((Cow::Borrowed("/sys/fs/cgroup/cpu,cpuacct"), p)), + // this can be expensive on systems with tons of mountpoints + // but we only get to this point when /proc/self/cgroups explicitly indicated + // this process belongs to a cpu-controller cgroup v1 and the defaults didn't work + find_mountpoint, + ]; for mount in mounts { + let Some((mount, group_path)) = mount(&group_path) else { continue }; + path.clear(); - path.push(mount); + path.push(mount.as_ref()); path.push(&group_path); // skip if we guessed the mount incorrectly @@ -508,16 +526,16 @@ mod cgroups { continue; } - while path.starts_with(mount) { + while path.starts_with(mount.as_ref()) { let mut parse_file = |name| { path.push(name); read_buf.clear(); - let mut f = File::open(&path).ok()?; - f.read_to_string(&mut read_buf).ok()?; + let f = File::open(&path); + path.pop(); // restore buffer before any early returns + f.ok()?.read_to_string(&mut read_buf).ok()?; let parsed = read_buf.trim().parse::().ok()?; - path.pop(); Some(parsed) }; @@ -531,10 +549,56 @@ mod cgroups { path.pop(); } + + // we passed the try_exists above so we should have traversed the correct hierarchy + // when reaching this line + break; } quota } + + /// Scan mountinfo for cgroup v1 mountpoint with a cpu controller + /// + /// If the cgroupfs is a bind mount then `group_path` is adjusted to skip + /// over the already-included prefix + fn find_mountpoint(group_path: &Path) -> Option<(Cow<'static, str>, &Path)> { + let mut reader = BufReader::new(File::open("/proc/self/mountinfo").ok()?); + let mut line = String::with_capacity(256); + loop { + line.clear(); + if reader.read_line(&mut line).ok()? == 0 { + break; + } + + let line = line.trim(); + let mut items = line.split(' '); + + let sub_path = items.nth(3)?; + let mount_point = items.next()?; + let mount_opts = items.next_back()?; + let filesystem_type = items.nth_back(1)?; + + if filesystem_type != "cgroup" || !mount_opts.split(',').any(|opt| opt == "cpu") { + // not a cgroup / not a cpu-controller + continue; + } + + let sub_path = Path::new(sub_path).strip_prefix("/").ok()?; + + if !group_path.starts_with(sub_path) { + // this is a bind-mount and the bound subdirectory + // does not contain the cgroup this process belongs to + continue; + } + + let trimmed_group_path = group_path.strip_prefix(sub_path).ok()?; + + return Some((Cow::Owned(mount_point.to_owned()), trimmed_group_path)); + } + + None + } } #[cfg(all(