fix: EXC-1838 Run hook after CanisterWasmMemoryLimitExceeded error is fixed

rs/execution_environment/src/execution/update.rs

@@ -51,7 +51,7 @@ pub fn execute_update(
     log_dirty_pages: FlagStatus,
     deallocation_sender: &DeallocationSender,
 ) -> ExecuteMessageResult {
-    let (clean_canister, prepaid_execution_cycles, resuming_aborted) =
+    let (mut clean_canister, prepaid_execution_cycles, resuming_aborted) =
         match prepaid_execution_cycles {
             Some(prepaid_execution_cycles) => (clean_canister, prepaid_execution_cycles, true),
             None => {
@@ -147,13 +147,31 @@ pub fn execute_update(
     let helper = match UpdateHelper::new(&clean_canister, &original, deallocation_sender) {
         Ok(helper) => helper,
         Err(err) => {
+            if err.code() == ErrorCode::CanisterWasmMemoryLimitExceeded
+                && original.call_or_task == CanisterCallOrTask::Task(CanisterTask::OnLowWasmMemory)
+            {
+                //`OnLowWasmMemoryHook` is taken from task_queue (i.e. `OnLowWasmMemoryHookStatus` is `Executed`),
+                // but its was not executed due to error `WasmMemoryLimitExceeded`. To ensure that the hook is executed
+                // when the error is resolved we need to set `OnLowWasmMemoryHookStatus` to `Ready`. Because of
+                // the way `OnLowWasmMemoryHookStatus::update` is implemented we first need to remove it from the
+                // task_queue (which calls `OnLowWasmMemoryHookStatus::update(false)`) followed with `enqueue`
+                // (which calls `OnLowWasmMemoryHookStatus::update(true)`) to ensure desired behavior.
+                clean_canister
+                    .system_state
+                    .task_queue
+                    .remove(ic_replicated_state::ExecutionTask::OnLowWasmMemory);
+                clean_canister
+                    .system_state
+                    .task_queue
+                    .enqueue(ic_replicated_state::ExecutionTask::OnLowWasmMemory);
+            }
             return finish_err(
                 clean_canister,
                 original.execution_parameters.instruction_limits.message(),
                 err,
                 original,
                 round,
-            )
+            );
         }
     };
 
@@ -341,26 +359,22 @@ impl UpdateHelper {
 
         validate_message(&canister, &original.method)?;
 
-        if let CanisterCallOrTask::Call(_) = original.call_or_task {
-            // TODO(RUN-957): Enforce the limit in heartbeat and timer after
-            // canister logging ships by removing the `if` above.
+        let wasm_memory_usage = canister
+            .execution_state
+            .as_ref()
+            .map_or(NumBytes::new(0), |es| {
+                num_bytes_try_from(es.wasm_memory.size).unwrap()
+            });
 
-            let wasm_memory_usage = canister
-                .execution_state
-                .as_ref()
-                .map_or(NumBytes::new(0), |es| {
-                    num_bytes_try_from(es.wasm_memory.size).unwrap()
-                });
+        if let Some(wasm_memory_limit) = clean_canister.system_state.wasm_memory_limit {
+            // A Wasm memory limit of 0 means unlimited.
+            if wasm_memory_limit.get() != 0 && wasm_memory_usage > wasm_memory_limit {
+                let err = HypervisorError::WasmMemoryLimitExceeded {
+                    bytes: wasm_memory_usage,
+                    limit: wasm_memory_limit,
+                };
 
-            if let Some(wasm_memory_limit) = clean_canister.system_state.wasm_memory_limit {
-                // A Wasm memory limit of 0 means unlimited.
-                if wasm_memory_limit.get() != 0 && wasm_memory_usage > wasm_memory_limit {
-                    let err = HypervisorError::WasmMemoryLimitExceeded {
-                        bytes: wasm_memory_usage,
-                        limit: wasm_memory_limit,
-                    };
-                    return Err(err.into_user_error(&canister.canister_id()));
-                }
+                return Err(err.into_user_error(&canister.canister_id()));
             }
         }
 

rs/execution_environment/src/execution_environment/tests/canister_task.rs

@@ -1443,3 +1443,103 @@ fn on_low_wasm_memory_is_executed_after_growing_stable_memory() {
         NumWasmPages::new(6)
     );
 }
+
+#[test]
+fn on_low_wasm_memory_hook_is_run_after_memory_surpass_limit() {
+    let mut test = ExecutionTestBuilder::new().with_manual_execution().build();
+
+    let update_grow_mem_size = 10;
+    let hook_grow_mem_size = 5;
+
+    let wat: String =
+        get_wat_with_update_and_hook_mem_grow(update_grow_mem_size, hook_grow_mem_size, true);
+
+    let canister_id = test.canister_from_wat(wat.as_str()).unwrap();
+
+    // Initially wasm_memory.size = 1
+    assert_eq!(
+        test.execution_state(canister_id).wasm_memory.size,
+        NumWasmPages::new(1)
+    );
+
+    test.ingress_raw(canister_id, "grow_mem", vec![]);
+
+    // First ingress messages gets executed.
+    // wasm_memory.size = 1 + 10 = 11
+    test.execute_slice(canister_id);
+
+    assert_eq!(
+        test.execution_state(canister_id).wasm_memory.size,
+        NumWasmPages::new(11)
+    );
+
+    // We update `wasm_memory_limit` to be smaller than `used_wasm_memory`.
+    test.canister_update_wasm_memory_limit_and_wasm_memory_threshold(
+        canister_id,
+        (10 * WASM_PAGE_SIZE_IN_BYTES as u64).into(),
+        (5 * WASM_PAGE_SIZE_IN_BYTES as u64).into(),
+    )
+    .unwrap();
+
+    // The update will also trigger `low_wasm_memory` hook.
+    assert_eq!(
+        test.state()
+            .canister_states
+            .get(&canister_id)
+            .unwrap()
+            .system_state
+            .task_queue
+            .peek_hook_status(),
+        OnLowWasmMemoryHookStatus::Ready
+    );
+
+    // Hook execution will not succeed since `used_wasm_memory` > `wasm_memory_limit`.
+    test.execute_slice(canister_id);
+
+    assert_eq!(
+        test.execution_state(canister_id).wasm_memory.size,
+        NumWasmPages::new(11)
+    );
+
+    // After execution of the hook fails, hook status will remain `Ready`.
+    assert_eq!(
+        test.state()
+            .canister_states
+            .get(&canister_id)
+            .unwrap()
+            .system_state
+            .task_queue
+            .peek_hook_status(),
+        OnLowWasmMemoryHookStatus::Ready
+    );
+
+    // We fix the error by setting `wasm_memory_limit` > `used_wasm_memory`.
+    // At the same time:
+    // `wasm_memory_limit` - `used_wasm_memory` < `wasm_memory_threshold`
+    // condition for `low_wasm_memory` hook remains satisfied.
+    // Hence, `low_wasm_memory` hook execution will follow.
+    test.canister_update_wasm_memory_limit_and_wasm_memory_threshold(
+        canister_id,
+        (20 * WASM_PAGE_SIZE_IN_BYTES as u64).into(),
+        (10 * WASM_PAGE_SIZE_IN_BYTES as u64).into(),
+    )
+    .unwrap();
+
+    test.execute_slice(canister_id);
+
+    assert_eq!(
+        test.execution_state(canister_id).wasm_memory.size,
+        NumWasmPages::new(16)
+    );
+
+    assert_eq!(
+        test.state()
+            .canister_states
+            .get(&canister_id)
+            .unwrap()
+            .system_state
+            .task_queue
+            .peek_hook_status(),
+        OnLowWasmMemoryHookStatus::Executed
+    );
+}