ax_task/

api.rs

//! Task APIs for multi-task configuration.

use alloc::{
    collections::BTreeMap,
    string::String,
    sync::{Arc, Weak},
};
use core::fmt;

use ax_memory_addr::VirtAddr;

#[cfg(feature = "lockdep")]
pub use crate::lockdep::{HeldLock, HeldLockStack};
pub(crate) use crate::run_queue::{current_run_queue, select_run_queue, select_wake_run_queue};
use crate::sync::PreemptIrqSaveState;
#[cfg_attr(doc, doc(cfg(feature = "task-ext")))]
#[cfg(feature = "task-ext")]
pub use crate::task::{AxTaskExt, TaskExt};
pub use crate::{
    interrupt::InterruptSnapshot,
    task::{CurrentTask, TaskId, TaskInner, TaskState},
    timers::{
        ClockEventControl, HardKernelTimerAction, HardKernelTimerCallback, KernelTimerAction,
        KernelTimerCallback, KernelTimerCancelOutcome, KernelTimerError, KernelTimerHandle,
        MonotonicDeadline, MonotonicInstant, RestartableKernelTimerCallback, TimerCpuId,
        arm_hard_kernel_timer, cancel_kernel_timer, disarm_hard_kernel_timer, init_timer_service,
        register_hard_restartable_kernel_timer, register_kernel_timer,
        register_restartable_kernel_timer, register_timer_callback,
    },
    wait_queue::WaitQueue,
};

/// The reference type of a task.
pub type AxTaskRef = Arc<AxTask>;

/// The weak reference type of a task.
pub type WeakAxTaskRef = Weak<AxTask>;

static TASK_REGISTRY: ax_lazyinit::LazyLock<crate::sync::SpinRwLock<BTreeMap<u64, WeakAxTaskRef>>> =
    ax_lazyinit::LazyLock::new(|| crate::sync::SpinRwLock::new(BTreeMap::new()));

/// The wrapper type for [`ax_cpumask::CpuMask`] with SMP configuration.
pub type AxCpuMask = ax_cpumask::CpuMask<{ crate::build_info::CPU_CAPACITY }>;

/// Returns the default stack size used by task creation helpers.
pub fn default_task_stack_size() -> usize {
    crate::build_info::DEFAULT_TASK_STACK_SIZE
}

cfg_if::cfg_if! {
    if #[cfg(feature = "sched-rr")] {
        const MAX_TIME_SLICE: usize = 5;
        pub(crate) type AxTask = ax_sched::RRTask<TaskInner, MAX_TIME_SLICE>;
        pub(crate) type Scheduler = ax_sched::RRScheduler<TaskInner, MAX_TIME_SLICE>;
    } else if #[cfg(feature = "sched-cfs")] {
        pub(crate) type AxTask = ax_sched::CFSTask<TaskInner>;
        pub(crate) type Scheduler = ax_sched::CFScheduler<TaskInner>;
    } else {
        // If no scheduler features are set, use FIFO as the default.
        pub(crate) type AxTask = ax_sched::FifoTask<TaskInner>;
        pub(crate) type Scheduler = ax_sched::FifoScheduler<TaskInner>;
    }
}

/// Gets the current task, or returns [`None`] if the current task is not
/// initialized.
pub fn current_may_uninit() -> Option<CurrentTask> {
    CurrentTask::try_get()
}

/// Reports whether the given fault address hits the current task's stack guard page.
#[cfg(feature = "stack-guard-page")]
pub fn diagnose_current_stack_guard_page_fault(fault_addr: VirtAddr) -> bool {
    current_may_uninit().is_some_and(|curr| curr.diagnose_stack_guard_page_fault(fault_addr))
}

/// Gets the current task.
///
/// # Panics
///
/// Panics if the current task is not initialized.
pub fn current() -> CurrentTask {
    CurrentTask::get()
}

/// Disables preemption for the current task when preemption is configured.
#[doc(hidden)]
pub fn disable_preempt() -> usize {
    #[cfg(feature = "preempt")]
    {
        crate::runtime_preempt::enter()
    }
    #[cfg(not(feature = "preempt"))]
    {
        0
    }
}

/// Enables preemption for the current task when preemption is configured.
#[doc(hidden)]
pub fn enable_preempt(token: usize) {
    #[cfg(feature = "preempt")]
    {
        crate::runtime_preempt::exit(token);
    }
    #[cfg(not(feature = "preempt"))]
    let _ = token;
}

/// Enables preemption at the final IRQ-return boundary.
#[doc(hidden)]
pub fn enable_preempt_from_irq_return(token: usize) {
    #[cfg(feature = "preempt")]
    {
        crate::runtime_preempt::exit_from_irq_return(token);
    }
    #[cfg(not(feature = "preempt"))]
    let _ = token;
}

/// Reports scheduler work to the runtime preemption safe-point adapter.
#[doc(hidden)]
pub fn runtime_preemption_pending() -> bool {
    #[cfg(feature = "preempt")]
    {
        current_may_uninit().is_some_and(|curr| curr.preemption_pending())
    }
    #[cfg(not(feature = "preempt"))]
    {
        false
    }
}

/// Runs the legacy scheduler action after the runtime claims its baton.
#[doc(hidden)]
pub fn runtime_preempt_current() {
    #[cfg(feature = "preempt")]
    crate::task::TaskInner::current_check_preempt_pending();
}

#[cfg(feature = "lockdep")]
#[doc(hidden)]
pub fn collect_current_task_held_locks(snapshot: &mut crate::sync::HeldLockSnapshot) {
    let _irq_guard = crate::sync::IrqSaveGuard::new();
    if let Some(curr) = current_may_uninit() {
        curr.with_held_locks(|stack| snapshot.extend(stack));
    }
}

#[cfg(feature = "lockdep")]
#[doc(hidden)]
pub fn push_current_task_held_lock(held: crate::sync::HeldLock) {
    let _irq_guard = crate::sync::IrqSaveGuard::new();
    if let Some(curr) = current_may_uninit() {
        curr.with_held_locks(|stack| stack.push(held));
    }
}

#[cfg(feature = "lockdep")]
#[doc(hidden)]
pub fn pop_current_task_held_lock(lock_addr: usize) {
    let _irq_guard = crate::sync::IrqSaveGuard::new();
    if let Some(curr) = current_may_uninit() {
        curr.with_held_locks(|stack| stack.pop_checked(lock_addr));
    }
}

#[cfg(feature = "lockdep")]
pub fn with_current_lockdep_stack<R>(f: impl FnOnce(&mut HeldLockStack) -> R) -> R {
    current().with_held_locks(f)
}

/// Initializes the task scheduler (for the primary CPU).
pub fn init_scheduler() {
    info!("Initialize scheduling...");

    #[cfg(feature = "host-test")]
    ax_hal::percpu::initialize_host_test_cpu();

    // Initialize the run queue.
    crate::run_queue::init();

    info!("  use {} scheduler.", Scheduler::scheduler_name());
}

pub(crate) fn cpu_mask_full() -> AxCpuMask {
    use ax_lazyinit::LazyLock;

    static CPU_MASK_FULL: LazyLock<AxCpuMask> = LazyLock::new(|| {
        let cpu_num = ax_hal::cpu_num();
        let mut cpumask = AxCpuMask::new();
        for cpu_id in 0..cpu_num {
            cpumask.set(cpu_id, true);
        }
        cpumask
    });

    *CPU_MASK_FULL
}

/// Initializes the task scheduler for secondary CPUs.
pub fn init_scheduler_secondary(stack_ptr: VirtAddr, stack_size: usize) {
    crate::run_queue::init_secondary(stack_ptr, stack_size);
}

/// Handles periodic timer ticks for the task manager.
///
/// For example, advance scheduler states, checks timed events, etc.
pub fn on_timer_tick() {
    on_timer_irq(true);
}

/// Handles a hardware timer interrupt.
pub fn on_timer_irq(scheduler_tick: bool) {
    crate::timers::check_events(scheduler_tick);
    if scheduler_tick {
        // Since irq and preemption are both disabled here,
        // we can get the current run queue without another context transition.
        current_run_queue::<crate::sync::RawState>().scheduler_timer_tick();
    }
}

#[doc(hidden)]
pub fn next_timer_deadline_nanos() -> Option<u64> {
    crate::timers::next_deadline_nanos()
}

/// Scheduler ticks CPU `cpu` has spent running a non-idle task since boot.
///
/// This is the load metric for an ondemand cpufreq governor: a monotonic per-CPU
/// counter bumped once per timer tick when the CPU is not idle. Sample the delta
/// over a window and divide by the elapsed ticks to get the busy fraction. Returns
/// 0 for an out-of-range `cpu`. The counter only advances inside the timer tick.
pub fn cpu_busy_ticks(cpu: usize) -> u64 {
    crate::run_queue::BUSY_TICKS
        .get(cpu)
        .map_or(0, |t| t.load(core::sync::atomic::Ordering::Relaxed))
}

/// Adds the given task to the run queue, returns the task reference.
pub fn spawn_task(task: TaskInner) -> AxTaskRef {
    spawn_task_with(task, |_| {})
}

/// Initializes the given task before adding it to the run queue.
///
/// The `initialize` callback receives the stable task reference before the task
/// becomes runnable. Use it to publish runtime-specific task metadata that the
/// task must be able to observe on its first instruction. The callback must not
/// wait for the new task to run because it has not been registered or queued.
///
/// # Panics
///
/// Panics if `initialize` panics.
pub fn spawn_task_with<F>(task: TaskInner, initialize: F) -> AxTaskRef
where
    F: FnOnce(&AxTaskRef),
{
    let task_ref = task.into_arc();
    initialize_task_before_schedule(&task_ref, initialize, |task_ref| {
        register_task(task_ref);
        select_run_queue::<PreemptIrqSaveState>(task_ref).add_task(task_ref.clone());
    });
    task_ref
}

fn initialize_task_before_schedule<T>(
    task: &T,
    initialize: impl FnOnce(&T),
    schedule: impl FnOnce(&T),
) {
    initialize(task);
    schedule(task);
}

/// Spawns a new task with the given parameters.
///
/// Returns the task reference.
pub fn spawn_raw<F>(f: F, name: String, stack_size: usize) -> AxTaskRef
where
    F: FnOnce() + Send + 'static,
{
    spawn_task(TaskInner::new(f, name, stack_size))
}

/// Spawns a new task with the given name and the default stack size.
///
/// Returns the task reference.
pub fn spawn_with_name<F>(f: F, name: String) -> AxTaskRef
where
    F: FnOnce() + Send + 'static,
{
    spawn_raw(f, name, default_task_stack_size())
}

/// Spawns a new task with the default parameters.
///
/// The default task name is an empty string. The default task stack size is
/// [`default_task_stack_size`].
///
/// Returns the task reference.
pub fn spawn<F>(f: F) -> AxTaskRef
where
    F: FnOnce() + Send + 'static,
{
    spawn_with_name(f, String::new())
}

/// Set the priority for current task.
///
/// The range of the priority is dependent on the underlying scheduler. For
/// example, in the [CFS] scheduler, the priority is the nice value, ranging from
/// -20 to 19.
///
/// Returns `true` if the priority is set successfully.
///
/// [CFS]: https://en.wikipedia.org/wiki/Completely_Fair_Scheduler
pub fn set_priority(prio: isize) -> bool {
    current_run_queue::<PreemptIrqSaveState>().set_current_priority(prio)
}

/// Set the affinity for the current task.
/// [`AxCpuMask`] is used to specify the CPU affinity.
/// Returns `true` if the affinity is set successfully.
///
/// TODO: support set the affinity for other tasks.
#[track_caller]
pub fn set_current_affinity(cpumask: AxCpuMask) -> bool {
    might_sleep();

    if cpumask.is_empty() {
        false
    } else {
        let curr = current().clone();

        curr.set_cpumask(cpumask);
        // After setting the affinity, we need to check if current cpu matches
        // the affinity. If not, we need to migrate the task to the correct CPU.
        #[cfg(feature = "smp")]
        if !cpumask.get(ax_hal::percpu::this_cpu_id()) {
            // Spawn a new migration task for migrating.
            let migration_task = TaskInner::new(
                move || crate::run_queue::migrate_entry(curr),
                "migration-task".into(),
                default_task_stack_size(),
            )
            .into_arc();

            // Migrate the current task to the correct CPU using the migration task.
            current_run_queue::<PreemptIrqSaveState>().migrate_current(migration_task);
        }
        true
    }
}

/// Current task gives up the CPU time voluntarily, and switches to another
/// ready task.
#[track_caller]
pub fn yield_now() {
    might_sleep();

    yield_now_unchecked();
}

/// Gives up the CPU from a kernel-internal path.
///
/// This bypasses the public `might_sleep()` guard and is intended only for
/// carefully reviewed scheduler or syscall paths that must yield while running
/// under internal kernel guards.
#[doc(hidden)]
pub(crate) fn yield_now_unchecked() {
    current_run_queue::<PreemptIrqSaveState>().yield_current()
}

/// Current task is going to sleep for the given duration.
#[track_caller]
pub fn sleep(dur: core::time::Duration) {
    sleep_until(ax_hal::time::monotonic_time() + dur);
}

/// Current task is going to sleep, it will be woken up at the given deadline.
/// The deadline is measured against the monotonic clock.
#[track_caller]
pub fn sleep_until(deadline: ax_hal::time::TimeValue) {
    might_sleep();
    current_run_queue::<PreemptIrqSaveState>().sleep_until(deadline);
}

/// Exits the current task.
#[track_caller]
pub fn exit(exit_code: i32) -> ! {
    might_sleep();

    current_run_queue::<PreemptIrqSaveState>().exit_current(exit_code)
}

fn current_irq_context() -> bool {
    #[cfg(not(feature = "host-test"))]
    {
        ax_hal::irq::in_irq_context()
    }
    #[cfg(feature = "host-test")]
    {
        false
    }
}

fn current_cpu_id() -> usize {
    #[cfg(not(feature = "host-test"))]
    {
        ax_hal::percpu::this_cpu_id()
    }
    #[cfg(feature = "host-test")]
    {
        0
    }
}

#[derive(Clone, Copy)]
struct AtomicContextReasons {
    irq_disabled: bool,
    irq_context: bool,
    preempt_disabled: bool,
}

impl AtomicContextReasons {
    const fn is_atomic(self) -> bool {
        self.irq_disabled || self.irq_context || self.preempt_disabled
    }
}

impl fmt::Display for AtomicContextReasons {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let mut wrote_any = false;
        f.write_str("[")?;
        if self.irq_disabled {
            f.write_str("irq_disabled")?;
            wrote_any = true;
        }
        if self.irq_context {
            if wrote_any {
                f.write_str(",")?;
            }
            f.write_str("irq_context")?;
            wrote_any = true;
        }
        if self.preempt_disabled {
            if wrote_any {
                f.write_str(",")?;
            }
            f.write_str("preempt_disabled")?;
            wrote_any = true;
        }
        if !wrote_any {
            f.write_str("none")?;
        }
        f.write_str("]")
    }
}

#[derive(Clone, Copy)]
struct AtomicContextSnapshot {
    irq_enabled: bool,
    irq_context: bool,
    preempt_count: usize,
    cpu_id: usize,
    task_id: Option<u64>,
    task_state: Option<TaskState>,
}

impl AtomicContextSnapshot {
    fn capture() -> Self {
        let current = current_may_uninit();
        let preempt_count = {
            #[cfg(feature = "preempt")]
            {
                let task_depth = current.as_ref().map_or(0, |curr| curr.preempt_count());
                #[cfg(not(feature = "host-test"))]
                {
                    task_depth
                }
                #[cfg(feature = "host-test")]
                {
                    task_depth + crate::sync::host_preempt_depth()
                }
            }
            #[cfg(not(feature = "preempt"))]
            {
                0
            }
        };

        Self {
            irq_enabled: ax_hal::asm::irqs_enabled(),
            irq_context: current_irq_context(),
            preempt_count,
            cpu_id: current_cpu_id(),
            task_id: current.as_ref().map(|curr| curr.id().as_u64()),
            task_state: current.as_ref().map(|curr| curr.state()),
        }
    }

    fn reasons(self) -> AtomicContextReasons {
        let irq_disabled = !self.irq_enabled;

        AtomicContextReasons {
            irq_disabled,
            irq_context: self.irq_context,
            preempt_disabled: self.preempt_count != 0,
        }
    }

    fn is_atomic(self) -> bool {
        self.reasons().is_atomic()
    }
}

/// Returns whether the current context is atomic, meaning sleeping or
/// rescheduling is not allowed.
///
/// This matches the intent of Linux's `might_sleep()`: catch misuse from
/// IRQ-disabled or preempt-disabled regions before a sleep-like action happens.
pub fn in_atomic_context() -> bool {
    AtomicContextSnapshot::capture().is_atomic()
}

/// Marks an operation as one that may sleep or reschedule.
///
/// Panics if it is executed in an atomic context.
#[track_caller]
pub fn might_sleep() {
    might_sleep_at(core::panic::Location::caller());
}

/// Checks a sleep-like operation and attributes failures to `caller`.
///
/// Runtime capability adapters use this entry point because their generated
/// cross-crate shim cannot preserve Rust's implicit `#[track_caller]` argument.
#[doc(hidden)]
pub fn might_sleep_at(caller: &'static core::panic::Location<'static>) {
    let snapshot = AtomicContextSnapshot::capture();
    if snapshot.is_atomic() {
        panic_atomic_sleep(snapshot, caller);
    }
}

#[cfg(not(feature = "lockdep"))]
fn panic_atomic_sleep(
    snapshot: AtomicContextSnapshot,
    caller: &'static core::panic::Location<'static>,
) -> ! {
    panic!(
        "sleeping or rescheduling is not allowed in atomic context: caller={}, reasons={}, \
         irq_enabled={}, irq_context={}, preempt_count={}, cpu_id={}, task_id={:?}, \
         task_state={:?}",
        caller,
        snapshot.reasons(),
        snapshot.irq_enabled,
        snapshot.irq_context,
        snapshot.preempt_count,
        snapshot.cpu_id,
        snapshot.task_id,
        snapshot.task_state
    );
}

#[cfg(feature = "lockdep")]
fn panic_atomic_sleep(
    snapshot: AtomicContextSnapshot,
    caller: &'static core::panic::Location<'static>,
) -> ! {
    let held_locks = crate::sync::current_task_held_lock_snapshot();
    panic!(
        "sleeping or rescheduling is not allowed in atomic context: caller={}, reasons={}, \
         irq_enabled={}, irq_context={}, preempt_count={}, cpu_id={}, task_id={:?}, \
         task_state={:?}, held_locks={}",
        caller,
        snapshot.reasons(),
        snapshot.irq_enabled,
        snapshot.irq_context,
        snapshot.preempt_count,
        snapshot.cpu_id,
        snapshot.task_id,
        snapshot.task_state,
        held_locks
    );
}

/// Wakes a task that may be sleeping, ensuring it can observe a newly-
/// delivered signal.
///
/// `TaskInner::interrupt()` sets the task's interrupt flag and fires the
/// interrupt waker, which unblocks the task via `AxWaker::wake_by_ref`. This
/// covers the common case where the task is blocked in `block_on` with
/// `interruptible` wrapping. For tasks blocked on raw `WaitQueue` objects
/// (which do not register an interrupt waker), this function provides an
/// escape hatch by additionally force-unblocking when the task appears to
/// be parked on a wait queue.
pub fn wake_task(task: &AxTaskRef) {
    // Fire the interrupt: sets the flag and wakes the interrupt_waker.
    // For tasks in block_on (the common case), AxWaker::wake_by_ref already
    // unblocks the task via the registered waker callback.
    task.interrupt();

    // For tasks blocked on a raw WaitQueue, interrupt_waker.wake() is a
    // no-op (no waker registered). Force-unblock by transitioning the task
    // from Blocked to Ready and placing it on the run queue of its
    // affinity CPU.
    //
    // SAFETY: unblock_task uses a CAS on the task state (Blocked → Ready),
    // so if the task is concurrently being woken by its WaitQueue, the CAS
    // fails and this is a harmless no-op. The stale entry in the WaitQueue
    // is benign: when WaitQueue::notify_one eventually pops it, the
    // subsequent unblock_task call will again CAS-fail (task already Ready
    // or Running).
    if task.state() == TaskState::Blocked {
        let mut rq = select_run_queue::<PreemptIrqSaveState>(task);
        rq.unblock_task(task.clone(), false);
    }
}

/// Registers a task for lookup by its scheduler task id.
///
/// This keeps a weak reference only; expired entries are ignored by lookup.
pub fn register_task(task: &AxTaskRef) {
    TASK_REGISTRY
        .write()
        .insert(task.id().as_u64(), Arc::downgrade(task));
}

/// Finds a task by its scheduler task id.
pub fn task_by_id(task_id: TaskId) -> Option<AxTaskRef> {
    TASK_REGISTRY
        .read()
        .get(&task_id.as_u64())
        .and_then(|task| task.upgrade())
}

/// Wakes a task by its scheduler task id.
pub fn wake_task_by_id(task_id: TaskId) -> bool {
    let Some(task) = task_by_id(task_id) else {
        return false;
    };
    wake_task(&task);
    true
}

/// The idle task routine.
///
/// It runs an infinite loop that keeps trying to hand over the CPU before
/// waiting for the next interrupt.
pub fn run_idle() -> ! {
    loop {
        yield_now_unchecked();
        trace!("idle task: waiting for IRQs...");
        #[cfg(not(feature = "host-test"))]
        ax_hal::asm::wait_for_irqs();
    }
}

#[cfg(test)]
mod tests {
    use core::cell::Cell;

    #[test]
    #[cfg(feature = "host-test")]
    fn host_atomic_context_query_does_not_require_cpu_local_state() {
        assert!(!super::in_atomic_context());
    }

    #[test]
    fn task_initialization_precedes_scheduling() {
        let initialized = Cell::new(false);

        super::initialize_task_before_schedule(
            &(),
            |_| initialized.set(true),
            |_| {
                assert!(
                    initialized.get(),
                    "task was scheduled before initialization"
                )
            },
        );

        assert!(initialized.get());
    }
}

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

    fn axtask_api_constants_hold_for_test() -> bool {
        // default_task_stack_size should return a non-zero value
        let stack_size = default_task_stack_size();
        assert!(stack_size > 0);
        assert!(stack_size.is_multiple_of(4096)); // Should be page-aligned

        true
    }

    fn axtask_api_type_aliases_hold_for_test() -> bool {
        // Test that type aliases exist and are usable
        // AxTaskRef = Arc<AxTask>
        // WeakAxTaskRef = Weak<AxTask>
        let _type_check: Option<super::AxTaskRef> = None;
        let _weak_check: Option<super::WeakAxTaskRef> = None;

        true
    }

    fn axtask_api_scheduler_name_hold_for_test() -> bool {
        // Test that Scheduler::scheduler_name() returns a non-empty string
        let name = super::Scheduler::scheduler_name();
        assert!(!name.is_empty());

        true
    }

    fn axtask_api_task_registry_functions_exist_hold_for_test() -> bool {
        let _lookup: fn(super::TaskId) -> Option<super::AxTaskRef> = super::task_by_id;
        let _wake: fn(super::TaskId) -> bool = super::wake_task_by_id;

        true
    }

    #[test]
    fn axtask_api_constants_hold() {
        assert!(axtask_api_constants_hold_for_test());
    }

    #[test]
    fn axtask_api_type_aliases_hold() {
        assert!(axtask_api_type_aliases_hold_for_test());
    }

    #[test]
    fn axtask_api_scheduler_name_hold() {
        assert!(axtask_api_scheduler_name_hold_for_test());
    }

    #[test]
    fn axtask_api_task_registry_functions_exist_hold() {
        assert!(axtask_api_task_registry_functions_exist_hold_for_test());
    }
}