refactor: EXC-1890 Use saturating arithmetics when evaluating hook co…

…nditions (#3879)
dfinity · Feb 10, 2025 · 63e785c · 63e785c
1 parent c309ff5
commit 63e785c
Show file tree

Hide file tree

Showing 4 changed files with 6 additions and 3 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/rs/replicated_state/BUILD.bazel b/rs/replicated_state/BUILD.bazel
@@ -34,6 +34,7 @@ DEPENDENCIES = [
     "@crate_index//:libc",
     "@crate_index//:maplit",
     "@crate_index//:nix",
+    "@crate_index//:num-traits",
     "@crate_index//:prometheus",
     "@crate_index//:prost",
     "@crate_index//:rand",

diff --git a/rs/replicated_state/Cargo.toml b/rs/replicated_state/Cargo.toml
@@ -36,6 +36,7 @@ lazy_static = { workspace = true }
 libc = { workspace = true }
 maplit = "1.0.2"
 nix = { workspace = true }
+num-traits = { workspace = true }
 phantom_newtype = { path = "../phantom_newtype" }
 prometheus = { workspace = true }
 prost = { workspace = true }

diff --git a/rs/replicated_state/src/canister_state/system_state/task_queue.rs b/rs/replicated_state/src/canister_state/system_state/task_queue.rs
@@ -5,6 +5,7 @@ use ic_protobuf::proxy::ProxyDecodeError;
 use ic_protobuf::state::canister_state_bits::v1 as pb;
 use ic_types::CanisterId;
 use ic_types::NumBytes;
+use num_traits::SaturatingSub;
 use serde::{Deserialize, Serialize};
 use std::collections::VecDeque;
 
@@ -383,8 +384,7 @@ pub fn is_low_wasm_memory_hook_condition_satisfied(
         memory_usage
     );
 
-    let memory_usage_without_wasm_memory =
-        NumBytes::new(memory_usage.get().saturating_sub(wasm_memory_usage.get()));
+    let memory_usage_without_wasm_memory = memory_usage.saturating_sub(&wasm_memory_usage);
 
     // If the canister has memory allocation, then it maximum allowed Wasm memory can be calculated
     // as min(memory_allocation - memory_usage_without_wasm_memory, wasm_memory_limit).
@@ -398,7 +398,7 @@ pub fn is_low_wasm_memory_hook_condition_satisfied(
                 memory_allocation
             );
             std::cmp::min(
-                memory_allocation - memory_usage_without_wasm_memory,
+                memory_allocation.saturating_sub(&memory_usage_without_wasm_memory),
                 wasm_memory_limit,
             )
         },