從 generator 的角度看 Rust 非同步程式碼

文｜Ruihang Xia

目前參與邊緣時序資料儲存引擎專案

本文 6992 字閱讀 18 分鐘

前言

作為 2018 edition 一個比較重要的特性 Rust 的非同步程式設計現在已經得到了廣泛的使用。使用的時候難免會好奇它是如何運作的，這篇文章嘗試從 generator 以及變數捕獲的方面進行探索，而後介紹了在嵌入式時序儲存引擎 ceresdb-helix 的研發過程中遇到的一個場景。

囿於作者水平內容難免存在一些錯漏之處，還煩請留言告知。

PART. 1 async/.await, coroutine and generator

async/.await 語法在 1.39 版本[1]進入 stable channel，它能夠很方便地編寫非同步程式碼：

、、、java
async fn asynchronous() {

// snipped

}

async fn foo() {

let x: usize = 233;
asynchronous().await;
println!("{}", x);

、、、

在上面的示例中，區域性變數 x 能夠直接在一次非同步過程（fn asynchoronous）之後使用，和寫同步程式碼一樣。而在這之前，非同步程式碼一般是通過類似 futures 0.1[2] 形式的組合子來使用，想要給接下來 (如 and_then()) 的非同步過程的使用的區域性變數需要被顯式手動地以閉包出入參的方式鏈式處理，體驗不是特別好。

async/.await 所做的實際上就是將程式碼變換一下，變成 generator/coroutine[3] 的形式去執行。一個 coroutine 過程可以被掛起，去做一些別的事情然後再繼續恢復執行，目前用起來就是 .await 的樣子。以上面的程式碼為例，在非同步過程 foo()中呼叫了另一個非同步過程 asynchronous() ，在第七行的 .await 時當前過程的執行被掛起，等到可以繼續執行的時候再被恢復。

而恢復執行可能需要之前的一些資訊，如在 foo()中我們在第八行用到了之前的資訊 x。也就是說 async 過程要有能力儲存一些內部區域性狀態，使得它們能夠在 .await 之後被繼續使用。換句話說要在 generator state 裡面儲存可能在 yield 之後被使用的區域性變數。這裡需要引入 pin[4] 機制解決可能出現的自引用問題，這部分不再贅述。

PART. 2 visualize generator via MIR

我們可以透過 MIR[5]來看一下前面提到的 generator 是什麼樣子的。MIR 是 Rust 的一箇中間表示，基於控制流圖 CFG[6]表示。CFG 能夠比較直觀地展示程式執行起來大概是什麼樣子，MIR 在有時不清楚你的 Rust 程式碼到底變成了什麼樣子的時候能夠起到一些幫助。

想要得到程式碼的 MIR 表示有幾種方法，假如現在手邊有一個可用的 Rust toolchain，可以像這樣傳遞一個環境變數給 rustc ，再使用 cargo 進行構建來產生 MIR：

RUSTFLAGS="--emit mir" cargo build

構建成功的話會在 target/debug/deps/ 目錄下生成一個 .mir 的檔案。或者也能通過 https://play.rust-lang.org/ 來獲取 MIR，在 Run 旁邊的溢位選單上選擇 MIR 就可以。

由 2021-08 nightly 的 toolchain 所產生的 MIR 大概是這個樣子的，有許多不認識的東西可以不用管，大概知道一下。

_0, _1 這些是變數
有許多語法和 Rust 差不多，如型別註解，函式定義及呼叫和註釋等就行了。

fn future_1() -> impl Future {
    let mut _0: impl std::future::Future; // return place in scope 0 at src/anchored.rs:27:21: 27:21
    let mut _1: [static generator@src/anchored.rs:27:21: 27:23]; // in scope 0 at src/anchored.rs:27:21: 27:23

    bb0: {
        discriminant(_1) = 0; // scope 0 at src/anchored.rs:27:21: 27:23
        _0 = from_generator::<[static generator@src/anchored.rs:27:21: 27:23]>(move _1) -> bb1; // scope 0 at src/anchored.rs:27:21: 27:23
                                         // mir::Constant
                                         // + span: src/anchored.rs:27:21: 27:23
                                         // + literal: Const { ty: fn([static generator@src/anchored.rs:27:21: 27:23]) -> impl std::future::Future {std::future::from_generator::<[static generator@src/anchored.rs:27:21: 27:23]>}, val: Value(Scalar(<ZST>)) }
    }

    bb1: {
        return; // scope 0 at src/anchored.rs:27:23: 27:23
    }
}

fn future_1::{closure#0}(_1: Pin<&mut [static generator@src/anchored.rs:27:21: 27:23]>, _2: ResumeTy) -> GeneratorState<(), ()> {
    debug _task_context => _4; // in scope 0 at src/anchored.rs:27:21: 27:23
    let mut _0: std::ops::GeneratorState<(), ()>; // return place in scope 0 at src/anchored.rs:27:21: 27:23
    let mut _3: (); // in scope 0 at src/anchored.rs:27:21: 27:23
    let mut _4: std::future::ResumeTy; // in scope 0 at src/anchored.rs:27:21: 27:23
    let mut _5: u32; // in scope 0 at src/anchored.rs:27:21: 27:23

    bb0: {
        _5 = discriminant((*(_1.0: &mut [static generator@src/anchored.rs:27:21: 27:23]))); // scope 0 at src/anchored.rs:27:21: 27:23
        switchInt(move _5) -> [0_u32: bb1, 1_u32: bb2, otherwise: bb3]; // scope 0 at src/anchored.rs:27:21: 27:23
    }

    bb1: {
        _4 = move _2; // scope 0 at src/anchored.rs:27:21: 27:23
        _3 = const (); // scope 0 at src/anchored.rs:27:21: 27:23
        ((_0 as Complete).0: ()) = move _3; // scope 0 at src/anchored.rs:27:23: 27:23
        discriminant(_0) = 1; // scope 0 at src/anchored.rs:27:23: 27:23
        discriminant((*(_1.0: &mut [static generator@src/anchored.rs:27:21: 27:23]))) = 1; // scope 0 at src/anchored.rs:27:23: 27:23
        return; // scope 0 at src/anchored.rs:27:23: 27:23
    }

    bb2: {
        assert(const false, "`async fn` resumed after completion") -> bb2; // scope 0 at src/anchored.rs:27:21: 27:23
    }

    bb3: {
        unreachable; // scope 0 at src/anchored.rs:27:21: 27:23
    }
}

這個 demo crate 中還有一些別的程式碼，不過對應上面的 MIR 的原始碼比較簡單：

async fn future_1() {}

只是一個簡單的空的非同步函式，可以看到生成的 MIR 會膨脹很多，如果內容稍微多一點的話通過文字形式不太好看。我們可以指定一下生成的 MIR 的格式，然後將它視覺化。

步驟大概如下：

RUSTFLAGS="--emit mir -Z dump-mir=F -Z dump-mir-dataflow -Z unpretty=mir-cfg" cargo build > mir.dot
dot -T svg -o mir.svg mir.dot

能夠在當前目錄下找到 mir.svg，開啟之後可以看到一個像流程圖的東西（另一幅差不多的圖省略掉了，有興趣的可以嘗試通過上面的方法自己生成一份）。

這裡將 MIR 按照基本單位 basic block (bb) 組織，原本的資訊都在，並且將各個 basic block 之間的跳轉關係畫了出來。從上面的圖中我們可以看到四個 basic blocks，其中一個是起點，另外三個是終點。首先起點的 bb0 switch（match in rust）了一個變數 _5，按照不同的值分支到不同的 blocks。能大概想象一下這樣的程式碼：

match _5 {
  0: jump(bb1),
    1: jump(bb2),
    _ => unreachable()
}

而 generator 的 state 可以當成就是那個 _5，不同的值就是這個 generator 的各個狀態。future_1 的狀態寫出來大概是這樣

enum Future1State {
    Start,
    Finished,
}

如果是 §1 中的 async fn foo()，可能還會多一個列舉值來表示那一次 yield。此時再想之前的問題，就能夠很自然地想到要跨越 generator 不同階段的變數需要如何儲存了。

enum FooState {
    Start,
    Yield(usize),
    Finished,
}

PART. 3 generator captured

讓我們把儲存在 generator state 中，能夠跨越 .await/yield 被後續階段使用的變數稱為被捕獲的變數。那麼能不能知道到底哪些變數實際上被捕獲了呢？讓我們試一試，首先寫一個稍微複雜一點的非同步函式：

async fn complex() {
    let x = 0;
    future_1().await;
    let y = 1;
    future_1().await;
    println!("{}, {}", x, y);
}

生成的 MIR 及 svg 比較複雜，擷取了一段放在了附錄中，可以嘗試自己生成一份完整的內容。

稍微瀏覽一下生成的內容，我們可以看到一個很長的型別總是出現，像是這樣子的東西：

[static generator@src/anchored.rs:27:20: 33:2]
// or
(((*(_1.0: &mut [static generator@src/anchored.rs:27:20: 33:2])) as variant#3).0: i32)

對照我們程式碼的位置可以發現這個型別中所帶的兩個檔案位置就是我們非同步函式 complex()的首尾兩個大括號，這個型別是一個跟我們這整個非同步函式相關的型別。

通過更進一步的探索我們大概能猜一下，上面程式碼片段中第一行的是一個實現了 Generator trait[7] 的匿名型別（struct），而 "as variant#3" 是 MIR 中的一個操作，Projection 的 Projection::Downcast，大概在這裡[8]生成。在這個 downcast 之後所做的 projection 的到的型別是我們認識的 i32。綜合其他類似的片段我們能夠推測這個匿名型別和前面描述的 generator state 是差不多的東西，而各個 variant 是不同的狀態元組，投影這個 N 元組能夠拿到被捕獲的區域性變數。

PART. 4 anchored

知道哪些變數會被捕獲能夠幫助我們理解自己的程式碼，也能夠基於這些資訊進行一些應用。

先提一下 Rust 型別系統中特殊的一種東西 auto trait[9] 。最常見的就是 Send 和 Sync，這種 auto trait 會自動為所有的型別實現，除非顯式地用 negative impl opt-out，並且 negative impl 會傳遞，如包含了 !Send 的 Rc 結構也是 !Send 的。通過 auto trait 和 negative impl 我們控制一些結構的型別，並讓編譯器幫忙檢查。

比如 anchored[10] crate 就是提供了通過 auto trait 和 generator 捕獲機制所實現的一個小工具，它能夠阻止非同步函式中指定的變數穿過 .await 點。比較有用的一個場景就是非同步過程中關於變數內部可變性的獲取。

通常來說，我們會通過非同步鎖如 tokio::sync::Mutex 來提供變數的內部可變性；如果這個變數不會穿過 .await point 即被 generator state 捕獲，那麼 std::sync::Mutex 這種同步鎖或者 RefCell 也能使用；如果想要更高的效能，避免這兩者執行時的開銷，那也能夠考慮 UnsafeCell 或其他 unsafe 手段，但是就有一點危險了。而通過 anchored 我們可以在這種場景下控制不安全因素，實現一個安全的方法來提供內部可變性，只要將變數通過 anchored::Anchored 這個 ZST 進行標記，再給整個 async fn 帶上一個 attribute 就能夠讓編譯器幫我們確認沒有東西錯誤地被捕獲並穿越了 .await、然後導致災難性的資料競爭。

就像這樣：

#[unanchored]
async fn foo(){
    {
        let bar = Anchored::new(Bar {});
    }
    async_fn().await;
}

而這種就會導致編譯錯誤：

#[unanchored]
async fn foo(){
    let bar = Anchored::new(Bar {});
    async_fn().await;
    drop(bar);
}

對於 std 的 Mutex, Ref 和 RefMut 等常見型別，clippy 提供了兩個 lints[11] ，它們也是通過分析 generator 的型別來實現的。並且與 anchored 一樣都有一個缺點，在除了像上面那樣明確使用單獨的 block 放置變數外，都會出現 false positive 的情況[12]。因為區域性變數在其他的形式下都會被記錄下來[13]，導致資訊被汙染。

anchored 目前還缺少一些 ergonomic 的介面，attribute macro 和 ecosystem 的其他工具互動的時候也存在一點問題，歡迎感興趣的小夥伴來了解一下 https://github.com/waynexia/a...

文件：https://docs.rs/anchored/0.1....

「參考」

[1]https://blog.rust-lang.org/20...

[2]https://docs.rs/futures/0.1.2...

[3]https://github.com/rust-lang/...

[4]https://doc.rust-lang.org/std...

[5]https://blog.rust-lang.org/20...

[6]https://en.wikipedia.org/wiki...

[7]https://doc.rust-lang.org/std...

[8]https://github.com/rust-lang/...

[9]https://doc.rust-lang.org/bet...