cgroup: support to set threaded mode in cgroup v2

Support to set threaded mode in cgroup v2. The premise of switching to threaded mode is that only the cgroup of cpuset, cpu and pids is supported.

Fixes: kata-containers#90

Signed-off-by: yaoyinnan <yaoyinnan@foxmail.com>

src/cgroup.rs

@@ -59,19 +59,30 @@ impl Default for Cgroup {
 }
 
 impl Cgroup {
+    pub fn v2(&self) -> bool {
+        self.hier.v2()
+    }
+
     /// Create this control group.
     fn create(&self) {
         if self.hier.v2() {
-            let _ret = create_v2_cgroup(self.hier.root(), &self.path);
+            let _ret = create_v2_cgroup(self.hier.root(), &self.path, None);
         } else {
             for subsystem in &self.subsystems {
                 subsystem.to_controller().create();
             }
         }
     }
 
-    pub fn v2(&self) -> bool {
-        self.hier.v2()
+    /// Create this control group with specified controllers.
+    fn create_with_specified_controllers(&self, specified_controllers: Option<Vec<String>>) {
+        if self.hier.v2() {
+            let _ret = create_v2_cgroup(self.hier.root(), &self.path, specified_controllers);
+        } else {
+            for subsystem in &self.subsystems {
+                subsystem.to_controller().create();
+            }
+        }
     }
 
     /// Create a new control group in the hierarchy `hier`, with name `path`.
@@ -83,11 +94,28 @@ impl Cgroup {
         cg
     }
 
+    /// Create a new control group in the hierarchy `hier`, with name `path`.
+    ///
+    /// Returns a handle to the control group that can be used to manipulate it.
+    pub fn new_with_specified_controllers<P: AsRef<Path>>(
+        hier: Box<dyn Hierarchy>,
+        path: P,
+        specified_controllers: Option<Vec<String>>,
+    ) -> Cgroup {
+        let cg = if let Some(sc) = specified_controllers.clone() {
+            Cgroup::load_with_specified_controllers(hier, path, sc)
+        } else {
+            Cgroup::load(hier, path)
+        };
+        cg.create_with_specified_controllers(specified_controllers);
+        cg
+    }
+
     /// Create a new control group in the hierarchy `hier`, with name `path` and `relative_paths`
     ///
     /// Returns a handle to the control group that can be used to manipulate it.
     ///
-    /// Note that this method is only meaningful for cgroup v1, call it is equivalent to call `new` in the v2 mode
+    /// Note that this method is only meaningful for cgroup v1, call it is equivalent to call `new` in the v2 mode.
     pub fn new_with_relative_paths<P: AsRef<Path>>(
         hier: Box<dyn Hierarchy>,
         path: P,
@@ -119,6 +147,32 @@ impl Cgroup {
         }
     }
 
+    /// Create a handle for a specified control group in the hierarchy `hier`, with name `path`.
+    ///
+    /// Returns a handle to the control group (that possibly does not exist until `create()` has
+    /// been called on the cgroup.
+    pub fn load_with_specified_controllers<P: AsRef<Path>>(
+        hier: Box<dyn Hierarchy>,
+        path: P,
+        specified_controllers: Vec<String>,
+    ) -> Cgroup {
+        let path = path.as_ref();
+        let mut subsystems = hier.subsystems();
+        if path.as_os_str() != "" {
+            subsystems = subsystems
+                .into_iter()
+                .filter(|x| specified_controllers.contains(&x.controller_name()))
+                .map(|x| x.enter(path))
+                .collect::<Vec<_>>();
+        }
+
+        Cgroup {
+            path: path.to_str().unwrap().to_string(),
+            subsystems,
+            hier,
+        }
+    }
+
     /// Create a handle for a control group in the hierarchy `hier`, with name `path` and `relative_paths`
     ///
     /// Returns a handle to the control group (that possibly does not exist until `create()` has
@@ -233,6 +287,14 @@ impl Cgroup {
         None
     }
 
+    /// Removes a proc from the control group.
+    ///
+    /// Note that this means that the task will be moved back to the root control group in the
+    /// hierarchy and any rules applied to that control group will _still_ apply to the proc.
+    pub fn remove_proc(&self, pid: CgroupPid) {
+        let _ = self.hier.root_control_group().add_task_by_tgid(pid);
+    }
+
     /// Removes a task from the control group.
     ///
     /// Note that this means that the task will be moved back to the root control group in the
@@ -241,6 +303,16 @@ impl Cgroup {
         let _ = self.hier.root_control_group().add_task(pid);
     }
 
+    /// Move to a proc to parent control group.
+    pub fn move_to_parent_proc(&self, pid: CgroupPid) {
+        let _ = self.hier.parent_control_group(&self.path).add_task_by_tgid(pid);
+    }
+
+    /// Move to a task to parent control group.
+    pub fn move_to_parent_task(&self, pid: CgroupPid) {
+        let _ = self.hier.parent_control_group(&self.path).add_task(pid);
+    }
+
     /// Attach a task to the control group.
     pub fn add_task(&self, pid: CgroupPid) -> Result<()> {
         if self.v2() {
@@ -249,7 +321,7 @@ impl Cgroup {
                 let c = subsystems[0].to_controller();
                 c.add_task(&pid)
             } else {
-                Ok(())
+                Err(Error::new(SubsystemsEmpty))
             }
         } else {
             self.subsystems()
@@ -265,6 +337,37 @@ impl Cgroup {
             .try_for_each(|sub| sub.to_controller().add_task_by_tgid(&pid))
     }
 
+    /// set cgroup.type
+    pub fn set_cgroup_type(&self, cgroup_type: &str) -> Result<()> {
+        if self.v2() {
+            let subsystems = self.subsystems();
+            if !subsystems.is_empty() {
+                let c = subsystems[0].to_controller();
+                c.set_cgroup_type(cgroup_type)
+            } else {
+                Err(Error::new(SubsystemsEmpty))
+            }
+        } else {
+            Err(Error::new(CgroupVersion))
+        }
+    }
+
+    /// get cgroup.type
+    pub fn get_cgroup_type(&self) -> Result<String> {
+        if self.v2() {
+            let subsystems = self.subsystems();
+            if !subsystems.is_empty() {
+                let c = subsystems[0].to_controller();
+                let cgroup_type = c.get_cgroup_type()?;
+                Ok(cgroup_type)
+            } else {
+                Err(Error::new(SubsystemsEmpty))
+            }
+        } else {
+            Err(Error::new(CgroupVersion))
+        }
+    }
+
     /// Set notify_on_release to the control group.
     pub fn set_notify_on_release(&self, enable: bool) -> Result<()> {
         self.subsystems()
@@ -281,6 +384,33 @@ impl Cgroup {
             .try_for_each(|sub| sub.to_controller().set_release_agent(path))
     }
 
+    /// Returns an Iterator that can be used to iterate over the procs that are currently in the
+    /// control group.
+    pub fn procs(&self) -> Vec<CgroupPid> {
+        // Collect the procs from all subsystems
+        let mut v = if self.v2() {
+            let subsystems = self.subsystems();
+            if !subsystems.is_empty() {
+                let c = subsystems[0].to_controller();
+                c.procs()
+            } else {
+                vec![]
+            }
+        } else {
+            self.subsystems()
+                .iter()
+                .map(|x| x.to_controller().procs())
+                .fold(vec![], |mut acc, mut x| {
+                    acc.append(&mut x);
+                    acc
+                })
+        };
+
+        v.sort();
+        v.dedup();
+        v
+    }
+
     /// Returns an Iterator that can be used to iterate over the tasks that are currently in the
     /// control group.
     pub fn tasks(&self) -> Vec<CgroupPid> {
@@ -319,7 +449,7 @@ fn enable_controllers(controllers: &[String], path: &Path) {
     }
 }
 
-fn supported_controllers() -> Vec<String> {
+fn get_supported_controllers() -> Vec<String> {
     let p = format!("{}/{}", UNIFIED_MOUNTPOINT, "cgroup.controllers");
     let ret = fs::read_to_string(p.as_str());
     ret.unwrap_or_default()
@@ -328,9 +458,22 @@ fn supported_controllers() -> Vec<String> {
         .collect::<Vec<String>>()
 }
 
-fn create_v2_cgroup(root: PathBuf, path: &str) -> Result<()> {
+fn create_v2_cgroup(
+    root: PathBuf,
+    path: &str,
+    specified_controllers: Option<Vec<String>>,
+) -> Result<()> {
     // controler list ["memory", "cpu"]
-    let controllers = supported_controllers();
+    let controllers = if let Some(s_controllers) = specified_controllers {
+        if verify_supported_controllers(s_controllers.as_ref()) {
+            s_controllers
+        } else {
+            return Err(Error::new(ErrorKind::SpecifiedControllers));
+        }
+    } else {
+        get_supported_controllers()
+    };
+
     let mut fp = root;
 
     // enable for root
@@ -358,6 +501,16 @@ fn create_v2_cgroup(root: PathBuf, path: &str) -> Result<()> {
     Ok(())
 }
 
+pub fn verify_supported_controllers(controllers: &[String]) -> bool {
+    let supported_controllers = get_supported_controllers();
+    for controller in controllers.iter() {
+        if !supported_controllers.contains(controller) {
+            return false;
+        }
+    }
+    true
+}
+
 pub fn get_cgroups_relative_paths() -> Result<HashMap<String, String>> {
     let path = "/proc/self/cgroup".to_string();
     get_cgroups_relative_paths_by_path(path)

src/cgroup_builder.rs

@@ -80,6 +80,7 @@ pub struct CgroupBuilder {
     name: String,
     /// Internal, unsupported field: use the associated builders instead.
     resources: Resources,
+    specified_controllers: Option<Vec<String>>,
 }
 
 impl CgroupBuilder {
@@ -90,6 +91,7 @@ impl CgroupBuilder {
         CgroupBuilder {
             name: name.to_owned(),
             resources: Resources::default(),
+            specified_controllers: None,
         }
     }
 
@@ -135,9 +137,20 @@ impl CgroupBuilder {
 
     /// Finalize the control group, consuming the builder and creating the control group.
     pub fn build(self, hier: Box<dyn Hierarchy>) -> Cgroup {
-        let cg = Cgroup::new(hier, self.name);
-        let _ret = cg.apply(&self.resources);
-        cg
+        if let Some(controllers) = self.specified_controllers {
+            let cg = Cgroup::new_with_specified_controllers(hier, self.name, Some(controllers));
+            let _ret = cg.apply(&self.resources);
+            cg
+        } else {
+            let cg = Cgroup::new(hier, self.name);
+            let _ret = cg.apply(&self.resources);
+            cg
+        }
+    }
+
+    pub fn set_specified_controllers(mut self, specified_controllers: Vec<String>) -> Self {
+        self.specified_controllers = Some(specified_controllers);
+        self
     }
 }
 

src/error.rs

@@ -51,6 +51,18 @@ pub enum ErrorKind {
     #[error("invalid bytes size")]
     InvalidBytesSize,
 
+    /// The specified controller is not in the list of supported controllers.
+    #[error("specified controller is not in the list of supported controllers")]
+    SpecifiedControllers,
+
+    /// Using method in wrong cgroup version.
+    #[error("using method in wrong cgroup version")]
+    CgroupVersion,
+
+    /// Subsystems is empty.
+    #[error("subsystems is empty")]
+    SubsystemsEmpty,
+
     /// An unknown error has occured.
     #[error("an unknown error")]
     Other,

src/hierarchies.rs

@@ -12,7 +12,7 @@
 use std::fs;
 use std::fs::File;
 use std::io::{BufRead, BufReader};
-use std::path::PathBuf;
+use std::path::{Path, PathBuf};
 
 use crate::blkio::BlkIoController;
 use crate::cpu::CpuController;
@@ -173,6 +173,12 @@ impl Hierarchy for V1 {
         Cgroup::load(auto(), "")
     }
 
+    fn parent_control_group(&self, path: &str) -> Cgroup {
+        let path = Path::new(path);
+        let parent_path = path.parent().unwrap().to_string_lossy().to_string();
+        Cgroup::load(auto(), &parent_path)
+    }
+
     fn root(&self) -> PathBuf {
         self.mountinfo
             .iter()
@@ -249,6 +255,12 @@ impl Hierarchy for V2 {
         Cgroup::load(auto(), "")
     }
 
+    fn parent_control_group(&self, path: &str) -> Cgroup {
+        let path = Path::new(path);
+        let parent_path = path.parent().unwrap().to_string_lossy().to_string();
+        Cgroup::load(auto(), &parent_path)
+    }
+
     fn root(&self) -> PathBuf {
         PathBuf::from(self.root.clone())
     }