diff --git a/runtime/src/bank.rs b/runtime/src/bank.rs index 19cd4aca171380..c5de09e82d0980 100644 --- a/runtime/src/bank.rs +++ b/runtime/src/bank.rs @@ -8070,32 +8070,46 @@ impl Bank { new_address: &Pubkey, datapoint_name: &'static str, ) { - if let Some(old_account) = self.get_account_with_fixed_root(old_address) { - if let Some(new_account) = self.get_account_with_fixed_root(new_address) { - datapoint_info!(datapoint_name, ("slot", self.slot, i64)); - - // Burn lamports in the old account - self.capitalization - .fetch_sub(old_account.lamports(), Relaxed); - - // Transfer new account to old account - self.store_account(old_address, &new_account); - - // Clear new account - self.store_account(new_address, &AccountSharedData::default()); - - // Unload a program from the bank's cache - self.loaded_programs_cache - .write() - .unwrap() - .remove_programs([*old_address].into_iter()); + let program_data_address = |address: &Pubkey| { + Pubkey::find_program_address(&[address.as_ref()], &bpf_loader_upgradeable::id()).0 + }; + let maybe_replace_program_account = + |old_address: &Pubkey, new_address: &Pubkey, is_program: bool| { + if let Some(old_account) = self.get_account_with_fixed_root(old_address) { + if let Some(new_account) = self.get_account_with_fixed_root(new_address) { + datapoint_info!(datapoint_name, ("slot", self.slot, i64)); + + // Burn lamports in the old account + self.capitalization + .fetch_sub(old_account.lamports(), Relaxed); + + // Transfer new account to old account + self.store_account(old_address, &new_account); + + // Clear new account + self.store_account(new_address, &AccountSharedData::default()); + + // Unload a program from the bank's cache + if is_program { + self.loaded_programs_cache + .write() + .unwrap() + .remove_programs([*old_address].into_iter()); + } - self.calculate_and_update_accounts_data_size_delta_off_chain( - old_account.data().len(), - new_account.data().len(), - ); - } - } + self.calculate_and_update_accounts_data_size_delta_off_chain( + old_account.data().len(), + new_account.data().len(), + ); + } + } + }; + maybe_replace_program_account( + &program_data_address(old_address), + &program_data_address(new_address), + false, /* is_program */ + ); + maybe_replace_program_account(old_address, new_address, true /* is_program */); } /// Get all the accounts for this bank and calculate stats diff --git a/runtime/src/bank/tests.rs b/runtime/src/bank/tests.rs index 70574666ef734c..0bc68b80f35887 100644 --- a/runtime/src/bank/tests.rs +++ b/runtime/src/bank/tests.rs @@ -8013,42 +8013,55 @@ fn test_compute_active_feature_set() { assert!(bank.feature_set.is_active(&test_feature)); } +fn set_up_account_with_bank(bank: &mut Bank, pubkey: &Pubkey, lamports: u64) -> AccountSharedData { + let new_account = AccountSharedData::from(Account { + lamports, + ..Account::default() + }); + bank.store_account_and_update_capitalization(pubkey, &new_account); + assert_eq!(bank.get_balance(pubkey), lamports); + new_account +} + #[test] fn test_program_replacement() { let mut bank = create_simple_test_bank(0); - // Setup original program account + // Setup original program accounts let old_address = Pubkey::new_unique(); + set_up_account_with_bank(&mut bank, &old_address, 100); + + let (old_data_address, _) = + Pubkey::find_program_address(&[old_address.as_ref()], &bpf_loader_upgradeable::id()); + set_up_account_with_bank(&mut bank, &old_data_address, 102); + + // Setup new program accounts let new_address = Pubkey::new_unique(); - bank.store_account_and_update_capitalization( - &old_address, - &AccountSharedData::from(Account { - lamports: 100, - ..Account::default() - }), - ); - assert_eq!(bank.get_balance(&old_address), 100); - - // Setup new program account - let new_program_account = AccountSharedData::from(Account { - lamports: 123, - ..Account::default() - }); - bank.store_account_and_update_capitalization(&new_address, &new_program_account); - assert_eq!(bank.get_balance(&new_address), 123); + let new_program_account = set_up_account_with_bank(&mut bank, &new_address, 123); + + let (new_data_address, _) = + Pubkey::find_program_address(&[new_address.as_ref()], &bpf_loader_upgradeable::id()); + let new_program_data_account = set_up_account_with_bank(&mut bank, &new_data_address, 125); let original_capitalization = bank.capitalization(); bank.replace_program_account(&old_address, &new_address, "bank-apply_program_replacement"); - // New program account is now empty + // New program accounts are now empty assert_eq!(bank.get_balance(&new_address), 0); + assert_eq!(bank.get_balance(&new_data_address), 0); // Old program account holds the new program account assert_eq!(bank.get_account(&old_address), Some(new_program_account)); - // Lamports in the old token account were burnt - assert_eq!(bank.capitalization(), original_capitalization - 100); + // Old program data account holds the new program data account + assert_eq!( + bank.get_account(&old_data_address), + Some(new_program_data_account) + ); + + // Lamports in the old token accounts were burnt + assert_eq!(bank.capitalization(), original_capitalization - 100 - 102); } fn min_rent_exempt_balance_for_sysvars(bank: &Bank, sysvar_ids: &[Pubkey]) -> u64 {