fix: Check hook status after an upgrade/reinstall/uninstall/install (#3195)

Problem:
@mraszyk pointed out
[here](dfinity/portal#3761 (comment)):

> The hook should not run after an upgrade/reinstall/uninstall/install
if the condition is not satisfied after the upgrade (although the
condition was satisfied before the upgrade and the hook did not execute
yet).

Solution:
This PR should enforce that the status of the hook condition is updated
each time we `upgrade/reinstall/uninstall/install` code.

Author: dragoljub-duric

rs/execution_environment/src/execution_environment.rs

@@ -3070,7 +3070,7 @@ impl ExecutionEnvironment {
         let execution_duration = since.elapsed().as_secs_f64();
         match dts_result {
             DtsInstallCodeResult::Finished {
-                canister,
+                mut canister,
                 mut message,
                 call_id,
                 instructions_used,
@@ -3108,6 +3108,7 @@ impl ExecutionEnvironment {
                         Err(err.into())
                     }
                 };
+                canister.update_on_low_wasm_memory_hook_condition();
                 state.put_canister_state(canister);
                 let refund = message.take_cycles();
                 // The message can be removed because a response was produced.

rs/execution_environment/src/execution_environment/tests.rs

@@ -713,7 +713,7 @@ fn get_canister_status_from_another_canister_when_memory_low() {
     let controller = test.universal_canister().unwrap();
     let binary = wat::parse_str("(module)").unwrap();
     let canister = test.create_canister(Cycles::new(1_000_000_000_000));
-    let memory_allocation = NumBytes::from(158);
+    let memory_allocation = NumBytes::from(300);
     test.install_canister_with_allocation(canister, binary, None, Some(memory_allocation.get()))
         .unwrap();
     let canister_status_args = Encode!(&CanisterIdRecord::from(canister)).unwrap();

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

@@ -2,6 +2,7 @@ use assert_matches::assert_matches;
 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::NextExecution;
 use ic_replicated_state::canister_state::WASM_PAGE_SIZE_IN_BYTES;
@@ -823,6 +824,7 @@ fn global_timer_produces_transient_error_on_out_of_cycles() {
 fn get_wat_with_update_and_hook_mem_grow(
     update_grow_mem_size: i32,
     hook_grow_mem_size: i32,
+    with_enchanced_ortogonal_persistence: bool,
 ) -> String {
     let mut wat = r#"
     (module
@@ -843,7 +845,17 @@ fn get_wat_with_update_and_hook_mem_grow(
     wat.push_str(
         r#")))
     )
-    (memory 1 20)
+    (memory 1 20)"#,
+    );
+    if with_enchanced_ortogonal_persistence {
+        wat.push_str(
+            r#"
+            (@custom "icp:private enhanced-orthogonal-persistence" "")
+            "#,
+        );
+    }
+    wat.push_str(
+        r#"
     )"#,
     );
     wat
@@ -856,7 +868,8 @@ fn on_low_wasm_memory_is_executed() {
     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);
+    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_wat(wat.as_str()).unwrap();
 
@@ -905,7 +918,8 @@ fn on_low_wasm_memory_is_executed_before_message() {
     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);
+    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_wat(wat.as_str()).unwrap();
 
@@ -955,14 +969,178 @@ fn on_low_wasm_memory_is_executed_before_message() {
     );
 }
 
+#[test]
+fn on_low_wasm_memory_is_executed_after_upgrade_if_condition_holds() {
+    let mut test = ExecutionTestBuilder::new().with_manual_execution().build();
+
+    let update_grow_mem_size = 7;
+    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();
+
+    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)
+    );
+
+    test.ingress_raw(canister_id, "grow_mem", vec![]);
+    test.ingress_raw(canister_id, "grow_mem", vec![]);
+
+    // First ingress messages gets executed.
+    // wasm_memory.size = 1 + 7 = 8
+    // wasm_capacity - used_wasm_memory < self.wasm_memory_threshold
+    // Hook condition is triggered.
+    test.execute_slice(canister_id);
+    assert_eq!(
+        test.execution_state(canister_id).wasm_memory.size,
+        NumWasmPages::new(8)
+    );
+
+    let result = test.upgrade_canister_v2(
+        canister_id,
+        wat::parse_str(wat).unwrap(),
+        CanisterUpgradeOptions {
+            skip_pre_upgrade: None,
+            wasm_memory_persistence: Some(WasmMemoryPersistence::Keep),
+        },
+    );
+    assert_eq!(result, Ok(()));
+
+    // Upgrade is executed, and the wasm_memory size is unchanged.
+    // Hook condition is triggered.
+    assert_eq!(
+        test.execution_state(canister_id).wasm_memory.size,
+        NumWasmPages::new(8)
+    );
+
+    // Though we have the second ingress message awaiting to be processed,
+    // hook will be executed first.
+    test.execute_slice(canister_id);
+    assert_eq!(
+        test.execution_state(canister_id).wasm_memory.size,
+        NumWasmPages::new(13)
+    );
+
+    // The second ingress message is executed after the hook.
+    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_not_executed_after_upgrade_if_condition_becomes_unsatisfied() {
+    let mut test = ExecutionTestBuilder::new().with_manual_execution().build();
+
+    let update_grow_mem_size = 3;
+    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, false);
+
+    let canister_id = test.canister_from_wat(wat.as_str()).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)
+    );
+
+    test.ingress_raw(canister_id, "grow_mem", vec![]);
+    test.ingress_raw(canister_id, "grow_mem", vec![]);
+    test.ingress_raw(canister_id, "grow_mem", vec![]);
+
+    // First ingress messages gets executed.
+    // wasm_memory.size = 1 + 3 = 4
+    // wasm_capacity - used_wasm_memory > self.wasm_memory_threshold
+    // Hook condition is not triggered.
+    test.execute_slice(canister_id);
+    assert_eq!(
+        test.execution_state(canister_id).wasm_memory.size,
+        NumWasmPages::new(4)
+    );
+
+    // Second ingress messages gets executed.
+    // wasm_memory.size = 4 + 3 = 7
+    // wasm_capacity - used_wasm_memory < self.wasm_memory_threshold
+    // Hook condition is triggered.
+    test.execute_slice(canister_id);
+    assert_eq!(
+        test.execution_state(canister_id).wasm_memory.size,
+        NumWasmPages::new(7)
+    );
+    println!("canister upgrade");
+
+    let result = test.upgrade_canister_v2(
+        canister_id,
+        wat::parse_str(wat).unwrap(),
+        CanisterUpgradeOptions {
+            skip_pre_upgrade: None,
+            wasm_memory_persistence: None,
+        },
+    );
+    assert_eq!(result, Ok(()));
+    println!("canister upgrade");
+
+    // Upgrade is executed, and the wasm_memory size reset to 1.
+    // Hook condition is not triggered.
+    assert_eq!(
+        test.execution_state(canister_id).wasm_memory.size,
+        NumWasmPages::new(1)
+    );
+
+    // Though the hook was initially scheduled, it is now removed
+    // from queue, and the third ingress message will be executed.
+    test.execute_slice(canister_id);
+    assert_eq!(
+        test.execution_state(canister_id).wasm_memory.size,
+        NumWasmPages::new(4)
+    );
+
+    // There are no messages left to be executed.
+    test.execute_slice(canister_id);
+    assert_eq!(
+        test.execution_state(canister_id).wasm_memory.size,
+        NumWasmPages::new(4)
+    );
+}
+
 #[test]
 fn on_low_wasm_memory_is_executed_once() {
     let mut test = ExecutionTestBuilder::new().build();
 
     let update_grow_mem_size = 7;
     let hook_grow_mem_size = 2;
 
-    let wat = get_wat_with_update_and_hook_mem_grow(update_grow_mem_size, hook_grow_mem_size);
+    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_wat(wat.as_str()).unwrap();
 

rs/replicated_state/src/canister_state.rs

@@ -375,6 +375,13 @@ impl CanisterState {
             + self.system_state.snapshots_memory_usage
     }
 
+    /// Returns the amount of Wasm memory currently used by the canister in bytes.
+    pub fn wasm_memory_usage(&self) -> NumBytes {
+        self.execution_state
+            .as_ref()
+            .map_or(NumBytes::from(0), |es| es.wasm_memory_usage())
+    }
+
     /// Returns the amount of execution memory (heap, stable, globals, Wasm)
     /// currently used by the canister in bytes.
     pub fn execution_memory_usage(&self) -> NumBytes {
@@ -564,6 +571,12 @@ impl CanisterState {
             .map_or(NumBytes::from(0), |es| es.heap_delta())
             + self.system_state.wasm_chunk_store.heap_delta()
     }
+
+    /// Updates status of `OnLowWasmMemory` hook.
+    pub fn update_on_low_wasm_memory_hook_condition(&mut self) {
+        self.system_state
+            .update_on_low_wasm_memory_hook_status(self.memory_usage(), self.wasm_memory_usage());
+    }
 }
 
 /// The result of `next_execution()` function.

rs/replicated_state/src/canister_state/execution_state.rs

@@ -561,13 +561,18 @@ impl ExecutionState {
         self.exports.has_method(method)
     }
 
+    /// Returns the Wasm memory currently used by the `ExecutionState`.
+    pub fn wasm_memory_usage(&self) -> NumBytes {
+        num_bytes_try_from(self.wasm_memory.size)
+            .expect("could not convert from wasm memory number of pages to bytes")
+    }
+
     /// Returns the memory currently used by the `ExecutionState`.
     pub fn memory_usage(&self) -> NumBytes {
         // We use 8 bytes per global.
         let globals_size_bytes = 8 * self.exported_globals.len() as u64;
         let wasm_binary_size_bytes = self.wasm_binary.binary.len() as u64;
-        num_bytes_try_from(self.wasm_memory.size)
-            .expect("could not convert from wasm memory number of pages to bytes")
+        self.wasm_memory_usage()
             + num_bytes_try_from(self.stable_memory.size)
                 .expect("could not convert from stable memory number of pages to bytes")
             + NumBytes::from(globals_size_bytes)