feat(RUN-991): Set the default Wasm memory limit (#1040)

This PR sets uninitialized `wasm_memory_limit` field in canister
settings to the default value, which is computed as follows:

- let `usage` be the Wasm memory size of the canister.
- let `halfway_to_max` be `(usage + 4GiB) / 2`.
- use `max(3GiB, halfway_to_max)` as the default value.

Note that if the field has already been initialized, then that value is
not overwritten.

For newly created canisters, the field is set to `3GiB`.

The initialization for existing canisters will happen at the end of the
first execution round after the replica upgrade.

rs/config/src/execution_environment.rs

@@ -140,6 +140,13 @@ pub const MAX_NUMBER_OF_SNAPSHOTS_PER_CANISTER: usize = 1;
 /// The worst case request latency used here should be equivalent to the request timeout in the adapter.
 pub const MAX_CANISTER_HTTP_REQUESTS_IN_FLIGHT: usize = 3000;
 
+/// The default value of `wasm_memory_limit` in the canister settings:
+/// - this value is used directly for newly created canisters.
+/// - existing canisters will get their field initialized as follows:
+///   - let `halfway_to_max = (memory_usage + 4GiB) / 2`
+///   - use the maximum of `default_wasm_memory_limit` and `halfway_to_max`.
+pub const DEFAULT_WASM_MEMORY_LIMIT: NumBytes = NumBytes::new(3 * GIB);
+
 #[derive(Clone, Debug, Deserialize, PartialEq, Eq, Serialize)]
 #[serde(default)]
 pub struct Config {
@@ -282,6 +289,13 @@ pub struct Config {
     pub ic00_sign_with_schnorr: FlagStatus,
 
     pub max_canister_http_requests_in_flight: usize,
+
+    /// The default value of `wasm_memory_limit` in the canister settings:
+    /// - this value is used directly for newly created canisters.
+    /// - existing canisters will get their field initialized as follows:
+    ///   - let `halfway_to_max = (memory_usage + 4GiB) / 2`
+    ///   - use the maximum of `default_wasm_memory_limit` and `halfway_to_max`.
+    pub default_wasm_memory_limit: NumBytes,
 }
 
 impl Default for Config {
@@ -356,6 +370,7 @@ impl Default for Config {
             ic00_schnorr_public_key: FlagStatus::Enabled,
             ic00_sign_with_schnorr: FlagStatus::Enabled,
             max_canister_http_requests_in_flight: MAX_CANISTER_HTTP_REQUESTS_IN_FLIGHT,
+            default_wasm_memory_limit: DEFAULT_WASM_MEMORY_LIMIT,
         }
     }
 }

rs/execution_environment/src/canister_manager.rs

@@ -125,6 +125,7 @@ pub(crate) struct CanisterMgrConfig {
     upload_wasm_chunk_instructions: NumInstructions,
     wasm_chunk_store_max_size: NumBytes,
     canister_snapshot_baseline_instructions: NumInstructions,
+    default_wasm_memory_limit: NumBytes,
 }
 
 impl CanisterMgrConfig {
@@ -145,6 +146,7 @@ impl CanisterMgrConfig {
         upload_wasm_chunk_instructions: NumInstructions,
         wasm_chunk_store_max_size: NumBytes,
         canister_snapshot_baseline_instructions: NumInstructions,
+        default_wasm_memory_limit: NumBytes,
     ) -> Self {
         Self {
             subnet_memory_capacity,
@@ -162,6 +164,7 @@ impl CanisterMgrConfig {
             upload_wasm_chunk_instructions,
             wasm_chunk_store_max_size,
             canister_snapshot_baseline_instructions,
+            default_wasm_memory_limit,
         }
     }
 }
@@ -762,6 +765,10 @@ impl CanisterManager {
             .reserved_cycles_limit
             .get_or_insert_with(|| self.cycles_account_manager.default_reserved_balance_limit());
 
+        settings
+            .wasm_memory_limit
+            .get_or_insert(self.config.default_wasm_memory_limit);
+
         // Validate settings before `create_canister_helper` applies them
         match self.validate_settings_for_canister_creation(
             settings,

rs/execution_environment/src/canister_manager/tests.rs

@@ -14,7 +14,9 @@ use assert_matches::assert_matches;
 use candid::Decode;
 use ic_base_types::{NumSeconds, PrincipalId};
 use ic_config::{
-    execution_environment::Config, flag_status::FlagStatus, subnet_config::SchedulerConfig,
+    execution_environment::{Config, DEFAULT_WASM_MEMORY_LIMIT},
+    flag_status::FlagStatus,
+    subnet_config::SchedulerConfig,
 };
 use ic_constants::SMALL_APP_SUBNET_MAX_SIZE;
 use ic_cycles_account_manager::{CyclesAccountManager, ResourceSaturation};
@@ -295,6 +297,7 @@ fn canister_manager_config(
         SchedulerConfig::application_subnet().upload_wasm_chunk_instructions,
         ic_config::embedders::Config::default().wasm_max_size,
         SchedulerConfig::application_subnet().canister_snapshot_baseline_instructions,
+        DEFAULT_WASM_MEMORY_LIMIT,
     )
 }
 

rs/execution_environment/src/execution_environment.rs

@@ -397,6 +397,7 @@ impl ExecutionEnvironment {
             upload_wasm_chunk_instructions,
             config.embedders_config.wasm_max_size,
             canister_snapshot_baseline_instructions,
+            config.default_wasm_memory_limit,
         );
         let metrics = ExecutionEnvironmentMetrics::new(metrics_registry);
         let canister_manager = CanisterManager::new(
@@ -3561,6 +3562,11 @@ impl ExecutionEnvironment {
         )
     }
 
+    /// Returns the default value of `wasm_memory_limit` in canister settings.
+    pub fn default_wasm_memory_limit(&self) -> NumBytes {
+        self.config.default_wasm_memory_limit
+    }
+
     /// For testing purposes only.
     #[doc(hidden)]
     pub fn hypervisor_for_testing(&self) -> &Hypervisor {

rs/execution_environment/src/scheduler.rs

@@ -25,8 +25,10 @@ use ic_replicated_state::{
     canister_state::{
         execution_state::NextScheduledMethod, system_state::CyclesUseCase, NextExecution,
     },
+    num_bytes_try_from,
     page_map::PageAllocatorFileDescriptor,
-    CanisterState, CanisterStatus, ExecutionTask, InputQueueType, NetworkTopology, ReplicatedState,
+    CanisterState, CanisterStatus, ExecutionTask, InputQueueType, NetworkTopology, NumWasmPages,
+    ReplicatedState,
 };
 use ic_system_api::InstructionLimits;
 use ic_types::{
@@ -36,6 +38,7 @@ use ic_types::{
     messages::{CanisterMessage, Ingress, MessageId, Response, StopCanisterContext, NO_DEADLINE},
     AccumulatedPriority, CanisterId, ComputeAllocation, Cycles, ExecutionRound, LongExecutionMode,
     MemoryAllocation, NumBytes, NumInstructions, NumSlices, Randomness, SubnetId, Time,
+    MAX_WASM_MEMORY_IN_BYTES,
 };
 use ic_types::{nominal_cycles::NominalCycles, NumMessages};
 use num_rational::Ratio;
@@ -1340,9 +1343,38 @@ impl SchedulerImpl {
                 }
             }
         }
+        self.initialize_wasm_memory_limit(state);
         self.check_dts_invariants(state, current_round_type);
     }
 
+    fn initialize_wasm_memory_limit(&self, state: &mut ReplicatedState) {
+        fn compute_default_wasm_memory_limit(default: NumBytes, usage: NumBytes) -> NumBytes {
+            // Returns the larger of the two:
+            // - the default value
+            // - the average between the current usage and the hard limit.
+            default.max(NumBytes::new(
+                MAX_WASM_MEMORY_IN_BYTES.saturating_add(usage.get()) / 2,
+            ))
+        }
+
+        let default_wasm_memory_limit = self.exec_env.default_wasm_memory_limit();
+        for (_id, canister) in state.canister_states.iter_mut() {
+            if canister.system_state.wasm_memory_limit.is_none() {
+                let num_wasm_pages = canister
+                    .execution_state
+                    .as_ref()
+                    .map_or_else(|| NumWasmPages::new(0), |es| es.wasm_memory.size);
+                if let Ok(wasm_memory_usage) = num_bytes_try_from(num_wasm_pages) {
+                    canister.system_state.wasm_memory_limit =
+                        Some(compute_default_wasm_memory_limit(
+                            default_wasm_memory_limit,
+                            wasm_memory_usage,
+                        ));
+                }
+            }
+        }
+    }
+
     /// Checks the deterministic time slicing invariant after round execution.
     fn check_dts_invariants(
         &self,

rs/execution_environment/tests/execution_test.rs

@@ -2,7 +2,7 @@ use assert_matches::assert_matches;
 use candid::Encode;
 use ic_base_types::PrincipalId;
 use ic_config::{
-    execution_environment::Config as HypervisorConfig,
+    execution_environment::{Config as HypervisorConfig, DEFAULT_WASM_MEMORY_LIMIT},
     flag_status::FlagStatus,
     subnet_config::{CyclesAccountManagerConfig, SubnetConfig},
 };
@@ -12,6 +12,7 @@ use ic_management_canister_types::{
     SignWithECDSAArgs, TakeCanisterSnapshotArgs, UpdateSettingsArgs, IC_00,
 };
 use ic_registry_subnet_type::SubnetType;
+use ic_replicated_state::NumWasmPages;
 use ic_state_machine_tests::{
     ErrorCode, IngressState, IngressStatus, MessageId, StateMachine, StateMachineBuilder,
     StateMachineConfig, UserError,
@@ -2485,11 +2486,9 @@ fn set_wasm_memory_limit_to_zero() {
 
 #[test]
 fn set_wasm_memory_limit_from_another_canister() {
-    let subnet_config = SubnetConfig::new(SubnetType::Application);
-    let env = StateMachine::new_with_config(StateMachineConfig::new(
-        subnet_config,
-        HypervisorConfig::default(),
-    ));
+    let env = StateMachineBuilder::new()
+        .with_subnet_type(SubnetType::Application)
+        .build();
 
     let initial_cycles = Cycles::new(1_000_000 * B);
     let canister1 = create_universal_canister_with_cycles(&env, None, initial_cycles);
@@ -2543,3 +2542,131 @@ fn set_wasm_memory_limit_from_another_canister() {
         .unwrap();
     assert_eq!(result, WasmResult::Reply(vec![1, 2, 3]));
 }
+
+#[test]
+fn canister_create_with_default_wasm_memory_limit() {
+    let env = StateMachineBuilder::new()
+        .with_subnet_type(SubnetType::Application)
+        .build();
+
+    let initial_cycles = Cycles::new(200 * B);
+    let canister_id = create_universal_canister_with_cycles(&env, None, initial_cycles);
+
+    let wasm_memory_limit = fetch_wasm_memory_limit(&env, canister_id);
+    assert_eq!(wasm_memory_limit, DEFAULT_WASM_MEMORY_LIMIT);
+}
+
+#[test]
+fn initialize_default_wasm_memory_limit_with_low_memory_usage() {
+    let env = StateMachineBuilder::new()
+        .with_subnet_type(SubnetType::Application)
+        .build();
+
+    let initial_cycles = Cycles::new(200 * B);
+    let canister_id = create_universal_canister_with_cycles(&env, None, initial_cycles);
+
+    let wasm_memory_limit = fetch_wasm_memory_limit(&env, canister_id);
+    assert_eq!(wasm_memory_limit, DEFAULT_WASM_MEMORY_LIMIT);
+
+    // Clear the Wasm memory limit.
+    {
+        let mut state = env.get_latest_state().as_ref().clone();
+        let canister = state.canister_state_mut(&canister_id).unwrap();
+        canister.system_state.wasm_memory_limit = None;
+        env.replace_canister_state(Arc::new(state), canister_id);
+    }
+
+    env.tick();
+
+    let wasm_memory_limit = fetch_wasm_memory_limit(&env, canister_id);
+    assert_eq!(wasm_memory_limit, DEFAULT_WASM_MEMORY_LIMIT);
+}
+
+#[test]
+fn initialize_default_wasm_memory_limit_with_high_memory_usage() {
+    let env = StateMachineBuilder::new()
+        .with_subnet_type(SubnetType::Application)
+        .build();
+
+    let wat = r#"(module
+        (import "ic0" "msg_reply" (func $msg_reply))
+        (func $grow_mem
+            (drop (memory.grow (i32.const 49152)))
+            (call $msg_reply)
+        )
+        (export "canister_update grow_mem" (func $grow_mem))
+        (memory $memory 0)
+    )"#;
+
+    let initial_cycles = Cycles::new(1_000_000 * B);
+    let canister_id = env
+        .install_canister_with_cycles(
+            wat::parse_str(wat).unwrap(),
+            vec![],
+            Some(
+                CanisterSettingsArgsBuilder::new()
+                    .with_wasm_memory_limit(10_000_000_000)
+                    .build(),
+            ),
+            initial_cycles,
+        )
+        .unwrap();
+
+    let wasm_memory_limit = fetch_wasm_memory_limit(&env, canister_id);
+    assert_eq!(wasm_memory_limit, NumBytes::new(10_000_000_000));
+
+    env.execute_ingress(canister_id, "grow_mem", vec![])
+        .unwrap();
+
+    // Check Wasm memory usage and clear the Wasm memory limit.
+    {
+        let mut state = env.get_latest_state().as_ref().clone();
+        let canister = state.canister_state_mut(&canister_id).unwrap();
+        canister.system_state.wasm_memory_limit = None;
+        assert_eq!(
+            canister.execution_state.as_ref().unwrap().wasm_memory.size,
+            NumWasmPages::new(49152)
+        );
+        env.replace_canister_state(Arc::new(state), canister_id);
+    }
+
+    env.tick();
+
+    let wasm_memory_limit = fetch_wasm_memory_limit(&env, canister_id);
+    // The initialized Wasm memory should be the average of the current memory
+    // usage and the hard limit of 4GiB.
+    assert_eq!(
+        wasm_memory_limit,
+        NumBytes::new((49152 + 65536) / 2 * WASM_PAGE_SIZE_IN_BYTES)
+    );
+}
+
+#[test]
+fn do_not_initialize_wasm_memory_limit_if_it_is_not_empty() {
+    let env = StateMachineBuilder::new()
+        .with_subnet_type(SubnetType::Application)
+        .build();
+
+    let wat = "(module)";
+    let initial_cycles = Cycles::new(1_000_000 * B);
+    let canister_id = env
+        .install_canister_with_cycles(
+            wat::parse_str(wat).unwrap(),
+            vec![],
+            Some(
+                CanisterSettingsArgsBuilder::new()
+                    .with_wasm_memory_limit(10_000_000_000)
+                    .build(),
+            ),
+            initial_cycles,
+        )
+        .unwrap();
+
+    let wasm_memory_limit = fetch_wasm_memory_limit(&env, canister_id);
+    assert_eq!(wasm_memory_limit, NumBytes::new(10_000_000_000));
+
+    env.tick();
+
+    let wasm_memory_limit = fetch_wasm_memory_limit(&env, canister_id);
+    assert_eq!(wasm_memory_limit, NumBytes::new(10_000_000_000));
+}

rs/replica_tests/tests/canister_lifecycle.rs

@@ -1,5 +1,6 @@
 use assert_matches::assert_matches;
 use candid::Encode;
+use ic_config::execution_environment::DEFAULT_WASM_MEMORY_LIMIT;
 use ic_config::subnet_config::CyclesAccountManagerConfig;
 use ic_config::Config;
 use ic_error_types::{ErrorCode, RejectCode};
@@ -719,7 +720,7 @@ fn can_get_canister_information() {
                 0u128,
                 0u128,
                 0u128,
-                Some(0),
+                Some(DEFAULT_WASM_MEMORY_LIMIT.get()),
             )
         );
 
@@ -777,7 +778,7 @@ fn can_get_canister_information() {
                     0u128,
                     0u128,
                     0u128,
-                    Some(0)
+                    Some(DEFAULT_WASM_MEMORY_LIMIT.get())
                 ),
                 CanisterStatusResultV2::decode(&res).unwrap(),
                 2 * BALANCE_EPSILON,

rs/tests/src/execution/big_stable_memory.rs

@@ -128,6 +128,7 @@ pub fn can_access_big_heap_and_big_stable_memory(env: TestEnv) {
             let mgr = ManagementCanister::create(&agent);
             let canister_id = mgr
                 .create_canister()
+                .with_optional_wasm_memory_limit(Some(4 * 1024 * 1024 * 1024))
                 .as_provisional_create_with_amount(None)
                 .with_effective_canister_id(node.effective_canister_id())
                 .call_and_wait()