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

…ity#667)

This PR is part of the OnLowWasmMemoryHook effort. Here we implement
checking for hook condition and propagating hook status to SystemState.
Condition for executing hook, as agreed in interface spec
[discussion](dfinity/interface-spec#286 (comment))
is:

1. In the case with memory_allocation `min(memory_allocation -
used_stable_memory, wasm_memory_limit) - used_wasm_memory`
2. Without memory allocation `wasm_memory_limit - used_wasm_memory` 

Hook status can be one of:

1. Condition is not satisfied
3. Ready for execution
4. Executed

The hook condition is checked whenever additional execution (stable or
Wasm) memory allocation is requested.

After this change, we will have all the required information of
SystemState and it will be necessary only to schedule hook execution in
the following round.

This PR will be followed with:

1. [EXC-1724](https://dfinity.atlassian.net/browse/EXC-1724) In which we
will refactor `SystemState::task_queue` to be new struct with additional
field `on_low_wasm_memory_hook_status`, to ensure that whenever
available first task that will be poped form task queue is
`OnLowWasmMemoryHook` (excluding paused executions). For that reason
`SystemState::on_low_wasm_memory_hook status` is private, and we use
`set` and `get`, so it can be easier encapsulated in the future
`TaskQueue` struct.
2. [EXC-1725](https://dfinity.atlassian.net/browse/EXC-1725) Schedule
and execute `OnLowWasmMemoryHook`


[EXC-1724]:
https://dfinity.atlassian.net/browse/EXC-1724?atlOrigin=eyJpIjoiNWRkNTljNzYxNjVmNDY3MDlhMDU5Y2ZhYzA5YTRkZjUiLCJwIjoiZ2l0aHViLWNvbS1KU1cifQ
[EXC-1725]:
https://dfinity.atlassian.net/browse/EXC-1725?atlOrigin=eyJpIjoiNWRkNTljNzYxNjVmNDY3MDlhMDU5Y2ZhYzA5YTRkZjUiLCJwIjoiZ2l0aHViLWNvbS1KU1cifQ

---------

Co-authored-by: Alin Sinpalean <[email protected]>

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

@@ -24,6 +24,7 @@ use ic_types::{
 };
 use ic_wasm_types::BinaryEncodedWasm;
 
+use ic_replicated_state::NumWasmPages;
 use lazy_static::lazy_static;
 use wasmtime::{Engine, Module, Store, StoreLimits, Val};
 
@@ -91,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

@@ -341,6 +341,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";
@@ -430,4 +437,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

@@ -700,6 +700,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>,
 }
@@ -872,3 +874,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,51 @@ 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),
+        }
+    }
 }
 
 /// A wrapper around the different canister statuses.
@@ -740,6 +785,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 +816,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 +845,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 +1605,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/state_layout/src/state_layout.rs

@@ -1,6 +1,3 @@
-use crate::error::LayoutError;
-use crate::utils::do_copy;
-
 use ic_base_types::{NumBytes, NumSeconds};
 use ic_config::flag_status::FlagStatus;
 use ic_logger::{error, info, warn, ReplicaLogger};
@@ -17,7 +14,10 @@ use ic_protobuf::{
 use ic_replicated_state::{
     canister_state::{
         execution_state::{NextScheduledMethod, WasmMetadata},
-        system_state::{wasm_chunk_store::WasmChunkStoreMetadata, CanisterHistory, CyclesUseCase},
+        system_state::{
+            wasm_chunk_store::WasmChunkStoreMetadata, CanisterHistory, CyclesUseCase,
+            OnLowWasmMemoryHookStatus,
+        },
     },
     page_map::{Shard, StorageLayout, StorageResult},
     CallContextManager, CanisterStatus, ExecutionTask, ExportedFunctions, Global, NumWasmPages,
@@ -42,6 +42,9 @@ use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::{Arc, Mutex};
 use std::time::Instant;
 
+use crate::error::LayoutError;
+use crate::utils::do_copy;
+
 #[cfg(test)]
 mod tests;
 
@@ -175,6 +178,7 @@ pub struct CanisterStateBits {
     pub wasm_memory_limit: Option<NumBytes>,
     pub next_snapshot_id: u64,
     pub snapshots_memory_usage: NumBytes,
+    pub on_low_wasm_memory_hook_status: OnLowWasmMemoryHookStatus,
 }
 
 /// This struct contains bits of the `CanisterSnapshot` that are not already
@@ -2281,6 +2285,12 @@ impl From<CanisterStateBits> for pb_canister_state_bits::CanisterStateBits {
             wasm_memory_limit: item.wasm_memory_limit.map(|v| v.get()),
             next_snapshot_id: item.next_snapshot_id,
             snapshots_memory_usage: item.snapshots_memory_usage.get(),
+            on_low_wasm_memory_hook_status: Some(
+                pb_canister_state_bits::OnLowWasmMemoryHookStatus::from(
+                    &item.on_low_wasm_memory_hook_status,
+                )
+                .into(),
+            ),
         }
     }
 }
@@ -2351,6 +2361,12 @@ impl TryFrom<pb_canister_state_bits::CanisterStateBits> for CanisterStateBits {
             );
         }
 
+        let on_low_wasm_memory_hook_status: Option<
+            pb_canister_state_bits::OnLowWasmMemoryHookStatus,
+        > = value.on_low_wasm_memory_hook_status.map(|v| {
+            pb_canister_state_bits::OnLowWasmMemoryHookStatus::try_from(v).unwrap_or_default()
+        });
+
         Ok(Self {
             controllers,
             last_full_execution_round: value.last_full_execution_round.into(),
@@ -2433,6 +2449,11 @@ impl TryFrom<pb_canister_state_bits::CanisterStateBits> for CanisterStateBits {
             wasm_memory_limit: value.wasm_memory_limit.map(NumBytes::from),
             next_snapshot_id: value.next_snapshot_id,
             snapshots_memory_usage: NumBytes::from(value.snapshots_memory_usage),
+            on_low_wasm_memory_hook_status: try_from_option_field(
+                on_low_wasm_memory_hook_status,
+                "CanisterStateBits::on_low_wasm_memory_hook_status",
+            )
+            .unwrap_or_default(),
         })
     }
 }

rs/state_layout/src/state_layout/tests.rs

@@ -4,8 +4,10 @@ use ic_management_canister_types::{
     CanisterChange, CanisterChangeDetails, CanisterChangeOrigin, CanisterInstallMode, IC_00,
 };
 use ic_replicated_state::{
-    canister_state::system_state::CanisterHistory,
-    metadata_state::subnet_call_context_manager::InstallCodeCallId, page_map::Shard, NumWasmPages,
+    canister_state::system_state::{CanisterHistory, OnLowWasmMemoryHookStatus},
+    metadata_state::subnet_call_context_manager::InstallCodeCallId,
+    page_map::Shard,
+    NumWasmPages,
 };
 use ic_test_utilities_logger::with_test_replica_logger;
 use ic_test_utilities_tmpdir::tmpdir;
@@ -58,6 +60,7 @@ fn default_canister_state_bits() -> CanisterStateBits {
         wasm_memory_limit: None,
         next_snapshot_id: 0,
         snapshots_memory_usage: NumBytes::from(0),
+        on_low_wasm_memory_hook_status: OnLowWasmMemoryHookStatus::default(),
     }
 }
 

rs/state_manager/src/checkpoint.rs

@@ -1,7 +1,3 @@
-use crate::{
-    CheckpointError, CheckpointMetrics, HasDowngrade, PageMapType, TipRequest,
-    CRITICAL_ERROR_CHECKPOINT_SOFT_INVARIANT_BROKEN, NUMBER_OF_CHECKPOINT_THREADS,
-};
 use crossbeam_channel::{unbounded, Sender};
 use ic_base_types::{subnet_id_try_from_protobuf, CanisterId, SnapshotId};
 use ic_config::flag_status::FlagStatus;
@@ -29,6 +25,11 @@ use std::convert::TryFrom;
 use std::sync::Arc;
 use std::time::{Duration, Instant};
 
+use crate::{
+    CheckpointError, CheckpointMetrics, HasDowngrade, PageMapType, TipRequest,
+    CRITICAL_ERROR_CHECKPOINT_SOFT_INVARIANT_BROKEN, NUMBER_OF_CHECKPOINT_THREADS,
+};
+
 #[cfg(test)]
 mod tests;
 
@@ -450,6 +451,7 @@ pub fn load_canister_state(
         canister_state_bits.next_snapshot_id,
         canister_state_bits.snapshots_memory_usage,
         metrics,
+        canister_state_bits.on_low_wasm_memory_hook_status,
     );
 
     let canister_state = CanisterState {

rs/state_manager/src/tip.rs

@@ -1077,6 +1077,9 @@ fn serialize_canister_to_tip(
             wasm_memory_limit: canister_state.system_state.wasm_memory_limit,
             next_snapshot_id: canister_state.system_state.next_snapshot_id,
             snapshots_memory_usage: canister_state.system_state.snapshots_memory_usage,
+            on_low_wasm_memory_hook_status: canister_state
+                .system_state
+                .get_on_low_wasm_memory_hook_status(),
         }
         .into(),
     )?;