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: #90

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

src/cgroup.rs

@@ -59,43 +59,76 @@ impl Default for Cgroup {
 }
 
 impl Cgroup {
+    pub fn v2(&self) -> bool {
+        self.hier.v2()
+    }
+
     /// Create this control group.
-    fn create(&self) {
+    fn create(&self) -> Result<()> {
         if self.hier.v2() {
-            let _ret = create_v2_cgroup(self.hier.root(), &self.path);
+            create_v2_cgroup(self.hier.root(), &self.path, None)
         } else {
             for subsystem in &self.subsystems {
                 subsystem.to_controller().create();
             }
+            Ok(())
         }
     }
 
-    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>>,
+    ) -> Result<()> {
+        if self.hier.v2() {
+            create_v2_cgroup(self.hier.root(), &self.path, specified_controllers)
+        } else {
+            for subsystem in &self.subsystems {
+                subsystem.to_controller().create();
+            }
+            Ok(())
+        }
     }
 
     /// 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<P: AsRef<Path>>(hier: Box<dyn Hierarchy>, path: P) -> Cgroup {
+    pub fn new<P: AsRef<Path>>(hier: Box<dyn Hierarchy>, path: P) -> Result<Cgroup> {
         let cg = Cgroup::load(hier, path);
-        cg.create();
-        cg
+        cg.create()?;
+        Ok(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>>,
+    ) -> Result<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)?;
+        Ok(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,
         relative_paths: HashMap<String, String>,
-    ) -> Cgroup {
+    ) -> Result<Cgroup> {
         let cg = Cgroup::load_with_relative_paths(hier, path, relative_paths);
-        cg.create();
-        cg
+        cg.create()?;
+        Ok(cg)
     }
 
     /// Create a handle for a control group in the hierarchy `hier`, with name `path`.
@@ -119,6 +152,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,12 +292,37 @@ 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) -> Result<()> {
+        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
     /// hierarchy and any rules applied to that control group will _still_ apply to the task.
-    pub fn remove_task(&self, pid: CgroupPid) {
-        let _ = self.hier.root_control_group().add_task(pid);
+    pub fn remove_task(&self, pid: CgroupPid) -> Result<()> {
+        self.hier.root_control_group().add_task(pid)
+    }
+
+    /// Move to a proc to parent control group.
+    pub fn move_proc_to_parent(&self, pid: CgroupPid) -> Result<()> {
+        self.hier
+            .parent_control_group(&self.path)
+            .add_task_by_tgid(pid)
+    }
+
+    /// Move to a task to parent control group.
+    pub fn move_task_to_parent(&self, pid: CgroupPid) -> Result<()> {
+        self.hier.parent_control_group(&self.path).add_task(pid)
+    }
+
+    /// Return a handle to the parent control group in the hierarchy.
+    pub fn parent_control_group(&self) -> Cgroup {
+        self.hier.parent_control_group(&self.path)
     }
 
     /// Attach a task to the control group.
@@ -249,7 +333,7 @@ impl Cgroup {
                 let c = subsystems[0].to_controller();
                 c.add_task(&pid)
             } else {
-                Ok(())
+                Err(Error::new(SubsystemsEmpty))
             }
         } else {
             self.subsystems()
@@ -260,9 +344,50 @@ impl Cgroup {
 
     /// Attach a task to the control group by thread group id.
     pub fn add_task_by_tgid(&self, pid: CgroupPid) -> Result<()> {
-        self.subsystems()
-            .iter()
-            .try_for_each(|sub| sub.to_controller().add_task_by_tgid(&pid))
+        if self.v2() {
+            let subsystems = self.subsystems();
+            if !subsystems.is_empty() {
+                let c = subsystems[0].to_controller();
+                c.add_task_by_tgid(&pid)
+            } else {
+                Err(Error::new(SubsystemsEmpty))
+            }
+        } else {
+            self.subsystems()
+                .iter()
+                .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.
@@ -281,6 +406,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 +471,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 +480,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 +523,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

@@ -57,11 +57,11 @@
 //!              .read(6, 1, 10)
 //!              .write(11, 1, 100)
 //!          .done()
-//!      .build(h);
+//!      .build(h).unwrap();
 //! ```
 
 use crate::{
-    BlkIoDeviceResource, BlkIoDeviceThrottleResource, Cgroup, DeviceResource, Hierarchy,
+    BlkIoDeviceResource, BlkIoDeviceThrottleResource, Cgroup, DeviceResource, Error, Hierarchy,
     HugePageResource, MaxValue, NetworkPriority, Resources,
 };
 
@@ -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,
         }
     }
 
@@ -134,10 +136,21 @@ 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
+    pub fn build(self, hier: Box<dyn Hierarchy>) -> Result<Cgroup, Error> {
+        if let Some(controllers) = self.specified_controllers {
+            let cg = Cgroup::new_with_specified_controllers(hier, self.name, Some(controllers))?;
+            cg.apply(&self.resources)?;
+            Ok(cg)
+        } else {
+            let cg = Cgroup::new(hier, self.name)?;
+            cg.apply(&self.resources)?;
+            Ok(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,