ferralk 0.9.4

//! Work-stealing queue primitives and the worker protocol behind them.
//!
//! The protocol types below use `loom`'s primitives when the crate is built
//! with `--cfg loom`, so the models in `loom_models` exercise the very code the
//! walker runs instead of a re-implementation of it.

use std::time::Duration;

#[cfg(loom)]
use loom::sync::{
    Condvar, Mutex, MutexGuard,
    atomic::{AtomicBool, AtomicUsize, Ordering},
};
#[cfg(not(loom))]
use std::sync::{
    Condvar, Mutex, MutexGuard,
    atomic::{AtomicBool, AtomicUsize, Ordering},
};

use crossbeam_deque::{Injector, Steal, Worker};

/// Backstop for cancellation requested through a `CancellationToken` that the
/// walk does not own, which therefore cannot notify the parked workers.
///
/// Progress never depends on it: hand-off and termination are driven by
/// notifications that are sent while `wake_lock` is held.
const CANCELLATION_POLL_INTERVAL: Duration = Duration::from_millis(10);

/// Low bits reserved for parked-worker registrations in `wake_state`.
///
/// Deriving the width from the walk's hard worker cap keeps the packed field's
/// no-carry invariant tied to the budget it protects.
const SLEEPER_BITS: u32 = usize::BITS - crate::MAX_WORKERS.leading_zeros();
const SLEEPER_MASK: usize = (1 << SLEEPER_BITS) - 1;
const WAKE_EPOCH_STEP: usize = 1 << SLEEPER_BITS;

pub(crate) struct Scheduler<T> {
    injector: Injector<T>,
}

impl<T> Scheduler<T> {
    pub(crate) fn new() -> Self {
        Self {
            injector: Injector::new(),
        }
    }

    pub(crate) fn push(&self, task: T) {
        self.injector.push(task);
    }

    pub(crate) fn worker(&self) -> Worker<T> {
        Worker::new_fifo()
    }

    /// Whether the shared queue holds nothing a worker could take.
    pub(crate) fn is_empty(&self) -> bool {
        self.injector.is_empty()
    }

    pub(crate) fn steal_into(&self, worker: &Worker<T>) -> Option<T> {
        loop {
            match self.injector.steal_batch_and_pop(worker) {
                Steal::Success(task) => return Some(task),
                Steal::Empty => return None,
                Steal::Retry => continue,
            }
        }
    }
}

/// Task accounting and parking protocol shared by the workers of one walk.
///
/// Tasks are counted here rather than queued: `begin_task` announces one,
/// the guard returned by `claim_task` releases it even while a panic unwinds,
/// and `wait_for_task` parks a worker until work or the end of the walk is
/// observable. Every notification is sent while `wake_lock` is held, so a
/// worker that has decided to park cannot miss one.
pub(crate) struct Coordinator {
    pending: CacheLine<AtomicUsize>,
    active_workers: CacheLine<AtomicUsize>,
    /// Publication counter checked by a waiter after its optimistic queue
    /// probe. It keeps that waiter from sleeping through work published while
    /// it was preparing to park.
    /// Publication epoch in the high bits and registered sleepers in the low
    /// bits. Keeping them in one atomic makes the producer's epoch bump and
    /// sleeper observation indivisible.
    wake_state: AtomicUsize,
    /// Whether the caller is still the only worker. This state belongs beside
    /// the parking protocol because it decides whether producing work must
    /// notify a waiter.
    lone: AtomicBool,
    wake_lock: Mutex<()>,
    wake: Condvar,
}

/// One value with a cache line to itself.
///
/// The counters this wraps are written once per directory by whichever worker
/// read it and loaded by every other worker deciding whether to park or to grow
/// the pool. Packed next to each other they share a line, so one worker's
/// increment invalidates the line the others are reading — false sharing
/// between counters that have nothing to say to each other.
///
/// 128 bytes rather than 64 on the two architectures that have a native
/// backend: Apple silicon and x86-64 both move cache lines in adjacent pairs,
/// so 64-byte separation leaves two values in the same coherence granule.
/// Elsewhere a line is 64 bytes and the padding follows.
#[cfg_attr(any(target_arch = "aarch64", target_arch = "x86_64"), repr(align(128)))]
#[cfg_attr(
    not(any(target_arch = "aarch64", target_arch = "x86_64")),
    repr(align(64))
)]
pub(crate) struct CacheLine<T>(pub(crate) T);

impl<T> std::ops::Deref for CacheLine<T> {
    type Target = T;

    fn deref(&self) -> &T {
        &self.0
    }
}

impl Coordinator {
    pub(crate) fn new() -> Self {
        Self {
            pending: CacheLine(AtomicUsize::new(0)),
            active_workers: CacheLine(AtomicUsize::new(0)),
            wake_state: AtomicUsize::new(0),
            lone: AtomicBool::new(true),
            wake_lock: Mutex::new(()),
            wake: Condvar::new(),
        }
    }

    /// Announces one task that some worker will claim later.
    pub(crate) fn begin_task(&self) {
        self.pending.fetch_add(1, Ordering::AcqRel);
    }

    /// Makes one task available to a worker.
    ///
    /// The producer publishes the task before checking `lone`. In lone mode
    /// this stays on the old lock-free fast path: no helper can exist yet. If
    /// widening wins the check, one helper is notified while holding the same
    /// lock a waiter uses to enter the condvar. If production wins it, widening
    /// starts helpers only after the task has been published, so their first
    /// queue probe observes it. Keeping both operations here lets loom drive
    /// the real fast path rather than a stand-in.
    pub(crate) fn schedule(&self, publish: impl FnOnce()) {
        self.begin_task();
        publish();
        if self.lone.load(Ordering::Acquire) {
            return;
        }
        self.announce_one();
    }

    /// Announces that helpers may exist. The caller invokes this before
    /// spawning the first helper, so a producer that still sees `lone` cannot
    /// strand work behind a parked worker.
    pub(crate) fn widen(&self) {
        self.lone.store(false, Ordering::Release);
    }

    /// Takes ownership of one announced task. The guard releases it on drop,
    /// including while a panic unwinds through the work.
    pub(crate) fn claim_task(&self) -> TaskGuard<'_> {
        TaskGuard { coordinator: self }
    }

    pub(crate) fn pending(&self) -> usize {
        self.pending.load(Ordering::Acquire)
    }

    /// Wakes every parked worker. Used for the states every worker has to see:
    /// the walk draining, being cancelled, aborting, or unwinding.
    pub(crate) fn wake_waiters(&self) {
        let guard = lock(&self.wake_lock);
        self.wake.notify_all();
        drop(guard);
    }

    /// Wakes one waiter after a batch steal makes the other entries in that
    /// batch visible in the thief's local deque. Those entries were announced
    /// before the steal, so no ordinary producer will wake for them.
    pub(crate) fn announce_batch_handoff(&self) {
        self.announce_one();
    }

    /// Publishes an observation point before waking one waiter. A worker that
    /// sees this epoch change leaves the mutex and retries its queue probes,
    /// so no potentially expensive steal scan runs while `wake_lock` is held.
    fn announce_one(&self) {
        let previous = self.wake_state.fetch_add(WAKE_EPOCH_STEP, Ordering::SeqCst);
        // The same RMW publishes a new epoch and observes the sleeper count.
        // A racing waiter either registered in `previous`, or its registration
        // observes this new epoch and it does not enter the condvar.
        if previous & SLEEPER_MASK == 0 {
            return;
        }
        let guard = lock(&self.wake_lock);
        self.wake.notify_one();
        drop(guard);
    }

    /// Parks until `try_take` yields a task, the walk is cancelled, or no task
    /// is outstanding any more.
    pub(crate) fn wait_for_task<T>(
        &self,
        should_stop: impl Fn() -> bool,
        mut try_take: impl FnMut() -> Option<T>,
    ) -> Option<T> {
        loop {
            if should_stop() {
                return None;
            }
            let observed_epoch = self.wake_state.load(Ordering::SeqCst) & !SLEEPER_MASK;
            if let Some(task) = try_take() {
                return Some(task);
            }
            let guard = lock(&self.wake_lock);
            // Re-checked under the lock: a notification sent from here on
            // cannot be missed, and anything announced earlier is visible.
            if self.pending.load(Ordering::Acquire) == 0 || should_stop() {
                return None;
            }
            // Register before the final epoch check. A producer that already
            // skipped the notification changed the epoch first; a later
            // producer observes this registration and takes the locked path.
            let previous = self.wake_state.fetch_add(1, Ordering::SeqCst);
            debug_assert_ne!(previous & SLEEPER_MASK, SLEEPER_MASK);
            if previous & !SLEEPER_MASK != observed_epoch {
                self.wake_state.fetch_sub(1, Ordering::SeqCst);
                drop(guard);
                continue;
            }
            let (guard, _) = self
                .wake
                .wait_timeout(guard, CANCELLATION_POLL_INTERVAL)
                .unwrap_or_else(|poisoned| poisoned.into_inner());
            self.wake_state.fetch_sub(1, Ordering::SeqCst);
            drop(guard);
        }
    }

    /// Reserves the slot of the thread that starts the walk.
    pub(crate) fn claim_caller_slot(&self) -> WorkerSlot<'_> {
        self.active_workers.fetch_add(1, Ordering::AcqRel);
        WorkerSlot { coordinator: self }
    }

    /// Reserves a slot for one more worker while queued tasks outnumber the
    /// workers already running and `budget` has room. The compare-exchange
    /// keeps concurrent growth attempts from overshooting the budget.
    pub(crate) fn claim_worker_slot(&self, budget: usize) -> Option<WorkerSlot<'_>> {
        let mut active = self.active_workers.load(Ordering::Acquire);
        loop {
            if active >= budget || self.pending() <= active {
                return None;
            }
            match self.active_workers.compare_exchange_weak(
                active,
                active + 1,
                Ordering::AcqRel,
                Ordering::Acquire,
            ) {
                Ok(_) => return Some(WorkerSlot { coordinator: self }),
                Err(observed) => active = observed,
            }
        }
    }

    /// Observation point for the loom models; the walk itself only ever asks
    /// for slots through [`Coordinator::claim_worker_slot`].
    #[cfg(loom)]
    pub(crate) fn active_workers(&self) -> usize {
        self.active_workers.load(Ordering::Acquire)
    }

    /// Observation point for loom: every exit from the wait protocol must
    /// retire its sleeper registration.
    #[cfg(loom)]
    pub(crate) fn sleepers(&self) -> usize {
        self.wake_state.load(Ordering::SeqCst) & SLEEPER_MASK
    }

    fn finish_task(&self) {
        if self.pending.fetch_sub(1, Ordering::AcqRel) == 1 {
            self.wake_waiters();
        }
    }
}

/// One in-flight task. Dropping it releases the task, so a panic inside the
/// work cannot strand the pending count and leave siblings waiting forever.
pub(crate) struct TaskGuard<'a> {
    coordinator: &'a Coordinator,
}

impl Drop for TaskGuard<'_> {
    fn drop(&mut self) {
        self.coordinator.finish_task();
    }
}

/// One running worker. Dropping it frees the slot for a later helper, whether
/// the worker finished, failed to start, or unwound.
pub(crate) struct WorkerSlot<'a> {
    coordinator: &'a Coordinator,
}

impl Drop for WorkerSlot<'_> {
    fn drop(&mut self) {
        self.coordinator
            .active_workers
            .fetch_sub(1, Ordering::AcqRel);
    }
}

fn lock<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
    mutex
        .lock()
        .unwrap_or_else(|poisoned| poisoned.into_inner())
}

#[cfg(all(test, not(loom)))]
mod tests {
    use super::Scheduler;

    #[test]
    fn injector_distributes_tasks_to_worker_local_queues() {
        let scheduler = Scheduler::new();
        scheduler.push(1);
        scheduler.push(2);
        let worker = scheduler.worker();

        let first = scheduler.steal_into(&worker).expect("first task");
        let second = worker.pop().unwrap_or_else(|| {
            scheduler
                .steal_into(&worker)
                .expect("second task after local batch")
        });
        assert_ne!(first, second);
        assert!(scheduler.steal_into(&worker).is_none());
    }
}

/// Loom models of the production protocol. They drive [`Coordinator`] itself,
/// with loom's primitives substituted underneath it, so every interleaving of
/// the real hand-off is checked.
///
/// Run with `RUSTFLAGS="--cfg loom" cargo test -p ferralk --lib loom_models`.
///
/// Note that loom's `Condvar::wait_timeout` never times out, so the poll
/// interval cannot paper over a missing notification here: a lost wakeup shows
/// up as a deadlock.
#[cfg(loom)]
mod loom_models {
    use std::{
        collections::VecDeque,
        panic::{AssertUnwindSafe, catch_unwind},
        sync::atomic::{AtomicUsize, Ordering},
    };

    use loom::{
        sync::{Arc, Mutex},
        thread,
    };

    use super::Coordinator;

    /// Reports how much of the state space a model actually covered, so a run
    /// shows its work instead of just passing silently.
    fn report(model: &str, executions: &AtomicUsize) {
        eprintln!(
            "loom explored {} executions of {model}",
            executions.load(Ordering::Relaxed)
        );
    }

    fn pop(queue: &Mutex<VecDeque<usize>>) -> Option<usize> {
        queue
            .lock()
            .unwrap_or_else(|poisoned| poisoned.into_inner())
            .pop_front()
    }

    fn push(queue: &Mutex<VecDeque<usize>>, task: usize) {
        queue
            .lock()
            .unwrap_or_else(|poisoned| poisoned.into_inner())
            .push_back(task);
    }

    /// A task announced while a worker is on its way into `wait_for_task` must
    /// still reach that worker, and the walk must end once it is done.
    ///
    /// The producer notifies through the single-worker path the walk uses when
    /// it queues one directory, so this covers `notify_one` as well as the
    /// hand-off. Loom's `wait_timeout` never times out, so a notification that
    /// failed to reach the consumer shows up here as a deadlock rather than as
    /// a slow test.
    #[test]
    fn a_queued_task_reaches_a_parking_worker() {
        static EXECUTIONS: AtomicUsize = AtomicUsize::new(0);
        loom::model(|| {
            EXECUTIONS.fetch_add(1, Ordering::Relaxed);
            let coordinator = Arc::new(Coordinator::new());
            let queue = Arc::new(Mutex::new(VecDeque::new()));
            // Helpers exist for this hand-off, so producing work must notify
            // one of them through the production scheduling path.
            coordinator.widen();
            coordinator.begin_task();

            let producer_coordinator = Arc::clone(&coordinator);
            let producer_queue = Arc::clone(&queue);
            let producer = thread::spawn(move || {
                let _root = producer_coordinator.claim_task();
                producer_coordinator.schedule(|| push(&producer_queue, 7));
            });

            let consumer_coordinator = Arc::clone(&coordinator);
            let consumer_queue = Arc::clone(&queue);
            let consumer = thread::spawn(move || {
                let task = consumer_coordinator.wait_for_task(|| false, || pop(&consumer_queue));
                assert_eq!(task, Some(7), "the queued task must reach the worker");
                drop(consumer_coordinator.claim_task());
                // With every task done the walk has to end instead of parking.
                assert_eq!(
                    consumer_coordinator.wait_for_task(|| false, || pop(&consumer_queue)),
                    None,
                    "a drained walk must let its workers finish"
                );
            });

            producer.join().expect("producer joins");
            consumer.join().expect("consumer joins");
            assert_eq!(coordinator.pending(), 0);
        });
        report("a_queued_task_reaches_a_parking_worker", &EXECUTIONS);
    }

    /// The fast path may observe no sleepers just as a worker registers under
    /// the mutex. Either the announcement sees that registration and notifies,
    /// or the worker sees the preceding epoch change and retries without
    /// sleeping. Loom's non-expiring timeout turns either missed edge into a
    /// deterministic deadlock.
    #[test]
    fn zero_sleeper_skip_races_with_parker_registration() {
        static EXECUTIONS: AtomicUsize = AtomicUsize::new(0);
        loom::model(|| {
            EXECUTIONS.fetch_add(1, Ordering::Relaxed);
            let coordinator = Arc::new(Coordinator::new());
            let queue = Arc::new(Mutex::new(VecDeque::new()));
            coordinator.begin_task();

            let producer_coordinator = Arc::clone(&coordinator);
            let producer_queue = Arc::clone(&queue);
            let producer = thread::spawn(move || {
                push(&producer_queue, 7);
                producer_coordinator.announce_batch_handoff();
            });

            let worker_coordinator = Arc::clone(&coordinator);
            let worker_queue = Arc::clone(&queue);
            let worker = thread::spawn(move || {
                let task = worker_coordinator.wait_for_task(|| false, || pop(&worker_queue));
                assert_eq!(task, Some(7), "the skipped notify must not lose work");
                drop(worker_coordinator.claim_task());
            });

            producer.join().expect("producer joins");
            worker.join().expect("worker joins");
            assert_eq!(coordinator.pending(), 0);
            assert_eq!(coordinator.sleepers(), 0);
        });
        report(
            "zero_sleeper_skip_races_with_parker_registration",
            &EXECUTIONS,
        );
    }

    /// A producer may publish work just as the caller turns its lone walk into
    /// a helper walk. The real `Coordinator` owns that decision: if the
    /// producer still sees lone mode, the helper starts only after the task is
    /// already published; otherwise it receives the normal one-worker wakeup.
    #[test]
    fn widening_does_not_lose_a_concurrent_scheduled_task() {
        static EXECUTIONS: AtomicUsize = AtomicUsize::new(0);
        loom::model(|| {
            EXECUTIONS.fetch_add(1, Ordering::Relaxed);
            let coordinator = Arc::new(Coordinator::new());
            let queue = Arc::new(Mutex::new(VecDeque::new()));
            // The caller is still processing its root while the first helper
            // becomes eligible to wait for a child task.
            coordinator.begin_task();

            let producer_coordinator = Arc::clone(&coordinator);
            let producer_queue = Arc::clone(&queue);
            let producer = thread::spawn(move || {
                let _root = producer_coordinator.claim_task();
                producer_coordinator.schedule(|| push(&producer_queue, 7));
            });

            let helper_coordinator = Arc::clone(&coordinator);
            let helper_queue = Arc::clone(&queue);
            let helper = thread::spawn(move || {
                helper_coordinator.widen();
                let task = helper_coordinator.wait_for_task(|| false, || pop(&helper_queue));
                assert_eq!(task, Some(7), "widening must not strand published work");
                drop(helper_coordinator.claim_task());
            });

            producer.join().expect("producer joins");
            helper.join().expect("helper joins");
            assert_eq!(coordinator.pending(), 0);
        });
        report(
            "widening_does_not_lose_a_concurrent_scheduled_task",
            &EXECUTIONS,
        );
    }

    /// A successful batch steal consumes one task but leaves siblings in the
    /// thief's local queue. The handoff notification must make one of those
    /// siblings observable to a worker that was already preparing to park.
    #[test]
    fn a_batch_handoff_reaches_a_parking_worker() {
        static EXECUTIONS: AtomicUsize = AtomicUsize::new(0);
        loom::model(|| {
            EXECUTIONS.fetch_add(1, Ordering::Relaxed);
            let coordinator = Arc::new(Coordinator::new());
            let queue = Arc::new(Mutex::new(VecDeque::new()));
            coordinator.widen();
            coordinator.begin_task();

            let thief_coordinator = Arc::clone(&coordinator);
            let thief_queue = Arc::clone(&queue);
            let thief = thread::spawn(move || {
                // This is the queue state after `steal_batch_and_pop`: the
                // first stolen task is already executing, and its sibling is
                // now only visible through the thief's local deque.
                push(&thief_queue, 7);
                thief_coordinator.announce_batch_handoff();
            });

            let worker_coordinator = Arc::clone(&coordinator);
            let worker_queue = Arc::clone(&queue);
            let worker = thread::spawn(move || {
                let task = worker_coordinator.wait_for_task(|| false, || pop(&worker_queue));
                assert_eq!(task, Some(7), "the batch sibling must reach a worker");
                drop(worker_coordinator.claim_task());
            });

            thief.join().expect("batch thief joins");
            worker.join().expect("worker joins");
            assert_eq!(coordinator.pending(), 0);
        });
        report("a_batch_handoff_reaches_a_parking_worker", &EXECUTIONS);
    }

    /// The task guard has to release its task while the panic unwinds, or the
    /// sibling worker parks on a count that never reaches zero.
    #[test]
    fn a_panicking_worker_releases_its_task() {
        static EXECUTIONS: AtomicUsize = AtomicUsize::new(0);
        // The injected panic is expected, so its report is filtered out instead
        // of printing once per explored execution. Everything else, loom's own
        // failures included, still reaches the hook that was installed before,
        // which is put back once the model is done.
        let previous = std::sync::Arc::new(std::panic::take_hook());
        let filtered = std::sync::Arc::clone(&previous);
        std::panic::set_hook(Box::new(move |info| {
            if !info.to_string().contains("injected task panic") {
                filtered(info);
            }
        }));

        loom::model(|| {
            EXECUTIONS.fetch_add(1, Ordering::Relaxed);
            let coordinator = Arc::new(Coordinator::new());
            coordinator.begin_task();

            let panicking_coordinator = Arc::clone(&coordinator);
            let panicking = thread::spawn(move || {
                let outcome = catch_unwind(AssertUnwindSafe(|| {
                    let _task = panicking_coordinator.claim_task();
                    panic!("injected task panic");
                }));
                assert!(outcome.is_err(), "the injected panic must be captured");
            });

            let sibling_coordinator = Arc::clone(&coordinator);
            let sibling = thread::spawn(move || {
                assert_eq!(
                    sibling_coordinator.wait_for_task(|| false, || None::<usize>),
                    None,
                    "the sibling must not park forever after a panic"
                );
            });

            panicking.join().expect("panicking worker joins");
            sibling.join().expect("sibling joins");
            assert_eq!(coordinator.pending(), 0);
        });

        let _ = std::panic::take_hook();
        std::panic::set_hook(Box::new(move |info| previous(info)));
        report("a_panicking_worker_releases_its_task", &EXECUTIONS);
    }

    /// Cancellation has to release parked workers as well, and the notification
    /// must not be lost between the re-check and the park.
    #[test]
    fn cancellation_releases_a_parking_worker() {
        static EXECUTIONS: AtomicUsize = AtomicUsize::new(0);
        loom::model(|| {
            EXECUTIONS.fetch_add(1, Ordering::Relaxed);
            let coordinator = Arc::new(Coordinator::new());
            let cancelled = Arc::new(Mutex::new(false));
            // A task stays outstanding, so only cancellation can end the wait.
            coordinator.begin_task();

            let canceller_coordinator = Arc::clone(&coordinator);
            let canceller_flag = Arc::clone(&cancelled);
            let canceller = thread::spawn(move || {
                *canceller_flag
                    .lock()
                    .unwrap_or_else(|poisoned| poisoned.into_inner()) = true;
                canceller_coordinator.wake_waiters();
            });

            let worker_coordinator = Arc::clone(&coordinator);
            let worker_flag = Arc::clone(&cancelled);
            let worker = thread::spawn(move || {
                let stopped = || {
                    *worker_flag
                        .lock()
                        .unwrap_or_else(|poisoned| poisoned.into_inner())
                };
                assert_eq!(
                    worker_coordinator.wait_for_task(stopped, || None::<usize>),
                    None,
                    "a cancelled walk must release its parked workers"
                );
            });

            canceller.join().expect("canceller joins");
            worker.join().expect("worker joins");
            drop(coordinator.claim_task());
        });
        report("cancellation_releases_a_parking_worker", &EXECUTIONS);
    }

    /// Concurrent growth attempts must not put more workers on the walk than
    /// the configured budget allows.
    #[test]
    fn concurrent_growth_stays_within_the_thread_budget() {
        static EXECUTIONS: AtomicUsize = AtomicUsize::new(0);
        loom::model(|| {
            EXECUTIONS.fetch_add(1, Ordering::Relaxed);
            let coordinator = Arc::new(Coordinator::new());
            let caller = coordinator.claim_caller_slot();
            // Backlog of three tasks against a budget of two workers.
            for _ in 0..3 {
                coordinator.begin_task();
            }

            let contender_coordinator = Arc::clone(&coordinator);
            let contender = thread::spawn(move || {
                let slot = contender_coordinator.claim_worker_slot(2);
                assert!(contender_coordinator.active_workers() <= 2);
                drop(slot);
            });

            let slot = coordinator.claim_worker_slot(2);
            assert!(coordinator.active_workers() <= 2);
            drop(slot);
            contender.join().expect("contender joins");

            drop(caller);
            assert_eq!(coordinator.active_workers(), 0);
            for _ in 0..3 {
                drop(coordinator.claim_task());
            }
            assert_eq!(coordinator.pending(), 0);
        });
        report(
            "concurrent_growth_stays_within_the_thread_budget",
            &EXECUTIONS,
        );
    }
}