use log::*;
use std::io::Write;
use std::path::PathBuf;
use crate::error::ErrorKind::*;
use crate::error::*;
use crate::{read_string_from, read_u64_from};
use crate::{
ControllIdentifier, ControllerInternal, Controllers, CpuResources, Resources, Subsystem,
};
#[derive(Debug, Clone)]
pub struct CpuSetController {
base: PathBuf,
path: PathBuf,
v2: bool,
}
pub struct CpuSet {
pub cpu_exclusive: bool,
pub cpus: Vec<(u64, u64)>,
pub effective_cpus: Vec<(u64, u64)>,
pub effective_mems: Vec<(u64, u64)>,
pub mem_exclusive: bool,
pub mem_hardwall: bool,
pub memory_migrate: bool,
pub memory_pressure: u64,
pub memory_pressure_enabled: Option<bool>,
pub memory_spread_page: bool,
pub memory_spread_slab: bool,
pub mems: Vec<(u64, u64)>,
pub sched_load_balance: bool,
pub sched_relax_domain_level: u64,
}
impl ControllerInternal for CpuSetController {
fn control_type(&self) -> Controllers {
Controllers::CpuSet
}
fn get_path(&self) -> &PathBuf {
&self.path
}
fn get_path_mut(&mut self) -> &mut PathBuf {
&mut self.path
}
fn get_base(&self) -> &PathBuf {
&self.base
}
fn is_v2(&self) -> bool {
self.v2
}
fn apply(&self, res: &Resources) -> Result<()> {
let res: &CpuResources = &res.cpu;
update!(self, set_cpus, res.cpus.as_ref());
update!(self, set_mems, res.mems.as_ref());
Ok(())
}
fn post_create(&self) {
if self.is_v2() {
return;
}
let current = self.get_path();
if current != self.get_base() {
match copy_from_parent(current.to_str().unwrap(), "cpuset.cpus") {
Ok(_) => (),
Err(err) => error!("error create_dir for cpuset.cpus {:?}", err),
}
match copy_from_parent(current.to_str().unwrap(), "cpuset.mems") {
Ok(_) => (),
Err(err) => error!("error create_dir for cpuset.mems {:?}", err),
}
}
}
}
fn find_no_empty_parent(from: &str, file: &str) -> Result<(String, Vec<PathBuf>)> {
let mut current_path = ::std::path::Path::new(from).to_path_buf();
let mut v = vec![];
loop {
let current_value =
match ::std::fs::read_to_string(current_path.clone().join(file).to_str().unwrap()) {
Ok(cpus) => String::from(cpus.trim()),
Err(e) => return Err(Error::with_cause(ReadFailed, e)),
};
if !current_value.is_empty() {
return Ok((current_value, v));
}
v.push(current_path.clone());
let parent = match current_path.parent() {
Some(p) => p,
None => return Ok(("".to_string(), v)),
};
current_path = parent.to_path_buf();
}
}
fn copy_from_parent(current: &str, file: &str) -> Result<()> {
let (value, parents) = find_no_empty_parent(current, file)?;
if value.is_empty() || parents.is_empty() {
return Ok(());
}
for p in parents.iter().rev() {
let mut pb = p.clone();
pb.push(file);
match ::std::fs::write(pb.to_str().unwrap(), value.as_bytes()) {
Ok(_) => (),
Err(e) => return Err(Error::with_cause(WriteFailed, e)),
}
}
Ok(())
}
impl ControllIdentifier for CpuSetController {
fn controller_type() -> Controllers {
Controllers::CpuSet
}
}
impl<'a> From<&'a Subsystem> for &'a CpuSetController {
fn from(sub: &'a Subsystem) -> &'a CpuSetController {
unsafe {
match sub {
Subsystem::CpuSet(c) => c,
_ => {
assert_eq!(1, 0);
let v = std::mem::MaybeUninit::uninit();
v.assume_init()
}
}
}
}
}
fn parse_range(s: String) -> Result<Vec<(u64, u64)>> {
let mut fin = Vec::new();
if s.is_empty() {
return Ok(fin);
}
let comma_split = s.split(',');
for sp in comma_split {
if sp.contains('-') {
let dash_split = sp.split('-').collect::<Vec<_>>();
if dash_split.len() != 2 {
return Err(Error::new(ParseError));
}
let first = dash_split[0].parse::<u64>();
let second = dash_split[1].parse::<u64>();
if first.is_err() || second.is_err() {
return Err(Error::new(ParseError));
}
fin.push((first.unwrap(), second.unwrap()));
} else {
let num = sp.parse::<u64>();
if num.is_err() {
return Err(Error::new(ParseError));
}
fin.push((num.clone().unwrap(), num.clone().unwrap()));
}
}
Ok(fin)
}
impl CpuSetController {
pub fn new(root: PathBuf, v2: bool) -> Self {
Self {
base: root.clone(),
path: root,
v2,
}
}
pub fn cpuset(&self) -> CpuSet {
CpuSet {
cpu_exclusive: {
self.open_path("cpuset.cpu_exclusive", false)
.and_then(read_u64_from)
.map(|x| x == 1)
.unwrap_or(false)
},
cpus: {
self.open_path("cpuset.cpus", false)
.and_then(read_string_from)
.and_then(parse_range)
.unwrap_or_default()
},
effective_cpus: {
self.open_path("cpuset.effective_cpus", false)
.and_then(read_string_from)
.and_then(parse_range)
.unwrap_or_default()
},
effective_mems: {
self.open_path("cpuset.effective_mems", false)
.and_then(read_string_from)
.and_then(parse_range)
.unwrap_or_default()
},
mem_exclusive: {
self.open_path("cpuset.mem_exclusive", false)
.and_then(read_u64_from)
.map(|x| x == 1)
.unwrap_or(false)
},
mem_hardwall: {
self.open_path("cpuset.mem_hardwall", false)
.and_then(read_u64_from)
.map(|x| x == 1)
.unwrap_or(false)
},
memory_migrate: {
self.open_path("cpuset.memory_migrate", false)
.and_then(read_u64_from)
.map(|x| x == 1)
.unwrap_or(false)
},
memory_pressure: {
self.open_path("cpuset.memory_pressure", false)
.and_then(read_u64_from)
.unwrap_or(0)
},
memory_pressure_enabled: {
self.open_path("cpuset.memory_pressure_enabled", false)
.and_then(read_u64_from)
.map(|x| x == 1)
.ok()
},
memory_spread_page: {
self.open_path("cpuset.memory_spread_page", false)
.and_then(read_u64_from)
.map(|x| x == 1)
.unwrap_or(false)
},
memory_spread_slab: {
self.open_path("cpuset.memory_spread_slab", false)
.and_then(read_u64_from)
.map(|x| x == 1)
.unwrap_or(false)
},
mems: {
self.open_path("cpuset.mems", false)
.and_then(read_string_from)
.and_then(parse_range)
.unwrap_or_default()
},
sched_load_balance: {
self.open_path("cpuset.sched_load_balance", false)
.and_then(read_u64_from)
.map(|x| x == 1)
.unwrap_or(false)
},
sched_relax_domain_level: {
self.open_path("cpuset.sched_relax_domain_level", false)
.and_then(read_u64_from)
.unwrap_or(0)
},
}
}
pub fn set_cpu_exclusive(&self, b: bool) -> Result<()> {
self.open_path("cpuset.cpu_exclusive", true)
.and_then(|mut file| {
if b {
file.write_all(b"1")
.map_err(|e| Error::with_cause(WriteFailed, e))
} else {
file.write_all(b"0")
.map_err(|e| Error::with_cause(WriteFailed, e))
}
})
}
pub fn set_mem_exclusive(&self, b: bool) -> Result<()> {
self.open_path("cpuset.mem_exclusive", true)
.and_then(|mut file| {
if b {
file.write_all(b"1")
.map_err(|e| Error::with_cause(WriteFailed, e))
} else {
file.write_all(b"0")
.map_err(|e| Error::with_cause(WriteFailed, e))
}
})
}
pub fn set_cpus(&self, cpus: &str) -> Result<()> {
self.open_path("cpuset.cpus", true).and_then(|mut file| {
file.write_all(cpus.as_ref())
.map_err(|e| Error::with_cause(WriteFailed, e))
})
}
pub fn set_mems(&self, mems: &str) -> Result<()> {
self.open_path("cpuset.mems", true).and_then(|mut file| {
file.write_all(mems.as_ref())
.map_err(|e| Error::with_cause(WriteFailed, e))
})
}
pub fn set_hardwall(&self, b: bool) -> Result<()> {
self.open_path("cpuset.mem_hardwall", true)
.and_then(|mut file| {
if b {
file.write_all(b"1")
.map_err(|e| Error::with_cause(WriteFailed, e))
} else {
file.write_all(b"0")
.map_err(|e| Error::with_cause(WriteFailed, e))
}
})
}
pub fn set_load_balancing(&self, b: bool) -> Result<()> {
self.open_path("cpuset.sched_load_balance", true)
.and_then(|mut file| {
if b {
file.write_all(b"1")
.map_err(|e| Error::with_cause(WriteFailed, e))
} else {
file.write_all(b"0")
.map_err(|e| Error::with_cause(WriteFailed, e))
}
})
}
pub fn set_rebalance_relax_domain_level(&self, i: i64) -> Result<()> {
self.open_path("cpuset.sched_relax_domain_level", true)
.and_then(|mut file| {
file.write_all(i.to_string().as_ref())
.map_err(|e| Error::with_cause(WriteFailed, e))
})
}
pub fn set_memory_migration(&self, b: bool) -> Result<()> {
self.open_path("cpuset.memory_migrate", true)
.and_then(|mut file| {
if b {
file.write_all(b"1")
.map_err(|e| Error::with_cause(WriteFailed, e))
} else {
file.write_all(b"0")
.map_err(|e| Error::with_cause(WriteFailed, e))
}
})
}
pub fn set_memory_spread_page(&self, b: bool) -> Result<()> {
self.open_path("cpuset.memory_spread_page", true)
.and_then(|mut file| {
if b {
file.write_all(b"1")
.map_err(|e| Error::with_cause(WriteFailed, e))
} else {
file.write_all(b"0")
.map_err(|e| Error::with_cause(WriteFailed, e))
}
})
}
pub fn set_memory_spread_slab(&self, b: bool) -> Result<()> {
self.open_path("cpuset.memory_spread_slab", true)
.and_then(|mut file| {
if b {
file.write_all(b"1")
.map_err(|e| Error::with_cause(WriteFailed, e))
} else {
file.write_all(b"0")
.map_err(|e| Error::with_cause(WriteFailed, e))
}
})
}
pub fn set_enable_memory_pressure(&self, b: bool) -> Result<()> {
if !self.path_exists("cpuset.memory_pressure_enabled") {
return Err(Error::new(InvalidOperation));
}
self.open_path("cpuset.memory_pressure_enabled", true)
.and_then(|mut file| {
if b {
file.write_all(b"1")
.map_err(|e| Error::with_cause(WriteFailed, e))
} else {
file.write_all(b"0")
.map_err(|e| Error::with_cause(WriteFailed, e))
}
})
}
}
#[cfg(test)]
mod tests {
use crate::cpuset;
#[test]
fn test_parse_range() {
let test_cases = vec![
"1,2,4-6,9".to_string(),
"".to_string(),
"1".to_string(),
"1-111".to_string(),
"1,2,3,4".to_string(),
"1-5,6-7,8-9".to_string(),
];
let expecteds = vec![
vec![(1, 1), (2, 2), (4, 6), (9, 9)],
vec![],
vec![(1, 1)],
vec![(1, 111)],
vec![(1, 1), (2, 2), (3, 3), (4, 4)],
vec![(1, 5), (6, 7), (8, 9)],
];
for (i, case) in test_cases.into_iter().enumerate() {
let range = cpuset::parse_range(case.clone());
println!("{:?} => {:?}", case, range);
assert!(range.is_ok());
assert_eq!(range.unwrap(), expecteds[i]);
}
}
}