Adds parameter recompile to Bank::load_program().

ledger-tool/src/program.rs

@@ -552,7 +552,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));
+        loaded_programs.replenish(key, bank.load_program(&key, false, None));
         debug!("Loaded program {}", key);
     }
     invoke_context.programs_loaded_for_tx_batch = &loaded_programs;

program-runtime/src/loaded_programs.rs

@@ -454,6 +454,8 @@ pub struct LoadedPrograms {
     /// More precisely, it starts with the recompilation phase a few hundred slots before the epoch boundary,
     /// and it ends with the first rerooting after the epoch boundary.
     pub upcoming_environments: Option<ProgramRuntimeEnvironments>,
+    /// List of loaded programs which should be recompiled before the next epoch (but don't have to).
+    pub programs_to_recompile: Vec<(Pubkey, Arc<LoadedProgram>)>,
     pub stats: Stats,
 }
 
@@ -626,6 +628,7 @@ impl LoadedPrograms {
             if let Some(upcoming_environments) = self.upcoming_environments.take() {
                 recompilation_phase_ends = true;
                 self.environments = upcoming_environments;
+                self.programs_to_recompile.clear();
             }
         }
         for second_level in self.entries.values_mut() {

runtime/src/bank.rs

@@ -1449,8 +1449,17 @@ impl Bank {
                 new.update_epoch_stakes(leader_schedule_epoch);
                 // Recompile loaded programs one at a time before the next epoch hits
                 let (_epoch, slot_index) = new.get_epoch_and_slot_index(new.slot());
-                let loaded_programs_cache = new.loaded_programs_cache.write().unwrap();
-                if slot_index.saturating_add(
+                let mut loaded_programs_cache = new.loaded_programs_cache.write().unwrap();
+                if !loaded_programs_cache.programs_to_recompile.is_empty() {
+                    if let Some((key, program_to_recompile)) =
+                        loaded_programs_cache.programs_to_recompile.pop()
+                    {
+                        drop(loaded_programs_cache);
+                        let recompiled = new.load_program(&key, false, 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)
                         .unwrap() as u64,
@@ -1488,6 +1497,11 @@ impl Bank {
                             Arc::new(program_runtime_environment_v2);
                     }
                     loaded_programs_cache.upcoming_environments = Some(upcoming_environments);
+                    loaded_programs_cache.programs_to_recompile = loaded_programs_cache
+                        .get_entries_sorted_by_tx_usage(
+                            changed_program_runtime_v1,
+                            changed_program_runtime_v2,
+                        );
                 }
             }
             if new.is_partitioned_rewards_code_enabled() {
@@ -4705,16 +4719,28 @@ impl Bank {
         ProgramAccountLoadResult::InvalidAccountData
     }
 
-    pub fn load_program(&self, pubkey: &Pubkey, reload: bool) -> Arc<LoadedProgram> {
+    pub fn load_program(
+        &self,
+        pubkey: &Pubkey,
+        reload: bool,
+        recompile: Option<Arc<LoadedProgram>>,
+    ) -> Arc<LoadedProgram> {
+        let effective_slot = if recompile.is_some() {
+            self.epoch_schedule()
+                .get_first_slot_in_epoch(self.epoch.saturating_add(1))
+        } else {
+            self.slot
+        };
         let loaded_programs_cache = self.loaded_programs_cache.read().unwrap();
-        let environments = &loaded_programs_cache.get_environments_for_epoch(self.epoch);
+        let environments = &loaded_programs_cache
+            .get_environments_for_epoch(self.get_epoch_and_slot_index(effective_slot).0);
 
         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,
@@ -4800,25 +4826,34 @@ 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()),
             )
         });
 
         let mut timings = ExecuteDetailsTimings::default();
         load_program_metrics.submit_datapoint(&mut timings);
+        if let Some(recompile) = recompile {
+            loaded_program.effective_slot = effective_slot;
+            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)
     }
 
@@ -5093,7 +5128,7 @@ impl Bank {
         let missing_programs: Vec<(Pubkey, Arc<LoadedProgram>)> = missing
             .iter()
             .map(|(key, count)| {
-                let program = self.load_program(key, false);
+                let program = self.load_program(key, false, None);
                 program.tx_usage_counter.store(*count, Ordering::Relaxed);
                 (*key, program)
             })
@@ -5103,7 +5138,7 @@ impl Bank {
         let unloaded_programs: Vec<(Pubkey, Arc<LoadedProgram>)> = unloaded
             .iter()
             .map(|(key, count)| {
-                let program = self.load_program(key, true);
+                let program = self.load_program(key, true, None);
                 program.tx_usage_counter.store(*count, Ordering::Relaxed);
                 (*key, program)
             })

runtime/src/bank/tests.rs

@@ -7225,7 +7225,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, false);
+    let program = bank.load_program(&key1, false, None);
     assert_matches!(program.program, LoadedProgramType::LegacyV1(_));
     assert_eq!(
         program.account_size,
@@ -7380,7 +7380,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(), false);
+    let loaded_program = bank.load_program(&program_keypair.pubkey(), false, None);
 
     // Invoke deployed program
     mock_process_instruction(