Skip to content

Commit b2c410e

Browse files
committed
scan mountinfo when hardcoded cgroupv1 mountpoints don't work
1 parent d823462 commit b2c410e

File tree

1 file changed

+83
-19
lines changed

1 file changed

+83
-19
lines changed

library/std/src/sys/unix/thread.rs

+83-19
Original file line numberDiff line numberDiff line change
@@ -381,19 +381,27 @@ pub fn available_parallelism() -> io::Result<NonZeroUsize> {
381381

382382
#[cfg(any(target_os = "android", target_os = "linux"))]
383383
mod cgroups {
384+
//! Currently not covered
385+
//! * cgroup v2 in non-standard mountpoints
386+
//! * paths containing control characters or spaces, since those would be escaped in procfs
387+
//! output and we don't unescape
388+
use crate::borrow::Cow;
384389
use crate::ffi::OsString;
385390
use crate::fs::{try_exists, File};
386391
use crate::io::Read;
392+
use crate::io::{BufRead, BufReader};
387393
use crate::os::unix::ffi::OsStringExt;
394+
use crate::path::Path;
388395
use crate::path::PathBuf;
389396
use crate::str::from_utf8;
390397

398+
#[derive(PartialEq)]
391399
enum Cgroup {
392400
V1,
393401
V2,
394402
}
395403

396-
/// Returns cgroup CPU quota in core-equivalents, rounded down, or usize::MAX if the quota cannot
404+
/// Returns cgroup CPU quota in core-equivalents, rounded down or usize::MAX if the quota cannot
397405
/// be determined or is not set.
398406
pub(super) fn quota() -> usize {
399407
let mut quota = usize::MAX;
@@ -407,27 +415,30 @@ mod cgroups {
407415
let mut buf = Vec::with_capacity(128);
408416
// find our place in the cgroup hierarchy
409417
File::open("/proc/self/cgroup").ok()?.read_to_end(&mut buf).ok()?;
410-
let (cgroup_path, version) = buf
411-
.split(|&c| c == b'\n')
412-
.filter_map(|line| {
418+
let (cgroup_path, version) =
419+
buf.split(|&c| c == b'\n').fold(None, |previous, line| {
413420
let mut fields = line.splitn(3, |&c| c == b':');
414421
// 2nd field is a list of controllers for v1 or empty for v2
415422
let version = match fields.nth(1) {
416-
Some(b"") => Some(Cgroup::V2),
423+
Some(b"") => Cgroup::V2,
417424
Some(controllers)
418425
if from_utf8(controllers)
419426
.is_ok_and(|c| c.split(",").any(|c| c == "cpu")) =>
420427
{
421-
Some(Cgroup::V1)
428+
Cgroup::V1
422429
}
423-
_ => None,
424-
}?;
430+
_ => return previous,
431+
};
432+
433+
// already-found v1 trumps v2 since it explicitly specifies its controllers
434+
if previous.is_some() && version == Cgroup::V2 {
435+
return previous;
436+
}
425437

426438
let path = fields.last()?;
427439
// skip leading slash
428440
Some((path[1..].to_owned(), version))
429-
})
430-
.next()?;
441+
})?;
431442
let cgroup_path = PathBuf::from(OsString::from_vec(cgroup_path));
432443

433444
quota = match version {
@@ -493,31 +504,38 @@ mod cgroups {
493504
let mut read_buf = String::with_capacity(20);
494505

495506
// Hardcode commonly used locations mentioned in the cgroups(7) manpage
496-
// since scanning mountinfo can be expensive on some systems.
497-
// This isn't exactly standardized since cgroupv1 was meant to allow flexibly
498-
// mixing and matching controller hierarchies.
499-
let mounts = ["/sys/fs/cgroup/cpu", "/sys/fs/cgroup/cpu,cpuacct"];
507+
// if that doesn't work scan mountinfo and adjust `group_path` for bind-mounts
508+
let mounts: &[fn(&Path) -> Option<(_, &Path)>] = &[
509+
|p| Some((Cow::Borrowed("/sys/fs/cgroup/cpu"), p)),
510+
|p| Some((Cow::Borrowed("/sys/fs/cgroup/cpu,cpuacct"), p)),
511+
// this can be expensive on systems with tons of mountpoints
512+
// but we only get to this point when /proc/self/cgroups explicitly indicated
513+
// this process belongs to a cpu-controller cgroup v1 and the defaults didn't work
514+
find_mountpoint,
515+
];
500516

501517
for mount in mounts {
518+
let Some((mount, group_path)) = mount(&group_path) else { continue };
519+
502520
path.clear();
503-
path.push(mount);
521+
path.push(mount.as_ref());
504522
path.push(&group_path);
505523

506524
// skip if we guessed the mount incorrectly
507525
if matches!(try_exists(&path), Err(_) | Ok(false)) {
508526
continue;
509527
}
510528

511-
while path.starts_with(mount) {
529+
while path.starts_with(mount.as_ref()) {
512530
let mut parse_file = |name| {
513531
path.push(name);
514532
read_buf.clear();
515533

516-
let mut f = File::open(&path).ok()?;
517-
f.read_to_string(&mut read_buf).ok()?;
534+
let f = File::open(&path);
535+
path.pop(); // restore buffer before any early returns
536+
f.ok()?.read_to_string(&mut read_buf).ok()?;
518537
let parsed = read_buf.trim().parse::<usize>().ok()?;
519538

520-
path.pop();
521539
Some(parsed)
522540
};
523541

@@ -531,10 +549,56 @@ mod cgroups {
531549

532550
path.pop();
533551
}
552+
553+
// we passed the try_exists above so we should have traversed the correct hierarchy
554+
// when reaching this line
555+
break;
534556
}
535557

536558
quota
537559
}
560+
561+
/// Scan mountinfo for cgroup v1 mountpoint with a cpu controller
562+
///
563+
/// If the cgroupfs is a bind mount then `group_path` is adjusted to skip
564+
/// over the already-included prefix
565+
fn find_mountpoint(group_path: &Path) -> Option<(Cow<'static, str>, &Path)> {
566+
let mut reader = BufReader::new(File::open("/proc/self/mountinfo").ok()?);
567+
let mut line = String::with_capacity(256);
568+
loop {
569+
line.clear();
570+
if reader.read_line(&mut line).ok()? == 0 {
571+
break;
572+
}
573+
574+
let line = line.trim();
575+
let mut items = line.split(' ');
576+
577+
let sub_path = items.nth(3)?;
578+
let mount_point = items.next()?;
579+
let mount_opts = items.next_back()?;
580+
let filesystem_type = items.nth_back(1)?;
581+
582+
if filesystem_type != "cgroup" || !mount_opts.split(',').any(|opt| opt == "cpu") {
583+
// not a cgroup / not a cpu-controller
584+
continue;
585+
}
586+
587+
let sub_path = Path::new(sub_path).strip_prefix("/").ok()?;
588+
589+
if !group_path.starts_with(sub_path) {
590+
// this is a bind-mount and the bound subdirectory
591+
// does not contain the cgroup this process belongs to
592+
continue;
593+
}
594+
595+
let trimmed_group_path = group_path.strip_prefix(sub_path).ok()?;
596+
597+
return Some((Cow::Owned(mount_point.to_owned()), trimmed_group_path));
598+
}
599+
600+
None
601+
}
538602
}
539603

540604
#[cfg(all(

0 commit comments

Comments
 (0)