Fix - Reloading of unload entries after finishing the recompilation p…

…hase (anza-xyz#1199)

* Adds reproducer to test_feature_activation_loaded_programs_epoch_transition.

* Differentiate entries by environment instead of adjusting the effective slot.

* Fixes tests of assign_program().

* Fixes env order in test_feature_activation_loaded_programs_recompilation_phase().

* Turns test_load_program_effective_slot() into test_load_program_environment().

* Adds comments inside ProgramCache::assign_program().

program-runtime/src/loaded_programs.rs

@@ -818,11 +818,33 @@ impl<FG: ForkGraph> ProgramCache<FG> {
             &entry.program,
             ProgramCacheEntryType::DelayVisibility
         ));
+        // This function always returns `true` during normal operation.
+        // Only during the recompilation phase this can return `false`
+        // for entries with `upcoming_environments`.
+        fn is_current_env(
+            environments: &ProgramRuntimeEnvironments,
+            env_opt: Option<&ProgramRuntimeEnvironment>,
+        ) -> bool {
+            env_opt
+                .map(|env| {
+                    Arc::ptr_eq(env, &environments.program_runtime_v1)
+                        || Arc::ptr_eq(env, &environments.program_runtime_v2)
+                })
+                .unwrap_or(true)
+        }
         let slot_versions = &mut self.entries.entry(key).or_default();
         match slot_versions.binary_search_by(|at| {
             at.effective_slot
                 .cmp(&entry.effective_slot)
                 .then(at.deployment_slot.cmp(&entry.deployment_slot))
+                .then(
+                    // This `.then()` has no effect during normal operation.
+                    // Only during the recompilation phase this does allow entries
+                    // which only differ in their environment to be interleaved in `slot_versions`.
+                    is_current_env(&self.environments, at.program.get_environment()).cmp(
+                        &is_current_env(&self.environments, entry.program.get_environment()),
+                    ),
+                )
         }) {
             Ok(index) => {
                 let existing = slot_versions.get_mut(index).unwrap();
@@ -1727,34 +1749,34 @@ mod tests {
     #[test_matrix(
         (
             ProgramCacheEntryType::Closed,
-            ProgramCacheEntryType::FailedVerification(Arc::new(BuiltinProgram::new_mock())),
+            ProgramCacheEntryType::FailedVerification(get_mock_env()),
             new_loaded_entry(get_mock_env()),
         ),
         (
-            ProgramCacheEntryType::FailedVerification(Arc::new(BuiltinProgram::new_mock())),
+            ProgramCacheEntryType::FailedVerification(get_mock_env()),
             ProgramCacheEntryType::Closed,
-            ProgramCacheEntryType::Unloaded(Arc::new(BuiltinProgram::new_mock())),
+            ProgramCacheEntryType::Unloaded(get_mock_env()),
             new_loaded_entry(get_mock_env()),
             ProgramCacheEntryType::Builtin(BuiltinProgram::new_mock()),
         )
     )]
     #[test_matrix(
         (
-            ProgramCacheEntryType::Unloaded(Arc::new(BuiltinProgram::new_mock())),
+            ProgramCacheEntryType::Unloaded(get_mock_env()),
         ),
         (
-            ProgramCacheEntryType::FailedVerification(Arc::new(BuiltinProgram::new_mock())),
+            ProgramCacheEntryType::FailedVerification(get_mock_env()),
             ProgramCacheEntryType::Closed,
-            ProgramCacheEntryType::Unloaded(Arc::new(BuiltinProgram::new_mock())),
+            ProgramCacheEntryType::Unloaded(get_mock_env()),
             ProgramCacheEntryType::Builtin(BuiltinProgram::new_mock()),
         )
     )]
     #[test_matrix(
         (ProgramCacheEntryType::Builtin(BuiltinProgram::new_mock()),),
         (
-            ProgramCacheEntryType::FailedVerification(Arc::new(BuiltinProgram::new_mock())),
+            ProgramCacheEntryType::FailedVerification(get_mock_env()),
             ProgramCacheEntryType::Closed,
-            ProgramCacheEntryType::Unloaded(Arc::new(BuiltinProgram::new_mock())),
+            ProgramCacheEntryType::Unloaded(get_mock_env()),
             new_loaded_entry(get_mock_env()),
         )
     )]

runtime/src/bank/tests.rs

@@ -12029,16 +12029,16 @@ fn test_feature_activation_loaded_programs_recompilation_phase() {
     goto_end_of_slot(bank.clone());
     let bank = new_from_parent_with_fork_next_slot(bank, bank_forks.as_ref());
     {
-        let program_cache = bank.transaction_processor.program_cache.read().unwrap();
+        let program_cache = bank.transaction_processor.program_cache.write().unwrap();
         let slot_versions = program_cache.get_slot_versions_for_tests(&program_keypair.pubkey());
         assert_eq!(slot_versions.len(), 2);
         assert!(Arc::ptr_eq(
             slot_versions[0].program.get_environment().unwrap(),
-            &current_env
+            &upcoming_env
         ));
         assert!(Arc::ptr_eq(
             slot_versions[1].program.get_environment().unwrap(),
-            &upcoming_env
+            &current_env
         ));
     }
 
@@ -12108,6 +12108,21 @@ fn test_feature_activation_loaded_programs_epoch_transition() {
     let bank = new_bank_from_parent_with_bank_forks(&bank_forks, bank, &Pubkey::default(), 33);
 
     // Load the program with the new environment.
+    let transaction = Transaction::new(&signers, message.clone(), bank.last_blockhash());
+    assert!(bank.process_transaction(&transaction).is_ok());
+
+    {
+        // Prune for rerooting and thus finishing the recompilation phase.
+        let mut program_cache = bank.transaction_processor.program_cache.write().unwrap();
+        program_cache.prune(bank.slot(), bank.epoch());
+
+        // Unload all (which is only the entry with the new environment)
+        program_cache.sort_and_unload(percentage::Percentage::from(0));
+    }
+
+    // Reload the unloaded program with the new environment.
+    goto_end_of_slot(bank.clone());
+    let bank = new_from_parent_with_fork_next_slot(bank, bank_forks.as_ref());
     let transaction = Transaction::new(&signers, message, bank.last_blockhash());
     assert!(bank.process_transaction(&transaction).is_ok());
 }

svm/src/program_loader.rs

@@ -127,7 +127,7 @@ pub fn load_program_with_pubkey<CB: TransactionProcessingCallback, FG: ForkGraph
     pubkey: &Pubkey,
     slot: Slot,
     effective_epoch: Epoch,
-    epoch_schedule: &EpochSchedule,
+    _epoch_schedule: &EpochSchedule,
     reload: bool,
 ) -> Option<Arc<ProgramCacheEntry>> {
     let environments = program_cache.get_environments_for_epoch(effective_epoch);
@@ -136,7 +136,7 @@ pub fn load_program_with_pubkey<CB: TransactionProcessingCallback, FG: ForkGraph
         ..LoadProgramMetrics::default()
     };
 
-    let mut loaded_program = match load_program_accounts(callbacks, pubkey)? {
+    let loaded_program = match load_program_accounts(callbacks, pubkey)? {
         ProgramAccountLoadResult::InvalidAccountData(owner) => Ok(
             ProgramCacheEntry::new_tombstone(slot, owner, ProgramCacheEntryType::Closed),
         ),
@@ -217,21 +217,6 @@ pub fn load_program_with_pubkey<CB: TransactionProcessingCallback, FG: ForkGraph
 
     let mut timings = ExecuteDetailsTimings::default();
     load_program_metrics.submit_datapoint(&mut timings);
-    if !Arc::ptr_eq(
-        &environments.program_runtime_v1,
-        &program_cache.environments.program_runtime_v1,
-    ) || !Arc::ptr_eq(
-        &environments.program_runtime_v2,
-        &program_cache.environments.program_runtime_v2,
-    ) {
-        // There can be two entries per program when the environment changes.
-        // One for the old environment before the epoch boundary and one for the new environment after the epoch boundary.
-        // These two entries have the same deployment slot, so they must differ in their effective slot instead.
-        // This is done by setting the effective slot of the entry for the new environment to the epoch boundary.
-        loaded_program.effective_slot = loaded_program
-            .effective_slot
-            .max(epoch_schedule.get_first_slot_in_epoch(effective_epoch));
-    }
     loaded_program.update_access_slot(slot);
     Some(Arc::new(loaded_program))
 }
@@ -842,36 +827,51 @@ mod tests {
     }
 
     #[test]
-    fn test_load_program_effective_slot() {
+    fn test_load_program_environment() {
         let key = Pubkey::new_unique();
         let mock_bank = MockBankCallback::default();
         let mut account_data = AccountSharedData::default();
-        account_data.set_owner(loader_v4::id());
+        account_data.set_owner(bpf_loader::id());
         let batch_processor = TransactionBatchProcessor::<TestForkGraph>::default();
 
-        batch_processor
-            .program_cache
-            .write()
-            .unwrap()
-            .upcoming_environments = Some(ProgramRuntimeEnvironments::default());
+        let upcoming_environments = ProgramRuntimeEnvironments::default();
+        let current_environments = {
+            let mut program_cache = batch_processor.program_cache.write().unwrap();
+            program_cache.upcoming_environments = Some(upcoming_environments.clone());
+            program_cache.environments.clone()
+        };
         mock_bank
             .account_shared_data
             .borrow_mut()
             .insert(key, account_data.clone());
 
         let program_cache = batch_processor.program_cache.read().unwrap();
 
-        let result = load_program_with_pubkey(
-            &mock_bank,
-            &program_cache,
-            &key,
-            200,
-            20,
-            &batch_processor.epoch_schedule,
-            false,
-        );
-
-        let slot = batch_processor.epoch_schedule.get_first_slot_in_epoch(20);
-        assert_eq!(result.unwrap().effective_slot, slot);
+        for is_upcoming_env in [false, true] {
+            let result = load_program_with_pubkey(
+                &mock_bank,
+                &program_cache,
+                &key,
+                200,
+                is_upcoming_env as u64,
+                &batch_processor.epoch_schedule,
+                false,
+            )
+            .unwrap();
+            assert_ne!(
+                is_upcoming_env,
+                Arc::ptr_eq(
+                    result.program.get_environment().unwrap(),
+                    &current_environments.program_runtime_v1,
+                )
+            );
+            assert_eq!(
+                is_upcoming_env,
+                Arc::ptr_eq(
+                    result.program.get_environment().unwrap(),
+                    &upcoming_environments.program_runtime_v1,
+                )
+            );
+        }
     }
 }