Introduce InstalledSchedulerPool trait (#33934)

* Introduce InstalledSchedulerPool

* Use type alias

* Remove log_prefix for now...

* Simplify return_to_pool()

* Simplify InstalledScheduler's context methods

* Reorder trait methods semantically

* Simplify Arc<Bank> handling

Cargo.toml

@@ -132,6 +132,7 @@ edition = "2021"
 
 [workspace.dependencies]
 Inflector = "0.11.4"
+aquamarine = "0.3.2"
 aes-gcm-siv = "0.10.3"
 ahash = "0.8.6"
 anyhow = "1.0.75"

ledger/src/blockstore_processor.rs

@@ -1945,7 +1945,7 @@ pub mod tests {
             genesis_utils::{
                 self, create_genesis_config_with_vote_accounts, ValidatorVoteKeypairs,
             },
-            installed_scheduler_pool::{MockInstalledScheduler, WaitReason},
+            installed_scheduler_pool::{MockInstalledScheduler, SchedulingContext, WaitReason},
         },
         solana_sdk::{
             account::{AccountSharedData, WritableAccount},
@@ -4527,11 +4527,17 @@ 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 txs = create_test_transactions(&mint_keypair, &genesis_config.hash());
 
         let mut mocked_scheduler = MockInstalledScheduler::new();
         let mut seq = mockall::Sequence::new();
+        mocked_scheduler
+            .expect_context()
+            .times(1)
+            .in_sequence(&mut seq)
+            .return_const(context);
         mocked_scheduler
             .expect_schedule_execution()
             .times(txs.len())
@@ -4542,6 +4548,11 @@ pub mod tests {
             .times(1)
             .in_sequence(&mut seq)
             .returning(|_| None);
+        mocked_scheduler
+            .expect_return_to_pool()
+            .times(1)
+            .in_sequence(&mut seq)
+            .returning(|| ());
         let bank = BankWithScheduler::new(bank, Some(Box::new(mocked_scheduler)));
 
         let batch = bank.prepare_sanitized_batch(&txs);

runtime/Cargo.toml

@@ -10,6 +10,7 @@ license = { workspace = true }
 edition = { workspace = true }
 
 [dependencies]
+aquamarine = { workspace = true }
 arrayref = { workspace = true }
 base64 = { workspace = true }
 bincode = { workspace = true }

runtime/src/bank_forks.rs

@@ -4,7 +4,9 @@ use {
     crate::{
         accounts_background_service::{AbsRequestSender, SnapshotRequest, SnapshotRequestKind},
         bank::{epoch_accounts_hash_utils, Bank, SquashTiming},
-        installed_scheduler_pool::BankWithScheduler,
+        installed_scheduler_pool::{
+            BankWithScheduler, InstalledSchedulerPoolArc, SchedulingContext,
+        },
         snapshot_config::SnapshotConfig,
     },
     log::*,
@@ -72,6 +74,7 @@ pub struct BankForks {
     last_accounts_hash_slot: Slot,
     in_vote_only_mode: Arc<AtomicBool>,
     highest_slot_at_startup: Slot,
+    scheduler_pool: Option<InstalledSchedulerPoolArc>,
 }
 
 impl Index<u64> for BankForks {
@@ -203,6 +206,7 @@ impl BankForks {
             last_accounts_hash_slot: root,
             in_vote_only_mode: Arc::new(AtomicBool::new(false)),
             highest_slot_at_startup: 0,
+            scheduler_pool: None,
         }));
 
         for bank in bank_forks.read().unwrap().banks.values() {
@@ -215,11 +219,26 @@ impl BankForks {
         bank_forks
     }
 
+    pub fn install_scheduler_pool(&mut self, pool: InstalledSchedulerPoolArc) {
+        info!("Installed new scheduler_pool into bank_forks: {:?}", pool);
+        assert!(
+            self.scheduler_pool.replace(pool).is_none(),
+            "Reinstalling scheduler pool isn't supported"
+        );
+    }
+
     pub fn insert(&mut self, mut bank: Bank) -> BankWithScheduler {
         bank.check_program_modification_slot =
             self.root.load(Ordering::Relaxed) < self.highest_slot_at_startup;
 
-        let bank = BankWithScheduler::new_without_scheduler(Arc::new(bank));
+        let bank = Arc::new(bank);
+        let bank = if let Some(scheduler_pool) = &self.scheduler_pool {
+            let context = SchedulingContext::new(bank.clone());
+            let scheduler = scheduler_pool.take_scheduler(context);
+            BankWithScheduler::new(bank, Some(scheduler))
+        } else {
+            BankWithScheduler::new_without_scheduler(bank)
+        };
         let prev = self.banks.insert(bank.slot(), bank.clone_with_scheduler());
         assert!(prev.is_none());
         let slot = bank.slot();