Prio scheduler crate

Cargo.toml

@@ -69,6 +69,7 @@ members = [
     "poh",
     "poh-bench",
     "poseidon",
+    "prio-graph-scheduler",
     "program-runtime",
     "program-test",
     "programs/address-lookup-table",
@@ -445,6 +446,7 @@ solana-package-metadata-macro = { path = "sdk/package-metadata-macro", version =
 solana-perf = { path = "perf", version = "=2.1.0" }
 solana-poh = { path = "poh", version = "=2.1.0" }
 solana-poseidon = { path = "poseidon", version = "=2.1.0" }
+solana-prio-graph-scheduler = { path = "prio-graph-scheduler", version = "=2.1.0" }
 solana-precompile-error = { path = "sdk/precompile-error", version = "=2.1.0" }
 solana-program = { path = "sdk/program", version = "=2.1.0", default-features = false }
 solana-program-error = { path = "sdk/program-error", version = "=2.1.0" }

core/Cargo.toml

@@ -90,6 +90,7 @@ solana-version = { workspace = true }
 solana-vote = { workspace = true }
 solana-vote-program = { workspace = true }
 solana-wen-restart = { workspace = true }
+solana-prio-graph-scheduler = { workspace = true }
 strum = { workspace = true, features = ["derive"] }
 strum_macros = { workspace = true }
 sys-info = { workspace = true }

core/src/banking_stage.rs

@@ -17,12 +17,8 @@ use {
     },
     crate::{
         banking_stage::{
-            consume_worker::ConsumeWorker,
-            packet_deserializer::PacketDeserializer,
-            transaction_scheduler::{
-                prio_graph_scheduler::PrioGraphScheduler,
-                scheduler_controller::SchedulerController, scheduler_error::SchedulerError,
-            },
+            consume_worker::ConsumeWorker, packet_deserializer::PacketDeserializer,
+            scheduler_controller::SchedulerController,
         },
         banking_trace::BankingPacketReceiver,
         tracer_packet_stats::TracerPacketStats,
@@ -36,6 +32,9 @@ use {
     solana_measure::measure_us,
     solana_perf::{data_budget::DataBudget, packet::PACKETS_PER_BATCH},
     solana_poh::poh_recorder::{PohRecorder, TransactionRecorder},
+    solana_prio_graph_scheduler::{
+        prio_graph_scheduler::PrioGraphScheduler, scheduler_error::SchedulerError,
+    },
     solana_runtime::{
         bank_forks::BankForks, prioritization_fee_cache::PrioritizationFeeCache,
         vote_sender_types::ReplayVoteSender,
@@ -73,8 +72,7 @@ mod packet_deserializer;
 mod packet_filter;
 mod packet_receiver;
 mod read_write_account_set;
-mod scheduler_messages;
-mod transaction_scheduler;
+mod scheduler_controller;
 
 // Fixed thread size seems to be fastest on GCP setup
 pub const NUM_THREADS: u32 = 6;

core/src/banking_stage/consume_worker.rs

@@ -2,11 +2,11 @@ use {
     super::{
         consumer::{Consumer, ExecuteAndCommitTransactionsOutput, ProcessTransactionBatchOutput},
         leader_slot_timing_metrics::LeaderExecuteAndCommitTimings,
-        scheduler_messages::{ConsumeWork, FinishedConsumeWork},
     },
     crossbeam_channel::{Receiver, RecvError, SendError, Sender},
     solana_measure::measure_us,
     solana_poh::leader_bank_notifier::LeaderBankNotifier,
+    solana_prio_graph_scheduler::scheduler_messages::{ConsumeWork, FinishedConsumeWork},
     solana_runtime::bank::Bank,
     solana_sdk::timing::AtomicInterval,
     solana_svm::transaction_error_metrics::TransactionErrorMetrics,
@@ -694,7 +694,6 @@ mod tests {
         crate::banking_stage::{
             committer::Committer,
             qos_service::QosService,
-            scheduler_messages::{MaxAge, TransactionBatchId, TransactionId},
             tests::{create_slow_genesis_config, sanitize_transactions, simulate_poh},
         },
         crossbeam_channel::unbounded,
@@ -703,6 +702,9 @@ mod tests {
             get_tmp_ledger_path_auto_delete, leader_schedule_cache::LeaderScheduleCache,
         },
         solana_poh::poh_recorder::{PohRecorder, WorkingBankEntry},
+        solana_prio_graph_scheduler::scheduler_messages::{
+            MaxAge, TransactionBatchId, TransactionId,
+        },
         solana_runtime::{
             bank_forks::BankForks, prioritization_fee_cache::PrioritizationFeeCache,
             vote_sender_types::ReplayVoteReceiver,

core/src/banking_stage/consumer.rs

@@ -7,7 +7,6 @@ use {
         },
         leader_slot_timing_metrics::LeaderExecuteAndCommitTimings,
         qos_service::QosService,
-        scheduler_messages::MaxAge,
         unprocessed_transaction_storage::{ConsumeScannerPayload, UnprocessedTransactionStorage},
         BankingStageStats,
     },
@@ -19,6 +18,7 @@ use {
         BankStart, PohRecorderError, RecordTransactionsSummary, RecordTransactionsTimings,
         TransactionRecorder,
     },
+    solana_prio_graph_scheduler::scheduler_messages::MaxAge,
     solana_runtime::{
         bank::{Bank, LoadAndExecuteTransactionsOutput},
         transaction_batch::TransactionBatch,

core/src/banking_stage/forward_packet_batches_by_accounts.rs

@@ -8,6 +8,7 @@ use {
     },
     solana_feature_set::FeatureSet,
     solana_perf::packet::Packet,
+    solana_prio_graph_scheduler::deserializable_packet::DeserializableTxPacket,
     solana_sdk::transaction::SanitizedTransaction,
     solana_svm_transaction::svm_message::SVMMessage,
     std::sync::Arc,

core/src/banking_stage/immutable_deserialized_packet.rs

@@ -2,6 +2,7 @@ use {
     super::packet_filter::PacketFilterFailure,
     solana_compute_budget::compute_budget_limits::ComputeBudgetLimits,
     solana_perf::packet::Packet,
+    solana_prio_graph_scheduler::deserializable_packet::DeserializableTxPacket,
     solana_runtime::bank::Bank,
     solana_runtime_transaction::instructions_processor::process_compute_budget_instructions,
     solana_sanitize::SanitizeError,
@@ -50,8 +51,9 @@ pub struct ImmutableDeserializedPacket {
     compute_unit_limit: u32,
 }
 
-impl ImmutableDeserializedPacket {
-    pub fn new(packet: Packet) -> Result<Self, DeserializedPacketError> {
+impl DeserializableTxPacket for ImmutableDeserializedPacket {
+    type DeserializeError = DeserializedPacketError;
+    fn new(packet: Packet) -> Result<Self, Self::DeserializeError> {
         let versioned_transaction: VersionedTransaction = packet.deserialize_slice(..)?;
         let sanitized_transaction = SanitizedVersionedTransaction::try_from(versioned_transaction)?;
         let message_bytes = packet_message(&packet)?;
@@ -86,34 +88,34 @@ impl ImmutableDeserializedPacket {
         })
     }
 
-    pub fn original_packet(&self) -> &Packet {
+    fn original_packet(&self) -> &Packet {
         &self.original_packet
     }
 
-    pub fn transaction(&self) -> &SanitizedVersionedTransaction {
+    fn transaction(&self) -> &SanitizedVersionedTransaction {
         &self.transaction
     }
 
-    pub fn message_hash(&self) -> &Hash {
+    fn message_hash(&self) -> &Hash {
         &self.message_hash
     }
 
-    pub fn is_simple_vote(&self) -> bool {
+    fn is_simple_vote(&self) -> bool {
         self.is_simple_vote
     }
 
-    pub fn compute_unit_price(&self) -> u64 {
+    fn compute_unit_price(&self) -> u64 {
         self.compute_unit_price
     }
 
-    pub fn compute_unit_limit(&self) -> u64 {
+    fn compute_unit_limit(&self) -> u64 {
         u64::from(self.compute_unit_limit)
     }
 
     // This function deserializes packets into transactions, computes the blake3 hash of transaction
     // messages.
     // Additionally, this returns the minimum deactivation slot of the resolved addresses.
-    pub fn build_sanitized_transaction(
+    fn build_sanitized_transaction(
         &self,
         votes_only: bool,
         bank: &Bank,
@@ -138,7 +140,9 @@ impl ImmutableDeserializedPacket {
         .ok()?;
         Some((tx, deactivation_slot))
     }
+}
 
+impl ImmutableDeserializedPacket {
     fn resolve_addresses_with_deactivation(
         transaction: &SanitizedVersionedTransaction,
         bank: &Bank,

core/src/banking_stage/latest_unprocessed_votes.rs

@@ -6,6 +6,7 @@ use {
     itertools::Itertools,
     rand::{thread_rng, Rng},
     solana_perf::packet::Packet,
+    solana_prio_graph_scheduler::deserializable_packet::DeserializableTxPacket,
     solana_runtime::{bank::Bank, epoch_stakes::EpochStakes},
     solana_sdk::{
         account::from_account,

core/src/banking_stage/packet_deserializer.rs

@@ -4,15 +4,10 @@ use {
     super::{
         immutable_deserialized_packet::{DeserializedPacketError, ImmutableDeserializedPacket},
         packet_filter::PacketFilterFailure,
-    },
-    crate::{
+    }, crate::{
         banking_trace::{BankingPacketBatch, BankingPacketReceiver},
         sigverify::SigverifyTracerPacketStats,
-    },
-    crossbeam_channel::RecvTimeoutError,
-    solana_perf::packet::PacketBatch,
-    solana_sdk::saturating_add_assign,
-    std::time::{Duration, Instant},
+    }, crossbeam_channel::RecvTimeoutError, solana_perf::packet::PacketBatch, solana_prio_graph_scheduler::deserializable_packet::DeserializableTxPacket, solana_sdk::saturating_add_assign, std::time::{Duration, Instant}
 };
 
 /// Results from deserializing packet batches.

core/src/banking_stage/packet_filter.rs

@@ -3,6 +3,7 @@ use {
     solana_builtins_default_costs::BUILTIN_INSTRUCTION_COSTS,
     solana_sdk::{ed25519_program, saturating_add_assign, secp256k1_program},
     thiserror::Error,
+    solana_prio_graph_scheduler::deserializable_packet::DeserializableTxPacket,
 };
 
 #[derive(Debug, Error, PartialEq)]

core/src/banking_stage/packet_receiver.rs

@@ -9,6 +9,7 @@ use {
     crate::{banking_trace::BankingPacketReceiver, tracer_packet_stats::TracerPacketStats},
     crossbeam_channel::RecvTimeoutError,
     solana_measure::{measure::Measure, measure_us},
+    solana_prio_graph_scheduler::deserializable_packet::DeserializableTxPacket,
     solana_sdk::{saturating_add_assign, timing::timestamp},
     std::{sync::atomic::Ordering, time::Duration},
 };

core/src/banking_stage/transaction_scheduler/scheduler_controller.rs → core/src/banking_stage/scheduler_controller.rs

@@ -2,31 +2,32 @@
 //!
 
 use {
-    super::{
-        prio_graph_scheduler::PrioGraphScheduler,
-        scheduler_error::SchedulerError,
-        scheduler_metrics::{
-            SchedulerCountMetrics, SchedulerLeaderDetectionMetrics, SchedulerTimingMetrics,
-        },
-        transaction_id_generator::TransactionIdGenerator,
-        transaction_state::SanitizedTransactionTTL,
-        transaction_state_container::TransactionStateContainer,
-    },
     crate::banking_stage::{
         consume_worker::ConsumeWorkerMetrics,
         consumer::Consumer,
         decision_maker::{BufferedPacketsDecision, DecisionMaker},
         forwarder::Forwarder,
         immutable_deserialized_packet::ImmutableDeserializedPacket,
         packet_deserializer::PacketDeserializer,
-        scheduler_messages::MaxAge,
         ForwardOption, LikeClusterInfo, TOTAL_BUFFERED_PACKETS,
     },
     arrayvec::ArrayVec,
     crossbeam_channel::RecvTimeoutError,
     solana_accounts_db::account_locks::validate_account_locks,
     solana_cost_model::cost_model::CostModel,
     solana_measure::measure_us,
+    solana_prio_graph_scheduler::deserializable_packet::DeserializableTxPacket,
+    solana_prio_graph_scheduler::{
+        id_generator::IdGenerator,
+        prio_graph_scheduler::PrioGraphScheduler,
+        scheduler_error::SchedulerError,
+        scheduler_messages::MaxAge,
+        scheduler_metrics::{
+            SchedulerCountMetrics, SchedulerLeaderDetectionMetrics, SchedulerTimingMetrics,
+        },
+        transaction_state::SanitizedTransactionTTL,
+        transaction_state_container::TransactionStateContainer,
+    },
     solana_runtime::{bank::Bank, bank_forks::BankForks},
     solana_runtime_transaction::instructions_processor::process_compute_budget_instructions,
     solana_sdk::{
@@ -53,12 +54,12 @@ pub(crate) struct SchedulerController<T: LikeClusterInfo> {
     packet_receiver: PacketDeserializer,
     bank_forks: Arc<RwLock<BankForks>>,
     /// Generates unique IDs for incoming transactions.
-    transaction_id_generator: TransactionIdGenerator,
+    transaction_id_generator: IdGenerator,
     /// Container for transaction state.
     /// Shared resource between `packet_receiver` and `scheduler`.
-    container: TransactionStateContainer,
+    container: TransactionStateContainer<ImmutableDeserializedPacket>,
     /// State for scheduling and communicating with worker threads.
-    scheduler: PrioGraphScheduler,
+    scheduler: PrioGraphScheduler<ImmutableDeserializedPacket>,
     /// Metrics tracking time for leader bank detection.
     leader_detection_metrics: SchedulerLeaderDetectionMetrics,
     /// Metrics tracking counts on transactions in different states
@@ -78,16 +79,18 @@ impl<T: LikeClusterInfo> SchedulerController<T> {
         decision_maker: DecisionMaker,
         packet_deserializer: PacketDeserializer,
         bank_forks: Arc<RwLock<BankForks>>,
-        scheduler: PrioGraphScheduler,
+        scheduler: PrioGraphScheduler<ImmutableDeserializedPacket>,
         worker_metrics: Vec<Arc<ConsumeWorkerMetrics>>,
         forwarder: Option<Forwarder<T>>,
     ) -> Self {
         Self {
             decision_maker,
             packet_receiver: packet_deserializer,
             bank_forks,
-            transaction_id_generator: TransactionIdGenerator::default(),
-            container: TransactionStateContainer::with_capacity(TOTAL_BUFFERED_PACKETS),
+            transaction_id_generator: IdGenerator::default(),
+            container: TransactionStateContainer::<ImmutableDeserializedPacket>::with_capacity(
+                TOTAL_BUFFERED_PACKETS,
+            ),
             scheduler,
             leader_detection_metrics: SchedulerLeaderDetectionMetrics::default(),
             count_metrics: SchedulerCountMetrics::default(),
@@ -717,9 +720,7 @@ mod tests {
         super::*,
         crate::{
             banking_stage::{
-                consumer::TARGET_NUM_TRANSACTIONS_PER_BATCH,
-                scheduler_messages::{ConsumeWork, FinishedConsumeWork, TransactionBatchId},
-                tests::create_slow_genesis_config,
+                consumer::TARGET_NUM_TRANSACTIONS_PER_BATCH, tests::create_slow_genesis_config,
             },
             banking_trace::BankingPacketBatch,
             sigverify::SigverifyTracerPacketStats,
@@ -733,6 +734,9 @@ mod tests {
         },
         solana_perf::packet::{to_packet_batches, PacketBatch, NUM_PACKETS},
         solana_poh::poh_recorder::{PohRecorder, Record, WorkingBankEntry},
+        solana_prio_graph_scheduler::scheduler_messages::{
+            ConsumeWork, FinishedConsumeWork, TransactionBatchId,
+        },
         solana_runtime::bank::Bank,
         solana_sdk::{
             compute_budget::ComputeBudgetInstruction, fee_calculator::FeeRateGovernor, hash::Hash,