feat(nns-root): Added take_canister_snapshot to Root canister.

rs/nervous_system/clients/src/lib.rs

@@ -5,6 +5,7 @@ pub mod delete_canister;
 pub mod ledger_client;
 pub mod management_canister_client;
 pub mod stop_canister;
+pub mod take_canister_snapshot;
 pub mod update_settings;
 
 mod request;

rs/nervous_system/clients/src/management_canister_client.rs

@@ -4,13 +4,16 @@ use crate::{
     canister_status::{CanisterStatusResultFromManagementCanister, canister_status},
     delete_canister::delete_canister,
     stop_canister::stop_canister,
+    take_canister_snapshot::take_canister_snapshot,
     update_settings::{UpdateSettings, update_settings},
 };
 use async_trait::async_trait;
 use candid::Encode;
 use ic_base_types::PrincipalId;
 use ic_error_types::RejectCode;
-use ic_management_canister_types_private::IC_00;
+use ic_management_canister_types_private::{
+    CanisterSnapshotResponse, IC_00, TakeCanisterSnapshotArgs,
+};
 use ic_nervous_system_proxied_canister_calls_tracker::ProxiedCanisterCallsTracker;
 use ic_nervous_system_runtime::Runtime;
 use std::{
@@ -54,6 +57,11 @@ pub trait ManagementCanisterClient {
         &self,
         canister_id_record: CanisterIdRecord,
     ) -> Result<(), (i32, String)>;
+
+    async fn take_canister_snapshot(
+        &self,
+        args: TakeCanisterSnapshotArgs,
+    ) -> Result<CanisterSnapshotResponse, (i32, String)>;
 }
 
 /// An example implementation of the ManagementCanisterClient trait.
@@ -171,6 +179,24 @@ impl<Rt: Runtime + Sync> ManagementCanisterClient for ManagementCanisterClientIm
 
         delete_canister::<Rt>(canister_id_record).await
     }
+
+    async fn take_canister_snapshot(
+        &self,
+        args: TakeCanisterSnapshotArgs,
+    ) -> Result<CanisterSnapshotResponse, (i32, String)> {
+        let _tracker = self.proxied_canister_calls_tracker.map(|tracker| {
+            let encoded_args = Encode!(&args).unwrap_or_default();
+            ProxiedCanisterCallsTracker::start_tracking(
+                tracker,
+                dfn_core::api::caller(),
+                IC_00,
+                "take_canister_snapshot",
+                &encoded_args,
+            )
+        });
+
+        take_canister_snapshot::<Rt>(args).await
+    }
 }
 
 /// A ManagementCanisterClient that wraps another ManagementCanisterClient.
@@ -288,6 +314,14 @@ where
         let _loan = self.try_borrow_slot()?;
         self.inner.delete_canister(canister_id_record).await
     }
+
+    async fn take_canister_snapshot(
+        &self,
+        args: TakeCanisterSnapshotArgs,
+    ) -> Result<CanisterSnapshotResponse, (i32, String)> {
+        let _loan = self.try_borrow_slot()?;
+        self.inner.take_canister_snapshot(args).await
+    }
 }
 
 /// Increments available_slot_count by used_slot_count when dropped.
@@ -341,6 +375,7 @@ pub enum MockManagementCanisterClientCall {
     CanisterMetadata(PrincipalId, String),
     StopCanister(CanisterIdRecord),
     DeleteCanister(CanisterIdRecord),
+    TakeCanisterSnapshot(TakeCanisterSnapshotArgs),
 }
 
 #[derive(Clone, Eq, PartialEq, Debug)]
@@ -351,6 +386,7 @@ pub enum MockManagementCanisterClientReply {
     CanisterMetadata(Result<Vec<u8>, (i32, String)>),
     StopCanister(Result<(), (i32, String)>),
     DeleteCanister(Result<(), (i32, String)>),
+    TakeCanisterSnapshot(Result<CanisterSnapshotResponse, (i32, String)>),
 }
 
 #[async_trait]
@@ -493,6 +529,32 @@ impl ManagementCanisterClient for MockManagementCanisterClient {
             ),
         }
     }
+
+    async fn take_canister_snapshot(
+        &self,
+        args: TakeCanisterSnapshotArgs,
+    ) -> Result<CanisterSnapshotResponse, (i32, String)> {
+        self.calls
+            .lock()
+            .unwrap()
+            .push_back(MockManagementCanisterClientCall::TakeCanisterSnapshot(args));
+
+        let reply = self
+            .replies
+            .lock()
+            .unwrap()
+            .pop_front()
+            .expect("Expected a MockManagementCanisterClientCall to be on the queue.");
+
+        match reply {
+            MockManagementCanisterClientReply::TakeCanisterSnapshot(result) => result,
+            err => panic!(
+                "Expected MockManagementCanisterClientReply::TakeCanisterSnapshot to be at \
+                the front of the queue. Had {:?}",
+                err
+            ),
+        }
+    }
 }
 
 impl Drop for MockManagementCanisterClient {

rs/nervous_system/clients/src/management_canister_client/tests.rs

@@ -84,6 +84,13 @@ async fn test_limit_outstanding_calls() {
             ) -> Result<(), (i32, String)> {
                 unimplemented!();
             }
+
+            async fn take_canister_snapshot(
+                &self,
+                _args: TakeCanisterSnapshotArgs,
+            ) -> Result<CanisterSnapshotResponse, (i32, String)> {
+                unimplemented!();
+            }
         }
 
         impl Drop for MockManagementCanisterClient {

rs/nervous_system/clients/src/take_canister_snapshot.rs

@@ -0,0 +1,14 @@
+use ic_management_canister_types_private::{
+    CanisterSnapshotResponse, IC_00, TakeCanisterSnapshotArgs,
+};
+use ic_nervous_system_runtime::Runtime;
+
+pub async fn take_canister_snapshot<Rt>(
+    args: TakeCanisterSnapshotArgs,
+) -> Result<CanisterSnapshotResponse, (i32, String)>
+where
+    Rt: Runtime,
+{
+    let (res,) = Rt::call_with_cleanup(IC_00, "take_canister_snapshot", (args,)).await?;
+    Ok(res)
+}

rs/nns/handlers/root/impl/canister/canister.rs

@@ -30,7 +30,8 @@ use ic_nns_handler_root::{
 };
 use ic_nns_handler_root_interface::{
     ChangeCanisterControllersRequest, ChangeCanisterControllersResponse,
-    UpdateCanisterSettingsRequest, UpdateCanisterSettingsResponse,
+    TakeCanisterSnapshotRequest, TakeCanisterSnapshotResponse, UpdateCanisterSettingsRequest,
+    UpdateCanisterSettingsResponse,
 };
 use std::cell::RefCell;
 
@@ -227,6 +228,20 @@ async fn update_canister_settings(
     .await
 }
 
+/// Takes a snapshot of a canister controlled by NNS Root. Only callable by NNS
+/// Governance.
+#[update]
+async fn take_canister_snapshot(
+    take_canister_snapshot_request: TakeCanisterSnapshotRequest,
+) -> TakeCanisterSnapshotResponse {
+    check_caller_is_governance();
+    canister_management::take_canister_snapshot(
+        take_canister_snapshot_request,
+        &mut new_management_canister_client(),
+    )
+    .await
+}
+
 /// Resources to serve for a given http_request
 /// Serve an HttpRequest made to this canister
 #[query(

rs/nns/handlers/root/impl/canister/root.did

@@ -151,6 +151,27 @@ type CallCanisterRequest = record {
   payload : blob;
 };
 
+type TakeCanisterSnapshotArgs = record {
+  canister_id : principal;
+  replace_snapshot : opt blob;
+};
+
+type TakeCanisterSnapshotResponse = variant {
+  Ok : TakeCanisterSnapshotOk;
+  Err : TakeCanisterSnapshotError;
+};
+
+type TakeCanisterSnapshotOk = record {
+  id : blob;
+  taken_at_timestamp : nat64;
+  total_size : nat64;
+};
+
+type TakeCanisterSnapshotError = record {
+  code : opt int32;
+  description : text;
+};
+
 service : () -> {
   canister_status : (CanisterIdRecord) -> (CanisterStatusResult);
   get_build_metadata : () -> (text) query;
@@ -166,4 +187,7 @@ service : () -> {
       UpdateCanisterSettingsResponse,
     );
   call_canister : (CallCanisterRequest) -> ();
+  take_canister_snapshot : (TakeCanisterSnapshotArgs) -> (
+      TakeCanisterSnapshotResponse,
+    );
 }

rs/nns/handlers/root/impl/src/canister_management.rs

@@ -1,12 +1,13 @@
 #![allow(deprecated)]
 use crate::PROXIED_CANISTER_CALLS_TRACKER;
-use ic_base_types::{CanisterId, PrincipalId};
+use ic_base_types::{CanisterId, PrincipalId, SnapshotId};
 use ic_cdk::{
     api::call::{RejectionCode, call_with_payment},
     call, caller, print,
 };
 use ic_management_canister_types_private::{
-    CanisterInstallMode::Install, CanisterSettingsArgsBuilder, CreateCanisterArgs, InstallCodeArgs,
+    CanisterInstallMode::Install, CanisterSettingsArgsBuilder, CanisterSnapshotResponse,
+    CreateCanisterArgs, InstallCodeArgs, TakeCanisterSnapshotArgs,
 };
 use ic_nervous_system_clients::{
     canister_id_record::CanisterIdRecord,
@@ -23,7 +24,8 @@ use ic_nns_common::{
     types::CallCanisterRequest,
 };
 use ic_nns_handler_root_interface::{
-    ChangeCanisterControllersRequest, ChangeCanisterControllersResponse,
+    ChangeCanisterControllersRequest, ChangeCanisterControllersResponse, TakeCanisterSnapshotError,
+    TakeCanisterSnapshotOk, TakeCanisterSnapshotRequest, TakeCanisterSnapshotResponse,
     UpdateCanisterSettingsError, UpdateCanisterSettingsRequest, UpdateCanisterSettingsResponse,
 };
 use ic_protobuf::{
@@ -251,3 +253,71 @@ pub async fn update_canister_settings(
         }
     }
 }
+
+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_resposne_to_take_canister_snapshot_ok(result);
+            TakeCanisterSnapshotResponse::Ok(result)
+        }
+
+        Err((code, description)) => TakeCanisterSnapshotResponse::Err(TakeCanisterSnapshotError {
+            code: Some(code),
+            description,
+        }),
+    }
+}
+
+fn convert_from_canister_snapshot_resposne_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,
+    }
+}

rs/nns/handlers/root/interface/BUILD.bazel

@@ -7,6 +7,7 @@ DEPENDENCIES = [
     "//rs/nervous_system/clients",
     "//rs/nns/constants",
     "//rs/types/base_types",
+    "//rs/types/management_canister_types",
     "@crate_index//:candid",
     "@crate_index//:ic-cdk",
     "@crate_index//:serde",

rs/nns/handlers/root/interface/Cargo.toml

@@ -8,6 +8,7 @@ async-trait = { workspace = true }
 candid = { workspace = true }
 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-nns-constants = { path = "../../../../nns/constants" }
 serde = { workspace = true }

rs/nns/handlers/root/interface/src/lib.rs

@@ -82,3 +82,28 @@ pub struct UpdateCanisterSettingsError {
     pub code: Option<i32>,
     pub description: String,
 }
+
+#[derive(Clone, Eq, PartialEq, Hash, Debug, CandidType, Deserialize)]
+pub struct TakeCanisterSnapshotRequest {
+    pub canister_id: PrincipalId,
+    pub replace_snapshot: Option</* snapshot ID */ Vec<u8>>,
+}
+
+#[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<u8>,
+    pub taken_at_timestamp: u64,
+    pub total_size: u64,
+}
+
+#[derive(Clone, Eq, PartialEq, Hash, Debug, CandidType, Deserialize)]
+pub struct TakeCanisterSnapshotError {
+    pub code: Option<i32>,
+    pub description: String,
+}