From 7445684e3320e8ad37adb565549f9f4fd41b7813 Mon Sep 17 00:00:00 2001 From: Justin Starry Date: Wed, 6 Mar 2024 01:01:07 +0000 Subject: [PATCH] refactor: unlock accounts --- accounts-db/src/accounts.rs | 88 +++++++++++++++++++------------- runtime/src/bank.rs | 9 +--- runtime/src/transaction_batch.rs | 10 +++- 3 files changed, 64 insertions(+), 43 deletions(-) diff --git a/accounts-db/src/accounts.rs b/accounts-db/src/accounts.rs index 371db9eb08c095..97378849cdd631 100644 --- a/accounts-db/src/accounts.rs +++ b/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, - 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) { + 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 diff --git a/runtime/src/bank.rs b/runtime/src/bank.rs index 3ea316f857a2bc..82dfdbcb756d19 100644 --- a/runtime/src/bank.rs +++ b/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) { + self.rc.accounts.unlock_accounts(txs) } pub fn remove_unrooted_slots(&self, slots: &[(Slot, BankId)]) { diff --git a/runtime/src/transaction_batch.rs b/runtime/src/transaction_batch.rs index 66711fd5a1acd5..ac2dc01ed3c9e5 100644 --- a/runtime/src/transaction_batch.rs +++ b/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)), + ) + } } }