RGBCube/uutils-coreutils
1
Fork 0
mirror of https://github.com/RGBCube/uutils-coreutils synced 2025-07-29 12:07:46 +00:00
Code Issues Activity

nproc: Improve the support of OMP_NUM_THREADS

Browse source
This commit is contained in:
Sylvestre Ledru 2022-03-13 18:13:48 +01:00
parent 2c6bbcf716
commit 24b4af768c
2 changed files with 32 additions and 12 deletions
Download patch file Download diff file Expand all files Collapse all files

20
src/uu/nproc/src/nproc.rs
View file

@ -32,7 +32,7 @@ const USAGE: &str = "{} [OPTIONS]...";
pub fn uumain(args: impl uucore::Args) -> UResult<()> {
let matches = uu_app().get_matches_from(args);
let mut ignore = match matches.value_of(OPT_IGNORE) {
let ignore = match matches.value_of(OPT_IGNORE) {
Some(numstr) => match numstr.trim().parse() {
Ok(num) => num,
Err(e) => {
@ -45,18 +45,18 @@ pub fn uumain(args: impl uucore::Args) -> UResult<()> {
None => 0,
};
if !matches.is_present(OPT_ALL) {
// OMP_NUM_THREADS doesn't have an impact on --all
ignore += match env::var("OMP_NUM_THREADS") {
Ok(threadstr) => threadstr.parse().unwrap_or(0),
Err(_) => 0,
};
}
let mut cores = if matches.is_present(OPT_ALL) {
num_cpus_all()
} else {
num_cpus::get()
// OMP_NUM_THREADS doesn't have an impact on --all
match env::var("OMP_NUM_THREADS") {
// 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
// fallback to the regular CPU detection
Err(_) => num_cpus::get(),
}
};
if cores <= ignore {

24
tests/by-util/test_nproc.rs
View file

@ -1,3 +1,4 @@
// spell-checker:ignore incorrectnumber
use crate::common::util::*;
#[test]
@ -15,11 +16,11 @@ fn test_nproc_all_omp() {
let result = TestScenario::new(util_name!())
.ucmd_keepenv()
.env("OMP_NUM_THREADS", "1")
.env("OMP_NUM_THREADS", "60")
.succeeds();
let nproc_omp: u8 = result.stdout_str().trim().parse().unwrap();
assert!(nproc - 1 == nproc_omp);
assert!(nproc_omp == 60);
let result = TestScenario::new(util_name!())
.ucmd_keepenv()
@ -28,6 +29,14 @@ fn test_nproc_all_omp() {
.succeeds();
let nproc_omp: u8 = result.stdout_str().trim().parse().unwrap();
assert!(nproc == nproc_omp);
// If the parsing fails, returns the number of CPU
let result = TestScenario::new(util_name!())
.ucmd_keepenv()
.env("OMP_NUM_THREADS", "incorrectnumber") // returns the number CPU
.succeeds();
let nproc_omp: u8 = result.stdout_str().trim().parse().unwrap();
assert!(nproc == nproc_omp);
}
#[test]
@ -52,3 +61,14 @@ fn test_nproc_ignore() {
assert!(nproc_total - 1 == nproc);
}
}
#[test]
fn test_nproc_ignore_all_omp() {
let result = TestScenario::new(util_name!())
.ucmd_keepenv()
.env("OMP_NUM_THREADS", "42")
.arg("--ignore=40")
.succeeds();
let nproc: u8 = result.stdout_str().trim().parse().unwrap();
assert!(nproc == 2);
}