⎇ 72758e75ce515f355f27615671c95ecaf8c2a8bc
crates/shared/ente-incarnate/src/cgroup.rs
Historia
Quién tocó qué
crudo
feat(shipote): quota enforce + cgroup memory.max + pipeline restart (fase L)
1//! Resolución y creación de cgroups v2 para el hijo.
2
3use crate::error::IncarnateError;
4use brahman_card::{CgroupSpec, ResourceLimits};
5use std::path::{Path, PathBuf};
6
7/// Cgroup actual del proceso que llama. Lo usamos como prefijo para paths
8/// declarados relativos en `CgroupSpec.path`.
9pub fn current_cgroup() -> Option<String> {
10 let s = std::fs::read_to_string("/proc/self/cgroup").ok()?;
11 s.lines()
12 .find_map(|l| l.strip_prefix("0::"))
13 .map(|s| s.trim().to_string())
14}
15
16/// Resuelve un path declarado contra la jerarquía real.
17pub fn resolve_cgroup_path(spec_path: &str) -> String {
18 if spec_path.is_empty() {
19 return String::new();
20 }
21 if spec_path.starts_with('/') {
22 return spec_path.to_string();
23 }
24 let trimmed = spec_path.trim_start_matches('/');
25 if let Some(cg) = current_cgroup() {
26 let base = if cg == "/" {
27 String::new()
28 } else {
29 cg.trim_end_matches('/').to_string()
30 };
31 format!("{base}/{trimmed}")
32 } else {
33 format!("/{trimmed}")
34 }
35}
36
37/// Crea el cgroup declarado y aplica weights. Devuelve el path absoluto
38/// resultante bajo `/sys/fs/cgroup`.
39pub fn ensure_cgroup(spec: &CgroupSpec) -> Result<PathBuf, IncarnateError> {
40 let rel = resolve_cgroup_path(&spec.path);
41 if rel.is_empty() {
42 return Err(IncarnateError::CgroupNotWritable {
43 path: PathBuf::from("(empty)"),
44 });
45 }
46 let abs = PathBuf::from(format!("/sys/fs/cgroup{}", rel));
47 std::fs::create_dir_all(&abs).map_err(|e| match e.kind() {
48 std::io::ErrorKind::PermissionDenied => IncarnateError::CgroupNotWritable { path: abs.clone() },
49 _ => IncarnateError::Io(e),
50 })?;
51 if let Some(w) = spec.cpu_weight {
52 let _ = std::fs::write(abs.join("cpu.weight"), format!("{w}\n"));
53 }
54 if let Some(w) = spec.io_weight {
55 // io.weight requiere "default <n>" en cgroup v2.
56 let _ = std::fs::write(abs.join("io.weight"), format!("default {w}\n"));
57 }
58 Ok(abs)
59}
60
61/// Escribe `memory.max` y `pids.max` al cgroup según `rlimits`. Falla
62/// silenciosamente si los archivos no son escribibles (cgroup no
63/// delegated). El kernel hace OOM kill cuando `memory.max` se excede,
64/// y bloquea forks cuando `pids.max` se alcanza.
65///
66/// `memory.max` acepta `max` o un número en bytes. `pids.max` igual.
67pub fn apply_rlimits_to_cgroup(cgroup_abs: &Path, rlimits: &ResourceLimits) -> Vec<String> {
68 let mut applied = Vec::new();
69 if let Some(mem) = rlimits.mem_bytes {
70 let path = cgroup_abs.join("memory.max");
71 match std::fs::write(&path, format!("{mem}\n")) {
72 Ok(_) => applied.push(format!("memory.max={mem}")),
73 Err(e) => tracing::warn!(?e, path = %path.display(), "memory.max write failed"),
74 }
75 }
76 if let Some(np) = rlimits.nproc {
77 let path = cgroup_abs.join("pids.max");
78 match std::fs::write(&path, format!("{np}\n")) {
79 Ok(_) => applied.push(format!("pids.max={np}")),
80 Err(e) => tracing::warn!(?e, path = %path.display(), "pids.max write failed"),
81 }
82 }
83 applied
84}
85
86/// Mueve `pid` a `cgroup_abs/cgroup.procs`.
87pub fn move_to_cgroup(cgroup_abs: &Path, pid: nix::unistd::Pid) -> Result<(), IncarnateError> {
88 let procs = cgroup_abs.join("cgroup.procs");
89 std::fs::write(&procs, format!("{}\n", pid.as_raw())).map_err(|e| match e.kind() {
90 std::io::ErrorKind::PermissionDenied => IncarnateError::CgroupNotWritable {
91 path: procs.clone(),
92 },
93 _ => IncarnateError::Io(e),
94 })
95}
96
97#[cfg(test)]
98mod tests {
99 use super::*;
100
101 #[test]
102 fn absolute_path_passthrough() {
103 assert_eq!(resolve_cgroup_path("/foo/bar"), "/foo/bar");
104 }
105
106 #[test]
107 fn empty_returns_empty() {
108 assert_eq!(resolve_cgroup_path(""), "");
109 }
110
111 #[test]
112 fn relative_path_prefixed() {
113 let r = resolve_cgroup_path("shipote/ws-1");
114 assert!(r.ends_with("/shipote/ws-1") || r == "/shipote/ws-1");
115 }
116}