Extend getConfirmedBlock rpc to return account pre- and post-balances (#7543)

automerge

runtime/src/bank.rs

@@ -161,6 +161,20 @@ pub struct TransactionResults {
     pub fee_collection_results: Vec<Result<()>>,
     pub processing_results: Vec<TransactionProcessResult>,
 }
+pub struct TransactionBalancesSet {
+    pub pre_balances: TransactionBalances,
+    pub post_balances: TransactionBalances,
+}
+impl TransactionBalancesSet {
+    pub fn new(pre_balances: TransactionBalances, post_balances: TransactionBalances) -> Self {
+        assert_eq!(pre_balances.len(), post_balances.len());
+        Self {
+            pre_balances,
+            post_balances,
+        }
+    }
+}
+pub type TransactionBalances = Vec<Vec<u64>>;
 
 #[derive(Clone, Debug, Eq, PartialEq)]
 pub enum HashAgeKind {
@@ -1030,6 +1044,18 @@ impl Bank {
         self.check_signatures(txs, iteration_order, age_results, &mut error_counters)
     }
 
+    pub fn collect_balances(&self, batch: &[Transaction]) -> TransactionBalances {
+        let mut balances: TransactionBalances = vec![];
+        for transaction in batch.iter() {
+            let mut transaction_balances: Vec<u64> = vec![];
+            for account_key in transaction.message.account_keys.iter() {
+                transaction_balances.push(self.get_balance(account_key));
+            }
+            balances.push(transaction_balances);
+        }
+        balances
+    }
+
     fn update_error_counters(error_counters: &ErrorCounters) {
         if 0 != error_counters.blockhash_not_found {
             inc_new_counter_error!(
@@ -1372,24 +1398,40 @@ impl Bank {
         &self,
         batch: &TransactionBatch,
         max_age: usize,
-    ) -> TransactionResults {
+        collect_balances: bool,
+    ) -> (TransactionResults, TransactionBalancesSet) {
+        let pre_balances = if collect_balances {
+            self.collect_balances(batch.transactions())
+        } else {
+            vec![]
+        };
         let (mut loaded_accounts, executed, _, tx_count, signature_count) =
             self.load_and_execute_transactions(batch, max_age);
 
-        self.commit_transactions(
+        let results = self.commit_transactions(
             batch.transactions(),
             batch.iteration_order(),
             &mut loaded_accounts,
             &executed,
             tx_count,
             signature_count,
+        );
+        let post_balances = if collect_balances {
+            self.collect_balances(batch.transactions())
+        } else {
+            vec![]
+        };
+        (
+            results,
+            TransactionBalancesSet::new(pre_balances, post_balances),
         )
     }
 
     #[must_use]
     pub fn process_transactions(&self, txs: &[Transaction]) -> Vec<Result<()>> {
         let batch = self.prepare_batch(txs, None);
-        self.load_execute_and_commit_transactions(&batch, MAX_RECENT_BLOCKHASHES)
+        self.load_execute_and_commit_transactions(&batch, MAX_RECENT_BLOCKHASHES, false)
+            .0
             .fee_collection_results
     }
 
@@ -1816,7 +1858,7 @@ mod tests {
         clock::DEFAULT_TICKS_PER_SLOT,
         epoch_schedule::MINIMUM_SLOTS_PER_EPOCH,
         genesis_config::create_genesis_config,
-        instruction::{Instruction, InstructionError},
+        instruction::{CompiledInstruction, Instruction, InstructionError},
         message::{Message, MessageHeader},
         nonce_instruction, nonce_state,
         poh_config::PohConfig,
@@ -3221,7 +3263,8 @@ mod tests {
 
         let lock_result = bank.prepare_batch(&pay_alice, None);
         let results_alice = bank
-            .load_execute_and_commit_transactions(&lock_result, MAX_RECENT_BLOCKHASHES)
+            .load_execute_and_commit_transactions(&lock_result, MAX_RECENT_BLOCKHASHES, false)
+            .0
             .fee_collection_results;
         assert_eq!(results_alice[0], Ok(()));
 
@@ -4731,4 +4774,94 @@ mod tests {
         /* Check fee charged */
         assert_eq!(bank.get_balance(&custodian_pubkey), 4_630_000);
     }
+
+    #[test]
+    fn test_collect_balances() {
+        let (genesis_config, _mint_keypair) = create_genesis_config(500);
+        let parent = Arc::new(Bank::new(&genesis_config));
+        let bank0 = Arc::new(new_from_parent(&parent));
+
+        let keypair = Keypair::new();
+        let pubkey0 = Pubkey::new_rand();
+        let pubkey1 = Pubkey::new_rand();
+        let program_id = Pubkey::new(&[2; 32]);
+        let keypair_account = Account::new(8, 0, &program_id);
+        let account0 = Account::new(11, 0, &program_id);
+        let program_account = Account::new(1, 10, &Pubkey::default());
+        bank0.store_account(&keypair.pubkey(), &keypair_account);
+        bank0.store_account(&pubkey0, &account0);
+        bank0.store_account(&program_id, &program_account);
+
+        let instructions = vec![CompiledInstruction::new(1, &(), vec![0])];
+        let tx0 = Transaction::new_with_compiled_instructions(
+            &[&keypair],
+            &[pubkey0],
+            Hash::default(),
+            vec![program_id],
+            instructions,
+        );
+        let instructions = vec![CompiledInstruction::new(1, &(), vec![0])];
+        let tx1 = Transaction::new_with_compiled_instructions(
+            &[&keypair],
+            &[pubkey1],
+            Hash::default(),
+            vec![program_id],
+            instructions,
+        );
+        let balances = bank0.collect_balances(&[tx0, tx1]);
+        assert_eq!(balances.len(), 2);
+        assert_eq!(balances[0], vec![8, 11, 1]);
+        assert_eq!(balances[1], vec![8, 0, 1]);
+    }
+
+    #[test]
+    fn test_pre_post_transaction_balances() {
+        let (mut genesis_config, _mint_keypair) = create_genesis_config(500);
+        let fee_calculator = FeeCalculator::new(1, 0);
+        genesis_config.fee_calculator = fee_calculator;
+        let parent = Arc::new(Bank::new(&genesis_config));
+        let bank0 = Arc::new(new_from_parent(&parent));
+
+        let keypair0 = Keypair::new();
+        let keypair1 = Keypair::new();
+        let pubkey0 = Pubkey::new_rand();
+        let pubkey1 = Pubkey::new_rand();
+        let pubkey2 = Pubkey::new_rand();
+        let keypair0_account = Account::new(8, 0, &Pubkey::default());
+        let keypair1_account = Account::new(9, 0, &Pubkey::default());
+        let account0 = Account::new(11, 0, &&Pubkey::default());
+        bank0.store_account(&keypair0.pubkey(), &keypair0_account);
+        bank0.store_account(&keypair1.pubkey(), &keypair1_account);
+        bank0.store_account(&pubkey0, &account0);
+
+        let blockhash = bank0.last_blockhash();
+
+        let tx0 = system_transaction::transfer(&keypair0, &pubkey0, 2, blockhash.clone());
+        let tx1 = system_transaction::transfer(&Keypair::new(), &pubkey1, 2, blockhash.clone());
+        let tx2 = system_transaction::transfer(&keypair1, &pubkey2, 12, blockhash.clone());
+        let txs = vec![tx0, tx1, tx2];
+
+        let lock_result = bank0.prepare_batch(&txs, None);
+        let (transaction_results, transaction_balances_set) =
+            bank0.load_execute_and_commit_transactions(&lock_result, MAX_RECENT_BLOCKHASHES, true);
+
+        assert_eq!(transaction_balances_set.pre_balances.len(), 3);
+        assert_eq!(transaction_balances_set.post_balances.len(), 3);
+
+        assert!(transaction_results.processing_results[0].0.is_ok());
+        assert_eq!(transaction_balances_set.pre_balances[0], vec![8, 11, 1]);
+        assert_eq!(transaction_balances_set.post_balances[0], vec![5, 13, 1]);
+
+        // Failed transactions still produce balance sets
+        // This is a TransactionError - not possible to charge fees
+        assert!(transaction_results.processing_results[1].0.is_err());
+        assert_eq!(transaction_balances_set.pre_balances[1], vec![0, 0, 1]);
+        assert_eq!(transaction_balances_set.post_balances[1], vec![0, 0, 1]);
+
+        // Failed transactions still produce balance sets
+        // This is an InstructionError - fees charged
+        assert!(transaction_results.processing_results[2].0.is_err());
+        assert_eq!(transaction_balances_set.pre_balances[2], vec![9, 0, 1]);
+        assert_eq!(transaction_balances_set.post_balances[2], vec![8, 0, 1]);
+    }
 }