. async/await 脱糖:状态机编译

追踪 Rust 编译器怎样按挂起点切分控制流、分析跨挂起点存活变量,并将 async 函数变换为实现了 Future 的枚举状态机.

在 Future 与 Waker 中,手写一个简单的 Double Double 组合子只需要十几行代码; 但如果一段异步逻辑包含多个顺序步骤、条件分支与循环,且后一步依赖前一步算出的局部变量,手写 Future Future 就会变成繁琐的状态维护工作. Rust 的 async async 与 .await .await 正是为消除这种样板代码而设计的语法糖: 编译器在后台将看似顺序执行的异步代码块,静态改写为一个实现了 Future Future 的枚举状态机.

1 挂起点如何切分控制流

先看一个具体的异步函数:它接收初始值 base base ,依次等待两个异步计算 step_one step_one 和 step_two step_two ,最后返回累加结果:

async fn fetch_and_add(base: u32) -> u32 {
    let tmp = base + 1;
    let a = step_one(tmp).await;
    let b = step_two(a).await;
    base + a + b
}
async fn fetch_and_add(base: u32) -> u32 {
    let tmp = base + 1;
    let a = step_one(tmp).await;
    let b = step_two(a).await;
    base + a + b
}

在这段代码中,只有两处可能发生挂起: step_one(tmp).await step_one(tmp).await 与 step_two(a).await step_two(a).await . 这两个 .await .await 表达式把整个函数的执行过程切成了三段不间断的同步代码块(加上首尾的未启动状态与已完成状态):

  1. 初始段(从入口到第一个 .await .await ):接收参数 base base ,计算 let tmp = base + 1 let tmp = base + 1 ,构造子状态机 fut1 = step_one(tmp) fut1 = step_one(tmp) ,并立即对其发起第一次 poll poll .
  2. 中间段(从第一个 .await .await 恢复到第二个 .await .await ):当 fut1 fut1 返回 Poll::Ready(a) Poll::Ready(a) 时,拿到 a a ,构造第二个子状态机 fut2 = step_two(a) fut2 = step_two(a) ,并对其发起 poll poll .
  3. 收尾段(从第二个 .await .await 恢复到函数返回):当 fut2 fut2 返回 Poll::Ready(b) Poll::Ready(b) 时,计算 base + a + b base + a + b ,返回 Poll::Ready(base + a + b) Poll::Ready(base + a + b) ,并将自身状态标记为已完成.
2 跨挂起点的存活变量分析

要把这三段代码装进一个枚举类型,编译器需要回答一个问题:每个枚举变体(状态)里需要保存哪些字段? 答案取决于变量的生命周期是否跨越了 .await .await 挂起点:

  • 变量 tmp tmp 在第一个 .await .await 之前算出,并在构造 step_one(tmp) step_one(tmp) 时立即按值消耗,第一个 .await .await 之后再也不会用到 tmp tmp . 因此, tmp tmp 只需要在执行第一次 poll poll 时暂存在工作线程的普通调用栈上,不需要存进状态机结构体里.
  • 参数 base base 从函数入口一直存活到最后一行 base + a + b base + a + b ,跨越了两个 .await .await 挂起点,因此必须在状态一与状态二中都保存 base base .
  • 变量 a a 在第一个 .await .await 结束时产生,在第二个 .await .await 结束后的 base + a + b base + a + b 中还要使用,跨越了第二个 .await .await ,因此必须在状态二中保存 a a .
  • 变量 b b 在第二个 .await .await 结束时产生,并在同一段同步代码中立即用于计算返回值,不跨越任何 .await .await ,因此也留在普通线程栈上即可.

把控制流切分与变量存活范围画在同一张图上,就能直接看出状态机每个阶段的内存需求:

3 手工还原编译器生成的状态机

按照上图的四个状态,我们可以把 fetch_and_add fetch_and_add 脱糖后的等价 Rust 代码完整写出来. 为了聚焦状态转移逻辑,这里假设 StepOneFut StepOneFut 与 StepTwoFut StepTwoFut 均满足 Unpin Unpin (含跨挂起点借用的非 Unpin Unpin 情形将在下一篇单独处理):

use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};

enum FetchAndAddFuture {
    /// 尚未被 poll 过的初始状态,只保存函数参数
    Start { base: u32 },
    /// 正在等待第一个 .await
    WaitingOne { base: u32, fut1: StepOneFut },
    /// 正在等待第二个 .await
    WaitingTwo { base: u32, a: u32, fut2: StepTwoFut },
    /// 已经返回过 Ready,不可再次 poll
    Done,
}

fn fetch_and_add(base: u32) -> FetchAndAddFuture {
    // 调用 async fn 只是构造初始状态,不执行函数体
    FetchAndAddFuture::Start { base }
}
use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};

enum FetchAndAddFuture {
    /// 尚未被 poll 过的初始状态,只保存函数参数
    Start { base: u32 },
    /// 正在等待第一个 .await
    WaitingOne { base: u32, fut1: StepOneFut },
    /// 正在等待第二个 .await
    WaitingTwo { base: u32, a: u32, fut2: StepTwoFut },
    /// 已经返回过 Ready,不可再次 poll
    Done,
}

fn fetch_and_add(base: u32) -> FetchAndAddFuture {
    // 调用 async fn 只是构造初始状态,不执行函数体
    FetchAndAddFuture::Start { base }
}

接下来看它的 Future::poll Future::poll 实现. 注意里面使用了一个 loop loop 循环:当 fut1 fut1 恰好在第一次被 poll poll 时就立刻返回 Poll::Ready(a) Poll::Ready(a) (例如命中了内存缓存,无需真正等待 I/O),状态机不需要返回到调用者手中,而是可以直接修改自身状态为 WaitingTwo WaitingTwo 并在同一次 poll poll 调用里顺着 loop loop 继续推进 fut2 fut2 :

impl Future for FetchAndAddFuture {
    type Output = u32;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<u32> {
        let this = &mut *self;
        loop {
            match this {
                FetchAndAddFuture::Start { base } => {
                    let base = *base;
                    // tmp 只在当前线程栈上存活,不进入枚举字段
                    let tmp = base + 1;
                    let fut1 = step_one(tmp);
                    *this = FetchAndAddFuture::WaitingOne { base, fut1 };
                }
                FetchAndAddFuture::WaitingOne { base, fut1 } => {
                    match Pin::new(fut1).poll(cx) {
                        Poll::Ready(a) => {
                            let base = *base;
                            let fut2 = step_two(a);
                            // 转移到下一状态,通过外层 loop 立即尝试 poll fut2
                            *this = FetchAndAddFuture::WaitingTwo { base, a, fut2 };
                        }
                        Poll::Pending => return Poll::Pending,
                    }
                }
                FetchAndAddFuture::WaitingTwo { base, a, fut2 } => {
                    match Pin::new(fut2).poll(cx) {
                        Poll::Ready(b) => {
                            let result = *base + *a + b;
                            *this = FetchAndAddFuture::Done;
                            return Poll::Ready(result);
                        }
                        Poll::Pending => return Poll::Pending,
                    }
                }
                FetchAndAddFuture::Done => {
                    panic!("async fn polled after completion");
                }
            }
        }
    }
}
impl Future for FetchAndAddFuture {
    type Output = u32;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<u32> {
        let this = &mut *self;
        loop {
            match this {
                FetchAndAddFuture::Start { base } => {
                    let base = *base;
                    // tmp 只在当前线程栈上存活,不进入枚举字段
                    let tmp = base + 1;
                    let fut1 = step_one(tmp);
                    *this = FetchAndAddFuture::WaitingOne { base, fut1 };
                }
                FetchAndAddFuture::WaitingOne { base, fut1 } => {
                    match Pin::new(fut1).poll(cx) {
                        Poll::Ready(a) => {
                            let base = *base;
                            let fut2 = step_two(a);
                            // 转移到下一状态,通过外层 loop 立即尝试 poll fut2
                            *this = FetchAndAddFuture::WaitingTwo { base, a, fut2 };
                        }
                        Poll::Pending => return Poll::Pending,
                    }
                }
                FetchAndAddFuture::WaitingTwo { base, a, fut2 } => {
                    match Pin::new(fut2).poll(cx) {
                        Poll::Ready(b) => {
                            let result = *base + *a + b;
                            *this = FetchAndAddFuture::Done;
                            return Poll::Ready(result);
                        }
                        Poll::Pending => return Poll::Pending,
                    }
                }
                FetchAndAddFuture::Done => {
                    panic!("async fn polled after completion");
                }
            }
        }
    }
}

对照最后那个 FetchAndAddFuture::Done FetchAndAddFuture::Done 分支: 这正是我们在 Future 与 Waker 中提到的“契约第一条”的代码体现——一旦状态机吐出 Poll::Ready(result) Poll::Ready(result) ,它就会把自身状态覆写为 Done Done ;若外部调用者违反协议再次 poll poll 它,状态机就会直接触发 panic! panic! .

4 状态机的内存布局与两项工程推论

观察 FetchAndAddFuture FetchAndAddFuture 这个 enum enum 的定义,Rust 的枚举采用带标签联合体布局: 整个状态机在内存中的大小,等于状态标签( discriminant discriminant )加上最大那个状态变体的大小; 不同状态变体(例如 WaitingOne WaitingOne 里的 fut1 fut1 与 WaitingTwo WaitingTwo 里的 fut2 fut2 )在物理内存上复用同一块字节区域.

只要记住“异步函数是捕获跨 .await .await 变量的枚举状态机”,日常开发中两个常见的现象就容易理解了:

4.1 为什么跨 .await .await 的大数组会撑爆状态机大小

如果在异步函数中声明了一个 let buf = [0u8; 65536]; let buf = [0u8; 65536]; ,并且这个 buf buf 的作用域跨越了某个 .await .await 点, 编译器就必须把这个 64 KB 的数组作为字段直接塞进状态机枚举的变体中. 当外层异步函数层层内联嵌套时,顶层 Future Future 的体积会迅速膨胀到数百 KB,甚至在将其从栈上移动到 tokio::spawn tokio::spawn 的瞬间引发线程栈溢出. 反之,如果用代码块 { let buf = ...; } { let buf = ...; } 把大变量的作用域限制在两个 .await .await 之间(使其在进入下一个 .await .await 前就析构),或者将其放入堆上的 Vec<u8> Vec<u8> / Box<[u8]> Box<[u8]> ,状态机里就只需要保存几个机器字长的胖指针.

4.2 为什么跨 .await .await 持有 MutexGuard MutexGuard 会破坏 Send Send 或引发死锁

在多线程工作窃取执行器(如 tokio::spawn tokio::spawn )中,一个任务在挂起前可能运行在线程 1 上,被唤醒后则可能被线程 2 偷走继续执行 poll poll . 因此, tokio::spawn tokio::spawn 要求传入的顶层 Future Future 满足 Send Send (可以安全地跨线程转移所有权).

编译器怎样判断 FetchAndAddFuture FetchAndAddFuture 是否实现 Send Send ? 规则直接按结构归纳:当且仅当状态机中所有状态变体保存的字段类型全都满足 Send Send 时,整个状态机才满足 Send Send .

标准库的同步互斥锁守卫 std::sync::MutexGuard<'_, T> std::sync::MutexGuard<'_, T> 专门没有实现 Send Send (因为 POSIX 线程库等底层操作系统原语要求加锁与解锁必须在同一个操作系统线程上完成). 如果你写出如下代码:

async fn bad_lock(mutex: &std::sync::Mutex<u32>) {
    let mut guard = mutex.lock().unwrap();
    *guard += 1;
    some_io().await; // guard 的作用域若延续到这里,就会进入状态机字段!
}
async fn bad_lock(mutex: &std::sync::Mutex<u32>) {
    let mut guard = mutex.lock().unwrap();
    *guard += 1;
    some_io().await; // guard 的作用域若延续到这里,就会进入状态机字段!
}

只要 guard guard 跨过了 some_io().await some_io().await ,编译器就必须把 guard: MutexGuard<'_, u32> guard: MutexGuard<'_, u32> 存进 WaitingIo WaitingIo 状态变体中,从而直接传染整个状态机,使其失去 Send Send ,导致外层的 tokio::spawn(bad_lock(...)) tokio::spawn(bad_lock(...)) 编译报错. 即使不报编译错误,在持锁期间 .await .await 挂起也会让当前工作线程在锁未释放的情况下切去运行别的任务;若刚好切到的另一个任务也要获取同一把同步锁,整个工作线程就会当场死锁. 因此,同步锁的守卫应当在 .await .await 之前显式 drop(guard) drop(guard) (或用局部代码块括起来),使其完全停留在当前那次 poll poll 的线程栈上,绝不进入状态机字段.

5 练习:状态机大小与字段分析
Exercises (2)

Exercise 3.1. 判断哪些局部变量会进入状态机

分析下面这段异步函数中, a a 、 x x 、 y y 、 z z 四个局部变量有哪些会被编译器放进挂起状态的枚举变体字段中:

async fn compute(input: u64) -> u64 {
    let x = {
        let a = [0u64; 1024];
        input + a[0]
    };
    let y = work().await;
    let z = y * 2;
    x + y + z
}
async fn compute(input: u64) -> u64 {
    let x = {
        let a = [0u64; 1024];
        input + a[0]
    };
    let y = work().await;
    let z = y * 2;
    x + y + z
}
Hint画出唯一那个 .await .await 的位置,逐一检查每个变量是在 .await .await 之前最后一次使用、之后才定义,还是在 .await .await 前定义且在 .await .await 后仍被读取或析构.
Solution

唯一需要保存在 Waiting Waiting 状态变体中的局部变量是 x x (连同正在被等待的子状态机 work().await work().await ).

  • a a :在局部代码块 { ... } { ... } 内定义并在 .await .await 之前就已经离开作用域销毁,只在第一次 poll poll 的调用栈上存在,不进入状态机.
  • x x :在 .await .await 之前算出,在 .await .await 之后的 x + y + z x + y + z 中仍要读取,跨越了挂起点,必须进入状态机.
  • y y 与 z z :都在 .await .await 完成之后才定义并在返回前用完,期间不再经过任何挂起点,只在最后一次 poll poll 的调用栈上存在,不进入状态机.

Exercise 3.2. 为什么直接递归的 async fn async fn 无法确定类型大小

用本节介绍的状态机内存布局解释:为什么不加 Box::pin Box::pin 的直接递归 async fn async fn 会触发编译器的“无限大小类型(recursive async fn async fn )”错误?
Hint回想 Term Term 语法树为什么不能直接写 Succ(Term) Succ(Term) 而必须写 Succ(Box<Term>) Succ(Box<Term>) ,再对照 async fn async fn 脱糖后的状态机字段.
Solution

假设 async fn countdown(n: u32) async fn countdown(n: u32) 直接在函数体内调用 countdown(n - 1).await countdown(n - 1).await . 编译器将其脱糖为状态机类型 CountdownFuture CountdownFuture 时,在挂起等待递归调用的状态变体里,必须内联保存一个被等待的子状态机字段 fut: CountdownFuture fut: CountdownFuture . 这导致 CountdownFuture CountdownFuture 的定义内部直接包含了未加指针间接层的 CountdownFuture CountdownFuture 自身,其类型大小满足方程 𝑆= tag +4+𝑆,在编译期无穷大.

解决方法与递归枚举相同:在递归挂起点处引入堆分配间接层(如 Box::pin(countdown(n - 1)).await Box::pin(countdown(n - 1)).await ),使状态机字段里保存的只是一个固定大小的 Pin<Box<...>> Pin<Box<...>> 指针. (注:自 Rust 1.77 起,只要递归调用点明确写出了 Box::pin(...) Box::pin(...) 等指针包装,编译器已支持直接编写递归 async fn async fn .)