Merge pull request #1045 from kubo39/nproc-counts-via-affinity-mask-if-available

nproc: counts CPU cores via affinity mask if available on Linux

Cargo.lock

@@ -679,7 +679,7 @@ version = "0.0.1"
 dependencies = [
  "getopts 0.2.14 (registry+https://github.com/rust-lang/crates.io-index)",
  "libc 0.2.18 (registry+https://github.com/rust-lang/crates.io-index)",
- "num_cpus 1.2.0 (registry+https://github.com/rust-lang/crates.io-index)",
+ "num_cpus 1.5.0 (registry+https://github.com/rust-lang/crates.io-index)",
  "uucore 0.0.1",
 ]
 
@@ -717,7 +717,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 
 [[package]]
 name = "num_cpus"
-version = "1.2.0"
+version = "1.5.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 dependencies = [
  "libc 0.2.18 (registry+https://github.com/rust-lang/crates.io-index)",
@@ -1389,7 +1389,7 @@ dependencies = [
 "checksum num-integer 0.1.33 (registry+https://github.com/rust-lang/crates.io-index)" = "21e4df1098d1d797d27ef0c69c178c3fab64941559b290fcae198e0825c9c8b5"
 "checksum num-iter 0.1.33 (registry+https://github.com/rust-lang/crates.io-index)" = "f7d1891bd7b936f12349b7d1403761c8a0b85a18b148e9da4429d5d102c1a41e"
 "checksum num-traits 0.1.37 (registry+https://github.com/rust-lang/crates.io-index)" = "e1cbfa3781f3fe73dc05321bed52a06d2d491eaa764c52335cf4399f046ece99"
-"checksum num_cpus 1.2.0 (registry+https://github.com/rust-lang/crates.io-index)" = "55aabf4e2d6271a2e4e4c0f2ea1f5b07cc589cc1a9e9213013b54a76678ca4f3"
+"checksum num_cpus 1.5.0 (registry+https://github.com/rust-lang/crates.io-index)" = "f6e850c7f35c3de263e6094e819f6b4b9c09190ff4438fc6dec1aef1568547bc"
 "checksum pretty-bytes 0.2.1 (registry+https://github.com/rust-lang/crates.io-index)" = "3095b93999fae14b4e0bb661c53875a441d9058b7b1a7ba2dfebc104d3776349"
 "checksum quick-error 1.1.0 (registry+https://github.com/rust-lang/crates.io-index)" = "0aad603e8d7fb67da22dbdf1f4b826ce8829e406124109e73cf1b2454b93a71c"
 "checksum rand 0.3.15 (registry+https://github.com/rust-lang/crates.io-index)" = "022e0636ec2519ddae48154b028864bdce4eaf7d35226ab8e65c611be97b189d"

src/nproc/Cargo.toml

@@ -10,7 +10,7 @@ path = "nproc.rs"
 [dependencies]
 getopts = "*"
 libc = "*"
-num_cpus = "*"
+num_cpus = "1.5"
 uucore = { path="../uucore" }
 
 [[bin]]

src/nproc/nproc.rs

@@ -12,12 +12,25 @@
 extern crate getopts;
 extern crate num_cpus;
 
+#[cfg(unix)]
+extern crate libc;
+
 #[macro_use]
 extern crate uucore;
 
 use std::io::Write;
 use std::env;
 
+#[cfg(target_os = "linux")]
+pub const _SC_NPROCESSORS_CONF: libc::c_int = 83;
+#[cfg(target_os = "macos")]
+pub const _SC_NPROCESSORS_CONF: libc::c_int = libc::_SC_NPROCESSORS_CONF;
+#[cfg(target_os = "freebsd")]
+pub const _SC_NPROCESSORS_CONF: libc::c_int = 57;
+#[cfg(target_os = "netbsd")]
+pub const _SC_NPROCESSORS_CONF: libc::c_int = 1001;
+
+
 static NAME: &'static str = "nproc";
 static VERSION: &'static str = env!("CARGO_PKG_VERSION");
 
@@ -75,7 +88,12 @@ Print the number of cores available to the current process.", NAME, VERSION);
         };
     }
 
-    let mut cores = num_cpus::get();
+    let mut cores = if matches.opt_present("all") {
+        num_cpus_all()
+    } else {
+        num_cpus::get()
+    };
+
     if cores <= ignore {
         cores = 1;
     } else {
@@ -84,3 +102,27 @@ Print the number of cores available to the current process.", NAME, VERSION);
     println!("{}", cores);
     0
 }
+
+#[cfg(any(target_os = "linux",
+          target_os = "macos",
+          target_os = "freebsd",
+          target_os = "netbsd"))]
+fn num_cpus_all() -> usize {
+    let nprocs = unsafe { libc::sysconf(_SC_NPROCESSORS_CONF) };
+    if nprocs == 1 {
+        // In some situation, /proc and /sys are not mounted, and sysconf returns 1.
+        // However, we want to guarantee that `nproc --all` >= `nproc`.
+        num_cpus::get()
+    } else {
+        if nprocs > 0 { nprocs as usize } else { 1 }
+    }
+}
+
+// Other platform(e.g., windows), num_cpus::get() directly.
+#[cfg(not(any(target_os = "linux",
+              target_os = "macos",
+              target_os = "freebsd",
+              target_os = "netbsd")))]
+fn num_cpus_all() -> usize {
+    num_cpus::get()
+}