refactor: unlock accounts

accounts-db/src/accounts.rs

@@ -80,11 +80,20 @@ impl AccountLocks {
             if *count == 0 {
                 occupied_entry.remove_entry();
             }
+        } else {
+            debug_assert!(
+                false,
+                "Attempted to remove a read-lock for a key that wasn't read-locked"
+            );
         }
     }
 
     fn unlock_write(&mut self, key: &Pubkey) {
-        self.write_locks.remove(key);
+        let removed = self.write_locks.remove(key);
+        debug_assert!(
+            removed,
+            "Attempted to remove a write-lock for a key that wasn't write-locked"
+        );
     }
 }
 
@@ -616,16 +625,12 @@ impl Accounts {
 
     /// Once accounts are unlocked, new transactions that modify that state can enter the pipeline
     #[allow(clippy::needless_collect)]
-    pub fn unlock_accounts<'a>(
-        &self,
-        txs: impl Iterator<Item = &'a SanitizedTransaction>,
-        results: &[Result<()>],
-    ) {
-        let keys: Vec<_> = txs
-            .zip(results)
-            .filter(|(_, res)| res.is_ok())
-            .map(|(tx, _)| tx.get_account_locks_unchecked())
-            .collect();
+    pub fn unlock_accounts<'a>(&self, txs: impl Iterator<Item = &'a SanitizedTransaction>) {
+        let keys: Vec<_> = txs.map(|tx| tx.get_account_locks_unchecked()).collect();
+        if keys.is_empty() {
+            return;
+        }
+
         let mut account_locks = self.account_locks.lock().unwrap();
         debug!("bank unlock accounts");
         keys.into_iter().for_each(|keys| {
@@ -812,6 +817,7 @@ mod tests {
         },
         std::{
             borrow::Cow,
+            iter,
             sync::atomic::{AtomicBool, AtomicU64, Ordering},
             thread, time,
         },
@@ -1100,7 +1106,7 @@ mod tests {
             let txs = vec![new_sanitized_tx(&[&keypair], message, Hash::default())];
             let results = accounts.lock_accounts(txs.iter(), MAX_TX_ACCOUNT_LOCKS);
             assert_eq!(results[0], Ok(()));
-            accounts.unlock_accounts(txs.iter(), &results);
+            accounts.unlock_accounts(txs.iter());
         }
 
         // Disallow over MAX_TX_ACCOUNT_LOCKS
@@ -1154,9 +1160,10 @@ mod tests {
             instructions,
         );
         let tx = new_sanitized_tx(&[&keypair0], message, Hash::default());
-        let results0 = accounts.lock_accounts([tx.clone()].iter(), MAX_TX_ACCOUNT_LOCKS);
-
-        assert!(results0[0].is_ok());
+        assert_eq!(
+            accounts.lock_accounts(iter::once(&tx), MAX_TX_ACCOUNT_LOCKS),
+            vec![Ok(())]
+        );
         assert_eq!(
             *accounts
                 .account_locks
@@ -1189,10 +1196,14 @@ mod tests {
         );
         let tx1 = new_sanitized_tx(&[&keypair1], message, Hash::default());
         let txs = vec![tx0, tx1];
-        let results1 = accounts.lock_accounts(txs.iter(), MAX_TX_ACCOUNT_LOCKS);
+        assert_eq!(
+            accounts.lock_accounts(txs.iter(), MAX_TX_ACCOUNT_LOCKS),
+            vec![
+                Ok(()), // Read-only account (keypair1) can be referenced multiple times
+                Err(TransactionError::AccountInUse), // Read-only account (keypair1) cannot also be locked as writable
+            ]
+        );
 
-        assert!(results1[0].is_ok()); // Read-only account (keypair1) can be referenced multiple times
-        assert!(results1[1].is_err()); // Read-only account (keypair1) cannot also be locked as writable
         assert_eq!(
             *accounts
                 .account_locks
@@ -1203,9 +1214,8 @@ mod tests {
                 .unwrap(),
             2
         );
-
-        accounts.unlock_accounts([tx].iter(), &results0);
-        accounts.unlock_accounts(txs.iter(), &results1);
+        accounts.unlock_accounts(iter::once(&tx));
+        accounts.unlock_accounts(iter::once(&txs[0]));
         let instructions = vec![CompiledInstruction::new(2, &(), vec![0, 1])];
         let message = Message::new_with_compiled_instructions(
             1,
@@ -1216,8 +1226,10 @@ mod tests {
             instructions,
         );
         let tx = new_sanitized_tx(&[&keypair1], message, Hash::default());
-        let results2 = accounts.lock_accounts([tx].iter(), MAX_TX_ACCOUNT_LOCKS);
-        assert!(results2[0].is_ok()); // Now keypair1 account can be locked as writable
+        assert_eq!(
+            accounts.lock_accounts(iter::once(&tx), MAX_TX_ACCOUNT_LOCKS),
+            vec![Ok(())] // Now keypair1 account can be locked as writable
+        );
 
         // Check that read-only lock with zero references is deleted
         assert!(accounts
@@ -1280,28 +1292,29 @@ mod tests {
             let results = accounts_clone
                 .clone()
                 .lock_accounts(txs.iter(), MAX_TX_ACCOUNT_LOCKS);
-            for result in results.iter() {
-                if result.is_ok() {
+            let locked_txs = txs.iter().zip(results).filter_map(|(tx, result)| {
+                result.ok().map(|_| {
                     counter_clone.clone().fetch_add(1, Ordering::SeqCst);
-                }
-            }
-            accounts_clone.unlock_accounts(txs.iter(), &results);
+                    tx
+                })
+            });
+            accounts_clone.unlock_accounts(locked_txs);
             if exit_clone.clone().load(Ordering::Relaxed) {
                 break;
             }
         });
         let counter_clone = counter;
         for _ in 0..5 {
-            let txs = vec![readonly_tx.clone()];
+            let tx = &readonly_tx;
             let results = accounts_arc
                 .clone()
-                .lock_accounts(txs.iter(), MAX_TX_ACCOUNT_LOCKS);
+                .lock_accounts(iter::once(tx), MAX_TX_ACCOUNT_LOCKS);
             if results[0].is_ok() {
                 let counter_value = counter_clone.clone().load(Ordering::SeqCst);
                 thread::sleep(time::Duration::from_millis(50));
                 assert_eq!(counter_value, counter_clone.clone().load(Ordering::SeqCst));
+                accounts_arc.unlock_accounts(iter::once(tx));
             }
-            accounts_arc.unlock_accounts(txs.iter(), &results);
             thread::sleep(time::Duration::from_millis(50));
         }
         exit.store(true, Ordering::Relaxed);
@@ -1442,9 +1455,14 @@ mod tests {
             MAX_TX_ACCOUNT_LOCKS,
         );
 
-        assert!(results[0].is_ok()); // Read-only account (keypair0) can be referenced multiple times
-        assert!(results[1].is_err()); // is not locked due to !qos_results[1].is_ok()
-        assert!(results[2].is_ok()); // Read-only account (keypair0) can be referenced multiple times
+        assert_eq!(
+            results,
+            vec![
+                Ok(()), // Read-only account (keypair0) can be referenced multiple times
+                Err(TransactionError::WouldExceedMaxBlockCostLimit), // is not locked due to !qos_results[1].is_ok()
+                Ok(()), // Read-only account (keypair0) can be referenced multiple times
+            ],
+        );
 
         // verify that keypair0 read-only lock twice (for tx0 and tx2)
         assert_eq!(
@@ -1466,7 +1484,7 @@ mod tests {
             .get(&keypair2.pubkey())
             .is_none());
 
-        accounts.unlock_accounts(txs.iter(), &results);
+        accounts.unlock_accounts([&txs[0], &txs[2]].into_iter());
 
         // check all locks to be removed
         assert!(accounts

runtime/src/bank.rs

@@ -4376,13 +4376,8 @@ impl Bank {
         account_overrides
     }
 
-    pub fn unlock_accounts(&self, batch: &mut TransactionBatch) {
-        if batch.needs_unlock() {
-            batch.set_needs_unlock(false);
-            self.rc
-                .accounts
-                .unlock_accounts(batch.sanitized_transactions().iter(), batch.lock_results())
-        }
+    pub fn unlock_accounts<'a>(&self, txs: impl Iterator<Item = &'a SanitizedTransaction>) {
+        self.rc.accounts.unlock_accounts(txs)
     }
 
     pub fn remove_unrooted_slots(&self, slots: &[(Slot, BankId)]) {

runtime/src/transaction_batch.rs

@@ -51,7 +51,15 @@ impl<'a, 'b> TransactionBatch<'a, 'b> {
 // Unlock all locked accounts in destructor.
 impl<'a, 'b> Drop for TransactionBatch<'a, 'b> {
     fn drop(&mut self) {
-        self.bank.unlock_accounts(self)
+        if self.needs_unlock() {
+            self.set_needs_unlock(false);
+            self.bank.unlock_accounts(
+                self.sanitized_transactions()
+                    .iter()
+                    .zip(self.lock_results())
+                    .filter_map(|(tx, res)| res.as_ref().ok().map(|_| tx)),
+            )
+        }
     }
 }