From 5fb07f58d6e088ac8030d27f8c897634e49d3eb5 Mon Sep 17 00:00:00 2001
From: Linwei Shang <linwei.shang@dfinity.org>
Date: Mon, 13 Jan 2025 14:05:51 -0500
Subject: [PATCH] feat: safe binding for msg_deadline (#545)

* feat: safe binding for msg_deadline

* Option<NonZeroU64> for msg_deadline
---
 e2e-tests/src/bin/api.rs | 40 +++++++++++++++++++++++++++++++++++++++-
 e2e-tests/tests/api.rs   |  7 +++++++
 ic-cdk/src/api.rs        | 30 +++++++++++++++++++++++++++++-
 3 files changed, 75 insertions(+), 2 deletions(-)

diff --git a/e2e-tests/src/bin/api.rs b/e2e-tests/src/bin/api.rs
index 9abe2f03..b5624069 100644
--- a/e2e-tests/src/bin/api.rs
+++ b/e2e-tests/src/bin/api.rs
@@ -1,5 +1,8 @@
+//! # NOTE
+//! The [`inspect_message`] function defined below mandates that all the update/query entrypoints must start with "call_".
+
 use candid::Principal;
-use ic_cdk::api::*;
+use ic_cdk::{api::*, call::ConfigurableCall};
 
 #[export_name = "canister_update call_msg_arg_data"]
 fn call_msg_arg_data() {
@@ -13,6 +16,41 @@ fn call_msg_caller() {
     msg_reply(vec![]);
 }
 
+/// This entrypoint will call [`call_msg_deadline`] with both best-effort and guaranteed responses.
+#[ic_cdk::update]
+async fn call_msg_deadline_caller() {
+    use ic_cdk::call::{Call, SendableCall};
+    // Call with best-effort responses.
+    let reply1 = Call::new(canister_self(), "call_msg_deadline")
+        .call_raw()
+        .await
+        .unwrap();
+    assert_eq!(reply1, vec![1]);
+    // Call with guaranteed responses.
+    let reply1 = Call::new(canister_self(), "call_msg_deadline")
+        .with_guaranteed_response()
+        .call_raw()
+        .await
+        .unwrap();
+    assert_eq!(reply1, vec![0]);
+}
+
+/// This entrypoint is to be called by [`call_msg_deadline_caller`].
+/// If the call was made with best-effort responses, `msg_deadline` should be `Some`, then return 1.
+/// If the call was made with guaranteed responses, `msg_deadline` should be `None`, then return 0.
+#[export_name = "canister_update call_msg_deadline"]
+fn call_msg_deadline() {
+    let reply = match msg_deadline() {
+        Some(v) => {
+            // `NonZeroU64::get()` converts the value to `u64`.
+            assert!(v.get() > 1);
+            1
+        }
+        None => 0,
+    };
+    msg_reply(vec![reply]);
+}
+
 #[export_name = "canister_update call_msg_reply"]
 fn call_msg_reply() {
     msg_reply(vec![42]);
diff --git a/e2e-tests/tests/api.rs b/e2e-tests/tests/api.rs
index 5b3e4af1..eb4b42d5 100644
--- a/e2e-tests/tests/api.rs
+++ b/e2e-tests/tests/api.rs
@@ -20,6 +20,13 @@ fn call_api() {
         .update_call(canister_id, sender, "call_msg_caller", vec![])
         .unwrap();
     assert_eq!(res, WasmResult::Reply(vec![]));
+    let res = pic
+        .update_call(canister_id, sender, "call_msg_deadline_caller", vec![])
+        .unwrap();
+    // Unlike the other entry points, `call_msg_dealine_caller` was implemented with the `#[update]` macro.
+    // So it returns the bytes of the Candid value `()` which is not the vec![]`.
+    // The assertion below is to check if the call was successful.
+    assert!(matches!(res, WasmResult::Reply(_)));
     // `msg_reject_code` and `msg_reject_msg` can't be tested here.
     // They are invoked in the reply/reject callback of inter-canister calls.
     // So the `call.rs` test covers them.
diff --git a/ic-cdk/src/api.rs b/ic-cdk/src/api.rs
index 961452c4..b5ed7e27 100644
--- a/ic-cdk/src/api.rs
+++ b/ic-cdk/src/api.rs
@@ -19,7 +19,7 @@
 //! For example, [`msg_arg_data`] wraps both `ic0::msg_arg_data_size` and `ic0::msg_arg_data_copy`.
 
 use candid::Principal;
-use std::convert::TryFrom;
+use std::{convert::TryFrom, num::NonZeroU64};
 
 pub mod call;
 pub mod management_canister;
@@ -83,6 +83,34 @@ pub fn msg_reject_msg() -> String {
     String::from_utf8_lossy(&bytes).into_owned()
 }
 
+/// Gets the deadline, in nanoseconds since 1970-01-01, after which the caller might stop waiting for a response.
+///
+/// For calls to update methods with best-effort responses and their callbacks,
+/// the deadline is computed based on the time the call was made,
+/// and the `timeout_seconds` parameter provided by the caller.
+/// In such cases, the deadline value will be converted to `NonZeroU64` and wrapped in `Some`.
+/// To get the deadline value as a `u64`, call `get()` on the `NonZeroU64` value.
+///
+/// ```rust,no_run
+/// use ic_cdk::api::msg_deadline;
+/// if let Some(deadline) = msg_deadline() {
+///     let deadline_value : u64 = deadline.get();
+/// }
+/// ```
+///
+/// For other calls (ingress messages and all calls to query and composite query methods,
+/// including calls in replicated mode), a `None` is returned.
+/// Please note that the raw `msg_deadline` system API returns 0 in such cases.
+/// This function is a wrapper around the raw system API that provides more semantic information through the return type.
+pub fn msg_deadline() -> Option<NonZeroU64> {
+    // SAFETY: ic0.msg_deadline is always safe to call.
+    let nano_seconds = unsafe { ic0::msg_deadline() };
+    match nano_seconds {
+        0 => None,
+        _ => Some(NonZeroU64::new(nano_seconds).unwrap()),
+    }
+}
+
 /// Replies to the sender with the data.
 pub fn msg_reply<T: AsRef<[u8]>>(data: T) {
     let buf = data.as_ref();