refactor: Provide a more efficient `ReplicatedState::message_memory_t…

…aken()` method

rs/execution_environment/src/execution_environment.rs

@@ -389,7 +389,7 @@ impl ExecutionEnvironment {
         &self.metrics.canister_not_found_error
     }
 
-    /// Look up the current amount of memory available on the subnet.
+    /// Computes the current amount of memory available on the subnet.
     pub fn subnet_available_memory(&self, state: &ReplicatedState) -> SubnetAvailableMemory {
         let memory_taken = state.memory_taken();
         SubnetAvailableMemory::new(
@@ -405,6 +405,15 @@ impl ExecutionEnvironment {
         )
     }
 
+    /// Computes the current amount of message memory available on the subnet.
+    ///
+    /// This is a more efficient alternative to `memory_taken()` for cases when only
+    /// the message memory usage is necessary.
+    pub fn subnet_available_message_memory(&self, state: &ReplicatedState) -> i64 {
+        self.config.subnet_message_memory_capacity.get() as i64
+            - state.message_memory_taken().get() as i64
+    }
+
     /// Executes a replicated message sent to a subnet.
     /// Returns the new replicated state and the number of left instructions.
     #[allow(clippy::cognitive_complexity)]

rs/execution_environment/src/scheduler.rs

@@ -1160,10 +1160,7 @@ impl SchedulerImpl {
     /// through message routing.
     pub fn induct_messages_on_same_subnet(&self, state: &mut ReplicatedState) {
         // Compute subnet available memory *before* taking out the canisters.
-        let mut subnet_available_memory = self
-            .exec_env
-            .subnet_available_memory(state)
-            .get_message_memory();
+        let mut subnet_available_memory = self.exec_env.subnet_available_message_memory(state);
 
         let mut canisters = state.take_canister_states();
 

rs/messaging/src/routing/stream_handler.rs

@@ -812,9 +812,7 @@ impl StreamHandlerImpl {
     /// Computes the subnet's available message memory, as the difference
     /// between the subnet's message memory capacity and its current usage.
     fn subnet_available_memory(&self, state: &ReplicatedState) -> i64 {
-        let memory_taken = state.memory_taken();
-        let message_memory_taken = memory_taken.messages();
-        self.subnet_message_memory_capacity.get() as i64 - message_memory_taken.get() as i64
+        self.subnet_message_memory_capacity.get() as i64 - state.message_memory_taken().get() as i64
     }
 
     /// Observes "time in backlog" (since learning about their existence from

rs/replicated_state/src/replicated_state.rs

@@ -588,7 +588,7 @@ impl ReplicatedState {
             .sum()
     }
 
-    /// Returns the memory taken by different types of memory resources.
+    /// Computes the memory taken by different types of memory resources.
     pub fn memory_taken(&self) -> MemoryTaken {
         let (
             raw_memory_taken,
@@ -633,6 +633,20 @@ impl ReplicatedState {
         }
     }
 
+    /// Computes the memory taken by messages.
+    ///
+    /// This is a more efficient alternative to `memory_taken()` for cases when only
+    /// the message memory usage is necessary.
+    pub fn message_memory_taken(&self) -> NumBytes {
+        let canisters_memory_usage: NumBytes = self
+            .canisters_iter()
+            .map(|canister| canister.system_state.message_memory_usage())
+            .sum();
+        let subnet_memory_usage = (self.subnet_queues.memory_usage() as u64).into();
+
+        canisters_memory_usage + subnet_memory_usage
+    }
+
     /// Returns the total memory taken by the ingress history in bytes.
     pub fn total_ingress_memory_taken(&self) -> NumBytes {
         self.metadata.ingress_history.memory_usage()

rs/replicated_state/tests/replicated_state.rs

@@ -158,6 +158,10 @@ impl ReplicatedStateFixture {
         self.state.memory_taken()
     }
 
+    fn message_memory_taken(&self) -> NumBytes {
+        self.state.message_memory_taken()
+    }
+
     fn remote_subnet_input_schedule(&self, canister: &CanisterId) -> &VecDeque<CanisterId> {
         self.state
             .canister_state(canister)
@@ -189,6 +193,10 @@ fn assert_message_memory_taken(queues_memory_usage: usize, fixture: &ReplicatedS
         queues_memory_usage as u64,
         fixture.memory_taken().messages().get()
     );
+    assert_eq!(
+        queues_memory_usage as u64,
+        fixture.message_memory_taken().get()
+    );
 }
 
 fn assert_canister_history_memory_taken(