feature: Limit compute usage of a chunk

The commit introduces aggregation of compute usage across all operations
performed during chunk application and limiting of this compute usage
to 1s.

This should not change the behavior of nodes in the short run because
compute costs match gas costs and this is validated by the
`debug_assert`.

There are two follow ups to this work:
- near#8859
- near#8860

chain/chain/src/test_utils/kv_runtime.rs

@@ -1183,6 +1183,7 @@ impl RuntimeAdapter for KeyValueRuntime {
                         logs: vec![],
                         receipt_ids: new_receipt_hashes,
                         gas_burnt: 0,
+                        compute_usage: 0,
                         tokens_burnt: 0,
                         executor_id: to.clone(),
                         metadata: ExecutionMetadata::V1,

chain/chain/src/types.rs

@@ -637,6 +637,7 @@ mod tests {
                 logs: vec!["outcome1".to_string()],
                 receipt_ids: vec![hash(&[1])],
                 gas_burnt: 100,
+                compute_usage: 200,
                 tokens_burnt: 10000,
                 executor_id: "alice".parse().unwrap(),
                 metadata: ExecutionMetadata::V1,
@@ -649,6 +650,7 @@ mod tests {
                 logs: vec!["outcome2".to_string()],
                 receipt_ids: vec![],
                 gas_burnt: 0,
+                compute_usage: 0,
                 tokens_burnt: 0,
                 executor_id: "bob".parse().unwrap(),
                 metadata: ExecutionMetadata::V1,

core/primitives/src/transaction.rs

@@ -9,6 +9,7 @@ use borsh::{BorshDeserialize, BorshSerialize};
 use near_crypto::{PublicKey, Signature};
 use near_o11y::pretty;
 use near_primitives_core::profile::{ProfileDataV2, ProfileDataV3};
+use near_primitives_core::types::Compute;
 use std::borrow::Borrow;
 use std::fmt;
 use std::hash::{Hash, Hasher};
@@ -415,6 +416,10 @@ pub struct ExecutionOutcome {
     pub receipt_ids: Vec<CryptoHash>,
     /// The amount of the gas burnt by the given transaction or receipt.
     pub gas_burnt: Gas,
+    /// The amount of compute time spent by the given transaction or receipt.
+    // TODO(#8859): Treat this field in the same way as `gas_burnt`.
+    #[borsh_skip]
+    pub compute_usage: Compute,
     /// The amount of tokens burnt corresponding to the burnt gas amount.
     /// This value doesn't always equal to the `gas_burnt` multiplied by the gas price, because
     /// the prepaid gas price might be lower than the actual gas price and it creates a deficit.
@@ -437,7 +442,7 @@ pub enum ExecutionMetadata {
     V1,
     /// V2: With ProfileData by legacy `Cost` enum
     V2(ProfileDataV2),
-    // V3: With ProfileData by gas parameters
+    /// V3: With ProfileData by gas parameters
     V3(ProfileDataV3),
 }
 
@@ -598,6 +603,7 @@ mod tests {
             logs: vec!["123".to_string(), "321".to_string()],
             receipt_ids: vec![],
             gas_burnt: 123,
+            compute_usage: 456,
             tokens_burnt: 1234000,
             executor_id: "alice".parse().unwrap(),
             metadata: ExecutionMetadata::V1,

runtime/near-vm-logic/src/logic.rs

@@ -15,9 +15,9 @@ use near_primitives_core::config::ExtCosts::*;
 use near_primitives_core::config::{ActionCosts, ExtCosts, VMConfig};
 use near_primitives_core::runtime::fees::{transfer_exec_fee, transfer_send_fee};
 use near_primitives_core::types::{
-    AccountId, Balance, EpochHeight, Gas, ProtocolVersion, StorageUsage,
+    AccountId, Balance, Compute, EpochHeight, Gas, GasDistribution, GasWeight, ProtocolVersion,
+    StorageUsage,
 };
-use near_primitives_core::types::{GasDistribution, GasWeight};
 use near_vm_errors::{FunctionCallError, InconsistentStateError};
 use near_vm_errors::{HostError, VMLogicError};
 use std::mem::size_of;
@@ -2792,13 +2792,15 @@ impl<'a> VMLogic<'a> {
 
         let mut profile = self.gas_counter.profile_data();
         profile.compute_wasm_instruction_cost(burnt_gas);
+        let compute_usage = profile.total_compute_usage(&self.config.ext_costs);
 
         VMOutcome {
             balance: self.current_account_balance,
             storage_usage: self.current_storage_usage,
             return_data: self.return_data,
             burnt_gas,
             used_gas,
+            compute_usage,
             logs: self.logs,
             profile,
             action_receipts: self.receipt_manager.action_receipts,
@@ -2917,6 +2919,7 @@ pub struct VMOutcome {
     pub return_data: ReturnData,
     pub burnt_gas: Gas,
     pub used_gas: Gas,
+    pub compute_usage: Compute,
     pub logs: Vec<String>,
     /// Data collected from making a contract call
     pub profile: ProfileDataV3,
@@ -2950,6 +2953,7 @@ impl VMOutcome {
             return_data: ReturnData::None,
             burnt_gas: 0,
             used_gas: 0,
+            compute_usage: 0,
             logs: Vec::new(),
             profile: ProfileDataV3::default(),
             action_receipts: Vec::new(),

runtime/runtime/src/actions.rs

@@ -1,6 +1,6 @@
 use crate::config::{
-    safe_add_gas, total_prepaid_exec_fees, total_prepaid_gas, total_prepaid_send_fees,
-    RuntimeConfig,
+    safe_add_compute, safe_add_gas, total_prepaid_exec_fees, total_prepaid_gas,
+    total_prepaid_send_fees, RuntimeConfig,
 };
 use crate::ext::{ExternalError, RuntimeExt};
 use crate::{metrics, ActionResult, ApplyState};
@@ -254,6 +254,7 @@ pub(crate) fn action_function_call(
     // return a real `gas_used` instead of the `gas_burnt` into `ActionResult` even for
     // `FunctionCall`s error.
     result.gas_used = safe_add_gas(result.gas_used, outcome.used_gas)?;
+    result.compute_usage = safe_add_compute(result.compute_usage, outcome.compute_usage)?;
     result.logs.extend(outcome.logs);
     result.profile.merge(&outcome.profile);
     if execution_succeeded {
@@ -687,6 +688,8 @@ pub(crate) fn apply_delegate_action(
     // gas_used is incremented because otherwise the gas will be refunded. Refund function checks only gas_used.
     result.gas_used = safe_add_gas(result.gas_used, prepaid_send_fees)?;
     result.gas_burnt = safe_add_gas(result.gas_burnt, prepaid_send_fees)?;
+    // TODO(#8806): Support compute costs for actions. For now they match burnt gas.
+    result.compute_usage = safe_add_compute(result.compute_usage, prepaid_send_fees)?;
     result.new_receipts.push(new_receipt);
 
     Ok(())

runtime/runtime/src/config.rs

@@ -14,7 +14,7 @@ use near_primitives::runtime::fees::{transfer_exec_fee, transfer_send_fee, Runti
 use near_primitives::transaction::{
     Action, AddKeyAction, DeployContractAction, FunctionCallAction, Transaction,
 };
-use near_primitives::types::{AccountId, Balance, Gas};
+use near_primitives::types::{AccountId, Balance, Compute, Gas};
 use near_primitives::version::{is_implicit_account_creation_enabled, ProtocolVersion};
 
 /// Describes the cost of converting this transaction into a receipt.
@@ -59,6 +59,10 @@ pub fn safe_add_balance(a: Balance, b: Balance) -> Result<Balance, IntegerOverfl
     a.checked_add(b).ok_or_else(|| IntegerOverflowError {})
 }
 
+pub fn safe_add_compute(a: Compute, b: Compute) -> Result<Compute, IntegerOverflowError> {
+    a.checked_add(b).ok_or_else(|| IntegerOverflowError {})
+}
+
 #[macro_export]
 macro_rules! safe_add_balance_apply {
     ($x: expr) => {$x};

runtime/runtime/src/lib.rs

@@ -1,7 +1,7 @@
 use crate::actions::*;
 use crate::balance_checker::check_balance;
 use crate::config::{
-    exec_fee, safe_add_balance, safe_add_gas, safe_gas_to_balance, total_deposit,
+    exec_fee, safe_add_balance, safe_add_compute, safe_add_gas, safe_gas_to_balance, total_deposit,
     total_prepaid_exec_fees, total_prepaid_gas, RuntimeConfig,
 };
 use crate::genesis::{GenesisStateApplier, StorageComputer};
@@ -35,7 +35,7 @@ use near_primitives::transaction::{
 };
 use near_primitives::trie_key::TrieKey;
 use near_primitives::types::{
-    validator_stake::ValidatorStake, AccountId, Balance, EpochInfoProvider, Gas,
+    validator_stake::ValidatorStake, AccountId, Balance, Compute, EpochInfoProvider, Gas,
     RawStateChangesWithTrieKey, ShardId, StateChangeCause, StateRoot,
 };
 use near_primitives::utils::{
@@ -125,6 +125,7 @@ pub struct ActionResult {
     pub gas_burnt: Gas,
     pub gas_burnt_for_function_call: Gas,
     pub gas_used: Gas,
+    pub compute_usage: Compute,
     pub result: Result<ReturnData, ActionError>,
     pub logs: Vec<LogEntry>,
     pub new_receipts: Vec<Receipt>,
@@ -147,6 +148,7 @@ impl ActionResult {
             next_result.gas_burnt_for_function_call,
         )?;
         self.gas_used = safe_add_gas(self.gas_used, next_result.gas_used)?;
+        self.compute_usage = safe_add_compute(self.compute_usage, next_result.compute_usage)?;
         self.profile.merge(&next_result.profile);
         self.result = next_result.result;
         self.logs.append(&mut next_result.logs);
@@ -171,6 +173,7 @@ impl Default for ActionResult {
             gas_burnt: 0,
             gas_burnt_for_function_call: 0,
             gas_used: 0,
+            compute_usage: 0,
             result: Ok(ReturnData::None),
             logs: vec![],
             new_receipts: vec![],
@@ -263,6 +266,8 @@ impl Runtime {
                         logs: vec![],
                         receipt_ids: vec![receipt.receipt_id],
                         gas_burnt: verification_result.gas_burnt,
+                        // TODO(#8806): Support compute costs for actions. For now they match burnt gas.
+                        compute_usage: verification_result.gas_burnt,
                         tokens_burnt: verification_result.burnt_amount,
                         executor_id: transaction.signer_id.clone(),
                         // TODO: profile data is only counted in apply_action, which only happened at process_receipt
@@ -295,16 +300,17 @@ impl Runtime {
         actions: &[Action],
         epoch_info_provider: &dyn EpochInfoProvider,
     ) -> Result<ActionResult, RuntimeError> {
-        // println!("enter apply_action");
-        let mut result = ActionResult::default();
         let exec_fees = exec_fee(
             &apply_state.config.fees,
             action,
             &receipt.receiver_id,
             apply_state.current_protocol_version,
         );
-        result.gas_burnt += exec_fees;
-        result.gas_used += exec_fees;
+        let mut result = ActionResult::default();
+        result.gas_used = exec_fees;
+        result.gas_burnt = exec_fees;
+        // TODO(#8806): Support compute costs for actions. For now they match burnt gas.
+        result.compute_usage = exec_fees;
         let account_id = &receipt.receiver_id;
         let is_the_only_action = actions.len() == 1;
         let is_refund = AccountId::is_system(&receipt.predecessor_id);
@@ -499,9 +505,11 @@ impl Runtime {
         let mut account = get_account(state_update, account_id)?;
         let mut actor_id = receipt.predecessor_id.clone();
         let mut result = ActionResult::default();
-        let exec_fee = apply_state.config.fees.fee(ActionCosts::new_action_receipt).exec_fee();
-        result.gas_used = exec_fee;
-        result.gas_burnt = exec_fee;
+        let exec_fees = apply_state.config.fees.fee(ActionCosts::new_action_receipt).exec_fee();
+        result.gas_used = exec_fees;
+        result.gas_burnt = exec_fees;
+        // TODO(#8806): Support compute costs for actions. For now they match burnt gas.
+        result.compute_usage = exec_fees;
         // Executing actions one by one
         for (action_index, action) in action_receipt.actions.iter().enumerate() {
             let action_hash = create_action_hash(
@@ -586,6 +594,7 @@ impl Runtime {
                 apply_state.current_protocol_version
             ) {
                 result.gas_burnt = 0;
+                result.compute_usage = 0;
                 result.gas_used = 0;
             }
 
@@ -749,6 +758,7 @@ impl Runtime {
                 logs: result.logs,
                 receipt_ids,
                 gas_burnt: result.gas_burnt,
+                compute_usage: result.compute_usage,
                 tokens_burnt,
                 executor_id: account_id.clone(),
                 metadata: ExecutionMetadata::V3(result.profile),
@@ -1253,6 +1263,7 @@ impl Runtime {
         // charge any gas for refund receipts, we still count the gas use towards the block gas
         // limit
         let mut total_gas_burnt = gas_used_for_migrations;
+        let mut total_compute_usage = gas_used_for_migrations;
 
         for signed_transaction in transactions {
             let (receipt, outcome_with_id) = self.process_transaction(
@@ -1268,6 +1279,15 @@ impl Runtime {
             }
 
             total_gas_burnt = safe_add_gas(total_gas_burnt, outcome_with_id.outcome.gas_burnt)?;
+            total_compute_usage =
+                safe_add_compute(total_compute_usage, outcome_with_id.outcome.compute_usage)?;
+
+            // TODO(#8032): Remove when compute costs are stabilized.
+            debug_assert_eq!(
+                total_compute_usage, total_gas_burnt,
+                "Compute usage must match burnt gas"
+            );
+
             outcomes.push(outcome_with_id);
         }
 
@@ -1277,7 +1297,8 @@ impl Runtime {
 
         let mut process_receipt = |receipt: &Receipt,
                                    state_update: &mut TrieUpdate,
-                                   total_gas_burnt: &mut Gas|
+                                   total_gas_burnt: &mut Gas,
+                                   total_compute_usage: &mut Compute|
          -> Result<_, RuntimeError> {
             let _span = tracing::debug_span!(
                 target: "runtime",
@@ -1302,12 +1323,21 @@ impl Runtime {
             if let Some(outcome_with_id) = result? {
                 *total_gas_burnt =
                     safe_add_gas(*total_gas_burnt, outcome_with_id.outcome.gas_burnt)?;
+                *total_compute_usage =
+                    safe_add_compute(*total_compute_usage, outcome_with_id.outcome.compute_usage)?;
+                // TODO(#8032): Remove when compute costs are stabilized.
+                debug_assert_eq!(
+                    total_compute_usage, total_gas_burnt,
+                    "Compute usage must match burnt gas"
+                );
                 outcomes.push(outcome_with_id);
             }
             Ok(())
         };
 
-        let gas_limit = apply_state.gas_limit.unwrap_or(Gas::max_value());
+        // TODO(#8859): Introduce a dedicated `compute_limit` for the chunk.
+        // For now compute limit always matches the gas limit.
+        let compute_limit = apply_state.gas_limit.unwrap_or(Gas::max_value());
 
         // We first process local receipts. They contain staking, local contract calls, etc.
         if let Some(prefetcher) = &mut prefetcher {
@@ -1316,18 +1346,23 @@ impl Runtime {
             _ = prefetcher.prefetch_receipts_data(&local_receipts);
         }
         for receipt in local_receipts.iter() {
-            if total_gas_burnt < gas_limit {
+            if total_compute_usage < compute_limit {
                 // NOTE: We don't need to validate the local receipt, because it's just validated in
                 // the `verify_and_charge_transaction`.
-                process_receipt(receipt, &mut state_update, &mut total_gas_burnt)?;
+                process_receipt(
+                    receipt,
+                    &mut state_update,
+                    &mut total_gas_burnt,
+                    &mut total_compute_usage,
+                )?;
             } else {
                 Self::delay_receipt(&mut state_update, &mut delayed_receipts_indices, receipt)?;
             }
         }
 
         // Then we process the delayed receipts. It's a backlog of receipts from the past blocks.
         while delayed_receipts_indices.first_index < delayed_receipts_indices.next_available_index {
-            if total_gas_burnt >= gas_limit {
+            if total_compute_usage >= compute_limit {
                 break;
             }
             let key = TrieKey::DelayedReceipt { index: delayed_receipts_indices.first_index };
@@ -1360,7 +1395,12 @@ impl Runtime {
             state_update.remove(key);
             // Math checked above: first_index is less than next_available_index
             delayed_receipts_indices.first_index += 1;
-            process_receipt(&receipt, &mut state_update, &mut total_gas_burnt)?;
+            process_receipt(
+                &receipt,
+                &mut state_update,
+                &mut total_gas_burnt,
+                &mut total_compute_usage,
+            )?;
             processed_delayed_receipts.push(receipt);
         }
 
@@ -1379,8 +1419,13 @@ impl Runtime {
                 apply_state.current_protocol_version,
             )
             .map_err(RuntimeError::ReceiptValidationError)?;
-            if total_gas_burnt < gas_limit {
-                process_receipt(receipt, &mut state_update, &mut total_gas_burnt)?;
+            if total_compute_usage < compute_limit {
+                process_receipt(
+                    receipt,
+                    &mut state_update,
+                    &mut total_gas_burnt,
+                    &mut total_compute_usage,
+                )?;
             } else {
                 Self::delay_receipt(&mut state_update, &mut delayed_receipts_indices, receipt)?;
             }

tools/state-viewer/src/contract_accounts.rs

@@ -677,6 +677,7 @@ mod tests {
                 logs: vec![],
                 receipt_ids,
                 gas_burnt: 100,
+                compute_usage: 200,
                 tokens_burnt: 2000,
                 executor_id: "someone.near".parse().unwrap(),
                 status: ExecutionStatus::SuccessValue(vec![]),