From 807c9aaa7ae2107dcfe048f51f59e681cac72679 Mon Sep 17 00:00:00 2001 From: Ashwin Sekar Date: Thu, 3 Oct 2024 19:14:27 +0000 Subject: [PATCH] bank: add lightweight bank hash cache to reduce bank forks usage --- core/src/replay_stage.rs | 45 +++++++++++++++++++++ runtime/src/bank.rs | 16 ++++++++ runtime/src/bank_forks.rs | 36 +++++++++++------ runtime/src/bank_hash_cache.rs | 71 ++++++++++++++++++++++++++++++++++ runtime/src/lib.rs | 1 + 5 files changed, 158 insertions(+), 11 deletions(-) create mode 100644 runtime/src/bank_hash_cache.rs diff --git a/core/src/replay_stage.rs b/core/src/replay_stage.rs index db0ee5aff30d53..b1f3944da0d9ca 100644 --- a/core/src/replay_stage.rs +++ b/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() + ); + } + } } diff --git a/runtime/src/bank.rs b/runtime/src/bank.rs index 77661a4bc43cf2..5fa512d1bf113a 100644 --- a/runtime/src/bank.rs +++ b/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, @@ -571,6 +572,7 @@ impl PartialEq for Bank { compute_budget: _, transaction_account_lock_limit: _, fee_structure: _, + 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. @@ -901,6 +903,9 @@ pub struct Bank { /// This _field_ was needed to be DCOU-ed to avoid 2 locks per bank freezing... #[cfg(feature = "dev-context-only-utils")] hash_overrides: Arc>, + + /// Lightweight cache of bank hashes + bank_hash_cache: Arc>, } struct VoteWithStakeDelegations { @@ -1021,6 +1026,7 @@ impl Bank { fee_structure: FeeStructure::default(), #[cfg(feature = "dev-context-only-utils")] hash_overrides: Arc::new(Mutex::new(HashOverrides::default())), + bank_hash_cache: Arc::>::default(), }; bank.transaction_processor = @@ -1282,6 +1288,7 @@ impl Bank { fee_structure: parent.fee_structure.clone(), #[cfg(feature = "dev-context-only-utils")] hash_overrides: parent.hash_overrides.clone(), + bank_hash_cache: parent.bank_hash_cache.clone(), }; let (_, ancestors_time_us) = measure_us!({ @@ -1660,6 +1667,7 @@ impl Bank { fee_structure: FeeStructure::default(), #[cfg(feature = "dev-context-only-utils")] hash_overrides: Arc::new(Mutex::new(HashOverrides::default())), + bank_hash_cache: Arc::>::default(), }; bank.transaction_processor = @@ -2852,6 +2860,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); } } @@ -6634,6 +6646,10 @@ impl Bank { self.check_program_modification_slot = check; } + pub(crate) fn set_bank_hash_cache(&mut self, bank_hash_cache: Arc>) { + self.bank_hash_cache = bank_hash_cache; + } + pub fn fee_structure(&self) -> &FeeStructure { &self.fee_structure } diff --git a/runtime/src/bank_forks.rs b/runtime/src/bank_forks.rs index aa3d78ea128ecd..3b8aaee1897407 100644 --- a/runtime/src/bank_forks.rs +++ b/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, highest_slot_at_startup: Slot, scheduler_pool: Option, + + /// Lightweight cache for just the latest bank hash + bank_hash_cache: Arc>, } impl Index for BankForks { @@ -88,7 +92,9 @@ impl Index for BankForks { } impl BankForks { - pub fn new_rw_arc(root_bank: Bank) -> Arc> { + pub fn new_rw_arc(mut root_bank: Bank) -> Arc> { + 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> { + self.bank_hash_cache.clone() + } + pub fn root_bank(&self) -> Arc { 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(); diff --git a/runtime/src/bank_hash_cache.rs b/runtime/src/bank_hash_cache.rs new file mode 100644 index 00000000000000..1b58804dd47410 --- /dev/null +++ b/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, +} + +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 { + 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 { + self.hashes.get(&slot).copied() + } + + /// Removes `slots` from the cache. Intended to be used with `BankForks::prune_non_rooted` + pub fn prune(&mut self, slots: I) + where + I: Iterator, + { + 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() + ); + } +} diff --git a/runtime/src/lib.rs b/runtime/src/lib.rs index bd11e97668eec0..ce7cc1eca344f2 100644 --- a/runtime/src/lib.rs +++ b/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;