feat(upgrader): prune resources in disaster recovery via upgrader

core/upgrader/api/spec.did

@@ -184,10 +184,17 @@ type RequestDisasterRecoveryRestoreInput = record {
   // A snapshot to be restored.
   snapshot_id : text;
 };
+// The resource to prune.
+type RequestDisasterRecoveryPruneInput = variant {
+  chunk_store;
+  snapshot : text;
+  state;
+};
 type RequestDisasterRecoveryInput = variant {
   InstallCode : RequestDisasterRecoveryInstallCodeInput;
   Snapshot : RequestDisasterRecoverySnapshotInput;
   Restore : RequestDisasterRecoveryRestoreInput;
+  Prune : RequestDisasterRecoveryPruneInput;
 };
 
 type InstallMode = variant {
@@ -262,10 +269,18 @@ type StationRecoveryRequestRestoreOperation = record {
   snapshot_id : text;
 };
 
+// The resource to prune.
+type StationRecoveryRequestPruneOperation = variant {
+  chunk_store;
+  snapshot : text;
+  state;
+};
+
 type StationRecoveryRequestOperation = variant {
   InstallCode : StationRecoveryRequestInstallCodeOperation;
   Snapshot : StationRecoveryRequestSnapshotOperation;
   Restore : StationRecoveryRequestRestoreOperation;
+  Prune : StationRecoveryRequestPruneOperation;
 };
 
 // Request to recover the station.

core/upgrader/api/src/lib.rs

@@ -176,11 +176,22 @@ pub struct RequestDisasterRecoveryRestoreInput {
     pub snapshot_id: String,
 }
 
+#[derive(Clone, Debug, CandidType, Deserialize)]
+pub enum RequestDisasterRecoveryPruneInput {
+    #[serde(rename = "snapshot")]
+    Snapshot(String),
+    #[serde(rename = "chunk_store")]
+    ChunkStore,
+    #[serde(rename = "state")]
+    State,
+}
+
 #[derive(Clone, Debug, CandidType, Deserialize)]
 pub enum RequestDisasterRecoveryInput {
     InstallCode(RequestDisasterRecoveryInstallCodeInput),
     Snapshot(RequestDisasterRecoverySnapshotInput),
     Restore(RequestDisasterRecoveryRestoreInput),
+    Prune(RequestDisasterRecoveryPruneInput),
 }
 
 #[derive(CandidType, Deserialize, Debug, Clone)]
@@ -236,11 +247,22 @@ pub struct StationRecoveryRequestRestoreOperation {
     pub snapshot_id: String,
 }
 
+#[derive(Clone, Debug, CandidType, Deserialize, PartialEq, Eq)]
+pub enum StationRecoveryRequestPruneOperation {
+    #[serde(rename = "snapshot")]
+    Snapshot(String),
+    #[serde(rename = "chunk_store")]
+    ChunkStore,
+    #[serde(rename = "state")]
+    State,
+}
+
 #[derive(Clone, Debug, CandidType, Deserialize, PartialEq, Eq)]
 pub enum StationRecoveryRequestOperation {
     InstallCode(StationRecoveryRequestInstallCodeOperation),
     Snapshot(StationRecoveryRequestSnapshotOperation),
     Restore(StationRecoveryRequestRestoreOperation),
+    Prune(StationRecoveryRequestPruneOperation),
 }
 
 #[derive(Clone, Debug, CandidType, Deserialize, PartialEq, Eq)]

core/upgrader/impl/src/mappers/disaster_recovery.rs

@@ -1,10 +1,11 @@
 use crate::model::{
     DisasterRecovery, DisasterRecoveryV0, RequestDisasterRecoveryInstallCodeLog,
-    RequestDisasterRecoveryOperationLog, RequestDisasterRecoveryRestoreLog,
-    RequestDisasterRecoverySnapshotLog, StationRecoveryRequest,
+    RequestDisasterRecoveryOperationLog, RequestDisasterRecoveryPruneLog,
+    RequestDisasterRecoveryRestoreLog, RequestDisasterRecoverySnapshotLog, StationRecoveryRequest,
     StationRecoveryRequestInstallCodeOperation,
     StationRecoveryRequestInstallCodeOperationFootprint, StationRecoveryRequestOperation,
-    StationRecoveryRequestOperationFootprint, StationRecoveryRequestRestoreOperation,
+    StationRecoveryRequestOperationFootprint, StationRecoveryRequestPruneOperation,
+    StationRecoveryRequestPruneOperationFootprint, StationRecoveryRequestRestoreOperation,
     StationRecoveryRequestRestoreOperationFootprint, StationRecoveryRequestSnapshotOperation,
     StationRecoveryRequestSnapshotOperationFootprint, StationRecoveryRequestV0,
 };
@@ -45,6 +46,22 @@ impl From<upgrader_api::RequestDisasterRecoveryInput> for StationRecoveryRequest
                         .expect("Failed to parse `snapshot_id`"),
                 })
             }
+            upgrader_api::RequestDisasterRecoveryInput::Prune(prune) => {
+                let prune_op = match prune {
+                    upgrader_api::RequestDisasterRecoveryPruneInput::Snapshot(snapshot_id) => {
+                        StationRecoveryRequestPruneOperation::Snapshot(
+                            hex::decode(snapshot_id).expect("Failed to parse `snapshot_id`"),
+                        )
+                    }
+                    upgrader_api::RequestDisasterRecoveryPruneInput::ChunkStore => {
+                        StationRecoveryRequestPruneOperation::ChunkStore
+                    }
+                    upgrader_api::RequestDisasterRecoveryPruneInput::State => {
+                        StationRecoveryRequestPruneOperation::State
+                    }
+                };
+                StationRecoveryRequestOperation::Prune(prune_op)
+            }
         }
     }
 }
@@ -76,6 +93,22 @@ impl From<&StationRecoveryRequestOperation> for StationRecoveryRequestOperationF
                     },
                 )
             }
+            StationRecoveryRequestOperation::Prune(ref prune) => {
+                let prune_op = match prune {
+                    StationRecoveryRequestPruneOperation::Snapshot(snapshot_id) => {
+                        StationRecoveryRequestPruneOperationFootprint::Snapshot(hex::encode(
+                            snapshot_id,
+                        ))
+                    }
+                    StationRecoveryRequestPruneOperation::ChunkStore => {
+                        StationRecoveryRequestPruneOperationFootprint::ChunkStore
+                    }
+                    StationRecoveryRequestPruneOperation::State => {
+                        StationRecoveryRequestPruneOperationFootprint::State
+                    }
+                };
+                StationRecoveryRequestOperationFootprint::Prune(prune_op)
+            }
         }
     }
 }
@@ -103,6 +136,20 @@ impl From<&StationRecoveryRequestOperation> for RequestDisasterRecoveryOperation
                     snapshot_id: hex::encode(&snapshot.snapshot_id),
                 })
             }
+            StationRecoveryRequestOperation::Prune(ref prune) => {
+                let prune_op = match prune {
+                    StationRecoveryRequestPruneOperation::Snapshot(snapshot_id) => {
+                        RequestDisasterRecoveryPruneLog::Snapshot(hex::encode(snapshot_id))
+                    }
+                    StationRecoveryRequestPruneOperation::ChunkStore => {
+                        RequestDisasterRecoveryPruneLog::ChunkStore
+                    }
+                    StationRecoveryRequestPruneOperation::State => {
+                        RequestDisasterRecoveryPruneLog::State
+                    }
+                };
+                RequestDisasterRecoveryOperationLog::Prune(prune_op)
+            }
         }
     }
 }
@@ -134,6 +181,22 @@ impl From<&StationRecoveryRequestOperation> for upgrader_api::StationRecoveryReq
                     },
                 )
             }
+            StationRecoveryRequestOperation::Prune(ref prune) => {
+                let prune_op = match prune {
+                    StationRecoveryRequestPruneOperation::Snapshot(snapshot_id) => {
+                        upgrader_api::StationRecoveryRequestPruneOperation::Snapshot(hex::encode(
+                            snapshot_id,
+                        ))
+                    }
+                    StationRecoveryRequestPruneOperation::ChunkStore => {
+                        upgrader_api::StationRecoveryRequestPruneOperation::ChunkStore
+                    }
+                    StationRecoveryRequestPruneOperation::State => {
+                        upgrader_api::StationRecoveryRequestPruneOperation::State
+                    }
+                };
+                upgrader_api::StationRecoveryRequestOperation::Prune(prune_op)
+            }
         }
     }
 }

core/upgrader/impl/src/model/disaster_recovery.rs

@@ -94,12 +94,21 @@ pub struct StationRecoveryRequestRestoreOperation {
     pub snapshot_id: Vec<u8>,
 }
 
+#[storable]
+#[derive(Clone, Debug, Eq, PartialEq)]
+pub enum StationRecoveryRequestPruneOperation {
+    Snapshot(Vec<u8>),
+    ChunkStore,
+    State,
+}
+
 #[storable]
 #[derive(Clone, Debug, Eq, PartialEq)]
 pub enum StationRecoveryRequestOperation {
     InstallCode(StationRecoveryRequestInstallCodeOperation),
     Snapshot(StationRecoveryRequestSnapshotOperation),
     Restore(StationRecoveryRequestRestoreOperation),
+    Prune(StationRecoveryRequestPruneOperation),
 }
 
 #[derive(Clone, Debug, Eq, PartialEq, Hash)]
@@ -120,11 +129,19 @@ pub struct StationRecoveryRequestRestoreOperationFootprint {
     pub snapshot_id: Vec<u8>,
 }
 
+#[derive(Clone, Debug, Eq, PartialEq, Hash)]
+pub enum StationRecoveryRequestPruneOperationFootprint {
+    Snapshot(String),
+    ChunkStore,
+    State,
+}
+
 #[derive(Clone, Debug, Eq, PartialEq, Hash)]
 pub enum StationRecoveryRequestOperationFootprint {
     InstallCode(StationRecoveryRequestInstallCodeOperationFootprint),
     Snapshot(StationRecoveryRequestSnapshotOperationFootprint),
     Restore(StationRecoveryRequestRestoreOperationFootprint),
+    Prune(StationRecoveryRequestPruneOperationFootprint),
 }
 
 #[storable]

core/upgrader/impl/src/model/logging.rs

@@ -46,11 +46,35 @@ pub struct RequestDisasterRecoveryRestoreLog {
     pub snapshot_id: String,
 }
 
+#[derive(Serialize)]
+pub enum RequestDisasterRecoveryPruneLog {
+    Snapshot(String),
+    ChunkStore,
+    State,
+}
+
+impl std::fmt::Display for RequestDisasterRecoveryPruneLog {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        match self {
+            RequestDisasterRecoveryPruneLog::Snapshot(snapshot_id) => {
+                write!(f, "snapshot_id {}", snapshot_id)
+            }
+            RequestDisasterRecoveryPruneLog::ChunkStore => {
+                write!(f, "chunk store")
+            }
+            RequestDisasterRecoveryPruneLog::State => {
+                write!(f, "state")
+            }
+        }
+    }
+}
+
 #[derive(Serialize)]
 pub enum RequestDisasterRecoveryOperationLog {
     InstallCode(RequestDisasterRecoveryInstallCodeLog),
     Snapshot(RequestDisasterRecoverySnapshotLog),
     Restore(RequestDisasterRecoveryRestoreLog),
+    Prune(RequestDisasterRecoveryPruneLog),
 }
 
 impl std::fmt::Display for RequestDisasterRecoveryOperationLog {
@@ -73,6 +97,9 @@ impl std::fmt::Display for RequestDisasterRecoveryOperationLog {
             RequestDisasterRecoveryOperationLog::Restore(snapshot) => {
                 write!(f, "Restore snapshot_id {}", snapshot.snapshot_id,)
             }
+            RequestDisasterRecoveryOperationLog::Prune(prune) => {
+                write!(f, "Prune {}", prune)
+            }
         }
     }
 }

core/upgrader/impl/src/services/disaster_recovery.rs

@@ -9,6 +9,7 @@ use crate::{
         RecoveryStatus, RequestDisasterRecoveryLog, RequestDisasterRecoveryOperationLog,
         SetAccountsAndAssetsLog, SetAccountsLog, SetCommitteeLog, StationRecoveryRequest,
         StationRecoveryRequestOperation, StationRecoveryRequestOperationFootprint,
+        StationRecoveryRequestPruneOperation,
     },
     services::LOGGER_SERVICE,
     set_disaster_recovery,
@@ -20,8 +21,9 @@ use lazy_static::lazy_static;
 use orbit_essentials::api::ServiceResult;
 use orbit_essentials::cdk::api::canister_version;
 use orbit_essentials::cdk::api::management_canister::main::{
-    load_canister_snapshot, take_canister_snapshot, LoadCanisterSnapshotArgs,
-    TakeCanisterSnapshotArgs,
+    clear_chunk_store, delete_canister_snapshot, load_canister_snapshot, take_canister_snapshot,
+    uninstall_code, CanisterIdRecord, ClearChunkStoreArgument, DeleteCanisterSnapshotArgs,
+    LoadCanisterSnapshotArgs, TakeCanisterSnapshotArgs,
 };
 use std::{
     collections::{HashMap, HashSet},
@@ -324,6 +326,41 @@ impl DisasterRecoveryService {
 
                 installer.start(station_canister_id).await?;
 
+                Ok(())
+            }
+            StationRecoveryRequestOperation::Prune(
+                StationRecoveryRequestPruneOperation::Snapshot(snapshot_id),
+            ) => {
+                let snapshot_args = DeleteCanisterSnapshotArgs {
+                    canister_id: station_canister_id,
+                    snapshot_id,
+                };
+                delete_canister_snapshot(snapshot_args)
+                    .await
+                    .map_err(|(_, err)| err)?;
+
+                Ok(())
+            }
+            StationRecoveryRequestOperation::Prune(
+                StationRecoveryRequestPruneOperation::ChunkStore,
+            ) => {
+                let canister_id_record = ClearChunkStoreArgument {
+                    canister_id: station_canister_id,
+                };
+                clear_chunk_store(canister_id_record)
+                    .await
+                    .map_err(|(_, err)| err)?;
+
+                Ok(())
+            }
+            StationRecoveryRequestOperation::Prune(StationRecoveryRequestPruneOperation::State) => {
+                let canister_id_record = CanisterIdRecord {
+                    canister_id: station_canister_id,
+                };
+                uninstall_code(canister_id_record)
+                    .await
+                    .map_err(|(_, err)| err)?;
+
                 Ok(())
             }
         }