From 64038d279f8b49974bed7d75041968efb9ea2a0c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Alexander=20Mei=C3=9Fner?= <AlexanderMeissner@gmx.net>
Date: Mon, 2 Oct 2023 23:15:05 +0200
Subject: [PATCH] Adds LoadedPrograms::get_environments_for_epoch().

---
 program-runtime/src/loaded_programs.rs | 24 ++++++++++++++++++++----
 runtime/src/bank.rs                    | 20 +++++++++++---------
 2 files changed, 31 insertions(+), 13 deletions(-)

diff --git a/program-runtime/src/loaded_programs.rs b/program-runtime/src/loaded_programs.rs
index dd649d532c6584..0ded17ee7877de 100644
--- a/program-runtime/src/loaded_programs.rs
+++ b/program-runtime/src/loaded_programs.rs
@@ -56,9 +56,12 @@ pub trait ForkGraph {
 
 /// Provides information about current working slot, and its ancestors
 pub trait WorkingSlot {
-    /// Returns the current slot value
+    /// Returns the current slot
     fn current_slot(&self) -> Slot;
 
+    /// Returns the epoch of the current slot
+    fn current_epoch(&self) -> Epoch;
+
     /// Returns true if the `other` slot is an ancestor of self, false otherwise
     fn is_ancestor(&self, other: Slot) -> bool;
 }
@@ -529,6 +532,11 @@ pub enum LoadedProgramMatchCriteria {
 }
 
 impl LoadedPrograms {
+    /// Returns the current environments depending on the given epoch
+    pub fn get_environments_for_epoch(&self, _epoch: Epoch) -> &ProgramRuntimeEnvironments {
+        &self.environments
+    }
+
     /// Refill the cache with a single entry. It's typically called during transaction loading,
     /// when the cache doesn't contain the entry corresponding to program `key`.
     /// The function dedupes the cache, in case some other thread replenished the entry in parallel.
@@ -716,6 +724,7 @@ impl LoadedPrograms {
         working_slot: &S,
         keys: impl Iterator<Item = (Pubkey, (LoadedProgramMatchCriteria, u64))>,
     ) -> ExtractedPrograms {
+        let environments = self.get_environments_for_epoch(working_slot.current_epoch());
         let mut missing = Vec::new();
         let mut unloaded = Vec::new();
         let found = keys
@@ -733,7 +742,7 @@ impl LoadedPrograms {
                                     return None;
                                 }
 
-                                if !Self::matches_environment(entry, &self.environments) {
+                                if !Self::matches_environment(entry, environments) {
                                     missing.push((key, count));
                                     return None;
                                 }
@@ -778,7 +787,7 @@ impl LoadedPrograms {
             loaded: LoadedProgramsForTxBatch {
                 entries: found,
                 slot: working_slot.current_slot(),
-                environments: self.environments.clone(),
+                environments: environments.clone(),
             },
             missing,
             unloaded,
@@ -920,7 +929,10 @@ mod tests {
         assert_matches::assert_matches,
         percentage::Percentage,
         solana_rbpf::vm::BuiltinProgram,
-        solana_sdk::{clock::Slot, pubkey::Pubkey},
+        solana_sdk::{
+            clock::{Epoch, Slot},
+            pubkey::Pubkey,
+        },
         std::{
             ops::ControlFlow,
             sync::{
@@ -1472,6 +1484,10 @@ mod tests {
             self.slot
         }
 
+        fn current_epoch(&self) -> Epoch {
+            0
+        }
+
         fn is_ancestor(&self, other: Slot) -> bool {
             self.fork
                 .iter()
diff --git a/runtime/src/bank.rs b/runtime/src/bank.rs
index 4bbd825579d0e7..5227db287c9f10 100644
--- a/runtime/src/bank.rs
+++ b/runtime/src/bank.rs
@@ -933,6 +933,10 @@ impl WorkingSlot for Bank {
         self.slot
     }
 
+    fn current_epoch(&self) -> Epoch {
+        self.epoch
+    }
+
     fn is_ancestor(&self, other: Slot) -> bool {
         self.ancestors.contains_key(&other)
     }
@@ -4675,12 +4679,8 @@ impl Bank {
     }
 
     pub fn load_program(&self, pubkey: &Pubkey, reload: bool) -> Arc<LoadedProgram> {
-        let environments = self
-            .loaded_programs_cache
-            .read()
-            .unwrap()
-            .environments
-            .clone();
+        let loaded_programs_cache = self.loaded_programs_cache.read().unwrap();
+        let environments = loaded_programs_cache.get_environments_for_epoch(self.epoch);
 
         let mut load_program_metrics = LoadProgramMetrics {
             program_id: pubkey.to_string(),
@@ -4773,20 +4773,22 @@ impl Bank {
                     })
                     .unwrap_or(LoadedProgram::new_tombstone(
                         self.slot,
-                        LoadedProgramType::FailedVerification(environments.program_runtime_v2),
+                        LoadedProgramType::FailedVerification(
+                            environments.program_runtime_v2.clone(),
+                        ),
                     ));
                 Ok(loaded_program)
             }
 
             ProgramAccountLoadResult::InvalidV4Program => Ok(LoadedProgram::new_tombstone(
                 self.slot,
-                LoadedProgramType::FailedVerification(environments.program_runtime_v2),
+                LoadedProgramType::FailedVerification(environments.program_runtime_v2.clone()),
             )),
         }
         .unwrap_or_else(|_| {
             LoadedProgram::new_tombstone(
                 self.slot,
-                LoadedProgramType::FailedVerification(environments.program_runtime_v1),
+                LoadedProgramType::FailedVerification(environments.program_runtime_v1.clone()),
             )
         });