ferralk 0.6.0

//! 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::{AtomicUsize, Ordering},
};
#[cfg(not(loom))]
use std::sync::{
    Condvar, Mutex, MutexGuard,
    atomic::{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);

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>,
    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);

/// How many parked workers one notification is meant for.
#[derive(Clone, Copy)]
enum Wake {
    One,
    All,
}

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_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);
    }

    /// 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) {
        self.wake_some(Wake::All);
    }

    /// Wakes at most one parked worker, for a producer that made exactly one
    /// task available. A worker that wakes and finds the task already taken by
    /// somebody who never parked simply parks again; the task is not lost,
    /// because whoever took it holds it.
    pub(crate) fn wake_one_waiter(&self) {
        self.wake_some(Wake::One);
    }

    /// The notification itself. Taking `wake_lock` around it is what makes the
    /// hand-off free of lost wakeups: a worker that observed an empty queue
    /// holds the lock until it parks, so it either sees the new state before
    /// parking or is notified afterwards.
    ///
    /// A parked-worker count that let a producer skip the lock as well was
    /// tried in #98 and dropped there. It was sound — producer and worker each
    /// ran a store-then-load with a `SeqCst` fence between the halves, and the
    /// loom models passed — but it multiplied the hand-off model's state space
    /// by fifteen, from 69,363 explored executions to 1,072,953, because three
    /// more atomic accesses inside `wait_for_task` are three more scheduling
    /// points per interleaving. That is a nineteen-fold loom job on every pull
    /// request, and no walk measurement could tell the two versions apart: this
    /// fires once per queued directory, a few thousand times across a 33 ms
    /// walk.
    fn wake_some(&self, wake: Wake) {
        drop(lock(&self.wake_lock));
        match wake {
            Wake::One => self.wake.notify_one(),
            Wake::All => self.wake.notify_all(),
        }
    }

    /// 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;
            }
            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;
            }
            if let Some(task) = try_take() {
                return Some(task);
            }
            let _ = self
                .wake
                .wait_timeout(guard, CANCELLATION_POLL_INTERVAL)
                .unwrap_or_else(|poisoned| poisoned.into_inner());
        }
    }

    /// 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)
    }

    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()));
            // The root task is outstanding before the helper starts.
            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.begin_task();
                push(&producer_queue, 7);
                // One task pushed, so this is the walk's own hand-off path.
                producer_coordinator.wake_one_waiter();
            });

            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 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,
        );
    }
}