sparreal_kernel/task/
tcb.rs

1use core::{
2    alloc::Layout,
3    fmt::Debug,
4    ops::{Deref, DerefMut},
5    ptr::NonNull,
6    sync::atomic::{AtomicUsize, Ordering},
7};
8
9use alloc::{boxed::Box, string::String};
10use log::trace;
11
12use crate::{platform, platform_if::PlatformImpl, task::schedule::*};
13
14use super::{TaskConfig, TaskError};
15
16#[repr(transparent)]
17#[derive(Clone, Copy)]
18pub struct TaskControlBlock(*mut u8);
19
20unsafe impl Send for TaskControlBlock {}
21
22impl From<*mut u8> for TaskControlBlock {
23    fn from(value: *mut u8) -> Self {
24        Self(value)
25    }
26}
27
28#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
29#[repr(transparent)]
30pub struct Pid(usize);
31
32impl Pid {
33    pub fn new() -> Self {
34        static ITER: AtomicUsize = AtomicUsize::new(0);
35        Self(ITER.fetch_add(1, Ordering::SeqCst))
36    }
37}
38
39impl Default for Pid {
40    fn default() -> Self {
41        Self::new()
42    }
43}
44
45impl Debug for Pid {
46    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
47        writeln!(f, "{:?}", self.0)
48    }
49}
50
51impl TaskControlBlock {
52    pub(super) fn new<F>(entry: F, config: TaskConfig) -> Result<Self, TaskError>
53    where
54        F: FnOnce() + Send + 'static,
55    {
56        let entry_box = Box::new(entry);
57
58        let buffer = NonNull::new(unsafe {
59            alloc::alloc::alloc_zeroed(
60                Layout::from_size_align(Self::tcb_size(config.stack_size), platform::page_size())
61                    .unwrap(),
62            )
63        })
64        .ok_or(TaskError::NoMemory)?;
65
66        let pid = Pid::new();
67
68        unsafe {
69            let task_data = &mut *(buffer.as_ptr() as *mut TaskControlBlockData);
70            task_data.pid = pid;
71            task_data.stack_size = config.stack_size;
72            task_data.priority = config.priority;
73            task_data.name = config.name;
74            task_data.state = TaskState::Idle;
75            task_data.entry = Some(entry_box);
76        }
77
78        let mut task = Self(buffer.as_ptr());
79        task.sp = task.stack_top() as usize;
80
81        unsafe {
82            task.sp -= PlatformImpl::cpu_context_size();
83            let ctx_ptr = task.sp as *mut u8;
84
85            PlatformImpl::cpu_context_set_sp(ctx_ptr, task.sp);
86            PlatformImpl::cpu_context_set_pc(ctx_ptr, task_entry as _);
87        }
88        Ok(task)
89    }
90
91    pub(super) fn new_main() -> Self {
92        let entry_box = Box::new(|| {});
93
94        let buffer = NonNull::new(unsafe {
95            alloc::alloc::alloc_zeroed(
96                Layout::from_size_align(Self::tcb_size(0), platform::page_size()).unwrap(),
97            )
98        })
99        .ok_or(TaskError::NoMemory)
100        .expect("main task no memory");
101
102        let pid = Pid::new();
103
104        unsafe {
105            let task_data = &mut *(buffer.as_ptr() as *mut TaskControlBlockData);
106            task_data.pid = pid;
107            task_data.stack_size = 0;
108            task_data.priority = 0;
109            task_data.name = "Main".into();
110            task_data.state = TaskState::Running;
111            task_data.entry = Some(entry_box);
112        }
113        Self(buffer.as_ptr())
114    }
115
116    fn tcb_size(stack_size: usize) -> usize {
117        size_of::<TaskControlBlockData>() + stack_size
118    }
119
120    fn stack_bottom(&self) -> *mut u8 {
121        unsafe { self.0.add(size_of::<TaskControlBlockData>()) }
122    }
123
124    fn stack_top(&self) -> *mut u8 {
125        unsafe { self.stack_bottom().add(self.stack_size) }
126    }
127
128    pub(super) unsafe fn drop(self) {
129        let size = Self::tcb_size(self.stack_size);
130
131        unsafe {
132            alloc::alloc::dealloc(
133                self.0,
134                Layout::from_size_align_unchecked(size, platform::page_size()),
135            )
136        };
137    }
138
139    fn addr(&self) -> *mut u8 {
140        self.0 as _
141    }
142
143    pub(super) fn switch_to(&self, next: &TaskControlBlock) {
144        trace!("switch {} -> {}", self.name, next.name);
145        set_current(next);
146        match self.state {
147            TaskState::Stopped => finished_push(*self),
148            _ => idle_push(*self),
149        }
150
151        unsafe {
152            PlatformImpl::cpu_context_switch(self.addr(), next.addr());
153        }
154    }
155}
156
157impl Deref for TaskControlBlock {
158    type Target = TaskControlBlockData;
159
160    fn deref(&self) -> &Self::Target {
161        unsafe { &*(self.0 as *mut TaskControlBlockData) }
162    }
163}
164
165impl DerefMut for TaskControlBlock {
166    fn deref_mut(&mut self) -> &mut Self::Target {
167        unsafe { &mut *(self.0 as *mut TaskControlBlockData) }
168    }
169}
170
171pub struct TaskControlBlockData {
172    pub pid: Pid,
173    pub name: String,
174    pub priority: usize,
175    pub stack_size: usize,
176    pub entry: Option<Box<dyn FnOnce()>>,
177    pub state: TaskState,
178    pub sp: usize,
179}
180
181#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
182pub enum TaskState {
183    Idle,
184    Running,
185    Suspend,
186    Stopped,
187}
188
189extern "C" fn task_entry() -> ! {
190    let mut task = current();
191
192    if let Some(entry) = task.entry.take() {
193        entry();
194        task.state = TaskState::Stopped;
195    }
196    schedule();
197    unreachable!("task exited!");
198}
199
200pub fn current() -> TaskControlBlock {
201    unsafe {
202        let ptr = PlatformImpl::get_current_tcb_addr();
203        TaskControlBlock::from(ptr)
204    }
205}
206
207pub fn set_current(tcb: &TaskControlBlock) {
208    unsafe {
209        PlatformImpl::set_current_tcb_addr(tcb.addr());
210    }
211}