vendor/num_cpus/src/linux.rs

   1 use std::collections::HashMap;
   2 use std::fs::File;
   3 use std::io::{BufRead, BufReader, Read};
   4 use std::mem;
   5 use std::path::{Path, PathBuf};
   6 use std::sync::atomic::{AtomicUsize, Ordering};
   7 use std::sync::Once;
   8
   9 use libc;
  10
  11 macro_rules! debug {
  12     ($($args:expr),*) => ({
  13         if false {
  14         //if true {
  15             println!($($args),*);
  16         }
  17     });
  18 }
  19
  20 macro_rules! some {
  21     ($e:expr) => {{
  22         match $e {
  23             Some(v) => v,
  24             None => {
  25                 debug!("NONE: {:?}", stringify!($e));
  26                 return None;
  27             }
  28         }
  29     }};
  30 }
  31
  32 pub fn get_num_cpus() -> usize {
  33     match cgroups_num_cpus() {
  34         Some(n) => n,
  35         None => logical_cpus(),
  36     }
  37 }
  38
  39 fn logical_cpus() -> usize {
  40     let mut set: libc::cpu_set_t = unsafe { mem::zeroed() };
  41     if unsafe { libc::sched_getaffinity(0, mem::size_of::<libc::cpu_set_t>(), &mut set) } == 0 {
  42         let mut count: u32 = 0;
  43         for i in 0..libc::CPU_SETSIZE as usize {
  44             if unsafe { libc::CPU_ISSET(i, &set) } {
  45                 count += 1
  46             }
  47         }
  48         count as usize
  49     } else {
  50         let cpus = unsafe { libc::sysconf(libc::_SC_NPROCESSORS_ONLN) };
  51         if cpus < 1 {
  52             1
  53         } else {
  54             cpus as usize
  55         }
  56     }
  57 }
  58
  59 pub fn get_num_physical_cpus() -> usize {
  60     let file = match File::open("/proc/cpuinfo") {
  61         Ok(val) => val,
  62         Err(_) => return get_num_cpus(),
  63     };
  64     let reader = BufReader::new(file);
  65     let mut map = HashMap::new();
  66     let mut physid: u32 = 0;
  67     let mut cores: usize = 0;
  68     let mut chgcount = 0;
  69     for line in reader.lines().filter_map(|result| result.ok()) {
  70         let mut it = line.split(':');
  71         let (key, value) = match (it.next(), it.next()) {
  72             (Some(key), Some(value)) => (key.trim(), value.trim()),
  73             _ => continue,
  74         };
  75         if key == "physical id" {
  76             match value.parse() {
  77                 Ok(val) => physid = val,
  78                 Err(_) => break,
  79             };
  80             chgcount += 1;
  81         }
  82         if key == "cpu cores" {
  83             match value.parse() {
  84                 Ok(val) => cores = val,
  85                 Err(_) => break,
  86             };
  87             chgcount += 1;
  88         }
  89         if chgcount == 2 {
  90             map.insert(physid, cores);
  91             chgcount = 0;
  92         }
  93     }
  94     let count = map.into_iter().fold(0, |acc, (_, cores)| acc + cores);
  95
  96     if count == 0 {
  97         get_num_cpus()
  98     } else {
  99         count
 100     }
 101 }
 102
 103 /// Cached CPUs calculated from cgroups.
 104 ///
 105 /// If 0, check logical cpus.
 106 // Allow deprecation warnings, we want to work on older rustc
 107 #[allow(warnings)]
 108 static CGROUPS_CPUS: AtomicUsize = ::std::sync::atomic::ATOMIC_USIZE_INIT;
 109
 110 fn cgroups_num_cpus() -> Option<usize> {
 111     #[allow(warnings)]
 112     static ONCE: Once = ::std::sync::ONCE_INIT;
 113
 114     ONCE.call_once(init_cgroups);
 115
 116     let cpus = CGROUPS_CPUS.load(Ordering::Acquire);
 117
 118     if cpus > 0 {
 119         Some(cpus)
 120     } else {
 121         None
 122     }
 123 }
 124
 125 fn init_cgroups() {
 126     // Should only be called once
 127     debug_assert!(CGROUPS_CPUS.load(Ordering::SeqCst) == 0);
 128
 129     if let Some(quota) = load_cgroups("/proc/self/cgroup", "/proc/self/mountinfo") {
 130         if quota == 0 {
 131             return;
 132         }
 133
 134         let logical = logical_cpus();
 135         let count = ::std::cmp::min(quota, logical);
 136
 137         CGROUPS_CPUS.store(count, Ordering::SeqCst);
 138     }
 139 }
 140
 141 fn load_cgroups<P1, P2>(cgroup_proc: P1, mountinfo_proc: P2) -> Option<usize>
 142 where
 143     P1: AsRef<Path>,
 144     P2: AsRef<Path>,
 145 {
 146     let subsys = some!(Subsys::load_cpu(cgroup_proc));
 147     let mntinfo = some!(MountInfo::load_cpu(mountinfo_proc));
 148     let cgroup = some!(Cgroup::translate(mntinfo, subsys));
 149     cgroup.cpu_quota()
 150 }
 151
 152 struct Cgroup {
 153     base: PathBuf,
 154 }
 155
 156 struct MountInfo {
 157     root: String,
 158     mount_point: String,
 159 }
 160
 161 struct Subsys {
 162     base: String,
 163 }
 164
 165 impl Cgroup {
 166     fn new(dir: PathBuf) -> Cgroup {
 167         Cgroup { base: dir }
 168     }
 169
 170     fn translate(mntinfo: MountInfo, subsys: Subsys) -> Option<Cgroup> {
 171         // Translate the subsystem directory via the host paths.
 172         debug!(
 173             "subsys = {:?}; root = {:?}; mount_point = {:?}",
 174             subsys.base, mntinfo.root, mntinfo.mount_point
 175         );
 176
 177         let rel_from_root = some!(Path::new(&subsys.base).strip_prefix(&mntinfo.root).ok());
 178
 179         debug!("rel_from_root: {:?}", rel_from_root);
 180
 181         // join(mp.MountPoint, relPath)
 182         let mut path = PathBuf::from(mntinfo.mount_point);
 183         path.push(rel_from_root);
 184         Some(Cgroup::new(path))
 185     }
 186
 187     fn cpu_quota(&self) -> Option<usize> {
 188         let quota_us = some!(self.quota_us());
 189         let period_us = some!(self.period_us());
 190
 191         // protect against dividing by zero
 192         if period_us == 0 {
 193             return None;
 194         }
 195
 196         // Ceil the division, since we want to be able to saturate
 197         // the available CPUs, and flooring would leave a CPU un-utilized.
 198
 199         Some((quota_us as f64 / period_us as f64).ceil() as usize)
 200     }
 201
 202     fn quota_us(&self) -> Option<usize> {
 203         self.param("cpu.cfs_quota_us")
 204     }
 205
 206     fn period_us(&self) -> Option<usize> {
 207         self.param("cpu.cfs_period_us")
 208     }
 209
 210     fn param(&self, param: &str) -> Option<usize> {
 211         let mut file = some!(File::open(self.base.join(param)).ok());
 212
 213         let mut buf = String::new();
 214         some!(file.read_to_string(&mut buf).ok());
 215
 216         buf.trim().parse().ok()
 217     }
 218 }
 219
 220 impl MountInfo {
 221     fn load_cpu<P: AsRef<Path>>(proc_path: P) -> Option<MountInfo> {
 222         let file = some!(File::open(proc_path).ok());
 223         let file = BufReader::new(file);
 224
 225         file.lines()
 226             .filter_map(|result| result.ok())
 227             .filter_map(MountInfo::parse_line)
 228             .next()
 229     }
 230
 231     fn parse_line(line: String) -> Option<MountInfo> {
 232         let mut fields = line.split(' ');
 233
 234         // 7 5 0:6 </> /sys/fs/cgroup/cpu,cpuacct rw,nosuid,nodev,noexec,relatime shared:7 - cgroup cgroup rw,cpu,cpuacct
 235         let mnt_root = some!(fields.nth(3));
 236         // 7 5 0:6 / </sys/fs/cgroup/cpu,cpuacct> rw,nosuid,nodev,noexec,relatime shared:7 - cgroup cgroup rw,cpu,cpuacct
 237         let mnt_point = some!(fields.next());
 238
 239         // Ignore all fields until the separator(-).
 240         // Note: there could be zero or more optional fields before hyphen.
 241         // See: https://man7.org/linux/man-pages/man5/proc.5.html
 242         // 7 5 0:6 / /sys/fs/cgroup/cpu,cpuacct rw,nosuid,nodev,noexec,relatime shared:7 <-> cgroup cgroup rw,cpu,cpuacct
 243         // Note: we cannot use `?` here because we need to support Rust 1.13.
 244         match fields.find(|&s| s == "-") {
 245             Some(_) => {}
 246             None => return None,
 247         };
 248
 249         // 7 5 0:6 / /sys/fs/cgroup/cpu,cpuacct rw,nosuid,nodev,noexec,relatime shared:7 - <cgroup> cgroup rw,cpu,cpuacct
 250         if fields.next() != Some("cgroup") {
 251             return None;
 252         }
 253
 254         // 7 5 0:6 / /sys/fs/cgroup/cpu,cpuacct rw,nosuid,nodev,noexec,relatime shared:7 - cgroup cgroup <rw,cpu,cpuacct>
 255         let super_opts = some!(fields.nth(1));
 256
 257         // We only care about the 'cpu' option
 258         if !super_opts.split(',').any(|opt| opt == "cpu") {
 259             return None;
 260         }
 261
 262         Some(MountInfo {
 263             root: mnt_root.to_owned(),
 264             mount_point: mnt_point.to_owned(),
 265         })
 266     }
 267 }
 268
 269 impl Subsys {
 270     fn load_cpu<P: AsRef<Path>>(proc_path: P) -> Option<Subsys> {
 271         let file = some!(File::open(proc_path).ok());
 272         let file = BufReader::new(file);
 273
 274         file.lines()
 275             .filter_map(|result| result.ok())
 276             .filter_map(Subsys::parse_line)
 277             .next()
 278     }
 279
 280     fn parse_line(line: String) -> Option<Subsys> {
 281         // Example format:
 282         // 11:cpu,cpuacct:/
 283         let mut fields = line.split(':');
 284
 285         let sub_systems = some!(fields.nth(1));
 286
 287         if !sub_systems.split(',').any(|sub| sub == "cpu") {
 288             return None;
 289         }
 290
 291         fields.next().map(|path| Subsys {
 292             base: path.to_owned(),
 293         })
 294     }
 295 }
 296
 297 #[cfg(test)]
 298 mod tests {
 299     use super::{Cgroup, MountInfo, Subsys};
 300     use std::path::{Path, PathBuf};
 301
 302     // `static_in_const` feature is not stable in Rust 1.13.
 303     static FIXTURES_PROC: &'static str = "fixtures/cgroups/proc/cgroups";
 304
 305     static FIXTURES_CGROUPS: &'static str = "fixtures/cgroups/cgroups";
 306
 307     macro_rules! join {
 308         ($base:expr, $($path:expr),+) => ({
 309             Path::new($base)
 310                 $(.join($path))+
 311         })
 312     }
 313
 314     #[test]
 315     fn test_load_mountinfo() {
 316         // test only one optional fields
 317         let path = join!(FIXTURES_PROC, "mountinfo");
 318
 319         let mnt_info = MountInfo::load_cpu(path).unwrap();
 320
 321         assert_eq!(mnt_info.root, "/");
 322         assert_eq!(mnt_info.mount_point, "/sys/fs/cgroup/cpu,cpuacct");
 323
 324         // test zero optional field
 325         let path = join!(FIXTURES_PROC, "mountinfo_zero_opt");
 326
 327         let mnt_info = MountInfo::load_cpu(path).unwrap();
 328
 329         assert_eq!(mnt_info.root, "/");
 330         assert_eq!(mnt_info.mount_point, "/sys/fs/cgroup/cpu,cpuacct");
 331
 332         // test multi optional fields
 333         let path = join!(FIXTURES_PROC, "mountinfo_multi_opt");
 334
 335         let mnt_info = MountInfo::load_cpu(path).unwrap();
 336
 337         assert_eq!(mnt_info.root, "/");
 338         assert_eq!(mnt_info.mount_point, "/sys/fs/cgroup/cpu,cpuacct");
 339     }
 340
 341     #[test]
 342     fn test_load_subsys() {
 343         let path = join!(FIXTURES_PROC, "cgroup");
 344
 345         let subsys = Subsys::load_cpu(path).unwrap();
 346
 347         assert_eq!(subsys.base, "/");
 348     }
 349
 350     #[test]
 351     fn test_cgroup_mount() {
 352         let cases = &[
 353             ("/", "/sys/fs/cgroup/cpu", "/", Some("/sys/fs/cgroup/cpu")),
 354             (
 355                 "/docker/01abcd",
 356                 "/sys/fs/cgroup/cpu",
 357                 "/docker/01abcd",
 358                 Some("/sys/fs/cgroup/cpu"),
 359             ),
 360             (
 361                 "/docker/01abcd",
 362                 "/sys/fs/cgroup/cpu",
 363                 "/docker/01abcd/",
 364                 Some("/sys/fs/cgroup/cpu"),
 365             ),
 366             (
 367                 "/docker/01abcd",
 368                 "/sys/fs/cgroup/cpu",
 369                 "/docker/01abcd/large",
 370                 Some("/sys/fs/cgroup/cpu/large"),
 371             ),
 372             // fails
 373             ("/docker/01abcd", "/sys/fs/cgroup/cpu", "/", None),
 374             ("/docker/01abcd", "/sys/fs/cgroup/cpu", "/docker", None),
 375             ("/docker/01abcd", "/sys/fs/cgroup/cpu", "/elsewhere", None),
 376             (
 377                 "/docker/01abcd",
 378                 "/sys/fs/cgroup/cpu",
 379                 "/docker/01abcd-other-dir",
 380                 None,
 381             ),
 382         ];
 383
 384         for &(root, mount_point, subsys, expected) in cases.iter() {
 385             let mnt_info = MountInfo {
 386                 root: root.into(),
 387                 mount_point: mount_point.into(),
 388             };
 389             let subsys = Subsys {
 390                 base: subsys.into(),
 391             };
 392
 393             let actual = Cgroup::translate(mnt_info, subsys).map(|c| c.base);
 394             let expected = expected.map(PathBuf::from);
 395             assert_eq!(actual, expected);
 396         }
 397     }
 398
 399     #[test]
 400     fn test_cgroup_cpu_quota() {
 401         let cgroup = Cgroup::new(join!(FIXTURES_CGROUPS, "good"));
 402         assert_eq!(cgroup.cpu_quota(), Some(6));
 403     }
 404
 405     #[test]
 406     fn test_cgroup_cpu_quota_divide_by_zero() {
 407         let cgroup = Cgroup::new(join!(FIXTURES_CGROUPS, "zero-period"));
 408         assert!(cgroup.quota_us().is_some());
 409         assert_eq!(cgroup.period_us(), Some(0));
 410         assert_eq!(cgroup.cpu_quota(), None);
 411     }
 412
 413     #[test]
 414     fn test_cgroup_cpu_quota_ceil() {
 415         let cgroup = Cgroup::new(join!(FIXTURES_CGROUPS, "ceil"));
 416         assert_eq!(cgroup.cpu_quota(), Some(2));
 417     }
 418 }