feat(registry): Add invariant check for sharded canister migrations (#5535)

This changes the routing table invariant check to make sure there is a
valid routing table built from shards, and that it is equal to the
single record routing table.

Additionally, it adds a benchmark for the routing table invariant check.

---------

Co-authored-by: IDX GitHub Automation <infra+github-automation@dfinity.org>

Author: max-dfinity

rs/registry/canister/canbench/canbench_results.yml

@@ -1,4 +1,11 @@
 benches:
+  measure_routing_table_invariant_checks_shards_and_unsharded:
+    total:
+      calls: 1
+      instructions: 563027802
+      heap_increase: 69
+      stable_memory_increase: 0
+    scopes: {}
   measure_snapshot_creation_with_1000_individual_entries:
     total:
       calls: 1

rs/registry/canister/src/flags.rs

@@ -1,6 +1,6 @@
 use std::cell::Cell;
 
-#[cfg(test)]
+#[cfg(any(test, feature = "canbench-rs"))]
 use ic_nervous_system_temporary::Temporary;
 
 thread_local! {
@@ -11,7 +11,7 @@ pub(crate) fn is_chunkifying_large_values_enabled() -> bool {
     IS_CHUNKIFYING_LARGE_VALUES_ENABLED.get()
 }
 
-#[cfg(test)]
+#[cfg(any(test, feature = "canbench-rs"))]
 pub fn temporarily_enable_chunkifying_large_values() -> Temporary {
     Temporary::new(&IS_CHUNKIFYING_LARGE_VALUES_ENABLED, true)
 }

rs/registry/canister/src/invariants/checks/benches.rs

@@ -1,9 +1,11 @@
+use crate::flags::temporarily_enable_chunkifying_large_values;
+use crate::invariants::routing_table::check_routing_table_invariants;
 use crate::registry::Registry;
 use canbench_rs::{bench, bench_fn, BenchResult};
 use ic_base_types::{CanisterId, PrincipalId, SubnetId};
 use ic_protobuf::registry::routing_table::v1::routing_table::Entry;
 use ic_protobuf::registry::routing_table::v1::RoutingTable;
-use ic_registry_keys::make_canister_ranges_key;
+use ic_registry_keys::{make_canister_ranges_key, make_routing_table_record_key};
 use ic_registry_routing_table::CanisterIdRange;
 use ic_registry_transport::upsert;
 use prost::Message;
@@ -90,3 +92,22 @@ fn measure_snapshot_creation_with_1_segment_of_1000_entries() -> BenchResult {
 fn measure_snapshot_creation_with_100_segments_of_1000_entries() -> BenchResult {
     benchmark_snapshot_creation_with_entries(100, 1000)
 }
+
+/// 20k entries costs 250m instructions to validate
+#[bench(raw)]
+fn measure_routing_table_invariant_checks_shards_and_unsharded() -> BenchResult {
+    let _feature = temporarily_enable_chunkifying_large_values();
+    let mut registry = setup_registry_with_rt_segments_with_x_entries_each(1000, 20);
+
+    let rt =
+        registry.get_routing_table_from_canister_range_records_or_panic(registry.latest_version());
+    let rt_mutation = upsert(
+        make_routing_table_record_key(),
+        RoutingTable::from(rt).encode_to_vec(),
+    );
+    registry.apply_mutations_for_test(vec![rt_mutation]);
+
+    let snapshot = registry.take_latest_snapshot();
+
+    bench_fn(|| check_routing_table_invariants(&snapshot))
+}

rs/registry/canister/src/invariants/routing_table.rs

@@ -2,10 +2,13 @@ use crate::invariants::common::{InvariantCheckError, RegistrySnapshot};
 
 use std::convert::TryFrom;
 
+use ic_base_types::CanisterId;
 use ic_protobuf::registry::routing_table::v1::{
     CanisterMigrations as pbCanisterMigrations, RoutingTable as pbRoutingTable,
 };
-use ic_registry_keys::{make_canister_migrations_record_key, make_routing_table_record_key};
+use ic_registry_keys::{
+    make_canister_migrations_record_key, make_canister_ranges_key, make_routing_table_record_key,
+};
 use ic_registry_routing_table::{CanisterMigrations, RoutingTable};
 use prost::Message;
 
@@ -19,14 +22,40 @@ pub(crate) fn check_routing_table_invariants(
 
 // Return routing table from snapshot
 fn get_routing_table(snapshot: &RegistrySnapshot) -> RoutingTable {
-    match snapshot.get(make_routing_table_record_key().as_bytes()) {
-        Some(routing_table_bytes) => {
-            let routing_table_proto =
-                pbRoutingTable::decode(routing_table_bytes.as_slice()).unwrap();
-            RoutingTable::try_from(routing_table_proto).unwrap()
-        }
-        None => panic!("No routing table in snapshot"),
+    // TODO(NNS1-3781): Remove this once we have sharded table supported by all clients.
+    let rt_from_routing_table_record =
+        match snapshot.get(make_routing_table_record_key().as_bytes()) {
+            Some(routing_table_bytes) => {
+                let routing_table_proto =
+                    pbRoutingTable::decode(routing_table_bytes.as_slice()).unwrap();
+                RoutingTable::try_from(routing_table_proto).unwrap()
+            }
+            None => panic!("No routing table in snapshot"),
+        };
+
+    // If there are shards, they should match the routing table record.
+    let shards = get_routing_table_shards(snapshot);
+    if !shards.is_empty() {
+        let rt_from_shards = RoutingTable::try_from(shards).unwrap();
+        assert_eq!(
+            rt_from_shards, rt_from_routing_table_record,
+            "Routing tables from shards and routing table record do not match."
+        );
+    }
+
+    rt_from_routing_table_record
+}
+
+fn get_routing_table_shards(snapshot: &RegistrySnapshot) -> Vec<pbRoutingTable> {
+    let start = make_canister_ranges_key(CanisterId::from_u64(0)).into_bytes();
+    let end = make_canister_ranges_key(CanisterId::from_u64(u64::MAX)).into_bytes();
+    let mut shards = vec![];
+    for (_, value) in snapshot.range(start..=end) {
+        let routing_table_proto = pbRoutingTable::decode(value.as_slice()).unwrap();
+        shards.push(routing_table_proto);
     }
+
+    shards
 }
 
 /// Iff `canister_migrations` is present, check that its invariants hold if reading and conversion succeed.
@@ -291,4 +320,35 @@ mod tests {
         assert!(check_routing_table_invariants(&snapshot).is_ok());
         assert!(check_canister_migrations_invariants(&snapshot).is_err());
     }
+
+    #[test]
+    #[should_panic(expected = "Routing tables from shards and routing table record do not match.")]
+    fn if_sharded_ranges_they_must_match_original_routing_table() {
+        let mut snapshot = RegistrySnapshot::new();
+
+        // The routing table before canister migration.
+        let routing_table = RoutingTable::try_from(btreemap! {
+            CanisterIdRange{ start: CanisterId::from(0x0), end: CanisterId::from(0xff) } => subnet_test_id(1),
+            CanisterIdRange{ start: CanisterId::from(0x100), end: CanisterId::from(0x1ff) } => subnet_test_id(2),
+            CanisterIdRange{ start: CanisterId::from(0x200), end: CanisterId::from(0x2ff) } => subnet_test_id(3),
+         }).unwrap();
+
+        let routing_table = PbRoutingTable::from(routing_table);
+        let key1 = make_routing_table_record_key();
+        let value1 = routing_table.encode_to_vec();
+
+        let rt_shard = RoutingTable::try_from(btreemap! {
+            CanisterIdRange{ start: CanisterId::from(0x0), end: CanisterId::from(0xff) } => subnet_test_id(1),
+            CanisterIdRange{ start: CanisterId::from(0x100), end: CanisterId::from(0x1ff) } => subnet_test_id(2),
+        }).unwrap();
+
+        let rt_shard = PbRoutingTable::from(rt_shard);
+        let key2 = make_canister_ranges_key(CanisterId::from_u64(0x0));
+        let value2 = rt_shard.encode_to_vec();
+
+        snapshot.insert(key1.into_bytes(), value1);
+        snapshot.insert(key2.into_bytes(), value2);
+
+        check_routing_table_invariants(&snapshot).unwrap();
+    }
 }

rs/registry/canister/unreleased_changelog.md

@@ -25,6 +25,9 @@ on the process that this file is part of, see
   [announced at the end of March in a forum][chunking] (and various other
   channels). This release marks the end of the "migration window" described in
   the aforelinked forum post.
+* The `check_routing_table_invariants` method now checks the new canister_ranges_
+  and ensures they match the `routing_table` record. The old invariant check will be
+  removed once `routing_table` is removed.
 
 [chunking]: https://forum.dfinity.org/t/breaking-registry-changes-for-large-records/42893?u=daniel-wong
 

rs/registry/routing_table/src/proto.rs

@@ -78,24 +78,41 @@ impl TryFrom<pb::RoutingTable> for RoutingTable {
     type Error = ProxyDecodeError;
 
     fn try_from(src: pb::RoutingTable) -> Result<Self, Self::Error> {
-        let mut map = BTreeMap::new();
-        for entry in src.entries {
-            let range = try_from_option_field(entry.range, "RoutingTable::Entry::range")?;
-            let subnet_id =
-                try_from_option_field(entry.subnet_id, "RoutingTable::Entry::subnet_id")?;
-            let subnet_id = subnet_id_try_from_protobuf(subnet_id)?;
-            if let Some(prev_subnet_id) = map.insert(range, subnet_id) {
-                return Err(ProxyDecodeError::DuplicateEntry {
-                    key: format!("{:?}", range),
-                    v1: prev_subnet_id.to_string(),
-                    v2: subnet_id.to_string(),
-                });
-            }
+        let entries_count = src.entries.len();
+        let map: BTreeMap<_, _> = src
+            .entries
+            .into_iter()
+            .map(|entry| {
+                let range = try_from_option_field(entry.range, "RoutingTable::Entry::range")?;
+                let subnet_id =
+                    try_from_option_field(entry.subnet_id, "RoutingTable::Entry::subnet_id")?;
+                let subnet_id = subnet_id_try_from_protobuf(subnet_id)?;
+
+                Ok((range, subnet_id))
+            })
+            .collect::<Result<_, Self::Error>>()?;
+
+        if map.len() != entries_count {
+            let diff = entries_count.saturating_sub(map.len());
+            return Err(ProxyDecodeError::Other(format!(
+                "There were {} duplicate entries in the routing table",
+                diff
+            )));
         }
+
         Ok(map.try_into()?)
     }
 }
 
+impl TryFrom<Vec<pb::RoutingTable>> for RoutingTable {
+    type Error = ProxyDecodeError;
+
+    fn try_from(src: Vec<pb::RoutingTable>) -> Result<Self, Self::Error> {
+        let entries = src.into_iter().flat_map(|table| table.entries).collect();
+        Self::try_from(pb::RoutingTable { entries })
+    }
+}
+
 impl From<CanisterMigrations> for pb::CanisterMigrations {
     fn from(src: CanisterMigrations) -> Self {
         Self::from(&src)