impl simd83 locking

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,60 @@ 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<()>> {
+        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 +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);
             }
         }
     }
@@ -69,7 +107,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 {
@@ -87,7 +125,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) {
@@ -106,19 +144,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) {

accounts-db/src/accounts.rs

@@ -35,6 +35,8 @@ use {
     },
 };
 
+const HANA_FEATURE_PLACEHOLDER: bool = false;
+
 pub type PubkeyAccountSlot = (Pubkey, AccountSharedData, Slot);
 
 struct TransactionAccountLocksIterator<'a, T: SVMMessage> {
@@ -598,15 +600,26 @@ impl Accounts {
         tx_account_locks_results: Vec<Result<TransactionAccountLocksIterator<impl SVMMessage>>>,
     ) -> Vec<Result<()>> {
         let account_locks = &mut self.account_locks.lock().unwrap();
-        tx_account_locks_results
-            .into_iter()
-            .map(|tx_account_locks_result| match tx_account_locks_result {
-                Ok(tx_account_locks) => {
-                    account_locks.try_lock_accounts(tx_account_locks.accounts_with_is_writable())
-                }
-                Err(err) => Err(err),
-            })
-            .collect()
+        if HANA_FEATURE_PLACEHOLDER {
+            let validated_batch_keys = tx_account_locks_results
+                .into_iter()
+                .map(|tx_account_locks_result| {
+                    tx_account_locks_result
+                        .map(|tx_account_locks| tx_account_locks.accounts_with_is_writable())
+                })
+                .collect();
+
+            account_locks.try_lock_transaction_batch(validated_batch_keys)
+        } else {
+            tx_account_locks_results
+                .into_iter()
+                .map(|tx_account_locks_result| match tx_account_locks_result {
+                    Ok(tx_account_locks) => account_locks
+                        .try_lock_accounts(tx_account_locks.accounts_with_is_writable()),
+                    Err(err) => Err(err),
+                })
+                .collect()
+        }
     }
 
     /// Once accounts are unlocked, new transactions that modify that state can enter the pipeline