From 5f74fc4f16698e9ec39c549e80b67bc787e5f4a6 Mon Sep 17 00:00:00 2001
From: Justin Starry <justin.starry@icloud.com>
Date: Wed, 10 Jan 2024 08:34:41 +0800
Subject: [PATCH] Update genesis processing to have a fallback collector id for
 tests (#34135)

* Update genesis processing to have a fallback collector id for tests

* DCOU-ify the collector id for tests parameter (#1902)

* wrap test_collector_id in DCOU

* rename param to collector_id_for_tests

* fix program test

* fix dcou

---------

Co-authored-by: Brooks <brooks@prumo.org>
---
 ledger/src/blockstore_processor.rs |  1 +
 program-test/src/lib.rs            |  1 +
 runtime/src/bank.rs                | 30 +++++++++++++++++++++---------
 runtime/src/bank/tests.rs          |  5 +++++
 4 files changed, 28 insertions(+), 9 deletions(-)

diff --git a/ledger/src/blockstore_processor.rs b/ledger/src/blockstore_processor.rs
index cc8a4e5cb..4fa5fa6f3 100644
--- a/ledger/src/blockstore_processor.rs
+++ b/ledger/src/blockstore_processor.rs
@@ -794,6 +794,7 @@ pub(crate) fn process_blockstore_for_bank_0(
         false,
         opts.accounts_db_config.clone(),
         accounts_update_notifier,
+        None,
         exit,
     );
     let bank0_slot = bank0.slot();
diff --git a/program-test/src/lib.rs b/program-test/src/lib.rs
index fb90a12c3..b27528043 100644
--- a/program-test/src/lib.rs
+++ b/program-test/src/lib.rs
@@ -834,6 +834,7 @@ impl ProgramTest {
             false,
             None,
             None,
+            None,
             Arc::default(),
         );
 
diff --git a/runtime/src/bank.rs b/runtime/src/bank.rs
index 3a7a1db44..271acb151 100644
--- a/runtime/src/bank.rs
+++ b/runtime/src/bank.rs
@@ -1040,6 +1040,7 @@ impl Bank {
         debug_do_not_add_builtins: bool,
         accounts_db_config: Option<AccountsDbConfig>,
         accounts_update_notifier: Option<AccountsUpdateNotifier>,
+        #[allow(unused)] collector_id_for_tests: Option<Pubkey>,
         exit: Arc<AtomicBool>,
     ) -> Self {
         let accounts_db = AccountsDb::new_with_config(
@@ -1058,7 +1059,11 @@ impl Bank {
         bank.runtime_config = runtime_config;
         bank.cluster_type = Some(genesis_config.cluster_type);
 
+        #[cfg(not(feature = "dev-context-only-utils"))]
         bank.process_genesis_config(genesis_config);
+        #[cfg(feature = "dev-context-only-utils")]
+        bank.process_genesis_config(genesis_config, collector_id_for_tests);
+
         bank.finish_init(
             genesis_config,
             additional_builtins,
@@ -3736,7 +3741,11 @@ impl Bank {
         self.parent_hash
     }
 
-    fn process_genesis_config(&mut self, genesis_config: &GenesisConfig) {
+    fn process_genesis_config(
+        &mut self,
+        genesis_config: &GenesisConfig,
+        #[cfg(feature = "dev-context-only-utils")] collector_id_for_tests: Option<Pubkey>,
+    ) {
         // Bootstrap validator collects fees until `new_from_parent` is called.
         self.fee_rate_governor = genesis_config.fee_rate_governor.clone();
 
@@ -3762,14 +3771,15 @@ impl Bank {
             self.accounts_data_size_initial += account.data().len() as u64;
         }
 
-        // Highest staked node is the first collector but if a genesis config
-        // doesn't define any staked nodes, we assume this genesis config is for
-        // testing and set the collector id to a unique pubkey.
-        self.collector_id = self
-            .stakes_cache
-            .stakes()
-            .highest_staked_node()
-            .unwrap_or_else(Pubkey::new_unique);
+        // After storing genesis accounts, the bank stakes cache will be warmed
+        // up and can be used to set the collector id to the highest staked
+        // node. If no staked nodes exist, allow fallback to an unstaked test
+        // collector id during tests.
+        let collector_id = self.stakes_cache.stakes().highest_staked_node();
+        #[cfg(feature = "dev-context-only-utils")]
+        let collector_id = collector_id.or(collector_id_for_tests);
+        self.collector_id =
+            collector_id.expect("genesis processing failed because no staked nodes exist");
 
         self.blockhash_queue.write().unwrap().genesis_hash(
             &genesis_config.hash(),
@@ -8184,6 +8194,7 @@ impl Bank {
             false,
             Some(ACCOUNTS_DB_CONFIG_FOR_TESTING),
             None,
+            Some(Pubkey::new_unique()),
             Arc::default(),
         )
     }
@@ -8206,6 +8217,7 @@ impl Bank {
             false,
             Some(ACCOUNTS_DB_CONFIG_FOR_BENCHMARKS),
             None,
+            Some(Pubkey::new_unique()),
             Arc::default(),
         )
     }
diff --git a/runtime/src/bank/tests.rs b/runtime/src/bank/tests.rs
index 517102bf4..8f14e5d92 100644
--- a/runtime/src/bank/tests.rs
+++ b/runtime/src/bank/tests.rs
@@ -9035,6 +9035,7 @@ where
         false,
         Some(ACCOUNTS_DB_CONFIG_FOR_TESTING),
         None,
+        None,
         Arc::default(),
     ));
     let vote_and_stake_accounts =
@@ -12664,6 +12665,7 @@ fn test_rewards_computation_and_partitioned_distribution_two_blocks() {
         false,
         Some(accounts_db_config),
         None,
+        None,
         Arc::default(),
     );
 
@@ -13383,6 +13385,7 @@ fn test_get_reward_distribution_num_blocks_cap() {
         false,
         Some(accounts_db_config),
         None,
+        Some(Pubkey::new_unique()),
         Arc::default(),
     );
 
@@ -13850,6 +13853,7 @@ fn test_rehash_with_skipped_rewrites() {
         false,
         Some(accounts_db_config),
         None,
+        Some(Pubkey::new_unique()),
         Arc::new(AtomicBool::new(false)),
     ));
     // This test is only meaningful while the bank hash contains rewrites.
@@ -13910,6 +13914,7 @@ fn test_rebuild_skipped_rewrites() {
         false,
         Some(accounts_db_config.clone()),
         None,
+        Some(Pubkey::new_unique()),
         Arc::new(AtomicBool::new(false)),
     ));
     // This test is only meaningful while the bank hash contains rewrites.