root_bank_cache (#27985)

* BankForks: atomic root

* root bank cache

* fix and clean up

* added some tests

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<AtomicSlot>,
+}
+
+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<Slot, Arc<Bank>>,
     descendants: HashMap<Slot, HashSet<Slot>>,
-    root: Slot,
+    root: Arc<AtomicSlot>,
+
     pub snapshot_config: Option<SnapshotConfig>,
 
     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<Slot, HashSet<Slot>> {
-        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<Slot>,
     ) -> (Vec<Arc<Bank>>, 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

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;

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<RwLock<BankForks>>,
+    cached_root_bank: Arc<Bank>,
+    root_slot: ReadOnlyAtomicSlot,
+}
+
+impl RootBankCache {
+    pub fn new(bank_forks: Arc<RwLock<BankForks>>) -> 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<Bank> {
+        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());
+        }
+    }
+}