Use RootBankCache

anza-xyz · Jan 22, 2025 · a3e78c2 · a3e78c2
1 parent cf20e54
commit a3e78c2
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 5 deletions.
diff --git a/core/src/banking_stage/unified_scheduler.rs b/core/src/banking_stage/unified_scheduler.rs
@@ -35,7 +35,7 @@ use {
     crate::banking_trace::Channels,
     agave_banking_stage_ingress_types::BankingPacketBatch,
     solana_poh::poh_recorder::PohRecorder,
-    solana_runtime::bank_forks::BankForks,
+    solana_runtime::{bank_forks::BankForks, root_bank_cache::RootBankCache},
     solana_unified_scheduler_pool::{BankingStageHelper, DefaultSchedulerPool},
     std::sync::{Arc, RwLock},
 };
@@ -49,7 +49,7 @@ pub(crate) fn ensure_banking_stage_setup(
     cluster_info: &impl LikeClusterInfo,
     poh_recorder: &Arc<RwLock<PohRecorder>>,
 ) {
-    let bank_forks = bank_forks.clone();
+    let mut root_bank_cache = RootBankCache::new(bank_forks.clone());
     let unified_receiver = channels.unified_receiver().clone();
     let decision_maker = DecisionMaker::new(cluster_info.id(), poh_recorder.clone());
     let transaction_recorder = poh_recorder.read().unwrap().new_recorder();
@@ -60,7 +60,7 @@ pub(crate) fn ensure_banking_stage_setup(
             if matches!(decision, BufferedPacketsDecision::Forward) {
                 return;
             }
-            let bank = bank_forks.read().unwrap().root_bank();
+            let bank = root_bank_cache.root_bank();
             for batch in batches.iter() {
                 // over-provision nevertheless some of packets could be invalid.
                 let task_id_base = helper.generate_task_ids(batch.len());

diff --git a/runtime/src/bank_forks.rs b/runtime/src/bank_forks.rs
@@ -33,6 +33,7 @@ use {
 
 pub const MAX_ROOT_DISTANCE_FOR_VOTE_ONLY: Slot = 400;
 pub type AtomicSlot = AtomicU64;
+#[derive(Clone)]
 pub struct ReadOnlyAtomicSlot {
     slot: Arc<AtomicSlot>,
 }

diff --git a/runtime/src/root_bank_cache.rs b/runtime/src/root_bank_cache.rs
@@ -11,6 +11,7 @@ use {
 };
 
 /// Cached root bank that only loads from bank forks if the root has been updated.
+#[derive(Clone)]
 pub struct RootBankCache {
     bank_forks: Arc<RwLock<BankForks>>,
     cached_root_bank: Weak<Bank>,

diff --git a/unified-scheduler-pool/src/lib.rs b/unified-scheduler-pool/src/lib.rs
@@ -160,7 +160,8 @@ struct BankingStageHandlerContext {
 }
 
 trait_set! {
-    pub trait BankingPacketHandler = DynClone + Fn(&BankingStageHelper, BankingPacketBatch) + Send + 'static;
+    pub trait BankingPacketHandler =
+        DynClone + FnMut(&BankingStageHelper, BankingPacketBatch) + Send + 'static;
 }
 // Make this `Clone`-able so that it can easily propagated to all the handler threads.
 clone_trait_object!(BankingPacketHandler);
@@ -1329,7 +1330,7 @@ impl<S: SpawnableScheduler<TH>, TH: TaskHandler> ThreadManager<S, TH> {
         };
 
         let handler_main_loop = || {
-            let handler_context = handler_context.clone();
+            let mut handler_context = handler_context.clone();
             let mut runnable_task_receiver = runnable_task_receiver.clone();
             let finished_blocked_task_sender = finished_blocked_task_sender.clone();
             let finished_idle_task_sender = finished_idle_task_sender.clone();