takeaway 0.1.4

//! 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.
///
/// # 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.
///
/// - 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`.
///
/// - 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`.
///
/// - 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.
///
/// - Uninitialized (4): The worker has not started yet.
///
///   The associated [`WakerSlot`] is locked by the associated worker and holds
///   [`None`].
///
///   Outgoing transitions:
///   - `initialize` to Running.
///
/// # 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: `store 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.
///   - Operation: `store 0`.
///
/// - `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.
///   - Operation: `store 3`.
///
/// - `wake-self`: The associated worker unlocks its waker after dropping it
///   and starts running.
///
///   - Transition: Locked to Running.
///   - Operation: `store 0`.
///
/// - `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`.
///
/// - `initialize`: The associated worker marks itself as running, verifying
///   that it is the only actor using the worker ID.
///
///   - Transition: Uninitialized to Running.
///   - Operation: `cmpxchg 4 -> 0`.
#[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(4),
        }
    }

    /// Initialize the worker.
    ///
    /// Mark the worker as running, verifying that this is the first (and only)
    /// worker with this ID.
    ///
    /// ## Transition
    ///
    /// `initialize`.
    ///
    /// ## Panics
    ///
    /// Panics if this control state has already been initialized.
    pub fn initialize(&self) {
        // Transition 'initialize': Uninitialized to Running.
        // - Possible prior states: anything.
        // - Possible future states: anything initialized.
        self.raw
            .compare_exchange(4, 0, Ordering::Relaxed, Ordering::Relaxed)
            .expect("Workers are never allocated the same IDs");
    }

    /// Mark the worker as asleep.
    ///
    /// Attempt to mark the worker as asleep, using the specified waker.
    ///
    /// ## 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) {
        // SAFETY:
        // - The state was last observed to be Running.
        // - Only 'set-self-asleep' is 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, the current state is Running.
        // - 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.
        // - Possible future states: anything.
        self.raw.store(3, Ordering::Release);
    }

    /// 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.
        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.
        //   - 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.
        // - Possible prior states: Locked/Contended.
        // - Possible future states: anything.
        let prev = self.raw.swap(0, Ordering::Release);

        // If the associated worker was contending the lock, unblock it.
        if prev == 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 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<()> {
        // Possible states: anything initialized.

        // Transition 'lock-self-asleep': Asleep to Locked.
        //
        // NOTE(memory ordering): In case of failure, the associated waker slot
        // may have been emptied, the write must become visible here.
        match self.raw.compare_exchange(
            3,
            1,
            Ordering::Relaxed,
            Ordering::Acquire,
        ) {
            Ok(3) => {
                // Possible states: Locked.
                ControlFlow::Continue(())
            }

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

                // Possible states: Running.

                ControlFlow::Break(())
            }

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

                // Possible states: Locked/Running.

                // 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(()),
                        1 => {}
                        _ => unsafe { unreachable_unchecked() },
                    }
                }

                // Transition 'wait-self-asleep': Locked to Contended.
                //
                // NOTE(memory ordering): In case of failure, the state may have
                // become Running, and the write emptying the waker slot must
                // become visible here.
                match self.raw.compare_exchange(
                    1,
                    2,
                    Ordering::Relaxed,
                    Ordering::Acquire,
                ) {
                    Ok(_) => {}
                    Err(0) => return ControlFlow::Break(()),
                    _ => unsafe { unreachable_unchecked() },
                }

                // Possible states: Contended/Running.

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

                    match self.raw.load(Ordering::Acquire) {
                        0 => return ControlFlow::Break(()),
                        2 => continue,
                        _ => unsafe { unreachable_unchecked() },
                    }
                }
            }

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

    /// Wake up the worker from 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) {
        // 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(()) => return,
        }

        // Possible states: Locked.

        // Drop the stored waker and transition to Running.

        // SAFETY:
        // - In the Locked state, 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.
        self.raw.store(0, Ordering::Release);
    }

    /// 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.
    ///
    /// ## Transition
    ///
    /// `lock-self-asleep`, `refresh-self-asleep` or `wake-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(()) => return Control::Running,
        }

        // Possible states: Locked.

        // 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 state, 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.
        self.raw.store(3, Ordering::Release);

        // Possible states: anything initialized.
        Control::Asleep
    }
}

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 {}