moirai-executor 0.9.1

//! ThreadScheduler core implementation.

use std::sync::{Arc, atomic::Ordering};

use moirai_core::{
    Priority,
    error::{ExecutorError, ExecutorResult},
};

use super::super::super::{class::WorkClass, job::ScheduledJob};
use super::super::types::{BoundedContendedWake, ThreadScheduler, get_current_worker_id};
use super::super::worker::{
    JOIN_FAST_SPIN_ATTEMPTS, is_quiescent, lock_mutex, wake_contended_workers, wake_worker,
};

/// A per-thread round-robin ticket for spreading queued submissions across
/// workers. Replaces a process-shared `AtomicUsize` that every producer thread
/// RMW'd on each submit: that counter's cache line bounced between all producing
/// cores under high submit rates. Each producer now rotates its own thread-local
/// sequence — contention-free, and still a uniform spread (the value is only a
/// load-balancing hint, not a synchronization point). A single producer still
/// rotates `0, 1, 2, …`, preserving the round-robin the runtime tests pin.
#[inline]
fn next_round_robin_ticket() -> usize {
    use std::cell::Cell;
    std::thread_local! {
        // clippy 1.97.0 FP: already const. ATLAS-MNEMOSYNE-CI-1.
        #[allow(clippy::missing_const_for_thread_local)]
        static TICKET: Cell<usize> = const { Cell::new(0) };
    }
    TICKET.with(|cell| {
        let ticket = cell.get();
        cell.set(ticket.wrapping_add(1));
        ticket
    })
}

impl<const BLOCKING_QUEUE_CAPACITY: usize, const SPIN_LIMIT: usize>
    ThreadScheduler<BLOCKING_QUEUE_CAPACITY, SPIN_LIMIT>
{
    /// Schedule a job for a compile-time work class.
    pub fn schedule<C, F>(
        &self,
        priority: Priority,
        locality_hint: Option<usize>,
        task: F,
    ) -> ExecutorResult<()>
    where
        C: WorkClass,
        F: FnOnce(usize) + Send + 'static,
    {
        let job = ScheduledJob::new(task);
        self.schedule_job::<C>(priority, locality_hint, job)
    }

    /// Lane identifier passed to a job the caller runs itself.
    ///
    /// Worker lanes are `0..worker_count()`; the calling thread is the lane
    /// past the last worker. It is a lane identity, never an index into the
    /// worker set.
    pub(crate) fn caller_lane_id(&self) -> usize {
        get_current_worker_id().unwrap_or_else(|| self.worker_count())
    }

    /// Schedule a job, dropping it if admission refuses it.
    ///
    /// Callers that can still run a refused job themselves use
    /// [`Self::admit_job`] instead.
    pub(crate) fn schedule_job<C>(
        &self,
        priority: Priority,
        locality_hint: Option<usize>,
        job: ScheduledJob,
    ) -> ExecutorResult<()>
    where
        C: WorkClass,
    {
        self.admit_job::<C>(priority, locality_hint, &mut Some(job))
    }

    /// Schedule a job, leaving it in `job` when admission refuses it.
    ///
    /// A refused job has not run and never will unless the caller runs it, so
    /// dropping it silently discards work. Admission failure and job loss are
    /// separate facts: a job a full queue turned away is still in `job` and
    /// still runnable, while one a lane consumed while failing is not, so the
    /// caller checks rather than assumes. `SchedulerScope::flush` runs what is
    /// left behind on the calling lane, which is how a scope keeps its "every
    /// spawned job runs before the scope returns" contract under backpressure.
    ///
    /// The job travels by slot rather than inside the error so the hot path
    /// keeps a small `Result`: a `ScheduledJob` carries fixed-size inline
    /// storage, and an `Err` variant holding one would widen every admission
    /// return, success included.
    pub(crate) fn admit_job<C>(
        &self,
        priority: Priority,
        locality_hint: Option<usize>,
        job: &mut Option<ScheduledJob>,
    ) -> ExecutorResult<()>
    where
        C: WorkClass,
    {
        if C::USES_BLOCKING_LANE {
            if self.inner.shutdown.load(Ordering::Acquire) {
                return Err(ExecutorError::ShuttingDown);
            }
            if let Some(lane) = self.inner.blocking_lane.get() {
                return lane.submit(
                    priority,
                    locality_hint,
                    job,
                    &self.inner.blocking_pending_tasks,
                );
            }

            let _lane_init = lock_mutex(&self.inner.blocking_lane_init);
            if self.inner.shutdown.load(Ordering::Acquire) {
                return Err(ExecutorError::ShuttingDown);
            }
            if self.inner.blocking_lane.get().is_none() {
                let candidate = super::super::blocking::BlockingLane::new(self.worker_count());
                candidate.start(Arc::clone(&self.inner), &self.inner.blocking_lane_prefix)?;
                if self.inner.blocking_lane.set(candidate).is_err() {
                    return Err(ExecutorError::ThreadPoolCreationFailed);
                }
            }
            let lane = self
                .inner
                .blocking_lane
                .get()
                .expect("invariant: blocking lane initialized before submission");
            return lane.submit(
                priority,
                locality_hint,
                job,
                &self.inner.blocking_pending_tasks,
            );
        }

        let pending_before_submit = self.inner.pending_tasks.load(Ordering::Acquire);
        let active_before_submit = self.inner.active_workers.load(Ordering::Acquire);
        let worker_index = self.select_worker_for_state::<C>(
            priority,
            locality_hint,
            pending_before_submit,
            active_before_submit,
        );
        // SeqCst (not Release): this increment is one half of a store-buffer /
        // Dekker handshake with a parking worker, which does
        // `idle_workers.fetch_or(mask, SeqCst)` then `pending_tasks.load(SeqCst)`
        // before parking, while the producer below does this increment then
        // `idle_workers.load(SeqCst)`. Correctness requires all four accesses to
        // share one SeqCst total order; a Release RMW is NOT in that order, which
        // permits the worker to read `pending == 0` AND the producer to read the
        // worker's idle bit as clear — a lost wakeup that stalls the task until an
        // unrelated submission. The increment must stay BEFORE the push so the
        // `execute_job` decrement (worker.rs) can never underflow from 0.
        // (On x86 `lock xadd` is already full-barrier, so this is free.)
        let previous_pending = self.inner.pending_tasks.fetch_add(1, Ordering::SeqCst);
        // Pair with the no-work worker's SeqCst shutdown/pending observation.
        // Either admission sees shutdown and rolls back, or a draining worker
        // sees this pending publication and remains available to execute it.
        if self.inner.shutdown.load(Ordering::SeqCst) {
            self.inner.pending_tasks.fetch_sub(1, Ordering::SeqCst);
            return Err(ExecutorError::ShuttingDown);
        }

        let admitting = job
            .take()
            .expect("invariant: job is present before admission");
        let rejected = if get_current_worker_id() == Some(worker_index) {
            self.inner.workers[worker_index]
                .lifo_slot
                .try_push(admitting)
                .and_then(|job| {
                    self.inner.workers[worker_index]
                        .queues
                        .try_push_external(priority, job)
                })
        } else {
            self.inner.workers[worker_index]
                .queues
                .try_push_external(priority, admitting)
        };

        if let Some(rejected) = rejected {
            self.inner.pending_tasks.fetch_sub(1, Ordering::SeqCst);
            // Back in the caller's slot: it never entered a queue, so whoever
            // owns the slot can still run it.
            *job = Some(rejected);
            return Err(ExecutorError::ResourceExhausted(format!(
                "worker {worker_index} scheduler admission queue is full"
            )));
        }

        // Try to wake up an idle worker via the lock-free wake lottery. The
        // bitset spans every worker (not just the first 64), so high-index
        // workers on large pools are reachable directly rather than stranded.
        let mut woken = false;
        if let Some(worker_to_wake) = self.inner.idle_workers.claim_one(self.inner.workers.len()) {
            wake_worker(&self.inner.workers[worker_to_wake]);
            woken = true;
        }

        if !woken {
            if previous_pending == 0 {
                wake_worker(&self.inner.workers[worker_index]);
            } else {
                let worker_count = self.inner.workers.len();
                if previous_pending < worker_count {
                    let _ = wake_contended_workers::<BoundedContendedWake>(
                        &*self.inner,
                        worker_index,
                        previous_pending,
                    );
                } else {
                    let wake_index = worker_index.wrapping_add(previous_pending) % worker_count;
                    wake_worker(&self.inner.workers[wake_index]);
                }
            }
        }
        Ok(())
    }

    /// Approximate number of queued jobs.
    pub fn pending_tasks(&self) -> usize {
        self.inner.pending_tasks.load(Ordering::Acquire)
            + self.inner.blocking_pending_tasks.load(Ordering::Acquire)
    }

    /// Number of workers currently executing jobs.
    pub fn active_workers(&self) -> usize {
        self.inner.active_workers.load(Ordering::Acquire)
            + self.inner.blocking_active_workers.load(Ordering::Acquire)
    }

    /// Returns true when queued or active work exists.
    pub fn has_work(&self) -> bool {
        !is_quiescent(&self.inner)
    }

    /// Wait until all queued and active work has completed without stopping workers.
    pub fn join(&self) -> ExecutorResult<()> {
        for _ in 0..JOIN_FAST_SPIN_ATTEMPTS {
            if is_quiescent(&self.inner) {
                return Ok(());
            }
            core::hint::spin_loop();
        }

        let mut guard = lock_mutex(&self.inner.wait_lock);
        // SeqCst: the joiner half of the quiescence Dekker handshake — this
        // registration must be in the same SeqCst total order as a completing
        // worker's `active_workers` decrement and its `join_waiters` load in
        // `notify_quiescent`. With AcqRel the store-buffer outcome hangs `join()`
        // (see `execute_job` and `tests/loom_join_quiescence.rs`). The re-check
        // below runs under `wait_lock`, closing the window against a concurrent
        // `notify_quiescent` that also takes the lock before signalling.
        self.inner.join_waiters.fetch_add(1, Ordering::SeqCst);
        while !is_quiescent(&self.inner) {
            guard = self
                .inner
                .wait_signal
                .wait(guard)
                .unwrap_or_else(|poisoned| poisoned.into_inner());
        }
        self.inner.join_waiters.fetch_sub(1, Ordering::AcqRel);
        Ok(())
    }

    /// Number of worker threads.
    pub fn worker_count(&self) -> usize {
        self.inner.workers.len()
    }

    #[cfg(test)]
    pub(crate) fn injector_capacity_per_worker(&self) -> usize {
        self.inner.workers[0].queues.injector_capacity()
    }

    /// Capture scheduler metrics.
    pub fn metrics(&self) -> super::super::types::ScheduleMetrics {
        super::super::types::ScheduleMetrics {
            worker_count: self.worker_count(),
            pending_tasks: self.pending_tasks(),
            active_workers: self.active_workers(),
            completed_tasks: self.inner.completed_tasks.load(Ordering::Acquire),
            failed_tasks: self.inner.failed_tasks.load(Ordering::Acquire),
        }
    }

    /// Number of indexed chunks run by their submitting caller after bounded
    /// scheduler admission rejected the corresponding worker job.
    ///
    /// This monotonic diagnostic counter is observational only. Its relaxed
    /// atomic ordering intentionally establishes no synchronization edge.
    #[must_use]
    pub fn admission_caller_runs(&self) -> u64 {
        self.inner.admission_caller_runs.load(Ordering::Relaxed)
    }

    pub(super) fn record_admission_caller_run(&self) {
        self.inner
            .admission_caller_runs
            .fetch_add(1, Ordering::Relaxed);
    }

    pub(crate) fn select_worker<C>(&self, priority: Priority, locality_hint: Option<usize>) -> usize
    where
        C: WorkClass,
    {
        self.select_worker_for_state::<C>(
            priority,
            locality_hint,
            self.inner.pending_tasks.load(Ordering::Acquire),
            self.inner.active_workers.load(Ordering::Acquire),
        )
    }

    pub(crate) fn select_worker_for_state<C>(
        &self,
        priority: Priority,
        locality_hint: Option<usize>,
        pending_tasks: usize,
        active_workers: usize,
    ) -> usize
    where
        C: WorkClass,
    {
        let worker_count = self.inner.workers.len();
        if let Some(hint) = locality_hint {
            return hint % worker_count;
        }

        if pending_tasks == 0 && active_workers <= 1 {
            return C::SERIAL_AFFINITY_OFFSET.wrapping_add(priority.index()) % worker_count;
        }

        let ticket = next_round_robin_ticket();
        ticket
            .wrapping_add(C::AFFINITY_OFFSET)
            .wrapping_add(priority.index())
            % worker_count
    }
}