Provision bench client accounts in genesis block (#4648)

* fixes to script

* shellcheck

* address review comments

bench-tps/src/cli.rs

@@ -17,6 +17,9 @@ pub struct Config {
     pub tx_count: usize,
     pub thread_batch_sleep_ms: usize,
     pub sustained: bool,
+    pub client_ids_and_stake_file: String,
+    pub write_to_client_file: bool,
+    pub read_from_client_file: bool,
 }
 
 impl Default for Config {
@@ -31,6 +34,9 @@ impl Default for Config {
             tx_count: 500_000,
             thread_batch_sleep_ms: 0,
             sustained: false,
+            client_ids_and_stake_file: String::new(),
+            write_to_client_file: false,
+            read_from_client_file: false,
         }
     }
 }
@@ -106,6 +112,20 @@ pub fn build_args<'a, 'b>() -> App<'a, 'b> {
                 .takes_value(true)
                 .help("Per-thread-per-iteration sleep in ms"),
         )
+        .arg(
+            Arg::with_name("write-client-keys")
+                .long("write-client-keys")
+                .value_name("FILENAME")
+                .takes_value(true)
+                .help("Generate client keys and stakes and write the list to YAML file"),
+        )
+        .arg(
+            Arg::with_name("read-client-keys")
+                .long("read-client-keys")
+                .value_name("FILENAME")
+                .takes_value(true)
+                .help("Read client keys and stakes from the YAML file"),
+        )
 }
 
 /// Parses a clap `ArgMatches` structure into a `Config`
@@ -163,5 +183,16 @@ pub fn extract_args<'a>(matches: &ArgMatches<'a>) -> Config {
 
     args.sustained = matches.is_present("sustained");
 
+    if let Some(s) = matches.value_of("write-client-keys") {
+        args.write_to_client_file = true;
+        args.client_ids_and_stake_file = s.to_string();
+    }
+
+    if let Some(s) = matches.value_of("read-client-keys") {
+        assert!(!args.write_to_client_file);
+        args.read_from_client_file = true;
+        args.client_ids_and_stake_file = s.to_string();
+    }
+
     args
 }

bench-tps/src/main.rs

@@ -1,8 +1,15 @@
 mod bench;
 mod cli;
 
-use crate::bench::{do_bench_tps, generate_and_fund_keypairs, Config, NUM_LAMPORTS_PER_ACCOUNT};
+use crate::bench::{
+    do_bench_tps, generate_and_fund_keypairs, generate_keypairs, Config, NUM_LAMPORTS_PER_ACCOUNT,
+};
 use solana::gossip_service::{discover_cluster, get_multi_client};
+use solana_sdk::signature::Keypair;
+use std::collections::HashMap;
+use std::fs::File;
+use std::io::prelude::*;
+use std::path::Path;
 use std::process::exit;
 
 fn main() {
@@ -22,8 +29,28 @@ fn main() {
         tx_count,
         thread_batch_sleep_ms,
         sustained,
+        client_ids_and_stake_file,
+        write_to_client_file,
+        read_from_client_file,
     } = cli_config;
 
+    if write_to_client_file {
+        let keypairs = generate_keypairs(&id, tx_count as u64 * 2);
+        let mut accounts = HashMap::new();
+        keypairs.iter().for_each(|keypair| {
+            accounts.insert(
+                serde_json::to_string(&keypair.to_bytes().to_vec()).unwrap(),
+                NUM_LAMPORTS_PER_ACCOUNT as u64,
+            );
+        });
+
+        let serialized = serde_yaml::to_string(&accounts).unwrap();
+        let path = Path::new(&client_ids_and_stake_file);
+        let mut file = File::create(path).unwrap();
+        file.write_all(&serialized.into_bytes()).unwrap();
+        return;
+    }
+
     println!("Connecting to the cluster");
     let (nodes, _replicators) =
         discover_cluster(&entrypoint_addr, num_nodes).unwrap_or_else(|err| {
@@ -41,13 +68,29 @@ fn main() {
         exit(1);
     }
 
-    let (keypairs, keypair_balance) = generate_and_fund_keypairs(
-        &client,
-        Some(drone_addr),
-        &id,
-        tx_count,
-        NUM_LAMPORTS_PER_ACCOUNT,
-    );
+    let (keypairs, keypair_balance) = if read_from_client_file {
+        let path = Path::new(&client_ids_and_stake_file);
+        let file = File::open(path).unwrap();
+
+        let accounts: HashMap<String, u64> = serde_yaml::from_reader(file).unwrap();
+        let mut keypairs = vec![];
+        let mut last_balance = 0;
+
+        accounts.into_iter().for_each(|(keypair, balance)| {
+            let bytes: Vec<u8> = serde_json::from_str(keypair.as_str()).unwrap();
+            keypairs.push(Keypair::from_bytes(&bytes).unwrap());
+            last_balance = balance;
+        });
+        (keypairs, last_balance)
+    } else {
+        generate_and_fund_keypairs(
+            &client,
+            Some(drone_addr),
+            &id,
+            tx_count,
+            NUM_LAMPORTS_PER_ACCOUNT,
+        )
+    };
 
     let config = Config {
         id,