From 6b338cf997961c40987c0776b77c618a8f5bca98 Mon Sep 17 00:00:00 2001
From: yaoyinnan <yaoyinnan@foxmail.com>
Date: Tue, 7 Feb 2023 22:40:06 +0800
Subject: [PATCH] Determine cgroup mode in add_task()

Determine the cgroup mode in add_task() to avoid the wrong operation of the caller writing threads to cgroup.threads in non-thread mode.

Fixes: #103

Signed-off-by: yaoyinnan <yaoyinnan@foxmail.com>
---
 src/cgroup.rs | 20 +++++++++++++++++++-
 src/error.rs  |  4 ++++
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/src/cgroup.rs b/src/cgroup.rs
index 60d9b24..988d9e1 100644
--- a/src/cgroup.rs
+++ b/src/cgroup.rs
@@ -16,6 +16,11 @@ use std::convert::From;
 use std::fs;
 use std::path::{Path, PathBuf};
 
+pub const CGROUP_MODE_DOMAIN: &str = "domain";
+pub const CGROUP_MODE_DOMAIN_THREADED: &str = "domain threaded";
+pub const CGROUP_MODE_DOMAIN_INVALID: &str = "domain invalid";
+pub const CGROUP_MODE_THREADED: &str = "threaded";
+
 /// A control group is the central structure to this crate.
 ///
 ///
@@ -346,7 +351,18 @@ impl Cgroup {
             let subsystems = self.subsystems();
             if !subsystems.is_empty() {
                 let c = subsystems[0].to_controller();
-                c.add_task(&tid)
+                let cgroup_type = self.get_cgroup_type()?;
+                // In cgroup v2, writing to the cgroup.threads file is only supported in thread mode.
+                if cgroup_type == *CGROUP_MODE_DOMAIN_THREADED
+                    || cgroup_type == *CGROUP_MODE_THREADED
+                {
+                    // It is used to move the threads of a process into a cgroup in thread mode.
+                    c.add_task(&tid)
+                } else {
+                    // When the cgroup type is domain or domain invalid,
+                    // cgroup.threads cannot be written.
+                    Err(Error::new(CgroupMode))
+                }
             } else {
                 Err(Error::new(SubsystemsEmpty))
             }
@@ -363,6 +379,8 @@ impl Cgroup {
             let subsystems = self.subsystems();
             if !subsystems.is_empty() {
                 let c = subsystems[0].to_controller();
+                // It is used to move a thread of the process to a cgroup,
+                // and other threads of the process will also move together.
                 c.add_task_by_tgid(&tgid)
             } else {
                 Err(Error::new(SubsystemsEmpty))
diff --git a/src/error.rs b/src/error.rs
index fd582e7..bc45a1a 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -59,6 +59,10 @@ pub enum ErrorKind {
     #[error("using method in wrong cgroup version")]
     CgroupVersion,
 
+    /// Using method in wrong cgroup mode.
+    #[error("using method in wrong cgroup mode.")]
+    CgroupMode,
+
     /// Subsystems is empty.
     #[error("subsystems is empty")]
     SubsystemsEmpty,