Remove public visibility of program cache from bank (#279)

anza-xyz · Mar 17, 2024 · cc3afa5 · cc3afa5
1 parent b27c80a
commit cc3afa5
Show file tree

Hide file tree

Showing 7 changed files with 43 additions and 39 deletions.
diff --git a/core/benches/banking_stage.rs b/core/benches/banking_stage.rs
@@ -398,10 +398,7 @@ fn simulate_process_entries(
     let bank_fork = BankForks::new_rw_arc(bank);
     let bank = bank_fork.read().unwrap().get_with_scheduler(slot).unwrap();
     bank.clone_without_scheduler()
-        .loaded_programs_cache
-        .write()
-        .unwrap()
-        .set_fork_graph(bank_fork.clone());
+        .set_fork_graph_in_program_cache(bank_fork.clone());
 
     for i in 0..(num_accounts / 2) {
         bank.transfer(initial_lamports, mint_keypair, &keypairs[i * 2].pubkey())

diff --git a/core/src/replay_stage.rs b/core/src/replay_stage.rs
@@ -1686,11 +1686,7 @@ impl ReplayStage {
             root_bank.clear_slot_signatures(slot);
 
             // Remove cached entries of the programs that were deployed in this slot.
-            root_bank
-                .loaded_programs_cache
-                .write()
-                .unwrap()
-                .prune_by_deployment_slot(slot);
+            root_bank.prune_program_cache_by_deployment_slot(slot);
 
             if let Some(bank_hash) = blockstore.get_bank_hash(slot) {
                 // If a descendant was successfully replayed and chained from a duplicate it must

diff --git a/ledger-tool/src/program.rs b/ledger-tool/src/program.rs
@@ -514,14 +514,9 @@ pub fn program(ledger_path: &Path, matches: &ArgMatches<'_>) {
     with_mock_invoke_context!(invoke_context, transaction_context, transaction_accounts);
 
     // Adding `DELAY_VISIBILITY_SLOT_OFFSET` to slots to accommodate for delay visibility of the program
-    let mut loaded_programs = LoadedProgramsForTxBatch::new(
-        bank.slot() + DELAY_VISIBILITY_SLOT_OFFSET,
-        bank.loaded_programs_cache
-            .read()
-            .unwrap()
-            .environments
-            .clone(),
-    );
+    let slot = bank.slot() + DELAY_VISIBILITY_SLOT_OFFSET;
+    let mut loaded_programs =
+        LoadedProgramsForTxBatch::new(slot, bank.get_runtime_environments_for_slot(slot));
     for key in cached_account_keys {
         loaded_programs.replenish(key, bank.load_program(&key, false, bank.epoch()));
         debug!("Loaded program {}", key);

diff --git a/ledger/src/blockstore_processor.rs b/ledger/src/blockstore_processor.rs
@@ -1674,11 +1674,7 @@ fn load_frozen_forks(
                 root = new_root_bank.slot();
 
                 leader_schedule_cache.set_root(new_root_bank);
-                new_root_bank
-                    .loaded_programs_cache
-                    .write()
-                    .unwrap()
-                    .prune(root, new_root_bank.epoch());
+                new_root_bank.prune_program_cache(root, new_root_bank.epoch());
                 let _ = bank_forks.write().unwrap().set_root(
                     root,
                     accounts_background_request_sender,

diff --git a/runtime/src/bank.rs b/runtime/src/bank.rs
@@ -98,7 +98,9 @@ use {
     solana_program_runtime::{
         compute_budget_processor::process_compute_budget_instructions,
         invoke_context::BuiltinFunctionWithContext,
-        loaded_programs::{LoadedProgram, LoadedProgramType, LoadedPrograms},
+        loaded_programs::{
+            LoadedProgram, LoadedProgramType, LoadedPrograms, ProgramRuntimeEnvironments,
+        },
         runtime_config::RuntimeConfig,
         timings::{ExecuteTimingType, ExecuteTimings},
     },
@@ -803,7 +805,7 @@ pub struct Bank {
 
     pub incremental_snapshot_persistence: Option<BankIncrementalSnapshotPersistence>,
 
-    pub loaded_programs_cache: Arc<RwLock<LoadedPrograms<BankForks>>>,
+    loaded_programs_cache: Arc<RwLock<LoadedPrograms<BankForks>>>,
 
     epoch_reward_status: EpochRewardStatus,
 
@@ -1467,6 +1469,36 @@ impl Bank {
         new
     }
 
+    pub fn set_fork_graph_in_program_cache(&self, fork_graph: Arc<RwLock<BankForks>>) {
+        self.loaded_programs_cache
+            .write()
+            .unwrap()
+            .set_fork_graph(fork_graph);
+    }
+
+    pub fn prune_program_cache(&self, new_root_slot: Slot, new_root_epoch: Epoch) {
+        self.loaded_programs_cache
+            .write()
+            .unwrap()
+            .prune(new_root_slot, new_root_epoch);
+    }
+
+    pub fn prune_program_cache_by_deployment_slot(&self, deployment_slot: Slot) {
+        self.loaded_programs_cache
+            .write()
+            .unwrap()
+            .prune_by_deployment_slot(deployment_slot);
+    }
+
+    pub fn get_runtime_environments_for_slot(&self, slot: Slot) -> ProgramRuntimeEnvironments {
+        let epoch = self.epoch_schedule.get_epoch(slot);
+        self.loaded_programs_cache
+            .read()
+            .unwrap()
+            .get_environments_for_epoch(epoch)
+            .clone()
+    }
+
     /// Epoch in which the new cooldown warmup rate for stake was activated
     pub fn new_warmup_cooldown_rate_epoch(&self) -> Option<Epoch> {
         self.feature_set

diff --git a/runtime/src/bank_forks.rs b/runtime/src/bank_forks.rs
@@ -126,12 +126,7 @@ impl BankForks {
             scheduler_pool: None,
         }));
 
-        root_bank
-            .loaded_programs_cache
-            .write()
-            .unwrap()
-            .set_fork_graph(bank_forks.clone());
-
+        root_bank.set_fork_graph_in_program_cache(bank_forks.clone());
         bank_forks
     }
 
@@ -451,11 +446,7 @@ impl BankForks {
 
     pub fn prune_program_cache(&self, root: Slot) {
         if let Some(root_bank) = self.banks.get(&root) {
-            root_bank
-                .loaded_programs_cache
-                .write()
-                .unwrap()
-                .prune(root, root_bank.epoch());
+            root_bank.prune_program_cache(root, root_bank.epoch());
         }
     }
 

diff --git a/unified-scheduler-pool/src/lib.rs b/unified-scheduler-pool/src/lib.rs
@@ -941,10 +941,7 @@ mod tests {
         let slot = bank.slot();
         let bank_fork = BankForks::new_rw_arc(bank);
         let bank = bank_fork.read().unwrap().get(slot).unwrap();
-        bank.loaded_programs_cache
-            .write()
-            .unwrap()
-            .set_fork_graph(bank_fork);
+        bank.set_fork_graph_in_program_cache(bank_fork);
         bank
     }