[#5536] tokio-util: fix panic in DelayQueue

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

@@ -140,11 +140,31 @@ impl<T: Stack> Level<T> {
 
         // TODO: This can probably be simplified w/ power of 2 math
         let level_start = now - (now % level_range);
-        let deadline = level_start + slot as u64 * slot_range;
-
+        let mut deadline = level_start + slot as u64 * slot_range;
+        if deadline < now {
+            // A timer is in a slot "prior" to the current time. This can occur
+            // because we do not have an infinite hierarchy of timer levels, and
+            // eventually a timer scheduled for a very distant time might end up
+            // being placed in a slot that is beyond the end of all of the
+            // arrays.
+            //
+            // To deal with this, we first limit timers to being scheduled no
+            // more than MAX_DURATION ticks in the future; that is, they're at
+            // most one rotation of the top level away. Then, we force timers
+            // that logically would go into the top+1 level, to instead go into
+            // the top level's slots.
+            //
+            // What this means is that the top level's slots act as a
+            // pseudo-ring buffer, and we rotate around them indefinitely. If we
+            // compute a deadline before now, and it's the top level, it
+            // therefore means we're actually looking at a slot in the future.
+            debug_assert_eq!(self.level, super::NUM_LEVELS - 1);
+
+            deadline += level_range;
+        }
         debug_assert!(
             deadline >= now,
-            "deadline={}; now={}; level={}; slot={}; occupied={:b}",
+            "deadline={:016X}; now={:016X}; level={}; slot={}; occupied={:b}",
             deadline,
             now,
             self.level,

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

@@ -254,8 +254,11 @@ fn level_for(elapsed: u64, when: u64) -> usize {
 
     // Mask in the trailing bits ignored by the level calculation in order to cap
     // the possible leading zeros
-    let masked = elapsed ^ when | SLOT_MASK;
-
+    let mut masked = elapsed ^ when | SLOT_MASK;
+    if masked >= MAX_DURATION {
+        // Fudge the timer into the top level
+        masked = MAX_DURATION - 1;
+    }
     let leading_zeros = masked.leading_zeros() as usize;
     let significant = 63 - leading_zeros;
     significant / 6

tokio-util/tests/time_delay_queue.rs

@@ -2,6 +2,7 @@
 #![warn(rust_2018_idioms)]
 #![cfg(feature = "full")]
 
+use futures::StreamExt;
 use tokio::time::{self, sleep, sleep_until, Duration, Instant};
 use tokio_test::{assert_pending, assert_ready, task};
 use tokio_util::time::DelayQueue;
@@ -786,6 +787,22 @@ async fn compact_change_deadline() {
     assert!(entry.is_none());
 }
 
+#[tokio::test(start_paused = true)]
+async fn item_expiry_greater_than_wheel() {
+    // This function tests that a delay queue that has existed for at least 2^36 milliseconds won't panic when a new item is inserted.
+    let mut queue = DelayQueue::new();
+    for _ in 0..2 {
+        tokio::time::advance(Duration::from_millis(1 << 35)).await;
+        queue.insert(0, Duration::from_millis(0));
+        queue.next().await;
+    }
+    // This should not panic
+    let no_panic = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
+        queue.insert(1, Duration::from_millis(1));
+    }));
+    assert!(no_panic.is_ok());
+}
+
 #[cfg_attr(target_os = "wasi", ignore = "FIXME: Does not seem to work with WASI")]
 #[tokio::test(start_paused = true)]
 async fn remove_after_compact() {