feat(recovery): [CON-1604] allow to override what state height to download

rs/recovery/BUILD.bazel

@@ -44,6 +44,7 @@ DEV_DEPENDENCIES = [
     "//rs/test_utilities/consensus",
     "//rs/test_utilities/tmpdir",
     "//rs/test_utilities/types",
+    "@crate_index//:assert_matches",
     "@crate_index//:tempfile",
 ]
 

rs/recovery/Cargo.toml

@@ -45,6 +45,7 @@ tokio = { workspace = true }
 url = { workspace = true }
 
 [dev-dependencies]
+assert_matches = { workspace = true }
 ic-crypto-tree-hash = { path = "../crypto/tree_hash" }
 ic-test-utilities-consensus = { path = "../test_utilities/consensus" }
 ic-test-utilities-tmpdir = { path = "../test_utilities/tmpdir" }

rs/recovery/src/app_subnet_recovery.rs

@@ -46,8 +46,8 @@ pub enum StepType {
     /// replica bug and not due to malicious actors, this step should not reveal any problems.
     MergeCertificationPools,
     /// In this step we will download all finalized consensus artifacts. For that we should use a
-    /// node, that is up to date with the highest finalization height because this node will contain
-    /// all required artifacts for the recovery.
+    /// node, that is up to date with the highest finalization and CUP height because this node
+    /// will contain all required artifacts for the recovery.
     DownloadConsensusPool,
     /// In this step we will download the subnet state from a node that is sufficiently up to date
     /// with the rest of the subnet, i.e. not behind by more than 1 DKG interval. To avoid
@@ -171,6 +171,10 @@ pub struct AppSubnetRecoveryArgs {
     #[clap(long)]
     pub keep_downloaded_state: Option<bool>,
 
+    /// Height of the checkpoint to download. If not provided, the latest checkpoint is used.
+    #[clap(long)]
+    pub download_state_height: Option<u64>,
+
     /// The method of uploading state. Possible values are either `local` (for a
     /// local recovery on the admin node) or the ipv6 address of the target node.
     /// Local recoveries allow us to skip a potentially expensive data transfer.
@@ -306,8 +310,8 @@ impl RecoveryIterator<StepType, StepTypeIter> for AppSubnetRecovery {
 
             StepType::DownloadConsensusPool => {
                 if self.params.download_pool_node.is_none() {
-                    // We could pick a node with highest finalization height automatically, but we
-                    // might have a preference between nodes of the same finalization height.
+                    // We could pick a node with highest finalization and CUP height automatically,
+                    // but we might have a preference between nodes of same heights.
                     print_height_info(
                         &self.logger,
                         &self.recovery.registry_helper,
@@ -336,6 +340,13 @@ impl RecoveryIterator<StepType, StepTypeIter> for AppSubnetRecovery {
                         "Preserve original downloaded state locally?",
                     ));
                 }
+
+                if self.params.download_state_height.is_none() {
+                    self.params.download_state_height = read_optional(
+                        &self.logger,
+                        "Enter the height of the checkpoint to download (leave empty for latest checkpoint):",
+                    );
+                }
             }
 
             StepType::ICReplay => {
@@ -453,9 +464,10 @@ impl RecoveryIterator<StepType, StepTypeIter> for AppSubnetRecovery {
             }
 
             StepType::DownloadState => match self.params.download_state_method {
-                Some(DataLocation::Local) => {
-                    Ok(Box::new(self.recovery.get_copy_local_state_step()))
-                }
+                Some(DataLocation::Local) => Ok(Box::new(
+                    self.recovery
+                        .get_copy_local_state_step(self.params.download_state_height)?,
+                )),
                 Some(DataLocation::Remote(node_ip)) => {
                     let (ssh_user, key_file) = if self.params.readonly_pub_key.is_some() {
                         (SshUser::Readonly, self.params.readonly_key_file.clone())
@@ -468,6 +480,7 @@ impl RecoveryIterator<StepType, StepTypeIter> for AppSubnetRecovery {
                         ssh_user,
                         key_file,
                         self.params.keep_downloaded_state == Some(true),
+                        self.params.download_state_height,
                     )?))
                 }
                 None => Err(RecoveryError::StepSkipped),

rs/recovery/src/cli.rs

@@ -212,7 +212,10 @@ fn print_summary(logger: &Logger, args: &RecoveryArgs, subnet_id: SubnetId) {
 
 pub fn print_height_info(logger: &Logger, registry_helper: &RegistryHelper, subnet_id: SubnetId) {
     info!(logger, "Collecting node heights from metrics...");
-    info!(logger, "Select a node with highest finalization height:");
+    info!(
+        logger,
+        "Select a node with highest finalization height and highest CUP height:"
+    );
     match get_available_nodes_heights_from_metrics(logger, registry_helper, subnet_id) {
         Ok(heights) => info!(logger, "{:#?}", heights),
         Err(err) => warn!(logger, "Failed to query height info: {:?}", err),