Buffer pre-session tasks in BP unified scheduler (anza-xyz#4949)

* Buffer pre-session tasks in BP unified scheduler

* Clean up comments and tests

* Add doc comments to SchedulingStateMachine

* Make SchedulingContext::bank() return Option<_>

* kick ci

* kick ci

Author: ryoqun

Cargo.lock

@@ -594,6 +594,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"

Cargo.toml

@@ -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());
 

ledger/src/blockstore_processor.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());
+            // Skip registering for block production. Both the tvu main loop in 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...
+            if matches!(mode, SchedulingMode::BlockVerification) {
+                scheduler_pool
+                    .register_timeout_listener(bank_with_scheduler.create_timeout_listener());
+            }
             bank_with_scheduler
         } else {
             BankWithScheduler::new_without_scheduler(bank)

programs/sbf/Cargo.lock

@@ -235,43 +235,54 @@ pub type SchedulerId = u64;
 /// expected to be used by a particular scheduler only for that duration of the time and to be
 /// disposed by the scheduler. Then, the scheduler may work on different banks with new
 /// `SchedulingContext`s.
+///
+/// There's a special construction only used for scheduler preallocation, which has no bank. Panics
+/// will be triggered when tried to be used normally across code-base.
 #[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
+    pub fn bank(&self) -> Option<&Arc<Bank>> {
+        self.bank.as_ref()
     }
 
-    pub fn slot(&self) -> Slot {
-        self.bank().slot()
+    pub fn slot(&self) -> Option<Slot> {
+        self.bank.as_ref().map(|bank| bank.slot())
     }
 }
 
@@ -424,11 +435,19 @@ pub struct BankWithSchedulerInner {
 pub type InstalledSchedulerRwLock = RwLock<SchedulerStatus>;
 
 impl BankWithScheduler {
+    /// Creates a new `BankWithScheduler` from bank and its associated scheduler.
+    ///
+    /// # Panics
+    ///
+    /// Panics if `scheduler`'s scheduling context is unmatched to given bank or for scheduler
+    /// preallocation.
     #[cfg_attr(feature = "dev-context-only-utils", qualifiers(pub))]
     pub(crate) fn new(bank: Arc<Bank>, scheduler: Option<InstalledSchedulerBox>) -> Self {
+        // Avoid the fatal situation in which bank is being associated with a scheduler associated
+        // to a different bank!
         if let Some(bank_in_context) = scheduler
             .as_ref()
-            .map(|scheduler| scheduler.context().bank())
+            .map(|scheduler| scheduler.context().bank().unwrap())
         {
             assert!(Arc::ptr_eq(&bank, bank_in_context));
         }
@@ -570,7 +589,9 @@ 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. So,
+                // unconditional context construction for verification is okay here.
+                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 +794,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();

runtime/src/bank_forks.rs

@@ -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 }

runtime/src/installed_scheduler_pool.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,49 @@ impl SchedulingStateMachine {
     /// indicating the scheduled task is blocked currently.
     ///
     /// Note that this function takes ownership of the task to allow for future optimizations.
+    #[cfg(any(test, doc))]
     #[must_use]
     pub fn schedule_task(&mut self, task: Task) -> Option<Task> {
+        self.schedule_or_buffer_task(task, false)
+    }
+
+    /// Adds given `task` to internal buffer, even if it's immediately schedulable otherwise.
+    ///
+    /// Put differently, buffering means to force the task to be blocked unconditionally after
+    /// normal scheduling processing.
+    ///
+    /// Thus, the task is internally retained inside this [`SchedulingStateMachine`], whether the
+    /// task is blocked or not. Eventually, the buffered task will be returned by one of later
+    /// invocations [`schedule_next_unblocked_task()`](Self::schedule_next_unblocked_task).
+    ///
+    /// Note that this function takes ownership of the task to allow for future optimizations.
+    pub fn buffer_task(&mut self, task: Task) {
+        self.schedule_or_buffer_task(task, true).unwrap_none();
+    }
+
+    /// Schedules or buffers given `task`, returning successful one unless buffering is forced.
+    ///
+    /// Refer to [`schedule_task()`](Self::schedule_task) and
+    /// [`buffer_task()`](Self::buffer_task) for the difference between _scheduling_ and
+    /// _buffering_ respectively.
+    ///
+    /// Note that this function takes ownership of the task to allow for future optimizations.
+    #[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| {
+            // locking succeeded, and then ...
+            if force_buffering {
+                // ... push to unblocked_task_queue, if buffering is forced.
+                self.unblocked_task_count.increment_self();
+                self.unblocked_task_queue.push_back(task);
+                None
+            } else {
+                // ... return the task back as schedulable to the caller as-is otherwise.
+                Some(task)
+            }
+        })
     }
 
     #[must_use]