From 6b338cf997961c40987c0776b77c618a8f5bca98 Mon Sep 17 00:00:00 2001 From: yaoyinnan Date: Tue, 7 Feb 2023 22:40:06 +0800 Subject: [PATCH 1/2] 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 --- 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, From 3dd07353243ace2fc35417681a013f37ce9f47d8 Mon Sep 17 00:00:00 2001 From: yaoyinnan Date: Tue, 7 Feb 2023 22:40:33 +0800 Subject: [PATCH 2/2] Add UT for add_task(). Add UT for add_task() for cgroup v1 and v2. Fixes: #103 Signed-off-by: yaoyinnan --- tests/cgroup.rs | 179 +++++++++++++++++++++++++++++++++--------------- 1 file changed, 125 insertions(+), 54 deletions(-) diff --git a/tests/cgroup.rs b/tests/cgroup.rs index 15d80fa..3c9bd7d 100644 --- a/tests/cgroup.rs +++ b/tests/cgroup.rs @@ -5,6 +5,10 @@ // //! Simple unit tests about the control groups system. +use cgroups_rs::cgroup::{ + CGROUP_MODE_DOMAIN, CGROUP_MODE_DOMAIN_INVALID, CGROUP_MODE_DOMAIN_THREADED, + CGROUP_MODE_THREADED, +}; use cgroups_rs::memory::MemController; use cgroups_rs::Controller; use cgroups_rs::{Cgroup, CgroupPid, Subsystem}; @@ -36,6 +40,127 @@ fn test_procs_iterator_cgroup() { cg.delete().unwrap(); } +#[test] +fn test_tasks_iterator_cgroup_v1() { + if cgroups_rs::hierarchies::is_cgroup2_unified_mode() { + return; + } + let h = cgroups_rs::hierarchies::auto(); + let pid = libc::pid_t::from(nix::unistd::getpid()) as u64; + let cg = Cgroup::new(h, String::from("test_tasks_iterator_cgroup_v1")).unwrap(); + { + // Add a task to the control group. + cg.add_task(CgroupPid::from(pid)).unwrap(); + + let mut tasks = cg.tasks().into_iter(); + // Verify that the task is indeed in the xcontrol group + assert_eq!(tasks.next(), Some(CgroupPid::from(pid))); + assert_eq!(tasks.next(), None); + + // Now, try removing it. + cg.remove_task(CgroupPid::from(pid)).unwrap(); + tasks = cg.tasks().into_iter(); + + // Verify that it was indeed removed. + assert_eq!(tasks.next(), None); + } + cg.delete().unwrap(); +} + +#[test] +fn test_tasks_iterator_cgroup_threaded_mode() { + if !cgroups_rs::hierarchies::is_cgroup2_unified_mode() { + return; + } + let pid = libc::pid_t::from(nix::unistd::getpid()) as u64; + let cg = Cgroup::new( + cgroups_rs::hierarchies::auto(), + String::from("test_tasks_iterator_cgroup_threaded_mode"), + ) + .unwrap(); + let cg_threaded_sub1 = Cgroup::new_with_specified_controllers( + cgroups_rs::hierarchies::auto(), + String::from("test_tasks_iterator_cgroup_threaded_mode/threaded_sub1"), + Some(vec![String::from("cpuset"), String::from("cpu")]), + ) + .unwrap(); + let cg_threaded_sub2 = Cgroup::new_with_specified_controllers( + cgroups_rs::hierarchies::auto(), + String::from("test_tasks_iterator_cgroup_threaded_mode/threaded_sub2"), + Some(vec![String::from("cpuset"), String::from("cpu")]), + ) + .unwrap(); + { + // Verify that cgroup type of the control group is domain mode. + assert_eq!(cg.get_cgroup_type().unwrap(), CGROUP_MODE_DOMAIN); + + // Set cgroup type of the sub-control group is thread mode. + cg_threaded_sub1 + .set_cgroup_type(CGROUP_MODE_THREADED) + .unwrap(); + // Verify that cgroup type of the sub-control group is thread mode. + assert_eq!( + cg_threaded_sub1.get_cgroup_type().unwrap(), + CGROUP_MODE_THREADED + ); + // Verify that the cgroup type of the sub-control group that does + // not set the cgroup type is domain invalid mode. + assert_eq!( + cg_threaded_sub2.get_cgroup_type().unwrap(), + CGROUP_MODE_DOMAIN_INVALID + ); + // Verify whether the cgroup type of the parent control group of + // the control group whose cgroup type is set to thread mode is + // domain thread mode. + assert_eq!(cg.get_cgroup_type().unwrap(), CGROUP_MODE_DOMAIN_THREADED); + + // Set cgroup type of the sub-control group is thread mode. + cg_threaded_sub2 + .set_cgroup_type(CGROUP_MODE_THREADED) + .unwrap(); + // Verify that cgroup type of the sub-control group is thread mode. + assert_eq!( + cg_threaded_sub2.get_cgroup_type().unwrap(), + CGROUP_MODE_THREADED + ); + + // Add a proc to the control group. + cg.add_task_by_tgid(CgroupPid::from(pid)).unwrap(); + + let mut procs = cg.procs().into_iter(); + // Verify that the task is indeed in the x control group + assert_eq!(procs.next(), Some(CgroupPid::from(pid))); + assert_eq!(procs.next(), None); + + // Add a task to the sub control group. + cg_threaded_sub1.add_task(CgroupPid::from(pid)).unwrap(); + + let mut tasks = cg_threaded_sub1.tasks().into_iter(); + // Verify that the task is indeed in the xcontrol group + assert_eq!(tasks.next(), Some(CgroupPid::from(pid))); + assert_eq!(tasks.next(), None); + + // Now, try move it to parent. + cg_threaded_sub1 + .move_task_to_parent(CgroupPid::from(pid)) + .unwrap(); + tasks = cg_threaded_sub1.tasks().into_iter(); + + // Verify that it was indeed removed. + assert_eq!(tasks.next(), None); + + // Now, try removing it. + cg.remove_task_by_tgid(CgroupPid::from(pid)).unwrap(); + procs = cg.procs().into_iter(); + + // Verify that it was indeed removed. + assert_eq!(procs.next(), None); + } + cg_threaded_sub1.delete().unwrap(); + cg_threaded_sub2.delete().unwrap(); + cg.delete().unwrap(); +} + #[test] fn test_kill_cgroup() { if !cgroups_rs::hierarchies::is_cgroup2_unified_mode() { @@ -146,57 +271,3 @@ fn test_cgroup_v2() { cg.delete().unwrap(); } - -#[test] -fn test_tasks_iterator_cgroup_threaded_mode() { - if !cgroups_rs::hierarchies::is_cgroup2_unified_mode() { - return; - } - let h = cgroups_rs::hierarchies::auto(); - let pid = libc::pid_t::from(nix::unistd::getpid()) as u64; - let cg = Cgroup::new(h, String::from("test_tasks_iterator_cgroup_threaded_mode")).unwrap(); - let h = cgroups_rs::hierarchies::auto(); - let specified_controllers = vec![String::from("cpuset"), String::from("cpu")]; - let cg_threaded = Cgroup::new_with_specified_controllers( - h, - String::from("test_tasks_iterator_cgroup_threaded_mode/threaded"), - Some(specified_controllers), - ) - .unwrap(); - cg_threaded.set_cgroup_type("threaded").unwrap(); - { - // Add a task to the control group. - cg.add_task_by_tgid(CgroupPid::from(pid)).unwrap(); - - let mut procs = cg.procs().into_iter(); - // Verify that the task is indeed in the xcontrol group - assert_eq!(procs.next(), Some(CgroupPid::from(pid))); - assert_eq!(procs.next(), None); - - // Add a task to the sub control group. - cg_threaded.add_task(CgroupPid::from(pid)).unwrap(); - - let mut tasks = cg_threaded.tasks().into_iter(); - // Verify that the task is indeed in the xcontrol group - assert_eq!(tasks.next(), Some(CgroupPid::from(pid))); - assert_eq!(tasks.next(), None); - - // Now, try move it to parent. - cg_threaded - .move_task_to_parent(CgroupPid::from(pid)) - .unwrap(); - tasks = cg_threaded.tasks().into_iter(); - - // Verify that it was indeed removed. - assert_eq!(tasks.next(), None); - - // Now, try removing it. - cg.remove_task_by_tgid(CgroupPid::from(pid)).unwrap(); - procs = cg.procs().into_iter(); - - // Verify that it was indeed removed. - assert_eq!(procs.next(), None); - } - cg_threaded.delete().unwrap(); - cg.delete().unwrap(); -}