accounts-db: relax intrabatch account locks

accounts-db/Cargo.toml

@@ -86,6 +86,7 @@ crate-type = ["lib"]
 name = "solana_accounts_db"
 
 [dev-dependencies]
+agave-reserved-account-keys = { workspace = true }
 assert_matches = { workspace = true }
 criterion = { workspace = true }
 libsecp256k1 = { workspace = true }
@@ -98,6 +99,7 @@ solana-compute-budget = { workspace = true }
 solana-instruction = { workspace = true }
 solana-logger = { workspace = true }
 solana-sdk-ids = { workspace = true }
+solana-signature = { workspace = true, features = ["rand"] }
 solana-slot-history = { workspace = true }
 static_assertions = { workspace = true }
 strum = { workspace = true, features = ["derive"] }

accounts-db/benches/bench_lock_accounts.rs

@@ -62,16 +62,21 @@ fn create_test_transactions(lock_count: usize, read_conflicts: bool) -> Vec<Sani
 fn bench_entry_lock_accounts(c: &mut Criterion) {
     let mut group = c.benchmark_group("bench_lock_accounts");
 
-    for (batch_size, lock_count, read_conflicts) in
-        iproduct!(BATCH_SIZES, LOCK_COUNTS, [false, true])
+    for (batch_size, lock_count, read_conflicts, relax_intrabatch_account_locks) in
+        iproduct!(BATCH_SIZES, LOCK_COUNTS, [false, true], [false, true])
     {
         let name = format!(
-            "batch_size_{batch_size}_locks_count_{lock_count}{}",
+            "batch_size_{batch_size}_locks_count_{lock_count}{}{}",
             if read_conflicts {
                 "_read_conflicts"
             } else {
                 ""
-            }
+            },
+            if relax_intrabatch_account_locks {
+                "_simd83"
+            } else {
+                "_old"
+            },
         );
 
         let accounts_db = AccountsDb::new_single_for_tests();
@@ -84,8 +89,11 @@ fn bench_entry_lock_accounts(c: &mut Criterion) {
         group.bench_function(name.as_str(), move |b| {
             b.iter(|| {
                 for batch in &transaction_batches {
-                    let results =
-                        accounts.lock_accounts(black_box(batch.iter()), MAX_TX_ACCOUNT_LOCKS);
+                    let results = accounts.lock_accounts(
+                        black_box(batch.iter()),
+                        MAX_TX_ACCOUNT_LOCKS,
+                        relax_intrabatch_account_locks,
+                    );
                     accounts.unlock_accounts(batch.iter().zip(&results));
                 }
             })

accounts-db/src/account_locks.rs

@@ -5,13 +5,13 @@ use {
     solana_message::AccountKeys,
     solana_pubkey::Pubkey,
     solana_transaction::sanitized::MAX_TX_ACCOUNT_LOCKS,
-    solana_transaction_error::TransactionError,
+    solana_transaction_error::{TransactionError, TransactionResult},
     std::{cell::RefCell, collections::hash_map},
 };
 
 #[derive(Debug, Default)]
 pub struct AccountLocks {
-    write_locks: AHashSet<Pubkey>,
+    write_locks: AHashMap<Pubkey, u64>,
     readonly_locks: AHashMap<Pubkey, u64>,
 }
 
@@ -20,31 +20,39 @@ impl AccountLocks {
     /// 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<Item = (&'a Pubkey, bool)> + 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);
-            }
-        }
+    /// NOTE this is the pre-SIMD83 logic and can be removed once SIMD83 is active.
+    pub fn try_lock_accounts(&mut self, keys: &[(&Pubkey, bool)]) -> TransactionResult<()> {
+        self.can_lock_accounts(keys)?;
+        self.lock_accounts(keys);
 
         Ok(())
     }
 
+    /// Lock accounts for all transactions in a batch which don't conflict
+    /// with existing locks. Returns a vector of `TransactionResult` indicating
+    /// success or failure for each transaction in the batch.
+    /// NOTE this is the SIMD83 logic; after the feature is active, it becomes
+    /// the only logic, and this note can be removed with the feature gate.
+    pub fn try_lock_transaction_batch<'a>(
+        &mut self,
+        validated_batch_keys: impl Iterator<Item = TransactionResult<Vec<(&'a Pubkey, bool)>>>,
+    ) -> Vec<TransactionResult<()>> {
+        let available_batch_keys: Vec<_> = validated_batch_keys
+            .map(|validated_keys| match validated_keys {
+                Ok(ref keys) => match self.can_lock_accounts(keys) {
+                    Ok(_) => validated_keys,
+                    Err(e) => Err(e),
+                },
+                Err(e) => Err(e),
+            })
+            .collect();
+
+        available_batch_keys
+            .into_iter()
+            .map(|available_keys| available_keys.map(|keys| self.lock_accounts(&keys)))
+            .collect()
+    }
+
     /// 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
@@ -59,14 +67,38 @@ impl AccountLocks {
         }
     }
 
+    fn can_lock_accounts(&self, keys: &[(&Pubkey, bool)]) -> TransactionResult<()> {
+        for (key, writable) in keys {
+            if *writable {
+                if !self.can_write_lock(key) {
+                    return Err(TransactionError::AccountInUse);
+                }
+            } else if !self.can_read_lock(key) {
+                return Err(TransactionError::AccountInUse);
+            }
+        }
+
+        Ok(())
+    }
+
+    fn lock_accounts(&mut self, keys: &[(&Pubkey, bool)]) {
+        for (key, writable) in keys {
+            if *writable {
+                self.lock_write(key);
+            } else {
+                self.lock_readonly(key);
+            }
+        }
+    }
+
     #[cfg_attr(feature = "dev-context-only-utils", qualifiers(pub))]
     fn is_locked_readonly(&self, key: &Pubkey) -> bool {
         self.readonly_locks.get(key).is_some_and(|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)
+        self.write_locks.get(key).is_some_and(|count| *count > 0)
     }
 
     fn can_read_lock(&self, key: &Pubkey) -> bool {
@@ -84,7 +116,7 @@ impl AccountLocks {
     }
 
     fn lock_write(&mut self, key: &Pubkey) {
-        self.write_locks.insert(*key);
+        *self.write_locks.entry(*key).or_default() += 1;
     }
 
     fn unlock_readonly(&mut self, key: &Pubkey) {
@@ -103,19 +135,26 @@ impl AccountLocks {
     }
 
     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"
-        );
+        if let hash_map::Entry::Occupied(mut occupied_entry) = self.write_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 write-lock for a key that wasn't write-locked"
+            );
+        }
     }
 }
 
 /// Validate account locks before locking.
 pub fn validate_account_locks(
     account_keys: AccountKeys,
     tx_account_lock_limit: usize,
-) -> Result<(), TransactionError> {
+) -> TransactionResult<()> {
     if account_keys.len() > tx_account_lock_limit {
         Err(TransactionError::TooManyAccountLocks)
     } else if has_duplicates(account_keys) {
@@ -165,23 +204,23 @@ mod tests {
         let key2 = Pubkey::new_unique();
 
         // Add write and read-lock.
-        let result = account_locks.try_lock_accounts([(&key1, true), (&key2, false)].into_iter());
+        let result = account_locks.try_lock_accounts(&[(&key1, true), (&key2, false)]);
         assert!(result.is_ok());
 
         // Try to add duplicate write-lock.
-        let result = account_locks.try_lock_accounts([(&key1, true)].into_iter());
+        let result = account_locks.try_lock_accounts(&[(&key1, true)]);
         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());
+        let result = account_locks.try_lock_accounts(&[(&key2, true)]);
         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());
+        let result = account_locks.try_lock_accounts(&[(&key1, false)]);
         assert_eq!(result, Err(TransactionError::AccountInUse));
 
         // Add read lock on read-locked account.
-        let result = account_locks.try_lock_accounts([(&key2, false)].into_iter());
+        let result = account_locks.try_lock_accounts(&[(&key2, false)]);
         assert!(result.is_ok());
 
         // Unlock write and read locks.