feat(gas-keys): charge compute for gas key nonce removal

core/store/src/utils/mod.rs

@@ -384,14 +384,22 @@ pub fn compute_gas_key_balance_sum(
     Ok(total)
 }
 
+pub struct RemoveAccountResult {
+    pub gas_key_nonce_count: usize,
+    pub gas_key_nonce_total_key_bytes: usize, // used to calculate compute cost
+}
+
 /// Removes account, code and all access keys and gas keys associated to it.
 pub fn remove_account(
     state_update: &mut TrieUpdate,
     account_id: &AccountId,
-) -> Result<(), StorageError> {
+) -> Result<RemoveAccountResult, StorageError> {
     state_update.remove(TrieKey::Account { account_id: account_id.clone() });
     state_update.remove(TrieKey::ContractCode { account_id: account_id.clone() });
 
+    let mut gas_key_nonce_count: usize = 0;
+    let mut gas_key_nonce_total_key_bytes: usize = 0;
+
     // Removing access keys and gas key nonces
     let lock = state_update.trie().lock_for_iter();
     let mut keys_to_remove: Vec<TrieKey> = Vec::new();
@@ -415,6 +423,8 @@ pub fn remove_account(
                     )
                 })?;
         if let Some(index) = nonce_index {
+            gas_key_nonce_count += 1;
+            gas_key_nonce_total_key_bytes += raw_key.len();
             keys_to_remove.push(TrieKey::GasKeyNonce {
                 account_id: account_id.clone(),
                 public_key: public_key.clone(),
@@ -452,7 +462,7 @@ pub fn remove_account(
     for key in data_keys {
         state_update.remove(TrieKey::ContractData { account_id: account_id.clone(), key });
     }
-    Ok(())
+    Ok(RemoveAccountResult { gas_key_nonce_count, gas_key_nonce_total_key_bytes })
 }
 
 pub fn get_genesis_state_roots(store: &Store) -> Option<Vec<StateRoot>> {

runtime/runtime/src/access_keys.rs

@@ -1,15 +1,16 @@
 use std::mem::size_of;
 
 use near_crypto::PublicKey;
-use near_parameters::RuntimeFeesConfig;
+use near_parameters::{RuntimeConfig, RuntimeFeesConfig};
 use near_primitives::account::{AccessKey, Account, GasKeyInfo};
 use near_primitives::action::{TransferToGasKeyAction, WithdrawFromGasKeyAction};
 use near_primitives::errors::{ActionErrorKind, IntegerOverflowError, RuntimeError};
 use near_primitives::transaction::{AddKeyAction, DeleteKeyAction};
-use near_primitives::types::{AccountId, BlockHeight, Compute, Nonce, NonceIndex, StorageUsage};
+use near_primitives::types::{AccountId, BlockHeight, Nonce, NonceIndex, StorageUsage};
 
-use crate::config::safe_add_compute;
+use crate::config::{safe_add_compute, storage_removes_compute};
 use crate::{ActionResult, ApplyState};
+use near_primitives::trie_key::gas_key_nonce_key_len;
 use near_store::{
     StorageError, TrieUpdate, get_access_key, remove_access_key, remove_gas_key_nonce,
     set_access_key, set_gas_key_nonce,
@@ -42,20 +43,14 @@ fn gas_key_storage_cost(
         + access_key_storage_usage(fee_config, public_key, access_key)
 }
 
-/// Returns the compute cost for deleting a single gas key nonce.
-fn gas_key_nonce_delete_compute_cost() -> Compute {
-    // TODO(gas-keys): properly handle GasKey fees
-    near_primitives::gas::Gas::ZERO.as_gas()
-}
-
 pub(crate) fn initial_nonce_value(block_height: BlockHeight) -> Nonce {
     // Set default nonce for newly created access key to avoid transaction hash collision.
     // See <https://github.com/near/nearcore/issues/3779>.
     (block_height - 1) * near_primitives::account::AccessKey::ACCESS_KEY_NONCE_RANGE_MULTIPLIER
 }
 
 pub(crate) fn action_delete_key(
-    fee_config: &RuntimeFeesConfig,
+    config: &RuntimeConfig,
     state_update: &mut TrieUpdate,
     account: &mut Account,
     result: &mut ActionResult,
@@ -66,7 +61,7 @@ pub(crate) fn action_delete_key(
     if let Some(access_key) = access_key {
         if let Some(gas_key_info) = access_key.gas_key_info() {
             delete_gas_key(
-                fee_config,
+                config,
                 state_update,
                 account,
                 result,
@@ -77,7 +72,7 @@ pub(crate) fn action_delete_key(
             )?;
         } else {
             delete_regular_key(
-                fee_config,
+                &config.fees,
                 state_update,
                 account,
                 account_id,
@@ -96,7 +91,7 @@ pub(crate) fn action_delete_key(
 }
 
 fn delete_gas_key(
-    fee_config: &RuntimeFeesConfig,
+    config: &RuntimeConfig,
     state_update: &mut TrieUpdate,
     account: &mut Account,
     result: &mut ActionResult,
@@ -116,15 +111,21 @@ fn delete_gas_key(
     }
     result.tokens_burnt =
         result.tokens_burnt.checked_add(gas_key_info.balance).ok_or(IntegerOverflowError)?;
+    let num_nonces = gas_key_info.num_nonces as usize;
     for i in 0..gas_key_info.num_nonces {
         remove_gas_key_nonce(state_update, account_id.clone(), public_key.clone(), i);
     }
-    let nonce_delete_compute_cost =
-        gas_key_nonce_delete_compute_cost() * gas_key_info.num_nonces as u64;
-    result.compute_usage = safe_add_compute(result.compute_usage, nonce_delete_compute_cost)?;
+    let nonce_key_len = gas_key_nonce_key_len(account_id, public_key);
+    let nonce_remove_compute = storage_removes_compute(
+        &config.wasm_config.ext_costs,
+        num_nonces,
+        nonce_key_len * num_nonces,
+        AccessKey::NONCE_VALUE_LEN * num_nonces,
+    );
+    result.compute_usage = safe_add_compute(result.compute_usage, nonce_remove_compute)?;
     remove_access_key(state_update, account_id.clone(), public_key.clone());
     account.set_storage_usage(account.storage_usage().saturating_sub(gas_key_storage_cost(
-        fee_config,
+        &config.fees,
         public_key,
         access_key,
         gas_key_info.num_nonces,
@@ -341,6 +342,7 @@ mod tests {
     use crate::ActionResult;
     use crate::ApplyState;
     use crate::actions_test_utils::{setup_account, test_delete_large_account};
+    use crate::config::storage_removes_compute;
     use crate::state_viewer::TrieViewer;
 
     use super::*;
@@ -362,6 +364,21 @@ mod tests {
     const TEST_NUM_NONCES: NonceIndex = 2;
     const TEST_GAS_KEY_BLOCK_HEIGHT: BlockHeight = 10;
 
+    fn expected_nonce_remove_compute(
+        account_id: &AccountId,
+        public_key: &PublicKey,
+        num_nonces: usize,
+    ) -> u64 {
+        let config = RuntimeConfig::test();
+        let nonce_key_len = gas_key_nonce_key_len(account_id, public_key);
+        storage_removes_compute(
+            &config.wasm_config.ext_costs,
+            num_nonces,
+            nonce_key_len * num_nonces,
+            AccessKey::NONCE_VALUE_LEN * num_nonces,
+        )
+    }
+
     fn create_apply_state(block_height: BlockHeight) -> ApplyState {
         ApplyState {
             apply_reason: ApplyChunkReason::UpdateTrackedShard,
@@ -544,7 +561,7 @@ mod tests {
         let mut result = ActionResult::default();
         let action = DeleteKeyAction { public_key: gas_key_public_key.clone() };
         action_delete_key(
-            &RuntimeFeesConfig::test(),
+            &RuntimeConfig::test(),
             &mut state_update,
             &mut account,
             &mut result,
@@ -570,6 +587,13 @@ mod tests {
                     .expect("failed to get gas key nonce");
             assert!(gas_key_nonce.is_none());
         }
+
+        let expected_compute = expected_nonce_remove_compute(
+            &account_id,
+            &gas_key_public_key,
+            TEST_NUM_NONCES as usize,
+        );
+        assert_eq!(result.compute_usage, expected_compute);
     }
 
     #[test]
@@ -591,7 +615,7 @@ mod tests {
         let mut result = ActionResult::default();
         let action = DeleteKeyAction { public_key: gas_key_public_key.clone() };
         action_delete_key(
-            &RuntimeFeesConfig::test(),
+            &RuntimeConfig::test(),
             &mut state_update,
             &mut account,
             &mut result,
@@ -603,6 +627,12 @@ mod tests {
 
         // Verify the balance was burned
         assert_eq!(result.tokens_burnt, deposit_amount);
+        let expected_compute = expected_nonce_remove_compute(
+            &account_id,
+            &gas_key_public_key,
+            TEST_NUM_NONCES as usize,
+        );
+        assert_eq!(result.compute_usage, expected_compute);
     }
 
     #[test]
@@ -619,7 +649,7 @@ mod tests {
         let mut result = ActionResult::default();
         let action = DeleteKeyAction { public_key: nonexistent_public_key.clone() };
         action_delete_key(
-            &RuntimeFeesConfig::test(),
+            &RuntimeConfig::test(),
             &mut state_update,
             &mut account,
             &mut result,
@@ -655,6 +685,12 @@ mod tests {
         assert!(action_result.result.is_ok());
         state_update.commit(StateChangeCause::InitialState);
 
+        let expected_compute: u64 = public_keys
+            .iter()
+            .map(|pk| expected_nonce_remove_compute(&account_id, pk, TEST_NUM_NONCES as usize))
+            .sum();
+        assert_eq!(action_result.compute_usage, expected_compute);
+
         let lock = state_update.trie().lock_for_iter();
         let trie_key_count = state_update
             .locked_iter(&trie_key_parsers::get_raw_prefix_for_access_keys(&account_id), &lock)
@@ -694,6 +730,11 @@ mod tests {
         let expected_burnt =
             deposit_amounts.iter().fold(Balance::ZERO, |acc, x| acc.checked_add(*x).unwrap());
         assert_eq!(action_result.tokens_burnt, expected_burnt);
+        let expected_compute: u64 = public_keys
+            .iter()
+            .map(|pk| expected_nonce_remove_compute(&account_id, pk, TEST_NUM_NONCES as usize))
+            .sum();
+        assert_eq!(action_result.compute_usage, expected_compute);
     }
 
     #[test]
@@ -1000,7 +1041,7 @@ mod tests {
         let mut result = ActionResult::default();
         let action = DeleteKeyAction { public_key: gas_key_public_key.clone() };
         action_delete_key(
-            &RuntimeFeesConfig::test(),
+            &RuntimeConfig::test(),
             &mut state_update,
             &mut account,
             &mut result,
@@ -1040,7 +1081,7 @@ mod tests {
         let mut result = ActionResult::default();
         let action = DeleteKeyAction { public_key: gas_key_public_key.clone() };
         action_delete_key(
-            &RuntimeFeesConfig::test(),
+            &RuntimeConfig::test(),
             &mut state_update,
             &mut account,
             &mut result,

runtime/runtime/src/actions.rs

@@ -1,6 +1,7 @@
 use crate::access_keys::initial_nonce_value;
 use crate::config::{
-    safe_add_compute, total_prepaid_exec_fees, total_prepaid_gas, total_prepaid_send_fees,
+    safe_add_compute, storage_removes_compute, total_prepaid_exec_fees, total_prepaid_gas,
+    total_prepaid_send_fees,
 };
 use crate::deterministic_account_id::create_deterministic_account;
 use crate::{ActionResult, ApplyState};
@@ -337,6 +338,7 @@ pub(crate) fn action_delete_account(
     result: &mut ActionResult,
     account_id: &AccountId,
     delete_account: &DeleteAccountAction,
+    config: &RuntimeConfig,
     current_protocol_version: ProtocolVersion,
 ) -> Result<(), StorageError> {
     let account_ref = account.as_ref().unwrap();
@@ -377,11 +379,22 @@ pub(crate) fn action_delete_account(
             .new_receipts
             .push(Receipt::new_balance_refund(&delete_account.beneficiary_id, account_balance));
     }
-    remove_account(state_update, account_id)?;
+    let remove_result = remove_account(state_update, account_id)?;
     result.tokens_burnt =
         result.tokens_burnt.checked_add(gas_key_balance_to_burn).ok_or_else(|| {
             StorageError::StorageInconsistentState("tokens_burnt overflow".to_string())
         })?;
+    if remove_result.gas_key_nonce_count > 0 {
+        let compute = storage_removes_compute(
+            &config.wasm_config.ext_costs,
+            remove_result.gas_key_nonce_count,
+            remove_result.gas_key_nonce_total_key_bytes,
+            AccessKey::NONCE_VALUE_LEN * remove_result.gas_key_nonce_count,
+        );
+        result.compute_usage = safe_add_compute(result.compute_usage, compute).map_err(|_| {
+            StorageError::StorageInconsistentState("compute_usage overflow".to_string())
+        })?;
+    }
     *actor_id = receipt.predecessor_id().clone();
     *account = None;
     Ok(())

runtime/runtime/src/actions_test_utils.rs

@@ -1,4 +1,5 @@
 use near_crypto::PublicKey;
+use near_parameters::RuntimeConfig;
 use near_primitives::account::{AccessKey, Account, AccountContract};
 use near_primitives::action::DeleteAccountAction;
 use near_primitives::hash::CryptoHash;
@@ -46,6 +47,7 @@ pub(crate) fn test_delete_large_account(
     let mut actor_id = account_id.clone();
     let mut action_result = ActionResult::default();
     let receipt = Receipt::new_balance_refund(&"alice.near".parse().unwrap(), Balance::ZERO);
+    let config = RuntimeConfig::test();
     let res = action_delete_account(
         state_update,
         &mut account,
@@ -54,6 +56,7 @@ pub(crate) fn test_delete_large_account(
         &mut action_result,
         account_id,
         &DeleteAccountAction { beneficiary_id: "bob".parse().unwrap() },
+        &config,
         PROTOCOL_VERSION,
     );
     assert!(res.is_ok());

runtime/runtime/src/config.rs

@@ -6,7 +6,8 @@ use near_primitives::action::DeployGlobalContractAction;
 use near_primitives::errors::IntegerOverflowError;
 // Just re-exporting RuntimeConfig for backwards compatibility.
 use near_parameters::{
-    ActionCosts, RuntimeConfig, RuntimeFeesConfig, transfer_exec_fee, transfer_send_fee,
+    ActionCosts, ExtCosts, ExtCostsConfig, RuntimeConfig, RuntimeFeesConfig, transfer_exec_fee,
+    transfer_send_fee,
 };
 pub use near_primitives::num_rational::Rational32;
 use near_primitives::transaction::{Action, DeployContractAction, Transaction};
@@ -68,6 +69,22 @@ fn gas_key_transfer_exec_fee(
         .unwrap()
 }
 
+/// Compute cost for `count` storage_remove operations with `total_key_bytes` total
+/// key bytes and `total_value_bytes` total value bytes across all removals.
+pub(crate) fn storage_removes_compute(
+    ext: &ExtCostsConfig,
+    count: usize,
+    total_key_bytes: usize,
+    total_value_bytes: usize,
+) -> Compute {
+    let count = count as u64;
+    let total_key_bytes = total_key_bytes as u64;
+    let total_value_bytes = total_value_bytes as u64;
+    ext.compute_cost(ExtCosts::storage_remove_base) * count
+        + ext.compute_cost(ExtCosts::storage_remove_key_byte) * total_key_bytes
+        + ext.compute_cost(ExtCosts::storage_remove_ret_value_byte) * total_value_bytes
+}
+
 /// Total sum of gas that needs to be burnt to send these actions.
 pub fn total_send_fees(
     config: &RuntimeConfig,

runtime/runtime/src/lib.rs

@@ -604,7 +604,7 @@ impl Runtime {
             Action::DeleteKey(delete_key) => {
                 metrics::ACTION_CALLED_COUNT.delete_key.inc();
                 action_delete_key(
-                    &apply_state.config.fees,
+                    &apply_state.config,
                     state_update,
                     account.as_mut().expect(EXPECT_ACCOUNT_EXISTS),
                     &mut result,
@@ -622,6 +622,7 @@ impl Runtime {
                     &mut result,
                     account_id,
                     delete_account,
+                    &apply_state.config,
                     apply_state.current_protocol_version,
                 )?;
             }