Added support for load_canister_snapshot to rs/nervous_system/clients. This is analogous to the changes in PR 8158.

Author: daniel-wong-dfinity-org

rs/nervous_system/clients/src/lib.rs

@@ -3,6 +3,7 @@ pub mod canister_metadata;
 pub mod canister_status;
 pub mod delete_canister;
 pub mod ledger_client;
+pub mod load_canister_snapshot;
 pub mod management_canister_client;
 pub mod stop_canister;
 pub mod take_canister_snapshot;

rs/nervous_system/clients/src/load_canister_snapshot.rs

@@ -0,0 +1,9 @@
+use ic_management_canister_types_private::{IC_00, LoadCanisterSnapshotArgs};
+use ic_nervous_system_runtime::Runtime;
+
+pub async fn load_canister_snapshot<Rt>(args: LoadCanisterSnapshotArgs) -> Result<(), (i32, String)>
+where
+    Rt: Runtime,
+{
+    Rt::call_with_cleanup(IC_00, "load_canister_snapshot", (args,)).await
+}

rs/nervous_system/clients/src/management_canister_client.rs

@@ -3,6 +3,7 @@ use crate::{
     canister_metadata::canister_metadata,
     canister_status::{CanisterStatusResultFromManagementCanister, canister_status},
     delete_canister::delete_canister,
+    load_canister_snapshot::load_canister_snapshot,
     stop_canister::stop_canister,
     take_canister_snapshot::take_canister_snapshot,
     update_settings::{UpdateSettings, update_settings},
@@ -12,7 +13,7 @@ use candid::Encode;
 use ic_base_types::PrincipalId;
 use ic_error_types::RejectCode;
 use ic_management_canister_types_private::{
-    CanisterSnapshotResponse, IC_00, TakeCanisterSnapshotArgs,
+    CanisterSnapshotResponse, IC_00, LoadCanisterSnapshotArgs, TakeCanisterSnapshotArgs,
 };
 use ic_nervous_system_proxied_canister_calls_tracker::ProxiedCanisterCallsTracker;
 use ic_nervous_system_runtime::Runtime;
@@ -62,6 +63,11 @@ pub trait ManagementCanisterClient {
         &self,
         args: TakeCanisterSnapshotArgs,
     ) -> Result<CanisterSnapshotResponse, (i32, String)>;
+
+    async fn load_canister_snapshot(
+        &self,
+        args: LoadCanisterSnapshotArgs,
+    ) -> Result<(), (i32, String)>;
 }
 
 /// An example implementation of the ManagementCanisterClient trait.
@@ -197,6 +203,24 @@ impl<Rt: Runtime + Sync> ManagementCanisterClient for ManagementCanisterClientIm
 
         take_canister_snapshot::<Rt>(args).await
     }
+
+    async fn load_canister_snapshot(
+        &self,
+        args: LoadCanisterSnapshotArgs,
+    ) -> Result<(), (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,
+                "load_canister_snapshot",
+                &encoded_args,
+            )
+        });
+
+        load_canister_snapshot::<Rt>(args).await
+    }
 }
 
 /// A ManagementCanisterClient that wraps another ManagementCanisterClient.
@@ -322,6 +346,14 @@ where
         let _loan = self.try_borrow_slot()?;
         self.inner.take_canister_snapshot(args).await
     }
+
+    async fn load_canister_snapshot(
+        &self,
+        args: LoadCanisterSnapshotArgs,
+    ) -> Result<(), (i32, String)> {
+        let _loan = self.try_borrow_slot()?;
+        self.inner.load_canister_snapshot(args).await
+    }
 }
 
 /// Increments available_slot_count by used_slot_count when dropped.
@@ -376,6 +408,7 @@ pub enum MockManagementCanisterClientCall {
     StopCanister(CanisterIdRecord),
     DeleteCanister(CanisterIdRecord),
     TakeCanisterSnapshot(TakeCanisterSnapshotArgs),
+    LoadCanisterSnapshot(LoadCanisterSnapshotArgs),
 }
 
 #[derive(Clone, Eq, PartialEq, Debug)]
@@ -387,6 +420,7 @@ pub enum MockManagementCanisterClientReply {
     StopCanister(Result<(), (i32, String)>),
     DeleteCanister(Result<(), (i32, String)>),
     TakeCanisterSnapshot(Result<CanisterSnapshotResponse, (i32, String)>),
+    LoadCanisterSnapshot(Result<(), (i32, String)>),
 }
 
 #[async_trait]
@@ -555,6 +589,32 @@ impl ManagementCanisterClient for MockManagementCanisterClient {
             ),
         }
     }
+
+    async fn load_canister_snapshot(
+        &self,
+        args: LoadCanisterSnapshotArgs,
+    ) -> Result<(), (i32, String)> {
+        self.calls
+            .lock()
+            .unwrap()
+            .push_back(MockManagementCanisterClientCall::LoadCanisterSnapshot(args));
+
+        let reply = self
+            .replies
+            .lock()
+            .unwrap()
+            .pop_front()
+            .expect("Expected a MockManagementCanisterClientCall to be on the queue.");
+
+        match reply {
+            MockManagementCanisterClientReply::LoadCanisterSnapshot(result) => result,
+            err => panic!(
+                "Expected MockManagementCanisterClientReply::LoadCanisterSnapshot to be at \
+                the front of the queue. Had {:?}",
+                err
+            ),
+        }
+    }
 }
 
 impl Drop for MockManagementCanisterClient {

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

@@ -5,6 +5,9 @@ use crate::canister_status::{
 };
 use candid::Nat;
 use ic_base_types::{CanisterId, PrincipalId};
+use ic_management_canister_types_private::{
+    CanisterSnapshotResponse, LoadCanisterSnapshotArgs, TakeCanisterSnapshotArgs,
+};
 use rand::{Rng, thread_rng};
 use std::time::Duration;
 
@@ -91,6 +94,13 @@ async fn test_limit_outstanding_calls() {
             ) -> Result<CanisterSnapshotResponse, (i32, String)> {
                 unimplemented!();
             }
+
+            async fn load_canister_snapshot(
+                &self,
+                _args: LoadCanisterSnapshotArgs,
+            ) -> Result<(), (i32, String)> {
+                unimplemented!();
+            }
         }
 
         impl Drop for MockManagementCanisterClient {