Wire scheduler up to adjust actual loaded accounts size cost for committed transactions (#1547)

* wire up scheduler to adjust actual loaded accounts size cost for committed transactions

* update tests

core/src/banking_stage/committer.rs

@@ -25,7 +25,10 @@ use {
 
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub enum CommitTransactionDetails {
-    Committed { compute_units: u64 },
+    Committed {
+        compute_units: u64,
+        loaded_accounts_data_size: usize,
+    },
     NotCommitted,
 }
 
@@ -104,12 +107,21 @@ impl Committer {
         let commit_transaction_statuses = tx_results
             .execution_results
             .iter()
-            .map(|execution_result| match execution_result.details() {
-                Some(details) => CommitTransactionDetails::Committed {
-                    compute_units: details.executed_units,
+            .zip(tx_results.loaded_accounts_stats.iter())
+            .map(
+                |(execution_result, loaded_accounts_stats)| match execution_result.details() {
+                    // reports actual execution CUs, and actual loaded accounts size for
+                    // transaction committed to block. qos_service uses these information to adjust
+                    // reserved block space.
+                    Some(details) => CommitTransactionDetails::Committed {
+                        compute_units: details.executed_units,
+                        loaded_accounts_data_size: loaded_accounts_stats
+                            .as_ref()
+                            .map_or(0, |stats| stats.loaded_accounts_data_size),
+                    },
+                    None => CommitTransactionDetails::NotCommitted,
                 },
-                None => CommitTransactionDetails::NotCommitted,
-            })
+            )
             .collect();
 
         let ((), find_and_send_votes_us) = measure_us!({

core/src/banking_stage/consumer.rs

@@ -1559,9 +1559,18 @@ mod tests {
             assert_eq!(retryable_transaction_indexes, vec![1]);
 
             let expected_block_cost = if !apply_cost_tracker_during_replay_enabled {
-                let actual_programs_execution_cost =
+                let (actual_programs_execution_cost, actual_loaded_accounts_data_size_cost) =
                     match commit_transactions_result.first().unwrap() {
-                        CommitTransactionDetails::Committed { compute_units } => *compute_units,
+                        CommitTransactionDetails::Committed {
+                            compute_units,
+                            loaded_accounts_data_size,
+                        } => (
+                            *compute_units,
+                            CostModel::calculate_loaded_accounts_data_size_cost(
+                                *loaded_accounts_data_size,
+                                &bank.feature_set,
+                            ),
+                        ),
                         CommitTransactionDetails::NotCommitted => {
                             unreachable!()
                         }
@@ -1570,6 +1579,8 @@ mod tests {
                 let mut cost = CostModel::calculate_cost(&transactions[0], &bank.feature_set);
                 if let TransactionCost::Transaction(ref mut usage_cost) = cost {
                     usage_cost.programs_execution_cost = actual_programs_execution_cost;
+                    usage_cost.loaded_accounts_data_size_cost =
+                        actual_loaded_accounts_data_size_cost;
                 }
 
                 block_cost + cost.sum()

core/src/banking_stage/qos_service.rs

@@ -194,14 +194,23 @@ impl QosService {
         let mut cost_tracker = bank.write_cost_tracker().unwrap();
         transaction_cost_results
             .zip(transaction_committed_status)
-            .for_each(|(tx_cost, transaction_committed_details)| {
+            .for_each(|(estimated_tx_cost, transaction_committed_details)| {
                 // Only transactions that the qos service included have to be
                 // checked for update
-                if let Ok(tx_cost) = tx_cost {
-                    if let CommitTransactionDetails::Committed { compute_units } =
-                        transaction_committed_details
+                if let Ok(estimated_tx_cost) = estimated_tx_cost {
+                    if let CommitTransactionDetails::Committed {
+                        compute_units,
+                        loaded_accounts_data_size,
+                    } = transaction_committed_details
                     {
-                        cost_tracker.update_execution_cost(tx_cost, *compute_units)
+                        cost_tracker.update_execution_cost(
+                            estimated_tx_cost,
+                            *compute_units,
+                            CostModel::calculate_loaded_accounts_data_size_cost(
+                                *loaded_accounts_data_size,
+                                &bank.feature_set,
+                            ),
+                        )
                     }
                 }
             });
@@ -723,13 +732,25 @@ mod tests {
         // calculate their costs, apply to cost_tracker
         let transaction_count = 5;
         let keypair = Keypair::new();
-        let transfer_tx = SanitizedTransaction::from_transaction_for_tests(
-            system_transaction::transfer(&keypair, &keypair.pubkey(), 1, Hash::default()),
-        );
+        let loaded_accounts_data_size: usize = 1_000_000;
+        let transaction = solana_sdk::transaction::Transaction::new_unsigned(solana_sdk::message::Message::new(
+            &[
+                solana_sdk::compute_budget::ComputeBudgetInstruction::set_loaded_accounts_data_size_limit(loaded_accounts_data_size as u32),
+                solana_sdk::system_instruction::transfer(&keypair.pubkey(), &solana_sdk::pubkey::Pubkey::new_unique(), 1),
+            ],
+            Some(&keypair.pubkey()),
+        ));
+        let transfer_tx = SanitizedTransaction::from_transaction_for_tests(transaction.clone());
         let txs: Vec<SanitizedTransaction> = (0..transaction_count)
             .map(|_| transfer_tx.clone())
             .collect();
-        let execute_units_adjustment = 10u64;
+        let execute_units_adjustment: u64 = 10;
+        let loaded_accounts_data_size_adjustment: usize = 32000;
+        let loaded_accounts_data_size_cost_adjustment =
+            CostModel::calculate_loaded_accounts_data_size_cost(
+                loaded_accounts_data_size_adjustment,
+                &bank.feature_set,
+            );
 
         // assert all tx_costs should be applied to cost_tracker if all execution_results are all committed
         {
@@ -755,9 +776,13 @@ mod tests {
                 .map(|tx_cost| CommitTransactionDetails::Committed {
                     compute_units: tx_cost.as_ref().unwrap().programs_execution_cost()
                         + execute_units_adjustment,
+                    loaded_accounts_data_size: loaded_accounts_data_size
+                        + loaded_accounts_data_size_adjustment,
                 })
                 .collect();
-            let final_txs_cost = total_txs_cost + execute_units_adjustment * transaction_count;
+            let final_txs_cost = total_txs_cost
+                + (execute_units_adjustment + loaded_accounts_data_size_cost_adjustment)
+                    * transaction_count;
 
             // All transactions are committed, no costs should be removed
             QosService::remove_costs(qos_cost_results.iter(), Some(&committed_status), &bank);
@@ -845,13 +870,25 @@ mod tests {
         // calculate their costs, apply to cost_tracker
         let transaction_count = 5;
         let keypair = Keypair::new();
-        let transfer_tx = SanitizedTransaction::from_transaction_for_tests(
-            system_transaction::transfer(&keypair, &keypair.pubkey(), 1, Hash::default()),
-        );
+        let loaded_accounts_data_size: usize = 1_000_000;
+        let transaction = solana_sdk::transaction::Transaction::new_unsigned(solana_sdk::message::Message::new(
+            &[
+                solana_sdk::compute_budget::ComputeBudgetInstruction::set_loaded_accounts_data_size_limit(loaded_accounts_data_size as u32),
+                solana_sdk::system_instruction::transfer(&keypair.pubkey(), &solana_sdk::pubkey::Pubkey::new_unique(), 1),
+            ],
+            Some(&keypair.pubkey()),
+        ));
+        let transfer_tx = SanitizedTransaction::from_transaction_for_tests(transaction.clone());
         let txs: Vec<SanitizedTransaction> = (0..transaction_count)
             .map(|_| transfer_tx.clone())
             .collect();
-        let execute_units_adjustment = 10u64;
+        let execute_units_adjustment: u64 = 10;
+        let loaded_accounts_data_size_adjustment: usize = 32000;
+        let loaded_accounts_data_size_cost_adjustment =
+            CostModel::calculate_loaded_accounts_data_size_cost(
+                loaded_accounts_data_size_adjustment,
+                &bank.feature_set,
+            );
 
         // assert only committed tx_costs are applied cost_tracker
         {
@@ -882,6 +919,8 @@ mod tests {
                         CommitTransactionDetails::Committed {
                             compute_units: tx_cost.as_ref().unwrap().programs_execution_cost()
                                 + execute_units_adjustment,
+                            loaded_accounts_data_size: loaded_accounts_data_size
+                                + loaded_accounts_data_size_adjustment,
                         }
                     }
                 })
@@ -896,8 +935,9 @@ mod tests {
             qos_cost_results.iter().enumerate().for_each(|(n, cost)| {
                 if n % 2 != 0 {
                     expected_final_txs_count += 1;
-                    expected_final_block_cost +=
-                        cost.as_ref().unwrap().sum() + execute_units_adjustment;
+                    expected_final_block_cost += cost.as_ref().unwrap().sum()
+                        + execute_units_adjustment
+                        + loaded_accounts_data_size_cost_adjustment;
                 }
             });
             assert_eq!(

cost-model/benches/cost_tracker.rs

@@ -58,7 +58,7 @@ fn bench_cost_tracker_non_contentious_transaction(bencher: &mut Bencher) {
             if cost_tracker.try_add(tx_cost).is_err() {
                 break;
             } // stop when hit limits
-            cost_tracker.update_execution_cost(tx_cost, 0); // update execution cost down to zero
+            cost_tracker.update_execution_cost(tx_cost, 0, 0); // update execution cost down to zero
         }
     });
 }
@@ -75,7 +75,7 @@ fn bench_cost_tracker_contentious_transaction(bencher: &mut Bencher) {
             if cost_tracker.try_add(tx_cost).is_err() {
                 break;
             } // stop when hit limits
-            cost_tracker.update_execution_cost(tx_cost, 0); // update execution cost down to zero
+            cost_tracker.update_execution_cost(tx_cost, 0, 0); // update execution cost down to zero
         }
     });
 }

cost-model/src/cost_model.rs

@@ -153,14 +153,10 @@ impl CostModel {
                     programs_execution_costs = u64::from(compute_budget_limits.compute_unit_limit);
                 }
 
-                if feature_set
-                    .is_active(&include_loaded_accounts_data_size_in_fee_calculation::id())
-                {
-                    loaded_accounts_data_size_cost = FeeStructure::calculate_memory_usage_cost(
-                        usize::try_from(compute_budget_limits.loaded_accounts_bytes).unwrap(),
-                        DEFAULT_HEAP_COST,
-                    )
-                }
+                loaded_accounts_data_size_cost = Self::calculate_loaded_accounts_data_size_cost(
+                    usize::try_from(compute_budget_limits.loaded_accounts_bytes).unwrap(),
+                    feature_set,
+                );
             }
             Err(_) => {
                 programs_execution_costs = 0;
@@ -172,6 +168,17 @@ impl CostModel {
         tx_cost.data_bytes_cost = data_bytes_len_total / INSTRUCTION_DATA_BYTES_COST;
     }
 
+    pub fn calculate_loaded_accounts_data_size_cost(
+        loaded_accounts_data_size: usize,
+        feature_set: &FeatureSet,
+    ) -> u64 {
+        if feature_set.is_active(&include_loaded_accounts_data_size_in_fee_calculation::id()) {
+            FeeStructure::calculate_memory_usage_cost(loaded_accounts_data_size, DEFAULT_HEAP_COST)
+        } else {
+            0
+        }
+    }
+
     fn calculate_account_data_size_on_deserialized_system_instruction(
         instruction: SystemInstruction,
     ) -> u64 {

cost-model/src/cost_tracker.rs

@@ -160,20 +160,25 @@ impl CostTracker {
         &mut self,
         estimated_tx_cost: &TransactionCost,
         actual_execution_units: u64,
+        actual_loaded_accounts_data_size_cost: u64,
     ) {
-        let estimated_execution_units = estimated_tx_cost.programs_execution_cost();
-        match actual_execution_units.cmp(&estimated_execution_units) {
+        let actual_load_and_execution_units =
+            actual_execution_units.saturating_add(actual_loaded_accounts_data_size_cost);
+        let estimated_load_and_execution_units = estimated_tx_cost
+            .programs_execution_cost()
+            .saturating_add(estimated_tx_cost.loaded_accounts_data_size_cost());
+        match actual_load_and_execution_units.cmp(&estimated_load_and_execution_units) {
             Ordering::Equal => (),
             Ordering::Greater => {
                 self.add_transaction_execution_cost(
                     estimated_tx_cost,
-                    actual_execution_units - estimated_execution_units,
+                    actual_load_and_execution_units - estimated_load_and_execution_units,
                 );
             }
             Ordering::Less => {
                 self.sub_transaction_execution_cost(
                     estimated_tx_cost,
-                    estimated_execution_units - actual_execution_units,
+                    estimated_load_and_execution_units - actual_load_and_execution_units,
                 );
             }
         }
@@ -857,51 +862,69 @@ mod tests {
 
     #[test]
     fn test_update_execution_cost() {
-        let acct1 = Pubkey::new_unique();
-        let acct2 = Pubkey::new_unique();
-        let acct3 = Pubkey::new_unique();
-        let cost = 100;
+        let estimated_programs_execution_cost = 100;
+        let estimated_loaded_accounts_data_size_cost = 200;
+        let number_writeble_accounts = 3;
+        let writable_accounts = std::iter::repeat_with(Pubkey::new_unique)
+            .take(number_writeble_accounts)
+            .collect();
 
         let tx_cost = TransactionCost::Transaction(UsageCostDetails {
-            writable_accounts: vec![acct1, acct2, acct3],
-            programs_execution_cost: cost,
+            writable_accounts,
+            programs_execution_cost: estimated_programs_execution_cost,
+            loaded_accounts_data_size_cost: estimated_loaded_accounts_data_size_cost,
             ..UsageCostDetails::default()
         });
+        // confirm tx_cost is only made up by programs_execution_cost and
+        // loaded_accounts_data_size_cost
+        let estimated_tx_cost = tx_cost.sum();
+        assert_eq!(
+            estimated_tx_cost,
+            estimated_programs_execution_cost + estimated_loaded_accounts_data_size_cost
+        );
 
-        let mut cost_tracker = CostTracker::default();
-
-        // Assert OK to add tx_cost
-        assert!(cost_tracker.try_add(&tx_cost).is_ok());
-        let (_costliest_account, costliest_account_cost) = cost_tracker.find_costliest_account();
-        assert_eq!(cost, cost_tracker.block_cost);
-        assert_eq!(cost, costliest_account_cost);
-        assert_eq!(1, cost_tracker.transaction_count);
-
-        // assert no-change if actual units is same as estimated units
-        let mut expected_cost = cost;
-        cost_tracker.update_execution_cost(&tx_cost, cost);
-        let (_costliest_account, costliest_account_cost) = cost_tracker.find_costliest_account();
-        assert_eq!(expected_cost, cost_tracker.block_cost);
-        assert_eq!(expected_cost, costliest_account_cost);
-        assert_eq!(1, cost_tracker.transaction_count);
-
-        // assert cost are adjusted down
-        let reduced_units = 3;
-        expected_cost -= reduced_units;
-        cost_tracker.update_execution_cost(&tx_cost, cost - reduced_units);
-        let (_costliest_account, costliest_account_cost) = cost_tracker.find_costliest_account();
-        assert_eq!(expected_cost, cost_tracker.block_cost);
-        assert_eq!(expected_cost, costliest_account_cost);
-        assert_eq!(1, cost_tracker.transaction_count);
+        let test_update_cost_tracker =
+            |execution_cost_adjust: i64, loaded_acounts_data_size_cost_adjust: i64| {
+                let mut cost_tracker = CostTracker::default();
+                assert!(cost_tracker.try_add(&tx_cost).is_ok());
+
+                let actual_programs_execution_cost =
+                    (estimated_programs_execution_cost as i64 + execution_cost_adjust) as u64;
+                let actual_loaded_accounts_data_size_cost =
+                    (estimated_loaded_accounts_data_size_cost as i64
+                        + loaded_acounts_data_size_cost_adjust) as u64;
+                let expected_cost = (estimated_tx_cost as i64
+                    + execution_cost_adjust
+                    + loaded_acounts_data_size_cost_adjust)
+                    as u64;
+
+                cost_tracker.update_execution_cost(
+                    &tx_cost,
+                    actual_programs_execution_cost,
+                    actual_loaded_accounts_data_size_cost,
+                );
 
-        // assert cost are adjusted up
-        let increased_units = 1;
-        expected_cost += increased_units;
-        cost_tracker.update_execution_cost(&tx_cost, cost + increased_units);
-        let (_costliest_account, costliest_account_cost) = cost_tracker.find_costliest_account();
-        assert_eq!(expected_cost, cost_tracker.block_cost);
-        assert_eq!(expected_cost, costliest_account_cost);
-        assert_eq!(1, cost_tracker.transaction_count);
+                assert_eq!(expected_cost, cost_tracker.block_cost);
+                assert_eq!(0, cost_tracker.vote_cost);
+                assert_eq!(
+                    number_writeble_accounts,
+                    cost_tracker.cost_by_writable_accounts.len()
+                );
+                for writable_account_cost in cost_tracker.cost_by_writable_accounts.values() {
+                    assert_eq!(expected_cost, *writable_account_cost);
+                }
+                assert_eq!(1, cost_tracker.transaction_count);
+            };
+
+        test_update_cost_tracker(0, 0);
+        test_update_cost_tracker(0, 9);
+        test_update_cost_tracker(0, -9);
+        test_update_cost_tracker(9, 0);
+        test_update_cost_tracker(9, 9);
+        test_update_cost_tracker(9, -9);
+        test_update_cost_tracker(-9, 0);
+        test_update_cost_tracker(-9, 9);
+        test_update_cost_tracker(-9, -9);
     }
 
     #[test]