feat: Support wasm_memory_persistence for canister upgrades

CHANGELOG.md

@@ -2,6 +2,12 @@
 
 # UNRELEASED
 
+### feat: support `wasm_memory_persistence` for canister upgrades
+
+Support Motoko enhanced orthogonal persistence with the canister upgrade option `wasm_memory_persistence`.
+* For all Motoko canisters with the installed version using enhanced orthogonal persistence: `Keep` the main memory.
+* For all other canisters, non-Motoko or Motoko with classical persistence: Use the `None` as default mode.
+
 ### feat: PocketIC state
 
 `dfx start --pocketic` no longer requires `--clean`, and can persist replica state between runs.

e2e/assets/wasm_memory_persistence/README.md

@@ -0,0 +1,3 @@
+`classical-actor.wasm` is built from `actor.mo` with `moc -o classical-actor.wasm actor.mo`.
+`enhanced-actor.wasm` is built from `actor.mo` with `moc -o enhanced-actor.wasm --enhanced-orthogonal-persistence actor.mo`.
+`actor.did` is obtained by `moc --idl actor.mo`.

e2e/assets/wasm_memory_persistence/actor.did

@@ -0,0 +1,3 @@
+service : {
+  getVersion: () -> (nat) query;
+}

e2e/assets/wasm_memory_persistence/actor.mo

@@ -0,0 +1,12 @@
+import Prim "mo:prim";
+
+actor {
+    stable var version = 0;
+
+    version += 1;
+    Prim.debugPrint("Deployed actor version " # debug_show (version));
+
+    public query func getVersion() : async Nat {
+        return version;
+    };
+};

e2e/assets/wasm_memory_persistence/dfx.json

@@ -0,0 +1,9 @@
+{
+  "canisters": {
+    "test": {
+      "type": "custom",
+      "candid": "actor.did",
+      "wasm": "classical-actor.wasm"
+    }
+  }
+}

e2e/tests-dfx/wasm_memory_persistence.bash

@@ -0,0 +1,66 @@
+#!/usr/bin/env bats
+
+load ../utils/_
+
+setup() {
+  standard_setup
+
+  dfx_new test
+}
+
+teardown() {
+  dfx_stop
+
+  standard_teardown
+}
+
+@test "migrate Motoko from classical persistence to classical persistence" {
+  install_asset wasm_memory_persistence
+  dfx_start
+  dfx deploy
+  dfx deploy --upgrade-unchanged
+  assert_command dfx canister call test getVersion '()'
+  assert_match "(2 : nat)"
+}
+
+@test "migrate Motoko from classical persistence to enhanced orthogonal persistence" {
+  install_asset wasm_memory_persistence
+  dfx_start
+  dfx deploy
+  jq '.canisters.test.wasm="enhanced-actor.wasm"' dfx.json | sponge dfx.json
+  dfx deploy
+  assert_command dfx canister call test getVersion '()'
+  assert_match "(2 : nat)"
+}
+
+@test "migrate Motoko from enhanced orthogonal persistence to enhanced orthogonal persistence" {
+  install_asset wasm_memory_persistence
+  dfx_start
+  jq '.canisters.test.wasm="enhanced-actor.wasm"' dfx.json | sponge dfx.json
+  dfx deploy
+  jq '.canisters.test.wasm="enhanced-actor.wasm"' dfx.json | sponge dfx.json
+  dfx deploy --upgrade-unchanged
+  assert_command dfx canister call test getVersion '()'
+  assert_match "(2 : nat)"
+}
+
+@test "failing Motoko downgrade from enhanced orthogonal persistence to classical persistence" {
+  install_asset wasm_memory_persistence
+  dfx_start
+  jq '.canisters.test.wasm="enhanced-actor.wasm"' dfx.json | sponge dfx.json
+  dfx deploy
+  jq '.canisters.test.wasm="classical-actor.wasm"' dfx.json | sponge dfx.json
+  assert_command_fail dfx deploy
+  assert_match "The \`wasm_memory_persistence: opt Keep\` upgrade option requires that the new canister module supports enhanced orthogonal persistence."
+}
+
+@test "re-install Motoko enhanced orthogonal persistence with classical persistence" {
+  install_asset wasm_memory_persistence
+  dfx_start
+  jq '.canisters.test.wasm="enhanced-actor.wasm"' dfx.json | sponge dfx.json
+  dfx deploy
+  jq '.canisters.test.wasm="classical-actor.wasm"' dfx.json | sponge dfx.json
+  echo yes | dfx canister install test --mode=reinstall
+  assert_command dfx canister call test getVersion '()'
+  assert_match "(1 : nat)"
+}

src/dfx/src/lib/operations/canister/install_canister.rs

@@ -21,7 +21,9 @@ use dfx_core::config::model::network_descriptor::NetworkDescriptor;
 use dfx_core::identity::CallSender;
 use fn_error_context::context;
 use ic_agent::Agent;
-use ic_utils::interfaces::management_canister::builders::InstallMode;
+use ic_utils::interfaces::management_canister::builders::{
+    CanisterUpgradeOptions, InstallMode, WasmMemoryPersistence,
+};
 use ic_utils::interfaces::ManagementCanister;
 use ic_utils::Argument;
 use itertools::Itertools;
@@ -63,9 +65,16 @@ pub async fn install_canister(
         "Previously installed module hash: {:?}",
         installed_module_hash.as_ref().map(hex::encode)
     );
+    let wasm_memory_persistence =
+        read_module_metadata(agent, canister_id, "enhanced-orthogonal-persistence")
+            .await
+            .map(|_| WasmMemoryPersistence::Keep);
     let mode = mode.unwrap_or_else(|| {
         if installed_module_hash.is_some() {
-            InstallMode::Upgrade(None)
+            InstallMode::Upgrade(Some(CanisterUpgradeOptions {
+                wasm_memory_persistence,
+                skip_pre_upgrade: None,
+            }))
         } else {
             InstallMode::Install
         }