From 4da0e28b487c76bb05d96022e3ec4764143313a2 Mon Sep 17 00:00:00 2001
From: "mergify[bot]" <37929162+mergify[bot]@users.noreply.github.com>
Date: Thu, 29 Feb 2024 14:35:12 -0800
Subject: [PATCH] v1.18: Add limit to looping in banking-stage (backport of
 #35342) (#35364)

---
 core/src/banking_stage/consumer.rs     |  3 ++-
 program-runtime/src/loaded_programs.rs |  9 ++++---
 runtime/src/bank.rs                    | 34 ++++++++++++++++++++------
 3 files changed, 35 insertions(+), 11 deletions(-)

diff --git a/core/src/banking_stage/consumer.rs b/core/src/banking_stage/consumer.rs
index 64b68889747633..093ee533aa45fb 100644
--- a/core/src/banking_stage/consumer.rs
+++ b/core/src/banking_stage/consumer.rs
@@ -574,7 +574,8 @@ impl Consumer {
                 transaction_status_sender_enabled,
                 &mut execute_and_commit_timings.execute_timings,
                 None, // account_overrides
-                self.log_messages_bytes_limit
+                self.log_messages_bytes_limit,
+                true,
             ));
         execute_and_commit_timings.load_execute_us = load_execute_us;
 
diff --git a/program-runtime/src/loaded_programs.rs b/program-runtime/src/loaded_programs.rs
index b676830e5bc4cc..dea50678890008 100644
--- a/program-runtime/src/loaded_programs.rs
+++ b/program-runtime/src/loaded_programs.rs
@@ -183,7 +183,7 @@ impl Stats {
             ("reloads", reloads, i64),
             ("insertions", insertions, i64),
             ("lost_insertions", lost_insertions, i64),
-            ("replacements", replacements, i64),
+            ("replace_entry", replacements, i64),
             ("one_hit_wonders", one_hit_wonders, i64),
             ("prunes_orphan", prunes_orphan, i64),
             ("prunes_environment", prunes_environment, i64),
@@ -568,6 +568,7 @@ pub struct LoadedProgramsForTxBatch {
     entries: HashMap<Pubkey, Arc<LoadedProgram>>,
     slot: Slot,
     pub environments: ProgramRuntimeEnvironments,
+    pub hit_max_limit: bool,
 }
 
 impl LoadedProgramsForTxBatch {
@@ -576,6 +577,7 @@ impl LoadedProgramsForTxBatch {
             entries: HashMap::new(),
             slot,
             environments,
+            hit_max_limit: false,
         }
     }
 
@@ -934,7 +936,7 @@ impl<FG: ForkGraph> LoadedPrograms<FG> {
         slot: Slot,
         key: Pubkey,
         loaded_program: Arc<LoadedProgram>,
-    ) {
+    ) -> bool {
         let second_level = self.entries.entry(key).or_default();
         debug_assert_eq!(
             second_level.cooperative_loading_lock,
@@ -955,8 +957,9 @@ impl<FG: ForkGraph> LoadedPrograms<FG> {
         {
             self.stats.lost_insertions.fetch_add(1, Ordering::Relaxed);
         }
-        self.assign_program(key, loaded_program);
+        let was_occupied = self.assign_program(key, loaded_program);
         self.loading_task_waiter.notify();
+        was_occupied
     }
 
     pub fn merge(&mut self, tx_batch_cache: &LoadedProgramsForTxBatch) {
diff --git a/runtime/src/bank.rs b/runtime/src/bank.rs
index 2ec0d09cc83000..98a8eda628da74 100644
--- a/runtime/src/bank.rs
+++ b/runtime/src/bank.rs
@@ -4289,6 +4289,7 @@ impl Bank {
             &mut timings,
             Some(&account_overrides),
             None,
+            true,
         );
 
         let post_simulation_accounts = loaded_transactions
@@ -4938,6 +4939,7 @@ impl Bank {
     fn replenish_program_cache(
         &self,
         program_accounts_map: &HashMap<Pubkey, (&Pubkey, u64)>,
+        limit_to_load_programs: bool,
     ) -> LoadedProgramsForTxBatch {
         let mut missing_programs: Vec<(Pubkey, (LoadedProgramMatchCriteria, u64))> =
             if self.check_program_modification_slot {
@@ -4983,11 +4985,14 @@ impl Bank {
                 }
                 // Submit our last completed loading task.
                 if let Some((key, program)) = program_to_store.take() {
-                    loaded_programs_cache.finish_cooperative_loading_task(
-                        self.slot(),
-                        key,
-                        program,
-                    );
+                    if loaded_programs_cache
+                        .finish_cooperative_loading_task(self.slot, key, program)
+                        && limit_to_load_programs
+                    {
+                        let mut ret = LoadedProgramsForTxBatch::default();
+                        ret.hit_max_limit = true;
+                        return ret;
+                    }
                 }
                 // Figure out which program needs to be loaded next.
                 let program_to_load = loaded_programs_cache.extract(
@@ -5072,7 +5077,7 @@ impl Bank {
         result
     }
 
-    #[allow(clippy::type_complexity)]
+    #[allow(clippy::too_many_arguments, clippy::type_complexity)]
     pub fn load_and_execute_transactions(
         &self,
         batch: &TransactionBatch,
@@ -5083,6 +5088,7 @@ impl Bank {
         timings: &mut ExecuteTimings,
         account_overrides: Option<&AccountOverrides>,
         log_messages_bytes_limit: Option<usize>,
+        limit_to_load_programs: bool,
     ) -> LoadAndExecuteTransactionsOutput {
         let sanitized_txs = batch.sanitized_transactions();
         debug!("processing transactions: {}", sanitized_txs.len());
@@ -5146,9 +5152,22 @@ impl Bank {
         }
 
         let programs_loaded_for_tx_batch = Rc::new(RefCell::new(
-            self.replenish_program_cache(&program_accounts_map),
+            self.replenish_program_cache(&program_accounts_map, limit_to_load_programs),
         ));
 
+        if programs_loaded_for_tx_batch.borrow().hit_max_limit {
+            return LoadAndExecuteTransactionsOutput {
+                loaded_transactions: vec![],
+                execution_results: vec![],
+                retryable_transaction_indexes,
+                executed_transactions_count: 0,
+                executed_non_vote_transactions_count: 0,
+                executed_with_successful_result_count: 0,
+                signature_count: 0,
+                error_counters,
+            };
+        }
+
         let mut load_time = Measure::start("accounts_load");
         let mut loaded_transactions = load_accounts(
             &self.rc.accounts.accounts_db,
@@ -6304,6 +6323,7 @@ impl Bank {
             timings,
             None,
             log_messages_bytes_limit,
+            false,
         );
 
         let (last_blockhash, lamports_per_signature) =