fix: [EXC-1831] Low Wasm memory hook will run once the canister's unf…

…rozen if it was scheduled before freezing (#3455)

Problem:
@mraszyk noticed
[here](dfinity/portal#3761 (comment))
that property:
”If the canister gets frozen immediately after the function is scheduled
for execution, the function will run once the canister's unfrozen if the
canister's wasm memory remains above the threshold t.”
is missing.

Solution:
When execution of the hook is stopped because of the freezing canister,
add the hook again to the task queue.

rs/execution_environment/src/execution/update.rs

@@ -65,8 +65,9 @@ pub fn execute_update(
                     .as_ref()
                     .map_or(false, |es| es.is_wasm64);
 
-                let prepaid_execution_cycles =
-                    match round.cycles_account_manager.prepay_execution_cycles(
+                let prepaid_execution_cycles = match round
+                    .cycles_account_manager
+                    .prepay_execution_cycles(
                         &mut canister.system_state,
                         memory_usage,
                         message_memory_usage,
@@ -76,19 +77,35 @@ pub fn execute_update(
                         reveal_top_up,
                         is_wasm64_execution.into(),
                     ) {
-                        Ok(cycles) => cycles,
-                        Err(err) => {
-                            return finish_call_with_error(
-                                UserError::new(ErrorCode::CanisterOutOfCycles, err),
-                                canister,
-                                call_or_task,
-                                NumInstructions::from(0),
-                                round.time,
-                                execution_parameters.subnet_type,
-                                round.log,
-                            );
+                    Ok(cycles) => cycles,
+                    Err(err) => {
+                        if 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 the freezing of the canister. To ensure that the hook is executed
+                            // when the canister is unfrozen 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.
+                            canister
+                                .system_state
+                                .task_queue
+                                .remove(ic_replicated_state::ExecutionTask::OnLowWasmMemory);
+                            canister
+                                .system_state
+                                .task_queue
+                                .enqueue(ic_replicated_state::ExecutionTask::OnLowWasmMemory);
                         }
-                    };
+                        return finish_call_with_error(
+                            UserError::new(ErrorCode::CanisterOutOfCycles, err),
+                            canister,
+                            call_or_task,
+                            NumInstructions::from(0),
+                            round.time,
+                            execution_parameters.subnet_type,
+                            round.log,
+                        );
+                    }
+                };
                 (canister, prepaid_execution_cycles, false)
             }
         };

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

@@ -1,9 +1,11 @@
 use assert_matches::assert_matches;
+use ic_base_types::NumSeconds;
 use ic_config::{execution_environment::Config as HypervisorConfig, subnet_config::SubnetConfig};
 use ic_error_types::RejectCode;
 use ic_management_canister_types::{CanisterSettingsArgsBuilder, CanisterStatusType};
 use ic_management_canister_types::{CanisterUpgradeOptions, WasmMemoryPersistence};
 use ic_registry_subnet_type::SubnetType;
+use ic_replicated_state::canister_state::system_state::OnLowWasmMemoryHookStatus;
 use ic_replicated_state::canister_state::NextExecution;
 use ic_replicated_state::canister_state::WASM_PAGE_SIZE_IN_BYTES;
 use ic_replicated_state::page_map::PAGE_SIZE;
@@ -861,6 +863,116 @@ fn get_wat_with_update_and_hook_mem_grow(
     wat
 }
 
+#[test]
+fn on_low_wasm_memory_hook_is_run_after_freezing() {
+    let mut test = ExecutionTestBuilder::new().with_manual_execution().build();
+
+    let update_grow_mem_size = 7;
+    let hook_grow_mem_size = 5;
+
+    let wat =
+        get_wat_with_update_and_hook_mem_grow(update_grow_mem_size, hook_grow_mem_size, false);
+
+    let canister_id = test
+        .canister_from_cycles_and_wat(Cycles::new(200_000_000_000_000), wat)
+        .unwrap();
+
+    test.canister_update_wasm_memory_limit_and_wasm_memory_threshold(
+        canister_id,
+        (20 * WASM_PAGE_SIZE_IN_BYTES as u64).into(),
+        (15 * WASM_PAGE_SIZE_IN_BYTES as u64).into(),
+    )
+    .unwrap();
+
+    // Here we have:
+    // wasm_capacity = wasm_memory_limit = 20 Wasm Pages
+    // wasm_memory_threshold = 15 Wasm Pages
+
+    // Initially wasm_memory.size = 1
+    assert_eq!(
+        test.execution_state(canister_id).wasm_memory.size,
+        NumWasmPages::new(1)
+    );
+
+    // Two ingress messages are sent.
+    test.ingress_raw(canister_id, "grow_mem", vec![]);
+    test.ingress_raw(canister_id, "grow_mem", vec![]);
+
+    // The first ingress message gets executed.
+    // wasm_memory.size = 1 + 7 = 8
+    test.execute_slice(canister_id);
+    assert_eq!(
+        test.execution_state(canister_id).wasm_memory.size,
+        NumWasmPages::new(8)
+    );
+
+    // wasm_capacity - used_wasm_memory < self.wasm_memory_threshold
+    // The hook condition is triggered.
+    // Hence hook should be executed next.
+    assert_eq!(
+        test.state()
+            .canister_states
+            .get(&canister_id)
+            .unwrap()
+            .system_state
+            .task_queue
+            .peek_hook_status(),
+        OnLowWasmMemoryHookStatus::Ready
+    );
+
+    // We update `freezing_threshold` making canister frozen.
+    test.update_freezing_threshold(canister_id, NumSeconds::new(100_000_000_000_000))
+        .unwrap();
+
+    // The execution of the hook is not finished due to freezing.
+    test.execute_slice(canister_id);
+    assert_eq!(
+        test.execution_state(canister_id).wasm_memory.size,
+        NumWasmPages::new(8)
+    );
+
+    // The hook status is still `Ready`.
+    assert_eq!(
+        test.state()
+            .canister_states
+            .get(&canister_id)
+            .unwrap()
+            .system_state
+            .task_queue
+            .peek_hook_status(),
+        OnLowWasmMemoryHookStatus::Ready
+    );
+
+    // We update `freezing_threshold` unfreezing canister.
+    test.update_freezing_threshold(canister_id, NumSeconds::new(100))
+        .unwrap();
+
+    // The hook is executed.
+    test.execute_slice(canister_id);
+    assert_eq!(
+        test.execution_state(canister_id).wasm_memory.size,
+        NumWasmPages::new(13)
+    );
+
+    assert_eq!(
+        test.state()
+            .canister_states
+            .get(&canister_id)
+            .unwrap()
+            .system_state
+            .task_queue
+            .peek_hook_status(),
+        OnLowWasmMemoryHookStatus::Executed
+    );
+
+    // The second ingress message is executed.
+    test.execute_slice(canister_id);
+    assert_eq!(
+        test.execution_state(canister_id).wasm_memory.size,
+        NumWasmPages::new(20)
+    );
+}
+
 #[test]
 fn on_low_wasm_memory_is_executed() {
     let mut test = ExecutionTestBuilder::new().build();