path: root/tokio/src/time/driver/entry.rs

//! Timer state structures.
//!
//! This module contains the heart of the intrusive timer implementation, and as
//! such the structures inside are full of tricky concurrency and unsafe code.
//!
//! # Ground rules
//!
//! The heart of the timer implementation here is the `TimerShared` structure,
//! shared between the `TimerEntry` and the driver. Generally, we permit access
//! to `TimerShared` ONLY via either 1) a mutable reference to `TimerEntry` or
//! 2) a held driver lock.
//!
//! It follows from this that any changes made while holding BOTH 1 and 2 will
//! be reliably visible, regardless of ordering. This is because of the acq/rel
//! fences on the driver lock ensuring ordering with 2, and rust mutable
//! reference rules for 1 (a mutable reference to an object can't be passed
//! between threads without an acq/rel barrier, and same-thread we have local
//! happens-before ordering).
//!
//! # State field
//!
//! Each timer has a state field associated with it. This field contains either
//! the current scheduled time, or a special flag value indicating its state.
//! This state can either indicate that the timer is on the 'pending' queue (and
//! thus will be fired with an `Ok(())` result soon) or that it has already been
//! fired/deregistered.
//!
//! This single state field allows for code that is firing the timer to
//! synchronize with any racing `reset` calls reliably.
//!
//! # Cached vs true timeouts
//!
//! To allow for the use case of a timeout that is periodically reset before
//! expiration to be as lightweight as possible, we support optimistically
//! lock-free timer resets, in the case where a timer is rescheduled to a later
//! point than it was originally scheduled for.
//!
//! This is accomplished by lazily rescheduling timers. That is, we update the
//! state field field with the true expiration of the timer from the holder of
//! the [`TimerEntry`]. When the driver services timers (ie, whenever it's
//! walking lists of timers), it checks this "true when" value, and reschedules
//! based on it.
//!
//! We do, however, also need to track what the expiration time was when we
//! originally registered the timer; this is used to locate the right linked
//! list when the timer is being cancelled. This is referred to as the "cached
//! when" internally.
//!
//! There is of course a race condition between timer reset and timer
//! expiration. If the driver fails to observe the updated expiration time, it
//! could trigger expiration of the timer too early. However, because
//! `mark_pending` performs a compare-and-swap, it will identify this race and
//! refuse to mark the timer as pending.

use crate::loom::cell::UnsafeCell;
use crate::loom::sync::atomic::Ordering;

use crate::sync::AtomicWaker;
use crate::time::Instant;
use crate::util::linked_list;

use super::Handle;

use std::cell::UnsafeCell as StdUnsafeCell;
use std::task::{Context, Poll, Waker};
use std::{marker::PhantomPinned, pin::Pin, ptr::NonNull};

type TimerResult = Result<(), crate::time::error::Error>;

const STATE_DEREGISTERED: u64 = u64::max_value();
const STATE_PENDING_FIRE: u64 = STATE_DEREGISTERED - 1;
const STATE_MIN_VALUE: u64 = STATE_PENDING_FIRE;

/// Not all platforms support 64-bit compare-and-swap. This hack replaces the
/// AtomicU64 with a mutex around a u64 on platforms that don't. This is slow,
/// unfortunately, but 32-bit platforms are a bit niche so it'll do for now.
///
/// Note: We use "x86 or 64-bit pointers" as the condition here because
/// target_has_atomic is not stable.
#[cfg(all(
    not(tokio_force_time_entry_locked),
    any(target_arch = "x86", target_pointer_width = "64")
))]
type AtomicU64 = crate::loom::sync::atomic::AtomicU64;

#[cfg(not(all(
    not(tokio_force_time_entry_locked),
    any(target_arch = "x86", target_pointer_width = "64")
)))]
#[derive(Debug)]
struct AtomicU64 {
    inner: crate::loom::sync::Mutex<u64>,
}

#[cfg(not(all(
    not(tokio_force_time_entry_locked),
    any(target_arch = "x86", target_pointer_width = "64")
)))]
impl AtomicU64 {
    fn new(v: u64) -> Self {
        Self {
            inner: crate::loom::sync::Mutex::new(v),
        }
    }

    fn load(&self, _order: Ordering) -> u64 {
        debug_assert_ne!(_order, Ordering::SeqCst); // we only provide AcqRel with the lock
        *self.inner.lock()
    }

    fn store(&self, v: u64, _order: Ordering) {
        debug_assert_ne!(_order, Ordering::SeqCst); // we only provide AcqRel with the lock
        *self.inner.lock() = v;
    }

    fn compare_exchange(
        &self,
        current: u64,
        new: u64,
        _success: Ordering,
        _failure: Ordering,
    ) -> Result<u64, u64> {
        debug_assert_ne!(_success, Ordering::SeqCst); // we only provide AcqRel with the lock
        debug_assert_ne!(_failure, Ordering::SeqCst);

        let mut lock = self.inner.lock();

        if *lock == current {
            *lock = new;
            Ok(current)
        } else {
            Err(*lock)
        }
    }

    fn compare_exchange_weak(
        &self,
        current: u64,
        new: u64,
        success: Ordering,
        failure: Ordering,
    ) -> Result<u64, u64> {
        self.compare_exchange(current, new, success, failure)
    }
}

/// This structure holds the current shared state of the timer - its scheduled
/// time (if registered), or otherwise the result of the timer completing, as
/// well as the registered waker.
///
/// Generally, the StateCell is only permitted to be accessed from two contexts:
/// Either a thread holding the corresponding &mut TimerEntry, or a thread
/// holding the timer driver lock. The write actions on the StateCell amount to
/// passing "ownership" of the StateCell between these contexts; moving a timer
/// from the TimerEntry to the driver requires _both_ holding the &mut
/// TimerEntry and the driver lock, while moving it back (firing the timer)
/// requires only the driver lock.
pub(super) struct StateCell {
    /// Holds either the scheduled expiration time for this timer, or (if the
    /// timer has been fired and is unregistered), [`u64::max_value()`].
    state: AtomicU64,
    /// If the timer is fired (an Acquire order read on state shows
    /// `u64::max_value()`), holds the result that should be returned from
    /// polling the timer. Otherwise, the contents are unspecified and reading
    /// without holding the driver lock is undefined behavior.
    result: UnsafeCell<TimerResult>,
    /// The currently-registered waker
    waker: CachePadded<AtomicWaker>,
}

impl Default for StateCell {
    fn default() -> Self {
        Self::new()
    }
}

impl std::fmt::Debug for StateCell {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "StateCell({:?})", self.read_state())
    }
}

impl StateCell {
    fn new() -> Self {
        Self {
            state: AtomicU64::new(STATE_DEREGISTERED),
            result: UnsafeCell::new(Ok(())),
            waker: CachePadded(AtomicWaker::new()),
        }
    }

    fn is_pending(&self) -> bool {
        self.state.load(Ordering::Relaxed) == STATE_PENDING_FIRE
    }

    /// Returns the current expiration time, or None if not currently scheduled.
    fn when(&self) -> Option<u64> {
        let cur_state = self.state.load(Ordering::Relaxed);

        if cur_state == u64::max_value() {
            None
        } else {
            Some(cur_state)
        }
    }

    /// If the timer is completed, returns the result of the timer. Otherwise,
    /// returns None and registers the waker.
    fn poll(&self, waker: &Waker) -> Poll<TimerResult> {
        // We must register first. This ensures that either `fire` will
        // observe the new waker, or we will observe a racing fire to have set
        // the state, or both.
        self.waker.0.register_by_ref(waker);

        self.read_state()
    }

    fn read_state(&self) -> Poll<TimerResult> {
        let cur_state = self.state.load(Ordering::Acquire);

        if cur_state == STATE_DEREGISTERED {
            // SAFETY: The driver has fired this timer; this involves writing
            // the result, and then writing (with release ordering) the state
            // field.
            Poll::Ready(unsafe { self.result.with(|p| *p) })
        } else {
            Poll::Pending
        }
    }

    /// Marks this timer as being moved to the pending list, if its scheduled
    /// time is not after `not_after`.
    ///
    /// If the timer is scheduled for a time after not_after, returns an Err
    /// containing the current scheduled time.
    ///
    /// SAFETY: Must hold the driver lock.
    unsafe fn mark_pending(&self, not_after: u64) -> Result<(), u64> {
        // Quick initial debug check to see if the timer is already fired. Since
        // firing the timer can only happen with the driver lock held, we know
        // we shouldn't be able to "miss" a transition to a fired state, even
        // with relaxed ordering.
        let mut cur_state = self.state.load(Ordering::Relaxed);

        loop {
            debug_assert!(cur_state < STATE_MIN_VALUE);

            if cur_state > not_after {
                break Err(cur_state);
            }

            match self.state.compare_exchange(
                cur_state,
                STATE_PENDING_FIRE,
                Ordering::AcqRel,
                Ordering::Acquire,
            ) {
                Ok(_) => {