refactor(registry): Use shards for get_subnet_for_canister (#5867)

This change allows get_subnet_for_canister to operate more efficiently,
only loading the portion of the routing table needed to return a
response, instead of all of the shards

Author: max-dfinity

rs/registry/canister/src/mutations/node_management/common.rs

@@ -258,41 +258,64 @@ pub fn node_exists_with_ipv4(registry: &Registry, ipv4_addr: &str) -> bool {
 /// Similar to `get_key_family` on the `RegistryClient`, return a list of
 /// tuples, (ID, value).  This strips the prefix from the key and returns the
 /// value as a decoded struct.
+///
+/// This function must return keys in order of their bytes, which should
+/// also be the same order as the string representations.
 pub(crate) fn get_key_family<T: prost::Message + Default>(
     registry: &Registry,
     prefix: &str,
 ) -> Vec<(String, T)> {
     get_key_family_iter(registry, prefix).collect()
 }
 
+/// This function must return keys in order of their bytes, which should
+/// also be the same order as the string representations.
 pub(crate) fn get_key_family_iter<'a, T: prost::Message + Default>(
     registry: &'a Registry,
     prefix: &'a str,
 ) -> impl Iterator<Item = (String, T)> + 'a {
     get_key_family_iter_at_version(registry, prefix, registry.latest_version())
 }
 
+/// This function must return keys in order of their bytes, which should
+/// also be the same order as the string representations.
 pub(crate) fn get_key_family_iter_at_version<'a, T: prost::Message + Default>(
     registry: &'a Registry,
     prefix: &'a str,
     version: u64,
 ) -> impl Iterator<Item = (String, T)> + 'a {
+    get_key_family_raw_iter_at_version(registry, prefix, version).filter_map(|(id, value)| {
+        let latest_value: Option<T> =
+            with_chunks(|chunks| decode_high_capacity_registry_value::<T, _>(value, chunks));
+
+        let latest_value = latest_value?;
+
+        Some((id, latest_value))
+    })
+}
+
+/// This function must return keys in order of their bytes, which should
+/// also be the same order as the string representations.
+pub(crate) fn get_key_family_raw_iter_at_version<'a>(
+    registry: &'a Registry,
+    prefix: &'a str,
+    version: u64,
+) -> impl Iterator<Item = (String, &'a HighCapacityRegistryValue)> + 'a {
     let prefix_bytes = prefix.as_bytes();
     let start = prefix_bytes.to_vec();
 
+    // Note, using the 'store' which is a BTreeMap is what guarantees the order of keys.
     registry
         .store
         .range(start..)
         .take_while(|(k, _)| k.starts_with(prefix_bytes))
         .filter_map(move |(key, values)| {
-            let value: &HighCapacityRegistryValue =
+            let latest_value: &HighCapacityRegistryValue =
                 values.iter().rev().find(|value| value.version <= version)?;
 
-            let latest_value: Option<T> =
-                with_chunks(|chunks| decode_high_capacity_registry_value::<T, _>(value, chunks));
-
-            // Skip deleted values.
-            let latest_value = latest_value?;
+            if !latest_value.is_present() {
+                return None; // Deleted or otherwise empty value.
+            }
 
             let id = key
                 .strip_prefix(prefix_bytes)

rs/registry/canister/src/mutations/routing_table.rs

@@ -1,8 +1,12 @@
-use crate::mutations::node_management::common::get_key_family_iter_at_version;
-use crate::{pb::v1::SubnetForCanister, registry::Registry};
+use crate::mutations::node_management::common::{
+    get_key_family_iter_at_version, get_key_family_raw_iter_at_version,
+};
+use crate::storage::with_chunks;
+use crate::{common::LOG_PREFIX, pb::v1::SubnetForCanister, registry::Registry};
 use dfn_core::CanisterId;
 use ic_base_types::{PrincipalId, SubnetId};
 use ic_protobuf::registry::routing_table::v1 as pb;
+use ic_registry_canister_chunkify::decode_high_capacity_registry_value;
 use ic_registry_keys::{
     make_canister_migrations_record_key, make_canister_ranges_key, make_routing_table_record_key,
     CANISTER_RANGES_PREFIX,
@@ -254,6 +258,41 @@ impl Registry {
         RoutingTable::try_from(pb::RoutingTable { entries }).unwrap()
     }
 
+    pub fn get_routing_table_shard_for_canister_id(
+        &self,
+        canister_id: CanisterId,
+        version: u64,
+    ) -> Result<RoutingTable, String> {
+        // get_key_family_* functions don't return deleted values.
+        let ranges: Vec<_> =
+            get_key_family_raw_iter_at_version(self, CANISTER_RANGES_PREFIX, version)
+                .map(|(k, v)| {
+                    let shard_canister_id_bytes = hex::decode(k).unwrap();
+                    (CanisterId::try_from(shard_canister_id_bytes).unwrap(), v)
+                })
+                .collect();
+
+        let maybe_range_to_decode = ranges
+            .into_iter()
+            .take_while(|(shard_canister_id, _)| shard_canister_id <= &canister_id)
+            .last()
+            .map(|(_, record)| record);
+
+        let Some(range_to_decode) = maybe_range_to_decode else {
+            return Err(format!("{}Could not find routing table shard", LOG_PREFIX));
+        };
+
+        match with_chunks(|chunks| {
+            decode_high_capacity_registry_value::<pb::RoutingTable, _>(range_to_decode, chunks)
+        }) {
+            Some(rt) => RoutingTable::try_from(rt).map_err(|e| e.to_string()),
+            None => Err(format!(
+                "{}Could not decode routing table shard",
+                LOG_PREFIX
+            )),
+        }
+    }
+
     /// Applies the given mutation to the routing table at the specified version.
     fn modify_routing_table(
         &self,
@@ -388,11 +427,14 @@ impl Registry {
         principal_id: &PrincipalId,
     ) -> Result<SubnetForCanister, GetSubnetForCanisterError> {
         let latest_version = self.latest_version();
-        let routing_table = self.get_routing_table_or_panic(latest_version);
-        let canister_id = CanisterId::try_from(*principal_id)
+
+        let canister_id = CanisterId::try_from_principal_id(*principal_id)
             .map_err(|_| GetSubnetForCanisterError::InvalidCanisterId)?;
+        let routing_table_segment = self
+            .get_routing_table_shard_for_canister_id(canister_id, latest_version)
+            .map_err(|_| GetSubnetForCanisterError::NoSubnetAssigned)?;
 
-        match routing_table
+        match routing_table_segment
             .lookup_entry(canister_id)
             .map(|(_, subnet_id)| subnet_id.get())
         {
@@ -460,6 +502,149 @@ mod tests {
         // GetSubnetForCanisterError::CanisterIdConversion currently not reachable - CanisterId::try_from() always succeeds
     }
 
+    // Helper functions to reduce test boilerplate
+    fn assert_subnet_for_canister(registry: &Registry, canister_id: u64, expected_subnet_idx: u64) {
+        let result = registry
+            .get_subnet_for_canister(&CanisterId::from(canister_id).get())
+            .unwrap_or_else(|_| panic!("No subnet found for canister {}", canister_id));
+
+        assert_eq!(
+            result.subnet_id.unwrap(),
+            PrincipalId::new_subnet_test_id(expected_subnet_idx),
+            "Canister {} should be on subnet with index {}",
+            canister_id,
+            expected_subnet_idx
+        );
+    }
+
+    // Helper function to assert that no subnet is assigned for a canister
+    fn assert_no_subnet_for_canister(registry: &Registry, canister_id: u64) {
+        assert_matches!(
+            registry
+                .get_subnet_for_canister(&CanisterId::from(canister_id).get())
+                .unwrap_err(),
+            GetSubnetForCanisterError::NoSubnetAssigned
+        );
+    }
+
+    #[test]
+    fn get_subnet_for_canister_with_multiple_shards() {
+        let mut registry = invariant_compliant_registry(0);
+        // Create enough entries to trigger multiple shards (more than MAX_RANGES_PER_CANISTER_RANGES)
+        let entries = make_rt_entry_definitions(0..25, 100); // 25 entries, each covering 100 canister IDs
+
+        let rt = rt_from_ranges(entries);
+
+        let mutations = routing_table_into_registry_mutation(&registry, rt);
+        registry.maybe_apply_mutation_internal(mutations);
+
+        // Test 5 random canister IDs from different ranges
+        // Test canister ID 50 (should be in range 0-99, subnet_id from index 0)
+        assert_subnet_for_canister(&registry, 50, 0);
+
+        // Test canister ID 550 (should be in range 500-599, subnet_id from index 5)
+        assert_subnet_for_canister(&registry, 550, 5);
+
+        // Test canister ID 1250 (should be in range 1200-1299, subnet_id from index 12)
+        assert_subnet_for_canister(&registry, 1250, 12);
+
+        // Test canister ID 1999 (should be in range 1900-1999, subnet_id from index 19)
+        assert_subnet_for_canister(&registry, 1999, 19);
+
+        // Test canister ID 2399 (should be in range 2300-2399, subnet_id from index 23)
+        assert_subnet_for_canister(&registry, 2399, 23);
+
+        // Test a canister ID outside all ranges
+        assert_no_subnet_for_canister(&registry, 3000);
+    }
+
+    #[test]
+    fn test_canister_lookup_at_shard_boundaries() {
+        let mut registry = invariant_compliant_registry(0);
+        let shards = shards(vec![
+            (0, vec![((0, 999), test_subnet(0))]),
+            (1000, vec![((1000, 1999), test_subnet(1))]),
+            (2000, vec![((2000, 2999), test_subnet(2))]),
+        ]);
+        apply_shards_to_registry(&mut registry, &shards);
+
+        assert_subnet_for_canister(&registry, 1000, 1);
+        assert_subnet_for_canister(&registry, 999, 0);
+        assert_subnet_for_canister(&registry, 1999, 1);
+        assert_subnet_for_canister(&registry, 2000, 2);
+        assert_no_subnet_for_canister(&registry, 3000);
+    }
+
+    #[test]
+    fn test_canister_lookup_with_gaps_in_shards() {
+        let mut registry = invariant_compliant_registry(0);
+        let shards = shards(vec![
+            (0, vec![((0, 500), test_subnet(0))]),
+            (2000, vec![((2000, 2500), test_subnet(2))]),
+        ]);
+        apply_shards_to_registry(&mut registry, &shards);
+
+        assert_subnet_for_canister(&registry, 250, 0);
+        assert_no_subnet_for_canister(&registry, 1500);
+        assert_subnet_for_canister(&registry, 2250, 2);
+    }
+
+    #[test]
+    fn test_canister_lookup_with_empty_shards() {
+        let mut registry = invariant_compliant_registry(0);
+        let shards = shards(vec![
+            (0, vec![((0, 500), test_subnet(0))]),
+            (1000, vec![]), // Empty shard
+            (2000, vec![((2000, 2500), test_subnet(2))]),
+        ]);
+        apply_shards_to_registry(&mut registry, &shards);
+
+        assert_no_subnet_for_canister(&registry, 1500);
+    }
+
+    #[test]
+    fn test_canister_lookup_below_first_shard() {
+        let mut registry = invariant_compliant_registry(0);
+        let shards = shards(vec![
+            (1000, vec![((1000, 1500), test_subnet(1))]),
+            (2000, vec![((2000, 2500), test_subnet(2))]),
+        ]);
+        apply_shards_to_registry(&mut registry, &shards);
+
+        assert_no_subnet_for_canister(&registry, 500);
+    }
+
+    #[test]
+    fn test_canister_lookup_with_gaps() {
+        let mut registry = invariant_compliant_registry(0);
+        let shards = shards(vec![
+            (0, vec![((100, 500), test_subnet(0))]),
+            (400, vec![((600, 800), test_subnet(1))]),
+            (850, vec![((900, 1200), test_subnet(2))]),
+        ]);
+        apply_shards_to_registry(&mut registry, &shards);
+
+        assert_subnet_for_canister(&registry, 300, 0);
+        assert_no_subnet_for_canister(&registry, 50);
+        assert_subnet_for_canister(&registry, 700, 1);
+    }
+
+    #[test]
+    fn test_canister_lookup_at_max_canister_id() {
+        let mut registry = invariant_compliant_registry(0);
+        let shards = shards(vec![
+            (0, vec![((0, 1000), test_subnet(0))]),
+            (
+                u64::MAX - 1000,
+                vec![((u64::MAX - 1000, u64::MAX), test_subnet(1))],
+            ),
+        ]);
+        apply_shards_to_registry(&mut registry, &shards);
+
+        assert_subnet_for_canister(&registry, u64::MAX, 1);
+        assert_no_subnet_for_canister(&registry, u64::MAX - 2000);
+    }
+
     #[test]
     fn test_routing_table_updates_and_deletes_entries_as_expected() {
         let mut registry = invariant_compliant_registry(0);
@@ -516,6 +701,11 @@ mod tests {
         );
     }
 
+    // Simple helper to create SubnetId from index
+    fn test_subnet(idx: u64) -> SubnetId {
+        SubnetId::new(PrincipalId::new_subnet_test_id(idx))
+    }
+
     /// Helper to build a RoutingTable from a Vec of ((start, end), subnet_id)
     fn rt_from_ranges(ranges: Vec<((u64, u64), SubnetId)>) -> RoutingTable {
         let pb_rt = shard_pb_rt(ranges);
@@ -587,7 +777,7 @@ mod tests {
                         i * entries_per_range,
                         i * entries_per_range + entries_per_range - 1,
                     ),
-                    SubnetId::new(PrincipalId::new_user_test_id(i)),
+                    SubnetId::new(PrincipalId::new_subnet_test_id(i)),
                 )
             })
             .collect()