feat: Low memory threshold

CHANGELOG.md

@@ -2,6 +2,12 @@
 
 # UNRELEASED
 
+### feat: `dfx canister [create|update-settings] --wasm-memory-threshold`
+
+This adds support for the WASM memory threshold, used in conjunction with `--wasm-memory-limit`.
+When the remaining memory until the limit falls below the threshold, the canister's
+`on_low_wasm_memory` handler is run.
+
 ### fix: `dfx deploy --by-proposal` no longer sends chunk data in ProposeCommitBatch
 
 Recently we made `dfx deploy` include some chunk data in CommitBatch, in order to streamline

Cargo.toml

@@ -22,11 +22,11 @@ license = "Apache-2.0"
 candid = "0.10.11"
 candid_parser = "0.1.4"
 dfx-core = { path = "src/dfx-core", version = "0.1.0" }
-ic-agent = "0.39"
+ic-agent = { version = "0.39", git = "https://github.com/dfinity/agent-rs", rev = "0a51f2a65dde7d9e1790c378bd60e1768e3be257" }
 ic-asset = { path = "src/canisters/frontend/ic-asset", version = "0.21.0" }
 ic-cdk = "0.13.1"
-ic-identity-hsm = "0.39"
-ic-utils = "0.39"
+ic-identity-hsm = { version = "0.39", git = "https://github.com/dfinity/agent-rs", rev = "0a51f2a65dde7d9e1790c378bd60e1768e3be257" }
+ic-utils = { version = "0.39", git = "https://github.com/dfinity/agent-rs", rev = "0a51f2a65dde7d9e1790c378bd60e1768e3be257" }
 
 aes-gcm = "0.10.3"
 anyhow = "1.0.56"

docs/cli-reference/dfx-canister.mdx

@@ -280,6 +280,7 @@ You can use the following options with the `dfx canister create` command.
 | `--memory-allocation <memory>`            | Specifies how much memory the canister is allowed to use in total. This should be a value in the range [0..12 GiB]. A setting of 0 means the canister will have access to memory on a “best-effort” basis: It will only be charged for the memory it uses, but at any point in time may stop running if it tries to allocate more memory when there isn’t space available on the subnet. |
 | `--reserved-cycles-limit <limit>`         | Specifies the upper limit for the canister's reserved cycles.                                                                                                                                                                                                                                                                                                                            |
 | `--wasm-memory-limit <limit>`             | Specifies a soft upper limit for the canister's heap memory.                                                                                                                                                                                                                                                                                                                             |
+| `--wasm-memory-threshold <threshold>`     | Specifies a threshold remaining amount of memory before the canister's low-memory hook runs.                                                                                                                                                                                                                                                                                             |
 | `--log-viewer <principal>`                | Specifies the principal as an allowed viewers. Can be specified more than once. Cannot be used with `--log-visibility`.                                                                                                                                                                                                                                                                  |
 | `--log-visibility <visibility>`           | Specifies who can read the canister's logs: "controllers" or "public". For custom allowed viewers, use `--log-viewer`.                                                                                                                                                                                                                                                                   |
 | `--no-wallet`                             | Performs the call with the user Identity as the Sender of messages. Bypasses the Wallet canister. Enabled by default.                                                                                                                                                                                                                                                                    |
@@ -1150,6 +1151,7 @@ You can specify the following options for the `dfx canister update-settings` com
 | `--remove-controller <principal>`         | Removes a principal from the list of controllers of the canister.                                                                                                                                                                                                                                                                                                                        |
 | `--remove-log-viewer <principal>`         | Removes a principal from the list of log viewers of the canister. Can be specified more than once to remove multiple log viewers.                                                                                                                                                                                                                                                        |
 | `--freezing-threshold <seconds>`          | Set the [freezing threshold](https://internetcomputer.org/docs/current/references/ic-interface-spec/#ic-create_canister) in seconds for a canister. This should be a value in the range [0..2^64^-1]. Very long thresholds require the `--confirm-very-long-freezing-threshold` option.                                                                                                  |
+| `--wasm-memory-threshold <threshold>`     | Specifies a threshold remaining amount of memory before the canister's low-memory hook runs.                                                                                                                                                                                                                                                                                             |
 | `-y`, `--yes`                             | Skips yes/no checks by answering 'yes'. Such checks can result in loss of control, so this is not recommended outside of CI.                                                                                                                                                                                                                                                             |
 
 ### Arguments

docs/dfx-json-schema.json

@@ -447,7 +447,8 @@
             "log_visibility": null,
             "memory_allocation": null,
             "reserved_cycles_limit": null,
-            "wasm_memory_limit": null
+            "wasm_memory_limit": null,
+            "wasm_memory_threshold": null
           },
           "allOf": [
             {
@@ -1020,6 +1021,19 @@
               "type": "null"
             }
           ]
+        },
+        "wasm_memory_threshold": {
+          "title": "Wasm Memory Threshold",
+          "description": "Specifies a threshold (in bytes) on the Wasm memory usage of the canister, as a distance from `wasm_memory_limit`.\n\nWhen the remaining memory before the limit drops below this threshold, its `on_low_wasm_memory` hook will be invoked. This enables it to self-optimize, or raise an alert, or otherwise attempt to prevent itself from reaching `wasm_memory_limit`.\n\nMust be a number of bytes between 0 and 2^48 (i.e. 256 TiB), inclusive. Can be specified as an integer, or as an SI unit string (e.g. \"4KB\", \"2 MiB\")",
+          "default": null,
+          "anyOf": [
+            {
+              "$ref": "#/definitions/Byte"
+            },
+            {
+              "type": "null"
+            }
+          ]
         }
       }
     },

e2e/assets/allocate_memory/src/e2e_project_backend/src/lib.rs

@@ -7,3 +7,8 @@ fn greet(s: String) -> String {
 fn greet_update(s: String) -> String {
     format!("Hello, {s}!")
 }
+
+#[ic_cdk::on_low_wasm_memory]
+fn on_low_wasm_memory() {
+    ic_cdk::println!("Low memory!");
+}

e2e/tests-dfx/update_settings.bash

@@ -61,6 +61,23 @@ teardown() {
   assert_contains "Canister exceeded its current Wasm memory limit of 8 bytes"
 }
 
+@test "set wasm memory threshold" {
+  dfx_new_rust
+  install_asset allocate_memory
+  dfx_start
+  assert_command dfx canister create e2e_project_backend --no-wallet --wasm-memory-threshold 2MiB --wasm-memory-limit 2MiB
+  assert_command dfx deploy e2e_project_backend
+  assert_command dfx canister status e2e_project_backend
+  assert_contains "Wasm memory threshold: 2_097_152 Bytes"
+  assert_command dfx canister update-settings e2e_project_backend --wasm-memory-threshold 1MiB
+  assert_command dfx canister status e2e_project_backend
+  assert_contains "Wasm memory threshold: 1_048_576 Bytes"
+  assert_command dfx canister call e2e_project_backend greet_update '("alice")'
+  sleep 1
+  assert_command dfx canister logs e2e_project_backend
+  assert_contains "Low memory!"
+}
+
 @test "set log visibility" {
   dfx_new
   dfx_start

src/dfx-core/src/config/model/dfinity.rs

@@ -16,11 +16,13 @@ use crate::error::dfx_config::GetRemoteCanisterIdError::GetRemoteCanisterIdFaile
 use crate::error::dfx_config::GetReservedCyclesLimitError::GetReservedCyclesLimitFailed;
 use crate::error::dfx_config::GetSpecifiedIdError::GetSpecifiedIdFailed;
 use crate::error::dfx_config::GetWasmMemoryLimitError::GetWasmMemoryLimitFailed;
+use crate::error::dfx_config::GetWasmMemoryThresholdError::GetWasmMemoryThresholdFailed;
 use crate::error::dfx_config::{
     AddDependenciesError, GetCanisterConfigError, GetCanisterNamesWithDependenciesError,
     GetComputeAllocationError, GetFreezingThresholdError, GetLogVisibilityError,
     GetMemoryAllocationError, GetPullCanistersError, GetRemoteCanisterIdError,
     GetReservedCyclesLimitError, GetSpecifiedIdError, GetWasmMemoryLimitError,
+    GetWasmMemoryThresholdError,
 };
 use crate::error::fs::CanonicalizePathError;
 use crate::error::load_dfx_config::LoadDfxConfigError;
@@ -477,6 +479,21 @@ pub struct InitializationValues {
     #[schemars(with = "Option<ByteSchema>")]
     pub wasm_memory_limit: Option<Byte>,
 
+    /// # Wasm Memory Threshold
+    ///
+    /// Specifies a threshold (in bytes) on the Wasm memory usage of the canister,
+    /// as a distance from `wasm_memory_limit`.
+    ///
+    /// When the remaining memory before the limit drops below this threshold, its
+    /// `on_low_wasm_memory` hook will be invoked. This enables it to self-optimize,
+    /// or raise an alert, or otherwise attempt to prevent itself from reaching
+    /// `wasm_memory_limit`.
+    ///
+    /// Must be a number of bytes between 0 and 2^48 (i.e. 256 TiB), inclusive.
+    /// Can be specified as an integer, or as an SI unit string (e.g. "4KB", "2 MiB")
+    #[schemars(with = "Option<ByteSchema>")]
+    pub wasm_memory_threshold: Option<Byte>,
+
     /// # Log Visibility
     /// Specifies who is allowed to read the canister's logs.
     ///
@@ -1004,6 +1021,17 @@ impl ConfigInterface {
             .wasm_memory_limit)
     }
 
+    pub fn get_wasm_memory_threshold(
+        &self,
+        canister_name: &str,
+    ) -> Result<Option<Byte>, GetWasmMemoryThresholdError> {
+        Ok(self
+            .get_canister_config(canister_name)
+            .map_err(|e| GetWasmMemoryThresholdFailed(canister_name.to_string(), e))?
+            .initialization_values
+            .wasm_memory_threshold)
+    }
+
     pub fn get_log_visibility(
         &self,
         canister_name: &str,

src/dfx-core/src/error/dfx_config.rs

@@ -58,6 +58,12 @@ pub enum GetWasmMemoryLimitError {
     GetWasmMemoryLimitFailed(String, #[source] GetCanisterConfigError),
 }
 
+#[derive(Error, Debug)]
+pub enum GetWasmMemoryThresholdError {
+    #[error("Failed to get Wasm memory threshold for canister '{0}'")]
+    GetWasmMemoryThresholdFailed(String, #[source] GetCanisterConfigError),
+}
+
 #[derive(Error, Debug)]
 pub enum GetLogVisibilityError {
     #[error("Failed to get log visibility for canister '{0}'")]