use log::*;
use std::collections::HashMap;
use std::fs::{self, File};
use std::io::{BufRead, BufReader, Read, Write};
use std::path::{Path, PathBuf};
use std::str::FromStr;
macro_rules! update_and_test {
($self: ident, $set_func:ident, $value:expr, $get_func:ident) => {
if let Some(v) = $value {
$self.$set_func(v)?;
if $self.$get_func()? != v {
return Err(Error::new(Other));
}
}
};
}
macro_rules! update {
($self: ident, $set_func:ident, $value:expr) => {
if let Some(v) = $value {
let _ = $self.$set_func(v);
}
};
}
pub mod blkio;
pub mod cgroup;
pub mod cgroup_builder;
pub mod cpu;
pub mod cpuacct;
pub mod cpuset;
pub mod devices;
pub mod error;
pub mod events;
pub mod freezer;
pub mod hierarchies;
pub mod hugetlb;
pub mod memory;
pub mod net_cls;
pub mod net_prio;
pub mod perf_event;
pub mod pid;
pub mod rdma;
pub mod systemd;
use crate::blkio::BlkIoController;
use crate::cpu::CpuController;
use crate::cpuacct::CpuAcctController;
use crate::cpuset::CpuSetController;
use crate::devices::DevicesController;
use crate::error::ErrorKind::*;
use crate::error::*;
use crate::freezer::FreezerController;
use crate::hugetlb::HugeTlbController;
use crate::memory::MemController;
use crate::net_cls::NetClsController;
use crate::net_prio::NetPrioController;
use crate::perf_event::PerfEventController;
use crate::pid::PidController;
use crate::rdma::RdmaController;
use crate::systemd::SystemdController;
#[doc(inline)]
pub use crate::cgroup::Cgroup;
#[derive(Debug, Clone)]
pub enum Subsystem {
Pid(PidController),
Mem(MemController),
CpuSet(CpuSetController),
CpuAcct(CpuAcctController),
Cpu(CpuController),
Devices(DevicesController),
Freezer(FreezerController),
NetCls(NetClsController),
BlkIo(BlkIoController),
PerfEvent(PerfEventController),
NetPrio(NetPrioController),
HugeTlb(HugeTlbController),
Rdma(RdmaController),
Systemd(SystemdController),
}
#[doc(hidden)]
#[derive(Eq, PartialEq, Debug, Clone)]
pub enum Controllers {
Pids,
Mem,
CpuSet,
CpuAcct,
Cpu,
Devices,
Freezer,
NetCls,
BlkIo,
PerfEvent,
NetPrio,
HugeTlb,
Rdma,
Systemd,
}
impl Controllers {
pub fn to_string(&self) -> String {
match self {
Controllers::Pids => return "pids".to_string(),
Controllers::Mem => return "memory".to_string(),
Controllers::CpuSet => return "cpuset".to_string(),
Controllers::CpuAcct => return "cpuacct".to_string(),
Controllers::Cpu => return "cpu".to_string(),
Controllers::Devices => return "devices".to_string(),
Controllers::Freezer => return "freezer".to_string(),
Controllers::NetCls => return "net_cls".to_string(),
Controllers::BlkIo => return "blkio".to_string(),
Controllers::PerfEvent => return "perf_event".to_string(),
Controllers::NetPrio => return "net_prio".to_string(),
Controllers::HugeTlb => return "hugetlb".to_string(),
Controllers::Rdma => return "rdma".to_string(),
Controllers::Systemd => return "name=systemd".to_string(),
}
}
}
mod sealed {
use super::*;
pub trait ControllerInternal {
fn apply(&self, res: &Resources) -> Result<()>;
fn control_type(&self) -> Controllers;
fn get_path(&self) -> &PathBuf;
fn get_path_mut(&mut self) -> &mut PathBuf;
fn get_base(&self) -> &PathBuf;
fn post_create(&self) {}
fn is_v2(&self) -> bool {
false
}
fn verify_path(&self) -> Result<()> {
if self.get_path().starts_with(self.get_base()) {
Ok(())
} else {
Err(Error::new(ErrorKind::InvalidPath))
}
}
fn open_path(&self, p: &str, w: bool) -> Result<File> {
let mut path = self.get_path().clone();
path.push(p);
self.verify_path()?;
if w {
match File::create(&path) {
Err(e) => return Err(Error::with_cause(ErrorKind::WriteFailed, e)),
Ok(file) => return Ok(file),
}
} else {
match File::open(&path) {
Err(e) => return Err(Error::with_cause(ErrorKind::ReadFailed, e)),
Ok(file) => return Ok(file),
}
}
}
fn get_max_value(&self, f: &str) -> Result<MaxValue> {
self.open_path(f, false).and_then(|mut file| {
let mut string = String::new();
let res = file.read_to_string(&mut string);
match res {
Ok(_) => parse_max_value(&string),
Err(e) => Err(Error::with_cause(ReadFailed, e)),
}
})
}
#[doc(hidden)]
fn path_exists(&self, p: &str) -> bool {
if let Err(_) = self.verify_path() {
return false;
}
std::path::Path::new(p).exists()
}
}
pub trait CustomizedAttribute: ControllerInternal {
fn set(&self, key: &str, value: &str) -> Result<()> {
self.open_path(key, true).and_then(|mut file| {
file.write_all(value.as_ref())
.map_err(|e| Error::with_cause(WriteFailed, e))
})
}
fn get(&self, key: &str) -> Result<String> {
self.open_path(key, false).and_then(|mut file: File| {
let mut string = String::new();
match file.read_to_string(&mut string) {
Ok(_) => Ok(string.trim().to_owned()),
Err(e) => Err(Error::with_cause(ReadFailed, e)),
}
})
}
}
}
pub(crate) use crate::sealed::{ControllerInternal, CustomizedAttribute};
pub trait Controller {
#[doc(hidden)]
fn control_type(&self) -> Controllers;
fn path(&self) -> &Path;
fn apply(&self, res: &Resources) -> Result<()>;
fn create(&self);
fn exists(&self) -> bool;
fn set_notify_on_release(&self, enable: bool) -> Result<()>;
fn set_release_agent(&self, path: &str) -> Result<()>;
fn delete(&self) -> Result<()>;
fn add_task(&self, pid: &CgroupPid) -> Result<()>;
fn add_task_by_tgid(&self, pid: &CgroupPid) -> Result<()>;
fn tasks(&self) -> Vec<CgroupPid>;
fn v2(&self) -> bool;
}
impl<T> Controller for T
where
T: ControllerInternal,
{
fn control_type(&self) -> Controllers {
ControllerInternal::control_type(self)
}
fn path(&self) -> &Path {
self.get_path()
}
fn apply(&self, res: &Resources) -> Result<()> {
ControllerInternal::apply(self, res)
}
fn create(&self) {
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(),
Err(e) => warn!("error create_dir: {:?} error: {:?}", self.get_path(), e),
}
}
fn set_notify_on_release(&self, enable: bool) -> Result<()> {
self.open_path("notify_on_release", true)
.and_then(|mut file| {
write!(file, "{}", enable as i32)
.map_err(|e| Error::with_cause(ErrorKind::WriteFailed, e))
})
}
fn set_release_agent(&self, path: &str) -> Result<()> {
self.open_path("release_agent", true).and_then(|mut file| {
file.write_all(path.as_bytes())
.map_err(|e| Error::with_cause(ErrorKind::WriteFailed, e))
})
}
fn exists(&self) -> bool {
self.get_path().exists()
}
fn delete(&self) -> Result<()> {
if !self.get_path().exists() {
return Ok(());
}
fs::remove_dir(self.get_path()).map_err(|e| Error::with_cause(ErrorKind::RemoveFailed, e))
}
fn add_task(&self, pid: &CgroupPid) -> Result<()> {
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))
})
}
fn add_task_by_tgid(&self, pid: &CgroupPid) -> Result<()> {
self.open_path("cgroup.procs", true).and_then(|mut file| {
file.write_all(pid.pid.to_string().as_ref())
.map_err(|e| Error::with_cause(ErrorKind::WriteFailed, e))
})
}
fn tasks(&self) -> Vec<CgroupPid> {
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();
for line in bf.lines() {
if let Ok(line) = line {
let n = line.trim().parse().unwrap_or(0u64);
v.push(n);
}
}
Ok(v.into_iter().map(CgroupPid::from).collect())
})
.unwrap_or(vec![])
}
fn v2(&self) -> bool {
self.is_v2()
}
}
#[doc(hidden)]
pub trait ControllIdentifier {
fn controller_type() -> Controllers;
}
pub trait Hierarchy: std::fmt::Debug + Send + Sync {
fn subsystems(&self) -> Vec<Subsystem>;
fn root(&self) -> PathBuf;
fn root_control_group(&self) -> Cgroup;
fn v2(&self) -> bool;
}
#[derive(Debug, Clone, Eq, PartialEq, Default)]
pub struct MemoryResources {
pub kernel_memory_limit: Option<i64>,
pub memory_hard_limit: Option<i64>,
pub memory_soft_limit: Option<i64>,
pub kernel_tcp_memory_limit: Option<i64>,
pub memory_swap_limit: Option<i64>,
pub swappiness: Option<u64>,
pub attrs: std::collections::HashMap<&'static str, String>,
}
#[derive(Debug, Clone, Eq, PartialEq, Default)]
pub struct PidResources {
pub maximum_number_of_processes: Option<MaxValue>,
}
#[derive(Debug, Clone, Eq, PartialEq, Default)]
pub struct CpuResources {
pub cpus: Option<String>,
pub mems: Option<String>,
pub shares: Option<u64>,
pub quota: Option<i64>,
pub period: Option<u64>,
pub realtime_runtime: Option<i64>,
pub realtime_period: Option<u64>,
pub attrs: std::collections::HashMap<&'static str, String>,
}
#[derive(Debug, Clone, Eq, PartialEq, Default)]
pub struct DeviceResource {
pub allow: bool,
pub devtype: crate::devices::DeviceType,
pub major: i64,
pub minor: i64,
pub access: Vec<crate::devices::DevicePermissions>,
}
#[derive(Debug, Clone, Eq, PartialEq, Default)]
pub struct DeviceResources {
pub devices: Vec<DeviceResource>,
}
#[derive(Debug, Clone, Eq, PartialEq, Default)]
pub struct NetworkPriority {
pub name: String,
pub priority: u64,
}
#[derive(Debug, Clone, Eq, PartialEq, Default)]
pub struct NetworkResources {
pub class_id: Option<u64>,
pub priorities: Vec<NetworkPriority>,
}
#[derive(Debug, Clone, Eq, PartialEq, Default)]
pub struct HugePageResource {
pub size: String,
pub limit: u64,
}
#[derive(Debug, Clone, Eq, PartialEq, Default)]
pub struct HugePageResources {
pub limits: Vec<HugePageResource>,
}
#[derive(Debug, Clone, Eq, PartialEq, Default)]
pub struct BlkIoDeviceResource {
pub major: u64,
pub minor: u64,
pub weight: Option<u16>,
pub leaf_weight: Option<u16>,
}
#[derive(Debug, Clone, Eq, PartialEq, Default)]
pub struct BlkIoDeviceThrottleResource {
pub major: u64,
pub minor: u64,
pub rate: u64,
}
#[derive(Debug, Clone, Eq, PartialEq, Default)]
pub struct BlkIoResources {
pub weight: Option<u16>,
pub leaf_weight: Option<u16>,
pub weight_device: Vec<BlkIoDeviceResource>,
pub throttle_read_bps_device: Vec<BlkIoDeviceThrottleResource>,
pub throttle_read_iops_device: Vec<BlkIoDeviceThrottleResource>,
pub throttle_write_bps_device: Vec<BlkIoDeviceThrottleResource>,
pub throttle_write_iops_device: Vec<BlkIoDeviceThrottleResource>,
}
#[derive(Debug, Clone, Eq, PartialEq, Default)]
pub struct Resources {
pub memory: MemoryResources,
pub pid: PidResources,
pub cpu: CpuResources,
pub devices: DeviceResources,
pub network: NetworkResources,
pub hugepages: HugePageResources,
pub blkio: BlkIoResources,
}
#[derive(Debug, PartialEq, Eq, PartialOrd, Ord)]
pub struct CgroupPid {
pub pid: u64,
}
impl From<u64> for CgroupPid {
fn from(u: u64) -> CgroupPid {
CgroupPid { pid: u }
}
}
impl<'a> From<&'a std::process::Child> for CgroupPid {
fn from(u: &std::process::Child) -> CgroupPid {
CgroupPid { pid: u.id() as u64 }
}
}
impl Subsystem {
fn enter(self, path: &Path) -> Self {
match self {
Subsystem::Pid(cont) => Subsystem::Pid({
let mut c = cont.clone();
c.get_path_mut().push(path);
c
}),
Subsystem::Mem(cont) => Subsystem::Mem({
let mut c = cont.clone();
c.get_path_mut().push(path);
c
}),
Subsystem::CpuSet(cont) => Subsystem::CpuSet({
let mut c = cont.clone();
c.get_path_mut().push(path);
c
}),
Subsystem::CpuAcct(cont) => Subsystem::CpuAcct({
let mut c = cont.clone();
c.get_path_mut().push(path);
c
}),
Subsystem::Cpu(cont) => Subsystem::Cpu({
let mut c = cont.clone();
c.get_path_mut().push(path);
c
}),
Subsystem::Devices(cont) => Subsystem::Devices({
let mut c = cont.clone();
c.get_path_mut().push(path);
c
}),
Subsystem::Freezer(cont) => Subsystem::Freezer({
let mut c = cont.clone();
c.get_path_mut().push(path);
c
}),
Subsystem::NetCls(cont) => Subsystem::NetCls({
let mut c = cont.clone();
c.get_path_mut().push(path);
c
}),
Subsystem::BlkIo(cont) => Subsystem::BlkIo({
let mut c = cont.clone();
c.get_path_mut().push(path);
c
}),
Subsystem::PerfEvent(cont) => Subsystem::PerfEvent({
let mut c = cont.clone();
c.get_path_mut().push(path);
c
}),
Subsystem::NetPrio(cont) => Subsystem::NetPrio({
let mut c = cont.clone();
c.get_path_mut().push(path);
c
}),
Subsystem::HugeTlb(cont) => Subsystem::HugeTlb({
let mut c = cont.clone();
c.get_path_mut().push(path);
c
}),
Subsystem::Rdma(cont) => Subsystem::Rdma({
let mut c = cont.clone();
c.get_path_mut().push(path);
c
}),
Subsystem::Systemd(cont) => Subsystem::Systemd({
let mut c = cont.clone();
c.get_path_mut().push(path);
c
}),
}
}
pub fn to_controller(&self) -> &dyn Controller {
match self {
Subsystem::Pid(cont) => cont,
Subsystem::Mem(cont) => cont,
Subsystem::CpuSet(cont) => cont,
Subsystem::CpuAcct(cont) => cont,
Subsystem::Cpu(cont) => cont,
Subsystem::Devices(cont) => cont,
Subsystem::Freezer(cont) => cont,
Subsystem::NetCls(cont) => cont,
Subsystem::BlkIo(cont) => cont,
Subsystem::PerfEvent(cont) => cont,
Subsystem::NetPrio(cont) => cont,
Subsystem::HugeTlb(cont) => cont,
Subsystem::Rdma(cont) => cont,
Subsystem::Systemd(cont) => cont,
}
}
pub fn controller_name(&self) -> String {
self.to_controller().control_type().to_string()
}
}
#[derive(Eq, PartialEq, Copy, Clone, Debug)]
pub enum MaxValue {
Max,
Value(i64),
}
impl Default for MaxValue {
fn default() -> Self {
MaxValue::Max
}
}
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<MaxValue> {
if s.trim() == "max" {
return Ok(MaxValue::Max);
}
match s.trim().parse() {
Ok(val) => Ok(MaxValue::Value(val)),
Err(e) => Err(Error::with_cause(ParseError, e)),
}
}
pub fn flat_keyed_to_vec(mut file: File) -> Result<Vec<(String, i64)>> {
let mut content = String::new();
file.read_to_string(&mut content)
.map_err(|e| Error::with_cause(ReadFailed, e))?;
let mut v = Vec::new();
for line in content.lines() {
let parts: Vec<&str> = line.split(' ').collect();
if parts.len() == 2 {
match parts[1].parse::<i64>() {
Ok(i) => {
v.push((parts[0].to_string(), i));
}
Err(_) => {}
}
}
}
Ok(v)
}
pub fn flat_keyed_to_hashmap(mut file: File) -> Result<HashMap<String, i64>> {
let mut content = String::new();
file.read_to_string(&mut content)
.map_err(|e| Error::with_cause(ReadFailed, e))?;
let mut h = HashMap::new();
for line in content.lines() {
let parts: Vec<&str> = line.split(' ').collect();
if parts.len() == 2 {
match parts[1].parse::<i64>() {
Ok(i) => {
h.insert(parts[0].to_string(), i);
}
Err(_) => {}
}
}
}
Ok(h)
}
pub fn nested_keyed_to_hashmap(mut file: File) -> Result<HashMap<String, HashMap<String, i64>>> {
let mut content = String::new();
file.read_to_string(&mut content)
.map_err(|e| Error::with_cause(ReadFailed, e))?;
let mut h = HashMap::new();
for line in content.lines() {
let parts: Vec<&str> = line.split(' ').collect();
if parts.len() == 0 {
continue;
}
let mut th = HashMap::new();
for item in parts[1..].into_iter() {
let fields: Vec<&str> = item.split('=').collect();
if fields.len() == 2 {
match fields[1].parse::<i64>() {
Ok(i) => {
th.insert(fields[0].to_string(), i);
}
Err(_) => {}
}
}
}
h.insert(parts[0].to_string(), th);
}
Ok(h)
}
fn read_from<T>(mut file: File) -> Result<T>
where
T: FromStr,
<T as FromStr>::Err: 'static + Send + Sync + std::error::Error,
{
let mut string = String::new();
match file.read_to_string(&mut string) {
Ok(_) => string
.trim()
.parse::<T>()
.map_err(|e| Error::with_cause(ParseError, e)),
Err(e) => Err(Error::with_cause(ReadFailed, e)),
}
}
fn read_string_from(mut file: File) -> Result<String> {
let mut string = String::new();
match file.read_to_string(&mut string) {
Ok(_) => Ok(string.trim().to_string()),
Err(e) => Err(Error::with_cause(ReadFailed, e)),
}
}
fn read_u64_from(file: File) -> Result<u64> {
read_from::<u64>(file)
}
fn read_i64_from(file: File) -> Result<i64> {
read_from::<i64>(file)
}