keygen: Add num-threads argument for grind (#14884)

* keygen: Add num-threads argument

* Review feedback

keygen/src/keygen.rs

@@ -1,9 +1,10 @@
 use bip39::{Language, Mnemonic, MnemonicType, Seed};
 use clap::{
-    crate_description, crate_name, value_t, values_t_or_exit, App, AppSettings, Arg, ArgMatches,
+    crate_description, crate_name, value_t, value_t_or_exit, values_t_or_exit, App, AppSettings,
-    SubCommand,
+    Arg, ArgMatches, SubCommand,
 };
 use solana_clap_utils::{
+    input_validators::is_parsable,
     keypair::{
         keypair_from_seed_phrase, prompt_passphrase, signer_from_path,
         SKIP_SEED_PHRASE_VALIDATION_ARG,
@@ -129,8 +130,8 @@ fn grind_validator_starts_and_ends_with(v: String) -> Result<(), String> {
     Ok(())
 }
 
-fn grind_print_info(grind_matches: &[GrindMatch]) {
+fn grind_print_info(grind_matches: &[GrindMatch], num_threads: usize) {
-    println!("Searching with {} threads for:", num_cpus::get());
+    println!("Searching with {} threads for:", num_threads);
     for gm in grind_matches {
         let mut msg = Vec::<String>::new();
         if gm.count.load(Ordering::Relaxed) > 1 {
@@ -159,6 +160,7 @@ fn grind_parse_args(
     starts_with_args: HashSet<String>,
     ends_with_args: HashSet<String>,
     starts_and_ends_with_args: HashSet<String>,
+    num_threads: usize,
 ) -> Vec<GrindMatch> {
     let mut grind_matches = Vec::<GrindMatch>::new();
     for sw in starts_with_args {
@@ -201,11 +203,12 @@ fn grind_parse_args(
             count: AtomicU64::new(args[2].parse::<u64>().unwrap()),
         });
     }
-    grind_print_info(&grind_matches);
+    grind_print_info(&grind_matches, num_threads);
     grind_matches
 }
 
 fn main() -> Result<(), Box<dyn error::Error>> {
+    let default_num_threads = num_cpus::get().to_string();
     let matches = App::new(crate_name!())
         .about(crate_description!())
         .version(solana_version::version!())
@@ -335,6 +338,15 @@ fn main() -> Result<(), Box<dyn error::Error>> {
                         .multiple(true)
                         .validator(grind_validator_starts_and_ends_with)
                         .help("Saves specified number of keypairs whos public key starts and ends with the indicated perfix and suffix\nExample: --starts-and-ends-with sol:ana:4\nPREFIX and SUFFIX type is Base58\nCOUNT type is u64"),
+                )
+                .arg(
+                    Arg::with_name("num_threads")
+                        .long("num-threads")
+                        .value_name("NUMBER")
+                        .takes_value(true)
+                        .validator(is_parsable::<usize>)
+                        .default_value(&default_num_threads)
+                        .help("Specify the number of grind threads"),
                 ),
         )
         .subcommand(
@@ -532,11 +544,14 @@ fn do_main(matches: &ArgMatches<'_>) -> Result<(), Box<dyn error::Error>> {
                 exit(1);
             }
 
+            let num_threads = value_t_or_exit!(matches.value_of("num_threads"), usize);
+
             let grind_matches = grind_parse_args(
                 ignore_case,
                 starts_with_args,
                 ends_with_args,
                 starts_and_ends_with_args,
+                num_threads,
             );
 
             let grind_matches_thread_safe = Arc::new(grind_matches);
@@ -545,7 +560,7 @@ fn do_main(matches: &ArgMatches<'_>) -> Result<(), Box<dyn error::Error>> {
             let start = Instant::now();
             let done = Arc::new(AtomicBool::new(false));
 
-            let thread_handles: Vec<_> = (0..num_cpus::get())
+            let thread_handles: Vec<_> = (0..num_threads)
                 .map(|_| {
                     let done = done.clone();
                     let attempts = attempts.clone();