From c800deaaccf1dd5caeb08d9416c4c570939c0c85 Mon Sep 17 00:00:00 2001 From: b-naber Date: Sun, 9 Jan 2022 12:41:30 +0100 Subject: [PATCH] util: add `shrink_to_fit` and `compact` methods to `DelayQueue` (#4170) --- tokio-util/Cargo.toml | 2 +- tokio-util/src/time/delay_queue.rs | 362 ++++++++++++++++++++++++--- tokio-util/src/time/wheel/mod.rs | 1 + tokio-util/src/time/wheel/stack.rs | 4 +- tokio-util/tests/time_delay_queue.rs | 173 +++++++++++++ 5 files changed, 503 insertions(+), 39 deletions(-) diff --git a/tokio-util/Cargo.toml b/tokio-util/Cargo.toml index 00fc018cd..f3c19c1e3 100644 --- a/tokio-util/Cargo.toml +++ b/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"] } diff --git a/tokio-util/src/time/delay_queue.rs b/tokio-util/src/time/delay_queue.rs index 6da83a574..697670d75 100644 --- a/tokio-util/src/time/delay_queue.rs +++ b/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 { /// Stores data associated with entries - slab: Slab>, + slab: SlabStorage, /// Lookup structure tracking all delays in the queue wheel: Wheel>, @@ -152,6 +157,216 @@ pub struct DelayQueue { waker: Option, } +#[derive(Default)] +struct SlabStorage { + inner: Slab>, + + // 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, + + // 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 SlabStorage { + pub(crate) fn with_capacity(capacity: usize) -> SlabStorage { + 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) -> 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 { + 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 { + 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 fmt::Debug for SlabStorage +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 Index for SlabStorage { + type Output = Data; + + 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 IndexMut for SlabStorage { + fn index_mut(&mut self, key: Key) -> &mut Data { + 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 { /// /// [`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 { /// Head of the stack - head: Option, + head: Option, _p: PhantomData T>, } @@ -201,10 +424,10 @@ struct Data { expired: bool, /// Next entry in the stack - next: Option, + next: Option, /// Previous entry in the stack - prev: Option, + prev: Option, } /// Maximum number of entries the queue can handle @@ -253,7 +476,7 @@ impl DelayQueue { pub fn with_capacity(capacity: usize) -> DelayQueue { 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 DelayQueue { } } - Key::new(key) + key } /// Attempts to pull out the next value of the delay queue, registering the @@ -369,13 +592,13 @@ impl DelayQueue { 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 DelayQueue { 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 DelayQueue { 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 DelayQueue { 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 DelayQueue { // 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 DelayQueue { } } + /// 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 { self.wheel @@ -750,7 +1017,7 @@ impl DelayQueue { /// 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>> { + fn poll_idx(&mut self, cx: &mut task::Context<'_>) -> Poll>> { use self::wheel::Stack; let expired = self.expired.pop(&mut self.slab); @@ -816,9 +1083,9 @@ impl futures_core::Stream for DelayQueue { } impl wheel::Stack for Stack { - type Owned = usize; - type Borrowed = usize; - type Store = Slab>; + type Owned = Key; + type Borrowed = Key; + type Store = SlabStorage; fn is_empty(&self) -> bool { self.head.is_none() @@ -837,28 +1104,29 @@ impl wheel::Stack for Stack { } store[item].next = old; - self.head = Some(item) + self.head = Some(item); } fn pop(&mut self, store: &mut Self::Store) -> Option { - 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 wheel::Stack for Stack { 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 for KeyInternal { + fn from(item: Key) -> Self { + KeyInternal::new(item.index) + } +} + +impl From for Key { + fn from(item: KeyInternal) -> Self { + Key::new(item.index) + } +} + impl Expired { /// Returns a reference to the inner value. pub fn get_ref(&self) -> &T { diff --git a/tokio-util/src/time/wheel/mod.rs b/tokio-util/src/time/wheel/mod.rs index 8fed0bf43..4191e401d 100644 --- a/tokio-util/src/time/wheel/mod.rs +++ b/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. diff --git a/tokio-util/src/time/wheel/stack.rs b/tokio-util/src/time/wheel/stack.rs index 6e55c38cc..c87adcafd 100644 --- a/tokio-util/src/time/wheel/stack.rs +++ b/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; /// Borrowed item - type Borrowed; + type Borrowed: Eq + Hash; /// Item storage, this allows a slab to be used instead of just the heap type Store; diff --git a/tokio-util/tests/time_delay_queue.rs b/tokio-util/tests/time_delay_queue.rs index 901afcaaa..1c30446af 100644 --- a/tokio-util/tests/time_delay_queue.rs +++ b/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) }