#![allow(clippy::unnecessary_unwrap)]
use log::*;
use std::collections::HashMap;
use std::fmt;
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 fmt::Display for Controllers {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Controllers::Pids => write!(f, "pids"),
Controllers::Mem => write!(f, "memory"),
Controllers::CpuSet => write!(f, "cpuset"),
Controllers::CpuAcct => write!(f, "cpuacct"),
Controllers::Cpu => write!(f, "cpu"),
Controllers::Devices => write!(f, "devices"),
Controllers::Freezer => write!(f, "freezer"),
Controllers::NetCls => write!(f, "net_cls"),
Controllers::BlkIo => write!(f, "blkio"),
Controllers::PerfEvent => write!(f, "perf_event"),
Controllers::NetPrio => write!(f, "net_prio"),
Controllers::HugeTlb => write!(f, "hugetlb"),
Controllers::Rdma => write!(f, "rdma"),
Controllers::Systemd => write!(f, "name=systemd"),
}
}
}
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) => Err(Error::with_cause(ErrorKind::WriteFailed, e)),
Ok(file) => Ok(file),
}
} else {
match File::open(&path) {
Err(e) => Err(Error::with_cause(ErrorKind::ReadFailed, e)),
Ok(file) => 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 self.verify_path().is_err() {
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()
.unwrap_or_else(|_| panic!("path should be valid: {:?}", self.path()));
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(());
}
remove_dir(self.get_path())
}
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)
.map(|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);
}
}
v.into_iter().map(CgroupPid::from).collect()
})
.unwrap_or_default()
}
fn v2(&self) -> bool {
self.is_v2()
}
}
fn remove_dir(dir: &PathBuf) -> Result<()> {
if fs::remove_dir(dir).is_ok() {
return Ok(());
}
if dir.exists() && dir.is_dir() {
for entry in fs::read_dir(dir).map_err(|e| Error::with_cause(ReadFailed, e))? {
let entry = entry.map_err(|e| Error::with_cause(ReadFailed, e))?;
let path = entry.path();
if path.is_dir() {
remove_dir(&path)?;
}
}
fs::remove_dir(dir).map_err(|e| Error::with_cause(RemoveFailed, e))?;
}
Ok(())
}
#[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(mut cont) => Subsystem::Pid({
cont.get_path_mut().push(path);
cont
}),
Subsystem::Mem(mut cont) => Subsystem::Mem({
cont.get_path_mut().push(path);
cont
}),
Subsystem::CpuSet(mut cont) => Subsystem::CpuSet({
cont.get_path_mut().push(path);
cont
}),
Subsystem::CpuAcct(mut cont) => Subsystem::CpuAcct({
cont.get_path_mut().push(path);
cont
}),
Subsystem::Cpu(mut cont) => Subsystem::Cpu({
cont.get_path_mut().push(path);
cont
}),
Subsystem::Devices(mut cont) => Subsystem::Devices({
cont.get_path_mut().push(path);
cont
}),
Subsystem::Freezer(mut cont) => Subsystem::Freezer({
cont.get_path_mut().push(path);
cont
}),
Subsystem::NetCls(mut cont) => Subsystem::NetCls({
cont.get_path_mut().push(path);
cont
}),
Subsystem::BlkIo(mut cont) => Subsystem::BlkIo({
cont.get_path_mut().push(path);
cont
}),
Subsystem::PerfEvent(mut cont) => Subsystem::PerfEvent({
cont.get_path_mut().push(path);
cont
}),
Subsystem::NetPrio(mut cont) => Subsystem::NetPrio({
cont.get_path_mut().push(path);
cont
}),
Subsystem::HugeTlb(mut cont) => Subsystem::HugeTlb({
cont.get_path_mut().push(path);
cont
}),
Subsystem::Rdma(mut cont) => Subsystem::Rdma({
cont.get_path_mut().push(path);
cont
}),
Subsystem::Systemd(mut cont) => Subsystem::Systemd({
cont.get_path_mut().push(path);
cont
}),
}
}
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,
}
}
}
impl fmt::Display for MaxValue {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
MaxValue::Max => write!(f, "max"),
MaxValue::Value(num) => write!(f, "{}", num.to_string()),
}
}
}
pub fn parse_max_value(s: &str) -> 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 {
if let Ok(i) = parts[1].parse::<i64>() {
v.push((parts[0].to_string(), i));
}
}
}
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 {
if let Ok(i) = parts[1].parse::<i64>() {
h.insert(parts[0].to_string(), i);
}
}
}
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.is_empty() {
continue;
}
let mut th = HashMap::new();
for item in parts[1..].iter() {
let fields: Vec<&str> = item.split('=').collect();
if fields.len() == 2 {
if let Ok(i) = fields[1].parse::<i64>() {
th.insert(fields[0].to_string(), i);
}
}
}
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)
}