From fad52dff34b026402770e1e2dcef80af4f1a2f44 Mon Sep 17 00:00:00 2001
From: Andrew Fitzgerald <apfitzge@gmail.com>
Date: Fri, 21 Jul 2023 14:10:00 -0700
Subject: [PATCH] DecisionMaker: Add more tests (#32505)

---
 core/src/banking_stage/decision_maker.rs | 115 +++++++++++++++++++----
 1 file changed, 98 insertions(+), 17 deletions(-)

diff --git a/core/src/banking_stage/decision_maker.rs b/core/src/banking_stage/decision_maker.rs
index ef705c0a2..e93051320 100644
--- a/core/src/banking_stage/decision_maker.rs
+++ b/core/src/banking_stage/decision_maker.rs
@@ -47,22 +47,10 @@ impl DecisionMaker {
             let poh_recorder = self.poh_recorder.read().unwrap();
             decision = Self::consume_or_forward_packets(
                 &self.my_pubkey,
-                || {
-                    poh_recorder.bank_start().filter(|bank_start| {
-                        bank_start.should_working_bank_still_be_processing_txs()
-                    })
-                },
-                || {
-                    poh_recorder.would_be_leader(
-                        (FORWARD_TRANSACTIONS_TO_LEADER_AT_SLOT_OFFSET - 1)
-                            * DEFAULT_TICKS_PER_SLOT,
-                    )
-                },
-                || {
-                    poh_recorder
-                        .would_be_leader(HOLD_TRANSACTIONS_SLOT_OFFSET * DEFAULT_TICKS_PER_SLOT)
-                },
-                || poh_recorder.leader_after_n_slots(FORWARD_TRANSACTIONS_TO_LEADER_AT_SLOT_OFFSET),
+                || Self::bank_start(&poh_recorder),
+                || Self::would_be_leader_shortly(&poh_recorder),
+                || Self::would_be_leader(&poh_recorder),
+                || Self::leader_pubkey(&poh_recorder),
             );
         }
 
@@ -101,6 +89,26 @@ impl DecisionMaker {
             BufferedPacketsDecision::Hold
         }
     }
+
+    fn bank_start(poh_recorder: &PohRecorder) -> Option<BankStart> {
+        poh_recorder
+            .bank_start()
+            .filter(|bank_start| bank_start.should_working_bank_still_be_processing_txs())
+    }
+
+    fn would_be_leader_shortly(poh_recorder: &PohRecorder) -> bool {
+        poh_recorder.would_be_leader(
+            (FORWARD_TRANSACTIONS_TO_LEADER_AT_SLOT_OFFSET - 1) * DEFAULT_TICKS_PER_SLOT,
+        )
+    }
+
+    fn would_be_leader(poh_recorder: &PohRecorder) -> bool {
+        poh_recorder.would_be_leader(HOLD_TRANSACTIONS_SLOT_OFFSET * DEFAULT_TICKS_PER_SLOT)
+    }
+
+    fn leader_pubkey(poh_recorder: &PohRecorder) -> Option<Pubkey> {
+        poh_recorder.leader_after_n_slots(FORWARD_TRANSACTIONS_TO_LEADER_AT_SLOT_OFFSET)
+    }
 }
 
 #[cfg(test)]
@@ -108,8 +116,15 @@ mod tests {
     use {
         super::*,
         core::panic,
+        solana_ledger::{blockstore::Blockstore, genesis_utils::create_genesis_config},
+        solana_poh::poh_recorder::create_test_recorder,
         solana_runtime::bank::Bank,
-        std::{sync::Arc, time::Instant},
+        solana_sdk::clock::NUM_CONSECUTIVE_LEADER_SLOTS,
+        std::{
+            env::temp_dir,
+            sync::{atomic::Ordering, Arc},
+            time::Instant,
+        },
     };
 
     #[test]
@@ -129,6 +144,72 @@ mod tests {
         assert!(BufferedPacketsDecision::Hold.bank_start().is_none());
     }
 
+    #[test]
+    fn test_make_consume_or_forward_decision() {
+        let genesis_config = create_genesis_config(2).genesis_config;
+        let bank = Arc::new(Bank::new_no_wallclock_throttle_for_tests(&genesis_config));
+        let ledger_path = temp_dir();
+        let blockstore = Arc::new(Blockstore::open(ledger_path.as_path()).unwrap());
+        let (exit, poh_recorder, poh_service, _entry_receiver) =
+            create_test_recorder(&bank, blockstore, None, None);
+
+        let my_pubkey = Pubkey::new_unique();
+        let decision_maker = DecisionMaker::new(my_pubkey, poh_recorder.clone());
+        poh_recorder.write().unwrap().reset(bank.clone(), None);
+        let bank = Arc::new(Bank::new_from_parent(&bank, &my_pubkey, bank.slot() + 1));
+
+        // Currently Leader - Consume
+        {
+            poh_recorder.write().unwrap().set_bank(bank.clone(), false);
+            let decision = decision_maker.make_consume_or_forward_decision();
+            assert!(matches!(decision, BufferedPacketsDecision::Consume(_)));
+        }
+
+        // Will be leader shortly - Hold
+        for next_leader_slot_offset in [0, 1].into_iter() {
+            let next_leader_slot = bank.slot() + next_leader_slot_offset;
+            poh_recorder.write().unwrap().reset(
+                bank.clone(),
+                Some((
+                    next_leader_slot,
+                    next_leader_slot + NUM_CONSECUTIVE_LEADER_SLOTS,
+                )),
+            );
+            let decision = decision_maker.make_consume_or_forward_decision();
+            assert!(
+                matches!(decision, BufferedPacketsDecision::Hold),
+                "next_leader_slot_offset: {next_leader_slot_offset}",
+            );
+        }
+
+        // Will be leader - ForwardAndHold
+        for next_leader_slot_offset in [2, 19].into_iter() {
+            let next_leader_slot = bank.slot() + next_leader_slot_offset;
+            poh_recorder.write().unwrap().reset(
+                bank.clone(),
+                Some((
+                    next_leader_slot,
+                    next_leader_slot + NUM_CONSECUTIVE_LEADER_SLOTS + 1,
+                )),
+            );
+            let decision = decision_maker.make_consume_or_forward_decision();
+            assert!(
+                matches!(decision, BufferedPacketsDecision::ForwardAndHold),
+                "next_leader_slot_offset: {next_leader_slot_offset}",
+            );
+        }
+
+        // Known leader, not me - Forward
+        {
+            poh_recorder.write().unwrap().reset(bank, None);
+            let decision = decision_maker.make_consume_or_forward_decision();
+            assert!(matches!(decision, BufferedPacketsDecision::Forward));
+        }
+
+        exit.store(true, Ordering::Relaxed);
+        poh_service.join().unwrap();
+    }
+
     #[test]
     fn test_should_process_or_forward_packets() {
         let my_pubkey = solana_sdk::pubkey::new_rand();