diff --git a/src/nproc/nproc.rs b/src/nproc/nproc.rs index 1212f6fa6..881e48159 100644 --- a/src/nproc/nproc.rs +++ b/src/nproc/nproc.rs @@ -9,6 +9,8 @@ * file that was distributed with this source code. */ +#![feature(macro_rules)] + extern crate getopts; use std::os; @@ -16,15 +18,23 @@ use std::os; static NAME : &'static str = "nproc"; static VERSION : &'static str = "0.0.0"; +#[path = "../common/util.rs"] +mod util; + pub fn uumain(args: Vec) -> int { let opts = [ + getopts::optflag("", "all", "print the number of cores available to the system"), + getopts::optopt("", "ignore", "ignore up to N cores", "N"), getopts::optflag("h", "help", "display this help and exit"), getopts::optflag("V", "version", "output version information and exit"), ]; let matches = match getopts::getopts(args.tail(), opts) { Ok(m) => m, - Err(err) => fail!("{}", err), + Err(err) => { + show_error!("{}", err); + return 1; + } }; if matches.opt_present("version") { @@ -38,14 +48,38 @@ pub fn uumain(args: Vec) -> int { println!("Usage:"); println!(" {} [OPTIONS] NAME...", NAME); println!(""); - print!("{}", getopts::usage("Print the number of cores available.", opts.as_slice()).as_slice()); - if matches.free.is_empty() { - return 1; - } + print!("{}", getopts::usage("Print the number of cores available to the current process.", opts)); return 0; } - println!("{}", os::num_cpus()); + let mut ignore = match matches.opt_str("ignore") { + Some(numstr) => match from_str(numstr.as_slice()) { + Some(num) => num, + None => { + show_error!("\"{}\" is not a valid number", numstr); + return 1; + } + }, + None => 0 + }; + + if !matches.opt_present("all") { + ignore += match os::getenv("OMP_NUM_THREADS") { + Some(threadstr) => match from_str(threadstr.as_slice()) { + Some(num) => num, + None => 0 + }, + None => 0 + }; + } + + let mut cores = os::num_cpus(); + if cores <= ignore { + cores = 1; + } else { + cores -= ignore; + } + println!("{}", cores); return 0 }