refactor: install the canister memories in create_execution_state

rs/execution_environment/benches/lib/src/common.rs

@@ -12,8 +12,8 @@ use ic_cycles_account_manager::ResourceSaturation;
 use ic_embedders::wasmtime_embedder::system_api::{ExecutionParameters, InstructionLimits};
 use ic_error_types::RejectCode;
 use ic_execution_environment::{
-    CompilationCostHandling, ExecutionEnvironment, ExecutionServicesForTesting, RoundLimits,
-    as_round_instructions,
+    CompilationCostHandling, ExecutionEnvironment, ExecutionServicesForTesting, MemorySource,
+    RoundLimits, as_round_instructions,
 };
 use ic_interfaces::execution_environment::{ExecutionMode, SubnetAvailableMemory};
 use ic_limits::SMALL_APP_SUBNET_MAX_SIZE;
@@ -142,6 +142,7 @@ where
             UNIX_EPOCH,
             &mut round_limits,
             CompilationCostHandling::CountFullAmount,
+            MemorySource::Fresh,
         )
         .1
         .expect("Failed to create execution state");

rs/execution_environment/benches/management_canister/create_execution_state.rs

@@ -1,6 +1,8 @@
 use criterion::{BatchSize, Criterion, criterion_group, criterion_main};
 use ic_config::execution_environment::{SUBNET_CALLBACK_SOFT_LIMIT, SUBNET_MEMORY_RESERVATION};
-use ic_execution_environment::{CompilationCostHandling, RoundLimits, as_round_instructions};
+use ic_execution_environment::{
+    CompilationCostHandling, MemorySource, RoundLimits, as_round_instructions,
+};
 use ic_interfaces::execution_environment::SubnetAvailableMemory;
 use ic_test_utilities_execution_environment::ExecutionTestBuilder;
 use ic_test_utilities_types::ids::canister_test_id;
@@ -39,6 +41,7 @@ fn run_benchmark(
                     ic_types::time::UNIX_EPOCH,
                     &mut round_limits,
                     compilation_cost_handling,
+                    MemorySource::Fresh,
                 );
             },
             BatchSize::SmallInput,

rs/execution_environment/src/canister_manager.rs

@@ -11,7 +11,7 @@ use crate::execution_environment::{
     RoundLimits,
 };
 use crate::execution_environment_metrics::ExecutionEnvironmentMetrics;
-use crate::hypervisor::Hypervisor;
+use crate::hypervisor::{Hypervisor, MemorySource};
 use crate::types::{IngressResponse, Response};
 use crate::util::{GOVERNANCE_CANISTER_ID, MIGRATION_CANISTER_ID};
 
@@ -2295,6 +2295,35 @@ impl CanisterManager {
             return Err(CanisterManagerError::LongExecutionAlreadyInProgress { canister_id });
         }
 
+        // Build the memories to restore up front, so that they can be handed to
+        // `create_execution_state` below instead of being installed after the
+        // fact. This is cheap: it only clones the storage of the snapshot's page
+        // maps and, for another canister's snapshot, allocates a fresh page
+        // allocator.
+        let (snapshot_wasm_memory, snapshot_stable_memory) = if canister_id == snapshot_canister_id
+        {
+            (
+                Memory::from(&execution_snapshot.wasm_memory),
+                Memory::from(&execution_snapshot.stable_memory),
+            )
+        } else {
+            let not_loadable = || CanisterManagerError::CanisterSnapshotNotLoadable {
+                canister_id,
+                snapshot_id,
+            };
+            let wasm_memory = Memory::try_from((
+                &execution_snapshot.wasm_memory,
+                Arc::clone(&self.fd_factory),
+            ))
+            .map_err(|_| not_loadable())?;
+            let stable_memory = Memory::try_from((
+                &execution_snapshot.stable_memory,
+                Arc::clone(&self.fd_factory),
+            ))
+            .map_err(|_| not_loadable())?;
+            (wasm_memory, stable_memory)
+        };
+
         // All basic checks have passed, prepay cycles for instructions.
         //
         // The Wasm execution mode is only used to pick the per-instruction rate
@@ -2339,6 +2368,10 @@ impl CanisterManager {
                     time,
                     round_limits,
                     compilation_cost_handling,
+                    MemorySource::Explicit {
+                        wasm_memory: snapshot_wasm_memory,
+                        stable_memory: snapshot_stable_memory,
+                    },
                 );
             debug_assert!(
                 instructions_for_execution <= prepaid_execution_instructions,
@@ -2423,38 +2456,6 @@ impl CanisterManager {
 
             new_execution_state.exported_globals = execution_snapshot.exported_globals.clone();
 
-            if canister_id == snapshot_canister_id {
-                new_execution_state.stable_memory = Memory::from(&execution_snapshot.stable_memory);
-                new_execution_state.wasm_memory = Memory::from(&execution_snapshot.wasm_memory);
-            } else {
-                let new_stable_memory = match Memory::try_from((
-                    &execution_snapshot.stable_memory,
-                    Arc::clone(&self.fd_factory),
-                )) {
-                    Ok(memory) => memory,
-                    Err(_) => {
-                        return Err(CanisterManagerError::CanisterSnapshotNotLoadable {
-                            canister_id,
-                            snapshot_id,
-                        });
-                    }
-                };
-                new_execution_state.stable_memory = new_stable_memory;
-
-                let new_wasm_memory = match Memory::try_from((
-                    &execution_snapshot.wasm_memory,
-                    Arc::clone(&self.fd_factory),
-                )) {
-                    Ok(memory) => memory,
-                    Err(_) => {
-                        return Err(CanisterManagerError::CanisterSnapshotNotLoadable {
-                            canister_id,
-                            snapshot_id,
-                        });
-                    }
-                };
-                new_execution_state.wasm_memory = new_wasm_memory;
-            }
             Some(new_execution_state)
         };
 

rs/execution_environment/src/execution/install.rs

@@ -8,10 +8,10 @@ use crate::canister_manager::types::{
 };
 use crate::execution::common::{ingress_status_with_processing_state, update_round_limits};
 use crate::execution::install_code::{
-    CanisterMemoryHandling, InstallCodeHelper, MemoryHandling, OriginalContext,
-    PausedInstallCodeHelper, finish_err,
+    InstallCodeHelper, OriginalContext, PausedInstallCodeHelper, finish_err,
 };
 use crate::execution_environment::{RoundContext, RoundLimits};
+use crate::hypervisor::MemorySource;
 use ic_base_types::PrincipalId;
 use ic_embedders::{
     wasm_executor::{CanisterStateChanges, PausedWasmExecution, WasmExecutionResult},
@@ -120,15 +120,12 @@ pub(crate) fn execute_install(
         round.time,
         round_limits,
         original.compilation_cost_handling,
+        // Install and re-install replace both the stable memory and the main
+        // memory with the initial memories of the new module.
+        MemorySource::Fresh,
     );
     helper.charge_for_compilation(instructions_from_compilation);
-    if let Err(err) = helper.replace_execution_state_and_allocations(
-        result,
-        CanisterMemoryHandling {
-            stable_memory_handling: MemoryHandling::Replace,
-            main_memory_handling: MemoryHandling::Replace,
-        },
-    ) {
+    if let Err(err) = helper.replace_execution_state_and_allocations(result) {
         let instructions_left = helper.instructions_left();
         return finish_err(
             clean_canister,

rs/execution_environment/src/execution/install_code.rs

@@ -42,39 +42,21 @@ use ic_replicated_state::canister_state::execution_state::WasmExecutionMode;
 #[cfg(test)]
 mod tests;
 
-/// Indicates whether the memory is kept or replaced with new (initial) memory.
-/// Applicable to both the stable memory and the main memory of a canister.
-#[derive(Copy, Clone, PartialEq, Debug)]
-pub(crate) enum MemoryHandling {
-    /// Retain the memory.
-    Keep,
-    /// Reset the memory.
-    Replace,
-}
-
-/// Specifies the handling of the canister's memories.
-/// * On install and re-install:
-///   - Replace both the stable memory and the main memory.
-/// * On upgrade:
-///   - For canisters with enhanced orthogonal persistence (Motoko):
-///     Retain both the main memory and the stable memory.
-///   - For all other canisters:
-///     Retain only the stable memory and erase the main memory.
-#[derive(Copy, Clone, PartialEq, Debug)]
-pub(crate) struct CanisterMemoryHandling {
-    pub stable_memory_handling: MemoryHandling,
-    pub main_memory_handling: MemoryHandling,
-}
-
 /// The main steps of `install_code` execution that may fail with an error or
 /// change the canister state.
 #[derive(Clone, Debug)]
 #[allow(clippy::large_enum_variant)]
 pub(crate) enum InstallCodeStep {
     ValidateInput,
     ReplaceExecutionStateAndAllocations {
+        /// The execution state as returned by
+        /// `Hypervisor::create_execution_state`, i.e. with the memories that
+        /// this `install_code` preserves (if any) already in place. Replaying
+        /// this step hence re-applies those memories as they were at the time
+        /// of the original execution rather than re-deriving them from the
+        /// replayed canister state; the two agree because replaying the
+        /// preceding steps is deterministic.
         maybe_execution_state: HypervisorResult<ExecutionState>,
-        memory_handling: CanisterMemoryHandling,
     },
     ClearCertifiedData,
     ClearLog,
@@ -591,21 +573,21 @@ impl InstallCodeHelper {
         Ok(())
     }
 
-    /// Replaces the execution state of the current canister with the freshly
-    /// created execution state. The stable memory and the main memory are
-    /// conditionally replaced based on the given `memory_handling`.
+    /// Replaces the execution state of the current canister with the newly
+    /// created execution state. Which memories the new execution state carries
+    /// (the initial ones of the new module or the preserved ones of the old
+    /// execution state) has already been decided by
+    /// `Hypervisor::create_execution_state`.
     ///
     /// It also updates the compute and memory allocations with the requested
     /// values in `original` context.
     pub fn replace_execution_state_and_allocations(
         &mut self,
         maybe_execution_state: HypervisorResult<ExecutionState>,
-        memory_handling: CanisterMemoryHandling,
     ) -> Result<(), CanisterManagerError> {
         self.steps
             .push(InstallCodeStep::ReplaceExecutionStateAndAllocations {
                 maybe_execution_state: maybe_execution_state.clone(),
-                memory_handling,
             });
 
         let old_memory_usage = self.canister.memory_usage();
@@ -616,23 +598,12 @@ impl InstallCodeHelper {
             .as_ref()
             .map_or(NumBytes::new(0), |es| es.metadata.memory_usage());
 
-        // Replace the execution state and maybe the stable memory.
-        let mut execution_state =
+        // Replace the execution state, dropping the old one.
+        let execution_state =
             maybe_execution_state.map_err(|err| (self.canister.canister_id(), err))?;
 
         let new_wasm_custom_sections_memory_used = execution_state.metadata.memory_usage();
 
-        if let Some(old) = self.canister.execution_state.take() {
-            match memory_handling.stable_memory_handling {
-                MemoryHandling::Keep => execution_state.stable_memory = old.stable_memory,
-                MemoryHandling::Replace => {}
-            }
-            match memory_handling.main_memory_handling {
-                MemoryHandling::Keep => execution_state.wasm_memory = old.wasm_memory,
-                MemoryHandling::Replace => {}
-            }
-        };
-
         self.canister.execution_state = Some(execution_state);
 
         let new_memory_usage = self.canister.memory_usage();
@@ -832,10 +803,7 @@ impl InstallCodeHelper {
             InstallCodeStep::ValidateInput => self.validate_input(original),
             InstallCodeStep::ReplaceExecutionStateAndAllocations {
                 maybe_execution_state,
-                memory_handling,
-            } => {
-                self.replace_execution_state_and_allocations(maybe_execution_state, memory_handling)
-            }
+            } => self.replace_execution_state_and_allocations(maybe_execution_state),
             InstallCodeStep::ClearCertifiedData => {
                 self.clear_certified_data();
                 Ok(())