Unset needs_unlock for rebatched transactions batches (#25095)

* Unset needs_unlock for rebatched transactions batches

* address review comments

ledger/src/blockstore_processor.rs

@@ -294,6 +294,21 @@ fn execute_batches_internal(
     first_err(&results)
 }
 
+fn rebatch_transactions<'a>(
+    lock_results: &'a [Result<()>],
+    bank: &'a Arc<Bank>,
+    sanitized_txs: &'a [SanitizedTransaction],
+    start: usize,
+    end: usize,
+) -> TransactionBatch<'a, 'a> {
+    let txs = &sanitized_txs[start..=end];
+    let results = &lock_results[start..=end];
+    let mut tx_batch = TransactionBatch::new(results.to_vec(), bank, Cow::from(txs));
+    tx_batch.set_needs_unlock(false);
+
+    tx_batch
+}
+
 fn execute_batches(
     bank: &Arc<Bank>,
     batches: &[TransactionBatch],
@@ -339,9 +354,8 @@ fn execute_batches(
             let next_index = index + 1;
             batch_cost = batch_cost.saturating_add(cost);
             if batch_cost >= target_batch_cost || next_index == sanitized_txs.len() {
-                let txs = &sanitized_txs[slice_start..=index];
-                let results = &lock_results[slice_start..=index];
-                let tx_batch = TransactionBatch::new(results.to_vec(), bank, Cow::from(txs));
+                let tx_batch =
+                    rebatch_transactions(&lock_results, bank, &sanitized_txs, slice_start, index);
                 slice_start = next_index;
                 tx_batches.push(tx_batch);
                 batch_cost = 0;
@@ -3922,6 +3936,49 @@ pub mod tests {
         assert_eq!(slot_2_bank.get_hash_age(&slot_2_hash), Some(0));
     }
 
+    #[test]
+    fn test_rebatch_transactions() {
+        let dummy_leader_pubkey = solana_sdk::pubkey::new_rand();
+        let GenesisConfigInfo {
+            genesis_config,
+            mint_keypair,
+            ..
+        } = create_genesis_config_with_leader(500, &dummy_leader_pubkey, 100);
+        let bank = Arc::new(Bank::new_for_tests(&genesis_config));
+
+        let pubkey = solana_sdk::pubkey::new_rand();
+        let keypair2 = Keypair::new();
+        let pubkey2 = solana_sdk::pubkey::new_rand();
+
+        let txs = vec![
+            SanitizedTransaction::from_transaction_for_tests(system_transaction::transfer(
+                &mint_keypair,
+                &pubkey,
+                1,
+                genesis_config.hash(),
+            )),
+            SanitizedTransaction::from_transaction_for_tests(system_transaction::transfer(
+                &keypair2,
+                &pubkey2,
+                1,
+                genesis_config.hash(),
+            )),
+        ];
+
+        let batch = bank.prepare_sanitized_batch(&txs);
+        assert!(batch.needs_unlock());
+
+        let batch2 = rebatch_transactions(
+            batch.lock_results(),
+            &bank,
+            batch.sanitized_transactions(),
+            0,
+            1,
+        );
+        assert!(batch.needs_unlock());
+        assert!(!batch2.needs_unlock());
+    }
+
     #[test]
     fn test_confirm_slot_entries_with_fix() {
         const HASHES_PER_TICK: u64 = 10;

runtime/src/bank.rs

@@ -3858,7 +3858,7 @@ impl Bank {
         let lock_result = transaction.get_account_locks(&self.feature_set).map(|_| ());
         let mut batch =
             TransactionBatch::new(vec![lock_result], self, Cow::Owned(vec![transaction]));
-        batch.needs_unlock = false;
+        batch.set_needs_unlock(false);
         batch
     }
 
@@ -3967,8 +3967,8 @@ impl Bank {
     }
 
     pub fn unlock_accounts(&self, batch: &mut TransactionBatch) {
-        if batch.needs_unlock {
-            batch.needs_unlock = false;
+        if batch.needs_unlock() {
+            batch.set_needs_unlock(false);
             self.rc
                 .accounts
                 .unlock_accounts(batch.sanitized_transactions().iter(), batch.lock_results())

runtime/src/transaction_batch.rs

@@ -9,7 +9,7 @@ pub struct TransactionBatch<'a, 'b> {
     lock_results: Vec<Result<()>>,
     bank: &'a Bank,
     sanitized_txs: Cow<'b, [SanitizedTransaction]>,
-    pub(crate) needs_unlock: bool,
+    needs_unlock: bool,
 }
 
 impl<'a, 'b> TransactionBatch<'a, 'b> {
@@ -38,6 +38,14 @@ impl<'a, 'b> TransactionBatch<'a, 'b> {
     pub fn bank(&self) -> &Bank {
         self.bank
     }
+
+    pub fn set_needs_unlock(&mut self, needs_unlock: bool) {
+        self.needs_unlock = needs_unlock;
+    }
+
+    pub fn needs_unlock(&self) -> bool {
+        self.needs_unlock
+    }
 }
 
 // Unlock all locked accounts in destructor.