Rust async code from the generator&#39;s point of view

Text｜Ruihang Xia

Currently participating in the edge time series data storage engine project

This article is 6992 words read 18 minutes

foreword

Asynchronous programming in Rust, one of the more important features of the 2018 edition, is now widely used. When using it, you will inevitably be curious about how it works. This article attempts to explore from the aspects of generator and variable capture, and then introduces a scenario encountered during the development of the embedded time series storage engine ceresdb-helix.

Due to the author's level of content, there are inevitably some errors and omissions, please leave a message to inform.

PART. 1 async/.await, coroutine and generator

The async/.await syntax entered the stable channel in version 1.39 [1], which makes it easy to write asynchronous code:

、、、java
async fn asynchronous() {

// snipped

}

async fn foo() {

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

、、、

In the above example, the local variable x can be used directly after an asynchronous process (fn asynchoronous), just like writing synchronous code. Prior to this, asynchronous code was generally used through combinators in the form of futures 0.1[2], and local variables that were intended to be used by subsequent asynchronous procedures (such as and_then()) needed to be explicitly manually closed. The way the package enters and leaves the parameters is chained, and the experience is not particularly good.

What async/.await does is actually transform the code into a generator/coroutine[3] form to execute. A coroutine process can be suspended, do something else and then resume execution, which is currently used like .await. Taking the above code as an example, another asynchronous process asynchronous() is called in the asynchronous process foo(), and the execution of the current process is suspended when the .await in the seventh line is resumed when it can continue to execute.

And resume execution may need some previous information, such as in foo() we use the previous information x in the eighth line. That is to say, async procedures have the ability to save some internal local state, so that they can continue to be used after .await. In other words save local variables in generator state that may be used after yield. The pin[4] mechanism needs to be introduced here to solve the possible self-reference problem, which will not be repeated in this part.

PART. 2 visualize generator via MIR

We can see what the generator mentioned above looks like through MIR[5]. MIR is an intermediate representation in Rust, based on the control flow graph CFG [6] representation. CFG can provide a more intuitive representation of what the program execution looks like, and MIR can be helpful when it is sometimes unclear what your Rust code actually looks like.

There are several ways to get the MIR representation of your code. If you have a working Rust toolchain at hand, you can pass an environment variable to rustc like this and build with cargo to generate the MIR:

RUSTFLAGS="--emit mir" cargo build

If the build is successful, a .mir file will be generated in the target/debug/deps/ directory. Or you can get MIR via https://play.rust-lang.org/ , just select MIR on the overflow menu next to Run.

The MIR generated by the toolchain of nightly in 2021-08 is probably like this. There are many things that you don't know that can be ignored. You should know about it.

_0, _1 these are variables
There are many syntaxes similar to Rust, such as type annotations, function definitions and calls and annotations.

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
    }
}

There are some other codes in this demo crate, but the source code corresponding to the above MIR is relatively simple:

async fn future_1() {}

Just a simple empty asynchronous function, you can see that the generated MIR will be inflated a lot, and if the content is a little more, it will not look good in text form. We can specify the format of the generated MIR and visualize it.

The steps are roughly as follows:

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

You can find mir.svg in the current directory. After opening it, you can see something like a flowchart (another similar picture is omitted. If you are interested, you can try to generate a copy yourself through the above method).

Here, the MIR is organized according to the basic unit basic block (bb), the original information is all there, and the jump relationship between each basic block is drawn. From the picture above we can see four basic blocks, one of which is the starting point and the other three are the ending point. First, the bb0 switch at the starting point (match in rust) has a variable _5, branching to different blocks according to different values. You can roughly imagine code like this:

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

The state of the generator can be regarded as the achievement _5, and the different values are the states of the generator. The status of future_1 is written like this

enum Future1State {
    Start,
    Finished,
}

In the case of async fn foo() in §1, there may be an additional enumeration value to indicate that yield. If you think about the previous question at this point, you can naturally think of how to save the variables that span different stages of the generator.

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

PART. 3 generator captured

Let's keep variables in generator state that can be used by subsequent stages across .await/yield called captured variables. So is it possible to know which variables are actually captured? Let's try it out, first by writing a slightly more complex async function:

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

The generated MIR and svg are more complicated, so I cut a section and put it in the appendix. You can try to generate a complete content by yourself.

Browsing the generated content a little, we can see that a long type is always present, something like this:

[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)

Comparing the location of our code, we can find that the two file locations in this type are the two curly brackets at the beginning and end of our asynchronous function complex(). This type is a type related to our entire asynchronous function.

Through further exploration, we can probably guess that the first line in the above code snippet is an anonymous type (struct) that implements the Generator trait[7], and "as variant#3" is an operation in MIR, Projection Projection::Downcast, probably generated here [8]. The type of projection made after this downcast is i32 as we know it. Combining other similar fragments, we can speculate that this anonymous type is similar to the generator state described above, and each variant is a different state tuple. Projecting this N-tuple can get the captured local variables.

PART. 4 anchored

Knowing which variables are captured can help us understand our code and make some applications based on this information.

Let's start by mentioning a special thing in Rust's type system, the auto trait[9] . The most common are Send and Sync, this auto trait is automatically implemented for all types unless explicitly opt-out with negative impl, and negative impl will be passed, such as the Rc structure containing !Send is also !Send. With the auto trait and negative impl we control the type of some structs and let the compiler check for help.

For example, the anchored[10] crate provides a small tool implemented through the auto trait and generator capture mechanism, which prevents the variables specified in the asynchronous function from passing through the .await point. A useful scenario is the acquisition of variable internal variability in an asynchronous process.

Usually, we provide internal mutability of variables through asynchronous locks such as tokio::sync::Mutex; if the variable does not pass through the .await point and is captured by the generator state, then std::sync::Mutex this A synchronization lock or RefCell can also be used; if you want higher performance and avoid the runtime overhead of both, you can also consider UnsafeCell or other unsafe means, but it is a little dangerous. With anchored, we can control the unsafe factors in this scenario and implement a safe method to provide internal variability. As long as the variable is marked with the ZST anchored::Anchored, and an attribute is added to the entire async fn. Let the compiler help us make sure that nothing is caught by mistake and traversing .await, leading to a catastrophic data race.

like this:

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

And this will result in a compilation error:

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

For common types such as Mutex, Ref and RefMut of std, clippy provides two lints[11], which are also realized by analyzing the type of generator. And it has the same disadvantage as anchored. In addition to explicitly using a separate block to place variables as above, there will be false positives [12]. Because local variables are recorded in other forms [13], information is polluted.

Anchored still lacks some ergonomic interfaces. There are also some problems when the attribute macro interacts with other tools of the ecosystem. Interested partners are welcome to learn about https://github.com/waynexia/anchored

Documentation: https://docs.rs/anchored/0.1.0/anchored/

"refer to"

[1]https://blog.rust-lang.org/2019/11/07/Async-await-stable.html

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

[3]https://github.com/rust-lang/rfcs/blob/master/text/2033-experimental-coroutines.md

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

[5]https://blog.rust-lang.org/2016/04/19/MIR.html

[6]https://en.wikipedia.org/wiki/Control-flow_graph

[7]https://doc.rust-lang.org/std/ops/trait.Generator.html

[8]https://github.com/rust-lang/rust/blob/b834c4c1bad7521af47f38f44a4048be0a1fe2ee/compiler/rustc_middle/src/mir/mod.rs#L1915

[9]https://doc.rust-lang.org/beta/unstable-book/language-features/auto-traits.html

[10]https://crates.io/crates/anchored