axtask/
api.rs

1//! Task APIs for multi-task configuration.
2
3use alloc::{string::String, sync::Arc};
4
5use kernel_guard::NoPreemptIrqSave;
6
7pub(crate) use crate::run_queue::{current_run_queue, select_run_queue};
8
9#[doc(cfg(feature = "multitask"))]
10pub use crate::task::{CurrentTask, TaskId, TaskInner};
11#[doc(cfg(feature = "multitask"))]
12pub use crate::task_ext::{TaskExtMut, TaskExtRef};
13#[doc(cfg(feature = "multitask"))]
14pub use crate::wait_queue::WaitQueue;
15
16/// The reference type of a task.
17pub type AxTaskRef = Arc<AxTask>;
18
19/// The wrapper type for [`cpumask::CpuMask`] with SMP configuration.
20pub type AxCpuMask = cpumask::CpuMask<{ axconfig::plat::MAX_CPU_NUM }>;
21
22cfg_if::cfg_if! {
23    if #[cfg(feature = "sched-rr")] {
24        const MAX_TIME_SLICE: usize = 5;
25        pub(crate) type AxTask = axsched::RRTask<TaskInner, MAX_TIME_SLICE>;
26        pub(crate) type Scheduler = axsched::RRScheduler<TaskInner, MAX_TIME_SLICE>;
27    } else if #[cfg(feature = "sched-cfs")] {
28        pub(crate) type AxTask = axsched::CFSTask<TaskInner>;
29        pub(crate) type Scheduler = axsched::CFScheduler<TaskInner>;
30    } else {
31        // If no scheduler features are set, use FIFO as the default.
32        pub(crate) type AxTask = axsched::FifoTask<TaskInner>;
33        pub(crate) type Scheduler = axsched::FifoScheduler<TaskInner>;
34    }
35}
36
37#[cfg(feature = "preempt")]
38struct KernelGuardIfImpl;
39
40#[cfg(feature = "preempt")]
41#[crate_interface::impl_interface]
42impl kernel_guard::KernelGuardIf for KernelGuardIfImpl {
43    fn disable_preempt() {
44        if let Some(curr) = current_may_uninit() {
45            curr.disable_preempt();
46        }
47    }
48
49    fn enable_preempt() {
50        if let Some(curr) = current_may_uninit() {
51            curr.enable_preempt(true);
52        }
53    }
54}
55
56/// Gets the current task, or returns [`None`] if the current task is not
57/// initialized.
58pub fn current_may_uninit() -> Option<CurrentTask> {
59    CurrentTask::try_get()
60}
61
62/// Gets the current task.
63///
64/// # Panics
65///
66/// Panics if the current task is not initialized.
67pub fn current() -> CurrentTask {
68    CurrentTask::get()
69}
70
71/// Initializes the task scheduler (for the primary CPU).
72pub fn init_scheduler() {
73    info!("Initialize scheduling...");
74
75    // Initialize the cpu count information.
76    init_cpu_mask_full();
77
78    // Initialize the run queue.
79    crate::run_queue::init();
80    #[cfg(feature = "irq")]
81    crate::timers::init();
82
83    info!("  use {} scheduler.", Scheduler::scheduler_name());
84}
85
86/// The full CPU mask of the system.
87static CPU_MASK_FULL: lazyinit::LazyInit<AxCpuMask> = lazyinit::LazyInit::new();
88
89/// Gets the cpu count information and initializes related data structures.
90fn init_cpu_mask_full() {
91    let cpu_num = axhal::cpu_num();
92    let mut cpumask = AxCpuMask::new();
93    for cpu_id in 0..cpu_num {
94        cpumask.set(cpu_id, true);
95    }
96
97    CPU_MASK_FULL.call_once(|| cpumask);
98}
99
100pub(crate) fn cpu_mask_full() -> AxCpuMask {
101    CPU_MASK_FULL
102        .get()
103        .expect("CPU mask not initialized")
104        .clone()
105}
106
107/// Initializes the task scheduler for secondary CPUs.
108pub fn init_scheduler_secondary() {
109    crate::run_queue::init_secondary();
110    #[cfg(feature = "irq")]
111    crate::timers::init();
112}
113
114/// Handles periodic timer ticks for the task manager.
115///
116/// For example, advance scheduler states, checks timed events, etc.
117#[cfg(feature = "irq")]
118#[doc(cfg(feature = "irq"))]
119pub fn on_timer_tick() {
120    use kernel_guard::NoOp;
121    crate::timers::check_events();
122    // Since irq and preemption are both disabled here,
123    // we can get current run queue with the default `kernel_guard::NoOp`.
124    current_run_queue::<NoOp>().scheduler_timer_tick();
125}
126
127/// Adds the given task to the run queue, returns the task reference.
128pub fn spawn_task(task: TaskInner) -> AxTaskRef {
129    let task_ref = task.into_arc();
130    select_run_queue::<NoPreemptIrqSave>(&task_ref).add_task(task_ref.clone());
131    task_ref
132}
133
134/// Spawns a new task with the given parameters.
135///
136/// Returns the task reference.
137pub fn spawn_raw<F>(f: F, name: String, stack_size: usize) -> AxTaskRef
138where
139    F: FnOnce() + Send + 'static,
140{
141    spawn_task(TaskInner::new(f, name, stack_size))
142}
143
144/// Spawns a new task with the default parameters.
145///
146/// The default task name is an empty string. The default task stack size is
147/// [`axconfig::TASK_STACK_SIZE`].
148///
149/// Returns the task reference.
150pub fn spawn<F>(f: F) -> AxTaskRef
151where
152    F: FnOnce() + Send + 'static,
153{
154    spawn_raw(f, "".into(), axconfig::TASK_STACK_SIZE)
155}
156
157/// Set the priority for current task.
158///
159/// The range of the priority is dependent on the underlying scheduler. For
160/// example, in the [CFS] scheduler, the priority is the nice value, ranging from
161/// -20 to 19.
162///
163/// Returns `true` if the priority is set successfully.
164///
165/// [CFS]: https://en.wikipedia.org/wiki/Completely_Fair_Scheduler
166pub fn set_priority(prio: isize) -> bool {
167    current_run_queue::<NoPreemptIrqSave>().set_current_priority(prio)
168}
169
170/// Set the affinity for the current task.
171/// [`AxCpuMask`] is used to specify the CPU affinity.
172/// Returns `true` if the affinity is set successfully.
173///
174/// TODO: support set the affinity for other tasks.
175pub fn set_current_affinity(cpumask: AxCpuMask) -> bool {
176    if cpumask.is_empty() {
177        false
178    } else {
179        let curr = current().clone();
180
181        curr.set_cpumask(cpumask);
182        // After setting the affinity, we need to check if current cpu matches
183        // the affinity. If not, we need to migrate the task to the correct CPU.
184        #[cfg(feature = "smp")]
185        if !cpumask.get(axhal::percpu::this_cpu_id()) {
186            const MIGRATION_TASK_STACK_SIZE: usize = 4096;
187            // Spawn a new migration task for migrating.
188            let migration_task = TaskInner::new(
189                move || crate::run_queue::migrate_entry(curr),
190                "migration-task".into(),
191                MIGRATION_TASK_STACK_SIZE,
192            )
193            .into_arc();
194
195            // Migrate the current task to the correct CPU using the migration task.
196            current_run_queue::<NoPreemptIrqSave>().migrate_current(migration_task);
197
198            assert!(
199                cpumask.get(axhal::percpu::this_cpu_id()),
200                "Migration failed"
201            );
202        }
203        true
204    }
205}
206
207/// Current task gives up the CPU time voluntarily, and switches to another
208/// ready task.
209pub fn yield_now() {
210    current_run_queue::<NoPreemptIrqSave>().yield_current()
211}
212
213/// Current task is going to sleep for the given duration.
214///
215/// If the feature `irq` is not enabled, it uses busy-wait instead.
216pub fn sleep(dur: core::time::Duration) {
217    sleep_until(axhal::time::wall_time() + dur);
218}
219
220/// Current task is going to sleep, it will be woken up at the given deadline.
221///
222/// If the feature `irq` is not enabled, it uses busy-wait instead.
223pub fn sleep_until(deadline: axhal::time::TimeValue) {
224    #[cfg(feature = "irq")]
225    current_run_queue::<NoPreemptIrqSave>().sleep_until(deadline);
226    #[cfg(not(feature = "irq"))]
227    axhal::time::busy_wait_until(deadline);
228}
229
230/// Exits the current task.
231pub fn exit(exit_code: i32) -> ! {
232    current_run_queue::<NoPreemptIrqSave>().exit_current(exit_code)
233}
234
235/// The idle task routine.
236///
237/// It runs an infinite loop that keeps calling [`yield_now()`].
238pub fn run_idle() -> ! {
239    loop {
240        yield_now();
241        debug!("idle task: waiting for IRQs...");
242        #[cfg(feature = "irq")]
243        axhal::asm::wait_for_irqs();
244    }
245}
axtask/api.rs

axtask/
api.rs
