time: update deadline on removal in DelayQueue (#4178)

tokio-util/src/time/delay_queue.rs

@@ -498,9 +498,20 @@ impl<T> DelayQueue<T> {
     /// # }
     /// ```
     pub fn remove(&mut self, key: &Key) -> Expired<T> {
+        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,

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)
 }