From 853c670254de995a82d64179a35fe241114c8982 Mon Sep 17 00:00:00 2001
From: Justin Starry <justin@solana.com>
Date: Wed, 6 Mar 2024 15:09:55 +0000
Subject: [PATCH] Add support for partial tx batch unlocking

---
 runtime/src/transaction_batch.rs | 83 +++++++++++++++++++++++++-------
 1 file changed, 65 insertions(+), 18 deletions(-)

diff --git a/runtime/src/transaction_batch.rs b/runtime/src/transaction_batch.rs
index 9d0ff5fb7ce007..e7fd3fd358b44f 100644
--- a/runtime/src/transaction_batch.rs
+++ b/runtime/src/transaction_batch.rs
@@ -46,6 +46,26 @@ impl<'a, 'b> TransactionBatch<'a, 'b> {
     pub fn needs_unlock(&self) -> bool {
         self.needs_unlock
     }
+
+    /// For every error result, if the corresponding transaction is
+    /// still locked, unlock the transaction and then record the new error.
+    pub fn unlock_failures(&mut self, transaction_results: Vec<Result<()>>) {
+        assert_eq!(self.lock_results.len(), transaction_results.len());
+        // Shouldn't happen but if a batch was marked as not needing an unlock,
+        // don't unlock failures.
+        if !self.needs_unlock() {
+            return;
+        }
+
+        let txs_and_results = transaction_results
+            .iter()
+            .enumerate()
+            .filter(|(index, result)| result.is_err() && self.lock_results[*index].is_ok())
+            .map(|(index, _)| (&self.sanitized_txs[index], &self.lock_results[index]));
+
+        self.bank.unlock_accounts(txs_and_results);
+        self.lock_results = transaction_results;
+    }
 }
 
 // Unlock all locked accounts in destructor.
@@ -67,12 +87,12 @@ mod tests {
     use {
         super::*,
         crate::genesis_utils::{create_genesis_config_with_leader, GenesisConfigInfo},
-        solana_sdk::{signature::Keypair, system_transaction},
+        solana_sdk::{signature::Keypair, system_transaction, transaction::TransactionError},
     };
 
     #[test]
     fn test_transaction_batch() {
-        let (bank, txs) = setup();
+        let (bank, txs) = setup(false);
 
         // Test getting locked accounts
         let batch = bank.prepare_sanitized_batch(&txs);
@@ -94,7 +114,7 @@ mod tests {
 
     #[test]
     fn test_simulation_batch() {
-        let (bank, txs) = setup();
+        let (bank, txs) = setup(false);
 
         // Prepare batch without locks
         let batch = bank.prepare_unlocked_batch_from_single_tx(&txs[0]);
@@ -109,7 +129,37 @@ mod tests {
         assert!(batch3.lock_results().iter().all(|x| x.is_ok()));
     }
 
-    fn setup() -> (Bank, Vec<SanitizedTransaction>) {
+    #[test]
+    fn test_unlock_failures() {
+        let (bank, txs) = setup(true);
+
+        // Test getting locked accounts
+        let mut batch = bank.prepare_sanitized_batch(&txs);
+        assert_eq!(
+            batch.lock_results,
+            vec![Ok(()), Err(TransactionError::AccountInUse), Ok(())]
+        );
+
+        let qos_results = vec![
+            Ok(()),
+            Err(TransactionError::AccountInUse),
+            Err(TransactionError::WouldExceedMaxBlockCostLimit),
+        ];
+        batch.unlock_failures(qos_results.clone());
+        assert_eq!(batch.lock_results, qos_results);
+
+        // Dropping the batch should unlock remaining locked transactions
+        drop(batch);
+
+        // The next batch should be able to lock all but the conflicting tx
+        let batch2 = bank.prepare_sanitized_batch(&txs);
+        assert_eq!(
+            batch2.lock_results,
+            vec![Ok(()), Err(TransactionError::AccountInUse), Ok(())]
+        );
+    }
+
+    fn setup(insert_conflicting_tx: bool) -> (Bank, Vec<SanitizedTransaction>) {
         let dummy_leader_pubkey = solana_sdk::pubkey::new_rand();
         let GenesisConfigInfo {
             genesis_config,
@@ -122,20 +172,17 @@ mod tests {
         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 mut txs = vec![SanitizedTransaction::from_transaction_for_tests(
+            system_transaction::transfer(&mint_keypair, &pubkey, 1, genesis_config.hash()),
+        )];
+        if insert_conflicting_tx {
+            txs.push(SanitizedTransaction::from_transaction_for_tests(
+                system_transaction::transfer(&mint_keypair, &pubkey2, 1, genesis_config.hash()),
+            ));
+        }
+        txs.push(SanitizedTransaction::from_transaction_for_tests(
+            system_transaction::transfer(&keypair2, &pubkey2, 1, genesis_config.hash()),
+        ));
 
         (bank, txs)
     }