Merge pull request #3286 from sylvestre/nproc-2

nproc: add the full support of OMP_THREAD_LIMIT

src/uu/nproc/src/nproc.rs

@@ -25,7 +25,9 @@ pub const _SC_NPROCESSORS_CONF: libc::c_int = 1001;
 static OPT_ALL: &str = "all";
 static OPT_IGNORE: &str = "ignore";
 
-static ABOUT: &str = "Print the number of cores available to the current process.";
+static ABOUT: &str = r#"Print the number of cores available to the current process.
+If the OMP_NUM_THREADS or OMP_THREAD_LIMIT environment variables are set, then
+they will determine the minimum and maximum returned value respectively."#;
 const USAGE: &str = "{} [OPTIONS]...";
 
 #[uucore::main]
@@ -45,6 +47,15 @@ pub fn uumain(args: impl uucore::Args) -> UResult<()> {
         None => 0,
     };
 
+    let limit = match env::var("OMP_THREAD_LIMIT") {
+        // Uses the OpenMP variable to limit the number of threads
+        // If the parsing fails, returns the max size (so, no impact)
+        Ok(threadstr) => threadstr.parse().unwrap_or(usize::MAX),
+        // the variable 'OMP_THREAD_LIMIT' doesn't exist
+        // fallback to the max
+        Err(_) => usize::MAX,
+    };
+
     let mut cores = if matches.is_present(OPT_ALL) {
         num_cpus_all()
     } else {
@@ -53,12 +64,12 @@ pub fn uumain(args: impl uucore::Args) -> UResult<()> {
             // Uses the OpenMP variable to force the number of threads
             // If the parsing fails, returns the number of CPU
             Ok(threadstr) => threadstr.parse().unwrap_or_else(|_| num_cpus::get()),
-            // the variable 'OMP_NUM_THREADS' doesn't exit
+            // the variable 'OMP_NUM_THREADS' doesn't exist
             // fallback to the regular CPU detection
             Err(_) => num_cpus::get(),
         }
     };
-
+    cores = std::cmp::min(limit, cores);
     if cores <= ignore {
         cores = 1;
     } else {

tests/by-util/test_nproc.rs

@@ -72,3 +72,37 @@ fn test_nproc_ignore_all_omp() {
     let nproc: u8 = result.stdout_str().trim().parse().unwrap();
     assert!(nproc == 2);
 }
+
+#[test]
+fn test_nproc_omp_limit() {
+    let result = TestScenario::new(util_name!())
+        .ucmd_keepenv()
+        .env("OMP_NUM_THREADS", "42")
+        .env("OMP_THREAD_LIMIT", "0")
+        .succeeds();
+    let nproc: u8 = result.stdout_str().trim().parse().unwrap();
+    assert!(nproc == 1);
+
+    let result = TestScenario::new(util_name!())
+        .ucmd_keepenv()
+        .env("OMP_NUM_THREADS", "42")
+        .env("OMP_THREAD_LIMIT", "2")
+        .succeeds();
+    let nproc: u8 = result.stdout_str().trim().parse().unwrap();
+    assert!(nproc == 2);
+
+    let result = TestScenario::new(util_name!())
+        .ucmd_keepenv()
+        .env("OMP_NUM_THREADS", "42")
+        .env("OMP_THREAD_LIMIT", "2bad")
+        .succeeds();
+    let nproc: u8 = result.stdout_str().trim().parse().unwrap();
+    assert!(nproc == 42);
+
+    let result = TestScenario::new(util_name!())
+        .ucmd_keepenv()
+        .env("OMP_THREAD_LIMIT", "1")
+        .succeeds();
+    let nproc: u8 = result.stdout_str().trim().parse().unwrap();
+    assert!(nproc == 1);
+}