diff --git a/rs/nervous_system/clients/src/lib.rs b/rs/nervous_system/clients/src/lib.rs index 10d7db52578c..6434cd2cc570 100644 --- a/rs/nervous_system/clients/src/lib.rs +++ b/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; diff --git a/rs/nervous_system/clients/src/management_canister_client.rs b/rs/nervous_system/clients/src/management_canister_client.rs index 15d857bd9a05..1e1ff6780a93 100644 --- a/rs/nervous_system/clients/src/management_canister_client.rs +++ b/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; } /// An example implementation of the ManagementCanisterClient trait. @@ -171,6 +179,24 @@ impl ManagementCanisterClient for ManagementCanisterClientIm delete_canister::(canister_id_record).await } + + async fn take_canister_snapshot( + &self, + args: TakeCanisterSnapshotArgs, + ) -> Result { + 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::(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 { + 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, (i32, String)>), StopCanister(Result<(), (i32, String)>), DeleteCanister(Result<(), (i32, String)>), + TakeCanisterSnapshot(Result), } #[async_trait] @@ -493,6 +529,32 @@ impl ManagementCanisterClient for MockManagementCanisterClient { ), } } + + async fn take_canister_snapshot( + &self, + args: TakeCanisterSnapshotArgs, + ) -> Result { + 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 { diff --git a/rs/nervous_system/clients/src/management_canister_client/tests.rs b/rs/nervous_system/clients/src/management_canister_client/tests.rs index 3cde2358dbea..fbefab470b1f 100644 --- a/rs/nervous_system/clients/src/management_canister_client/tests.rs +++ b/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 { + unimplemented!(); + } } impl Drop for MockManagementCanisterClient { diff --git a/rs/nervous_system/clients/src/take_canister_snapshot.rs b/rs/nervous_system/clients/src/take_canister_snapshot.rs new file mode 100644 index 000000000000..4262e23ba1c1 --- /dev/null +++ b/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( + args: TakeCanisterSnapshotArgs, +) -> Result +where + Rt: Runtime, +{ + let (res,) = Rt::call_with_cleanup(IC_00, "take_canister_snapshot", (args,)).await?; + Ok(res) +}