takeaway 0.1.4

//! Workers.

use alloc::{rc::Rc, vec::Vec};
use core::{
    ops::ControlFlow,
    task::{Poll, Waker},
};

use crate::{
    Enqueuer, Queue, Task, TaskPriority, batch::Batch, control::Control,
};

//----------- Worker -----------------------------------------------------------

/// A thread-local view of the task queue.
///
/// This is the primary interface to `takeaway`.  It manages a set of pending
/// tasks, communicating with the global [`Queue`] to distribute tasks across
/// the system, and provides methods for adding and retrieving tasks.
///
/// # Usage
///
/// You can obtain a [`Worker`] from [`Worker::new()`].  To enqueue tasks,
/// use [`Worker::enqueue()`].  To retrieve tasks, use [`Worker::next()`].  In
/// either case, lower-level interfaces are provided for more control.
///
/// ```no_run
/// # use takeaway::{Queue, Worker};
/// #
/// # async fn worker(queue: &Queue<()>, id: usize) {
/// // Construct the 'Worker'.
/// let mut worker = Worker::new(queue, id);
///
/// // Enqueue some initial tasks.
/// worker.enqueue_one(todo!());
///
/// // Process tasks.
/// while let Some(task) = worker.next().await {
///     // Execute the task.
///     //...
///
///     // Enqueue sub-tasks as necessary.
///     worker.enqueue_one(todo!());
/// }
///
/// // The task queue has shut down.
/// # }
/// ```
pub struct Worker<'q, T: Task> {
    /// The ID of this worker.
    id: usize,

    /// The global task queue.
    queue: &'q Queue<T>,

    /// The associated enqueuer.
    enqueuer: Rc<Enqueuer<T>>,

    /// Postponed tasks.
    ///
    /// These tasks are not part of the current batch, and are sorted in
    /// ascending order of priority.
    postponed: Vec<T>,

    /// The current batch of tasks.
    batch: Batch<T>,

    /// The control state of the worker.
    control: Control,
}

impl<'q, T: Task> Worker<'q, T> {
    /// Construct a new [`Worker`].
    ///
    /// The worker will be assigned the specified ID.
    ///
    /// ## Panics
    ///
    /// - Panics if `id >= queue.config().num_workers().get()`.
    /// - Panics if another worker has been assigned this ID already.
    pub fn new(queue: &'q Queue<T>, id: usize) -> Self {
        assert!(id < queue.config.num_workers.get());

        // Lock this ID.
        queue.control[id].initialize();

        Self {
            id,
            queue,
            enqueuer: Rc::new(Enqueuer::new()),
            postponed: Vec::new(),
            // SAFETY: This worker uses 'id' and 'global'.
            batch: unsafe { Batch::new(id, &queue.config) },
            control: Control::Running,
        }
    }

    /// The ID of this worker.
    ///
    /// This is the value assigned to the worker in [`Worker::new()`].
    #[inline]
    pub const fn id(&self) -> usize {
        self.id
    }

    /// The associated global queue.
    #[inline]
    pub const fn queue(&self) -> &'q Queue<T> {
        self.queue
    }

    /// The associated [`Enqueuer`].
    #[inline]
    pub const fn enqueuer(&self) -> &Rc<Enqueuer<T>> {
        &self.enqueuer
    }

    /// Enqueue a set of tasks.
    ///
    /// The tasks will be added to a thread-local [`Vec`], where they will
    /// remain until they are used in the [`Worker`]'s batch.
    ///
    /// This is a shorthand for `self.enqueuer().extend(tasks)`.  If the worker
    /// is borrowed mutably, the [`Enqueuer`] can be used directly by cloning
    /// the [`Rc`] returned by [`Worker::enqueuer()`].
    #[inline]
    pub fn enqueue(&self, tasks: impl IntoIterator<Item = T>) {
        self.enqueuer.extend(tasks)
    }

    /// Enqueue a single task.
    ///
    /// The task will be added to a thread-local [`Vec`], where it will remain
    /// until it is used in the [`Worker`]'s batch.
    ///
    /// This is a shorthand for `self.enqueuer().add(task)`.
    #[inline]
    pub fn enqueue_one(&self, task: T) {
        self.enqueuer.add(task);
    }

    /// Retrieve a single task asynchronously.
    ///
    /// If [`None`] is returned, the task queue has been shut down.
    pub async fn next(&mut self) -> Option<T> {
        core::future::poll_fn(|cx| self.poll(cx.waker())).await
    }

    /// Poll for a single task.
    ///
    /// If [`None`] is returned, the task queue has been shut down.
    #[inline]
    pub fn poll(&mut self, waker: &Waker) -> Poll<Option<T>> {
        // Get a task from the batch, or refresh it.
        if let Some(task) = unsafe { self.batch.next(self.id, self.queue) } {
            debug_assert_eq!(self.control, Control::Running);
            return Poll::Ready(Some(task));
        }

        self.poll_slow(waker)
    }

    #[cold]
    fn poll_slow(&mut self, waker: &Waker) -> Poll<Option<T>> {
        if self.control == Control::Running {
            self.refresh_and_poll(waker)
        } else if self.control == Control::Asleep {
            let control = &self.queue.control[self.id];
            let waker_slot = &self.queue.waker_slot[self.id];

            // First, check for recently enqueued tasks.
            if self.enqueuer.take_waker().is_none() {
                // Mark the worker as awake.
                // SAFETY: The state is currently Asleep.
                unsafe { control.set_self_awake(waker_slot) };
                self.control = Control::Running;
            } else {
                // Poll the control state.
                // SAFETY: The state is currently Asleep.
                self.control =
                    unsafe { control.poll_asleep(waker, waker_slot) };
            }

            // React to the updated state.
            if matches!(self.control, Control::Asleep) {
                self.enqueuer.set_waker(waker.clone());

                Poll::Pending
            } else {
                if self.queue.queue_control.mark_awake().is_break() {
                    self.control = Control::ShuttingDown;
                    return Poll::Ready(None);
                }

                self.poll(waker)
            }
        } else if self.control == Control::Initialized {
            // Mark the worker as running.
            self.control = Control::Running;

            self.poll(waker)
        } else {
            // The queue is shutting down.
            Poll::Ready(None)
        }
    }

    /// Refresh the batch and retrieve the next task to execute.
    fn refresh_and_poll(&mut self, waker: &Waker) -> Poll<Option<T>> {
        if self.queue.queue_control.has_shut_down() {
            return Poll::Ready(None);
        }

        self.collect();

        let task = self.postponed.pop();

        unsafe {
            self.batch.fill(
                task.is_some() as usize,
                &mut self.postponed,
                self.id,
                self.queue,
            )
        };

        match task {
            Some(task) => Poll::Ready(Some(task)),

            None => {
                // Mark the worker as asleep.
                let control = &self.queue.control[self.id];
                let waker_slot = &self.queue.waker_slot[self.id];
                // SAFETY: 'sleep_state' and 'control' correspond.
                unsafe { control.set_self_asleep(waker_slot, waker.clone()) };
                self.enqueuer.set_waker(waker.clone());
                self.control = match self.queue.queue_control.mark_asleep() {
                    ControlFlow::Continue(0) if self.queue.config.oneshot => {
                        self.queue.shutdown();
                        Control::ShuttingDown
                    }
                    ControlFlow::Continue(_) => Control::Asleep,
                    ControlFlow::Break(_) => Control::ShuttingDown,
                };

                match self.control {
                    Control::Asleep => Poll::Pending,
                    Control::ShuttingDown => Poll::Ready(None),
                    _ => unreachable!(),
                }
            }
        }
    }

    /// Collect tasks into the local queue.
    fn collect(&mut self) {
        // Drain the current batch of tasks.
        unsafe { self.batch.drain(&mut self.postponed, self.id, self.queue) };

        // Approximate the highest-available priority thus far.
        //
        // Both the local and public queues are sorted in ascending order of
        // priority, and the last elements of the two have comparable priorities
        // (since they were adjacent when previously sorted), so the last task
        // in the local queue right now, after the appending the public queue,
        // is almost exactly the highest-priority task known.  This ignores any
        // recently-enqueued tasks, but we only need an approximation here.
        let max_priority = self.postponed.last().map(|t| t.priority());

        // Collect recently-enqueued tasks.
        self.enqueuer.drain(&mut self.postponed);

        // Steal higher-priority tasks from the global queue.
        let pq = self.batch.pub_queue.id(&self.queue.config);
        // SAFETY:
        // - We just took ownership of the public queue 'pq'.
        // - 'pq' does not contain any tasks.
        // - We set our public priority to 'None'.
        if let Some(pq) = unsafe { self.queue.steal(max_priority, pq) } {
            // Update the public queue details, even if nobody can see them.
            let stealer = &self.queue.stealer[self.id];
            // SAFETY:
            // - Our public priority is 'None'.
            // - 'stealer' currently contains 'last_pubqueue'.
            unsafe { stealer.set(self.batch.pub_queue, pq) };
            self.batch.pub_queue = pq;

            // Locate the public queue contents.
            let pq = self.batch.pub_queue.id(&self.queue.config);
            let len = self.batch.pub_queue.len(&self.queue.config);
            let pq = &self.queue.pq_contents[pq];

            // Move the tasks out of the public queue.
            // SAFETY:
            // - 'len <= pq_size' as per 'PubQueue::len()'.
            // - We just stole 'pq', and it had 'len' elements at the time of
            //   the theft, and we have not changed it.
            // - This worker owns 'pq' so nobody will write to it.
            unsafe { pq.move_to(len.get(), &mut self.postponed) };
        }

        // Sort the now-filled local queue by priority.
        let batch_size = self.queue.config.batch_size.get();
        if <T::Priority as TaskPriority>::TRIVIAL {
            // Don't sort the queue at all.
        } else if self.postponed.len() > batch_size * 2 {
            // Only sort the used elements of the queue.
            let sort_offset = self.postponed.len() - batch_size * 2;
            self.postponed
                .select_nth_unstable_by_key(sort_offset, |task| {
                    task.priority()
                });
            self.postponed[sort_offset..]
                .sort_unstable_by_key(|task| task.priority());
        } else {
            // The queue is not saturated; sort it completely.
            self.postponed.sort_unstable_by_key(|task| task.priority());
        }

        debug_assert!(self.batch.counter.is_empty());
        debug_assert!(self.batch.local.is_empty());
        debug_assert!(self.batch.pq_priority.is_none());
    }
}

impl<T: Task> Drop for Worker<'_, T> {
    fn drop(&mut self) {
        // Collect tasks from the public queue, as the global 'Queue' does not
        // know whether they exist.
        unsafe { self.batch.drain(&mut self.postponed, self.id, self.queue) };
    }
}

//----------- Tests ------------------------------------------------------------

#[cfg(test)]
mod test {
    use core::{
        convert::Infallible,
        iter,
        num::{NonZeroU8, NonZeroUsize},
        task::{Poll, Waker},
    };

    use crate::{Config, Task, Worker};

    #[test]
    fn new() {
        let workers = 8.try_into().unwrap();
        let batch_size = NonZeroUsize::MIN;
        let queue = Config::new(workers)
            .with_batch_size(batch_size)
            .build::<Infallible>();
        for id in 0..8 {
            let _ = Worker::new(&queue, id);
        }
    }

    #[test]
    #[should_panic]
    fn new_too_many() {
        let workers = 4.try_into().unwrap();
        let batch_size = NonZeroUsize::MIN;
        let queue = Config::new(workers)
            .with_batch_size(batch_size)
            .build::<Infallible>();
        for id in 0..5 {
            let _ = Worker::new(&queue, id);
        }
    }

    #[test]
    fn next_single() {
        struct Foo(u8);

        impl Task for Foo {
            type Priority = ();

            fn priority(&self) -> Self::Priority {}
        }

        let workers = 1.try_into().unwrap();
        let batch_size = NonZeroUsize::MIN;
        let queue = Config::new(workers).with_batch_size(batch_size).build();
        let mut worker = Worker::new(&queue, 0);

        // Enqueue the tasks.
        worker.enqueuer().extend((0..4).map(Foo));

        // Execute the tasks.
        let mut seen = [false; 4];
        let waker = Waker::noop();
        while let Poll::Ready(Some(task)) = worker.poll(waker) {
            let Foo(index) = task;
            assert!(!seen[index as usize]);
            seen[index as usize] = true;
        }

        assert_eq!(seen, [true; 4]);
    }

    #[test]
    fn priority() {
        struct Foo(NonZeroU8);

        impl Task for Foo {
            type Priority = NonZeroU8;

            fn priority(&self) -> Self::Priority {
                self.0
            }
        }

        let workers = 1.try_into().unwrap();
        let batch_size = NonZeroUsize::MIN;
        let queue = Config::new(workers).with_batch_size(batch_size).build();
        let mut worker = Worker::new(&queue, 0);

        // Enqueue the tasks.
        worker
            .enqueuer()
            .extend((1..5).rev().map(|i| Foo(NonZeroU8::new(i).unwrap())));

        // Execute the tasks.
        let waker = Waker::noop();
        assert!(
            iter::from_fn(|| match worker.poll(waker) {
                Poll::Ready(Some(t)) => Some(t),
                _ => None,
            })
            .map(|t| t.0.get())
            .eq((1..5).rev())
        );
    }
}