util: fix panics on updating DelayQueue entries (#3270)

tokio-util/src/time/delay_queue.rs

@@ -36,14 +36,14 @@ use std::task::{self, Poll, Waker};
 /// # `Stream` implementation
 ///
 /// Items are retrieved from the queue via [`DelayQueue::poll_expired`]. If no delays have
-/// expired, no items are returned. In this case, `NotReady` is returned and the
+/// expired, no items are returned. In this case, `Pending` is returned and the
 /// current task is registered to be notified once the next item's delay has
 /// expired.
 ///
 /// If no items are in the queue, i.e. `is_empty()` returns `true`, then `poll`
 /// returns `Ready(None)`. This indicates that the stream has reached an end.
 /// However, if a new item is inserted *after*, `poll` will once again start
-/// returning items or `NotReady.
+/// returning items or `Pending.
 ///
 /// Items are returned ordered by their expirations. Items that are configured
 /// to expire first will be returned first. There are no ordering guarantees
@@ -538,7 +538,7 @@ impl<T> DelayQueue<T> {
     ///
     ///     delay_queue.reset_at(&key, Instant::now() + Duration::from_secs(10));
     ///
-    ///     // "foo"is now scheduled to be returned in 10 seconds
+    ///     // "foo" is now scheduled to be returned in 10 seconds
     /// # }
     /// ```
     pub fn reset_at(&mut self, key: &Key, when: Instant) {
@@ -548,6 +548,8 @@ impl<T> DelayQueue<T> {
         let when = self.normalize_deadline(when);
 
         self.slab[key.index].when = when;
+        self.slab[key.index].expired = false;
+
         self.insert_idx(when, key.index);
 
         let next_deadline = self.next_deadline();
@@ -711,7 +713,7 @@ impl<T> DelayQueue<T> {
     /// Returns `true` if there are no items in the queue.
     ///
     /// Note that this function returns `false` even if all items have not yet
-    /// expired and a call to `poll` will return `NotReady`.
+    /// expired and a call to `poll` will return `Pending`.
     ///
     /// # Examples
     ///

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

@@ -233,14 +233,13 @@ fn slot_for(duration: u64, level: usize) -> usize {
     ((duration >> (level * 6)) % LEVEL_MULT as u64) as usize
 }
 
-/*
 #[cfg(all(test, not(loom)))]
 mod test {
     use super::*;
 
     #[test]
     fn test_slot_for() {
-        for pos in 1..64 {
+        for pos in 0..64 {
             assert_eq!(pos as usize, slot_for(pos, 0));
         }
 
@@ -252,4 +251,3 @@ mod test {
         }
     }
 }
-*/

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

@@ -116,9 +116,17 @@ where
         Ok(())
     }
 
-    /// Remove `item` from thee timing wheel.
+    /// Remove `item` from the timing wheel.
     pub(crate) fn remove(&mut self, item: &T::Borrowed, store: &mut T::Store) {
         let when = T::when(item, store);
+
+        assert!(
+            self.elapsed <= when,
+            "elapsed={}; when={}",
+            self.elapsed,
+            when
+        );
+
         let level = self.level_for(when);
 
         self.levels[level].remove_entry(when, item, store);
@@ -240,9 +248,11 @@ where
 }
 
 fn level_for(elapsed: u64, when: u64) -> usize {
-    let masked = elapsed ^ when;
+    const SLOT_MASK: u64 = (1 << 6) - 1;
 
-    assert!(masked != 0, "elapsed={}; when={}", elapsed, when);
+    // 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 leading_zeros = masked.leading_zeros() as usize;
     let significant = 63 - leading_zeros;
@@ -255,7 +265,7 @@ mod test {
 
     #[test]
     fn test_level_for() {
-        for pos in 1..64 {
+        for pos in 0..64 {
             assert_eq!(
                 0,
                 level_for(0, pos),