1//! A queue of delayed elements.
2//!
3//! See [`DelayQueue`] for more details.
4//!
5//! [`DelayQueue`]: struct@DelayQueue
6
7use crate::time::wheel::{self, Wheel};
8
9use futures_core::ready;
10use tokio::time::{sleep_until, Duration, Instant, Sleep};
11
12use core::ops::{Index, IndexMut};
13use slab::Slab;
14use std::cmp;
15use std::collections::HashMap;
16use std::convert::From;
17use std::fmt;
18use std::fmt::Debug;
19use std::future::Future;
20use std::marker::PhantomData;
21use std::pin::Pin;
22use std::task::{self, Poll, Waker};
23
24/// A queue of delayed elements.
25///
26/// Once an element is inserted into the `DelayQueue`, it is yielded once the
27/// specified deadline has been reached.
28///
29/// # Usage
30///
31/// Elements are inserted into `DelayQueue` using the [`insert`] or
32/// [`insert_at`] methods. A deadline is provided with the item and a [`Key`] is
33/// returned. The key is used to remove the entry or to change the deadline at
34/// which it should be yielded back.
35///
36/// Once delays have been configured, the `DelayQueue` is used via its
37/// [`Stream`] implementation. [`poll_expired`] is called. If an entry has reached its
38/// deadline, it is returned. If not, `Poll::Pending` is returned indicating that the
39/// current task will be notified once the deadline has been reached.
40///
41/// # `Stream` implementation
42///
43/// Items are retrieved from the queue via [`DelayQueue::poll_expired`]. If no delays have
44/// expired, no items are returned. In this case, `Poll::Pending` is returned and the
45/// current task is registered to be notified once the next item's delay has
46/// expired.
47///
48/// If no items are in the queue, i.e. `is_empty()` returns `true`, then `poll`
49/// returns `Poll::Ready(None)`. This indicates that the stream has reached an end.
50/// However, if a new item is inserted *after*, `poll` will once again start
51/// returning items or `Poll::Pending`.
52///
53/// Items are returned ordered by their expirations. Items that are configured
54/// to expire first will be returned first. There are no ordering guarantees
55/// for items configured to expire at the same instant. Also note that delays are
56/// rounded to the closest millisecond.
57///
58/// # Implementation
59///
60/// The [`DelayQueue`] is backed by a separate instance of a timer wheel similar to that used internally
61/// by Tokio's standalone timer utilities such as [`sleep`]. Because of this, it offers the same
62/// performance and scalability benefits.
63///
64/// State associated with each entry is stored in a [`slab`]. This amortizes the cost of allocation,
65/// and allows reuse of the memory allocated for expired entries.
66///
67/// Capacity can be checked using [`capacity`] and allocated preemptively by using
68/// the [`reserve`] method.
69///
70/// # Usage
71///
72/// Using `DelayQueue` to manage cache entries.
73///
74/// ```rust,no_run
75/// use tokio_util::time::{DelayQueue, delay_queue};
76///
77/// use futures::ready;
78/// use std::collections::HashMap;
79/// use std::task::{Context, Poll};
80/// use std::time::Duration;
81/// # type CacheKey = String;
82/// # type Value = String;
83///
84/// struct Cache {
85/// entries: HashMap<CacheKey, (Value, delay_queue::Key)>,
86/// expirations: DelayQueue<CacheKey>,
87/// }
88///
89/// const TTL_SECS: u64 = 30;
90///
91/// impl Cache {
92/// fn insert(&mut self, key: CacheKey, value: Value) {
93/// let delay = self.expirations
94/// .insert(key.clone(), Duration::from_secs(TTL_SECS));
95///
96/// self.entries.insert(key, (value, delay));
97/// }
98///
99/// fn get(&self, key: &CacheKey) -> Option<&Value> {
100/// self.entries.get(key)
101/// .map(|&(ref v, _)| v)
102/// }
103///
104/// fn remove(&mut self, key: &CacheKey) {
105/// if let Some((_, cache_key)) = self.entries.remove(key) {
106/// self.expirations.remove(&cache_key);
107/// }
108/// }
109///
110/// fn poll_purge(&mut self, cx: &mut Context<'_>) -> Poll<()> {
111/// while let Some(entry) = ready!(self.expirations.poll_expired(cx)) {
112/// self.entries.remove(entry.get_ref());
113/// }
114///
115/// Poll::Ready(())
116/// }
117/// }
118/// ```
119///
120/// [`insert`]: method@Self::insert
121/// [`insert_at`]: method@Self::insert_at
122/// [`Key`]: struct@Key
123/// [`Stream`]: https://docs.rs/futures/0.1/futures/stream/trait.Stream.html
124/// [`poll_expired`]: method@Self::poll_expired
125/// [`Stream::poll_expired`]: method@Self::poll_expired
126/// [`DelayQueue`]: struct@DelayQueue
127/// [`sleep`]: fn@tokio::time::sleep
128/// [`slab`]: slab
129/// [`capacity`]: method@Self::capacity
130/// [`reserve`]: method@Self::reserve
131#[derive(Debug)]
132pub struct DelayQueue<T> {
133 /// Stores data associated with entries
134 slab: SlabStorage<T>,
135
136 /// Lookup structure tracking all delays in the queue
137 wheel: Wheel<Stack<T>>,
138
139 /// Delays that were inserted when already expired. These cannot be stored
140 /// in the wheel
141 expired: Stack<T>,
142
143 /// Delay expiring when the *first* item in the queue expires
144 delay: Option<Pin<Box<Sleep>>>,
145
146 /// Wheel polling state
147 wheel_now: u64,
148
149 /// Instant at which the timer starts
150 start: Instant,
151
152 /// Waker that is invoked when we potentially need to reset the timer.
153 /// Because we lazily create the timer when the first entry is created, we
154 /// need to awaken any poller that polled us before that point.
155 waker: Option<Waker>,
156}
157
158#[derive(Default)]
159struct SlabStorage<T> {
160 inner: Slab<Data<T>>,
161
162 // A `compact` call requires a re-mapping of the `Key`s that were changed
163 // during the `compact` call of the `slab`. Since the keys that were given out
164 // cannot be changed retroactively we need to keep track of these re-mappings.
165 // The keys of `key_map` correspond to the old keys that were given out and
166 // the values to the `Key`s that were re-mapped by the `compact` call.
167 key_map: HashMap<Key, KeyInternal>,
168
169 // Index used to create new keys to hand out.
170 next_key_index: usize,
171
172 // Whether `compact` has been called, necessary in order to decide whether
173 // to include keys in `key_map`.
174 compact_called: bool,
175}
176
177impl<T> SlabStorage<T> {
178 pub(crate) fn with_capacity(capacity: usize) -> SlabStorage<T> {
179 SlabStorage {
180 inner: Slab::with_capacity(capacity),
181 key_map: HashMap::new(),
182 next_key_index: 0,
183 compact_called: false,
184 }
185 }
186
187 // Inserts data into the inner slab and re-maps keys if necessary
188 pub(crate) fn insert(&mut self, val: Data<T>) -> Key {
189 let mut key = KeyInternal::new(self.inner.insert(val));
190 let key_contained = self.key_map.contains_key(&key.into());
191
192 if key_contained {
193 // It's possible that a `compact` call creates capacity in `self.inner` in
194 // such a way that a `self.inner.insert` call creates a `key` which was
195 // previously given out during an `insert` call prior to the `compact` call.
196 // If `key` is contained in `self.key_map`, we have encountered this exact situation,
197 // We need to create a new key `key_to_give_out` and include the relation
198 // `key_to_give_out` -> `key` in `self.key_map`.
199 let key_to_give_out = self.create_new_key();
200 assert!(!self.key_map.contains_key(&key_to_give_out.into()));
201 self.key_map.insert(key_to_give_out.into(), key);
202 key = key_to_give_out;
203 } else if self.compact_called {
204 // Include an identity mapping in `self.key_map` in order to allow us to
205 // panic if a key that was handed out is removed more than once.
206 self.key_map.insert(key.into(), key);
207 }
208
209 key.into()
210 }
211
212 // Re-map the key in case compact was previously called.
213 // Note: Since we include identity mappings in key_map after compact was called,
214 // we have information about all keys that were handed out. In the case in which
215 // compact was called and we try to remove a Key that was previously removed
216 // we can detect invalid keys if no key is found in `key_map`. This is necessary
217 // in order to prevent situations in which a previously removed key
218 // corresponds to a re-mapped key internally and which would then be incorrectly
219 // removed from the slab.
220 //
221 // Example to illuminate this problem:
222 //
223 // Let's assume our `key_map` is {1 -> 2, 2 -> 1} and we call remove(1). If we
224 // were to remove 1 again, we would not find it inside `key_map` anymore.
225 // If we were to imply from this that no re-mapping was necessary, we would
226 // incorrectly remove 1 from `self.slab.inner`, which corresponds to the
227 // handed-out key 2.
228 pub(crate) fn remove(&mut self, key: &Key) -> Data<T> {
229 let remapped_key = if self.compact_called {
230 match self.key_map.remove(key) {
231 Some(key_internal) => key_internal,
232 None => panic!("invalid key"),
233 }
234 } else {
235 (*key).into()
236 };
237
238 self.inner.remove(remapped_key.index)
239 }
240
241 pub(crate) fn shrink_to_fit(&mut self) {
242 self.inner.shrink_to_fit();
243 self.key_map.shrink_to_fit();
244 }
245
246 pub(crate) fn compact(&mut self) {
247 if !self.compact_called {
248 for (key, _) in self.inner.iter() {
249 self.key_map.insert(Key::new(key), KeyInternal::new(key));
250 }
251 }
252
253 let mut remapping = HashMap::new();
254 self.inner.compact(|_, from, to| {
255 remapping.insert(from, to);
256 true
257 });
258
259 // At this point `key_map` contains a mapping for every element.
260 for internal_key in self.key_map.values_mut() {
261 if let Some(new_internal_key) = remapping.get(&internal_key.index) {
262 *internal_key = KeyInternal::new(*new_internal_key);
263 }
264 }
265
266 if self.key_map.capacity() > 2 * self.key_map.len() {
267 self.key_map.shrink_to_fit();
268 }
269
270 self.compact_called = true;
271 }
272
273 // Tries to re-map a `Key` that was given out to the user to its
274 // corresponding internal key.
275 fn remap_key(&self, key: &Key) -> Option<KeyInternal> {
276 let key_map = &self.key_map;
277 if self.compact_called {
278 key_map.get(key).copied()
279 } else {
280 Some((*key).into())
281 }
282 }
283
284 fn create_new_key(&mut self) -> KeyInternal {
285 while self.key_map.contains_key(&Key::new(self.next_key_index)) {
286 self.next_key_index = self.next_key_index.wrapping_add(1);
287 }
288
289 KeyInternal::new(self.next_key_index)
290 }
291
292 pub(crate) fn len(&self) -> usize {
293 self.inner.len()
294 }
295
296 pub(crate) fn capacity(&self) -> usize {
297 self.inner.capacity()
298 }
299
300 pub(crate) fn clear(&mut self) {
301 self.inner.clear();
302 self.key_map.clear();
303 self.compact_called = false;
304 }
305
306 pub(crate) fn reserve(&mut self, additional: usize) {
307 self.inner.reserve(additional);
308
309 if self.compact_called {
310 self.key_map.reserve(additional);
311 }
312 }
313
314 pub(crate) fn is_empty(&self) -> bool {
315 self.inner.is_empty()
316 }
317
318 pub(crate) fn contains(&self, key: &Key) -> bool {
319 let remapped_key = self.remap_key(key);
320
321 match remapped_key {
322 Some(internal_key) => self.inner.contains(internal_key.index),
323 None => false,
324 }
325 }
326}
327
328impl<T> fmt::Debug for SlabStorage<T>
329where
330 T: fmt::Debug,
331{
332 fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
333 if fmt.alternate() {
334 fmt.debug_map().entries(self.inner.iter()).finish()
335 } else {
336 fmt.debug_struct("Slab")
337 .field("len", &self.len())
338 .field("cap", &self.capacity())
339 .finish()
340 }
341 }
342}
343
344impl<T> Index<Key> for SlabStorage<T> {
345 type Output = Data<T>;
346
347 fn index(&self, key: Key) -> &Self::Output {
348 let remapped_key = self.remap_key(&key);
349
350 match remapped_key {
351 Some(internal_key) => &self.inner[internal_key.index],
352 None => panic!("Invalid index {}", key.index),
353 }
354 }
355}
356
357impl<T> IndexMut<Key> for SlabStorage<T> {
358 fn index_mut(&mut self, key: Key) -> &mut Data<T> {
359 let remapped_key = self.remap_key(&key);
360
361 match remapped_key {
362 Some(internal_key) => &mut self.inner[internal_key.index],
363 None => panic!("Invalid index {}", key.index),
364 }
365 }
366}
367
368/// An entry in `DelayQueue` that has expired and been removed.
369///
370/// Values are returned by [`DelayQueue::poll_expired`].
371///
372/// [`DelayQueue::poll_expired`]: method@DelayQueue::poll_expired
373#[derive(Debug)]
374pub struct Expired<T> {
375 /// The data stored in the queue
376 data: T,
377
378 /// The expiration time
379 deadline: Instant,
380
381 /// The key associated with the entry
382 key: Key,
383}
384
385/// Token to a value stored in a `DelayQueue`.
386///
387/// Instances of `Key` are returned by [`DelayQueue::insert`]. See [`DelayQueue`]
388/// documentation for more details.
389///
390/// [`DelayQueue`]: struct@DelayQueue
391/// [`DelayQueue::insert`]: method@DelayQueue::insert
392#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
393pub struct Key {
394 index: usize,
395}
396
397// Whereas `Key` is given out to users that use `DelayQueue`, internally we use
398// `KeyInternal` as the key type in order to make the logic of mapping between keys
399// as a result of `compact` calls clearer.
400#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
401struct KeyInternal {
402 index: usize,
403}
404
405#[derive(Debug)]
406struct Stack<T> {
407 /// Head of the stack
408 head: Option<Key>,
409 _p: PhantomData<fn() -> T>,
410}
411
412#[derive(Debug)]
413struct Data<T> {
414 /// The data being stored in the queue and will be returned at the requested
415 /// instant.
416 inner: T,
417
418 /// The instant at which the item is returned.
419 when: u64,
420
421 /// Set to true when stored in the `expired` queue
422 expired: bool,
423
424 /// Next entry in the stack
425 next: Option<Key>,
426
427 /// Previous entry in the stack
428 prev: Option<Key>,
429}
430
431/// Maximum number of entries the queue can handle
432const MAX_ENTRIES: usize = (1 << 30) - 1;
433
434impl<T> DelayQueue<T> {
435 /// Creates a new, empty, `DelayQueue`.
436 ///
437 /// The queue will not allocate storage until items are inserted into it.
438 ///
439 /// # Examples
440 ///
441 /// ```rust
442 /// # use tokio_util::time::DelayQueue;
443 /// let delay_queue: DelayQueue<u32> = DelayQueue::new();
444 /// ```
445 pub fn new() -> DelayQueue<T> {
446 DelayQueue::with_capacity(0)
447 }
448
449 /// Creates a new, empty, `DelayQueue` with the specified capacity.
450 ///
451 /// The queue will be able to hold at least `capacity` elements without
452 /// reallocating. If `capacity` is 0, the queue will not allocate for
453 /// storage.
454 ///
455 /// # Examples
456 ///
457 /// ```rust
458 /// # use tokio_util::time::DelayQueue;
459 /// # use std::time::Duration;
460 ///
461 /// # #[tokio::main]
462 /// # async fn main() {
463 /// let mut delay_queue = DelayQueue::with_capacity(10);
464 ///
465 /// // These insertions are done without further allocation
466 /// for i in 0..10 {
467 /// delay_queue.insert(i, Duration::from_secs(i));
468 /// }
469 ///
470 /// // This will make the queue allocate additional storage
471 /// delay_queue.insert(11, Duration::from_secs(11));
472 /// # }
473 /// ```
474 pub fn with_capacity(capacity: usize) -> DelayQueue<T> {
475 DelayQueue {
476 wheel: Wheel::new(),
477 slab: SlabStorage::with_capacity(capacity),
478 expired: Stack::default(),
479 delay: None,
480 wheel_now: 0,
481 start: Instant::now(),
482 waker: None,
483 }
484 }
485
486 /// Inserts `value` into the queue set to expire at a specific instant in
487 /// time.
488 ///
489 /// This function is identical to `insert`, but takes an `Instant` instead
490 /// of a `Duration`.
491 ///
492 /// `value` is stored in the queue until `when` is reached. At which point,
493 /// `value` will be returned from [`poll_expired`]. If `when` has already been
494 /// reached, then `value` is immediately made available to poll.
495 ///
496 /// The return value represents the insertion and is used as an argument to
497 /// [`remove`] and [`reset`]. Note that [`Key`] is a token and is reused once
498 /// `value` is removed from the queue either by calling [`poll_expired`] after
499 /// `when` is reached or by calling [`remove`]. At this point, the caller
500 /// must take care to not use the returned [`Key`] again as it may reference
501 /// a different item in the queue.
502 ///
503 /// See [type] level documentation for more details.
504 ///
505 /// # Panics
506 ///
507 /// This function panics if `when` is too far in the future.
508 ///
509 /// # Examples
510 ///
511 /// Basic usage
512 ///
513 /// ```rust
514 /// use tokio::time::{Duration, Instant};
515 /// use tokio_util::time::DelayQueue;
516 ///
517 /// # #[tokio::main]
518 /// # async fn main() {
519 /// let mut delay_queue = DelayQueue::new();
520 /// let key = delay_queue.insert_at(
521 /// "foo", Instant::now() + Duration::from_secs(5));
522 ///
523 /// // Remove the entry
524 /// let item = delay_queue.remove(&key);
525 /// assert_eq!(*item.get_ref(), "foo");
526 /// # }
527 /// ```
528 ///
529 /// [`poll_expired`]: method@Self::poll_expired
530 /// [`remove`]: method@Self::remove
531 /// [`reset`]: method@Self::reset
532 /// [`Key`]: struct@Key
533 /// [type]: #
534 #[track_caller]
535 pub fn insert_at(&mut self, value: T, when: Instant) -> Key {
536 assert!(self.slab.len() < MAX_ENTRIES, "max entries exceeded");
537
538 // Normalize the deadline. Values cannot be set to expire in the past.
539 let when = self.normalize_deadline(when);
540
541 // Insert the value in the store
542 let key = self.slab.insert(Data {
543 inner: value,
544 when,
545 expired: false,
546 next: None,
547 prev: None,
548 });
549
550 self.insert_idx(when, key);
551
552 // Set a new delay if the current's deadline is later than the one of the new item
553 let should_set_delay = if let Some(ref delay) = self.delay {
554 let current_exp = self.normalize_deadline(delay.deadline());
555 current_exp > when
556 } else {
557 true
558 };
559
560 if should_set_delay {
561 if let Some(waker) = self.waker.take() {
562 waker.wake();
563 }
564
565 let delay_time = self.start + Duration::from_millis(when);
566 if let Some(ref mut delay) = &mut self.delay {
567 delay.as_mut().reset(delay_time);
568 } else {
569 self.delay = Some(Box::pin(sleep_until(delay_time)));
570 }
571 }
572
573 key
574 }
575
576 /// Attempts to pull out the next value of the delay queue, registering the
577 /// current task for wakeup if the value is not yet available, and returning
578 /// `None` if the queue is exhausted.
579 pub fn poll_expired(&mut self, cx: &mut task::Context<'_>) -> Poll<Option<Expired<T>>> {
580 if !self
581 .waker
582 .as_ref()
583 .map(|w| w.will_wake(cx.waker()))
584 .unwrap_or(false)
585 {
586 self.waker = Some(cx.waker().clone());
587 }
588
589 let item = ready!(self.poll_idx(cx));
590 Poll::Ready(item.map(|key| {
591 let data = self.slab.remove(&key);
592 debug_assert!(data.next.is_none());
593 debug_assert!(data.prev.is_none());
594
595 Expired {
596 key,
597 data: data.inner,
598 deadline: self.start + Duration::from_millis(data.when),
599 }
600 }))
601 }
602
603 /// Inserts `value` into the queue set to expire after the requested duration
604 /// elapses.
605 ///
606 /// This function is identical to `insert_at`, but takes a `Duration`
607 /// instead of an `Instant`.
608 ///
609 /// `value` is stored in the queue until `timeout` duration has
610 /// elapsed after `insert` was called. At that point, `value` will
611 /// be returned from [`poll_expired`]. If `timeout` is a `Duration` of
612 /// zero, then `value` is immediately made available to poll.
613 ///
614 /// The return value represents the insertion and is used as an
615 /// argument to [`remove`] and [`reset`]. Note that [`Key`] is a
616 /// token and is reused once `value` is removed from the queue
617 /// either by calling [`poll_expired`] after `timeout` has elapsed
618 /// or by calling [`remove`]. At this point, the caller must not
619 /// use the returned [`Key`] again as it may reference a different
620 /// item in the queue.
621 ///
622 /// See [type] level documentation for more details.
623 ///
624 /// # Panics
625 ///
626 /// This function panics if `timeout` is greater than the maximum
627 /// duration supported by the timer in the current `Runtime`.
628 ///
629 /// # Examples
630 ///
631 /// Basic usage
632 ///
633 /// ```rust
634 /// use tokio_util::time::DelayQueue;
635 /// use std::time::Duration;
636 ///
637 /// # #[tokio::main]
638 /// # async fn main() {
639 /// let mut delay_queue = DelayQueue::new();
640 /// let key = delay_queue.insert("foo", Duration::from_secs(5));
641 ///
642 /// // Remove the entry
643 /// let item = delay_queue.remove(&key);
644 /// assert_eq!(*item.get_ref(), "foo");
645 /// # }
646 /// ```
647 ///
648 /// [`poll_expired`]: method@Self::poll_expired
649 /// [`remove`]: method@Self::remove
650 /// [`reset`]: method@Self::reset
651 /// [`Key`]: struct@Key
652 /// [type]: #
653 #[track_caller]
654 pub fn insert(&mut self, value: T, timeout: Duration) -> Key {
655 self.insert_at(value, Instant::now() + timeout)
656 }
657
658 #[track_caller]
659 fn insert_idx(&mut self, when: u64, key: Key) {
660 use self::wheel::{InsertError, Stack};
661
662 // Register the deadline with the timer wheel
663 match self.wheel.insert(when, key, &mut self.slab) {
664 Ok(_) => {}
665 Err((_, InsertError::Elapsed)) => {
666 self.slab[key].expired = true;
667 // The delay is already expired, store it in the expired queue
668 self.expired.push(key, &mut self.slab);
669 }
670 Err((_, err)) => panic!("invalid deadline; err={:?}", err),
671 }
672 }
673
674 /// Returns the deadline of the item associated with `key`.
675 ///
676 /// Since the queue operates at millisecond granularity, the returned
677 /// deadline may not exactly match the value that was given when initially
678 /// inserting the item into the queue.
679 ///
680 /// # Panics
681 ///
682 /// This function panics if `key` is not contained by the queue.
683 ///
684 /// # Examples
685 ///
686 /// Basic usage
687 ///
688 /// ```rust
689 /// use tokio_util::time::DelayQueue;
690 /// use std::time::Duration;
691 ///
692 /// # #[tokio::main]
693 /// # async fn main() {
694 /// let mut delay_queue = DelayQueue::new();
695 ///
696 /// let key1 = delay_queue.insert("foo", Duration::from_secs(5));
697 /// let key2 = delay_queue.insert("bar", Duration::from_secs(10));
698 ///
699 /// assert!(delay_queue.deadline(&key1) < delay_queue.deadline(&key2));
700 /// # }
701 /// ```
702 #[track_caller]
703 pub fn deadline(&self, key: &Key) -> Instant {
704 self.start + Duration::from_millis(self.slab[*key].when)
705 }
706
707 /// Removes the key from the expired queue or the timer wheel
708 /// depending on its expiration status.
709 ///
710 /// # Panics
711 ///
712 /// Panics if the key is not contained in the expired queue or the wheel.
713 #[track_caller]
714 fn remove_key(&mut self, key: &Key) {
715 use crate::time::wheel::Stack;
716
717 // Special case the `expired` queue
718 if self.slab[*key].expired {
719 self.expired.remove(key, &mut self.slab);
720 } else {
721 self.wheel.remove(key, &mut self.slab);
722 }
723 }
724
725 /// Removes the item associated with `key` from the queue.
726 ///
727 /// There must be an item associated with `key`. The function returns the
728 /// removed item as well as the `Instant` at which it will the delay will
729 /// have expired.
730 ///
731 /// # Panics
732 ///
733 /// The function panics if `key` is not contained by the queue.
734 ///
735 /// # Examples
736 ///
737 /// Basic usage
738 ///
739 /// ```rust
740 /// use tokio_util::time::DelayQueue;
741 /// use std::time::Duration;
742 ///
743 /// # #[tokio::main]
744 /// # async fn main() {
745 /// let mut delay_queue = DelayQueue::new();
746 /// let key = delay_queue.insert("foo", Duration::from_secs(5));
747 ///
748 /// // Remove the entry
749 /// let item = delay_queue.remove(&key);
750 /// assert_eq!(*item.get_ref(), "foo");
751 /// # }
752 /// ```
753 #[track_caller]
754 pub fn remove(&mut self, key: &Key) -> Expired<T> {
755 let prev_deadline = self.next_deadline();
756
757 self.remove_key(key);
758 let data = self.slab.remove(key);
759
760 let next_deadline = self.next_deadline();
761 if prev_deadline != next_deadline {
762 match (next_deadline, &mut self.delay) {
763 (None, _) => self.delay = None,
764 (Some(deadline), Some(delay)) => delay.as_mut().reset(deadline),
765 (Some(deadline), None) => self.delay = Some(Box::pin(sleep_until(deadline))),
766 }
767 }
768
769 Expired {
770 key: Key::new(key.index),
771 data: data.inner,
772 deadline: self.start + Duration::from_millis(data.when),
773 }
774 }
775
776 /// Attempts to remove the item associated with `key` from the queue.
777 ///
778 /// Removes the item associated with `key`, and returns it along with the
779 /// `Instant` at which it would have expired, if it exists.
780 ///
781 /// Returns `None` if `key` is not in the queue.
782 ///
783 /// # Examples
784 ///
785 /// Basic usage
786 ///
787 /// ```rust
788 /// use tokio_util::time::DelayQueue;
789 /// use std::time::Duration;
790 ///
791 /// # #[tokio::main(flavor = "current_thread")]
792 /// # async fn main() {
793 /// let mut delay_queue = DelayQueue::new();
794 /// let key = delay_queue.insert("foo", Duration::from_secs(5));
795 ///
796 /// // The item is in the queue, `try_remove` returns `Some(Expired("foo"))`.
797 /// let item = delay_queue.try_remove(&key);
798 /// assert_eq!(item.unwrap().into_inner(), "foo");
799 ///
800 /// // The item is not in the queue anymore, `try_remove` returns `None`.
801 /// let item = delay_queue.try_remove(&key);
802 /// assert!(item.is_none());
803 /// # }
804 /// ```
805 pub fn try_remove(&mut self, key: &Key) -> Option<Expired<T>> {
806 if self.slab.contains(key) {
807 Some(self.remove(key))
808 } else {
809 None
810 }
811 }
812
813 /// Sets the delay of the item associated with `key` to expire at `when`.
814 ///
815 /// This function is identical to `reset` but takes an `Instant` instead of
816 /// a `Duration`.
817 ///
818 /// The item remains in the queue but the delay is set to expire at `when`.
819 /// If `when` is in the past, then the item is immediately made available to
820 /// the caller.
821 ///
822 /// # Panics
823 ///
824 /// This function panics if `when` is too far in the future or if `key` is
825 /// not contained by the queue.
826 ///
827 /// # Examples
828 ///
829 /// Basic usage
830 ///
831 /// ```rust
832 /// use tokio::time::{Duration, Instant};
833 /// use tokio_util::time::DelayQueue;
834 ///
835 /// # #[tokio::main]
836 /// # async fn main() {
837 /// let mut delay_queue = DelayQueue::new();
838 /// let key = delay_queue.insert("foo", Duration::from_secs(5));
839 ///
840 /// // "foo" is scheduled to be returned in 5 seconds
841 ///
842 /// delay_queue.reset_at(&key, Instant::now() + Duration::from_secs(10));
843 ///
844 /// // "foo" is now scheduled to be returned in 10 seconds
845 /// # }
846 /// ```
847 #[track_caller]
848 pub fn reset_at(&mut self, key: &Key, when: Instant) {
849 self.remove_key(key);
850
851 // Normalize the deadline. Values cannot be set to expire in the past.
852 let when = self.normalize_deadline(when);
853
854 self.slab[*key].when = when;
855 self.slab[*key].expired = false;
856
857 self.insert_idx(when, *key);
858
859 let next_deadline = self.next_deadline();
860 if let (Some(ref mut delay), Some(deadline)) = (&mut self.delay, next_deadline) {
861 // This should awaken us if necessary (ie, if already expired)
862 delay.as_mut().reset(deadline);
863 }
864 }
865
866 /// Shrink the capacity of the slab, which `DelayQueue` uses internally for storage allocation.
867 /// This function is not guaranteed to, and in most cases, won't decrease the capacity of the slab
868 /// to the number of elements still contained in it, because elements cannot be moved to a different
869 /// index. To decrease the capacity to the size of the slab use [`compact`].
870 ///
871 /// This function can take O(n) time even when the capacity cannot be reduced or the allocation is
872 /// shrunk in place. Repeated calls run in O(1) though.
873 ///
874 /// [`compact`]: method@Self::compact
875 pub fn shrink_to_fit(&mut self) {
876 self.slab.shrink_to_fit();
877 }
878
879 /// Shrink the capacity of the slab, which `DelayQueue` uses internally for storage allocation,
880 /// to the number of elements that are contained in it.
881 ///
882 /// This methods runs in O(n).
883 ///
884 /// # Examples
885 ///
886 /// Basic usage
887 ///
888 /// ```rust
889 /// use tokio_util::time::DelayQueue;
890 /// use std::time::Duration;
891 ///
892 /// # #[tokio::main]
893 /// # async fn main() {
894 /// let mut delay_queue = DelayQueue::with_capacity(10);
895 ///
896 /// let key1 = delay_queue.insert(5, Duration::from_secs(5));
897 /// let key2 = delay_queue.insert(10, Duration::from_secs(10));
898 /// let key3 = delay_queue.insert(15, Duration::from_secs(15));
899 ///
900 /// delay_queue.remove(&key2);
901 ///
902 /// delay_queue.compact();
903 /// assert_eq!(delay_queue.capacity(), 2);
904 /// # }
905 /// ```
906 pub fn compact(&mut self) {
907 self.slab.compact();
908 }
909
910 /// Gets the [`Key`] that [`poll_expired`] will pull out of the queue next, without
911 /// pulling it out or waiting for the deadline to expire.
912 ///
913 /// Entries that have already expired may be returned in any order, but it is
914 /// guaranteed that this method returns them in the same order as when items
915 /// are popped from the `DelayQueue`.
916 ///
917 /// # Examples
918 ///
919 /// Basic usage
920 ///
921 /// ```rust
922 /// use tokio_util::time::DelayQueue;
923 /// use std::time::Duration;
924 ///
925 /// # #[tokio::main]
926 /// # async fn main() {
927 /// let mut delay_queue = DelayQueue::new();
928 ///
929 /// let key1 = delay_queue.insert("foo", Duration::from_secs(10));
930 /// let key2 = delay_queue.insert("bar", Duration::from_secs(5));
931 /// let key3 = delay_queue.insert("baz", Duration::from_secs(15));
932 ///
933 /// assert_eq!(delay_queue.peek().unwrap(), key2);
934 /// # }
935 /// ```
936 ///
937 /// [`Key`]: struct@Key
938 /// [`poll_expired`]: method@Self::poll_expired
939 pub fn peek(&self) -> Option<Key> {
940 use self::wheel::Stack;
941
942 self.expired.peek().or_else(|| self.wheel.peek())
943 }
944
945 /// Returns the next time to poll as determined by the wheel.
946 ///
947 /// Note that this does not include deadlines in the `expired` queue.
948 fn next_deadline(&self) -> Option<Instant> {
949 self.wheel
950 .poll_at()
951 .map(|poll_at| self.start + Duration::from_millis(poll_at))
952 }
953
954 /// Sets the delay of the item associated with `key` to expire after
955 /// `timeout`.
956 ///
957 /// This function is identical to `reset_at` but takes a `Duration` instead
958 /// of an `Instant`.
959 ///
960 /// The item remains in the queue but the delay is set to expire after
961 /// `timeout`. If `timeout` is zero, then the item is immediately made
962 /// available to the caller.
963 ///
964 /// # Panics
965 ///
966 /// This function panics if `timeout` is greater than the maximum supported
967 /// duration or if `key` is not contained by the queue.
968 ///
969 /// # Examples
970 ///
971 /// Basic usage
972 ///
973 /// ```rust
974 /// use tokio_util::time::DelayQueue;
975 /// use std::time::Duration;
976 ///
977 /// # #[tokio::main]
978 /// # async fn main() {
979 /// let mut delay_queue = DelayQueue::new();
980 /// let key = delay_queue.insert("foo", Duration::from_secs(5));
981 ///
982 /// // "foo" is scheduled to be returned in 5 seconds
983 ///
984 /// delay_queue.reset(&key, Duration::from_secs(10));
985 ///
986 /// // "foo"is now scheduled to be returned in 10 seconds
987 /// # }
988 /// ```
989 #[track_caller]
990 pub fn reset(&mut self, key: &Key, timeout: Duration) {
991 self.reset_at(key, Instant::now() + timeout);
992 }
993
994 /// Clears the queue, removing all items.
995 ///
996 /// After calling `clear`, [`poll_expired`] will return `Ok(Ready(None))`.
997 ///
998 /// Note that this method has no effect on the allocated capacity.
999 ///
1000 /// [`poll_expired`]: method@Self::poll_expired
1001 ///
1002 /// # Examples
1003 ///
1004 /// ```rust
1005 /// use tokio_util::time::DelayQueue;
1006 /// use std::time::Duration;
1007 ///
1008 /// # #[tokio::main]
1009 /// # async fn main() {
1010 /// let mut delay_queue = DelayQueue::new();
1011 ///
1012 /// delay_queue.insert("foo", Duration::from_secs(5));
1013 ///
1014 /// assert!(!delay_queue.is_empty());
1015 ///
1016 /// delay_queue.clear();
1017 ///
1018 /// assert!(delay_queue.is_empty());
1019 /// # }
1020 /// ```
1021 pub fn clear(&mut self) {
1022 self.slab.clear();
1023 self.expired = Stack::default();
1024 self.wheel = Wheel::new();
1025 self.delay = None;
1026 }
1027
1028 /// Returns the number of elements the queue can hold without reallocating.
1029 ///
1030 /// # Examples
1031 ///
1032 /// ```rust
1033 /// use tokio_util::time::DelayQueue;
1034 ///
1035 /// let delay_queue: DelayQueue<i32> = DelayQueue::with_capacity(10);
1036 /// assert_eq!(delay_queue.capacity(), 10);
1037 /// ```
1038 pub fn capacity(&self) -> usize {
1039 self.slab.capacity()
1040 }
1041
1042 /// Returns the number of elements currently in the queue.
1043 ///
1044 /// # Examples
1045 ///
1046 /// ```rust
1047 /// use tokio_util::time::DelayQueue;
1048 /// use std::time::Duration;
1049 ///
1050 /// # #[tokio::main]
1051 /// # async fn main() {
1052 /// let mut delay_queue: DelayQueue<i32> = DelayQueue::with_capacity(10);
1053 /// assert_eq!(delay_queue.len(), 0);
1054 /// delay_queue.insert(3, Duration::from_secs(5));
1055 /// assert_eq!(delay_queue.len(), 1);
1056 /// # }
1057 /// ```
1058 pub fn len(&self) -> usize {
1059 self.slab.len()
1060 }
1061
1062 /// Reserves capacity for at least `additional` more items to be queued
1063 /// without allocating.
1064 ///
1065 /// `reserve` does nothing if the queue already has sufficient capacity for
1066 /// `additional` more values. If more capacity is required, a new segment of
1067 /// memory will be allocated and all existing values will be copied into it.
1068 /// As such, if the queue is already very large, a call to `reserve` can end
1069 /// up being expensive.
1070 ///
1071 /// The queue may reserve more than `additional` extra space in order to
1072 /// avoid frequent reallocations.
1073 ///
1074 /// # Panics
1075 ///
1076 /// Panics if the new capacity exceeds the maximum number of entries the
1077 /// queue can contain.
1078 ///
1079 /// # Examples
1080 ///
1081 /// ```
1082 /// use tokio_util::time::DelayQueue;
1083 /// use std::time::Duration;
1084 ///
1085 /// # #[tokio::main]
1086 /// # async fn main() {
1087 /// let mut delay_queue = DelayQueue::new();
1088 ///
1089 /// delay_queue.insert("hello", Duration::from_secs(10));
1090 /// delay_queue.reserve(10);
1091 ///
1092 /// assert!(delay_queue.capacity() >= 11);
1093 /// # }
1094 /// ```
1095 #[track_caller]
1096 pub fn reserve(&mut self, additional: usize) {
1097 assert!(
1098 self.slab.capacity() + additional <= MAX_ENTRIES,
1099 "max queue capacity exceeded"
1100 );
1101 self.slab.reserve(additional);
1102 }
1103
1104 /// Returns `true` if there are no items in the queue.
1105 ///
1106 /// Note that this function returns `false` even if all items have not yet
1107 /// expired and a call to `poll` will return `Poll::Pending`.
1108 ///
1109 /// # Examples
1110 ///
1111 /// ```
1112 /// use tokio_util::time::DelayQueue;
1113 /// use std::time::Duration;
1114 ///
1115 /// # #[tokio::main]
1116 /// # async fn main() {
1117 /// let mut delay_queue = DelayQueue::new();
1118 /// assert!(delay_queue.is_empty());
1119 ///
1120 /// delay_queue.insert("hello", Duration::from_secs(5));
1121 /// assert!(!delay_queue.is_empty());
1122 /// # }
1123 /// ```
1124 pub fn is_empty(&self) -> bool {
1125 self.slab.is_empty()
1126 }
1127
1128 /// Polls the queue, returning the index of the next slot in the slab that
1129 /// should be returned.
1130 ///
1131 /// A slot should be returned when the associated deadline has been reached.
1132 fn poll_idx(&mut self, cx: &mut task::Context<'_>) -> Poll<Option<Key>> {
1133 use self::wheel::Stack;
1134
1135 let expired = self.expired.pop(&mut self.slab);
1136
1137 if expired.is_some() {
1138 return Poll::Ready(expired);
1139 }
1140
1141 loop {
1142 if let Some(ref mut delay) = self.delay {
1143 if !delay.is_elapsed() {
1144 ready!(Pin::new(&mut *delay).poll(cx));
1145 }
1146
1147 let now = crate::time::ms(delay.deadline() - self.start, crate::time::Round::Down);
1148
1149 self.wheel_now = now;
1150 }
1151
1152 // We poll the wheel to get the next value out before finding the next deadline.
1153 let wheel_idx = self.wheel.poll(self.wheel_now, &mut self.slab);
1154
1155 self.delay = self.next_deadline().map(|when| Box::pin(sleep_until(when)));
1156
1157 if let Some(idx) = wheel_idx {
1158 return Poll::Ready(Some(idx));
1159 }
1160
1161 if self.delay.is_none() {
1162 return Poll::Ready(None);
1163 }
1164 }
1165 }
1166
1167 fn normalize_deadline(&self, when: Instant) -> u64 {
1168 let when = if when < self.start {
1169 0
1170 } else {
1171 crate::time::ms(when - self.start, crate::time::Round::Up)
1172 };
1173
1174 cmp::max(when, self.wheel.elapsed())
1175 }
1176}
1177
1178// We never put `T` in a `Pin`...
1179impl<T> Unpin for DelayQueue<T> {}
1180
1181impl<T> Default for DelayQueue<T> {
1182 fn default() -> DelayQueue<T> {
1183 DelayQueue::new()
1184 }
1185}
1186
1187impl<T> futures_core::Stream for DelayQueue<T> {
1188 // DelayQueue seems much more specific, where a user may care that it
1189 // has reached capacity, so return those errors instead of panicking.
1190 type Item = Expired<T>;
1191
1192 fn poll_next(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Option<Self::Item>> {
1193 DelayQueue::poll_expired(self.get_mut(), cx)
1194 }
1195}
1196
1197impl<T> wheel::Stack for Stack<T> {
1198 type Owned = Key;
1199 type Borrowed = Key;
1200 type Store = SlabStorage<T>;
1201
1202 fn is_empty(&self) -> bool {
1203 self.head.is_none()
1204 }
1205
1206 fn push(&mut self, item: Self::Owned, store: &mut Self::Store) {
1207 // Ensure the entry is not already in a stack.
1208 debug_assert!(store[item].next.is_none());
1209 debug_assert!(store[item].prev.is_none());
1210
1211 // Remove the old head entry
1212 let old = self.head.take();
1213
1214 if let Some(idx) = old {
1215 store[idx].prev = Some(item);
1216 }
1217
1218 store[item].next = old;
1219 self.head = Some(item);
1220 }
1221
1222 fn pop(&mut self, store: &mut Self::Store) -> Option<Self::Owned> {
1223 if let Some(key) = self.head {
1224 self.head = store[key].next;
1225
1226 if let Some(idx) = self.head {
1227 store[idx].prev = None;
1228 }
1229
1230 store[key].next = None;
1231 debug_assert!(store[key].prev.is_none());
1232
1233 Some(key)
1234 } else {
1235 None
1236 }
1237 }
1238
1239 fn peek(&self) -> Option<Self::Owned> {
1240 self.head
1241 }
1242
1243 #[track_caller]
1244 fn remove(&mut self, item: &Self::Borrowed, store: &mut Self::Store) {
1245 let key = *item;
1246 assert!(store.contains(item));
1247
1248 // Ensure that the entry is in fact contained by the stack
1249 debug_assert!({
1250 // This walks the full linked list even if an entry is found.
1251 let mut next = self.head;
1252 let mut contains = false;
1253
1254 while let Some(idx) = next {
1255 let data = &store[idx];
1256
1257 if idx == *item {
1258 debug_assert!(!contains);
1259 contains = true;
1260 }
1261
1262 next = data.next;
1263 }
1264
1265 contains
1266 });
1267
1268 if let Some(next) = store[key].next {
1269 store[next].prev = store[key].prev;
1270 }
1271
1272 if let Some(prev) = store[key].prev {
1273 store[prev].next = store[key].next;
1274 } else {
1275 self.head = store[key].next;
1276 }
1277
1278 store[key].next = None;
1279 store[key].prev = None;
1280 }
1281
1282 fn when(item: &Self::Borrowed, store: &Self::Store) -> u64 {
1283 store[*item].when
1284 }
1285}
1286
1287impl<T> Default for Stack<T> {
1288 fn default() -> Stack<T> {
1289 Stack {
1290 head: None,
1291 _p: PhantomData,
1292 }
1293 }
1294}
1295
1296impl Key {
1297 pub(crate) fn new(index: usize) -> Key {
1298 Key { index }
1299 }
1300}
1301
1302impl KeyInternal {
1303 pub(crate) fn new(index: usize) -> KeyInternal {
1304 KeyInternal { index }
1305 }
1306}
1307
1308impl From<Key> for KeyInternal {
1309 fn from(item: Key) -> Self {
1310 KeyInternal::new(item.index)
1311 }
1312}
1313
1314impl From<KeyInternal> for Key {
1315 fn from(item: KeyInternal) -> Self {
1316 Key::new(item.index)
1317 }
1318}
1319
1320impl<T> Expired<T> {
1321 /// Returns a reference to the inner value.
1322 pub fn get_ref(&self) -> &T {
1323 &self.data
1324 }
1325
1326 /// Returns a mutable reference to the inner value.
1327 pub fn get_mut(&mut self) -> &mut T {
1328 &mut self.data
1329 }
1330
1331 /// Consumes `self` and returns the inner value.
1332 pub fn into_inner(self) -> T {
1333 self.data
1334 }
1335
1336 /// Returns the deadline that the expiration was set to.
1337 pub fn deadline(&self) -> Instant {
1338 self.deadline
1339 }
1340
1341 /// Returns the key that the expiration is indexed by.
1342 pub fn key(&self) -> Key {
1343 self.key
1344 }
1345}
1346