Wire unified scheduler into banking experimentally

Cargo.toml

@@ -312,11 +312,13 @@ curve25519-dalek = { version = "4.1.3", features = ["digest", "rand_core"] }
 dashmap = "5.5.3"
 derivation-path = { version = "0.2.0", default-features = false }
 derive-where = "1.2.7"
+derive_more = { version = "1.0.0", features = ["full"] }
 dialoguer = "0.10.4"
 digest = "0.10.7"
 dir-diff = "0.3.3"
 dirs-next = "2.0.0"
 dlopen2 = "0.5.0"
+dyn-clone = "1.0.17"
 eager = "0.1.0"
 ed25519-dalek = "=1.0.1"
 ed25519-dalek-bip32 = "0.2.0"
@@ -660,6 +662,7 @@ tokio-util = "0.7"
 toml = "0.8.12"
 tonic = "0.9.2"
 tonic-build = "0.9.2"
+trait-set = "0.3.0"
 trees = "0.4.2"
 tungstenite = "0.20.1"
 uriparse = "0.6.4"

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

@@ -9,7 +9,10 @@ use {
     rayon::prelude::*,
     solana_client::connection_cache::ConnectionCache,
     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,
     },
@@ -38,6 +41,8 @@ use {
     },
     solana_streamer::socket::SocketAddrSpace,
     solana_tpu_client::tpu_client::DEFAULT_TPU_CONNECTION_POOL_SIZE,
+    solana_unified_scheduler_logic::SchedulingMode,
+    solana_unified_scheduler_pool::{DefaultSchedulerPool, SupportedSchedulingMode},
     std::{
         sync::{atomic::Ordering, Arc, RwLock},
         thread::sleep,
@@ -442,20 +447,41 @@ fn main() {
             BANKING_TRACE_DIR_DEFAULT_BYTE_LIMIT,
         )))
         .unwrap();
+    let prioritization_fee_cache = Arc::new(PrioritizationFeeCache::new(0u64));
+    let cluster_info = {
+        let keypair = Arc::new(Keypair::new());
+        let node = Node::new_localhost_with_pubkey(&keypair.pubkey());
+        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(),
+            poh_recorder.read().unwrap().new_recorder(),
+        );
+        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);
-    let cluster_info = {
-        let keypair = Arc::new(Keypair::new());
-        let node = Node::new_localhost_with_pubkey(&keypair.pubkey());
-        ClusterInfo::new(node.info, keypair, SocketAddrSpace::Unspecified)
-    };
-    let cluster_info = Arc::new(cluster_info);
+    } = banking_tracer_channels;
     let tpu_disable_quic = matches.is_present("tpu_disable_quic");
     let connection_cache = if tpu_disable_quic {
         ConnectionCache::with_udp(
@@ -469,7 +495,7 @@ fn main() {
         )
     };
     let banking_stage = BankingStage::new_num_threads(
-        block_production_method,
+        block_production_method.clone(),
         &cluster_info,
         &poh_recorder,
         non_vote_receiver,
@@ -481,10 +507,22 @@ fn main() {
         None,
         Arc::new(connection_cache),
         bank_forks.clone(),
-        &Arc::new(PrioritizationFeeCache::new(0u64)),
+        &prioritization_fee_cache,
         false,
     );
 
+    // 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.
@@ -545,10 +583,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

@@ -26,6 +26,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 }
@@ -92,6 +93,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 }
@@ -106,7 +108,6 @@ tokio = { workspace = true, features = ["full"] }
 trees = { workspace = true }
 
 [dev-dependencies]
-assert_matches = { workspace = true }
 fs_extra = { workspace = true }
 serde_json = { workspace = true }
 serial_test = { workspace = true }
@@ -123,7 +124,6 @@ solana-program-runtime = { workspace = true, features = ["metrics"] }
 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",
 ] }