Refund transient VM linear-memory budget at VM teardown

The memory budget has always measured cumulative memory -- the sum of all
allocations -- and does not credit back transient memory that is later
freed. For most invocations cumulative memory stays close to peak live
memory, so this is a fine proxy. But each VM instantiation charges the
contract's WASM linear memory (~1 MiB with default toolchain settings),
which is freed at VM teardown; under deep or wide cross-contract calling,
many VMs are created and destroyed and cumulative memory diverges sharply
from peak, significantly overcharging for memory never held at once.

Track the mem_bytes charged for each VM store's linear memory in its
VmStoreData (initial allocation + runtime grows), and refund it to the
budget's memory dimension when the store -- and thus its linear memory --
is freed at Drop. Scope is VM linear memory only; CPU is unaffected. The
refund credits only the real budget total, never the shadow budget.

Adds Budget::refund_mem and BudgetDimension::refund, and updates the
affected metering assertions.

Author: jayz22

soroban-env-host/src/budget.rs

@@ -1332,6 +1332,13 @@ impl Budget {
             .get_memory_cost(ty, 1, input)
     }
 
+    /// Refunds `amount` to the memory dimension (transient linear memory freed
+    /// at VM teardown). Saturating; never refunds more than was charged.
+    pub(crate) fn refund_mem(&self, amount: u64) -> Result<(), HostError> {
+        self.0.try_borrow_mut_or_err()?.mem_bytes.refund(amount);
+        Ok(())
+    }
+
     /// Runs a user provided closure in shadow mode -- all metering is done
     /// through the shadow budget.
     ///

soroban-env-host/src/budget/dimension.rs

@@ -140,6 +140,18 @@ impl BudgetDimension {
         self.shadow_total_count = 0;
     }
 
+    /// Credits `amount` back to this dimension (used to refund transient
+    /// allocations, e.g. a VM's linear memory freed at teardown). Saturating,
+    /// so it can never underflow or refund more than was charged.
+    ///
+    /// Only the real total is refunded: VM linear memory is charged to the real
+    /// budget (never the shadow budget), so decrementing `shadow_total_count`
+    /// here would corrupt shadow accounting (e.g. extend the debug-logging
+    /// budget).
+    pub(crate) fn refund(&mut self, amount: u64) {
+        self.total_count = self.total_count.saturating_sub(amount);
+    }
+
     pub(crate) fn check_budget_limit(&self, is_shadow: IsShadowMode) -> Result<(), HostError> {
         let over_limit = if is_shadow.0 {
             self.shadow_total_count > self.shadow_limit

soroban-env-host/src/budget/wasmi_helper.rs

@@ -25,8 +25,10 @@ pub(crate) const WASMI_LIMITS_CONFIG: WasmiLimits = WasmiLimits {
 };
 
 // The `ResourceLimiter` is implemented for each VM store's data (rather than
-// for `Host`, which is shared by all stores): wasmi invokes the limiter on the
-// growing store's own data.
+// for `Host`, which is shared by all stores) so that the linear-memory
+// charges accumulate structurally in the store that owns the growing memory:
+// wasmi invokes the limiter on the growing store's own data, both during
+// instantiation and for runtime `memory.grow`.
 impl ResourceLimiter for VmStoreData {
     fn memory_growing(
         &mut self,
@@ -60,11 +62,36 @@ impl ResourceLimiter for VmStoreData {
                     .map_err(|_| errors::MemoryError::OutOfBoundsGrowth)?;
             }
 
+            // VM linear memory is charged only to the real budget and refunded
+            // only from the real budget (see `Vm::execute_with_mem_refund`);
+            // guest memory only grows during instantiation or execution,
+            // neither of which runs in shadow mode. This is just a defensive
+            // check.
+            if self
+                .host
+                .as_budget()
+                .is_in_shadow_mode()
+                .map_err(|_| errors::MemoryError::OutOfBoundsGrowth)?
+            {
+                return Err(errors::MemoryError::OutOfBoundsGrowth);
+            }
+
             self.host
                 .as_budget()
                 .charge(ContractCostType::MemAlloc, Some(delta))
-                .map(|_| true)
-                .map_err(|_| errors::MemoryError::OutOfBoundsGrowth)
+                .map_err(|_| errors::MemoryError::OutOfBoundsGrowth)?;
+
+            // Accumulate the charge in this store's data so that the VM's
+            // usage scope can refund it when the VM is torn down (see
+            // `Vm::execute_with_mem_refund`).
+            let charged = self
+                .host
+                .as_budget()
+                .get_memory_cost(ContractCostType::MemAlloc, Some(delta))
+                .map_err(|_| errors::MemoryError::OutOfBoundsGrowth)?;
+            self.linear_mem_charged = self.linear_mem_charged.saturating_add(charged);
+
+            Ok(true)
         } else {
             Err(errors::MemoryError::OutOfBoundsGrowth)
         }

soroban-env-host/src/host/frame.rs

@@ -753,17 +753,19 @@ impl Host {
             }
         };
         let vm = self.instantiate_vm(id, wasm_hash)?;
-        let relative_objects = Vec::new();
-        self.with_frame(
-            Frame::ContractVM {
-                vm: Rc::clone(&vm),
-                fn_name: *func,
-                args: args_vec,
-                instance,
-                relative_objects,
-            },
-            || vm.invoke_function_raw(self, func, args, treat_missing_function_as_noop),
-        )
+        vm.execute_with_mem_refund(self, |vm| {
+            let relative_objects = Vec::new();
+            self.with_frame(
+                Frame::ContractVM {
+                    vm: Rc::clone(vm),
+                    fn_name: *func,
+                    args: args_vec,
+                    instance,
+                    relative_objects,
+                },
+                || vm.invoke_function_raw(self, func, args, treat_missing_function_as_noop),
+            )
+        })
     }
 
     fn instantiate_vm(&self, id: &ContractId, wasm_hash: &Hash) -> Result<Rc<Vm>, HostError> {
@@ -891,12 +893,12 @@ impl Host {
         match &instance.executable {
             ContractExecutable::Wasm(wasm_hash) => {
                 let vm = self.instantiate_vm(contract_id, wasm_hash)?;
-                Ok(vm.module.proto_version)
+                vm.execute_with_mem_refund(self, |vm| Ok(vm.module.proto_version))
             }
             ContractExecutable::ExternalRef(external_ref) => {
                 let wasm_hash = self.resolve_external_ref_wasm_hash(external_ref)?;
                 let vm = self.instantiate_vm(contract_id, &wasm_hash)?;
-                Ok(vm.module.proto_version)
+                vm.execute_with_mem_refund(self, |vm| Ok(vm.module.proto_version))
             }
             ContractExecutable::StellarAsset => self.get_ledger_protocol_version(),
         }

soroban-env-host/src/host/invocation_metering.rs

@@ -1092,7 +1092,7 @@ mod test {
         expect![[r#"
             InvocationResources {
                 instructions: 4146007,
-                mem_bytes: 2861644,
+                mem_bytes: 764460,
                 disk_read_entries: 0,
                 memory_read_entries: 2,
                 write_entries: 2,
@@ -1110,7 +1110,7 @@ mod test {
                 invocation: CreateContractEntryPoint,
                 resources: SubInvocationResources {
                     instructions: 4146007,
-                    mem_bytes: 2861644,
+                    mem_bytes: 764460,
                     disk_read_entries: 0,
                     memory_read_entries: 2,
                     write_entries: 2,
@@ -1146,7 +1146,7 @@ mod test {
         expect![[r#"
             InvocationResources {
                 instructions: 316541,
-                mem_bytes: 1135019,
+                mem_bytes: 86427,
                 disk_read_entries: 0,
                 memory_read_entries: 3,
                 write_entries: 0,
@@ -1173,7 +1173,7 @@ mod test {
                 ),
                 resources: SubInvocationResources {
                     instructions: 316541,
-                    mem_bytes: 1135019,
+                    mem_bytes: 86427,
                     disk_read_entries: 0,
                     memory_read_entries: 3,
                     write_entries: 0,
@@ -1207,7 +1207,7 @@ mod test {
         expect![[r#"
             InvocationResources {
                 instructions: 320592,
-                mem_bytes: 1135562,
+                mem_bytes: 86970,
                 disk_read_entries: 0,
                 memory_read_entries: 3,
                 write_entries: 1,
@@ -1233,7 +1233,7 @@ mod test {
         expect![[r#"
             InvocationResources {
                 instructions: 315840,
-                mem_bytes: 1134867,
+                mem_bytes: 86275,
                 disk_read_entries: 0,
                 memory_read_entries: 3,
                 write_entries: 0,
@@ -1259,7 +1259,7 @@ mod test {
         expect![[r#"
             InvocationResources {
                 instructions: 322503,
-                mem_bytes: 1135918,
+                mem_bytes: 87326,
                 disk_read_entries: 0,
                 memory_read_entries: 3,
                 write_entries: 1,
@@ -1285,7 +1285,7 @@ mod test {
         expect![[r#"
             InvocationResources {
                 instructions: 316380,
-                mem_bytes: 1134935,
+                mem_bytes: 86343,
                 disk_read_entries: 0,
                 memory_read_entries: 3,
                 write_entries: 0,
@@ -1311,7 +1311,7 @@ mod test {
         expect![[r#"
             InvocationResources {
                 instructions: 317605,
-                mem_bytes: 1135287,
+                mem_bytes: 86695,
                 disk_read_entries: 0,
                 memory_read_entries: 3,
                 write_entries: 0,
@@ -1337,7 +1337,7 @@ mod test {
         expect![[r#"
             InvocationResources {
                 instructions: 318007,
-                mem_bytes: 1135287,
+                mem_bytes: 86695,
                 disk_read_entries: 0,
                 memory_read_entries: 3,
                 write_entries: 0,
@@ -1363,7 +1363,7 @@ mod test {
         expect![[r#"
             InvocationResources {
                 instructions: 317444,
-                mem_bytes: 1135355,
+                mem_bytes: 86763,
                 disk_read_entries: 0,
                 memory_read_entries: 3,
                 write_entries: 0,
@@ -1396,7 +1396,7 @@ mod test {
         expect![[r#"
             InvocationResources {
                 instructions: 320615,
-                mem_bytes: 1135822,
+                mem_bytes: 87230,
                 disk_read_entries: 2,
                 memory_read_entries: 1,
                 write_entries: 2,
@@ -1428,7 +1428,7 @@ mod test {
         expect![[r#"
             InvocationResources {
                 instructions: 323152,
-                mem_bytes: 1136269,
+                mem_bytes: 87677,
                 disk_read_entries: 3,
                 memory_read_entries: 0,
                 write_entries: 3,
@@ -1874,7 +1874,7 @@ mod test {
                             ),
                             data: String(
                                 ScString(
-                                    StringM(invocation resource limits are exceeded: instructions: 1788420 > 10, memory bytes: 1163561 > 20, contract events size bytes: 800 > 50, contract data key 'ContractData(LedgerKeyContractData { contract: Contract(ContractId(Hash(2e0ff7a55065f3a896723a964c3d9862a4722bfc77229fe4875f390ef2a0027e))), key: LedgerKeyContractInstance, durability: Persistent })' size: 48 > 25, contract data entry with key 'ContractData(LedgerKeyContractData { contract: Contract(ContractId(Hash(2e0ff7a55065f3a896723a964c3d9862a4722bfc77229fe4875f390ef2a0027e))), key: LedgerKeyContractInstance, durability: Persistent })' size: 104 > 35, contract code entry with key 'ContractCode(LedgerKeyContractCode { hash: Hash(1a2ee5a89dd2162a20e716b3e2de70a2d662480a9565350378661871bcf8e398) })' size: 1840 > 45),
+                                    StringM(invocation resource limits are exceeded: instructions: 1788420 > 10, memory bytes: 114969 > 20, contract events size bytes: 800 > 50, contract data key 'ContractData(LedgerKeyContractData { contract: Contract(ContractId(Hash(2e0ff7a55065f3a896723a964c3d9862a4722bfc77229fe4875f390ef2a0027e))), key: LedgerKeyContractInstance, durability: Persistent })' size: 48 > 25, contract data entry with key 'ContractData(LedgerKeyContractData { contract: Contract(ContractId(Hash(2e0ff7a55065f3a896723a964c3d9862a4722bfc77229fe4875f390ef2a0027e))), key: LedgerKeyContractInstance, durability: Persistent })' size: 104 > 35, contract code entry with key 'ContractCode(LedgerKeyContractCode { hash: Hash(1a2ee5a89dd2162a20e716b3e2de70a2d662480a9565350378661871bcf8e398) })' size: 1840 > 45),
                                 ),
                             ),
                         },
@@ -1927,7 +1927,7 @@ mod test {
                             ),
                             data: String(
                                 ScString(
-                                    StringM(invocation resource limits are exceeded: instructions: 320868 > 10, memory bytes: 1135974 > 20, total footprint ledger entries: 5 > 4, disk read ledger entries: 2 > 1, disk read bytes: 3132 > 30, write ledger entries: 2 > 0, write bytes: 3132 > 40, contract data key 'ContractData(LedgerKeyContractData { contract: Contract(ContractId(Hash(ba863dea340f907c97f640ecbe669125e9f8f3b63ed1f4ed0f30073b869e5441))), key: Symbol(ScSymbol(StringM(key_1))), durability: Persistent })' size: 60 > 25, contract data key 'ContractData(LedgerKeyContractData { contract: Contract(ContractId(Hash(ba863dea340f907c97f640ecbe669125e9f8f3b63ed1f4ed0f30073b869e5441))), key: LedgerKeyContractInstance, durability: Persistent })' size: 48 > 25, contract data entry with key 'ContractData(LedgerKeyContractData { contract: Contract(ContractId(Hash(ba863dea340f907c97f640ecbe669125e9f8f3b63ed1f4ed0f30073b869e5441))), key: LedgerKeyContractInstance, durability: Persistent })' size: 104 > 35, contract code entry with key 'ContractCode(LedgerKeyContractCode { hash: Hash(fc644715caaead746e6145f4331ff75c427c965c20d2995a9942b01247515962) })' size: 3028 > 45),
+                                    StringM(invocation resource limits are exceeded: instructions: 320868 > 10, memory bytes: 87382 > 20, total footprint ledger entries: 5 > 4, disk read ledger entries: 2 > 1, disk read bytes: 3132 > 30, write ledger entries: 2 > 0, write bytes: 3132 > 40, contract data key 'ContractData(LedgerKeyContractData { contract: Contract(ContractId(Hash(ba863dea340f907c97f640ecbe669125e9f8f3b63ed1f4ed0f30073b869e5441))), key: Symbol(ScSymbol(StringM(key_1))), durability: Persistent })' size: 60 > 25, contract data key 'ContractData(LedgerKeyContractData { contract: Contract(ContractId(Hash(ba863dea340f907c97f640ecbe669125e9f8f3b63ed1f4ed0f30073b869e5441))), key: LedgerKeyContractInstance, durability: Persistent })' size: 48 > 25, contract data entry with key 'ContractData(LedgerKeyContractData { contract: Contract(ContractId(Hash(ba863dea340f907c97f640ecbe669125e9f8f3b63ed1f4ed0f30073b869e5441))), key: LedgerKeyContractInstance, durability: Persistent })' size: 104 > 35, contract code entry with key 'ContractCode(LedgerKeyContractCode { hash: Hash(fc644715caaead746e6145f4331ff75c427c965c20d2995a9942b01247515962) })' size: 3028 > 45),
                                 ),
                             ),
                         },

soroban-env-host/src/host/lifecycle.rs

@@ -313,22 +313,26 @@ impl Host {
             // native test contracts behave like wasm. They will never be
             // instantiated, this is just to exercise their storage logic.
         } else {
-            let _check_vm = Vm::new(
+            let check_vm = Vm::new(
                 self,
                 ContractId(Hash(hash_bytes.metered_clone(self)?)),
                 wasm_bytes_m.as_slice(),
             )?;
-            // At this point we do a secondary parse on what we've checked to be a valid
-            // module in order to extract a refined cost model, which we'll store in the
-            // code entry's ext field, for future parsing and instantiations.
-            _check_vm.module.cost_inputs.charge_for_parsing(self)?;
-            ext = crate::xdr::ContractCodeEntryExt::V1(crate::xdr::ContractCodeEntryV1 {
-                ext: ExtensionPoint::V0,
-                cost_inputs: crate::vm::ParsedModule::extract_refined_contract_cost_inputs(
-                    self,
-                    wasm_bytes_m.as_slice(),
-                )?,
-            });
+            ext = check_vm.execute_with_mem_refund(self, |check_vm| {
+                // At this point we do a secondary parse on what we've checked to be a valid
+                // module in order to extract a refined cost model, which we'll store in the
+                // code entry's ext field, for future parsing and instantiations.
+                check_vm.module.cost_inputs.charge_for_parsing(self)?;
+                Ok(crate::xdr::ContractCodeEntryExt::V1(
+                    crate::xdr::ContractCodeEntryV1 {
+                        ext: ExtensionPoint::V0,
+                        cost_inputs: crate::vm::ParsedModule::extract_refined_contract_cost_inputs(
+                            self,
+                            wasm_bytes_m.as_slice(),
+                        )?,
+                    },
+                ))
+            })?;
         }
 
         let hash_obj = self.add_host_object(self.scbytes_from_slice(hash_bytes.as_slice())?)?;