From effead29d19b6f07ff2c546dcec76671ca98e8d6 Mon Sep 17 00:00:00 2001
From: Jens Reidel <adrian@travitia.xyz>
Date: Sun, 16 Apr 2023 09:43:03 +0200
Subject: [PATCH] time: add `DelayQueue::peek` (#5569)

Signed-off-by: Jens Reidel <adrian@travitia.xyz>
---
 tokio-util/src/time/delay_queue.rs   | 39 ++++++++++++++++++++++++++++
 tokio-util/src/time/wheel/level.rs   |  4 +++
 tokio-util/src/time/wheel/mod.rs     | 10 +++++++
 tokio-util/src/time/wheel/stack.rs   |  3 +++
 tokio-util/tests/time_delay_queue.rs | 38 +++++++++++++++++++++++++++
 5 files changed, 94 insertions(+)

diff --git a/tokio-util/src/time/delay_queue.rs b/tokio-util/src/time/delay_queue.rs
index a87103941..9136d900b 100644
--- a/tokio-util/src/time/delay_queue.rs
+++ b/tokio-util/src/time/delay_queue.rs
@@ -874,6 +874,41 @@ impl<T> DelayQueue<T> {
         self.slab.compact();
     }
 
+    /// Gets the [`Key`] that [`poll_expired`] will pull out of the queue next, without
+    /// pulling it out or waiting for the deadline to expire.
+    ///
+    /// Entries that have already expired may be returned in any order, but it is
+    /// guaranteed that this method returns them in the same order as when items
+    /// are popped from the `DelayQueue`.
+    ///
+    /// # Examples
+    ///
+    /// Basic usage
+    ///
+    /// ```rust
+    /// use tokio_util::time::DelayQueue;
+    /// use std::time::Duration;
+    ///
+    /// # #[tokio::main]
+    /// # async fn main() {
+    /// let mut delay_queue = DelayQueue::new();
+    ///
+    /// let key1 = delay_queue.insert("foo", Duration::from_secs(10));
+    /// let key2 = delay_queue.insert("bar", Duration::from_secs(5));
+    /// let key3 = delay_queue.insert("baz", Duration::from_secs(15));
+    ///
+    /// assert_eq!(delay_queue.peek().unwrap(), key2);
+    /// # }
+    /// ```
+    ///
+    /// [`Key`]: struct@Key
+    /// [`poll_expired`]: method@Self::poll_expired
+    pub fn peek(&self) -> Option<Key> {
+        use self::wheel::Stack;
+
+        self.expired.peek().or_else(|| self.wheel.peek())
+    }
+
     /// Returns the next time to poll as determined by the wheel
     fn next_deadline(&mut self) -> Option<Instant> {
         self.wheel
@@ -1166,6 +1201,10 @@ impl<T> wheel::Stack for Stack<T> {
         }
     }
 
+    fn peek(&self) -> Option<Self::Owned> {
+        self.head
+    }
+
     #[track_caller]
     fn remove(&mut self, item: &Self::Borrowed, store: &mut Self::Store) {
         let key = *item;
diff --git a/tokio-util/src/time/wheel/level.rs b/tokio-util/src/time/wheel/level.rs
index 8ea30af30..9cdcb7336 100644
--- a/tokio-util/src/time/wheel/level.rs
+++ b/tokio-util/src/time/wheel/level.rs
@@ -206,6 +206,10 @@ impl<T: Stack> Level<T> {
 
         ret
     }
+
+    pub(crate) fn peek_entry_slot(&self, slot: usize) -> Option<T::Owned> {
+        self.slot[slot].peek()
+    }
 }
 
 impl<T> fmt::Debug for Level<T> {
diff --git a/tokio-util/src/time/wheel/mod.rs b/tokio-util/src/time/wheel/mod.rs
index ffa05ab71..4058b615a 100644
--- a/tokio-util/src/time/wheel/mod.rs
+++ b/tokio-util/src/time/wheel/mod.rs
@@ -139,6 +139,12 @@ where
         self.next_expiration().map(|expiration| expiration.deadline)
     }
 
+    /// Next key that will expire
+    pub(crate) fn peek(&self) -> Option<T::Owned> {
+        self.next_expiration()
+            .and_then(|expiration| self.peek_entry(&expiration))
+    }
+
     /// Advances the timer up to the instant represented by `now`.
     pub(crate) fn poll(&mut self, now: u64, store: &mut T::Store) -> Option<T::Owned> {
         loop {
@@ -244,6 +250,10 @@ where
         self.levels[expiration.level].pop_entry_slot(expiration.slot, store)
     }
 
+    fn peek_entry(&self, expiration: &Expiration) -> Option<T::Owned> {
+        self.levels[expiration.level].peek_entry_slot(expiration.slot)
+    }
+
     fn level_for(&self, when: u64) -> usize {
         level_for(self.elapsed, when)
     }
diff --git a/tokio-util/src/time/wheel/stack.rs b/tokio-util/src/time/wheel/stack.rs
index c87adcafd..7d32f27d2 100644
--- a/tokio-util/src/time/wheel/stack.rs
+++ b/tokio-util/src/time/wheel/stack.rs
@@ -22,6 +22,9 @@ pub(crate) trait Stack: Default {
     /// Pop an item from the stack
     fn pop(&mut self, store: &mut Self::Store) -> Option<Self::Owned>;
 
+    /// Peek into the stack.
+    fn peek(&self) -> Option<Self::Owned>;
+
     fn remove(&mut self, item: &Self::Borrowed, store: &mut Self::Store);
 
     fn when(item: &Self::Borrowed, store: &Self::Store) -> u64;
diff --git a/tokio-util/tests/time_delay_queue.rs b/tokio-util/tests/time_delay_queue.rs
index c8b74d1ec..d8773a3d6 100644
--- a/tokio-util/tests/time_delay_queue.rs
+++ b/tokio-util/tests/time_delay_queue.rs
@@ -823,6 +823,44 @@ async fn remove_after_compact_poll() {
     assert!(panic.is_err());
 }
 
+#[tokio::test(start_paused = true)]
+async fn peek() {
+    let mut queue = task::spawn(DelayQueue::new());
+
+    let now = Instant::now();
+
+    let key = queue.insert_at("foo", now + ms(5));
+    let key2 = queue.insert_at("bar", now);
+    let key3 = queue.insert_at("baz", now + ms(10));
+
+    assert_eq!(queue.peek(), Some(key2));
+
+    sleep(ms(6)).await;
+
+    assert_eq!(queue.peek(), Some(key2));
+
+    let entry = assert_ready_some!(poll!(queue));
+    assert_eq!(entry.get_ref(), &"bar");
+
+    assert_eq!(queue.peek(), Some(key));
+
+    let entry = assert_ready_some!(poll!(queue));
+    assert_eq!(entry.get_ref(), &"foo");
+
+    assert_eq!(queue.peek(), Some(key3));
+
+    assert_pending!(poll!(queue));
+
+    sleep(ms(5)).await;
+
+    assert_eq!(queue.peek(), Some(key3));
+
+    let entry = assert_ready_some!(poll!(queue));
+    assert_eq!(entry.get_ref(), &"baz");
+
+    assert!(queue.peek().is_none());
+}
+
 fn ms(n: u64) -> Duration {
     Duration::from_millis(n)
 }