本章主要内容为分时多任务的切换，重点代码在 os/src/task 目录下。此处代码在我看来写得有些繁琐，不够清晰，正好做实验也要改，本节先不考虑实验目标，直接修改代码。
-
TrapContext此结构体在
os/src/trap/context.rs，原始版本：use riscv::register::sstatus::{Sstatus, self, SPP}; #[repr(C)] pub struct TrapContext { pub x: [usize; 32], pub sstatus: Sstatus, pub sepc: usize, } impl TrapContext { pub fn set_sp(&mut self, sp: usize) { self.x[2] = sp; } pub fn app_init_context(entry: usize, sp: usize) -> Self { let mut sstatus = sstatus::read(); sstatus.set_spp(SPP::User); let mut cx = Self { x: [0; 32], sstatus, sepc: entry, }; cx.set_sp(sp); cx } }
set_sp方法在后面章节好像用到了，但是本章只在后面的app_init_context里用了，因此删除方法，内联进去。同时，在
os/src/trap/mod.rs的trap_handler中有如下写法：Trap::Exception(Exception::UserEnvCall) => { cx.sepc += 4; cx.x[10] = syscall(cx.x[17], [cx.x[10], cx.x[11], cx.x[12]]) as usize; }
这个操作完全属于
TrapContext而且相当典型，应该封装为方法。并且，修改后，通用寄存器和sstatus不再需要从外部访问了，可以删去pub。修改后的TrapContext：use riscv::register::sstatus::{self, Sstatus, SPP}; #[repr(C)] pub struct TrapContext { x: [usize; 32], sstatus: Sstatus, pub sepc: usize, } impl TrapContext { pub fn app_init_context(entry: usize, sp: usize) -> Self { let mut sstatus = sstatus::read(); sstatus.set_spp(SPP::User); let mut cx = Self { x: [0; 32], sstatus, sepc: entry, }; cx.x[2] = sp; cx } pub fn ecall(&mut self) { self.sepc += 4; self.x[10] = crate::syscall::syscall(self.x[17], [self.x[10], self.x[11], self.x[12]]) as _; } }
-
TaskContext注意到
zero_init()方法用于零初始化。pub fn zero_init() -> Self { Self { ra: 0, sp: 0, s: [0; 12], } }
对于这种返回简单常量的方法尽量直接写常量，让语义更明确。改为：
pub const ZERO: Self = Self { ra: 0, sp: 0, s: [0; 12], };
-
TaskManager::run_first_taskTaskManager::run_first_task用于从启动控制流切换到第一个用户任务流。fn run_first_task(&self) -> ! { let mut inner = self.inner.exclusive_access(); let task0 = &mut inner.tasks[0]; task0.task_status = TaskStatus::Running; let next_task_cx_ptr = &task0.task_cx as *const TaskContext; drop(inner); let mut _unused = TaskContext::zero_init(); // before this, we should drop local variables that must be dropped manually unsafe { __switch( &mut _unused as *mut TaskContext, next_task_cx_ptr, ); } panic!("unreachable in run_first_task!"); }
对于变量
_unused，教程是这么说的：因此，我们显式声明了一个
_unused变量，并将它的地址作为第一个参数传给__switch，这样保存一些寄存器之后的启动栈栈顶的位置将会保存在此变量中。然而无论是此变量还是启动栈我们之后均不会涉及到，一旦应用开始运行，我们就开始在应用的用户栈和内核栈之间开始切换了。这里声明此变量的意义仅仅是为了避免覆盖到其他数据。然而，如果我们观察
__switch的（汇编）实现，会发现：.altmacro .macro SAVE_SN n sd s\n, (\n+2)*8(a0) .endm .macro LOAD_SN n ld s\n, (\n+2)*8(a1) .endm .section .text .globl __switch __switch: # __switch( # current_task_cx_ptr: *mut TaskContext, # next_task_cx_ptr: *const TaskContext # ) # save kernel stack of current task sd sp, 8(a0) # save ra & s0~s11 of current execution sd ra, 0(a0) .set n, 0 .rept 12 SAVE_SN %n .set n, n + 1 .endr # restore ra & s0~s11 of next execution ld ra, 0(a1) .set n, 0 .rept 12 LOAD_SN %n .set n, n + 1 .endr # restore kernel stack of next task ld sp, 8(a1) ret
其被
# restore ra & s0~s11 of next execution注释行分隔为上下两部分，上半部分用于保存当前上下文，下半部分用于切换到新任务上下文，这两部分实际上是无关的。因此，我们可以把这两部分分开，顺便整理一下：.altmacro .macro SAVE_SN n sd s\n, (\n+2)*8(a0) .endm .macro LOAD_SN n ld s\n, (\n+2)*8(a1) .endm .section .text # __switch( # next_task_cx_ptr: *const TaskContext # current_task_cx_ptr: *mut TaskContext, # ) .globl __switch # __init( # next_task_cx_ptr: *const TaskContext # ) .globl __init __switch: sd ra, 0(a1) sd sp, 8(a1) .set n, 0 .rept 12 SAVE_SN %n .set n, n + 1 .endr __init: .set n, 0 .rept 12 LOAD_SN %n .set n, n + 1 .endr ld sp, 8(a0) ld ra, 0(a0) ret
__switch还是__switch，__init则跳过保存上下文的步骤，直接加载新任务。**注意！**为了让只有一个参数的
__init能单独调用，我们需要将__switch的两个参数颠倒一下，让新任务上下文统一在第一个参数的位置。**特别注意！！！**颠倒参数寄存器不要忘了把上面宏里的寄存器也颠倒过来！为了这个我折腾了一个下午！
接下来，在
os/src/task/switch.rs里修改__switch签名并增加__init：core::arch::global_asm!(include_str!("switch.S")); use super::TaskContext; extern "C" { pub fn __switch(next_task_cx_ptr: *const TaskContext, current_task_cx_ptr: *mut TaskContext); pub fn __init(next_task_cx_ptr: *const TaskContext); }
最后将
TaskManager::run_first_task改为：fn run_first_task(&self) -> ! { let next_task_cx_ptr = { let mut inner = self.inner.exclusive_access(); let task0 = unsafe { inner.tasks.get_unchecked_mut(0) }; task0.task_status = TaskStatus::Running; &task0.task_cx as _ }; unsafe { __init(next_task_cx_ptr) }; unreachable!("run_first_task"); }
是不是清爽了不少？实际执行的指令数也减少了 14 个，四舍五入就是 0 个……啊不是，是一个亿啊！而且这个
__init在下面还会用到。 -
TaskManager::run_next_task先看原版：
fn run_next_task(&self) { if let Some(next) = self.find_next_task() { let mut inner = self.inner.exclusive_access(); let current = inner.current_task; inner.tasks[next].task_status = TaskStatus::Running; inner.current_task = next; let current_task_cx_ptr = &mut inner.tasks[current].task_cx as *mut TaskContext; let next_task_cx_ptr = &inner.tasks[next].task_cx as *const TaskContext; drop(inner); // before this, we should drop local variables that must be dropped manually unsafe { __switch( current_task_cx_ptr, next_task_cx_ptr, ); } // go back to user mode } else { panic!("All applications completed!"); } }
对
TaskManager::run_next_task的改造分为两步。首先，先想办法利用我们的
__init，减少 14 条访存指令，落实节能减排，加快上下文切换。显然，只有可能恢复到当前控制流时才需要保存当前控制流的上下文。因此，对于当前任务退出导致的上下文切换，可以使用__init跳过保存。任务退出时，会先由
TaskManager::mark_current_exited将状态修改为TaskStatus::Exited，因此我们只需要识别这个状态：fn run_next_task(&self) { let current_task_cx_ptr = { let mut inner = self.inner.exclusive_access(); let current = inner.current_task; let context = unsafe { inner.tasks.get_unchecked_mut(current) }; if context.task_status == TaskStatus::Exited { core::ptr::null_mut() } else { &mut inner.tasks[current].task_cx as *mut TaskContext } }; if let Some(next) = self.find_next_task() { let next_task_cx_ptr = { let mut inner = self.inner.exclusive_access(); inner.tasks[next].task_status = TaskStatus::Running; inner.current_task = next; &inner.tasks[next].task_cx as *const _ }; if current_task_cx_ptr.is_null() { unsafe { __init(next_task_cx_ptr) }; unreachable!(); } else { unsafe { __switch(next_task_cx_ptr, current_task_cx_ptr) }; // go back to user mode } } else { panic!("All applications completed!"); } }
代码看起来变复杂了，实际上大部分都是高级优化下会内联掉的起别名语句，实际开销最大的可能是多出来的 if 导致的条件跳转。
接下来，我们考察一个任务都能有哪些状态？
pub enum TaskStatus { UnInit, Ready, Running, Exited, }
创建内核任务表时，由于任务表是一个定长数组，必须有一个合法的初始值，因此需要
UnInit状态；用户程序加载到内存时，有值的任务更新到Ready状态；正在运行的任务处于Running状态；任务未完成但被换出时，状态变为Ready；任务退出时，状态变为Exited。然而，这些状态都是有必要的吗？首先，
UnInit只是一个填充，由于存储了实际任务数量，无意义的格子根本不会被访问到，所以理论上填充什么都不会影响到程序运行。考虑到安全性，我们可以用Exited填充没有初始化的格子，这样即使后面程序写错也不会导致切换到意想不到的位置。更好的方案是直接用
core::mem::MaybeUninit。其次，一个任务切出后，要么还能继续运行，要么不能继续运行，我们只需要区分这两种状态。而任务切出前，一定是正在运行而处于
Running状 态。而从将Running状态修改为Ready或Exited是由两个mark...方法控制的。因此，实际上mark_current_suspended是 不需要的，Running状态和Exited状态已经可以供run_next_task区分了。因此，我们可以直接删除
UnInit枚举项和TaskManager::mark_current_suspended。顺便把这一堆也简化掉：pub fn run_first_task() { TASK_MANAGER.run_first_task(); } fn run_next_task() { TASK_MANAGER.run_next_task(); } fn mark_current_suspended() { TASK_MANAGER.mark_current_suspended(); } fn mark_current_exited() { TASK_MANAGER.mark_current_exited(); } pub fn suspend_current_and_run_next() { mark_current_suspended(); run_next_task(); } pub fn exit_current_and_run_next() { mark_current_exited(); run_next_task(); }
现在：
pub fn run_first_task() { TASK_MANAGER.run_first_task(); } pub fn suspend_current_and_run_next() { TASK_MANAGER.run_next_task(); } pub fn exit_current_and_run_next() { TASK_MANAGER.mark_current_exited(); TASK_MANAGER.run_next_task(); }
现在
1 reference的run_first_task变得突兀起来了……这是必需的吗？只要仔细观察
TaskManager::run_first_task和TaskManager::run_next_task的异同，就能发现这两个方法实际上并没有本质区别。经过刚才的修改，TaskManager::run_first_task不再需要一个任意的_unused保存用不到的控制流了，而通过判断当前任务的状态，TaskManager::run_next_task也可以选择不保存上下文。在启动期间，所有任务被初始化为Ready状态，而从一个任务切出时，任务只能处于Running或Exited状态，而只有Running状态才需要保存上下文。继续观察
suspend_current_and_run_next和exit_current_and_run_next，我们发现exit_current_and_run_next只是要先标记当前任务退出，这个工作也可以留在TaskManager::run_next_task里。因此，TaskManager::run_next_task可以进一步修改为：fn run_next_task(&self, current_exited: bool) { let mut inner = self.inner.exclusive_access(); // 更新当前任务状态，准备切换 let index = inner.current; let context = unsafe { inner.tasks.get_unchecked_mut(index) }; let current_task_cx_ptr = match context.task_status { TaskStatus::Ready => { // 当前任务就绪状态，仅出现在第一个任务 // 实际此时没有任务在运行，因此不需要保存上下文 core::ptr::null_mut() } TaskStatus::Running if current_exited => { // 当前任务已退出，不需要保存上下文 context.task_status = TaskStatus::Exited; core::ptr::null_mut() } TaskStatus::Running => { // 当前任务耗尽时间片 context.task_status = TaskStatus::Ready; &mut inner.tasks[index].task_cx as *mut TaskContext } TaskStatus::Exited => unreachable!(), }; drop(inner); // 弹出下一个任务 let next = self.find_next_task().expect("All applications completed!"); let mut inner = self.inner.exclusive_access(); inner.current = next; // 更新当前任务状态 let context = unsafe { inner.tasks.get_unchecked_mut(next.index) }; let next_task_cx_ptr = &context.task_cx as *const _; context.task_status = TaskStatus::Running; // 释放 `inner` 准备切换控制流 drop(inner); // 当前上下文不需要保存 if current_task_cx_ptr.is_null() { unsafe { __init(next_task_cx_ptr) }; unreachable!("Task exited!"); } else { unsafe { __switch(next_task_cx_ptr, current_task_cx_ptr) }; // go back to user mode } }
现在，
TaskManager有且只有这一个方法了。暴露出去的函数修改为：
pub fn run_next_task() { TASK_MANAGER.run_next_task(false); } pub fn exit_current_and_run_next() { TASK_MANAGER.run_next_task(true); }
现在不再需要一个特殊的
run_first_task，启动时也可以调用run_next_task。但是如果不修改TaskManager的初始化，序号为 1 的任务现在将最先启动，现象会改变，但这对实验没有影响。
终于可以开始做实验了。不过如果你已经完成了上面的修改，应该对整个任务的切换过程已经烂熟于心，实验只不过是个锦上添花的事情罢了。
这里就不抄题了，为了最终实现的优雅，一个优先队列是很有必要的。现在还没有堆分配，因此不能直接使用 BinaryHeap，我们只好自己造一个，还好定容单线程版本是很容易实现的：
use core::{fmt::Debug, mem::MaybeUninit};
/// 具有指定容量的优先队列（小顶堆）
pub struct PriorityQueue<T, const N: usize> {
len: usize,
val: [MaybeUninit<T>; N],
}
impl<T, const N: usize> Default for PriorityQueue<T, N> {
fn default() -> Self {
Self {
len: 0,
val: MaybeUninit::uninit_array(),
}
}
}
impl<T, const N: usize> PriorityQueue<T, N> {
#[inline]
fn get(&self, i: usize) -> &T {
unsafe { self.val.get_unchecked(i).assume_init_ref() }
}
}
impl<T: Ord, const N: usize> PriorityQueue<T, N> {
pub fn push(&mut self, t: T) {
if self.len == N {
panic!("Out of capacity!");
}
self.val[self.len] = MaybeUninit::new(t);
// 从底部上浮
let mut i = self.len;
while i > 0 {
// 父节点的序号
let j = i >> 1;
// 交换或退出
if self.get(i) < self.get(j) {
self.val.swap(i, j);
i = j;
} else {
break;
}
}
self.len += 1;
}
pub fn pop(&mut self) -> Option<T> {
if self.len == 0 {
return None;
}
self.len -= 1;
self.val.swap(0, self.len);
let n = self.len;
// 从顶部下沉
let mut i = 0;
loop {
// 较大子节点的序号
let j = i * 2 + 1;
let j = if n <= j {
break;
} else if n == j + 1 || self.get(j) < self.get(j + 1) {
j
} else {
j + 1
};
// 交换或退出
if self.get(i) > self.get(j) {
self.val.swap(i, j);
i = j;
} else {
break;
}
}
unsafe {
Some(
core::mem::replace(
//
self.val.get_unchecked_mut(n),
MaybeUninit::uninit(),
)
.assume_init(),
)
}
}
}
impl<T: Debug, const N: usize> Debug for PriorityQueue<T, N> {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
write!(f, "[")?;
if self.len > 0 {
write!(f, "{:?}", self.get(0))?;
for i in 1..self.len {
write!(f, ", {:?}", self.get(i))?;
}
}
write!(f, "]",)
}
}MaybeUninit::uninit_array() 目前仍然 unstable，需要在 os/src/main.rs 顶上添加：
#![feature(maybe_uninit_uninit_array)]对于 TaskControlBlock，我的处理方式是保持其主体部分在数组中不变，另取一个带索引的 Stride 类型在队列中排序。因为我们目前还没有堆，因此 Rust 的移动语义在性能上起不到什么作用。小顶堆中大量的交换操作都意味着拷贝整个数据结构两次，因此这个结构最好做得小一些。修改 os/src/task/task.rs 如下：
use super::TaskContext;
#[derive(Copy, Clone)]
pub struct TaskControlBlock {
pub priority: usize,
pub task_status: TaskStatus,
pub task_cx: TaskContext,
}
#[derive(Copy, Clone, PartialEq)]
pub enum TaskStatus {
Ready,
Running,
Exited,
}
pub const BIG_STRIDE: usize = isize::MAX as _;
#[derive(PartialEq, Eq, Clone, Copy, Debug)]
pub struct TaskStride {
pub stride: usize,
pub index: usize,
}
impl PartialOrd for TaskStride {
fn partial_cmp(&self, other: &Self) -> Option<core::cmp::Ordering> {
Some(self.cmp(other))
}
}
impl Ord for TaskStride {
fn cmp(&self, other: &Self) -> core::cmp::Ordering {
(self.stride.wrapping_sub(other.stride) as isize).cmp(&0)
}
}优先级在 TCB 中，给 TaskStride 实现 Ord 供优先队列使用。由于我们已经将 TaskManager::run_first_task 和 TaskManager::run_next_task 合并，现在也可以统一通过优先队列调度了。即使任务在初始化时就具有不同的优先级也能正确调度。
构造优先队列：
struct TaskManagerInner {
queue: PriorityQueue<TaskStride, MAX_APP_NUM>,
tasks: [TaskControlBlock; MAX_APP_NUM],
current: TaskStride,
}
lazy_static! {
pub static ref TASK_MANAGER: TaskManager = {
let num_app = get_num_app();
let mut queue = PriorityQueue::default();
let mut tasks = [TaskControlBlock {
priority: 2,
task_cx: TaskContext::ZERO,
task_status: TaskStatus::Exited,
}; MAX_APP_NUM];
for i in 0..num_app {
queue.push(TaskStride {
stride: if i == 0 { 0 } else { 1 },
index: i,
});
tasks[i].task_cx = TaskContext::goto_restore(init_app_cx(i));
tasks[i].task_status = TaskStatus::Ready;
}
TaskManager {
num_app,
inner: unsafe {
UPSafeCell::new(TaskManagerInner {
queue,
tasks,
current: TaskStride {
stride: 0,
index: 0,
},
})
},
}
};
}为了使第一个调用的是 0 号任务，让其他任务的初始 stride 大于 0。初始的 TaskManagerInner::current 可以随便写。
切换任务时更新 stride，顺便把 set_priority 暴露出去：
impl TaskManager {
fn run_next_task(&self, current_exited: bool) {
let mut inner = self.inner.exclusive_access();
// 更新当前任务状态，准备切换
let TaskStride { stride, index } = inner.current;
let context = unsafe { inner.tasks.get_unchecked_mut(index) };
let current_task_cx_ptr = match context.task_status {
TaskStatus::Ready => {
// 当前任务就绪状态，仅出现在第一个任务
// 实际此时没有任务在运行，因此不需要保存上下文
core::ptr::null_mut()
}
TaskStatus::Running if current_exited => {
// 当前任务已退出，不需要保存上下文
context.task_status = TaskStatus::Exited;
core::ptr::null_mut()
}
TaskStatus::Running => {
// 当前任务耗尽时间片
context.task_status = TaskStatus::Ready;
let stride = stride.wrapping_add(BIG_STRIDE / context.priority);
inner.queue.push(TaskStride { stride, index });
&mut inner.tasks[index].task_cx as *mut TaskContext
}
TaskStatus::Exited => unreachable!(),
};
// 弹出下一个任务
let next = inner.queue.pop().expect("All applications completed!");
inner.current = next;
// 更新当前任务状态
let context = unsafe { inner.tasks.get_unchecked_mut(next.index) };
let next_task_cx_ptr = &context.task_cx as *const _;
context.task_status = TaskStatus::Running;
// 释放 `inner` 准备切换控制流
drop(inner);
// 当前上下文不需要保存
if current_task_cx_ptr.is_null() {
unsafe { __init(next_task_cx_ptr) };
unreachable!("Task exited!");
} else {
unsafe { __switch(next_task_cx_ptr, current_task_cx_ptr) };
// go back to user mode
}
}
fn set_priority(&self, value: isize) {
let mut inner = self.inner.exclusive_access();
let TaskStride { stride: _, index } = inner.current;
inner.tasks[index].priority = value as _;
}
}
pub fn set_priority(value: isize) {
TASK_MANAGER.set_priority(value);
}再把系统调用那边改一下，先判断参数正确性然后调 set_priority：
pub fn sys_setpriority(prio: isize) -> isize {
if prio >= 2 {
set_priority(prio);
prio
} else {
-1
}
}到这里实验就做完了。