Reclaim tokens in client TPS demo

src/bin/client-demo.rs

@@ -83,14 +83,22 @@ fn generate_and_send_txs(
     txs: i64,
     last_id: &mut Hash,
     threads: usize,
+    reclaim: bool,
 ) {
-    println!("Signing transactions... {}", txs,);
+    println!("Signing transactions... {}", txs / 2,);
     let signing_start = Instant::now();
 
-    let transactions: Vec<_> = keypairs
-        .par_iter()
-        .map(|keypair| Transaction::new(&id.keypair(), keypair.pubkey(), 1, *last_id))
-        .collect();
+    let transactions: Vec<_> = if !reclaim {
+        keypairs
+            .par_iter()
+            .map(|keypair| Transaction::new(&id.keypair(), keypair.pubkey(), 1, *last_id))
+            .collect()
+    } else {
+        keypairs
+            .par_iter()
+            .map(|keypair| Transaction::new(keypair, id.pubkey(), 1, *last_id))
+            .collect()
+    };
 
     let duration = signing_start.elapsed();
     let ns = duration.as_secs() * 1_000_000_000 + u64::from(duration.subsec_nanos());
@@ -103,7 +111,11 @@ fn generate_and_send_txs(
         duration_as_ms(&duration),
     );
 
-    println!("Transfering {} transactions in {} batches", txs, threads);
+    println!(
+        "Transfering {} transactions in {} batches",
+        txs / 2,
+        threads
+    );
     let transfer_start = Instant::now();
     let sz = transactions.len() / threads;
     let chunks: Vec<_> = transactions.chunks(sz).collect();
@@ -140,7 +152,7 @@ fn main() {
     env_logger::init();
     let mut threads = 4usize;
     let mut num_nodes = 1usize;
-    let mut time_sec = 60;
+    let mut time_sec = 90;
 
     let matches = App::new("solana-client-demo")
         .arg(
@@ -224,18 +236,22 @@ fn main() {
     let mut client = mk_client(&leader);
 
     let starting_balance = client.poll_get_balance(&id.pubkey()).unwrap();
-
     let txs: i64 = 500_000;
-    println!("Airdropping {:?} tokens", txs);
-    let _airdrop = request_airdrop(&drone_addr, &id, txs as u64).unwrap();
-    // TODO: return airdrop Result from Drone
-    sleep(Duration::from_millis(100));
-    let balance = client.poll_get_balance(&id.pubkey()).unwrap();
-    println!("Your balance is: {:?}", balance);
-
-    if balance < txs || (starting_balance == balance) {
-        println!("TPS airdrop limit reached; wait 60sec to retry");
-        exit(1);
+
+    if starting_balance < txs {
+        let airdrop_amount = txs - starting_balance;
+        println!("Airdropping {:?} tokens", airdrop_amount);
+        let _airdrop = request_airdrop(&drone_addr, &id, airdrop_amount as u64).unwrap();
+        // TODO: return airdrop Result from Drone
+        sleep(Duration::from_millis(100));
+
+        let balance = client.poll_get_balance(&id.pubkey()).unwrap();
+        println!("Your balance is: {:?}", balance);
+
+        if balance < txs || (starting_balance == balance) {
+            println!("TPS airdrop limit reached; wait 60sec to retry");
+            exit(1);
+        }
     }
 
     println!("Get last ID...");
@@ -247,7 +263,7 @@ fn main() {
     let rnd = GenKeys::new(seed);
 
     println!("Creating keypairs...");
-    let keypairs = rnd.gen_n_keypairs(txs);
+    let keypairs = rnd.gen_n_keypairs(txs / 2);
 
     let first_count = client.transaction_count();
     println!("initial count {}", first_count);
@@ -275,8 +291,23 @@ fn main() {
     let clients = (0..threads).map(|_| mk_client(&leader)).collect();
 
     // generate and send transactions for the specified duration
-    let time = Duration::new(time_sec, 0);
-    let now = Instant::now();
+    let time = Duration::new(time_sec / 2, 0);
+    let mut now = Instant::now();
+    while now.elapsed() < time {
+        generate_and_send_txs(
+            &mut client,
+            &clients,
+            &id,
+            &keypairs,
+            &leader,
+            txs,
+            &mut last_id,
+            threads,
+            false,
+        );
+    }
+    last_id = client.get_last_id();
+    now = Instant::now();
     while now.elapsed() < time {
         generate_and_send_txs(
             &mut client,
@@ -287,6 +318,7 @@ fn main() {
             txs,
             &mut last_id,
             threads,
+            true,
         );
     }