. Future 与 Waker:基于轮询的异步契约

剖析标准库中 Future、Poll、Context 与 Waker 的类型定义、轮询协议不变量,以及 RawWaker 手工虚表的零成本擦除机制.

在 为什么需要异步 中我们看到,Rust 用“外部主动拉取状态机 + 事件源回调通知重排”代替了层层嵌套的回调闭包. 这套机制在标准库 core::future core::future 与 core::task core::task 中只占极小的篇幅: 标准库既不包含事件循环,也不包含网络或定时器实现,它只定义了驱动者(执行器)与被驱动者( Future Future )之间的通信契约.

1 Future Future 与 Poll Poll 的定义

打开标准库 core::future::Future core::future::Future 与 core::task::Poll core::task::Poll ,核心定义只有十几行:

pub trait Future {
    type Output;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output>;
}

pub enum Poll<T> {
    Ready(T),
    Pending,
}
pub trait Future {
    type Output;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output>;
}

pub enum Poll<T> {
    Ready(T),
    Pending,
}

这里暂时把第一个参数 self: Pin<&mut Self> self: Pin<&mut Self> 当作普通的可变借用 &mut Self &mut Self 来读(关于为什么必须用 Pin Pin 包装,将在 Pin 与自引用结构 中专门拆解). 除去 Pin Pin 之后, poll poll 的签名表达了一件非常朴素的事: 调用者提供可变借用与任务上下文 cx cx ,让状态机尝试往前推进; 状态机要么已经算完并返回 Poll::Ready(val) Poll::Ready(val) ,要么因为等待某个条件而无法继续,返回 Poll::Pending Poll::Pending .

为什么不能省掉 cx: &mut Context<'_> cx: &mut Context<'_> ,只让执行器在循环里一遍遍调用 poll() poll() ? 如果执行器不知道一个返回了 Pending Pending 的任务什么时候才会就绪,它就只能采取忙等待,在死循环里不停地询问每一个任务“好了吗?”. 这会把 CPU 核心跑满到 100%,退化成低效的自旋.

为了避免忙等待, poll poll 协议规定了一条双向契约:

  1. 调用者的承诺:当 poll poll 返回 Poll::Ready(val) Poll::Ready(val) 后,调用者认为该 Future Future 已完成,不得再次对它调用 poll poll (除非该类型额外实现了 FusedFuture FusedFuture 并显式允许).
  2. 实现者的承诺:当 poll poll 返回 Poll::Pending Poll::Pending 时,该 Future Future 必须确保已经将 cx.waker() cx.waker() (或其克隆)注册到了某个未来会触发的事件源上;当阻挡任务前进的条件解除时,事件源有责任调用该 Waker Waker 的 wake() wake() 方法,通知执行器“现在值得再 poll poll 我一次了”.

如果一个叶子 Future Future 既返回了 Poll::Pending Poll::Pending ,又没有在任何地方保存或触发 cx.waker() cx.waker() ,执行器就会把它移出就绪队列并永久挂起——任务会在没有任何报错的情况下“静默死锁”.

2 手写一个叶子 Future Future 与组合子

为了看清 poll poll 是如何一层层传导的,我们不借助 async async / await await 语法糖,先手动实现两个 Future Future : 一个是模拟等待倒计时的叶子状态机 Countdown Countdown ,另一个是将结果乘二的外层组合子 Double<F> Double<F> .

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

/// 每次被 poll 时减一,减到 0 时就绪.
pub struct Countdown {
    pub remaining: u32,
}

impl Future for Countdown {
    type Output = u32;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<u32> {
        if self.remaining == 0 {
            Poll::Ready(21)
        } else {
            self.remaining -= 1;
            // 关键:返回 Pending 前必须安排唤醒!
            // 这里演示立即通知执行器下一轮继续 poll 自己;
            // 真实的异步 I/O 会把 cx.waker().clone() 存进反应器.
            cx.waker().wake_by_ref();
            Poll::Pending
        }
    }
}

/// 包装另一个 Future,在其就绪后将输出乘 2.
pub struct Double<F> {
    pub inner: F,
}

impl<F: Future<Output = u32> + Unpin> Future for Double<F> {
    type Output = u32;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<u32> {
        // 向下委托 poll 调用,把同一个 Context 传给内层状态机
        match Pin::new(&mut self.inner).poll(cx) {
            Poll::Ready(val) => Poll::Ready(val * 2),
            Poll::Pending => Poll::Pending,
        }
    }
}
use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};

/// 每次被 poll 时减一,减到 0 时就绪.
pub struct Countdown {
    pub remaining: u32,
}

impl Future for Countdown {
    type Output = u32;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<u32> {
        if self.remaining == 0 {
            Poll::Ready(21)
        } else {
            self.remaining -= 1;
            // 关键:返回 Pending 前必须安排唤醒!
            // 这里演示立即通知执行器下一轮继续 poll 自己;
            // 真实的异步 I/O 会把 cx.waker().clone() 存进反应器.
            cx.waker().wake_by_ref();
            Poll::Pending
        }
    }
}

/// 包装另一个 Future,在其就绪后将输出乘 2.
pub struct Double<F> {
    pub inner: F,
}

impl<F: Future<Output = u32> + Unpin> Future for Double<F> {
    type Output = u32;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<u32> {
        // 向下委托 poll 调用,把同一个 Context 传给内层状态机
        match Pin::new(&mut self.inner).poll(cx) {
            Poll::Ready(val) => Poll::Ready(val * 2),
            Poll::Pending => Poll::Pending,
        }
    }
}

注意 Double::poll Double::poll 的写法:它本身并不知道底层到底在等计时器、套接字还是通道,也不需要直接操作 cx.waker() cx.waker() ; 它只负责把 cx cx 原封不动地透传给内层的 self.inner.poll(cx) self.inner.poll(cx) . 标准库还提供了 std::task::ready!(expr) std::task::ready!(expr) 宏(历史代码中常见 futures::ready! futures::ready! ),专门用来简化上述 match match 模式:

let val = std::task::ready!(Pin::new(&mut self.inner).poll(cx));
Poll::Ready(val * 2)
let val = std::task::ready!(Pin::new(&mut self.inner).poll(cx));
Poll::Ready(val * 2)

当多层 Future Future 组合在一起时,外层是一个大的结构体,内层是嵌套在其中的字段. 一次顶层的 poll poll 调用,就是沿着这棵结构体树自顶向下的一连串同步方法调用,直到抵达最底层的叶子 Future Future :

3 为什么 Waker Waker 使用手工虚表 RawWaker RawWaker

理解了 Waker Waker 的职责后,一个工程问题随之而来: Rust 生态中有许多不同的执行器—— tokio tokio 的多线程工作窃取调度器、 async-std async-std 、 smol smol ,乃至嵌入式芯片上跑在中断里的微型轮询器. 不同执行器里的“任务句柄”长相完全不同:有的指向 Arc<Task> Arc<Task> ,有的只是一个静态数组的下标,有的甚至不需要任何数据指针. 而同一个底层库(比如某个驱动硬件的 Future Future )必须能够在不重新编译、不带执行器泛型参数 Future<E: Executor> Future<E: Executor> 的前提下,被任何执行器驱动.

如果在 Future Future 特型上加泛型参数,整个异步生态就会分裂成互不兼容的孤岛; 如果直接把 Waker Waker 定义成 Box<dyn WakeTrait> Box<dyn WakeTrait> ,每次克隆 Waker Waker 就可能触发一次堆内存分配,或者强制依赖 alloc alloc 库(无法用于无堆的裸机 no_std no_std 环境).

为此,标准库在 core::task core::task 中采用了一套手工构造的胖指针(类型擦除虚表)设计,由 Waker Waker 、 RawWaker RawWaker 与 RawWakerVTable RawWakerVTable 组成:

pub struct Waker {
    waker: RawWaker,
}

#[derive(PartialEq, Debug)]
pub struct RawWaker {
    data: *const (),
    vtable: &'static RawWakerVTable,
}

pub struct RawWakerVTable {
    clone: unsafe fn(*const ()) -> RawWaker,
    wake: unsafe fn(*const ()),
    wake_by_ref: unsafe fn(*const ()),
    drop: unsafe fn(*const ()),
}
pub struct Waker {
    waker: RawWaker,
}

#[derive(PartialEq, Debug)]
pub struct RawWaker {
    data: *const (),
    vtable: &'static RawWakerVTable,
}

pub struct RawWakerVTable {
    clone: unsafe fn(*const ()) -> RawWaker,
    wake: unsafe fn(*const ()),
    wake_by_ref: unsafe fn(*const ()),
    drop: unsafe fn(*const ()),
}

Waker Waker 在内存中恰好占用两个机器字长(64 位平台上为 16 字节): 一个擦除了具体类型的裸指针 data: *const () data: *const () ,加上一个指向编译期静态常量表 RawWakerVTable RawWakerVTable 的引用.

这套四函数虚表精确覆盖了 Waker Waker 的全部生命周期操作,且每一项都留有性能空间:

  • clone clone :复制一份 Waker Waker .若 data data 底层是 Arc<Task> Arc<Task> , clone clone 只需调用 Arc::increment_strong_count Arc::increment_strong_count 原地增加原子引用计数,并返回指向同一 data data 和同一 vtable vtable 的新 RawWaker RawWaker ,全程没有新的堆内存分配;若在裸机上 data data 只是全局中断标志的地址, clone clone 甚至什么都不用做,直接复制指针即可.
  • wake wake 与 wake_by_ref wake_by_ref :前者按值消耗 Waker Waker (可以省去先增引用计数再减引用计数的原子操作),后者通过借用触发唤醒.
  • drop drop : Waker Waker 析构时调用,例如对底层的 Arc<Task> Arc<Task> 递减引用计数.

标准库在拥有堆分配器的 std::task::Wake std::task::Wake 特型中,已经为 Arc<T: Wake> Arc<T: Wake> 封装好了这套 RawWakerVTable RawWakerVTable 样板代码; 普通编写执行器的人只需为自己的任务结构体实现 Wake Wake 特型,再调用 Waker::from(arc_task) Waker::from(arc_task) 即可安全地拿到 Waker Waker ,无需手写 unsafe unsafe 裸指针转换.

4 常见陷阱:伪唤醒与 Waker Waker 更新

在实际编写底层 Future Future 时,有两个与 Waker Waker 相关的边界条件容易出错:

  1. 容忍伪唤醒: 仅仅因为 Future::poll Future::poll 被再次调用,不代表等待的事件一定已经完成. 比如一个组合子同时等待两个子 Future Future (如 select! select! 或 join! join! ),其中左支触发了同一个 Waker Waker ,执行器会把整个顶层任务重新 poll poll 一遍,此时右支也会跟着被 poll poll ,但右支等待的 I/O 根本还没就绪. 因此, poll poll 的实现必须始终先检查底层资源的真实就绪状态(例如再次尝试非阻塞读取并检查是否仍返回 WouldBlock WouldBlock ),绝不能假设“第二次进 poll poll 就一定成功”.
  2. 替换旧的 Waker Waker : 一个 Future Future 在两次 poll poll 之间,可能会被迁移到另一个执行器或另一个任务上下文中驱动,导致第二次传入的 cx.waker() cx.waker() 与第一次保存的不是同一个唤醒目标. 如果底层事件源在多次返回 Pending Pending 的过程中只保留了第一次 poll poll 时的 Waker Waker ,后续的唤醒就会发往错误的旧任务. 正确的做法是在每次返回 Pending Pending 前,用最新的 cx.waker() cx.waker() 更新已存储的 Waker Waker (可借助 Waker::will_wake Waker::will_wake 先检查新旧 Waker Waker 是否指向同一目标,避免无谓的原子引用计数更新):
if !stored_waker.will_wake(cx.waker()) {
    stored_waker = cx.waker().clone();
}
if !stored_waker.will_wake(cx.waker()) {
    stored_waker = cx.waker().clone();
}
5 练习:轮询协议
Exercises (2)

Exercise 2.1. 诊断“静默挂起”的 Future Future

阅读下面这个试图“第一次挂起、第二次返回 42 42 ”的 Future Future 实现,说明把它交给真实执行器运行时会发生什么,并指出原因与修复方式:

pub struct BrokenOnce {
    pub ready: bool,
}

impl Future for BrokenOnce {
    type Output = i32;

    fn poll(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<i32> {
        if self.ready {
            Poll::Ready(42)
        } else {
            self.ready = true;
            Poll::Pending
        }
    }
}
pub struct BrokenOnce {
    pub ready: bool,
}

impl Future for BrokenOnce {
    type Output = i32;

    fn poll(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<i32> {
        if self.ready {
            Poll::Ready(42)
        } else {
            self.ready = true;
            Poll::Pending
        }
    }
}
Hint检查当 self.ready self.ready 为 false false 时, BrokenOnce::poll BrokenOnce::poll 是否向 cx.waker() cx.waker() 发出了任何通知或将其保存下来.
Solution

第一次调用 poll poll 时, self.ready self.ready 被设为 true true 并返回 Poll::Pending Poll::Pending , 但函数既没有调用 cx.waker().wake_by_ref() cx.waker().wake_by_ref() ,也没有把 cx.waker().clone() cx.waker().clone() 交给任何外部事件源. 在遵守标准协议的执行器(如 tokio tokio 或后文我们将实现的队列执行器)中,执行器收到 Pending Pending 后会将任务移出就绪队列,等待 Waker Waker 被触发时再将其放回. 由于没有任何地方会触发这个 Waker Waker ,该任务永远不会获得第二次 poll poll 的机会,程序会一直挂起(或阻塞在 block_on block_on 的条件变量上).

修复方法:若本意是“让出一次执行权后立即恢复”(类似 yield_now yield_now ),应在返回 Poll::Pending Poll::Pending 前加上 cx.waker().wake_by_ref(); cx.waker().wake_by_ref(); .

Exercise 2.2. 组合子的 poll poll 次数与唤醒传递

设顶层任务为 Double { inner: Countdown { remaining: 2 } } Double { inner: Countdown { remaining: 2 } } . 将它交给执行器运行直到完成,顶层的 Double::poll Double::poll 一共会被调用几次?每次返回什么?
Hint沿着 Double::poll -> Countdown::poll Double::poll -> Countdown::poll 逐步追踪每次调用时 remaining remaining 的变化和返回值.
Solution

共需要 3 次顶层 poll poll 调用:

  1. 第一次 poll poll : remaining remaining 从 2 2 减为 1 1 ,调用 wake_by_ref() wake_by_ref() ,返回 Poll::Pending Poll::Pending ; Double Double 收到 Pending Pending 后直接向上返回 Poll::Pending Poll::Pending .
  2. 第二次 poll poll : remaining remaining 从 1 1 减为 0 0 ,调用 wake_by_ref() wake_by_ref() ,返回 Poll::Pending Poll::Pending ; Double Double 向上返回 Poll::Pending Poll::Pending .
  3. 第三次 poll poll : remaining == 0 remaining == 0 , Countdown Countdown 返回 Poll::Ready(21) Poll::Ready(21) ; Double Double 将其乘 2,返回 Poll::Ready(42) Poll::Ready(42) .

在这个过程中, Double Double 自己完全没有触碰 cx cx ,但 Countdown Countdown 通过透传下来的 cx cx 直接唤醒了顶层任务.