From 99df0851f6ddfecf415c820aca40825cbfb309fb Mon Sep 17 00:00:00 2001
From: Andrew Fitzgerald <apfitzge@gmail.com>
Date: Tue, 7 May 2024 16:33:06 -0500
Subject: [PATCH] load balance by CUs (#1221)

(cherry picked from commit 41e2c969d7fad64c517a0c187780c7fe95901b15)
---
 .../transaction_scheduler/prio_graph_scheduler.rs        | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/core/src/banking_stage/transaction_scheduler/prio_graph_scheduler.rs b/core/src/banking_stage/transaction_scheduler/prio_graph_scheduler.rs
index 16fa19cca44895..8fa6559d166310 100644
--- a/core/src/banking_stage/transaction_scheduler/prio_graph_scheduler.rs
+++ b/core/src/banking_stage/transaction_scheduler/prio_graph_scheduler.rs
@@ -195,6 +195,8 @@ impl PrioGraphScheduler {
                     |thread_set| {
                         Self::select_thread(
                             thread_set,
+                            &batches.total_cus,
+                            self.in_flight_tracker.cus_in_flight_per_thread(),
                             &batches.transactions,
                             self.in_flight_tracker.num_in_flight_per_thread(),
                         )
@@ -419,6 +421,8 @@ impl PrioGraphScheduler {
     /// Panics if the `thread_set` is empty.
     fn select_thread(
         thread_set: ThreadSet,
+        batch_cus_per_thread: &[u64],
+        in_flight_cus_per_thread: &[u64],
         batches_per_thread: &[Vec<SanitizedTransaction>],
         in_flight_per_thread: &[usize],
     ) -> ThreadId {
@@ -427,11 +431,12 @@ impl PrioGraphScheduler {
             .map(|thread_id| {
                 (
                     thread_id,
+                    batch_cus_per_thread[thread_id] + in_flight_cus_per_thread[thread_id],
                     batches_per_thread[thread_id].len() + in_flight_per_thread[thread_id],
                 )
             })
-            .min_by(|a, b| a.1.cmp(&b.1))
-            .map(|(thread_id, _)| thread_id)
+            .min_by(|a, b| a.1.cmp(&b.1).then_with(|| a.2.cmp(&b.2)))
+            .map(|(thread_id, _, _)| thread_id)
             .unwrap()
     }