diff options
Diffstat (limited to 'vendor/num_cpus/src/linux.rs')
| -rw-r--r-- | vendor/num_cpus/src/linux.rs | 418 |
1 files changed, 418 insertions, 0 deletions
diff --git a/vendor/num_cpus/src/linux.rs b/vendor/num_cpus/src/linux.rs new file mode 100644 index 000000000..36f472717 --- /dev/null +++ b/vendor/num_cpus/src/linux.rs @@ -0,0 +1,418 @@ +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); + + if let Some(quota) = load_cgroups("/proc/self/cgroup", "/proc/self/mountinfo") { + if quota == 0 { + return; + } + + let logical = logical_cpus(); + let count = ::std::cmp::min(quota, logical); + + CGROUPS_CPUS.store(count, Ordering::SeqCst); + } +} + +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(' '); + + // 7 5 0:6 </> /sys/fs/cgroup/cpu,cpuacct rw,nosuid,nodev,noexec,relatime shared:7 - cgroup cgroup rw,cpu,cpuacct + let mnt_root = some!(fields.nth(3)); + // 7 5 0:6 / </sys/fs/cgroup/cpu,cpuacct> rw,nosuid,nodev,noexec,relatime shared:7 - cgroup cgroup rw,cpu,cpuacct + let mnt_point = some!(fields.next()); + + // Ignore all fields until the separator(-). + // Note: there could be zero or more optional fields before hyphen. + // See: https://man7.org/linux/man-pages/man5/proc.5.html + // 7 5 0:6 / /sys/fs/cgroup/cpu,cpuacct rw,nosuid,nodev,noexec,relatime shared:7 <-> cgroup cgroup rw,cpu,cpuacct + // Note: we cannot use `?` here because we need to support Rust 1.13. + match fields.find(|&s| s == "-") { + Some(_) => {} + None => return None, + }; + + // 7 5 0:6 / /sys/fs/cgroup/cpu,cpuacct rw,nosuid,nodev,noexec,relatime shared:7 - <cgroup> cgroup rw,cpu,cpuacct + if fields.next() != Some("cgroup") { + return None; + } + + // 7 5 0:6 / /sys/fs/cgroup/cpu,cpuacct rw,nosuid,nodev,noexec,relatime shared:7 - cgroup cgroup <rw,cpu,cpuacct> + 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 super::{Cgroup, MountInfo, Subsys}; + use std::path::{Path, PathBuf}; + + // `static_in_const` feature is not stable in Rust 1.13. + 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() { + // test only one optional fields + 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 zero optional field + let path = join!(FIXTURES_PROC, "mountinfo_zero_opt"); + + 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 multi optional fields + let path = join!(FIXTURES_PROC, "mountinfo_multi_opt"); + + 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(PathBuf::from); + 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)); + } +} |