Possibly Breaking Change: Snapshot API for all Canisters

src/canister/individual_user_template/can.did

@@ -624,7 +624,7 @@ service : (IndividualUserTemplateInitArgs) -> {
   http_request : (HttpRequest) -> (HttpResponse) query;
   list_namespace_keys : (nat64) -> (Result_18) query;
   list_namespaces : (nat64, nat64) -> (vec NamespaceForFrontend) query;
-  load_snapshot : (nat64) -> ();
+  load_snapshot : () -> ();
   once_reenqueue_timers_for_pending_bet_outcomes : () -> (Result_19);
   read_key_value_pair : (nat64, text) -> (Result_5) query;
   receive_and_save_snaphot : (nat64, blob) -> ();

src/canister/individual_user_template/src/api/snapshot/get_snapshot.rs

@@ -62,7 +62,7 @@ fn receive_and_save_snaphot(offset: u64, state_bytes: Vec<u8>) {
 }
 
 #[update(guard = "is_reclaim_canister_id")]
-fn load_snapshot(length: u64) {
+fn load_snapshot() {
     let state_bytes =
         SNAPSHOT_DATA.with(|snapshot_data_ref_cell| snapshot_data_ref_cell.borrow().clone());
 

src/canister/platform_orchestrator/Cargo.toml

@@ -17,6 +17,8 @@ ic-cdk-macros = { workspace = true }
 ic-stable-structures = { workspace = true }
 ciborium = { workspace = true }
 shared_utils = { workspace = true }
+serde_json = { workspace = true }
+serde_json_any_key = "2.0.0"
 
 [dev-dependencies]
 test_utils = { workspace = true }

src/canister/platform_orchestrator/can.did

@@ -82,13 +82,15 @@ type WasmType = variant {
 };
 service : (PlatformOrchestratorInitArgs) -> {
   add_principal_as_global_admin : (principal) -> ();
+  clear_snapshot : () -> ();
   collect_creator_dao_stats_in_the_network : () -> ();
   delete_all_sns_creator_token_in_the_network : () -> ();
   delete_all_sns_creator_token_of_an_individual_canister : (principal) -> (
       Result,
     );
   deposit_cycles_to_canister : (principal, nat) -> (Result_1);
   deregister_subnet_orchestrator : (principal, bool) -> ();
+  download_snapshot : (nat64, nat64) -> (blob) query;
   get_all_available_subnet_orchestrators : () -> (vec principal) query;
   get_all_global_admins : () -> (vec principal) query;
   get_all_subnet_orchestrators : () -> (vec principal) query;
@@ -101,6 +103,7 @@ service : (PlatformOrchestratorInitArgs) -> {
   get_subnets_upgrade_status_report : () -> (SubnetUpgradeReport) query;
   get_version : () -> (text) query;
   http_request : (HttpRequest) -> (HttpResponse) query;
+  load_snapshot : () -> ();
   make_individual_canister_logs_private : (principal) -> (Result);
   make_individual_canister_logs_public : (principal) -> (Result);
   make_subnet_orchestrator_logs_private : (principal) -> (Result);
@@ -115,6 +118,7 @@ service : (PlatformOrchestratorInitArgs) -> {
   populate_known_principal_for_all_subnet : () -> ();
   provision_empty_canisters_in_a_subnet : (principal, nat64) -> (Result);
   provision_subnet_orchestrator_canister : (principal) -> (Result_2);
+  receive_and_save_snaphot : (nat64, blob) -> ();
   receive_creator_dao_stats_from_subnet_orchestrator : (
       principal,
       vec principal,
@@ -126,6 +130,7 @@ service : (PlatformOrchestratorInitArgs) -> {
   remove_subnet_orchestrators_from_available_list : (principal) -> (Result_1);
   report_subnet_upgrade_status : (UpgradeStatus) -> (Result);
   reset_canisters_ml_feed_cache : () -> (Result_1);
+  save_snapshot_json : () -> (nat32);
   set_reserved_cycle_limit_for_subnet_orchestrator : (principal, nat) -> (
       Result,
     );

src/canister/platform_orchestrator/src/api/mod.rs

@@ -4,3 +4,4 @@ pub mod cycle_management;
 pub mod generic_proposal;
 pub mod monitoring;
 pub mod stats;
+pub mod snapshot;

src/canister/platform_orchestrator/src/api/snapshot/mod.rs

@@ -0,0 +1,67 @@
+use ic_cdk_macros::{query, update};
+
+use crate::{data_model::CanisterData, CANISTER_DATA, SNAPSHOT_DATA};
+use shared_utils::common::utils::permissions::is_reclaim_canister_id;
+
+#[update(guard = "is_reclaim_canister_id")]
+fn save_snapshot_json() -> u32 {
+    let mut state_bytes = vec![];
+
+    CANISTER_DATA.with(|canister_data_ref_cell| {
+        let serde_str = serde_json::to_string(&*canister_data_ref_cell.borrow()).unwrap();
+        state_bytes = serde_str.as_bytes().to_vec();
+    });
+
+    let len = state_bytes.len() as u32;
+
+    SNAPSHOT_DATA.with(|snapshot_data_ref_cell| {
+        *snapshot_data_ref_cell.borrow_mut() = state_bytes;
+    });
+
+    len
+}
+
+#[query(guard = "is_reclaim_canister_id")]
+fn download_snapshot(offset: u64, length: u64) -> Vec<u8> {
+    let state_bytes = SNAPSHOT_DATA.with(|snapshot_data_ref_cell| {
+        let snapshot = snapshot_data_ref_cell.borrow();
+
+        snapshot[offset as usize..(offset + length) as usize].to_vec()
+    });
+
+    state_bytes
+}
+
+#[update(guard = "is_reclaim_canister_id")]
+fn receive_and_save_snaphot(offset: u64, state_bytes: Vec<u8>) {
+    SNAPSHOT_DATA.with(|snapshot_data_ref_cell| {
+        let mut snapshot = snapshot_data_ref_cell.borrow_mut();
+        if snapshot.len() < (offset + state_bytes.len() as u64) as usize {
+            snapshot.resize((offset + state_bytes.len() as u64) as usize, 0);
+        }
+        snapshot.splice(
+            offset as usize..(offset + state_bytes.len() as u64) as usize,
+            state_bytes,
+        );
+    });
+}
+
+#[update(guard = "is_reclaim_canister_id")]
+fn load_snapshot() {
+    let state_bytes =
+        SNAPSHOT_DATA.with(|snapshot_data_ref_cell| snapshot_data_ref_cell.borrow().clone());
+
+    let canister_data_snapshot: CanisterData =
+        serde_json::from_str(std::str::from_utf8(&state_bytes).unwrap()).unwrap();
+
+    CANISTER_DATA.with(|canister_data_ref_cell| {
+        *canister_data_ref_cell.borrow_mut() = canister_data_snapshot;
+    });
+}
+
+#[update(guard = "is_reclaim_canister_id")]
+fn clear_snapshot() {
+    SNAPSHOT_DATA.with(|snapshot_data_ref_cell| {
+        *snapshot_data_ref_cell.borrow_mut() = vec![];
+    });
+}

src/canister/platform_orchestrator/src/lib.rs

@@ -28,6 +28,7 @@ mod utils;
 
 thread_local! {
     pub static CANISTER_DATA: RefCell<CanisterData> = RefCell::default();
+    pub static SNAPSHOT_DATA: RefCell<Vec<u8>> = RefCell::default();
 }
 
 export_candid!();

src/canister/post_cache/Cargo.toml

@@ -15,6 +15,7 @@ ic-cdk-timers = { workspace = true }
 shared_utils = { workspace = true }
 ic-cdk-macros = { workspace = true }
 serde = { workspace = true }
+serde_json = { workspace = true }
 
 [dev-dependencies]
 test_utils = { workspace = true }

src/canister/post_cache/src/lib.rs

@@ -20,6 +20,7 @@ mod util;
 
 thread_local! {
     static CANISTER_DATA: RefCell<CanisterData> = RefCell::default();
+    static SNAPSHOT_DATA: RefCell<Vec<u8>> = RefCell::default();
 }
 
 export_candid!();

src/canister/user_index/can.did

@@ -127,6 +127,7 @@ type UserIndexInitArgs = record {
 service : (UserIndexInitArgs) -> {
   allot_empty_canister : () -> (Result);
   are_signups_enabled : () -> (bool) query;
+  clear_snapshot : () -> ();
   collect_creator_dao_stats_in_the_network : () -> (Result_1);
   create_pool_of_individual_user_available_canisters : (text, blob) -> (
       Result_2,
@@ -135,6 +136,7 @@ service : (UserIndexInitArgs) -> {
   delete_all_sns_creator_token_of_an_individual_canister : (principal) -> (
       Result_3,
     );
+  download_snapshot : (nat64, nat64) -> (blob) query;
   get_current_list_of_all_well_known_principal_values : () -> (
       vec record { KnownPrincipalType; principal },
     ) query;
@@ -166,13 +168,15 @@ service : (UserIndexInitArgs) -> {
     ) query;
   http_request : (HttpRequest) -> (HttpResponse) query;
   issue_rewards_for_referral : (principal, principal, principal) -> (Result_2);
+  load_snapshot : () -> ();
   make_individual_canister_logs_private : (principal) -> (Result_3);
   make_individual_canister_logs_public : (principal) -> (Result_3);
   notify_specific_individual_canister_to_upgrade_creator_dao_governance_canisters : (
       principal,
       blob,
     ) -> (Result_3);
   provision_empty_canisters : (nat64) -> ();
+  receive_and_save_snaphot : (nat64, blob) -> ();
   receive_creator_dao_stats_from_individual_canister : (vec principal) -> (
       Result_3,
     );
@@ -185,6 +189,7 @@ service : (UserIndexInitArgs) -> {
   reset_user_canisters_ml_feed_cache : () -> (text);
   reset_user_individual_canisters : (vec principal) -> (Result_2);
   return_cycles_to_platform_orchestrator_canister : () -> (Result_2);
+  save_snapshot_json : () -> (nat32);
   set_permission_to_upgrade_individual_canisters : (bool) -> (text);
   start_upgrades_for_individual_canisters : (text, blob) -> (text);
   toggle_signups_enabled : () -> (Result_3);

src/canister/user_index/src/api/mod.rs

@@ -8,3 +8,4 @@ pub mod upgrade_individual_user_template;
 pub mod user_record;
 pub mod user_signup;
 pub mod well_known_principal;
+pub mod snapshot;

src/canister/user_index/src/api/snapshot/mod.rs

@@ -0,0 +1,67 @@
+use ic_cdk_macros::{query, update};
+
+use crate::{data_model::CanisterData, CANISTER_DATA, SNAPSHOT_DATA};
+use shared_utils::common::utils::permissions::is_reclaim_canister_id;
+
+#[update(guard = "is_reclaim_canister_id")]
+fn save_snapshot_json() -> u32 {
+    let mut state_bytes = vec![];
+
+    CANISTER_DATA.with(|canister_data_ref_cell| {
+        let serde_str = serde_json::to_string(&*canister_data_ref_cell.borrow()).unwrap();
+        state_bytes = serde_str.as_bytes().to_vec();
+    });
+
+    let len = state_bytes.len() as u32;
+
+    SNAPSHOT_DATA.with(|snapshot_data_ref_cell| {
+        *snapshot_data_ref_cell.borrow_mut() = state_bytes;
+    });
+
+    len
+}
+
+#[query(guard = "is_reclaim_canister_id")]
+fn download_snapshot(offset: u64, length: u64) -> Vec<u8> {
+    let state_bytes = SNAPSHOT_DATA.with(|snapshot_data_ref_cell| {
+        let snapshot = snapshot_data_ref_cell.borrow();
+
+        snapshot[offset as usize..(offset + length) as usize].to_vec()
+    });
+
+    state_bytes
+}
+
+#[update(guard = "is_reclaim_canister_id")]
+fn receive_and_save_snaphot(offset: u64, state_bytes: Vec<u8>) {
+    SNAPSHOT_DATA.with(|snapshot_data_ref_cell| {
+        let mut snapshot = snapshot_data_ref_cell.borrow_mut();
+        if snapshot.len() < (offset + state_bytes.len() as u64) as usize {
+            snapshot.resize((offset + state_bytes.len() as u64) as usize, 0);
+        }
+        snapshot.splice(
+            offset as usize..(offset + state_bytes.len() as u64) as usize,
+            state_bytes,
+        );
+    });
+}
+
+#[update(guard = "is_reclaim_canister_id")]
+fn load_snapshot() {
+    let state_bytes =
+        SNAPSHOT_DATA.with(|snapshot_data_ref_cell| snapshot_data_ref_cell.borrow().clone());
+
+    let canister_data_snapshot: CanisterData =
+        serde_json::from_str(std::str::from_utf8(&state_bytes).unwrap()).unwrap();
+
+    CANISTER_DATA.with(|canister_data_ref_cell| {
+        *canister_data_ref_cell.borrow_mut() = canister_data_snapshot;
+    });
+}
+
+#[update(guard = "is_reclaim_canister_id")]
+fn clear_snapshot() {
+    SNAPSHOT_DATA.with(|snapshot_data_ref_cell| {
+        *snapshot_data_ref_cell.borrow_mut() = vec![];
+    });
+}

src/canister/user_index/src/lib.rs

@@ -25,6 +25,7 @@ mod util;
 
 thread_local! {
     static CANISTER_DATA: RefCell<CanisterData> = RefCell::default();
+    static SNAPSHOT_DATA: RefCell<Vec<u8>> = RefCell::default();
 }
 
 export_candid!();