1use std::path::{Path, PathBuf};
38
39const CGROUP_ROOT: &str = "/sys/fs/cgroup";
41
42const V1_UNLIMITED: u64 = u64::MAX / 4;
49
50#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
52pub struct Cap {
53 pub cgroup: Option<u64>,
56 pub host: Option<u64>,
58}
59
60#[must_use]
66pub fn cap() -> Cap {
67 static ONCE: std::sync::OnceLock<Cap> = std::sync::OnceLock::new();
68 *ONCE.get_or_init(Cap::read)
69}
70
71impl Cap {
72 #[must_use]
74 pub fn read() -> Cap {
75 Cap {
76 cgroup: cgroup_limit(Path::new(CGROUP_ROOT), Path::new("/proc/self/cgroup")),
77 host: host_memory(),
78 }
79 }
80
81 #[must_use]
87 pub fn limit(&self) -> Option<u64> {
88 match (self.cgroup, self.host) {
89 (Some(c), Some(h)) => Some(c.min(h)),
90 (Some(c), None) => Some(c),
91 (None, h) => h,
92 }
93 }
94
95 #[must_use]
100 pub fn budget(&self) -> u64 {
101 self.limit().map_or(0, |b| b / 4)
102 }
103}
104
105fn parse_limit(text: &str) -> Option<u64> {
110 let text = text.trim();
111 if text == "max" {
112 return None;
113 }
114 let n: u64 = text.parse().ok()?;
115 if n >= V1_UNLIMITED { None } else { Some(n) }
116}
117
118fn parse_self_cgroup(text: &str) -> Option<&str> {
124 text.lines()
125 .find_map(|line| line.strip_prefix("0::"))
126 .map(str::trim)
127}
128
129fn cgroup_limit(root: &Path, self_cgroup: &Path) -> Option<u64> {
135 v2_limit(root, self_cgroup).or_else(|| {
138 let text = std::fs::read_to_string(root.join("memory/memory.limit_in_bytes")).ok()?;
139 parse_limit(&text)
140 })
141}
142
143fn v2_limit(root: &Path, self_cgroup: &Path) -> Option<u64> {
146 let rel = std::fs::read_to_string(self_cgroup).ok()?;
147 let rel = parse_self_cgroup(&rel)?.trim_start_matches('/');
148
149 let mut dir: PathBuf = root.join(rel);
150 let mut best: Option<u64> = None;
151 loop {
152 if let Ok(text) = std::fs::read_to_string(dir.join("memory.max"))
153 && let Some(n) = parse_limit(&text)
154 {
155 best = Some(best.map_or(n, |b: u64| b.min(n)));
156 }
157 if dir == root {
158 break;
159 }
160 match dir.parent() {
163 Some(p) if p.starts_with(root) || p == root => dir = p.to_path_buf(),
164 _ => break,
165 }
166 }
167 best
168}
169
170#[cfg(target_os = "linux")]
172fn host_memory() -> Option<u64> {
173 let (pages, size) = unsafe {
176 (
177 libc::sysconf(libc::_SC_PHYS_PAGES),
178 libc::sysconf(libc::_SC_PAGESIZE),
179 )
180 };
181 if pages > 0 && size > 0 {
182 Some(pages as u64 * size as u64)
183 } else {
184 None
185 }
186}
187
188#[cfg(target_vendor = "apple")]
190fn host_memory() -> Option<u64> {
191 let mut out: u64 = 0;
192 let mut len = size_of::<u64>();
193 let rc = unsafe {
197 libc::sysctlbyname(
198 c"hw.memsize".as_ptr(),
199 (&raw mut out).cast(),
200 &raw mut len,
201 std::ptr::null_mut(),
202 0,
203 )
204 };
205 if rc == 0 && out > 0 { Some(out) } else { None }
206}
207
208#[cfg(not(any(target_os = "linux", target_vendor = "apple")))]
210fn host_memory() -> Option<u64> {
211 None
212}
213
214#[cfg(test)]
215mod tests {
216 use super::*;
217
218 struct Tree(PathBuf);
221
222 impl Tree {
223 fn new(name: &str) -> Tree {
224 let dir = std::env::temp_dir().join(format!("yo-cap-{name}-{}", std::process::id()));
225 let _ = std::fs::remove_dir_all(&dir);
226 std::fs::create_dir_all(&dir).expect("could not make a temporary directory");
227 Tree(dir)
228 }
229
230 fn write(&self, rel: &str, text: &str) -> PathBuf {
231 let at = self.0.join(rel);
232 std::fs::create_dir_all(at.parent().expect("a file has a parent"))
233 .expect("could not make a directory");
234 std::fs::write(&at, text).expect("could not write");
235 at
236 }
237
238 fn root(&self) -> PathBuf {
239 self.0.join("cgroup")
240 }
241 }
242
243 impl Drop for Tree {
244 fn drop(&mut self) {
245 let _ = std::fs::remove_dir_all(&self.0);
246 }
247 }
248
249 #[test]
250 fn max_means_there_is_no_limit() {
251 assert_eq!(parse_limit("max\n"), None);
252 assert_eq!(parse_limit(" max "), None);
253 }
254
255 #[test]
256 fn a_number_near_the_top_is_cgroup_v1_saying_no_limit() {
257 assert_eq!(parse_limit("9223372036854771712"), None);
258 assert_eq!(parse_limit(&u64::MAX.to_string()), None);
259 }
260
261 #[test]
262 fn a_real_number_is_bytes() {
263 assert_eq!(parse_limit("2147483648\n"), Some(2 * 1024 * 1024 * 1024));
264 }
265
266 #[test]
267 fn nonsense_is_no_limit_rather_than_a_panic() {
268 assert_eq!(parse_limit(""), None);
269 assert_eq!(parse_limit("-1"), None);
270 assert_eq!(parse_limit("2gb"), None);
271 }
272
273 #[test]
274 fn the_v2_line_is_the_one_with_no_controllers() {
275 let text = "12:pids:/user.slice\n1:name=systemd:/user.slice\n0::/user.slice/app.scope\n";
276 assert_eq!(parse_self_cgroup(text), Some("/user.slice/app.scope"));
277 }
278
279 #[test]
280 fn inside_a_cgroup_namespace_the_path_is_just_the_root() {
281 assert_eq!(parse_self_cgroup("0::/\n"), Some("/"));
282 }
283
284 #[test]
285 fn a_v1_only_machine_has_no_v2_line() {
286 assert_eq!(parse_self_cgroup("6:memory:/\n3:cpu:/\n"), None);
287 }
288
289 #[test]
290 fn the_limit_is_read_from_the_leaf() {
291 let t = Tree::new("leaf");
292 let me = t.write("proc", "0::/a/b\n");
293 t.write("cgroup/a/b/memory.max", "1073741824\n");
294 assert_eq!(cgroup_limit(&t.root(), &me), Some(1024 * 1024 * 1024));
295 }
296
297 #[test]
298 fn an_ancestor_with_a_tighter_limit_wins() {
299 let t = Tree::new("ancestor");
304 let me = t.write("proc", "0::/pods/one\n");
305 t.write("cgroup/memory.max", "max\n");
306 t.write("cgroup/pods/memory.max", "536870912\n");
307 t.write("cgroup/pods/one/memory.max", "4294967296\n");
308 assert_eq!(cgroup_limit(&t.root(), &me), Some(512 * 1024 * 1024));
309 }
310
311 #[test]
312 fn a_tree_that_says_max_all_the_way_up_has_no_limit() {
313 let t = Tree::new("nolimit");
314 let me = t.write("proc", "0::/a\n");
315 t.write("cgroup/memory.max", "max\n");
316 t.write("cgroup/a/memory.max", "max\n");
317 assert_eq!(cgroup_limit(&t.root(), &me), None);
318 }
319
320 #[test]
321 fn cgroup_v1_is_read_when_v2_has_nothing_to_say() {
322 let t = Tree::new("v1");
323 let me = t.write("proc", "6:memory:/\n");
324 t.write("cgroup/memory/memory.limit_in_bytes", "268435456\n");
325 assert_eq!(cgroup_limit(&t.root(), &me), Some(256 * 1024 * 1024));
326 }
327
328 #[test]
329 fn a_machine_with_no_cgroups_at_all_reports_none() {
330 let t = Tree::new("nothing");
331 assert_eq!(
332 cgroup_limit(&t.root(), &t.0.join("not-here")),
333 None,
334 "a missing file is a machine without cgroups, not an error"
335 );
336 }
337
338 #[test]
339 fn the_tighter_of_the_two_is_the_one_that_counts() {
340 let big = 64 * 1024 * 1024 * 1024;
341 let small = 2 * 1024 * 1024 * 1024;
342 assert_eq!(
343 Cap {
344 cgroup: Some(small),
345 host: Some(big)
346 }
347 .limit(),
348 Some(small)
349 );
350 assert_eq!(
353 Cap {
354 cgroup: Some(big),
355 host: Some(small)
356 }
357 .limit(),
358 Some(small)
359 );
360 }
361
362 #[test]
363 fn the_budget_is_a_quarter_and_zero_when_there_is_nothing_to_take_a_quarter_of() {
364 let cap = Cap {
365 cgroup: Some(4 * 1024 * 1024 * 1024),
366 host: None,
367 };
368 assert_eq!(cap.budget(), 1024 * 1024 * 1024);
369 assert_eq!(Cap::default().budget(), 0);
370 }
371
372 #[test]
373 fn asking_the_real_machine_answers_something_sensible() {
374 let cap = Cap::read();
375 if let Some(h) = cap.host {
378 assert!(h >= 64 * 1024 * 1024, "a host with {h} bytes is not real");
379 }
380 if let Some(l) = cap.limit() {
381 assert_eq!(cap.budget(), l / 4);
382 }
383 }
384}