From 4e628a8b4ca0d2eb5736e01fdca604707b7e294f Mon Sep 17 00:00:00 2001
From: Adam Spofford <93943719+adamspofford-dfinity@users.noreply.github.com>
Date: Mon, 4 Nov 2024 15:01:09 -0800
Subject: [PATCH] fix: Update CanisterInstallMode signature (#527)

---
 .github/workflows/ci.yml                      |  2 +-
 .../src/api/management_canister/main/types.rs | 36 +++++++++++++++++--
 2 files changed, 34 insertions(+), 4 deletions(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 814e425b4..6a5341580 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -43,7 +43,7 @@ jobs:
     runs-on: ${{ matrix.os }}
     strategy:
       matrix:
-        os: [ubuntu-latest, macos-12]
+        os: [ubuntu-latest, macos-13-large]
     steps:
       - name: Checkout repository
         uses: actions/checkout@v4
diff --git a/src/ic-cdk/src/api/management_canister/main/types.rs b/src/ic-cdk/src/api/management_canister/main/types.rs
index ec68304d5..9eea37e81 100644
--- a/src/ic-cdk/src/api/management_canister/main/types.rs
+++ b/src/ic-cdk/src/api/management_canister/main/types.rs
@@ -191,10 +191,10 @@ pub enum CanisterInstallMode {
     Reinstall,
     /// Upgrade an existing canister.
     #[serde(rename = "upgrade")]
-    Upgrade(Option<SkipPreUpgrade>),
+    Upgrade(Option<UpgradeFlags>),
 }
 
-/// If set to true, the pre_upgrade step will be skipped during the canister upgrade
+/// Flags for canister installation with [`CanisterInstallMode::Upgrade`].
 #[derive(
     CandidType,
     Serialize,
@@ -209,7 +209,37 @@ pub enum CanisterInstallMode {
     Copy,
     Default,
 )]
-pub struct SkipPreUpgrade(pub Option<bool>);
+pub struct UpgradeFlags {
+    /// If set to `true`, the `pre_upgrade` step will be skipped during the canister upgrade
+    pub skip_pre_upgrade: Option<bool>,
+    /// If set to `Keep`, the WASM heap memory will be preserved instead of cleared.
+    pub wasm_memory_persistence: Option<WasmPersistenceMode>,
+}
+
+/// WASM memory persistence setting for [`UpgradeFlags`].
+#[derive(
+    CandidType,
+    Serialize,
+    Deserialize,
+    Debug,
+    PartialEq,
+    Eq,
+    PartialOrd,
+    Ord,
+    Hash,
+    Clone,
+    Copy,
+    Default,
+)]
+pub enum WasmPersistenceMode {
+    /// Preserve heap memory (only officially supported by Motoko)
+    #[serde(rename = "keep")]
+    Keep,
+    /// Clear heap memory
+    #[serde(rename = "replace")]
+    #[default]
+    Replace,
+}
 
 /// WASM module.
 pub type WasmModule = Vec<u8>;