cgroup: make it possible to move a control group to root

Signed-off-by: Levente Kurusa <lkurusa@acm.org>

src/cgroup.rs

@@ -17,12 +17,15 @@ use std::convert::From;
 /// > specialized behaviour.
 ///
 /// This crate is an attempt at providing a Rust-native way of managing these cgroups.
-pub struct Cgroup {
+pub struct Cgroup<'b> {
     /// The list of subsystems that control this cgroup
     subsystems: Vec<Subsystem>,
+
+    /// The hierarchy.
+    hier: &'b Hierarchy,
 }
 
-impl Cgroup {
+impl<'b> Cgroup<'b> {
 
     /// Create this control group.
     fn create(self: &Self) {
@@ -38,15 +41,29 @@ impl Cgroup {
     /// Note that if the handle goes out of scope and is dropped, the control group is _not_
     /// destroyed.
     pub fn new(hier: &Hierarchy, path: String) -> Cgroup {
+        let cg = Cgroup::load(hier, path);
+        cg.create();
+        cg
+    }
+
+    /// Create a handle for a 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.
+    ///
+    /// Note that if the handle goes out of scope and is dropped, the control group is _not_
+    /// destroyed.
+    pub fn load(hier: &Hierarchy, path: String) -> Cgroup {
         let mut subsystems = hier.subsystems();
-        subsystems = subsystems.into_iter().map(|x| x.enter(&path)).collect::<Vec<_>>();
+        if path != ""  {
+            subsystems = subsystems.into_iter().map(|x| x.enter(&path)).collect::<Vec<_>>();
+        }
 
         let cg = Cgroup {
-            //name: path,
             subsystems: subsystems,
+            hier: hier,
         };
 
-        cg.create();
         cg
     }
 
@@ -114,8 +131,29 @@ impl Cgroup {
         None
     }
 
+    /// 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: &Self, pid: CgroupPid) {
+        self.hier.root_control_group().add_task(pid);
+    }
+
     /// Attach a task to the control group.
     pub fn add_task(self: &Self, pid: CgroupPid) {
         self.subsystems().iter().for_each(|sub| sub.to_controller().add_task(&pid));
     }
+
+    /// Returns an Iterator that can be used to iterate over the tasks that are currently in the
+    /// control group.
+    pub fn tasks(self: &Self) -> Vec<CgroupPid> {
+        /* Collect the tasks from all subsystems */
+        let mut v = self.subsystems().iter()
+            .map(|x| x.to_controller().tasks())
+            .fold(vec![], |mut acc, mut x| { acc.append(&mut x); acc });
+        v.sort();
+        v.dedup();
+        v
+    }
+
 }

src/hierarchies.rs

@@ -23,6 +23,8 @@ use ::net_prio::NetPrioController;
 use ::hugetlb::HugeTlbController;
 use ::rdma::RdmaController;
 
+use ::cgroup::Cgroup;
+
 
 /// The standard, original cgroup implementation. Often referred to as "cgroupv1".
 pub struct V1 {
@@ -75,6 +77,10 @@ impl Hierarchy for V1 {
         subs
     }
 
+    fn root_control_group(self: &Self) -> Cgroup {
+        Cgroup::load(self, "".to_string())
+    }
+
     fn check_support(self: &Self, sub: Controllers) -> bool {
         let root = self.root().read_dir().unwrap();
         for entry in root {

src/lib.rs

@@ -1,6 +1,6 @@
 use std::path::PathBuf;
 use std::fs::File;
-use std::io::Write;
+use std::io::{BufRead, BufReader, Write};
 
 pub mod hierarchies;
 pub mod pid;
@@ -176,6 +176,21 @@ pub trait Controller {
             file.write_all(pid.pid.to_string().as_ref()).ok()
         });
     }
+
+    /// Get the list of tasks that this controller has.
+    fn tasks(self: &Self) -> Vec<CgroupPid> {
+        self.open_path("tasks", false).and_then(|file| {
+            let bf = BufReader::new(file);
+            let mut v = Vec::new();
+            for line in bf.lines() {
+                if let Ok(line) = line {
+                    let n = line.trim().parse().unwrap_or(0u64);
+                    v.push(n);
+                }
+            }
+            Some(v.into_iter().map(CgroupPid::from).collect())
+        }).unwrap_or(vec![])
+    }
 }
 
 #[doc(hidden)]
@@ -188,11 +203,17 @@ pub trait ControllIdentifier {
 pub trait Hierarchy {
     /// Returns what subsystems are supported by the hierarchy.
     fn subsystems(self: &Self) -> Vec<Subsystem>;
+
     /// Returns the root directory of the hierarchy.
     fn root(self: &Self) -> PathBuf;
+
+    /// Return a handle to the root control group in the hierarchy.
+    fn root_control_group(self: &Self) -> Cgroup;
+
     /// Checks whether a certain subsystem is supported in the hierarchy.
     ///
     /// This is an internal function and should not be used.
+    #[doc(hidden)]
     fn check_support(self: &Self, sub: Controllers) -> bool;
 }
 
@@ -391,6 +412,7 @@ pub struct Resources {
 
 /// A structure representing a `pid`. Currently implementations exist for `u64` and
 /// `std::process::Child`.
+#[derive(Debug, PartialEq, Eq, PartialOrd, Ord)]
 pub struct CgroupPid {
     /// The process identifier
     pub pid: u64,