diff --git a/core/src/banking_stage.rs b/core/src/banking_stage.rs index 4a10dbe696..2800dc55b8 100644 --- a/core/src/banking_stage.rs +++ b/core/src/banking_stage.rs @@ -509,7 +509,7 @@ impl BankingStage { // expires. let txs = batch.transactions(); let pre_balances = if transaction_status_sender.is_some() { - bank.collect_balances(txs) + bank.collect_balances(batch) } else { vec![] }; @@ -545,7 +545,7 @@ impl BankingStage { .processing_results; if let Some(sender) = transaction_status_sender { - let post_balances = bank.collect_balances(txs); + let post_balances = bank.collect_balances(batch); send_transaction_status_batch( bank.clone(), batch.transactions(), diff --git a/runtime/src/bank.rs b/runtime/src/bank.rs index d5913057b5..842c6b3a61 100644 --- a/runtime/src/bank.rs +++ b/runtime/src/bank.rs @@ -1354,9 +1354,9 @@ impl Bank { } } - pub fn collect_balances(&self, batch: &[Transaction]) -> TransactionBalances { + pub fn collect_balances(&self, batch: &TransactionBatch) -> TransactionBalances { let mut balances: TransactionBalances = vec![]; - for transaction in batch.iter() { + for transaction in OrderedIterator::new(batch.transactions(), batch.iteration_order()) { let mut transaction_balances: Vec = vec![]; for account_key in transaction.message.account_keys.iter() { transaction_balances.push(self.get_balance(account_key)); @@ -2120,7 +2120,7 @@ impl Bank { collect_balances: bool, ) -> (TransactionResults, TransactionBalancesSet) { let pre_balances = if collect_balances { - self.collect_balances(batch.transactions()) + self.collect_balances(batch) } else { vec![] }; @@ -2136,7 +2136,7 @@ impl Bank { signature_count, ); let post_balances = if collect_balances { - self.collect_balances(batch.transactions()) + self.collect_balances(batch) } else { vec![] }; @@ -6800,10 +6800,20 @@ mod tests { vec![program_id], instructions, ); - let balances = bank0.collect_balances(&[tx0, tx1]); + let txs = vec![tx0, tx1]; + let iteration_order: Vec = vec![0, 1]; + let batch = bank0.prepare_batch(&txs, Some(iteration_order)); + let balances = bank0.collect_balances(&batch); assert_eq!(balances.len(), 2); assert_eq!(balances[0], vec![8, 11, 1]); assert_eq!(balances[1], vec![8, 0, 1]); + + let iteration_order: Vec = vec![1, 0]; + let batch = bank0.prepare_batch(&txs, Some(iteration_order)); + let balances = bank0.collect_balances(&batch); + assert_eq!(balances.len(), 2); + assert_eq!(balances[0], vec![8, 0, 1]); + assert_eq!(balances[1], vec![8, 11, 1]); } #[test]