util: add shrink_to_fit and compact methods to DelayQueue (#4170)

tokio-util/Cargo.toml

@@ -44,7 +44,7 @@ futures-io = { version = "0.3.0", optional = true }
 futures-util = { version = "0.3.0", optional = true }
 log = "0.4"
 pin-project-lite = "0.2.0"
-slab = { version = "0.4.1", optional = true } # Backs `DelayQueue`
+slab = { version = "0.4.4", optional = true } # Backs `DelayQueue`
 
 [dev-dependencies]
 tokio = { version = "1.0.0", path = "../tokio", features = ["full"] }

tokio-util/src/time/delay_queue.rs

@@ -9,8 +9,13 @@ use crate::time::wheel::{self, Wheel};
 use futures_core::ready;
 use tokio::time::{error::Error, sleep_until, Duration, Instant, Sleep};
 
+use core::ops::{Index, IndexMut};
 use slab::Slab;
 use std::cmp;
+use std::collections::HashMap;
+use std::convert::From;
+use std::fmt;
+use std::fmt::Debug;
 use std::future::Future;
 use std::marker::PhantomData;
 use std::pin::Pin;
@@ -128,7 +133,7 @@ use std::task::{self, Poll, Waker};
 #[derive(Debug)]
 pub struct DelayQueue<T> {
     /// Stores data associated with entries
-    slab: Slab<Data<T>>,
+    slab: SlabStorage<T>,
 
     /// Lookup structure tracking all delays in the queue
     wheel: Wheel<Stack<T>>,
@@ -152,6 +157,216 @@ pub struct DelayQueue<T> {
     waker: Option<Waker>,
 }
 
+#[derive(Default)]
+struct SlabStorage<T> {
+    inner: Slab<Data<T>>,
+
+    // A `compact` call requires a re-mapping of the `Key`s that were changed
+    // during the `compact` call of the `slab`. Since the keys that were given out
+    // cannot be changed retroactively we need to keep track of these re-mappings.
+    // The keys of `key_map` correspond to the old keys that were given out and
+    // the values to the `Key`s that were re-mapped by the `compact` call.
+    key_map: HashMap<Key, KeyInternal>,
+
+    // Index used to create new keys to hand out.
+    next_key_index: usize,
+
+    // Whether `compact` has been called, necessary in order to decide whether
+    // to include keys in `key_map`.
+    compact_called: bool,
+}
+
+impl<T> SlabStorage<T> {
+    pub(crate) fn with_capacity(capacity: usize) -> SlabStorage<T> {
+        SlabStorage {
+            inner: Slab::with_capacity(capacity),
+            key_map: HashMap::new(),
+            next_key_index: 0,
+            compact_called: false,
+        }
+    }
+
+    // Inserts data into the inner slab and re-maps keys if necessary
+    pub(crate) fn insert(&mut self, val: Data<T>) -> Key {
+        let mut key = KeyInternal::new(self.inner.insert(val));
+        let key_contained = self.key_map.contains_key(&key.into());
+
+        if key_contained {
+            // It's possible that a `compact` call creates capacitiy in `self.inner` in
+            // such a way that a `self.inner.insert` call creates a `key` which was
+            // previously given out during an `insert` call prior to the `compact` call.
+            // If `key` is contained in `self.key_map`, we have encountered this exact situation,
+            // We need to create a new key `key_to_give_out` and include the relation
+            // `key_to_give_out` -> `key` in `self.key_map`.
+            let key_to_give_out = self.create_new_key();
+            assert!(!self.key_map.contains_key(&key_to_give_out.into()));
+            self.key_map.insert(key_to_give_out.into(), key);
+            key = key_to_give_out;
+        } else if self.compact_called {
+            // Include an identity mapping in `self.key_map` in order to allow us to
+            // panic if a key that was handed out is removed more than once.
+            self.key_map.insert(key.into(), key);
+        }
+
+        key.into()
+    }
+
+    // Re-map the key in case compact was previously called.
+    // Note: Since we include identity mappings in key_map after compact was called,
+    // we have information about all keys that were handed out. In the case in which
+    // compact was called and we try to remove a Key that was previously removed
+    // we can detect invalid keys if no key is found in `key_map`. This is necessary
+    // in order to prevent situations in which a previously removed key
+    // corresponds to a re-mapped key internally and which would then be incorrectly
+    // removed from the slab.
+    //
+    // Example to illuminate this problem:
+    //
+    // Let's assume our `key_map` is {1 -> 2, 2 -> 1} and we call remove(1). If we
+    // were to remove 1 again, we would not find it inside `key_map` anymore.
+    // If we were to imply from this that no re-mapping was necessary, we would
+    // incorrectly remove 1 from `self.slab.inner`, which corresponds to the
+    // handed-out key 2.
+    pub(crate) fn remove(&mut self, key: &Key) -> Data<T> {
+        let remapped_key = if self.compact_called {
+            match self.key_map.remove(key) {
+                Some(key_internal) => key_internal,
+                None => panic!("invalid key"),
+            }
+        } else {
+            (*key).into()
+        };
+
+        self.inner.remove(remapped_key.index)
+    }
+
+    pub(crate) fn shrink_to_fit(&mut self) {
+        self.inner.shrink_to_fit();
+        self.key_map.shrink_to_fit();
+    }
+
+    pub(crate) fn compact(&mut self) {
+        if !self.compact_called {
+            for (key, _) in self.inner.iter() {
+                self.key_map.insert(Key::new(key), KeyInternal::new(key));
+            }
+        }
+
+        let mut remapping = HashMap::new();
+        self.inner.compact(|_, from, to| {
+            remapping.insert(from, to);
+            true
+        });
+
+        // At this point `key_map` contains a mapping for every element.
+        for internal_key in self.key_map.values_mut() {
+            if let Some(new_internal_key) = remapping.get(&internal_key.index) {
+                *internal_key = KeyInternal::new(*new_internal_key);
+            }
+        }
+
+        if self.key_map.capacity() > 2 * self.key_map.len() {
+            self.key_map.shrink_to_fit();
+        }
+
+        self.compact_called = true;
+    }
+
+    // Tries to re-map a `Key` that was given out to the user to its
+    // corresponding internal key.
+    fn remap_key(&self, key: &Key) -> Option<KeyInternal> {
+        let key_map = &self.key_map;
+        if self.compact_called {
+            key_map.get(&*key).copied()
+        } else {
+            Some((*key).into())
+        }
+    }
+
+    fn create_new_key(&mut self) -> KeyInternal {
+        while self.key_map.contains_key(&Key::new(self.next_key_index)) {
+            self.next_key_index = self.next_key_index.wrapping_add(1);
+        }
+
+        KeyInternal::new(self.next_key_index)
+    }
+
+    pub(crate) fn len(&self) -> usize {
+        self.inner.len()
+    }
+
+    pub(crate) fn capacity(&self) -> usize {
+        self.inner.capacity()
+    }
+
+    pub(crate) fn clear(&mut self) {
+        self.inner.clear();
+        self.key_map.clear();
+        self.compact_called = false;
+    }
+
+    pub(crate) fn reserve(&mut self, additional: usize) {
+        self.inner.reserve(additional);
+
+        if self.compact_called {
+            self.key_map.reserve(additional);
+        }
+    }
+
+    pub(crate) fn is_empty(&self) -> bool {
+        self.inner.is_empty()
+    }
+
+    pub(crate) fn contains(&self, key: &Key) -> bool {
+        let remapped_key = self.remap_key(key);
+
+        match remapped_key {
+            Some(internal_key) => self.inner.contains(internal_key.index),
+            None => false,
+        }
+    }
+}
+
+impl<T> fmt::Debug for SlabStorage<T>
+where
+    T: fmt::Debug,
+{
+    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
+        if fmt.alternate() {
+            fmt.debug_map().entries(self.inner.iter()).finish()
+        } else {
+            fmt.debug_struct("Slab")
+                .field("len", &self.len())
+                .field("cap", &self.capacity())
+                .finish()
+        }
+    }
+}
+
+impl<T> Index<Key> for SlabStorage<T> {
+    type Output = Data<T>;
+
+    fn index(&self, key: Key) -> &Self::Output {
+        let remapped_key = self.remap_key(&key);
+
+        match remapped_key {
+            Some(internal_key) => &self.inner[internal_key.index],
+            None => panic!("Invalid index {}", key.index),
+        }
+    }
+}
+
+impl<T> IndexMut<Key> for SlabStorage<T> {
+    fn index_mut(&mut self, key: Key) -> &mut Data<T> {
+        let remapped_key = self.remap_key(&key);
+
+        match remapped_key {
+            Some(internal_key) => &mut self.inner[internal_key.index],
+            None => panic!("Invalid index {}", key.index),
+        }
+    }
+}
+
 /// An entry in `DelayQueue` that has expired and been removed.
 ///
 /// Values are returned by [`DelayQueue::poll_expired`].
@@ -176,15 +391,23 @@ pub struct Expired<T> {
 ///
 /// [`DelayQueue`]: struct@DelayQueue
 /// [`DelayQueue::insert`]: method@DelayQueue::insert
-#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 pub struct Key {
     index: usize,
 }
 
+// Whereas `Key` is given out to users that use `DelayQueue`, internally we use
+// `KeyInternal` as the key type in order to make the logic of mapping between keys
+// as a result of `compact` calls clearer.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+struct KeyInternal {
+    index: usize,
+}
+
 #[derive(Debug)]
 struct Stack<T> {
     /// Head of the stack
-    head: Option<usize>,
+    head: Option<Key>,
     _p: PhantomData<fn() -> T>,
 }
 
@@ -201,10 +424,10 @@ struct Data<T> {
     expired: bool,
 
     /// Next entry in the stack
-    next: Option<usize>,
+    next: Option<Key>,
 
     /// Previous entry in the stack
-    prev: Option<usize>,
+    prev: Option<Key>,
 }
 
 /// Maximum number of entries the queue can handle
@@ -253,7 +476,7 @@ impl<T> DelayQueue<T> {
     pub fn with_capacity(capacity: usize) -> DelayQueue<T> {
         DelayQueue {
             wheel: Wheel::new(),
-            slab: Slab::with_capacity(capacity),
+            slab: SlabStorage::with_capacity(capacity),
             expired: Stack::default(),
             delay: None,
             wheel_now: 0,
@@ -348,7 +571,7 @@ impl<T> DelayQueue<T> {
             }
         }
 
-        Key::new(key)
+        key
     }
 
     /// Attempts to pull out the next value of the delay queue, registering the
@@ -369,13 +592,13 @@ impl<T> DelayQueue<T> {
 
         let item = ready!(self.poll_idx(cx));
         Poll::Ready(item.map(|result| {
-            result.map(|idx| {
-                let data = self.slab.remove(idx);
+            result.map(|key| {
+                let data = self.slab.remove(&key);
                 debug_assert!(data.next.is_none());
                 debug_assert!(data.prev.is_none());
 
                 Expired {
-                    key: Key::new(idx),
+                    key,
                     data: data.inner,
                     deadline: self.start + Duration::from_millis(data.when),
                 }
@@ -437,7 +660,7 @@ impl<T> DelayQueue<T> {
         self.insert_at(value, Instant::now() + timeout)
     }
 
-    fn insert_idx(&mut self, when: u64, key: usize) {
+    fn insert_idx(&mut self, when: u64, key: Key) {
         use self::wheel::{InsertError, Stack};
 
         // Register the deadline with the timer wheel
@@ -462,10 +685,10 @@ impl<T> DelayQueue<T> {
         use crate::time::wheel::Stack;
 
         // Special case the `expired` queue
-        if self.slab[key.index].expired {
-            self.expired.remove(&key.index, &mut self.slab);
+        if self.slab[*key].expired {
+            self.expired.remove(key, &mut self.slab);
         } else {
-            self.wheel.remove(&key.index, &mut self.slab);
+            self.wheel.remove(key, &mut self.slab);
         }
     }
 
@@ -501,7 +724,7 @@ impl<T> DelayQueue<T> {
         let prev_deadline = self.next_deadline();
 
         self.remove_key(key);
-        let data = self.slab.remove(key.index);
+        let data = self.slab.remove(key);
 
         let next_deadline = self.next_deadline();
         if prev_deadline != next_deadline {
@@ -559,10 +782,10 @@ impl<T> DelayQueue<T> {
         // Normalize the deadline. Values cannot be set to expire in the past.
         let when = self.normalize_deadline(when);
 
-        self.slab[key.index].when = when;
-        self.slab[key.index].expired = false;
+        self.slab[*key].when = when;
+        self.slab[*key].expired = false;
 
-        self.insert_idx(when, key.index);
+        self.insert_idx(when, *key);
 
         let next_deadline = self.next_deadline();
         if let (Some(ref mut delay), Some(deadline)) = (&mut self.delay, next_deadline) {
@@ -571,6 +794,50 @@ impl<T> DelayQueue<T> {
         }
     }
 
+    /// Shrink the capacity of the slab, which `DelayQueue` uses internally for storage allocation.
+    /// This function is not guaranteed to, and in most cases, won't decrease the capacity of the slab
+    /// to the number of elements still contained in it, because elements cannot be moved to a different
+    /// index. To decrease the capacity to the size of the slab use [`compact`].
+    ///
+    /// This function can take O(n) time even when the capacity cannot be reduced or the allocation is
+    /// shrunk in place. Repeated calls run in O(1) though.
+    ///
+    /// [`compact`]: method@Self::compact
+    pub fn shrink_to_fit(&mut self) {
+        self.slab.shrink_to_fit();
+    }
+
+    /// Shrink the capacity of the slab, which `DelayQueue` uses internally for storage allocation,
+    /// to the number of elements that are contained in it.
+    ///
+    /// This methods runs in O(n).
+    ///
+    /// # Examples
+    ///
+    /// Basic usage
+    ///
+    /// ```rust
+    /// use tokio_util::time::DelayQueue;
+    /// use std::time::Duration;
+    ///
+    /// # #[tokio::main]
+    /// # async fn main() {
+    /// let mut delay_queue = DelayQueue::with_capacity(10);
+    ///
+    /// let key1 = delay_queue.insert(5, Duration::from_secs(5));
+    /// let key2 = delay_queue.insert(10, Duration::from_secs(10));
+    /// let key3 = delay_queue.insert(15, Duration::from_secs(15));
+    ///
+    /// delay_queue.remove(&key2);
+    ///
+    /// delay_queue.compact();
+    /// assert_eq!(delay_queue.capacity(), 2);
+    /// # }
+    /// ```
+    pub fn compact(&mut self) {
+        self.slab.compact();
+    }
+
     /// Returns the next time to poll as determined by the wheel
     fn next_deadline(&mut self) -> Option<Instant> {
         self.wheel
@@ -750,7 +1017,7 @@ impl<T> DelayQueue<T> {
     /// should be returned.
     ///
     /// A slot should be returned when the associated deadline has been reached.
-    fn poll_idx(&mut self, cx: &mut task::Context<'_>) -> Poll<Option<Result<usize, Error>>> {
+    fn poll_idx(&mut self, cx: &mut task::Context<'_>) -> Poll<Option<Result<Key, Error>>> {
         use self::wheel::Stack;
 
         let expired = self.expired.pop(&mut self.slab);
@@ -816,9 +1083,9 @@ impl<T> futures_core::Stream for DelayQueue<T> {
 }
 
 impl<T> wheel::Stack for Stack<T> {
-    type Owned = usize;
-    type Borrowed = usize;
-    type Store = Slab<Data<T>>;
+    type Owned = Key;
+    type Borrowed = Key;
+    type Store = SlabStorage<T>;
 
     fn is_empty(&self) -> bool {
         self.head.is_none()
@@ -837,28 +1104,29 @@ impl<T> wheel::Stack for Stack<T> {
         }
 
         store[item].next = old;
-        self.head = Some(item)
+        self.head = Some(item);
     }
 
     fn pop(&mut self, store: &mut Self::Store) -> Option<Self::Owned> {
-        if let Some(idx) = self.head {
-            self.head = store[idx].next;
+        if let Some(key) = self.head {
+            self.head = store[key].next;
 
             if let Some(idx) = self.head {
                 store[idx].prev = None;
             }
 
-            store[idx].next = None;
-            debug_assert!(store[idx].prev.is_none());
+            store[key].next = None;
+            debug_assert!(store[key].prev.is_none());
 
-            Some(idx)
+            Some(key)
         } else {
             None
         }
     }
 
     fn remove(&mut self, item: &Self::Borrowed, store: &mut Self::Store) {
-        assert!(store.contains(*item));
+        let key = *item;
+        assert!(store.contains(item));
 
         // Ensure that the entry is in fact contained by the stack
         debug_assert!({
@@ -867,29 +1135,31 @@ impl<T> wheel::Stack for Stack<T> {
             let mut contains = false;
 
             while let Some(idx) = next {
+                let data = &store[idx];
+
                 if idx == *item {
                     debug_assert!(!contains);
                     contains = true;
                 }
 
-                next = store[idx].next;
+                next = data.next;
             }
 
             contains
         });
 
-        if let Some(next) = store[*item].next {
-            store[next].prev = store[*item].prev;
+        if let Some(next) = store[key].next {
+            store[next].prev = store[key].prev;
         }
 
-        if let Some(prev) = store[*item].prev {
-            store[prev].next = store[*item].next;
+        if let Some(prev) = store[key].prev {
+            store[prev].next = store[key].next;
         } else {
-            self.head = store[*item].next;
+            self.head = store[key].next;
         }
 
-        store[*item].next = None;
-        store[*item].prev = None;
+        store[key].next = None;
+        store[key].prev = None;
     }
 
     fn when(item: &Self::Borrowed, store: &Self::Store) -> u64 {
@@ -912,6 +1182,24 @@ impl Key {
     }
 }
 
+impl KeyInternal {
+    pub(crate) fn new(index: usize) -> KeyInternal {
+        KeyInternal { index }
+    }
+}
+
+impl From<Key> for KeyInternal {
+    fn from(item: Key) -> Self {
+        KeyInternal::new(item.index)
+    }
+}
+
+impl From<KeyInternal> for Key {
+    fn from(item: KeyInternal) -> Self {
+        Key::new(item.index)
+    }
+}
+
 impl<T> Expired<T> {
     /// Returns a reference to the inner value.
     pub fn get_ref(&self) -> &T {

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

@@ -6,6 +6,7 @@ mod stack;
 pub(crate) use self::stack::Stack;
 
 use std::borrow::Borrow;
+use std::fmt::Debug;
 use std::usize;
 
 /// Timing wheel implementation.

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

@@ -1,4 +1,6 @@
 use std::borrow::Borrow;
+use std::cmp::Eq;
+use std::hash::Hash;
 
 /// Abstracts the stack operations needed to track timeouts.
 pub(crate) trait Stack: Default {
@@ -6,7 +8,7 @@ pub(crate) trait Stack: Default {
     type Owned: Borrow<Self::Borrowed>;
 
     /// Borrowed item
-    type Borrowed;
+    type Borrowed: Eq + Hash;
 
     /// Item storage, this allows a slab to be used instead of just the heap
     type Store;

tokio-util/tests/time_delay_queue.rs

@@ -109,6 +109,7 @@ async fn multi_delay_at_start() {
 
     let start = Instant::now();
     for elapsed in 0..1200 {
+        println!("elapsed: {:?}", elapsed);
         let elapsed = elapsed + 1;
         tokio::time::sleep_until(start + ms(elapsed)).await;
 
@@ -128,10 +129,12 @@ async fn multi_delay_at_start() {
             assert_pending!(poll!(queue));
         }
     }
+    println!("finished multi_delay_start");
 }
 
 #[tokio::test]
 async fn insert_in_past_fires_immediately() {
+    println!("running insert_in_past_fires_immediately");
     time::pause();
 
     let mut queue = task::spawn(DelayQueue::new());
@@ -142,6 +145,7 @@ async fn insert_in_past_fires_immediately() {
     queue.insert_at("foo", now);
 
     assert_ready!(poll!(queue));
+    println!("finished insert_in_past_fires_immediately");
 }
 
 #[tokio::test]
@@ -640,6 +644,175 @@ async fn delay_queue_poll_expired_when_empty() {
     assert!(assert_ready!(poll!(delay_queue)).is_none());
 }
 
+#[tokio::test(start_paused = true)]
+async fn compact_expire_empty() {
+    let mut queue = task::spawn(DelayQueue::new());
+
+    let now = Instant::now();
+
+    queue.insert_at("foo1", now + ms(10));
+    queue.insert_at("foo2", now + ms(10));
+
+    sleep(ms(10)).await;
+
+    let mut res = vec![];
+    while res.len() < 2 {
+        let entry = assert_ready_ok!(poll!(queue));
+        res.push(entry.into_inner());
+    }
+
+    queue.compact();
+
+    assert_eq!(queue.len(), 0);
+    assert_eq!(queue.capacity(), 0);
+}
+
+#[tokio::test(start_paused = true)]
+async fn compact_remove_empty() {
+    let mut queue = task::spawn(DelayQueue::new());
+
+    let now = Instant::now();
+
+    let key1 = queue.insert_at("foo1", now + ms(10));
+    let key2 = queue.insert_at("foo2", now + ms(10));
+
+    queue.remove(&key1);
+    queue.remove(&key2);
+
+    queue.compact();
+
+    assert_eq!(queue.len(), 0);
+    assert_eq!(queue.capacity(), 0);
+}
+
+#[tokio::test(start_paused = true)]
+// Trigger a re-mapping of keys in the slab due to a `compact` call and
+// test removal of re-mapped keys
+async fn compact_remove_remapped_keys() {
+    let mut queue = task::spawn(DelayQueue::new());
+
+    let now = Instant::now();
+
+    queue.insert_at("foo1", now + ms(10));
+    queue.insert_at("foo2", now + ms(10));
+
+    // should be assigned indices 3 and 4
+    let key3 = queue.insert_at("foo3", now + ms(20));
+    let key4 = queue.insert_at("foo4", now + ms(20));
+
+    sleep(ms(10)).await;
+
+    let mut res = vec![];
+    while res.len() < 2 {
+        let entry = assert_ready_ok!(poll!(queue));
+        res.push(entry.into_inner());
+    }
+
+    // items corresponding to `foo3` and `foo4` will be assigned
+    // new indices here
+    queue.compact();
+
+    queue.insert_at("foo5", now + ms(10));
+
+    // test removal of re-mapped keys
+    let expired3 = queue.remove(&key3);
+    let expired4 = queue.remove(&key4);
+
+    assert_eq!(expired3.into_inner(), "foo3");
+    assert_eq!(expired4.into_inner(), "foo4");
+
+    queue.compact();
+    assert_eq!(queue.len(), 1);
+    assert_eq!(queue.capacity(), 1);
+}
+
+#[tokio::test(start_paused = true)]
+async fn compact_change_deadline() {
+    let mut queue = task::spawn(DelayQueue::new());
+
+    let mut now = Instant::now();
+
+    queue.insert_at("foo1", now + ms(10));
+    queue.insert_at("foo2", now + ms(10));
+
+    // should be assigned indices 3 and 4
+    queue.insert_at("foo3", now + ms(20));
+    let key4 = queue.insert_at("foo4", now + ms(20));
+
+    sleep(ms(10)).await;
+
+    let mut res = vec![];
+    while res.len() < 2 {
+        let entry = assert_ready_ok!(poll!(queue));
+        res.push(entry.into_inner());
+    }
+
+    // items corresponding to `foo3` and `foo4` should be assigned
+    // new indices
+    queue.compact();
+
+    now = Instant::now();
+
+    queue.insert_at("foo5", now + ms(10));
+    let key6 = queue.insert_at("foo6", now + ms(10));
+
+    queue.reset_at(&key4, now + ms(20));
+    queue.reset_at(&key6, now + ms(20));
+
+    // foo3 and foo5 will expire
+    sleep(ms(10)).await;
+
+    while res.len() < 4 {
+        let entry = assert_ready_ok!(poll!(queue));
+        res.push(entry.into_inner());
+    }
+
+    sleep(ms(10)).await;
+
+    while res.len() < 6 {
+        let entry = assert_ready_ok!(poll!(queue));
+        res.push(entry.into_inner());
+    }
+
+    let entry = assert_ready!(poll!(queue));
+    assert!(entry.is_none());
+}
+
+#[tokio::test(start_paused = true)]
+async fn remove_after_compact() {
+    let now = Instant::now();
+    let mut queue = DelayQueue::new();
+
+    let foo_key = queue.insert_at("foo", now + ms(10));
+    queue.insert_at("bar", now + ms(20));
+    queue.remove(&foo_key);
+    queue.compact();
+
+    let panic = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
+        queue.remove(&foo_key);
+    }));
+    assert!(panic.is_err());
+}
+
+#[tokio::test(start_paused = true)]
+async fn remove_after_compact_poll() {
+    let now = Instant::now();
+    let mut queue = task::spawn(DelayQueue::new());
+
+    let foo_key = queue.insert_at("foo", now + ms(10));
+    queue.insert_at("bar", now + ms(20));
+
+    sleep(ms(10)).await;
+    assert_eq!(assert_ready_ok!(poll!(queue)).key(), foo_key);
+
+    queue.compact();
+
+    let panic = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
+        queue.remove(&foo_key);
+    }));
+    assert!(panic.is_err());
+}
+
 fn ms(n: u64) -> Duration {
     Duration::from_millis(n)
 }