| use std::collections::HashMap; |
| use std::fs::File; |
| use std::io::{BufRead, BufReader, Read}; |
| use std::mem; |
| use std::path::{Path, PathBuf}; |
| use std::sync::atomic::{AtomicUsize, Ordering}; |
| use std::sync::Once; |
| |
| use libc; |
| |
| macro_rules! debug { |
| ($($args:expr),*) => ({ |
| if false { |
| //if true { |
| println!($($args),*); |
| } |
| }); |
| } |
| |
| macro_rules! some { |
| ($e:expr) => ({ |
| match $e { |
| Some(v) => v, |
| None => { |
| debug!("NONE: {:?}", stringify!($e)); |
| return None; |
| } |
| } |
| }) |
| } |
| |
| pub fn get_num_cpus() -> usize { |
| match cgroups_num_cpus() { |
| Some(n) => n, |
| None => logical_cpus(), |
| } |
| } |
| |
| fn logical_cpus() -> usize { |
| let mut set: libc::cpu_set_t = unsafe { mem::zeroed() }; |
| if unsafe { libc::sched_getaffinity(0, mem::size_of::<libc::cpu_set_t>(), &mut set) } == 0 { |
| let mut count: u32 = 0; |
| for i in 0..libc::CPU_SETSIZE as usize { |
| if unsafe { libc::CPU_ISSET(i, &set) } { |
| count += 1 |
| } |
| } |
| count as usize |
| } else { |
| let cpus = unsafe { libc::sysconf(libc::_SC_NPROCESSORS_ONLN) }; |
| if cpus < 1 { |
| 1 |
| } else { |
| cpus as usize |
| } |
| } |
| } |
| |
| pub fn get_num_physical_cpus() -> usize { |
| let file = match File::open("/proc/cpuinfo") { |
| Ok(val) => val, |
| Err(_) => return get_num_cpus(), |
| }; |
| let reader = BufReader::new(file); |
| let mut map = HashMap::new(); |
| let mut physid: u32 = 0; |
| let mut cores: usize = 0; |
| let mut chgcount = 0; |
| for line in reader.lines().filter_map(|result| result.ok()) { |
| let mut it = line.split(':'); |
| let (key, value) = match (it.next(), it.next()) { |
| (Some(key), Some(value)) => (key.trim(), value.trim()), |
| _ => continue, |
| }; |
| if key == "physical id" { |
| match value.parse() { |
| Ok(val) => physid = val, |
| Err(_) => break, |
| }; |
| chgcount += 1; |
| } |
| if key == "cpu cores" { |
| match value.parse() { |
| Ok(val) => cores = val, |
| Err(_) => break, |
| }; |
| chgcount += 1; |
| } |
| if chgcount == 2 { |
| map.insert(physid, cores); |
| chgcount = 0; |
| } |
| } |
| let count = map.into_iter().fold(0, |acc, (_, cores)| acc + cores); |
| |
| if count == 0 { |
| get_num_cpus() |
| } else { |
| count |
| } |
| } |
| |
| /// Cached CPUs calculated from cgroups. |
| /// |
| /// If 0, check logical cpus. |
| // Allow deprecation warnings, we want to work on older rustc |
| #[allow(warnings)] |
| static CGROUPS_CPUS: AtomicUsize = ::std::sync::atomic::ATOMIC_USIZE_INIT; |
| |
| fn cgroups_num_cpus() -> Option<usize> { |
| #[allow(warnings)] |
| static ONCE: Once = ::std::sync::ONCE_INIT; |
| |
| ONCE.call_once(init_cgroups); |
| |
| let cpus = CGROUPS_CPUS.load(Ordering::Acquire); |
| |
| if cpus > 0 { |
| Some(cpus) |
| } else { |
| None |
| } |
| } |
| |
| fn init_cgroups() { |
| // Should only be called once |
| debug_assert!(CGROUPS_CPUS.load(Ordering::SeqCst) == 0); |
| |
| match load_cgroups("/proc/self/cgroup", "/proc/self/mountinfo") { |
| Some(quota) => { |
| if quota == 0 { |
| return; |
| } |
| |
| let logical = logical_cpus(); |
| let count = ::std::cmp::min(quota, logical); |
| |
| CGROUPS_CPUS.store(count, Ordering::SeqCst); |
| } |
| None => return, |
| } |
| } |
| |
| fn load_cgroups<P1, P2>(cgroup_proc: P1, mountinfo_proc: P2) -> Option<usize> |
| where |
| P1: AsRef<Path>, |
| P2: AsRef<Path>, |
| { |
| let subsys = some!(Subsys::load_cpu(cgroup_proc)); |
| let mntinfo = some!(MountInfo::load_cpu(mountinfo_proc)); |
| let cgroup = some!(Cgroup::translate(mntinfo, subsys)); |
| cgroup.cpu_quota() |
| } |
| |
| struct Cgroup { |
| base: PathBuf, |
| } |
| |
| struct MountInfo { |
| root: String, |
| mount_point: String, |
| } |
| |
| struct Subsys { |
| base: String, |
| } |
| |
| impl Cgroup { |
| fn new(dir: PathBuf) -> Cgroup { |
| Cgroup { |
| base: dir, |
| } |
| } |
| |
| fn translate(mntinfo: MountInfo, subsys: Subsys) -> Option<Cgroup> { |
| // Translate the subsystem directory via the host paths. |
| debug!( |
| "subsys = {:?}; root = {:?}; mount_point = {:?}", |
| subsys.base, |
| mntinfo.root, |
| mntinfo.mount_point |
| ); |
| |
| let rel_from_root = some!(Path::new(&subsys.base).strip_prefix(&mntinfo.root).ok()); |
| |
| debug!("rel_from_root: {:?}", rel_from_root); |
| |
| // join(mp.MountPoint, relPath) |
| let mut path = PathBuf::from(mntinfo.mount_point); |
| path.push(rel_from_root); |
| Some(Cgroup::new(path)) |
| } |
| |
| fn cpu_quota(&self) -> Option<usize> { |
| let quota_us = some!(self.quota_us()); |
| let period_us = some!(self.period_us()); |
| |
| // protect against dividing by zero |
| if period_us == 0 { |
| return None; |
| } |
| |
| // Ceil the division, since we want to be able to saturate |
| // the available CPUs, and flooring would leave a CPU un-utilized. |
| |
| Some((quota_us as f64 / period_us as f64).ceil() as usize) |
| } |
| |
| fn quota_us(&self) -> Option<usize> { |
| self.param("cpu.cfs_quota_us") |
| } |
| |
| fn period_us(&self) -> Option<usize> { |
| self.param("cpu.cfs_period_us") |
| } |
| |
| fn param(&self, param: &str) -> Option<usize> { |
| let mut file = some!(File::open(self.base.join(param)).ok()); |
| |
| let mut buf = String::new(); |
| some!(file.read_to_string(&mut buf).ok()); |
| |
| buf.trim().parse().ok() |
| } |
| } |
| |
| impl MountInfo { |
| fn load_cpu<P: AsRef<Path>>(proc_path: P) -> Option<MountInfo> { |
| let file = some!(File::open(proc_path).ok()); |
| let file = BufReader::new(file); |
| |
| file.lines() |
| .filter_map(|result| result.ok()) |
| .filter_map(MountInfo::parse_line) |
| .next() |
| } |
| |
| fn parse_line(line: String) -> Option<MountInfo> { |
| let mut fields = line.split(' '); |
| |
| let mnt_root = some!(fields.nth(3)); |
| let mnt_point = some!(fields.nth(0)); |
| |
| if fields.nth(3) != Some("cgroup") { |
| return None; |
| } |
| |
| let super_opts = some!(fields.nth(1)); |
| |
| // We only care about the 'cpu' option |
| if !super_opts.split(',').any(|opt| opt == "cpu") { |
| return None; |
| } |
| |
| Some(MountInfo { |
| root: mnt_root.to_owned(), |
| mount_point: mnt_point.to_owned(), |
| }) |
| } |
| } |
| |
| impl Subsys { |
| fn load_cpu<P: AsRef<Path>>(proc_path: P) -> Option<Subsys> { |
| let file = some!(File::open(proc_path).ok()); |
| let file = BufReader::new(file); |
| |
| file.lines() |
| .filter_map(|result| result.ok()) |
| .filter_map(Subsys::parse_line) |
| .next() |
| } |
| |
| fn parse_line(line: String) -> Option<Subsys> { |
| // Example format: |
| // 11:cpu,cpuacct:/ |
| let mut fields = line.split(':'); |
| |
| let sub_systems = some!(fields.nth(1)); |
| |
| if !sub_systems.split(',').any(|sub| sub == "cpu") { |
| return None; |
| } |
| |
| fields.next().map(|path| Subsys { base: path.to_owned() }) |
| } |
| } |
| |
| #[cfg(test)] |
| mod tests { |
| use std::path::{Path, PathBuf}; |
| use super::{Cgroup, MountInfo, Subsys}; |
| |
| |
| static FIXTURES_PROC: &'static str = "fixtures/cgroups/proc/cgroups"; |
| |
| static FIXTURES_CGROUPS: &'static str = "fixtures/cgroups/cgroups"; |
| |
| macro_rules! join { |
| ($base:expr, $($path:expr),+) => ({ |
| Path::new($base) |
| $(.join($path))+ |
| }) |
| } |
| |
| #[test] |
| fn test_load_mountinfo() { |
| let path = join!(FIXTURES_PROC, "mountinfo"); |
| |
| let mnt_info = MountInfo::load_cpu(path).unwrap(); |
| |
| assert_eq!(mnt_info.root, "/"); |
| assert_eq!(mnt_info.mount_point, "/sys/fs/cgroup/cpu,cpuacct"); |
| } |
| |
| #[test] |
| fn test_load_subsys() { |
| let path = join!(FIXTURES_PROC, "cgroup"); |
| |
| let subsys = Subsys::load_cpu(path).unwrap(); |
| |
| assert_eq!(subsys.base, "/"); |
| } |
| |
| #[test] |
| fn test_cgroup_mount() { |
| let cases = &[ |
| ( |
| "/", |
| "/sys/fs/cgroup/cpu", |
| "/", |
| Some("/sys/fs/cgroup/cpu"), |
| ), |
| ( |
| "/docker/01abcd", |
| "/sys/fs/cgroup/cpu", |
| "/docker/01abcd", |
| Some("/sys/fs/cgroup/cpu"), |
| ), |
| ( |
| "/docker/01abcd", |
| "/sys/fs/cgroup/cpu", |
| "/docker/01abcd/", |
| Some("/sys/fs/cgroup/cpu"), |
| ), |
| ( |
| "/docker/01abcd", |
| "/sys/fs/cgroup/cpu", |
| "/docker/01abcd/large", |
| Some("/sys/fs/cgroup/cpu/large"), |
| ), |
| |
| // fails |
| |
| ( |
| "/docker/01abcd", |
| "/sys/fs/cgroup/cpu", |
| "/", |
| None, |
| ), |
| ( |
| "/docker/01abcd", |
| "/sys/fs/cgroup/cpu", |
| "/docker", |
| None, |
| ), |
| ( |
| "/docker/01abcd", |
| "/sys/fs/cgroup/cpu", |
| "/elsewhere", |
| None, |
| ), |
| ( |
| "/docker/01abcd", |
| "/sys/fs/cgroup/cpu", |
| "/docker/01abcd-other-dir", |
| None, |
| ), |
| ]; |
| |
| for &(root, mount_point, subsys, expected) in cases.iter() { |
| let mnt_info = MountInfo { |
| root: root.into(), |
| mount_point: mount_point.into(), |
| }; |
| let subsys = Subsys { |
| base: subsys.into(), |
| }; |
| |
| let actual = Cgroup::translate(mnt_info, subsys).map(|c| c.base); |
| let expected = expected.map(|s| PathBuf::from(s)); |
| assert_eq!(actual, expected); |
| } |
| } |
| |
| #[test] |
| fn test_cgroup_cpu_quota() { |
| let cgroup = Cgroup::new(join!(FIXTURES_CGROUPS, "good")); |
| assert_eq!(cgroup.cpu_quota(), Some(6)); |
| } |
| |
| #[test] |
| fn test_cgroup_cpu_quota_divide_by_zero() { |
| let cgroup = Cgroup::new(join!(FIXTURES_CGROUPS, "zero-period")); |
| assert!(cgroup.quota_us().is_some()); |
| assert_eq!(cgroup.period_us(), Some(0)); |
| assert_eq!(cgroup.cpu_quota(), None); |
| } |
| |
| #[test] |
| fn test_cgroup_cpu_quota_ceil() { |
| let cgroup = Cgroup::new(join!(FIXTURES_CGROUPS, "ceil")); |
| assert_eq!(cgroup.cpu_quota(), Some(2)); |
| } |
| } |