[timers] Fix missing wakeups in timer service

Exactly as (#5085) but for the timer service. Changes to the scheduling were not correctly waking up
the wakers for `poll_next_timer`. This was ok in the existing code because we were reconstructing the timer stream in every `LeaderState::run` call.
Now that this is changing in the next PR, this has to be fixed.

> For full disclaimer, this was caught and fixed by codex.

Author: MohamedBassem

crates/timer/src/service/mod.rs

@@ -10,17 +10,19 @@
 
 #![allow(clippy::enum_variant_names)]
 
-use pin_project::pin_project;
-use restate_types::timer::TimerKey;
 use std::collections::HashSet;
 use std::fmt::Debug;
 use std::future;
 use std::future::Future;
 use std::pin::Pin;
 use std::task::{Context, Poll, Waker, ready};
+
+use pin_project::pin_project;
 use tokio_util::sync::ReusableBoxFuture;
 use tracing::trace;
 
+use restate_types::timer::TimerKey;
+
 pub mod clock;
 #[cfg(test)]
 mod tests;
@@ -56,6 +58,7 @@ enum ProcessTimersState<TimerKey, SleepFuture> {
     ReadNextTimer,
     AwaitTimer {
         timer_key: TimerKey,
+        waker: Waker,
         #[pin]
         sleep: SleepFuture,
     },
@@ -234,12 +237,15 @@ where
                         ProcessTimersStateProj::ReadNextTimer => {
                             // nothing to do because peek timer will be read next
                         }
-                        ProcessTimersStateProj::AwaitTimer { timer_key, .. } => {
+                        ProcessTimersStateProj::AwaitTimer {
+                            timer_key, waker, ..
+                        } => {
                             // we might wait for a later timer if the newly added timer fires earlier
                             if new_timer_key < *timer_key {
                                 trace!(
                                     "Reset process timer state to ReadNextTimer because added timer fires earlier."
                                 );
+                                waker.wake_by_ref();
                                 process_timers_state.set(ProcessTimersState::ReadNextTimer);
                             }
                         }
@@ -300,8 +306,11 @@ where
                     ProcessTimersStateProj::ReadNextTimer => {
                         // nothing to do
                     }
-                    ProcessTimersStateProj::AwaitTimer { timer_key, .. } => {
+                    ProcessTimersStateProj::AwaitTimer {
+                        timer_key, waker, ..
+                    } => {
                         if key == *timer_key {
+                            waker.wake_by_ref();
                             process_timers_state.set(ProcessTimersState::ReadNextTimer);
 
                             trace!("Skip awaiting removed timer '{key:?}'. Read next timer.");
@@ -344,7 +353,8 @@ where
 
         loop {
             match state.as_mut().project() {
-                StateProj::Idle(_) => {
+                StateProj::Idle(waker) => {
+                    waker.clone_from(cx.waker());
                     return Poll::Pending;
                 }
                 StateProj::LoadTimers { removed_timers } => {
@@ -423,6 +433,7 @@ where
                                 );
                                 process_timers_state.set(ProcessTimersState::AwaitTimer {
                                     timer_key: timer_key.clone(),
+                                    waker: cx.waker().clone(),
                                     sleep,
                                 });
                             } else {
@@ -455,7 +466,8 @@ where
                             state.set(State::LoadTimers { removed_timers });
                         }
                     }
-                    ProcessTimersStateProj::AwaitTimer { sleep, .. } => {
+                    ProcessTimersStateProj::AwaitTimer { sleep, waker, .. } => {
+                        waker.clone_from(cx.waker());
                         ready!(sleep.poll(cx));
                         process_timers_state.set(ProcessTimersState::TriggerTimer);
                     }

crates/timer/src/service/tests.rs

@@ -8,22 +8,27 @@
 // the Business Source License, use of this software will be governed
 // by the Apache License, Version 2.0.
 
-use crate::service::clock::TokioClock;
-use crate::service::clock::tests::ManualClock;
-use crate::{Timer, TimerReader, TimerService};
-use futures_util::FutureExt;
-use restate_test_util::let_assert;
-use restate_types::time::MillisSinceEpoch;
-use restate_types::timer::TimerKey;
 use std::cmp::Ordering;
 use std::collections::BTreeMap;
 use std::fmt::Debug;
 use std::pin::Pin;
+use std::sync::atomic::{AtomicBool, Ordering as AtomicOrdering};
 use std::sync::{Arc, Mutex};
+use std::task::{Context, Poll, Wake, Waker};
 use std::time::{Duration, SystemTime};
+
+use futures_util::FutureExt;
 use test_log::test;
 use tokio::sync::oneshot;
 
+use restate_test_util::let_assert;
+use restate_types::time::MillisSinceEpoch;
+use restate_types::timer::TimerKey;
+
+use crate::service::clock::TokioClock;
+use crate::service::clock::tests::ManualClock;
+use crate::{Timer, TimerReader, TimerService};
+
 #[derive(Debug, Clone)]
 struct MockTimerReader<T>
 where
@@ -159,6 +164,14 @@ impl TimerKey for TimerValue {
     }
 }
 
+struct TestWaker(AtomicBool);
+
+impl Wake for TestWaker {
+    fn wake(self: Arc<Self>) {
+        self.0.store(true, AtomicOrdering::Relaxed);
+    }
+}
+
 #[test(tokio::test)]
 async fn no_timer_is_dropped() {
     let timer_reader = MockTimerReader::new();
@@ -324,6 +337,55 @@ async fn earlier_timers_replace_older_ones() {
     }
 }
 
+#[test]
+fn adding_earlier_timer_wakes_service() {
+    let clock = ManualClock::new(MillisSinceEpoch::UNIX_EPOCH);
+    let timer_reader = MockTimerReader::<TimerValue>::new();
+    timer_reader.add_timer(TimerValue::new(1, 10.into()));
+
+    let service = TimerService::new(clock, Some(1), timer_reader);
+    tokio::pin!(service);
+
+    let test_waker = Arc::new(TestWaker(AtomicBool::new(false)));
+    let waker = Waker::from(Arc::clone(&test_waker));
+    let mut cx = Context::from_waker(&waker);
+
+    assert_eq!(service.as_mut().poll_next_timer(&mut cx), Poll::Pending);
+
+    service
+        .as_mut()
+        .add_timer(TimerValue::new(0, MillisSinceEpoch::UNIX_EPOCH));
+
+    assert!(test_waker.0.load(AtomicOrdering::Relaxed));
+    assert_eq!(
+        service.as_mut().poll_next_timer(&mut cx),
+        Poll::Ready(TimerValue::new(0, MillisSinceEpoch::UNIX_EPOCH))
+    );
+}
+
+#[test]
+fn removing_awaited_timer_wakes_service() {
+    let clock = ManualClock::new(MillisSinceEpoch::UNIX_EPOCH);
+    let timer_reader = MockTimerReader::<TimerValue>::new();
+    let timer = TimerValue::new(0, 10.into());
+    timer_reader.add_timer(timer);
+
+    let service = TimerService::new(clock, Some(1), timer_reader.clone());
+    tokio::pin!(service);
+
+    let test_waker = Arc::new(TestWaker(AtomicBool::new(false)));
+    let waker = Waker::from(Arc::clone(&test_waker));
+    let mut cx = Context::from_waker(&waker);
+
+    assert_eq!(service.as_mut().poll_next_timer(&mut cx), Poll::Pending);
+
+    timer_reader.remove_timer(timer);
+    service.as_mut().remove_timer(timer);
+
+    assert!(test_waker.0.load(AtomicOrdering::Relaxed));
+    assert_eq!(service.as_mut().poll_next_timer(&mut cx), Poll::Pending);
+}
+
 async fn yield_to_timer_service<
     Timer: crate::Timer + Debug + 'static,
     Clock: crate::Clock,