RGBCube/uutils-coreutils
1
Fork 0
mirror of https://github.com/RGBCube/uutils-coreutils synced 2025-07-27 11:07:44 +00:00
Code Issues Activity

Merge pull request #3286 from sylvestre/nproc-2

Browse source 
nproc: add the full support of OMP_THREAD_LIMIT
This commit is contained in:
Terts Diepraam 2022-03-21 18:11:41 +01:00 • committed by GitHub
commit e9a6cf043f
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 48 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -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 {

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

@ -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);
}