async-await Rust: 200 多行程式碼實現一個極簡 runtime

2023-07-12 18:00:41

What I cannot create, I do not understand

Rust 中的 runtime 到底是咋回事, 為了徹底搞懂它, 我在儘量不借助第三方 crate 的情況下實現了一個玩具 runtime, 之所以說是玩具,因為它沒有複雜的排程演演算法(只有一個全域性 task queue)

程式碼除了 mpmc(multi-producer, multi-consumer) 使用第三方 crate crossbeam 之外, 其餘程式碼一律手擼

可以這麼玩

fn main() {
    let toy = Toy::new();

    for i in 1..=20 {
        toy.spawn(async move {
            let ret = FakeIO::new(Duration::from_secs(i)).await;
            println!("{:?}: {:?}", thread::current().id(), ret);
        })
    }

    toy.run(4); // 4 threads
}

其中 FakeIO 也是足夠單純

pub struct FakeIO {
    finished: Arc<AtomicBool>,
    duration: Duration,
}

impl Future for FakeIO {
    type Output = Duration;

    fn poll(
        self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Self::Output> {
        if self.finished.load(Ordering::Acquire) {
            return Poll::Ready(self.duration);
        }

        let finished = self.finished.clone();
        let waker = cx.waker().clone();
        let duration = self.duration;

        thread::spawn(move || {
            thread::sleep(duration);

            finished.store(true, Ordering::Release);

            waker.wake();
        });

        Poll::Pending
    }
}

資料結構

資料結構就下面幾個(參考了 tokio 的設計)

struct Task {
    raw: RawTask,
}

unsafe impl Send for Task {}
unsafe impl Sync for Task {}

struct RawTask {
    ptr: NonNull<Header>, // pointer to Cell<T> where T: Future
}

struct Header {
    // todo: maybe replace the Mutex<State> with AtomicUsize
    state: Mutex<State>,
    vtable: &'static Vtable,
    sender: crossbeam::channel::Sender<Task>,
}

#[derive(Default)]
struct State {
    running: bool,
    notified: bool,
    completed: bool,
}

/// #[repr(C)] make sure `*mut Cell<T>` can cast to valid `*mut Header`, and backwards. 
/// In the default situation, the data layout may not be the same as the order in which the fields are specified in the declaration of the type
/// 預設情況下 Rust 的資料佈局不一定會按照 field 的宣告順序排列
/// [The Default Representation](https://doc.rust-lang.org/reference/type-layout.html?#the-default-representation)
///
/// [playground link](https://play.rust-lang.org/?version=stable&mode=debug&edition=2021&gist=39ac84782d121970598b91201b168f82)
/// 
/// you can easilly view the data layout with this crate https://github.com/hangj/layout-rs
#[repr(C)]
struct Cell<T: Future> {
    header: Header,
    future: T,
    output: Option<T::Output>,
}

struct Vtable {
    poll_task: unsafe fn(NonNull<Header>),
    clone_task: unsafe fn(NonNull<Header>) -> NonNull<Header>,
    drop_task: unsafe fn(NonNull<Header>),
}

其中值得注意的是:

  • RawTask 內的 ptr 實際上指向的是 NonNull<Cell<T: Future>>
  • Cell<T: Future> 被標記了 #repr(C), 原因已在註釋中說明
  • vtable 的設計參考了 Waker 中的 vtable, 相當於利用泛型函數儲存了型別資訊, 便於後面從裸指標恢復到原始型別

點選「閱讀原文」直達 toy-runtime 倉庫

Have fun!