SIMD#83 - Allow conflicting entries in the account locking function

accounts-db/Cargo.toml

@@ -70,6 +70,8 @@ serde_bytes = { workspace = true }
 solana-accounts-db = { path = ".", features = ["dev-context-only-utils"] }
 solana-compute-budget = { workspace = true }
 solana-logger = { workspace = true }
+solana-ledger = { workspace = true }
+solana-runtime = { workspace = true, features = ["dev-context-only-utils"] }
 solana-sdk = { workspace = true, features = ["dev-context-only-utils"] }
 static_assertions = { workspace = true }
 strum = { workspace = true, features = ["derive"] }
@@ -114,3 +116,4 @@ harness = false
 
 [lints]
 workspace = true
+

accounts-db/benches/bench_lock_accounts.rs

@@ -20,10 +20,10 @@ const LOCK_COUNTS: [usize; 2] = [2, 64];
 // total transactions per run
 const TOTAL_TRANSACTIONS: usize = 1024;
 
-fn create_test_transactions(lock_count: usize, read_conflicts: bool) -> Vec<SanitizedTransaction> {
-    // keys available to be shared between transactions, depending on mode
-    // currently, we test batches with no conflicts and batches with reader/reader conflicts
-    // in the future with SIMD83, we will also test reader/writer and writer/writer conflicts
+fn create_test_transactions(
+    lock_count: usize,
+    self_conflicting_batch: bool,
+) -> Vec<SanitizedTransaction> {
     let shared_pubkeys: Vec<_> = (0..lock_count).map(|_| Pubkey::new_unique()).collect();
     let mut transactions = vec![];
 
@@ -37,11 +37,15 @@ fn create_test_transactions(lock_count: usize, read_conflicts: bool) -> Vec<Sani
             let account_meta = if i == 0 {
                 AccountMeta::new(Pubkey::new_unique(), true)
             } else if i % 2 == 0 {
-                AccountMeta::new(Pubkey::new_unique(), false)
-            } else if read_conflicts {
-                AccountMeta::new_readonly(shared_pubkeys[i], false)
+                if !self_conflicting_batch {
+                    AccountMeta::new(Pubkey::new_unique(), false)
+                } else {
+                    AccountMeta::new_readonly(shared_pubkeys[i], false)
+                }
+            } else if self_conflicting_batch {
+                AccountMeta::new(shared_pubkeys[i], false)
             } else {
-                AccountMeta::new_readonly(Pubkey::new_unique(), false)
+                AccountMeta::new_readonly(shared_pubkeys[i], false)
             };
 
             account_metas.push(account_meta);
@@ -59,32 +63,32 @@ 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
+    for (batch_size, lock_count, self_conflicting_batch) in
         iproduct!(BATCH_SIZES, LOCK_COUNTS, [false, true])
     {
-        let name = format!(
-            "batch_size_{batch_size}_locks_count_{lock_count}{}",
-            if read_conflicts {
-                "_read_conflicts"
-            } else {
-                ""
-            }
-        );
+        let mut group = if self_conflicting_batch {
+            c.benchmark_group("bench_lock_accounts_with_conflicting_batches")
+        } else {
+            c.benchmark_group("bench_lock_accounts")
+        };
+
+        let name = format!("batch_size_{batch_size}_locks_count_{lock_count}");
 
         let accounts_db = AccountsDb::new_single_for_tests();
         let accounts = Accounts::new(Arc::new(accounts_db));
 
-        let transactions = create_test_transactions(lock_count, read_conflicts);
+        let transactions = create_test_transactions(lock_count, self_conflicting_batch);
         group.throughput(Throughput::Elements(transactions.len() as u64));
         let transaction_batches: Vec<_> = transactions.chunks(batch_size).collect();
 
         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,
+                        self_conflicting_batch,
+                    );
                     accounts.unlock_accounts(batch.iter().zip(&results));
                 }
             })

accounts-db/src/account_locks.rs

@@ -1,18 +1,23 @@
 #[cfg(feature = "dev-context-only-utils")]
 use qualifier_attr::qualifiers;
 use {
+    crate::accounts::BatchAccountLocks,
     ahash::{AHashMap, AHashSet},
     solana_sdk::{
         message::AccountKeys,
         pubkey::Pubkey,
         transaction::{TransactionError, MAX_TX_ACCOUNT_LOCKS},
     },
-    std::{cell::RefCell, collections::hash_map},
+    std::{
+        cell::{RefCell, RefMut},
+        collections::hash_map,
+    },
 };
 
 #[derive(Debug, Default)]
 pub struct AccountLocks {
-    write_locks: AHashSet<Pubkey>,
+    // A key can have multiple outstanding write locks in the case of a self-conflicting batch.
+    write_locks: AHashMap<Pubkey, u64>,
     readonly_locks: AHashMap<Pubkey, u64>,
 }
 
@@ -21,6 +26,7 @@ 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.
+    /// This function will become obsolete after self conflicting batches are allowed.
     pub fn try_lock_accounts<'a>(
         &mut self,
         keys: impl Iterator<Item = (&'a Pubkey, bool)> + Clone,
@@ -46,6 +52,44 @@ impl AccountLocks {
         Ok(())
     }
 
+    pub fn try_lock_accounts_with_conflicting_batches<'a>(
+        &mut self,
+        keys: impl Iterator<Item = (&'a Pubkey, bool)> + Clone,
+        batch_account_locks: &mut RefMut<BatchAccountLocks>,
+    ) -> (Result<(), TransactionError>, bool) {
+        let mut self_conflicting_batch = false;
+
+        for (key, writable) in keys.clone() {
+            if writable {
+                if !self.can_write_lock(key) {
+                    if !(batch_account_locks.writables.contains(key)
+                        || batch_account_locks.readables.contains(key))
+                    {
+                        return (Err(TransactionError::AccountInUse), false);
+                    }
+                    self_conflicting_batch = true;
+                }
+            } else if !self.can_read_lock(key) {
+                if !batch_account_locks.writables.contains(key) {
+                    return (Err(TransactionError::AccountInUse), false);
+                }
+                self_conflicting_batch = true;
+            }
+        }
+
+        for (key, writable) in keys {
+            if writable {
+                batch_account_locks.insert_write_lock(key);
+                self.lock_write(key);
+            } else {
+                batch_account_locks.insert_read_lock(key);
+                self.lock_readonly(key);
+            }
+        }
+
+        (Ok(()), self_conflicting_batch)
+    }
+
     /// 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
@@ -69,7 +113,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 +131,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,11 +150,18 @@ 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"
+            );
+        }
     }
 }
 
@@ -158,7 +209,11 @@ fn has_duplicates(account_keys: AccountKeys) -> bool {
 
 #[cfg(test)]
 mod tests {
-    use {super::*, solana_sdk::message::v0::LoadedAddresses};
+    use {super::*, solana_sdk::message::v0::LoadedAddresses, std::cell::RefCell};
+
+    thread_local! {
+        static BATCH_ACCOUNT_LOCKS: RefCell<BatchAccountLocks> = RefCell::new(BatchAccountLocks::with_capacity(64*128));
+    }
 
     #[test]
     fn test_account_locks() {
@@ -288,4 +343,64 @@ mod tests {
         let account_keys = AccountKeys::new(&keys, None);
         assert!(has_duplicates(account_keys));
     }
+
+    #[test]
+    fn test_account_locks_with_conflicting_batches_and_unlock() {
+        let mut account_locks = AccountLocks::default();
+
+        let key1 = Pubkey::new_unique();
+        let key2 = Pubkey::new_unique();
+        BATCH_ACCOUNT_LOCKS.with(|batch_account_locks| {
+            let mut batch_account_locks = batch_account_locks.borrow_mut();
+            // Add write and read-lock.
+            let (result, _) = account_locks.try_lock_accounts_with_conflicting_batches(
+                [(&key1, true), (&key2, false)].into_iter(),
+                &mut batch_account_locks,
+            );
+            assert!(result.is_ok());
+
+            // Try to add duplicate write-lock, allowed in conflicting batch.
+            let (result, _) = account_locks.try_lock_accounts_with_conflicting_batches(
+                [(&key1, true)].into_iter(),
+                &mut batch_account_locks,
+            );
+            assert!(result.is_ok());
+
+            // Try to add write lock on read-locked account, allowed in conflicting batch.
+            let (result, _) = account_locks.try_lock_accounts_with_conflicting_batches(
+                [(&key2, true)].into_iter(),
+                &mut batch_account_locks,
+            );
+            assert!(result.is_ok());
+
+            // Try to add read lock on write-locked account, allowed in conflicting batch.
+            let (result, _) = account_locks.try_lock_accounts_with_conflicting_batches(
+                [(&key1, false)].into_iter(),
+                &mut batch_account_locks,
+            );
+            assert!(result.is_ok());
+
+            // Add read lock on read-locked account.
+            let (result, _) = account_locks.try_lock_accounts_with_conflicting_batches(
+                [(&key2, false)].into_iter(),
+                &mut batch_account_locks,
+            );
+            assert!(result.is_ok());
+
+            // Unlock write and read locks.
+            account_locks.unlock_accounts([(&key1, true), (&key2, false)].into_iter());
+
+            // More remaining write-locks and Read-lock.
+            assert!(account_locks.is_locked_write(&key1));
+            assert!(account_locks.is_locked_readonly(&key2));
+
+            // Unlock remaining write locks
+            account_locks.unlock_accounts([(&key1, true)].into_iter());
+            assert!(!account_locks.is_locked_write(&key1));
+
+            // Unlock read lock.
+            account_locks.unlock_accounts([(&key2, false)].into_iter());
+            assert!(!account_locks.is_locked_readonly(&key2));
+        })
+    }
 }