【Async OS】最小化异步操作系统

前提

需要了解操作系统基础知识，实现简单的操作系统内核框架，并且要实现堆内存分配，因为必须要使用 Box 和 Pin，还要实现打印功能用于测试，还需要实现时间获取，用于模拟异步延迟。

Async in Rust

Rust 提供了 Future trait 用于实现异步操作，其结构定义如下：

pub trait Future {
    type Output;
    fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output>;
}

poll 方法接受两个参数，Pin<&mut Self> 其实和 &mut Self 类似，只是需要 Pin 来固定内存地址，cx 参数是用于传入一个唤醒器，在异步任务完成时可以通过唤醒器发出信号。poll 方法返回一个 Poll 枚举：

pub enum Poll<T> {
    Ready(T),
    Pending,
}

大致工作原理其实很简单，调用异步函数的 poll 方法，如果返回的是 Pending 表示值还不可用，CPU 可以先去执行其他任务，稍候再试。返回 Ready 则表示任务已完成，可以接着执行往下的程序。

运行时

知道基本原理后，基于异步来构建操作系统的思路就很清晰了，即遇到 Pending 就切换到另一个任务，直到所有任务都完成。

先创建一个 Runtime 结构体，包含 tasks 队列，用于存储需要执行的任务。spawn 方法用于将任务添加到队列，run 方法用于执行异步任务，其逻辑是先取出队列中一个任务，通过 loop 不断尝试执行异步任务，如果任务 Pending，则先去执行另一个任务，以此实现非阻塞，直到队列任务全部为空。

poll 方法需要传入 Waker 参数，用于在异步任务完成后发出信号，因为目前是 loop 盲等的机制，并没有实现真正的唤醒，所以先采用一个虚假唤醒器。

pub struct Runtime {
    tasks: VecDeque<Task>,
}

impl Runtime {
    pub fn new() -> Self {
        Runtime {
            tasks: VecDeque::new(),
        }
    }

    pub fn spawn(&mut self, future: impl Future<Output = ()> + Send + Sync + 'static) {
        self.tasks.push_back(Task::new(future))
    }
    
    pub fn run(&mut self) {
        while let Some(mut task) = self.tasks.pop_front() {
            let waker = dummy_waker();
            let mut context = Context::from_waker(&waker);
            loop {
                match task.poll(&mut context) {
                    Poll::Ready(val) => break val,
                    Poll::Pending => {
                        self.tasks.push_back(task);
                        break;
                    }
                };
            }
        }
    }
}

任务

任务的结构也很简单，含有一个内部 future，在 poll 时执行 future 的 poll。内部 future 定义很长但是不复杂，基于的 future 结构是 Future<Output = ()>，然后需要一个 'static 生命周期，同时需要 Send 和 Sync 实现跨线程共享，虽然现在没用到，但是 Rust 编译器可不同意你不写。dyn 声明动态类型也是必须要，同时还需要使用 Box 包裹来使编译器确定闭包大小，Pin 用于固定内存位置不可以动。

struct Task {
    future: Pin<Box<dyn Future<Output = ()> + Send + Sync + 'static>>,
}

impl Task {
    pub fn new(future: impl Future<Output = ()> + Send + Sync + 'static) -> Task {
        Task {
            future: Box::pin(future),
        }
    }

    fn poll(&mut self, cx: &mut Context) -> Poll<()> {
        self.future.as_mut().poll(cx)
    }
}

虚假唤醒器

无需了解过多，真正实现唤醒器的时候才需深入了解。

fn dummy_waker() -> Waker {
    unsafe { Waker::from_raw(dummy_raw_waker()) }
}

fn dummy_raw_waker() -> RawWaker {
    fn no_op(_: *const ()) {}
    fn clone(_: *const ()) -> RawWaker {
        dummy_raw_waker()
    }

    let vtable = &RawWakerVTable::new(clone, no_op, no_op, no_op);
    RawWaker::new(0 as *const (), vtable)
}

Delay

基础框架实现完后，还需要实现一个延迟任务，用于模拟耗时操作。我打算实现一个 delay 方法模拟 sleep，以测试任务 sleep 的时候运行时会切换到下一个任务。代码很简单，DelayFuture 结构体包含一个 target_time 和一个 waker，target_time 表示延迟到什么时候，waker 用于在延迟完成后发出信号。poll 方法中判断当前时间是否大于 target_time，如果大于则返回 Ready，否则返回 Pending。

pub async fn delay(ms: usize) {
    DelayFuture::new(ms).await;
}

struct DelayFuture {
    target_time: usize,
    waker: Option<Waker>,
}

impl DelayFuture {
    fn new(ms: usize) -> Self {
        DelayFuture {
            target_time: get_time_ms() + ms,
            waker: None,
        }
    }
}

impl Future for DelayFuture {
    type Output = ();

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        if get_time_ms() >= self.target_time {
            Poll::Ready(())
        } else {
            Poll::Pending
        }
    }
}

测试

rust_main 操作系统的 rust 入口，task1 和 task2 是两个异步任务，task1 先打印 start task 1，然后延迟 200ms，再打印 end task 1。task2 也是类似，只是延迟时间更长。运行时先执行 task1，然后切换到 task2，再切换到 task1，最后切换到 task2。

#[no_mangle]
fn rust_main() {
    let mut rt = Runtime::new();
    rt.spawn(task1());
    rt.spawn(task2());
    rt.run();
}

async fn task1() {
    println!("start task 1");
    delay(200).await;
    println!("end task 1");
}

async fn task2() {
    println!("start task 2");
    delay(500).await;
    println!("end task 2");
}

打印结果：

start task 1
start task 2
end task 1
end task 2

总结

Rust 的异步编程还是很方便的，但是目前的实现还是很粗糙，比如没有实现真正的唤醒器，也没有实现真正的异步操作，只是模拟了异步延迟。下一步是实现真正的唤醒器，然后尝试实现异步的 I/O。