From 8e4e6e8737ec230bf936aaaabd6b7814ed63bce7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Alexander=20Mei=C3=9Fner?= <AlexanderMeissner@gmx.net>
Date: Wed, 21 Feb 2024 23:45:15 +0000
Subject: [PATCH] Moves recompilation specifics out of load_program().

---
 ledger-tool/src/program.rs       |  2 +-
 runtime/src/bank.rs              | 13 ++++++++++---
 svm/src/transaction_processor.rs | 20 +++-----------------
 3 files changed, 14 insertions(+), 21 deletions(-)

diff --git a/ledger-tool/src/program.rs b/ledger-tool/src/program.rs
index b56affd4c905c2..af50d59bca0255 100644
--- a/ledger-tool/src/program.rs
+++ b/ledger-tool/src/program.rs
@@ -523,7 +523,7 @@ pub fn program(ledger_path: &Path, matches: &ArgMatches<'_>) {
             .clone(),
     );
     for key in cached_account_keys {
-        loaded_programs.replenish(key, bank.load_program(&key, false, None));
+        loaded_programs.replenish(key, bank.load_program(&key, false, bank.epoch()));
         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 7e051019c99871..73ca586c78c2ec 100644
--- a/runtime/src/bank.rs
+++ b/runtime/src/bank.rs
@@ -1360,8 +1360,15 @@ impl Bank {
                 if let Some((key, program_to_recompile)) =
                     loaded_programs_cache.programs_to_recompile.pop()
                 {
+                    let effective_epoch = loaded_programs_cache.latest_root_epoch.saturating_add(1);
                     drop(loaded_programs_cache);
-                    let recompiled = new.load_program(&key, false, Some(program_to_recompile));
+                    let recompiled = new.load_program(&key, false, effective_epoch);
+                    recompiled
+                        .tx_usage_counter
+                        .fetch_add(program_to_recompile.tx_usage_counter.load(Relaxed), Relaxed);
+                    recompiled
+                        .ix_usage_counter
+                        .fetch_add(program_to_recompile.ix_usage_counter.load(Relaxed), Relaxed);
                     let mut loaded_programs_cache = new.loaded_programs_cache.write().unwrap();
                     loaded_programs_cache.assign_program(key, recompiled);
                 }
@@ -7480,10 +7487,10 @@ impl Bank {
         &self,
         pubkey: &Pubkey,
         reload: bool,
-        recompile: Option<Arc<LoadedProgram>>,
+        effective_epoch: Epoch,
     ) -> Arc<LoadedProgram> {
         self.transaction_processor
-            .load_program(self, pubkey, reload, recompile)
+            .load_program(self, pubkey, reload, effective_epoch)
     }
 }
 
diff --git a/svm/src/transaction_processor.rs b/svm/src/transaction_processor.rs
index d877bf3d09b54c..d322bfa72a5be3 100644
--- a/svm/src/transaction_processor.rs
+++ b/svm/src/transaction_processor.rs
@@ -51,10 +51,7 @@ use {
         collections::{hash_map::Entry, HashMap},
         fmt::{Debug, Formatter},
         rc::Rc,
-        sync::{
-            atomic::{AtomicU64, Ordering},
-            Arc, RwLock,
-        },
+        sync::{atomic::Ordering, Arc, RwLock},
     },
 };
 
@@ -419,7 +416,7 @@ impl<FG: ForkGraph> TransactionBatchProcessor<FG> {
 
             if let Some((key, count)) = program_to_load {
                 // Load, verify and compile one program.
-                let program = self.load_program(callback, &key, false, None);
+                let program = self.load_program(callback, &key, false, self.epoch);
                 program.tx_usage_counter.store(count, Ordering::Relaxed);
                 program_to_store = Some((key, program));
             } else if missing_programs.is_empty() {
@@ -654,14 +651,9 @@ impl<FG: ForkGraph> TransactionBatchProcessor<FG> {
         callbacks: &CB,
         pubkey: &Pubkey,
         reload: bool,
-        recompile: Option<Arc<LoadedProgram>>,
+        effective_epoch: Epoch,
     ) -> Arc<LoadedProgram> {
         let loaded_programs_cache = self.loaded_programs_cache.read().unwrap();
-        let effective_epoch = if recompile.is_some() {
-            loaded_programs_cache.latest_root_epoch.saturating_add(1)
-        } else {
-            self.epoch
-        };
         let environments = loaded_programs_cache.get_environments_for_epoch(effective_epoch);
         let mut load_program_metrics = LoadProgramMetrics {
             program_id: pubkey.to_string(),
@@ -742,12 +734,6 @@ impl<FG: ForkGraph> TransactionBatchProcessor<FG> {
         loaded_program.effective_slot = loaded_program
             .effective_slot
             .max(self.epoch_schedule.get_first_slot_in_epoch(effective_epoch));
-        if let Some(recompile) = recompile {
-            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));
-        }
         loaded_program.update_access_slot(self.slot);
         Arc::new(loaded_program)
     }