feat: [MR-636] Add size limits as fields to the stream builder (#3885)

The driver behind this change is to make smaller streams possible in
state machine tests and also unit or prop tests. But for state machine
tests especially, having to fill up an entire stream with 10k filler
requests just to test something at capacity takes a long time.
Since this change makes the stream limit constants public, removing the
duplicate
[MAX_STREAM_MESSAGES](https://sourcegraph.com/github.com/dfinity/ic/-/blob/rs/xnet/payload_builder/src/lib.rs?L229)
in the payload builder as a small additional change makes sense here
(even though it's not strictly related).

---------

Co-authored-by: IDX GitHub Automation <infra+github-automation@dfinity.org>
Co-authored-by: Alin Sinpalean <58422065+alin-at-dfinity@users.noreply.github.com>

Author: 3 people

.github/CODEOWNERS

@@ -113,6 +113,7 @@ go_deps.bzl               @dfinity/idx
 /rs/config/                                             @dfinity/consensus
 /rs/config/src/embedders.rs                             @dfinity/execution
 /rs/config/src/execution_environment.rs                 @dfinity/execution
+/rs/config/src/message_routing.rs                       @dfinity/ic-message-routing-owners
 /rs/config/src/state_manager.rs                         @dfinity/ic-message-routing-owners
 /rs/config/src/subnet_config.rs                         @dfinity/execution
 /rs/consensus/                                          @dfinity/consensus

rs/config/src/message_routing.rs

@@ -1,5 +1,17 @@
 use serde::{Deserialize, Serialize};
 
+/// Desired byte size of an outgoing stream.
+///
+/// At most `MAX_STREAM_MESSAGES` are enqueued into a stream; but only until its
+/// `count_bytes()` is greater than or equal to `TARGET_STREAM_SIZE_BYTES`.
+pub const TARGET_STREAM_SIZE_BYTES: usize = 10 * 1024 * 1024;
+
+/// Maximum number of messages in a stream.
+///
+/// At most `MAX_STREAM_MESSAGES` are enqueued into a stream; but only until its
+/// `count_bytes()` is greater than or equal to `TARGET_STREAM_SIZE_BYTES`.
+pub const MAX_STREAM_MESSAGES: usize = 10_000;
+
 #[derive(Clone, Eq, PartialEq, Debug, Deserialize, Serialize)]
 #[serde(default)]
 /// Message Routing replica config.

rs/messaging/src/message_routing.rs

@@ -4,6 +4,7 @@ use crate::{
 };
 use ic_config::embedders::BestEffortResponsesFeature;
 use ic_config::execution_environment::{BitcoinConfig, Config as HypervisorConfig};
+use ic_config::message_routing::{MAX_STREAM_MESSAGES, TARGET_STREAM_SIZE_BYTES};
 use ic_cycles_account_manager::CyclesAccountManager;
 use ic_interfaces::{crypto::ErrorReproducibility, execution_environment::ChainKeySettings};
 use ic_interfaces::{
@@ -616,6 +617,7 @@ pub(crate) type NodePublicKeys = BTreeMap<NodeId, Vec<u8>>;
 pub(crate) type ApiBoundaryNodes = BTreeMap<NodeId, ApiBoundaryNodeEntry>;
 
 impl BatchProcessorImpl {
+    #[allow(clippy::too_many_arguments)]
     fn new(
         state_manager: Arc<dyn StateManager<State = ReplicatedState>>,
         certified_stream_store: Arc<dyn CertifiedStreamStore>,
@@ -624,6 +626,8 @@ impl BatchProcessorImpl {
         hypervisor_config: HypervisorConfig,
         cycles_account_manager: Arc<CyclesAccountManager>,
         subnet_id: SubnetId,
+        max_stream_messages: usize,
+        target_stream_size_bytes: usize,
         metrics: MessageRoutingMetrics,
         metrics_registry: &MetricsRegistry,
         log: ReplicaLogger,
@@ -657,6 +661,8 @@ impl BatchProcessorImpl {
         ));
         let stream_builder = Box::new(routing::stream_builder::StreamBuilderImpl::new(
             subnet_id,
+            max_stream_messages,
+            target_stream_size_bytes,
             metrics_registry,
             &metrics,
             time_in_stream_metrics,
@@ -1491,6 +1497,11 @@ impl MessageRoutingImpl {
             hypervisor_config,
             cycles_account_manager,
             subnet_id,
+            // Do NOT replace these constants. Stream limits must remain constant on mainnet,
+            // otherwise the payload builder might mistakenly identify subnets as dishonest.
+            // Changes must be carefully considered.
+            MAX_STREAM_MESSAGES,
+            TARGET_STREAM_SIZE_BYTES,
             metrics.clone(),
             metrics_registry,
             log.clone(),
@@ -1512,6 +1523,8 @@ impl MessageRoutingImpl {
     ) -> Self {
         let stream_builder = Box::new(routing::stream_builder::StreamBuilderImpl::new(
             subnet_id,
+            MAX_STREAM_MESSAGES,
+            TARGET_STREAM_SIZE_BYTES,
             metrics_registry,
             &MessageRoutingMetrics::new(metrics_registry),
             Arc::new(Mutex::new(LatencyMetrics::new_time_in_stream(
@@ -1597,8 +1610,7 @@ impl MessageRouting for MessageRoutingImpl {
     }
 }
 
-/// An MessageRouting implementation that processes batches synchronously. Primarily used for
-/// testing.
+/// An MessageRouting implementation that processes batches synchronously. Used for state machine tests.
 pub struct SyncMessageRouting {
     batch_processor: Arc<Mutex<dyn BatchProcessor>>,
     state_manager: Arc<dyn StateManager<State = ReplicatedState>>,
@@ -1616,6 +1628,8 @@ impl SyncMessageRouting {
         hypervisor_config: HypervisorConfig,
         cycles_account_manager: Arc<CyclesAccountManager>,
         subnet_id: SubnetId,
+        max_stream_messages: usize,
+        target_stream_size_bytes: usize,
         metrics_registry: &MetricsRegistry,
         log: ReplicaLogger,
         registry: Arc<dyn RegistryClient>,
@@ -1631,6 +1645,8 @@ impl SyncMessageRouting {
             hypervisor_config,
             cycles_account_manager,
             subnet_id,
+            max_stream_messages,
+            target_stream_size_bytes,
             metrics,
             metrics_registry,
             log.clone(),

rs/messaging/src/routing/stream_builder.rs

@@ -53,18 +53,6 @@ struct StreamBuilderMetrics {
     pub critical_error_induct_response_failed: IntCounter,
 }
 
-/// Desired byte size of an outgoing stream.
-///
-/// At most `MAX_STREAM_MESSAGES` are enqueued into a stream; but only until its
-/// `count_bytes()` is greater than or equal to `TARGET_STREAM_SIZE_BYTES`.
-const TARGET_STREAM_SIZE_BYTES: usize = 10 * 1024 * 1024;
-
-/// Maximum number of messages in a stream.
-///
-/// At most `MAX_STREAM_MESSAGES` are enqueued into a stream; but only until its
-/// `count_bytes()` is greater than or equal to `TARGET_STREAM_SIZE_BYTES`.
-const MAX_STREAM_MESSAGES: usize = 10_000;
-
 const METRIC_STREAM_MESSAGES: &str = "mr_stream_messages";
 const METRIC_STREAM_BYTES: &str = "mr_stream_bytes";
 const METRIC_STREAM_BEGIN: &str = "mr_stream_begin";
@@ -173,8 +161,14 @@ pub(crate) trait StreamBuilder: Send {
     fn build_streams(&self, state: ReplicatedState) -> ReplicatedState;
 }
 
+/// Routes messages from canister output queues into streams, up to the specified limits.
+///
+/// At most `max_stream_messages` are enqueued into a stream; but only until its
+/// `count_bytes()` is greater than or equal to `target_stream_size_bytes`.
 pub(crate) struct StreamBuilderImpl {
     subnet_id: SubnetId,
+    max_stream_messages: usize,
+    target_stream_size_bytes: usize,
     metrics: StreamBuilderMetrics,
     time_in_stream_metrics: Arc<Mutex<LatencyMetrics>>,
 
@@ -187,6 +181,8 @@ pub(crate) struct StreamBuilderImpl {
 impl StreamBuilderImpl {
     pub(crate) fn new(
         subnet_id: SubnetId,
+        max_stream_messages: usize,
+        target_stream_size_bytes: usize,
         metrics_registry: &MetricsRegistry,
         message_routing_metrics: &MessageRoutingMetrics,
         time_in_stream_metrics: Arc<Mutex<LatencyMetrics>>,
@@ -195,6 +191,8 @@ impl StreamBuilderImpl {
     ) -> Self {
         Self {
             subnet_id,
+            max_stream_messages,
+            target_stream_size_bytes,
             metrics: StreamBuilderMetrics::new(metrics_registry, message_routing_metrics),
             time_in_stream_metrics,
             best_effort_responses,
@@ -270,15 +268,8 @@ impl StreamBuilderImpl {
             .observe(msg.payload_size_bytes().get() as f64);
     }
 
-    /// Implementation of `StreamBuilder::build_streams()` that takes a
-    /// `target_stream_size_bytes` argument to limit how many messages will be
-    /// routed into each stream.
-    fn build_streams_impl(
-        &self,
-        mut state: ReplicatedState,
-        max_stream_messages: usize,
-        target_stream_size_bytes: usize,
-    ) -> ReplicatedState {
+    /// Implementation of `StreamBuilder::build_streams()`.
+    fn build_streams_impl(&self, mut state: ReplicatedState) -> ReplicatedState {
         /// Pops the previously peeked message.
         ///
         /// Panics:
@@ -294,23 +285,20 @@ impl StreamBuilderImpl {
             message
         }
 
-        /// Tests whether a stream is over the message count limit, byte limit or (if
-        /// directed at a system subnet) over `2 * SYSTEM_SUBNET_STREAM_MSG_LIMIT`.
-        fn is_at_limit(
-            stream: &btree_map::Entry<SubnetId, Stream>,
-            max_stream_messages: usize,
-            target_stream_size_bytes: usize,
-            is_local_message: bool,
-            destination_subnet_type: SubnetType,
-        ) -> bool {
+        // Tests whether a stream is over the message count limit, byte limit or (if
+        // directed at a system subnet) over `2 * SYSTEM_SUBNET_STREAM_MSG_LIMIT`.
+        let is_at_limit = |stream: &btree_map::Entry<SubnetId, Stream>,
+                           is_local_message: bool,
+                           destination_subnet_type: SubnetType|
+         -> bool {
             let stream = match stream {
                 btree_map::Entry::Occupied(occupied_entry) => occupied_entry.get(),
                 btree_map::Entry::Vacant(_) => return false,
             };
             let stream_messages_len = stream.messages().len();
 
-            if stream_messages_len >= max_stream_messages
-                || stream.count_bytes() >= target_stream_size_bytes
+            if stream_messages_len >= self.max_stream_messages
+                || stream.count_bytes() >= self.target_stream_size_bytes
             {
                 // At limit if message count or byte size limits (enforced across all outgoing
                 // streams) are hit.
@@ -323,7 +311,7 @@ impl StreamBuilderImpl {
             !is_local_message
                 && destination_subnet_type == SubnetType::System
                 && stream_messages_len >= 2 * SYSTEM_SUBNET_STREAM_MSG_LIMIT
-        }
+        };
 
         let mut streams = state.take_streams();
         let routing_table = state.routing_table();
@@ -373,8 +361,6 @@ impl StreamBuilderImpl {
                     let dst_stream_entry = streams.entry(dst_subnet_id);
                     if is_at_limit(
                         &dst_stream_entry,
-                        max_stream_messages,
-                        target_stream_size_bytes,
                         self.subnet_id == dst_subnet_id,
                         destination_subnet_type,
                     ) {
@@ -588,6 +574,6 @@ impl StreamBuilderImpl {
 
 impl StreamBuilder for StreamBuilderImpl {
     fn build_streams(&self, state: ReplicatedState) -> ReplicatedState {
-        self.build_streams_impl(state, MAX_STREAM_MESSAGES, TARGET_STREAM_SIZE_BYTES)
+        self.build_streams_impl(state)
     }
 }

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

@@ -2,6 +2,7 @@ use crate::message_routing::MessageRoutingMetrics;
 
 use super::*;
 use ic_base_types::NumSeconds;
+use ic_config::message_routing::{MAX_STREAM_MESSAGES, TARGET_STREAM_SIZE_BYTES};
 use ic_error_types::RejectCode;
 use ic_management_canister_types_private::Method;
 use ic_registry_routing_table::{CanisterIdRange, RoutingTable};
@@ -336,9 +337,22 @@ fn build_streams_local_canisters() {
 }
 
 #[test]
-fn build_streams_impl_at_limit_leaves_state_untouched() {
+fn build_streams_impl_at_message_limit_leaves_state_untouched() {
+    build_streams_impl_at_limit_leaves_state_untouched_impl(0, usize::MAX);
+}
+
+#[test]
+fn build_streams_impl_at_memory_limit_leaves_state_untouched() {
+    build_streams_impl_at_limit_leaves_state_untouched_impl(usize::MAX, 0);
+}
+
+fn build_streams_impl_at_limit_leaves_state_untouched_impl(
+    max_stream_messages: usize,
+    target_stream_size_bytes: usize,
+) {
     with_test_replica_logger(|log| {
-        let (stream_builder, mut provided_state, metrics_registry) = new_fixture(&log);
+        let (stream_builder, mut provided_state, metrics_registry) =
+            new_fixture_with_limits(&log, max_stream_messages, target_stream_size_bytes);
         provided_state.metadata.network_topology.routing_table = Arc::new(RoutingTable::try_from(
             btreemap! {
                 CanisterIdRange{ start: CanisterId::from(0), end: CanisterId::from(0xfff) } => REMOTE_SUBNET,
@@ -360,10 +374,7 @@ fn build_streams_impl_at_limit_leaves_state_untouched() {
         let expected_state = provided_state.clone();
 
         // Act.
-        let result_state = stream_builder.build_streams_impl(provided_state.clone(), usize::MAX, 0);
-        assert_eq!(result_state, expected_state);
-
-        let result_state = stream_builder.build_streams_impl(provided_state, 0, usize::MAX);
+        let result_state = stream_builder.build_streams_impl(provided_state.clone());
         assert_eq!(result_state, expected_state);
 
         assert_eq!(
@@ -413,18 +424,26 @@ fn build_streams_impl_respects_limits(
     expected_messages: u64,
 ) {
     with_test_replica_logger(|log| {
-        let (stream_builder, mut provided_state, metrics_registry) = new_fixture(&log);
+        let msgs = generate_messages_for_test(/* senders = */ 2, /* receivers = */ 2);
+        let msg_count = msgs.len();
+        // All messages returned by `generate_messages_for_test` are of the same size
+        let msg_size = msgs.first().unwrap().count_bytes() as u64;
+
+        // Target stream size: stream struct plus `max_stream_messages_by_byte_size - 1`
+        // messages plus 1 byte. Since this is a target / soft limit, it should ensure
+        // that exactly `max_stream_messages_by_byte_size` messages (or
+        // `max_stream_messages_by_byte_size * msg_size` bytes) are routed.
+        let target_stream_size_bytes =
+            size_of::<Stream>() + (max_stream_messages_by_byte_size - 1) * msg_size as usize + 1;
+
+        let (stream_builder, mut provided_state, metrics_registry) =
+            new_fixture_with_limits(&log, max_stream_messages, target_stream_size_bytes);
         provided_state.metadata.network_topology.routing_table = Arc::new(RoutingTable::try_from(
             btreemap! {
                 CanisterIdRange{ start: CanisterId::from(0), end: CanisterId::from(0xfff) } => REMOTE_SUBNET,
             },
         ).unwrap());
 
-        let msgs = generate_messages_for_test(/* senders = */ 2, /* receivers = */ 2);
-        let msg_count = msgs.len();
-        // All messages returned by `generate_messages_for_test` are of the same size
-        let msg_size = msgs.first().unwrap().count_bytes() as u64;
-
         assert!(
             msg_count > expected_messages as usize,
             "Invalid test setup: msg_count ({}) must be greater than routed_messages ({})",
@@ -459,19 +478,8 @@ fn build_streams_impl_respects_limits(
             streams.insert(REMOTE_SUBNET, expected_stream);
         });
 
-        // Target stream size: stream struct plus `max_stream_messages_by_byte_size - 1`
-        // messages plus 1 byte. Since this is a target / soft limit, it should ensure
-        // that exactly `max_stream_messages_by_byte_size` messages (or
-        // `max_stream_messages_by_byte_size * msg_size` bytes) are routed.
-        let target_stream_size_bytes =
-            size_of::<Stream>() + (max_stream_messages_by_byte_size - 1) * msg_size as usize + 1;
-
         // Act.
-        let result_state = stream_builder.build_streams_impl(
-            provided_state,
-            max_stream_messages,
-            target_stream_size_bytes,
-        );
+        let result_state = stream_builder.build_streams_impl(provided_state);
 
         assert_eq!(expected_state.canister_states, result_state.canister_states);
         assert_eq!(expected_state.metadata, result_state.metadata);
@@ -998,13 +1006,19 @@ fn build_streams_with_oversized_payloads() {
 }
 
 /// Sets up the `StreamHandlerImpl`, `ReplicatedState` and `MetricsRegistry` to
-/// be used by a test.
-fn new_fixture(log: &ReplicaLogger) -> (StreamBuilderImpl, ReplicatedState, MetricsRegistry) {
+/// be used by a test using specific stream limits.
+fn new_fixture_with_limits(
+    log: &ReplicaLogger,
+    max_stream_messages: usize,
+    target_stream_size_bytes: usize,
+) -> (StreamBuilderImpl, ReplicatedState, MetricsRegistry) {
     let mut state = ReplicatedState::new(LOCAL_SUBNET, SubnetType::Application);
     state.metadata.batch_time = Time::from_nanos_since_unix_epoch(5);
     let metrics_registry = MetricsRegistry::new();
     let stream_builder = StreamBuilderImpl::new(
         LOCAL_SUBNET,
+        max_stream_messages,
+        target_stream_size_bytes,
         &metrics_registry,
         &MessageRoutingMetrics::new(&metrics_registry),
         Arc::new(Mutex::new(LatencyMetrics::new_time_in_stream(
@@ -1017,6 +1031,12 @@ fn new_fixture(log: &ReplicaLogger) -> (StreamBuilderImpl, ReplicatedState, Metr
     (stream_builder, state, metrics_registry)
 }
 
+/// Sets up the `StreamHandlerImpl`, `ReplicatedState` and `MetricsRegistry` to
+/// be used by a test using default stream limits.
+fn new_fixture(log: &ReplicaLogger) -> (StreamBuilderImpl, ReplicatedState, MetricsRegistry) {
+    new_fixture_with_limits(log, MAX_STREAM_MESSAGES, TARGET_STREAM_SIZE_BYTES)
+}
+
 /// Simulates routing the given requests into a `StreamIndexedQueue` with the
 /// given `start` index, until `byte_limit` is reached or exceeded.
 ///