scx_cosmos/
cgroup.rs

1// SPDX-License-Identifier: GPL-2.0
2
3use std::collections::HashSet;
4use std::fs::File;
5use std::io::{BufRead, BufReader};
6use std::path::{Component, Path, PathBuf};
7
8use anyhow::{bail, Context, Result};
9
10/// Read process membership from the host's unified cgroup hierarchy.
11pub(crate) struct CgroupReader {
12    mount_point: PathBuf,
13}
14
15impl CgroupReader {
16    /// Locate the host cgroup v2 mount from mountinfo.
17    pub(crate) fn discover() -> Result<Self> {
18        let file = File::open("/proc/self/mountinfo").context("open /proc/self/mountinfo")?;
19        let mount_point =
20            parse_cgroup2_mount(BufReader::new(file))?.context("root cgroup v2 mount not found")?;
21        Ok(Self { mount_point })
22    }
23
24    /// Return the exact, non-root cgroup directory for a process.
25    pub(crate) fn process_cgroup(&self, tgid: u32) -> Result<Option<PathBuf>> {
26        let path = format!("/proc/{tgid}/cgroup");
27        let file = File::open(&path).with_context(|| format!("open {path}"))?;
28        let cgroup_path = parse_process_cgroup(BufReader::new(file))?;
29        let Some(relative) = relative_cgroup_path(&cgroup_path)? else {
30            return Ok(None);
31        };
32        Ok(Some(self.mount_point.join(relative)))
33    }
34
35    /// Return all process IDs in an exact cgroup (not its descendants).
36    pub(crate) fn processes(&self, cgroup: &Path) -> Result<HashSet<u32>> {
37        let type_path = cgroup.join("cgroup.type");
38        let cgroup_type = std::fs::read_to_string(&type_path)
39            .with_context(|| format!("read {}", type_path.display()))?;
40        if cgroup_type.trim_end() != "domain" {
41            bail!(
42                "cgroup {} is not an exact domain: unsupported cgroup type {:?}",
43                cgroup.display(),
44                cgroup_type.trim_end()
45            );
46        }
47
48        let path = cgroup.join("cgroup.procs");
49        let file = File::open(&path).with_context(|| format!("open {}", path.display()))?;
50        parse_cgroup_procs(BufReader::new(file))
51    }
52}
53
54fn parse_process_cgroup(reader: impl BufRead) -> Result<PathBuf> {
55    let mut unified = None;
56
57    for line in reader.lines() {
58        let line = line.context("read process cgroup entry")?;
59        let Some(path) = line.strip_prefix("0::") else {
60            continue;
61        };
62        if unified.is_some() {
63            bail!("multiple cgroup v2 entries");
64        }
65        if path.ends_with(" (deleted)") {
66            bail!("cgroup was deleted");
67        }
68        unified = Some(PathBuf::from(path));
69    }
70
71    unified.context("cgroup v2 entry not found")
72}
73
74fn relative_cgroup_path(path: &Path) -> Result<Option<PathBuf>> {
75    if !path.is_absolute() {
76        bail!("cgroup path is not absolute: {}", path.display());
77    }
78    let relative = path
79        .strip_prefix(Path::new("/"))
80        .context("strip cgroup root")?;
81    if relative.as_os_str().is_empty() {
82        return Ok(None);
83    }
84    if relative
85        .components()
86        .any(|component| !matches!(component, Component::Normal(_)))
87    {
88        bail!("invalid cgroup path: {}", path.display());
89    }
90    Ok(Some(relative.to_path_buf()))
91}
92
93fn parse_cgroup_procs(reader: impl BufRead) -> Result<HashSet<u32>> {
94    let mut processes = HashSet::new();
95
96    for line in reader.lines() {
97        let line = line.context("read cgroup.procs entry")?;
98        let value = line.trim();
99        if value.is_empty() {
100            continue;
101        }
102        let tgid = value
103            .parse::<u32>()
104            .with_context(|| format!("invalid process ID in cgroup.procs: {value}"))?;
105        if tgid == 0 {
106            bail!("invalid process ID in cgroup.procs: 0");
107        }
108        processes.insert(tgid);
109    }
110
111    Ok(processes)
112}
113
114fn parse_cgroup2_mount(reader: impl BufRead) -> Result<Option<PathBuf>> {
115    for line in reader.lines() {
116        let line = line.context("read mountinfo entry")?;
117        let Some((mount_fields, fs_fields)) = line.split_once(" - ") else {
118            continue;
119        };
120        let mount_fields: Vec<_> = mount_fields.split_whitespace().collect();
121        let mut fs_fields = fs_fields.split_whitespace();
122        if mount_fields.len() < 5 || fs_fields.next() != Some("cgroup2") {
123            continue;
124        }
125        // Require the hierarchy root so process cgroup paths resolve beneath
126        // this mount without namespace-specific path translation.
127        if mount_fields[3] != "/" {
128            continue;
129        }
130        return Ok(Some(decode_mount_path(mount_fields[4])?));
131    }
132
133    Ok(None)
134}
135
136fn decode_mount_path(value: &str) -> Result<PathBuf> {
137    let mut input = value.chars();
138    let mut output = String::with_capacity(value.len());
139
140    while let Some(character) = input.next() {
141        if character != '\\' {
142            output.push(character);
143            continue;
144        }
145
146        let mut digits = [0u16; 3];
147        for digit in &mut digits {
148            let Some(character @ '0'..='7') = input.next() else {
149                bail!("invalid mountinfo path escape: {value}");
150            };
151            *digit = u16::from(character as u8 - b'0');
152        }
153        let decoded = (digits[0] << 6) | (digits[1] << 3) | digits[2];
154        if decoded > u16::from(u8::MAX) {
155            bail!("mountinfo path escape is out of range: {value}");
156        }
157        let decoded = decoded as u8;
158        if !decoded.is_ascii() {
159            bail!("non-ASCII mountinfo path escape is unsupported: {value}");
160        }
161        output.push(char::from(decoded));
162    }
163
164    Ok(PathBuf::from(output))
165}
166
167#[cfg(test)]
168mod tests {
169    use std::io::Cursor;
170
171    use super::*;
172
173    #[test]
174    fn parses_unified_cgroup_entry() {
175        let input = b"7:cpu:/legacy\n0::/services/inference\n";
176        assert_eq!(
177            parse_process_cgroup(Cursor::new(input)).unwrap(),
178            PathBuf::from("/services/inference")
179        );
180    }
181
182    #[test]
183    fn rejects_deleted_or_missing_cgroup() {
184        assert!(parse_process_cgroup(Cursor::new(b"0::/gone (deleted)\n")).is_err());
185        assert!(parse_process_cgroup(Cursor::new(b"7:cpu:/legacy\n")).is_err());
186        assert!(parse_process_cgroup(Cursor::new(b"0::/one\n0::/two\n")).is_err());
187    }
188
189    #[test]
190    fn rejects_root_and_unsafe_paths() {
191        assert_eq!(relative_cgroup_path(Path::new("/")).unwrap(), None);
192        assert!(relative_cgroup_path(Path::new("relative")).is_err());
193        assert!(relative_cgroup_path(Path::new("/safe/../unsafe")).is_err());
194    }
195
196    #[test]
197    fn parses_cgroup2_mount_and_escapes() {
198        let input = b"10 1 0:1 / /old rw - cgroup cgroup rw\n\
199                      39 30 0:33 / /sys/fs/cgroup\\040root rw - cgroup2 cgroup2 rw\n";
200        assert_eq!(
201            parse_cgroup2_mount(Cursor::new(input)).unwrap(),
202            Some(PathBuf::from("/sys/fs/cgroup root"))
203        );
204    }
205
206    #[test]
207    fn decodes_utf8_mount_path() {
208        assert_eq!(
209            decode_mount_path("/sys/fs/cgroup-é\\040root").unwrap(),
210            PathBuf::from("/sys/fs/cgroup-é root")
211        );
212        assert!(decode_mount_path("/sys/fs/cgroup\\377root").is_err());
213    }
214
215    #[test]
216    fn ignores_non_root_cgroup2_mount() {
217        let input = b"39 30 0:33 /slice /sys/fs/cgroup rw - cgroup2 cgroup2 rw\n";
218        assert_eq!(parse_cgroup2_mount(Cursor::new(input)).unwrap(), None);
219    }
220
221    #[test]
222    fn parses_and_deduplicates_cgroup_procs() {
223        let input = b"10\n20\n10\n";
224        assert_eq!(
225            parse_cgroup_procs(Cursor::new(input)).unwrap(),
226            HashSet::from([10, 20])
227        );
228        assert!(parse_cgroup_procs(Cursor::new(b"not-a-pid\n")).is_err());
229    }
230}
scx_cosmos/cgroup.rs

scx_cosmos/
cgroup.rs
