diff --git a/rs/execution_environment/src/execution_environment.rs b/rs/execution_environment/src/execution_environment.rs index b4a44e45bd4..622f0553d48 100644 --- a/rs/execution_environment/src/execution_environment.rs +++ b/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)] diff --git a/rs/execution_environment/src/scheduler.rs b/rs/execution_environment/src/scheduler.rs index 6d864440e5c..0fe5edaf1ce 100644 --- a/rs/execution_environment/src/scheduler.rs +++ b/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(); diff --git a/rs/messaging/src/routing/stream_handler.rs b/rs/messaging/src/routing/stream_handler.rs index 1909d0fceb3..33879b32710 100644 --- a/rs/messaging/src/routing/stream_handler.rs +++ b/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 diff --git a/rs/replicated_state/src/replicated_state.rs b/rs/replicated_state/src/replicated_state.rs index 0054493f24e..013658bd7b9 100644 --- a/rs/replicated_state/src/replicated_state.rs +++ b/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() diff --git a/rs/replicated_state/tests/replicated_state.rs b/rs/replicated_state/tests/replicated_state.rs index d3635affb26..c0c6a6d32cc 100644 --- a/rs/replicated_state/tests/replicated_state.rs +++ b/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 { 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(