BAM Validator: Alternative Implementation (Using Banking Stage for Unified Execution)

.github/workflows/release.yml

@@ -90,4 +90,4 @@ jobs:
 #              title: 'Bump version to ${{ steps.parse_info.outputs.next_version }}',
 #              head: 'version-bump-${{ steps.parse_info.outputs.next_version }}',
 #              base: '${{ steps.parse_info.outputs.target_branch }}'
-#            })
+#            })

.gitmodules

@@ -7,3 +7,6 @@
 [submodule "jito-protos/protos"]
 	path = jito-protos/protos
 	url = https://github.com/jito-labs/mev-protos.git
+[submodule "jito-protos/bam-protos"]
+	path = jito-protos/bam-protos
+	url = git@github.com:jito-labs/bam-protos.git

accounts-db/benches/bench_lock_accounts.rs

@@ -85,7 +85,7 @@ fn bench_entry_lock_accounts(c: &mut Criterion) {
                 for batch in &transaction_batches {
                     let results =
                         accounts.lock_accounts(black_box(batch.iter()), MAX_TX_ACCOUNT_LOCKS);
-                    accounts.unlock_accounts(batch.iter().zip(&results));
+                    accounts.unlock_accounts(batch.iter().zip(&results), false);
                 }
             })
         });

accounts-db/src/accounts.rs

@@ -25,8 +25,9 @@ use {
     },
     solana_transaction_context::TransactionAccount,
     std::{
+        cell::RefCell,
         cmp::Reverse,
-        collections::{BinaryHeap, HashSet},
+        collections::{BinaryHeap, HashMap, HashSet},
         ops::RangeBounds,
         sync::{
             atomic::{AtomicUsize, Ordering},
@@ -35,6 +36,11 @@ use {
     },
 };
 
+thread_local! {
+    /// Reusable deduped locks for batched account locking.
+    static REUSABLE_DEDUPED_LOCKS: RefCell<HashMap<Pubkey, bool>> = RefCell::new(HashMap::new());
+}
+
 pub type PubkeyAccountSlot = (Pubkey, AccountSharedData, Slot);
 
 pub struct TransactionAccountLocksIterator<'a, T: SVMMessage> {
@@ -570,7 +576,7 @@ impl Accounts {
                     .map(|_| TransactionAccountLocksIterator::new(tx))
             })
             .collect();
-        self.lock_accounts_inner(tx_account_locks_results, None, None)
+        self.lock_accounts_inner(tx_account_locks_results, None, None, false)
     }
 
     #[must_use]
@@ -581,6 +587,7 @@ impl Accounts {
         tx_account_lock_limit: usize,
         additional_read_locks: Option<&HashSet<Pubkey>>,
         additional_write_locks: Option<&HashSet<Pubkey>>,
+        batched: bool,
     ) -> Vec<Result<()>> {
         // Validate the account locks, then get iterator if successful validation.
         let tx_account_locks_results: Vec<Result<_>> = txs
@@ -595,6 +602,7 @@ impl Accounts {
             tx_account_locks_results,
             additional_read_locks,
             additional_write_locks,
+            batched,
         )
     }
 
@@ -604,7 +612,15 @@ impl Accounts {
         tx_account_locks_results: Vec<Result<TransactionAccountLocksIterator<impl SVMMessage>>>,
         additional_read_locks: Option<&HashSet<Pubkey>>,
         additional_write_locks: Option<&HashSet<Pubkey>>,
+        batched: bool,
     ) -> Vec<Result<()>> {
+        if batched {
+            return self.lock_accounts_batched(
+                tx_account_locks_results,
+                additional_read_locks,
+                additional_write_locks,
+            );
+        }
         let account_locks = &mut self.account_locks.lock().unwrap();
         tx_account_locks_results
             .into_iter()
@@ -619,15 +635,97 @@ impl Accounts {
             .collect()
     }
 
+    /// Lock accounts in batches, deduplicating locks across transactions.
+    /// Returns a vec with all transaction locks associated with the
+    /// first transaction result, and `Ok(())` for all other transactions (If successful).
+    fn lock_accounts_batched(
+        &self,
+        tx_account_locks_results: Vec<Result<TransactionAccountLocksIterator<impl SVMMessage>>>,
+        additional_read_locks: Option<&HashSet<Pubkey>>,
+        additional_write_locks: Option<&HashSet<Pubkey>>,
+    ) -> Vec<Result<()>> {
+        if let Some(err) = tx_account_locks_results.iter().find_map(|res| {
+            if let Err(err) = res {
+                Some(err.clone())
+            } else {
+                None
+            }
+        }) {
+            return tx_account_locks_results
+                .into_iter()
+                .map(|_| Err(err.clone()))
+                .collect();
+        }
+
+        let len = tx_account_locks_results.len();
+        let result = REUSABLE_DEDUPED_LOCKS.with_borrow_mut(|deduped_locks| {
+            deduped_locks.clear();
+            Self::get_deduped_batch_locks(
+                deduped_locks,
+                tx_account_locks_results.into_iter().map(|res| res.unwrap()),
+            );
+            let account_locks = &mut self.account_locks.lock().unwrap();
+            account_locks.try_lock_accounts(
+                deduped_locks
+                    .iter()
+                    .map(|(pubkey, is_writable)| (pubkey, *is_writable)),
+                additional_read_locks,
+                additional_write_locks,
+            )
+        });
+
+        (0..len).map(|_| result.clone()).collect()
+    }
+
+    /// Deduplicate the locks across all transactions in a batch; promoting read-only locks to writable
+    /// if a writable lock exists for the same pubkey.
+    fn get_deduped_batch_locks<'a>(
+        deduped_locks: &mut HashMap<Pubkey, bool>,
+        tx_account_locks: impl Iterator<
+            Item = TransactionAccountLocksIterator<'a, impl SVMMessage + 'a>,
+        >,
+    ) {
+        for tx_account_locks in tx_account_locks {
+            for (pubkey, is_writable) in tx_account_locks.accounts_with_is_writable() {
+                if let Some(existing) = deduped_locks.get(pubkey) {
+                    if *existing && !is_writable {
+                        continue;
+                    }
+                }
+                deduped_locks.insert(*pubkey, is_writable);
+            }
+        }
+    }
+
     /// Once accounts are unlocked, new transactions that modify that state can enter the pipeline
     pub fn unlock_accounts<'a, Tx: SVMMessage + 'a>(
         &self,
         txs_and_results: impl Iterator<Item = (&'a Tx, &'a Result<()>)> + Clone,
+        batched: bool,
     ) {
         if !txs_and_results.clone().any(|(_, res)| res.is_ok()) {
             return;
         }
 
+        if batched {
+            REUSABLE_DEDUPED_LOCKS.with_borrow_mut(|deduped_locks| {
+                deduped_locks.clear();
+                Self::get_deduped_batch_locks(
+                    deduped_locks,
+                    txs_and_results
+                        .clone()
+                        .map(|(tx, _)| TransactionAccountLocksIterator::new(tx)),
+                );
+                let mut account_locks = self.account_locks.lock().unwrap();
+                account_locks.unlock_accounts(
+                    deduped_locks
+                        .iter()
+                        .map(|(pubkey, is_writable)| (pubkey, *is_writable)),
+                );
+            });
+            return;
+        }
+
         let mut account_locks = self.account_locks.lock().unwrap();
         debug!("bank unlock accounts");
         for (tx, res) in txs_and_results {
@@ -991,7 +1089,7 @@ mod tests {
             let txs = vec![new_sanitized_tx(&[&keypair], message, Hash::default())];
             let results = accounts.lock_accounts(txs.iter(), MAX_TX_ACCOUNT_LOCKS);
             assert_eq!(results, vec![Ok(())]);
-            accounts.unlock_accounts(txs.iter().zip(&results));
+            accounts.unlock_accounts(txs.iter().zip(&results), false);
         }
 
         // Disallow over MAX_TX_ACCOUNT_LOCKS
@@ -1089,8 +1187,8 @@ mod tests {
             .unwrap()
             .is_locked_readonly(&keypair1.pubkey()));
 
-        accounts.unlock_accounts(iter::once(&tx).zip(&results0));
-        accounts.unlock_accounts(txs.iter().zip(&results1));
+        accounts.unlock_accounts(iter::once(&tx).zip(&results0), false);
+        accounts.unlock_accounts(txs.iter().zip(&results1), false);
         let instructions = vec![CompiledInstruction::new(2, &(), vec![0, 1])];
         let message = Message::new_with_compiled_instructions(
             1,
@@ -1171,7 +1269,7 @@ mod tests {
                     counter_clone.clone().fetch_add(1, Ordering::Release);
                 }
             }
-            accounts_clone.unlock_accounts(txs.iter().zip(&results));
+            accounts_clone.unlock_accounts(txs.iter().zip(&results), false);
             if exit_clone.clone().load(Ordering::Relaxed) {
                 break;
             }
@@ -1187,7 +1285,7 @@ mod tests {
                 thread::sleep(time::Duration::from_millis(50));
                 assert_eq!(counter_value, counter_clone.clone().load(Ordering::Acquire));
             }
-            accounts_arc.unlock_accounts(txs.iter().zip(&results));
+            accounts_arc.unlock_accounts(txs.iter().zip(&results), false);
             thread::sleep(time::Duration::from_millis(50));
         }
         exit.store(true, Ordering::Relaxed);
@@ -1321,6 +1419,7 @@ mod tests {
             MAX_TX_ACCOUNT_LOCKS,
             None,
             None,
+            false,
         );
 
         assert_eq!(
@@ -1638,4 +1737,89 @@ mod tests {
             ));
         }
     }
+
+    #[test]
+    fn test_batched_locking() {
+        let keypair0 = Keypair::new();
+        let keypair1 = Keypair::new();
+        let keypair2 = Keypair::new();
+        let keypair3 = Keypair::new();
+
+        let account0 = AccountSharedData::new(1, 0, &Pubkey::default());
+        let account1 = AccountSharedData::new(2, 0, &Pubkey::default());
+        let account2 = AccountSharedData::new(3, 0, &Pubkey::default());
+        let account3 = AccountSharedData::new(4, 0, &Pubkey::default());
+
+        let accounts_db = AccountsDb::new_single_for_tests();
+        let accounts = Accounts::new(Arc::new(accounts_db));
+        accounts.store_for_tests(0, &keypair0.pubkey(), &account0);
+        accounts.store_for_tests(0, &keypair1.pubkey(), &account1);
+        accounts.store_for_tests(0, &keypair2.pubkey(), &account2);
+        accounts.store_for_tests(0, &keypair3.pubkey(), &account3);
+
+        let instructions = vec![CompiledInstruction::new(2, &(), vec![0, 1])];
+        let message = Message::new_with_compiled_instructions(
+            1,
+            0,
+            2,
+            vec![keypair1.pubkey(), keypair0.pubkey(), native_loader::id()],
+            Hash::default(),
+            instructions,
+        );
+        let tx0 = new_sanitized_tx(&[&keypair1], message, Hash::default());
+        let instructions = vec![CompiledInstruction::new(2, &(), vec![0, 1])];
+        let message = Message::new_with_compiled_instructions(
+            1,
+            0,
+            2,
+            vec![keypair2.pubkey(), keypair0.pubkey(), native_loader::id()],
+            Hash::default(),
+            instructions,
+        );
+        let tx1 = new_sanitized_tx(&[&keypair2], message, Hash::default());
+        let instructions = vec![CompiledInstruction::new(2, &(), vec![0, 1])];
+        let message = Message::new_with_compiled_instructions(
+            1,
+            0,
+            2,
+            vec![keypair3.pubkey(), keypair0.pubkey(), native_loader::id()],
+            Hash::default(),
+            instructions,
+        );
+        let tx2 = new_sanitized_tx(&[&keypair3], message, Hash::default());
+        let txs = vec![tx0, tx1, tx2];
+
+        let qos_results = vec![Ok(()), Ok(()), Ok(())];
+
+        let results = accounts.lock_accounts_with_results(
+            txs.iter(),
+            qos_results.into_iter(),
+            MAX_TX_ACCOUNT_LOCKS,
+            None,
+            None,
+            true,
+        );
+
+        assert_eq!(
+            results,
+            vec![
+                Ok(()), // Read-only account (keypair0) can be referenced multiple times
+                Ok(()), // Read-only account (keypair0) can be referenced multiple times
+                Ok(()), // Read-only account (keypair0) can be referenced multiple times
+            ],
+        );
+
+        // verify that keypair0 read-only locked
+        assert!(accounts
+            .account_locks
+            .lock()
+            .unwrap()
+            .is_locked_readonly(&keypair0.pubkey()));
+        // verify that keypair2 (for tx1) is write-locked (2 txns referencing it)
+        assert!(accounts
+            .account_locks
+            .lock()
+            .unwrap()
+            .is_locked_write(&keypair2.pubkey()));
+    }
 }

banking-bench/src/main.rs

@@ -507,6 +507,8 @@ fn main() {
         HashSet::default(),
         BundleAccountLocker::default(),
         |_| 0,
+        None,
+        None,
     );
 
     // This is so that the signal_receiver does not go out of scope after the closure.

core/Cargo.toml

@@ -80,6 +80,7 @@ serde = { workspace = true }
 serde_bytes = { workspace = true }
 serde_derive = { workspace = true }
 slab = { workspace = true }
+smallvec = { workspace = true }
 solana-accounts-db = { workspace = true }
 solana-bloom = { workspace = true }
 solana-builtins-default-costs = { workspace = true }
@@ -132,6 +133,7 @@ solana-version = { workspace = true }
 solana-vote = { workspace = true }
 solana-vote-program = { workspace = true }
 solana-wen-restart = { workspace = true }
+spl-memo = { workspace = true, features = ["no-entrypoint"] }
 strum = { workspace = true, features = ["derive"] }
 strum_macros = { workspace = true }
 sys-info = { workspace = true }

core/benches/banking_stage.rs

@@ -339,6 +339,8 @@ fn bench_banking(
         HashSet::default(),
         BundleAccountLocker::default(),
         |_| 0,
+        None,
+        None,
     );
 
     let chunk_len = verified.len() / CHUNKS;

core/benches/consumer.rs

@@ -177,6 +177,7 @@ fn bench_process_and_record_transactions(bencher: &mut Bencher, batch_size: usiz
                 transaction_iter.next().unwrap(),
                 0,
                 &|_| 0,
+                false,
             );
             assert!(summary
                 .execute_and_commit_transactions_output