os/task/
task.rs

1//!Implementation of [`TaskControlBlock`]
2use super::TaskContext;
3use super::{KernelStack, PidHandle, pid_alloc};
4use crate::config::TRAP_CONTEXT;
5use crate::mm::{KERNEL_SPACE, MemorySet, PhysPageNum, VirtAddr};
6use crate::sync::UPSafeCell;
7use crate::trap::{TrapContext, trap_handler};
8use alloc::sync::{Arc, Weak};
9use alloc::vec::Vec;
10use core::cell::RefMut;
11
12pub struct TaskControlBlock {
13    // immutable
14    pub pid: PidHandle,
15    pub kernel_stack: KernelStack,
16    // mutable
17    inner: UPSafeCell<TaskControlBlockInner>,
18}
19
20pub struct TaskControlBlockInner {
21    pub trap_cx_ppn: PhysPageNum,
22    #[allow(unused)]
23    pub base_size: usize,
24    pub task_cx: TaskContext,
25    pub task_status: TaskStatus,
26    pub memory_set: MemorySet,
27    pub parent: Option<Weak<TaskControlBlock>>,
28    pub children: Vec<Arc<TaskControlBlock>>,
29    pub exit_code: i32,
30}
31
32impl TaskControlBlockInner {
33    /*
34    pub fn get_task_cx_ptr2(&self) -> *const usize {
35        &self.task_cx_ptr as *const usize
36    }
37    */
38    pub fn get_trap_cx(&self) -> &'static mut TrapContext {
39        self.trap_cx_ppn.get_mut()
40    }
41    pub fn get_user_token(&self) -> usize {
42        self.memory_set.token()
43    }
44    fn get_status(&self) -> TaskStatus {
45        self.task_status
46    }
47    pub fn is_zombie(&self) -> bool {
48        self.get_status() == TaskStatus::Zombie
49    }
50}
51
52impl TaskControlBlock {
53    pub fn inner_exclusive_access(&self) -> RefMut<'_, TaskControlBlockInner> {
54        self.inner.exclusive_access()
55    }
56    pub fn new(elf_data: &[u8]) -> Self {
57        // memory_set with elf program headers/trampoline/trap context/user stack
58        let (memory_set, user_sp, entry_point) = MemorySet::from_elf(elf_data);
59        let trap_cx_ppn = memory_set
60            .translate(VirtAddr::from(TRAP_CONTEXT).into())
61            .unwrap()
62            .ppn();
63        // alloc a pid and a kernel stack in kernel space
64        let pid_handle = pid_alloc();
65        let kernel_stack = KernelStack::new(&pid_handle);
66        let kernel_stack_top = kernel_stack.get_top();
67        // push a task context which goes to trap_return to the top of kernel stack
68        let task_control_block = Self {
69            pid: pid_handle,
70            kernel_stack,
71            inner: unsafe {
72                UPSafeCell::new(TaskControlBlockInner {
73                    trap_cx_ppn,
74                    base_size: user_sp,
75                    task_cx: TaskContext::goto_trap_return(kernel_stack_top),
76                    task_status: TaskStatus::Ready,
77                    memory_set,
78                    parent: None,
79                    children: Vec::new(),
80                    exit_code: 0,
81                })
82            },
83        };
84        // prepare TrapContext in user space
85        let trap_cx = task_control_block.inner_exclusive_access().get_trap_cx();
86        *trap_cx = TrapContext::app_init_context(
87            entry_point,
88            user_sp,
89            KERNEL_SPACE.exclusive_access().token(),
90            kernel_stack_top,
91            trap_handler as usize,
92        );
93        task_control_block
94    }
95    pub fn exec(&self, elf_data: &[u8]) {
96        // memory_set with elf program headers/trampoline/trap context/user stack
97        let (memory_set, user_sp, entry_point) = MemorySet::from_elf(elf_data);
98        let trap_cx_ppn = memory_set
99            .translate(VirtAddr::from(TRAP_CONTEXT).into())
100            .unwrap()
101            .ppn();
102
103        // **** access inner exclusively
104        let mut inner = self.inner_exclusive_access();
105        // substitute memory_set
106        inner.memory_set = memory_set;
107        // update trap_cx ppn
108        inner.trap_cx_ppn = trap_cx_ppn;
109        // initialize base_size
110        inner.base_size = user_sp;
111        // initialize trap_cx
112        let trap_cx = inner.get_trap_cx();
113        *trap_cx = TrapContext::app_init_context(
114            entry_point,
115            user_sp,
116            KERNEL_SPACE.exclusive_access().token(),
117            self.kernel_stack.get_top(),
118            trap_handler as usize,
119        );
120        // **** release inner automatically
121    }
122    pub fn fork(self: &Arc<Self>) -> Arc<Self> {
123        // ---- access parent PCB exclusively
124        let mut parent_inner = self.inner_exclusive_access();
125        // copy user space(include trap context)
126        let memory_set = MemorySet::from_existed_user(&parent_inner.memory_set);
127        let trap_cx_ppn = memory_set
128            .translate(VirtAddr::from(TRAP_CONTEXT).into())
129            .unwrap()
130            .ppn();
131        // alloc a pid and a kernel stack in kernel space
132        let pid_handle = pid_alloc();
133        let kernel_stack = KernelStack::new(&pid_handle);
134        let kernel_stack_top = kernel_stack.get_top();
135        let task_control_block = Arc::new(TaskControlBlock {
136            pid: pid_handle,
137            kernel_stack,
138            inner: unsafe {
139                UPSafeCell::new(TaskControlBlockInner {
140                    trap_cx_ppn,
141                    base_size: parent_inner.base_size,
142                    task_cx: TaskContext::goto_trap_return(kernel_stack_top),
143                    task_status: TaskStatus::Ready,
144                    memory_set,
145                    parent: Some(Arc::downgrade(self)),
146                    children: Vec::new(),
147                    exit_code: 0,
148                })
149            },
150        });
151        // add child
152        parent_inner.children.push(task_control_block.clone());
153        // modify kernel_sp in trap_cx
154        // **** access children PCB exclusively
155        let trap_cx = task_control_block.inner_exclusive_access().get_trap_cx();
156        trap_cx.kernel_sp = kernel_stack_top;
157        // return
158        task_control_block
159        // ---- release parent PCB automatically
160        // **** release children PCB automatically
161    }
162    pub fn getpid(&self) -> usize {
163        self.pid.0
164    }
165}
166
167#[derive(Copy, Clone, PartialEq)]
168pub enum TaskStatus {
169    Ready,
170    Running,
171    Zombie,
172}