feat: Stop canisters before upgrading

CHANGELOG.md

@@ -1,6 +1,7 @@
 # Unreleased
 
 * feat: Added support for creating canisters on cloud engine subnets. Note that local networks cannot yet create these subnets.
+* feat: Upgrading canisters now stops them before the upgrade and starts them again afterwards
 * feat: `icp canister logs` supports filtering by timestamp (`--since`, `--until`) and log index (`--since-index`, `--until-index`)
 * feat: Support `log_memory_limit` canister setting in `icp canister settings update` and `icp canister settings sync`
 * feat: Leaving off the method name parameter in `icp canister call` prompts you with an interactive list of methods

crates/icp-cli/src/commands/canister/install.rs

@@ -13,7 +13,7 @@ use crate::{
     commands::args,
     operations::{
         candid_compat::{CandidCompatibility, check_candid_compatibility},
-        install::{install_canister, resolve_install_mode},
+        install::{install_canister, resolve_install_mode_and_status},
     },
 };
 
@@ -120,8 +120,9 @@ pub(crate) async fn exec(ctx: &Context, args: &InstallArgs) -> Result<(), anyhow
         .transpose()?;
 
     let canister_display = args.cmd_args.canister.to_string();
-    let install_mode =
-        resolve_install_mode(&agent, &canister_display, &canister_id, &args.mode).await?;
+    let (install_mode, status) =
+        resolve_install_mode_and_status(&agent, &canister_display, &canister_id, &args.mode)
+            .await?;
 
     // Candid interface compatibility check for upgrades
     if !args.yes && matches!(install_mode, CanisterInstallMode::Upgrade(_)) {
@@ -158,6 +159,7 @@ pub(crate) async fn exec(ctx: &Context, args: &InstallArgs) -> Result<(), anyhow
         &canister_display,
         &wasm,
         install_mode,
+        status,
         init_args_bytes.as_deref(),
     )
     .await?;

crates/icp-cli/src/commands/deploy.rs

@@ -20,7 +20,7 @@ use crate::{
         build::build_many_with_progress_bar,
         candid_compat::check_candid_compatibility_many,
         create::CreateOperation,
-        install::{install_many, resolve_install_mode},
+        install::{install_many, resolve_install_mode_and_status},
         settings::sync_settings_many,
         sync::sync_many,
     },
@@ -242,7 +242,8 @@ pub(crate) async fn exec(ctx: &Context, args: &DeployArgs) -> Result<(), anyhow:
                 .await
                 .map_err(|e| anyhow!(e))?;
 
-            let mode = resolve_install_mode(&agent, name, &cid, &args.mode).await?;
+            let (mode, status) =
+                resolve_install_mode_and_status(&agent, name, &cid, &args.mode).await?;
 
             let env = ctx.get_environment(&environment_selection).await?;
             let (_canister_path, canister_info) =
@@ -254,7 +255,7 @@ pub(crate) async fn exec(ctx: &Context, args: &DeployArgs) -> Result<(), anyhow:
                 .map(|ia| ia.to_bytes())
                 .transpose()?;
 
-            Ok::<_, anyhow::Error>((name.clone(), cid, mode, init_args_bytes))
+            Ok::<_, anyhow::Error>((name.clone(), cid, mode, status, init_args_bytes))
         }
     }))
     .await?;
@@ -265,7 +266,7 @@ pub(crate) async fn exec(ctx: &Context, args: &DeployArgs) -> Result<(), anyhow:
             agent.clone(),
             canisters
                 .iter()
-                .map(|(name, cid, mode, _)| (&**name, *cid, *mode)),
+                .map(|(name, cid, mode, _, _)| (&**name, *cid, *mode)),
             ctx.artifacts.clone(),
             Arc::new(ctx.term.clone()),
             ctx.debug,

crates/icp-cli/src/operations/install.rs

@@ -1,7 +1,7 @@
 use futures::{StreamExt, stream::FuturesOrdered};
 use ic_agent::{Agent, AgentError, export::Principal};
 use ic_management_canister_types::{
-    CanisterId, ChunkHash, UpgradeFlags, UploadChunkArgs, WasmMemoryPersistence,
+    CanisterId, CanisterStatusType, ChunkHash, UpgradeFlags, UploadChunkArgs, WasmMemoryPersistence,
 };
 use ic_utils::interfaces::{
     ManagementCanister, management_canister::builders::CanisterInstallMode,
@@ -10,7 +10,7 @@ use icp::context::TermWriter;
 use sha2::{Digest, Sha256};
 use snafu::{ResultExt, Snafu};
 use std::sync::Arc;
-use tracing::debug;
+use tracing::{debug, warn};
 
 use crate::progress::{ProgressManager, ProgressManagerSettings};
 
@@ -23,6 +23,18 @@ pub enum InstallOperationError {
 
     #[snafu(display("agent error: {source}"))]
     Agent { source: AgentError },
+
+    #[snafu(display("Failed to stop canister '{canister_name}' before upgrade"))]
+    StopCanister {
+        canister_name: String,
+        source: AgentError,
+    },
+
+    #[snafu(display("Failed to start canister '{canister_name}' after upgrade"))]
+    StartCanister {
+        canister_name: String,
+        source: AgentError,
+    },
 }
 
 #[derive(Debug, Snafu)]
@@ -41,28 +53,26 @@ struct InstallFailure {
 /// Resolve a mode string ("auto", "install", "reinstall", "upgrade") into
 /// a [`CanisterInstallMode`]. For "auto", queries `canister_status` to
 /// determine whether the canister already has code installed.
-pub(crate) async fn resolve_install_mode(
+pub(crate) async fn resolve_install_mode_and_status(
     agent: &Agent,
     canister_name: &str,
     canister_id: &Principal,
     mode: &str,
-) -> Result<CanisterInstallMode, ResolveInstallModeError> {
+) -> Result<(CanisterInstallMode, CanisterStatusType), ResolveInstallModeError> {
+    let mgmt = ManagementCanister::create(agent);
+    let (status,) = mgmt
+        .canister_status(canister_id)
+        .await
+        .context(ResolveInstallModeSnafu { canister_name })?;
     match mode {
-        "auto" => {
-            let mgmt = ManagementCanister::create(agent);
-            let (status,) = mgmt
-                .canister_status(canister_id)
-                .await
-                .context(ResolveInstallModeSnafu { canister_name })?;
-            Ok(if status.module_hash.is_some() {
-                CanisterInstallMode::Upgrade(None)
-            } else {
-                CanisterInstallMode::Install
-            })
-        }
-        "install" => Ok(CanisterInstallMode::Install),
-        "reinstall" => Ok(CanisterInstallMode::Reinstall),
-        "upgrade" => Ok(CanisterInstallMode::Upgrade(None)),
+        "auto" => Ok(if status.module_hash.is_some() {
+            (CanisterInstallMode::Upgrade(None), status.status)
+        } else {
+            (CanisterInstallMode::Install, status.status)
+        }),
+        "install" => Ok((CanisterInstallMode::Install, status.status)),
+        "reinstall" => Ok((CanisterInstallMode::Reinstall, status.status)),
+        "upgrade" => Ok((CanisterInstallMode::Upgrade(None), status.status)),
         _ => panic!("invalid install mode: {mode}"),
     }
 }
@@ -80,6 +90,7 @@ pub(crate) async fn install_canister(
     canister_name: &str,
     wasm: &[u8],
     mode: CanisterInstallMode,
+    status: CanisterStatusType,
     init_args: Option<&[u8]>,
 ) -> Result<(), InstallOperationError> {
     let mode = match mode {
@@ -106,7 +117,16 @@ pub(crate) async fn install_canister(
         canister_name, mode
     );
 
-    do_install_operation(agent, canister_id, canister_name, wasm, mode, init_args).await
+    do_install_operation(
+        agent,
+        canister_id,
+        canister_name,
+        wasm,
+        mode,
+        status,
+        init_args,
+    )
+    .await
 }
 
 async fn do_install_operation(
@@ -115,6 +135,7 @@ async fn do_install_operation(
     canister_name: &str,
     wasm: &[u8],
     mode: CanisterInstallMode,
+    status: CanisterStatusType,
     init_args: Option<&[u8]>,
 ) -> Result<(), InstallOperationError> {
     let mgmt = ManagementCanister::create(agent);
@@ -143,9 +164,12 @@ async fn do_install_operation(
             builder = builder.with_raw_arg(args.into());
         }
 
-        builder
-            .await
-            .map_err(|source| InstallOperationError::Agent { source })?;
+        stop_and_start_if_upgrade(&mgmt, canister_id, canister_name, mode, status, async {
+            builder
+                .await
+                .map_err(|source| InstallOperationError::Agent { source })
+        })
+        .await?;
     } else {
         // Large wasm: use chunked installation
         debug!("Installing wasm for {canister_name} using chunked installation");
@@ -197,31 +221,91 @@ async fn do_install_operation(
             builder = builder.with_raw_arg(args.to_vec());
         }
 
-        builder
-            .await
-            .map_err(|source| InstallOperationError::Agent { source })?;
+        let install_res =
+            stop_and_start_if_upgrade(&mgmt, canister_id, canister_name, mode, status, async {
+                builder
+                    .await
+                    .map_err(|source| InstallOperationError::Agent { source })
+            })
+            .await;
 
         // Clear chunk store after successful installation to free up storage
-        mgmt.clear_chunk_store(canister_id)
+        let clear_res = mgmt
+            .clear_chunk_store(canister_id)
             .await
-            .map_err(|source| InstallOperationError::Agent { source })?;
+            .map_err(|source| InstallOperationError::Agent { source });
+
+        if let Err(clear_error) = clear_res {
+            if let Err(install_error) = install_res {
+                warn!("Failed to clear chunk store after failed install: {clear_error}");
+                return Err(install_error);
+            } else {
+                return Err(clear_error);
+            }
+        }
+        install_res?;
     }
 
     Ok(())
 }
 
+async fn stop_and_start_if_upgrade(
+    mgmt: &ManagementCanister<'_>,
+    canister_id: &Principal,
+    canister_name: &str,
+    mode: CanisterInstallMode,
+    status: CanisterStatusType,
+    f: impl Future<Output = Result<(), InstallOperationError>>,
+) -> Result<(), InstallOperationError> {
+    let should_guard = matches!(
+        mode,
+        CanisterInstallMode::Upgrade(_) | CanisterInstallMode::Reinstall
+    ) && matches!(status, CanisterStatusType::Running);
+    // Stop the canister before proceeding
+    if should_guard {
+        mgmt.stop_canister(canister_id)
+            .await
+            .context(StopCanisterSnafu { canister_name })?;
+    }
+    // Install the canister
+    let install_result = f.await;
+    // Restart the canister whether or not the installation succeeded
+    if should_guard {
+        let start_result = mgmt.start_canister(canister_id).await;
+        if let Err(start_error) = start_result {
+            // If both install and start failed, report the install error since it's more likely to be the root cause
+            if let Err(install_error) = install_result {
+                warn!("Failed to start canister after failed upgrade: {start_error}");
+                return Err(install_error);
+            } else {
+                return Err(start_error).context(StartCanisterSnafu { canister_name });
+            }
+        }
+    }
+
+    install_result
+}
+
 /// Installs code to multiple canisters and displays progress bars.
 pub(crate) async fn install_many(
     agent: Agent,
-    canisters: impl IntoIterator<Item = (String, Principal, CanisterInstallMode, Option<Vec<u8>>)>,
+    canisters: impl IntoIterator<
+        Item = (
+            String,
+            Principal,
+            CanisterInstallMode,
+            CanisterStatusType,
+            Option<Vec<u8>>,
+        ),
+    >,
     artifacts: Arc<dyn icp::store_artifact::Access>,
     term: Arc<TermWriter>,
     debug: bool,
 ) -> Result<(), InstallManyError> {
     let mut futs = FuturesOrdered::new();
     let progress_manager = ProgressManager::new(ProgressManagerSettings { hidden: debug });
 
-    for (name, cid, mode, init_args) in canisters {
+    for (name, cid, mode, status, init_args) in canisters {
         let pb = progress_manager.create_progress_bar(&name);
         let agent = agent.clone();
         let install_fn = {
@@ -238,7 +322,16 @@ pub(crate) async fn install_many(
                     }
                 })?;
 
-                install_canister(&agent, &cid, &name, &wasm, mode, init_args.as_deref()).await
+                install_canister(
+                    &agent,
+                    &cid,
+                    &name,
+                    &wasm,
+                    mode,
+                    status,
+                    init_args.as_deref(),
+                )
+                .await
             }
         };