timeout: prevent lost wakeup in TimerStub::poll

pingora-timeout/src/timer.rs

@@ -72,10 +72,30 @@ pub struct TimerStub(Arc<Notify>, Arc<AtomicBool>);
 impl TimerStub {
     /// Wait for the timer to expire.
     pub async fn poll(self) {
+        self.poll_impl(|| {}).await;
+    }
+
+    // The hook makes the check-to-registration window deterministic in tests.
+    async fn poll_impl<F>(self, before_waiter: F)
+    where
+        F: FnOnce(),
+    {
+        if self.1.load(Ordering::SeqCst) {
+            return;
+        }
+
+        // Register the waiter before checking the flag again. `notify_waiters`
+        // does not retain a permit for waiters created after the notification.
+        before_waiter();
+        let notified = self.0.notified();
+        tokio::pin!(notified);
+        notified.as_mut().enable();
+
         if self.1.load(Ordering::SeqCst) {
             return;
         }
-        self.0.notified().await;
+
+        notified.await;
     }
 }
 
@@ -324,4 +344,16 @@ mod tests {
         t1.poll().await;
         assert_eq!(now.elapsed().as_secs(), 2);
     }
+
+    #[tokio::test]
+    async fn test_timer_stub_poll_does_not_miss_fire() {
+        let timer = Timer::new();
+        let stub = timer.subscribe();
+
+        // Fire after the first flag check but before the waiter is registered.
+        // The second check must observe the fire instead of sleeping forever.
+        tokio::time::timeout(Duration::from_secs(1), stub.poll_impl(|| timer.fire()))
+            .await
+            .expect("TimerStub::poll missed a fire");
+    }
 }