use std::path::{Path, PathBuf};
const CGROUP_ROOT: &str = "/sys/fs/cgroup";
const V1_UNLIMITED: u64 = u64::MAX / 4;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct Cap {
pub cgroup: Option<u64>,
pub host: Option<u64>,
}
#[must_use]
pub fn cap() -> Cap {
static ONCE: std::sync::OnceLock<Cap> = std::sync::OnceLock::new();
*ONCE.get_or_init(Cap::read)
}
impl Cap {
#[must_use]
pub fn read() -> Cap {
Cap {
cgroup: cgroup_limit(Path::new(CGROUP_ROOT), Path::new("/proc/self/cgroup")),
host: host_memory(),
}
}
#[must_use]
pub fn limit(&self) -> Option<u64> {
match (self.cgroup, self.host) {
(Some(c), Some(h)) => Some(c.min(h)),
(Some(c), None) => Some(c),
(None, h) => h,
}
}
#[must_use]
pub fn budget(&self) -> u64 {
self.limit().map_or(0, |b| b / 4)
}
}
fn parse_limit(text: &str) -> Option<u64> {
let text = text.trim();
if text == "max" {
return None;
}
let n: u64 = text.parse().ok()?;
if n >= V1_UNLIMITED { None } else { Some(n) }
}
fn parse_self_cgroup(text: &str) -> Option<&str> {
text.lines()
.find_map(|line| line.strip_prefix("0::"))
.map(str::trim)
}
fn cgroup_limit(root: &Path, self_cgroup: &Path) -> Option<u64> {
v2_limit(root, self_cgroup).or_else(|| {
let text = std::fs::read_to_string(root.join("memory/memory.limit_in_bytes")).ok()?;
parse_limit(&text)
})
}
fn v2_limit(root: &Path, self_cgroup: &Path) -> Option<u64> {
let rel = std::fs::read_to_string(self_cgroup).ok()?;
let rel = parse_self_cgroup(&rel)?.trim_start_matches('/');
let mut dir: PathBuf = root.join(rel);
let mut best: Option<u64> = None;
loop {
if let Ok(text) = std::fs::read_to_string(dir.join("memory.max"))
&& let Some(n) = parse_limit(&text)
{
best = Some(best.map_or(n, |b: u64| b.min(n)));
}
if dir == root {
break;
}
match dir.parent() {
Some(p) if p.starts_with(root) || p == root => dir = p.to_path_buf(),
_ => break,
}
}
best
}
#[cfg(all(target_os = "linux", not(miri)))]
fn host_memory() -> Option<u64> {
let (pages, size) = unsafe {
(
libc::sysconf(libc::_SC_PHYS_PAGES),
libc::sysconf(libc::_SC_PAGESIZE),
)
};
if pages > 0 && size > 0 {
Some(pages as u64 * size as u64)
} else {
None
}
}
#[cfg(all(target_vendor = "apple", not(miri)))]
fn host_memory() -> Option<u64> {
let mut out: u64 = 0;
let mut len = size_of::<u64>();
let rc = unsafe {
libc::sysctlbyname(
c"hw.memsize".as_ptr(),
(&raw mut out).cast(),
&raw mut len,
std::ptr::null_mut(),
0,
)
};
if rc == 0 && out > 0 { Some(out) } else { None }
}
#[cfg(any(miri, not(any(target_os = "linux", target_vendor = "apple"))))]
fn host_memory() -> Option<u64> {
None
}
#[cfg(test)]
mod tests {
use super::*;
struct Tree(PathBuf);
impl Tree {
fn new(name: &str) -> Tree {
let dir = std::env::temp_dir().join(format!("yo-cap-{name}-{}", std::process::id()));
let _ = std::fs::remove_dir_all(&dir);
std::fs::create_dir_all(&dir).expect("could not make a temporary directory");
Tree(dir)
}
fn write(&self, rel: &str, text: &str) -> PathBuf {
let at = self.0.join(rel);
std::fs::create_dir_all(at.parent().expect("a file has a parent"))
.expect("could not make a directory");
std::fs::write(&at, text).expect("could not write");
at
}
fn root(&self) -> PathBuf {
self.0.join("cgroup")
}
}
impl Drop for Tree {
fn drop(&mut self) {
let _ = std::fs::remove_dir_all(&self.0);
}
}
#[test]
fn max_means_there_is_no_limit() {
assert_eq!(parse_limit("max\n"), None);
assert_eq!(parse_limit(" max "), None);
}
#[test]
fn a_number_near_the_top_is_cgroup_v1_saying_no_limit() {
assert_eq!(parse_limit("9223372036854771712"), None);
assert_eq!(parse_limit(&u64::MAX.to_string()), None);
}
#[test]
fn a_real_number_is_bytes() {
assert_eq!(parse_limit("2147483648\n"), Some(2 * 1024 * 1024 * 1024));
}
#[test]
fn nonsense_is_no_limit_rather_than_a_panic() {
assert_eq!(parse_limit(""), None);
assert_eq!(parse_limit("-1"), None);
assert_eq!(parse_limit("2gb"), None);
}
#[test]
fn the_v2_line_is_the_one_with_no_controllers() {
let text = "12:pids:/user.slice\n1:name=systemd:/user.slice\n0::/user.slice/app.scope\n";
assert_eq!(parse_self_cgroup(text), Some("/user.slice/app.scope"));
}
#[test]
fn inside_a_cgroup_namespace_the_path_is_just_the_root() {
assert_eq!(parse_self_cgroup("0::/\n"), Some("/"));
}
#[test]
fn a_v1_only_machine_has_no_v2_line() {
assert_eq!(parse_self_cgroup("6:memory:/\n3:cpu:/\n"), None);
}
#[test]
fn the_limit_is_read_from_the_leaf() {
let t = Tree::new("leaf");
let me = t.write("proc", "0::/a/b\n");
t.write("cgroup/a/b/memory.max", "1073741824\n");
assert_eq!(cgroup_limit(&t.root(), &me), Some(1024 * 1024 * 1024));
}
#[test]
fn an_ancestor_with_a_tighter_limit_wins() {
let t = Tree::new("ancestor");
let me = t.write("proc", "0::/pods/one\n");
t.write("cgroup/memory.max", "max\n");
t.write("cgroup/pods/memory.max", "536870912\n");
t.write("cgroup/pods/one/memory.max", "4294967296\n");
assert_eq!(cgroup_limit(&t.root(), &me), Some(512 * 1024 * 1024));
}
#[test]
fn a_tree_that_says_max_all_the_way_up_has_no_limit() {
let t = Tree::new("nolimit");
let me = t.write("proc", "0::/a\n");
t.write("cgroup/memory.max", "max\n");
t.write("cgroup/a/memory.max", "max\n");
assert_eq!(cgroup_limit(&t.root(), &me), None);
}
#[test]
fn cgroup_v1_is_read_when_v2_has_nothing_to_say() {
let t = Tree::new("v1");
let me = t.write("proc", "6:memory:/\n");
t.write("cgroup/memory/memory.limit_in_bytes", "268435456\n");
assert_eq!(cgroup_limit(&t.root(), &me), Some(256 * 1024 * 1024));
}
#[test]
fn a_machine_with_no_cgroups_at_all_reports_none() {
let t = Tree::new("nothing");
assert_eq!(
cgroup_limit(&t.root(), &t.0.join("not-here")),
None,
"a missing file is a machine without cgroups, not an error"
);
}
#[test]
fn the_tighter_of_the_two_is_the_one_that_counts() {
let big = 64 * 1024 * 1024 * 1024;
let small = 2 * 1024 * 1024 * 1024;
assert_eq!(
Cap {
cgroup: Some(small),
host: Some(big)
}
.limit(),
Some(small)
);
assert_eq!(
Cap {
cgroup: Some(big),
host: Some(small)
}
.limit(),
Some(small)
);
}
#[test]
fn the_budget_is_a_quarter_and_zero_when_there_is_nothing_to_take_a_quarter_of() {
let cap = Cap {
cgroup: Some(4 * 1024 * 1024 * 1024),
host: None,
};
assert_eq!(cap.budget(), 1024 * 1024 * 1024);
assert_eq!(Cap::default().budget(), 0);
}
#[test]
fn asking_the_real_machine_answers_something_sensible() {
let cap = Cap::read();
if let Some(h) = cap.host {
assert!(h >= 64 * 1024 * 1024, "a host with {h} bytes is not real");
}
if let Some(l) = cap.limit() {
assert_eq!(cap.budget(), l / 4);
}
}
}