diff --git a/pingora-timeout/src/timer.rs b/pingora-timeout/src/timer.rs index c6c587e0d..9281fc197 100644 --- a/pingora-timeout/src/timer.rs +++ b/pingora-timeout/src/timer.rs @@ -72,10 +72,30 @@ pub struct TimerStub(Arc, Arc); 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(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"); + } }