Fix simulation for auto-restored contract code. (#1590)

### What

Fix simulation for auto-restored contract code.

When a contract code entry is auto-restored, it will likely miss the
module cache (if it's been in the hot archive), so we should charge the
recompilation to the budget (in the same way as we already do for the
entry that has been just created via upload).

### Why

Improving simulation accuracy.

### Known limitations

N/A

Author: dmkozh

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

@@ -720,12 +720,13 @@ impl Host {
         //     - User invoked us with bad input, eg. calling a
         //       contract that wasn't provided in footprint/storage.
         //
-        //     - User uploaded the wasm _in this transaction_ so we
-        //       didn't cache it when starting the transaction (and
-        //       couldn't add it: additions use a shared wasmi engine
-        //       that owns all the cache entries, and that engine is
-        //       locked while we're running; uploads use a throwaway
-        //       engine for validation purposes).
+        //     - User uploaded the wasm in this ledger so we didn't
+        //       cache it when starting the ledger (and couldn't add
+        //       it: the module cache and the wasmi engine used to
+        //       build modules are both locked and shared across
+        //       threads during execution, we don't want to perturb
+        //       it even if we could; uploads use a throwaway engine
+        //       for validation purposes).
         //
         //   3. Even more pathological: the module cache was built,
         //      and contained the module, but someone _removed_ the
@@ -742,23 +743,34 @@ impl Host {
         // with a storage error.
 
         #[cfg(any(test, feature = "recording_mode"))]
-        // In recording mode: if a contract was present in the initial snapshot image, it is part of
-        // the set of contracts that would have been built into a module cache in enforcing mode;
-        // we want to suppress the cost of parsing those (simulating them as module cache hits).
-        //
-        // If a contract is _not_ in the initial snapshot image, it's because someone just uploaded
-        // it during execution. Those would be cache misses in enforcing mode, and so it is right to
-        // continue to charge for them as such (charging the parse cost on each call) in recording.
+        // In recording mode:
+        //   - We have no _real_ module cache.
+        //   - We have a choice of whether simulate a cache-hit.
+        //     - We will "simulate a hit" by doing a fresh parse but charging it
+        //       to the shadow budget, not the real budget.
+        //   - We _want_ to simulate a miss any time _would_ be a corresponding
+        //     (charged-for) miss in enforcing mode, because otherwise we're
+        //     under-charging and the tx we're simulating will fail in execution.
+        //   - One case we know for sure will cause a miss: if the module
+        //     literally isn't in the snapshot at all. This happens when someone
+        //     uploads a contract and tries running it in the same ledger.
+        //   - Other cases we're _not sure_: a module might be expired and evicted
+        //     (thus removed from module cache) between simulation time and
+        //     enforcement time.
+        //   - But we can and do make an approximation here:
+        //     - If the module is _expired_ we assume it's on its way to eviction
+        //       soon and simulate a miss, risking overcharging.
+        //     - If the module is _not expired_ we assume it'll be survive until
+        //       execution, simulate a hit, and risk undercharging.
         if self.in_storage_recording_mode()? {
             if let Some((parsed_module, wasmi_linker)) =
                 self.budget_ref().with_observable_shadow_mode(|| {
                     use crate::vm::ParsedModule;
                     let wasm_key = self.contract_code_ledger_key(wasm_hash)?;
-                    if self
-                        .try_borrow_storage()?
-                        .get_snapshot_value(self, &wasm_key)?
-                        .is_some()
-                    {
+                    let is_key_live_in_snapshot = self
+                        .try_borrow_storage_mut()?
+                        .is_key_live_in_snapshot(self, &wasm_key)?;
+                    if is_key_live_in_snapshot {
                         let (code, costs) = self.retrieve_wasm_from_storage(&wasm_hash)?;
                         let parsed_module =
                             ParsedModule::new_with_isolated_engine(self, code.as_slice(), costs)?;

soroban-env-host/src/storage.rs

@@ -499,13 +499,30 @@ impl Storage {
     }
 
     #[cfg(any(test, feature = "recording_mode"))]
-    pub(crate) fn get_snapshot_value(
+    /// Returns `true` if the key exists in the snapshot and is live w.r.t
+    /// the current ledger sequence.
+    pub(crate) fn is_key_live_in_snapshot(
         &self,
         host: &Host,
         key: &Rc<LedgerKey>,
-    ) -> Result<Option<EntryWithLiveUntil>, HostError> {
+    ) -> Result<bool, HostError> {
         match &self.mode {
-            FootprintMode::Recording(snapshot) => snapshot.get(key),
+            FootprintMode::Recording(snapshot) => {
+                let snapshot_value = snapshot.get(key)?;
+                if let Some((_, live_until_ledger)) = snapshot_value {
+                    if let Some(live_until_ledger) = live_until_ledger {
+                        let current_ledger_sequence =
+                            host.with_ledger_info(|li| Ok(li.sequence_number))?;
+                        Ok(live_until_ledger >= current_ledger_sequence)
+                    } else {
+                        // Non-Soroban entries are always live.
+                        Ok(true)
+                    }
+                } else {
+                    // Key is not in the snapshot.
+                    Ok(false)
+                }
+            }
             FootprintMode::Enforcing => Err(host.err(
                 ScErrorType::Storage,
                 ScErrorCode::InternalError,

soroban-env-host/src/test/e2e_tests.rs

@@ -44,6 +44,7 @@ use soroban_test_wasms::{
     NO_ARGUMENT_CONSTRUCTOR_TEST_CONTRACT_P22, SIMPLE_ACCOUNT_CONTRACT, SUM_I32,
     UPDATEABLE_CONTRACT,
 };
+use std::collections::HashSet;
 use std::rc::Rc;
 
 // It's tricky to get exactly the same instruction consumption
@@ -307,8 +308,15 @@ fn invoke_host_function_helper_with_restored_entries(
                 .unwrap()
         })
         .collect();
-
-    let module_cache = build_module_cache_for_entries(ledger_info, ledger_entries_with_ttl)?;
+    let mut restored_contracts = HashSet::new();
+    for restored_id in restored_entry_ids {
+        let key = &resources.footprint.read_write[*restored_id as usize];
+        if let LedgerKey::ContractCode(code) = key {
+            restored_contracts.insert(code.hash.clone());
+        }
+    }
+    let module_cache =
+        build_module_cache_for_entries(ledger_info, ledger_entries_with_ttl, &restored_contracts)?;
 
     let budget = Budget::default();
     budget
@@ -372,12 +380,18 @@ fn invoke_host_function_helper(
 fn build_module_cache_for_entries(
     ledger_info: &LedgerInfo,
     ledger_entries_with_ttl: Vec<(LedgerEntry, Option<u32>)>,
+    restored_contracts: &HashSet<Hash>,
 ) -> Result<ModuleCache, HostError> {
     let ctx = E2eTestCompilationContext::new()?;
     let cache = ModuleCache::new(&ctx)?;
     for (e, _) in ledger_entries_with_ttl.iter() {
         if let LedgerEntryData::ContractCode(cd) = &e.data {
             let contract_id = Hash(sha256_hash_from_bytes_raw(&cd.code, ctx.as_budget())?);
+            // Restored contracts are not yet in the module cache and need to be
+            // compiled during execution.
+            if restored_contracts.contains(&contract_id) {
+                continue;
+            }
             let code_cost_inputs = match &cd.ext {
                 ContractCodeEntryExt::V0 => VersionedContractCodeCostInputs::V0 {
                     wasm_bytes: cd.code.len(),
@@ -442,7 +456,7 @@ fn invoke_host_function_using_simulation_with_signers(
     host_fn: &HostFunction,
     source_account: &AccountId,
     ledger_info: &LedgerInfo,
-    ledger_entries_with_ttl: Vec<(LedgerEntry, Option<u32>)>,
+    mut ledger_entries_with_ttl: Vec<(LedgerEntry, Option<u32>)>,
     prng_seed: &[u8; 32],
     dummy_host: &Host,
     signers: &Vec<TestSigner>,
@@ -543,7 +557,21 @@ fn invoke_host_function_using_simulation_with_signers(
         * (1.0 + RECORDING_MODE_INSTRUCTIONS_RANGE))
         as u32;
 
-    let enforcing_result = invoke_host_function_helper(
+    let restored_live_until_ledger =
+        ledger_info.sequence_number + ledger_info.min_persistent_entry_ttl - 1;
+    for restored_id in &recording_result.restored_rw_entry_ids {
+        let key = recording_result.resources.footprint.read_write[*restored_id as usize].clone();
+        // Update TTL of the entry corresponding to the key
+        if let Some((_, ttl)) = ledger_entries_with_ttl
+            .iter_mut()
+            .find(|(le, _)| ledger_entry_to_ledger_key(le, &Budget::default()).unwrap() == key)
+        {
+            *ttl = Some(restored_live_until_ledger);
+        } else {
+            panic!("restored entry not found in the original entries");
+        }
+    }
+    let enforcing_result = invoke_host_function_helper_with_restored_entries(
         enable_diagnostics,
         host_fn,
         &recording_result.resources,
@@ -552,6 +580,7 @@ fn invoke_host_function_using_simulation_with_signers(
         ledger_info,
         ledger_entries_with_ttl,
         prng_seed,
+        recording_result.restored_rw_entry_ids.as_slice(),
     )?;
 
     assert_eq!(
@@ -2246,26 +2275,27 @@ fn test_auto_restore_with_extension_in_recording_mode() {
         &val,
         ContractDataDurability::Persistent,
     );
+    let le_with_ttl = vec![
+        (
+            cd.wasm_entry.clone(),
+            Some(ledger_info.sequence_number - 100_000),
+        ),
+        (
+            cd.contract_entry.clone(),
+            Some(ledger_info.sequence_number - 1),
+        ),
+        (
+            persistent_data_entry.clone(),
+            Some(ledger_info.sequence_number - 10),
+        ),
+    ];
     let res = invoke_host_function_recording_helper(
         true,
         &host_fn,
         &cd.deployer,
         RecordingInvocationAuthMode::Recording(true),
         &ledger_info,
-        vec![
-            (
-                cd.wasm_entry.clone(),
-                Some(ledger_info.sequence_number - 100_000),
-            ),
-            (
-                cd.contract_entry.clone(),
-                Some(ledger_info.sequence_number - 1),
-            ),
-            (
-                persistent_data_entry.clone(),
-                Some(ledger_info.sequence_number - 10),
-            ),
-        ],
+        le_with_ttl.clone(),
         &prng_seed(),
         None,
     )
@@ -2345,11 +2375,21 @@ fn test_auto_restore_with_extension_in_recording_mode() {
                 .try_into()
                 .unwrap(),
             },
-            instructions: 1027906,
+            instructions: 1562621,
             disk_read_bytes: data_entry_size + wasm_entry_size + instance_entry_size,
             write_bytes: data_entry_size + wasm_entry_size + instance_entry_size,
         }
     );
+
+    let _ = invoke_host_function_using_simulation(
+        true,
+        &host_fn,
+        &cd.deployer,
+        &ledger_info,
+        le_with_ttl,
+        &prng_seed(),
+    )
+    .unwrap();
 }
 
 #[test]
@@ -2376,26 +2416,27 @@ fn test_auto_restore_with_overwrite_in_recording_mode() {
         &key,
         ContractDataDurability::Persistent,
     );
+    let le_with_ttl = vec![
+        (
+            cd.wasm_entry.clone(),
+            Some(ledger_info.sequence_number + 100_000),
+        ),
+        (
+            cd.contract_entry.clone(),
+            Some(ledger_info.sequence_number - 100),
+        ),
+        (
+            persistent_data_entry.clone(),
+            Some(ledger_info.sequence_number - 1),
+        ),
+    ];
     let res = invoke_host_function_recording_helper(
         true,
         &host_fn,
         &cd.deployer,
         RecordingInvocationAuthMode::Recording(true),
         &ledger_info,
-        vec![
-            (
-                cd.wasm_entry.clone(),
-                Some(ledger_info.sequence_number + 100_000),
-            ),
-            (
-                cd.contract_entry.clone(),
-                Some(ledger_info.sequence_number - 100),
-            ),
-            (
-                persistent_data_entry.clone(),
-                Some(ledger_info.sequence_number - 1),
-            ),
-        ],
+        le_with_ttl.clone(),
         &prng_seed(),
         None,
     )
@@ -2477,6 +2518,16 @@ fn test_auto_restore_with_overwrite_in_recording_mode() {
             write_bytes: data_entry_size + instance_entry_size,
         }
     );
+
+    let _ = invoke_host_function_using_simulation(
+        true,
+        &host_fn,
+        &cd.deployer,
+        &ledger_info,
+        le_with_ttl,
+        &prng_seed(),
+    )
+    .unwrap();
 }
 
 #[test]
@@ -2594,7 +2645,7 @@ fn test_auto_restore_with_expired_temp_entry_in_recording_mode() {
                     .try_into()
                     .unwrap(),
             },
-            instructions: 1026761,
+            instructions: 1561476,
             disk_read_bytes: wasm_entry_size + instance_entry_size,
             write_bytes: wasm_entry_size + instance_entry_size,
         }
@@ -2714,7 +2765,7 @@ fn test_auto_restore_with_recreated_temp_entry_in_recording_mode() {
                 read_only: vec![cd.contract_key.clone()].try_into().unwrap(),
                 read_write: vec![data_key, cd.wasm_key.clone()].try_into().unwrap(),
             },
-            instructions: 1028934,
+            instructions: 1563649,
             disk_read_bytes: wasm_entry_size,
             write_bytes: wasm_entry_size + temp_entry_size,
         }

soroban-simulation/src/test/simulation.rs

@@ -557,7 +557,7 @@ fn test_simulate_invoke_contract_with_autorestore() {
     assert!(res.contract_events.is_empty());
     assert!(!res.diagnostic_events.is_empty());
 
-    let expected_instructions = 9936120;
+    let expected_instructions = 10942546;
     let wasm_entry_size = contracts[0]
         .wasm_entry
         .to_xdr(Limits::none())
@@ -589,11 +589,11 @@ fn test_simulate_invoke_contract_with_autorestore() {
                 disk_read_bytes: wasm_entry_size + contract_1_size,
                 write_bytes: wasm_entry_size + contract_1_size,
             },
-            resource_fee: 6230340,
+            resource_fee: 6231346,
         })
     );
     assert_eq!(res.simulated_instructions, expected_instructions);
-    assert_eq!(res.simulated_memory, 4968053);
+    assert_eq!(res.simulated_memory, 5471262);
     assert_eq!(
         res.modified_entries,
         vec![