diff --git a/ntex-util/src/services/keepalive.rs b/ntex-util/src/services/keepalive.rs index bb74af381..e6a49bbca 100644 --- a/ntex-util/src/services/keepalive.rs +++ b/ntex-util/src/services/keepalive.rs @@ -119,21 +119,26 @@ where if expire <= now() { Err((self.f)()) } else { - ctx.poll_once(|cx| match self.sleep.poll_elapsed(cx) { - Poll::Ready(()) => { - let now = now(); - let expire = self.expire.get() + self.dur; - if expire <= now { - Err((self.f)()) - } else { - let expire = expire - now; - self.sleep - .reset(Millis(expire.as_millis().try_into().unwrap_or(u32::MAX))); - let _ = self.sleep.poll_elapsed(cx); - Ok(()) + ctx.poll_once(|cx| { + loop { + match self.sleep.poll_elapsed(cx) { + Poll::Ready(()) => { + let now = now(); + let expire = self.expire.get() + self.dur; + if expire <= now { + return Err((self.f)()); + } + let expire = expire - now; + + // sleep must be reset to non zero duration, + // otherwise it stays in elapsed state and waker + // never gets registered + let expire: u32 = expire.as_millis().try_into().unwrap_or(u32::MAX); + self.sleep.reset(Millis(expire.max(1))); + } + Poll::Pending => return Ok(()), } } - Poll::Pending => Ok(()), }) } } @@ -202,4 +207,21 @@ mod tests { assert_eq!(res, Ok(())); assert_eq!(svc.ready().await, Err(TestErr)); } + + #[ntex::test] + async fn test_ka_sub_millis() { + let svc = std::rc::Rc::new(KeepAliveService::new(Millis(100), || TestErr)); + + // less than millisecond is left before expiration + svc.expire + .set(now().checked_sub(svc.dur).unwrap() + time::Duration::from_micros(500)); + svc.sleep.elapse(); + + let p = Pipeline::::with((), svc.clone()).bind(); + assert_eq!(p.ready().await, Ok(())); + + // timer has to be re-armed, otherwise waker never gets registered + // and service readiness never resolves + assert!(!svc.sleep.is_elapsed()); + } }