perf: Replace RoutingTable::route by NetworkTopology::route (#6361)

This commit removes `RoutingTable::route`, which was used to route
messages to the correct subnet.

The difference to the still existing `RoutingTable::lookup_entry` was
that `route` would additionally consider the case where the target is a
subnet id instead of a canister id. However, previously this was
implemented by scanning over all canister ranges and checking if the
input occurs as the subnet id for any range. There are two issues with
this, which this commit both solves:

1. As it's a linear scan, route was `O(n)`, where n is the number of
ranges in the routing table. Currently `n` is small, but will grow with
the upcoming canister migration feature.

2. There is a theoretical edge case of subnets with no ranges assigned
to them. It was not possible to route to such a subnet.

This commit solves both problems by instead implementing `route` in the
`NetworkTopology` struct. `NetworkTopology` contains both a
`RoutingTable` and a `BTreeMap` of subnets. Instead of a linear scan, we
now only need to do a quick lookup in the `BTreeMap`. Furthermore,
subnets without any ranges will also be found this way.

This commit is somewhat large due to the necessary test changes. In
production code, the switch is very simple.

Author: schneiderstefan

rs/embedders/src/wasmtime_embedder/system_api/routing.rs

@@ -45,7 +45,6 @@ fn route_canister_id(
     network_topology: &NetworkTopology,
 ) -> Result<PrincipalId, ResolveDestinationError> {
     network_topology
-        .routing_table
         .route(canister_id.get())
         .map(|subnet_id| subnet_id.get())
         .ok_or(ResolveDestinationError::SubnetNotFound(canister_id, method))

rs/execution_environment/src/canister_manager.rs

@@ -1284,12 +1284,7 @@ impl CanisterManager {
             return Err(CanisterManagerError::CanisterAlreadyExists(new_canister_id));
         }
 
-        if state
-            .metadata
-            .network_topology
-            .routing_table
-            .route(specified_id)
-            == Some(state.metadata.own_subnet_id)
+        if state.metadata.network_topology.route(specified_id) == Some(state.metadata.own_subnet_id)
         {
             Ok(new_canister_id)
         } else {

rs/execution_environment/src/execution_environment.rs

@@ -3290,7 +3290,7 @@ impl ExecutionEnvironment {
 
         let topology = &state.metadata.network_topology;
         // If the request isn't from the NNS, then we need to charge for it.
-        let source_subnet = topology.routing_table.route(request.sender.get());
+        let source_subnet = topology.route(request.sender.get());
         if source_subnet != Some(state.metadata.network_topology.nns_subnet_id) {
             let cost_schedule = state.get_own_cost_schedule();
             let signature_fee = self.calculate_signature_fee(&args, subnet_size, cost_schedule);

rs/messaging/src/routing/stream_builder.rs

@@ -305,10 +305,8 @@ impl StreamBuilderImpl {
         };
 
         let mut streams = state.take_streams();
-        let routing_table = state.routing_table();
-        let subnet_types: BTreeMap<_, _> = state
-            .metadata
-            .network_topology
+        let network_topology = state.metadata.network_topology.clone();
+        let subnet_types: BTreeMap<_, _> = network_topology
             .subnets
             .iter()
             .map(|(subnet_id, topology)| (*subnet_id, topology.subnet_type))
@@ -342,7 +340,7 @@ impl StreamBuilderImpl {
             }
             last_output_size = output_size;
 
-            match routing_table.route(msg.receiver().get()) {
+            match network_topology.route(msg.receiver().get()) {
                 // Destination subnet found.
                 Some(dst_subnet_id) => {
                     let dst_stream_entry = streams.entry(dst_subnet_id);

rs/messaging/src/routing/stream_builder/tests.rs

@@ -606,6 +606,11 @@ fn build_streams_with_messages_targeted_to_other_subnets() {
                 CanisterIdRange{ start: CanisterId::from(0), end: CanisterId::from(0xfff) } => REMOTE_SUBNET,
             },
         ).unwrap());
+        provided_state
+            .metadata
+            .network_topology
+            .subnets
+            .insert(REMOTE_SUBNET, Default::default());
 
         // Set up the provided_canister_states.
         let provided_canister_states = canister_states_with_outputs(msgs.clone());

rs/messaging/src/routing/stream_handler.rs

@@ -570,7 +570,6 @@ impl StreamHandlerImpl {
             let new_destination = state
                 .metadata
                 .network_topology
-                .routing_table
                 .route(response.receiver().get())
                 .expect("Canister disappeared from registry. Registry in an inconsistent state.");
             info!(
@@ -827,11 +826,7 @@ impl StreamHandlerImpl {
         available_guaranteed_response_memory: &mut i64,
     ) -> InductionResult {
         // Subnet that should have received the message according to the routing table.
-        let receiver_host_subnet = state
-            .metadata
-            .network_topology
-            .routing_table
-            .route(msg.receiver().get());
+        let receiver_host_subnet = state.metadata.network_topology.route(msg.receiver().get());
 
         let payload_size = msg.payload_size_bytes().get();
         let msg_cycles = msg.cycles();
@@ -966,11 +961,7 @@ impl StreamHandlerImpl {
         state: &ReplicatedState,
     ) -> bool {
         // Remote subnet that should have sent the message according to the routing table.
-        let expected_subnet_id = state
-            .metadata
-            .network_topology
-            .routing_table
-            .route(msg.sender().get());
+        let expected_subnet_id = state.metadata.network_topology.route(msg.sender().get());
 
         match expected_subnet_id {
             // The actual originating subnet and the routing table entry for the sender are in agreement.
@@ -1030,11 +1021,7 @@ impl StreamHandlerImpl {
     ) -> bool {
         debug_assert_eq!(
             Some(actual_receiver_subnet_id),
-            state
-                .metadata
-                .network_topology
-                .routing_table
-                .route(msg.receiver().get())
+            state.metadata.network_topology.route(msg.receiver().get())
         );
 
         // Reroute if `msg.receiver()` is being migrated from `self.subnet_id` to

rs/messaging/src/routing/stream_handler/tests.rs

@@ -4085,20 +4085,38 @@ fn with_test_setup_and_config(
 
         // Ensure the routing table maps `LOCAL_CANISTER` to `LOCAL_SUBNET`,
         // `REMOTE_CANISTER` to `REMOTE_SUBNET` and `UNKNOWN_CANISTER` to `None`.
-        let routing_table = Arc::new(RoutingTable::try_from(btreemap! {
-            CanisterIdRange{ start: CanisterId::from(0x0), end: CanisterId::from(0xff) } => LOCAL_SUBNET,
-            CanisterIdRange{ start: CanisterId::from(0x100), end: CanisterId::from(0x1ff) } => REMOTE_SUBNET,
-        }).unwrap());
+        let local_range = CanisterIdRange {
+            start: CanisterId::from(0x0),
+            end: CanisterId::from(0xff),
+        };
+        let remote_range = CanisterIdRange {
+            start: CanisterId::from(0x100),
+            end: CanisterId::from(0x1ff),
+        };
+        let routing_table = Arc::new(
+            RoutingTable::try_from(btreemap! {
+                local_range => LOCAL_SUBNET,
+                remote_range => REMOTE_SUBNET,
+            })
+            .unwrap(),
+        );
         assert_eq!(
-            Some(LOCAL_SUBNET),
-            routing_table.route(LOCAL_CANISTER.get())
+            Some((local_range, LOCAL_SUBNET)),
+            routing_table.lookup_entry(*LOCAL_CANISTER)
         );
         assert_eq!(
-            Some(REMOTE_SUBNET),
-            routing_table.route(REMOTE_CANISTER.get())
+            Some((remote_range, REMOTE_SUBNET)),
+            routing_table.lookup_entry(*REMOTE_CANISTER)
         );
-        assert!(routing_table.route(UNKNOWN_CANISTER.get()).is_none());
+        assert!(routing_table.lookup_entry(*UNKNOWN_CANISTER).is_none());
         state.metadata.network_topology.routing_table = routing_table;
+        for subnet in [LOCAL_SUBNET, REMOTE_SUBNET] {
+            state
+                .metadata
+                .network_topology
+                .subnets
+                .insert(subnet, Default::default());
+        }
 
         // Generate testing canister using `LOCAL_CANISTER` as the canister ID.
         let mut canister_state = CanisterStateBuilder::new()

rs/pocket_ic_server/src/pocket_ic.rs

@@ -676,8 +676,11 @@ impl PocketIcSubnets {
     }
 
     fn route(&self, canister_id: CanisterId) -> Option<Arc<StateMachine>> {
-        let subnet_id = self.routing_table.route(canister_id.get());
-        subnet_id.map(|subnet_id| self.get(subnet_id).unwrap())
+        self.get(SubnetId::from(canister_id.get())).or_else(|| {
+            self.routing_table
+                .lookup_entry(canister_id)
+                .and_then(|(_, subnet_id)| self.get(subnet_id))
+        })
     }
 
     fn time(&self) -> SystemTime {

rs/registry/routing_table/src/lib.rs

@@ -1,7 +1,7 @@
 mod proto;
 
 use candid::CandidType;
-use ic_base_types::{CanisterId, CanisterIdError, PrincipalId, SubnetId};
+use ic_base_types::{CanisterId, CanisterIdError, SubnetId};
 use ic_protobuf::proxy::ProxyDecodeError;
 use serde::{Deserialize, Serialize};
 use std::convert::TryInto;
@@ -527,30 +527,6 @@ impl RoutingTable {
         Ok(())
     }
 
-    /// Returns the `SubnetId` that the given `principal_id` is assigned to or
-    /// `None` if an assignment cannot be found.
-    pub fn route(&self, principal_id: PrincipalId) -> Option<SubnetId> {
-        // Check if the given `principal_id` is a subnet.
-        // Note that the following assumes that all known subnets are in the routing
-        // table, even if they're empty (i.e. no canister exists on them). In the
-        // future, if this assumption does not hold, the list of existing
-        // subnets should be taken from the rest of the registry (which should
-        // be the absolute source of truth).
-        if let Some(subnet_id) = self.0.values().find(|x| x.get() == principal_id) {
-            return Some(*subnet_id);
-        }
-
-        // If the `principal_id` was not a subnet, it must be a `CanisterId` (otherwise
-        // we can't route to it).
-        match CanisterId::try_from(principal_id) {
-            Ok(canister_id) => {
-                lookup_in_ranges(&self.0, canister_id).map(|(_range, subnet_id)| subnet_id)
-            }
-            // Cannot route to any subnet as we couldn't convert to a `CanisterId`.
-            Err(_) => None,
-        }
-    }
-
     /// Returns the corresponding `CanisterIdRange` and `SubnetId` that the given `canister_id` is assigned to
     /// or `None` if an assignment cannot be found.
     pub fn lookup_entry(&self, canister_id: CanisterId) -> Option<(CanisterIdRange, SubnetId)> {