From c623dc3fbab8fa0b5a4f6b844debcfddf93ef9a2 Mon Sep 17 00:00:00 2001 From: bin liu Date: Wed, 2 Sep 2020 22:03:11 +0800 Subject: [PATCH] add basic v2 cpu/memory functions Signed-off-by: bin liu --- src/cgroup.rs | 62 +++++++++++++++++++++++++++++++++----- src/cpu.rs | 47 ++++++++++++++++++++++++++++- src/hierarchies.rs | 75 ++++++++++++++++++++++++++++++---------------- src/lib.rs | 41 ++++++++++++++++++------- src/memory.rs | 57 ++++++++++++++++++++++++++++++----- src/pid.rs | 4 +-- tests/cgroup.rs | 25 +++++++++++++++- 7 files changed, 256 insertions(+), 55 deletions(-) diff --git a/src/cgroup.rs b/src/cgroup.rs index dd424ef..0cf1ba4 100644 --- a/src/cgroup.rs +++ b/src/cgroup.rs @@ -4,6 +4,7 @@ use crate::error::*; use crate::{CgroupPid, ControllIdentifier, Controller, Hierarchy, Resources, Subsystem}; +use std::collections::HashMap; use std::convert::From; use std::fs; use std::path::{Path, PathBuf}; @@ -41,6 +42,10 @@ impl<'b> Cgroup<'b> { } } + pub fn v2(&self) -> bool { + self.hier.v2() + } + /// 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. @@ -79,6 +84,42 @@ impl<'b> Cgroup<'b> { cg } + pub fn new_with_prefix>(hier: Box<&'b dyn Hierarchy>, path: P, prefixes: HashMap) -> Cgroup<'b> { + let cg = Cgroup::load_with_prefix(hier, path, prefixes); + cg.create(); + cg + } + + pub fn load_with_prefix>(hier: Box<&'b dyn Hierarchy>, path: P, prefixes: HashMap) -> Cgroup<'b> { + let path = path.as_ref(); + let mut subsystems = hier.subsystems(); + if path.as_os_str() != "" { + subsystems = subsystems + .into_iter() + .map(|x| { + let cn = x.controller_name(); + if prefixes.contains_key(&cn) { + let prefix = prefixes.get(&cn).unwrap(); + let valid_path = prefix.trim_start_matches("/").to_string(); + let mut p = PathBuf::from(valid_path); + p.push(path); + x.enter(p.as_ref()) + }else { + x.enter(path) + } + }) + .collect::>(); + } + + let cg = Cgroup { + subsystems: subsystems, + hier: hier, + path: path.to_str().unwrap().to_string(), + }; + + cg + } + /// The list of subsystems that this control group supports. pub fn subsystems(&self) -> &Vec { &self.subsystems @@ -175,6 +216,16 @@ impl<'b> Cgroup<'b> { pub const UNIFIED_MOUNTPOINT: &'static str = "/sys/fs/cgroup"; +fn enable_controllers(controllers: &Vec, path: &PathBuf) { + let mut f = path.clone(); + f.push("cgroup.subtree_control"); + for c in controllers{ + let body = format!("+{}", c); + // FIXME set mode to 0644 + let _rest = fs::write(f.as_path(), body.as_bytes()); + } +} + fn supported_controllers(p: &PathBuf) -> Vec{ let p = format!("{}/{}", UNIFIED_MOUNTPOINT, "cgroup.controllers"); let ret = fs::read_to_string(p.as_str()); @@ -186,6 +237,9 @@ fn create_v2_cgroup(root: PathBuf, path: &str) -> Result<()> { let controllers = supported_controllers(&root); let mut fp = root; + // enable for root + enable_controllers(&controllers, &fp); + // path: "a/b/c" let elements = path.split("/").collect::>(); let last_index = elements.len() - 1 ; @@ -203,13 +257,7 @@ fn create_v2_cgroup(root: PathBuf, path: &str) -> Result<()> { if i < last_index { // enable controllers for substree - let mut f = fp.clone(); - f.push("cgroup.subtree_control"); - for c in &controllers{ - let body = format!("+{}", c); - // FIXME set mode to 0644 - let _rest = fs::write(f.as_path(), body.as_bytes()); - } + enable_controllers(&controllers, &fp); } } diff --git a/src/cpu.rs b/src/cpu.rs index ae48c1c..d816a58 100644 --- a/src/cpu.rs +++ b/src/cpu.rs @@ -149,7 +149,12 @@ impl CpuController { /// `shares` to `200` ensures that control group `B` receives twice as much as CPU bandwidth. /// (Assuming both `A` and `B` are of the same parent) pub fn set_shares(&self, shares: u64) -> Result<()> { - self.open_path("cpu.shares", true).and_then(|mut file| { + let mut file = "cpu.shares"; + if self.v2 { + file = "cpu.weight"; + } + // NOTE: .CpuShares is not used here. Conversion is the caller's responsibility. + self.open_path(file, true).and_then(|mut file| { file.write_all(shares.to_string().as_ref()) .map_err(|e| Error::with_cause(WriteFailed, e)) }) @@ -194,4 +199,44 @@ impl CpuController { self.open_path("cpu.cfs_quota_us", false) .and_then(read_u64_from) } + + pub fn set_cfs_quota_and_period(&self, quota: u64, period: u64) -> Result<()> { + if !self.v2 { + self.set_cfs_quota(quota)?; + return self.set_cfs_period(period); + } + let mut line = "max".to_string(); + if quota > 0 { + line = quota.to_string(); + } + + let mut p = period; + if period == 0 { + // This default value is documented in + // https://www.kernel.org/doc/html/latest/admin-guide/cgroup-v2.html + p = 100000 + } + line = format!("{} {}", line, p); + self.open_path("cpu.max", true) + .and_then(|mut file| { + file.write_all(line.as_ref()) + .map_err(|e| Error::with_cause(WriteFailed, e)) + }) + } + + pub fn set_rt_runtime(&self, us: i64) -> Result<()> { + self.open_path("cpu.rt_runtime_us", true) + .and_then(|mut file| { + file.write_all(us.to_string().as_ref()) + .map_err(|e| Error::with_cause(WriteFailed, e)) + }) + } + + pub fn set_rt_period_us(&self, us: u64) -> Result<()> { + self.open_path("cpu.rt_period_us", true) + .and_then(|mut file| { + file.write_all(us.to_string().as_ref()) + .map_err(|e| Error::with_cause(WriteFailed, e)) + }) + } } diff --git a/src/hierarchies.rs b/src/hierarchies.rs index f99fd3f..4853b7c 100644 --- a/src/hierarchies.rs +++ b/src/hierarchies.rs @@ -70,7 +70,7 @@ impl Hierarchy for V1 { subs.push(Subsystem::NetCls(NetClsController::new(self.root()))); } if self.check_support(Controllers::BlkIo) { - subs.push(Subsystem::BlkIo(BlkIoController::new(self.root(), true))); + subs.push(Subsystem::BlkIo(BlkIoController::new(self.root(), false))); } if self.check_support(Controllers::PerfEvent) { subs.push(Subsystem::PerfEvent(PerfEventController::new(self.root()))); @@ -125,7 +125,6 @@ impl Hierarchy for V2 { } let controllers = ret.unwrap().trim().to_string(); - println!("controllers: {:?}", controllers); let controller_list: Vec<&str> = controllers.split(' ').collect(); @@ -140,30 +139,30 @@ impl Hierarchy for V2 { } } - if self.check_support(Controllers::CpuAcct) { - subs.push(Subsystem::CpuAcct(CpuAcctController::new(self.root()))); - } - if self.check_support(Controllers::Devices) { - subs.push(Subsystem::Devices(DevicesController::new(self.root()))); - } - if self.check_support(Controllers::Freezer) { - subs.push(Subsystem::Freezer(FreezerController::new(self.root()))); - } - if self.check_support(Controllers::NetCls) { - subs.push(Subsystem::NetCls(NetClsController::new(self.root()))); - } - if self.check_support(Controllers::PerfEvent) { - subs.push(Subsystem::PerfEvent(PerfEventController::new(self.root()))); - } - if self.check_support(Controllers::NetPrio) { - subs.push(Subsystem::NetPrio(NetPrioController::new(self.root()))); - } - if self.check_support(Controllers::HugeTlb) { - subs.push(Subsystem::HugeTlb(HugeTlbController::new(self.root(), true))); - } - if self.check_support(Controllers::Rdma) { - subs.push(Subsystem::Rdma(RdmaController::new(self.root()))); - } + // if self.check_support(Controllers::CpuAcct) { + // subs.push(Subsystem::CpuAcct(CpuAcctController::new(self.root()))); + // } + // if self.check_support(Controllers::Devices) { + // subs.push(Subsystem::Devices(DevicesController::new(self.root()))); + // } + // if self.check_support(Controllers::Freezer) { + // subs.push(Subsystem::Freezer(FreezerController::new(self.root()))); + // } + // if self.check_support(Controllers::NetCls) { + // subs.push(Subsystem::NetCls(NetClsController::new(self.root()))); + // } + // if self.check_support(Controllers::PerfEvent) { + // subs.push(Subsystem::PerfEvent(PerfEventController::new(self.root()))); + // } + // if self.check_support(Controllers::NetPrio) { + // subs.push(Subsystem::NetPrio(NetPrioController::new(self.root()))); + // } + // if self.check_support(Controllers::HugeTlb) { + // subs.push(Subsystem::HugeTlb(HugeTlbController::new(self.root(), true))); + // } + // if self.check_support(Controllers::Rdma) { + // subs.push(Subsystem::Rdma(RdmaController::new(self.root()))); + // } subs } @@ -205,6 +204,7 @@ impl V2 { pub const UNIFIED_MOUNTPOINT: &'static str = "/sys/fs/cgroup"; +#[cfg(all(target_os = "linux", not(target_env = "musl")))] pub fn is_cgroup2_unified_mode() -> bool { let path = Path::new(UNIFIED_MOUNTPOINT); let fs_stat = statfs::statfs(path); @@ -212,9 +212,32 @@ pub fn is_cgroup2_unified_mode() -> bool { return false } + // FIXME notwork, nix will not compile CGROUP2_SUPER_MAGIC because not(target_env = "musl") fs_stat.unwrap().filesystem_type() == statfs::CGROUP2_SUPER_MAGIC } +pub const INIT_CGROUP_PATHS: &'static str = "/proc/1/cgroup"; + +#[cfg(all(target_os = "linux", target_env = "musl"))] +pub fn is_cgroup2_unified_mode() -> bool { + let lines = fs::read_to_string(INIT_CGROUP_PATHS); + if lines.is_err() { + return false + } + + for line in lines.unwrap().lines(){ + let fields: Vec<&str> = line.split(':').collect(); + if fields.len() != 3 { + continue; + } + if fields[0] != "0" { + return false; + } + } + + true +} + pub fn auto() -> Box { if is_cgroup2_unified_mode() { Box::new(V2::new()) diff --git a/src/lib.rs b/src/lib.rs index 419263b..00e8e42 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -231,7 +231,7 @@ impl Controller for T where T: ControllerInternal { /// Create this controller fn create(&self) { - self.verify_path().expect("path should be valid"); + self.verify_path().expect(format!("path should be valid: {:?}", self.path()).as_str()); match ::std::fs::create_dir_all(self.get_path()) { Ok(_) => self.post_create(), @@ -253,7 +253,11 @@ impl Controller for T where T: ControllerInternal { /// Attach a task to this controller. fn add_task(&self, pid: &CgroupPid) -> Result<()> { - self.open_path("tasks", true).and_then(|mut file| { + let mut file = "tasks"; + if self.is_v2() { + file = "cgroup.procs"; + } + self.open_path(file, true).and_then(|mut file| { file.write_all(pid.pid.to_string().as_ref()) .map_err(|e| Error::with_cause(ErrorKind::WriteFailed, e)) }) @@ -261,7 +265,11 @@ impl Controller for T where T: ControllerInternal { /// Get the list of tasks that this controller has. fn tasks(&self) -> Vec { - self.open_path("tasks", false) + let mut file = "tasks"; + if self.is_v2() { + file = "cgroup.procs"; + } + self.open_path(file, false) .and_then(|file| { let bf = BufReader::new(file); let mut v = Vec::new(); @@ -608,6 +616,10 @@ impl Subsystem { Subsystem::Rdma(cont) => cont, } } + + fn controller_name(&self) -> String { + self.to_controller().control_type().to_string() + } } @@ -627,6 +639,22 @@ impl Default for MaxValue { } } +impl MaxValue { + fn to_i64(&self) -> i64 { + match self { + MaxValue::Max => -1, + MaxValue::Value(num) => *num, + } + } + + fn to_string(&self) -> String { + match self { + MaxValue::Max => "max".to_string(), + MaxValue::Value(num) => num.to_string(), + } + } +} + pub fn parse_max_value(s: &String) -> Result { if s.trim() == "max" { return Ok(MaxValue::Max) @@ -636,10 +664,3 @@ pub fn parse_max_value(s: &String) -> Result { Err(e) => Err(Error::with_cause(ParseError, e)), } } - -pub fn max_value_to_string(m: MaxValue) -> String { - match m { - MaxValue::Max => "max".to_string(), - MaxValue::Value(num) => num.to_string(), - } -} \ No newline at end of file diff --git a/src/memory.rs b/src/memory.rs index 8793459..32f2b07 100644 --- a/src/memory.rs +++ b/src/memory.rs @@ -13,11 +13,9 @@ use crate::error::*; use crate::events; use crate::{ - ControllIdentifier, ControllerInternal, Controllers, MemoryResources, Resources, Subsystem, + ControllIdentifier, ControllerInternal, Controllers, MaxValue, MemoryResources, Resources, Subsystem, }; -use crate::{MaxValue, max_value_to_string, parse_max_value}; - /// A controller that allows controlling the `memory` subsystem of a Cgroup. /// /// In essence, using the memory controller, the user can gather statistics about the memory usage @@ -481,8 +479,7 @@ impl MemController { let v = value.0; let f = value.1; if v.is_some() { - let v = v.unwrap(); - let v = max_value_to_string(v); + let v = v.unwrap().to_string(); self.open_path(f, true) .and_then(|mut file| { file.write_all(v.as_ref()) @@ -504,12 +501,44 @@ impl MemController { Ok(m) } + fn memory_stat_v2(&self) -> Memory { + let set = self.get_mem().unwrap(); + + Memory { + fail_cnt: 0, + limit_in_bytes: set.max.unwrap().to_i64(), + usage_in_bytes: self + .open_path("memory.current", false) + .and_then(read_u64_from) + .unwrap_or(0), + max_usage_in_bytes: 0, + move_charge_at_immigrate: 0, + numa_stat: NumaStat::default(), + oom_control: OomControl::default(), + soft_limit_in_bytes: set.high.unwrap().to_i64(), + stat: self + .open_path("memory.stat", false) + .and_then(read_string_from) + .and_then(parse_memory_stat) + .unwrap_or(MemoryStat::default()), + swappiness: self + .open_path("memory.swap.current", false) + .and_then(read_u64_from) + .unwrap_or(0), + use_hierarchy: 0, + } + } + /// Gathers overall statistics (and the current state of) about the memory usage of the control /// group's tasks. /// /// See the individual fields for more explanation, and as always, remember to consult the /// kernel Documentation and/or sources. pub fn memory_stat(&self) -> Memory { + if self.v2 { + return self.memory_stat_v2(); + } + Memory { fail_cnt: self .open_path("memory.failcnt", false) @@ -670,7 +699,11 @@ impl MemController { /// Set the memory usage limit of the control group, in bytes. pub fn set_limit(&self, limit: i64) -> Result<()> { - self.open_path("memory.limit_in_bytes", true) + let mut file = "memory.limit_in_bytes"; + if self.v2 { + file = "memory.max"; + } + self.open_path(file, true) .and_then(|mut file| { file.write_all(limit.to_string().as_ref()) .map_err(|e| Error::with_cause(WriteFailed, e)) @@ -688,7 +721,11 @@ impl MemController { /// Set the memory+swap limit of the control group, in bytes. pub fn set_memswap_limit(&self, limit: i64) -> Result<()> { - self.open_path("memory.memsw.limit_in_bytes", true) + let mut file = "memory.memsw.limit_in_bytes"; + if self.v2 { + file = "memory.swap.max"; + } + self.open_path(file, true) .and_then(|mut file| { file.write_all(limit.to_string().as_ref()) .map_err(|e| Error::with_cause(WriteFailed, e)) @@ -709,7 +746,11 @@ impl MemController { /// This limit is enforced when the system is nearing OOM conditions. Contrast this with the /// hard limit, which is _always_ enforced. pub fn set_soft_limit(&self, limit: i64) -> Result<()> { - self.open_path("memory.soft_limit_in_bytes", true) + let mut file = "memory.soft_limit_in_bytes"; + if self.v2 { + file = "memory.low" + } + self.open_path(file, true) .and_then(|mut file| { file.write_all(limit.to_string().as_ref()) .map_err(|e| Error::with_cause(WriteFailed, e)) diff --git a/src/pid.rs b/src/pid.rs index 3be3da0..4eb6754 100644 --- a/src/pid.rs +++ b/src/pid.rs @@ -10,7 +10,7 @@ use crate::error::*; use crate::error::ErrorKind::*; use crate::{ - ControllIdentifier, ControllerInternal, Controllers, MaxValue, max_value_to_string, parse_max_value, PidResources, Resources, Subsystem, + ControllIdentifier, ControllerInternal, Controllers, MaxValue, parse_max_value, PidResources, Resources, Subsystem, }; /// A controller that allows controlling the `pids` subsystem of a Cgroup. @@ -150,7 +150,7 @@ impl PidController { /// extra processes to a control group disregards the limit. pub fn set_pid_max(&self, max_pid: MaxValue) -> Result<()> { self.open_path("pids.max", true).and_then(|mut file| { - let string_to_write = max_value_to_string(max_pid); + let string_to_write = max_pid.to_string(); match file.write_all(string_to_write.as_ref()) { Ok(_) => Ok(()), Err(e) => Err(Error::with_cause(WriteFailed, e)), diff --git a/tests/cgroup.rs b/tests/cgroup.rs index 177eea3..1f5efbd 100644 --- a/tests/cgroup.rs +++ b/tests/cgroup.rs @@ -1,5 +1,7 @@ //! Simple unit tests about the control groups system. -use cgroups::{Cgroup, CgroupPid, Hierarchy}; +use cgroups::{Cgroup, CgroupPid, Hierarchy, Subsystem}; +use cgroups::memory::{MemController, SetMemory}; +use std::collections::HashMap; #[test] fn test_tasks_iterator() { @@ -24,3 +26,24 @@ fn test_tasks_iterator() { } cg.delete(); } + + +#[test] +fn test_cgroup_with_prefix() { + let h = cgroups::hierarchies::auto(); + let h = Box::new(&*h); + let mut prefixes = HashMap::new(); + prefixes.insert("memory".to_string(), "/memory/abc/def".to_string()); + let cg = Cgroup::new_with_prefix(h, String::from("test_cgroup_with_prefix"), prefixes); + { + let subsystems = cg.subsystems(); + println!("mem path: {:?}", &subsystems); + subsystems.into_iter().for_each(|sub| match sub { + Subsystem::Pid(c) => {println!("path {:?}", c);}, + // base: "/sys/fs/cgroup", path: "/sys/fs/cgroup/memory/abc/def/test_cgroup_with_prefix" + Subsystem::Mem(c) => {println!("path {:?}", c);}, + _ => {}, + }); + } + cg.delete(); +}