accounts-db: relax intrabatch account locks

accounts-db/Cargo.toml

@@ -97,7 +97,9 @@ solana-accounts-db = { path = ".", features = ["dev-context-only-utils"] }
 solana-compute-budget = { workspace = true }
 solana-instruction = { workspace = true }
 solana-logger = { workspace = true }
+solana-reserved-account-keys = { 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,10 +20,61 @@ 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.
+    /// NOTE this is the pre-SIMD83 logic and can be removed once SIMD83 is active.
     pub fn try_lock_accounts<'a>(
         &mut self,
         keys: impl Iterator<Item = (&'a Pubkey, bool)> + Clone,
-    ) -> Result<(), TransactionError> {
+    ) -> TransactionResult<()> {
+        self.can_lock_accounts(keys.clone())?;
+        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<impl Iterator<Item = (&'a Pubkey, bool)> + Clone>,
+        >,
+    ) -> Vec<TransactionResult<()>> {
+        let available_batch_keys: Vec<_> = validated_batch_keys
+            .map(|validated_keys| {
+                validated_keys
+                    .clone()
+                    .and_then(|keys| self.can_lock_accounts(keys))
+                    .and(validated_keys)
+            })
+            .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
+    /// an account that wasn't locked in the way requested.
+    pub fn unlock_accounts<'a>(&mut self, keys: impl Iterator<Item = (&'a Pubkey, bool)>) {
+        for (k, writable) in keys {
+            if writable {
+                self.unlock_write(k);
+            } else {
+                self.unlock_readonly(k);
+            }
+        }
+    }
+
+    fn can_lock_accounts<'a>(
+        &self,
+        keys: impl Iterator<Item = (&'a Pubkey, bool)> + Clone,
+    ) -> TransactionResult<()> {
         for (key, writable) in keys.clone() {
             if writable {
                 if !self.can_write_lock(key) {
@@ -34,27 +85,15 @@ impl AccountLocks {
             }
         }
 
-        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<Item = (&'a Pubkey, bool)>) {
-        for (k, writable) in keys {
+    fn lock_accounts<'a>(&mut self, keys: impl Iterator<Item = (&'a Pubkey, bool)> + Clone) {
+        for (key, writable) in keys {
             if writable {
-                self.unlock_write(k);
+                self.lock_write(key);
             } else {
-                self.unlock_readonly(k);
+                self.lock_readonly(key);
             }
         }
     }
@@ -66,7 +105,7 @@ impl AccountLocks {
 
     #[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 +123,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 +142,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) {