From fa55d49af8ad3a5390180d5ca542e0886389ab7f Mon Sep 17 00:00:00 2001
From: Pankaj Garg <pankaj@solana.com>
Date: Mon, 8 May 2023 11:20:57 -0700
Subject: [PATCH] address review comments

---
 program-runtime/src/loaded_programs.rs | 16 ++++++++++++--
 programs/bpf_loader/src/lib.rs         | 29 +++++++++++++++-----------
 runtime/src/bank.rs                    | 22 +++++++------------
 runtime/src/bank/tests.rs              |  2 +-
 4 files changed, 39 insertions(+), 30 deletions(-)

diff --git a/program-runtime/src/loaded_programs.rs b/program-runtime/src/loaded_programs.rs
index b7497803077e8d..ddce49a32bd849 100644
--- a/program-runtime/src/loaded_programs.rs
+++ b/program-runtime/src/loaded_programs.rs
@@ -276,7 +276,7 @@ pub struct LoadedPrograms {
 pub struct LoadedProgramsForTxBatch {
     /// Pubkey is the address of a program.
     /// LoadedProgram is the corresponding program entry valid for the slot in which a transaction is being executed.
-    pub entries: HashMap<Pubkey, Arc<LoadedProgram>>,
+    entries: HashMap<Pubkey, Arc<LoadedProgram>>,
     slot: Slot,
 }
 
@@ -321,9 +321,15 @@ impl LoadedProgramsForTxBatch {
         self.slot
     }
 
-    pub fn set_slot(&mut self, slot: Slot) {
+    pub fn set_slot_for_tests(&mut self, slot: Slot) {
         self.slot = slot;
     }
+
+    pub fn merge(&mut self, other: &Self) {
+        other.entries.iter().for_each(|(key, entry)| {
+            self.replenish(*key, entry.clone());
+        })
+    }
 }
 
 pub enum LoadedProgramMatchCriteria {
@@ -493,6 +499,12 @@ impl LoadedPrograms {
         )
     }
 
+    pub fn merge(&mut self, tx_batch_cache: &LoadedProgramsForTxBatch) {
+        tx_batch_cache.entries.iter().for_each(|(key, entry)| {
+            self.replenish(*key, entry.clone());
+        })
+    }
+
     /// Unloads programs which were used infrequently
     pub fn sort_and_unload(&mut self, shrink_to: PercentageInteger) {
         let sorted_candidates: Vec<(Pubkey, Arc<LoadedProgram>)> = self
diff --git a/programs/bpf_loader/src/lib.rs b/programs/bpf_loader/src/lib.rs
index 61489fe258b0bb..af505be8307351 100644
--- a/programs/bpf_loader/src/lib.rs
+++ b/programs/bpf_loader/src/lib.rs
@@ -1273,17 +1273,22 @@ fn process_loader_upgradeable_instruction(
                                 instruction_context,
                                 &log_collector,
                             )?;
-                            let clock = invoke_context.get_sysvar_cache().get_clock()?;
-                            invoke_context
-                                .programs_modified_by_tx
-                                .borrow_mut()
-                                .replenish(
-                                    program_key,
-                                    Arc::new(LoadedProgram::new_tombstone(
-                                        clock.slot,
-                                        LoadedProgramType::Closed,
-                                    )),
-                                );
+                            if invoke_context
+                                .feature_set
+                                .is_active(&delay_visibility_of_program_deployment::id())
+                            {
+                                let clock = invoke_context.get_sysvar_cache().get_clock()?;
+                                invoke_context
+                                    .programs_modified_by_tx
+                                    .borrow_mut()
+                                    .replenish(
+                                        program_key,
+                                        Arc::new(LoadedProgram::new_tombstone(
+                                            clock.slot,
+                                            LoadedProgramType::Closed,
+                                        )),
+                                    );
+                            }
                         }
                         _ => {
                             ic_logger_msg!(log_collector, "Invalid Program account");
@@ -1706,7 +1711,7 @@ pub mod test_utils {
                     false,
                 ) {
                     let mut cache = invoke_context.programs_modified_by_tx.borrow_mut();
-                    cache.set_slot(DELAY_VISIBILITY_SLOT_OFFSET);
+                    cache.set_slot_for_tests(DELAY_VISIBILITY_SLOT_OFFSET);
                     cache.replenish(*pubkey, Arc::new(loaded_program));
                 }
             }
diff --git a/runtime/src/bank.rs b/runtime/src/bank.rs
index aa35ef65f10a50..db652c5b4316d4 100644
--- a/runtime/src/bank.rs
+++ b/runtime/src/bank.rs
@@ -4706,13 +4706,9 @@ impl Bank {
                         // Update batch specific cache of the loaded programs with the modifications
                         // made by the transaction, if it executed successfully.
                         if details.status.is_ok() {
-                            programs_modified_by_tx.borrow().entries.iter().for_each(
-                                |(key, entry)| {
-                                    programs_loaded_for_tx_batch
-                                        .borrow_mut()
-                                        .replenish(*key, entry.clone());
-                                },
-                            )
+                            programs_loaded_for_tx_batch
+                                .borrow_mut()
+                                .merge(&programs_modified_by_tx.borrow());
                         }
                     }
                     result
@@ -5204,14 +5200,10 @@ impl Bank {
             } = execution_result
             {
                 if details.status.is_ok() {
-                    let mut cache = self.loaded_programs_cache.write().unwrap();
-                    programs_modified_by_tx
-                        .borrow()
-                        .entries
-                        .iter()
-                        .for_each(|(key, entry)| {
-                            cache.replenish(*key, entry.clone());
-                        })
+                    self.loaded_programs_cache
+                        .write()
+                        .unwrap()
+                        .merge(&programs_modified_by_tx.borrow());
                 }
             }
         }
diff --git a/runtime/src/bank/tests.rs b/runtime/src/bank/tests.rs
index 8ad66a5ec71e9a..5a26dba22bdcce 100644
--- a/runtime/src/bank/tests.rs
+++ b/runtime/src/bank/tests.rs
@@ -7772,7 +7772,7 @@ fn test_bpf_loader_upgradeable_deploy_with_max_len() {
         solana_bpf_loader_program::process_instruction,
         |invoke_context| {
             let mut cache = invoke_context.programs_modified_by_tx.borrow_mut();
-            cache.set_slot(bank.slot() + DELAY_VISIBILITY_SLOT_OFFSET);
+            cache.set_slot_for_tests(bank.slot() + DELAY_VISIBILITY_SLOT_OFFSET);
             cache.replenish(program_keypair.pubkey(), loaded_program.clone());
         },
         |_invoke_context| {},