Wire unified scheduler into banking experimentally

banking-bench/Cargo.toml

@@ -28,6 +28,8 @@ solana-runtime = { workspace = true, features = ["dev-context-only-utils"] }
 solana-sdk = { workspace = true }
 solana-streamer = { workspace = true }
 solana-tpu-client = { workspace = true }
+solana-unified-scheduler-logic = { workspace = true }
+solana-unified-scheduler-pool = { workspace = true }
 solana-version = { workspace = true }
 
 [features]

banking-bench/src/main.rs

@@ -8,7 +8,10 @@ use {
     rand::{thread_rng, Rng},
     rayon::prelude::*,
     solana_core::{
-        banking_stage::{update_bank_forks_and_poh_recorder_for_new_tpu_bank, BankingStage},
+        banking_stage::{
+            unified_scheduler::ensure_banking_stage_setup,
+            update_bank_forks_and_poh_recorder_for_new_tpu_bank, BankingStage,
+        },
         banking_trace::{BankingTracer, Channels, BANKING_TRACE_DIR_DEFAULT_BYTE_LIMIT},
         validator::{BlockProductionMethod, TransactionStructure},
     },
@@ -36,6 +39,8 @@ use {
         transaction::Transaction,
     },
     solana_streamer::socket::SocketAddrSpace,
+    solana_unified_scheduler_logic::SchedulingMode,
+    solana_unified_scheduler_pool::{DefaultSchedulerPool, SupportedSchedulingMode},
     std::{
         sync::{atomic::Ordering, Arc, RwLock},
         thread::sleep,
@@ -452,16 +457,35 @@ fn main() {
         ClusterInfo::new(node.info, keypair, SocketAddrSpace::Unspecified)
     };
     let cluster_info = Arc::new(cluster_info);
+    let banking_tracer_channels = if matches!(
+        block_production_method,
+        BlockProductionMethod::UnifiedScheduler
+    ) {
+        let pool = DefaultSchedulerPool::new(
+            SupportedSchedulingMode::Either(SchedulingMode::BlockProduction),
+            None,
+            None,
+            None,
+            Some(replay_vote_sender.clone()),
+            prioritization_fee_cache.clone(),
+        );
+        let channels = banking_tracer.create_channels_for_scheduler_pool(&pool);
+        ensure_banking_stage_setup(&pool, &bank_forks, &channels, &cluster_info, &poh_recorder);
+        bank_forks.write().unwrap().install_scheduler_pool(pool);
+        channels
+    } else {
+        banking_tracer.create_channels(false)
+    };
     let Channels {
         non_vote_sender,
         non_vote_receiver,
         tpu_vote_sender,
         tpu_vote_receiver,
         gossip_vote_sender,
         gossip_vote_receiver,
-    } = banking_tracer.create_channels(false);
+    } = banking_tracer_channels;
     let banking_stage = BankingStage::new_num_threads(
-        block_production_method,
+        block_production_method.clone(),
         transaction_struct,
         &cluster_info,
         &poh_recorder,
@@ -476,6 +500,18 @@ fn main() {
         &prioritization_fee_cache,
     );
 
+    // This bench processes transactions, starting from the very first bank, so special-casing is
+    // needed for unified scheduler.
+    if matches!(
+        block_production_method,
+        BlockProductionMethod::UnifiedScheduler
+    ) {
+        bank = bank_forks
+            .write()
+            .unwrap()
+            .reinstall_block_production_scheduler_into_working_genesis_bank();
+    }
+
     // This is so that the signal_receiver does not go out of scope after the closure.
     // If it is dropped before poh_service, then poh_service will error when
     // calling send() on the channel.
@@ -536,10 +572,11 @@ fn main() {
             tx_total_us += now.elapsed().as_micros() as u64;
 
             let mut poh_time = Measure::start("poh_time");
-            poh_recorder
+            let cleared_bank = poh_recorder
                 .write()
                 .unwrap()
                 .reset(bank.clone(), Some((bank.slot(), bank.slot() + 1)));
+            assert_matches!(cleared_bank, None);
             poh_time.stop();
 
             let mut new_bank_time = Measure::start("new_bank");

core/Cargo.toml

@@ -28,6 +28,7 @@ bytes = { workspace = true }
 chrono = { workspace = true, features = ["default", "serde"] }
 crossbeam-channel = { workspace = true }
 dashmap = { workspace = true, features = ["rayon", "raw-api"] }
+derive_more = { workspace = true }
 etcd-client = { workspace = true, features = ["tls"] }
 futures = { workspace = true }
 histogram = { workspace = true }
@@ -94,6 +95,7 @@ solana-tls-utils = { workspace = true }
 solana-tpu-client = { workspace = true }
 solana-transaction-status = { workspace = true }
 solana-turbine = { workspace = true }
+solana-unified-scheduler-logic = { workspace = true }
 solana-unified-scheduler-pool = { workspace = true }
 solana-version = { workspace = true }
 solana-vote = { workspace = true }
@@ -125,7 +127,6 @@ solana-rpc = { workspace = true, features = ["dev-context-only-utils"] }
 solana-sdk = { workspace = true, features = ["dev-context-only-utils"] }
 solana-stake-program = { workspace = true }
 solana-system-program = { workspace = true }
-solana-unified-scheduler-logic = { workspace = true }
 solana-unified-scheduler-pool = { workspace = true, features = [
     "dev-context-only-utils",
 ] }

core/src/banking_simulation.rs

@@ -2,6 +2,7 @@
 use {
     crate::{
         banking_stage::{
+            unified_scheduler::ensure_banking_stage_setup,
             update_bank_forks_and_poh_recorder_for_new_tpu_bank, BankingStage, LikeClusterInfo,
         },
         banking_trace::{
@@ -45,6 +46,7 @@ use {
     },
     solana_streamer::socket::SocketAddrSpace,
     solana_turbine::broadcast_stage::{BroadcastStage, BroadcastStageType},
+    solana_unified_scheduler_pool::DefaultSchedulerPool,
     std::{
         collections::BTreeMap,
         fmt::Display,
@@ -691,6 +693,7 @@ impl BankingSimulator {
         blockstore: Arc<Blockstore>,
         block_production_method: BlockProductionMethod,
         transaction_struct: TransactionStructure,
+        unified_scheduler_pool: Option<Arc<DefaultSchedulerPool>>,
     ) -> (SenderLoop, SimulatorLoop, SimulatorThreads) {
         let parent_slot = self.parent_slot().unwrap();
         let mut packet_batches_by_time = self.banking_trace_events.packet_batches_by_time;
@@ -777,14 +780,27 @@ impl BankingSimulator {
         let cluster_info_for_banking = Arc::new(DummyClusterInfo {
             id: simulated_leader.into(),
         });
+        let banking_tracer_channels = if let Some(pool) = unified_scheduler_pool {
+            let channels = retracer.create_channels_for_scheduler_pool(&pool);
+            ensure_banking_stage_setup(
+                &pool,
+                &bank_forks,
+                &channels,
+                &cluster_info_for_banking,
+                &poh_recorder,
+            );
+            channels
+        } else {
+            retracer.create_channels(false)
+        };
         let Channels {
             non_vote_sender,
             non_vote_receiver,
             tpu_vote_sender,
             tpu_vote_receiver,
             gossip_vote_sender,
             gossip_vote_receiver,
-        } = retracer.create_channels(false);
+        } = banking_tracer_channels;
 
         let (replay_vote_sender, _replay_vote_receiver) = unbounded();
         let (retransmit_slots_sender, retransmit_slots_receiver) = unbounded();
@@ -907,13 +923,15 @@ impl BankingSimulator {
         blockstore: Arc<Blockstore>,
         block_production_method: BlockProductionMethod,
         transaction_struct: TransactionStructure,
+        unified_scheduler_pool: Option<Arc<DefaultSchedulerPool>>,
     ) -> Result<(), SimulateError> {
         let (sender_loop, simulator_loop, simulator_threads) = self.prepare_simulation(
             genesis_config,
             bank_forks,
             blockstore,
             block_production_method,
             transaction_struct,
+            unified_scheduler_pool,
         );
 
         sender_loop.log_starting();

core/src/banking_stage.rs

@@ -39,6 +39,7 @@ use {
         vote_sender_types::ReplayVoteSender,
     },
     solana_sdk::{pubkey::Pubkey, timing::AtomicInterval},
+    solana_unified_scheduler_logic::SchedulingMode,
     std::{
         cmp, env,
         ops::Deref,
@@ -391,6 +392,9 @@ impl BankingStage {
                     prioritization_fee_cache,
                 )
             }
+            BlockProductionMethod::UnifiedScheduler => Self {
+                bank_thread_hdls: vec![],
+            },
         }
     }
 
@@ -735,11 +739,15 @@ pub(crate) fn update_bank_forks_and_poh_recorder_for_new_tpu_bank(
     tpu_bank: Bank,
     track_transaction_indexes: bool,
 ) {
-    let tpu_bank = bank_forks.write().unwrap().insert(tpu_bank);
+    let tpu_bank = bank_forks
+        .write()
+        .unwrap()
+        .insert_with_scheduling_mode(SchedulingMode::BlockProduction, tpu_bank);
     poh_recorder
         .write()
         .unwrap()
-        .set_bank(tpu_bank, track_transaction_indexes);
+        .set_bank(tpu_bank.clone_with_scheduler(), track_transaction_indexes);
+    tpu_bank.unblock_block_production();
 }
 
 #[cfg(test)]

core/src/banking_stage/decision_maker.rs

@@ -7,8 +7,9 @@ use {
         },
         pubkey::Pubkey,
     },
+    solana_unified_scheduler_pool::{BankingStageMonitor, BankingStageStatus},
     std::{
-        sync::{Arc, RwLock},
+        sync::{atomic::Ordering::Relaxed, Arc, RwLock},
         time::{Duration, Instant},
     },
 };
@@ -31,9 +32,10 @@ impl BufferedPacketsDecision {
     }
 }
 
-#[derive(Clone)]
+#[derive(Clone, derive_more::Debug)]
 pub struct DecisionMaker {
     my_pubkey: Pubkey,
+    #[debug("{poh_recorder:p}")]
     poh_recorder: Arc<RwLock<PohRecorder>>,
 
     cached_decision: Option<BufferedPacketsDecision>,
@@ -136,6 +138,21 @@ impl DecisionMaker {
     }
 }
 
+impl BankingStageMonitor for DecisionMaker {
+    fn status(&mut self) -> BankingStageStatus {
+        if self.poh_recorder.read().unwrap().is_exited.load(Relaxed) {
+            BankingStageStatus::Exited
+        } else if matches!(
+            self.make_consume_or_forward_decision(),
+            BufferedPacketsDecision::Forward,
+        ) {
+            BankingStageStatus::Inactive
+        } else {
+            BankingStageStatus::Active
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use {

core/src/banking_stage/packet_deserializer.rs

@@ -199,7 +199,6 @@ impl PacketDeserializer {
         })
     }
 
-    #[allow(dead_code)]
     pub(crate) fn deserialize_packets_with_indexes(
         packet_batch: &PacketBatch,
     ) -> impl Iterator<Item = (ImmutableDeserializedPacket, usize)> + '_ {

core/src/banking_stage/unified_scheduler.rs

@@ -32,7 +32,7 @@ use {
         packet_deserializer::PacketDeserializer,
         LikeClusterInfo,
     },
-    crate::banking_trace::Channels,
+    crate::{banking_stage::BankingStage, banking_trace::Channels},
     agave_banking_stage_ingress_types::BankingPacketBatch,
     solana_poh::poh_recorder::PohRecorder,
     solana_runtime::{bank_forks::BankForks, root_bank_cache::RootBankCache},
@@ -49,16 +49,22 @@ pub(crate) fn ensure_banking_stage_setup(
     cluster_info: &impl LikeClusterInfo,
     poh_recorder: &Arc<RwLock<PohRecorder>>,
 ) {
+    if !pool.block_production_supported() {
+        return;
+    }
+
     let mut root_bank_cache = RootBankCache::new(bank_forks.clone());
     let unified_receiver = channels.unified_receiver().clone();
     let mut decision_maker = DecisionMaker::new(cluster_info.id(), poh_recorder.clone());
+    let block_producing_scheduler_handler_threads = BankingStage::num_threads() as usize;
+    let banking_stage_monitor = Box::new(decision_maker.clone());
     let transaction_recorder = poh_recorder.read().unwrap().new_recorder();
 
     let banking_packet_handler = Box::new(
         move |helper: &BankingStageHelper, batches: BankingPacketBatch| {
             let decision = decision_maker.make_consume_or_forward_decision();
             if matches!(decision, BufferedPacketsDecision::Forward) {
-                return;
+                return Ok(());
             }
             let bank = root_bank_cache.root_bank();
             for batch in batches.iter() {
@@ -80,14 +86,17 @@ pub(crate) fn ensure_banking_stage_setup(
                     let index = task_id_base + packet_index;
 
                     let task = helper.create_new_task(transaction, index);
-                    helper.send_new_task(task);
+                    helper.send_new_task(task)?
                 }
             }
+            Ok(())
         },
     );
 
     pool.register_banking_stage(
         unified_receiver,
+        block_producing_scheduler_handler_threads,
+        banking_stage_monitor,
         banking_packet_handler,
         transaction_recorder,
     );