takeaway 0.1.0

//! The global task queue.

use alloc::boxed::Box;
use core::sync::atomic::{AtomicBool, AtomicUsize, Ordering};

use crate::config::Config;
use crate::control::{AtomicControl, WakerSlot};
use crate::pubqueue::{PQContents, Priority, PubQueue, Stealer};
use crate::task::Task;
use crate::util::CachePadded;

//----------- Queue ------------------------------------------------------------

/// Shared state for a task queue.
///
/// This holds the global state for `takeaway`, and is a shared platform workers
/// use to communicate with each other.  Most importantly, workers will publish
/// information about their tasks here, making them available to steal.
pub struct Queue<T: Task> {
    /// The configuration of the queue.
    pub(crate) config: Config,

    /// The number of started workers.
    ///
    /// Automatic shutdowns (for one-shot mode) are only initiated once all
    /// workers have started (i.e. [`crate::Worker::next()`] has been called for
    /// the first time).
    pub(crate) started: Box<CachePadded<AtomicUsize>>,

    /// The number of pending tasks.
    ///
    /// In one-shot mode, this variable tracks the number of pending tasks
    /// across the entire system.  Under the right conditions, if it reaches
    /// zero, the system has run out of tasks and will shut down.
    pub(crate) pending: Box<CachePadded<AtomicUsize>>,

    /// Whether a shutdown has been initiated.
    pub(crate) shutdown: Box<CachePadded<AtomicBool>>,

    /// The theft-priority of every worker thread.
    ///
    /// If a worker thread has shared a public queue for stealing, this element
    /// is set to a non-zero value.  It can be set to zero by a different thread
    /// to steal the public queue, or the worker thread itself to regain control
    /// of its own public queue.
    pub(crate) priority: Box<[CachePadded<Priority<T>>]>,

    /// The stealer of every worker thread.
    ///
    /// A stealer acts as a bidirectional channel between the owning thread and
    /// a thief thread, holding the public queue in between and allowing either
    /// one to take control of it.
    pub(crate) stealer: Box<[CachePadded<Stealer>]>,

    /// The contents of every public queue.
    ///
    /// There are exactly as many public queues as worker threads, but a worker
    /// thread does not necessarily own the public queue at the same index.
    pub(crate) pq_contents: Box<[Box<PQContents<T>>]>,

    /// The control state of every worker.
    pub(crate) control: Box<[CachePadded<AtomicControl>]>,

    /// Storage for a waker for every worker.
    pub(crate) waker_slot: Box<[CachePadded<WakerSlot>]>,
}

impl<T: Task> Queue<T> {
    /// Construct a new [`Queue`].
    ///
    /// This is equivalent to [`Config::build()`]; pick whichever one is more
    /// idiomatic in your context.
    pub fn new(config: Config) -> Self {
        let num_workers = config.num_workers.get();

        let started = Box::new(CachePadded::new(AtomicUsize::new(0)));

        let pending = Box::new(CachePadded::new(AtomicUsize::new(0)));

        let shutdown = Box::new(CachePadded::new(AtomicBool::new(false)));

        let priority = (0..num_workers)
            .map(|_| CachePadded::new(Priority::default()))
            .collect();

        let stealer = (0..num_workers)
            // SAFETY: 'id' is a valid ID, i.e. 'id < workers'.
            .map(|id| CachePadded::new(unsafe { Stealer::new(id, &config) }))
            .collect();

        let pq_contents = (0..num_workers)
            .map(|_| PQContents::new_boxed(&config))
            .collect();

        let control = (0..num_workers)
            .map(|_| CachePadded::new(AtomicControl::new()))
            .collect();

        let waker_slot = (0..num_workers)
            .map(|_| CachePadded::new(WakerSlot::new()))
            .collect();

        Self {
            config,
            started,
            pending,
            shutdown,
            priority,
            stealer,
            pq_contents,
            control,
            waker_slot,
        }
    }

    /// The configuration of this queue.
    #[inline]
    pub const fn config(&self) -> &Config {
        &self.config
    }

    /// Shut down the task queue.
    ///
    /// By default, [`takeaway`](crate) assumes that there is an infinite stream
    /// of tasks to process.  [`Worker::next()`] blocks (asynchronously) until a
    /// task is available to execute.  However, once [`Queue::shutdown()`] is
    /// called, [`Worker`]s will stop blocking and report that no more tasks are
    /// available.
    ///
    /// [`Worker`]: crate::Worker
    /// [`Worker::next()`]: crate::Worker::next()
    ///
    /// Calling this multiple times has no effect.
    pub fn shutdown(&self) {
        // Mark the state as shut down.
        self.shutdown.store(true, Ordering::Relaxed);

        // Wake up all sleeping workers.
        for (control, waker_slot) in self.control.iter().zip(&self.waker_slot) {
            // SAFETY: 'sleep_state' and 'waker' correspond.
            unsafe { control.shutdown(waker_slot) };
        }
    }

    /// Whether the task queue has been shut down.
    pub fn has_shut_down(&self) -> bool {
        self.shutdown.load(Ordering::Relaxed)
    }

    /// Whether all workers have started.
    pub(crate) fn all_started(&self) -> bool {
        self.started.load(Ordering::Relaxed) == self.config.num_workers.get()
    }

    /// Steal a public queue of high-priority tasks.
    ///
    /// If a public queue containing tasks of a high priority (all above the
    /// given minimum) can be found, it will be stolen and returned.
    ///
    /// ## Safety
    ///
    /// A worker thread `i` can call `steal(min, iq)` if `i` owns `iq`, `iq` is
    /// empty, and `priority[i]` is `None`.
    ///
    /// Additional invariants:
    /// - `replacement < num_workers`
    pub(crate) unsafe fn steal(
        &self,
        min_priority: Option<T::Priority>,
        replacement: usize,
    ) -> Option<PubQueue> {
        debug_assert!(replacement < self.config.num_workers.get());

        // Try stealing from every worker thread in turn.
        // SAFETY: As per the caller, `priority[i]` is `None`.
        let num_workers = self.config.num_workers.get();
        let index = (0..num_workers).find(|&index| unsafe {
            self.priority[index].steal_if_above(min_priority).is_some()
        })?;

        // Steal the public queue from this thread.
        // SAFETY:
        // - As per the caller, `i` owns `replacement`.
        // - `i` is now stealing from `j`.
        Some(unsafe { self.stealer[index].steal(replacement, &self.config) })
    }

    /// Wake up a thread.
    ///
    /// When new tasks are published, this method can be called to wake up any
    /// sleeping threads so they can steal those tasks.
    pub(crate) fn wake(&self) {
        // Look over all the sleep states.
        for (control, waker_slot) in self.control.iter().zip(&self.waker_slot) {
            // SAFETY: 'sleep_state' and 'waker' correspond.
            if unsafe { control.try_wake(waker_slot) } {
                // This worker has been woken up.  Stop.
                break;
            }
        }
    }
}