Define PohRecorder set_bank related test helper methods (#33626)

Define PohRecorder set_bank related test methods

Cargo.lock

@@ -6587,6 +6587,7 @@ dependencies = [
  "solana-measure",
  "solana-metrics",
  "solana-perf",
+ "solana-poh",
  "solana-runtime",
  "solana-sdk",
  "thiserror",

banking-bench/Cargo.toml

@@ -21,12 +21,15 @@ solana-ledger = { workspace = true }
 solana-logger = { workspace = true }
 solana-measure = { workspace = true }
 solana-perf = { workspace = true }
-solana-poh = { workspace = true }
+solana-poh = { workspace = true, features = ["dev-context-only-utils"] }
 solana-runtime = { workspace = true }
 solana-sdk = { workspace = true }
 solana-streamer = { workspace = true }
 solana-tpu-client = { workspace = true }
 solana-version = { workspace = true }
 
+[features]
+dev-context-only-utils = []
+
 [package.metadata.docs.rs]
 targets = ["x86_64-unknown-linux-gnu"]

banking-bench/src/main.rs

@@ -549,7 +549,10 @@ fn main() {
                 );
 
                 assert!(poh_recorder.read().unwrap().bank().is_none());
-                poh_recorder.write().unwrap().set_bank(bank.clone(), false);
+                poh_recorder
+                    .write()
+                    .unwrap()
+                    .set_bank_for_test(bank.clone());
                 assert!(poh_recorder.read().unwrap().bank().is_some());
                 debug!(
                     "new_bank_time: {}us insert_time: {}us poh_time: {}us",

core/Cargo.toml

@@ -89,6 +89,7 @@ serial_test = { workspace = true }
 # See order-crates-for-publishing.py for using this unusual `path = "."`
 solana-core = { path = ".", features = ["dev-context-only-utils"] }
 solana-logger = { workspace = true }
+solana-poh = { workspace = true, features = ["dev-context-only-utils"] }
 solana-program-runtime = { workspace = true }
 solana-runtime = { workspace = true, features = ["dev-context-only-utils"] }
 solana-sdk = { workspace = true, features = ["dev-context-only-utils"] }

core/src/banking_stage.rs

@@ -1071,7 +1071,10 @@ mod tests {
 
             let poh_simulator = simulate_poh(record_receiver, &poh_recorder);
 
-            poh_recorder.write().unwrap().set_bank(bank.clone(), false);
+            poh_recorder
+                .write()
+                .unwrap()
+                .set_bank_for_test(bank.clone());
             let pubkey = solana_sdk::pubkey::new_rand();
             let keypair2 = Keypair::new();
             let pubkey2 = solana_sdk::pubkey::new_rand();

core/src/banking_stage/consume_worker.rs

@@ -283,7 +283,10 @@ mod tests {
             ..
         } = &test_frame;
         let worker_thread = std::thread::spawn(move || worker.run());
-        poh_recorder.write().unwrap().set_bank(bank.clone(), false);
+        poh_recorder
+            .write()
+            .unwrap()
+            .set_bank_for_test(bank.clone());
 
         let pubkey1 = Pubkey::new_unique();
 
@@ -325,7 +328,10 @@ mod tests {
             ..
         } = &test_frame;
         let worker_thread = std::thread::spawn(move || worker.run());
-        poh_recorder.write().unwrap().set_bank(bank.clone(), false);
+        poh_recorder
+            .write()
+            .unwrap()
+            .set_bank_for_test(bank.clone());
 
         let pubkey1 = Pubkey::new_unique();
         let pubkey2 = Pubkey::new_unique();
@@ -370,7 +376,10 @@ mod tests {
             ..
         } = &test_frame;
         let worker_thread = std::thread::spawn(move || worker.run());
-        poh_recorder.write().unwrap().set_bank(bank.clone(), false);
+        poh_recorder
+            .write()
+            .unwrap()
+            .set_bank_for_test(bank.clone());
 
         let pubkey1 = Pubkey::new_unique();
         let pubkey2 = Pubkey::new_unique();

core/src/banking_stage/consumer.rs

@@ -805,7 +805,10 @@ mod tests {
         let recorder = poh_recorder.new_recorder();
         let poh_recorder = Arc::new(RwLock::new(poh_recorder));
 
-        poh_recorder.write().unwrap().set_bank(bank.clone(), false);
+        poh_recorder
+            .write()
+            .unwrap()
+            .set_bank_for_test(bank.clone());
 
         let poh_simulator = simulate_poh(record_receiver, &poh_recorder);
 
@@ -966,7 +969,10 @@ mod tests {
 
             let poh_simulator = simulate_poh(record_receiver, &poh_recorder);
 
-            poh_recorder.write().unwrap().set_bank(bank.clone(), false);
+            poh_recorder
+                .write()
+                .unwrap()
+                .set_bank_for_test(bank.clone());
             let (replay_vote_sender, _replay_vote_receiver) = unbounded();
             let committer = Committer::new(
                 None,
@@ -1093,7 +1099,10 @@ mod tests {
 
             let poh_simulator = simulate_poh(record_receiver, &poh_recorder);
 
-            poh_recorder.write().unwrap().set_bank(bank.clone(), false);
+            poh_recorder
+                .write()
+                .unwrap()
+                .set_bank_for_test(bank.clone());
             let (replay_vote_sender, _replay_vote_receiver) = unbounded();
             let committer = Committer::new(
                 None,
@@ -1179,7 +1188,10 @@ mod tests {
 
             let poh_simulator = simulate_poh(record_receiver, &poh_recorder);
 
-            poh_recorder.write().unwrap().set_bank(bank.clone(), false);
+            poh_recorder
+                .write()
+                .unwrap()
+                .set_bank_for_test(bank.clone());
             let (replay_vote_sender, _replay_vote_receiver) = unbounded();
             let committer = Committer::new(
                 None,
@@ -1328,7 +1340,10 @@ mod tests {
             let recorder = poh_recorder.new_recorder();
             let poh_recorder = Arc::new(RwLock::new(poh_recorder));
 
-            poh_recorder.write().unwrap().set_bank(bank.clone(), false);
+            poh_recorder
+                .write()
+                .unwrap()
+                .set_bank_for_test(bank.clone());
 
             let poh_simulator = simulate_poh(record_receiver, &poh_recorder);
 
@@ -1628,7 +1643,10 @@ mod tests {
 
             let poh_simulator = simulate_poh(record_receiver, &poh_recorder);
 
-            poh_recorder.write().unwrap().set_bank(bank.clone(), false);
+            poh_recorder
+                .write()
+                .unwrap()
+                .set_bank_for_test(bank.clone());
 
             let shreds = entries_to_test_shreds(
                 &entries,
@@ -1765,7 +1783,10 @@ mod tests {
 
             let poh_simulator = simulate_poh(record_receiver, &poh_recorder);
 
-            poh_recorder.write().unwrap().set_bank(bank.clone(), false);
+            poh_recorder
+                .write()
+                .unwrap()
+                .set_bank_for_test(bank.clone());
 
             let shreds = entries_to_test_shreds(
                 &entries,
@@ -1864,7 +1885,7 @@ mod tests {
             assert_eq!(buffered_packet_batches.len(), num_conflicting_transactions);
             // When the working bank in poh_recorder is Some, all packets should be processed.
             // Multi-Iterator will process them 1-by-1 if all txs are conflicting.
-            poh_recorder.write().unwrap().set_bank(bank, false);
+            poh_recorder.write().unwrap().set_bank_for_test(bank);
             let bank_start = poh_recorder.read().unwrap().bank_start().unwrap();
             let banking_stage_stats = BankingStageStats::default();
             consumer.consume_buffered_packets(
@@ -1942,7 +1963,7 @@ mod tests {
             assert_eq!(buffered_packet_batches.len(), num_conflicting_transactions);
             // When the working bank in poh_recorder is Some, all packets should be processed.
             // Multi-Iterator will process them 1-by-1 if all txs are conflicting.
-            poh_recorder.write().unwrap().set_bank(bank, false);
+            poh_recorder.write().unwrap().set_bank_for_test(bank);
             let bank_start = poh_recorder.read().unwrap().bank_start().unwrap();
             consumer.consume_buffered_packets(
                 &bank_start,
@@ -1995,7 +2016,10 @@ mod tests {
             // When the working bank in poh_recorder is Some, all packets should be processed
             // except except for retryable errors. Manually take the lock of a transaction to
             // simulate another thread processing a transaction with that lock.
-            poh_recorder.write().unwrap().set_bank(bank.clone(), false);
+            poh_recorder
+                .write()
+                .unwrap()
+                .set_bank_for_test(bank.clone());
             let bank_start = poh_recorder.read().unwrap().bank_start().unwrap();
 
             let lock_account = transactions[0].message.account_keys[1];
@@ -2116,7 +2140,7 @@ mod tests {
             assert_eq!(buffered_packet_batches.len(), num_conflicting_transactions);
             // When the working bank in poh_recorder is Some, all packets should be processed.
             // Multi-Iterator will process them 1-by-1 if all txs are conflicting.
-            poh_recorder.write().unwrap().set_bank(bank, false);
+            poh_recorder.write().unwrap().set_bank_for_test(bank);
             let bank_start = poh_recorder.read().unwrap().bank_start().unwrap();
             let banking_stage_stats = BankingStageStats::default();
             consumer.consume_buffered_packets(

core/src/banking_stage/decision_maker.rs

@@ -164,7 +164,10 @@ mod tests {
 
         // Currently Leader - Consume
         {
-            poh_recorder.write().unwrap().set_bank(bank.clone(), false);
+            poh_recorder
+                .write()
+                .unwrap()
+                .set_bank_for_test(bank.clone());
             let decision = decision_maker.make_consume_or_forward_decision();
             assert_matches!(decision, BufferedPacketsDecision::Consume(_));
         }

poh/Cargo.toml

@@ -27,6 +27,10 @@ bincode = { workspace = true }
 rand = { workspace = true }
 solana-logger = { workspace = true }
 solana-perf = { workspace = true }
+solana-poh = { path = ".", features = ["dev-context-only-utils"] }
+
+[features]
+dev-context-only-utils = []
 
 [lib]
 crate-type = ["lib"]

poh/src/poh_recorder.rs

@@ -642,6 +642,16 @@ impl PohRecorder {
         let _ = self.flush_cache(false);
     }
 
+    #[cfg(feature = "dev-context-only-utils")]
+    pub fn set_bank_for_test(&mut self, bank: Arc<Bank>) {
+        self.set_bank(bank, false)
+    }
+
+    #[cfg(test)]
+    pub fn set_bank_with_transaction_index_for_test(&mut self, bank: Arc<Bank>) {
+        self.set_bank(bank, true)
+    }
+
     // Flush cache will delay flushing the cache for a bank until it past the WorkingBank::min_tick_height
     // On a record flush will flush the cache at the WorkingBank::min_tick_height, since a record
     // occurs after the min_tick_height was generated
@@ -1219,7 +1229,7 @@ mod tests {
                 Arc::new(AtomicBool::default()),
             );
 
-            poh_recorder.set_bank(bank, false);
+            poh_recorder.set_bank_for_test(bank);
             assert!(poh_recorder.working_bank.is_some());
             poh_recorder.clear_bank();
             assert!(poh_recorder.working_bank.is_none());
@@ -1253,7 +1263,7 @@ mod tests {
             let bank1 = Arc::new(Bank::new_from_parent(bank0, &Pubkey::default(), 1));
 
             // Set a working bank
-            poh_recorder.set_bank(bank1.clone(), false);
+            poh_recorder.set_bank_for_test(bank1.clone());
 
             // Tick until poh_recorder.tick_height == working bank's min_tick_height
             let num_new_ticks = bank1.tick_height() - poh_recorder.tick_height();
@@ -1322,7 +1332,7 @@ mod tests {
             );
             assert_eq!(poh_recorder.tick_height, bank.max_tick_height() + 1);
 
-            poh_recorder.set_bank(bank.clone(), false);
+            poh_recorder.set_bank_for_test(bank.clone());
             poh_recorder.tick();
 
             assert_eq!(poh_recorder.tick_height, bank.max_tick_height() + 2);
@@ -1363,7 +1373,7 @@ mod tests {
 
             bank0.fill_bank_with_ticks_for_tests();
             let bank1 = Arc::new(Bank::new_from_parent(bank0, &Pubkey::default(), 1));
-            poh_recorder.set_bank(bank1.clone(), false);
+            poh_recorder.set_bank_for_test(bank1.clone());
             // Let poh_recorder tick up to bank1.tick_height() - 1
             for _ in 0..bank1.tick_height() - 1 {
                 poh_recorder.tick()
@@ -1404,7 +1414,7 @@ mod tests {
                 Arc::new(AtomicBool::default()),
             );
 
-            poh_recorder.set_bank(bank.clone(), false);
+            poh_recorder.set_bank_for_test(bank.clone());
             let tx = test_tx();
             let h1 = hash(b"hello world!");
 
@@ -1448,7 +1458,7 @@ mod tests {
 
             bank0.fill_bank_with_ticks_for_tests();
             let bank1 = Arc::new(Bank::new_from_parent(bank0, &Pubkey::default(), 1));
-            poh_recorder.set_bank(bank1.clone(), false);
+            poh_recorder.set_bank_for_test(bank1.clone());
 
             // Record up to exactly min tick height
             let min_tick_height = poh_recorder.working_bank.as_ref().unwrap().min_tick_height;
@@ -1502,7 +1512,7 @@ mod tests {
                 Arc::new(AtomicBool::default()),
             );
 
-            poh_recorder.set_bank(bank.clone(), false);
+            poh_recorder.set_bank_for_test(bank.clone());
             let num_ticks_to_max = bank.max_tick_height() - poh_recorder.tick_height;
             for _ in 0..num_ticks_to_max {
                 poh_recorder.tick();
@@ -1542,7 +1552,7 @@ mod tests {
                 Arc::new(AtomicBool::default()),
             );
 
-            poh_recorder.set_bank(bank.clone(), true);
+            poh_recorder.set_bank_with_transaction_index_for_test(bank.clone());
             poh_recorder.tick();
             assert_eq!(
                 poh_recorder
@@ -1616,7 +1626,7 @@ mod tests {
 
             bank0.fill_bank_with_ticks_for_tests();
             let bank1 = Arc::new(Bank::new_from_parent(bank0, &Pubkey::default(), 1));
-            poh_recorder.set_bank(bank1, false);
+            poh_recorder.set_bank_for_test(bank1);
 
             // Check we can make two ticks without hitting min_tick_height
             let remaining_ticks_to_min =
@@ -1764,7 +1774,7 @@ mod tests {
                 Arc::new(AtomicBool::default()),
             );
 
-            poh_recorder.set_bank(bank.clone(), false);
+            poh_recorder.set_bank_for_test(bank.clone());
             assert_eq!(bank.slot(), 0);
             poh_recorder.reset(bank, Some((4, 4)));
             assert!(poh_recorder.working_bank.is_none());
@@ -1796,7 +1806,7 @@ mod tests {
                     None,
                     Arc::new(AtomicBool::default()),
                 );
-            poh_recorder.set_bank(bank, false);
+            poh_recorder.set_bank_for_test(bank);
             poh_recorder.clear_bank();
             assert!(receiver.try_recv().is_ok());
         }
@@ -1831,7 +1841,7 @@ mod tests {
                 Arc::new(AtomicBool::default()),
             );
 
-            poh_recorder.set_bank(bank.clone(), false);
+            poh_recorder.set_bank_for_test(bank.clone());
 
             // Simulate ticking much further than working_bank.max_tick_height
             let max_tick_height = poh_recorder.working_bank.as_ref().unwrap().max_tick_height;
@@ -2126,7 +2136,7 @@ mod tests {
             // Move the bank up a slot (so that max_tick_height > slot 0's tick_height)
             let bank = Arc::new(Bank::new_from_parent(bank, &Pubkey::default(), 1));
             // If we set the working bank, the node should be leader within next 2 slots
-            poh_recorder.set_bank(bank.clone(), false);
+            poh_recorder.set_bank_for_test(bank.clone());
             assert!(poh_recorder.would_be_leader(2 * bank.ticks_per_slot()));
         }
     }
@@ -2160,7 +2170,7 @@ mod tests {
             for _ in 0..(bank.ticks_per_slot() * 3) {
                 poh_recorder.tick();
             }
-            poh_recorder.set_bank(bank.clone(), false);
+            poh_recorder.set_bank_for_test(bank.clone());
             assert!(!bank.is_hash_valid_for_age(&genesis_hash, 0));
             assert!(bank.is_hash_valid_for_age(&genesis_hash, 1));
         }

poh/src/poh_service.rs

@@ -498,7 +498,7 @@ mod tests {
                 hashes_per_batch,
                 record_receiver,
             );
-            poh_recorder.write().unwrap().set_bank(bank, false);
+            poh_recorder.write().unwrap().set_bank_for_test(bank);
 
             // get some events
             let mut hashes = 0;

scripts/check-dev-context-only-utils.sh

@@ -29,6 +29,7 @@ source ci/rust-version.sh nightly
 # reason to bend dev-context-only-utils's original intention and that listed
 # package isn't part of released binaries.
 declare tainted_packages=(
+  solana-banking-bench
   solana-ledger-tool
 )