feature: Get compute usage from profile (#8783)

This seemed like the simplest approach that introduced minimal runtime overhead and code changes (as opposed to introducing a dedicated field in the GasCounter). Happy to consider the alternatives though.

Part of #8032

Some follow-up work is planned in #8795

core/primitives-core/src/config.rs

@@ -286,7 +286,7 @@ impl ExtCostsConfig {
 
     /// Convenience constructor to use in tests where the exact gas cost does
     /// not need to correspond to a specific protocol version.
-    pub fn test() -> ExtCostsConfig {
+    pub fn test_with_undercharging_factor(factor: u64) -> ExtCostsConfig {
         let costs = enum_map! {
             ExtCosts::base => SAFETY_MULTIPLIER * 88256037,
             ExtCosts::contract_loading_base => SAFETY_MULTIPLIER * 11815321,
@@ -352,10 +352,15 @@ impl ExtCostsConfig {
             ExtCosts::alt_bn128_g1_sum_base => 3_000_000_000,
             ExtCosts::alt_bn128_g1_sum_element => 5_000_000_000,
         }
-        .map(|_, value| ParameterCost { gas: value, compute: value });
+        .map(|_, value| ParameterCost { gas: value, compute: value * factor });
         ExtCostsConfig { costs }
     }
 
+    /// `test_with_undercharging_factor` with a factor of 1.
+    pub fn test() -> ExtCostsConfig {
+        Self::test_with_undercharging_factor(1)
+    }
+
     fn free() -> ExtCostsConfig {
         ExtCostsConfig {
             costs: enum_map! {
@@ -482,10 +487,14 @@ pub enum ActionCosts {
 }
 
 impl ExtCosts {
-    pub fn value(self, config: &ExtCostsConfig) -> Gas {
+    pub fn gas(self, config: &ExtCostsConfig) -> Gas {
         config.gas_cost(self)
     }
 
+    pub fn compute(self, config: &ExtCostsConfig) -> Compute {
+        config.compute_cost(self)
+    }
+
     pub fn param(&self) -> Parameter {
         match self {
             ExtCosts::base => Parameter::WasmBase,

core/primitives-core/src/profile.rs

@@ -1,7 +1,7 @@
 pub use profile_v2::ProfileDataV2;
 
-use crate::config::{ActionCosts, ExtCosts};
-use crate::types::Gas;
+use crate::config::{ActionCosts, ExtCosts, ExtCostsConfig};
+use crate::types::{Compute, Gas};
 use borsh::{BorshDeserialize, BorshSerialize};
 use enum_map::{enum_map, Enum, EnumMap};
 use std::fmt;
@@ -105,6 +105,34 @@ impl ProfileDataV3 {
     pub fn action_gas(&self) -> Gas {
         self.actions_profile.as_slice().iter().copied().fold(0, Gas::saturating_add)
     }
+
+    /// Returns total compute usage of host calls.
+    pub fn total_compute_usage(&self, ext_costs_config: &ExtCostsConfig) -> Compute {
+        let ext_compute_cost = self
+            .wasm_ext_profile
+            .iter()
+            .map(|(key, value)| {
+                // Technically, gas cost might be zero while the compute cost is non-zero. To
+                // handle this case, we would need to explicitly count number of calls, not just
+                // the total gas usage.
+                // We don't have such costs at the moment, so this case is not implemented.
+                debug_assert!(key.gas(ext_costs_config) > 0 || key.compute(ext_costs_config) == 0);
+
+                if *value == 0 {
+                    return *value;
+                }
+                // If the `value` is non-zero, the gas cost also must be non-zero.
+                debug_assert!(key.gas(ext_costs_config) != 0);
+                debug_assert!(*value % key.gas(ext_costs_config) == 0);
+                // TODO(#8795): Consider storing the count of calls and avoid division here.
+                (*value / key.gas(ext_costs_config)).saturating_mul(key.compute(ext_costs_config))
+            })
+            .fold(0, Compute::saturating_add);
+
+        // We currently only support compute costs for host calls. In the future we might add
+        // them for actions as well.
+        ext_compute_cost.saturating_add(self.action_gas()).saturating_add(self.get_wasm_cost())
+    }
 }
 
 impl BorshDeserialize for ProfileDataV3 {
@@ -254,6 +282,26 @@ mod test {
         assert_eq!(profile_data.get_ext_cost(ExtCosts::storage_read_base), 33);
     }
 
+    #[test]
+    fn test_total_compute_usage() {
+        let ext_costs_config = ExtCostsConfig::test_with_undercharging_factor(3);
+        let mut profile_data = ProfileDataV3::default();
+        profile_data.add_ext_cost(
+            ExtCosts::storage_read_base,
+            2 * ExtCosts::storage_read_base.gas(&ext_costs_config),
+        );
+        profile_data.add_ext_cost(
+            ExtCosts::storage_write_base,
+            5 * ExtCosts::storage_write_base.gas(&ext_costs_config),
+        );
+        profile_data.add_action_cost(ActionCosts::function_call_base, 100);
+
+        assert_eq!(
+            profile_data.total_compute_usage(&ext_costs_config),
+            3 * profile_data.host_gas() + profile_data.action_gas()
+        );
+    }
+
     #[test]
     fn test_borsh_ser_deser() {
         let mut profile_data = ProfileDataV3::default();

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

@@ -79,7 +79,7 @@ impl GasCounter {
                 gas_limit: min(max_gas_burnt, prepaid_gas),
                 opcode_cost: Gas::from(opcode_cost),
             },
-            max_gas_burnt: max_gas_burnt,
+            max_gas_burnt,
             promises_gas: 0,
             prepaid_gas,
             is_view,
@@ -218,9 +218,8 @@ impl GasCounter {
 
     /// A helper function to pay a multiple of a cost.
     pub fn pay_per(&mut self, cost: ExtCosts, num: u64) -> Result<()> {
-        let use_gas = num
-            .checked_mul(cost.value(&self.ext_costs_config))
-            .ok_or(HostError::IntegerOverflow)?;
+        let use_gas =
+            num.checked_mul(cost.gas(&self.ext_costs_config)).ok_or(HostError::IntegerOverflow)?;
 
         self.inc_ext_costs_counter(cost, num);
         self.update_profile_host(cost, use_gas);
@@ -229,7 +228,7 @@ impl GasCounter {
 
     /// A helper function to pay base cost gas.
     pub fn pay_base(&mut self, cost: ExtCosts) -> Result<()> {
-        let base_fee = cost.value(&self.ext_costs_config);
+        let base_fee = cost.gas(&self.ext_costs_config);
         self.inc_ext_costs_counter(cost, 1);
         self.update_profile_host(cost, base_fee);
         self.burn_gas(base_fee)
@@ -284,6 +283,9 @@ mod tests {
     use crate::{ExtCostsConfig, HostError};
     use near_primitives_core::types::Gas;
 
+    /// Max prepaid amount of gas.
+    const MAX_GAS: u64 = 300_000_000_000_000;
+
     fn make_test_counter(max_burnt: Gas, prepaid: Gas, is_view: bool) -> super::GasCounter {
         super::GasCounter::new(ExtCostsConfig::test(), max_burnt, 1, prepaid, is_view)
     }
@@ -344,4 +346,24 @@ mod tests {
         test(8, 5, false, Err(HostError::GasExceeded));
         test(8, 5, true, Ok(()));
     }
+
+    #[test]
+    fn test_profile_compute_cost_is_accurate() {
+        let mut counter = make_test_counter(MAX_GAS, MAX_GAS, false);
+        counter.pay_base(near_primitives::config::ExtCosts::storage_write_base).unwrap();
+        counter.pay_per(near_primitives::config::ExtCosts::storage_write_value_byte, 10).unwrap();
+        counter.pay_wasm_gas(20).unwrap();
+        counter
+            .pay_action_accumulated(
+                100,
+                100,
+                near_primitives::config::ActionCosts::new_data_receipt_byte,
+            )
+            .unwrap();
+
+        let mut profile = counter.profile_data().clone();
+        profile.compute_wasm_instruction_cost(counter.burnt_gas());
+
+        assert_eq!(profile.total_compute_usage(&ExtCostsConfig::test()), counter.burnt_gas());
+    }
 }