accounts-db: relax intrabatch account locks

accounts-db/benches/bench_lock_accounts.rs

@@ -83,8 +83,9 @@ fn bench_entry_lock_accounts(c: &mut Criterion) {
         group.bench_function(name.as_str(), move |b| {
             b.iter(|| {
                 for batch in &transaction_batches {
+                    // HANA TODO bench both. but ideally in a way where we can compare them? idk
                     let results =
-                        accounts.lock_accounts(black_box(batch.iter()), MAX_TX_ACCOUNT_LOCKS);
+                        accounts.lock_accounts(black_box(batch.iter()), MAX_TX_ACCOUNT_LOCKS, true);
                     accounts.unlock_accounts(batch.iter().zip(&results));
                 }
             })

accounts-db/src/account_locks.rs

@@ -5,14 +5,14 @@ use {
     solana_sdk::{
         message::AccountKeys,
         pubkey::Pubkey,
-        transaction::{TransactionError, MAX_TX_ACCOUNT_LOCKS},
+        transaction::{Result, TransactionError, MAX_TX_ACCOUNT_LOCKS},
     },
     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>,
 }
 
@@ -21,10 +21,65 @@ 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> {
+    ) -> Result<()> {
+        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: Vec<Result<impl Iterator<Item = (&'a Pubkey, bool)> + Clone>>,
+    ) -> Vec<Result<()>> {
+        // HANA TODO the vec allocation here is unfortunate but hard to avoid
+        // we cannot do this in one closure because of borrow rules
+        // play around with alternate strategies, according to benches this may be up to
+        // 50% slower for small batches and few locks, but for large batches and many locks
+        // it is around 20% faster. so it's not horrible but ideally we improve all-case perf
+        let available_batch_keys: Vec<_> = validated_batch_keys
+            .into_iter()
+            .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,
+    ) -> Result<()> {
         for (key, writable) in keys.clone() {
             if writable {
                 if !self.can_write_lock(key) {
@@ -35,27 +90,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);
             }
         }
     }
@@ -67,7 +110,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).map_or(false, |count| *count > 0)
     }
 
     fn can_read_lock(&self, key: &Pubkey) -> bool {
@@ -85,7 +128,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) {
@@ -104,19 +147,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> {
+) -> Result<()> {
     if account_keys.len() > tx_account_lock_limit {
         Err(TransactionError::TooManyAccountLocks)
     } else if has_duplicates(account_keys) {