From bb66b25614e1e592974039cd36b3dbdc63b57793 Mon Sep 17 00:00:00 2001
From: Andrew Liu <24590316+aliu@users.noreply.github.com>
Date: Tue, 21 Nov 2023 00:32:58 -0800
Subject: [PATCH] util: add `DelayQueue::deadline`

---
 tokio-util/src/time/delay_queue.rs | 39 ++++++++++++++++++++++++++++--
 1 file changed, 37 insertions(+), 2 deletions(-)

diff --git a/tokio-util/src/time/delay_queue.rs b/tokio-util/src/time/delay_queue.rs
index 9136d900b59..2b33e36188d 100644
--- a/tokio-util/src/time/delay_queue.rs
+++ b/tokio-util/src/time/delay_queue.rs
@@ -671,6 +671,39 @@ impl<T> DelayQueue<T> {
         }
     }
 
+    /// Returns the deadline of the item associated with `key`.
+    ///
+    /// Since the queue operates at millisecond granularity, the returned
+    /// deadline may not exactly match the value that was given when initially
+    /// inserting the item into the queue.
+    ///
+    /// # Panics
+    ///
+    /// This function panics if `key` is not contained by the queue.
+    ///
+    /// # 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(5));
+    /// let key2 = delay_queue.insert("bar", Duration::from_secs(10));
+    ///
+    /// assert!(delay_queue.deadline(&key1) < delay_queue.deadline(&key2));
+    /// # }
+    /// ```
+    #[track_caller]
+    pub fn deadline(&self, key: &Key) -> Instant {
+        self.start + Duration::from_millis(self.slab[*key].when)
+    }
+
     /// Removes the key from the expired queue or the timer wheel
     /// depending on its expiration status.
     ///
@@ -909,8 +942,10 @@ impl<T> DelayQueue<T> {
         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> {
+    /// Returns the next time to poll as determined by the wheel.
+    ///
+    /// Note that this does not include deadlines in the `expired` queue.
+    fn next_deadline(&self) -> Option<Instant> {
         self.wheel
             .poll_at()
             .map(|poll_at| self.start + Duration::from_millis(poll_at))