time: return the earliest key from DelayQueue::peek (#8402)

tokio-util/src/time/delay_queue.rs

@@ -951,7 +951,7 @@ impl<T> DelayQueue<T> {
     pub fn peek(&self) -> Option<Key> {
         use self::wheel::Stack;
 
-        self.expired.peek().or_else(|| self.wheel.peek())
+        self.expired.peek().or_else(|| self.wheel.peek(&self.slab))
     }
 
     /// Returns the next time to poll as determined by the wheel.
@@ -1260,6 +1260,27 @@ impl<T> wheel::Stack for Stack<T> {
         self.head
     }
 
+    fn peek_earliest(&self, store: &Self::Store) -> Option<Self::Owned> {
+        let head = self.head?;
+        let mut earliest = (head, store[head].when);
+        let mut curr = store[head].next;
+
+        while let Some(key) = curr {
+            let data = &store[key];
+
+            // The comparison is strict so that the first entry seen wins a tie,
+            // which agrees with `pop` when every entry in the slot shares a
+            // deadline.
+            if data.when < earliest.1 {
+                earliest = (key, data.when);
+            }
+
+            curr = data.next;
+        }
+
+        Some(earliest.0)
+    }
+
     #[track_caller]
     fn remove(&mut self, item: &Self::Borrowed, store: &mut Self::Store) {
         let key = *item;

tokio-util/src/time/wheel/level.rs

@@ -148,8 +148,8 @@ impl<T: Stack> Level<T> {
         ret
     }
 
-    pub(crate) fn peek_entry_slot(&self, slot: usize) -> Option<T::Owned> {
+    pub(crate) fn peek_entry_slot(&self, slot: usize, store: &T::Store) -> Option<T::Owned> {
-        self.slot[slot].peek()
+        self.slot[slot].peek_earliest(store)
     }
 }
 

tokio-util/src/time/wheel/mod.rs

@@ -140,9 +140,9 @@ where
     }
 
     /// Next key that will expire
-    pub(crate) fn peek(&self) -> Option<T::Owned> {
+    pub(crate) fn peek(&self, store: &T::Store) -> Option<T::Owned> {
         self.next_expiration()
-            .and_then(|expiration| self.peek_entry(&expiration))
+            .and_then(|expiration| self.peek_entry(&expiration, store))
     }
 
     /// Advances the timer up to the instant represented by `now`.
@@ -250,8 +250,8 @@ where
         self.levels[expiration.level].pop_entry_slot(expiration.slot, store)
     }
 
-    fn peek_entry(&self, expiration: &Expiration) -> Option<T::Owned> {
+    fn peek_entry(&self, expiration: &Expiration, store: &T::Store) -> Option<T::Owned> {
-        self.levels[expiration.level].peek_entry_slot(expiration.slot)
+        self.levels[expiration.level].peek_entry_slot(expiration.slot, store)
     }
 
     fn level_for(&self, when: u64) -> usize {

tokio-util/src/time/wheel/stack.rs

@@ -25,6 +25,13 @@ pub(crate) trait Stack: Default {
     /// Peek into the stack.
     fn peek(&self) -> Option<Self::Owned>;
 
+    /// Peek at the item in the stack with the earliest deadline.
+    ///
+    /// Unlike `peek`, this does not have to agree with `pop`: a slot in a level
+    /// above zero spans a range of deadlines, so its entries are only ordered
+    /// once they cascade down.
+    fn peek_earliest(&self, store: &Self::Store) -> Option<Self::Owned>;
+
     fn remove(&mut self, item: &Self::Borrowed, store: &mut Self::Store);
 
     fn when(item: &Self::Borrowed, store: &Self::Store) -> u64;

tokio-util/tests/time_delay_queue.rs

@@ -902,6 +902,41 @@ async fn peek() {
     assert!(queue.peek().is_none());
 }
 
+#[tokio::test(start_paused = true)]
+async fn peek_entries_sharing_a_wheel_slot() {
+    // Only the lowest level of the timer wheel has one deadline per slot. A
+    // level-one slot spans 64ms, so entries with different deadlines share it
+    // and are only ordered once the slot cascades down.
+    let mut queue = task::spawn(DelayQueue::new());
+
+    let now = Instant::now();
+
+    // 64ms..=127ms all fall in the same level-one slot. Entries are pushed onto
+    // the front of a slot, so inserting the later deadline second puts it at the
+    // head.
+    let early = queue.insert_at("early", now + ms(100));
+    let late = queue.insert_at("late", now + ms(120));
+
+    assert_eq!(queue.peek(), Some(early));
+
+    sleep(ms(105)).await;
+
+    assert_eq!(queue.peek(), Some(early));
+
+    let entry = assert_ready_some!(poll!(queue));
+    assert_eq!(entry.key(), early);
+    assert_eq!(entry.get_ref(), &"early");
+
+    assert_eq!(queue.peek(), Some(late));
+
+    sleep(ms(20)).await;
+
+    let entry = assert_ready_some!(poll!(queue));
+    assert_eq!(entry.key(), late);
+
+    assert!(queue.peek().is_none());
+}
+
 #[tokio::test(start_paused = true)]
 async fn wake_after_remove_last() {
     let mut queue = task::spawn(DelayQueue::new());