Skip to content

Commit

Permalink
scan mountinfo when hardcoded cgroupv1 mountpoints don't work
Browse files Browse the repository at this point in the history
  • Loading branch information
the8472 committed Jun 24, 2022
1 parent d823462 commit b2c410e
Showing 1 changed file with 83 additions and 19 deletions.
102 changes: 83 additions & 19 deletions library/std/src/sys/unix/thread.rs
Original file line number Diff line number Diff line change
Expand Up @@ -381,19 +381,27 @@ pub fn available_parallelism() -> io::Result<NonZeroUsize> {

#[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;
Expand All @@ -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 {
Expand Down Expand Up @@ -493,31 +504,38 @@ 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
if matches!(try_exists(&path), Err(_) | Ok(false)) {
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::<usize>().ok()?;

path.pop();
Some(parsed)
};

Expand All @@ -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(
Expand Down

0 comments on commit b2c410e

Please sign in to comment.