moirai-executor 0.9.1

use std::{
    cell::UnsafeCell,
    mem::MaybeUninit,
    ptr::NonNull,
    sync::atomic::{AtomicBool, AtomicU8, AtomicU32, AtomicU64, Ordering},
    time::{Duration, Instant},
};

use moirai_core::Priority;

use super::super::task::TaskMetadata;
use super::fan_out;

/// Inverse of [`Priority::index`]: `PRIORITY_FROM_INDEX[p.index()] == p` for
/// every variant (asserted by `priority_index_round_trips` in the registry tests).
pub(crate) const PRIORITY_FROM_INDEX: [Priority; Priority::Critical.index() + 1] = [
    Priority::Low,
    Priority::Normal,
    Priority::High,
    Priority::Critical,
];

pub(crate) const NO_WORKER: u32 = u32::MAX;
pub(crate) const TIMESTAMP_NOT_RECORDED: u64 = u64::MAX;
pub(crate) const TASK_STATE_BLOCK_SIZE: usize = 1024;

/// Which settled blocks a retirement sweep may release.
#[derive(Debug, Clone, Copy)]
pub(super) enum Retirement {
    /// Only blocks whose every task completed at or before the instant.
    CompletedBefore(Instant),
    /// Any block whose every task completed, whatever its age.
    Forced,
}

/// One fixed-size block of task-state slots.
///
/// A slot is written once, by the registration that owns its id, and then
/// published by a release store to its `published` flag; every observer reads
/// the flag with acquire semantics before touching the state, so a lookup racing
/// a registration sees either an absent slot or a complete state. A published
/// state never moves and is never replaced. The only way a state's storage is
/// released is retiring the whole block, which the registry does once every slot
/// has completed and released its lifecycle token. An owned token's block `Arc`
/// keeps the allocation alive; scheduler-bounded tokens require their registry
/// to outlive the job and make their final access to the state when they retire.
///
/// The flags live apart from the states: a flag beside its state would pad every
/// 64-byte state to 72 bytes, and the retirement scan reads flags without
/// touching state lines.
pub(super) struct TaskStateBlock {
    published: Box<[AtomicBool]>,
    states: Box<[UnsafeCell<MaybeUninit<TaskState>>]>,
}

// SAFETY: a state is written only before its `published` flag is set, by the one
// registration that owns its slot, and only read after an acquire load observes
// the flag. `TaskState` is `Send + Sync`, so sharing published states across
// threads is sound.
unsafe impl Sync for TaskStateBlock {}

/// The waker slot of a [`TaskState`]: an `Option<Waker>` behind a tiny
/// self-contained spin lock.
///
/// [`std::sync::Mutex`] was rejected for this slot: `Mutex<Option<Waker>>` is
/// 32 bytes on macOS, where std reaches its `pthread_mutex_t` through a lazily
/// boxed pointer, against 24 on the futex targets (Linux, Windows), and with it
/// the per-task state was 80 bytes on macOS. The lock here is one flag byte
/// plus padding, and the critical sections clone a waker, run
/// [`Waker::will_wake`](std::task::Waker::will_wake) over a short member list,
/// or take the waker out — all
/// bounded, never blocking, so spinning with a yield is safe under any
/// scheduler.
pub(super) struct WakerSlot {
    locked: AtomicBool,
    slot: UnsafeCell<Option<std::task::Waker>>,
}

impl WakerSlot {
    const fn new() -> Self {
        Self {
            locked: AtomicBool::new(false),
            slot: UnsafeCell::new(None),
        }
    }

    /// Run `body` with exclusive access to the slot.
    fn with_slot<R>(&self, body: impl FnOnce(&mut Option<std::task::Waker>) -> R) -> R {
        let mut spins = 0_usize;
        while self
            .locked
            .compare_exchange_weak(false, true, Ordering::Acquire, Ordering::Relaxed)
            .is_err()
        {
            // Bounded spin first, then yield: the critical sections are short,
            // but a preempted holder must never deadlock the contender.
            if spins < 16 {
                std::hint::spin_loop();
            } else {
                std::thread::yield_now();
            }
            spins += 1;
        }
        // SAFETY: the `locked` flag is held by this call, so this borrow is the
        // only slot access; it ends before the flag is released.
        let slot = unsafe { &mut *self.slot.get() };
        let result = body(slot);
        self.locked.store(false, Ordering::Release);
        result
    }
}

// SAFETY: the slot is reached only through `with_slot`, which serializes access
// with the `locked` flag; a `Waker` is `Send`, so moving one between threads
// through the slot is sound.
unsafe impl Send for WakerSlot {}
unsafe impl Sync for WakerSlot {}

/// Shared lifecycle state for one task.
pub(crate) struct TaskState {
    pub(crate) created_at: Instant,
    pub(super) started_after_ns: AtomicU64,
    pub(super) completed_after_ns: AtomicU64,
    /// Worker index at start, or [`NO_WORKER`]. `u32` keeps the struct within
    /// its pinned budget on every supported pointer width; no scheduler this
    /// executor builds approaches 2^32 workers.
    pub(super) worker_id: AtomicU32,
    /// Waker slot behind its own tiny spin lock — see [`WakerSlot`] for why
    /// this is not a [`std::sync::Mutex`].
    pub(super) waker: WakerSlot,
    /// True while a lifecycle token can still access this slot; clearing it is
    /// the token's last access ([`TaskState::retire_token`]).
    token_active: AtomicBool,
    /// Spawn priority stored as its [`Priority::index`] discriminant.
    pub(super) priority: AtomicU8,
    /// Set by `cancel_task`; observed cooperatively at job start.
    pub(super) cancel_requested: AtomicBool,
    /// Set when the task ended without a result: a cancel request was honored,
    /// or its lifecycle token was dropped before completing. The task's handle
    /// reports `TaskError::Cancelled` in both cases.
    pub(super) cancelled: AtomicBool,
}

// A registry block holds 1,024 of these plus one flag byte each; the size is
// pinned because retained memory per task is this figure: 64 bytes on every
// 64-bit target (`Instant` is 16 bytes, 8-aligned, on each), 56 on wasm32.
// `worker_id` is `u32` and the waker slot carries its own spin lock (see
// `WakerSlot`) to stay within it.
const _: () = assert!(size_of::<TaskState>() <= 64);

impl std::fmt::Debug for TaskState {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("TaskState")
            .field("created_at", &self.created_at)
            .field("started_after_ns", &self.started_after_ns)
            .field("completed_after_ns", &self.completed_after_ns)
            .field("worker_id", &self.worker_id)
            .field(
                "waker_registered",
                &self.waker.with_slot(|slot| slot.is_some()),
            )
            .finish()
    }
}

impl TaskState {
    #[inline]
    pub(super) fn new() -> Self {
        Self {
            created_at: Instant::now(),
            started_after_ns: AtomicU64::new(TIMESTAMP_NOT_RECORDED),
            completed_after_ns: AtomicU64::new(TIMESTAMP_NOT_RECORDED),
            worker_id: AtomicU32::new(NO_WORKER),
            waker: WakerSlot::new(),
            token_active: AtomicBool::new(true),
            // Lossless enum-to-int cast: Priority discriminants are 0..=3.
            priority: AtomicU8::new(Priority::Normal as u8),
            cancel_requested: AtomicBool::new(false),
            cancelled: AtomicBool::new(false),
        }
    }

    #[inline]
    pub(super) fn set_priority(&self, priority: Priority) {
        // Lossless enum-to-int cast: Priority discriminants are 0..=3.
        self.priority.store(priority as u8, Ordering::Relaxed);
    }

    #[inline]
    pub(super) fn priority(&self) -> Priority {
        // Invariant: the slot only ever stores `priority as u8` (0..=3), so the
        // lookup cannot go out of bounds.
        PRIORITY_FROM_INDEX[usize::from(self.priority.load(Ordering::Relaxed))]
    }

    /// Flag the task for cooperative cancellation.
    #[inline]
    pub(super) fn request_cancel(&self) {
        self.cancel_requested.store(true, Ordering::Release);
    }

    #[inline]
    pub(super) fn cancel_requested(&self) -> bool {
        self.cancel_requested.load(Ordering::Acquire)
    }

    #[inline]
    pub(super) fn is_cancelled(&self) -> bool {
        self.cancelled.load(Ordering::Acquire)
    }

    #[inline]
    pub(super) fn token_active(&self) -> bool {
        self.token_active.load(Ordering::Acquire)
    }

    /// Clear the token-active flag, releasing a lease's claim on the state.
    ///
    /// This is a lease's last access. Once the flag is clear a retention sweep
    /// may free the whole block, and it does not wait for the store to return.
    /// A `&TaskState` passed to a method stays valid until that method returns,
    /// so retiring through one lets the block be freed under a live reference;
    /// this takes the state by pointer, projects the flag by address, and
    /// borrows only the one-byte atomic.
    ///
    /// # Safety
    ///
    /// `state` points to a live `TaskState` whose flag the calling lease has
    /// not yet cleared, and the caller makes no further access to it.
    #[inline]
    pub(super) unsafe fn retire_token(state: NonNull<Self>) {
        // SAFETY: the caller guarantees `state` is live; `&raw const` projects
        // the field address without creating a reference to the `TaskState`.
        let flag = unsafe { &raw const (*state.as_ptr()).token_active };
        // SAFETY: the flag is a live `AtomicBool` for the same reason. The
        // borrow covers this one byte and ends with the store.
        unsafe { (*flag).store(false, Ordering::Release) };
    }

    /// Publish that a cancel request was honored: the task completes without
    /// its body having run, and any registered waiter is woken.
    pub(super) fn mark_cancelled(&self) {
        self.cancelled.store(true, Ordering::Release);
        self.mark_completed();
    }

    #[inline]
    pub(super) fn mark_started(&self, worker_id: usize) -> u64 {
        let started_after_ns = elapsed_nanos_since(self.created_at);
        self.started_after_ns
            .store(started_after_ns, Ordering::Release);
        // A real scheduler never reaches `u32::MAX` workers; on a hypothetical
        // wider pool, report "no worker" rather than mis-attribute one.
        self.worker_id.store(
            u32::try_from(worker_id).unwrap_or(NO_WORKER),
            Ordering::Release,
        );
        started_after_ns
    }

    #[inline]
    pub(super) fn mark_completed_since(&self, started_after_ns: u64) -> Duration {
        // `Instant` documents saturation for rare platform monotonicity
        // violations. Preserve that contract across thread/core migration by
        // clamping the published completion offset to the recorded start.
        let completed_after_ns = elapsed_nanos_since(self.created_at).max(started_after_ns);
        self.completed_after_ns
            .store(completed_after_ns, Ordering::Release);

        // The lock is released before the waker runs: `with_slot` takes the
        // waker out and releases the spin lock, and only then does it wake —
        // so a waker that waits on this task cannot deadlock against the
        // slot. Same discipline as `moirai-async`'s sync primitives.
        let waker = self.waker.with_slot(|slot| slot.take());
        if let Some(waker) = waker {
            waker.wake();
        }

        Duration::from_nanos(completed_after_ns - started_after_ns)
    }

    pub(super) fn mark_completed(&self) {
        let started_after_ns = self.started_after_ns.load(Ordering::Acquire);
        let started_after_ns = if started_after_ns == TIMESTAMP_NOT_RECORDED {
            elapsed_nanos_since(self.created_at)
        } else {
            started_after_ns
        };
        self.mark_completed_since(started_after_ns);
    }

    /// Register a waker to be notified when the task completes.
    ///
    /// Wakers that do not `will_wake` the held one accumulate in the slot
    /// ([`fan_out`]), so every waiter is woken by completion.
    pub(super) fn register_waker(&self, waker: &std::task::Waker) {
        self.waker.with_slot(|slot| fan_out::register(slot, waker));
        // Store first, then re-check completion, mirroring the ordering
        // `mark_completed_since` publishes: it stores the completion offset
        // before taking the waker. A task that completed before this store
        // has already taken the absent waker and will never take again, so
        // the one just stored would be held for the life of the slot —
        // along with whatever it owns, typically an `Arc` to async task
        // state. Reclaiming it here is race-free in both directions: if
        // completion lands after the store, it takes and wakes; if it
        // landed before, this take wins and wakes instead. Only one take
        // can succeed, and a spurious wake is always permitted.
        if self.is_completed() {
            let stranded = self.waker.with_slot(|slot| slot.take());
            if let Some(stranded) = stranded {
                stranded.wake();
            }
        }
    }

    pub(super) fn is_completed(&self) -> bool {
        self.completed_after_ns.load(Ordering::Acquire) != TIMESTAMP_NOT_RECORDED
    }

    pub(super) fn completed_at(&self) -> Option<Instant> {
        instant_from_offset(
            self.created_at,
            self.completed_after_ns.load(Ordering::Acquire),
        )
    }

    pub(super) fn snapshot(&self, id: u64) -> TaskMetadata {
        let worker_id = match self.worker_id.load(Ordering::Acquire) {
            NO_WORKER => None,
            // Lossless widening: `usize` is at least 32 bits on every
            // supported target.
            worker_id => Some(worker_id as usize),
        };

        TaskMetadata {
            id,
            created_at: self.created_at,
            started_at: instant_from_offset(
                self.created_at,
                self.started_after_ns.load(Ordering::Acquire),
            ),
            completed_at: self.completed_at(),
            worker_id,
            priority: self.priority(),
            cancelled: self.is_cancelled(),
        }
    }
}

impl TaskStateBlock {
    pub(super) fn new() -> Self {
        let published = std::iter::repeat_with(|| AtomicBool::new(false))
            .take(TASK_STATE_BLOCK_SIZE)
            .collect();
        let states = std::iter::repeat_with(|| UnsafeCell::new(MaybeUninit::uninit()))
            .take(TASK_STATE_BLOCK_SIZE)
            .collect();

        Self { published, states }
    }

    /// Shared view of the state at `slot`, if the slot is registered and in range.
    pub(super) fn get(&self, slot: usize) -> Option<&TaskState> {
        if !self.published.get(slot)?.load(Ordering::Acquire) {
            return None;
        }
        // SAFETY: the acquire load observed the release store that follows the
        // slot's one write, so the state is initialized; nothing writes it again.
        Some(unsafe { (*self.states[slot].get()).assume_init_ref() })
    }

    /// Register a fresh state at `slot`, returning its stable address.
    ///
    /// # Safety
    ///
    /// No other call to `insert` has been or will be made for `slot`. The
    /// registry meets this by issuing every task id exactly once from its
    /// atomic counter and registering each id once.
    pub(super) unsafe fn insert(&self, slot: usize) -> NonNull<TaskState> {
        debug_assert!(
            !self.published[slot].load(Ordering::Relaxed),
            "a task id registers exactly once"
        );
        let cell = self.states[slot].get();
        // SAFETY: the caller is the slot's only registrant and the slot is not
        // yet published, so no reader touches the cell and this is the only
        // access to it. The address is stable: the boxed slice never moves or
        // shrinks and a published state is never replaced.
        unsafe { (*cell).write(TaskState::new()) };
        // SAFETY: `UnsafeCell::get` never returns null. The address derives
        // from the cell itself, as `get` does, and not from the `&mut` that
        // `write` returns: a pointer reborrowed from that transient reference
        // sits above it in the borrow stack, and the shared reads and atomic
        // stores that later readers make through the cell would invalidate it
        // before the lease uses it.
        let address = unsafe { NonNull::new_unchecked(cell.cast::<TaskState>()) };
        self.published[slot].store(true, Ordering::Release);
        address
    }

    /// Whether every task this block held has completed and released its
    /// lifecycle token, so the block can retire.
    ///
    /// Every slot must be registered: a vacant slot is an id that was issued
    /// but whose registration has not run yet, and retiring the block under it
    /// would lose that registration. `first_slot_unissued` exempts slot 0 of
    /// block 0, the one id the registry never issues. `retire` bounds the
    /// completion age; slots are examined newest first, so a block still inside
    /// its retention window is rejected after the first slot or two.
    pub(super) fn is_settled(&self, first_slot_unissued: bool, retire: Retirement) -> bool {
        (0..TASK_STATE_BLOCK_SIZE)
            .rev()
            .all(|slot| match self.get(slot) {
                None => first_slot_unissued && slot == 0,
                Some(state) => {
                    !state.token_active()
                        && state.completed_at().is_some_and(|completed| match retire {
                            Retirement::Forced => true,
                            Retirement::CompletedBefore(cutoff) => completed <= cutoff,
                        })
                }
            })
    }

    /// Iterate shared views of the registered states in this block.
    pub(super) fn states(&self) -> impl Iterator<Item = &TaskState> {
        (0..TASK_STATE_BLOCK_SIZE).filter_map(|slot| self.get(slot))
    }
}

impl Drop for TaskStateBlock {
    fn drop(&mut self) {
        for (published, state) in self.published.iter_mut().zip(self.states.iter_mut()) {
            if *published.get_mut() {
                // SAFETY: the flag is set only after the state is written, and
                // `&mut self` excludes every other access.
                unsafe { state.get_mut().assume_init_drop() };
            }
        }
    }
}

impl std::fmt::Debug for TaskStateBlock {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("TaskStateBlock")
            .field("slots", &self.published.len())
            .field("occupied", &self.states().count())
            .finish()
    }
}

#[inline]
pub(crate) fn elapsed_nanos_since(origin: Instant) -> u64 {
    let elapsed = origin.elapsed().as_nanos();
    elapsed.min(u128::from(TIMESTAMP_NOT_RECORDED - 1)) as u64
}

pub(crate) fn instant_from_offset(origin: Instant, offset_ns: u64) -> Option<Instant> {
    if offset_ns == TIMESTAMP_NOT_RECORDED {
        None
    } else {
        origin.checked_add(Duration::from_nanos(offset_ns))
    }
}

pub(crate) fn task_location(id: u64) -> (usize, usize) {
    let index = usize::try_from(id).expect("task ID must fit in usize");
    (index / TASK_STATE_BLOCK_SIZE, index % TASK_STATE_BLOCK_SIZE)
}

#[cfg(test)]
mod tests {
    use super::{TaskState, elapsed_nanos_since};
    use std::sync::atomic::Ordering;
    use std::time::Duration;

    #[test]
    fn completion_clamps_to_recorded_start_offset() {
        let state = TaskState::new();
        let future_start = elapsed_nanos_since(state.created_at).saturating_add(1_000_000);
        state
            .started_after_ns
            .store(future_start, Ordering::Release);

        let elapsed = state.mark_completed_since(future_start);
        let snapshot = state.snapshot(7);

        assert_eq!(elapsed, Duration::ZERO);
        assert_eq!(snapshot.started_at, snapshot.completed_at);
        assert_eq!(snapshot.execution_duration(), Some(Duration::ZERO));
    }
}