Fix BankForks::new_rw_arc memory leak (anza-xyz#1893)

bench-tps/src/bench.rs

@@ -1223,7 +1223,7 @@ pub fn fund_keypairs<T: 'static + TpsClient + Send + Sync + ?Sized>(
 mod tests {
     use {
         super::*,
-        solana_runtime::{bank::Bank, bank_client::BankClient},
+        solana_runtime::{bank::Bank, bank_client::BankClient, bank_forks::BankForks},
         solana_sdk::{
             commitment_config::CommitmentConfig,
             feature_set::FeatureSet,
@@ -1234,16 +1234,18 @@ mod tests {
         },
     };
 
-    fn bank_with_all_features(genesis_config: &GenesisConfig) -> Arc<Bank> {
+    fn bank_with_all_features(
+        genesis_config: &GenesisConfig,
+    ) -> (Arc<Bank>, Arc<RwLock<BankForks>>) {
         let mut bank = Bank::new_for_tests(genesis_config);
         bank.feature_set = Arc::new(FeatureSet::all_enabled());
-        bank.wrap_with_bank_forks_for_tests().0
+        bank.wrap_with_bank_forks_for_tests()
     }
 
     #[test]
     fn test_bench_tps_bank_client() {
         let (genesis_config, id) = create_genesis_config(sol_to_lamports(10_000.0));
-        let bank = bank_with_all_features(&genesis_config);
+        let (bank, _bank_forks) = bank_with_all_features(&genesis_config);
         let client = Arc::new(BankClient::new_shared(bank));
 
         let config = Config {
@@ -1264,7 +1266,7 @@ mod tests {
     #[test]
     fn test_bench_tps_fund_keys() {
         let (genesis_config, id) = create_genesis_config(sol_to_lamports(10_000.0));
-        let bank = bank_with_all_features(&genesis_config);
+        let (bank, _bank_forks) = bank_with_all_features(&genesis_config);
         let client = Arc::new(BankClient::new_shared(bank));
         let keypair_count = 20;
         let lamports = 20;
@@ -1289,7 +1291,7 @@ mod tests {
         let (mut genesis_config, id) = create_genesis_config(sol_to_lamports(10_000.0));
         let fee_rate_governor = FeeRateGovernor::new(11, 0);
         genesis_config.fee_rate_governor = fee_rate_governor;
-        let bank = bank_with_all_features(&genesis_config);
+        let (bank, _bank_forks) = bank_with_all_features(&genesis_config);
         let client = Arc::new(BankClient::new_shared(bank));
         let keypair_count = 20;
         let lamports = 20;
@@ -1307,7 +1309,7 @@ mod tests {
     #[test]
     fn test_bench_tps_create_durable_nonce() {
         let (genesis_config, id) = create_genesis_config(sol_to_lamports(10_000.0));
-        let bank = bank_with_all_features(&genesis_config);
+        let (bank, _bank_forks) = bank_with_all_features(&genesis_config);
         let client = Arc::new(BankClient::new_shared(bank));
         let keypair_count = 10;
         let lamports = 10_000_000;

core/benches/banking_stage.rs

@@ -398,7 +398,7 @@ fn simulate_process_entries(
     let bank_fork = BankForks::new_rw_arc(bank);
     let bank = bank_fork.read().unwrap().get_with_scheduler(slot).unwrap();
     bank.clone_without_scheduler()
-        .set_fork_graph_in_program_cache(bank_fork.clone());
+        .set_fork_graph_in_program_cache(Arc::downgrade(&bank_fork));
 
     for i in 0..(num_accounts / 2) {
         bank.transfer(initial_lamports, mint_keypair, &keypairs[i * 2].pubkey())

core/benches/consumer.rs

@@ -19,7 +19,7 @@ use {
         poh_recorder::{create_test_recorder, PohRecorder},
         poh_service::PohService,
     },
-    solana_runtime::bank::Bank,
+    solana_runtime::{bank::Bank, bank_forks::BankForks},
     solana_sdk::{
         account::{Account, ReadableAccount},
         signature::Keypair,
@@ -89,6 +89,7 @@ fn create_consumer(poh_recorder: &RwLock<PohRecorder>) -> Consumer {
 
 struct BenchFrame {
     bank: Arc<Bank>,
+    _bank_forks: Arc<RwLock<BankForks>>,
     ledger_path: TempDir,
     exit: Arc<AtomicBool>,
     poh_recorder: Arc<RwLock<PohRecorder>>,
@@ -115,7 +116,7 @@ fn setup() -> BenchFrame {
     bank.write_cost_tracker()
         .unwrap()
         .set_limits(u64::MAX, u64::MAX, u64::MAX);
-    let bank = bank.wrap_with_bank_forks_for_tests().0;
+    let (bank, bank_forks) = bank.wrap_with_bank_forks_for_tests();
 
     let ledger_path = TempDir::new().unwrap();
     let blockstore = Arc::new(
@@ -126,6 +127,7 @@ fn setup() -> BenchFrame {
 
     BenchFrame {
         bank,
+        _bank_forks: bank_forks,
         ledger_path,
         exit,
         poh_recorder,
@@ -147,6 +149,7 @@ fn bench_process_and_record_transactions(bencher: &mut Bencher, batch_size: usiz
 
     let BenchFrame {
         bank,
+        _bank_forks,
         ledger_path: _ledger_path,
         exit,
         poh_recorder,

core/src/banking_stage.rs

@@ -1064,7 +1064,7 @@ mod tests {
             drop(poh_recorder);
 
             let mut blockhash = start_hash;
-            let bank = Bank::new_no_wallclock_throttle_for_tests(&genesis_config).0;
+            let (bank, _bank_forks) = Bank::new_no_wallclock_throttle_for_tests(&genesis_config);
             bank.process_transaction(&fund_tx).unwrap();
             //receive entries + ticks
             loop {
@@ -1208,7 +1208,7 @@ mod tests {
                 .map(|(_bank, (entry, _tick_height))| entry)
                 .collect();
 
-            let bank = Bank::new_no_wallclock_throttle_for_tests(&genesis_config).0;
+            let (bank, _bank_forks) = Bank::new_no_wallclock_throttle_for_tests(&genesis_config);
             for entry in entries {
                 bank.process_entry_transactions(entry.transactions)
                     .iter()
@@ -1232,7 +1232,7 @@ mod tests {
             mint_keypair,
             ..
         } = create_genesis_config(10_000);
-        let bank = Bank::new_no_wallclock_throttle_for_tests(&genesis_config).0;
+        let (bank, _bank_forks) = Bank::new_no_wallclock_throttle_for_tests(&genesis_config);
         let ledger_path = get_tmp_ledger_path_auto_delete!();
         {
             let blockstore = Blockstore::open(ledger_path.path())

core/src/banking_stage/consume_worker.rs

@@ -710,7 +710,8 @@ mod tests {
         },
         solana_poh::poh_recorder::{PohRecorder, WorkingBankEntry},
         solana_runtime::{
-            prioritization_fee_cache::PrioritizationFeeCache, vote_sender_types::ReplayVoteReceiver,
+            bank_forks::BankForks, prioritization_fee_cache::PrioritizationFeeCache,
+            vote_sender_types::ReplayVoteReceiver,
         },
         solana_sdk::{
             genesis_config::GenesisConfig, poh_config::PohConfig, pubkey::Pubkey,
@@ -729,6 +730,7 @@ mod tests {
         mint_keypair: Keypair,
         genesis_config: GenesisConfig,
         bank: Arc<Bank>,
+        _bank_forks: Arc<RwLock<BankForks>>,
         _ledger_path: TempDir,
         _entry_receiver: Receiver<WorkingBankEntry>,
         poh_recorder: Arc<RwLock<PohRecorder>>,
@@ -745,7 +747,7 @@ mod tests {
             mint_keypair,
             ..
         } = create_slow_genesis_config(10_000);
-        let bank = Bank::new_no_wallclock_throttle_for_tests(&genesis_config).0;
+        let (bank, bank_forks) = Bank::new_no_wallclock_throttle_for_tests(&genesis_config);
 
         let ledger_path = get_tmp_ledger_path_auto_delete!();
         let blockstore = Blockstore::open(ledger_path.path())
@@ -788,6 +790,7 @@ mod tests {
                 mint_keypair,
                 genesis_config,
                 bank,
+                _bank_forks: bank_forks,
                 _ledger_path: ledger_path,
                 _entry_receiver: entry_receiver,
                 poh_recorder,

core/src/banking_stage/consumer.rs

@@ -851,7 +851,7 @@ mod tests {
         solana_poh::poh_recorder::{PohRecorder, Record, WorkingBankEntry},
         solana_program_runtime::timings::ProgramTiming,
         solana_rpc::transaction_status_service::TransactionStatusService,
-        solana_runtime::prioritization_fee_cache::PrioritizationFeeCache,
+        solana_runtime::{bank_forks::BankForks, prioritization_fee_cache::PrioritizationFeeCache},
         solana_sdk::{
             account::AccountSharedData,
             account_utils::StateMut,
@@ -989,6 +989,7 @@ mod tests {
     ) -> (
         Vec<Transaction>,
         Arc<Bank>,
+        Arc<RwLock<BankForks>>,
         Arc<RwLock<PohRecorder>>,
         Receiver<WorkingBankEntry>,
         GenesisConfigInfo,
@@ -1003,7 +1004,7 @@ mod tests {
         } = &genesis_config_info;
         let blockstore =
             Blockstore::open(ledger_path).expect("Expected to be able to open database ledger");
-        let bank = Bank::new_no_wallclock_throttle_for_tests(genesis_config).0;
+        let (bank, bank_forks) = Bank::new_no_wallclock_throttle_for_tests(genesis_config);
         let exit = Arc::new(AtomicBool::default());
         let (poh_recorder, entry_receiver, record_receiver) = PohRecorder::new(
             bank.tick_height(),
@@ -1032,6 +1033,7 @@ mod tests {
         (
             transactions,
             bank,
+            bank_forks,
             poh_recorder,
             entry_receiver,
             genesis_config_info,
@@ -1059,7 +1061,7 @@ mod tests {
             mint_keypair,
             ..
         } = create_slow_genesis_config(10_000);
-        let bank = Bank::new_no_wallclock_throttle_for_tests(&genesis_config).0;
+        let (bank, _bank_forks) = Bank::new_no_wallclock_throttle_for_tests(&genesis_config);
         let pubkey = solana_sdk::pubkey::new_rand();
 
         let transactions = sanitize_transactions(vec![system_transaction::transfer(
@@ -1187,7 +1189,7 @@ mod tests {
             mint_keypair,
             ..
         } = create_slow_genesis_config(10_000);
-        let bank = Bank::new_no_wallclock_throttle_for_tests(&genesis_config).0;
+        let (bank, _bank_forks) = Bank::new_no_wallclock_throttle_for_tests(&genesis_config);
         let pubkey = Pubkey::new_unique();
 
         // setup nonce account with a durable nonce different from the current
@@ -1343,7 +1345,7 @@ mod tests {
             mint_keypair,
             ..
         } = create_slow_genesis_config(10_000);
-        let bank = Bank::new_no_wallclock_throttle_for_tests(&genesis_config).0;
+        let (bank, _bank_forks) = Bank::new_no_wallclock_throttle_for_tests(&genesis_config);
         let pubkey = solana_sdk::pubkey::new_rand();
 
         let transactions = {
@@ -1427,7 +1429,7 @@ mod tests {
         } = create_slow_genesis_config(10_000);
         let mut bank = Bank::new_for_tests(&genesis_config);
         bank.ns_per_slot = u128::MAX;
-        let bank = bank.wrap_with_bank_forks_for_tests().0;
+        let (bank, _bank_forks) = bank.wrap_with_bank_forks_for_tests();
         let pubkey = solana_sdk::pubkey::new_rand();
 
         let ledger_path = get_tmp_ledger_path_auto_delete!();
@@ -1583,7 +1585,7 @@ mod tests {
             mint_keypair,
             ..
         } = create_slow_genesis_config(10_000);
-        let bank = Bank::new_no_wallclock_throttle_for_tests(&genesis_config).0;
+        let (bank, _bank_forks) = Bank::new_no_wallclock_throttle_for_tests(&genesis_config);
         let pubkey = solana_sdk::pubkey::new_rand();
         let pubkey1 = solana_sdk::pubkey::new_rand();
 
@@ -1660,7 +1662,7 @@ mod tests {
             mint_keypair,
             ..
         } = create_slow_genesis_config(lamports);
-        let bank = Bank::new_no_wallclock_throttle_for_tests(&genesis_config).0;
+        let (bank, _bank_forks) = Bank::new_no_wallclock_throttle_for_tests(&genesis_config);
         // set cost tracker limits to MAX so it will not filter out TXs
         bank.write_cost_tracker()
             .unwrap()
@@ -1721,7 +1723,7 @@ mod tests {
             mint_keypair,
             ..
         } = create_slow_genesis_config(10_000);
-        let bank = Bank::new_no_wallclock_throttle_for_tests(&genesis_config).0;
+        let (bank, _bank_forks) = Bank::new_no_wallclock_throttle_for_tests(&genesis_config);
         // set cost tracker limits to MAX so it will not filter out TXs
         bank.write_cost_tracker()
             .unwrap()
@@ -1780,7 +1782,7 @@ mod tests {
             mint_keypair,
             ..
         } = create_slow_genesis_config(10_000);
-        let bank = Bank::new_no_wallclock_throttle_for_tests(&genesis_config).0;
+        let (bank, _bank_forks) = Bank::new_no_wallclock_throttle_for_tests(&genesis_config);
 
         let pubkey = solana_sdk::pubkey::new_rand();
 
@@ -1861,7 +1863,7 @@ mod tests {
         } = create_slow_genesis_config(solana_sdk::native_token::sol_to_lamports(1000.0));
         genesis_config.rent.lamports_per_byte_year = 50;
         genesis_config.rent.exemption_threshold = 2.0;
-        let bank = Bank::new_no_wallclock_throttle_for_tests(&genesis_config).0;
+        let (bank, _bank_forks) = Bank::new_no_wallclock_throttle_for_tests(&genesis_config);
         let pubkey = solana_sdk::pubkey::new_rand();
         let pubkey1 = solana_sdk::pubkey::new_rand();
         let keypair1 = Keypair::new();
@@ -2130,7 +2132,7 @@ mod tests {
     fn test_consume_buffered_packets() {
         let ledger_path = get_tmp_ledger_path_auto_delete!();
         {
-            let (transactions, bank, poh_recorder, _entry_receiver, _, poh_simulator) =
+            let (transactions, bank, _bank_forks, poh_recorder, _entry_receiver, _, poh_simulator) =
                 setup_conflicting_transactions(ledger_path.path());
             let recorder: TransactionRecorder = poh_recorder.read().unwrap().new_recorder();
             let num_conflicting_transactions = transactions.len();
@@ -2203,8 +2205,15 @@ mod tests {
     fn test_consume_buffered_packets_sanitization_error() {
         let ledger_path = get_tmp_ledger_path_auto_delete!();
         {
-            let (mut transactions, bank, poh_recorder, _entry_receiver, _, poh_simulator) =
-                setup_conflicting_transactions(ledger_path.path());
+            let (
+                mut transactions,
+                bank,
+                _bank_forks,
+                poh_recorder,
+                _entry_receiver,
+                _,
+                poh_simulator,
+            ) = setup_conflicting_transactions(ledger_path.path());
             let duplicate_account_key = transactions[0].message.account_keys[0];
             transactions[0]
                 .message
@@ -2259,7 +2268,7 @@ mod tests {
     fn test_consume_buffered_packets_retryable() {
         let ledger_path = get_tmp_ledger_path_auto_delete!();
         {
-            let (transactions, bank, poh_recorder, _entry_receiver, _, poh_simulator) =
+            let (transactions, bank, _bank_forks, poh_recorder, _entry_receiver, _, poh_simulator) =
                 setup_conflicting_transactions(ledger_path.path());
             let recorder = poh_recorder.read().unwrap().new_recorder();
             let num_conflicting_transactions = transactions.len();
@@ -2355,8 +2364,15 @@ mod tests {
     fn test_consume_buffered_packets_batch_priority_guard() {
         let ledger_path = get_tmp_ledger_path_auto_delete!();
         {
-            let (_, bank, poh_recorder, _entry_receiver, genesis_config_info, poh_simulator) =
-                setup_conflicting_transactions(ledger_path.path());
+            let (
+                _,
+                bank,
+                _bank_forks,
+                poh_recorder,
+                _entry_receiver,
+                genesis_config_info,
+                poh_simulator,
+            ) = setup_conflicting_transactions(ledger_path.path());
             let recorder = poh_recorder.read().unwrap().new_recorder();
 
             // Setup transactions:

core/src/banking_stage/decision_maker.rs

@@ -148,7 +148,7 @@ mod tests {
     #[test]
     fn test_make_consume_or_forward_decision() {
         let genesis_config = create_genesis_config(2).genesis_config;
-        let bank = Bank::new_no_wallclock_throttle_for_tests(&genesis_config).0;
+        let (bank, _bank_forks) = Bank::new_no_wallclock_throttle_for_tests(&genesis_config);
         let ledger_path = temp_dir();
         let blockstore = Arc::new(Blockstore::open(ledger_path.as_path()).unwrap());
         let (exit, poh_recorder, poh_service, _entry_receiver) =

core/src/banking_stage/read_write_account_set.rs

@@ -84,7 +84,7 @@ mod tests {
     use {
         super::ReadWriteAccountSet,
         solana_ledger::genesis_utils::GenesisConfigInfo,
-        solana_runtime::{bank::Bank, genesis_utils::create_genesis_config},
+        solana_runtime::{bank::Bank, bank_forks::BankForks, genesis_utils::create_genesis_config},
         solana_sdk::{
             account::AccountSharedData,
             address_lookup_table::{
@@ -101,7 +101,10 @@ mod tests {
             signer::Signer,
             transaction::{MessageHash, SanitizedTransaction, VersionedTransaction},
         },
-        std::{borrow::Cow, sync::Arc},
+        std::{
+            borrow::Cow,
+            sync::{Arc, RwLock},
+        },
     };
 
     fn create_test_versioned_message(
@@ -171,9 +174,9 @@ mod tests {
         )
     }
 
-    fn create_test_bank() -> Arc<Bank> {
+    fn create_test_bank() -> (Arc<Bank>, Arc<RwLock<BankForks>>) {
         let GenesisConfigInfo { genesis_config, .. } = create_genesis_config(10_000);
-        Bank::new_no_wallclock_throttle_for_tests(&genesis_config).0
+        Bank::new_no_wallclock_throttle_for_tests(&genesis_config)
     }
 
     // Helper function (could potentially use test_case in future).
@@ -182,7 +185,7 @@ mod tests {
     // conflict_index = 2 means write lock conflict with address table key
     // conflict_index = 3 means read lock conflict with address table key
     fn test_check_and_take_locks(conflict_index: usize, add_write: bool, expectation: bool) {
-        let bank = create_test_bank();
+        let (bank, _bank_forks) = create_test_bank();
         let (bank, table_address) = create_test_address_lookup_table(bank, 2);
         let tx = create_test_sanitized_transaction(
             &Keypair::new(),

core/src/banking_stage/unprocessed_transaction_storage.rs

@@ -1069,7 +1069,7 @@ mod tests {
             mint_keypair,
             ..
         } = create_genesis_config(10);
-        let current_bank = Bank::new_with_bank_forks_for_tests(&genesis_config).0;
+        let (current_bank, _bank_forks) = Bank::new_with_bank_forks_for_tests(&genesis_config);
 
         let simple_transactions: Vec<Transaction> = (0..256)
             .map(|_id| {