bank: add lightweight bank hash cache to reduce bank forks usage

core/src/replay_stage.rs

@@ -9276,4 +9276,49 @@ pub(crate) mod tests {
             &mut PurgeRepairSlotCounter::default(),
         );
     }
+
+    #[test]
+    fn test_bank_hash_cache_root() {
+        /*
+            Build fork structure:
+
+                 slot 0
+                   |
+                 slot 1
+                 /    \
+            slot 2    |
+               |    slot 3
+            slot 4    |
+                    slot 5
+                      |
+                    slot 6
+        */
+        let (vote_simulator, _blockstore) = setup_default_forks(1, None);
+        let VoteSimulator { bank_forks, .. } = vote_simulator;
+        let bank_hash_cache = bank_forks.read().unwrap().bank_hash_cache();
+
+        for slot in 0..=6 {
+            assert_eq!(
+                bank_hash_cache.read().unwrap().bank_hash(slot).unwrap(),
+                bank_forks.read().unwrap().bank_hash(slot).unwrap(),
+                "bank hash mismatch {slot}"
+            );
+        }
+
+        bank_forks
+            .write()
+            .unwrap()
+            .set_root(3, &AbsRequestSender::default(), None)
+            .unwrap();
+        for slot in [0, 1, 2, 4] {
+            assert!(bank_forks.read().unwrap().get(slot).is_none());
+            assert!(bank_hash_cache.read().unwrap().bank_hash(slot).is_none());
+        }
+        for slot in [3, 5, 6] {
+            assert_eq!(
+                bank_hash_cache.read().unwrap().bank_hash(slot).unwrap(),
+                bank_forks.read().unwrap().bank_hash(slot).unwrap()
+            );
+        }
+    }
 }

runtime/src/bank.rs

@@ -42,6 +42,7 @@ use {
             partitioned_epoch_rewards::{EpochRewardStatus, StakeRewards, VoteRewardsAccounts},
         },
         bank_forks::BankForks,
+        bank_hash_cache::BankHashCache,
         epoch_stakes::{split_epoch_stakes, EpochStakes, NodeVoteAccounts, VersionedEpochStakes},
         installed_scheduler_pool::{BankWithScheduler, InstalledSchedulerRwLock},
         rent_collector::RentCollectorWithMetrics,
@@ -579,6 +580,7 @@ impl PartialEq for Bank {
             fee_structure: _,
             accounts_lt_hash: _,
             cache_for_accounts_lt_hash: _,
+            bank_hash_cache: _,
             // Ignore new fields explicitly if they do not impact PartialEq.
             // Adding ".." will remove compile-time checks that if a new field
             // is added to the struct, this PartialEq is accordingly updated.
@@ -920,6 +922,9 @@ pub struct Bank {
     /// The accounts lt hash needs both the initial and final state of each
     /// account that was modified in this slot.  Cache the initial state here.
     cache_for_accounts_lt_hash: RwLock<AHashMap<Pubkey, InitialStateOfAccount>>,
+
+    /// Lightweight cache of bank hashes
+    bank_hash_cache: Arc<RwLock<BankHashCache>>,
 }
 
 struct VoteWithStakeDelegations {
@@ -1042,6 +1047,7 @@ impl Bank {
             hash_overrides: Arc::new(Mutex::new(HashOverrides::default())),
             accounts_lt_hash: Mutex::new(AccountsLtHash(LtHash([0xBAD1; LtHash::NUM_ELEMENTS]))),
             cache_for_accounts_lt_hash: RwLock::new(AHashMap::new()),
+            bank_hash_cache: Arc::<RwLock<BankHashCache>>::default(),
         };
 
         bank.transaction_processor =
@@ -1316,6 +1322,7 @@ impl Bank {
             hash_overrides: parent.hash_overrides.clone(),
             accounts_lt_hash: Mutex::new(parent.accounts_lt_hash.lock().unwrap().clone()),
             cache_for_accounts_lt_hash: RwLock::new(AHashMap::new()),
+            bank_hash_cache: parent.bank_hash_cache.clone(),
         };
 
         let (_, ancestors_time_us) = measure_us!({
@@ -1696,6 +1703,7 @@ impl Bank {
             hash_overrides: Arc::new(Mutex::new(HashOverrides::default())),
             accounts_lt_hash: Mutex::new(AccountsLtHash(LtHash([0xBAD2; LtHash::NUM_ELEMENTS]))),
             cache_for_accounts_lt_hash: RwLock::new(AHashMap::new()),
+            bank_hash_cache: Arc::<RwLock<BankHashCache>>::default(),
         };
 
         bank.transaction_processor =
@@ -2904,6 +2912,10 @@ impl Bank {
             self.freeze_started.store(true, Relaxed);
             *hash = self.hash_internal_state();
             self.rc.accounts.accounts_db.mark_slot_frozen(self.slot());
+            self.bank_hash_cache
+                .write()
+                .unwrap()
+                .freeze(self.slot, *hash);
         }
     }
 
@@ -6703,6 +6715,10 @@ impl Bank {
         self.check_program_modification_slot = check;
     }
 
+    pub(crate) fn set_bank_hash_cache(&mut self, bank_hash_cache: Arc<RwLock<BankHashCache>>) {
+        self.bank_hash_cache = bank_hash_cache;
+    }
+
     pub fn fee_structure(&self) -> &FeeStructure {
         &self.fee_structure
     }

runtime/src/bank_forks.rs

@@ -4,6 +4,7 @@ use {
     crate::{
         accounts_background_service::{AbsRequestSender, SnapshotRequest, SnapshotRequestKind},
         bank::{epoch_accounts_hash_utils, Bank, SquashTiming},
+        bank_hash_cache::BankHashCache,
         installed_scheduler_pool::{
             BankWithScheduler, InstalledSchedulerPoolArc, SchedulingContext,
         },
@@ -78,6 +79,9 @@ pub struct BankForks {
     in_vote_only_mode: Arc<AtomicBool>,
     highest_slot_at_startup: Slot,
     scheduler_pool: Option<InstalledSchedulerPoolArc>,
+
+    /// Lightweight cache for just the latest bank hash
+    bank_hash_cache: Arc<RwLock<BankHashCache>>,
 }
 
 impl Index<u64> for BankForks {
@@ -88,7 +92,9 @@ impl Index<u64> for BankForks {
 }
 
 impl BankForks {
-    pub fn new_rw_arc(root_bank: Bank) -> Arc<RwLock<Self>> {
+    pub fn new_rw_arc(mut root_bank: Bank) -> Arc<RwLock<Self>> {
+        let bank_hash_cache = Arc::new(RwLock::new(BankHashCache::default()));
+        root_bank.set_bank_hash_cache(bank_hash_cache.clone());
         let root_bank = Arc::new(root_bank);
         let root_slot = root_bank.slot();
 
@@ -100,16 +106,14 @@ impl BankForks {
 
         let parents = root_bank.parents();
         for parent in parents {
-            if banks
-                .insert(
-                    parent.slot(),
-                    BankWithScheduler::new_without_scheduler(parent.clone()),
-                )
-                .is_some()
-            {
-                // All ancestors have already been inserted by another fork
-                break;
-            }
+            banks.insert(
+                parent.slot(),
+                BankWithScheduler::new_without_scheduler(parent.clone()),
+            );
+            bank_hash_cache
+                .write()
+                .unwrap()
+                .freeze(parent.slot(), parent.hash());
         }
 
         let mut descendants = HashMap::<_, HashSet<_>>::new();
@@ -128,6 +132,7 @@ impl BankForks {
             in_vote_only_mode: Arc::new(AtomicBool::new(false)),
             highest_slot_at_startup: 0,
             scheduler_pool: None,
+            bank_hash_cache,
         }));
 
         root_bank.set_fork_graph_in_program_cache(Arc::downgrade(&bank_forks));
@@ -207,6 +212,10 @@ impl BankForks {
         self.get(slot).map(|bank| bank.hash())
     }
 
+    pub fn bank_hash_cache(&self) -> Arc<RwLock<BankHashCache>> {
+        self.bank_hash_cache.clone()
+    }
+
     pub fn root_bank(&self) -> Arc<Bank> {
         self[self.root()].clone()
     }
@@ -223,6 +232,7 @@ impl BankForks {
         if self.root.load(Ordering::Relaxed) < self.highest_slot_at_startup {
             bank.set_check_program_modification_slot(true);
         }
+        bank.set_bank_hash_cache(self.bank_hash_cache());
 
         let bank = Arc::new(bank);
         let bank = if let Some(scheduler_pool) = &self.scheduler_pool {
@@ -445,6 +455,10 @@ impl BankForks {
         let dropped_banks_len = removed_banks.len();
 
         let mut drop_parent_banks_time = Measure::start("set_root::drop_banks");
+        self.bank_hash_cache
+            .write()
+            .unwrap()
+            .prune(removed_banks.iter().map(|bank| bank.slot()));
         drop(parents);
         drop_parent_banks_time.stop();
 

runtime/src/bank_hash_cache.rs

@@ -0,0 +1,71 @@
+//! Lightweight cache tracking the bank hashes of replayed banks.
+//! This can be useful to avoid read-locking bank forks just to query the bank hash.
+
+use {
+    solana_sdk::{clock::Slot, hash::Hash},
+    std::collections::BTreeMap,
+};
+
+#[derive(Default, Debug)]
+pub struct BankHashCache {
+    hashes: BTreeMap<Slot, Hash>,
+}
+
+impl BankHashCache {
+    /// Insert new frozen bank, returning the hash of the previously dumped bank if it exists
+    pub fn freeze(&mut self, slot: Slot, hash: Hash) -> Option<Hash> {
+        self.hashes.insert(slot, hash)
+    }
+
+    /// Returns the replayed bank hash of `slot`.
+    /// If `slot` has been dumped, returns the previously replayed hash.
+    pub fn bank_hash(&self, slot: Slot) -> Option<Hash> {
+        self.hashes.get(&slot).copied()
+    }
+
+    /// Removes `slots` from the cache. Intended to be used with `BankForks::prune_non_rooted`
+    pub fn prune<I>(&mut self, slots: I)
+    where
+        I: Iterator<Item = Slot>,
+    {
+        for slot in slots {
+            self.hashes.remove(&slot);
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use {
+        crate::{
+            bank::Bank,
+            bank_forks::BankForks,
+            genesis_utils::{create_genesis_config, GenesisConfigInfo},
+        },
+        solana_sdk::pubkey::Pubkey,
+    };
+
+    #[test]
+    fn test_bank_hash_cache() {
+        let GenesisConfigInfo { genesis_config, .. } = create_genesis_config(10_000);
+        let bank0 = Bank::new_for_tests(&genesis_config);
+        let slot = bank0.slot();
+        let bank_forks = BankForks::new_rw_arc(bank0);
+        let bank_hash_cache = bank_forks.read().unwrap().bank_hash_cache();
+        bank_forks.read().unwrap()[slot].freeze();
+        assert_eq!(
+            bank_hash_cache.read().unwrap().bank_hash(slot).unwrap(),
+            bank_forks.read().unwrap()[slot].hash()
+        );
+
+        let bank0 = bank_forks.read().unwrap().get(slot).unwrap();
+        let slot = 10;
+        let bank10 = Bank::new_from_parent(bank0, &Pubkey::new_unique(), slot);
+        bank_forks.write().unwrap().insert(bank10);
+        bank_forks.read().unwrap()[slot].freeze();
+        assert_eq!(
+            bank_hash_cache.read().unwrap().bank_hash(slot).unwrap(),
+            bank_forks.read().unwrap()[slot].hash()
+        );
+    }
+}

runtime/src/lib.rs

@@ -9,6 +9,7 @@ pub mod accounts_background_service;
 pub mod bank;
 pub mod bank_client;
 pub mod bank_forks;
+pub mod bank_hash_cache;
 pub mod bank_utils;
 pub mod commitment;
 pub mod epoch_stakes;