ax_task/

wait_queue.rs

use alloc::collections::VecDeque;

use crate::{
    AxTaskRef, CurrentTask, current_run_queue, select_wake_run_queue,
    sync::{PreemptIrqSaveState, SpinLock, SpinLockIrqSaveGuard},
};

/// A queue to store sleeping tasks.
///
/// # Examples
///
/// ```
/// use core::sync::atomic::{AtomicU32, Ordering};
///
/// use ax_task::WaitQueue;
///
/// static VALUE: AtomicU32 = AtomicU32::new(0);
/// static WQ: WaitQueue = WaitQueue::new();
///
/// ax_task::init_scheduler();
/// // spawn a new task that updates `VALUE` and notifies the main task
/// ax_task::spawn(|| {
///     assert_eq!(VALUE.load(Ordering::Acquire), 0);
///     VALUE.fetch_add(1, Ordering::Release);
///     WQ.notify_one(true); // wake up the main task
/// });
///
/// WQ.wait(); // block until `notify()` is called
/// assert_eq!(VALUE.load(Ordering::Acquire), 1);
/// ```
pub struct WaitQueue {
    queue: SpinLock<VecDeque<AxTaskRef>>,
}

pub(crate) type WaitQueueGuard<'a> = SpinLockIrqSaveGuard<'a, VecDeque<AxTaskRef>>;

impl Default for WaitQueue {
    fn default() -> Self {
        Self::new()
    }
}

impl WaitQueue {
    /// Creates an empty wait queue.
    pub const fn new() -> Self {
        Self {
            queue: SpinLock::new(VecDeque::new()),
        }
    }

    /// Returns whether this queue currently contains no blocked tasks.
    pub fn is_empty(&self) -> bool {
        self.queue.lock_irqsave().is_empty()
    }

    /// Cancel events by removing the task from the wait queue.
    /// If `from_timer_list` is true, try to remove the task from the timer list.
    fn cancel_events(&self, curr: CurrentTask, _from_timer_list: bool) {
        // A task can be woken by only one event (timer or `notify()`), so remove it from the other queue.
        if curr.in_wait_queue() {
            // wake up by timer (timeout).
            self.queue.lock_irqsave().retain(|t| !curr.ptr_eq(t));
            curr.set_in_wait_queue(false);
        }

        // Try to cancel a timer event from timer lists.
        // Just mark task's current timer ticket ID as expired.
        if _from_timer_list {
            curr.timer_ticket_expired();
            // Note:
            //  this task is still not removed from timer list of target CPU,
            //  which may cause some redundant timer events because it still needs to
            //  go through the process of expiring an event from the timer list and invoking the callback.
            //  (it can be considered a lazy-removal strategy, it will be ignored when it is about to take effect.)
        }
    }

    /// Blocks the current task and put it into the wait queue, until other task
    /// notifies it.
    #[track_caller]
    pub fn wait(&self) {
        crate::api::might_sleep();
        current_run_queue::<PreemptIrqSaveState>().blocked_resched(self.queue.lock_irqsave());
        self.cancel_events(crate::current(), false);
    }

    /// Blocks the current task and put it into the wait queue, until the given
    /// `condition` becomes true.
    ///
    /// Note that even other tasks notify this task, it will not wake up until
    /// the condition becomes true.
    #[track_caller]
    pub fn wait_until<F>(&self, condition: F)
    where
        F: Fn() -> bool,
    {
        crate::api::might_sleep();
        let curr = crate::current();
        loop {
            let mut rq = current_run_queue::<PreemptIrqSaveState>();
            let wq = self.queue.lock_irqsave();
            if condition() {
                break;
            }

            rq.blocked_resched(wq);
            // Preemption may occur here.
        }
        self.cancel_events(curr, false);
    }

    /// Blocks the current task and put it into the wait queue, until other tasks
    /// notify it, or the given duration has elapsed.
    #[track_caller]
    pub fn wait_timeout(&self, dur: core::time::Duration) -> bool {
        crate::api::might_sleep();
        let mut rq = current_run_queue::<PreemptIrqSaveState>();
        let curr = crate::current();
        let deadline = ax_hal::time::monotonic_time() + dur;
        debug!(
            "task wait_timeout: {} deadline={:?}",
            curr.id_name(),
            deadline
        );
        let timeout = loop {
            crate::timers::set_alarm_wakeup(deadline, curr.clone());
            rq.blocked_resched(self.queue.lock_irqsave());

            // Still in the wait queue means the timer path woke us. Re-check
            // the monotonic deadline so an early wake cannot truncate sleeps.
            if !curr.in_wait_queue() {
                break false;
            }
            if ax_hal::time::monotonic_time() >= deadline {
                break true;
            }
        };

        // Always try to remove the task from the timer list.
        self.cancel_events(curr, true);
        timeout
    }

    /// Blocks the current task and put it into the wait queue, until the given
    /// `condition` becomes true, or the given duration has elapsed.
    ///
    /// Note that even other tasks notify this task, it will not wake up until
    /// the above conditions are met.
    #[track_caller]
    pub fn wait_timeout_until<F>(&self, dur: core::time::Duration, condition: F) -> bool
    where
        F: Fn() -> bool,
    {
        crate::api::might_sleep();
        let curr = crate::current();
        let deadline = ax_hal::time::monotonic_time() + dur;
        debug!(
            "task wait_timeout: {}, deadline={:?}",
            curr.id_name(),
            deadline
        );
        let mut timeout = true;
        loop {
            let mut rq = current_run_queue::<PreemptIrqSaveState>();
            if ax_hal::time::monotonic_time() >= deadline {
                break;
            }
            let wq = self.queue.lock_irqsave();
            if condition() {
                timeout = false;
                break;
            }

            crate::timers::set_alarm_wakeup(deadline, curr.clone());
            rq.blocked_resched(wq);
            // Preemption may occur here.
        }
        // Always try to remove the task from the timer list.
        self.cancel_events(curr, true);
        timeout
    }

    /// Wakes up one task in the wait queue, usually the first one.
    /// If `resched` is true, the current task will be preempted when the
    /// preemption is enabled.
    pub fn notify_one(&self, resched: bool) -> bool {
        let task = self.pop_front();
        if let Some(task) = task {
            unblock_one_task(task, resched);
            return true;
        }
        false
    }

    /// Wakes up one task from IRQ context.
    ///
    /// This method is intended for low-level deferred notification paths. It
    /// only unblocks the worker and marks the current task for rescheduling
    /// after IRQ/preemption guards are released; it must not be used as a
    /// substitute for publishing the condition that the waiter will observe.
    pub fn notify_one_from_irq(&self) -> bool {
        self.notify_one(true)
    }

    /// Wakes up one task in the wait queue and runs a callback on it.
    ///
    /// The callback `func` is invoked while holding the wait-queue lock and
    /// before the selected task is unblocked. It receives the task's ID as a
    /// `u64` when a task is available, or `0` if the wait queue is empty.
    /// This can be used for lock handoff or other bookkeeping associated with
    /// the waking task.
    ///
    /// If `resched` is true, the current task will be preempted when the
    /// preemption is enabled.
    pub fn notify_one_with<F>(&self, resched: bool, func: F) -> bool
    where
        F: Fn(u64),
    {
        let task = {
            let mut wq = self.queue.lock_irqsave();
            match wq.pop_front() {
                Some(task) => {
                    func(task.id().as_u64());
                    task.set_in_wait_queue(false);
                    Some(task)
                }
                None => {
                    func(0);
                    None
                }
            }
        };

        if let Some(task) = task {
            unblock_one_task(task, resched);
            return true;
        }
        false
    }

    /// Wakes all tasks in the wait queue.
    ///
    /// If `resched` is true, the current task will yield.
    pub fn notify_all(&self, resched: bool) {
        while self.notify_one(resched) {
            // loop until the wait queue is empty
        }
    }

    /// Wakes all tasks from IRQ context.
    ///
    /// This method is intended for low-level deferred notification paths. It
    /// only unblocks workers and marks the current task for rescheduling after
    /// IRQ/preemption guards are released; it must not be used as a substitute
    /// for publishing the condition that waiters will observe.
    pub fn notify_all_from_irq(&self) {
        while self.notify_one_from_irq() {
            // loop until the wait queue is empty
        }
    }

    fn pop_front(&self) -> Option<AxTaskRef> {
        let mut wq = self.queue.lock_irqsave();
        let task = wq.pop_front()?;
        task.set_in_wait_queue(false);
        Some(task)
    }
}

fn unblock_one_task(task: AxTaskRef, resched: bool) {
    // Select run queue by the CPU set of the task.
    select_wake_run_queue::<PreemptIrqSaveState>(&task).unblock_task(task, resched)
}

#[cfg(test)]
mod coverage_tests {
    use super::*;

    fn wait_queue_new_and_default_hold_for_test() -> bool {
        // Test WaitQueue::new() and Default
        let wq = WaitQueue::new();
        let wq_default = WaitQueue::default();

        // Both should create valid WaitQueue instances
        // We can't easily test the internal state without locking,
        // but we can verify the struct exists and can be created
        let _wq_ref = &wq;
        let _wq_default_ref = &wq_default;

        true
    }

    #[test]
    fn wait_queue_new_and_default_hold() {
        assert!(wait_queue_new_and_default_hold_for_test());
    }
}