PohRecorder: SharedWorkingBank (#7280)

Author: apfitzge

Cargo.lock

@@ -549,7 +549,7 @@ impl BankingStage {
                     log_messages_bytes_limit,
                 ),
                 finished_work_sender.clone(),
-                poh_recorder.read().unwrap().new_leader_bank_notifier(),
+                poh_recorder.read().unwrap().shared_working_bank(),
             );
 
             worker_metrics.push(consume_worker.metrics_handle());

core/src/banking_stage.rs

@@ -6,7 +6,7 @@ use {
     },
     crossbeam_channel::{Receiver, RecvError, SendError, Sender},
     solana_measure::measure_us,
-    solana_poh::leader_bank_notifier::LeaderBankNotifier,
+    solana_poh::poh_recorder::SharedWorkingBank,
     solana_runtime::bank::Bank,
     solana_runtime_transaction::transaction_with_meta::TransactionWithMeta,
     solana_svm::transaction_error_metrics::TransactionErrorMetrics,
@@ -16,7 +16,7 @@ use {
             atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering},
             Arc,
         },
-        time::Duration,
+        time::{Duration, Instant},
     },
     thiserror::Error,
 };
@@ -34,7 +34,7 @@ pub(crate) struct ConsumeWorker<Tx> {
     consumer: Consumer,
     consumed_sender: Sender<FinishedConsumeWork<Tx>>,
 
-    leader_bank_notifier: Arc<LeaderBankNotifier>,
+    shared_working_bank: SharedWorkingBank,
     metrics: Arc<ConsumeWorkerMetrics>,
 }
 
@@ -44,13 +44,13 @@ impl<Tx: TransactionWithMeta> ConsumeWorker<Tx> {
         consume_receiver: Receiver<ConsumeWork<Tx>>,
         consumer: Consumer,
         consumed_sender: Sender<FinishedConsumeWork<Tx>>,
-        leader_bank_notifier: Arc<LeaderBankNotifier>,
+        shared_working_bank: SharedWorkingBank,
     ) -> Self {
         Self {
             consume_receiver,
             consumer,
             consumed_sender,
-            leader_bank_notifier,
+            shared_working_bank,
             metrics: Arc::new(ConsumeWorkerMetrics::new(id)),
         }
     }
@@ -67,7 +67,7 @@ impl<Tx: TransactionWithMeta> ConsumeWorker<Tx> {
     }
 
     fn consume_loop(&self, work: ConsumeWork<Tx>) -> Result<(), ConsumeWorkerError<Tx>> {
-        let (maybe_consume_bank, get_bank_us) = measure_us!(self.get_consume_bank());
+        let (maybe_consume_bank, get_bank_us) = measure_us!(self.working_bank_with_timeout());
         let Some(mut bank) = maybe_consume_bank else {
             self.metrics
                 .timing_metrics
@@ -82,10 +82,12 @@ impl<Tx: TransactionWithMeta> ConsumeWorker<Tx> {
 
         for work in try_drain_iter(work, &self.consume_receiver) {
             if bank.is_complete() || {
-                // check if the bank got interrupted before completion
-                self.get_consume_bank_id() != Some(bank.bank_id())
+                // if working bank has changed, then try to get a new bank.
+                self.working_bank()
+                    .map(|working_bank| Arc::ptr_eq(&working_bank, &bank))
+                    .unwrap_or(true)
             } {
-                let (maybe_new_bank, get_bank_us) = measure_us!(self.get_consume_bank());
+                let (maybe_new_bank, get_bank_us) = measure_us!(self.working_bank_with_timeout());
                 if let Some(new_bank) = maybe_new_bank {
                     self.metrics
                         .timing_metrics
@@ -130,16 +132,23 @@ impl<Tx: TransactionWithMeta> ConsumeWorker<Tx> {
         Ok(())
     }
 
-    /// Try to get a bank for consuming.
-    fn get_consume_bank(&self) -> Option<Arc<Bank>> {
-        self.leader_bank_notifier
-            .get_or_wait_for_in_progress(Duration::from_millis(50))
-            .upgrade()
+    /// Get the current poh working bank with a timeout - if the Bank is
+    /// not available within the timeout, return None.
+    fn working_bank_with_timeout(&self) -> Option<Arc<Bank>> {
+        const TIMEOUT: Duration = Duration::from_millis(50);
+        let now = Instant::now();
+        while now.elapsed() < TIMEOUT {
+            if let Some(bank) = self.working_bank() {
+                return Some(bank);
+            }
+        }
+
+        None
     }
 
-    /// Try to get the id for the bank that should be used for consuming
-    fn get_consume_bank_id(&self) -> Option<u64> {
-        self.leader_bank_notifier.get_current_bank_id()
+    /// Get the current poh working bank without a timeout.
+    fn working_bank(&self) -> Option<Arc<Bank>> {
+        self.shared_working_bank.load()
     }
 
     /// Retry current batch and all outstanding batches.
@@ -845,7 +854,7 @@ mod tests {
             consume_receiver,
             consumer,
             consumed_sender,
-            poh_recorder.read().unwrap().new_leader_bank_notifier(),
+            poh_recorder.read().unwrap().shared_working_bank(),
         );
 
         (

core/src/banking_stage/consume_worker.rs

@@ -20,6 +20,7 @@ name = "solana_poh"
 dev-context-only-utils = []
 
 [dependencies]
+arc-swap = { workspace = true }
 core_affinity = { workspace = true }
 crossbeam-channel = { workspace = true }
 log = { workspace = true }