From b519477bf0aaed93d3732081a7ee83f7e1c2ae56 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Alexander=20Mei=C3=9Fner?= <AlexanderMeissner@gmx.net>
Date: Tue, 25 Jul 2023 16:44:56 +0200
Subject: [PATCH] Adds parameter recompile to Bank::load_program().

---
 ledger-tool/src/program.rs |  2 +-
 runtime/src/bank.rs        | 44 ++++++++++++++++++++++++++++----------
 runtime/src/bank/tests.rs  |  4 ++--
 3 files changed, 36 insertions(+), 14 deletions(-)

diff --git a/ledger-tool/src/program.rs b/ledger-tool/src/program.rs
index f97c61f6038a55..b2454c2972342b 100644
--- a/ledger-tool/src/program.rs
+++ b/ledger-tool/src/program.rs
@@ -541,7 +541,7 @@ pub fn program(ledger_path: &Path, matches: &ArgMatches<'_>) {
     let mut loaded_programs =
         LoadedProgramsForTxBatch::new(bank.slot() + DELAY_VISIBILITY_SLOT_OFFSET);
     for key in cached_account_keys {
-        loaded_programs.replenish(key, bank.load_program(&key));
+        loaded_programs.replenish(key, bank.load_program(&key, None));
         debug!("Loaded program {}", key);
     }
     invoke_context.programs_loaded_for_tx_batch = &loaded_programs;
diff --git a/runtime/src/bank.rs b/runtime/src/bank.rs
index 5817a5a721f788..430a960779a23d 100644
--- a/runtime/src/bank.rs
+++ b/runtime/src/bank.rs
@@ -1470,7 +1470,12 @@ impl Bank {
                 if let Some((_env, ref mut programs_to_recompile)) =
                     &mut loaded_programs_cache.upcoming_program_runtime_environment_v1
                 {
-                    
+                    if let Some((key, program_to_recompile)) = programs_to_recompile.pop() {
+                        drop(loaded_programs_cache);
+                        let recompiled = new.load_program(&key, Some(program_to_recompile));
+                        let mut loaded_programs_cache = new.loaded_programs_cache.write().unwrap();
+                        loaded_programs_cache.replenish(key, recompiled);
+                    }
                 } else if slot_index.saturating_add(
                     solana_program_runtime::loaded_programs::MAX_LOADED_ENTRY_COUNT
                         .checked_div(2)
@@ -4680,20 +4685,28 @@ impl Bank {
         ProgramAccountLoadResult::InvalidAccountData
     }
 
-    pub fn load_program(&self, pubkey: &Pubkey) -> Arc<LoadedProgram> {
-        let program_runtime_environment_v1 = self
-            .loaded_programs_cache
-            .read()
-            .unwrap()
-            .program_runtime_environment_v1
-            .clone();
+    pub fn load_program(
+        &self,
+        pubkey: &Pubkey,
+        recompile: Option<Arc<LoadedProgram>>,
+    ) -> Arc<LoadedProgram> {
+        let loaded_programs_cache = self.loaded_programs_cache.read().unwrap();
+        let program_runtime_environment_v1 = if recompile.is_some() {
+            loaded_programs_cache
+                .upcoming_program_runtime_environment_v1
+                .as_ref()
+                .map(|(env, _programs_to_recompile)| env)
+                .unwrap()
+        } else {
+            &loaded_programs_cache.program_runtime_environment_v1
+        };
 
         let mut load_program_metrics = LoadProgramMetrics {
             program_id: pubkey.to_string(),
             ..LoadProgramMetrics::default()
         };
 
-        let loaded_program = match self.load_program_accounts(pubkey) {
+        let mut loaded_program = match self.load_program_accounts(pubkey) {
             ProgramAccountLoadResult::AccountNotFound => Ok(LoadedProgram::new_tombstone(
                 self.slot,
                 LoadedProgramType::Closed,
@@ -4745,12 +4758,21 @@ impl Bank {
         .unwrap_or_else(|_| {
             LoadedProgram::new_tombstone(
                 self.slot,
-                LoadedProgramType::FailedVerification(program_runtime_environment_v1),
+                LoadedProgramType::FailedVerification(program_runtime_environment_v1.clone()),
             )
         });
 
         let mut timings = ExecuteDetailsTimings::default();
         load_program_metrics.submit_datapoint(&mut timings);
+        if let Some(recompile) = recompile {
+            loaded_program.effective_slot = self
+                .epoch_schedule()
+                .get_first_slot_in_epoch(self.epoch.saturating_add(1));
+            loaded_program.tx_usage_counter =
+                AtomicU64::new(recompile.tx_usage_counter.load(Ordering::Relaxed));
+            loaded_program.ix_usage_counter =
+                AtomicU64::new(recompile.ix_usage_counter.load(Ordering::Relaxed));
+        }
         Arc::new(loaded_program)
     }
 
@@ -4988,7 +5010,7 @@ impl Bank {
         let missing_programs: Vec<(Pubkey, Arc<LoadedProgram>)> = missing_programs
             .iter()
             .map(|(key, count)| {
-                let program = self.load_program(key);
+                let program = self.load_program(key, None);
                 program.tx_usage_counter.store(*count, Ordering::Relaxed);
                 (*key, program)
             })
diff --git a/runtime/src/bank/tests.rs b/runtime/src/bank/tests.rs
index 62331f4762e061..706d81042a2e76 100644
--- a/runtime/src/bank/tests.rs
+++ b/runtime/src/bank/tests.rs
@@ -7275,7 +7275,7 @@ fn test_bank_load_program() {
     programdata_account.set_rent_epoch(1);
     bank.store_account_and_update_capitalization(&key1, &program_account);
     bank.store_account_and_update_capitalization(&programdata_key, &programdata_account);
-    let program = bank.load_program(&key1);
+    let program = bank.load_program(&key1, None);
     assert!(matches!(program.program, LoadedProgramType::LegacyV1(_)));
     assert_eq!(
         program.account_size,
@@ -7430,7 +7430,7 @@ fn test_bpf_loader_upgradeable_deploy_with_max_len() {
         assert_eq!(*elf.get(i).unwrap(), *byte);
     }
 
-    let loaded_program = bank.load_program(&program_keypair.pubkey());
+    let loaded_program = bank.load_program(&program_keypair.pubkey(), None);
 
     // Invoke deployed program
     mock_process_instruction(