From d0027ddb6472e97b237e8002ef9f4b933ec67d2b Mon Sep 17 00:00:00 2001 From: Andrew Fitzgerald Date: Mon, 5 Aug 2024 10:14:55 -0500 Subject: [PATCH] Separate AccountLocks + Refactor (#2390) --- accounts-db/src/account_locks.rs | 155 ++++++++++++++++++++++++ accounts-db/src/accounts.rs | 195 +++++-------------------------- accounts-db/src/lib.rs | 1 + 3 files changed, 185 insertions(+), 166 deletions(-) create mode 100644 accounts-db/src/account_locks.rs diff --git a/accounts-db/src/account_locks.rs b/accounts-db/src/account_locks.rs new file mode 100644 index 00000000000000..dcd512e7d72984 --- /dev/null +++ b/accounts-db/src/account_locks.rs @@ -0,0 +1,155 @@ +#[cfg(feature = "dev-context-only-utils")] +use qualifier_attr::qualifiers; +use { + ahash::{AHashMap, AHashSet}, + solana_sdk::{pubkey::Pubkey, transaction::TransactionError}, + std::collections::hash_map, +}; + +#[derive(Debug, Default)] +pub struct AccountLocks { + write_locks: AHashSet, + readonly_locks: AHashMap, +} + +impl AccountLocks { + /// Lock the account keys in `keys` for a transaction. + /// The bool in the tuple indicates if the account is writable. + /// Returns an error if any of the accounts are already locked in a way + /// that conflicts with the requested lock. + pub fn try_lock_accounts<'a>( + &mut self, + keys: impl Iterator + Clone, + ) -> Result<(), TransactionError> { + for (key, writable) in keys.clone() { + if writable { + if !self.can_write_lock(key) { + return Err(TransactionError::AccountInUse); + } + } else if !self.can_read_lock(key) { + return Err(TransactionError::AccountInUse); + } + } + + for (key, writable) in keys { + if writable { + self.lock_write(key); + } else { + self.lock_readonly(key); + } + } + + Ok(()) + } + + /// Unlock the account keys in `keys` after a transaction. + /// The bool in the tuple indicates if the account is writable. + /// In debug-mode this function will panic if an attempt is made to unlock + /// an account that wasn't locked in the way requested. + pub fn unlock_accounts<'a>(&mut self, keys: impl Iterator) { + for (k, writable) in keys { + if writable { + self.unlock_write(k); + } else { + self.unlock_readonly(k); + } + } + } + + #[cfg_attr(feature = "dev-context-only-utils", qualifiers(pub))] + fn is_locked_readonly(&self, key: &Pubkey) -> bool { + self.readonly_locks + .get(key) + .map_or(false, |count| *count > 0) + } + + #[cfg_attr(feature = "dev-context-only-utils", qualifiers(pub))] + fn is_locked_write(&self, key: &Pubkey) -> bool { + self.write_locks.contains(key) + } + + fn can_read_lock(&self, key: &Pubkey) -> bool { + // If the key is not write-locked, it can be read-locked + !self.is_locked_write(key) + } + + fn can_write_lock(&self, key: &Pubkey) -> bool { + // If the key is not read-locked or write-locked, it can be write-locked + !self.is_locked_readonly(key) && !self.is_locked_write(key) + } + + fn lock_readonly(&mut self, key: &Pubkey) { + *self.readonly_locks.entry(*key).or_default() += 1; + } + + fn lock_write(&mut self, key: &Pubkey) { + self.write_locks.insert(*key); + } + + fn unlock_readonly(&mut self, key: &Pubkey) { + if let hash_map::Entry::Occupied(mut occupied_entry) = self.readonly_locks.entry(*key) { + let count = occupied_entry.get_mut(); + *count -= 1; + 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) { + let removed = self.write_locks.remove(key); + debug_assert!( + removed, + "Attempted to remove a write-lock for a key that wasn't write-locked" + ); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_account_locks() { + let mut account_locks = AccountLocks::default(); + + let key1 = Pubkey::new_unique(); + let key2 = Pubkey::new_unique(); + + // Add write and read-lock. + let result = account_locks.try_lock_accounts([(&key1, true), (&key2, false)].into_iter()); + assert!(result.is_ok()); + + // Try to add duplicate write-lock. + let result = account_locks.try_lock_accounts([(&key1, true)].into_iter()); + assert_eq!(result, Err(TransactionError::AccountInUse)); + + // Try to add write lock on read-locked account. + let result = account_locks.try_lock_accounts([(&key2, true)].into_iter()); + assert_eq!(result, Err(TransactionError::AccountInUse)); + + // Try to add read lock on write-locked account. + let result = account_locks.try_lock_accounts([(&key1, false)].into_iter()); + assert_eq!(result, Err(TransactionError::AccountInUse)); + + // Add read lock on read-locked account. + let result = account_locks.try_lock_accounts([(&key2, false)].into_iter()); + assert!(result.is_ok()); + + // Unlock write and read locks. + account_locks.unlock_accounts([(&key1, true), (&key2, false)].into_iter()); + + // No more remaining write-locks. Read-lock remains. + assert!(!account_locks.is_locked_write(&key1)); + assert!(account_locks.is_locked_readonly(&key2)); + + // Unlock read lock. + account_locks.unlock_accounts([(&key2, false)].into_iter()); + assert!(!account_locks.is_locked_readonly(&key2)); + } +} diff --git a/accounts-db/src/accounts.rs b/accounts-db/src/accounts.rs index bb8656648c23c9..d781bb45e11b8a 100644 --- a/accounts-db/src/accounts.rs +++ b/accounts-db/src/accounts.rs @@ -1,5 +1,6 @@ use { crate::{ + account_locks::AccountLocks, accounts_db::{ AccountStorageEntry, AccountsAddRootTiming, AccountsDb, LoadHint, LoadedAccount, ScanAccountStorageData, ScanStorageResult, VerifyAccountsHashAndLamportsConfig, @@ -8,7 +9,6 @@ use { ancestors::Ancestors, storable_accounts::StorableAccounts, }, - ahash::{AHashMap, AHashSet}, dashmap::DashMap, log::*, solana_sdk::{ @@ -18,13 +18,13 @@ use { message::v0::{LoadedAddresses, MessageAddressTableLookup}, pubkey::Pubkey, slot_hashes::SlotHashes, - transaction::{Result, SanitizedTransaction, TransactionError}, + transaction::{Result, SanitizedTransaction}, transaction_context::TransactionAccount, }, solana_svm_transaction::svm_message::SVMMessage, std::{ cmp::Reverse, - collections::{hash_map, BinaryHeap, HashSet}, + collections::{BinaryHeap, HashSet}, ops::RangeBounds, sync::{ atomic::{AtomicUsize, Ordering}, @@ -35,58 +35,6 @@ use { pub type PubkeyAccountSlot = (Pubkey, AccountSharedData, Slot); -#[derive(Debug, Default)] -pub struct AccountLocks { - write_locks: AHashSet, - readonly_locks: AHashMap, -} - -impl AccountLocks { - fn is_locked_readonly(&self, key: &Pubkey) -> bool { - self.readonly_locks - .get(key) - .map_or(false, |count| *count > 0) - } - - fn is_locked_write(&self, key: &Pubkey) -> bool { - self.write_locks.contains(key) - } - - fn insert_new_readonly(&mut self, key: &Pubkey) { - assert!(self.readonly_locks.insert(*key, 1).is_none()); - } - - fn lock_readonly(&mut self, key: &Pubkey) -> bool { - self.readonly_locks.get_mut(key).map_or(false, |count| { - *count += 1; - true - }) - } - - fn unlock_readonly(&mut self, key: &Pubkey) { - if let hash_map::Entry::Occupied(mut occupied_entry) = self.readonly_locks.entry(*key) { - let count = occupied_entry.get_mut(); - *count -= 1; - 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) { - let removed = self.write_locks.remove(key); - debug_assert!( - removed, - "Attempted to remove a write-lock for a key that wasn't write-locked" - ); - } -} - struct TransactionAccountLocksIterator<'a, T: SVMMessage> { transaction: &'a T, } @@ -561,48 +509,6 @@ impl Accounts { self.accounts_db.store_uncached(slot, &[(pubkey, account)]); } - fn lock_account<'a>( - &self, - account_locks: &mut AccountLocks, - keys: impl Iterator + Clone, - ) -> Result<()> { - for (k, writable) in keys.clone() { - if writable { - if account_locks.is_locked_write(k) || account_locks.is_locked_readonly(k) { - debug!("Writable account in use: {:?}", k); - return Err(TransactionError::AccountInUse); - } - } else if account_locks.is_locked_write(k) { - debug!("Read-only account in use: {:?}", k); - return Err(TransactionError::AccountInUse); - } - } - - for (k, writable) in keys { - if writable { - account_locks.write_locks.insert(*k); - } else if !account_locks.lock_readonly(k) { - account_locks.insert_new_readonly(k); - } - } - - Ok(()) - } - - fn unlock_account<'a>( - &self, - account_locks: &mut AccountLocks, - keys: impl Iterator, - ) { - for (k, writable) in keys { - if writable { - account_locks.unlock_write(k); - } else { - account_locks.unlock_readonly(k); - } - } - } - /// This function will prevent multiple threads from modifying the same account state at the /// same time #[must_use] @@ -653,7 +559,7 @@ impl Accounts { .into_iter() .map(|tx_account_locks_result| match tx_account_locks_result { Ok(tx_account_locks) => { - self.lock_account(account_locks, tx_account_locks.accounts_with_is_writable()) + account_locks.try_lock_accounts(tx_account_locks.accounts_with_is_writable()) } Err(err) => Err(err), }) @@ -674,10 +580,7 @@ impl Accounts { for (tx, res) in txs_and_results { if res.is_ok() { let tx_account_locks = TransactionAccountLocksIterator::new(tx.message()); - self.unlock_account( - &mut account_locks, - tx_account_locks.accounts_with_is_writable(), - ); + account_locks.unlock_accounts(tx_account_locks.accounts_with_is_writable()); } } } @@ -714,7 +617,7 @@ mod tests { message::{Message, MessageHeader}, native_loader, signature::{signers::Signers, Keypair, Signer}, - transaction::{Transaction, MAX_TX_ACCOUNT_LOCKS}, + transaction::{Transaction, TransactionError, MAX_TX_ACCOUNT_LOCKS}, }, std::{ borrow::Cow, @@ -1046,16 +949,11 @@ mod tests { let results0 = accounts.lock_accounts([tx.clone()].iter(), MAX_TX_ACCOUNT_LOCKS); assert_eq!(results0, vec![Ok(())]); - assert_eq!( - *accounts - .account_locks - .lock() - .unwrap() - .readonly_locks - .get(&keypair1.pubkey()) - .unwrap(), - 1 - ); + assert!(accounts + .account_locks + .lock() + .unwrap() + .is_locked_readonly(&keypair1.pubkey())); let instructions = vec![CompiledInstruction::new(2, &(), vec![0, 1])]; let message = Message::new_with_compiled_instructions( @@ -1086,16 +984,11 @@ mod tests { Err(TransactionError::AccountInUse), // Read-only account (keypair1) cannot also be locked as writable ], ); - assert_eq!( - *accounts - .account_locks - .lock() - .unwrap() - .readonly_locks - .get(&keypair1.pubkey()) - .unwrap(), - 2 - ); + assert!(accounts + .account_locks + .lock() + .unwrap() + .is_locked_readonly(&keypair1.pubkey())); accounts.unlock_accounts(iter::once(&tx).zip(&results0)); accounts.unlock_accounts(txs.iter().zip(&results1)); @@ -1120,8 +1013,7 @@ mod tests { .account_locks .lock() .unwrap() - .readonly_locks - .contains_key(&keypair1.pubkey())); + .is_locked_readonly(&keypair1.pubkey())); } #[test] @@ -1235,29 +1127,22 @@ mod tests { assert!(results0[0].is_ok()); // Instruction program-id account demoted to readonly - assert_eq!( - *accounts - .account_locks - .lock() - .unwrap() - .readonly_locks - .get(&native_loader::id()) - .unwrap(), - 1 - ); + assert!(accounts + .account_locks + .lock() + .unwrap() + .is_locked_readonly(&native_loader::id())); // Non-program accounts remain writable assert!(accounts .account_locks .lock() .unwrap() - .write_locks - .contains(&keypair0.pubkey())); + .is_locked_write(&keypair0.pubkey())); assert!(accounts .account_locks .lock() .unwrap() - .write_locks - .contains(&keypair1.pubkey())); + .is_locked_write(&keypair1.pubkey())); } impl Accounts { @@ -1346,40 +1231,18 @@ mod tests { ], ); - // verify that keypair0 read-only lock twice (for tx0 and tx2) - assert_eq!( - *accounts - .account_locks - .lock() - .unwrap() - .readonly_locks - .get(&keypair0.pubkey()) - .unwrap(), - 2 - ); - // verify that keypair2 (for tx1) is not write-locked - assert!(!accounts - .account_locks - .lock() - .unwrap() - .write_locks - .contains(&keypair2.pubkey())); - - accounts.unlock_accounts(txs.iter().zip(&results)); - - // check all locks to be removed + // verify that keypair0 read-only locked assert!(accounts .account_locks .lock() .unwrap() - .readonly_locks - .is_empty()); - assert!(accounts + .is_locked_readonly(&keypair0.pubkey())); + // verify that keypair2 (for tx1) is not write-locked + assert!(!accounts .account_locks .lock() .unwrap() - .write_locks - .is_empty()); + .is_locked_write(&keypair2.pubkey())); } #[test] diff --git a/accounts-db/src/lib.rs b/accounts-db/src/lib.rs index da4f3a9549db67..3cc8f686eff2e0 100644 --- a/accounts-db/src/lib.rs +++ b/accounts-db/src/lib.rs @@ -5,6 +5,7 @@ extern crate lazy_static; pub mod account_info; +mod account_locks; pub mod account_storage; pub mod accounts; mod accounts_cache;