takeaway 0.1.4

//! The global task queue.

use alloc::boxed::Box;
use core::ops::ControlFlow;
use core::sync::atomic::{AtomicIsize, 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 control state of the queue.
    ///
    /// This indicates the number of sleeping threads and whether a shutdown
    /// has been initiated.
    pub(crate) queue_control: Box<CachePadded<AtomicQueueControl>>,

    /// 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 queue_control =
            Box::new(CachePadded::new(AtomicQueueControl::new(num_workers)));

        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,
            queue_control,
            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.queue_control.shutdown();

        // 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.try_wake(waker_slot) };
        }
    }

    /// Whether the task queue has been shut down.
    pub fn has_shut_down(&self) -> bool {
        self.queue_control.has_shut_down()
    }

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

//----------- AtomicQueueControl -----------------------------------------------

/// The control state of a [`Queue`].
#[repr(transparent)]
pub struct AtomicQueueControl {
    raw: AtomicIsize,
}

impl AtomicQueueControl {
    /// Construct a new [`AtomicQueueControl`].
    pub const fn new(num_workers: usize) -> Self {
        assert!(num_workers <= isize::MAX as usize);
        let num_workers = num_workers as isize;
        Self {
            raw: AtomicIsize::new(num_workers),
        }
    }

    /// Mark a worker as asleep.
    ///
    /// Returns [`ControlFlow::Break`] if a shutdown has been initiated.
    /// Otherwise returns [`ControlFlow::Continue`] with the number of
    /// still-running threads.
    pub fn mark_asleep(&self) -> ControlFlow<(), usize> {
        match self.raw.fetch_update(
            Ordering::Release,
            Ordering::Relaxed,
            |state| (state >= 0).then(|| state.strict_sub(1)),
        ) {
            Ok(state) => ControlFlow::Continue(state as usize - 1),
            Err(_) => ControlFlow::Break(()),
        }
    }

    /// Mark a worker as awake.
    ///
    /// Returns [`ControlFlow::Break`] if a shutdown has been initiated.
    pub fn mark_awake(&self) -> ControlFlow<()> {
        match self.raw.fetch_update(
            Ordering::Release,
            Ordering::Relaxed,
            |state| (state >= 0).then(|| state.strict_add(1)),
        ) {
            Ok(_) => ControlFlow::Continue(()),
            Err(_) => ControlFlow::Break(()),
        }
    }

    /// Initiate a shutdown.
    pub fn shutdown(&self) {
        self.raw.store(-1, Ordering::Relaxed);
    }

    /// Check whether a shutdown has been initiated.
    pub fn has_shut_down(&self) -> bool {
        self.raw.load(Ordering::Relaxed) < 0
    }
}