From 7ee45d6ddee608a682183ca0a9f7d16a9e1b4e6f Mon Sep 17 00:00:00 2001 From: Daniel Wong Date: Wed, 14 Jan 2026 17:18:45 +0100 Subject: [PATCH 01/12] Moved LoadCanisterSnapshot* types from rs/handlers/root/interface to rs/nervous_system/root. --- Cargo.lock | 2 ++ rs/nervous_system/root/BUILD.bazel | 1 + rs/nervous_system/root/Cargo.toml | 1 + rs/nervous_system/root/src/lib.rs | 1 + .../root/src/load_canister_snapshot.rs | 24 +++++++++++++++++ rs/nns/handlers/root/interface/BUILD.bazel | 1 + rs/nns/handlers/root/interface/Cargo.toml | 1 + rs/nns/handlers/root/interface/src/lib.rs | 27 +++++-------------- 8 files changed, 37 insertions(+), 21 deletions(-) create mode 100644 rs/nervous_system/root/src/load_canister_snapshot.rs diff --git a/Cargo.lock b/Cargo.lock index 285b8cf6ae84..f09a39240ba1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11795,6 +11795,7 @@ dependencies = [ "async-trait", "candid", "dfn_core", + "ic-base-types", "ic-cdk", "ic-crypto-sha2", "ic-management-canister-types-private", @@ -12351,6 +12352,7 @@ dependencies = [ "ic-cdk", "ic-management-canister-types-private", "ic-nervous-system-clients", + "ic-nervous-system-root", "ic-nns-constants", "serde", ] diff --git a/rs/nervous_system/root/BUILD.bazel b/rs/nervous_system/root/BUILD.bazel index 5ac732a318cb..a3a8f8144115 100644 --- a/rs/nervous_system/root/BUILD.bazel +++ b/rs/nervous_system/root/BUILD.bazel @@ -9,6 +9,7 @@ DEPENDENCIES = [ "//rs/nervous_system/lock", "//rs/nervous_system/runtime", "//rs/rust_canisters/dfn_core", + "//rs/types/base_types", "//rs/types/management_canister_types", "@crate_index//:candid", "@crate_index//:ic-cdk", diff --git a/rs/nervous_system/root/Cargo.toml b/rs/nervous_system/root/Cargo.toml index bba27ef51b03..6891f16b1e23 100644 --- a/rs/nervous_system/root/Cargo.toml +++ b/rs/nervous_system/root/Cargo.toml @@ -11,6 +11,7 @@ documentation.workspace = true [dependencies] candid = { workspace = true } dfn_core = { path = "../../rust_canisters/dfn_core" } +ic-base-types = { path = "../../types/base_types" } ic-cdk = { workspace = true } ic-crypto-sha2 = { path = "../../crypto/sha2" } ic-management-canister-types-private = { path = "../../types/management_canister_types" } diff --git a/rs/nervous_system/root/src/lib.rs b/rs/nervous_system/root/src/lib.rs index 98761b0a15f2..c5d861f0cc30 100644 --- a/rs/nervous_system/root/src/lib.rs +++ b/rs/nervous_system/root/src/lib.rs @@ -1,3 +1,4 @@ pub mod change_canister; +pub mod load_canister_snapshot; pub const LOG_PREFIX: &str = "[Root Canister] "; diff --git a/rs/nervous_system/root/src/load_canister_snapshot.rs b/rs/nervous_system/root/src/load_canister_snapshot.rs new file mode 100644 index 000000000000..bef93f788c4d --- /dev/null +++ b/rs/nervous_system/root/src/load_canister_snapshot.rs @@ -0,0 +1,24 @@ +use candid::CandidType; +use ic_base_types::PrincipalId; +use serde::Deserialize; + +#[derive(Clone, Eq, PartialEq, Hash, Debug, CandidType, Deserialize)] +pub struct LoadCanisterSnapshotRequest { + pub canister_id: PrincipalId, + pub snapshot_id: Vec, +} + +#[derive(Clone, Eq, PartialEq, Hash, Debug, CandidType, Deserialize)] +pub enum LoadCanisterSnapshotResponse { + Ok(LoadCanisterSnapshotOk), + Err(LoadCanisterSnapshotError), +} + +#[derive(Clone, Eq, PartialEq, Hash, Debug, CandidType, Deserialize)] +pub struct LoadCanisterSnapshotOk {} + +#[derive(Clone, Eq, PartialEq, Hash, Debug, CandidType, Deserialize)] +pub struct LoadCanisterSnapshotError { + pub code: Option, + pub description: String, +} diff --git a/rs/nns/handlers/root/interface/BUILD.bazel b/rs/nns/handlers/root/interface/BUILD.bazel index 1d91d0e8daf3..a33e56b39888 100644 --- a/rs/nns/handlers/root/interface/BUILD.bazel +++ b/rs/nns/handlers/root/interface/BUILD.bazel @@ -5,6 +5,7 @@ package(default_visibility = ["//visibility:public"]) DEPENDENCIES = [ # Keep sorted. "//rs/nervous_system/clients", + "//rs/nervous_system/root", "//rs/nns/constants", "//rs/types/base_types", "//rs/types/management_canister_types", diff --git a/rs/nns/handlers/root/interface/Cargo.toml b/rs/nns/handlers/root/interface/Cargo.toml index 3e1f39450458..161c97270591 100644 --- a/rs/nns/handlers/root/interface/Cargo.toml +++ b/rs/nns/handlers/root/interface/Cargo.toml @@ -10,5 +10,6 @@ ic-base-types = { path = "../../../../types/base_types" } ic-cdk = { workspace = true } ic-management-canister-types-private = { path = "../../../../types/management_canister_types" } ic-nervous-system-clients = { path = "../../../../nervous_system/clients" } +ic-nervous-system-root = { path = "../../../../nervous_system/root" } ic-nns-constants = { path = "../../../../nns/constants" } serde = { workspace = true } diff --git a/rs/nns/handlers/root/interface/src/lib.rs b/rs/nns/handlers/root/interface/src/lib.rs index 3f2f0a03ac69..1d9180b3fb71 100644 --- a/rs/nns/handlers/root/interface/src/lib.rs +++ b/rs/nns/handlers/root/interface/src/lib.rs @@ -3,6 +3,12 @@ use ic_base_types::PrincipalId; use ic_nervous_system_clients::update_settings::CanisterSettings; use serde::Deserialize; +#[doc(inline)] +pub use ic_nervous_system_root::load_canister_snapshot::{ + LoadCanisterSnapshotError, LoadCanisterSnapshotOk, LoadCanisterSnapshotRequest, + LoadCanisterSnapshotResponse, +}; + pub mod client; /// The request structure to the `change_canister_controllers` API. @@ -107,24 +113,3 @@ pub struct TakeCanisterSnapshotError { pub code: Option, pub description: String, } - -#[derive(Clone, Eq, PartialEq, Hash, Debug, CandidType, Deserialize)] -pub struct LoadCanisterSnapshotRequest { - pub canister_id: PrincipalId, - pub snapshot_id: Vec, -} - -#[derive(Clone, Eq, PartialEq, Hash, Debug, CandidType, Deserialize)] -pub enum LoadCanisterSnapshotResponse { - Ok(LoadCanisterSnapshotOk), - Err(LoadCanisterSnapshotError), -} - -#[derive(Clone, Eq, PartialEq, Hash, Debug, CandidType, Deserialize)] -pub struct LoadCanisterSnapshotOk {} - -#[derive(Clone, Eq, PartialEq, Hash, Debug, CandidType, Deserialize)] -pub struct LoadCanisterSnapshotError { - pub code: Option, - pub description: String, -} From 96f36142eb848b53aec7bd97b416c9dfef5e51ac Mon Sep 17 00:00:00 2001 From: Daniel Wong Date: Wed, 14 Jan 2026 17:26:19 +0100 Subject: [PATCH 02/12] Analogous to the previous commit, except this time, we are dealing with TakeCanisterSnapshot: moved request and response types from rs/handlers/root/interface to rs/nervous_system/root. --- rs/nervous_system/root/src/lib.rs | 1 + .../root/src/take_canister_snapshot.rs | 28 +++++++++++++++++ rs/nns/handlers/root/interface/src/lib.rs | 31 ++++--------------- 3 files changed, 35 insertions(+), 25 deletions(-) create mode 100644 rs/nervous_system/root/src/take_canister_snapshot.rs diff --git a/rs/nervous_system/root/src/lib.rs b/rs/nervous_system/root/src/lib.rs index c5d861f0cc30..34b38965126b 100644 --- a/rs/nervous_system/root/src/lib.rs +++ b/rs/nervous_system/root/src/lib.rs @@ -1,4 +1,5 @@ pub mod change_canister; pub mod load_canister_snapshot; +pub mod take_canister_snapshot; pub const LOG_PREFIX: &str = "[Root Canister] "; diff --git a/rs/nervous_system/root/src/take_canister_snapshot.rs b/rs/nervous_system/root/src/take_canister_snapshot.rs new file mode 100644 index 000000000000..3b3603b4229a --- /dev/null +++ b/rs/nervous_system/root/src/take_canister_snapshot.rs @@ -0,0 +1,28 @@ +use candid::CandidType; +use ic_base_types::PrincipalId; +use serde::Deserialize; + +#[derive(Clone, Eq, PartialEq, Hash, Debug, CandidType, Deserialize)] +pub struct TakeCanisterSnapshotRequest { + pub canister_id: PrincipalId, + pub replace_snapshot: Option>, +} + +#[derive(Clone, Eq, PartialEq, Hash, Debug, CandidType, Deserialize)] +pub enum TakeCanisterSnapshotResponse { + Ok(TakeCanisterSnapshotOk), + Err(TakeCanisterSnapshotError), +} + +#[derive(Clone, Eq, PartialEq, Hash, Debug, CandidType, Deserialize)] +pub struct TakeCanisterSnapshotOk { + pub id: Vec, + pub taken_at_timestamp: u64, + pub total_size: u64, +} + +#[derive(Clone, Eq, PartialEq, Hash, Debug, CandidType, Deserialize)] +pub struct TakeCanisterSnapshotError { + pub code: Option, + pub description: String, +} diff --git a/rs/nns/handlers/root/interface/src/lib.rs b/rs/nns/handlers/root/interface/src/lib.rs index 1d9180b3fb71..2617b48327d9 100644 --- a/rs/nns/handlers/root/interface/src/lib.rs +++ b/rs/nns/handlers/root/interface/src/lib.rs @@ -3,6 +3,12 @@ use ic_base_types::PrincipalId; use ic_nervous_system_clients::update_settings::CanisterSettings; use serde::Deserialize; +#[doc(inline)] +pub use ic_nervous_system_root::take_canister_snapshot::{ + TakeCanisterSnapshotError, TakeCanisterSnapshotOk, TakeCanisterSnapshotRequest, + TakeCanisterSnapshotResponse, +}; + #[doc(inline)] pub use ic_nervous_system_root::load_canister_snapshot::{ LoadCanisterSnapshotError, LoadCanisterSnapshotOk, LoadCanisterSnapshotRequest, @@ -88,28 +94,3 @@ pub struct UpdateCanisterSettingsError { pub code: Option, pub description: String, } - -#[derive(Clone, Eq, PartialEq, Hash, Debug, CandidType, Deserialize)] -pub struct TakeCanisterSnapshotRequest { - pub canister_id: PrincipalId, - pub replace_snapshot: Option>, -} - -#[derive(Clone, Eq, PartialEq, Hash, Debug, CandidType, Deserialize)] -pub enum TakeCanisterSnapshotResponse { - Ok(TakeCanisterSnapshotOk), - Err(TakeCanisterSnapshotError), -} - -#[derive(Clone, Eq, PartialEq, Hash, Debug, CandidType, Deserialize)] -pub struct TakeCanisterSnapshotOk { - pub id: Vec, - pub taken_at_timestamp: u64, - pub total_size: u64, -} - -#[derive(Clone, Eq, PartialEq, Hash, Debug, CandidType, Deserialize)] -pub struct TakeCanisterSnapshotError { - pub code: Option, - pub description: String, -} From ecfbfbefa04d6ac48866b44f50f43c20fba7a441 Mon Sep 17 00:00:00 2001 From: Daniel Wong Date: Wed, 14 Jan 2026 17:36:43 +0100 Subject: [PATCH 03/12] Move implementation of load_canister_snapshot from rs/handlers/root/impl to rs/nervous_system/root. --- .../root/src/load_canister_snapshot.rs | 55 ++++++++++++++++++- .../root/impl/src/canister_management.rs | 46 ++-------------- 2 files changed, 59 insertions(+), 42 deletions(-) diff --git a/rs/nervous_system/root/src/load_canister_snapshot.rs b/rs/nervous_system/root/src/load_canister_snapshot.rs index bef93f788c4d..2049b4ce9986 100644 --- a/rs/nervous_system/root/src/load_canister_snapshot.rs +++ b/rs/nervous_system/root/src/load_canister_snapshot.rs @@ -1,5 +1,8 @@ use candid::CandidType; -use ic_base_types::PrincipalId; +use ic_base_types::{CanisterId, PrincipalId, SnapshotId}; +use ic_management_canister_types_private::LoadCanisterSnapshotArgs; +use ic_nervous_system_clients::management_canister_client::ManagementCanisterClient; +use ic_nervous_system_runtime::Runtime; use serde::Deserialize; #[derive(Clone, Eq, PartialEq, Hash, Debug, CandidType, Deserialize)] @@ -22,3 +25,53 @@ pub struct LoadCanisterSnapshotError { pub code: Option, pub description: String, } + +pub async fn load_canister_snapshot( + load_canister_snapshot_request: LoadCanisterSnapshotRequest, + management_canister_client: &mut impl ManagementCanisterClient, +) -> LoadCanisterSnapshotResponse +where + Rt: Runtime, +{ + let LoadCanisterSnapshotRequest { + canister_id, + snapshot_id, + } = load_canister_snapshot_request; + + let snapshot_id = match SnapshotId::try_from(snapshot_id) { + Ok(ok) => ok, + Err(err) => { + return LoadCanisterSnapshotResponse::Err(LoadCanisterSnapshotError { + code: None, + description: format!("Invalid snapshot ID: {err}"), + }); + } + }; + + let canister_id = match CanisterId::try_from(canister_id) { + Ok(ok) => ok, + Err(err) => { + return LoadCanisterSnapshotResponse::Err(LoadCanisterSnapshotError { + code: None, + description: format!("Invalid canister ID: {err}"), + }); + } + }; + + let load_canister_snapshot_args = LoadCanisterSnapshotArgs::new( + canister_id, + snapshot_id, + management_canister_client.canister_version(), + ); + + match management_canister_client + .load_canister_snapshot(load_canister_snapshot_args) + .await + { + Ok(()) => LoadCanisterSnapshotResponse::Ok(LoadCanisterSnapshotOk {}), + Err((code, description)) => LoadCanisterSnapshotResponse::Err(LoadCanisterSnapshotError { + code: Some(code), + description, + }), + } +} diff --git a/rs/nns/handlers/root/impl/src/canister_management.rs b/rs/nns/handlers/root/impl/src/canister_management.rs index 8d26dfe2811e..244bef7e0562 100644 --- a/rs/nns/handlers/root/impl/src/canister_management.rs +++ b/rs/nns/handlers/root/impl/src/canister_management.rs @@ -328,45 +328,9 @@ pub async fn load_canister_snapshot( load_canister_snapshot_request: LoadCanisterSnapshotRequest, management_canister_client: &mut impl ManagementCanisterClient, ) -> LoadCanisterSnapshotResponse { - let LoadCanisterSnapshotRequest { - canister_id, - snapshot_id, - } = load_canister_snapshot_request; - - let snapshot_id = match SnapshotId::try_from(snapshot_id) { - Ok(ok) => ok, - Err(err) => { - return LoadCanisterSnapshotResponse::Err(LoadCanisterSnapshotError { - code: None, - description: format!("Invalid snapshot ID: {err}"), - }); - } - }; - - let canister_id = match CanisterId::try_from(canister_id) { - Ok(ok) => ok, - Err(err) => { - return LoadCanisterSnapshotResponse::Err(LoadCanisterSnapshotError { - code: None, - description: format!("Invalid canister ID: {err}"), - }); - } - }; - - let load_canister_snapshot_args = LoadCanisterSnapshotArgs::new( - canister_id, - snapshot_id, - management_canister_client.canister_version(), - ); - - match management_canister_client - .load_canister_snapshot(load_canister_snapshot_args) - .await - { - Ok(()) => LoadCanisterSnapshotResponse::Ok(LoadCanisterSnapshotOk {}), - Err((code, description)) => LoadCanisterSnapshotResponse::Err(LoadCanisterSnapshotError { - code: Some(code), - description, - }), - } + ic_nervous_system_root::load_canister_snapshot::load_canister_snapshot::( + load_canister_snapshot_request, + management_canister_client, + ) + .await } From 2e8b85d2b80508b2ede862df82f21c9b7ac98016 Mon Sep 17 00:00:00 2001 From: Daniel Wong Date: Wed, 14 Jan 2026 17:39:32 +0100 Subject: [PATCH 04/12] Moved implementation of take_canister_snapshot from rs/nns/handlers/root to rs/nervous_system/root, analogous to the previous commit. --- .../root/src/take_canister_snapshot.rs | 78 ++++++++++++++++++- .../root/impl/src/canister_management.rs | 65 ++-------------- 2 files changed, 82 insertions(+), 61 deletions(-) diff --git a/rs/nervous_system/root/src/take_canister_snapshot.rs b/rs/nervous_system/root/src/take_canister_snapshot.rs index 3b3603b4229a..36208b771da9 100644 --- a/rs/nervous_system/root/src/take_canister_snapshot.rs +++ b/rs/nervous_system/root/src/take_canister_snapshot.rs @@ -1,5 +1,10 @@ use candid::CandidType; -use ic_base_types::PrincipalId; +use ic_base_types::{CanisterId, PrincipalId, SnapshotId}; +use ic_management_canister_types_private::{ + CanisterSnapshotResponse, TakeCanisterSnapshotArgs, +}; +use ic_nervous_system_clients::management_canister_client::ManagementCanisterClient; +use ic_nervous_system_runtime::Runtime; use serde::Deserialize; #[derive(Clone, Eq, PartialEq, Hash, Debug, CandidType, Deserialize)] @@ -26,3 +31,74 @@ pub struct TakeCanisterSnapshotError { pub code: Option, pub description: String, } + +pub async fn take_canister_snapshot( + take_canister_snapshot_request: TakeCanisterSnapshotRequest, + management_canister_client: &mut impl ManagementCanisterClient, +) -> TakeCanisterSnapshotResponse +where + Rt: Runtime, +{ + let TakeCanisterSnapshotRequest { + canister_id, + replace_snapshot, + } = take_canister_snapshot_request; + + let replace_snapshot = match replace_snapshot { + None => None, + Some(snapshot_id) => { + let snapshot_id = match SnapshotId::try_from(&snapshot_id) { + Ok(ok) => ok, + Err(err) => { + return TakeCanisterSnapshotResponse::Err(TakeCanisterSnapshotError { + code: None, + description: format!("Invalid snapshot ID ({snapshot_id:02X?}): {err}"), + }); + } + }; + + Some(snapshot_id) + } + }; + + let take_canister_snapshot_args = TakeCanisterSnapshotArgs { + canister_id, + replace_snapshot, + uninstall_code: None, + sender_canister_version: management_canister_client.canister_version(), + }; + + match management_canister_client + .take_canister_snapshot(take_canister_snapshot_args) + .await + { + Ok(result) => { + let result = + convert_from_canister_snapshot_response_to_take_canister_snapshot_ok(result); + TakeCanisterSnapshotResponse::Ok(result) + } + + Err((code, description)) => TakeCanisterSnapshotResponse::Err(TakeCanisterSnapshotError { + code: Some(code), + description, + }), + } +} + +fn convert_from_canister_snapshot_response_to_take_canister_snapshot_ok( + response: CanisterSnapshotResponse, +) -> TakeCanisterSnapshotOk { + let CanisterSnapshotResponse { + id, + taken_at_timestamp, + total_size, + } = response; + + let id = id.to_vec(); + + TakeCanisterSnapshotOk { + id, + taken_at_timestamp, + total_size, + } +} diff --git a/rs/nns/handlers/root/impl/src/canister_management.rs b/rs/nns/handlers/root/impl/src/canister_management.rs index 244bef7e0562..eb7a982cd391 100644 --- a/rs/nns/handlers/root/impl/src/canister_management.rs +++ b/rs/nns/handlers/root/impl/src/canister_management.rs @@ -260,69 +260,14 @@ pub async fn take_canister_snapshot( take_canister_snapshot_request: TakeCanisterSnapshotRequest, management_canister_client: &mut impl ManagementCanisterClient, ) -> TakeCanisterSnapshotResponse { - let TakeCanisterSnapshotRequest { - canister_id, - replace_snapshot, - } = take_canister_snapshot_request; - - let replace_snapshot = match replace_snapshot { - None => None, - Some(snapshot_id) => { - let snapshot_id = match SnapshotId::try_from(&snapshot_id) { - Ok(ok) => ok, - Err(err) => { - return TakeCanisterSnapshotResponse::Err(TakeCanisterSnapshotError { - code: None, - description: format!("Invalid snapshot ID ({snapshot_id:02X?}): {err}"), - }); - } - }; - - Some(snapshot_id) - } - }; - - let take_canister_snapshot_args = TakeCanisterSnapshotArgs { - canister_id, - replace_snapshot, - uninstall_code: None, - sender_canister_version: management_canister_client.canister_version(), - }; - - match management_canister_client - .take_canister_snapshot(take_canister_snapshot_args) - .await - { - Ok(result) => { - let result = - convert_from_canister_snapshot_response_to_take_canister_snapshot_ok(result); - TakeCanisterSnapshotResponse::Ok(result) - } - - Err((code, description)) => TakeCanisterSnapshotResponse::Err(TakeCanisterSnapshotError { - code: Some(code), - description, - }), - } + ic_nervous_system_root::take_canister_snapshot::take_canister_snapshot::( + take_canister_snapshot_request, + management_canister_client, + ) + .await } -fn convert_from_canister_snapshot_response_to_take_canister_snapshot_ok( - response: CanisterSnapshotResponse, -) -> TakeCanisterSnapshotOk { - let CanisterSnapshotResponse { - id, - taken_at_timestamp, - total_size, - } = response; - - let id = id.to_vec(); - TakeCanisterSnapshotOk { - id, - taken_at_timestamp, - total_size, - } -} pub async fn load_canister_snapshot( load_canister_snapshot_request: LoadCanisterSnapshotRequest, From 4acc1b3984a643d43ebdef83971b06bf34e28f3c Mon Sep 17 00:00:00 2001 From: Daniel Wong Date: Wed, 14 Jan 2026 17:45:22 +0100 Subject: [PATCH 05/12] Clean up some imports that are no longer needed due to previous two commits. --- rs/nns/handlers/root/impl/src/canister_management.rs | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/rs/nns/handlers/root/impl/src/canister_management.rs b/rs/nns/handlers/root/impl/src/canister_management.rs index eb7a982cd391..c546b44c0af7 100644 --- a/rs/nns/handlers/root/impl/src/canister_management.rs +++ b/rs/nns/handlers/root/impl/src/canister_management.rs @@ -1,13 +1,12 @@ #![allow(deprecated)] use crate::PROXIED_CANISTER_CALLS_TRACKER; -use ic_base_types::{CanisterId, PrincipalId, SnapshotId}; +use ic_base_types::{CanisterId, PrincipalId}; use ic_cdk::{ api::call::{RejectionCode, call_with_payment}, call, caller, print, }; use ic_management_canister_types_private::{ - CanisterInstallMode::Install, CanisterSettingsArgsBuilder, CanisterSnapshotResponse, - CreateCanisterArgs, InstallCodeArgs, LoadCanisterSnapshotArgs, TakeCanisterSnapshotArgs, + CanisterInstallMode::Install, CanisterSettingsArgsBuilder, CreateCanisterArgs, InstallCodeArgs, }; use ic_nervous_system_clients::{ canister_id_record::CanisterIdRecord, @@ -24,9 +23,8 @@ use ic_nns_common::{ types::CallCanisterRequest, }; use ic_nns_handler_root_interface::{ - ChangeCanisterControllersRequest, ChangeCanisterControllersResponse, LoadCanisterSnapshotError, - LoadCanisterSnapshotOk, LoadCanisterSnapshotRequest, LoadCanisterSnapshotResponse, - TakeCanisterSnapshotError, TakeCanisterSnapshotOk, TakeCanisterSnapshotRequest, + ChangeCanisterControllersRequest, ChangeCanisterControllersResponse, + LoadCanisterSnapshotRequest, LoadCanisterSnapshotResponse, TakeCanisterSnapshotRequest, TakeCanisterSnapshotResponse, UpdateCanisterSettingsError, UpdateCanisterSettingsRequest, UpdateCanisterSettingsResponse, }; @@ -267,8 +265,6 @@ pub async fn take_canister_snapshot( .await } - - pub async fn load_canister_snapshot( load_canister_snapshot_request: LoadCanisterSnapshotRequest, management_canister_client: &mut impl ManagementCanisterClient, From f972cfe2252e9b8057c9c40441adac35a76b62bb Mon Sep 17 00:00:00 2001 From: Daniel Wong Date: Tue, 20 Jan 2026 11:36:38 +0100 Subject: [PATCH 06/12] Added exclusively_stop_and_start_canister to rs/nervous_system/root. This is a private helper function that will be used in (near) future commits. Its main jobs are 1. Stop canister before main operation, and start it again after the main operation. 2. Acquire a lock on the canister so that different operations do not stomp on each other (and release after) This pattern is already used during canister upgrades. Therefore, one of the places where this will be used is canister upgrades. The other two places where this will be used in the very near future are taking and loading canister snapshots. --- rs/nervous_system/root/src/lib.rs | 2 + rs/nervous_system/root/src/private.rs | 244 ++++++++++++++++++++++++++ 2 files changed, 246 insertions(+) create mode 100644 rs/nervous_system/root/src/private.rs diff --git a/rs/nervous_system/root/src/lib.rs b/rs/nervous_system/root/src/lib.rs index 34b38965126b..05e60022187b 100644 --- a/rs/nervous_system/root/src/lib.rs +++ b/rs/nervous_system/root/src/lib.rs @@ -2,4 +2,6 @@ pub mod change_canister; pub mod load_canister_snapshot; pub mod take_canister_snapshot; +pub(crate) mod private; + pub const LOG_PREFIX: &str = "[Root Canister] "; diff --git a/rs/nervous_system/root/src/private.rs b/rs/nervous_system/root/src/private.rs new file mode 100644 index 000000000000..0cf22c64afc9 --- /dev/null +++ b/rs/nervous_system/root/src/private.rs @@ -0,0 +1,244 @@ +use crate::{LOG_PREFIX, change_canister::{stop_canister, start_canister}}; +use dfn_core::api::CanisterId; +use ic_nervous_system_lock::acquire_for; +use ic_nervous_system_runtime::CdkRuntime; +use std::{ + cell::RefCell, + collections::BTreeMap, + fmt::{self, Debug, Display, Formatter}, + future::Future, +}; + +// Records the acquisition of locks that protect a canister. +thread_local! { + pub(crate) static CANISTER_LOCK_ACQUISITIONS: RefCell> = const { RefCell::new(BTreeMap::new()) }; +} + +#[derive(Debug, Clone, PartialEq, Eq,)] +pub(crate) struct Reject { + pub code: i32, + pub message: String, +} + +impl From<(i32, String)> for Reject { + fn from(src: (i32, String)) -> Self { + let (code, message) = src; + + Self { + code, + message, + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +#[allow(clippy::enum_variant_names)] // Because all names end with "Failed". +pub(crate) enum ExclusivelyStopAndStartCanisterError { + LockAcquisitionFailed { + ongoing_operation_description: String, + }, + + StopBeforeMainOperationFailed { + stop_reject: Reject, + restart_result: Result<(), Reject>, + }, + + StartAfterMainOperationFailed { + start_reject: Reject, + main_operation_result: String, + }, +} + +pub(crate) use ExclusivelyStopAndStartCanisterError::{ + LockAcquisitionFailed, StopBeforeMainOperationFailed, StartAfterMainOperationFailed, +}; + +impl Display for ExclusivelyStopAndStartCanisterError { + fn fmt(&self, formatter: &mut Formatter) -> fmt::Result { + let result = match self { + LockAcquisitionFailed { ongoing_operation_description } => { + format!( + "Another operation is currently in progress \ + (on the same canister): {ongoing_operation_description}.", + ) + } + + StopBeforeMainOperationFailed { stop_reject, restart_result } => { + let main_reason = format!( + "Failed to stop the canister beforehand: {stop_reject:?}" + ); + + let restart_status = match restart_result { + Ok(()) => "fortunately, re-starting the canister succeeded".to_string(), + Err(restart_err) => { + format!( + "unfortunately, re-starting the canister failed: {restart_err:?}", + ) + } + }; + + format!("{main_reason}; {restart_status}.") + } + + StartAfterMainOperationFailed { start_reject, main_operation_result } => { + let start_failure = format!( + "Unable to start the canister again after the main operation: \ + {start_reject:?}", + ); + + format!("{start_failure}; main_operation_result: {main_operation_result}.") + } + }; + + write!(formatter, "{result}") + } +} + +/// Does a couple of things around main_operation: +/// +/// * Before: Acquire the canister's lock (hence "exclusive"), and stop the +/// canister (assuming stop_before is true, which it normally would be). +/// +/// * After: the reverse: (re-)start the canister, and release the lock. +/// +/// If stopping the canister fails, tries to start it again, and return Err +/// WITHOUT proceeding. +/// +/// The "after" operations are performed regardless of the outcome of the main +/// operation. Releasing the lock is performed regardless of whether re-starting +/// succeeds. +pub(crate) async fn exclusively_stop_and_start_canister( + canister_id: CanisterId, + operation_description: &str, + stop_before: bool, + main_operation: MainOperation, +) -> Result +where + MainOperation: FnOnce() -> Fut, + Fut: Future, + R: Debug, +{ + // Try to acquire lock for this canister; fail immediately if the canister + // is already locked (indicating that some other operation is currently in + // progress). + let _release_on_drop = acquire_for( + &CANISTER_LOCK_ACQUISITIONS, + canister_id, + operation_description.to_string(), + ) + .map_err(|ongoing_operation_description| LockAcquisitionFailed { ongoing_operation_description })?; + + if stop_before { + stop_before_main_operation(canister_id, operation_description).await?; + } + + let main_operation_result = main_operation().await; + + if stop_before { + return restart_after_main_operation( + canister_id, + operation_description, + main_operation_result, + ) + .await; + } + + Ok(main_operation_result) +} + +async fn stop_before_main_operation( + canister_id: CanisterId, + operation_description: &str, +) -> Result<(), ExclusivelyStopAndStartCanisterError> { + let stop_reject = match stop_canister::(canister_id).await { + Ok(()) => { + return Ok(()); + } + Err(err) => Reject::from(err), + }; + + // Recover from failure to stop by (re-)starting the canister. + + println!( + "{LOG_PREFIX}WARNING: Failed to stop canister {canister_id} \ + ({stop_reject:?}) while preparing to perform {operation_description}, \ + which will now not be attempted. But first, attempting to \ + (re-)start the canister.", + ); + + let restart_result = start_canister::(canister_id) + .await + .map_err(|start_err| { + let start_err = Reject::from(start_err); + + println!( + "{LOG_PREFIX}WARNING: After failing to stop canister {canister_id} \ + in preparation to perform {operation_description}, re-starting the canister \ + ALSO failed: {start_err:?}." + ); + + start_err + }); + + if restart_result.is_ok() { + println!( + "{LOG_PREFIX}It's not all bad news for {operation_description}. \ + So, even though stopping canister {canister_id} beforehand \ + did not work, (re-)starting it DID INDEED work. \ + (We will still NOT proceed with the main operation though.)", + ); + } + + Err(StopBeforeMainOperationFailed { + stop_reject, + restart_result, + }) + + /* + let final_err = format!( + "Failed to stop canister {canister_id} while preparing to perform a {operation}: \ + {code} {message}. Tried re-starting the canister after that, and the result of \ + re-starting was {restart_status}." + ); + + Err((code, final_err)) + */ +} + +async fn restart_after_main_operation( + canister_id: CanisterId, + operation_description: &str, + main_operation_result: R, +) -> Result +where + R: Debug +{ + match start_canister::(canister_id).await { + Ok(()) => { + Ok(main_operation_result) + } + + Err(err) => { + // This would be pretty terrible, especially if Governance is the + // target canister... + + let start_reject = Reject::from(err); + let main_operation_result = format!("{main_operation_result:?}"); + + println!( + "{LOG_PREFIX}ERROR: Failed to re-start canister {canister_id} after \ + performing {operation_description}: {start_reject:?}." + ); + + Err(StartAfterMainOperationFailed { + start_reject, + main_operation_result, + }) + } + } +} From 266b4be45dee046cffcd31f7943618e562029867 Mon Sep 17 00:00:00 2001 From: Daniel Wong Date: Tue, 20 Jan 2026 15:51:30 +0100 Subject: [PATCH 07/12] The `load_canister_snapshot` function (in rs/nervous_system/root) now locks the canister, and stops it before the main operation (i.e. calling the Management canister's `load_canister_snapshot` method), making it safer. (Ofc, the canister is also started after the main operation too, and the lock is released at the very end.) This is implemented with the new helper: `exclusively_stop_and_start_canister`. Next, we'll do the same thing for `take_canister_snapshot`. --- .../root/src/load_canister_snapshot.rs | 51 ++++++++++++++----- 1 file changed, 39 insertions(+), 12 deletions(-) diff --git a/rs/nervous_system/root/src/load_canister_snapshot.rs b/rs/nervous_system/root/src/load_canister_snapshot.rs index 2049b4ce9986..eca8c5df66f0 100644 --- a/rs/nervous_system/root/src/load_canister_snapshot.rs +++ b/rs/nervous_system/root/src/load_canister_snapshot.rs @@ -1,3 +1,4 @@ +use crate::private::{exclusively_stop_and_start_canister}; use candid::CandidType; use ic_base_types::{CanisterId, PrincipalId, SnapshotId}; use ic_management_canister_types_private::LoadCanisterSnapshotArgs; @@ -33,6 +34,8 @@ pub async fn load_canister_snapshot( where Rt: Runtime, { + let operation_description = format!("{:?}", load_canister_snapshot_request); + let LoadCanisterSnapshotRequest { canister_id, snapshot_id, @@ -58,20 +61,44 @@ where } }; - let load_canister_snapshot_args = LoadCanisterSnapshotArgs::new( + let result = exclusively_stop_and_start_canister::<_, _, _>( canister_id, - snapshot_id, - management_canister_client.canister_version(), - ); + &operation_description, + true, // stop_before + || async { + let load_canister_snapshot_args = LoadCanisterSnapshotArgs::new( + canister_id, + snapshot_id, + management_canister_client.canister_version(), + ); + + management_canister_client + .load_canister_snapshot(load_canister_snapshot_args) + .await + }, + ) + .await; - match management_canister_client - .load_canister_snapshot(load_canister_snapshot_args) - .await - { + let result = match result { + Ok(ok) => ok, + Err(err) => { + let description = format!("{err}"); + return LoadCanisterSnapshotResponse::Err( + LoadCanisterSnapshotError { + code: None, + description, + } + ); + } + }; + + match result { Ok(()) => LoadCanisterSnapshotResponse::Ok(LoadCanisterSnapshotOk {}), - Err((code, description)) => LoadCanisterSnapshotResponse::Err(LoadCanisterSnapshotError { - code: Some(code), - description, - }), + Err((code, description)) => { + LoadCanisterSnapshotResponse::Err(LoadCanisterSnapshotError { + code: Some(code), + description, + }) + } } } From bb594cb5e042ec9a218574236149f48137c6bea5 Mon Sep 17 00:00:00 2001 From: Daniel Wong Date: Tue, 20 Jan 2026 16:16:04 +0100 Subject: [PATCH 08/12] Do some extra safety stuff within take_canister_snapshot (like the previous commit): acquire canister lock, and stop the canister (ofc, after the main operation, start the canister back up again, and release the lock).| (This is like the previous commit, which enhanced load_canister_snapshot.) --- .../root/src/take_canister_snapshot.rs | 50 ++++++++++++++++--- 1 file changed, 42 insertions(+), 8 deletions(-) diff --git a/rs/nervous_system/root/src/take_canister_snapshot.rs b/rs/nervous_system/root/src/take_canister_snapshot.rs index 36208b771da9..fcf88f35cf72 100644 --- a/rs/nervous_system/root/src/take_canister_snapshot.rs +++ b/rs/nervous_system/root/src/take_canister_snapshot.rs @@ -1,3 +1,4 @@ +use crate::private::exclusively_stop_and_start_canister; use candid::CandidType; use ic_base_types::{CanisterId, PrincipalId, SnapshotId}; use ic_management_canister_types_private::{ @@ -39,6 +40,8 @@ pub async fn take_canister_snapshot( where Rt: Runtime, { + let operation_description = format!("{:?}", take_canister_snapshot_request); + let TakeCanisterSnapshotRequest { canister_id, replace_snapshot, @@ -61,17 +64,48 @@ where } }; - let take_canister_snapshot_args = TakeCanisterSnapshotArgs { + let canister_id = match CanisterId::try_from(canister_id) { + Ok(id) => id, + Err(e) => { + return TakeCanisterSnapshotResponse::Err(TakeCanisterSnapshotError { + code: None, + description: format!("Invalid canister ID: {:?}", e), + }); + } + }; + + let result = exclusively_stop_and_start_canister::<_, _, _>( canister_id, - replace_snapshot, - uninstall_code: None, - sender_canister_version: management_canister_client.canister_version(), + &operation_description, + true, // stop_before + || async { + let canister_id = PrincipalId::from(canister_id); + + let take_canister_snapshot_args = TakeCanisterSnapshotArgs { + canister_id, + replace_snapshot, + uninstall_code: None, + sender_canister_version: management_canister_client.canister_version(), + }; + + management_canister_client + .take_canister_snapshot(take_canister_snapshot_args) + .await + }, + ) + .await; + + let result = match result { + Ok(ok) => ok, + Err(err) => { + return TakeCanisterSnapshotResponse::Err(TakeCanisterSnapshotError { + code: None, + description: format!("{err}"), + }); + } }; - match management_canister_client - .take_canister_snapshot(take_canister_snapshot_args) - .await - { + match result { Ok(result) => { let result = convert_from_canister_snapshot_response_to_take_canister_snapshot_ok(result); From 2ed3469d59cc0e5a39b605a35e2f3da970851a4a Mon Sep 17 00:00:00 2001 From: Daniel Wong Date: Tue, 20 Jan 2026 16:43:41 +0100 Subject: [PATCH 09/12] Stop letting callers choose between DfnRuntime and CdkRuntime when calling (take|load)_canister_snapshot. Nobody ever used DfcRuntime in these cases; the ability to choose was just following an existing but depricated pattern. --- .../root/src/load_canister_snapshot.rs | 8 ++---- .../root/src/take_canister_snapshot.rs | 8 ++---- .../handlers/root/impl/canister/canister.rs | 4 +-- .../root/impl/src/canister_management.rs | 25 +------------------ 4 files changed, 7 insertions(+), 38 deletions(-) diff --git a/rs/nervous_system/root/src/load_canister_snapshot.rs b/rs/nervous_system/root/src/load_canister_snapshot.rs index eca8c5df66f0..e03a456f638e 100644 --- a/rs/nervous_system/root/src/load_canister_snapshot.rs +++ b/rs/nervous_system/root/src/load_canister_snapshot.rs @@ -3,7 +3,6 @@ use candid::CandidType; use ic_base_types::{CanisterId, PrincipalId, SnapshotId}; use ic_management_canister_types_private::LoadCanisterSnapshotArgs; use ic_nervous_system_clients::management_canister_client::ManagementCanisterClient; -use ic_nervous_system_runtime::Runtime; use serde::Deserialize; #[derive(Clone, Eq, PartialEq, Hash, Debug, CandidType, Deserialize)] @@ -27,13 +26,10 @@ pub struct LoadCanisterSnapshotError { pub description: String, } -pub async fn load_canister_snapshot( +pub async fn load_canister_snapshot( load_canister_snapshot_request: LoadCanisterSnapshotRequest, management_canister_client: &mut impl ManagementCanisterClient, -) -> LoadCanisterSnapshotResponse -where - Rt: Runtime, -{ +) -> LoadCanisterSnapshotResponse { let operation_description = format!("{:?}", load_canister_snapshot_request); let LoadCanisterSnapshotRequest { diff --git a/rs/nervous_system/root/src/take_canister_snapshot.rs b/rs/nervous_system/root/src/take_canister_snapshot.rs index fcf88f35cf72..31278923bdf7 100644 --- a/rs/nervous_system/root/src/take_canister_snapshot.rs +++ b/rs/nervous_system/root/src/take_canister_snapshot.rs @@ -5,7 +5,6 @@ use ic_management_canister_types_private::{ CanisterSnapshotResponse, TakeCanisterSnapshotArgs, }; use ic_nervous_system_clients::management_canister_client::ManagementCanisterClient; -use ic_nervous_system_runtime::Runtime; use serde::Deserialize; #[derive(Clone, Eq, PartialEq, Hash, Debug, CandidType, Deserialize)] @@ -33,13 +32,10 @@ pub struct TakeCanisterSnapshotError { pub description: String, } -pub async fn take_canister_snapshot( +pub async fn take_canister_snapshot( take_canister_snapshot_request: TakeCanisterSnapshotRequest, management_canister_client: &mut impl ManagementCanisterClient, -) -> TakeCanisterSnapshotResponse -where - Rt: Runtime, -{ +) -> TakeCanisterSnapshotResponse { let operation_description = format!("{:?}", take_canister_snapshot_request); let TakeCanisterSnapshotRequest { diff --git a/rs/nns/handlers/root/impl/canister/canister.rs b/rs/nns/handlers/root/impl/canister/canister.rs index 2f2874c50749..3db4cd6e7b8e 100644 --- a/rs/nns/handlers/root/impl/canister/canister.rs +++ b/rs/nns/handlers/root/impl/canister/canister.rs @@ -235,7 +235,7 @@ async fn take_canister_snapshot( take_canister_snapshot_request: TakeCanisterSnapshotRequest, ) -> TakeCanisterSnapshotResponse { check_caller_is_governance(); - canister_management::take_canister_snapshot( + ic_nervous_system_root::take_canister_snapshot::take_canister_snapshot( take_canister_snapshot_request, &mut new_management_canister_client(), ) @@ -249,7 +249,7 @@ async fn load_canister_snapshot( load_canister_snapshot_request: LoadCanisterSnapshotRequest, ) -> LoadCanisterSnapshotResponse { check_caller_is_governance(); - canister_management::load_canister_snapshot( + ic_nervous_system_root::load_canister_snapshot::load_canister_snapshot( load_canister_snapshot_request, &mut new_management_canister_client(), ) diff --git a/rs/nns/handlers/root/impl/src/canister_management.rs b/rs/nns/handlers/root/impl/src/canister_management.rs index c546b44c0af7..d473ab284d0c 100644 --- a/rs/nns/handlers/root/impl/src/canister_management.rs +++ b/rs/nns/handlers/root/impl/src/canister_management.rs @@ -24,8 +24,7 @@ use ic_nns_common::{ }; use ic_nns_handler_root_interface::{ ChangeCanisterControllersRequest, ChangeCanisterControllersResponse, - LoadCanisterSnapshotRequest, LoadCanisterSnapshotResponse, TakeCanisterSnapshotRequest, - TakeCanisterSnapshotResponse, UpdateCanisterSettingsError, UpdateCanisterSettingsRequest, + UpdateCanisterSettingsError, UpdateCanisterSettingsRequest, UpdateCanisterSettingsResponse, }; use ic_protobuf::{ @@ -253,25 +252,3 @@ pub async fn update_canister_settings( } } } - -pub async fn take_canister_snapshot( - take_canister_snapshot_request: TakeCanisterSnapshotRequest, - management_canister_client: &mut impl ManagementCanisterClient, -) -> TakeCanisterSnapshotResponse { - ic_nervous_system_root::take_canister_snapshot::take_canister_snapshot::( - take_canister_snapshot_request, - management_canister_client, - ) - .await -} - -pub async fn load_canister_snapshot( - load_canister_snapshot_request: LoadCanisterSnapshotRequest, - management_canister_client: &mut impl ManagementCanisterClient, -) -> LoadCanisterSnapshotResponse { - ic_nervous_system_root::load_canister_snapshot::load_canister_snapshot::( - load_canister_snapshot_request, - management_canister_client, - ) - .await -} From ae6e2ac4408e04b5bcdcf4a322ef2ff91b03ba56 Mon Sep 17 00:00:00 2001 From: Daniel Wong Date: Tue, 20 Jan 2026 19:10:01 +0100 Subject: [PATCH 10/12] More or less a refactoring of change_canister (in rs/nervous_system/root). Mainly, changed the implementation of change_canister to use the new exclusively_stop_and_start_canister helper function. Using the helper ensures that code upgrades and snapshottting do not interfere with one another. This also simplifies the implementation of change_canister by factoring out common implementation code. The other change to change_canister is that the Runtime generic parameter is removed. This is fine because all callers were using CdkRuntime anyway. --- rs/nervous_system/root/src/change_canister.rs | 104 ++++++++---------- .../handlers/root/impl/canister/canister.rs | 2 +- .../handlers/root/impl/src/root_proposals.rs | 3 +- rs/sns/root/canister/canister.rs | 2 +- 4 files changed, 46 insertions(+), 65 deletions(-) diff --git a/rs/nervous_system/root/src/change_canister.rs b/rs/nervous_system/root/src/change_canister.rs index 18217c46752d..c2b9aa679fef 100644 --- a/rs/nervous_system/root/src/change_canister.rs +++ b/rs/nervous_system/root/src/change_canister.rs @@ -1,5 +1,17 @@ +//! Basically, this wraps the `install_code` method of the Management canister. +//! +//! (Therefore, a better name for this would have been `install_code`.) +//! +//! There are a couple of additional things that this does though: +//! +//! 1. Lock the canister. +//! 2. Stop the canister before installing (this can be skipped). +//! +//! Of course, after calling the Management canister's `install_code` method, +//! this start the canister back up again, and unlocks it. + #![allow(deprecated)] -use crate::LOG_PREFIX; +use crate::{LOG_PREFIX, private::exclusively_stop_and_start_canister}; use candid::{CandidType, Deserialize, Encode, Principal}; use dfn_core::api::CanisterId; #[cfg(target_arch = "wasm32")] @@ -10,10 +22,8 @@ use ic_management_canister_types_private::{ InstallCodeArgs, }; use ic_nervous_system_clients::canister_id_record::CanisterIdRecord; -use ic_nervous_system_lock::acquire_for; use ic_nervous_system_runtime::Runtime; use serde::Serialize; -use std::{cell::RefCell, collections::BTreeMap}; /// The structure allows reconstructing a potentially large WASM from chunks needed to upgrade or /// reinstall some target canister. @@ -216,70 +226,41 @@ pub struct StopOrStartCanisterRequest { pub action: CanisterAction, } -// Thread-local storage for per-canister locks -// Key: CanisterId, Value: ChangeCanisterRequest (for debugging/logging) -thread_local! { - static CANISTER_CHANGE_LOCKS: RefCell> = - const {RefCell::new(BTreeMap::new()) }; -} - -pub async fn change_canister(request: ChangeCanisterRequest) -> Result<(), String> -where - Rt: Runtime, -{ +pub async fn change_canister(request: ChangeCanisterRequest) -> Result<(), String> { let canister_id = request.canister_id; let stop_before_installing = request.stop_before_installing; - // Try to acquire lock for this canister - fail immediately if locked - let _guard = match acquire_for(&CANISTER_CHANGE_LOCKS, canister_id, request.clone()) { - Ok(guard) => guard, - Err(conflicting_request) => { - return Err(format!( - "Canister {canister_id} is currently locked by another change operation. Conflicting request: {conflicting_request:?}" - )); - } - }; - - if stop_before_installing { - let stop_result = stop_canister::(canister_id).await; - if stop_result.is_err() { - println!("{LOG_PREFIX}change_canister: Failed to stop canister, trying to restart..."); - return match start_canister::(canister_id).await { - Ok(_) => Err(format!( - "Failed to stop canister {canister_id:?}. After failing to stop, attempted to start it, and succeeded in that." - )), - Err(_) => { - println!("{LOG_PREFIX}change_canister: Failed to restart canister."); - Err(format!( - "Failed to stop canister {canister_id:?}. After failing to stop, attempted to start it, and failed in that." - )) - } - }; - } - } - let request_str = format!("{request:?}"); // Ship code to the canister. // - // Note that there's no guarantee that the canister to install/reinstall/upgrade - // is actually stopped here, even if stop_before_installing is true. This is - // because there could be a concurrent request to restart it. This could be - // guaranteed with a "stopped precondition" in the management canister, or - // with some locking here. - let res = install_code(request).await; - // For once, we don't want to unwrap the result here. The reason is that, if the - // installation failed (e.g., the wasm was rejected because it's invalid), - // then we want to restart the canister. So we just keep the res to be - // unwrapped later. - - // Restart the canister, if needed - if stop_before_installing { - start_canister::(canister_id).await.unwrap(); - } + let result = exclusively_stop_and_start_canister::<_, _, _>( + canister_id, + &request_str, + stop_before_installing, + || async { + // Note that there's no guarantee that the canister to install/reinstall/upgrade + // is actually stopped here, even if stop_before_installing is true. This is + // because there could be a concurrent request to restart it. This could be + // guaranteed with a "stopped precondition" in the management canister, or + // with some locking here. + install_code(request).await + // For once, we don't want to unwrap the result here. The reason is that, if the + // installation failed (e.g., the wasm was rejected because it's invalid), + // then we want to restart the canister. So we just keep the res to be + // unwrapped later. + }, + ) + .await; - // Check the result of the install_code - res.map_err(|(rejection_code, message)| { + let result = match result { + Ok(ok) => ok, + Err(err) => { + return Err(format!("{err}")); + } + }; + + result.map_err(|(rejection_code, message)| { format!( "Attempt to call install_code with request {request_str} failed with code \ {rejection_code:?}: {message}" @@ -480,10 +461,11 @@ mod tests { }; // Manually insert a lock for this canister to simulate a concurrent operation - CANISTER_CHANGE_LOCKS.with(|locks| { + // Manually insert a lock for this canister to simulate a concurrent operation + crate::private::CANISTER_LOCK_ACQUISITIONS.with(|locks| { locks .borrow_mut() - .insert(canister_id, conflicting_request.clone()); + .insert(canister_id, format!("{conflicting_request:?}")); }); // Now try to call change_canister on the same canister - this should fail diff --git a/rs/nns/handlers/root/impl/canister/canister.rs b/rs/nns/handlers/root/impl/canister/canister.rs index 3db4cd6e7b8e..54508786462a 100644 --- a/rs/nns/handlers/root/impl/canister/canister.rs +++ b/rs/nns/handlers/root/impl/canister/canister.rs @@ -147,7 +147,7 @@ fn change_nns_canister(request: ChangeCanisterRequest) { // Because change_canister is async, and because we can't directly use // `await`, we need to use the `spawn` trick. let future = async move { - let change_canister_result = change_canister::(request).await; + let change_canister_result = change_canister(request).await; match change_canister_result { Ok(()) => { println!("{LOG_PREFIX}change_canister: Canister change completed successfully."); diff --git a/rs/nns/handlers/root/impl/src/root_proposals.rs b/rs/nns/handlers/root/impl/src/root_proposals.rs index ccfba9fa3bd2..50ec5a680e69 100644 --- a/rs/nns/handlers/root/impl/src/root_proposals.rs +++ b/rs/nns/handlers/root/impl/src/root_proposals.rs @@ -11,7 +11,6 @@ use ic_nervous_system_root::{ LOG_PREFIX, change_canister::{ChangeCanisterRequest, change_canister}, }; -use ic_nervous_system_runtime::CdkRuntime; use ic_nns_common::registry::get_value; use ic_nns_constants::{GOVERNANCE_CANISTER_ID, REGISTRY_CANISTER_ID}; use ic_protobuf::registry::{ @@ -419,7 +418,7 @@ pub async fn vote_on_root_proposal_to_upgrade_governance_canister( println!("{message}"); return Err(message); } - let _ = change_canister::(payload).await; + let _ = change_canister(payload).await; Ok(()) } else if proposal.is_byzantine_majority_no() { PROPOSALS.with(|proposals| proposals.borrow_mut().remove(&proposer)); diff --git a/rs/sns/root/canister/canister.rs b/rs/sns/root/canister/canister.rs index bc8dd82f4841..37bd4e7ad446 100644 --- a/rs/sns/root/canister/canister.rs +++ b/rs/sns/root/canister/canister.rs @@ -233,7 +233,7 @@ fn change_canister(request: ChangeCanisterRequest) { // spawn to do the real work in the background. CanisterRuntime::spawn_future(async move { let change_canister_result = - ic_nervous_system_root::change_canister::change_canister::(request) + ic_nervous_system_root::change_canister::change_canister(request) .await; // We don't want to panic in here, or the log messages will be lost when // the state rolls back. From 545518b2651370cfe19db3446e8d303620f6e1eb Mon Sep 17 00:00:00 2001 From: Daniel Wong Date: Tue, 20 Jan 2026 19:25:45 +0100 Subject: [PATCH 11/12] lint --- .../root/src/load_canister_snapshot.rs | 22 ++++----- rs/nervous_system/root/src/private.rs | 47 ++++++++++--------- .../root/src/take_canister_snapshot.rs | 4 +- .../root/impl/src/canister_management.rs | 3 +- rs/sns/root/canister/canister.rs | 3 +- 5 files changed, 38 insertions(+), 41 deletions(-) diff --git a/rs/nervous_system/root/src/load_canister_snapshot.rs b/rs/nervous_system/root/src/load_canister_snapshot.rs index e03a456f638e..cb5d62d4e1d6 100644 --- a/rs/nervous_system/root/src/load_canister_snapshot.rs +++ b/rs/nervous_system/root/src/load_canister_snapshot.rs @@ -1,4 +1,4 @@ -use crate::private::{exclusively_stop_and_start_canister}; +use crate::private::exclusively_stop_and_start_canister; use candid::CandidType; use ic_base_types::{CanisterId, PrincipalId, SnapshotId}; use ic_management_canister_types_private::LoadCanisterSnapshotArgs; @@ -79,22 +79,18 @@ pub async fn load_canister_snapshot( Ok(ok) => ok, Err(err) => { let description = format!("{err}"); - return LoadCanisterSnapshotResponse::Err( - LoadCanisterSnapshotError { - code: None, - description, - } - ); + return LoadCanisterSnapshotResponse::Err(LoadCanisterSnapshotError { + code: None, + description, + }); } }; match result { Ok(()) => LoadCanisterSnapshotResponse::Ok(LoadCanisterSnapshotOk {}), - Err((code, description)) => { - LoadCanisterSnapshotResponse::Err(LoadCanisterSnapshotError { - code: Some(code), - description, - }) - } + Err((code, description)) => LoadCanisterSnapshotResponse::Err(LoadCanisterSnapshotError { + code: Some(code), + description, + }), } } diff --git a/rs/nervous_system/root/src/private.rs b/rs/nervous_system/root/src/private.rs index 0cf22c64afc9..1150ccc8ed36 100644 --- a/rs/nervous_system/root/src/private.rs +++ b/rs/nervous_system/root/src/private.rs @@ -1,4 +1,7 @@ -use crate::{LOG_PREFIX, change_canister::{stop_canister, start_canister}}; +use crate::{ + LOG_PREFIX, + change_canister::{start_canister, stop_canister}, +}; use dfn_core::api::CanisterId; use ic_nervous_system_lock::acquire_for; use ic_nervous_system_runtime::CdkRuntime; @@ -19,7 +22,7 @@ thread_local! { >> = const { RefCell::new(BTreeMap::new()) }; } -#[derive(Debug, Clone, PartialEq, Eq,)] +#[derive(Debug, Clone, PartialEq, Eq)] pub(crate) struct Reject { pub code: i32, pub message: String, @@ -29,10 +32,7 @@ impl From<(i32, String)> for Reject { fn from(src: (i32, String)) -> Self { let (code, message) = src; - Self { - code, - message, - } + Self { code, message } } } @@ -55,37 +55,42 @@ pub(crate) enum ExclusivelyStopAndStartCanisterError { } pub(crate) use ExclusivelyStopAndStartCanisterError::{ - LockAcquisitionFailed, StopBeforeMainOperationFailed, StartAfterMainOperationFailed, + LockAcquisitionFailed, StartAfterMainOperationFailed, StopBeforeMainOperationFailed, }; impl Display for ExclusivelyStopAndStartCanisterError { fn fmt(&self, formatter: &mut Formatter) -> fmt::Result { let result = match self { - LockAcquisitionFailed { ongoing_operation_description } => { + LockAcquisitionFailed { + ongoing_operation_description, + } => { format!( "Another operation is currently in progress \ (on the same canister): {ongoing_operation_description}.", ) } - StopBeforeMainOperationFailed { stop_reject, restart_result } => { - let main_reason = format!( - "Failed to stop the canister beforehand: {stop_reject:?}" - ); + StopBeforeMainOperationFailed { + stop_reject, + restart_result, + } => { + let main_reason = + format!("Failed to stop the canister beforehand: {stop_reject:?}"); let restart_status = match restart_result { Ok(()) => "fortunately, re-starting the canister succeeded".to_string(), Err(restart_err) => { - format!( - "unfortunately, re-starting the canister failed: {restart_err:?}", - ) + format!("unfortunately, re-starting the canister failed: {restart_err:?}",) } }; format!("{main_reason}; {restart_status}.") } - StartAfterMainOperationFailed { start_reject, main_operation_result } => { + StartAfterMainOperationFailed { + start_reject, + main_operation_result, + } => { let start_failure = format!( "Unable to start the canister again after the main operation: \ {start_reject:?}", @@ -131,7 +136,9 @@ where canister_id, operation_description.to_string(), ) - .map_err(|ongoing_operation_description| LockAcquisitionFailed { ongoing_operation_description })?; + .map_err(|ongoing_operation_description| LockAcquisitionFailed { + ongoing_operation_description, + })?; if stop_before { stop_before_main_operation(canister_id, operation_description).await?; @@ -216,12 +223,10 @@ async fn restart_after_main_operation( main_operation_result: R, ) -> Result where - R: Debug + R: Debug, { match start_canister::(canister_id).await { - Ok(()) => { - Ok(main_operation_result) - } + Ok(()) => Ok(main_operation_result), Err(err) => { // This would be pretty terrible, especially if Governance is the diff --git a/rs/nervous_system/root/src/take_canister_snapshot.rs b/rs/nervous_system/root/src/take_canister_snapshot.rs index 31278923bdf7..282b75248be4 100644 --- a/rs/nervous_system/root/src/take_canister_snapshot.rs +++ b/rs/nervous_system/root/src/take_canister_snapshot.rs @@ -1,9 +1,7 @@ use crate::private::exclusively_stop_and_start_canister; use candid::CandidType; use ic_base_types::{CanisterId, PrincipalId, SnapshotId}; -use ic_management_canister_types_private::{ - CanisterSnapshotResponse, TakeCanisterSnapshotArgs, -}; +use ic_management_canister_types_private::{CanisterSnapshotResponse, TakeCanisterSnapshotArgs}; use ic_nervous_system_clients::management_canister_client::ManagementCanisterClient; use serde::Deserialize; diff --git a/rs/nns/handlers/root/impl/src/canister_management.rs b/rs/nns/handlers/root/impl/src/canister_management.rs index d473ab284d0c..0cbbc800a547 100644 --- a/rs/nns/handlers/root/impl/src/canister_management.rs +++ b/rs/nns/handlers/root/impl/src/canister_management.rs @@ -24,8 +24,7 @@ use ic_nns_common::{ }; use ic_nns_handler_root_interface::{ ChangeCanisterControllersRequest, ChangeCanisterControllersResponse, - UpdateCanisterSettingsError, UpdateCanisterSettingsRequest, - UpdateCanisterSettingsResponse, + UpdateCanisterSettingsError, UpdateCanisterSettingsRequest, UpdateCanisterSettingsResponse, }; use ic_protobuf::{ registry::nns::v1::{NnsCanisterRecord, NnsCanisterRecords}, diff --git a/rs/sns/root/canister/canister.rs b/rs/sns/root/canister/canister.rs index 37bd4e7ad446..477566d05b44 100644 --- a/rs/sns/root/canister/canister.rs +++ b/rs/sns/root/canister/canister.rs @@ -233,8 +233,7 @@ fn change_canister(request: ChangeCanisterRequest) { // spawn to do the real work in the background. CanisterRuntime::spawn_future(async move { let change_canister_result = - ic_nervous_system_root::change_canister::change_canister(request) - .await; + ic_nervous_system_root::change_canister::change_canister(request).await; // We don't want to panic in here, or the log messages will be lost when // the state rolls back. match change_canister_result { From a058b3549c66148d3d4b74367ab820b8bdf6a780 Mon Sep 17 00:00:00 2001 From: Daniel Wong Date: Tue, 20 Jan 2026 19:51:49 +0100 Subject: [PATCH 12/12] Fix test. Problem was that I removed a generic parameter. This unlocked the opportunity to delete MockRuntime. --- rs/nervous_system/root/src/change_canister.rs | 70 +++---------------- 1 file changed, 8 insertions(+), 62 deletions(-) diff --git a/rs/nervous_system/root/src/change_canister.rs b/rs/nervous_system/root/src/change_canister.rs index c2b9aa679fef..d6ac7496d194 100644 --- a/rs/nervous_system/root/src/change_canister.rs +++ b/rs/nervous_system/root/src/change_canister.rs @@ -385,66 +385,7 @@ where #[cfg(test)] mod tests { use super::*; - use async_trait::async_trait; - use candid::utils::{ArgumentDecoder, ArgumentEncoder}; use dfn_core::api::CanisterId; - use std::future::Future; - - // Mock runtime that returns errors for all inter-canister calls - // This allows us to test the locking behavior without actually making calls - struct MockRuntime; - - #[async_trait] - impl Runtime for MockRuntime { - async fn call_without_cleanup( - _id: CanisterId, - _method: &str, - _args: In, - ) -> Result - where - In: ArgumentEncoder + Send, - Out: for<'a> ArgumentDecoder<'a>, - { - Err(( - 1, - "MockRuntime: call_without_cleanup not implemented".to_string(), - )) - } - - async fn call_with_cleanup( - _id: CanisterId, - _method: &str, - _args: In, - ) -> Result - where - In: ArgumentEncoder + Send, - Out: for<'a> ArgumentDecoder<'a>, - { - Err(( - 1, - "MockRuntime: call_with_cleanup not implemented".to_string(), - )) - } - - async fn call_bytes_with_cleanup( - _id: CanisterId, - _method: &str, - _args: &[u8], - ) -> Result, (i32, String)> { - Err(( - 1, - "MockRuntime: call_bytes_with_cleanup not implemented".to_string(), - )) - } - - fn spawn_future>(_future: F) { - // Do nothing - we don't need to actually spawn - } - - fn canister_version() -> u64 { - 1 - } - } #[tokio::test] async fn test_change_canister_fails_when_lock_exists() { @@ -478,12 +419,17 @@ mod tests { arg: vec![16, 17, 18], }; - let result = change_canister::(new_request).await; + let result = change_canister(new_request).await; // Should return an error indicating the canister is locked assert!(result.is_err()); let error_msg = result.unwrap_err(); - assert!(error_msg.contains("currently locked by another change operation")); - assert!(error_msg.contains(&format!("{canister_id}"))); + for key_word in ["another operation", "in progress"] { + assert!( + error_msg.to_lowercase().contains(key_word), + "{key_word:?} not in {error_msg:?}" + ); + } + assert!(error_msg.contains(&format!("{canister_id}")), "{error_msg}"); } }