path: root/tokio/src/task/state.rs

use crate::loom::sync::atomic::AtomicUsize;

use std::fmt;
use std::sync::atomic::Ordering::{AcqRel, Acquire, Release};
use std::usize;

pub(super) struct State {
    val: AtomicUsize,
}

/// Current state value
#[derive(Copy, Clone)]
pub(super) struct Snapshot(usize);

/// The task is currently being run.
const RUNNING: usize = 0b00_0001;

/// The task has been notified by a waker.
const NOTIFIED: usize = 0b00_0010;

/// The task is complete.
///
/// Once this bit is set, it is never unset
const COMPLETE: usize = 0b00_0100;

/// The primary task handle has been dropped.
const RELEASED: usize = 0b00_1000;

/// The join handle is still around
const JOIN_INTEREST: usize = 0b01_0000;

/// A join handle waker has been set
const JOIN_WAKER: usize = 0b10_0000;

/// The task has been forcibly canceled.
const CANCELLED: usize = 0b100_0000;

/// All bits
const LIFECYCLE_MASK: usize =
    RUNNING | NOTIFIED | COMPLETE | RELEASED | JOIN_INTEREST | JOIN_WAKER | CANCELLED;

/// Bits used by the waker ref count portion of the state.
///
/// Ref counts only cover **wakers**. Other handles are tracked with other state
/// bits.
const WAKER_COUNT_MASK: usize = usize::MAX - LIFECYCLE_MASK;

/// Number of positions to shift the ref count
const WAKER_COUNT_SHIFT: usize = WAKER_COUNT_MASK.count_zeros() as usize;

/// One ref count
const WAKER_ONE: usize = 1 << WAKER_COUNT_SHIFT;

/// Initial state
const INITIAL_STATE: usize = NOTIFIED;

/// All transitions are performed via RMW operations. This establishes an
/// unambiguous modification order.
impl State {
    /// Starts with a ref count of 2
    pub(super) fn new_joinable() -> State {
        State {
            val: AtomicUsize::new(INITIAL_STATE | JOIN_INTEREST),
        }
    }

    /// Load the current state, establishes `Acquire` ordering.
    pub(super) fn load(&self) -> Snapshot {
        Snapshot(self.val.load(Acquire))
    }

    /// Transition a task to the `Running` state.
    ///
    /// Returns a snapshot of the state **after** the transition.
    pub(super) fn transition_to_running(&self) -> Snapshot {
        const DELTA: usize = RUNNING | NOTIFIED;

        let prev = Snapshot(self.val.fetch_xor(DELTA, Acquire));
        debug_assert!(prev.is_notified());

        if prev.is_running() {
            // We were signalled to cancel
            //
            // Apply the state
            let prev = self.val.fetch_or(CANCELLED, AcqRel);
            return Snapshot(prev | CANCELLED);
        }

        debug_assert!(!prev.is_running());

        let next = Snapshot(prev.0 ^ DELTA);

        debug_assert!(next.is_running());
        debug_assert!(!next.is_notified());

        next
    }

    /// Transition the task from `Running` -> `Idle`.
    ///
    /// Returns a snapshot of the state **after** the transition.
    pub(super) fn transition_to_idle(&self) -> Snapshot {
        const DELTA: usize = RUNNING;

        let prev = Snapshot(self.val.fetch_xor(DELTA, AcqRel));

        if !prev.is_running() {
            // We were signaled to cancel.
            //
            // Apply the state
            let prev = self.val.fetch_or(CANCELLED, AcqRel);
            return Snapshot(prev | CANCELLED);
        }

        let next = Snapshot(prev.0 ^ DELTA);

        debug_assert!(!next.is_running());

        next
    }

    /// Transition the task from `Running` -> `Complete`.
    ///
    /// Returns a snapshot of the state **after** the transition.
    pub(super) fn transition_to_complete(&self) -> Snapshot {
        const DELTA: usize = RUNNING | COMPLETE;

        let prev = Snapshot(self.val.fetch_xor(DELTA, AcqRel));

        debug_assert!(!prev.is_complete());

        let next = Snapshot(prev.0 ^ DELTA);

        debug_assert!(next.is_complete());

        next
    }

    /// Transition the task from `Running` -> `Released`.
    ///
    /// Returns a snapshot of the state **after** the transition.
    pub(super) fn transition_to_released(&self) -> Snapshot {
        const DELTA: usize = RUNNING | COMPLETE | RELEASED;

        let prev = Snapshot(self.val.fetch_xor(DELTA, AcqRel));

        debug_assert!(prev.is_running());
        debug_assert!(!prev.is_complete());
        debug_assert!(!prev.is_released());

        let next = Snapshot(prev.0 ^ DELTA);

        debug_assert!(!next.is_running());
        debug_assert!(next.is_complete());
        debug_assert!(next.is_released());

        next
    }

    /// Transition the task to the canceled state.
    ///
    /// Returns the snapshot of the state **after** the transition **if** the
    /// transition was made successfully
    ///
    /// # States
    ///
    /// - Notifed: task may be in a queue, caller must not release.
    /// - Running: cannot drop. The poll handle will handle releasing.
    /// - Other prior states do not require cancellation.
    ///
    /// If the task has been notified, then it may still be in a queue. The
    /// caller must not release the task.
    pub(super) fn transition_to_canceled_from_queue(&self) -> Snapshot {
        let prev = Snapshot(self.val.fetch_or(CANCELLED, AcqRel));

        debug_assert!(!prev.is_complete());
        debug_assert!(!prev.is_running() || prev.is_notified());

        Snapshot(prev.0 | CANCELLED)
    }

    pub(super) fn transition_to_canceled_from_list(&self) -> Option<Snapshot> {
        let mut prev = self.load();

        loop {
            if !prev.is_active() {
                return None;
            }

            let mut next = prev;

            // Use the running flag to signal cancellation
            if prev.is_running() {
                next.0 -= RUNNING;
                next.0 |= NOTIFIED;
            } else if prev.is_notified() {
                next.0 += RUNNING;
                next.0 |= NOTIFIED;
            } else {
                next.0 |= CANCELLED;
            }

            let res = self.val.compare_exchange(prev.0, next.0, AcqRel, Acquire);

            match res {
                Ok(_) if next.is_canceled() => return Some(next),
                Ok(_) => return None,
                Err(actual) => {
                    prev = Snapshot(actual);
                }
            }
        }
    }

    /// Final transition to `Released`. Called when primary task handle is
    /// dropped. This is roughly a "ref decrement" operation.
    ///
    /// Returns a snapshot of the state **after** the transition.
    pub(super) fn release_task(&self) -> Snapshot {
        use crate::loom::sync::atomic;

        const DELTA: usize = RELEASED;

        let prev = Snapshot(self.val.fetch_or(DELTA, Release));

        debug_assert!(!prev.is_released());
        debug_assert!(prev.is_terminal(), "state = {:?}", prev);

        let next = Snapshot(prev.0 | DELTA);

        debug_assert!(next.is_released());

        if next.is_final_ref() || (next.has_join_waker() && !next.is_join_interested()) {
            // The final reference to the task was dropped, the caller must free the
            // memory. Establish an acquire ordering.
            atomic::fence(Acquire);
        }

        next
    }

    /// Transition the state to `Scheduled`.
    ///
    /// Returns `true` if the task needs to be submitted to the pool for
    /// execution
    pub(super) fn transition_to_notified(&self) -> bool {
        const MASK: usize = RUNNING | NOTIFIED | COMPLETE | CANCELLED;

        let prev = self.val.fetch_or(NOTIFIED, Release);
        prev & MASK == 0
    }

    /// Optimistically try to swap the state assuming the join handle is
    /// __immediately__ dropped on spawn
    pub(super) fn drop_join_handle_fast(&self) -> bool {
        use std::sync::atomic::Ordering::Relaxed;

        // Relaxed is acceptable as if this function is called and succeeds,
        // then nothing has been done w/ the join handle.
        //
        // The moment the join handle is used (polled), the `JOIN_WAKER` flag is
        // set, at which point the CAS will fail.
        //
        // Given this, there is no risk if this operation is reordered.
        self.va