diff --git a/runtime/src/bank_forks.rs b/runtime/src/bank_forks.rs index 3252ffef44..09c60eee74 100644 --- a/runtime/src/bank_forks.rs +++ b/runtime/src/bank_forks.rs @@ -12,12 +12,25 @@ use { std::{ collections::{hash_map::Entry, HashMap, HashSet}, ops::Index, - sync::{atomic::AtomicBool, Arc}, + sync::{ + atomic::{AtomicBool, AtomicU64, Ordering}, + Arc, + }, time::Instant, }, }; pub const MAX_ROOT_DISTANCE_FOR_VOTE_ONLY: Slot = 400; +pub type AtomicSlot = AtomicU64; +pub struct ReadOnlyAtomicSlot { + slot: Arc, +} + +impl ReadOnlyAtomicSlot { + pub fn get(&self) -> Slot { + self.slot.load(Ordering::Relaxed) + } +} #[derive(Debug, Default, Copy, Clone)] struct SetRootMetrics { @@ -45,7 +58,8 @@ struct SetRootTimings { pub struct BankForks { banks: HashMap>, descendants: HashMap>, - root: Slot, + root: Arc, + pub snapshot_config: Option, pub accounts_hash_interval_slots: Slot, @@ -84,7 +98,7 @@ impl BankForks { /// Create a map of bank slot id to the set of ancestors for the bank slot. pub fn ancestors(&self) -> HashMap> { - let root = self.root; + let root = self.root(); self.banks .iter() .map(|(slot, bank)| { @@ -160,7 +174,7 @@ impl BankForks { } } Self { - root, + root: Arc::new(AtomicSlot::new(root)), banks, descendants, snapshot_config: None, @@ -219,7 +233,8 @@ impl BankForks { highest_confirmed_root: Option, ) -> (Vec>, SetRootMetrics) { let old_epoch = self.root_bank().epoch(); - self.root = root; + self.root.store(root, Ordering::Relaxed); + let root_bank = self .banks .get(&root) @@ -427,7 +442,14 @@ impl BankForks { } pub fn root(&self) -> Slot { - self.root + self.root.load(Ordering::Relaxed) + } + + /// Gets a read-only wrapper to an atomic slot holding the root slot. + pub fn get_atomic_root(&self) -> ReadOnlyAtomicSlot { + ReadOnlyAtomicSlot { + slot: self.root.clone(), + } } /// After setting a new root, prune the banks that are no longer on rooted paths diff --git a/runtime/src/lib.rs b/runtime/src/lib.rs index 9a86795eb7..18a6b6b367 100644 --- a/runtime/src/lib.rs +++ b/runtime/src/lib.rs @@ -55,6 +55,7 @@ mod read_only_accounts_cache; pub mod rent_collector; mod rent_paying_accounts_by_partition; mod rolling_bit_field; +pub mod root_bank_cache; pub mod runtime_config; pub mod secondary_index; pub mod serde_snapshot; diff --git a/runtime/src/root_bank_cache.rs b/runtime/src/root_bank_cache.rs new file mode 100644 index 0000000000..6723e26fe5 --- /dev/null +++ b/runtime/src/root_bank_cache.rs @@ -0,0 +1,84 @@ +//! A wrapper around a root `Bank` that only loads from bank forks if the root has been updated. +//! This can be useful to avoid read-locking the bank forks until the root has been updated. +//! + +use { + crate::{ + bank::Bank, + bank_forks::{BankForks, ReadOnlyAtomicSlot}, + }, + std::sync::{Arc, RwLock}, +}; + +/// Cached root bank that only loads from bank forks if the root has been updated. +pub struct RootBankCache { + bank_forks: Arc>, + cached_root_bank: Arc, + root_slot: ReadOnlyAtomicSlot, +} + +impl RootBankCache { + pub fn new(bank_forks: Arc>) -> Self { + let (cached_root_bank, root_slot) = { + let lock = bank_forks.read().unwrap(); + (lock.root_bank(), lock.get_atomic_root()) + }; + Self { + bank_forks, + cached_root_bank, + root_slot, + } + } + + pub fn root_bank(&mut self) -> Arc { + let current_root_slot = self.root_slot.get(); + if self.cached_root_bank.slot() != current_root_slot { + let lock = self.bank_forks.read().unwrap(); + self.cached_root_bank = lock.root_bank(); + } + self.cached_root_bank.clone() + } +} + +#[cfg(test)] +mod tests { + use { + super::*, + crate::{ + accounts_background_service::AbsRequestSender, + bank_forks::BankForks, + genesis_utils::{create_genesis_config, GenesisConfigInfo}, + }, + solana_sdk::pubkey::Pubkey, + }; + + #[test] + fn test_root_bank_cache() { + let GenesisConfigInfo { genesis_config, .. } = create_genesis_config(10_000); + let bank = Bank::new_for_tests(&genesis_config); + let bank_forks = Arc::new(RwLock::new(BankForks::new(bank))); + + let mut root_bank_cache = RootBankCache::new(bank_forks.clone()); + + let bank = bank_forks.read().unwrap().root_bank(); + assert_eq!(bank, root_bank_cache.root_bank()); + + { + let child_bank = Bank::new_from_parent(&bank, &Pubkey::default(), 1); + bank_forks.write().unwrap().insert(child_bank); + + // cached slot is still 0 since we have not set root + assert_eq!(bank.slot(), root_bank_cache.cached_root_bank.slot()); + } + { + bank_forks + .write() + .unwrap() + .set_root(1, &AbsRequestSender::default(), None); + let bank = bank_forks.read().unwrap().root_bank(); + assert!(bank.slot() != root_bank_cache.cached_root_bank.slot()); + assert!(bank != root_bank_cache.cached_root_bank); + assert_eq!(bank, root_bank_cache.root_bank()); + } + } +}