Buffer pre-session tasks in BP unified scheduler

Cargo.toml

@@ -596,6 +596,7 @@ tower = "0.5.2"
 trait-set = "0.3.0"
 trees = "0.4.2"
 tungstenite = "0.20.1"
+unwrap_none = "0.1.2"
 uriparse = "0.6.4"
 url = "2.5.4"
 vec_extract_if_polyfill = "0.1.0"

ledger/src/blockstore_processor.rs

@@ -5032,7 +5032,7 @@ pub mod tests {
             ..
         } = create_genesis_config_with_leader(500, &dummy_leader_pubkey, 100);
         let bank = Arc::new(Bank::new_for_tests(&genesis_config));
-        let context = SchedulingContext::new(bank.clone());
+        let context = SchedulingContext::for_verification(bank.clone());
 
         let txs = create_test_transactions(&mint_keypair, &genesis_config.hash());
 

runtime/src/bank_forks.rs

@@ -246,7 +246,14 @@ impl BankForks {
             let context = SchedulingContext::new_with_mode(mode, bank.clone());
             let scheduler = scheduler_pool.take_scheduler(context);
             let bank_with_scheduler = BankWithScheduler::new(bank, Some(scheduler));
-            scheduler_pool.register_timeout_listener(bank_with_scheduler.create_timeout_listener());
+            if matches!(mode, SchedulingMode::BlockVerification) {
+                // Skip registering for block production. Both the replay stage and PohRecorder
+                // don't support _concurrent block production_ at all. It's strongly assumed that
+                // block is produced in singleton way and it's actually desired, while ignoring the
+                // opportunity cost of (hopefully rare!) fork switching...
+                scheduler_pool
+                    .register_timeout_listener(bank_with_scheduler.create_timeout_listener());
+            }
             bank_with_scheduler
         } else {
             BankWithScheduler::new_without_scheduler(bank)

runtime/src/installed_scheduler_pool.rs

@@ -238,40 +238,48 @@ pub type SchedulerId = u64;
 #[derive(Clone, Debug)]
 pub struct SchedulingContext {
     mode: SchedulingMode,
-    bank: Arc<Bank>,
+    bank: Option<Arc<Bank>>,
 }
 
 impl SchedulingContext {
-    pub fn new(bank: Arc<Bank>) -> Self {
-        // mode will be configurable later
+    pub fn for_preallocation() -> Self {
         Self {
-            mode: SchedulingMode::BlockVerification,
-            bank,
+            mode: SchedulingMode::BlockProduction,
+            bank: None,
         }
     }
 
     pub fn new_with_mode(mode: SchedulingMode, bank: Arc<Bank>) -> Self {
-        Self { mode, bank }
+        Self {
+            mode,
+            bank: Some(bank),
+        }
+    }
+
+    #[cfg_attr(feature = "dev-context-only-utils", qualifiers(pub))]
+    fn for_verification(bank: Arc<Bank>) -> Self {
+        Self::new_with_mode(SchedulingMode::BlockVerification, bank)
     }
 
     #[cfg(feature = "dev-context-only-utils")]
     pub fn for_production(bank: Arc<Bank>) -> Self {
-        Self {
-            mode: SchedulingMode::BlockProduction,
-            bank,
-        }
+        Self::new_with_mode(SchedulingMode::BlockProduction, bank)
+    }
+
+    pub fn is_preallocated(&self) -> bool {
+        self.bank.is_none()
     }
 
     pub fn mode(&self) -> SchedulingMode {
         self.mode
     }
 
     pub fn bank(&self) -> &Arc<Bank> {
-        &self.bank
+        self.bank.as_ref().unwrap()
     }
 
-    pub fn slot(&self) -> Slot {
-        self.bank().slot()
+    pub fn slot(&self) -> Option<Slot> {
+        self.bank.as_ref().map(|bank| bank.slot())
     }
 }
 
@@ -570,7 +578,8 @@ impl BankWithSchedulerInner {
                 let pool = pool.clone();
                 drop(scheduler);
 
-                let context = SchedulingContext::new(self.bank.clone());
+                // Schedulers can be stale only if its mode is block-verification.
+                let context = SchedulingContext::for_verification(self.bank.clone());
                 let mut scheduler = self.scheduler.write().unwrap();
                 trace!("with_active_scheduler: {:?}", scheduler);
                 scheduler.transition_from_stale_to_active(|pool, result_with_timings| {
@@ -773,7 +782,7 @@ mod tests {
         mock.expect_context()
             .times(1)
             .in_sequence(&mut seq.lock().unwrap())
-            .return_const(SchedulingContext::new(bank));
+            .return_const(SchedulingContext::for_verification(bank));
 
         for wait_reason in is_dropped_flags {
             let seq_cloned = seq.clone();

unified-scheduler-logic/Cargo.toml

@@ -15,6 +15,7 @@ solana-pubkey = { workspace = true }
 solana-runtime-transaction = { workspace = true }
 solana-transaction = { workspace = true }
 static_assertions = { workspace = true }
+unwrap_none = { workspace = true }
 
 [dev-dependencies]
 solana-instruction = { workspace = true }

unified-scheduler-logic/src/lib.rs

@@ -103,9 +103,10 @@ use {
     solana_transaction::sanitized::SanitizedTransaction,
     static_assertions::const_assert_eq,
     std::{collections::VecDeque, mem, sync::Arc},
+    unwrap_none::UnwrapNone,
 };
 
-#[derive(Clone, Copy, Debug)]
+#[derive(Clone, Copy, Debug, PartialEq)]
 pub enum SchedulingMode {
     BlockVerification,
     BlockProduction,
@@ -668,11 +669,29 @@ impl SchedulingStateMachine {
     /// indicating the scheduled task is blocked currently.
     ///
     /// Note that this function takes ownership of the task to allow for future optimizations.
+    #[cfg(test)]
     #[must_use]
     pub fn schedule_task(&mut self, task: Task) -> Option<Task> {
+        self.schedule_or_buffer_task(task, false)
+    }
+
+    pub fn buffer_task(&mut self, task: Task) {
+        self.schedule_or_buffer_task(task, true).unwrap_none();
+    }
+
+    #[must_use]
+    pub fn schedule_or_buffer_task(&mut self, task: Task, force_buffering: bool) -> Option<Task> {
         self.total_task_count.increment_self();
         self.active_task_count.increment_self();
-        self.try_lock_usage_queues(task)
+        self.try_lock_usage_queues(task).and_then(|task| {
+            if !force_buffering {
+                Some(task)
+            } else {
+                self.unblocked_task_count.increment_self();
+                self.unblocked_task_queue.push_back(task);
+                None
+            }
+        })
     }
 
     #[must_use]