Merge branch 'alin/MR-427-state_tool-split' into 'master'

feat: [MR-427] Implement state splitting in `state-tool`

Implement a split command with mutually exclusive `--retain` and `--drop`
arguments, so the same canister ID ranges can be passed when splitting off both
subnets A' (`--retain`) and B (`--drop`). 

See merge request dfinity-lab/public/ic!12534

Author: alin-at-dfinity

Cargo.lock

@@ -8,6 +8,7 @@ DEPENDENCIES = [
     "//rs/monitoring/logger",
     "//rs/monitoring/metrics",
     "//rs/protobuf",
+    "//rs/registry/routing_table",
     "//rs/registry/subnet_type",
     "//rs/replicated_state",
     "//rs/state_layout",

rs/state_tool/BUILD.bazel

@@ -16,6 +16,7 @@ ic-config = { path = "../config" }
 ic-logger = { path = "../monitoring/logger" }
 ic-metrics = { path = "../monitoring/metrics" }
 ic-protobuf = { path = "../protobuf" }
+ic-registry-routing-table = { path = "../registry/routing_table" }
 ic-registry-subnet-type = { path = "../registry/subnet_type" }
 ic-replicated-state = { path = "../replicated_state" }
 ic-state-layout = { path = "../state_layout" }

rs/state_tool/Cargo.toml

@@ -6,5 +6,6 @@ pub mod decode;
 pub mod import_state;
 pub mod list;
 pub mod manifest;
+pub mod split;
 mod utils;
 pub mod verify_manifest;

rs/state_tool/src/commands.rs

@@ -0,0 +1,102 @@
+//! Prunes a replicated state, as part of a subnet split.
+
+use ic_logger::replica_logger::no_op_logger;
+use ic_metrics::MetricsRegistry;
+use ic_registry_routing_table::{difference, CanisterIdRange, CanisterIdRanges, WellFormedError};
+use ic_state_manager::split::split;
+use ic_types::{CanisterId, PrincipalId};
+use std::{iter::once, path::PathBuf};
+
+/// Loads the latest checkpoint under the given root; splits off the state of
+/// `subnet_id`, retaining or dropping the provided canister ID ranges (exactly
+/// one of which must be non-empty); and writes back the split state as a new
+/// checkpoint, under the same root.
+pub fn do_split(
+    root: PathBuf,
+    subnet_id: PrincipalId,
+    retain: Vec<CanisterIdRange>,
+    drop: Vec<CanisterIdRange>,
+) -> Result<(), String> {
+    let canister_id_ranges = resolve(retain, drop).map_err(|e| format!("{:?}", e))?;
+    let metrics_registry = MetricsRegistry::new();
+    let log = no_op_logger();
+
+    split(root, subnet_id, canister_id_ranges, &metrics_registry, log)
+}
+
+/// Converts a pair of `retain` and `drop` range vectors (exactly one of which
+/// is expected to be non-empty) into a well-formed `CanisterIdRanges` covering
+/// all canisters to be retained. Returns an error if the provided inputs are
+/// not well formed.
+///
+/// Panics if none or both of the inputs are empty.
+fn resolve(
+    retain: Vec<CanisterIdRange>,
+    drop: Vec<CanisterIdRange>,
+) -> Result<CanisterIdRanges, WellFormedError> {
+    if !retain.is_empty() && drop.is_empty() {
+        // Validate and return `retain`.
+        CanisterIdRanges::try_from(retain)
+    } else if retain.is_empty() && !drop.is_empty() {
+        // Validate `drop` and return the diff between all possible canisters and it.
+        let all_canister_ids = CanisterIdRange {
+            start: CanisterId::from_u64(0),
+            end: CanisterId::from_u64(u64::MAX),
+        };
+        difference(
+            once(&all_canister_ids),
+            CanisterIdRanges::try_from(drop)?.iter(),
+        )
+    } else {
+        panic!("Expecting exactly one of `retain` and `drop` to be non-empty");
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn make_range(start: u64, end: u64) -> CanisterIdRange {
+        CanisterIdRange {
+            start: CanisterId::from_u64(start),
+            end: CanisterId::from_u64(end),
+        }
+    }
+
+    #[test]
+    fn test_resolve_retain() {
+        let retain = make_range(3, 4);
+        assert_eq!(
+            CanisterIdRanges::try_from(vec![retain]),
+            resolve(vec![retain], vec![])
+        );
+    }
+
+    #[test]
+    fn test_resolve_drop() {
+        let drop = make_range(3, 4);
+        assert_eq!(
+            CanisterIdRanges::try_from(vec![make_range(0, 2), make_range(5, u64::MAX)]),
+            resolve(vec![], vec![drop])
+        );
+    }
+
+    #[test]
+    fn test_resolve_not_well_formed() {
+        let retain = make_range(4, 3);
+        resolve(vec![retain], vec![]).unwrap_err();
+    }
+
+    #[test]
+    #[should_panic]
+    fn test_resolve_both_non_empty() {
+        let range = make_range(3, 4);
+        resolve(vec![range], vec![range]).ok();
+    }
+
+    #[test]
+    #[should_panic]
+    fn test_resolve_both_empty() {
+        resolve(vec![], vec![]).ok();
+    }
+}

rs/state_tool/src/commands/split.rs

@@ -5,6 +5,8 @@
 //! checkpoint manifests, import state trees).
 
 use clap::Parser;
+use ic_registry_routing_table::CanisterIdRange;
+use ic_types::PrincipalId;
 use std::path::PathBuf;
 
 mod commands;
@@ -99,7 +101,7 @@ enum Opt {
         /// their path.
         ///
         /// To make sure that accidentally passing something that matches
-        /// unwanted file paths, the list of processed files is explititly
+        /// unwanted file paths, the list of processed files is explicitly
         /// printed.
         #[clap(long = "canister")]
         canister: String,
@@ -142,6 +144,28 @@ enum Opt {
         #[clap(long = "bytes")]
         bytes: String,
     },
+
+    /// Prunes a replicated state, as part of a subnet split.
+    #[clap(name = "split")]
+    #[clap(group(
+        clap::ArgGroup::new("ranges")
+            .required(true)
+            .args(&["retain", "drop"]),
+    ))]
+    Split {
+        /// Path to the state layout.
+        #[clap(long, required = true)]
+        root: PathBuf,
+        /// The ID of the subnet being split off.
+        #[clap(long, required = true)]
+        subnet_id: PrincipalId,
+        /// Canister ID ranges to retain (assigned to the subnet in the routing table).
+        #[clap(long, multiple_values(true))]
+        retain: Vec<CanisterIdRange>,
+        /// Canister ID ranges to drop (assigned to other subnet in the routing table).
+        #[clap(long, multiple_values(true))]
+        drop: Vec<CanisterIdRange>,
+    },
 }
 
 fn main() {
@@ -175,6 +199,12 @@ fn main() {
         Opt::PrincipalFromBytes { bytes } => {
             commands::convert_ids::do_principal_from_byte_string(bytes)
         }
+        Opt::Split {
+            root,
+            subnet_id,
+            retain,
+            drop,
+        } => commands::split::do_split(root, subnet_id, retain, drop),
     };
 
     if let Err(e) = result {