future explained(0)

Future explained(0) – state machine 原理

generator(生成器)是 Rust 中 async/await 的基础,state machine 又是生成器的基础。

和闭包类似,不同的编译器会给 generator 生成不同的语义。但有一点是约定俗成的,也就是生成器最大的特点:程序的执行流程可以在生成器和调用者之间来回切换。

当我们需要暂时从生成器中返回的时候,可以使用 yield 关键字,当调用者希望再次进入生成器的时候,就调用 resume 方法,这时程序就从上次 yield 返回的那个点之后继续执行,还能获取到上下文信息,就好像能够按下了暂停按钮,又按下了继续按钮,程序执行完全没有感知。

其实不仅仅是 Rust,LLVM 作为一个广泛使用的后端,也支持了 yield 语法(因此 c++ 最新版本也可以实现 stackless 协程了)。接下来我们就分别探究一下二者实现的思路,以求对于 async/await 有一个更加深入的认识。

1. LLVM 实现

首先来看 LLVM 支持的方式:通过 generator + yield 两个关键字,可以将原先的一个普通函数,改写为生成器。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
#include <stdio.h>

generator int foo()
{
yield 1;
yield 2;
yield 3;
}

int main()
{
foreach (int i in foo())
printf("Value: %d\n", i);

return 0;
}

我们希望的执行流是:

  1. 下一次重新执行该 generator,可以从上一次 yield 点开始执行
  2. 可以保存栈的上下文

可以想象,我们需要保存很多状态,比如必须要 yield 的位置,这样才能在 resume 的时候执行正确位置的代码。Where? What?

  1. 我们肯定不会将状态保存在函数栈帧中,因为一旦函数返回,该栈帧中的所有局部变量都失效了。也就是说,必须要将上下文(context)保存在函数外的某个结构体内(static 局部变量也不行,因为如果重复调用就炸了)。

  2. 接下来考虑,需要保存什么样的状态,才能保证对于执行流和上下文没有感知?

    • 为了从 yield 下一条指令开始执行,肯定需要保存该指令的地址;
    • 为了保证栈帧中的局部变量不会失效,一定要将其保存存在调用栈之外的某个结构中(其实就是闭包)

实际上 LLVM 正是这样做的,通过 context 保存下一次 resume 需要跳转执行的位置,以及调用栈中的局部变量。每次调用这个 generator,都需要从外部传入相应的 context 上下文变量,将某些局部变量的读写都映射为针对 context 中相应字段的读写。

1
2
Tips:
上下文变量还需要一个特殊的 setup 函数来处理,每一个 generator 都唯一对应 context,从而保证状态的确定性。

来简单看一下 LLVM 生成的 IR,围绕着 context 这个最关键的Generator 状态展开,关键位置都做了标记。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
// 上下文(状态)
%foo_context = type {
i8*, ; 0: block (pc 指针位置)
i32 ; 1: value (result)
}

// foo 初始化,必须在调用 foo 之前设置
define void @foo_setup(%foo_context* %context) nounwind {
; set up 'block'
%1 = getelementptr %foo_context* %context, i32 0, i32 0
store i8* blockaddress(@foo_yield, %.yield1), i8** %1

ret void
}

; The boolean returned indicates if a result was available or not.
; Once no more results are available, the caller is expected to not call
; the iterator again.
define i1 @foo_yield(%foo_context* %context) nounwind {
; 将 context.state 中的值作为下标,跳转到对应的位置
%1 = getelementptr %foo_context* %context, i32 0, i32 0
%2 = load i8** %1
indirectbr i8* %2, [ label %.yield1, label %.yield2, label %.yield3, label %.done ]

.yield1:
; 将返回值保存在 foo_context.result 中
%3 = getelementptr %foo_context* %context, i32 0, i32 1
store i32 1, i32* %3

; 将 'block' 块的 index 保存在 foo_context.state 中
; make 'block' point to next block to execute
%4 = getelementptr %foo_context* %context, i32 0, i32 0
store i8* blockaddress(@foo_yield, %.yield2), i8** %4

; 返回
ret i1 1

.yield2:
; 将返回值保存在 foo_context.result 中
%5 = getelementptr %foo_context* %context, i32 0, i32 1
store i32 2, i32* %5

; 将 'block' 块的 index 保存在 foo_context.state 中
%6 = getelementptr %foo_context* %context, i32 0, i32 0
store i8* blockaddress(@foo_yield, %.yield3), i8** %6

; 返回
ret i1 1

.yield3:
; 将返回值保存在 foo_context.result 中
%7 = getelementptr %foo_context* %context, i32 0, i32 1
store i32 3, i32* %7

; 将 'block' 块的 index 保存在 foo_context.state 中
%8 = getelementptr %foo_context* %context, i32 0, i32 0
store i8* blockaddress(@foo_yield, %.done), i8** %8

; 返回
ret i1 1

; 结束,返回默认值 0
.done:
ret i1 0
}

define void @main() nounwind {
; 为 generator 分配 context 空间
%context = alloca %foo_context
call void @foo_setup(%foo_context* %context)
br label %.head

.head:
; 每次调用,都需要传入 context 参数
; foreach (int i in foo())
%1 = call i1 @foo_yield(%foo_context* %context)

; 省略 ...

.tail:
ret void
}

上面的 generator 中,实际上没有保存什么局部变量,但是如果 generator 中存在跨越 yield 的局部变量,就需要将这些内容保存到 context 的某个字段中,这里不过多演示。

2. Rust 实现

虽然 Rust 底层依赖于 LLVM,LLVM 也提供了 generator 的能力,但 Rust 单独实现了一套类似的 generator。实际上 generator 也不算太难,简单来说,等同于一个闭包,只不过这个闭包不是一次性执行完成,而是分为多个阶段,每一个阶段都会对应一个状态,状态之间会进行转移。That’s all!

下面主要会借助 Generator 和 Yield,探讨一下目前 async/await 实现的底层原理。

2.1 Yield & Generator

首先来看一个 yield 关键字的 demo。

由于迭代器是一个实验性质的特性,因此要打开 generators,generator_trait 这两个 feature。 在我们的代码中,可以通过 yield + 闭包来声明 Generator。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
#![feature(generators, generator_trait)]

use std::ops::{Generator, GeneratorState};
use std::pin::Pin;

fn main() {
let mut g = || {
yield 1;
return "foo"
};

assert_eq!(Pin::new(&mut g).resume(()), GeneratorState::Yielded(1));
assert_eq!(Pin::new(&mut g).resume(()), GeneratorState::Complete("foo"));
}

之所以用闭包实现,按照 LLVM 的实现机制以及上面的一小段简介,不难推断出原因:没错,需要保存上下文状态信息!那么,现在最重要的问题就是,这个上下文究竟是怎么保存的?编译器在背后究竟做了什么黑魔法

就像 Future 的核心是 Future traitGenerator 的核心是 Generator trait,我们先来看一下这几个 trait 的定义。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
#[derive(Clone, Copy, PartialEq, PartialOrd, Eq, Ord, Debug, Hash)]
#[lang = "generator_state"]
pub enum GeneratorState<Y, R> {
Yielded(Y),
Complete(R),
}

#[lang = "generator"]
#[fundamental]
pub trait Generator<R = ()> {
type Yield;
type Return;
fn resume(self: Pin<&mut Self>, arg: R) -> GeneratorState<Self::Yield, Self::Return>;
}

简单和 Future trait 进行一下对比:

  1. Future只有一个类型参数 Output,表示最终返回的结果,Poll 其实就相当于 Option;Generator 有两个类型参数 Yield 和 Return,可以在返回零次或多次 Yield 后最终返回Return,GeneratorState 其实就相当于 Result;
  2. Future 的 poll 中有一个额外参数 ctx,一般是用来注册回调的;Generator 的 resume 中有一个额外参数 arg,也就是每次进入的时候都可以传一个参数,默认为 (),我理解是赋予了这个 Generator 扩展的能力。

Generator trait 在 99% 的情况下,不需要我们手动实现,而是编译器帮助我们生成:

参照 LLVM 的实现机制,对程序做控制流分析,将整个 generator 视为一个状态机,根据 yield, return 这两个关键字划分为多个子代码块。在 LLVM 中,实际上是通过一个指针保存了每次 resume 需要跳转的位置,但有一个更加优雅的方式:enum + match

用一个 enum 保存每一个代码块执行中的状态,每次 yield 或者 return 之前,会将这段子代码块中的所有指令全部执行完成,然后修改 enum 中保存的状态,也就是状态转移,以及上下文中的局部变量,也都被捕获进来。

每一次 resume,其实就是需要匹配 enum 中的状态,去执行对应部分的代码,最后修改状态,是状态转移,直到返回 Complete。

按照这种思路,可以将上面 main 函数的闭包,翻译为下面的结构。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
#![feature(arbitrary_self_types, generators, generator_trait)]

use std::ops::{Generator, GeneratorState};
use std::pin::Pin;

fn main() {
let ret = "foo";
let mut generator = {
enum __Generator {
Start(&'static str),
Yield1(&'static str),
Done,
}

impl Generator for __Generator {
type Yield = i32;
type Return = &'static str;

fn resume(mut self: Pin<&mut Self>, resume: ()) -> GeneratorState<i32, &'static str> {
use std::mem;
match mem::replace(&mut *self, __Generator::Done) {
__Generator::Start(s) => {
*self = __Generator::Yield1(s);
GeneratorState::Yielded(1)
}

__Generator::Yield1(s) => {
*self = __Generator::Done;
GeneratorState::Complete(s)
}

__Generator::Done => {
panic!("generator resumed after completion")
}
}
}
}

__Generator::Start(ret)
};

Pin::new(&mut generator).resume(());
Pin::new(&mut generator).resume(());
}

enum __Generator 就是生成的匿名结构,编译器会为其实现 Generate trait,就是上面分析所说的状态机核心逻辑。

总的来说,就是通过 闭包 + 状态机,实现了整个 Generator。

有了 Generator 的基础,再去看 async/await 的实现,就会轻松不少。我们只需要去关心,究竟是如何将 async/await 转换成闭包,并且用 状态匹配 + 状态转移 的方式去驱动执行。

1
2
Tips:
需要提示一点,generator 在 Rust 中仍然是实验性的功能,只不过目前的 async/await 底层是借助了这一点来实现,未来不保证会不会修 改为别的方案,因为自引用结构的确有一些 annoying

2.2 HIR & MIR

Rust 针对 async/await 的处理,主要分为 HIR 和 MIR 两个阶段。

1
2
3
4
5
6
7
8
Tips:
Rust 面向 LLVM 编程,但是 rustc 编译器也做了非常多的优化,将其划分为 HIR,MIR

1. HIR. 全称是 "high-level (H) intermediate representation (IR)",可以将其理解为 AST 的另一中表示方式,主要是将语法糖 做了转换。

2. MIR. The "mid-level (M) intermediate representation (IR)"。HIR 的简化版本,将高级的表达形式,转换成更加低级的结 构。rust 中的生命周期在这里会被抹除。

3. LIR. Rust 中相当于 LLVM-IR。"low-level (L) intermediate representation (IR)" 是一种非常接近机器码的表示方式,会 对 MIR 做进一步的精简。

2.2.1 HIR

在 HIR 中,其实是针对 async fn test(f: F) where F:Future { f.await } 做了脱糖处理。

比如,在 HIR 的表达式解析模块中,会将 async 代码块构造为一个实现了 future trait 的 generator:

1
2
3
std::future::from_generator(static move? |_task_context| -> <ret_ty> {
<body>
})

会将 <expr>.await desugar 变为一个 loop + match。如果返回 Pending,下一句将会执行 yield

1
2
3
4
5
6
7
8
9
10
11
12
match <expr> {
mut pinned => loop {
match unsafe { ::std::future::Future::poll(
<::std::pin::Pin>::new_unchecked(&mut pinned),
::std::future::get_context(task_context),
) } {
::std::task::Poll::Ready(result) => break result,
::std::task::Poll::Pending => {}
}
task_context = yield ();
}
}

也就是说,会将 async 代码块翻译为 from_generator 生成的匿名结构体,这个结构体由编译器实现 Future trait。

你可能会好奇,GeneratorState 和 Poll 并不兼容啊,前者返回 GeneratorState,后者返回 Poll,这是为什么呢?别担心,在 GeFuture::poll 函数中,特别做了处理:

GenFuture 包装了一个 Generator 的时候,GenFuture::poll 其实会调用内部的 generator 的 resume 方法,返回的 Yield 对应 Pending,Complete 对应 Ready。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
struct GenFuture<T: Generator<Yield = ()>>(T);

impl<T: Generator<Yield = ()>> !Unpin for GenFuture<T> {}

#[unstable(feature = "gen_future", issue = "50547")]
impl<T: Generator<Yield = ()>> Future for GenFuture<T> {
type Output = T::Return;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
// Safe because we're !Unpin + !Drop mapping to a ?Unpin value
let gen = unsafe { Pin::map_unchecked_mut(self, |s| &mut s.0) };
set_task_context(cx, || match gen.resume() {
GeneratorState::Yielded(()) => Poll::Pending,
GeneratorState::Complete(x) => Poll::Ready(x),
})
}
}

pub fn from_generator<T: Generator<Yield = ()>>(x: T) -> impl Future<Output = T::Return> {
GenFuture(x)
}

这一部分源码在 rust expr codeGenFuture code 中可以找到,

到这一步为止,HIR 对于 async 代码块的处理就完成了,接下来是最重要的 MIR 处理。

2.2.2 MIR

MIR 中有一个 transform 模块专门做了这件事情。

其实主要的处理逻辑我们已经看过了,就是将 yield 转换成匿名 enum,并插入多个 yield point即可。这里并不过多展开。


本博客所有文章除特别声明外,均采用 CC BY-SA 4.0 协议 ,转载请注明出处!