feat(gstd): Add handle_reply to MessageFuture (#4046)

Cargo.toml

@@ -29,6 +29,7 @@ members = [
     "examples/async-critical",
     "examples/async-custom-entry",
     "examples/async-init",
+    "examples/async-reply-hook",
     "examples/async-signal-entry",
     "examples/async-tester",
     "examples/autoreply",
@@ -430,6 +431,7 @@ demo-async = { path = "examples/async" }
 demo-async-critical = { path = "examples/async-critical" }
 demo-async-custom-entry = { path = "examples/async-custom-entry" }
 demo-async-init = { path = "examples/async-init" }
+demo-async-reply-hook = { path = "examples/async-reply-hook" }
 demo-async-recursion = { path = "examples/async-recursion" }
 demo-async-signal-entry = { path = "examples/async-signal-entry" }
 demo-async-tester = { path = "examples/async-tester" }

examples/async-reply-hook/Cargo.toml

@@ -0,0 +1,21 @@
+[package]
+name = "demo-async-reply-hook"
+version = "0.1.0"
+authors.workspace = true
+edition.workspace = true
+license.workspace = true
+homepage.workspace = true
+repository.workspace = true
+
+[dependencies]
+gstd.workspace = true
+parity-scale-codec.workspace = true
+futures.workspace = true
+
+[build-dependencies]
+gear-wasm-builder.workspace = true
+
+[features]
+debug = ["gstd/debug"]
+default = ["std"]
+std = []

examples/async-reply-hook/build.rs

@@ -0,0 +1,21 @@
+// This file is part of Gear.
+
+// Copyright (C) 2023-2024 Gear Technologies Inc.
+// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
+
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with this program. If not, see <https://www.gnu.org/licenses/>.
+
+fn main() {
+    gear_wasm_builder::build();
+}

examples/async-reply-hook/src/lib.rs

@@ -0,0 +1,30 @@
+// This file is part of Gear.
+
+// Copyright (C) 2023-2024 Gear Technologies Inc.
+// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
+
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with this program. If not, see <https://www.gnu.org/licenses/>.
+
+#![no_std]
+
+#[cfg(feature = "std")]
+mod code {
+    include!(concat!(env!("OUT_DIR"), "/wasm_binary.rs"));
+}
+
+#[cfg(feature = "std")]
+pub use code::WASM_BINARY_OPT as WASM_BINARY;
+
+#[cfg(target_arch = "wasm32")]
+mod wasm;

examples/async-reply-hook/src/wasm.rs

@@ -0,0 +1,88 @@
+// This file is part of Gear.
+
+// Copyright (C) 2023-2024 Gear Technologies Inc.
+// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
+
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with this program. If not, see <https://www.gnu.org/licenses/>.
+
+use gstd::{critical, debug, exec, msg, prelude::*, ActorId};
+
+#[gstd::async_main]
+async fn main() {
+    let source = msg::source();
+
+    // Case 1: Message without reply hook
+    let m1 = gstd::msg::send_bytes_for_reply(source, b"for_reply_1", 0, 0)
+        .expect("Failed to send message");
+
+    // Case 2: Message with reply hook but we don't reply to it
+    let m2 = gstd::msg::send_bytes_for_reply(source, b"for_reply_2", 0, 1_000_000_000)
+        .expect("Failed to send message")
+        .up_to(Some(5))
+        .expect("Failed to set timeout")
+        .handle_reply(|| {
+            // This should be called in gas / 100 blocks, but the program exits by that time
+            unreachable!("This should not be called");
+        })
+        .expect("Failed to set reply hook");
+
+    // Case 3: Message with reply hook and we reply to it
+    let for_reply_3 = gstd::rc::Rc::new(core::cell::RefCell::new(false));
+    let for_reply_3_clone = for_reply_3.clone();
+    let m3 = gstd::msg::send_bytes_for_reply(source, b"for_reply_3", 0, 1_000_000_000)
+        .expect("Failed to send message")
+        .handle_reply(move || {
+            debug!("reply message_id: {:?}", msg::id());
+            debug!("reply payload: {:?}", msg::load_bytes());
+
+            assert_eq!(msg::load_bytes().unwrap(), [3]);
+
+            msg::send_bytes(msg::source(), b"saw_reply_3", 0);
+            for_reply_3_clone.replace(true);
+        })
+        .expect("Failed to set reply hook");
+
+    // Case 4: We reply to message after timeout
+    let m4 = gstd::msg::send_bytes_for_reply(source, b"for_reply_4", 0, 1_000_000_000)
+        .expect("Failed to send message")
+        .up_to(Some(5))
+        .expect("Failed to set timeout")
+        .handle_reply(|| {
+            debug!("reply message_id: {:?}", msg::id());
+            debug!("reply payload: {:?}", msg::load_bytes());
+
+            assert_eq!(msg::load_bytes().unwrap(), [4]);
+
+            msg::send_bytes(msg::source(), b"saw_reply_4", 0);
+        })
+        .expect("Failed to set reply hook");
+
+    m1.await.expect("Received error reply");
+
+    assert_eq!(
+        m2.await.expect_err("Should receive timeout"),
+        gstd::errors::Error::Timeout(8, 8)
+    );
+
+    m3.await.expect("Received error reply");
+    // check for_reply_3 handle_reply executed
+    assert!(for_reply_3.replace(false));
+
+    assert_eq!(
+        m4.await.expect_err("Should receive timeout"),
+        gstd::errors::Error::Timeout(8, 8)
+    );
+
+    msg::send_bytes(source, b"completed", 0);
+}

examples/distributor/src/wasm.rs

@@ -171,7 +171,7 @@ extern "C" fn handle() {
 
 #[no_mangle]
 extern "C" fn handle_reply() {
-    gstd::record_reply();
+    gstd::handle_reply_with_hook();
 }
 
 #[no_mangle]

gstd/codegen/src/lib.rs

@@ -196,7 +196,7 @@ fn generate_handle_reply_if_required(mut code: TokenStream, attr: Option<Path>)
         let handle_reply: TokenStream = quote!(
             #[no_mangle]
             extern "C" fn handle_reply() {
-                gstd::record_reply();
+                gstd::handle_reply_with_hook();
                 #attr ();
             }
         )
@@ -498,7 +498,7 @@ pub fn wait_for_reply(attr: TokenStream, item: TokenStream) -> TokenStream {
             // Registering signal.
             crate::async_runtime::signals().register_signal(waiting_reply_to);
 
-            Ok(crate::msg::MessageFuture { waiting_reply_to })
+            Ok(crate::msg::MessageFuture { waiting_reply_to, reply_deposit })
         }
 
         #[doc = #for_reply_as_docs]
@@ -514,7 +514,7 @@ pub fn wait_for_reply(attr: TokenStream, item: TokenStream) -> TokenStream {
             // Registering signal.
             crate::async_runtime::signals().register_signal(waiting_reply_to);
 
-            Ok(crate::msg::CodecMessageFuture::<D> { waiting_reply_to, _marker: Default::default() })
+            Ok(crate::msg::CodecMessageFuture::<D> { waiting_reply_to, reply_deposit, _marker: Default::default() })
         }
     }
     .into()
@@ -591,7 +591,7 @@ pub fn wait_create_program_for_reply(attr: TokenStream, item: TokenStream) -> To
             // Registering signal.
             crate::async_runtime::signals().register_signal(waiting_reply_to);
 
-            Ok(crate::msg::CreateProgramFuture { waiting_reply_to, program_id })
+            Ok(crate::msg::CreateProgramFuture { waiting_reply_to, program_id, reply_deposit })
         }
 
         #[doc = #for_reply_as_docs]
@@ -607,7 +607,7 @@ pub fn wait_create_program_for_reply(attr: TokenStream, item: TokenStream) -> To
             // Registering signal.
             crate::async_runtime::signals().register_signal(waiting_reply_to);
 
-            Ok(crate::msg::CodecCreateProgramFuture::<D> { waiting_reply_to, program_id, _marker: Default::default() })
+            Ok(crate::msg::CodecCreateProgramFuture::<D> { waiting_reply_to, program_id, reply_deposit, _marker: Default::default() })
         }
     }
     .into()

gstd/src/async_runtime/mod.rs

@@ -18,11 +18,13 @@
 
 mod futures;
 mod locks;
+mod reply_hooks;
 mod signals;
 mod waker;
 
 pub use self::futures::message_loop;
 pub(crate) use locks::Lock;
+pub(crate) use reply_hooks::HooksMap;
 pub(crate) use signals::ReplyPoll;
 
 use self::futures::FuturesMap;
@@ -49,9 +51,20 @@ pub(crate) fn locks() -> &'static mut LocksMap {
     unsafe { LOCKS.get_or_insert_with(LocksMap::default) }
 }
 
+static mut REPLY_HOOKS: Option<HooksMap> = None;
+
+pub(crate) fn reply_hooks() -> &'static mut HooksMap {
+    unsafe { REPLY_HOOKS.get_or_insert_with(HooksMap::new) }
+}
+
 /// Default reply handler.
-pub fn record_reply() {
+pub fn handle_reply_with_hook() {
     signals().record_reply();
+
+    // Execute reply hook (if it was registered)
+    let replied_to =
+        crate::msg::reply_to().expect("`gstd::handle_reply_with_hook()` called in wrong context");
+    reply_hooks().execute_and_remove(replied_to);
 }
 
 /// Default signal handler.
@@ -64,4 +77,5 @@ pub fn handle_signal() {
 
     futures().remove(&msg_id);
     locks().remove_message_entry(msg_id);
+    reply_hooks().remove(msg_id)
 }

gstd/src/async_runtime/reply_hooks.rs

@@ -0,0 +1,49 @@
+// This file is part of Gear.
+
+// Copyright (C) 2023-2024 Gear Technologies Inc.
+// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
+
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with this program. If not, see <https://www.gnu.org/licenses/>.
+
+use crate::MessageId;
+use alloc::boxed::Box;
+use hashbrown::HashMap;
+
+pub(crate) struct HooksMap(HashMap<MessageId, Box<dyn FnOnce()>>);
+
+impl HooksMap {
+    pub fn new() -> Self {
+        Self(HashMap::new())
+    }
+
+    /// Register hook to be executed when a reply for message_id is received.
+    pub(crate) fn register<F: FnOnce() + 'static>(&mut self, mid: MessageId, f: F) {
+        if self.0.contains_key(&mid) {
+            panic!("handle_reply: reply hook for this message_id is already registered");
+        }
+        self.0.insert(mid, Box::new(f));
+    }
+
+    /// Execute hook for message_id (if registered)
+    pub(crate) fn execute_and_remove(&mut self, message_id: MessageId) {
+        if let Some(f) = self.0.remove(&message_id) {
+            f();
+        }
+    }
+
+    /// Clear hook for message_id without executing it.
+    pub(crate) fn remove(&mut self, message_id: MessageId) {
+        self.0.remove(&message_id);
+    }
+}

gstd/src/common/errors.rs

@@ -172,6 +172,10 @@ pub enum UsageError {
     ZeroSystemReservationAmount,
     /// This error occurs when providing zero duration to mutex lock function
     ZeroMxLockDuration,
+    /// This error occurs when handle_reply is called without (or with zero)
+    /// reply deposit
+    /// (see [`MessageFuture::handle_reply`](crate::msg::MessageFuture::handle_reply)).
+    ZeroReplyDeposit,
 }
 
 impl fmt::Display for UsageError {
@@ -182,6 +186,9 @@ impl fmt::Display for UsageError {
                 write!(f, "System reservation amount can not be zero in config")
             }
             UsageError::ZeroMxLockDuration => write!(f, "Mutex lock duration can not be zero"),
+            UsageError::ZeroReplyDeposit => {
+                write!(f, "Reply deposit can not be zero when setting reply hook")
+            }
         }
     }
 }

gstd/src/lib.rs

@@ -159,7 +159,7 @@ mod reservations;
 pub mod sync;
 pub mod util;
 
-pub use async_runtime::{handle_signal, message_loop, record_reply};
+pub use async_runtime::{handle_reply_with_hook, handle_signal, message_loop};
 pub use common::{errors, primitives_ext::*};
 pub use config::Config;
 pub use gcore::{