diff --git a/tokio-util/src/time/delay_queue.rs b/tokio-util/src/time/delay_queue.rs index b53d18046..d8f02ae4e 100644 --- a/tokio-util/src/time/delay_queue.rs +++ b/tokio-util/src/time/delay_queue.rs @@ -951,7 +951,7 @@ impl DelayQueue { pub fn peek(&self) -> Option { use self::wheel::Stack; - self.expired.peek().or_else(|| self.wheel.peek()) + self.expired.peek().or_else(|| self.wheel.peek(&self.slab)) } /// Returns the next time to poll as determined by the wheel. @@ -1260,6 +1260,27 @@ impl wheel::Stack for Stack { self.head } + fn peek_earliest(&self, store: &Self::Store) -> Option { + let head = self.head?; + let mut earliest = (head, store[head].when); + let mut curr = store[head].next; + + while let Some(key) = curr { + let data = &store[key]; + + // The comparison is strict so that the first entry seen wins a tie, + // which agrees with `pop` when every entry in the slot shares a + // deadline. + if data.when < earliest.1 { + earliest = (key, data.when); + } + + curr = data.next; + } + + Some(earliest.0) + } + #[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 f3a08d8bf..0f62c22fc 100644 --- a/tokio-util/src/time/wheel/level.rs +++ b/tokio-util/src/time/wheel/level.rs @@ -148,8 +148,8 @@ impl Level { ret } - pub(crate) fn peek_entry_slot(&self, slot: usize) -> Option { - self.slot[slot].peek() + pub(crate) fn peek_entry_slot(&self, slot: usize, store: &T::Store) -> Option { + self.slot[slot].peek_earliest(store) } } diff --git a/tokio-util/src/time/wheel/mod.rs b/tokio-util/src/time/wheel/mod.rs index e93c63626..726cc659f 100644 --- a/tokio-util/src/time/wheel/mod.rs +++ b/tokio-util/src/time/wheel/mod.rs @@ -140,9 +140,9 @@ where } /// Next key that will expire - pub(crate) fn peek(&self) -> Option { + pub(crate) fn peek(&self, store: &T::Store) -> Option { self.next_expiration() - .and_then(|expiration| self.peek_entry(&expiration)) + .and_then(|expiration| self.peek_entry(&expiration, store)) } /// Advances the timer up to the instant represented by `now`. @@ -250,8 +250,8 @@ where self.levels[expiration.level].pop_entry_slot(expiration.slot, store) } - fn peek_entry(&self, expiration: &Expiration) -> Option { - self.levels[expiration.level].peek_entry_slot(expiration.slot) + fn peek_entry(&self, expiration: &Expiration, store: &T::Store) -> Option { + self.levels[expiration.level].peek_entry_slot(expiration.slot, store) } fn level_for(&self, when: u64) -> usize { diff --git a/tokio-util/src/time/wheel/stack.rs b/tokio-util/src/time/wheel/stack.rs index 7d32f27d2..d07dea219 100644 --- a/tokio-util/src/time/wheel/stack.rs +++ b/tokio-util/src/time/wheel/stack.rs @@ -25,6 +25,13 @@ pub(crate) trait Stack: Default { /// Peek into the stack. fn peek(&self) -> Option; + /// Peek at the item in the stack with the earliest deadline. + /// + /// Unlike `peek`, this does not have to agree with `pop`: a slot in a level + /// above zero spans a range of deadlines, so its entries are only ordered + /// once they cascade down. + fn peek_earliest(&self, store: &Self::Store) -> Option; + 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 a2db75c00..8cc775922 100644 --- a/tokio-util/tests/time_delay_queue.rs +++ b/tokio-util/tests/time_delay_queue.rs @@ -902,6 +902,41 @@ async fn peek() { assert!(queue.peek().is_none()); } +#[tokio::test(start_paused = true)] +async fn peek_entries_sharing_a_wheel_slot() { + // Only the lowest level of the timer wheel has one deadline per slot. A + // level-one slot spans 64ms, so entries with different deadlines share it + // and are only ordered once the slot cascades down. + let mut queue = task::spawn(DelayQueue::new()); + + let now = Instant::now(); + + // 64ms..=127ms all fall in the same level-one slot. Entries are pushed onto + // the front of a slot, so inserting the later deadline second puts it at the + // head. + let early = queue.insert_at("early", now + ms(100)); + let late = queue.insert_at("late", now + ms(120)); + + assert_eq!(queue.peek(), Some(early)); + + sleep(ms(105)).await; + + assert_eq!(queue.peek(), Some(early)); + + let entry = assert_ready_some!(poll!(queue)); + assert_eq!(entry.key(), early); + assert_eq!(entry.get_ref(), &"early"); + + assert_eq!(queue.peek(), Some(late)); + + sleep(ms(20)).await; + + let entry = assert_ready_some!(poll!(queue)); + assert_eq!(entry.key(), late); + + assert!(queue.peek().is_none()); +} + #[tokio::test(start_paused = true)] async fn wake_after_remove_last() { let mut queue = task::spawn(DelayQueue::new());