takeaway 0.1.0

//! Control states.
//!
//! Every worker thread has a control state, indicating whether it is running,
//! asleep, or shutting down.

use core::{
    cell::UnsafeCell,
    hint::unreachable_unchecked,
    ops::ControlFlow,
    sync::atomic::{AtomicU32, Ordering},
    task::Waker,
};

use crate::util::Backoff;

//----------- Control ----------------------------------------------------------

/// The control state of a worker.
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub enum Control {
    /// The worker has been initialized but is not yet running.
    Initialized,

    /// The worker is actively polling for tasks.
    Running,

    /// The worker ran out of tasks and is waiting.
    Asleep,

    /// The worker is shutting down.
    ShuttingDown,
}

//----------- AtomicControl ----------------------------------------------------

/// An atomic [`Control`].
///
/// This variable provides several functions:
/// - Ensuring worker IDs are uniquely allocated.
/// - Allowing the worker to sleep and be woken up reliably.
/// - Notifying the worker that the system is shutting down.
///
/// # States
///
/// - Running (0): The worker is actively processing tasks.
///
///   The associated [`WakerSlot`] is locked by the associated worker and
///   holds [`None`].
///
///   Outgoing transitions:
///   - `set-self-asleep` to Asleep.
///   - `mark-shutdown` to Running + Shutting Down.
///
/// - Locked (1): The worker is transitioning between Running and Asleep.
///
///   The associated [`WakerSlot`] is locked by the actor that transitioned
///   the variable the Locked state.
///
///   Outgoing transitions:
///   - `refresh-self-asleep` to Asleep, following `lock-self-asleep`.
///   - `wait-self-asleep` to Contended, following `try-wake`.
///   - `wake` to Running, following `try-wake`.
///   - `wake-self` to Running, following `lock-self-asleep`.
///   - `mark-shutdown` to Locked + Shutting Down.
///
/// - Contended (2): The worker is waiting for an ongoing wakeup to complete.
///
///   The associated [`WakerSlot`] is locked by the actor that transitioned
///   the variable the Locked state.
///
///   Outgoing transitions:
///   - `wake` to Running, following `try-wake` and `wait-self-asleep`.
///   - `mark-shutdown` to Contended + Shutting Down.
///
/// - Asleep (3): The worker ran out of tasks, and is waiting to be woken up.
///
///   The associated [`WakerSlot`] is unlocked and holds [`Some`].
///
///   Outgoing transitions:
///   - `lock-self-asleep` to Locked.
///   - `try-wake` to Locked.
///   - `mark-shutdown` to Asleep + Shutting Down.
///
/// - Running + Shutting Down (4): The worker is shutting down.
///
///   Coming from `mark-shutdown`: the associated [`WakerSlot`] is locked by the
///   associated worker and holds [`None`].
///
///   Coming from `wake`: the associated [`WakerSlot`] is not locked and holds
///   [`None`].
///
///   Outgoing transitions: none.
///
/// - Locked + Shutting Down (5): Locked, but a shutdown is occurring.
///
///   The associated [`WakerSlot`] is locked by the actor that transitioned
///   the variable the Locked state.
///
///   Outgoing transitions:
///   - `refresh-self-asleep` to Asleep + Shutting Down, following
///     `lock-self-asleep` and `mark-shutdown`.
///   - `wake` to Running + Shutting Down, following `try-wake` and
///     `mark-shutdown`.
///   - `wake-self` to Running + Shutting Down, following `lock-self-asleep` and
///     `mark-shutdown`.
///
/// - Contended + Shutting Down (6): Contended, but a shutdown is occurring.
///
///   The associated [`WakerSlot`] is locked by the actor that transitioned
///   the variable the Locked state.
///
///   Outgoing transitions:
///   - `wake` to Running + Shutting Down, following `try-wake`,
///     `wait-self-asleep`, and `mark-shutdown`.
///
/// - Asleep + Shutting Down (7): The worker is waiting to be woken up, and is
///   about to be woken up because a shutdown is occurring.
///
///   If coming from `refresh-self-asleep`, the associated [`WakerSlot`] is
///   locked by the associated worker and holds [`Some`].
///
///   If coming from `mark-shutdown`, the associated [`WakerSlot`] is locked by
///   the worker performing the shutdown and holds [`Some`].
///
///   Outgoing transitions:
///   - `wake-shutdown` to Running + Shutting Down, following `mark-shutdown`.
///   - `shutdown-self` to Running + Shutting Down, following
///     `lock-self-asleep`, `mark-shutdown`, and `refresh-self-asleep`.
///
/// - Uninitialized (8): The worker has not started yet.
///
///   The associated [`WakerSlot`] is locked by the associated worker and holds
///   [`None`].
///
///   Outgoing transitions:
///   - `mark-shutdown` to Uninitialized + Shutting Down.
///   - `initialize` to Running.
///
/// - Uninitialized + Shutting Down (12): The worker has not started yet, but
///   the system is shutting down.
///
///   The associated [`WakerSlot`] is locked by the associated worker and holds
///   [`None`].
///
///   Outgoing transitions:
///   - `initialize` to Running + Shutting Down.
///
/// # Transitions
///
/// - `set-self-asleep`: The associated worker marks itself as asleep, because
///   it has run out of tasks and was not able to steal any more.
///
///   - Transition: Running to Asleep.
///   - Operation: `cmpxchg 0 -> 3`.
///
/// - `try-wake`: The calling worker tries to wake up the associated worker so
///   it can steal tasks (such as those it just published).
///
///   - Transition: Asleep to Locked.
///   - Operation: `cmpxchg 3 -> 1`.
///
/// - `wake`: The calling worker marks the associated worker as running, after
///   locking and taking its waker using `try-wake`.
///
///   - Transition: Locked/Contended to Running, modulo Shutting Down.
///   - Operation: `and 4`.
///
/// - `lock-self-asleep`: The associated worker locks its waker in order to
///   refresh it while it is sleeping.
///
///   - Transition: Asleep to Locked.
///   - Operation: `cmpxchg 3 -> 1`.
///
/// - `refresh-self-asleep`: The associated worker unlocks its waker after
///   refreshing it and returns to sleeping.
///
///   - Transition: Locked to Asleep, modulo Shutting Down.
///   - Operation: `or 3`.
///
/// - `wake-self`: The associated worker unlocks its waker after dropping it
///   and starts running.
///
///   - Transition: Locked to Running, modulo Shutting Down.
///   - Operation: `and 4`.
///
/// - `shutdown-self`: The associated worker shuts itself down after receiving
///   notification of shutdown.
///
///   - Transition: Asleep + Shutting Down to Running + Shutting Down.
///   - Operation: `store 4`.
///
/// - `wait-self-asleep`: The associated worker blocks on the lock on its waker
///   after spinning and waiting for it to finish.
///
///   - Transition: Locked to Contended.
///   - Operation: `cmpxchg 1 -> 2`.
///
/// - `mark-shutdown`: The calling worker notifies the associated worker that a
///   shutdown is occurring, synchronizing with the associated worker's control
///   state changes.
///
///   - Transition: Anything to that + Shutting Down.
///   - Operation: `or 4`.
///
/// - `wake-shutdown`: The calling worker wakes up the associated worker after
///   notifying it of shutdown using `mark-shutdown`.
///
///   - Transition: Asleep + Shutting Down to Running + Shutting Down.
///   - Operation: `store 4`.
///
/// - `initialize`: The associated worker marks itself as running, verifying
///   that it is the only actor using the worker ID.
///
///   - Transition: Uninitialized to Running, modulo Shutting Down.
///   - Operation: `and 7`.
#[repr(transparent)]
pub struct AtomicControl {
    raw: AtomicU32,
}

impl AtomicControl {
    /// Construct a new [`ControlState`].
    ///
    /// The state is initialized to Uninitialized.
    pub const fn new() -> Self {
        Self {
            raw: AtomicU32::new(8),
        }
    }

    /// Initialize the worker.
    ///
    /// Mark the worker as running, verifying that this is the first (and only)
    /// worker with this ID.  Returns `true` if a notification of shutdown has
    /// been received.
    ///
    /// ## Transition
    ///
    /// `initialize`.
    ///
    /// ## Panics
    ///
    /// Panics if this control state has already been initialized.
    pub fn initialize(&self) -> bool {
        // Transition 'initialize': Uninitialized to Running, modulo Shutting Down.
        // - Possible prior states: anything.
        // - Possible future states: anything initialized.
        let old = self.raw.fetch_and(7, Ordering::Relaxed);

        assert!(old & 8 != 0, "Two workers were allocated the same IDs");

        old & 4 != 0
    }

    /// Mark the worker as asleep.
    ///
    /// Attempt to mark the worker as asleep, using the specified waker.  If a
    /// notification of shutdown is received, `true` is returned.
    ///
    /// ## Transition
    ///
    /// `set-self-asleep`.
    ///
    /// ## Safety
    ///
    /// - `self` and `waker_slot` must correspond.
    /// - The caller must be the associated worker.
    /// - The state was last observed to be Running.
    pub unsafe fn set_self_asleep(
        &self,
        waker_slot: &WakerSlot,
        waker: Waker,
    ) -> Control {
        // SAFETY:
        // - The state was last observed to be Running.
        // - Only 'set-self-asleep' and 'mark-shutdown' are possible.
        //   - 'set-self-asleep' can only be performed by the associated worker,
        //     i.e. this actor, but the state was last observed to be Running,
        //     not Asleep.
        //   - Thus 'mark-shutdown' is the only possible transition.
        // - Thus, the current state is Running (possibly + Shutting Down).
        //   - If the state is Running, the associated worker (i.e. this actor)
        //     locks the waker slot, which holds 'None'.
        //   - If the state is Running + Shutting Down, it was reached from
        //     'mark-shutdown', which means the associated worker (i.e. this
        //     actor) locks the waker slot, which holds 'None'.
        // - Thus, the associated worker (i.e. this actor) locks the waker slot,
        //   which holds 'None'.
        unsafe { waker_slot.insert(waker) };

        // Transition 'set-self-asleep': Running to Asleep.
        // - Prior possible states: Running (possibly + Shutting Down).
        // - Possible future states: anything.
        match self.raw.compare_exchange(
            0,
            3,
            Ordering::Release,
            Ordering::Relaxed,
        ) {
            Ok(0) => Control::Asleep,
            Err(4) => Control::ShuttingDown,
            _ => unsafe { unreachable_unchecked() },
        }
    }

    /// Try to wake this worker.
    ///
    /// A fast, best-effort attempt is made to wake the worker if it is asleep.
    /// If the worker is actually woken up, `true` is returned.
    ///
    /// ## Transition
    ///
    /// `try-wake`; if it is successful, then `wake`.
    ///
    /// ## Safety
    ///
    /// `self` and `waker_slot` must correspond.
    pub unsafe fn try_wake(&self, waker_slot: &WakerSlot) -> bool {
        // Transition 'try-wake': Asleep to Locked.
        // - Possible prior states: anything.
        // - Possible future states: Locked/Contended, possibly + Shutting Down.
        match self.raw.compare_exchange(
            3,
            1,
            Ordering::Acquire,
            Ordering::Relaxed,
        ) {
            Ok(_) => {}
            Err(_) => return false,
        }

        // SAFETY:
        // - As per the caller, 'self' and 'waker_slot' correspond.
        // - 'self' was Asleep.
        //   - At the time, 'waker_slot' held a 'Some'.
        // - 'self' is Locked/Contended (possibly + Shutting Down).
        //   - In all these states, the actor that set the state to Locked has
        //     locked the waker slot.
        // - This actor set the state to Locked.
        // - Thus, this actor has locked the waker slot.
        //   - Immediately before it was locked, it held a 'Some'.
        // - The actor has not yet modified the waker slot.
        //   - Thus the waker slot still holds a 'Some'.
        let waker = unsafe { waker_slot.extract() };

        // Transition 'wake': Locked/Contended to Running, modulo Shutting Down.
        // - Possible prior states: Locked/Contended (possibly + Shutting Down).
        // - Possible future states: anything.
        let prev = self.raw.fetch_and(4, Ordering::Release);

        // If the associated worker was contending the lock, unblock it.
        if prev % 4 == 2 {
            atomic_wait::wake_one(&self.raw);
        }

        // Activate the waker.
        waker.wake();
        true
    }

    /// Lock the control state during sleep.
    ///
    /// This is called by the associated worker, while it is asleep, in order
    /// to check for external wakeups.
    ///
    /// If [`ControlFlow::Break`] is returned, the control state could not be
    /// locked, because a wakeup or shutdown is occurring.
    ///
    /// If [`ControlFlow::Continue`] is returned, the control state is locked
    /// and the waker slot can be modified freely.  It will contain a [`Some`].
    ///
    /// ## Transition
    ///
    /// `lock-self-asleep`, and possibly `wait-self-asleep`.
    ///
    /// ## Safety
    ///
    /// - The caller must be the associated worker.
    /// - The state was last observed to be Asleep.
    unsafe fn lock_self_asleep(&self) -> ControlFlow<Control> {
        // Possible states: anything initialized.

        // Transition 'lock-self-asleep': Asleep to Locked.
        match self.raw.compare_exchange(
            3,
            1,
            Ordering::Relaxed,
            Ordering::Relaxed,
        ) {
            Ok(3) => {
                // Possible states: Locked (possibly + Shutting Down).
                ControlFlow::Continue(())
            }

            Err(0) => {
                // The worker must have been woken up externally.

                // Possible states: Running (possibly + Shutting Down).

                ControlFlow::Break(Control::Running)
            }

            Err(1) => {
                // The worker is being woken up externally.

                // Possible states: Locked/Running (possibly + Shutting Down).

                // Spin and wait.
                let backoff = Backoff::new();
                while !backoff.is_completed() {
                    backoff.snooze();

                    // TODO: Use 'Relaxed' here and introduce a secondary load?
                    match self.raw.load(Ordering::Acquire) {
                        0 => return ControlFlow::Break(Control::Running),
                        1 => {}
                        4 | 5 => {
                            return ControlFlow::Break(Control::ShuttingDown);
                        }
                        _ => unsafe { unreachable_unchecked() },
                    }
                }

                // Transition 'wait-self-asleep': Locked to Contended.
                match self.raw.compare_exchange(
                    1,
                    2,
                    Ordering::Acquire,
                    Ordering::Relaxed,
                ) {
                    Ok(_) => {}
                    Err(0) => return ControlFlow::Break(Control::Running),
                    Err(4 | 5) => {
                        return ControlFlow::Break(Control::ShuttingDown);
                    }
                    _ => unsafe { unreachable_unchecked() },
                }

                // Possible states: Contended/Running (possibly + Shutting Down).

                loop {
                    atomic_wait::wait(&self.raw, 2);

                    match self.raw.load(Ordering::Acquire) {
                        0 => return ControlFlow::Break(Control::Running),
                        2 => {}
                        4 | 6 => {
                            return ControlFlow::Break(Control::ShuttingDown);
                        }
                        _ => unsafe { unreachable_unchecked() },
                    }
                }
            }

            Err(4 | 5 | 7) => ControlFlow::Break(Control::ShuttingDown),

            _ => unsafe { unreachable_unchecked() },
        }
    }

    /// Wake up the worker during sleep.
    ///
    /// ## Transition
    ///
    /// - `lock-self-asleep` (possibly with `wait-self-asleep`).
    /// - `wake-self`.
    ///
    /// ## Safety
    ///
    /// - `self` and `waker_slot` must correspond.
    /// - The caller must be the associated worker.
    /// - The state was last observed to be Asleep.
    pub unsafe fn set_self_awake(&self, waker_slot: &WakerSlot) -> Control {
        // Try to lock the control state.

        // SAFETY:
        // - As per the caller, the caller is the associated worker.
        // - As per the caller, the state was last observed to be Asleep.
        match unsafe { self.lock_self_asleep() } {
            ControlFlow::Continue(()) => {}
            ControlFlow::Break(control) => return control,
        }

        // Possible states: Locked (possibly + Shutting Down).

        // Drop the stored waker and transition to Running.

        // SAFETY:
        // - In the Locked (possibly + Shutting Down) states, the waker slot is
        //   locked by the actor that transitioned the control state to Locked.
        // - Here, the current actor did so using 'lock-self-asleep', so the
        //   waker slot is locked by this actor.
        let _ = unsafe { waker_slot.extract() };

        // Transition 'wake-self': Locked to Running, modulo Shutting Down.
        let prev = self.raw.fetch_and(4, Ordering::Release);

        if prev & 4 == 0 {
            // Possible states: Running, possibly + Shutting Down.
            Control::Running
        } else {
            // Possible states: Running + Shutting Down.
            Control::ShuttingDown
        }
    }

    /// Poll the control state while running.
    ///
    /// This is called by the associated worker to check for a shutdown while
    /// it's running.  It returns the updated [`Control`], which is either
    /// running or shutting down.
    ///
    /// ## Safety
    ///
    /// - The caller must be the associated worker.
    /// - The state was last observed to be Running.
    pub unsafe fn poll_running(&self) -> Control {
        match self.raw.load(Ordering::Relaxed) {
            0 => Control::Running,
            4 => Control::ShuttingDown,
            _ => unsafe { unreachable_unchecked() },
        }
    }

    /// Poll the control state during sleep.
    ///
    /// This is called by the associated worker in order to check its control
    /// state when it is asleep and polled.  It returns [`Poll::Ready`] once the
    /// worker is marked as awake, and `true` if a shutdown has been received.
    ///
    /// ## Transition
    ///
    /// `lock-self-asleep`, `refresh-self-asleep` or `wake-self`,
    /// `shutdown-self`, and/or `wait-self-asleep`.
    ///
    /// ## Safety
    ///
    /// - `self` and `waker_slot` must correspond.
    /// - The caller must be the associated worker.
    /// - The state was last observed to be Asleep.
    pub unsafe fn poll_asleep(
        &self,
        waker: &Waker,
        waker_slot: &WakerSlot,
    ) -> Control {
        // Try to lock the control state.

        // SAFETY:
        // - As per the caller, the caller is the associated worker.
        // - As per the caller, the state was last observed to be Asleep.
        match unsafe { self.lock_self_asleep() } {
            ControlFlow::Continue(()) => {}
            ControlFlow::Break(control) => return control,
        }

        // Possible states: Locked (possibly + Shutting Down).

        // Replace the stored waker and transition to Asleep.

        // NOTE: Nobody's going to contend this lock, so we can perform a
        // potentially expensive operation here.
        let waker = waker.clone();

        // SAFETY:
        // - In the Locked (possibly + Shutting Down) states, the waker slot is
        //   locked by the actor that transitioned the control state to Locked.
        // - Here, the current actor did so using 'lock-self-asleep', so the
        //   waker slot is locked by this actor.
        let _ = unsafe { waker_slot.replace(waker) };

        // Transition 'refresh-self-asleep': Locked to Asleep, modulo
        //   Shutting Down.
        let prev = self.raw.fetch_or(3, Ordering::Release);

        if prev & 4 == 0 {
            // Possible states: anything initialized.
            Control::Asleep
        } else {
            // Possible states: Asleep + Shutting Down.

            // SAFETY:
            // - The current state is 'Asleep + Shutting Down' and it was
            //   reached via 'refresh-self-asleep'.
            // - As such, the waker slot is locked be the associated worker.
            let _ = unsafe { waker_slot.extract() };

            // Transition 'shutdown-self': (Asleep to Running) + Shutting Down.
            self.raw.store(4, Ordering::Relaxed);

            // Possible states: Running + Shutting Down.
            Control::ShuttingDown
        }
    }

    /// Shut down this worker.
    ///
    /// This is called by a worker initiating a shutdown, on the control state
    /// of every worker (including itself).  If the worker is asleep, it will be
    /// woken up.
    ///
    /// ## Transition
    ///
    /// `mark-shutdown` and `wake-shutdown`.
    ///
    /// ## Safety
    ///
    /// - `self` and `waker_slot` must correspond.
    pub unsafe fn shutdown(&self, waker_slot: &WakerSlot) {
        // Possible states: anything.

        // Transition 'mark-shutdown': anything to itself + Shutting Down.
        let prev = self.raw.fetch_or(4, Ordering::AcqRel);

        // Possible states: 'prev' + Shutting Down.

        // If the worker was asleep, wake it up.
        if prev == 3 {
            // Possible states: Asleep + Shutting Down.

            // SAFETY: TODO
            let waker = unsafe { waker_slot.extract() };

            // Transition 'wake-shutdown': (Asleep to Running) + Shutting Down.
            // - Possible prior states: Asleep + Shutting Down.
            // - Possible future states: Running + Shutting Down.
            self.raw.store(4, Ordering::Relaxed);

            // Wake up the worker.
            waker.wake();
        }
    }
}

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

//----------- WakerSlot --------------------------------------------------------

/// Storage for a thread's waker.
#[repr(transparent)]
pub struct WakerSlot {
    raw: UnsafeCell<Option<Waker>>,
}

impl WakerSlot {
    /// Construct a new, empty [`WakerSlot`].
    pub const fn new() -> Self {
        Self {
            raw: UnsafeCell::new(None),
        }
    }

    /// Insert a waker in this slot.
    ///
    /// ## Safety
    ///
    /// This slot must be empty.  For the duration of this method, no other
    /// threads can read or write this slot.
    pub unsafe fn insert(&self, waker: Waker) {
        debug_assert!(unsafe { &*self.raw.get() }.is_none());

        unsafe { self.raw.get().write(Some(waker)) };
    }

    /// Replace the waker in this slot.
    ///
    /// ## Safety
    ///
    /// There must be a waker stored in this slot.  For the duration of this
    /// method, no other threads can read or write this slot.
    pub unsafe fn replace(&self, waker: Waker) -> Waker {
        debug_assert!(unsafe { &*self.raw.get() }.is_some());

        let slot = unsafe { (*self.raw.get()).as_mut().unwrap_unchecked() };
        core::mem::replace(slot, waker)
    }

    /// Extract the waker in this slot.
    ///
    /// ## Safety
    ///
    /// There must be a waker stored in this slot.  For the duration of this
    /// method, no other threads can read or write this slot.
    pub unsafe fn extract(&self) -> Waker {
        unsafe { (*self.raw.get()).take().unwrap_unchecked() }
    }
}

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

unsafe impl Send for WakerSlot {}
unsafe impl Sync for WakerSlot {}