cgroup.rs

  1//! cgroups v2 session management for Linux process tracking.
  2//!
  3//! Each terminal session gets its own cgroup, providing an inescapable
  4//! mechanism for tracking and killing all descendant processes.
  5
  6use std::fs;
  7use std::io;
  8use std::path::{Path, PathBuf};
  9
 10/// A cgroup v2 session for tracking processes spawned by a terminal.
 11///
 12/// All processes in the cgroup can be killed atomically, regardless of
 13/// `setsid()`, `setpgid()`, or reparenting.
 14pub struct CgroupSession {
 15    cgroup_path: PathBuf,
 16    owns_cgroup: bool,
 17}
 18
 19impl CgroupSession {
 20    /// Create a new cgroup under the user's systemd slice.
 21    ///
 22    /// The cgroup is created at:
 23    /// `/sys/fs/cgroup/user.slice/user-<uid>.slice/user@<uid>.service/zed-terminal-<uuid>.scope`
 24    pub fn new() -> io::Result<Self> {
 25        let uid = unsafe { libc::getuid() };
 26        let uuid = uuid::Uuid::new_v4();
 27        let scope_name = format!("zed-terminal-{uuid}.scope");
 28
 29        // Try the systemd user slice first
 30        let user_slice = format!("/sys/fs/cgroup/user.slice/user-{uid}.slice/user@{uid}.service");
 31
 32        let cgroup_path = if Path::new(&user_slice).exists() {
 33            let path = PathBuf::from(&user_slice).join(&scope_name);
 34            fs::create_dir(&path)?;
 35            path
 36        } else {
 37            // Fallback: try directly under the current process's cgroup
 38            let self_cgroup = read_self_cgroup()?;
 39            let parent = PathBuf::from("/sys/fs/cgroup").join(self_cgroup.trim_start_matches('/'));
 40            if !parent.exists() {
 41                return Err(io::Error::other(
 42                    "cgroups v2 not available: cannot find parent cgroup",
 43                ));
 44            }
 45            let path = parent.join(&scope_name);
 46            fs::create_dir(&path)?;
 47            path
 48        };
 49
 50        Ok(Self {
 51            cgroup_path,
 52            owns_cgroup: true,
 53        })
 54    }
 55
 56    /// Reconstruct a CgroupSession from a path string (used by the child process).
 57    /// Does NOT create the cgroup — assumes the parent already created it.
 58    pub fn from_path(path: &str) -> Self {
 59        Self {
 60            cgroup_path: PathBuf::from(path),
 61            owns_cgroup: false,
 62        }
 63    }
 64
 65    /// Returns the cgroup filesystem path.
 66    pub fn path(&self) -> &Path {
 67        &self.cgroup_path
 68    }
 69
 70    /// Returns the cgroup path as a string for serialization.
 71    pub fn path_string(&self) -> String {
 72        self.cgroup_path.to_string_lossy().into_owned()
 73    }
 74
 75    /// Move a process into this cgroup by writing its PID to cgroup.procs.
 76    pub fn add_process(&self, pid: libc::pid_t) -> io::Result<()> {
 77        let procs_path = self.cgroup_path.join("cgroup.procs");
 78        fs::write(&procs_path, pid.to_string().as_bytes())
 79    }
 80
 81    /// Kill all processes in the cgroup.
 82    ///
 83    /// Tries the atomic `cgroup.kill` interface first (kernel 5.14+),
 84    /// falling back to reading cgroup.procs and killing each PID.
 85    pub fn kill_all(&self) -> io::Result<()> {
 86        // Step 1: Freeze the cgroup to prevent new forks
 87        let freeze_path = self.cgroup_path.join("cgroup.freeze");
 88        if freeze_path.exists() {
 89            if let Err(err) = fs::write(&freeze_path, b"1") {
 90                log::debug!("Failed to freeze cgroup: {err}");
 91            }
 92        }
 93
 94        // Step 2: Try atomic kill via cgroup.kill (kernel 5.14+)
 95        let kill_path = self.cgroup_path.join("cgroup.kill");
 96        if kill_path.exists() {
 97            if fs::write(&kill_path, b"1").is_ok() {
 98                return Ok(());
 99            }
100        }
101
102        // Step 3: Fallback — read cgroup.procs and kill each PID
103        let procs_path = self.cgroup_path.join("cgroup.procs");
104        loop {
105            let content = fs::read_to_string(&procs_path)?;
106            let pids: Vec<libc::pid_t> = content
107                .lines()
108                .filter_map(|line| line.trim().parse().ok())
109                .collect();
110
111            if pids.is_empty() {
112                break;
113            }
114
115            for pid in &pids {
116                unsafe {
117                    libc::kill(*pid, libc::SIGKILL);
118                }
119            }
120
121            // Brief sleep to let processes die before re-scanning
122            std::thread::sleep(std::time::Duration::from_millis(10));
123        }
124
125        Ok(())
126    }
127
128    /// Remove the cgroup directory. Must be called after all processes are dead.
129    pub fn cleanup(&self) {
130        if let Err(err) = fs::remove_dir(&self.cgroup_path) {
131            log::warn!(
132                "Failed to remove cgroup directory {:?}: {err}",
133                self.cgroup_path
134            );
135        }
136    }
137
138    /// Kill all processes and clean up the cgroup.
139    pub fn kill_all_and_cleanup(&self) {
140        if let Err(err) = self.kill_all() {
141            log::warn!("Failed to kill cgroup processes: {err}");
142        }
143        self.cleanup();
144    }
145}
146
147impl Drop for CgroupSession {
148    fn drop(&mut self) {
149        if self.owns_cgroup {
150            self.kill_all_and_cleanup();
151        }
152    }
153}
154
155/// Read the current process's cgroup path from /proc/self/cgroup.
156/// For cgroups v2, the format is "0::/path".
157fn read_self_cgroup() -> io::Result<String> {
158    let content = fs::read_to_string("/proc/self/cgroup")?;
159    for line in content.lines() {
160        if let Some(path) = line.strip_prefix("0::") {
161            return Ok(path.to_string());
162        }
163    }
164    Err(io::Error::other(
165        "cgroups v2 not found in /proc/self/cgroup",
166    ))
167}