feat: Implement low wasm memory hook (#528)

Author: adamspofford-dfinity

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<TokenStream, Error> {
+    dfn_macro(MethodType::OnLowWasmMemory, attr, item)
+}
+
 #[cfg(test)]
 mod test {
     use super::*;

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,
+    )
+}

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

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)
 

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;

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() {}