feat: [EXC-1669] Propagate hook execution status to SystemState

rs/canister_sandbox/src/sandbox_server.rs

@@ -153,7 +153,10 @@ mod tests {
     use ic_limits::SMALL_APP_SUBNET_MAX_SIZE;
     use ic_logger::replica_logger::no_op_logger;
     use ic_registry_subnet_type::SubnetType;
-    use ic_replicated_state::{Global, NumWasmPages, PageIndex, PageMap};
+    use ic_replicated_state::{
+        canister_state::system_state::OnLowWasmMemoryHookStatus, Global, NumWasmPages, PageIndex,
+        PageMap,
+    };
     use ic_system_api::{
         sandbox_safe_system_state::{CanisterStatusView, SandboxSafeSystemState},
         ApiType, ExecutionParameters, InstructionLimits,
@@ -228,6 +231,7 @@ mod tests {
             RequestMetadata::new(0, Time::from_nanos_since_unix_epoch(0)),
             caller,
             0,
+            OnLowWasmMemoryHookStatus::ConditionNotSatisfied,
         )
     }
 

rs/embedders/src/wasm_executor.rs

@@ -597,6 +597,7 @@ pub fn process(
         embedder.config().feature_flags.canister_backtrace,
         embedder.config().max_sum_exported_function_name_lengths,
         stable_memory.clone(),
+        wasm_memory.size,
         out_of_instructions_handler,
         logger,
     );

rs/embedders/src/wasmtime_embedder/wasmtime_embedder_tests.rs

@@ -92,6 +92,7 @@ fn test_wasmtime_system_api() {
         EmbeddersConfig::default().feature_flags.canister_backtrace,
         EmbeddersConfig::default().max_sum_exported_function_name_lengths,
         Memory::new_for_testing(),
+        NumWasmPages::from(0),
         Rc::new(DefaultOutOfInstructionsHandler::default()),
         no_op_logger(),
     );

rs/embedders/tests/wasmtime_random_memory_writes.rs

@@ -111,6 +111,7 @@ fn test_api_for_update(
         EmbeddersConfig::default().feature_flags.canister_backtrace,
         EmbeddersConfig::default().max_sum_exported_function_name_lengths,
         Memory::new_for_testing(),
+        NumWasmPages::from(0),
         Rc::new(DefaultOutOfInstructionsHandler::new(instruction_limit)),
         log,
     )

rs/execution_environment/src/canister_manager.rs

@@ -28,14 +28,13 @@ use ic_management_canister_types::{
 };
 use ic_registry_provisional_whitelist::ProvisionalWhitelist;
 use ic_registry_subnet_type::SubnetType;
-use ic_replicated_state::canister_state::system_state::ReservationError;
 use ic_replicated_state::{
     canister_snapshots::{CanisterSnapshot, CanisterSnapshotError},
     canister_state::{
         execution_state::Memory,
         system_state::{
             wasm_chunk_store::{self, WasmChunkStore},
-            CyclesUseCase,
+            CyclesUseCase, ReservationError,
         },
         NextExecution,
     },

rs/protobuf/def/state/canister_state_bits/v1/canister_state_bits.proto

@@ -340,6 +340,13 @@ enum LongExecutionMode {
   LONG_EXECUTION_MODE_PRIORITIZED = 2;
 }
 
+enum OnLowWasmMemoryHookStatus {
+  ON_LOW_WASM_MEMORY_HOOK_STATUS_UNSPECIFIED = 0;
+  ON_LOW_WASM_MEMORY_HOOK_STATUS_CONDITION_NOT_SATISFIED = 1;
+  ON_LOW_WASM_MEMORY_HOOK_STATUS_READY = 2;
+  ON_LOW_WASM_MEMORY_HOOK_STATUS_EXECUTED = 3;
+}
+
 message CanisterStateBits {
   reserved 1;
   reserved "controller";
@@ -429,4 +436,5 @@ message CanisterStateBits {
   int64 priority_credit = 48;
   LongExecutionMode long_execution_mode = 49;
   optional uint64 wasm_memory_threshold = 50;
+  optional OnLowWasmMemoryHookStatus on_low_wasm_memory_hook_status = 53;
 }

rs/protobuf/src/gen/state/state.canister_state_bits.v1.rs

@@ -698,6 +698,8 @@ pub struct CanisterStateBits {
     pub long_execution_mode: i32,
     #[prost(uint64, optional, tag = "50")]
     pub wasm_memory_threshold: ::core::option::Option<u64>,
+    #[prost(enumeration = "OnLowWasmMemoryHookStatus", optional, tag = "53")]
+    pub on_low_wasm_memory_hook_status: ::core::option::Option<i32>,
     #[prost(oneof = "canister_state_bits::CanisterStatus", tags = "11, 12, 13")]
     pub canister_status: ::core::option::Option<canister_state_bits::CanisterStatus>,
 }
@@ -870,3 +872,39 @@ impl LongExecutionMode {
         }
     }
 }
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)]
+#[repr(i32)]
+pub enum OnLowWasmMemoryHookStatus {
+    Unspecified = 0,
+    ConditionNotSatisfied = 1,
+    Ready = 2,
+    Executed = 3,
+}
+impl OnLowWasmMemoryHookStatus {
+    /// String value of the enum field names used in the ProtoBuf definition.
+    ///
+    /// The values are not transformed in any way and thus are considered stable
+    /// (if the ProtoBuf definition does not change) and safe for programmatic use.
+    pub fn as_str_name(&self) -> &'static str {
+        match self {
+            OnLowWasmMemoryHookStatus::Unspecified => "ON_LOW_WASM_MEMORY_HOOK_STATUS_UNSPECIFIED",
+            OnLowWasmMemoryHookStatus::ConditionNotSatisfied => {
+                "ON_LOW_WASM_MEMORY_HOOK_STATUS_CONDITION_NOT_SATISFIED"
+            }
+            OnLowWasmMemoryHookStatus::Ready => "ON_LOW_WASM_MEMORY_HOOK_STATUS_READY",
+            OnLowWasmMemoryHookStatus::Executed => "ON_LOW_WASM_MEMORY_HOOK_STATUS_EXECUTED",
+        }
+    }
+    /// Creates an enum from field names used in the ProtoBuf definition.
+    pub fn from_str_name(value: &str) -> ::core::option::Option<Self> {
+        match value {
+            "ON_LOW_WASM_MEMORY_HOOK_STATUS_UNSPECIFIED" => Some(Self::Unspecified),
+            "ON_LOW_WASM_MEMORY_HOOK_STATUS_CONDITION_NOT_SATISFIED" => {
+                Some(Self::ConditionNotSatisfied)
+            }
+            "ON_LOW_WASM_MEMORY_HOOK_STATUS_READY" => Some(Self::Ready),
+            "ON_LOW_WASM_MEMORY_HOOK_STATUS_EXECUTED" => Some(Self::Executed),
+            _ => None,
+        }
+    }
+}

rs/replicated_state/src/canister_state/system_state.rs

@@ -371,6 +371,96 @@ pub struct SystemState {
     /// This amount contributes to the total `memory_usage` of the canister as
     /// reported by `CanisterState::memory_usage`.
     pub snapshots_memory_usage: NumBytes,
+
+    /// Status of low_on_wasm_memory hook execution.
+    on_low_wasm_memory_hook_status: OnLowWasmMemoryHookStatus,
+}
+
+/// A wrapper around the different statuses of `OnLowWasmMemory` hook execution.
+#[derive(Clone, Copy, Eq, PartialEq, Debug, Default, Deserialize, Serialize)]
+pub enum OnLowWasmMemoryHookStatus {
+    #[default]
+    ConditionNotSatisfied,
+    Ready,
+    Executed,
+}
+
+impl From<&OnLowWasmMemoryHookStatus> for pb::OnLowWasmMemoryHookStatus {
+    fn from(item: &OnLowWasmMemoryHookStatus) -> Self {
+        use OnLowWasmMemoryHookStatus::*;
+
+        match *item {
+            ConditionNotSatisfied => Self::ConditionNotSatisfied,
+            Ready => Self::Ready,
+            Executed => Self::Executed,
+        }
+    }
+}
+
+impl TryFrom<pb::OnLowWasmMemoryHookStatus> for OnLowWasmMemoryHookStatus {
+    type Error = ProxyDecodeError;
+
+    fn try_from(value: pb::OnLowWasmMemoryHookStatus) -> Result<Self, Self::Error> {
+        match value {
+            pb::OnLowWasmMemoryHookStatus::Unspecified => Err(ProxyDecodeError::ValueOutOfRange {
+                typ: "OnLowWasmMemoryHookStatus",
+                err: format!(
+                    "Unexpected value of status of on low wasm memory hook: {:?}",
+                    value
+                ),
+            }),
+            pb::OnLowWasmMemoryHookStatus::ConditionNotSatisfied => {
+                Ok(OnLowWasmMemoryHookStatus::ConditionNotSatisfied)
+            }
+            pb::OnLowWasmMemoryHookStatus::Ready => Ok(OnLowWasmMemoryHookStatus::Ready),
+            pb::OnLowWasmMemoryHookStatus::Executed => Ok(OnLowWasmMemoryHookStatus::Executed),
+        }
+    }
+}
+
+impl OnLowWasmMemoryHookStatus {
+    pub fn update(
+        &mut self,
+        wasm_memory_threshold: NumBytes,
+        memory_allocation: Option<NumBytes>,
+        wasm_memory_limit: Option<NumBytes>,
+        used_stable_memory: NumBytes,
+        used_wasm_memory: NumBytes,
+    ) {
+        // If wasm memory limit is not set, the default is 4 GiB. Wasm memory
+        // limit is ignored for query methods, response callback handlers,
+        // global timers, heartbeats, and canister pre_upgrade.
+        let wasm_memory_limit =
+            wasm_memory_limit.unwrap_or_else(|| NumBytes::new(4 * 1024 * 1024 * 1024));
+
+        // If the canister has memory allocation, then it maximum allowed Wasm memory
+        // can be calculated as min(memory_allocation - used_stable_memory, wasm_memory_limit).
+        let wasm_capacity = memory_allocation.map_or_else(
+            || wasm_memory_limit,
+            |memory_allocation| {
+                debug_assert!(
+                    used_stable_memory <= memory_allocation,
+                    "Used stable memory: {:?} is larger than memory allocation: {:?}.",
+                    used_stable_memory,
+                    memory_allocation
+                );
+                std::cmp::min(memory_allocation - used_stable_memory, wasm_memory_limit)
+            },
+        );
+
+        // Conceptually we can think that the remaining Wasm memory is
+        // equal to `wasm_capacity - used_wasm_memory` and that should
+        // be compared with `wasm_memory_threshold` when checking for
+        // the condition for the hook. However, since `wasm_memory_limit`
+        // is ignored in some executions as stated above it is possible
+        // that `used_wasm_memory` is greater than `wasm_capacity` to
+        // avoid overflowing subtraction we adopted inequality.
+        if wasm_capacity >= used_wasm_memory + wasm_memory_threshold {
+            *self = OnLowWasmMemoryHookStatus::ConditionNotSatisfied;
+        } else if *self != OnLowWasmMemoryHookStatus::Executed {
+            *self = OnLowWasmMemoryHookStatus::Ready;
+        }
+    }
 }
 
 /// A wrapper around the different canister statuses.
@@ -740,6 +830,7 @@ impl SystemState {
             wasm_memory_limit: None,
             next_snapshot_id: 0,
             snapshots_memory_usage: NumBytes::from(0),
+            on_low_wasm_memory_hook_status: OnLowWasmMemoryHookStatus::default(),
         }
     }
 
@@ -770,6 +861,7 @@ impl SystemState {
         next_snapshot_id: u64,
         snapshots_memory_usage: NumBytes,
         metrics: &dyn CheckpointLoadingMetrics,
+        on_low_wasm_memory_hook_status: OnLowWasmMemoryHookStatus,
     ) -> Self {
         let system_state = Self {
             controllers,
@@ -798,6 +890,7 @@ impl SystemState {
             wasm_memory_limit,
             next_snapshot_id,
             snapshots_memory_usage,
+            on_low_wasm_memory_hook_status,
         };
         system_state.check_invariants().unwrap_or_else(|msg| {
             metrics.observe_broken_soft_invariant(msg);
@@ -1557,6 +1650,14 @@ impl SystemState {
             _ => None,
         }
     }
+
+    pub fn set_on_low_wasm_memory_hook_status(&mut self, status: OnLowWasmMemoryHookStatus) {
+        self.on_low_wasm_memory_hook_status = status;
+    }
+
+    pub fn get_on_low_wasm_memory_hook_status(&self) -> OnLowWasmMemoryHookStatus {
+        self.on_low_wasm_memory_hook_status
+    }
 }
 
 /// Implements memory limits verification for pushing a canister-to-canister

rs/replicated_state/src/canister_state/tests.rs

@@ -6,7 +6,7 @@ use crate::canister_state::execution_state::CustomSection;
 use crate::canister_state::execution_state::CustomSectionType;
 use crate::canister_state::execution_state::WasmMetadata;
 use crate::canister_state::system_state::{
-    CallContextManager, CanisterHistory, CanisterStatus, CyclesUseCase,
+    CallContextManager, CanisterHistory, CanisterStatus, CyclesUseCase, OnLowWasmMemoryHookStatus,
     MAX_CANISTER_HISTORY_CHANGES,
 };
 use crate::metadata_state::subnet_call_context_manager::InstallCodeCallId;
@@ -1100,3 +1100,86 @@ fn reverts_stopping_status_after_split() {
 
     assert_eq!(expected_state, canister_state);
 }
+
+const GIB: u64 = 1 << 30;
+
+fn helper_is_condition_satisfied_for_on_low_wasm_memory_hook(
+    wasm_memory_threshold: u64,
+    memory_allocation: Option<u64>,
+    wasm_memory_limit: Option<u64>,
+    used_stable_memory: u64,
+    used_wasm_memory: u64,
+) -> bool {
+    let wasm_memory_limit = wasm_memory_limit.unwrap_or(4 * GIB);
+
+    let wasm_capacity = memory_allocation.map_or(wasm_memory_limit, |memory_allocation| {
+        std::cmp::min(memory_allocation - used_stable_memory, wasm_memory_limit)
+    });
+
+    wasm_capacity < used_wasm_memory + wasm_memory_threshold
+}
+
+fn helper_test_on_low_wasm_memory_hook(
+    start_status: OnLowWasmMemoryHookStatus,
+    status_if_condition_satisfied: OnLowWasmMemoryHookStatus,
+) {
+    for wasm_memory_threshold in [0, GIB, 2 * GIB, 3 * GIB, 4 * GIB] {
+        for memory_allocation in [None, Some(GIB), Some(2 * GIB), Some(3 * GIB), Some(4 * GIB)] {
+            for wasm_memory_limit in [None, Some(GIB), Some(2 * GIB), Some(3 * GIB), Some(4 * GIB)]
+            {
+                for used_stable_memory in [0, GIB] {
+                    for used_wasm_memory in [0, GIB, 2 * GIB, 3 * GIB, 4 * GIB] {
+                        let mut hook_status = start_status;
+
+                        hook_status.update(
+                            wasm_memory_threshold.into(),
+                            memory_allocation.map(|m| m.into()),
+                            wasm_memory_limit.map(|m| m.into()),
+                            used_stable_memory.into(),
+                            used_wasm_memory.into(),
+                        );
+
+                        assert_eq!(
+                            hook_status,
+                            if helper_is_condition_satisfied_for_on_low_wasm_memory_hook(
+                                wasm_memory_threshold,
+                                memory_allocation,
+                                wasm_memory_limit,
+                                used_stable_memory,
+                                used_wasm_memory
+                            ) {
+                                status_if_condition_satisfied
+                            } else {
+                                OnLowWasmMemoryHookStatus::ConditionNotSatisfied
+                            }
+                        );
+                    }
+                }
+            }
+        }
+    }
+}
+
+#[test]
+fn test_on_low_wasm_memory_hook_start_status_condition_not_satisfied() {
+    helper_test_on_low_wasm_memory_hook(
+        OnLowWasmMemoryHookStatus::ConditionNotSatisfied,
+        OnLowWasmMemoryHookStatus::Ready,
+    );
+}
+
+#[test]
+fn test_on_low_wasm_memory_hook_start_status_ready() {
+    helper_test_on_low_wasm_memory_hook(
+        OnLowWasmMemoryHookStatus::Ready,
+        OnLowWasmMemoryHookStatus::Ready,
+    );
+}
+
+#[test]
+fn test_on_low_wasm_memory_hook_start_status_executed() {
+    helper_test_on_low_wasm_memory_hook(
+        OnLowWasmMemoryHookStatus::Executed,
+        OnLowWasmMemoryHookStatus::Executed,
+    );
+}