diff --git a/src/cgroup.rs b/src/cgroup.rs index 6255593..3b4a379 100644 --- a/src/cgroup.rs +++ b/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, + + /// 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::>(); + if path != "" { + subsystems = subsystems.into_iter().map(|x| x.enter(&path)).collect::>(); + } 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 { + /* 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 + } + } diff --git a/src/hierarchies.rs b/src/hierarchies.rs index 832443b..500ae72 100644 --- a/src/hierarchies.rs +++ b/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 { diff --git a/src/lib.rs b/src/lib.rs index 5c27b80..3a4e5fe 100644 --- a/src/lib.rs +++ b/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 { + 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; + /// 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,