Add handling to close accounts to many-accounts bench (#16199)

* gitignore farf

* Improve cli args

* Use derived addresses for accounts

* Add parameter to close every nth account created

accounts-cluster-bench/.gitignore

@@ -0,0 +1 @@
+/farf/

accounts-cluster-bench/src/main.rs

@@ -21,7 +21,6 @@ use solana_sdk::{
     transaction::Transaction,
 };
 use solana_transaction_status::parse_token::spl_token_v2_0_instruction;
-use spl_token_v2_0::solana_program::pubkey::Pubkey as SplPubkey;
 use std::{
     net::SocketAddr,
     process::exit,
@@ -251,28 +250,38 @@ impl TransactionExecutor {
     }
 }
 
-fn make_message(
+struct SeedTracker {
+    max_created: Arc<AtomicU64>,
+    max_closed: Arc<AtomicU64>,
+}
+
+fn make_create_message(
     keypair: &Keypair,
+    base_keypair: &Keypair,
+    max_created_seed: Arc<AtomicU64>,
     num_instructions: usize,
     balance: u64,
     maybe_space: Option<u64>,
     mint: Option<Pubkey>,
-) -> (Message, Vec<Keypair>) {
+) -> Message {
     let space = maybe_space.unwrap_or_else(|| thread_rng().gen_range(0, 1000));
 
-    let (instructions, new_keypairs): (Vec<_>, Vec<_>) = (0..num_instructions)
+    let instructions: Vec<_> = (0..num_instructions)
         .into_iter()
         .map(|_| {
-            let new_keypair = Keypair::new();
-
             let program_id = if mint.is_some() {
                 inline_spl_token_v2_0::id()
             } else {
                 system_program::id()
             };
-            let mut instructions = vec![system_instruction::create_account(
+            let seed = max_created_seed.fetch_add(1, Ordering::Relaxed).to_string();
+            let to_pubkey =
+                Pubkey::create_with_seed(&base_keypair.pubkey(), &seed, &program_id).unwrap();
+            let mut instructions = vec![system_instruction::create_account_with_seed(
                 &keypair.pubkey(),
-                &new_keypair.pubkey(),
+                &to_pubkey,
+                &base_keypair.pubkey(),
+                &seed,
                 balance,
                 space,
                 &program_id,
@@ -281,25 +290,69 @@ fn make_message(
                 instructions.push(spl_token_v2_0_instruction(
                     spl_token_v2_0::instruction::initialize_account(
                         &spl_token_v2_0::id(),
-                        &spl_token_v2_0_pubkey(&new_keypair.pubkey()),
+                        &spl_token_v2_0_pubkey(&to_pubkey),
                         &spl_token_v2_0_pubkey(&mint_address),
-                        &SplPubkey::new_unique(),
+                        &spl_token_v2_0_pubkey(&base_keypair.pubkey()),
                     )
                     .unwrap(),
                 ));
             }
 
-            (instructions, new_keypair)
+            instructions
         })
-        .unzip();
+        .collect();
     let instructions: Vec<_> = instructions.into_iter().flatten().collect();
 
-    (
-        Message::new(&instructions, Some(&keypair.pubkey())),
-        new_keypairs,
-    )
+    Message::new(&instructions, Some(&keypair.pubkey()))
 }
 
+fn make_close_message(
+    keypair: &Keypair,
+    base_keypair: &Keypair,
+    max_closed_seed: Arc<AtomicU64>,
+    num_instructions: usize,
+    balance: u64,
+    spl_token: bool,
+) -> Message {
+    let instructions: Vec<_> = (0..num_instructions)
+        .into_iter()
+        .map(|_| {
+            let program_id = if spl_token {
+                inline_spl_token_v2_0::id()
+            } else {
+                system_program::id()
+            };
+            let seed = max_closed_seed.fetch_add(1, Ordering::Relaxed).to_string();
+            let address =
+                Pubkey::create_with_seed(&base_keypair.pubkey(), &seed, &program_id).unwrap();
+            if spl_token {
+                spl_token_v2_0_instruction(
+                    spl_token_v2_0::instruction::close_account(
+                        &spl_token_v2_0::id(),
+                        &spl_token_v2_0_pubkey(&address),
+                        &spl_token_v2_0_pubkey(&keypair.pubkey()),
+                        &spl_token_v2_0_pubkey(&base_keypair.pubkey()),
+                        &[],
+                    )
+                    .unwrap(),
+                )
+            } else {
+                system_instruction::transfer_with_seed(
+                    &address,
+                    &base_keypair.pubkey(),
+                    seed,
+                    &program_id,
+                    &keypair.pubkey(),
+                    balance,
+                )
+            }
+        })
+        .collect();
+
+    Message::new(&instructions, Some(&keypair.pubkey()))
+}
+
+#[allow(clippy::too_many_arguments)]
 fn run_accounts_bench(
     entrypoint_addr: SocketAddr,
     faucet_addr: SocketAddr,
@@ -307,6 +360,7 @@ fn run_accounts_bench(
     iterations: usize,
     maybe_space: Option<u64>,
     batch_size: usize,
+    close_nth: u64,
     maybe_lamports: Option<u64>,
     num_instructions: usize,
     mint: Option<Pubkey>,
@@ -322,7 +376,8 @@ fn run_accounts_bench(
     let mut count = 0;
     let mut recent_blockhash = client.get_recent_blockhash().expect("blockhash");
     let mut tx_sent_count = 0;
-    let mut total_account_count = 0;
+    let mut total_accounts_created = 0;
+    let mut total_accounts_closed = 0;
     let mut balance = client.get_balance(&keypair.pubkey()).unwrap_or(0);
     let mut last_balance = Instant::now();
 
@@ -334,6 +389,12 @@ fn run_accounts_bench(
             .expect("min balance")
     });
 
+    let base_keypair = Keypair::new();
+    let seed_tracker = SeedTracker {
+        max_created: Arc::new(AtomicU64::default()),
+        max_closed: Arc::new(AtomicU64::default()),
+    };
+
     info!("Starting balance: {}", balance);
 
     let executor = TransactionExecutor::new(entrypoint_addr);
@@ -344,8 +405,15 @@ fn run_accounts_bench(
             last_blockhash = Instant::now();
         }
 
-        let (message, _keypairs) =
-            make_message(keypair, num_instructions, min_balance, maybe_space, mint);
+        let message = make_create_message(
+            keypair,
+            &base_keypair,
+            seed_tracker.max_created.clone(),
+            num_instructions,
+            min_balance,
+            maybe_space,
+            mint,
+        );
         let fee = recent_blockhash.1.calculate_fee(&message);
         let lamports = min_balance + fee;
 
@@ -370,27 +438,55 @@ fn run_accounts_bench(
         if sigs_len < batch_size {
             let num_to_create = batch_size - sigs_len;
             info!("creating {} new", num_to_create);
-            let (txs, _new_keypairs): (Vec<_>, Vec<_>) = (0..num_to_create)
+            let txs: Vec<_> = (0..num_to_create)
                 .into_par_iter()
                 .map(|_| {
-                    let (message, new_keypairs) =
-                        make_message(keypair, num_instructions, min_balance, maybe_space, mint);
-                    let signers: Vec<&Keypair> = new_keypairs
-                        .iter()
-                        .chain(std::iter::once(keypair))
-                        .collect();
-                    (
-                        Transaction::new(&signers, message, recent_blockhash.0),
-                        new_keypairs,
-                    )
+                    let message = make_create_message(
+                        keypair,
+                        &base_keypair,
+                        seed_tracker.max_created.clone(),
+                        num_instructions,
+                        min_balance,
+                        maybe_space,
+                        mint,
+                    );
+                    let signers: Vec<&Keypair> = vec![keypair, &base_keypair];
+                    Transaction::new(&signers, message, recent_blockhash.0)
                 })
-                .unzip();
+                .collect();
             balance = balance.saturating_sub(lamports * txs.len() as u64);
             info!("txs: {}", txs.len());
             let new_ids = executor.push_transactions(txs);
             info!("ids: {}", new_ids.len());
             tx_sent_count += new_ids.len();
-            total_account_count += num_instructions * new_ids.len();
+            total_accounts_created += num_instructions * new_ids.len();
+
+            if close_nth > 0 {
+                let expected_closed = total_accounts_created as u64 / close_nth;
+                if expected_closed > total_accounts_closed {
+                    let txs: Vec<_> = (0..expected_closed - total_accounts_closed)
+                        .into_par_iter()
+                        .map(|_| {
+                            let message = make_close_message(
+                                keypair,
+                                &base_keypair,
+                                seed_tracker.max_closed.clone(),
+                                1,
+                                min_balance,
+                                mint.is_some(),
+                            );
+                            let signers: Vec<&Keypair> = vec![keypair, &base_keypair];
+                            Transaction::new(&signers, message, recent_blockhash.0)
+                        })
+                        .collect();
+                    balance = balance.saturating_sub(fee * txs.len() as u64);
+                    info!("close txs: {}", txs.len());
+                    let new_ids = executor.push_transactions(txs);
+                    info!("close ids: {}", new_ids.len());
+                    tx_sent_count += new_ids.len();
+                    total_accounts_closed += new_ids.len() as u64;
+                }
+            }
         } else {
             let _ = executor.drain_cleared();
         }
@@ -398,8 +494,8 @@ fn run_accounts_bench(
         count += 1;
         if last_log.elapsed().as_millis() > 3000 {
             info!(
-                "total_accounts: {} tx_sent_count: {} loop_count: {} balance: {}",
-                total_account_count, tx_sent_count, count, balance
+                "total_accounts_created: {} total_accounts_closed: {} tx_sent_count: {} loop_count: {} balance: {}",
+                total_accounts_created, total_accounts_closed, tx_sent_count, count, balance
             );
             last_log = Instant::now();
         }
@@ -455,14 +551,21 @@ fn main() {
         )
         .arg(
             Arg::with_name("batch_size")
-                .long("batch_size")
+                .long("batch-size")
                 .takes_value(true)
                 .value_name("BYTES")
-                .help("Size of accounts to create"),
+                .help("Number of transactions to send per batch"),
+        )
+        .arg(
+            Arg::with_name("close_nth")
+                .long("close-frequency")
+                .takes_value(true)
+                .value_name("BYTES")
+                .help("Send close transactions after this many accounts created"),
         )
         .arg(
             Arg::with_name("num_instructions")
-                .long("num_instructions")
+                .long("num-instructions")
                 .takes_value(true)
                 .value_name("NUM")
                 .help("Number of accounts to create on each transaction"),
@@ -508,6 +611,7 @@ fn main() {
     let space = value_t!(matches, "space", u64).ok();
     let lamports = value_t!(matches, "lamports", u64).ok();
     let batch_size = value_t!(matches, "batch_size", usize).unwrap_or(4);
+    let close_nth = value_t!(matches, "close_nth", u64).unwrap_or(0);
     let iterations = value_t!(matches, "iterations", usize).unwrap_or(10);
     let num_instructions = value_t!(matches, "num_instructions", usize).unwrap_or(1);
     if num_instructions == 0 || num_instructions > 500 {
@@ -551,6 +655,7 @@ fn main() {
         iterations,
         space,
         batch_size,
+        close_nth,
         lamports,
         num_instructions,
         mint,
@@ -585,6 +690,7 @@ pub mod test {
         let iterations = 10;
         let maybe_space = None;
         let batch_size = 100;
+        let close_nth = 100;
         let maybe_lamports = None;
         let num_instructions = 2;
         let mut start = Measure::start("total accounts run");
@@ -595,6 +701,7 @@ pub mod test {
             iterations,
             maybe_space,
             batch_size,
+            close_nth,
             maybe_lamports,
             num_instructions,
             None,