diff options
Diffstat (limited to 'third_party/rust/tokio/src/time')
21 files changed, 4183 insertions, 0 deletions
diff --git a/third_party/rust/tokio/src/time/clock.rs b/third_party/rust/tokio/src/time/clock.rs new file mode 100644 index 0000000000..4ac24af3d0 --- /dev/null +++ b/third_party/rust/tokio/src/time/clock.rs @@ -0,0 +1,164 @@ +//! Source of time abstraction. +//! +//! By default, `std::time::Instant::now()` is used. However, when the +//! `test-util` feature flag is enabled, the values returned for `now()` are +//! configurable. + +cfg_not_test_util! { + use crate::time::{Duration, Instant}; + + #[derive(Debug, Clone)] + pub(crate) struct Clock {} + + pub(crate) fn now() -> Instant { + Instant::from_std(std::time::Instant::now()) + } + + impl Clock { + pub(crate) fn new() -> Clock { + Clock {} + } + + pub(crate) fn now(&self) -> Instant { + now() + } + + pub(crate) fn is_paused(&self) -> bool { + false + } + + pub(crate) fn advance(&self, _dur: Duration) { + unreachable!(); + } + } +} + +cfg_test_util! { + use crate::time::{Duration, Instant}; + use std::sync::{Arc, Mutex}; + use crate::runtime::context; + + /// A handle to a source of time. + #[derive(Debug, Clone)] + pub(crate) struct Clock { + inner: Arc<Mutex<Inner>>, + } + + #[derive(Debug)] + struct Inner { + /// Instant to use as the clock's base instant. + base: std::time::Instant, + + /// Instant at which the clock was last unfrozen + unfrozen: Option<std::time::Instant>, + } + + /// Pause time + /// + /// The current value of `Instant::now()` is saved and all subsequent calls + /// to `Instant::now()` will return the saved value. This is useful for + /// running tests that are dependent on time. + /// + /// # Panics + /// + /// Panics if time is already frozen or if called from outside of the Tokio + /// runtime. + pub fn pause() { + let clock = context::clock().expect("time cannot be frozen from outside the Tokio runtime"); + clock.pause(); + } + + /// Resume time + /// + /// Clears the saved `Instant::now()` value. Subsequent calls to + /// `Instant::now()` will return the value returned by the system call. + /// + /// # Panics + /// + /// Panics if time is not frozen or if called from outside of the Tokio + /// runtime. + pub fn resume() { + let clock = context::clock().expect("time cannot be frozen from outside the Tokio runtime"); + let mut inner = clock.inner.lock().unwrap(); + + if inner.unfrozen.is_some() { + panic!("time is not frozen"); + } + + inner.unfrozen = Some(std::time::Instant::now()); + } + + /// Advance time + /// + /// Increments the saved `Instant::now()` value by `duration`. Subsequent + /// calls to `Instant::now()` will return the result of the increment. + /// + /// # Panics + /// + /// Panics if time is not frozen or if called from outside of the Tokio + /// runtime. + pub async fn advance(duration: Duration) { + let clock = context::clock().expect("time cannot be frozen from outside the Tokio runtime"); + clock.advance(duration); + crate::task::yield_now().await; + } + + /// Return the current instant, factoring in frozen time. + pub(crate) fn now() -> Instant { + if let Some(clock) = context::clock() { + clock.now() + } else { + Instant::from_std(std::time::Instant::now()) + } + } + + impl Clock { + /// Return a new `Clock` instance that uses the current execution context's + /// source of time. + pub(crate) fn new() -> Clock { + let now = std::time::Instant::now(); + + Clock { + inner: Arc::new(Mutex::new(Inner { + base: now, + unfrozen: Some(now), + })), + } + } + + pub(crate) fn pause(&self) { + let mut inner = self.inner.lock().unwrap(); + + let elapsed = inner.unfrozen.as_ref().expect("time is already frozen").elapsed(); + inner.base += elapsed; + inner.unfrozen = None; + } + + pub(crate) fn is_paused(&self) -> bool { + let inner = self.inner.lock().unwrap(); + inner.unfrozen.is_none() + } + + pub(crate) fn advance(&self, duration: Duration) { + let mut inner = self.inner.lock().unwrap(); + + if inner.unfrozen.is_some() { + panic!("time is not frozen"); + } + + inner.base += duration; + } + + pub(crate) fn now(&self) -> Instant { + let inner = self.inner.lock().unwrap(); + + let mut ret = inner.base; + + if let Some(unfrozen) = inner.unfrozen { + ret += unfrozen.elapsed(); + } + + Instant::from_std(ret) + } + } +} diff --git a/third_party/rust/tokio/src/time/delay.rs b/third_party/rust/tokio/src/time/delay.rs new file mode 100644 index 0000000000..8088c9955c --- /dev/null +++ b/third_party/rust/tokio/src/time/delay.rs @@ -0,0 +1,99 @@ +use crate::time::driver::Registration; +use crate::time::{Duration, Instant}; + +use std::future::Future; +use std::pin::Pin; +use std::task::{self, Poll}; + +/// Waits until `deadline` is reached. +/// +/// No work is performed while awaiting on the delay to complete. The delay +/// operates at millisecond granularity and should not be used for tasks that +/// require high-resolution timers. +/// +/// # Cancellation +/// +/// Canceling a delay is done by dropping the returned future. No additional +/// cleanup work is required. +pub fn delay_until(deadline: Instant) -> Delay { + let registration = Registration::new(deadline, Duration::from_millis(0)); + Delay { registration } +} + +/// Waits until `duration` has elapsed. +/// +/// Equivalent to `delay_until(Instant::now() + duration)`. An asynchronous +/// analog to `std::thread::sleep`. +/// +/// No work is performed while awaiting on the delay to complete. The delay +/// operates at millisecond granularity and should not be used for tasks that +/// require high-resolution timers. +/// +/// # Cancellation +/// +/// Canceling a delay is done by dropping the returned future. No additional +/// cleanup work is required. +pub fn delay_for(duration: Duration) -> Delay { + delay_until(Instant::now() + duration) +} + +/// Future returned by [`delay_until`](delay_until) and +/// [`delay_for`](delay_for). +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct Delay { + /// The link between the `Delay` instance and the timer that drives it. + /// + /// This also stores the `deadline` value. + registration: Registration, +} + +impl Delay { + pub(crate) fn new_timeout(deadline: Instant, duration: Duration) -> Delay { + let registration = Registration::new(deadline, duration); + Delay { registration } + } + + /// Returns the instant at which the future will complete. + pub fn deadline(&self) -> Instant { + self.registration.deadline() + } + + /// Returns `true` if the `Delay` has elapsed + /// + /// A `Delay` is elapsed when the requested duration has elapsed. + pub fn is_elapsed(&self) -> bool { + self.registration.is_elapsed() + } + + /// Resets the `Delay` instance to a new deadline. + /// + /// Calling this function allows changing the instant at which the `Delay` + /// future completes without having to create new associated state. + /// + /// This function can be called both before and after the future has + /// completed. + pub fn reset(&mut self, deadline: Instant) { + self.registration.reset(deadline); + } +} + +impl Future for Delay { + type Output = (); + + fn poll(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> { + // `poll_elapsed` can return an error in two cases: + // + // - AtCapacity: this is a pathlogical case where far too many + // delays have been scheduled. + // - Shutdown: No timer has been setup, which is a mis-use error. + // + // Both cases are extremely rare, and pretty accurately fit into + // "logic errors", so we just panic in this case. A user couldn't + // really do much better if we passed the error onwards. + match ready!(self.registration.poll_elapsed(cx)) { + Ok(()) => Poll::Ready(()), + Err(e) => panic!("timer error: {}", e), + } + } +} diff --git a/third_party/rust/tokio/src/time/delay_queue.rs b/third_party/rust/tokio/src/time/delay_queue.rs new file mode 100644 index 0000000000..59f901a95d --- /dev/null +++ b/third_party/rust/tokio/src/time/delay_queue.rs @@ -0,0 +1,887 @@ +//! A queue of delayed elements. +//! +//! See [`DelayQueue`] for more details. +//! +//! [`DelayQueue`]: struct@DelayQueue + +use crate::time::wheel::{self, Wheel}; +use crate::time::{delay_until, Delay, Duration, Error, Instant}; + +use slab::Slab; +use std::cmp; +use std::future::Future; +use std::marker::PhantomData; +use std::pin::Pin; +use std::task::{self, Poll}; + +/// A queue of delayed elements. +/// +/// Once an element is inserted into the `DelayQueue`, it is yielded once the +/// specified deadline has been reached. +/// +/// # Usage +/// +/// Elements are inserted into `DelayQueue` using the [`insert`] or +/// [`insert_at`] methods. A deadline is provided with the item and a [`Key`] is +/// returned. The key is used to remove the entry or to change the deadline at +/// which it should be yielded back. +/// +/// Once delays have been configured, the `DelayQueue` is used via its +/// [`Stream`] implementation. [`poll`] is called. If an entry has reached its +/// deadline, it is returned. If not, `Poll::Pending` indicating that the +/// current task will be notified once the deadline has been reached. +/// +/// # `Stream` implementation +/// +/// Items are retrieved from the queue via [`Stream::poll`]. If no delays have +/// expired, no items are returned. In this case, `NotReady` is returned and the +/// current task is registered to be notified once the next item's delay has +/// expired. +/// +/// If no items are in the queue, i.e. `is_empty()` returns `true`, then `poll` +/// returns `Ready(None)`. This indicates that the stream has reached an end. +/// However, if a new item is inserted *after*, `poll` will once again start +/// returning items or `NotReady. +/// +/// Items are returned ordered by their expirations. Items that are configured +/// to expire first will be returned first. There are no ordering guarantees +/// for items configured to expire the same instant. Also note that delays are +/// rounded to the closest millisecond. +/// +/// # Implementation +/// +/// The `DelayQueue` is backed by the same hashed timing wheel implementation as +/// [`Timer`] as such, it offers the same performance benefits. See [`Timer`] +/// for further implementation notes. +/// +/// State associated with each entry is stored in a [`slab`]. This allows +/// amortizing the cost of allocation. Space created for expired entries is +/// reused when inserting new entries. +/// +/// Capacity can be checked using [`capacity`] and allocated preemptively by using +/// the [`reserve`] method. +/// +/// # Usage +/// +/// Using `DelayQueue` to manage cache entries. +/// +/// ```rust,no_run +/// use tokio::time::{delay_queue, DelayQueue, Error}; +/// +/// use futures::ready; +/// use std::collections::HashMap; +/// use std::task::{Context, Poll}; +/// use std::time::Duration; +/// # type CacheKey = String; +/// # type Value = String; +/// +/// struct Cache { +/// entries: HashMap<CacheKey, (Value, delay_queue::Key)>, +/// expirations: DelayQueue<CacheKey>, +/// } +/// +/// const TTL_SECS: u64 = 30; +/// +/// impl Cache { +/// fn insert(&mut self, key: CacheKey, value: Value) { +/// let delay = self.expirations +/// .insert(key.clone(), Duration::from_secs(TTL_SECS)); +/// +/// self.entries.insert(key, (value, delay)); +/// } +/// +/// fn get(&self, key: &CacheKey) -> Option<&Value> { +/// self.entries.get(key) +/// .map(|&(ref v, _)| v) +/// } +/// +/// fn remove(&mut self, key: &CacheKey) { +/// if let Some((_, cache_key)) = self.entries.remove(key) { +/// self.expirations.remove(&cache_key); +/// } +/// } +/// +/// fn poll_purge(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Error>> { +/// while let Some(res) = ready!(self.expirations.poll_expired(cx)) { +/// let entry = res?; +/// self.entries.remove(entry.get_ref()); +/// } +/// +/// Poll::Ready(Ok(())) +/// } +/// } +/// ``` +/// +/// [`insert`]: #method.insert +/// [`insert_at`]: #method.insert_at +/// [`Key`]: struct@Key +/// [`Stream`]: https://docs.rs/futures/0.1/futures/stream/trait.Stream.html +/// [`poll`]: #method.poll +/// [`Stream::poll`]: #method.poll +/// [`Timer`]: ../struct.Timer.html +/// [`slab`]: https://docs.rs/slab +/// [`capacity`]: #method.capacity +/// [`reserve`]: #method.reserve +#[derive(Debug)] +pub struct DelayQueue<T> { + /// Stores data associated with entries + slab: Slab<Data<T>>, + + /// Lookup structure tracking all delays in the queue + wheel: Wheel<Stack<T>>, + + /// Delays that were inserted when already expired. These cannot be stored + /// in the wheel + expired: Stack<T>, + + /// Delay expiring when the *first* item in the queue expires + delay: Option<Delay>, + + /// Wheel polling state + poll: wheel::Poll, + + /// Instant at which the timer starts + start: Instant, +} + +/// An entry in `DelayQueue` that has expired and removed. +/// +/// Values are returned by [`DelayQueue::poll`]. +/// +/// [`DelayQueue::poll`]: method@DelayQueue::poll +#[derive(Debug)] +pub struct Expired<T> { + /// The data stored in the queue + data: T, + + /// The expiration time + deadline: Instant, + + /// The key associated with the entry + key: Key, +} + +/// Token to a value stored in a `DelayQueue`. +/// +/// Instances of `Key` are returned by [`DelayQueue::insert`]. See [`DelayQueue`] +/// documentation for more details. +/// +/// [`DelayQueue`]: struct@DelayQueue +/// [`DelayQueue::insert`]: method@DelayQueue::insert +#[derive(Debug, Clone)] +pub struct Key { + index: usize, +} + +#[derive(Debug)] +struct Stack<T> { + /// Head of the stack + head: Option<usize>, + _p: PhantomData<fn() -> T>, +} + +#[derive(Debug)] +struct Data<T> { + /// The data being stored in the queue and will be returned at the requested + /// instant. + inner: T, + + /// The instant at which the item is returned. + when: u64, + + /// Set to true when stored in the `expired` queue + expired: bool, + + /// Next entry in the stack + next: Option<usize>, + + /// Previous entry in the stack + prev: Option<usize>, +} + +/// Maximum number of entries the queue can handle +const MAX_ENTRIES: usize = (1 << 30) - 1; + +impl<T> DelayQueue<T> { + /// Creates a new, empty, `DelayQueue` + /// + /// The queue will not allocate storage until items are inserted into it. + /// + /// # Examples + /// + /// ```rust + /// # use tokio::time::DelayQueue; + /// let delay_queue: DelayQueue<u32> = DelayQueue::new(); + /// ``` + pub fn new() -> DelayQueue<T> { + DelayQueue::with_capacity(0) + } + + /// Creates a new, empty, `DelayQueue` with the specified capacity. + /// + /// The queue will be able to hold at least `capacity` elements without + /// reallocating. If `capacity` is 0, the queue will not allocate for + /// storage. + /// + /// # Examples + /// + /// ```rust + /// # use tokio::time::DelayQueue; + /// # use std::time::Duration; + /// + /// # #[tokio::main] + /// # async fn main() { + /// let mut delay_queue = DelayQueue::with_capacity(10); + /// + /// // These insertions are done without further allocation + /// for i in 0..10 { + /// delay_queue.insert(i, Duration::from_secs(i)); + /// } + /// + /// // This will make the queue allocate additional storage + /// delay_queue.insert(11, Duration::from_secs(11)); + /// # } + /// ``` + pub fn with_capacity(capacity: usize) -> DelayQueue<T> { + DelayQueue { + wheel: Wheel::new(), + slab: Slab::with_capacity(capacity), + expired: Stack::default(), + delay: None, + poll: wheel::Poll::new(0), + start: Instant::now(), + } + } + + /// Inserts `value` into the queue set to expire at a specific instant in + /// time. + /// + /// This function is identical to `insert`, but takes an `Instant` instead + /// of a `Duration`. + /// + /// `value` is stored in the queue until `when` is reached. At which point, + /// `value` will be returned from [`poll`]. If `when` has already been + /// reached, then `value` is immediately made available to poll. + /// + /// The return value represents the insertion and is used at an argument to + /// [`remove`] and [`reset`]. Note that [`Key`] is token and is reused once + /// `value` is removed from the queue either by calling [`poll`] after + /// `when` is reached or by calling [`remove`]. At this point, the caller + /// must take care to not use the returned [`Key`] again as it may reference + /// a different item in the queue. + /// + /// See [type] level documentation for more details. + /// + /// # Panics + /// + /// This function panics if `when` is too far in the future. + /// + /// # Examples + /// + /// Basic usage + /// + /// ```rust + /// use tokio::time::{DelayQueue, Duration, Instant}; + /// + /// # #[tokio::main] + /// # async fn main() { + /// let mut delay_queue = DelayQueue::new(); + /// let key = delay_queue.insert_at( + /// "foo", Instant::now() + Duration::from_secs(5)); + /// + /// // Remove the entry + /// let item = delay_queue.remove(&key); + /// assert_eq!(*item.get_ref(), "foo"); + /// # } + /// ``` + /// + /// [`poll`]: #method.poll + /// [`remove`]: #method.remove + /// [`reset`]: #method.reset + /// [`Key`]: struct@Key + /// [type]: # + pub fn insert_at(&mut self, value: T, when: Instant) -> Key { + assert!(self.slab.len() < MAX_ENTRIES, "max entries exceeded"); + + // Normalize the deadline. Values cannot be set to expire in the past. + let when = self.normalize_deadline(when); + + // Insert the value in the store + let key = self.slab.insert(Data { + inner: value, + when, + expired: false, + next: None, + prev: None, + }); + + self.insert_idx(when, key); + + // Set a new delay if the current's deadline is later than the one of the new item + let should_set_delay = if let Some(ref delay) = self.delay { + let current_exp = self.normalize_deadline(delay.deadline()); + current_exp > when + } else { + true + }; + + if should_set_delay { + let delay_time = self.start + Duration::from_millis(when); + if let Some(ref mut delay) = &mut self.delay { + delay.reset(delay_time); + } else { + self.delay = Some(delay_until(delay_time)); + } + } + + Key::new(key) + } + + /// Attempts to pull out the next value of the delay queue, registering the + /// current task for wakeup if the value is not yet available, and returning + /// None if the queue is exhausted. + pub fn poll_expired( + &mut self, + cx: &mut task::Context<'_>, + ) -> Poll<Option<Result<Expired<T>, Error>>> { + let item = ready!(self.poll_idx(cx)); + Poll::Ready(item.map(|result| { + result.map(|idx| { + let data = self.slab.remove(idx); + debug_assert!(data.next.is_none()); + debug_assert!(data.prev.is_none()); + + Expired { + key: Key::new(idx), + data: data.inner, + deadline: self.start + Duration::from_millis(data.when), + } + }) + })) + } + + /// Inserts `value` into the queue set to expire after the requested duration + /// elapses. + /// + /// This function is identical to `insert_at`, but takes a `Duration` + /// instead of an `Instant`. + /// + /// `value` is stored in the queue until `when` is reached. At which point, + /// `value` will be returned from [`poll`]. If `when` has already been + /// reached, then `value` is immediately made available to poll. + /// + /// The return value represents the insertion and is used at an argument to + /// [`remove`] and [`reset`]. Note that [`Key`] is token and is reused once + /// `value` is removed from the queue either by calling [`poll`] after + /// `when` is reached or by calling [`remove`]. At this point, the caller + /// must take care to not use the returned [`Key`] again as it may reference + /// a different item in the queue. + /// + /// See [type] level documentation for more details. + /// + /// # Panics + /// + /// This function panics if `timeout` is greater than the maximum supported + /// duration. + /// + /// # Examples + /// + /// Basic usage + /// + /// ```rust + /// use tokio::time::DelayQueue; + /// use std::time::Duration; + /// + /// # #[tokio::main] + /// # async fn main() { + /// let mut delay_queue = DelayQueue::new(); + /// let key = delay_queue.insert("foo", Duration::from_secs(5)); + /// + /// // Remove the entry + /// let item = delay_queue.remove(&key); + /// assert_eq!(*item.get_ref(), "foo"); + /// # } + /// ``` + /// + /// [`poll`]: #method.poll + /// [`remove`]: #method.remove + /// [`reset`]: #method.reset + /// [`Key`]: struct@Key + /// [type]: # + pub fn insert(&mut self, value: T, timeout: Duration) -> Key { + self.insert_at(value, Instant::now() + timeout) + } + + fn insert_idx(&mut self, when: u64, key: usize) { + use self::wheel::{InsertError, Stack}; + + // Register the deadline with the timer wheel + match self.wheel.insert(when, key, &mut self.slab) { + Ok(_) => {} + Err((_, InsertError::Elapsed)) => { + self.slab[key].expired = true; + // The delay is already expired, store it in the expired queue + self.expired.push(key, &mut self.slab); + } + Err((_, err)) => panic!("invalid deadline; err={:?}", err), + } + } + + /// Removes the item associated with `key` from the queue. + /// + /// There must be an item associated with `key`. The function returns the + /// removed item as well as the `Instant` at which it will the delay will + /// have expired. + /// + /// # Panics + /// + /// The function panics if `key` is not contained by the queue. + /// + /// # Examples + /// + /// Basic usage + /// + /// ```rust + /// use tokio::time::DelayQueue; + /// use std::time::Duration; + /// + /// # #[tokio::main] + /// # async fn main() { + /// let mut delay_queue = DelayQueue::new(); + /// let key = delay_queue.insert("foo", Duration::from_secs(5)); + /// + /// // Remove the entry + /// let item = delay_queue.remove(&key); + /// assert_eq!(*item.get_ref(), "foo"); + /// # } + /// ``` + pub fn remove(&mut self, key: &Key) -> Expired<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); + } else { + self.wheel.remove(&key.index, &mut self.slab); + } + + let data = self.slab.remove(key.index); + + Expired { + key: Key::new(key.index), + data: data.inner, + deadline: self.start + Duration::from_millis(data.when), + } + } + + /// Sets the delay of the item associated with `key` to expire at `when`. + /// + /// This function is identical to `reset` but takes an `Instant` instead of + /// a `Duration`. + /// + /// The item remains in the queue but the delay is set to expire at `when`. + /// If `when` is in the past, then the item is immediately made available to + /// the caller. + /// + /// # Panics + /// + /// This function panics if `when` is too far in the future or if `key` is + /// not contained by the queue. + /// + /// # Examples + /// + /// Basic usage + /// + /// ```rust + /// use tokio::time::{DelayQueue, Duration, Instant}; + /// + /// # #[tokio::main] + /// # async fn main() { + /// let mut delay_queue = DelayQueue::new(); + /// let key = delay_queue.insert("foo", Duration::from_secs(5)); + /// + /// // "foo" is scheduled to be returned in 5 seconds + /// + /// delay_queue.reset_at(&key, Instant::now() + Duration::from_secs(10)); + /// + /// // "foo"is now scheduled to be returned in 10 seconds + /// # } + /// ``` + pub fn reset_at(&mut self, key: &Key, when: Instant) { + self.wheel.remove(&key.index, &mut self.slab); + + // 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.insert_idx(when, key.index); + + let next_deadline = self.next_deadline(); + if let (Some(ref mut delay), Some(deadline)) = (&mut self.delay, next_deadline) { + delay.reset(deadline); + } + } + + /// Returns the next time poll as determined by the wheel + fn next_deadline(&mut self) -> Option<Instant> { + self.wheel + .poll_at() + .map(|poll_at| self.start + Duration::from_millis(poll_at)) + } + + /// Sets the delay of the item associated with `key` to expire after + /// `timeout`. + /// + /// This function is identical to `reset_at` but takes a `Duration` instead + /// of an `Instant`. + /// + /// The item remains in the queue but the delay is set to expire after + /// `timeout`. If `timeout` is zero, then the item is immediately made + /// available to the caller. + /// + /// # Panics + /// + /// This function panics if `timeout` is greater than the maximum supported + /// duration or if `key` is not contained by the queue. + /// + /// # Examples + /// + /// Basic usage + /// + /// ```rust + /// use tokio::time::DelayQueue; + /// use std::time::Duration; + /// + /// # #[tokio::main] + /// # async fn main() { + /// let mut delay_queue = DelayQueue::new(); + /// let key = delay_queue.insert("foo", Duration::from_secs(5)); + /// + /// // "foo" is scheduled to be returned in 5 seconds + /// + /// delay_queue.reset(&key, Duration::from_secs(10)); + /// + /// // "foo"is now scheduled to be returned in 10 seconds + /// # } + /// ``` + pub fn reset(&mut self, key: &Key, timeout: Duration) { + self.reset_at(key, Instant::now() + timeout); + } + + /// Clears the queue, removing all items. + /// + /// After calling `clear`, [`poll`] will return `Ok(Ready(None))`. + /// + /// Note that this method has no effect on the allocated capacity. + /// + /// [`poll`]: #method.poll + /// + /// # Examples + /// + /// ```rust + /// use tokio::time::DelayQueue; + /// use std::time::Duration; + /// + /// # #[tokio::main] + /// # async fn main() { + /// let mut delay_queue = DelayQueue::new(); + /// + /// delay_queue.insert("foo", Duration::from_secs(5)); + /// + /// assert!(!delay_queue.is_empty()); + /// + /// delay_queue.clear(); + /// + /// assert!(delay_queue.is_empty()); + /// # } + /// ``` + pub fn clear(&mut self) { + self.slab.clear(); + self.expired = Stack::default(); + self.wheel = Wheel::new(); + self.delay = None; + } + + /// Returns the number of elements the queue can hold without reallocating. + /// + /// # Examples + /// + /// ```rust + /// use tokio::time::DelayQueue; + /// + /// let delay_queue: DelayQueue<i32> = DelayQueue::with_capacity(10); + /// assert_eq!(delay_queue.capacity(), 10); + /// ``` + pub fn capacity(&self) -> usize { + self.slab.capacity() + } + + /// Returns the number of elements currently in the queue. + /// + /// # Examples + /// + /// ```rust + /// use tokio::time::DelayQueue; + /// use std::time::Duration; + /// + /// # #[tokio::main] + /// # async fn main() { + /// let mut delay_queue: DelayQueue<i32> = DelayQueue::with_capacity(10); + /// assert_eq!(delay_queue.len(), 0); + /// delay_queue.insert(3, Duration::from_secs(5)); + /// assert_eq!(delay_queue.len(), 1); + /// # } + /// ``` + pub fn len(&self) -> usize { + self.slab.len() + } + + /// Reserves capacity for at least `additional` more items to be queued + /// without allocating. + /// + /// `reserve` does nothing if the queue already has sufficient capacity for + /// `additional` more values. If more capacity is required, a new segment of + /// memory will be allocated and all existing values will be copied into it. + /// As such, if the queue is already very large, a call to `reserve` can end + /// up being expensive. + /// + /// The queue may reserve more than `additional` extra space in order to + /// avoid frequent reallocations. + /// + /// # Panics + /// + /// Panics if the new capacity exceeds the maximum number of entries the + /// queue can contain. + /// + /// # Examples + /// + /// ``` + /// use tokio::time::DelayQueue; + /// use std::time::Duration; + /// + /// # #[tokio::main] + /// # async fn main() { + /// let mut delay_queue = DelayQueue::new(); + /// + /// delay_queue.insert("hello", Duration::from_secs(10)); + /// delay_queue.reserve(10); + /// + /// assert!(delay_queue.capacity() >= 11); + /// # } + /// ``` + pub fn reserve(&mut self, additional: usize) { + self.slab.reserve(additional); + } + + /// Returns `true` if there are no items in the queue. + /// + /// Note that this function returns `false` even if all items have not yet + /// expired and a call to `poll` will return `NotReady`. + /// + /// # Examples + /// + /// ``` + /// use tokio::time::DelayQueue; + /// use std::time::Duration; + /// + /// # #[tokio::main] + /// # async fn main() { + /// let mut delay_queue = DelayQueue::new(); + /// assert!(delay_queue.is_empty()); + /// + /// delay_queue.insert("hello", Duration::from_secs(5)); + /// assert!(!delay_queue.is_empty()); + /// # } + /// ``` + pub fn is_empty(&self) -> bool { + self.slab.is_empty() + } + + /// Polls the queue, returning the index of the next slot in the slab that + /// 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>>> { + use self::wheel::Stack; + + let expired = self.expired.pop(&mut self.slab); + + if expired.is_some() { + return Poll::Ready(expired.map(Ok)); + } + + loop { + if let Some(ref mut delay) = self.delay { + if !delay.is_elapsed() { + ready!(Pin::new(&mut *delay).poll(cx)); + } + + let now = crate::time::ms(delay.deadline() - self.start, crate::time::Round::Down); + + self.poll = wheel::Poll::new(now); + } + + // We poll the wheel to get the next value out before finding the next deadline. + let wheel_idx = self.wheel.poll(&mut self.poll, &mut self.slab); + + self.delay = self.next_deadline().map(delay_until); + + if let Some(idx) = wheel_idx { + return Poll::Ready(Some(Ok(idx))); + } + + if self.delay.is_none() { + return Poll::Ready(None); + } + } + } + + fn normalize_deadline(&self, when: Instant) -> u64 { + let when = if when < self.start { + 0 + } else { + crate::time::ms(when - self.start, crate::time::Round::Up) + }; + + cmp::max(when, self.wheel.elapsed()) + } +} + +// We never put `T` in a `Pin`... +impl<T> Unpin for DelayQueue<T> {} + +impl<T> Default for DelayQueue<T> { + fn default() -> DelayQueue<T> { + DelayQueue::new() + } +} + +#[cfg(feature = "stream")] +impl<T> futures_core::Stream for DelayQueue<T> { + // DelayQueue seems much more specific, where a user may care that it + // has reached capacity, so return those errors instead of panicking. + type Item = Result<Expired<T>, Error>; + + fn poll_next(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Option<Self::Item>> { + DelayQueue::poll_expired(self.get_mut(), cx) + } +} + +impl<T> wheel::Stack for Stack<T> { + type Owned = usize; + type Borrowed = usize; + type Store = Slab<Data<T>>; + + fn is_empty(&self) -> bool { + self.head.is_none() + } + + fn push(&mut self, item: Self::Owned, store: &mut Self::Store) { + // Ensure the entry is not already in a stack. + debug_assert!(store[item].next.is_none()); + debug_assert!(store[item].prev.is_none()); + + // Remove the old head entry + let old = self.head.take(); + + if let Some(idx) = old { + store[idx].prev = Some(item); + } + + store[item].next = old; + 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(idx) = self.head { + store[idx].prev = None; + } + + store[idx].next = None; + debug_assert!(store[idx].prev.is_none()); + + Some(idx) + } else { + None + } + } + + fn remove(&mut self, item: &Self::Borrowed, store: &mut Self::Store) { + assert!(store.contains(*item)); + + // Ensure that the entry is in fact contained by the stack + debug_assert!({ + // This walks the full linked list even if an entry is found. + let mut next = self.head; + let mut contains = false; + + while let Some(idx) = next { + if idx == *item { + debug_assert!(!contains); + contains = true; + } + + next = store[idx].next; + } + + contains + }); + + if let Some(next) = store[*item].next { + store[next].prev = store[*item].prev; + } + + if let Some(prev) = store[*item].prev { + store[prev].next = store[*item].next; + } else { + self.head = store[*item].next; + } + + store[*item].next = None; + store[*item].prev = None; + } + + fn when(item: &Self::Borrowed, store: &Self::Store) -> u64 { + store[*item].when + } +} + +impl<T> Default for Stack<T> { + fn default() -> Stack<T> { + Stack { + head: None, + _p: PhantomData, + } + } +} + +impl Key { + pub(crate) fn new(index: usize) -> Key { + Key { index } + } +} + +impl<T> Expired<T> { + /// Returns a reference to the inner value. + pub fn get_ref(&self) -> &T { + &self.data + } + + /// Returns a mutable reference to the inner value. + pub fn get_mut(&mut self) -> &mut T { + &mut self.data + } + + /// Consumes `self` and returns the inner value. + pub fn into_inner(self) -> T { + self.data + } + + /// Returns the deadline that the expiration was set to. + pub fn deadline(&self) -> Instant { + self.deadline + } +} diff --git a/third_party/rust/tokio/src/time/driver/atomic_stack.rs b/third_party/rust/tokio/src/time/driver/atomic_stack.rs new file mode 100644 index 0000000000..7e5a83fa52 --- /dev/null +++ b/third_party/rust/tokio/src/time/driver/atomic_stack.rs @@ -0,0 +1,124 @@ +use crate::time::driver::Entry; +use crate::time::Error; + +use std::ptr; +use std::sync::atomic::AtomicPtr; +use std::sync::atomic::Ordering::SeqCst; +use std::sync::Arc; + +/// A stack of `Entry` nodes +#[derive(Debug)] +pub(crate) struct AtomicStack { + /// Stack head + head: AtomicPtr<Entry>, +} + +/// Entries that were removed from the stack +#[derive(Debug)] +pub(crate) struct AtomicStackEntries { + ptr: *mut Entry, +} + +/// Used to indicate that the timer has shutdown. +const SHUTDOWN: *mut Entry = 1 as *mut _; + +impl AtomicStack { + pub(crate) fn new() -> AtomicStack { + AtomicStack { + head: AtomicPtr::new(ptr::null_mut()), + } + } + + /// Pushes an entry onto the stack. + /// + /// Returns `true` if the entry was pushed, `false` if the entry is already + /// on the stack, `Err` if the timer is shutdown. + pub(crate) fn push(&self, entry: &Arc<Entry>) -> Result<bool, Error> { + // First, set the queued bit on the entry + let queued = entry.queued.fetch_or(true, SeqCst); + + if queued { + // Already queued, nothing more to do + return Ok(false); + } + + let ptr = Arc::into_raw(entry.clone()) as *mut _; + + let mut curr = self.head.load(SeqCst); + + loop { + if curr == SHUTDOWN { + // Don't leak the entry node + let _ = unsafe { Arc::from_raw(ptr) }; + + return Err(Error::shutdown()); + } + + // Update the `next` pointer. This is safe because setting the queued + // bit is a "lock" on this field. + unsafe { + *(entry.next_atomic.get()) = curr; + } + + let actual = self.head.compare_and_swap(curr, ptr, SeqCst); + + if actual == curr { + break; + } + + curr = actual; + } + + Ok(true) + } + + /// Takes all entries from the stack + pub(crate) fn take(&self) -> AtomicStackEntries { + let ptr = self.head.swap(ptr::null_mut(), SeqCst); + AtomicStackEntries { ptr } + } + + /// Drains all remaining nodes in the stack and prevent any new nodes from + /// being pushed onto the stack. + pub(crate) fn shutdown(&self) { + // Shutdown the processing queue + let ptr = self.head.swap(SHUTDOWN, SeqCst); + + // Let the drop fn of `AtomicStackEntries` handle draining the stack + drop(AtomicStackEntries { ptr }); + } +} + +// ===== impl AtomicStackEntries ===== + +impl Iterator for AtomicStackEntries { + type Item = Arc<Entry>; + + fn next(&mut self) -> Option<Self::Item> { + if self.ptr.is_null() { + return None; + } + + // Convert the pointer to an `Arc<Entry>` + let entry = unsafe { Arc::from_raw(self.ptr) }; + + // Update `self.ptr` to point to the next element of the stack + self.ptr = unsafe { *entry.next_atomic.get() }; + + // Unset the queued flag + let res = entry.queued.fetch_and(false, SeqCst); + debug_assert!(res); + + // Return the entry + Some(entry) + } +} + +impl Drop for AtomicStackEntries { + fn drop(&mut self) { + for entry in self { + // Flag the entry as errored + entry.error(); + } + } +} diff --git a/third_party/rust/tokio/src/time/driver/entry.rs b/third_party/rust/tokio/src/time/driver/entry.rs new file mode 100644 index 0000000000..20cc824019 --- /dev/null +++ b/third_party/rust/tokio/src/time/driver/entry.rs @@ -0,0 +1,345 @@ +use crate::loom::sync::atomic::AtomicU64; +use crate::sync::AtomicWaker; +use crate::time::driver::{Handle, Inner}; +use crate::time::{Duration, Error, Instant}; + +use std::cell::UnsafeCell; +use std::ptr; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering::SeqCst; +use std::sync::{Arc, Weak}; +use std::task::{self, Poll}; +use std::u64; + +/// Internal state shared between a `Delay` instance and the timer. +/// +/// This struct is used as a node in two intrusive data structures: +/// +/// * An atomic stack used to signal to the timer thread that the entry state +/// has changed. The timer thread will observe the entry on this stack and +/// perform any actions as necessary. +/// +/// * A doubly linked list used **only** by the timer thread. Each slot in the +/// timer wheel is a head pointer to the list of entries that must be +/// processed during that timer tick. +#[derive(Debug)] +pub(crate) struct Entry { + /// Only accessed from `Registration`. + time: CachePadded<UnsafeCell<Time>>, + + /// Timer internals. Using a weak pointer allows the timer to shutdown + /// without all `Delay` instances having completed. + /// + /// When `None`, the entry has not yet been linked with a timer instance. + inner: Weak<Inner>, + + /// Tracks the entry state. This value contains the following information: + /// + /// * The deadline at which the entry must be "fired". + /// * A flag indicating if the entry has already been fired. + /// * Whether or not the entry transitioned to the error state. + /// + /// When an `Entry` is created, `state` is initialized to the instant at + /// which the entry must be fired. When a timer is reset to a different + /// instant, this value is changed. + state: AtomicU64, + + /// Task to notify once the deadline is reached. + waker: AtomicWaker, + + /// True when the entry is queued in the "process" stack. This value + /// is set before pushing the value and unset after popping the value. + /// + /// TODO: This could possibly be rolled up into `state`. + pub(super) queued: AtomicBool, + + /// Next entry in the "process" linked list. + /// + /// Access to this field is coordinated by the `queued` flag. + /// + /// Represents a strong Arc ref. + pub(super) next_atomic: UnsafeCell<*mut Entry>, + + /// When the entry expires, relative to the `start` of the timer + /// (Inner::start). This is only used by the timer. + /// + /// A `Delay` instance can be reset to a different deadline by the thread + /// that owns the `Delay` instance. In this case, the timer thread will not + /// immediately know that this has happened. The timer thread must know the + /// last deadline that it saw as it uses this value to locate the entry in + /// its wheel. + /// + /// Once the timer thread observes that the instant has changed, it updates + /// the wheel and sets this value. The idea is that this value eventually + /// converges to the value of `state` as the timer thread makes updates. + when: UnsafeCell<Option<u64>>, + + /// Next entry in the State's linked list. + /// + /// This is only accessed by the timer + pub(super) next_stack: UnsafeCell<Option<Arc<Entry>>>, + + /// Previous entry in the State's linked list. + /// + /// This is only accessed by the timer and is used to unlink a canceled + /// entry. + /// + /// This is a weak reference. + pub(super) prev_stack: UnsafeCell<*const Entry>, +} + +/// Stores the info for `Delay`. +#[derive(Debug)] +pub(crate) struct Time { + pub(crate) deadline: Instant, + pub(crate) duration: Duration, +} + +/// Flag indicating a timer entry has elapsed +const ELAPSED: u64 = 1 << 63; + +/// Flag indicating a timer entry has reached an error state +const ERROR: u64 = u64::MAX; + +// ===== impl Entry ===== + +impl Entry { + pub(crate) fn new(handle: &Handle, deadline: Instant, duration: Duration) -> Arc<Entry> { + let inner = handle.inner().unwrap(); + let entry: Entry; + + // Increment the number of active timeouts + if inner.increment().is_err() { + entry = Entry::new2(deadline, duration, Weak::new(), ERROR) + } else { + let when = inner.normalize_deadline(deadline); + let state = if when <= inner.elapsed() { + ELAPSED + } else { + when + }; + entry = Entry::new2(deadline, duration, Arc::downgrade(&inner), state); + } + + let entry = Arc::new(entry); + if inner.queue(&entry).is_err() { + entry.error(); + } + + entry + } + + /// Only called by `Registration` + pub(crate) fn time_ref(&self) -> &Time { + unsafe { &*self.time.0.get() } + } + + /// Only called by `Registration` + #[allow(clippy::mut_from_ref)] // https://github.com/rust-lang/rust-clippy/issues/4281 + pub(crate) unsafe fn time_mut(&self) -> &mut Time { + &mut *self.time.0.get() + } + + /// The current entry state as known by the timer. This is not the value of + /// `state`, but lets the timer know how to converge its state to `state`. + pub(crate) fn when_internal(&self) -> Option<u64> { + unsafe { *self.when.get() } + } + + pub(crate) fn set_when_internal(&self, when: Option<u64>) { + unsafe { + *self.when.get() = when; + } + } + + /// Called by `Timer` to load the current value of `state` for processing + pub(crate) fn load_state(&self) -> Option<u64> { + let state = self.state.load(SeqCst); + + if is_elapsed(state) { + None + } else { + Some(state) + } + } + + pub(crate) fn is_elapsed(&self) -> bool { + let state = self.state.load(SeqCst); + is_elapsed(state) + } + + pub(crate) fn fire(&self, when: u64) { + let mut curr = self.state.load(SeqCst); + + loop { + if is_elapsed(curr) || curr > when { + return; + } + + let next = ELAPSED | curr; + let actual = self.state.compare_and_swap(curr, next, SeqCst); + + if curr == actual { + break; + } + + curr = actual; + } + + self.waker.wake(); + } + + pub(crate) fn error(&self) { + // Only transition to the error state if not currently elapsed + let mut curr = self.state.load(SeqCst); + + loop { + if is_elapsed(curr) { + return; + } + + let next = ERROR; + + let actual = self.state.compare_and_swap(curr, next, SeqCst); + + if curr == actual { + break; + } + + curr = actual; + } + + self.waker.wake(); + } + + pub(crate) fn cancel(entry: &Arc<Entry>) { + let state = entry.state.fetch_or(ELAPSED, SeqCst); + + if is_elapsed(state) { + // Nothing more to do + return; + } + + // If registered with a timer instance, try to upgrade the Arc. + let inner = match entry.upgrade_inner() { + Some(inner) => inner, + None => return, + }; + + let _ = inner.queue(entry); + } + + pub(crate) fn poll_elapsed(&self, cx: &mut task::Context<'_>) -> Poll<Result<(), Error>> { + let mut curr = self.state.load(SeqCst); + + if is_elapsed(curr) { + return Poll::Ready(if curr == ERROR { + Err(Error::shutdown()) + } else { + Ok(()) + }); + } + + self.waker.register_by_ref(cx.waker()); + + curr = self.state.load(SeqCst); + + if is_elapsed(curr) { + return Poll::Ready(if curr == ERROR { + Err(Error::shutdown()) + } else { + Ok(()) + }); + } + + Poll::Pending + } + + /// Only called by `Registration` + pub(crate) fn reset(entry: &mut Arc<Entry>) { + let inner = match entry.upgrade_inner() { + Some(inner) => inner, + None => return, + }; + + let deadline = entry.time_ref().deadline; + let when = inner.normalize_deadline(deadline); + let elapsed = inner.elapsed(); + + let mut curr = entry.state.load(SeqCst); + let mut notify; + + loop { + // In these two cases, there is no work to do when resetting the + // timer. If the `Entry` is in an error state, then it cannot be + // used anymore. If resetting the entry to the current value, then + // the reset is a noop. + if curr == ERROR || curr == when { + return; + } + + let next; + + if when <= elapsed { + next = ELAPSED; + notify = !is_elapsed(curr); + } else { + next = when; + notify = true; + } + + let actual = entry.state.compare_and_swap(curr, next, SeqCst); + + if curr == actual { + break; + } + + curr = actual; + } + + if notify { + let _ = inner.queue(entry); + } + } + + fn new2(deadline: Instant, duration: Duration, inner: Weak<Inner>, state: u64) -> Self { + Self { + time: CachePadded(UnsafeCell::new(Time { deadline, duration })), + inner, + waker: AtomicWaker::new(), + state: AtomicU64::new(state), + queued: AtomicBool::new(false), + next_atomic: UnsafeCell::new(ptr::null_mut()), + when: UnsafeCell::new(None), + next_stack: UnsafeCell::new(None), + prev_stack: UnsafeCell::new(ptr::null_mut()), + } + } + + fn upgrade_inner(&self) -> Option<Arc<Inner>> { + self.inner.upgrade() + } +} + +fn is_elapsed(state: u64) -> bool { + state & ELAPSED == ELAPSED +} + +impl Drop for Entry { + fn drop(&mut self) { + let inner = match self.upgrade_inner() { + Some(inner) => inner, + None => return, + }; + + inner.decrement(); + } +} + +unsafe impl Send for Entry {} +unsafe impl Sync for Entry {} + +#[cfg_attr(target_arch = "x86_64", repr(align(128)))] +#[cfg_attr(not(target_arch = "x86_64"), repr(align(64)))] +#[derive(Debug)] +struct CachePadded<T>(T); diff --git a/third_party/rust/tokio/src/time/driver/handle.rs b/third_party/rust/tokio/src/time/driver/handle.rs new file mode 100644 index 0000000000..38b1761c8e --- /dev/null +++ b/third_party/rust/tokio/src/time/driver/handle.rs @@ -0,0 +1,38 @@ +use crate::runtime::context; +use crate::time::driver::Inner; +use std::fmt; +use std::sync::{Arc, Weak}; + +/// Handle to time driver instance. +#[derive(Clone)] +pub(crate) struct Handle { + inner: Weak<Inner>, +} + +impl Handle { + /// Creates a new timer `Handle` from a shared `Inner` timer state. + pub(crate) fn new(inner: Weak<Inner>) -> Self { + Handle { inner } + } + + /// Tries to get a handle to the current timer. + /// + /// # Panics + /// + /// This function panics if there is no current timer set. + pub(crate) fn current() -> Self { + context::time_handle() + .expect("there is no timer running, must be called from the context of Tokio runtime") + } + + /// Tries to return a strong ref to the inner + pub(crate) fn inner(&self) -> Option<Arc<Inner>> { + self.inner.upgrade() + } +} + +impl fmt::Debug for Handle { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "Handle") + } +} diff --git a/third_party/rust/tokio/src/time/driver/mod.rs b/third_party/rust/tokio/src/time/driver/mod.rs new file mode 100644 index 0000000000..4616816f3f --- /dev/null +++ b/third_party/rust/tokio/src/time/driver/mod.rs @@ -0,0 +1,391 @@ +//! Time driver + +mod atomic_stack; +use self::atomic_stack::AtomicStack; + +mod entry; +pub(super) use self::entry::Entry; + +mod handle; +pub(crate) use self::handle::Handle; + +mod registration; +pub(crate) use self::registration::Registration; + +mod stack; +use self::stack::Stack; + +use crate::loom::sync::atomic::{AtomicU64, AtomicUsize}; +use crate::park::{Park, Unpark}; +use crate::time::{wheel, Error}; +use crate::time::{Clock, Duration, Instant}; + +use std::sync::atomic::Ordering::{Acquire, Relaxed, Release, SeqCst}; + +use std::sync::Arc; +use std::usize; +use std::{cmp, fmt}; + +/// Time implementation that drives [`Delay`], [`Interval`], and [`Timeout`]. +/// +/// A `Driver` instance tracks the state necessary for managing time and +/// notifying the [`Delay`] instances once their deadlines are reached. +/// +/// It is expected that a single instance manages many individual [`Delay`] +/// instances. The `Driver` implementation is thread-safe and, as such, is able +/// to handle callers from across threads. +/// +/// After creating the `Driver` instance, the caller must repeatedly call +/// [`turn`]. The time driver will perform no work unless [`turn`] is called +/// repeatedly. +/// +/// The driver has a resolution of one millisecond. Any unit of time that falls +/// between milliseconds are rounded up to the next millisecond. +/// +/// When an instance is dropped, any outstanding [`Delay`] instance that has not +/// elapsed will be notified with an error. At this point, calling `poll` on the +/// [`Delay`] instance will result in `Err` being returned. +/// +/// # Implementation +/// +/// THe time driver is based on the [paper by Varghese and Lauck][paper]. +/// +/// A hashed timing wheel is a vector of slots, where each slot handles a time +/// slice. As time progresses, the timer walks over the slot for the current +/// instant, and processes each entry for that slot. When the timer reaches the +/// end of the wheel, it starts again at the beginning. +/// +/// The implementation maintains six wheels arranged in a set of levels. As the +/// levels go up, the slots of the associated wheel represent larger intervals +/// of time. At each level, the wheel has 64 slots. Each slot covers a range of +/// time equal to the wheel at the lower level. At level zero, each slot +/// represents one millisecond of time. +/// +/// The wheels are: +/// +/// * Level 0: 64 x 1 millisecond slots. +/// * Level 1: 64 x 64 millisecond slots. +/// * Level 2: 64 x ~4 second slots. +/// * Level 3: 64 x ~4 minute slots. +/// * Level 4: 64 x ~4 hour slots. +/// * Level 5: 64 x ~12 day slots. +/// +/// When the timer processes entries at level zero, it will notify all the +/// `Delay` instances as their deadlines have been reached. For all higher +/// levels, all entries will be redistributed across the wheel at the next level +/// down. Eventually, as time progresses, entries will [`Delay`] instances will +/// either be canceled (dropped) or their associated entries will reach level +/// zero and be notified. +#[derive(Debug)] +pub(crate) struct Driver<T> { + /// Shared state + inner: Arc<Inner>, + + /// Timer wheel + wheel: wheel::Wheel<Stack>, + + /// Thread parker. The `Driver` park implementation delegates to this. + park: T, + + /// Source of "now" instances + clock: Clock, +} + +/// Timer state shared between `Driver`, `Handle`, and `Registration`. +pub(crate) struct Inner { + /// The instant at which the timer started running. + start: Instant, + + /// The last published timer `elapsed` value. + elapsed: AtomicU64, + + /// Number of active timeouts + num: AtomicUsize, + + /// Head of the "process" linked list. + process: AtomicStack, + + /// Unparks the timer thread. + unpark: Box<dyn Unpark>, +} + +/// Maximum number of timeouts the system can handle concurrently. +const MAX_TIMEOUTS: usize = usize::MAX >> 1; + +// ===== impl Driver ===== + +impl<T> Driver<T> +where + T: Park, +{ + /// Creates a new `Driver` instance that uses `park` to block the current + /// thread and `now` to get the current `Instant`. + /// + /// Specifying the source of time is useful when testing. + pub(crate) fn new(park: T, clock: Clock) -> Driver<T> { + let unpark = Box::new(park.unpark()); + + Driver { + inner: Arc::new(Inner::new(clock.now(), unpark)), + wheel: wheel::Wheel::new(), + park, + clock, + } + } + + /// Returns a handle to the timer. + /// + /// The `Handle` is how `Delay` instances are created. The `Delay` instances + /// can either be created directly or the `Handle` instance can be passed to + /// `with_default`, setting the timer as the default timer for the execution + /// context. + pub(crate) fn handle(&self) -> Handle { + Handle::new(Arc::downgrade(&self.inner)) + } + + /// Converts an `Expiration` to an `Instant`. + fn expiration_instant(&self, when: u64) -> Instant { + self.inner.start + Duration::from_millis(when) + } + + /// Runs timer related logic + fn process(&mut self) { + let now = crate::time::ms( + self.clock.now() - self.inner.start, + crate::time::Round::Down, + ); + let mut poll = wheel::Poll::new(now); + + while let Some(entry) = self.wheel.poll(&mut poll, &mut ()) { + let when = entry.when_internal().expect("invalid internal entry state"); + + // Fire the entry + entry.fire(when); + + // Track that the entry has been fired + entry.set_when_internal(None); + } + + // Update the elapsed cache + self.inner.elapsed.store(self.wheel.elapsed(), SeqCst); + } + + /// Processes the entry queue + /// + /// This handles adding and canceling timeouts. + fn process_queue(&mut self) { + for entry in self.inner.process.take() { + match (entry.when_internal(), entry.load_state()) { + (None, None) => { + // Nothing to do + } + (Some(_), None) => { + // Remove the entry + self.clear_entry(&entry); + } + (None, Some(when)) => { + // Queue the entry + self.add_entry(entry, when); + } + (Some(_), Some(next)) => { + self.clear_entry(&entry); + self.add_entry(entry, next); + } + } + } + } + + fn clear_entry(&mut self, entry: &Arc<Entry>) { + self.wheel.remove(entry, &mut ()); + entry.set_when_internal(None); + } + + /// Fires the entry if it needs to, otherwise queue it to be processed later. + /// + /// Returns `None` if the entry was fired. + fn add_entry(&mut self, entry: Arc<Entry>, when: u64) { + use crate::time::wheel::InsertError; + + entry.set_when_internal(Some(when)); + + match self.wheel.insert(when, entry, &mut ()) { + Ok(_) => {} + Err((entry, InsertError::Elapsed)) => { + // The entry's deadline has elapsed, so fire it and update the + // internal state accordingly. + entry.set_when_internal(None); + entry.fire(when); + } + Err((entry, InsertError::Invalid)) => { + // The entry's deadline is invalid, so error it and update the + // internal state accordingly. + entry.set_when_internal(None); + entry.error(); + } + } + } +} + +impl<T> Park for Driver<T> +where + T: Park, +{ + type Unpark = T::Unpark; + type Error = T::Error; + + fn unpark(&self) -> Self::Unpark { + self.park.unpark() + } + + fn park(&mut self) -> Result<(), Self::Error> { + self.process_queue(); + + match self.wheel.poll_at() { + Some(when) => { + let now = self.clock.now(); + let deadline = self.expiration_instant(when); + + if deadline > now { + let dur = deadline - now; + + if self.clock.is_paused() { + self.park.park_timeout(Duration::from_secs(0))?; + self.clock.advance(dur); + } else { + self.park.park_timeout(dur)?; + } + } else { + self.park.park_timeout(Duration::from_secs(0))?; + } + } + None => { + self.park.park()?; + } + } + + self.process(); + + Ok(()) + } + + fn park_timeout(&mut self, duration: Duration) -> Result<(), Self::Error> { + self.process_queue(); + + match self.wheel.poll_at() { + Some(when) => { + let now = self.clock.now(); + let deadline = self.expiration_instant(when); + + if deadline > now { + let duration = cmp::min(deadline - now, duration); + + if self.clock.is_paused() { + self.park.park_timeout(Duration::from_secs(0))?; + self.clock.advance(duration); + } else { + self.park.park_timeout(duration)?; + } + } else { + self.park.park_timeout(Duration::from_secs(0))?; + } + } + None => { + self.park.park_timeout(duration)?; + } + } + + self.process(); + + Ok(()) + } +} + +impl<T> Drop for Driver<T> { + fn drop(&mut self) { + use std::u64; + + // Shutdown the stack of entries to process, preventing any new entries + // from being pushed. + self.inner.process.shutdown(); + + // Clear the wheel, using u64::MAX allows us to drain everything + let mut poll = wheel::Poll::new(u64::MAX); + + while let Some(entry) = self.wheel.poll(&mut poll, &mut ()) { + entry.error(); + } + } +} + +// ===== impl Inner ===== + +impl Inner { + fn new(start: Instant, unpark: Box<dyn Unpark>) -> Inner { + Inner { + num: AtomicUsize::new(0), + elapsed: AtomicU64::new(0), + process: AtomicStack::new(), + start, + unpark, + } + } + + fn elapsed(&self) -> u64 { + self.elapsed.load(SeqCst) + } + + #[cfg(all(test, loom))] + fn num(&self, ordering: std::sync::atomic::Ordering) -> usize { + self.num.load(ordering) + } + + /// Increments the number of active timeouts + fn increment(&self) -> Result<(), Error> { + let mut curr = self.num.load(Relaxed); + loop { + if curr == MAX_TIMEOUTS { + return Err(Error::at_capacity()); + } + + match self + .num + .compare_exchange_weak(curr, curr + 1, Release, Relaxed) + { + Ok(_) => return Ok(()), + Err(next) => curr = next, + } + } + } + + /// Decrements the number of active timeouts + fn decrement(&self) { + let prev = self.num.fetch_sub(1, Acquire); + debug_assert!(prev <= MAX_TIMEOUTS); + } + + fn queue(&self, entry: &Arc<Entry>) -> Result<(), Error> { + if self.process.push(entry)? { + // The timer is notified so that it can process the timeout + self.unpark.unpark(); + } + + Ok(()) + } + + fn normalize_deadline(&self, deadline: Instant) -> u64 { + if deadline < self.start { + return 0; + } + + crate::time::ms(deadline - self.start, crate::time::Round::Up) + } +} + +impl fmt::Debug for Inner { + fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt.debug_struct("Inner").finish() + } +} + +#[cfg(all(test, loom))] +mod tests; diff --git a/third_party/rust/tokio/src/time/driver/registration.rs b/third_party/rust/tokio/src/time/driver/registration.rs new file mode 100644 index 0000000000..b77357e735 --- /dev/null +++ b/third_party/rust/tokio/src/time/driver/registration.rs @@ -0,0 +1,53 @@ +use crate::time::driver::{Entry, Handle}; +use crate::time::{Duration, Error, Instant}; + +use std::sync::Arc; +use std::task::{self, Poll}; + +/// Registration with a timer. +/// +/// The association between a `Delay` instance and a timer is done lazily in +/// `poll` +#[derive(Debug)] +pub(crate) struct Registration { + entry: Arc<Entry>, +} + +impl Registration { + pub(crate) fn new(deadline: Instant, duration: Duration) -> Registration { + let handle = Handle::current(); + + Registration { + entry: Entry::new(&handle, deadline, duration), + } + } + + pub(crate) fn deadline(&self) -> Instant { + self.entry.time_ref().deadline + } + + pub(crate) fn reset(&mut self, deadline: Instant) { + unsafe { + self.entry.time_mut().deadline = deadline; + } + + Entry::reset(&mut self.entry); + } + + pub(crate) fn is_elapsed(&self) -> bool { + self.entry.is_elapsed() + } + + pub(crate) fn poll_elapsed(&self, cx: &mut task::Context<'_>) -> Poll<Result<(), Error>> { + // Keep track of task budget + ready!(crate::coop::poll_proceed(cx)); + + self.entry.poll_elapsed(cx) + } +} + +impl Drop for Registration { + fn drop(&mut self) { + Entry::cancel(&self.entry); + } +} diff --git a/third_party/rust/tokio/src/time/driver/stack.rs b/third_party/rust/tokio/src/time/driver/stack.rs new file mode 100644 index 0000000000..3e2924f265 --- /dev/null +++ b/third_party/rust/tokio/src/time/driver/stack.rs @@ -0,0 +1,121 @@ +use crate::time::driver::Entry; +use crate::time::wheel; + +use std::ptr; +use std::sync::Arc; + +/// A doubly linked stack +#[derive(Debug)] +pub(crate) struct Stack { + head: Option<Arc<Entry>>, +} + +impl Default for Stack { + fn default() -> Stack { + Stack { head: None } + } +} + +impl wheel::Stack for Stack { + type Owned = Arc<Entry>; + type Borrowed = Entry; + type Store = (); + + fn is_empty(&self) -> bool { + self.head.is_none() + } + + fn push(&mut self, entry: Self::Owned, _: &mut Self::Store) { + // Get a pointer to the entry to for the prev link + let ptr: *const Entry = &*entry as *const _; + + // Remove the old head entry + let old = self.head.take(); + + unsafe { + // Ensure the entry is not already in a stack. + debug_assert!((*entry.next_stack.get()).is_none()); + debug_assert!((*entry.prev_stack.get()).is_null()); + + if let Some(ref entry) = old.as_ref() { + debug_assert!({ + // The head is not already set to the entry + ptr != &***entry as *const _ + }); + + // Set the previous link on the old head + *entry.prev_stack.get() = ptr; + } + + // Set this entry's next pointer + *entry.next_stack.get() = old; + } + + // Update the head pointer + self.head = Some(entry); + } + + /// Pops an item from the stack + fn pop(&mut self, _: &mut ()) -> Option<Arc<Entry>> { + let entry = self.head.take(); + + unsafe { + if let Some(entry) = entry.as_ref() { + self.head = (*entry.next_stack.get()).take(); + + if let Some(entry) = self.head.as_ref() { + *entry.prev_stack.get() = ptr::null(); + } + + *entry.prev_stack.get() = ptr::null(); + } + } + + entry + } + + fn remove(&mut self, entry: &Entry, _: &mut ()) { + unsafe { + // Ensure that the entry is in fact contained by the stack + debug_assert!({ + // This walks the full linked list even if an entry is found. + let mut next = self.head.as_ref(); + let mut contains = false; + + while let Some(n) = next { + if entry as *const _ == &**n as *const _ { + debug_assert!(!contains); + contains = true; + } + + next = (*n.next_stack.get()).as_ref(); + } + + contains + }); + + // Unlink `entry` from the next node + let next = (*entry.next_stack.get()).take(); + + if let Some(next) = next.as_ref() { + (*next.prev_stack.get()) = *entry.prev_stack.get(); + } + + // Unlink `entry` from the prev node + + if let Some(prev) = (*entry.prev_stack.get()).as_ref() { + *prev.next_stack.get() = next; + } else { + // It is the head + self.head = next; + } + + // Unset the prev pointer + *entry.prev_stack.get() = ptr::null(); + } + } + + fn when(item: &Entry, _: &()) -> u64 { + item.when_internal().expect("invalid internal state") + } +} diff --git a/third_party/rust/tokio/src/time/driver/tests/mod.rs b/third_party/rust/tokio/src/time/driver/tests/mod.rs new file mode 100644 index 0000000000..88ff5525da --- /dev/null +++ b/third_party/rust/tokio/src/time/driver/tests/mod.rs @@ -0,0 +1,55 @@ +use crate::park::Unpark; +use crate::time::driver::Inner; +use crate::time::Instant; + +use loom::thread; + +use std::sync::atomic::Ordering; +use std::sync::Arc; + +struct MockUnpark; + +impl Unpark for MockUnpark { + fn unpark(&self) {} +} + +#[test] +fn balanced_incr_and_decr() { + const OPS: usize = 5; + + fn incr(inner: Arc<Inner>) { + for _ in 0..OPS { + inner.increment().expect("increment should not have failed"); + thread::yield_now(); + } + } + + fn decr(inner: Arc<Inner>) { + let mut ops_performed = 0; + while ops_performed < OPS { + if inner.num(Ordering::Relaxed) > 0 { + ops_performed += 1; + inner.decrement(); + } + thread::yield_now(); + } + } + + loom::model(|| { + let unpark = Box::new(MockUnpark); + let instant = Instant::now(); + + let inner = Arc::new(Inner::new(instant, unpark)); + + let incr_inner = inner.clone(); + let decr_inner = inner.clone(); + + let incr_hndle = thread::spawn(move || incr(incr_inner)); + let decr_hndle = thread::spawn(move || decr(decr_inner)); + + incr_hndle.join().expect("should never fail"); + decr_hndle.join().expect("should never fail"); + + assert_eq!(inner.num(Ordering::SeqCst), 0); + }) +} diff --git a/third_party/rust/tokio/src/time/error.rs b/third_party/rust/tokio/src/time/error.rs new file mode 100644 index 0000000000..0667b97ac1 --- /dev/null +++ b/third_party/rust/tokio/src/time/error.rs @@ -0,0 +1,72 @@ +use self::Kind::*; +use std::error; +use std::fmt; + +/// Errors encountered by the timer implementation. +/// +/// Currently, there are two different errors that can occur: +/// +/// * `shutdown` occurs when a timer operation is attempted, but the timer +/// instance has been dropped. In this case, the operation will never be able +/// to complete and the `shutdown` error is returned. This is a permanent +/// error, i.e., once this error is observed, timer operations will never +/// succeed in the future. +/// +/// * `at_capacity` occurs when a timer operation is attempted, but the timer +/// instance is currently handling its maximum number of outstanding delays. +/// In this case, the operation is not able to be performed at the current +/// moment, and `at_capacity` is returned. This is a transient error, i.e., at +/// some point in the future, if the operation is attempted again, it might +/// succeed. Callers that observe this error should attempt to [shed load]. One +/// way to do this would be dropping the future that issued the timer operation. +/// +/// [shed load]: https://en.wikipedia.org/wiki/Load_Shedding +#[derive(Debug)] +pub struct Error(Kind); + +#[derive(Debug)] +enum Kind { + Shutdown, + AtCapacity, +} + +impl Error { + /// Creates an error representing a shutdown timer. + pub fn shutdown() -> Error { + Error(Shutdown) + } + + /// Returns `true` if the error was caused by the timer being shutdown. + pub fn is_shutdown(&self) -> bool { + match self.0 { + Kind::Shutdown => true, + _ => false, + } + } + + /// Creates an error representing a timer at capacity. + pub fn at_capacity() -> Error { + Error(AtCapacity) + } + + /// Returns `true` if the error was caused by the timer being at capacity. + pub fn is_at_capacity(&self) -> bool { + match self.0 { + Kind::AtCapacity => true, + _ => false, + } + } +} + +impl error::Error for Error {} + +impl fmt::Display for Error { + fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { + use self::Kind::*; + let descr = match self.0 { + Shutdown => "the timer is shutdown, must be called from the context of Tokio runtime", + AtCapacity => "timer is at capacity and cannot create a new entry", + }; + write!(fmt, "{}", descr) + } +} diff --git a/third_party/rust/tokio/src/time/instant.rs b/third_party/rust/tokio/src/time/instant.rs new file mode 100644 index 0000000000..f2cb4bc97d --- /dev/null +++ b/third_party/rust/tokio/src/time/instant.rs @@ -0,0 +1,199 @@ +#![allow(clippy::trivially_copy_pass_by_ref)] + +use std::fmt; +use std::ops; +use std::time::Duration; + +/// A measurement of the system clock, useful for talking to +/// external entities like the file system or other processes. +#[derive(Clone, Copy, Eq, PartialEq, PartialOrd, Ord, Hash)] +pub struct Instant { + std: std::time::Instant, +} + +impl Instant { + /// Returns an instant corresponding to "now". + /// + /// # Examples + /// + /// ``` + /// use tokio::time::Instant; + /// + /// let now = Instant::now(); + /// ``` + pub fn now() -> Instant { + variant::now() + } + + /// Create a `tokio::time::Instant` from a `std::time::Instant`. + pub fn from_std(std: std::time::Instant) -> Instant { + Instant { std } + } + + /// Convert the value into a `std::time::Instant`. + pub fn into_std(self) -> std::time::Instant { + self.std + } + + /// Returns the amount of time elapsed from another instant to this one. + /// + /// # Panics + /// + /// This function will panic if `earlier` is later than `self`. + pub fn duration_since(&self, earlier: Instant) -> Duration { + self.std.duration_since(earlier.std) + } + + /// Returns the amount of time elapsed from another instant to this one, or + /// None if that instant is later than this one. + /// + /// # Examples + /// + /// ``` + /// use tokio::time::{Duration, Instant, delay_for}; + /// + /// #[tokio::main] + /// async fn main() { + /// let now = Instant::now(); + /// delay_for(Duration::new(1, 0)).await; + /// let new_now = Instant::now(); + /// println!("{:?}", new_now.checked_duration_since(now)); + /// println!("{:?}", now.checked_duration_since(new_now)); // None + /// } + /// ``` + pub fn checked_duration_since(&self, earlier: Instant) -> Option<Duration> { + self.std.checked_duration_since(earlier.std) + } + + /// Returns the amount of time elapsed from another instant to this one, or + /// zero duration if that instant is earlier than this one. + /// + /// # Examples + /// + /// ``` + /// use tokio::time::{Duration, Instant, delay_for}; + /// + /// #[tokio::main] + /// async fn main() { + /// let now = Instant::now(); + /// delay_for(Duration::new(1, 0)).await; + /// let new_now = Instant::now(); + /// println!("{:?}", new_now.saturating_duration_since(now)); + /// println!("{:?}", now.saturating_duration_since(new_now)); // 0ns + /// } + /// ``` + pub fn saturating_duration_since(&self, earlier: Instant) -> Duration { + self.std.saturating_duration_since(earlier.std) + } + + /// Returns the amount of time elapsed since this instant was created. + /// + /// # Panics + /// + /// This function may panic if the current time is earlier than this + /// instant, which is something that can happen if an `Instant` is + /// produced synthetically. + /// + /// # Examples + /// + /// ``` + /// use tokio::time::{Duration, Instant, delay_for}; + /// + /// #[tokio::main] + /// async fn main() { + /// let instant = Instant::now(); + /// let three_secs = Duration::from_secs(3); + /// delay_for(three_secs).await; + /// assert!(instant.elapsed() >= three_secs); + /// } + /// ``` + pub fn elapsed(&self) -> Duration { + Instant::now() - *self + } + + /// Returns `Some(t)` where `t` is the time `self + duration` if `t` can be + /// represented as `Instant` (which means it's inside the bounds of the + /// underlying data structure), `None` otherwise. + pub fn checked_add(&self, duration: Duration) -> Option<Instant> { + self.std.checked_add(duration).map(Instant::from_std) + } + + /// Returns `Some(t)` where `t` is the time `self - duration` if `t` can be + /// represented as `Instant` (which means it's inside the bounds of the + /// underlying data structure), `None` otherwise. + pub fn checked_sub(&self, duration: Duration) -> Option<Instant> { + self.std.checked_sub(duration).map(Instant::from_std) + } +} + +impl From<std::time::Instant> for Instant { + fn from(time: std::time::Instant) -> Instant { + Instant::from_std(time) + } +} + +impl From<Instant> for std::time::Instant { + fn from(time: Instant) -> std::time::Instant { + time.into_std() + } +} + +impl ops::Add<Duration> for Instant { + type Output = Instant; + + fn add(self, other: Duration) -> Instant { + Instant::from_std(self.std + other) + } +} + +impl ops::AddAssign<Duration> for Instant { + fn add_assign(&mut self, rhs: Duration) { + *self = *self + rhs; + } +} + +impl ops::Sub for Instant { + type Output = Duration; + + fn sub(self, rhs: Instant) -> Duration { + self.std - rhs.std + } +} + +impl ops::Sub<Duration> for Instant { + type Output = Instant; + + fn sub(self, rhs: Duration) -> Instant { + Instant::from_std(self.std - rhs) + } +} + +impl ops::SubAssign<Duration> for Instant { + fn sub_assign(&mut self, rhs: Duration) { + *self = *self - rhs; + } +} + +impl fmt::Debug for Instant { + fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { + self.std.fmt(fmt) + } +} + +#[cfg(not(feature = "test-util"))] +mod variant { + use super::Instant; + + pub(super) fn now() -> Instant { + Instant::from_std(std::time::Instant::now()) + } +} + +#[cfg(feature = "test-util")] +mod variant { + use super::Instant; + + pub(super) fn now() -> Instant { + crate::time::clock::now() + } +} diff --git a/third_party/rust/tokio/src/time/interval.rs b/third_party/rust/tokio/src/time/interval.rs new file mode 100644 index 0000000000..090e2d1f05 --- /dev/null +++ b/third_party/rust/tokio/src/time/interval.rs @@ -0,0 +1,139 @@ +use crate::future::poll_fn; +use crate::time::{delay_until, Delay, Duration, Instant}; + +use std::future::Future; +use std::pin::Pin; +use std::task::{Context, Poll}; + +/// Creates new `Interval` that yields with interval of `duration`. The first +/// tick completes immediately. +/// +/// An interval will tick indefinitely. At any time, the `Interval` value can be +/// dropped. This cancels the interval. +/// +/// This function is equivalent to `interval_at(Instant::now(), period)`. +/// +/// # Panics +/// +/// This function panics if `period` is zero. +/// +/// # Examples +/// +/// ``` +/// use tokio::time::{self, Duration}; +/// +/// #[tokio::main] +/// async fn main() { +/// let mut interval = time::interval(Duration::from_millis(10)); +/// +/// interval.tick().await; +/// interval.tick().await; +/// interval.tick().await; +/// +/// // approximately 20ms have elapsed. +/// } +/// ``` +pub fn interval(period: Duration) -> Interval { + assert!(period > Duration::new(0, 0), "`period` must be non-zero."); + + interval_at(Instant::now(), period) +} + +/// Creates new `Interval` that yields with interval of `period` with the +/// first tick completing at `at`. +/// +/// An interval will tick indefinitely. At any time, the `Interval` value can be +/// dropped. This cancels the interval. +/// +/// # Panics +/// +/// This function panics if `period` is zero. +/// +/// # Examples +/// +/// ``` +/// use tokio::time::{interval_at, Duration, Instant}; +/// +/// #[tokio::main] +/// async fn main() { +/// let start = Instant::now() + Duration::from_millis(50); +/// let mut interval = interval_at(start, Duration::from_millis(10)); +/// +/// interval.tick().await; +/// interval.tick().await; +/// interval.tick().await; +/// +/// // approximately 70ms have elapsed. +/// } +/// ``` +pub fn interval_at(start: Instant, period: Duration) -> Interval { + assert!(period > Duration::new(0, 0), "`period` must be non-zero."); + + Interval { + delay: delay_until(start), + period, + } +} + +/// Stream returned by [`interval`](interval) and [`interval_at`](interval_at). +#[derive(Debug)] +pub struct Interval { + /// Future that completes the next time the `Interval` yields a value. + delay: Delay, + + /// The duration between values yielded by `Interval`. + period: Duration, +} + +impl Interval { + #[doc(hidden)] // TODO: document + pub fn poll_tick(&mut self, cx: &mut Context<'_>) -> Poll<Instant> { + // Wait for the delay to be done + ready!(Pin::new(&mut self.delay).poll(cx)); + + // Get the `now` by looking at the `delay` deadline + let now = self.delay.deadline(); + + // The next interval value is `duration` after the one that just + // yielded. + let next = now + self.period; + self.delay.reset(next); + + // Return the current instant + Poll::Ready(now) + } + + /// Completes when the next instant in the interval has been reached. + /// + /// # Examples + /// + /// ``` + /// use tokio::time; + /// + /// use std::time::Duration; + /// + /// #[tokio::main] + /// async fn main() { + /// let mut interval = time::interval(Duration::from_millis(10)); + /// + /// interval.tick().await; + /// interval.tick().await; + /// interval.tick().await; + /// + /// // approximately 20ms have elapsed. + /// } + /// ``` + #[allow(clippy::should_implement_trait)] // TODO: rename (tokio-rs/tokio#1261) + pub async fn tick(&mut self) -> Instant { + poll_fn(|cx| self.poll_tick(cx)).await + } +} + +#[cfg(feature = "stream")] +impl crate::stream::Stream for Interval { + type Item = Instant; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Instant>> { + Poll::Ready(Some(ready!(self.poll_tick(cx)))) + } +} diff --git a/third_party/rust/tokio/src/time/mod.rs b/third_party/rust/tokio/src/time/mod.rs new file mode 100644 index 0000000000..7070d6b257 --- /dev/null +++ b/third_party/rust/tokio/src/time/mod.rs @@ -0,0 +1,130 @@ +//! Utilities for tracking time. +//! +//! This module provides a number of types for executing code after a set period +//! of time. +//! +//! * `Delay` is a future that does no work and completes at a specific `Instant` +//! in time. +//! +//! * `Interval` is a stream yielding a value at a fixed period. It is +//! initialized with a `Duration` and repeatedly yields each time the duration +//! elapses. +//! +//! * `Timeout`: Wraps a future or stream, setting an upper bound to the amount +//! of time it is allowed to execute. If the future or stream does not +//! complete in time, then it is canceled and an error is returned. +//! +//! * `DelayQueue`: A queue where items are returned once the requested delay +//! has expired. +//! +//! These types are sufficient for handling a large number of scenarios +//! involving time. +//! +//! These types must be used from within the context of the `Runtime`. +//! +//! # Examples +//! +//! Wait 100ms and print "Hello World!" +//! +//! ``` +//! use tokio::time::delay_for; +//! +//! use std::time::Duration; +//! +//! +//! #[tokio::main] +//! async fn main() { +//! delay_for(Duration::from_millis(100)).await; +//! println!("100 ms have elapsed"); +//! } +//! ``` +//! +//! Require that an operation takes no more than 300ms. Note that this uses the +//! `timeout` function on the `FutureExt` trait. This trait is included in the +//! prelude. +//! +//! ``` +//! use tokio::time::{timeout, Duration}; +//! +//! async fn long_future() { +//! // do work here +//! } +//! +//! # async fn dox() { +//! let res = timeout(Duration::from_secs(1), long_future()).await; +//! +//! if res.is_err() { +//! println!("operation timed out"); +//! } +//! # } +//! ``` + +mod clock; +pub(crate) use self::clock::Clock; +#[cfg(feature = "test-util")] +pub use clock::{advance, pause, resume}; + +pub mod delay_queue; +#[doc(inline)] +pub use delay_queue::DelayQueue; + +mod delay; +pub use delay::{delay_for, delay_until, Delay}; + +pub(crate) mod driver; + +mod error; +pub use error::Error; + +mod instant; +pub use self::instant::Instant; + +mod interval; +pub use interval::{interval, interval_at, Interval}; + +mod timeout; +#[doc(inline)] +pub use timeout::{timeout, timeout_at, Elapsed, Timeout}; + +cfg_stream! { + mod throttle; + pub use throttle::{throttle, Throttle}; +} + +mod wheel; + +#[cfg(test)] +#[cfg(not(loom))] +mod tests; + +// Re-export for convenience +pub use std::time::Duration; + +// ===== Internal utils ===== + +enum Round { + Up, + Down, +} + +/// Convert a `Duration` to milliseconds, rounding up and saturating at +/// `u64::MAX`. +/// +/// The saturating is fine because `u64::MAX` milliseconds are still many +/// million years. +#[inline] +fn ms(duration: Duration, round: Round) -> u64 { + const NANOS_PER_MILLI: u32 = 1_000_000; + const MILLIS_PER_SEC: u64 = 1_000; + + // Round up. + let millis = match round { + Round::Up => (duration.subsec_nanos() + NANOS_PER_MILLI - 1) / NANOS_PER_MILLI, + Round::Down => duration.subsec_millis(), + }; + + duration + .as_secs() + .saturating_mul(MILLIS_PER_SEC) + .saturating_add(u64::from(millis)) +} diff --git a/third_party/rust/tokio/src/time/tests/mod.rs b/third_party/rust/tokio/src/time/tests/mod.rs new file mode 100644 index 0000000000..4710d470f7 --- /dev/null +++ b/third_party/rust/tokio/src/time/tests/mod.rs @@ -0,0 +1,22 @@ +mod test_delay; + +use crate::time::{self, Instant}; +use std::time::Duration; + +fn assert_send<T: Send>() {} +fn assert_sync<T: Sync>() {} + +#[test] +fn registration_is_send_and_sync() { + use crate::time::driver::Registration; + + assert_send::<Registration>(); + assert_sync::<Registration>(); +} + +#[test] +#[should_panic] +fn delay_is_eager() { + let when = Instant::now() + Duration::from_millis(100); + let _ = time::delay_until(when); +} diff --git a/third_party/rust/tokio/src/time/tests/test_delay.rs b/third_party/rust/tokio/src/time/tests/test_delay.rs new file mode 100644 index 0000000000..f843434be4 --- /dev/null +++ b/third_party/rust/tokio/src/time/tests/test_delay.rs @@ -0,0 +1,447 @@ +#![warn(rust_2018_idioms)] + +use crate::park::{Park, Unpark}; +use crate::time::driver::{Driver, Entry, Handle}; +use crate::time::Clock; +use crate::time::{Duration, Instant}; + +use tokio_test::task; +use tokio_test::{assert_ok, assert_pending, assert_ready_ok}; + +use std::sync::Arc; + +macro_rules! poll { + ($e:expr) => { + $e.enter(|cx, e| e.poll_elapsed(cx)) + }; +} + +#[test] +fn frozen_utility_returns_correct_advanced_duration() { + let clock = Clock::new(); + clock.pause(); + let start = clock.now(); + + clock.advance(ms(10)); + assert_eq!(clock.now() - start, ms(10)); +} + +#[test] +fn immediate_delay() { + let (mut driver, clock, handle) = setup(); + let start = clock.now(); + + let when = clock.now(); + let mut e = task::spawn(delay_until(&handle, when)); + + assert_ready_ok!(poll!(e)); + + assert_ok!(driver.park_timeout(Duration::from_millis(1000))); + + // The time has not advanced. The `turn` completed immediately. + assert_eq!(clock.now() - start, ms(1000)); +} + +#[test] +fn delayed_delay_level_0() { + let (mut driver, clock, handle) = setup(); + let start = clock.now(); + + for &i in &[1, 10, 60] { + // Create a `Delay` that elapses in the future + let mut e = task::spawn(delay_until(&handle, start + ms(i))); + + // The delay has not elapsed. + assert_pending!(poll!(e)); + + assert_ok!(driver.park()); + assert_eq!(clock.now() - start, ms(i)); + + assert_ready_ok!(poll!(e)); + } +} + +#[test] +fn sub_ms_delayed_delay() { + let (mut driver, clock, handle) = setup(); + + for _ in 0..5 { + let deadline = clock.now() + ms(1) + Duration::new(0, 1); + + let mut e = task::spawn(delay_until(&handle, deadline)); + + assert_pending!(poll!(e)); + + assert_ok!(driver.park()); + assert_ready_ok!(poll!(e)); + + assert!(clock.now() >= deadline); + + clock.advance(Duration::new(0, 1)); + } +} + +#[test] +fn delayed_delay_wrapping_level_0() { + let (mut driver, clock, handle) = setup(); + let start = clock.now(); + + assert_ok!(driver.park_timeout(ms(5))); + assert_eq!(clock.now() - start, ms(5)); + + let mut e = task::spawn(delay_until(&handle, clock.now() + ms(60))); + + assert_pending!(poll!(e)); + + assert_ok!(driver.park()); + assert_eq!(clock.now() - start, ms(64)); + assert_pending!(poll!(e)); + + assert_ok!(driver.park()); + assert_eq!(clock.now() - start, ms(65)); + + assert_ready_ok!(poll!(e)); +} + +#[test] +fn timer_wrapping_with_higher_levels() { + let (mut driver, clock, handle) = setup(); + let start = clock.now(); + + // Set delay to hit level 1 + let mut e1 = task::spawn(delay_until(&handle, clock.now() + ms(64))); + assert_pending!(poll!(e1)); + + // Turn a bit + assert_ok!(driver.park_timeout(ms(5))); + + // Set timeout such that it will hit level 0, but wrap + let mut e2 = task::spawn(delay_until(&handle, clock.now() + ms(60))); + assert_pending!(poll!(e2)); + + // This should result in s1 firing + assert_ok!(driver.park()); + assert_eq!(clock.now() - start, ms(64)); + + assert_ready_ok!(poll!(e1)); + assert_pending!(poll!(e2)); + + assert_ok!(driver.park()); + assert_eq!(clock.now() - start, ms(65)); + + assert_ready_ok!(poll!(e1)); +} + +#[test] +fn delay_with_deadline_in_past() { + let (mut driver, clock, handle) = setup(); + let start = clock.now(); + + // Create `Delay` that elapsed immediately. + let mut e = task::spawn(delay_until(&handle, clock.now() - ms(100))); + + // Even though the delay expires in the past, it is not ready yet + // because the timer must observe it. + assert_ready_ok!(poll!(e)); + + // Turn the timer, it runs for the elapsed time + assert_ok!(driver.park_timeout(ms(1000))); + + // The time has not advanced. The `turn` completed immediately. + assert_eq!(clock.now() - start, ms(1000)); +} + +#[test] +fn delayed_delay_level_1() { + let (mut driver, clock, handle) = setup(); + let start = clock.now(); + + // Create a `Delay` that elapses in the future + let mut e = task::spawn(delay_until(&handle, clock.now() + ms(234))); + + // The delay has not elapsed. + assert_pending!(poll!(e)); + + // Turn the timer, this will wake up to cascade the timer down. + assert_ok!(driver.park_timeout(ms(1000))); + assert_eq!(clock.now() - start, ms(192)); + + // The delay has not elapsed. + assert_pending!(poll!(e)); + + // Turn the timer again + assert_ok!(driver.park_timeout(ms(1000))); + assert_eq!(clock.now() - start, ms(234)); + + // The delay has elapsed. + assert_ready_ok!(poll!(e)); + + let (mut driver, clock, handle) = setup(); + let start = clock.now(); + + // Create a `Delay` that elapses in the future + let mut e = task::spawn(delay_until(&handle, clock.now() + ms(234))); + + // The delay has not elapsed. + assert_pending!(poll!(e)); + + // Turn the timer with a smaller timeout than the cascade. + assert_ok!(driver.park_timeout(ms(100))); + assert_eq!(clock.now() - start, ms(100)); + + assert_pending!(poll!(e)); + + // Turn the timer, this will wake up to cascade the timer down. + assert_ok!(driver.park_timeout(ms(1000))); + assert_eq!(clock.now() - start, ms(192)); + + // The delay has not elapsed. + assert_pending!(poll!(e)); + + // Turn the timer again + assert_ok!(driver.park_timeout(ms(1000))); + assert_eq!(clock.now() - start, ms(234)); + + // The delay has elapsed. + assert_ready_ok!(poll!(e)); +} + +#[test] +fn concurrently_set_two_timers_second_one_shorter() { + let (mut driver, clock, handle) = setup(); + let start = clock.now(); + + let mut e1 = task::spawn(delay_until(&handle, clock.now() + ms(500))); + let mut e2 = task::spawn(delay_until(&handle, clock.now() + ms(200))); + + // The delay has not elapsed + assert_pending!(poll!(e1)); + assert_pending!(poll!(e2)); + + // Delay until a cascade + assert_ok!(driver.park()); + assert_eq!(clock.now() - start, ms(192)); + + // Delay until the second timer. + assert_ok!(driver.park()); + assert_eq!(clock.now() - start, ms(200)); + + // The shorter delay fires + assert_ready_ok!(poll!(e2)); + assert_pending!(poll!(e1)); + + assert_ok!(driver.park()); + assert_eq!(clock.now() - start, ms(448)); + + assert_pending!(poll!(e1)); + + // Turn again, this time the time will advance to the second delay + assert_ok!(driver.park()); + assert_eq!(clock.now() - start, ms(500)); + + assert_ready_ok!(poll!(e1)); +} + +#[test] +fn short_delay() { + let (mut driver, clock, handle) = setup(); + let start = clock.now(); + + // Create a `Delay` that elapses in the future + let mut e = task::spawn(delay_until(&handle, clock.now() + ms(1))); + + // The delay has not elapsed. + assert_pending!(poll!(e)); + + // Turn the timer, but not enough time will go by. + assert_ok!(driver.park()); + + // The delay has elapsed. + assert_ready_ok!(poll!(e)); + + // The time has advanced to the point of the delay elapsing. + assert_eq!(clock.now() - start, ms(1)); +} + +#[test] +fn sorta_long_delay_until() { + const MIN_5: u64 = 5 * 60 * 1000; + + let (mut driver, clock, handle) = setup(); + let start = clock.now(); + + // Create a `Delay` that elapses in the future + let mut e = task::spawn(delay_until(&handle, clock.now() + ms(MIN_5))); + + // The delay has not elapsed. + assert_pending!(poll!(e)); + + let cascades = &[262_144, 262_144 + 9 * 4096, 262_144 + 9 * 4096 + 15 * 64]; + + for &elapsed in cascades { + assert_ok!(driver.park()); + assert_eq!(clock.now() - start, ms(elapsed)); + + assert_pending!(poll!(e)); + } + + assert_ok!(driver.park()); + assert_eq!(clock.now() - start, ms(MIN_5)); + + // The delay has elapsed. + assert_ready_ok!(poll!(e)); +} + +#[test] +fn very_long_delay() { + const MO_5: u64 = 5 * 30 * 24 * 60 * 60 * 1000; + + let (mut driver, clock, handle) = setup(); + let start = clock.now(); + + // Create a `Delay` that elapses in the future + let mut e = task::spawn(delay_until(&handle, clock.now() + ms(MO_5))); + + // The delay has not elapsed. + assert_pending!(poll!(e)); + + let cascades = &[ + 12_884_901_888, + 12_952_010_752, + 12_959_875_072, + 12_959_997_952, + ]; + + for &elapsed in cascades { + assert_ok!(driver.park()); + assert_eq!(clock.now() - start, ms(elapsed)); + + assert_pending!(poll!(e)); + } + + // Turn the timer, but not enough time will go by. + assert_ok!(driver.park()); + + // The time has advanced to the point of the delay elapsing. + assert_eq!(clock.now() - start, ms(MO_5)); + + // The delay has elapsed. + assert_ready_ok!(poll!(e)); +} + +#[test] +fn unpark_is_delayed() { + // A special park that will take much longer than the requested duration + struct MockPark(Clock); + + struct MockUnpark; + + impl Park for MockPark { + type Unpark = MockUnpark; + type Error = (); + + fn unpark(&self) -> Self::Unpark { + MockUnpark + } + + fn park(&mut self) -> Result<(), Self::Error> { + panic!("parking forever"); + } + + fn park_timeout(&mut self, duration: Duration) -> Result<(), Self::Error> { + assert_eq!(duration, ms(0)); + self.0.advance(ms(436)); + Ok(()) + } + } + + impl Unpark for MockUnpark { + fn unpark(&self) {} + } + + let clock = Clock::new(); + clock.pause(); + let start = clock.now(); + let mut driver = Driver::new(MockPark(clock.clone()), clock.clone()); + let handle = driver.handle(); + + let mut e1 = task::spawn(delay_until(&handle, clock.now() + ms(100))); + let mut e2 = task::spawn(delay_until(&handle, clock.now() + ms(101))); + let mut e3 = task::spawn(delay_until(&handle, clock.now() + ms(200))); + + assert_pending!(poll!(e1)); + assert_pending!(poll!(e2)); + assert_pending!(poll!(e3)); + + assert_ok!(driver.park()); + + assert_eq!(clock.now() - start, ms(500)); + + assert_ready_ok!(poll!(e1)); + assert_ready_ok!(poll!(e2)); + assert_ready_ok!(poll!(e3)); +} + +#[test] +fn set_timeout_at_deadline_greater_than_max_timer() { + const YR_1: u64 = 365 * 24 * 60 * 60 * 1000; + const YR_5: u64 = 5 * YR_1; + + let (mut driver, clock, handle) = setup(); + let start = clock.now(); + + for _ in 0..5 { + assert_ok!(driver.park_timeout(ms(YR_1))); + } + + let mut e = task::spawn(delay_until(&handle, clock.now() + ms(1))); + assert_pending!(poll!(e)); + + assert_ok!(driver.park_timeout(ms(1000))); + assert_eq!(clock.now() - start, ms(YR_5) + ms(1)); + + assert_ready_ok!(poll!(e)); +} + +fn setup() -> (Driver<MockPark>, Clock, Handle) { + let clock = Clock::new(); + clock.pause(); + let driver = Driver::new(MockPark(clock.clone()), clock.clone()); + let handle = driver.handle(); + + (driver, clock, handle) +} + +fn delay_until(handle: &Handle, when: Instant) -> Arc<Entry> { + Entry::new(&handle, when, ms(0)) +} + +struct MockPark(Clock); + +struct MockUnpark; + +impl Park for MockPark { + type Unpark = MockUnpark; + type Error = (); + + fn unpark(&self) -> Self::Unpark { + MockUnpark + } + + fn park(&mut self) -> Result<(), Self::Error> { + panic!("parking forever"); + } + + fn park_timeout(&mut self, duration: Duration) -> Result<(), Self::Error> { + self.0.advance(duration); + Ok(()) + } +} + +impl Unpark for MockUnpark { + fn unpark(&self) {} +} + +fn ms(n: u64) -> Duration { + Duration::from_millis(n) +} diff --git a/third_party/rust/tokio/src/time/throttle.rs b/third_party/rust/tokio/src/time/throttle.rs new file mode 100644 index 0000000000..435bef6381 --- /dev/null +++ b/third_party/rust/tokio/src/time/throttle.rs @@ -0,0 +1,117 @@ +//! Slow down a stream by enforcing a delay between items. + +use crate::stream::Stream; +use crate::time::{Delay, Duration, Instant}; + +use std::future::Future; +use std::marker::Unpin; +use std::pin::Pin; +use std::task::{self, Poll}; + +use pin_project_lite::pin_project; + +/// Slows down a stream by enforcing a delay between items. +/// They will be produced not more often than the specified interval. +/// +/// # Example +/// +/// Create a throttled stream. +/// ```rust,norun +/// use std::time::Duration; +/// use tokio::stream::StreamExt; +/// use tokio::time::throttle; +/// +/// # async fn dox() { +/// let mut item_stream = throttle(Duration::from_secs(2), futures::stream::repeat("one")); +/// +/// loop { +/// // The string will be produced at most every 2 seconds +/// println!("{:?}", item_stream.next().await); +/// } +/// # } +/// ``` +pub fn throttle<T>(duration: Duration, stream: T) -> Throttle<T> +where + T: Stream, +{ + let delay = if duration == Duration::from_millis(0) { + None + } else { + Some(Delay::new_timeout(Instant::now() + duration, duration)) + }; + + Throttle { + delay, + duration, + has_delayed: true, + stream, + } +} + +pin_project! { + /// Stream for the [`throttle`](throttle) function. + #[derive(Debug)] + #[must_use = "streams do nothing unless polled"] + pub struct Throttle<T> { + // `None` when duration is zero. + delay: Option<Delay>, + duration: Duration, + + // Set to true when `delay` has returned ready, but `stream` hasn't. + has_delayed: bool, + + // The stream to throttle + #[pin] + stream: T, + } +} + +// XXX: are these safe if `T: !Unpin`? +impl<T: Unpin> Throttle<T> { + /// Acquires a reference to the underlying stream that this combinator is + /// pulling from. + pub fn get_ref(&self) -> &T { + &self.stream + } + + /// Acquires a mutable reference to the underlying stream that this combinator + /// is pulling from. + /// + /// Note that care must be taken to avoid tampering with the state of the stream + /// which may otherwise confuse this combinator. + pub fn get_mut(&mut self) -> &mut T { + &mut self.stream + } + + /// Consumes this combinator, returning the underlying stream. + /// + /// Note that this may discard intermediate state of this combinator, so care + /// should be taken to avoid losing resources when this is called. + pub fn into_inner(self) -> T { + self.stream + } +} + +impl<T: Stream> Stream for Throttle<T> { + type Item = T::Item; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Option<Self::Item>> { + if !self.has_delayed && self.delay.is_some() { + ready!(Pin::new(self.as_mut().project().delay.as_mut().unwrap()).poll(cx)); + *self.as_mut().project().has_delayed = true; + } + + let value = ready!(self.as_mut().project().stream.poll_next(cx)); + + if value.is_some() { + let dur = self.duration; + if let Some(ref mut delay) = self.as_mut().project().delay { + delay.reset(Instant::now() + dur); + } + + *self.as_mut().project().has_delayed = false; + } + + Poll::Ready(value) + } +} diff --git a/third_party/rust/tokio/src/time/timeout.rs b/third_party/rust/tokio/src/time/timeout.rs new file mode 100644 index 0000000000..401856a881 --- /dev/null +++ b/third_party/rust/tokio/src/time/timeout.rs @@ -0,0 +1,185 @@ +//! Allows a future to execute for a maximum amount of time. +//! +//! See [`Timeout`] documentation for more details. +//! +//! [`Timeout`]: struct@Timeout + +use crate::time::{delay_until, Delay, Duration, Instant}; + +use std::fmt; +use std::future::Future; +use std::pin::Pin; +use std::task::{self, Poll}; + +/// Require a `Future` to complete before the specified duration has elapsed. +/// +/// If the future completes before the duration has elapsed, then the completed +/// value is returned. Otherwise, an error is returned and the future is +/// canceled. +/// +/// # Cancelation +/// +/// Cancelling a timeout is done by dropping the future. No additional cleanup +/// or other work is required. +/// +/// The original future may be obtained by calling [`Timeout::into_inner`]. This +/// consumes the `Timeout`. +/// +/// # Examples +/// +/// Create a new `Timeout` set to expire in 10 milliseconds. +/// +/// ```rust +/// use tokio::time::timeout; +/// use tokio::sync::oneshot; +/// +/// use std::time::Duration; +/// +/// # async fn dox() { +/// let (tx, rx) = oneshot::channel(); +/// # tx.send(()).unwrap(); +/// +/// // Wrap the future with a `Timeout` set to expire in 10 milliseconds. +/// if let Err(_) = timeout(Duration::from_millis(10), rx).await { +/// println!("did not receive value within 10 ms"); +/// } +/// # } +/// ``` +pub fn timeout<T>(duration: Duration, future: T) -> Timeout<T> +where + T: Future, +{ + let delay = Delay::new_timeout(Instant::now() + duration, duration); + Timeout::new_with_delay(future, delay) +} + +/// Require a `Future` to complete before the specified instant in time. +/// +/// If the future completes before the instant is reached, then the completed +/// value is returned. Otherwise, an error is returned. +/// +/// # Cancelation +/// +/// Cancelling a timeout is done by dropping the future. No additional cleanup +/// or other work is required. +/// +/// The original future may be obtained by calling [`Timeout::into_inner`]. This +/// consumes the `Timeout`. +/// +/// # Examples +/// +/// Create a new `Timeout` set to expire in 10 milliseconds. +/// +/// ```rust +/// use tokio::time::{Instant, timeout_at}; +/// use tokio::sync::oneshot; +/// +/// use std::time::Duration; +/// +/// # async fn dox() { +/// let (tx, rx) = oneshot::channel(); +/// # tx.send(()).unwrap(); +/// +/// // Wrap the future with a `Timeout` set to expire 10 milliseconds into the +/// // future. +/// if let Err(_) = timeout_at(Instant::now() + Duration::from_millis(10), rx).await { +/// println!("did not receive value within 10 ms"); +/// } +/// # } +/// ``` +pub fn timeout_at<T>(deadline: Instant, future: T) -> Timeout<T> +where + T: Future, +{ + let delay = delay_until(deadline); + + Timeout { + value: future, + delay, + } +} + +/// Future returned by [`timeout`](timeout) and [`timeout_at`](timeout_at). +#[must_use = "futures do nothing unless you `.await` or poll them"] +#[derive(Debug)] +pub struct Timeout<T> { + value: T, + delay: Delay, +} + +/// Error returned by `Timeout`. +#[derive(Debug, PartialEq)] +pub struct Elapsed(()); + +impl Elapsed { + // Used on StreamExt::timeout + #[allow(unused)] + pub(crate) fn new() -> Self { + Elapsed(()) + } +} + +impl<T> Timeout<T> { + pub(crate) fn new_with_delay(value: T, delay: Delay) -> Timeout<T> { + Timeout { value, delay } + } + + /// Gets a reference to the underlying value in this timeout. + pub fn get_ref(&self) -> &T { + &self.value + } + + /// Gets a mutable reference to the underlying value in this timeout. + pub fn get_mut(&mut self) -> &mut T { + &mut self.value + } + + /// Consumes this timeout, returning the underlying value. + pub fn into_inner(self) -> T { + self.value + } +} + +impl<T> Future for Timeout<T> +where + T: Future, +{ + type Output = Result<T::Output, Elapsed>; + + fn poll(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> { + // First, try polling the future + + // Safety: we never move `self.value` + unsafe { + let p = self.as_mut().map_unchecked_mut(|me| &mut me.value); + if let Poll::Ready(v) = p.poll(cx) { + return Poll::Ready(Ok(v)); + } + } + + // Now check the timer + // Safety: X_X! + unsafe { + match self.map_unchecked_mut(|me| &mut me.delay).poll(cx) { + Poll::Ready(()) => Poll::Ready(Err(Elapsed(()))), + Poll::Pending => Poll::Pending, + } + } + } +} + +// ===== impl Elapsed ===== + +impl fmt::Display for Elapsed { + fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { + "deadline has elapsed".fmt(fmt) + } +} + +impl std::error::Error for Elapsed {} + +impl From<Elapsed> for std::io::Error { + fn from(_err: Elapsed) -> std::io::Error { + std::io::ErrorKind::TimedOut.into() + } +} diff --git a/third_party/rust/tokio/src/time/wheel/level.rs b/third_party/rust/tokio/src/time/wheel/level.rs new file mode 100644 index 0000000000..49f9bfb9cf --- /dev/null +++ b/third_party/rust/tokio/src/time/wheel/level.rs @@ -0,0 +1,255 @@ +use crate::time::wheel::Stack; + +use std::fmt; + +/// Wheel for a single level in the timer. This wheel contains 64 slots. +pub(crate) struct Level<T> { + level: usize, + + /// Bit field tracking which slots currently contain entries. + /// + /// Using a bit field to track slots that contain entries allows avoiding a + /// scan to find entries. This field is updated when entries are added or + /// removed from a slot. + /// + /// The least-significant bit represents slot zero. + occupied: u64, + + /// Slots + slot: [T; LEVEL_MULT], +} + +/// Indicates when a slot must be processed next. +#[derive(Debug)] +pub(crate) struct Expiration { + /// The level containing the slot. + pub(crate) level: usize, + + /// The slot index. + pub(crate) slot: usize, + + /// The instant at which the slot needs to be processed. + pub(crate) deadline: u64, +} + +/// Level multiplier. +/// +/// Being a power of 2 is very important. +const LEVEL_MULT: usize = 64; + +impl<T: Stack> Level<T> { + pub(crate) fn new(level: usize) -> Level<T> { + // Rust's derived implementations for arrays require that the value + // contained by the array be `Copy`. So, here we have to manually + // initialize every single slot. + macro_rules! s { + () => { + T::default() + }; + }; + + Level { + level, + occupied: 0, + slot: [ + // It does not look like the necessary traits are + // derived for [T; 64]. + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + s!(), + ], + } + } + + /// Finds the slot that needs to be processed next and returns the slot and + /// `Instant` at which this slot must be processed. + pub(crate) fn next_expiration(&self, now: u64) -> Option<Expiration> { + // Use the `occupied` bit field to get the index of the next slot that + // needs to be processed. + let slot = match self.next_occupied_slot(now) { + Some(slot) => slot, + None => return None, + }; + + // From the slot index, calculate the `Instant` at which it needs to be + // processed. This value *must* be in the future with respect to `now`. + + let level_range = level_range(self.level); + let slot_range = slot_range(self.level); + + // TODO: This can probably be simplified w/ power of 2 math + let level_start = now - (now % level_range); + let deadline = level_start + slot as u64 * slot_range; + + debug_assert!( + deadline >= now, + "deadline={}; now={}; level={}; slot={}; occupied={:b}", + deadline, + now, + self.level, + slot, + self.occupied + ); + + Some(Expiration { + level: self.level, + slot, + deadline, + }) + } + + fn next_occupied_slot(&self, now: u64) -> Option<usize> { + if self.occupied == 0 { + return None; + } + + // Get the slot for now using Maths + let now_slot = (now / slot_range(self.level)) as usize; + let occupied = self.occupied.rotate_right(now_slot as u32); + let zeros = occupied.trailing_zeros() as usize; + let slot = (zeros + now_slot) % 64; + + Some(slot) + } + + pub(crate) fn add_entry(&mut self, when: u64, item: T::Owned, store: &mut T::Store) { + let slot = slot_for(when, self.level); + + self.slot[slot].push(item, store); + self.occupied |= occupied_bit(slot); + } + + pub(crate) fn remove_entry(&mut self, when: u64, item: &T::Borrowed, store: &mut T::Store) { + let slot = slot_for(when, self.level); + + self.slot[slot].remove(item, store); + + if self.slot[slot].is_empty() { + // The bit is currently set + debug_assert!(self.occupied & occupied_bit(slot) != 0); + + // Unset the bit + self.occupied ^= occupied_bit(slot); + } + } + + pub(crate) fn pop_entry_slot(&mut self, slot: usize, store: &mut T::Store) -> Option<T::Owned> { + let ret = self.slot[slot].pop(store); + + if ret.is_some() && self.slot[slot].is_empty() { + // The bit is currently set + debug_assert!(self.occupied & occupied_bit(slot) != 0); + + self.occupied ^= occupied_bit(slot); + } + + ret + } +} + +impl<T> fmt::Debug for Level<T> { + fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt.debug_struct("Level") + .field("occupied", &self.occupied) + .finish() + } +} + +fn occupied_bit(slot: usize) -> u64 { + 1 << slot +} + +fn slot_range(level: usize) -> u64 { + LEVEL_MULT.pow(level as u32) as u64 +} + +fn level_range(level: usize) -> u64 { + LEVEL_MULT as u64 * slot_range(level) +} + +/// Convert a duration (milliseconds) and a level to a slot position +fn slot_for(duration: u64, level: usize) -> usize { + ((duration >> (level * 6)) % LEVEL_MULT as u64) as usize +} + +/* +#[cfg(all(test, not(loom)))] +mod test { + use super::*; + + #[test] + fn test_slot_for() { + for pos in 1..64 { + assert_eq!(pos as usize, slot_for(pos, 0)); + } + + for level in 1..5 { + for pos in level..64 { + let a = pos * 64_usize.pow(level as u32); + assert_eq!(pos as usize, slot_for(a as u64, level)); + } + } + } +} +*/ diff --git a/third_party/rust/tokio/src/time/wheel/mod.rs b/third_party/rust/tokio/src/time/wheel/mod.rs new file mode 100644 index 0000000000..a2ef27fc6c --- /dev/null +++ b/third_party/rust/tokio/src/time/wheel/mod.rs @@ -0,0 +1,314 @@ +mod level; +pub(crate) use self::level::Expiration; +use self::level::Level; + +mod stack; +pub(crate) use self::stack::Stack; + +use std::borrow::Borrow; +use std::usize; + +/// Timing wheel implementation. +/// +/// This type provides the hashed timing wheel implementation that backs `Timer` +/// and `DelayQueue`. +/// +/// The structure is generic over `T: Stack`. This allows handling timeout data +/// being stored on the heap or in a slab. In order to support the latter case, +/// the slab must be passed into each function allowing the implementation to +/// lookup timer entries. +/// +/// See `Timer` documentation for some implementation notes. +#[derive(Debug)] +pub(crate) struct Wheel<T> { + /// The number of milliseconds elapsed since the wheel started. + elapsed: u64, + + /// Timer wheel. + /// + /// Levels: + /// + /// * 1 ms slots / 64 ms range + /// * 64 ms slots / ~ 4 sec range + /// * ~ 4 sec slots / ~ 4 min range + /// * ~ 4 min slots / ~ 4 hr range + /// * ~ 4 hr slots / ~ 12 day range + /// * ~ 12 day slots / ~ 2 yr range + levels: Vec<Level<T>>, +} + +/// Number of levels. Each level has 64 slots. By using 6 levels with 64 slots +/// each, the timer is able to track time up to 2 years into the future with a +/// precision of 1 millisecond. +const NUM_LEVELS: usize = 6; + +/// The maximum duration of a delay +const MAX_DURATION: u64 = (1 << (6 * NUM_LEVELS)) - 1; + +#[derive(Debug)] +pub(crate) enum InsertError { + Elapsed, + Invalid, +} + +/// Poll expirations from the wheel +#[derive(Debug, Default)] +pub(crate) struct Poll { + now: u64, + expiration: Option<Expiration>, +} + +impl<T> Wheel<T> +where + T: Stack, +{ + /// Create a new timing wheel + pub(crate) fn new() -> Wheel<T> { + let levels = (0..NUM_LEVELS).map(Level::new).collect(); + + Wheel { elapsed: 0, levels } + } + + /// Return the number of milliseconds that have elapsed since the timing + /// wheel's creation. + pub(crate) fn elapsed(&self) -> u64 { + self.elapsed + } + + /// Insert an entry into the timing wheel. + /// + /// # Arguments + /// + /// * `when`: is the instant at which the entry should be fired. It is + /// represented as the number of milliseconds since the creation + /// of the timing wheel. + /// + /// * `item`: The item to insert into the wheel. + /// + /// * `store`: The slab or `()` when using heap storage. + /// + /// # Return + /// + /// Returns `Ok` when the item is successfully inserted, `Err` otherwise. + /// + /// `Err(Elapsed)` indicates that `when` represents an instant that has + /// already passed. In this case, the caller should fire the timeout + /// immediately. + /// + /// `Err(Invalid)` indicates an invalid `when` argument as been supplied. + pub(crate) fn insert( + &mut self, + when: u64, + item: T::Owned, + store: &mut T::Store, + ) -> Result<(), (T::Owned, InsertError)> { + if when <= self.elapsed { + return Err((item, InsertError::Elapsed)); + } else if when - self.elapsed > MAX_DURATION { + return Err((item, InsertError::Invalid)); + } + + // Get the level at which the entry should be stored + let level = self.level_for(when); + + self.levels[level].add_entry(when, item, store); + + debug_assert!({ + self.levels[level] + .next_expiration(self.elapsed) + .map(|e| e.deadline >= self.elapsed) + .unwrap_or(true) + }); + + Ok(()) + } + + /// Remove `item` from thee timing wheel. + pub(crate) fn remove(&mut self, item: &T::Borrowed, store: &mut T::Store) { + let when = T::when(item, store); + let level = self.level_for(when); + + self.levels[level].remove_entry(when, item, store); + } + + /// Instant at which to poll + pub(crate) fn poll_at(&self) -> Option<u64> { + self.next_expiration().map(|expiration| expiration.deadline) + } + + pub(crate) fn poll(&mut self, poll: &mut Poll, store: &mut T::Store) -> Option<T::Owned> { + loop { + if poll.expiration.is_none() { + poll.expiration = self.next_expiration().and_then(|expiration| { + if expiration.deadline > poll.now { + None + } else { + Some(expiration) + } + }); + } + + match poll.expiration { + Some(ref expiration) => { + if let Some(item) = self.poll_expiration(expiration, store) { + return Some(item); + } + + self.set_elapsed(expiration.deadline); + } + None => { + self.set_elapsed(poll.now); + return None; + } + } + + poll.expiration = None; + } + } + + /// Returns the instant at which the next timeout expires. + fn next_expiration(&self) -> Option<Expiration> { + // Check all levels + for level in 0..NUM_LEVELS { + if let Some(expiration) = self.levels[level].next_expiration(self.elapsed) { + // There cannot be any expirations at a higher level that happen + // before this one. + debug_assert!(self.no_expirations_before(level + 1, expiration.deadline)); + + return Some(expiration); + } + } + + None + } + + /// Used for debug assertions + fn no_expirations_before(&self, start_level: usize, before: u64) -> bool { + let mut res = true; + + for l2 in start_level..NUM_LEVELS { + if let Some(e2) = self.levels[l2].next_expiration(self.elapsed) { + if e2.deadline < before { + res = false; + } + } + } + + res + } + + pub(crate) fn poll_expiration( + &mut self, + expiration: &Expiration, + store: &mut T::Store, + ) -> Option<T::Owned> { + while let Some(item) = self.pop_entry(expiration, store) { + if expiration.level == 0 { + debug_assert_eq!(T::when(item.borrow(), store), expiration.deadline); + + return Some(item); + } else { + let when = T::when(item.borrow(), store); + + let next_level = expiration.level - 1; + + self.levels[next_level].add_entry(when, item, store); + } + } + + None + } + + fn set_elapsed(&mut self, when: u64) { + assert!( + self.elapsed <= when, + "elapsed={:?}; when={:?}", + self.elapsed, + when + ); + + if when > self.elapsed { + self.elapsed = when; + } + } + + fn pop_entry(&mut self, expiration: &Expiration, store: &mut T::Store) -> Option<T::Owned> { + self.levels[expiration.level].pop_entry_slot(expiration.slot, store) + } + + fn level_for(&self, when: u64) -> usize { + level_for(self.elapsed, when) + } +} + +fn level_for(elapsed: u64, when: u64) -> usize { + let masked = elapsed ^ when; + + assert!(masked != 0, "elapsed={}; when={}", elapsed, when); + + let leading_zeros = masked.leading_zeros() as usize; + let significant = 63 - leading_zeros; + significant / 6 +} + +impl Poll { + pub(crate) fn new(now: u64) -> Poll { + Poll { + now, + expiration: None, + } + } +} + +#[cfg(all(test, not(loom)))] +mod test { + use super::*; + + #[test] + fn test_level_for() { + for pos in 1..64 { + assert_eq!( + 0, + level_for(0, pos), + "level_for({}) -- binary = {:b}", + pos, + pos + ); + } + + for level in 1..5 { + for pos in level..64 { + let a = pos * 64_usize.pow(level as u32); + assert_eq!( + level, + level_for(0, a as u64), + "level_for({}) -- binary = {:b}", + a, + a + ); + + if pos > level { + let a = a - 1; + assert_eq!( + level, + level_for(0, a as u64), + "level_for({}) -- binary = {:b}", + a, + a + ); + } + + if pos < 64 { + let a = a + 1; + assert_eq!( + level, + level_for(0, a as u64), + "level_for({}) -- binary = {:b}", + a, + a + ); + } + } + } + } +} diff --git a/third_party/rust/tokio/src/time/wheel/stack.rs b/third_party/rust/tokio/src/time/wheel/stack.rs new file mode 100644 index 0000000000..6e55c38ccd --- /dev/null +++ b/third_party/rust/tokio/src/time/wheel/stack.rs @@ -0,0 +1,26 @@ +use std::borrow::Borrow; + +/// Abstracts the stack operations needed to track timeouts. +pub(crate) trait Stack: Default { + /// Type of the item stored in the stack + type Owned: Borrow<Self::Borrowed>; + + /// Borrowed item + type Borrowed; + + /// Item storage, this allows a slab to be used instead of just the heap + type Store; + + /// Returns `true` if the stack is empty + fn is_empty(&self) -> bool; + + /// Push an item onto the stack + fn push(&mut self, item: Self::Owned, store: &mut Self::Store); + + /// Pop an item from the stack + fn pop(&mut self, store: &mut Self::Store) -> Option<Self::Owned>; + + fn remove(&mut self, item: &Self::Borrowed, store: &mut Self::Store); + + fn when(item: &Self::Borrowed, store: &Self::Store) -> u64; +} |