1use alloc::{boxed::Box, string::String, sync::Arc};
2use core::ops::Deref;
3use core::sync::atomic::{AtomicBool, AtomicI32, AtomicU8, AtomicU64, Ordering};
4use core::{alloc::Layout, cell::UnsafeCell, fmt, ptr::NonNull};
5
6#[cfg(feature = "preempt")]
7use core::sync::atomic::AtomicUsize;
8
9use kspin::SpinNoIrq;
10use memory_addr::{VirtAddr, align_up_4k};
11
12use axhal::arch::TaskContext;
13#[cfg(feature = "tls")]
14use axhal::tls::TlsArea;
15
16use crate::task_ext::AxTaskExt;
17use crate::{AxCpuMask, AxTask, AxTaskRef, WaitQueue};
18
19#[derive(Debug, Clone, Copy, Eq, PartialEq)]
21pub struct TaskId(u64);
22
23#[repr(u8)]
25#[derive(Debug, Clone, Copy, Eq, PartialEq)]
26pub(crate) enum TaskState {
27 Running = 1,
29 Ready = 2,
31 Blocked = 3,
34 Exited = 4,
36}
37
38pub struct TaskInner {
40 id: TaskId,
41 name: String,
42 is_idle: bool,
43 is_init: bool,
44
45 entry: Option<*mut dyn FnOnce()>,
46 state: AtomicU8,
47
48 cpumask: SpinNoIrq<AxCpuMask>,
50
51 in_wait_queue: AtomicBool,
53
54 #[cfg(feature = "smp")]
56 on_cpu: AtomicBool,
57
58 #[cfg(feature = "irq")]
62 timer_ticket_id: AtomicU64,
63
64 #[cfg(feature = "preempt")]
65 need_resched: AtomicBool,
66 #[cfg(feature = "preempt")]
67 preempt_disable_count: AtomicUsize,
68
69 exit_code: AtomicI32,
70 wait_for_exit: WaitQueue,
71
72 kstack: Option<TaskStack>,
73 ctx: UnsafeCell<TaskContext>,
74 task_ext: AxTaskExt,
75
76 #[cfg(feature = "tls")]
77 tls: TlsArea,
78}
79
80impl TaskId {
81 fn new() -> Self {
82 static ID_COUNTER: AtomicU64 = AtomicU64::new(1);
83 Self(ID_COUNTER.fetch_add(1, Ordering::Relaxed))
84 }
85
86 pub const fn as_u64(&self) -> u64 {
88 self.0
89 }
90}
91
92impl From<u8> for TaskState {
93 #[inline]
94 fn from(state: u8) -> Self {
95 match state {
96 1 => Self::Running,
97 2 => Self::Ready,
98 3 => Self::Blocked,
99 4 => Self::Exited,
100 _ => unreachable!(),
101 }
102 }
103}
104
105unsafe impl Send for TaskInner {}
106unsafe impl Sync for TaskInner {}
107
108impl TaskInner {
109 pub fn new<F>(entry: F, name: String, stack_size: usize) -> Self
111 where
112 F: FnOnce() + Send + 'static,
113 {
114 let mut t = Self::new_common(TaskId::new(), name);
115 debug!("new task: {}", t.id_name());
116 let kstack = TaskStack::alloc(align_up_4k(stack_size));
117
118 #[cfg(feature = "tls")]
119 let tls = VirtAddr::from(t.tls.tls_ptr() as usize);
120 #[cfg(not(feature = "tls"))]
121 let tls = VirtAddr::from(0);
122
123 t.entry = Some(Box::into_raw(Box::new(entry)));
124 t.ctx_mut().init(task_entry as usize, kstack.top(), tls);
125 t.kstack = Some(kstack);
126 if t.name == "idle" {
127 t.is_idle = true;
128 }
129 t
130 }
131
132 pub const fn id(&self) -> TaskId {
134 self.id
135 }
136
137 pub fn name(&self) -> &str {
139 self.name.as_str()
140 }
141
142 pub fn id_name(&self) -> alloc::string::String {
144 alloc::format!("Task({}, {:?})", self.id.as_u64(), self.name)
145 }
146
147 pub fn join(&self) -> Option<i32> {
151 self.wait_for_exit
152 .wait_until(|| self.state() == TaskState::Exited);
153 Some(self.exit_code.load(Ordering::Acquire))
154 }
155
156 pub unsafe fn task_ext_ptr(&self) -> *mut u8 {
166 self.task_ext.as_ptr()
167 }
168
169 pub fn init_task_ext<T: Sized>(&mut self, data: T) -> Option<&T> {
174 if self.task_ext.is_empty() {
175 self.task_ext.write(data).map(|data| &*data)
176 } else {
177 None
178 }
179 }
180
181 #[inline]
183 pub const fn ctx_mut(&mut self) -> &mut TaskContext {
184 self.ctx.get_mut()
185 }
186
187 #[inline]
189 pub const fn kernel_stack_top(&self) -> Option<VirtAddr> {
190 match &self.kstack {
191 Some(s) => Some(s.top()),
192 None => None,
193 }
194 }
195
196 #[inline]
200 pub fn cpumask(&self) -> AxCpuMask {
201 *self.cpumask.lock()
202 }
203
204 #[inline]
209 pub fn set_cpumask(&self, cpumask: AxCpuMask) {
210 *self.cpumask.lock() = cpumask
211 }
212}
213
214impl TaskInner {
216 fn new_common(id: TaskId, name: String) -> Self {
217 Self {
218 id,
219 name,
220 is_idle: false,
221 is_init: false,
222 entry: None,
223 state: AtomicU8::new(TaskState::Ready as u8),
224 cpumask: SpinNoIrq::new(AxCpuMask::full()),
226 in_wait_queue: AtomicBool::new(false),
227 #[cfg(feature = "irq")]
228 timer_ticket_id: AtomicU64::new(0),
229 #[cfg(feature = "smp")]
230 on_cpu: AtomicBool::new(false),
231 #[cfg(feature = "preempt")]
232 need_resched: AtomicBool::new(false),
233 #[cfg(feature = "preempt")]
234 preempt_disable_count: AtomicUsize::new(0),
235 exit_code: AtomicI32::new(0),
236 wait_for_exit: WaitQueue::new(),
237 kstack: None,
238 ctx: UnsafeCell::new(TaskContext::new()),
239 task_ext: AxTaskExt::empty(),
240 #[cfg(feature = "tls")]
241 tls: TlsArea::alloc(),
242 }
243 }
244
245 pub(crate) fn new_init(name: String) -> Self {
254 let mut t = Self::new_common(TaskId::new(), name);
255 t.is_init = true;
256 #[cfg(feature = "smp")]
257 t.set_on_cpu(true);
258 if t.name == "idle" {
259 t.is_idle = true;
260 }
261 t
262 }
263
264 pub(crate) fn into_arc(self) -> AxTaskRef {
265 Arc::new(AxTask::new(self))
266 }
267
268 #[inline]
269 pub(crate) fn state(&self) -> TaskState {
270 self.state.load(Ordering::Acquire).into()
271 }
272
273 #[inline]
274 pub(crate) fn set_state(&self, state: TaskState) {
275 self.state.store(state as u8, Ordering::Release)
276 }
277
278 #[inline]
282 pub(crate) fn transition_state(&self, current_state: TaskState, new_state: TaskState) -> bool {
283 self.state
284 .compare_exchange(
285 current_state as u8,
286 new_state as u8,
287 Ordering::AcqRel,
288 Ordering::Acquire,
289 )
290 .is_ok()
291 }
292
293 #[inline]
294 pub(crate) fn is_running(&self) -> bool {
295 matches!(self.state(), TaskState::Running)
296 }
297
298 #[inline]
299 pub(crate) fn is_ready(&self) -> bool {
300 matches!(self.state(), TaskState::Ready)
301 }
302
303 #[inline]
304 pub(crate) const fn is_init(&self) -> bool {
305 self.is_init
306 }
307
308 #[inline]
309 pub(crate) const fn is_idle(&self) -> bool {
310 self.is_idle
311 }
312
313 #[inline]
314 pub(crate) fn in_wait_queue(&self) -> bool {
315 self.in_wait_queue.load(Ordering::Acquire)
316 }
317
318 #[inline]
319 pub(crate) fn set_in_wait_queue(&self, in_wait_queue: bool) {
320 self.in_wait_queue.store(in_wait_queue, Ordering::Release);
321 }
322
323 #[inline]
325 #[cfg(feature = "irq")]
326 pub(crate) fn timer_ticket(&self) -> u64 {
327 self.timer_ticket_id.load(Ordering::Acquire)
328 }
329
330 #[inline]
332 #[cfg(feature = "irq")]
333 pub(crate) fn set_timer_ticket(&self, timer_ticket_id: u64) {
334 assert!(timer_ticket_id != 0);
337 self.timer_ticket_id
338 .store(timer_ticket_id, Ordering::Release);
339 }
340
341 #[inline]
344 #[cfg(feature = "irq")]
345 pub(crate) fn timer_ticket_expired(&self) {
346 self.timer_ticket_id.store(0, Ordering::Release);
347 }
348
349 #[inline]
350 #[cfg(feature = "preempt")]
351 pub(crate) fn set_preempt_pending(&self, pending: bool) {
352 self.need_resched.store(pending, Ordering::Release)
353 }
354
355 #[inline]
356 #[cfg(feature = "preempt")]
357 pub(crate) fn can_preempt(&self, current_disable_count: usize) -> bool {
358 self.preempt_disable_count.load(Ordering::Acquire) == current_disable_count
359 }
360
361 #[inline]
362 #[cfg(feature = "preempt")]
363 pub(crate) fn disable_preempt(&self) {
364 self.preempt_disable_count.fetch_add(1, Ordering::Relaxed);
365 }
366
367 #[inline]
368 #[cfg(feature = "preempt")]
369 pub(crate) fn enable_preempt(&self, resched: bool) {
370 if self.preempt_disable_count.fetch_sub(1, Ordering::Relaxed) == 1 && resched {
371 Self::current_check_preempt_pending();
373 }
374 }
375
376 #[cfg(feature = "preempt")]
377 fn current_check_preempt_pending() {
378 use kernel_guard::NoPreemptIrqSave;
379 let curr = crate::current();
380 if curr.need_resched.load(Ordering::Acquire) && curr.can_preempt(0) {
381 let mut rq = crate::current_run_queue::<NoPreemptIrqSave>();
384 if curr.need_resched.load(Ordering::Acquire) {
385 rq.preempt_resched()
386 }
387 }
388 }
389
390 pub(crate) fn notify_exit(&self, exit_code: i32) {
392 self.exit_code.store(exit_code, Ordering::Release);
393 self.wait_for_exit.notify_all(false);
394 }
395
396 #[inline]
397 pub(crate) const unsafe fn ctx_mut_ptr(&self) -> *mut TaskContext {
398 self.ctx.get()
399 }
400
401 #[cfg(feature = "smp")]
408 #[inline]
409 pub(crate) fn on_cpu(&self) -> bool {
410 self.on_cpu.load(Ordering::Acquire)
411 }
412
413 #[cfg(feature = "smp")]
415 #[inline]
416 pub(crate) fn set_on_cpu(&self, on_cpu: bool) {
417 self.on_cpu.store(on_cpu, Ordering::Release)
418 }
419}
420
421impl fmt::Debug for TaskInner {
422 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
423 f.debug_struct("TaskInner")
424 .field("id", &self.id)
425 .field("name", &self.name)
426 .field("state", &self.state())
427 .finish()
428 }
429}
430
431impl Drop for TaskInner {
432 fn drop(&mut self) {
433 debug!("task drop: {}", self.id_name());
434 }
435}
436
437struct TaskStack {
438 ptr: NonNull<u8>,
439 layout: Layout,
440}
441
442impl TaskStack {
443 pub fn alloc(size: usize) -> Self {
444 let layout = Layout::from_size_align(size, 16).unwrap();
445 Self {
446 ptr: NonNull::new(unsafe { alloc::alloc::alloc(layout) }).unwrap(),
447 layout,
448 }
449 }
450
451 pub const fn top(&self) -> VirtAddr {
452 unsafe { core::mem::transmute(self.ptr.as_ptr().add(self.layout.size())) }
453 }
454}
455
456impl Drop for TaskStack {
457 fn drop(&mut self) {
458 unsafe { alloc::alloc::dealloc(self.ptr.as_ptr(), self.layout) }
459 }
460}
461
462use core::mem::ManuallyDrop;
463
464pub struct CurrentTask(ManuallyDrop<AxTaskRef>);
468
469impl CurrentTask {
470 pub(crate) fn try_get() -> Option<Self> {
471 let ptr: *const super::AxTask = axhal::cpu::current_task_ptr();
472 if !ptr.is_null() {
473 Some(Self(unsafe { ManuallyDrop::new(AxTaskRef::from_raw(ptr)) }))
474 } else {
475 None
476 }
477 }
478
479 pub(crate) fn get() -> Self {
480 Self::try_get().expect("current task is uninitialized")
481 }
482
483 pub fn as_task_ref(&self) -> &AxTaskRef {
485 &self.0
486 }
487
488 pub(crate) fn clone(&self) -> AxTaskRef {
489 self.0.deref().clone()
490 }
491
492 pub(crate) fn ptr_eq(&self, other: &AxTaskRef) -> bool {
493 Arc::ptr_eq(&self.0, other)
494 }
495
496 pub(crate) unsafe fn init_current(init_task: AxTaskRef) {
497 assert!(init_task.is_init());
498 #[cfg(feature = "tls")]
499 axhal::arch::write_thread_pointer(init_task.tls.tls_ptr() as usize);
500 let ptr = Arc::into_raw(init_task);
501 unsafe {
502 axhal::cpu::set_current_task_ptr(ptr);
503 }
504 }
505
506 pub(crate) unsafe fn set_current(prev: Self, next: AxTaskRef) {
507 let Self(arc) = prev;
508 ManuallyDrop::into_inner(arc); let ptr = Arc::into_raw(next);
510 unsafe {
511 axhal::cpu::set_current_task_ptr(ptr);
512 }
513 }
514}
515
516impl Deref for CurrentTask {
517 type Target = TaskInner;
518 fn deref(&self) -> &Self::Target {
519 self.0.deref()
520 }
521}
522
523extern "C" fn task_entry() -> ! {
524 #[cfg(feature = "smp")]
525 unsafe {
526 crate::run_queue::clear_prev_task_on_cpu();
528 }
529 #[cfg(feature = "irq")]
531 axhal::arch::enable_irqs();
532 let task = crate::current();
533 if let Some(entry) = task.entry {
534 unsafe { Box::from_raw(entry)() };
535 }
536 crate::exit(0);
537}