tawasuyu

tawasuyu/brahman

Rust 326TOML 123JSON 17Markdown 144606 Definitions

⎇ 72758e75ce515f355f27615671c95ecaf8c2a8bc crates/shared/ente-incarnate/src/cgroup.rs Rust · 116 lines · 4.0 KB History Who touched what raw

feat(shipote): quota enforce + cgroup memory.max + pipeline restart (fase L) 4c9d1b4c1 · sergio ·

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}