Track slot allocation fee on state metadata

aptos-move/aptos-gas-meter/src/meter.rs

@@ -495,8 +495,12 @@ where
             .map_err(|e| e.finish(Location::Undefined))
     }
 
-    fn storage_fee_per_write(&self, key: &StateKey, op: &WriteOp) -> Fee {
-        self.vm_gas_params().txn.storage_fee_per_write(key, op)
+    fn storage_fee_for_state_slot(&self, op: &WriteOp) -> Fee {
+        self.vm_gas_params().txn.storage_fee_for_slot(op)
+    }
+
+    fn storage_fee_for_state_bytes(&self, key: &StateKey, op: &WriteOp) -> Fee {
+        self.vm_gas_params().txn.storage_fee_for_bytes(key, op)
     }
 
     fn storage_fee_per_event(&self, event: &ContractEvent) -> Fee {

aptos-move/aptos-gas-meter/src/traits.rs

@@ -6,7 +6,7 @@ use aptos_gas_schedule::VMGasParameters;
 use aptos_types::{
     contract_event::ContractEvent, state_store::state_key::StateKey, write_set::WriteOp,
 };
-use aptos_vm_types::storage::StorageGasParameters;
+use aptos_vm_types::{change_set::VMChangeSet, storage::StorageGasParameters};
 use move_binary_format::errors::{Location, PartialVMResult, VMResult};
 use move_core_types::gas_algebra::{InternalGas, InternalGasUnit, NumBytes};
 use move_vm_types::gas::GasMeter as MoveGasMeter;
@@ -103,8 +103,11 @@ pub trait AptosGasMeter: MoveGasMeter {
     /// storage costs.
     fn charge_io_gas_for_write(&mut self, key: &StateKey, op: &WriteOp) -> VMResult<()>;
 
-    /// Calculates the storage fee for a write operation.
-    fn storage_fee_per_write(&self, key: &StateKey, op: &WriteOp) -> Fee;
+    /// Calculates the storage fee for a state slot allocation.
+    fn storage_fee_for_state_slot(&self, op: &WriteOp) -> Fee;
+
+    /// Calculates the storage fee for state bytes.
+    fn storage_fee_for_state_bytes(&self, key: &StateKey, op: &WriteOp) -> Fee;
 
     /// Calculates the storage fee for an event.
     fn storage_fee_per_event(&self, event: &ContractEvent) -> Fee;
@@ -116,16 +119,16 @@ pub trait AptosGasMeter: MoveGasMeter {
     fn storage_fee_for_transaction_storage(&self, txn_size: NumBytes) -> Fee;
 
     /// Charges the storage fees for writes, events & txn storage in a lump sum, minimizing the
-    /// loss of precision.
+    /// loss of precision. Refundable portion of the charge is recorded on the WriteOp itself,
+    /// due to which mutable references are required on the parameter list wherever proper.
     ///
     /// The contract requires that this function behaves in a way that is consistent to
     /// the ones defining the costs.
     /// Due to this reason, you should normally not override the default implementation,
     /// unless you are doing something special, such as injecting additional logging logic.
-    fn charge_storage_fee_for_all<'a>(
+    fn process_storage_fee_for_all(
         &mut self,
-        write_ops: impl IntoIterator<Item = (&'a StateKey, &'a WriteOp)>,
-        events: impl IntoIterator<Item = &'a ContractEvent>,
+        change_set: &mut VMChangeSet,
         txn_size: NumBytes,
         gas_unit_price: FeePerGasUnit,
     ) -> VMResult<()> {
@@ -142,10 +145,16 @@ pub trait AptosGasMeter: MoveGasMeter {
         }
 
         // Calculate the storage fees.
-        let write_fee = write_ops.into_iter().fold(Fee::new(0), |acc, (key, op)| {
-            acc + self.storage_fee_per_write(key, op)
-        });
-        let event_fee = events.into_iter().fold(Fee::new(0), |acc, event| {
+        let write_fee = change_set
+            .write_set_iter_mut()
+            .fold(Fee::new(0), |acc, (key, op)| {
+                let slot_fee = self.storage_fee_for_state_slot(op);
+                let bytes_fee = self.storage_fee_for_state_bytes(key, op);
+                Self::maybe_record_storage_deposit(op, slot_fee);
+
+                acc + slot_fee + bytes_fee
+            });
+        let event_fee = change_set.events().iter().fold(Fee::new(0), |acc, event| {
             acc + self.storage_fee_per_event(event)
         });
         let event_discount = self.storage_discount_for_events(event_fee);
@@ -161,6 +170,28 @@ pub trait AptosGasMeter: MoveGasMeter {
         Ok(())
     }
 
+    // The slot fee is refundable, we record it on the WriteOp itself and it'll end up in
+    // the state DB.
+    fn maybe_record_storage_deposit(write_op: &mut WriteOp, slot_fee: Fee) {
+        use WriteOp::*;
+
+        match write_op {
+            CreationWithMetadata {
+                ref mut metadata,
+                data: _,
+            } => {
+                if !slot_fee.is_zero() {
+                    metadata.set_deposit(slot_fee.into())
+                }
+            },
+            Creation(..)
+            | Modification(..)
+            | Deletion
+            | ModificationWithMetadata { .. }
+            | DeletionWithMetadata { .. } => {},
+        }
+    }
+
     // Below are getters reexported from the gas algebra.
     // Gas meter instances should not reimplement these themselves.
 

aptos-move/aptos-gas-profiling/Cargo.toml

@@ -22,6 +22,7 @@ aptos-gas-algebra = { workspace = true }
 aptos-gas-meter = { workspace = true }
 aptos-package-builder = { workspace = true }
 aptos-types = { workspace = true }
+aptos-vm-types = { workspace = true }
 
 move-binary-format = { workspace = true }
 move-core-types = { workspace = true }

aptos-move/aptos-gas-profiling/src/profiler.rs

@@ -10,6 +10,7 @@ use aptos_gas_meter::AptosGasMeter;
 use aptos_types::{
     contract_event::ContractEvent, state_store::state_key::StateKey, write_set::WriteOp,
 };
+use aptos_vm_types::change_set::VMChangeSet;
 use move_binary_format::{
     errors::{Location, PartialVMResult, VMResult},
     file_format::CodeOffset,
@@ -479,7 +480,9 @@ where
     delegate! {
         fn algebra(&self) -> &Self::Algebra;
 
-        fn storage_fee_per_write(&self, key: &StateKey, op: &WriteOp) -> Fee;
+        fn storage_fee_for_state_slot(&self, op: &WriteOp) -> Fee;
+
+        fn storage_fee_for_state_bytes(&self, key: &StateKey, op: &WriteOp) -> Fee;
 
         fn storage_fee_per_event(&self, event: &ContractEvent) -> Fee;
 
@@ -511,10 +514,9 @@ where
         res
     }
 
-    fn charge_storage_fee_for_all<'a>(
+    fn process_storage_fee_for_all(
         &mut self,
-        write_ops: impl IntoIterator<Item = (&'a StateKey, &'a WriteOp)>,
-        events: impl IntoIterator<Item = &'a ContractEvent>,
+        change_set: &mut VMChangeSet,
         txn_size: NumBytes,
         gas_unit_price: FeePerGasUnit,
     ) -> VMResult<()> {
@@ -533,8 +535,12 @@ where
         // Writes
         let mut write_fee = Fee::new(0);
         let mut write_set_storage = vec![];
-        for (key, op) in write_ops.into_iter() {
-            let fee = self.storage_fee_per_write(key, op);
+        for (key, op) in change_set.write_set_iter_mut() {
+            let slot_fee = self.storage_fee_for_state_slot(op);
+            let bytes_fee = self.storage_fee_for_state_bytes(key, op);
+            Self::maybe_record_storage_deposit(op, slot_fee);
+
+            let fee = slot_fee + bytes_fee;
             write_set_storage.push(WriteStorage {
                 key: key.clone(),
                 op_type: write_op_type(op),
@@ -546,7 +552,7 @@ where
         // Events
         let mut event_fee = Fee::new(0);
         let mut event_fees = vec![];
-        for event in events {
+        for event in change_set.events().iter() {
             let fee = self.storage_fee_per_event(event);
             event_fees.push(EventStorage {
                 ty: event.type_tag().clone(),

aptos-move/aptos-gas-schedule/src/gas_schedule/transaction.rs

@@ -189,27 +189,31 @@ impl TransactionGasParameters {
         }
     }
 
-    /// New formula to charge storage fee for a write, measured in APT.
-    pub fn storage_fee_per_write(&self, key: &StateKey, op: &WriteOp) -> Fee {
+    pub fn storage_fee_for_slot(&self, op: &WriteOp) -> Fee {
         use WriteOp::*;
 
-        let excess_fee = |key: &StateKey, data: &[u8]| -> Fee {
-            let size = NumBytes::new(key.size() as u64) + NumBytes::new(data.len() as u64);
-            match size.checked_sub(self.free_write_bytes_quota) {
-                Some(excess) => excess * self.storage_fee_per_excess_state_byte,
-                None => 0.into(),
-            }
-        };
-
         match op {
-            Creation(data) | CreationWithMetadata { data, .. } => {
-                self.storage_fee_per_state_slot_create * NumSlots::new(1) + excess_fee(key, data)
+            Creation(..) | CreationWithMetadata { .. } => {
+                self.storage_fee_per_state_slot_create * NumSlots::new(1)
             },
-            Modification(data) | ModificationWithMetadata { data, .. } => excess_fee(key, data),
-            Deletion | DeletionWithMetadata { .. } => 0.into(),
+            Modification(..)
+            | ModificationWithMetadata { .. }
+            | Deletion
+            | DeletionWithMetadata { .. } => 0.into(),
         }
     }
 
+    pub fn storage_fee_for_bytes(&self, key: &StateKey, op: &WriteOp) -> Fee {
+        if let Some(data) = op.bytes() {
+            let size = NumBytes::new(key.size() as u64) + NumBytes::new(data.len() as u64);
+            if let Some(excess) = size.checked_sub(self.free_write_bytes_quota) {
+                return excess * self.storage_fee_per_excess_state_byte;
+            }
+        }
+
+        0.into()
+    }
+
     /// New formula to charge storage fee for an event, measured in APT.
     pub fn storage_fee_per_event(&self, event: &ContractEvent) -> Fee {
         NumBytes::new(event.size() as u64) * self.storage_fee_per_event_byte

aptos-move/aptos-memory-usage-tracker/src/lib.rs

@@ -463,7 +463,9 @@ where
     delegate! {
         fn algebra(&self) -> &Self::Algebra;
 
-        fn storage_fee_per_write(&self, key: &StateKey, op: &WriteOp) -> Fee;
+        fn storage_fee_for_state_slot(&self, op: &WriteOp) -> Fee;
+
+        fn storage_fee_for_state_bytes(&self, key: &StateKey, op: &WriteOp) -> Fee;
 
         fn storage_fee_per_event(&self, event: &ContractEvent) -> Fee;
 

aptos-move/aptos-vm-types/src/change_set.rs

@@ -159,6 +159,13 @@ impl VMChangeSet {
             .chain(self.aggregator_write_set.iter())
     }
 
+    pub fn write_set_iter_mut(&mut self) -> impl Iterator<Item = (&StateKey, &mut WriteOp)> {
+        self.resource_write_set
+            .iter_mut()
+            .chain(self.module_write_set.iter_mut())
+            .chain(self.aggregator_write_set.iter_mut())
+    }
+
     pub fn resource_write_set(&self) -> &BTreeMap<StateKey, WriteOp> {
         &self.resource_write_set
     }

aptos-move/aptos-vm/src/aptos_vm.rs

@@ -491,15 +491,14 @@ impl AptosVM {
         change_set_configs: &ChangeSetConfigs,
         txn_data: &TransactionMetadata,
     ) -> Result<RespawnedSession<'r, 'l>, VMStatus> {
-        let change_set = session.finish(&mut (), change_set_configs)?;
+        let mut change_set = session.finish(&mut (), change_set_configs)?;
 
         for (key, op) in change_set.write_set_iter() {
             gas_meter.charge_io_gas_for_write(key, op)?;
         }
 
-        gas_meter.charge_storage_fee_for_all(
-            change_set.write_set_iter(),
-            change_set.events(),
+        gas_meter.process_storage_fee_for_all(
+            &mut change_set,
             txn_data.transaction_size,
             txn_data.gas_unit_price,
         )?;

aptos-move/aptos-vm/src/aptos_vm_impl.rs

@@ -17,7 +17,7 @@ use aptos_gas_schedule::{
     AptosGasParameters, FromOnChainGasSchedule, MiscGasParameters, NativeGasParameters,
 };
 use aptos_logger::{enabled, prelude::*, Level};
-use aptos_state_view::StateView;
+use aptos_state_view::{StateView, StateViewId};
 use aptos_types::{
     account_config::CORE_CODE_ADDRESS,
     chain_id::ChainId,
@@ -642,11 +642,9 @@ impl<'a> AptosVMInternals<'a> {
     }
 
     /// Returns the internal gas schedule if it has been loaded, or an error if it hasn't.
-    pub fn gas_params(
-        self,
-        log_context: &AdapterLogSchema,
-    ) -> Result<&'a AptosGasParameters, VMStatus> {
-        self.0.get_gas_parameters(log_context)
+    pub fn gas_params(self) -> Result<&'a AptosGasParameters, VMStatus> {
+        let log_context = AdapterLogSchema::new(StateViewId::Miscellaneous, 0);
+        self.0.get_gas_parameters(&log_context)
     }
 
     /// Returns the version of Move Runtime.

aptos-move/aptos-vm/src/move_vm_ext/session.rs

@@ -306,6 +306,8 @@ impl<'r, 'l> SessionExt<'r, 'l> {
         let mut new_slot_metadata: Option<StateValueMetadata> = None;
         if is_storage_slot_metadata_enabled {
             if let Some(current_time) = current_time {
+                // The deposit on the metadata is a placeholder (0), it will be updated later when
+                // storage fee is charged.
                 new_slot_metadata = Some(StateValueMetadata::new(0, current_time));
             }
         }

aptos-move/e2e-move-tests/src/harness.rs

@@ -30,6 +30,7 @@ use aptos_types::{
         TransactionPayload, TransactionStatus,
     },
 };
+use aptos_vm::AptosVM;
 use move_core_types::{
     language_storage::{StructTag, TypeTag},
     move_resource::MoveStructType,
@@ -603,6 +604,10 @@ impl MoveHarness {
         );
     }
 
+    pub fn new_vm(&self) -> AptosVM {
+        AptosVM::new(self.executor.data_store())
+    }
+
     pub fn set_default_gas_unit_price(&mut self, gas_unit_price: u64) {
         self.default_gas_unit_price = gas_unit_price;
     }

aptos-move/e2e-move-tests/src/tests/state_metadata.rs

@@ -45,10 +45,22 @@ fn test_metadata_tracking() {
         &account1,
         aptos_cached_packages::aptos_stdlib::aptos_account_transfer(address3, 100),
     );
+
+    let slot_fee = harness
+        .new_vm()
+        .internals()
+        .gas_params()
+        .unwrap()
+        .vm
+        .txn
+        .storage_fee_per_state_slot_create
+        .into();
+    assert!(slot_fee > 0);
+
     // Observe that metadata is tracked for address3 resources
     assert_eq!(
         harness.read_resource_metadata(&address3, coin_store.clone()),
-        Some(Some(StateValueMetadata::new(0, &timestamp))),
+        Some(Some(StateValueMetadata::new(slot_fee, &timestamp,))),
     );
 
     // Bump the timestamp and modify the resources, observe that metadata doesn't change.
@@ -67,6 +79,6 @@ fn test_metadata_tracking() {
     );
     assert_eq!(
         harness.read_resource_metadata(&address3, coin_store),
-        Some(Some(StateValueMetadata::new(0, &timestamp))),
+        Some(Some(StateValueMetadata::new(slot_fee, &timestamp))),
     );
 }

types/src/state_store/state_value.rs

@@ -42,6 +42,12 @@ impl StateValueMetadata {
             creation_time_usecs: creation_time_usecs.microseconds,
         }
     }
+
+    pub fn set_deposit(&mut self, amount: u64) {
+        match self {
+            StateValueMetadata::V0 { deposit, .. } => *deposit = amount,
+        }
+    }
 }
 
 #[derive(Clone, Debug, CryptoHasher)]