From 3792fd9b53f29dc7b47f7bf7a53cc13f72e5fffd Mon Sep 17 00:00:00 2001 From: Adam Spofford Date: Fri, 1 Nov 2024 09:13:11 -0700 Subject: [PATCH 1/2] Implement low wasm memory hook --- src/ic-cdk-macros/src/export.rs | 24 +++++++++++++++++------- src/ic-cdk-macros/src/lib.rs | 10 ++++++++++ src/ic-cdk/src/macros.rs | 21 +++++++++++++++++++++ src/ic-cdk/tests/pass/blank_methods.rs | 7 ++++++- 4 files changed, 54 insertions(+), 8 deletions(-) diff --git a/src/ic-cdk-macros/src/export.rs b/src/ic-cdk-macros/src/export.rs index 6d087d398..320c94d77 100644 --- a/src/ic-cdk-macros/src/export.rs +++ b/src/ic-cdk-macros/src/export.rs @@ -33,18 +33,20 @@ enum MethodType { Query, Heartbeat, InspectMessage, + OnLowWasmMemory, } impl MethodType { pub fn is_lifecycle(&self) -> bool { - matches!( - self, + match self { MethodType::Init - | MethodType::PreUpgrade - | MethodType::PostUpgrade - | MethodType::Heartbeat - | MethodType::InspectMessage - ) + | MethodType::PreUpgrade + | MethodType::PostUpgrade + | MethodType::Heartbeat + | MethodType::InspectMessage + | MethodType::OnLowWasmMemory => true, + MethodType::Update | MethodType::Query => false, + } } } @@ -58,6 +60,7 @@ impl std::fmt::Display for MethodType { MethodType::Update => f.write_str("update"), MethodType::Heartbeat => f.write_str("heartbeat"), MethodType::InspectMessage => f.write_str("inspect_message"), + MethodType::OnLowWasmMemory => f.write_str("on_low_wasm_memory"), } } } @@ -309,6 +312,13 @@ pub(crate) fn ic_inspect_message( dfn_macro(MethodType::InspectMessage, attr, item) } +pub(crate) fn ic_on_low_wasm_memory( + attr: TokenStream, + item: TokenStream, +) -> Result { + dfn_macro(MethodType::OnLowWasmMemory, attr, item) +} + #[cfg(test)] mod test { use super::*; diff --git a/src/ic-cdk-macros/src/lib.rs b/src/ic-cdk-macros/src/lib.rs index d9e75e690..00a0d566d 100644 --- a/src/ic-cdk-macros/src/lib.rs +++ b/src/ic-cdk-macros/src/lib.rs @@ -91,3 +91,13 @@ pub fn heartbeat(attr: TokenStream, item: TokenStream) -> TokenStream { pub fn inspect_message(attr: TokenStream, item: TokenStream) -> TokenStream { handle_debug_and_errors(export::ic_inspect_message, "ic_inspect_message", attr, item) } + +#[proc_macro_attribute] +pub fn on_low_wasm_memory(attr: TokenStream, item: TokenStream) -> TokenStream { + handle_debug_and_errors( + export::ic_on_low_wasm_memory, + "ic_on_low_wasm_memory", + attr, + item, + ) +} diff --git a/src/ic-cdk/src/macros.rs b/src/ic-cdk/src/macros.rs index 8c3162875..e2f88c5b9 100644 --- a/src/ic-cdk/src/macros.rs +++ b/src/ic-cdk/src/macros.rs @@ -313,3 +313,24 @@ pub use ic_cdk_macros::heartbeat; /// } /// ``` pub use ic_cdk_macros::inspect_message; + +/// Register the `canister_on_low_wasm_memory` entry point of a canister. +/// +/// This attribute macro will export the function `canister_on_low_wasm_memory` +/// in the canister module. +/// +/// The function under this attribute must have no return value. +/// +/// Each canister can only have one `canister_on_low_wasm_memory` entry point. +/// +/// # Example +/// +/// ```rust +/// # use ic_cdk::on_low_wasm_memory; +/// #[on_low_wasm_memory] +/// fn low_memory_handler() { +/// // ... +/// # unimplemented!() +/// } +/// ``` +pub use ic_cdk_macros::on_low_wasm_memory; diff --git a/src/ic-cdk/tests/pass/blank_methods.rs b/src/ic-cdk/tests/pass/blank_methods.rs index bd3922901..32e4440e6 100644 --- a/src/ic-cdk/tests/pass/blank_methods.rs +++ b/src/ic-cdk/tests/pass/blank_methods.rs @@ -1,4 +1,6 @@ -use ic_cdk::{heartbeat, init, inspect_message, post_upgrade, pre_upgrade, query, update}; +use ic_cdk::{ + heartbeat, init, inspect_message, on_low_wasm_memory, post_upgrade, pre_upgrade, query, update, +}; #[init] fn init() {} @@ -30,4 +32,7 @@ fn heartbeat() {} #[inspect_message] fn inspect_message() {} +#[on_low_wasm_memory] +fn on_low_wasm_memory() {} + fn main() {} From 7e5a8b05f128e636b41b560edacc286dd6e673d2 Mon Sep 17 00:00:00 2001 From: Adam Spofford Date: Fri, 1 Nov 2024 09:43:14 -0700 Subject: [PATCH 2/2] readme/changelog --- src/ic-cdk/CHANGELOG.md | 1 + src/ic-cdk/README.md | 1 + 2 files changed, 2 insertions(+) diff --git a/src/ic-cdk/CHANGELOG.md b/src/ic-cdk/CHANGELOG.md index 2c58efa9d..55d130c1d 100644 --- a/src/ic-cdk/CHANGELOG.md +++ b/src/ic-cdk/CHANGELOG.md @@ -12,6 +12,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added +- Attribute `#[on_low_wasm_memory]` for low-memory hook. (#528) - Support Threshold Schnorr signing management canister API. (#518) ## [0.16.0] - 2024-08-27 diff --git a/src/ic-cdk/README.md b/src/ic-cdk/README.md index aaa8b802d..80d7ffcf2 100644 --- a/src/ic-cdk/README.md +++ b/src/ic-cdk/README.md @@ -58,6 +58,7 @@ These macros are directly related to the [Internet Computer Specification](https * [`post_upgrade`](https://docs.rs/ic-cdk/latest/ic_cdk/attr.post_upgrade.html) * [`inspect_message`](https://docs.rs/ic-cdk/latest/ic_cdk/attr.inspect_message.html) * [`heartbeat`](https://docs.rs/ic-cdk/latest/ic_cdk/attr.heartbeat.html) +* [`on_low_wasm_memory`](https://docs.rs/ic-cdk/latest/ic_cdk/attr.on_low_wasm_memory.html) * [`update`](https://docs.rs/ic-cdk/latest/ic_cdk/attr.update.html) * [`query`](https://docs.rs/ic-cdk/latest/ic_cdk/attr.query.html)