feat: allow specifying wasm_memory_persistence (#274)

dfinity · Jan 14, 2025 · e221e8a · e221e8a
1 parent 8052e49
commit e221e8a
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 9 deletions.
diff --git a/service/pool/IC.mo b/service/pool/IC.mo
@@ -19,6 +19,16 @@ module {
     wasm_memory_limit : Nat;
     log_visibility : log_visibility;
   };
+  public type canister_install_mode = {
+    #reinstall;
+    #upgrade : ?{
+      wasm_memory_persistence : ?{
+          #Keep;
+          #Replace;
+      };
+    };
+    #install;
+  };
   public type snapshot = {
       id : snapshot_id;
       total_size : Nat64;
@@ -48,7 +58,7 @@ module {
   public type install_chunked_code_args = {
     arg : Blob;
     wasm_module_hash : Blob;
-    mode : { #reinstall; #upgrade; #install };
+    mode : canister_install_mode;
     chunk_hashes_list : [{hash : Blob}];
     target_canister : canister_id;
     store_canister : ?canister_id;
@@ -88,7 +98,7 @@ module {
     install_code : shared {
         arg : Blob;
         wasm_module : wasm_module;
-        mode : { #reinstall; #upgrade; #install };
+        mode : canister_install_mode;
         canister_id : canister_id;
     } -> async ();
     list_canister_snapshots : shared { canister_id : canister_id } -> async [snapshot];

diff --git a/service/pool/Main.mo b/service/pool/Main.mo
@@ -85,7 +85,7 @@ shared (creator) actor class Self(opt_params : ?Types.InitParams) = this {
         await f1;
         await* f2;
     };
-    private func getExpiredCanisterInfo(origin : Logs.Origin) : async* (Types.CanisterInfo, {#install; #reinstall}) {
+    private func getExpiredCanisterInfo(origin : Logs.Origin) : async* (Types.CanisterInfo, ICType.canister_install_mode) {
         let res = switch (pool.getExpiredCanisterId()) {
             case (#newId) {
               try {
@@ -208,19 +208,19 @@ shared (creator) actor class Self(opt_params : ?Types.InitParams) = this {
         statsByOrigin.addInstall({ origin = "external"; tags = [] });
     };
     // Combine create_canister and install_code into a single update call. Returns the current available canister id.
-    public shared ({ caller }) func deployCanister(opt_info: ?Types.CanisterInfo, args: ?Types.DeployArgs) : async (Types.CanisterInfo, {#install; #upgrade; #reinstall}) {
+    public shared ({ caller }) func deployCanister(opt_info: ?Types.CanisterInfo, args: ?Types.DeployArgs) : async (Types.CanisterInfo, ICType.canister_install_mode) {
         if (not Principal.isController(caller)) {
             throw Error.reject "Only called by controller";
         };
         if (Option.isSome(params.stored_module) and Option.isSome(args)) {
             throw Error.reject "args should be null when stored_module is set";
         };
         let origin = { origin = "admin"; tags = [] };
-        let (info, mode) = switch (opt_info) {
+        let (info, default_mode) = switch (opt_info) {
         case null { await* getExpiredCanisterInfo(origin) };
         case (?info) {
                  if (pool.find info) {
-                     (info, #upgrade)
+                     (info, #upgrade(null))
                  } else {
                      if (pool.findId(info.id)) {
                          await* getExpiredCanisterInfo(origin)
@@ -231,6 +231,14 @@ shared (creator) actor class Self(opt_params : ?Types.InitParams) = this {
                  };
              };
         };
+        let mode = switch (args) {
+            case (?args) {
+                Option.get(args.mode, default_mode);
+            };
+            case (null) {
+                default_mode;
+            };
+        };
         switch (args) {
         case (?args) {
                  let wasm = if (Option.get(args.bypass_wasm_transform, false)) {
@@ -369,7 +377,7 @@ shared (creator) actor class Self(opt_params : ?Types.InitParams) = this {
             };
             switch (args.mode) {
             case (#install) { tags.add("mode:install") };
-            case (#upgrade) { tags.add("mode:upgrade") };
+            case (#upgrade(_)) { tags.add("mode:upgrade") };
             case (#reinstall) { tags.add("mode:reinstall") };
             };
             let origin = { origin = install_config.origin.origin; tags = Buffer.toArray(tags) };
@@ -601,7 +609,7 @@ shared (creator) actor class Self(opt_params : ?Types.InitParams) = this {
     public shared ({ caller }) func install_code({
         arg : Blob;
         wasm_module : ICType.wasm_module;
-        mode : { #reinstall; #upgrade; #install };
+        mode : ICType.canister_install_mode;
         canister_id : ICType.canister_id;
     }) : async () {
         switch (sanitizeInputs(caller, canister_id)) {

diff --git a/service/pool/Types.mo b/service/pool/Types.mo
@@ -48,13 +48,14 @@ module {
     public type InstallArgs = {
         arg : Blob;
         wasm_module : Blob;
-        mode : { #reinstall; #upgrade; #install };
+        mode : ICType.canister_install_mode;
         canister_id : Principal;
     };
     public type DeployArgs = {
         arg : Blob;
         wasm_module : Blob;
         bypass_wasm_transform : ?Bool;
+        mode : ?ICType.canister_install_mode;
     };
     public type InstallConfig = {
         profiling: Bool;