diff --git a/keygen/src/keygen.rs b/keygen/src/keygen.rs index 0d40834ae9..c9536ea30e 100644 --- a/keygen/src/keygen.rs +++ b/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, - SubCommand, + crate_description, crate_name, value_t, value_t_or_exit, values_t_or_exit, App, AppSettings, + 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]) { - println!("Searching with {} threads for:", num_cpus::get()); +fn grind_print_info(grind_matches: &[GrindMatch], num_threads: usize) { + println!("Searching with {} threads for:", num_threads); for gm in grind_matches { let mut msg = Vec::::new(); if gm.count.load(Ordering::Relaxed) > 1 { @@ -159,6 +160,7 @@ fn grind_parse_args( starts_with_args: HashSet, ends_with_args: HashSet, starts_and_ends_with_args: HashSet, + num_threads: usize, ) -> Vec { let mut grind_matches = Vec::::new(); for sw in starts_with_args { @@ -201,11 +203,12 @@ fn grind_parse_args( count: AtomicU64::new(args[2].parse::().unwrap()), }); } - grind_print_info(&grind_matches); + grind_print_info(&grind_matches, num_threads); grind_matches } fn main() -> Result<(), Box> { + 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> { .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::) + .default_value(&default_num_threads) + .help("Specify the number of grind threads"), ), ) .subcommand( @@ -532,11 +544,14 @@ fn do_main(matches: &ArgMatches<'_>) -> Result<(), Box> { 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> { 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();