Refactor accounts.rs with Justin's comments to improve lock accounts (#21406)

with results code path.
- fix a bug that could unlock accounts that weren't locked
- add test to the refactored function
- skip enumerating transaction accounts if qos results is an error
- add #[must_use] annotation
- avoid clone error in results
- add qos error code to unlock_accounts match statement
- remove unnecessary AbiExample

core/src/banking_stage.rs

@@ -940,7 +940,8 @@ impl BankingStage {
         // Once accounts are locked, other threads cannot encode transactions that will modify the
         // same account state
         let mut lock_time = Measure::start("lock_time");
-        let batch = bank.prepare_sanitized_batch_with_results(txs, transactions_qos_results.iter());
+        let batch =
+            bank.prepare_sanitized_batch_with_results(txs, transactions_qos_results.into_iter());
         lock_time.stop();
 
         // retryable_txs includes AccountInUse, WouldExceedMaxBlockCostLimit and

core/src/qos_service.rs

@@ -24,6 +24,10 @@ use {
 };
 
 pub struct QosService {
+    // cost_model instance is owned by validator, shared between replay_stage and
+    // banking_stage. replay_stage writes the latest on-chain program timings to
+    // it; banking_stage's qos_service reads that information to calculate
+    // transaction cost, hence RwLock wrapped.
     cost_model: Arc<RwLock<CostModel>>,
     metrics: Arc<QosServiceMetrics>,
     reporting_thread: Option<JoinHandle<()>>,

runtime/src/accounts.rs

@@ -969,22 +969,27 @@ impl Accounts {
             .collect()
     }
 
+    #[must_use]
     #[allow(clippy::needless_collect)]
     pub fn lock_accounts_with_results<'a>(
         &self,
         txs: impl Iterator<Item = &'a SanitizedTransaction>,
-        results: impl Iterator<Item = &'a Result<()>>,
+        results: impl Iterator<Item = Result<()>>,
         demote_program_write_locks: bool,
     ) -> Vec<Result<()>> {
-        let keys: Vec<_> = txs
-            .map(|tx| tx.get_account_locks(demote_program_write_locks))
+        let key_results: Vec<_> = txs
+            .zip(results)
+            .map(|(tx, result)| match result {
+                Ok(()) => Ok(tx.get_account_locks(demote_program_write_locks)),
+                Err(e) => Err(e),
+            })
             .collect();
         let account_locks = &mut self.account_locks.lock().unwrap();
-        keys.into_iter()
-            .zip(results)
-            .map(|(keys, result)| match result {
-                Ok(()) => self.lock_account(account_locks, keys.writable, keys.readonly),
-                Err(e) => Err(e.clone()),
+        key_results
+            .into_iter()
+            .map(|key_result| match key_result {
+                Ok(keys) => self.lock_account(account_locks, keys.writable, keys.readonly),
+                Err(e) => Err(e),
             })
             .collect()
     }
@@ -1003,6 +1008,8 @@ impl Accounts {
                 Err(TransactionError::AccountInUse) => None,
                 Err(TransactionError::SanitizeFailure) => None,
                 Err(TransactionError::AccountLoadedTwice) => None,
+                Err(TransactionError::WouldExceedMaxBlockCostLimit) => None,
+                Err(TransactionError::WouldExceedMaxAccountCostLimit) => None,
                 _ => Some(tx.get_account_locks(demote_program_write_locks)),
             })
             .collect();
@@ -2391,6 +2398,117 @@ mod tests {
             .contains(&keypair1.pubkey()));
     }
 
+    #[test]
+    fn test_accounts_locks_with_results() {
+        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 = Accounts::new_with_config_for_tests(
+            Vec::new(),
+            &ClusterType::Development,
+            AccountSecondaryIndexes::default(),
+            false,
+            AccountShrinkThreshold::default(),
+        );
+        accounts.store_slow_uncached(0, &keypair0.pubkey(), &account0);
+        accounts.store_slow_uncached(0, &keypair1.pubkey(), &account1);
+        accounts.store_slow_uncached(0, &keypair2.pubkey(), &account2);
+        accounts.store_slow_uncached(0, &keypair3.pubkey(), &account3);
+
+        let demote_program_write_locks = true;
+
+        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(()),
+            Err(TransactionError::WouldExceedMaxBlockCostLimit),
+            Ok(()),
+        ];
+
+        let results = accounts.lock_accounts_with_results(
+            txs.iter(),
+            qos_results.into_iter(),
+            demote_program_write_locks,
+        );
+
+        assert!(results[0].is_ok()); // Read-only account (keypair0) can be referenced multiple times
+        assert!(results[1].is_err()); // is not locked due to !qos_results[1].is_ok()
+        assert!(results[2].is_ok()); // Read-only account (keypair0) can be referenced multiple times
+
+        // verify that keypair0 read-only lock twice (for tx0 and tx2)
+        assert_eq!(
+            *accounts
+                .account_locks
+                .lock()
+                .unwrap()
+                .readonly_locks
+                .get(&keypair0.pubkey())
+                .unwrap(),
+            2
+        );
+        // verify that keypair2 (for tx1) is not write-locked
+        assert!(accounts
+            .account_locks
+            .lock()
+            .unwrap()
+            .write_locks
+            .get(&keypair2.pubkey())
+            .is_none());
+
+        accounts.unlock_accounts(txs.iter(), &results, demote_program_write_locks);
+
+        // check all locks to be removed
+        assert!(accounts
+            .account_locks
+            .lock()
+            .unwrap()
+            .readonly_locks
+            .is_empty());
+        assert!(accounts
+            .account_locks
+            .lock()
+            .unwrap()
+            .write_locks
+            .is_empty());
+    }
+
     #[test]
     fn test_collect_accounts_to_store() {
         let keypair0 = Keypair::new();

runtime/src/bank.rs

@@ -3356,7 +3356,7 @@ impl Bank {
     pub fn prepare_sanitized_batch_with_results<'a, 'b>(
         &'a self,
         transactions: &'b [SanitizedTransaction],
-        transaction_results: impl Iterator<Item = &'b Result<()>>,
+        transaction_results: impl Iterator<Item = Result<()>>,
     ) -> TransactionBatch<'a, 'b> {
         // this lock_results could be: Ok, AccountInUse, WouldExceedBlockMaxLimit or WouldExceedAccountMaxLimit
         let lock_results = self.rc.accounts.lock_accounts_with_results(

runtime/src/cost_model.rs

@@ -14,17 +14,15 @@ use std::collections::HashMap;
 const MAX_WRITABLE_ACCOUNTS: usize = 256;
 
 // costs are stored in number of 'compute unit's
-#[derive(AbiExample, Debug)]
+#[derive(Debug)]
 pub struct TransactionCost {
     pub writable_accounts: Vec<Pubkey>,
     pub signature_cost: u64,
     pub write_lock_cost: u64,
     pub data_bytes_cost: u64,
     pub execution_cost: u64,
-    // `cost_weight` is a multiplier to be applied to tx cost, that
-    // allows to increase/decrease tx cost linearly based on algo.
-    // for example, vote tx could have weight zero to bypass cost
-    // limit checking during block packing.
+    // `cost_weight` is a multiplier could be applied to transaction cost,
+    // if set to zero allows the transaction to bypass cost limit check.
     pub cost_weight: u32,
 }