From 07878325c355c535390f6ce4455134c4aa5c96b5 Mon Sep 17 00:00:00 2001
From: Danny Canter <danny@dcantah.dev>
Date: Sun, 8 Jan 2023 02:04:46 -0800
Subject: [PATCH] cgroup2: Add Kill method

Fixes: #98

Add a Kill method for cgroups v2. This uses the new functionality in
kernels 5.14+ where writing to a cgroup.kill file will send a SIGKILL
to every process running in the cgroup. This would be useful for
kata to avoid freezing+manually sending SIGKILL+thawing process it
does currently to emulate runc's behavior.

Signed-off-by: Danny Canter <danny@dcantah.dev>
---
 src/cgroup.rs   | 22 ++++++++++++++++++++++
 tests/cgroup.rs | 46 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 68 insertions(+)

diff --git a/src/cgroup.rs b/src/cgroup.rs
index a11d12e..60d9b24 100644
--- a/src/cgroup.rs
+++ b/src/cgroup.rs
@@ -318,6 +318,28 @@ impl Cgroup {
         self.hier.parent_control_group(&self.path)
     }
 
+    /// Kill every process in the control group. Only supported for v2 cgroups and on
+    /// kernels 5.14+. This will fail with InvalidOperation if the 'cgroup.kill' file does
+    /// not exist.
+    pub fn kill(&self) -> Result<()> {
+        if !self.v2() {
+            return Err(Error::new(CgroupVersion));
+        }
+
+        let val = "1";
+        let file_name = "cgroup.kill";
+        let p = self.hier.root().join(self.path.clone()).join(file_name);
+
+        // If cgroup.kill doesn't exist they're not on 5.14+ so lets
+        // surface some error the caller can check against.
+        if !p.exists() {
+            return Err(Error::new(InvalidOperation));
+        }
+
+        fs::write(p, val)
+            .map_err(|e| Error::with_cause(WriteFailed(file_name.to_string(), val.to_string()), e))
+    }
+
     /// Attach a task to the control group.
     pub fn add_task(&self, tid: CgroupPid) -> Result<()> {
         if self.v2() {
diff --git a/tests/cgroup.rs b/tests/cgroup.rs
index 7bef852..15d80fa 100644
--- a/tests/cgroup.rs
+++ b/tests/cgroup.rs
@@ -8,6 +8,9 @@
 use cgroups_rs::memory::MemController;
 use cgroups_rs::Controller;
 use cgroups_rs::{Cgroup, CgroupPid, Subsystem};
+use std::process::Command;
+use std::thread::sleep;
+use std::time::Duration;
 
 #[test]
 fn test_procs_iterator_cgroup() {
@@ -33,6 +36,49 @@ fn test_procs_iterator_cgroup() {
     cg.delete().unwrap();
 }
 
+#[test]
+fn test_kill_cgroup() {
+    if !cgroups_rs::hierarchies::is_cgroup2_unified_mode() {
+        return;
+    }
+    let h = cgroups_rs::hierarchies::auto();
+    let cg = Cgroup::new(h, String::from("test_kill_cgroup")).unwrap();
+    {
+        // Spawn a proc, don't want to getpid(2) here.
+        let mut child = Command::new("sleep").arg("infinity").spawn().unwrap();
+        cg.add_task_by_tgid(CgroupPid::from(child.id() as u64))
+            .unwrap();
+
+        let cg_procs = cg.procs();
+        assert_eq!(cg_procs.len(), 1_usize);
+
+        // Now kill and wait on the proc.
+        cg.kill().unwrap();
+
+        let mut tries = 0;
+        let status: Option<std::process::ExitStatus> = loop {
+            match child.try_wait() {
+                Ok(Some(status)) => {
+                    break Some(status);
+                }
+                Ok(None) => {
+                    if tries > 3 {
+                        break None;
+                    }
+                    sleep(Duration::from_millis(100));
+                    tries += 1;
+                }
+                Err(e) => {
+                    child.kill().unwrap();
+                    panic!("error attempting to wait: {}", e);
+                }
+            }
+        };
+        assert!(!status.is_none());
+    }
+    cg.delete().unwrap();
+}
+
 #[test]
 fn test_cgroup_with_relative_paths() {
     if cgroups_rs::hierarchies::is_cgroup2_unified_mode() {