fix: reply after the timer body (#679)

* Move reply after task await

* Add test and fix

* ensure test fails on main

* improve the test

---------

Co-authored-by: Linwei Shang <[email protected]>

Author: adamspofford-dfinity

e2e-tests/src/bin/timers.rs

@@ -81,6 +81,17 @@ fn start_repeating() {
     REPEATING.with(|repeating| repeating.set(id));
 }
 
+#[update]
+fn start_repeating_async() {
+    let id = set_timer_interval(Duration::from_secs(1), async || {
+        Call::bounded_wait(canister_self(), "add_event_method")
+            .with_arg("repeat")
+            .await
+            .unwrap();
+    });
+    REPEATING.with(|repeating| repeating.set(id));
+}
+
 #[update]
 fn set_self_cancelling_periodic_timer() {
     let id = set_timer_interval(Duration::from_secs(0), async || {

e2e-tests/tests/timers.rs

@@ -4,6 +4,15 @@ use std::time::Duration;
 mod test_utilities;
 use test_utilities::{cargo_build_canister, pic_base, update};
 
+/// Advance the time by a given number of seconds.
+/// "tick" the IC every second so that it progress and produce a block so that timers get triggered.
+fn advance_seconds(pic: &PocketIc, seconds: u32) {
+    for _ in 0..seconds {
+        pic.advance_time(Duration::from_secs(1));
+        pic.tick();
+    }
+}
+
 #[test]
 fn test_timers() {
     let wasm = cargo_build_canister("timers");
@@ -86,13 +95,6 @@ fn test_scheduling_many_timers() {
     assert_eq!(timers_to_schedule, executed_timers);
 }
 
-fn advance_seconds(pic: &PocketIc, seconds: u32) {
-    for _ in 0..seconds {
-        pic.advance_time(Duration::from_secs(1));
-        pic.tick();
-    }
-}
-
 #[test]
 fn test_set_global_timers() {
     let wasm = cargo_build_canister("timers");
@@ -153,3 +155,25 @@ fn test_async_timers() {
         3
     );
 }
+
+#[test]
+fn test_periodic_timers_repeat_when_task_make_calls() {
+    let wasm = cargo_build_canister("timers");
+    let pic = pic_base().build();
+
+    let canister_id = pic.create_canister();
+    pic.add_cycles(canister_id, 2_000_000_000_000);
+    pic.install_canister(canister_id, wasm, vec![], None);
+
+    update::<(), ()>(&pic, canister_id, "start_repeating_async", ()).unwrap();
+    // start_repeating_async sets a repeating timer with a 1s interval
+    // advance time by 3 seconds should trigger 3 repeats
+    advance_seconds(&pic, 3);
+    update::<(), ()>(&pic, canister_id, "stop_repeating", ()).unwrap();
+
+    let (events,): (Vec<String>,) = query_candid(&pic, canister_id, "get_events", ()).unwrap();
+    assert_eq!(
+        events[..],
+        ["method repeat", "method repeat", "method repeat"]
+    );
+}

ic-cdk-timers/src/lib.rs

@@ -255,9 +255,9 @@ unsafe extern "C" fn timer_scope_callback(env: usize) {
                 return;
             }
             x => {
-                let reject_data_size = ic0::msg_arg_data_size();
+                let reject_data_size = ic0::msg_reject_msg_size();
                 let mut reject_data = Vec::with_capacity(reject_data_size);
-                ic0::msg_arg_data_copy_uninit(
+                ic0::msg_reject_msg_copy_uninit(
                     &mut reject_data.spare_capacity_mut()[..reject_data_size],
                     0,
                 );
@@ -475,7 +475,10 @@ extern "C" fn timer_executor() {
         if let Some(task) = task {
             match task {
                 Task::Once(fut) => {
-                    ic_cdk_executor::spawn_protected(fut);
+                    ic_cdk_executor::spawn_protected(async {
+                        fut.await;
+                        ic0::msg_reply();
+                    });
                     TASKS.with_borrow_mut(|tasks| tasks.remove(task_id));
                 }
                 Task::Repeated { func, interval } => {
@@ -496,11 +499,12 @@ extern "C" fn timer_executor() {
 
                         let mut guard = RepeatGuard(Some(func), task_id, interval);
                         guard.0.as_mut().unwrap().call_mut().await;
+                        ic0::msg_reply();
                     });
                 }
             }
+        } else {
+            ic0::msg_reply();
         }
-        ic0::msg_reply_data_append(&[]);
-        ic0::msg_reply();
     });
 }