diff --git a/tokio-util/src/time/delay_queue.rs b/tokio-util/src/time/delay_queue.rs index e9ab25a1679..6da83a574d0 100644 --- a/tokio-util/src/time/delay_queue.rs +++ b/tokio-util/src/time/delay_queue.rs @@ -498,9 +498,20 @@ impl DelayQueue { /// # } /// ``` pub fn remove(&mut self, key: &Key) -> Expired { + let prev_deadline = self.next_deadline(); + self.remove_key(key); let data = self.slab.remove(key.index); + let next_deadline = self.next_deadline(); + if prev_deadline != next_deadline { + match (next_deadline, &mut self.delay) { + (None, _) => self.delay = None, + (Some(deadline), Some(delay)) => delay.as_mut().reset(deadline), + (Some(deadline), None) => self.delay = Some(Box::pin(sleep_until(deadline))), + } + } + Expired { key: Key::new(key.index), data: data.inner, diff --git a/tokio-util/tests/time_delay_queue.rs b/tokio-util/tests/time_delay_queue.rs index 2758a22d5f2..901afcaaa94 100644 --- a/tokio-util/tests/time_delay_queue.rs +++ b/tokio-util/tests/time_delay_queue.rs @@ -630,6 +630,16 @@ async fn insert_in_past_after_poll_fires_immediately() { assert_eq!(entry, "bar"); } +#[tokio::test] +async fn delay_queue_poll_expired_when_empty() { + let mut delay_queue = task::spawn(DelayQueue::new()); + let key = delay_queue.insert(0, std::time::Duration::from_secs(10)); + assert_pending!(poll!(delay_queue)); + + delay_queue.remove(&key); + assert!(assert_ready!(poll!(delay_queue)).is_none()); +} + fn ms(n: u64) -> Duration { Duration::from_millis(n) }