CHAPTER 01

第 1 章:异步哲学与核心模型:Future、Waker 与协作式调度

官方源: tokio-rs/tokio · Commit @e800714a · 全书进度: 第 1 / 14 章

第 1 章:异步哲学与核心模型:Future、Waker 与协作式调度

异步编程在 Rust 中不是一个库,而是一套语言级的协议。Tokio 之所以能成为生产级运行时,不是因为它发明了 Future,而是因为它精确地实现了这套协议中每一个契约的边界条件。本章不急于跳进 Tokio 的调度器代码,而是先把「三件套」——Future、Waker、Executor——的职责边界和反向控制流讲透。理解了这三者如何咬合,后续章节中 Runtime 的组装、work-stealing 调度、I/O 驱动才有落脚点。

1.1 从阻塞到拉取:为什么 Rust 选择 poll 而非回调

直觉模型

想象你在餐厅点了一份需要现做的菜。回调式异步(如 Node.js 早期风格)相当于你留下手机号,厨师做好后主动打给你——控制权在厨师手里,你的代码只是被动响应。拉取式异步(Rust 的选择)相当于你拿到一张取餐凭证,你自己决定什么时候去窗口问「好了吗」:没好就去做别的事,好了就取走。

这个区别看似微小,却决定了整个系统的形态。回调式模型中,每个异步操作都必须携带一个「完成后做什么」的闭包,闭包层层嵌套形成回调地狱,且取消操作极其困难——你无法「撤回」一个已经注册的回调。拉取式模型中,Future 只是一个状态机,poll 是纯粹的查询动作,不推进就不消耗资源,取消就是 drop,干净利落。

拉取式模型的核心契约

Rust 标准库定义的 Future trait 只有两个要素:一个 poll 方法,一个 Output 关联类型。Tokio 并没有重新定义这个 trait,而是直接复用标准库的实现。这一点在源码中有明确体现:

rust
// tokio/src/future/mod.rs
cfg_not_trace! {
    cfg_rt! {
        pub(crate) use std::future::Future;
    }
}

📎 tokio/src/future/mod.rs:24-28

这段代码揭示了一个重要事实:在未启用 tracing 特性时,Tokio 内部的 Future 就是 std::future::Future 的别名,没有任何包装。只有在启用 tracing 时,才会用 InstrumentedFuture 替换:

rust
cfg_trace! {
    mod trace;
    #[allow(unused_imports)]
    pub(crate) use trace::InstrumentedFuture as Future;
}

📎 tokio/src/future/mod.rs:18-22

〔设计推断与架构权衡〕
这种「默认零开销、按需插桩」的设计是 Tokio 的一贯哲学:核心路径不引入任何额外抽象层,可观测性作为可选特性叠加。InstrumentedFuture 的存在说明 Tokio 团队认为 tracing 的插桩成本不应由所有用户承担。

poll 契约的三个隐含约束

poll 方法的签名是 fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output>。这个签名里藏着三条契约,违反任何一条都会导致未定义行为或逻辑错误:

契约一:Pin 保证自引用安全。 Pin<&mut Self> 意味着 Future 一旦被 poll,其内存地址就不能再移动。这是因为 async 块编译后会生成包含自引用的状态机——局部变量可能持有指向同一状态机内其他字段的引用。如果允许移动,这些引用就会悬空。

契约二:Pending 必须已注册唤醒。 当 poll 返回 Poll::Pending 时,Future 必须已经通过 cx.waker() 获取并保存了 Waker,或者已经将 Waker 注册到了某个事件源。否则执行器将永远不知道该 Future 何时可以再次被 poll,导致任务永久挂起。

契约三:Ready 之后不应再 poll。 一旦 poll 返回 Poll::Ready,再次 poll 同一个 Future 是逻辑错误(虽然不会导致 UB,但行为未定义)。执行器有责任在收到 Ready 后不再调度该任务。

这三条契约中,契约二是最容易出错的地方,也是 Waker 存在的根本原因。

1.2 Waker:反向控制流的载体

直觉模型

Waker 是餐厅给你的「震动取餐器」。你不需要站在窗口反复问「好了吗」——那会浪费你的时间。你只需要在第一次去窗口时把取餐器交给厨师(注册 Waker),然后安心做别的事。菜好了,厨师按下按钮,取餐器震动(调用 wake),你收到信号后再去窗口取餐(重新 poll)。

如果没有 Waker,执行器只有两种选择:要么忙轮询所有任务(浪费 CPU),要么永远不 poll 已返回 Pending 的任务(任务饿死)。Waker 是打破这个僵局的唯一机制。

Waker 的内存布局与虚表设计

Waker 是标准库类型,但它的设计直接影响了 Tokio 的任务结构。Waker 本质上是一个胖指针:一个 RawWaker 结构体,包含一个数据指针和一个虚表指针。

rust
// 标准库中的定义(非 Tokio 源码,此处为背景说明)
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 本身不关心「唤醒」具体意味着什么。它只是四个函数指针的载体。Tokio 可以提供一个 Waker,其 wake 函数把任务重新推入调度队列;而另一个运行时(比如 futures crate 的 block_on)可以提供完全不同的 Waker 实现。这种「数据 + 虚表」的模式使得 Waker 可以在不同运行时之间传递而不丢失语义。

wake 和 wake_by_ref 的区别至关重要:wake 消耗 Waker 的所有权(调用后 Waker 被 drop),而 wake_by_ref 只借用。执行器通常实现 wake_by_ref 为「将任务标记为就绪并入队」,而 wake 则在此基础上额外处理引用计数的递减。Tokio 的任务结构中,Waker 的数据指针指向任务的引用计数头,每次 clone 增加计数,drop 减少计数,计数归零时释放任务内存。

唤醒的完整时序

下面这张时序图展示了一个 TCP 读取操作从发起到被唤醒的完整链路。注意 Waker 是如何从任务上下文一路传递到 I/O 驱动的:

mermaid
sequenceDiagram
    participant App as 应用任务
    participant Exec as 调度器 Worker
    participant Future as TcpStream::read Future
    participant Reactor as I/O 驱动 (epoll)
    participant Kernel as 操作系统内核

    App->>Future: poll(cx) 携带 Waker
    Future->>Reactor: 注册可读兴趣 + 保存 Waker
    Reactor->>Kernel: epoll_ctl(ADD, fd, EPOLLIN)
    Future-->>Exec: 返回 Poll::Pending
    Note over Exec: 任务挂起,Worker 去执行其他任务
    Kernel-->>Reactor: epoll_wait 返回 fd 就绪
    Reactor->>Reactor: 查找 fd 对应的 Waker
    Reactor->>Exec: waker.wake_by_ref()
    Note over Exec: 任务重新入队
    Exec->>Future: 再次 poll(cx)
    Future->>Kernel: read(fd, buf) 非阻塞读取
    Kernel-->>Future: 返回数据
    Future-->>App: 返回 Poll::Ready(n)

这张图的关键在于:Waker 是唯一能从 Reactor 反向触达 Executor 的通道。Reactor 不持有任务的任何其他信息,它只知道「当这个 fd 就绪时,调用这个 Waker」。这种解耦使得 I/O 驱动可以独立于调度器实现,两者只通过 Waker 这个窄接口通信。

虚假唤醒:契约的灰色地带

Tokio 的文档明确承认虚假唤醒的存在:

Normally, tasks are scheduled only if they have been woken by calling wake on their waker. However, this is not guaranteed, and Tokio may schedule tasks that have not been woken under some circumstances.

📎 tokio/src/runtime/mod.rs:306-309

〔设计推断与架构权衡〕
这意味着 poll 的实现必须能够容忍「没有被唤醒就被再次 poll」的情况。一个正确的 Future 在返回 Pending 后,即使没有任何事件发生,再次被 poll 时也应该返回 Pending 而不是 panic 或产生错误结果。这个约束看似宽松,实际上对状态机的设计提出了要求:不能假设「两次 poll 之间一定有事件发生」。

1.3 Executor:从 Future 到任务的封装

直觉模型

Executor 是餐厅的调度员。他手里有一摞订单(任务队列),决定哪个订单先做、谁来做。当取餐器震动时,他把对应订单重新排进队列。没有调度员,厨师们就不知道该做哪道菜,也不知道该在什么时候切换工作。

但 Executor 的职责远不止「轮询 Future」。它必须解决三个核心问题:任务的生命周期管理(创建、调度、完成、取消)、公平性保证(防止某个任务饿死其他任务)、资源驱动集成(I/O 和定时器事件如何转化为唤醒)。

任务的内存布局:从 Future 到 Task

当调用 tokio::spawn 时,传入的 Future 并不会被直接放入队列。它会被包装成一个 Task 结构,包含引用计数头、调度元数据和 Future 本身。这个包装过程有一个关键的优化决策:

rust
/// Boundary value to prevent stack overflow caused by a large-sized
/// Future being placed in the stack.
pub(crate) const BOX_FUTURE_THRESHOLD: usize = if cfg!(debug_assertions)  {
    2048
} else {
    16384
};

pub(crate) struct AutoBox(std::marker::PhantomData);

impl AutoBox {
    /// `true` if a value of type `T` is larger than [`BOX_FUTURE_THRESHOLD`].
    pub(crate) const SHOULD_BOX: bool = std::mem::size_of::() > BOX_FUTURE_THRESHOLD;
}

📎 tokio/src/runtime/mod.rs:649-673

这段代码解决了一个非常具体的问题:如果 Future 太大(超过 16KB,debug 模式下 2KB),直接内联到 Task 结构中会导致栈溢出或内存浪费。AutoBox 通过编译期常量 SHOULD_BOX 来决定是否将 Future 装箱。

〔设计推断与架构权衡〕
注释中特别强调了「用关联常量而非运行时 if」的原因:如果用运行时判断,编译器会为每个 T 同时实例化两条分支的代码(一条处理 T,一条处理 Pin<Box<T>>),导致代码膨胀。而用常量分支,单态化收集器会剪掉不可达的分支,只为实际使用的类型生成代码。这是一个典型的「用类型系统替代运行时判断」的优化。

调度公平性:31 与 61 的魔法数字

Tokio 的调度器文档中定义了一个形式化的公平性保证:

If the total number of tasks does not grow without bound, and no task is blocking the thread, then it is guaranteed that tasks are scheduled fairly.

📎 tokio/src/runtime/mod.rs:279-281

这个保证的实现依赖于两个关键参数。对于 current-thread 运行时:

The runtime will prefer to choose the next task to schedule from the local queue, and will only pick a task from the global queue if the local queue is empty, or if it has picked a task from the local queue 31 times in a row.

📎 tokio/src/runtime/mod.rs:328-333

The runtime will check for new IO or timer events whenever there are no tasks ready to be scheduled, or when it has scheduled 61 tasks in a row.

📎 tokio/src/runtime/mod.rs:335-337

这两个数字(31 和 61)不是随意选择的。31 是 2 的 5 次方减 1,可以用位运算快速判断;61 则是为了确保 I/O 事件不会被无限延迟——即使任务队列永远非空,每 61 次调度后也必须检查一次 I/O。

〔设计推断与架构权衡〕
为什么是 31 而不是 32?因为计数器从 0 开始,每调度一次加 1,当计数器达到 31 时触发全局队列检查。用 counter & 31 == 31 判断比 counter % 32 == 0 更高效(虽然现代编译器会自动优化)。61 的选择则更微妙:它需要足够大以避免频繁的 epoll_wait 系统调用开销,又需要足够小以保证 I/O 延迟在可接受范围内。

多线程运行时的 LIFO 槽优化

多线程运行时在公平性之上还增加了一个性能优化——LIFO 槽:

The multi thread runtime uses the lifo slot optimization: Whenever a task wakes up another task, the other task is added to the worker thread's lifo slot instead of being added to a queue.

📎 tokio/src/runtime/mod.rs:373-377

这个优化的直觉是:当一个任务唤醒另一个任务时,被唤醒的任务很可能与当前任务有数据依赖(比如生产者-消费者模式)。把它放在 LIFO 槽中,当前任务完成后立即执行它,可以利用 CPU 缓存的热数据。

但 LIFO 槽有一个防滥用机制:

if a worker thread uses the lifo slot three times in a row, it is temporarily disabled until the worker thread has scheduled a task that didn't come from the lifo slot.

📎 tokio/src/runtime/mod.rs:380-382

〔设计推断与架构权衡〕
这个「三次连续使用后禁用」的规则是为了防止两个任务互相唤醒形成活锁。如果任务 A 唤醒任务 B,B 又唤醒 A,没有这个限制的话,LIFO 槽会被这两个任务永久占用,其他任务永远得不到调度。三次的限制给了其他任务一个插入的机会。

任务取消:abort 的真实语义

JoinHandle::abort 的行为经常被误解。文档明确指出:

Be aware that calls to JoinHandle::abort just schedule the task for cancellation, and will return before the cancellation has completed.

📎 tokio/src/task/mod.rs:146-148

这意味着 abort 不是同步的。它只是设置一个标志位,任务会在下一个 .await 点检查这个标志并自行终止。如果任务正在执行一段没有 .await 的 CPU 密集代码,abort 不会立即生效。

更微妙的是:

Note that aborting a task does not guarantee that it fails with a cancelled error, since it may complete normally first.

📎 tokio/src/task/mod.rs:134-138

〔设计推断与架构权衡〕
这个语义的设计动机是:取消是一个「尽力而为」的操作。Tokio 不强制杀死任务(Rust 没有安全的强制终止机制),而是协作式地请求任务自行退出。这与 spawn_blocking 任务不可取消的设计是一致的——阻塞任务没有 .await 点,无法检查取消标志。

1.4 设计思考:三件套的边界与代价

为什么 Future 不包含 Executor

Rust 的 Future trait 刻意不包含「如何调度自己」的信息。这是一个深思熟虑的解耦决策。如果 Future 知道自己的 Executor,那么:

1. 同一个 Future 无法在不同运行时上执行(比如从 Tokio 迁移到 async-std)

2. 测试时无法用简单的 block_on 驱动

3. 组合器(如 select!、join!)无法跨运行时工作

Waker 的存在正是为了在保持这种解耦的同时,仍然允许 Future 通知 Executor。Waker 是一个「能力令牌」——Future 只知道「我可以调用这个来请求重新调度」,但不知道调度具体如何发生。

协作式调度的代价

Tokio 的任务是协作式的:任务只有在 .await 点才会让出执行权。这意味着:

code that spends a long time without reaching an .await will prevent other tasks from running.

📎 tokio/src/lib.rs:178-179

〔设计推断与架构权衡〕
这是协作式调度的根本代价。操作系统可以在任意指令边界抢占线程,但 Tokio 只能在 .await 点切换任务。如果一个任务执行了一个 10 秒的 CPU 密集循环且中间没有 .await,那么同一个 worker 线程上的其他所有任务都会被阻塞 10 秒。Tokio 的应对策略是提供 spawn_blocking 和 block_in_place,把这类工作转移到专用线程池。但这是用户的责任,运行时无法自动检测。

公平性保证的边界条件

Tokio 的公平性保证有两个前提条件:任务总数有上界,且没有任务阻塞线程。这两个条件在实际生产环境中经常被违反:

  • 如果任务不断 spawn 新任务且不回收,任务总数无上界,公平性保证失效
  • 如果某个任务执行了阻塞系统调用(比如同步文件 I/O),它阻塞了整个 worker 线程
〔设计推断与架构权衡〕
这就是为什么 Tokio 文档反复强调「不要在异步任务中执行阻塞操作」。公平性保证不是运行时的硬性保证,而是「在正确使用的前提下」的保证。运行时不检测违规行为,因为检测本身需要开销。

1.5 本章小结

本章建立了理解 Tokio 的三个基石:

Future 是拉取式的状态机。 poll 是纯粹的查询动作,返回 Pending 时必须已注册唤醒,返回 Ready 后不应再被 poll。Tokio 直接复用 std::future::Future,不做额外包装(除非启用 tracing)。

Waker 是反向控制流的唯一通道。 它通过「数据指针 + 虚表」的设计实现了运行时无关性。wake 消耗所有权,wake_by_ref 只借用。虚假唤醒是允许的,Future 必须容忍。

Executor 负责生命周期、公平性和资源集成。 它把 Future 包装成 Task,通过 AutoBox 在编译期决定是否装箱,通过 31/61 这两个魔法数字平衡本地队列和全局队列的调度,通过 LIFO 槽优化数据依赖场景的性能。

这三个组件通过窄接口解耦:Future 只知道 poll,Waker 只知道 wake,Executor 只知道「轮询直到 Pending 或 Ready」。正是这种解耦使得 Tokio 可以在不修改 Future 定义的前提下,实现 work-stealing 调度、I/O 驱动集成、协作式预算等高级特性。

本章思考与自测

Q1: 如果将 AutoBox::SHOULD_BOX 的判断从编译期常量改为运行时 if size_of::<T>() > THRESHOLD,会对编译产物产生什么影响?为什么 Tokio 的注释特别强调这一点?

参考解析:根据 📎 tokio/src/runtime/mod.rs:657-667 的注释,如果用运行时 if,编译器会为每个 T 同时实例化两条分支的代码——一条处理 T 直接内联的情况,一条处理 Pin<Box<T>> 的情况。这意味着每个 spawn 的 Future 类型都会生成两份任务驱动代码(task harness),导致二进制体积翻倍。而用关联常量 SHOULD_BOX,由于它在 T 确定后就是编译期常量,单态化收集器会剪掉不可达的分支,只为实际使用的路径生成代码。这是一个「用类型系统替代运行时判断」的典型优化,代价是 AutoBox 必须是一个泛型结构体而非普通函数。

Q2: 假设一个任务在 poll 中返回了 Pending,但忘记注册 Waker。在 current-thread 运行时和 multi-thread 运行时下,这个任务分别会发生什么?Tokio 有没有机制检测这种情况?

参考解析:根据 📎 tokio/src/runtime/mod.rs:306-309,Tokio 允许虚假唤醒,这意味着任务可能在没有被唤醒的情况下被重新调度。但这不意味着忘记注册 Waker 是安全的。在 current-thread 运行时下,如果本地队列和全局队列都为空,运行时会进入 park 状态等待 I/O 或定时器事件。忘记注册 Waker 的任务永远不会被重新入队,导致永久挂起。在 multi-thread 运行时下,情况类似,但如果有其他任务持续唤醒,该任务可能因为虚假唤醒而被偶然重新调度——但这不可依赖。Tokio 没有运行时检测机制来发现「返回 Pending 但未注册 Waker」的情况,因为这需要在每次 poll 后检查 Waker 是否被使用,开销太大。这是 Future 实现者的责任。

Q3: LIFO 槽的「三次连续使用后禁用」规则是为了防止什么具体场景?如果去掉这个限制,在什么样的任务依赖模式下会导致其他任务饿死?

参考解析:根据 📎 tokio/src/runtime/mod.rs:380-382,LIFO 槽在连续使用三次后会被临时禁用,直到调度了一个非 LIFO 来源的任务。这个规则防止的场景是:两个任务互相唤醒形成紧密循环。例如任务 A 处理完一批数据后唤醒任务 B,任务 B 处理完后立即唤醒任务 A。如果没有三次限制,A 和 B 会永远占据 LIFO 槽,worker 线程会在这两个任务之间无限切换,本地队列和全局队列中的其他任务永远得不到执行机会。三次的限制确保了每处理三轮「互相唤醒」后,至少有一个其他任务被调度,打破了活锁。这个数字的选择是经验性的:太小会降低 LIFO 优化的收益,太大会增加其他任务的延迟。

至此,Future、Waker 与 Executor 三者的职责边界与协作机制已经清晰:Future 定义计算,Waker 负责唤醒,Executor 驱动执行。但单个组件无法独立工作,它们必须被组装进一个统一的运行时环境。下一章,我们将追踪 Runtime::new 与 Builder::build 的完整装配链路,看调度器、I/O 驱动、时间驱动和阻塞线程池如何被注入同一个 Runtime 实例,并揭示 current_thread 与 multi_thread 两种形态在装配阶段的根本差异。

AI 赋能代码库精读 · 本地优先架构

读完了本章?为你自己的私有项目生成专属架构全景书

基于 Tauri 2 + Rust 本地原生引擎,100% 源码离线隐私安全,零代码上传云端。像阅读一本传世专著一样拆解你的复杂系统。

⚡ Tauri 2 · Rust 原生引擎 · 100% 离线私密安全 · 适配超百万行代码库

CHAPTER 02

第 2 章:Runtime 装配工厂:Builder 如何把线程池与驱动拼装成一个运行时

官方源: tokio-rs/tokio · Commit @e800714a · 全书进度: 第 2 / 14 章

第 2 章:Runtime 装配工厂:Builder 如何把线程池与驱动拼装成一个运行时

从 Builder 到 Runtime:一次装配的完整旅程

上一章我们把 Future、Waker、Executor 三者的职责边界讲清楚了。但一个真实可用的运行时远不止「一个 Executor」——它还需要 I/O 事件循环、定时器、阻塞线程池,并且这些组件必须共享同一套句柄、同一份生命周期。本章追踪 Builder::build 的完整装配链路,回答一个核心问题:一个 Runtime 内部到底有哪些组件,它们如何被拼装并共享句柄。

Tokio 的装配入口是 Builder。它本身是一个纯配置容器,所有字段都是「意图声明」,不持有任何运行时资源。真正的资源创建发生在 build() 调用时。

直觉模型:Builder 是「装修图纸」,Runtime 是「交房后的房子」

Builder 就像一张装修图纸:你在上面标注「要几个房间(worker_threads)」「要不要通水(enable_io)」「要不要通电(enable_time)」「外包帮工上限(max_blocking_threads)」。图纸本身不产生任何实体。直到调用 build(),施工队才按图施工,把调度器、驱动、线程池这些「房间」真正建起来,并交付一个 Runtime 实例。

若没有 Builder 这一层,用户就必须手动 new 出每个组件、手动接线、手动处理失败回滚——任何一处顺序错误都会导致句柄悬空或资源泄漏。Builder 的价值在于:把「配置」与「构造」彻底分离,让构造过程可以集中做校验、失败清理和句柄共享。

内存布局:Builder 的字段分区

Builder 的字段可以按职责分成四组。第一组是形态与开关:kind 决定调度器形态,enable_io / enable_time 决定是否创建对应驱动。

📎 tokio/src/runtime/builder.rs:55-68

rust
pub struct Builder {
    kind: Kind,
    name: Option,
    enable_io: bool,
    nevents: usize,
    nevents_busy: Option,
    enable_time: bool,
    start_paused: bool,
    // ...
}

第二组是线程池参数:worker_threads 是 Option<usize>,None 表示「延迟到 build 时按 CPU 核数自动探测」;max_blocking_threads 默认 512。

📎 tokio/src/runtime/builder.rs:73-79

rust
worker_threads: Option,
max_blocking_threads: usize,

第三组是回调钩子,全部是 Option<Arc<dyn Fn ...>>。注意它们用 Arc 而非 Box,因为这些回调要被克隆进每个 worker 线程的 Config。

📎 tokio/src/runtime/builder.rs:87-97

rust
pub(super) after_start: Option,
pub(super) before_stop: Option,
pub(super) before_park: Option,
pub(super) after_unpark: Option,

第四组是调度启发式与随机种子:global_queue_interval、event_interval、disable_lifo_slot、seed_generator。

📎 tokio/src/runtime/builder.rs:116-134

rust
pub(super) global_queue_interval: Option,
pub(super) event_interval: u32,
pub(super) disable_lifo_slot: bool,
pub(super) seed_generator: RngSeedGenerator,

这里有一个值得注意的设计:Kind 是一个 Copy 的小枚举,只有两个变体。

📎 tokio/src/runtime/builder.rs:261-265

rust
#[derive(Clone, Copy)]
pub(crate) enum Kind {
    CurrentThread,
    #[cfg(feature = "rt-multi-thread")]
    MultiThread,
}

MultiThread 变体被 rt-multi-thread feature 门控。这意味着在只启用 rt feature 的构建里,Kind 只有一个变体,build() 的 match 会被编译器优化成单分支——用类型系统而非运行时判断来消除多线程调度器的代码体积。

默认值的哲学:为什么 I/O 和 time 默认关闭

Builder::new 是所有构造的公共入口。它把 enable_io 和 enable_time 都设为 false。

📎 tokio/src/runtime/builder.rs:309-318

rust
// I/O defaults to "off"
enable_io: false,
nevents: 1024,
nevents_busy: None,

// Time defaults to "off"
enable_time: false,

// The clock starts not-paused
start_paused: false,
〔设计推断与架构权衡〕
这个默认值选择是刻意的:创建 I/O 驱动需要向操作系统申请 epoll/kqueue 句柄,创建 time 驱动需要启动定时器基础设施。如果用户只是想要一个纯计算的任务调度器(比如跑 CPU 密集的 async 逻辑),强制创建这些驱动就是纯粹的浪费。#[tokio::main] 宏之所以「开箱即用」,是因为它内部调用了 enable_all()。

enable_all() 的实现揭示了 feature 门控如何影响「全开」的语义。

📎 tokio/src/runtime/builder.rs:398-419

rust
pub fn enable_all(&mut self) -> &mut Self {
    #[cfg(any(
        feature = "net",
        all(unix, feature = "process"),
        all(unix, feature = "signal")
    ))]
    self.enable_io();

    #[cfg(all(
        tokio_unstable,
        feature = "io-uring",
        // ...
    ))]
    self.enable_io_uring();

    #[cfg(feature = "time")]
    self.enable_time();

    self
}

注意 enable_io() 只在启用了 net、process 或 signal feature 时才被调用。如果用户只启用了 time feature,enable_all() 不会打开 I/O 驱动——因为编译产物里根本没有 I/O 驱动代码。

装配主路径:build() 的分流

build() 是装配的起点,它按 kind 分流到两条完全不同的路径。

📎 tokio/src/runtime/builder.rs:1146-1152

rust
pub fn build(&mut self) -> io::Result {
    match &self.kind {
        Kind::CurrentThread => self.build_current_thread_runtime(),
        #[cfg(feature = "rt-multi-thread")]
        Kind::MultiThread => self.build_threaded_runtime(),
    }
}

这两条路径的差异远不止「一个线程 vs 多个线程」。下面分别展开。

路径一:current_thread 的装配

build_current_thread_runtime 本身很薄,它委托给 build_current_thread_runtime_components,然后把返回的三元组包进 Runtime。

📎 tokio/src/runtime/builder.rs:1725-1736

rust
fn build_current_thread_runtime(&mut self) -> io::Result {
    use crate::runtime::runtime::Scheduler;

    let (scheduler, handle, blocking_pool) =
        self.build_current_thread_runtime_components(None)?;

    Ok(Runtime::from_parts(
        Scheduler::CurrentThread(scheduler),
        handle,
        blocking_pool,
    ))
}

真正的装配逻辑在 build_current_thread_runtime_components。它的执行顺序至关重要:

📎 tokio/src/runtime/builder.rs:1760-1766

rust
let mut cfg = self.get_cfg();
cfg.timer_flavor = TimerFlavor::Traditional;
let (driver, driver_handle) = driver::Driver::new(cfg)?;

// Blocking pool
let blocking_pool = blocking::create_blocking_pool(self, self.max_blocking_threads, 0);
let blocking_spawner = blocking_pool.spawner().clone();

第一步创建 driver,返回一对 (driver, driver_handle)。注意这里 ? 直接向上传播错误——如果 I/O 驱动初始化失败(比如 epoll 创建失败),整个 build 返回 Err,此时 blocking pool 还没创建,无需清理。

第二步创建 blocking pool,并立刻取出它的 spawner 克隆。这个 spawner 会被注入调度器,让调度器有能力把阻塞任务投递到线程池。

第三步生成两个独立的 RNG 种子生成器。

📎 tokio/src/runtime/builder.rs:1768-1770

rust
let seed_generator_1 = self.seed_generator.next_generator();
let seed_generator_2 = self.seed_generator.next_generator();
〔设计推断与架构权衡〕
为什么需要两个?seed_generator_1 被放进 Config,供调度器内部使用(比如 select! 的随机分支顺序);seed_generator_2 传给 CurrentThread::new,供任务侧使用。分离两个生成器可以避免调度器内部消费随机数影响用户可见的随机序列,从而保证 rng_seed 的可复现性。

第四步是核心:把 driver、driver_handle、blocking_spawner、种子和 Config 一起交给 CurrentThread::new。

📎 tokio/src/runtime/builder.rs:1776-1807

rust
let (scheduler, handle) = CurrentThread::new(
    driver,
    driver_handle,
    blocking_spawner,
    seed_generator_2,
    Config {
        before_park: self.before_park.clone(),
        after_unpark: self.after_unpark.clone(),
        // ...
        global_queue_interval: self.global_queue_interval,
        event_interval: self.event_interval,
        // ...
        enable_eager_driver_handoff: false,
        seed_generator: seed_generator_1,
        // ...
    },
    local_tid,
    self.name.clone(),
);

这里有一个关键细节:enable_eager_driver_handoff 被硬编码为 false。

📎 tokio/src/runtime/builder.rs:1795-1798

rust
// This setting never makes sense for a current thread runtime,
// as it only configures how the I/O driver is stolen across
// workers.
enable_eager_driver_handoff: false,
〔设计推断与架构权衡〕
这个注释点明了该选项的本质:它描述的是「多个 worker 之间如何抢占 I/O 驱动」,而 current_thread 只有一个线程,不存在抢占,所以强制关闭。这是「配置项语义与形态强相关」的典型例子——同一个 Builder 字段在不同形态下含义不同。

最后,CurrentThread::new 返回的 handle 被包进 scheduler::Handle::CurrentThread,再包进公开的 Handle。

📎 tokio/src/runtime/builder.rs:1816-1822

rust
let handle = Handle {
    inner: scheduler::Handle::CurrentThread(handle),
};

Ok((scheduler, handle, blocking_pool))

路径二:multi_thread 的装配

build_threaded_runtime 的骨架与 current_thread 类似,但有三处本质差异。第一处是 worker 线程数的确定:

📎 tokio/src/runtime/builder.rs:2185

rust
let worker_threads = self.worker_threads.unwrap_or_else(num_cpus);

None 在这里被解析为 num_cpus()。这就是「延迟自动探测」的落地点——探测发生在 build 时而非 Builder::new 时,因为 CPU 亲和性可能在两者之间变化。

第二处差异在 blocking pool 的容量计算:

📎 tokio/src/runtime/builder.rs:2189-2192

rust
let blocking_pool =
    blocking::create_blocking_pool(self, self.max_blocking_threads + worker_threads, worker_threads);
let blocking_spawner = blocking_pool.spawner().clone();

注意 max_blocking_threads + worker_threads。对比 current_thread 路径传入的是 self.max_blocking_threads 和 0。

📎 tokio/src/runtime/builder.rs:1765

rust
let blocking_pool = blocking::create_blocking_pool(self, self.max_blocking_threads, 0);
〔设计推断与架构权衡〕
这个差异揭示了 blocking pool 容量语义:multi_thread 下,max_blocking_threads 是「额外的」阻塞线程上限,实际总线程上限要加上 worker 线程数。第三个参数(current_thread 传 0,multi_thread 传 worker_threads)很可能是「预留线程数」或「初始线程数」的提示。这个设计让 max_blocking_threads 的语义在两种形态下保持一致:它描述的是「超出核心 worker 之外还能额外开多少阻塞线程」。

第三处差异是 MultiThread::new 返回三元组而非二元组:

📎 tokio/src/runtime/builder.rs:2198-2226

rust
let (scheduler, handle, launch) = MultiThread::new(
    worker_threads,
    driver,
    driver_handle,
    blocking_spawner,
    seed_generator_2,
    Config {
        // ...
        enable_eager_driver_handoff: self.enable_eager_driver_handoff,
        // ...
    },
    self.timer_flavor,
    self.name.clone(),
);

多出来的 launch 是一个「启动句柄」。MultiThread::new 只负责构造调度器结构,并不立即启动 worker 线程。真正的启动发生在后面:

📎 tokio/src/runtime/builder.rs:2228-2234

rust
let handle = Handle { inner: scheduler::Handle::MultiThread(handle) };

// Spawn the thread pool workers
let _enter = handle.enter();
launch.launch();

Ok(Runtime::from_parts(Scheduler::MultiThread(scheduler), handle, blocking_pool))

handle.enter() 进入运行时上下文,然后 launch.launch() 才真正 spawn 出所有 worker 线程。这个「先构造、后启动」的两阶段设计非常关键。

〔设计推断与架构权衡〕
为什么不能边构造边启动?因为 worker 线程一旦启动就会立刻开始 poll 任务,而任务可能引用 handle。如果 handle 还没构造完,就会出现「worker 拿着半成品句柄」的竞态。两阶段设计保证了:所有 worker 线程启动时,完整的 Handle 已经就绪。_enter 守卫确保 worker 线程在启动瞬间就处于正确的运行时上下文中。

装配流程图

下面这张图把两条路径的装配顺序、关键分支和错误路径画在一起。注意 driver::Driver::new 失败时直接返回 Err,此时 blocking pool 尚未创建。

mermaid
flowchart TD
    start["Builder::build()"] --> match_kind{"self.kind?"}

    match_kind -->|CurrentThread| ct_cfg["get_cfg() + timer_flavor=Traditional"]
    match_kind -->|MultiThread| mt_workers["worker_threads = self.worker_threads.unwrap_or_else(num_cpus)"]

    ct_cfg --> ct_driver["driver::Driver::new(cfg)?"]
    mt_workers --> mt_driver["driver::Driver::new(self.get_cfg())?"]

    ct_driver -->|Err| ret_err["return Err(io::Error)"]
    mt_driver -->|Err| ret_err

    ct_driver -->|Ok driver, driver_handle| ct_pool["create_blocking_pool(self, max_blocking_threads, 0)"]
    mt_driver -->|Ok driver, driver_handle| mt_pool["create_blocking_pool(self, max_blocking_threads + worker_threads, worker_threads)"]

    ct_pool --> ct_seed["next_generator() x2"]
    mt_pool --> mt_seed["next_generator() x2"]

    ct_seed --> ct_new["CurrentThread::new(driver, driver_handle, blocking_spawner, ...)"]
    mt_seed --> mt_new["MultiThread::new(worker_threads, driver, ...) -> (scheduler, handle, launch)"]

    ct_new --> ct_wrap["Handle { inner: CurrentThread(handle) }"]
    mt_new --> mt_wrap["Handle { inner: MultiThread(handle) }"]

    ct_wrap --> ct_rt["Runtime::from_parts(Scheduler::CurrentThread, handle, blocking_pool)"]
    mt_wrap --> mt_enter["handle.enter()"]
    mt_enter --> mt_launch["launch.launch() 启动 worker 线程"]
    mt_launch --> mt_rt["Runtime::from_parts(Scheduler::MultiThread, handle, blocking_pool)"]

句柄共享:Handle 如何成为跨组件的「通行证」

装配完成后,Runtime 持有 scheduler、handle、blocking_pool 三件套。其中 handle 是共享的核心。它的内部是一个枚举:

📎 tokio/src/runtime/scheduler/mod.rs:29-41

rust
#[derive(Debug, Clone)]
pub(crate) enum Handle {
    #[cfg(feature = "rt")]
    CurrentThread(Arc),

    #[cfg(feature = "rt-multi-thread")]
    MultiThread(Arc),

    #[cfg(not(feature = "rt"))]
    #[allow(dead_code)]
    Disabled,
}

注意两个变体都包着 Arc。这意味着 Handle 的克隆是廉价的引用计数递增,可以被自由地分发到任意线程。Handle 提供了统一的访问接口,把形态差异封装在 match 内部。例如 driver():

📎 tokio/src/runtime/scheduler/mod.rs:53-64

rust
pub(crate) fn driver(&self) -> &driver::Handle {
    match *self {
        #[cfg(feature = "rt")]
        Handle::CurrentThread(ref h) => &h.driver,

        #[cfg(feature = "rt-multi-thread")]
        Handle::MultiThread(ref h) => &h.driver,

        #[cfg(not(feature = "rt"))]
        Handle::Disabled => unreachable!(),
    }
}

blocking_spawner() 用了 match_flavor! 宏来消除重复:

📎 tokio/src/runtime/scheduler/mod.rs:96-98

rust
pub(crate) fn blocking_spawner(&self) -> &blocking::Spawner {
    match_flavor!(self, Handle(h) => &h.blocking_spawner)
}

这个宏展开后就是上面 driver() 那样的 match。它的价值在于:当新增一个需要按形态分发的访问器时,只需一行 match_flavor!,而不必手写两遍 match 分支。

公开的 Handle 是内部 scheduler::Handle 的薄包装:

📎 tokio/src/runtime/handle.rs:13-15

rust
pub struct Handle {
    pub(crate) inner: scheduler::Handle,
}

用户拿到的 Handle 可以跨线程克隆、可以 spawn、可以 block_on。spawn 的实现展示了 AutoBox 的编译期分支:

📎 tokio/src/runtime/handle.rs:197-208

rust
pub fn spawn(&self, future: F) -> JoinHandle
where
    F: Future + Send + 'static,
    F::Output: Send + 'static,
{
    let fut_size = mem::size_of::();
    if AutoBox::::SHOULD_BOX {
        self.spawn_named(Box::pin(future), SpawnMeta::new_unnamed(fut_size))
    } else {
        self.spawn_named(future, SpawnMeta::new_unnamed(fut_size))
    }
}

AutoBox::<F>::SHOULD_BOX 是一个关联常量,由 size_of::<F>() 与阈值比较得出。

📎 tokio/src/runtime/mod.rs:668-673

rust
pub(crate) struct AutoBox(std::marker::PhantomData);

impl AutoBox {
    pub(crate) const SHOULD_BOX: bool = std::mem::size_of::() > BOX_FUTURE_THRESHOLD;
}
〔设计推断与架构权衡〕
注释里解释了为什么用关联常量而非运行时 if:如果用运行时判断,spawn_named 会被单态化两次(一次针对 F,一次针对 Pin<Box<F>>),导致每个 spawn 的 future 都生成两份任务 harness,代码体积翻倍。用常量分支后,单态化收集器只保留实际走到的那个分支。

设计思考:装配顺序、错误恢复与生产踩坑

顺序即契约。装配顺序 driver -> blocking_pool -> scheduler 不是随意的。driver 最先创建,因为它是唯一可能因 OS 资源不足而失败、且失败后无需清理其他组件的步骤。blocking_pool 在 driver 之后、scheduler 之前,因为 scheduler 需要 blocking_spawner。如果 blocking_pool 创建失败(实际上它不太会失败),driver 会被 drop 自动清理。

current_thread 的 local_tid 分支。build_local 走的是 build_current_thread_local_runtime,它把当前线程 ID 传进去:

📎 tokio/src/runtime/builder.rs:1738-1751

rust
fn build_current_thread_local_runtime(&mut self) -> io::Result {
    use crate::runtime::local_runtime::LocalRuntimeScheduler;

    let tid = std::thread::current().id();

    let (scheduler, handle, blocking_pool) =
        self.build_current_thread_runtime_components(Some(tid))?;

    Ok(LocalRuntime::from_parts(
        LocalRuntimeScheduler::CurrentThread(scheduler),
        handle,
        blocking_pool,
    ))
}

这个 tid 被存进 Handle,后续 can_spawn_local_on_local_runtime 用它校验「spawn_local 是否在 owner 线程上调用」:

📎 tokio/src/runtime/scheduler/mod.rs:140-147

rust
pub(crate) fn can_spawn_local_on_local_runtime(&self) -> bool {
    match self {
        Handle::CurrentThread(h) => h.local_tid.is_some_and(|x| std::thread::current().id() == x),

        #[cfg(feature = "rt-multi-thread")]
        Handle::MultiThread(_) => false,
    }
}
〔设计推断与架构权衡〕
这是 LocalRuntime 安全性的基石:!Send 的 future 只能在其 owner 线程上被 poll,而 local_tid 就是这个约束的运行时检查点。如果去掉这个检查,跨线程 spawn_local 会导致 !Send 数据被并发访问,引发 UB。

生产踩坑一:worker_threads(0) 会 panic。worker_threads 方法有断言:

📎 tokio/src/runtime/builder.rs:582-586

rust
pub fn worker_threads(&mut self, val: usize) -> &mut Self {
    assert!(val > 0, "Worker threads cannot be set to 0");
    self.worker_threads = Some(val);
    self
}

这个断言在配置阶段就失败,而不是等到 build。好处是错误定位更早,坏处是如果线程数来自配置文件的动态值,用户必须在调用前自己校验。

生产踩坑二:max_blocking_threads 设太小会挂起。文档明确警告:

📎 tokio/src/runtime/builder.rs:600-601

rust
/// It's recommended to not set this limit too low in order to avoid hanging on operations
/// requiring [`spawn_blocking`].
〔设计推断与架构权衡〕
因为 blocking pool 的队列没有背压——任务会一直堆积直到有线程可用。如果所有阻塞线程都在等待某个「需要新阻塞线程才能完成」的操作,就会死锁。文档里「the queue does not apply any backpressure, it could potentially grow unbounded」正是这个风险的注脚。

生产踩坑三:UnhandledPanic::ShutdownRuntime 只支持 current_thread。

📎 tokio/src/runtime/builder.rs:1374-1381

rust
pub fn unhandled_panic(&mut self, behavior: UnhandledPanic) -> &mut Self {
    if !matches!(self.kind, Kind::CurrentThread) && matches!(behavior, UnhandledPanic::ShutdownRuntime) {
        panic!("UnhandledPanic::ShutdownRuntime is only supported in current thread runtime");
    }

    self.unhandled_panic = behavior;
    self
}
〔设计推断与架构权衡〕
这个限制的原因是:multi_thread 下「立即关闭运行时」需要协调所有 worker 线程的停止,实现复杂度高且语义模糊(正在 poll 的其他任务怎么办?)。current_thread 只有一个线程,关闭语义清晰。

本章小结

本章追踪了 Builder::build 的完整装配链路。核心结论:

1. Builder 是纯配置容器,build() 才创建资源。装配顺序 driver -> blocking_pool -> scheduler 由错误恢复需求决定。

2. current_thread 与 multi_thread 的差异不止线程数:blocking pool 容量计算不同(max_blocking_threads vs max_blocking_threads + worker_threads),multi_thread 多一个 launch 两阶段启动,enable_eager_driver_handoff 在 current_thread 下被强制关闭。

3. Handle 是跨组件共享的核心,内部用 Arc 包裹形态特定的句柄,通过 match 或 match_flavor! 宏统一访问。

4. AutoBox 用关联常量在编译期决定是否装箱 future,避免代码体积翻倍。

5. local_tid 是 LocalRuntime 安全性的运行时检查点。

下一章,我们将进入任务的生命周期:spawn 如何把一个 Future 变成可调度实体,JoinHandle 如何与任务状态机交互,以及任务在 PENDING / RUNNING / COMPLETE 之间的状态迁移。

本章思考与自测

Q1: 如果把 build_threaded_runtime 中 create_blocking_pool 的容量参数从 self.max_blocking_threads + worker_threads 改成 self.max_blocking_threads,在什么场景下会导致阻塞任务饿死?为什么 current_thread 路径可以传 self.max_blocking_threads?

参考解析:根据 📎 tokio/src/runtime/builder.rs:2189-2192,multi_thread 路径传入 self.max_blocking_threads + worker_threads,而 current_thread 路径 📎 tokio/src/runtime/builder.rs:1765 传入 self.max_blocking_threads。差异的根源在于:multi_thread 下,worker 线程本身也会执行阻塞任务(例如 block_in_place 会把 worker 线程临时转成阻塞线程),所以阻塞线程的总预算必须包含 worker 线程数。如果改成只传 self.max_blocking_threads,当 max_blocking_threads 设得较小(比如 1)且已有 worker 线程在 block_in_place 中占用预算时,新的 spawn_blocking 任务将无线程可用,堆积在无背压队列中,导致依赖这些阻塞任务的 async 任务永久挂起。current_thread 只有一个线程且不支持 block_in_place 的 worker 转换语义,所以不需要加上 worker 数。

Q2: MultiThread::new 返回 launch 句柄,真正启动 worker 线程的是 launch.launch()。如果去掉 handle.enter() 这一行直接调用 launch.launch(),会发生什么?

参考解析:根据 📎 tokio/src/runtime/builder.rs:2230-2232,启动前有 let _enter = handle.enter(); 然后才 launch.launch()。handle.enter() 的作用是设置线程本地上下文(thread-local),让当前线程「看起来」处于运行时内部。worker 线程启动后会立即开始 poll 任务,而任务代码可能调用 Handle::current()、tokio::spawn 等依赖上下文的 API。如果去掉 _enter,worker 线程在启动瞬间的上下文设置可能不完整(取决于 launch 内部是否自行设置),最坏情况下 worker 线程上执行的初始化代码调用 Handle::current() 会 panic(CONTEXT_MISSING_ERROR)。即使 launch 内部为每个 worker 设置了上下文,_enter 也保证了「启动动作本身」发生在正确的上下文中,避免启动过程中的竞态。

Q3: AutoBox::<F>::SHOULD_BOX 用关联常量而非运行时 if size_of::<F>() > THRESHOLD。假设改成运行时判断,除了代码体积翻倍,还会在什么情况下导致性能退化?

参考解析:根据 📎 tokio/src/runtime/mod.rs:657-673 的注释,运行时 if 会让 spawn_named 对每个 T 单态化两次(T 和 Pin<Box<T>> 各一次)。除了代码体积翻倍,性能退化体现在:1) 指令缓存(i-cache)压力增大,因为两套 harness 代码都要驻留;2) 编译器无法对「实际只走一个分支」做优化,运行时分支预测虽然通常准确,但分支本身和两套代码的寄存器分配差异会累积;3) 更隐蔽的是,Pin<Box<T>> 路径会强制堆分配,如果运行时判断因为某种原因(比如 size_of 在泛型上下文中未完全常量折叠)误判,小 future 也会被装箱,每次 spawn 多一次堆分配。关联常量让单态化收集器在编译期就剪掉未走的分支,零运行时开销。

AI 赋能代码库精读 · 本地优先架构

读完了本章?为你自己的私有项目生成专属架构全景书

基于 Tauri 2 + Rust 本地原生引擎,100% 源码离线隐私安全,零代码上传云端。像阅读一本传世专著一样拆解你的复杂系统。

⚡ Tauri 2 · Rust 原生引擎 · 100% 离线私密安全 · 适配超百万行代码库

CHAPTER 03

第 3 章:一粒任务的诞生:spawn 如何把一个 Future 变成可调度实体

官方源: tokio-rs/tokio · Commit @e800714a · 全书进度: 第 3 / 14 章

第 3 章:一粒任务的诞生:spawn 如何把一个 Future 变成可调度实体

上一章我们完成了 Runtime 的装配:I/O driver、time driver、blocking pool 与调度器被注入同一个 Runtime 实例,Handle 成为跨线程访问这些组件的共享句柄。但装配好的运行时此时还只是一个空壳——它拥有驱动任务的引擎,却没有任何任务可驱动。本章要回答的问题正是:当你敲下 tokio::spawn(async { ... }) 的那一刻,那个 async 块究竟经历了什么,才从一段普通的 Rust 代码变成一个「可被调度器接管、可被唤醒、可被 join」的实体。这是「任务的一生」的上半场,我们聚焦于诞生:从 Handle::spawn 出发,穿过 new_task 的引用计数分配,落到 Cell<T, S> 的内存布局,最终看清任务如何被投递到某个 worker 的本地队列或全局注入队列。下半场(第 4 章)才会进入调度循环与 poll/wake 闭环。

3.1 Future 不是任务:一次 spawn 到底创造了什么

直觉模型

把 Future 想象成一张「菜谱」,把任务想象成「厨房里正在被烹饪的一道菜」。菜谱本身是静态的、可复制的、没有任何执行状态;只有当厨房(调度器)决定「现在做这道菜」,给它分配一个灶台(worker)、一个订单号(TaskId)、一个出餐口(JoinHandle),它才成为一道「在制菜品」。若没有这层包装,调度器就无从知道「这道菜做到哪一步了」「谁在等它」「做好了通知谁」——它只能看到一张菜谱,无法管理。

数据结构与内存布局

Tokio 用 Task<S> 表示「被运行时拥有的任务引用」,它是对 RawTask 的透明包装:

rust
#[repr(transparent)]
pub(crate) struct Task {
    raw: RawTask,
    _p: PhantomData,
}

📎 tokio/src/runtime/task/mod.rs:233-238

#[repr(transparent)] 意味着 Task<S> 与 RawTask 在内存上完全一致,没有额外开销。PhantomData<S> 只是编译期的类型标记,标记这个任务属于哪个调度器类型 S。

真正承载任务全部状态的是 Cell<T, S>,它的布局是整个任务模块的基石:

rust
#[repr(C)]
pub(super) struct Cell {
    pub(super) header: Header,
    pub(super) core: Core,
    pub(super) trailer: Trailer,
}

📎 tokio/src/runtime/task/core.rs:126-136

三个字段按「热-温-冷」排列。Header 是热数据(每次调度、每次状态转换都要访问),Core 是温数据(poll 时访问),Trailer 是冷数据(仅在创建与销毁时访问)。注释明确写道:Header 必须是第一个字段,因为任务结构体会同时被 *mut Cell 和 *mut Header 引用 📎 tokio/src/runtime/task/core.rs:37-43。

更关键的是缓存行对齐。Cell 上挂着一长串 #[cfg_attr(..., repr(align(...)))],按目标架构选择对齐字节数:x86_64/aarch64/powerpc64 用 128 字节,arm/mips/sparc/hexagon 用 32 字节,m68k 用 16 字节,s390x 用 256 字节,其余默认 64 字节 📎 tokio/src/runtime/task/core.rs:64-125。注释解释了为什么 x86_64 要用 128 而非 64:从 Intel Sandy Bridge 起,空间预取器会一次拉取成对的 64 字节缓存行,所以必须对齐到 128 字节才能避免伪共享 📎 tokio/src/runtime/task/core.rs:45-53。

〔设计推断与架构权衡〕
这个对齐策略的代价是每个任务至少浪费一个缓存行的空间。但任务状态位(state)会被多个 worker 线程高频读写——一个线程在 poll 时设置 RUNNING 位,另一个线程在唤醒时读 NOTIFIED 位——若两个任务的状态位落在同一缓存行,每次状态转换都会触发缓存行在核心间来回弹跳(cache line ping-pong),性能损失远超内存浪费。Tokio 选择用空间换时间。

Header 本身被约束在 8 个指针大小以内:

rust
#[test]
#[cfg(not(loom))]
fn header_lte_cache_line() {
    assert!(std::mem::size_of::() ());
}

📎 tokio/src/runtime/task/core.rs:591-593

这个测试确保 Header 不会超过 64 字节(8 × 8),从而在 64 字节缓存行的架构上能完整落入一行。Header 的字段包括:state: State(原子状态位)、queue_next: UnsafeCell<Option<NonNull<Header>>>(注入队列的链表指针)、vtable: &'static Vtable(函数指针表)、owner_id: UnsafeCell<Option<NonZeroU64>>(所属 OwnedTasks 列表的 ID)、scheduled_at: UnsafeCell<ScheduleLatencyInstant>(调度延迟测量)📎 tokio/src/runtime/task/core.rs:169-198。

Core<T, S> 持有调度器句柄 scheduler: S、任务 ID task_id: Id,以及最核心的 stage: CoreStage<T> 📎 tokio/src/runtime/task/core.rs:148-165。Stage 是一个三态枚举:

rust
#[repr(C)]
pub(super) enum Stage {
    Running(T),
    Finished(super::Result),
    Consumed,
}

📎 tokio/src/runtime/task/core.rs:225-229

这正是「Future 与 Output 复用同一块内存」的关键:任务运行期间 Stage::Running 持有 future,完成后原地替换为 Stage::Finished(output),被 JoinHandle 取走后变为 Stage::Consumed。#[repr(C)] 注释指向一个 Miri issue,说明这个布局对 unsafe 代码的正确性有硬性要求 📎 tokio/src/runtime/task/core.rs:225-229。

Trailer 存放冷数据:owned: linked_list::Pointers<Header>(OwnedTasks 链表指针)、waker: UnsafeCell<Option<Waker>>(等待任务完成的消费者 waker)、hooks: TaskHarnessScheduleHooks 📎 tokio/src/runtime/task/core.rs:205-213。

Step-by-Step:从 spawn 到入队

我们代入一个具体场景:在 multi_thread 运行时中,worker 线程 A 执行 tokio::spawn(async { 42 })。

第一步:构造任务三件套。 new_task 是任务诞生的唯一入口:

rust
fn new_task(
    task: T,
    scheduler: S,
    id: Id,
    spawned_at: SpawnLocation,
) -> (Task, Notified, JoinHandle)

📎 tokio/src/runtime/task/mod.rs:336-346

它调用 RawTask::new::<T, S> 分配 Cell,然后从同一个 raw 指针派生出三个引用:Task(owned 引用,通常立即放入 OwnedTasks)、Notified(通知引用,交给调度器)、JoinHandle(结果读取句柄)📎 tokio/src/runtime/task/mod.rs:347-363。注意三者共享同一个 raw,各自持有一个引用计数。

第二步:分配 Cell 并写入初始状态。 Cell::new 在堆上分配整个结构:

rust
let result = Box::new(Cell {
    trailer: Trailer::new(scheduler.hooks()),
    header: new_header(state, vtable, ...),
    core: Core {
        scheduler,
        stage: CoreStage {
            stage: UnsafeCell::new(Stage::Running(future)),
        },
        task_id,
        ...
    },
});

📎 tokio/src/runtime/task/core.rs:261-278

vtable 由 raw::vtable::<T, S>() 生成,是一张针对具体 T 和 S 单态化的函数指针表 📎 tokio/src/runtime/task/core.rs:260。future 被直接移入 Stage::Running,没有额外装箱。

第三步:debug 断言验证布局。 在 debug_assertions 下,Cell::new 会调用 check 函数,用 Header::get_trailer、Header::get_scheduler、Header::get_id_ptr 等基于 vtable 偏移量的指针运算,逐一断言「通过 header 反查到的字段地址」与「实际字段地址」一致 📎 tokio/src/runtime/task/core.rs:280-321。这是对 vtable 偏移量正确性的运行时自检。

第四步:投递到调度器。 调度器拿到 Notified<S> 后,调用 Schedule::schedule 📎 tokio/src/runtime/task/mod.rs:315。在 multi_thread 下,这会走 push_back_or_overflow,把任务推入当前 worker 的本地队列,队列满时溢出到注入队列。

下面这张图刻画了从 new_task 到入队的控制流与分支:

mermaid
flowchart TD
    spawn_call["Handle::spawn(future)"] --> new_task["new_task::(future, scheduler, id)"]
    new_task --> raw_new["RawTask::new::"]
    raw_new --> cell_new["Cell::new: Box::new(Cell{header, core, trailer})"]
    cell_new --> vtable["raw::vtable::() 生成函数指针表"]
    cell_new --> stage["Stage::Running(future) 移入"]
    cell_new --> debug_check{"debug_assertions?"}
    debug_check -->|是| check_layout["check(): 断言 trailer/scheduler/id 偏移量"]
    debug_check -->|否| skip_check["跳过"]
    check_layout --> triple["派生 (Task, Notified, JoinHandle)"]
    skip_check --> triple
    triple --> owned["Task 存入 OwnedTasks"]
    triple --> sched["Notified 交给 Schedule::schedule"]
    sched --> push{"本地队列有容量?"}
    push -->|是| local_push["push_back_finish: 写入 buffer[tail & MASK]"]
    push -->|否| overflow_check{"steal == real?"}
    overflow_check -->|否, 有并发窃取| inject_only["overflow.push(task) 仅注入"]
    overflow_check -->|是| push_overflow["push_overflow: CAS 认领后半批"]
    push_overflow --> cas_ok{"CAS 成功?"}
    cas_ok -->|是| inject_batch["overflow.push_batch(后半批 + 当前 task)"]
    cas_ok -->|否| retry["返回 Err(task), 重试 push_back_or_overflow"]
    retry --> push

这张图揭示了几个关键分支:debug 断言只在调试构建生效;本地队列满时并非直接溢出,而是先判断是否有并发窃取者(steal != real),若有则只把当前任务推入注入队列,因为窃取者腾出的空间很快可用。

设计思考:为什么是三个引用而不是一个

new_task 返回三个引用,而非一个。这是引用计数设计的核心:Task 代表「运行时拥有这个任务」,Notified 代表「这个任务已被通知、待调度」,JoinHandle 代表「有人关心它的结果」。三者生命周期独立——JoinHandle 可以被 drop(任务继续运行,结果丢弃),Notified 在 poll 后消失,Task 在任务完成并从 OwnedTasks 移除后释放。若只有一个引用,就无法表达「任务还在跑但没人 join」这种状态。

UnownedTask 是另一个重要分支:它持有两个引用计数,用于 blocking 任务(不存入 OwnedTasks)📎 tokio/src/runtime/task/mod.rs:286-295。unowned 函数通过 mem::forget(task) 和 mem::forget(notified) 把两个引用合并进 UnownedTask 📎 tokio/src/runtime/task/mod.rs:388-397。这个「两个引用」的设计动机是:blocking 任务没有 OwnedTasks 列表来持有 owned 引用,所以需要额外一个引用计数来保证任务在运行期间不被释放。

3.2 状态位:一个 usize 如何编码任务的全部生命周期

直觉模型

把任务状态想象成一张「体检报告单」,上面有若干独立的勾选框:是否正在被 poll、是否已完成、是否被通知、是否被取消、是否有人 join。Tokio 没有用多个布尔字段,而是把这些勾选位压进一个 AtomicUsize。这样每次状态转换只需一次 CAS,而非多次加锁。若没有这个设计,任务状态转换会变成多把锁的嵌套,死锁风险与开销都会飙升。

位域布局

State 的位域在模块文档中有完整定义 📎 tokio/src/runtime/task/mod.rs:32-53:

  • RUNNING:任务是否正在被 poll 或取消。这一位同时充当任务的锁 📎 tokio/src/runtime/task/mod.rs:37-38。
  • COMPLETE:future 已完全完成并被 drop。一旦置位永不清除,且永不与 RUNNING 同时置位 📎 tokio/src/runtime/task/mod.rs:40-41。
  • NOTIFIED:当前是否存在一个 Notified 对象 📎 tokio/src/runtime/task/mod.rs:43。
  • CANCELLED:任务应尽快被取消 📎 tokio/src/runtime/task/mod.rs:45-46。
  • JOIN_INTEREST:存在 JoinHandle 📎 tokio/src/runtime/task/mod.rs:48。
  • JOIN_WAKER:作为 join handle waker 的访问控制位 📎 tokio/src/runtime/task/mod.rs:50-51。

剩余位用于引用计数 📎 tokio/src/runtime/task/mod.rs:53。

RUNNING 位充当锁这一点值得展开。模块文档的 Safety 章节指出:对 future 的任何可变访问都必须在修改 RUNNING 位获得锁之后进行,从而保证独占访问 📎 tokio/src/runtime/task/mod.rs:130-133。这意味着 poll 一个任务时,线程先 CAS 设置 RUNNING,成功后独占 future;若失败说明别的线程正在 poll,本次 poll 直接返回。这把「poll 的互斥」与「状态转换」合并成一次原子操作,避免了单独的互斥锁。

JOIN_WAKER 的访问控制协议

JOIN_WAKER 位是整个状态机中最精妙的部分。它解决的问题是:waker 字段(在 Trailer 中)会被两个线程并发访问——运行时在任务完成时读它来唤醒 join 者,JoinHandle 在 poll 时写它来注册 waker。模块文档给出了 7 条规则 📎 tokio/src/runtime/task/mod.rs:75-120:

1. JOIN_WAKER 初始为 0。

2. 为 0 时,JoinHandle 对 waker 字段有独占(可变)访问权。

3. 为 1 时,JoinHandle 只有共享(只读)访问权。

4. 为 1 且 COMPLETE 为 1 时,运行时对 waker 字段有共享(只读)访问权。

5. JoinHandle 要写 waker,必须:(i) 成功把 JOIN_WAKER 置 0 以获得独占权,(ii) 写入 waker,(iii) 成功把 JOIN_WAKER 置 1。

6. JoinHandle 只能在 COMPLETE 为 0 时改 JOIN_WAKER;运行时只能在 COMPLETE 为 1 时改。

7. 若 JOIN_INTEREST 为 0 且 COMPLETE 为 1,运行时对 waker 字段有独占访问权(用于 drop waker)。

规则 6 隐含了竞态:步骤 (i) 或 (iii) 可能失败。若 (i) 失败,放弃写 waker;若 (iii) 失败(另一线程在此期间置了 COMPLETE),则清空 waker 字段 📎 tokio/src/runtime/task/mod.rs:110-120。这套协议的本质是:用一个原子位在「写者」和「读者」之间动态转移所有权,避免为 waker 字段单独加锁。

引用计数的两种递减

Task 的 drop 递减一次引用计数,UnownedTask 的 drop 递减两次:

rust
impl Drop for Task {
    fn drop(&mut self) {
        if self.header().state.ref_dec() {
            self.raw.dealloc();
        }
    }
}

📎 tokio/src/runtime/task/mod.rs:580-586

rust
impl Drop for UnownedTask {
    fn drop(&mut self) {
        if self.raw.header().state.ref_dec_twice() {
            self.raw.dealloc();
        }
    }
}

📎 tokio/src/runtime/task/mod.rs:590-596

ref_dec 返回 true 表示这是最后一个引用,此时才真正释放 Cell 内存。ref_dec_twice 是 UnownedTask 持有两个计数的直接体现。

设计思考:为什么状态位与引用计数共用一个原子

〔设计推断与架构权衡〕
把状态位和引用计数放在同一个 AtomicUsize 里,是为了让「递减引用计数」与「设置状态位」这两个动作能在一次 CAS 中完成。模块文档在 Schedule::release 的注释中明确提到:「任务模块会批量处理 ref-dec 与其他选项的设置」📎 tokio/src/runtime/task/mod.rs:302-304。如果状态位和引用计数分属两个原子变量,那么「释放最后一个引用」与「标记完成」之间就会出现窗口,需要额外的同步。合并后,ref_dec 可以原子地完成「减计数 + 检查是否归零」,避免了 ABA 类问题。

3.3 JoinHandle:结果如何跨越任务边界回传

直觉模型

JoinHandle 就像餐厅给你的「取餐凭证」。任务(厨房)完成时,把菜品(output)放到出餐口(Stage::Finished),然后按响你的取餐器(waker)。你拿着凭证来取,凭证本身不持有菜品,只是指向出餐口的指针。若你把凭证丢了(drop JoinHandle),菜品会被直接倒掉(output 被 drop),但厨房不会因此停工。

数据结构

JoinHandle<T> 同样是对 RawTask 的透明包装:

rust
pub struct JoinHandle {
    raw: RawTask,
    _p: PhantomData,
}

📎 tokio/src/runtime/task/join.rs:163-166

PhantomData<T> 标记输出类型。JoinHandle<T> 在 T: Send 时才是 Send/Sync 📎 tokio/src/runtime/task/join.rs:169-170,这保证了非 Send 输出不会被跨线程移动。

Step-by-Step:await 一个 JoinHandle

JoinHandle 实现了 Future,其 poll 是结果回传的核心:

rust
fn poll(self: Pin, cx: &mut Context) -> Poll {
    ready!(crate::trace::trace_leaf());
    let mut ret = Poll::Pending;
    let coop = ready!(crate::task::coop::poll_proceed(cx));
    unsafe {
        self.raw.try_read_output(&mut ret, cx.waker());
    }
    if ret.is_ready() {
        coop.made_progress();
    }
    ret
}

📎 tokio/src/runtime/task/join.rs:327-354

注意几个细节:trace_leaf 用于 tracing 插桩;coop::poll_proceed 消耗协作预算(第 12 章详述);try_read_output 通过 vtable 擦除泛型,把返回值放在栈上、用 *mut () 传入 📎 tokio/src/runtime/task/join.rs:327-354。这个「返回值放栈上」的技巧是因为 vtable 函数无法泛型化返回类型 T,只能通过裸指针回写。

〔设计推断与架构权衡〕
try_read_output 内部逻辑(在 raw.rs 中,本章未提供源码):先检查 COMPLETE 位,若已置位则调用 take_output 取走 Stage::Finished 中的结果;否则把 cx.waker() 注册到 Trailer::waker 字段,返回 Pending。注册过程正是走 3.2 节的 JOIN_WAKER 协议。

结果的所有权转移

模块文档的「Non-Send output」章节精确描述了结果的所有权规则 📎 tokio/src/runtime/task/mod.rs:151-170:

  • 任务完成时,output 被放入 Stage,然后执行「设置 COMPLETE」的转换,并读取此刻的 JOIN_INTEREST 值。
  • 若 JOIN_INTEREST 为 0(无 JoinHandle),output 立即被 drop 📎 tokio/src/runtime/task/mod.rs:157-158。
  • 若 JOIN_INTEREST 为 1,JoinHandle 负责清理 output 📎 tokio/src/runtime/task/mod.rs:160-161。

对非 Send output,文档给出了三步论证:output 在 poll future 的线程上创建;JoinHandle<Output> 在 Output 非 Send 时也非 Send,所以它也在 spawn 线程上;因此 JoinHandle 取走或 drop output 时不会跨线程移动 📎 tokio/src/runtime/task/mod.rs:164-170。

JoinHandle 的 drop:快慢两条路径

rust
impl Drop for JoinHandle {
    fn drop(&mut self) {
        if self.raw.state().drop_join_handle_fast().is_ok() {
            return;
        }
        self.raw.drop_join_handle_slow();
    }
}

📎 tokio/src/runtime/task/join.rs:358-364

drop_join_handle_fast 尝试用一次 CAS 完成「清除 JOIN_INTEREST 位 + 递减引用计数」。若失败(例如任务正在完成,状态位被占用),则走 drop_join_handle_slow 的慢路径。这是典型的「乐观快路径 + 悲观慢路径」模式。

设计思考:为什么 JoinHandle 不直接持有 output

〔设计推断与架构权衡〕
若 JoinHandle 直接持有 output,那么 output 必须在任务完成时被移动到 JoinHandle 所在线程。但 JoinHandle 可能被移动到任意线程(只要 T: Send),而 output 的产生线程是 poll 线程。直接持有会导致「output 在 poll 线程产生,却要在 join 线程 drop」的跨线程移动,对非 Send output 直接违反类型系统。Tokio 选择让 output 留在 Cell 中(Stage::Finished),JoinHandle 只持有指向 Cell 的 RawTask,取结果时通过 take_output 原地取走。这样 output 的 drop 发生在 JoinHandle 所在线程,但前提是该线程与 poll 线程相同(非 Send 场景下成立)。

3.4 本地队列:work-stealing 的生产者-消费者结构

直觉模型

每个 worker 有一个「私人待办清单」(本地队列),容量 256。worker 自己从头部取任务(LIFO,利用缓存局部性),其他 worker 从尾部窃取任务(FIFO,取走最老的、最可能已完成的任务)。若没有本地队列,所有任务都挤在全局队列,每次取任务都要竞争全局锁,多核扩展性会崩溃。

内存布局:head 与 tail 的分离

rust
pub(crate) struct Inner {
    head: AtomicUnsignedLong,
    tail: AtomicUnsignedShort,
    buffer: Box>>; LOCAL_QUEUE_CAPACITY]>,
}

📎 tokio/src/runtime/scheduler/multi_thread/queue.rs:36-57

head 是 AtomicUnsignedLong(64 位,若平台支持 u64),tail 是 AtomicUnsignedShort(32 位)。注释解释了为什么索引比实际需要更宽:为了 ABA 缓解,以及区分「满」和「空」缓冲区 📎 tokio/src/runtime/scheduler/multi_thread/queue.rs:37-49。

head 内部打包了两个 UnsignedShort:低位是「真实头部」(real head),高位是「窃取者正在处理的第一个位置」(steal head)。当两者相等时,没有活跃的窃取者 📎 tokio/src/runtime/scheduler/multi_thread/queue.rs:39-49。这个双值打包是 work-stealing 队列的核心技巧:窃取者先 CAS 更新 steal 值来「认领」一批任务,完成后把 steal 值追上 real 值,表示窃取结束。

LOCAL_QUEUE_CAPACITY 在非 loom 下是 256,loom 下缩到 4 以便测试更多边界 📎 tokio/src/runtime/scheduler/multi_thread/queue.rs:62-69。MASK = LOCAL_QUEUE_CAPACITY - 1,用于环形缓冲区索引 📎 tokio/src/runtime/scheduler/multi_thread/queue.rs:71。

Step-by-Step:push_back_or_overflow 的完整分支

这是本地队列最复杂的函数,我们逐分支解析:

rust
pub(crate) fn push_back_or_overflow>(
    &mut self,
    mut task: task::Notified,
    overflow: &O,
    stats: &mut Stats,
) {
    let tail = loop {
        let head = self.inner.head.load(Acquire);
        let (steal, real) = unpack(head);
        let tail = unsafe { self.inner.tail.unsync_load() };

        if tail.wrapping_sub(steal)  return,
                Err(v) => { task = v; }
            }
        }
    };
    self.push_back_finish(task, tail);
}

📎 tokio/src/runtime/scheduler/multi_thread/queue.rs:188-223

三条分支:

1. 有容量(tail - steal < CAPACITY):break tail,跳出循环后调用 push_back_finish 写入缓冲区。

2. 无容量但有并发窃取者(steal != real):窃取者会腾出空间,所以只把当前任务推入注入队列,立即返回 📎 tokio/src/runtime/scheduler/multi_thread/queue.rs:204-208。

3. 无容量且无窃取者:调用 push_overflow 把后半批任务溢出到注入队列 📎 tokio/src/runtime/scheduler/multi_thread/queue.rs:209-219。若 CAS 失败(输给并发窃取者),push_overflow 返回 Err(task),循环重试。

push_back_finish 写入任务并更新 tail:

rust
fn push_back_finish(&self, task: task::Notified, tail: UnsignedShort) {
    let idx = tail as usize & MASK;
    self.inner.buffer[idx].with_mut(|ptr| {
        unsafe { ptr::write((*ptr).as_mut_ptr(), task); }
    });
    self.inner.tail.store(tail.wrapping_add(1), Release);
}

📎 tokio/src/runtime/scheduler/multi_thread/queue.rs:226-244

Release 序保证写入的任务对窃取者可见。

push_overflow:为什么溢出后半批

rust
const NUM_TASKS_TAKEN: UnsignedShort = (LOCAL_QUEUE_CAPACITY / 2) as UnsignedShort;

📎 tokio/src/runtime/scheduler/multi_thread/queue.rs:265

溢出时取走 128 个任务。注释详细解释了为什么取后半批而非前半批 📎 tokio/src/runtime/scheduler/multi_thread/queue.rs:295-306:从注入队列取任务时,总是放在前半部分。所以若一个任务在后半部分,就能确定它不是刚从注入队列取来的。这保证了「从注入队列取出的任务不会被立刻放回注入队列」(至少在被 poll 一次之前)。

CAS 认领后半批:

rust
if self.inner.head.compare_exchange_weak(
    pack(head, head), pack(tail, tail), Release, Relaxed
).is_err() {
    return Err(task);
}

📎 tokio/src/runtime/scheduler/multi_thread/queue.rs:283-293

把 head 从 (head, head) 更新到 (tail, tail),即同时推进 steal 和 real 到 tail,认领全部任务。成功后把 tail 回退到 tail + NUM_TASKS_TAKEN,表示前半批仍留在本地队列 📎 tokio/src/runtime/scheduler/multi_thread/queue.rs:314-316。

pop 与 steal_into:取任务的两条路径

pop 是 worker 自己取任务(从头部,LIFO):

rust
pub(crate) fn pop(&mut self) -> Option> {
    let mut head = self.inner.head.load(Acquire);
    let idx = loop {
        let (steal, real) = unpack(head);
        let tail = unsafe { self.inner.tail.unsync_load() };
        if real == tail { return None; }
        let next_real = real.wrapping_add(1);
        let next = if steal == real {
            pack(next_real, next_real)
        } else {
            assert_ne!(steal, next_real);
            pack(steal, next_real)
        };
        let res = self.inner.head.compare_exchange_weak(head, next, AcqRel, Acquire);
        match res {
            Ok(_) => break real as usize & MASK,
            Err(actual) => head = actual,
        }
    };
    Some(self.inner.buffer[idx].with(|ptr| unsafe { ptr::read(ptr).assume_init() }))
}

📎 tokio/src/runtime/scheduler/multi_thread/queue.rs:361-399

关键分支:若 steal == real(无窃取者),同时推进两者;否则只推进 real,保留 steal 不动 📎 tokio/src/runtime/scheduler/multi_thread/queue.rs:377-384。assert_ne!(steal, next_real) 确保不会把 real 推进到 steal 的位置,否则会破坏窃取者的认领状态。

steal_into 是窃取路径,先检查目标队列是否有足够空间:

rust
if dst_tail.wrapping_sub(steal) > LOCAL_QUEUE_CAPACITY as UnsignedShort / 2 {
    return None;
}

📎 tokio/src/runtime/scheduler/multi_thread/queue.rs:431-435

目标队列超过半满就不窃取,避免窃取后立刻又溢出。

steal_into2 是窃取的核心,计算窃取数量:

rust
let n = src_tail.wrapping_sub(src_head_real);
let n = n - n / 2;

📎 tokio/src/runtime/scheduler/multi_thread/queue.rs:487-488

窃取一半(向上取整)。然后 CAS 更新 head 的 steal 值来认领:

rust
let steal_to = src_head_real.wrapping_add(n);
next_packed = pack(src_head_steal, steal_to);
let res = self.0.head.compare_exchange_weak(prev_packed, next_packed, AcqRel, Acquire);

📎 tokio/src/runtime/scheduler/multi_thread/queue.rs:496-506

注意这里只更新了 real 值(pack(src_head_steal, steal_to) 中 steal 保持不变),把 real 推进到 steal_to。这表示「这些任务已被认领,其他窃取者不能再碰」。窃取完成后,再把 steal 追上 real:

rust
loop {
    let head = unpack(prev_packed).1;
    next_packed = pack(head, head);
    let res = self.0.head.compare_exchange_weak(prev_packed, next_packed, AcqRel, Acquire);
    match res {
        Ok(_) => return n,
        Err(actual) => prev_packed = actual,
    }
}

📎 tokio/src/runtime/scheduler/multi_thread/queue.rs:548-561

下面这张时序图刻画了「生产者 push、消费者 pop、窃取者 steal」三方并发交互:

mermaid
sequenceDiagram
    participant P as "Worker A (生产者)"
    participant Q as "Local 队列 Inner"
    participant C as "Worker A (消费者 pop)"
    participant S as "Worker B (窃取者)"

    P->>Q: "load head (Acquire)"
    P->>Q: "unsync_load tail"
    Note over P: "tail - steal >Q: "push_back_finish: buffer[idx] = task"
    P->>Q: "store tail+1 (Release)"

    C->>Q: "load head (Acquire)"
    C->>Q: "unsync_load tail"
    Note over C: "real == tail? 空则返回 None"
    C->>Q: "CAS head: pack(real+1, real+1)"
    Q-->>C: "Ok, 读取 buffer[real & MASK]"

    S->>Q: "load head (Acquire)"
    S->>Q: "load tail (Acquire)"
    Note over S: "src_head_steal != src_head_real? 返回 0"
    S->>Q: "CAS head: pack(steal, real+n) 认领一半"
    Q-->>S: "Ok, 拷贝 n 个任务到 dst"
    S->>Q: "CAS head: pack(real+n, real+n) 完成窃取"
    Q-->>S: "返回 n"

设计思考:为什么本地队列是 LIFO 而窃取是 FIFO

〔设计推断与架构权衡〕
worker 自己从头部取(LIFO),因为最近推入的任务最可能还在 CPU 缓存中,且最可能是「刚被唤醒、数据还热」的任务。窃取者从尾部取(FIFO),因为最老的任务最可能已经完成大部分工作,窃取它能最快减轻受害者负载。这种「LIFO 本地 + FIFO 窃取」的组合是 work-stealing 调度的经典设计,兼顾了缓存局部性与负载均衡。

至此,任务已经完成了从 Future 到可调度实体的蜕变:它被分配了引用计数、放进了 Cell 的内存布局,并成功投递到 worker 的本地队列或全局注入队列。但任务被放入队列只是开始,真正让它运转起来的是 worker 线程的调度循环。下一章我们将进入「任务的一生」下半场,追踪 worker 如何从队列中取出任务、调用 Future::poll,并在返回 Pending 时通过 Waker 注册唤醒,最终触发 schedule 重新入队——「唤醒 → 入队 → 再 poll」这一闭环的完整调用路径,以及 work-stealing 策略与 LIFO 槽位优化,都将在那里揭晓。

AI 赋能代码库精读 · 本地优先架构

读完了本章?为你自己的私有项目生成专属架构全景书

基于 Tauri 2 + Rust 本地原生引擎,100% 源码离线隐私安全,零代码上传云端。像阅读一本传世专著一样拆解你的复杂系统。

⚡ Tauri 2 · Rust 原生引擎 · 100% 离线私密安全 · 适配超百万行代码库

CHAPTER 04

第 4 章:调度循环的心跳:poll 循环与工作窃取 (Work-Stealing) 算法

官方源: tokio-rs/tokio · Commit @e800714a · 全书进度: 第 4 / 14 章

第 4 章:调度循环的心跳:poll 循环与工作窃取 (Work-Stealing) 算法

从队列到执行:worker 主循环的骨架

上一章我们把任务送进了 Local 队列或全局注入队列。但队列只是「待办清单」,真正让任务跑起来的,是 worker 线程里那个永不停止的循环。这一章我们追踪 Context::run——它是整个多线程调度器的心脏。

先建立直觉:worker 线程就像一个厨师,面前有一摞自己的订单(run_queue),旁边还有一个公共订单架(inject)。厨师先看自己手边最近的一张(lifo_slot),没有就从自己那摞拿,再没有就去公共架抓一把,还不行就去别的厨师那摞里偷几张。全都空了他才去休息,但休息时耳朵还竖着——一有订单进来就立刻醒来。

若没有这个循环,任务被入队后就永远躺在队列里,Future::poll 永远不会被调用,整个运行时就是一堆死数据。

Core 的内存布局与状态字段

worker 的可变状态全部装在 Core 里,它被 Box 分配在堆上,通过 AtomicCell<Core> 在 Worker 与线程本地 Context 之间传递。

Core 的关键字段如下 📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:113-167:

  • tick: u32:每次循环自增,用于周期性触发维护(maintenance)和全局队列检查。
  • lifo_slot: Option<Notified>:LIFO 槽位,这是本章最精妙的设计。当 worker 自己调度一个任务时,它不进 run_queue,而是放进这个槽位,下次取任务时优先从这里拿。
  • lifo_enabled: bool:LIFO 槽位的开关,用于防止 ping-pong 场景下的饥饿。
  • run_queue: queue::Local<Arc<Handle>>:本地队列,上一章剖析过的 Local 结构。
  • is_searching: bool:worker 是否正在搜索可窃取的任务。
  • is_shutdown: bool / is_traced: bool:关闭与追踪标志。
  • park: Option<Parker>:park 器,用 Option 包裹是为了在借用检查器下方便地取出/放回。
  • global_queue_interval: u32:多久检查一次全局队列。
  • rand: FastRand:快速随机数生成器,用于随机选择窃取起点。
〔设计推断与架构权衡〕
注意 lifo_slot 是 Option<Notified> 而非队列——它只存一个任务。这个设计动机在源码注释里说得很清楚 📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:117-121:worker 自己调度的任务存进这个槽位,worker 会在检查 run_queue 之前先检查它,效果是「最后被调度的任务下一个运行」(LIFO)。这是为了改善局部性,对消息传递模式特别有效,能降低延迟。

为什么 LIFO 能降低延迟?考虑一个典型的消息传递场景:任务 A 处理完消息后唤醒任务 B,B 处理完又唤醒 A。如果 A 唤醒 B 后 B 立刻运行,B 需要的数据很可能还在 CPU 缓存里(因为 A 刚碰过)。如果 B 被塞到队列尾部,等前面几十个任务跑完,缓存早被冲掉了。

但 LIFO 有饥饿风险。源码用 MAX_LIFO_POLLS_PER_TICK = 3 来限制 📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:263-263:每个 tick 最多优先 LIFO 槽位 3 次,超过就禁用,让其他任务有机会执行。

主循环 walkthrough:一次完整的调度周期

我们代入一个具体场景:worker 0 刚从 park 中醒来,run_queue 里有 5 个任务,lifo_slot 里有 1 个任务,全局队列有 3 个任务。

主循环入口是 Context::run 📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:570-642。它先重置 lifo_enabled(因为 core 可能被 block_in_place 偷走过,状态需要归位)📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:571-573,然后进入 while !core.is_shutdown 循环。

每轮循环做四件事:

第一步:tick 与维护。 core.tick() 自增计数器 📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:587。接着 self.maintenance(core) 检查 tick % event_interval == 0,若是则调用 park_yield 以 0 超时驱动 I/O 和定时器 📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:809-826。

第二步:取任务。 core.next_task(&self.worker) 是核心取任务逻辑 📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:1090-1156。它分两条路径:

  • 当 tick % global_queue_interval == 0 时,优先从全局队列取,取不到再取本地 📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:1091-1098。这是为了防止全局队列里的任务被饿死。
  • 否则优先取本地任务 📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:1090-1156。

本地取任务由 next_local_task 完成 📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:1158-1160:

rust
fn next_local_task(&mut self) -> Option {
    self.lifo_slot.take().or_else(|| self.run_queue.pop())
}

先取 LIFO 槽位,再取队列头部(LIFO 弹出)。这就是上一章说的「本地 LIFO」。

如果本地为空但全局队列非空,worker 会批量从全局队列拉取任务 📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:1110-1154。批量大小 n 的计算很讲究:min(inject.len() / remotes.len() + 1, cap),其中 cap 又取 min(remaining_slots, max_capacity / 2)。源码注释解释了为什么限制在队列容量的一半 📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:1120-1131:确保拉取的任务落在本地队列的前半部分,这样即使后续发生溢出,这些任务也不会被推回全局队列(溢出只影响后半部分)。

第三步:运行任务。 拿到任务后调用 run_task 📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:647-796。这是本章最复杂的函数,我们下一节专门展开。

第四步:窃取或 park。 如果 next_task 返回 None,说明本地和全局都没活了,调用 steal_work 📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:1167-1195。窃取失败则进入 park 或 park_yield 📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:613-621。

整个控制流如下:

mermaid
flowchart TD
    start["Context::run 进入循环"] --> tick["core.tick() 自增"]
    tick --> maint{"tick % event_interval == 0?"}
    maint -->|是| park_yield["park_yield 驱动 I/O 与定时器"]
    maint -->|否| next
    park_yield --> next["core.next_task()"]
    next --> has_task{"取到任务?"}
    has_task -->|是| run_task["run_task 执行 poll"]
    run_task --> cont{"core 还在?"}
    cont -->|是| tick
    cont -->|否| ret["return 退出"]
    has_task -->|否| steal["core.steal_work()"]
    steal --> steal_ok{"窃取成功?"}
    steal_ok -->|是| run_task
    steal_ok -->|否| defer_check{"defer 非空?"}
    defer_check -->|是| py["park_yield"]
    defer_check -->|否| pk["park 阻塞等待"]
    py --> tick
    pk --> tick

run_task:poll 与 LIFO 槽位的闭环

run_task 是任务真正被 poll 的地方,也是「唤醒 → 入队 → 再 poll」闭环的收口点。

进入函数后第一件事是 assert_owner 📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:648,把 Notified 转换成 Task,同时断言当前线程确实是这个任务的 owner(debug 断言)。

接着 transition_from_searching 📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:652——如果 worker 之前在搜索状态,现在找到任务了,要退出搜索状态,并可能唤醒其他 parked worker。

然后是关键的 budget 包裹 📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:695-795:

rust
coop::budget(|| {
    task.run();
    let mut lifo_polls = 0;
    loop {
        let mut core = match self.core.borrow_mut().take() {
            Some(core) => core,
            None => return ControlFlow::Break(()),
        };
        let task = match core.lifo_slot.take() {
            Some(task) => task,
            None => {
                self.reset_lifo_enabled(&mut core);
                core.stats.end_poll();
                return ControlFlow::Continue(core);
            }
        };
        if !coop::has_budget_remaining() {
            core.run_queue.push_back_or_overflow(task, ...);
            return ControlFlow::Continue(core);
        }
        lifo_polls += 1;
        if lifo_polls >= MAX_LIFO_POLLS_PER_TICK {
            core.lifo_enabled = false;
        }
        let task = self.worker.handle.shared.owned.assert_owner(task);
        *self.core.borrow_mut() = Some(core);
        task.run();
    }
})

这段代码揭示了 LIFO 槽位的完整闭环:task.run() 执行 Future::poll,poll 过程中如果任务唤醒了自己或别的任务,schedule_local 会把新任务放进 lifo_slot 📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:1396-1408。poll 返回后,循环立刻检查 lifo_slot,如果有任务就继续跑——不回到主循环,直接在同一个 budget 内连续 poll。

这就是「唤醒 → 入队 → 再 poll」在 LIFO 路径上的体现:唤醒时任务被放进 lifo_slot,poll 返回后立即被取出再 poll,形成紧密的闭环。

注意 self.core.borrow_mut().take() 的 None 分支 📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:716-724:如果 core 被偷走了(比如任务里调用了 block_in_place),worker 必须返回 ControlFlow::Break(()),让 Context::run 退出。这是 block_in_place 与调度循环的交互点。

唤醒路径:Waker 如何触发重新入队

当 Future::poll 返回 Pending 时,任务需要注册一个 Waker,等事件就绪时被唤醒。Tokio 的 Waker 实现极其精简——它就是一个指向任务 Header 的裸指针加一张 vtable。

waker_ref 构造 WakerRef 📎 tokio/src/runtime/task/waker.rs:11-34,用 ManuallyDrop 包裹 Waker 避免 drop 时减引用计数。vtable 是静态的 📎 tokio/src/runtime/task/waker.rs:119-119:

rust
static WAKER_VTABLE: RawWakerVTable =
    RawWakerVTable::new(clone_waker, wake_by_val, wake_by_ref, drop_waker);

四个函数都只是把裸指针还原成 Header,然后调用 RawTask 的对应方法 📎 tokio/src/runtime/task/waker.rs:70-116。比如 wake_by_ref 最终调用 raw.wake_by_ref() 📎 tokio/src/runtime/task/waker.rs:106-116。

wake_by_ref 的语义是:把任务状态从 PENDING 转为 SCHEDULED,如果转换成功(即之前确实是 PENDING),就调用 Schedule::schedule 把任务重新入队。

对于多线程调度器,schedule 的实现在 Handle::schedule_task 📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:1353-1376:

rust
pub(super) fn schedule_task(&self, task: Notified, is_yield: bool) {
    with_current(|maybe_cx| {
        if let Some(cx) = maybe_cx {
            if self.ptr_eq(&cx.worker.handle) {
                if let Some(core) = cx.core.borrow_mut().as_mut() {
                    self.schedule_local(core, task, is_yield);
                    return;
                }
            }
        }
        self.push_remote_task(task);
        self.notify_parked_remote();
    });
}

逻辑分两支:

  • 如果当前线程就是这个调度器的 worker,且持有 core,走 schedule_local 📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:1385-1417——放进 LIFO 槽位或本地队列。
  • 否则(从外部线程唤醒,或 core 被偷走),走 push_remote_task 推入全局注入队列,并 notify_parked_remote 唤醒一个 parked worker 📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:1379-1383。

schedule_local 内部又分两支 📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:1385-1417:如果是 yield 或 LIFO 已禁用,推入 run_queue 尾部;否则放进 lifo_slot,并把原来槽位里的任务挤到队列尾部。

mermaid
sequenceDiagram
    participant Future as "Future::poll"
    participant Waker as "Waker(wake_by_ref)"
    participant RawTask as "RawTask::wake_by_ref"
    participant Handle as "Handle::schedule_task"
    participant Core as "Core(schedule_local)"
    participant Inject as "InjectQueue"
    participant Parker as "Unparker"

    Future->>Waker: "返回 Pending, 注册 waker"
    Note over Future: "事件就绪(如 epoll)"
    Waker->>RawTask: "raw.wake_by_ref()"
    RawTask->>RawTask: "state: PENDING -> SCHEDULED"
    RawTask->>Handle: "schedule(Notified)"
    alt 当前线程是同一 worker 且持有 core
        Handle->>Core: "schedule_local: 放入 lifo_slot"
    else 外部线程或 core 被偷走
        Handle->>Inject: "push_remote_task"
        Handle->>Parker: "notify_parked_remote().unpark()"
    end

park 与 unpark:状态机与唤醒的原子性

worker 没活干时要 park,但 park/unpark 是最容易出竞态的地方。Tokio 用 AtomicUsize 状态机加 Condvar 兜底来解决。

Inner 的字段 📎 tokio/src/runtime/scheduler/multi_thread/park.rs:31-43:state: AtomicUsize、mutex: Mutex<()>、condvar: Condvar、shared: Arc<Shared>。状态常量有四个 📎 tokio/src/runtime/scheduler/multi_thread/park.rs:36-45:

  • EMPTY = 0:未 park。
  • PARKED_CONDVAR = 1:在 condvar 上 park。
  • PARKED_DRIVER = 2:在 I/O driver 上 park。
  • NOTIFIED = 3:已被唤醒。

这是一个显式状态机,我们用它画状态图(这是本章唯一符合 stateDiagram-v2 准入条件的地方——源码里确实有这四个状态常量):

mermaid
stateDiagram-v2
    [*] --> Empty
    Empty --> ParkedCondvar : "park_condvar() CAS(EMPTY->PARKED_CONDVAR)"
    Empty --> ParkedDriver : "park_driver() CAS(EMPTY->PARKED_DRIVER)"
    Empty --> Notified : "unpark() swap(NOTIFIED)"
    ParkedCondvar --> Empty : "condvar 唤醒后 CAS(NOTIFIED->EMPTY)"
    ParkedCondvar --> Empty : "超时 swap(EMPTY)"
    ParkedDriver --> Empty : "driver 返回后 swap(EMPTY)"
    Notified --> Empty : "park() CAS(NOTIFIED->EMPTY) 消费通知"
    Notified --> Notified : "再次 unpark() swap(NOTIFIED)"

unpark 的实现 📎 tokio/src/runtime/scheduler/multi_thread/park.rs:277-290 用 swap 而非 CAS,源码注释解释了原因 📎 tokio/src/runtime/scheduler/multi_thread/park.rs:277-290:必须执行 release 操作让 park 线程观察到 unpark 之前的写入,所以即使 state 已经是 NOTIFIED 也要写一次。

park 先尝试消费已有的通知 📎 tokio/src/runtime/scheduler/multi_thread/park.rs:132-149:如果 CAS NOTIFIED -> EMPTY 成功,说明之前已被唤醒,直接返回不阻塞。否则尝试拿 driver 锁,拿到就在 driver 上 park,拿不到就用 condvar 兜底 📎 tokio/src/runtime/scheduler/multi_thread/park.rs:143-148。

park_condvar 里有个经典的双重检查 📎 tokio/src/runtime/scheduler/multi_thread/park.rs:162-180:先 CAS EMPTY -> PARKED_CONDVAR,如果失败且是 NOTIFIED,说明在设置状态前就被唤醒了,此时必须 swap(EMPTY) 来同步 unpark 的写入 📎 tokio/src/runtime/scheduler/multi_thread/park.rs:167-177。注释特别强调:即使知道是 NOTIFIED 也必须读一次,因为 unpark 可能在我们读 NOTIFIED 之后又被调用了一次。

unpark_condvar 的注释 📎 tokio/src/runtime/scheduler/multi_thread/park.rs:292-307 点出了 condvar 的经典陷阱:parked 线程设置 PARKED 状态和真正 wait 之间有窗口期,如果在这期间 notify 会被忽略。解决方案是 park 线程此时持有 mutex,unpark 线程先 drop(self.mutex.lock()) 获取锁(从而等待 park 线程释放),再 notify_one。

设计思考:为什么 LIFO 槽位是单槽而非队列

〔设计推断与架构权衡〕
单槽设计是刻意的权衡。如果用队列,每次唤醒都要入队、每次取任务都要出队,开销更大;而且队列会积累多个任务,破坏「最近唤醒的最先跑」这个局部性假设。单槽的语义是「只记住最近一个」,被挤出的任务进普通队列——这恰好符合局部性收益递减的规律:最近一个任务最热,第二个次之,第三个往后收益就很小了。

MAX_LIFO_POLLS_PER_TICK = 3 这个魔数 📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:263-263 也是经验值。源码注释说「跑几次 LIFO 槽位似乎足以受益于局部性,超过 3 次可能过度加权」。这防止了 A 唤醒 B、B 唤醒 A 的 ping-pong 场景把其他任务饿死。

另一个值得注意的设计是 steal_work 的「半数搜索」策略 📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:1158-1160:只有当不到一半的 worker 在搜索时,新 worker 才真正尝试窃取。这避免了所有 worker 同时疯狂窃取导致的 CAS 争用。transition_to_searching 通过 idle.transition_worker_to_searching() 来协调 📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:1197-1203。

窃取从随机起点开始 📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:1172-1174,遍历所有 remote,跳过自己 📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:1179-1182,调用 steal_into 尝试窃取。全部失败后回退到全局队列 📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:1197-1203。

本章小结

worker 主循环 Context::run 是调度器的心脏:每轮 tick 后先取任务(LIFO 槽位 → 本地队列 → 全局队列),取到就 run_task 执行 poll,取不到就窃取,窃取失败就 park。run_task 内部的 LIFO 循环把「唤醒 → 入队 → 再 poll」压缩在同一个 budget 内,形成低延迟闭环。Waker 是裸指针加静态 vtable,wake_by_ref 通过状态转换触发 schedule,根据当前线程是否是同一 worker 决定走本地队列还是全局队列。park/unpark 用四状态原子机加 condvar 兜底,解决了唤醒丢失的经典竞态。

下一章我们将离开调度器,进入 I/O 世界:Reactor 如何把 epoll 事件翻译成 Waker 唤醒,让 AsyncFd 的 Pending 变成 Ready。

本章思考与自测

Q1: 如果把 next_local_task 改成先取 run_queue 再取 lifo_slot,在消息传递密集的场景下会有什么后果?

参考解析:next_local_task 当前实现是 self.lifo_slot.take().or_else(|| self.run_queue.pop()) 📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:1158-1160,先取 LIFO 槽位。如果反过来先取 run_queue,那么刚被唤醒、数据还热的任务会被排到队列里其他任务之后执行。在 A→B→A 的消息传递模式下,B 被唤醒后不会立即运行,而是等队列里其他任务跑完,此时 A 写入的数据可能已被挤出 CPU 缓存,局部性收益丧失。更严重的是,lifo_slot 里的任务会一直等到 run_queue 清空才被执行,延迟显著上升。源码注释 📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:117-121 明确指出这个顺序是为了「改善局部性,受益于消息传递模式并降低延迟」。

Q2: park_condvar 中,如果去掉 Err(NOTIFIED) 分支里的 self.state.swap(EMPTY, SeqCst),只保留 return,会有什么问题?

参考解析:源码在 Err(NOTIFIED) 分支里执行 let old = self.state.swap(EMPTY, SeqCst) 📎 tokio/src/runtime/scheduler/multi_thread/park.rs:167-177。注释解释 📎 tokio/src/runtime/scheduler/multi_thread/park.rs:168-173:unpark 可能在我们读到 NOTIFIED 之后又被调用了一次,必须执行一次 acquire 操作与那个 unpark 同步,才能观察到它之前的所有写入。如果只 return 不 swap,state 会停留在 NOTIFIED,下一次 park 时 CAS NOTIFIED -> EMPTY 会成功并立即返回(消费了一个已经过期的通知),但更糟的是 unpark 的 release 写入没有被同步,park 线程可能看不到 unpark 之前写入的数据,导致内存可见性问题。这是典型的「丢失唤醒 + 内存序」双重 bug。

Q3: run_task 中,当 self.core.borrow_mut().take() 返回 None 时为什么返回 ControlFlow::Break(()) 而不是 Continue?

参考解析:self.core.borrow_mut().take() 返回 None 意味着 core 已经被偷走 📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:716-724。core 被偷走的唯一途径是任务内部调用了 block_in_place,它会通过 maybe_move_runtime 把 core 从 cx.core 取出并交给新线程 📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:473-497。此时当前线程已经不再持有调度能力,如果返回 Continue,Context::run 会继续循环并调用 core.next_task() 等需要 core 的方法,但 core 已经不在 self.core 里了,会导致 panic 或状态不一致。返回 Break 让 Context::run 直接 return 📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:594-597,把控制权交还给 run 函数,由它处理后续(比如 cx.defer.wake() 📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:564)。注释也说明 📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:719-721:此时不能调用 reset_lifo_enabled,因为 core 被偷走了,偷走者会在 Context::run 顶部处理。

AI 赋能代码库精读 · 本地优先架构

读完了本章?为你自己的私有项目生成专属架构全景书

基于 Tauri 2 + Rust 本地原生引擎,100% 源码离线隐私安全,零代码上传云端。像阅读一本传世专著一样拆解你的复杂系统。

⚡ Tauri 2 · Rust 原生引擎 · 100% 离线私密安全 · 适配超百万行代码库

CHAPTER 05

第 5 章:I/O 驱动与就绪通知:Reactor 如何把 epoll 事件转化为 Waker 唤醒

官方源: tokio-rs/tokio · Commit @e800714a · 全书进度: 第 5 / 14 章

第 5 章:I/O 驱动与就绪通知:Reactor 如何把 epoll 事件转化为 Waker 唤醒

上一章我们追踪了 worker 线程的主循环:任务被 poll,返回 Pending 时把 Waker 存进某个地方,事件就绪后 Waker 被触发,任务重新入队。但「某个地方」到底是哪里?Waker 怎么在 epoll 事件到来时被找回来?这正是 Reactor 要回答的问题。先建立直觉模型:把整个 I/O 就绪通知机制想象成一家餐厅的取餐叫号系统——顾客(任务)点完餐后不会站在窗口死等,而是拿一个震动器(Waker)回座位;后厨(内核 epoll)做好餐后,前台(Reactor)根据订单号(Token)找到对应的震动器并按下按钮。若没有这套系统,每个任务只能轮询 socket,CPU 会被烧光;或者用阻塞线程等待,一个连接一个线程,规模上不去。Tokio 的 Reactor 由三个文件构成三层结构,职责严格分离:driver.rs 是事件循环本体,持有 mio::Poll,负责调用 poll() 阻塞等待内核事件,并把事件翻译成对 ScheduledIo 的读写;registration.rs 是面向用户的注册句柄,TcpStream 内部持有的就是它,提供 poll_read_ready / poll_write_ready 等 API;scheduled_io.rs 是每个 fd 的状态槽,存储读写就绪位与 Waker 列表,是事件与任务之间的桥梁。模块组装关系可参考 tokio/src/runtime/io/mod.rs:5-16:driver 导出 Driver、Handle、ReadyEvent,registration 导出 Registration,scheduled_io 导出 ScheduledIo。下图锚定了本章要追踪的完整数据流:TcpStream → Registration → ScheduledIo → Handle/Driver → 内核 → 回到 ScheduledIo → Waker。接下来我们逐层拆解。

驱动层:Driver 与 Handle 的职责切分

直觉模型

Driver 是唯一拥有 mio::Poll 的实体,它只能在单个线程里被 &mut 访问——这是事件循环的独占性要求。而 Handle 是可克隆、可跨线程共享的注册入口,任何线程想注册新 fd 都通过它。若没有这个切分,要么把 mio::Poll 加锁(每次注册都竞争),要么让所有注册都回到 driver 线程(引入跨线程消息队列)。Tokio 选择让 Handle 直接持有 mio::Registry 的克隆,注册操作可以并发进行,只有真正的事件等待才需要独占。

内存布局与字段

先看 Driver 的字段 📎 tokio/src/runtime/io/driver.rs:25-38:

  • signal_ready: bool:Unix 信号事件是否到达,用于 signal 驱动。
  • events: mio::Events:主事件缓冲区,跨 turn 调用复用,避免每次分配。
  • events_busy: Option<mio::Events>:非阻塞 poll 专用缓冲区,仅当 max_io_events_per_busy_tick 被设置时存在。
  • poll: mio::Poll:内核事件队列的封装。

再看 Handle 📎 tokio/src/runtime/io/driver.rs:41-75:

  • registry: mio::Registry:mio::Poll::registry() 的克隆,用于 register/deregister。
  • registrations: RegistrationSet:所有活跃注册的集合,负责分配 Token 与 ScheduledIo。
  • synced: Mutex<registration_set::Synced>:保护 RegistrationSet 的同步状态。
  • waker: mio::Waker:用于从任意线程唤醒阻塞在 turn 里的 driver。
  • metrics: IoDriverMetrics:统计 fd 数量、就绪事件数。

这里有个关键设计:events_busy 的存在 📎 tokio/src/runtime/io/driver.rs:25-38 是为了解决非阻塞 poll 会吞掉事件的问题。注释 📎 tokio/src/runtime/io/driver.rs:189-190 说得很清楚:非阻塞 poll 取走的事件如果留在主缓冲区里,下次 poll 就看不到了;用独立缓冲区,未处理的事件仍留在内核队列,下次 poll 会重新返回。

Step-by-Step:一次 turn 的执行

turn 是 driver 的核心函数 📎 tokio/src/runtime/io/driver.rs:184-261。假设 worker 线程发现没有任务可跑,调用 park → turn(handle, None) 阻塞等待:

第一步:断言未 shutdown 📎 tokio/src/runtime/io/driver.rs:185,并释放待清理的注册 📎 tokio/src/runtime/io/driver.rs:187。release_pending_registrations 检查 needs_release(),若有则调用 registrations.release() 📎 tokio/src/runtime/io/driver.rs:336-340。

第二步:选择事件缓冲区 📎 tokio/src/runtime/io/driver.rs:191-194。若 max_wait 是零且 events_busy 存在,用 busy 缓冲区;否则用主缓冲区。

第三步:调用 self.poll.poll(events, max_wait) 📎 tokio/src/runtime/io/driver.rs:198。这是真正阻塞在 epoll_wait 的地方。错误处理很克制:Interrupted 直接忽略(信号打断是正常的)📎 tokio/src/runtime/io/driver.rs:200,WASI 下的 InvalidInput 也忽略 📎 tokio/src/runtime/io/driver.rs:201-205,其他错误直接 panic 📎 tokio/src/runtime/io/driver.rs:206。

第四步:遍历事件 📎 tokio/src/runtime/io/driver.rs:211-233。对每个 event:

  • 若 token == TOKEN_WAKEUP(值为 0)📎 tokio/src/runtime/io/driver.rs:214,什么都不做——这是 unpark 用来打断阻塞的。
  • 若 token == TOKEN_SIGNAL(值为 1)📎 tokio/src/runtime/io/driver.rs:216,置 signal_ready = true。
  • 否则是普通 I/O 事件 📎 tokio/src/runtime/io/driver.rs:218-231:把 mio::Ready 转成 Tokio 的 Ready,用 EXPOSE_IO.from_exposed_addr(token.0) 把 token 还原成 *const ScheduledIo 指针,然后 set_readiness(Tick::Set, |curr| curr | ready) 累积就绪位,再 io.wake(ready) 触发对应方向的 Waker。

这里 EXPOSE_IO 是一个 PtrExposeDomain<ScheduledIo> 📎 tokio/src/runtime/io/mod.rs:21-22,它把指针「暴露」成一个 usize 作为 mio::Token。安全性注释 📎 tokio/src/runtime/io/driver.rs:222-225 说明了为什么这个 unsafe 转换是安全的:指针在从 mio 注销且 driver 不再并发 poll 之前不会被释放,且 driver 持有 Arc<ScheduledIo> 的所有权。

第五步:处理 io_uring 完成队列(仅 Linux + tokio_unstable)📎 tokio/src/runtime/io/driver.rs:235-258,包括 CQ 溢出时的 flush 循环。

第六步:累加 metrics 📎 tokio/src/runtime/io/driver.rs:265-267。

mermaid
flowchart TD
    start["turn(handle, max_wait)"] --> assert["debug_assert!(!is_shutdown)"]
    assert --> release["release_pending_registrations()"]
    release --> pick{"max_wait == 0且 events_busy 存在?"}
    pick -->|是| busy["events = events_busy"]
    pick -->|否| main["events = events"]
    busy --> poll["poll.poll(events, max_wait)"]
    main --> poll
    poll --> pollres{"poll 返回?"}
    pollres -->|"Ok / Interrupted"| iter["遍历 events.iter()"]
    pollres -->|"其他 Err"| panic["panic!(unexpected error)"]
    iter --> tok{"event.token()?"}
    tok -->|"TOKEN_WAKEUP"| skip["忽略,仅用于打断阻塞"]
    tok -->|"TOKEN_SIGNAL"| sig["signal_ready = true"]
    tok -->|"普通 fd token"| cast["EXPOSE_IO.from_exposed_addr(token.0)"]
    cast --> setr["io.set_readiness(Tick::Set, curr | ready)"]
    setr --> wake["io.wake(ready)"]
    wake --> iter
    skip --> iter
    sig --> iter
    iter --> uring["dispatch_completions() (io-uring)"]
    uring --> metrics["metrics.incr_ready_count_by(ready_count)"]

设计思考:为什么 Handle 要持有 mio::Waker

unpark 📎 tokio/src/runtime/io/driver.rs:280-283 调用 self.waker.wake()。这个 mio::Waker 在 Driver::new 时用 TOKEN_WAKEUP 注册 📎 tokio/src/runtime/io/driver.rs:124。当 driver 阻塞在 poll.poll() 里时,另一个线程调用 unpark 会往 epoll 里塞一个 TOKEN_WAKEUP 事件,poll 立即返回,遍历时看到这个 token 直接跳过 📎 tokio/src/runtime/io/driver.rs:214-215。

〔设计推断与架构权衡〕
这个机制在 deregister_source 里被用到 📎 tokio/src/runtime/io/driver.rs:315-334:注销一个 source 后,如果 registrations.deregister 返回 true(表示这是最后一个引用),就 unpark()。为什么? 因为 driver 可能正阻塞在 poll 里等待这个 fd 的事件,而 fd 已经被注销,内核不会再产生事件;必须主动唤醒 driver,让它重新检查注册集合并可能退出阻塞。否则 driver 会一直睡到 max_wait 超时,延迟 shutdown。

另一个细节:deregister_source 先调用 self.registry.deregister(source) 📎 tokio/src/runtime/io/driver.rs:322,再清理 registrations 📎 tokio/src/runtime/io/driver.rs:315-334。注释 📎 tokio/src/runtime/io/driver.rs:320-321 说「Cleanup ALWAYS happens」——即使 OS 层 deregister 失败,也要清理内部状态,最后才返回 OS 错误 📎 tokio/src/runtime/io/driver.rs:336-340。这是典型的资源清理优先于错误传播模式。

注册层:Registration 如何把 Waker 存进 ScheduledIo

直觉模型

Registration 是任务与 fd 之间的契约。它持有两个东西:一个 scheduler::Handle(用于在需要时访问 runtime),一个 Arc<ScheduledIo>(fd 的状态槽)。当任务调用 poll_read_ready 时,Registration 把 Waker 交给 ScheduledIo 保管;当 driver 收到事件时,从 ScheduledIo 里取出 Waker 唤醒。

内存布局与字段

Registration 只有两个字段 📎 tokio/src/runtime/io/registration.rs:46-54:

  • handle: scheduler::Handle:runtime 句柄,注释 📎 tokio/src/runtime/io/registration.rs:46-54 说「TODO: this can probably be moved into ScheduledIo」,说明作者认为这个字段位置可以优化。
  • shared: Arc<ScheduledIo>:共享状态,Arc 保证 driver 和任务都能访问。
〔设计推断与架构权衡〕
注意 Registration 手动实现了 Send 和 Sync 📎 tokio/src/runtime/io/registration.rs:57-58。为什么需要 unsafe impl? 因为 scheduler::Handle 内部可能包含非 Send/Sync 的字段(比如 Rc),但 Registration 的使用场景要求它能跨线程。文档注释 📎 tokio/src/runtime/io/registration.rs:28-33 给出了关键约束:调用者必须保证最多两个任务并发使用同一个 Registration,一个读、一个写。违反这个约束虽然内存安全,但会导致通知丢失和任务挂起。

Step-by-Step:poll_read_ready 的调用链

假设任务在 TcpStream::poll_read 里发现 socket 没数据,需要注册读兴趣。调用链是 TcpStream::poll_read_priv → PollEvented::poll_read → Registration::poll_read_io → poll_io → poll_ready。

poll_ready 是核心 📎 tokio/src/runtime/io/registration.rs:155-171:

第一步:trace_leaf() 📎 tokio/src/runtime/io/registration.rs:160,用于 tracing 埋点。

第二步:coop::poll_proceed(cx) 📎 tokio/src/runtime/io/registration.rs:155-171。这是第 12 章要讲的协作式预算机制。如果预算耗尽,返回 Pending 并注册一个特殊的 Waker,让任务在下一轮被重新调度。

第三步:self.shared.poll_readiness(cx, direction) 📎 tokio/src/runtime/io/registration.rs:155-171。这是真正与 ScheduledIo 交互的地方:检查当前就绪位,若已就绪立即返回 Ready;否则把 cx.waker() 存进 ScheduledIo 的对应方向槽位,返回 Pending。

第四步:检查 ev.is_shutdown 📎 tokio/src/runtime/io/registration.rs:155-171。若 runtime 正在关闭,返回 RUNTIME_SHUTTING_DOWN_ERROR。

第五步:coop.made_progress() 📎 tokio/src/runtime/io/registration.rs:169,标记预算消耗,返回就绪事件。

poll_io 在 poll_ready 之上加了一层重试循环 📎 tokio/src/runtime/io/registration.rs:173-192:

rust
loop {
    let ev = ready!(self.poll_ready(cx, direction))?;
    match f() {
        Ok(ret) => return Poll::Ready(Ok(ret)),
        Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
            self.clear_readiness(ev);
        }
        Err(e) => return Poll::Ready(Err(e)),
    }
}

这里体现了 readiness 是提示而非保证 的核心思想:poll_ready 说可读,但真正 read() 时可能返回 WouldBlock(比如另一个线程抢先读走了数据)。此时必须 clear_readiness(ev) 📎 tokio/src/runtime/io/registration.rs:187 清掉就绪位,然后循环重新等待。若不清,任务会陷入「以为可读 → read 失败 → 又以为可读」的忙循环。

设计思考:try_io 与 async_io 的分工

try_io 📎 tokio/src/runtime/io/registration.rs:194-213 是同步版本:先 ready_event(interest) 检查就绪位,若为空直接返回 WouldBlock 📎 tokio/src/runtime/io/registration.rs:194-213;否则执行 f(),若 f() 返回 WouldBlock 则清就绪位 📎 tokio/src/runtime/io/registration.rs:207-210。它不注册 Waker,适合 try_read 这类「试一下就走」的场景。

async_io 📎 tokio/src/runtime/io/registration.rs:225-245 是异步版本:readiness(interest).await 会注册 Waker 并等待,然后执行 f(),WouldBlock 时清就绪位并循环。注意它在循环里还调用了 coop::poll_proceed 📎 tokio/src/runtime/io/registration.rs:233,防止在大量 WouldBlock 重试中耗尽预算。

生产踩坑:Drop 里的 Waker 清理

Registration::drop 📎 tokio/src/runtime/io/registration.rs:253-262 调用 self.shared.clear_wakers()。注释 📎 tokio/src/runtime/io/registration.rs:253-262 解释了原因:ScheduledIo 里存的 Waker 可能持有 Arc<driver::Inner>,而 driver::Inner 又持有 ScheduledIo,形成循环引用。清理 Waker 是打破循环的手段。但注释也承认这是「imperfect solution」——如果 Registration 本身被存进了 Waker,循环仍然存在。这是 tokio-rs/tokio#3481 讨论的问题。

〔设计推断与架构权衡〕
生产环境中的表现是:如果大量连接被 drop 但 runtime 未退出,内存不会立即回收,直到下一次 clear_wakers 或 runtime shutdown。对于长连接服务,这通常不是问题;但对于短连接高频创建/销毁的场景,需要关注 ScheduledIo 的回收时机。

从 TcpStream::read 到 Waker 唤醒的完整链路

直觉模型

现在把三层串起来。用户在 TcpStream 上调用 .read().await,实际执行的是 AsyncRead::poll_read → PollEvented::poll_read → Registration::poll_read_io。当数据没到时,Waker 被存进 ScheduledIo;当 epoll 报告可读时,driver 从 ScheduledIo 取出 Waker 并唤醒,任务被重新调度,再次 poll 时 poll_readiness 发现就绪位已置,直接返回 Ready,read() 成功。

Step-by-Step:一次完整的读等待

阶段一:注册兴趣。TcpStream::new 📎 tokio/src/net/tcp/stream.rs:166-169 调用 PollEvented::new(connected),后者内部调用 Registration::new_with_interest_and_handle 📎 tokio/src/runtime/io/registration.rs:73-81,进而 handle.driver().io().add_source(io, interest) 📎 tokio/src/runtime/io/registration.rs:73-81。

add_source 📎 tokio/src/runtime/io/driver.rs:288-312 做三件事:

1. registrations.allocate(&mut synced.lock()) 分配一个 ScheduledIo,拿到 token 📎 tokio/src/runtime/io/driver.rs:293-294。

2. self.registry.register(source, token, interest.to_mio()) 向内核注册 📎 tokio/src/runtime/io/driver.rs:298。若失败,必须把刚分配的 ScheduledIo 从集合里移除 📎 tokio/src/runtime/io/driver.rs:300-303,否则泄漏。

3. metrics.incr_fd_count() 计数 📎 tokio/src/runtime/io/driver.rs:309。

阶段二:等待就绪。任务 poll TcpStream::poll_read 📎 tokio/src/net/tcp/stream.rs:1492-1498 → poll_read_priv 📎 tokio/src/net/tcp/stream.rs:1451-1458 → PollEvented::poll_read → Registration::poll_read_io 📎 tokio/src/runtime/io/registration.rs:133-139 → poll_io → poll_ready → ScheduledIo::poll_readiness。此时若未就绪,Waker 存入 ScheduledIo 的读槽位,返回 Pending。

阶段三:事件到达。driver 的 turn 从 poll.poll() 拿到事件 📎 tokio/src/runtime/io/driver.rs:198,遍历时对每个 fd 事件执行 io.set_readiness(Tick::Set, |curr| curr | ready) 和 io.wake(ready) 📎 tokio/src/runtime/io/driver.rs:228-229。wake 内部取出对应方向的 Waker 并调用 wake()。

阶段四:任务重调度。Waker::wake() 把任务重新入队到 worker 的本地队列(上一章讲过)。worker 再次 poll 该任务,poll_readiness 发现就绪位已置,返回 Ready,read() 成功。

mermaid
sequenceDiagram
    participant Task as "任务 (worker 线程)"
    participant Reg as "Registration"
    participant SIO as "ScheduledIo"
    participant Drv as "Driver (I/O 线程)"
    participant OS as "epoll/kqueue"

    Task->>Reg: "poll_read_ready(cx)"
    Reg->>SIO: "poll_readiness(cx, Read)"
    SIO-->>Reg: "Pending (Waker 已存入读槽位)"
    Reg-->>Task: "Poll::Pending"
    Note over Task: 任务让出,worker 去跑别的任务
    Drv->>OS: "poll.poll(events, max_wait)"
    OS-->>Drv: "event(token=fd_ptr, READABLE)"
    Drv->>SIO: "set_readiness(Tick::Set, curr | READABLE)"
    Drv->>SIO: "wake(READABLE)"
    SIO->>Task: "Waker::wake() 重新入队"
    Note over Task: worker 再次 poll 该任务
    Task->>Reg: "poll_read_ready(cx)"
    Reg->>SIO: "poll_readiness(cx, Read)"
    SIO-->>Reg: "Ready(ReadyEvent{ready: READABLE})"
    Reg-->>Task: "Poll::Ready(Ok(ev))"
    Task->>Task: "read() 成功返回数据"

重要分支:assume_ready 优化

TcpStream::new_accepted 📎 tokio/src/net/tcp/stream.rs:174-181 是一个值得注意的优化。accept 返回的 socket 天然可写,且通常已经持有对端的第一批字节。如果等 driver 的第一个事件,在高负载下这个事件可能排在所有已建立连接的事件后面,造成延迟。所以 new_accepted 直接调用 assume_ready(Ready::READABLE | Ready::WRITABLE) 📎 tokio/src/net/tcp/stream.rs:174-181。

assume_ready 的注释 📎 tokio/src/runtime/io/registration.rs:103-105 说:「A wrong guess costs one WouldBlock, which clears the readiness again.」——猜错了代价只是一次 WouldBlock,poll_io 的循环会清掉就绪位并重新等待。这是一个乐观猜测 + 快速纠错的设计。

设计思考:为什么 I/O 驱动与调度器解耦

〔设计推断与架构权衡〕
从源码结构看,Driver 和 worker 线程是分离的:Driver 被放在 runtime 的某个专用位置(通常是 block_on 线程或专门的 I/O 线程),而 worker 线程只持有 Handle。这种解耦带来几个好处:

1. 注册无锁化:Handle 持有 mio::Registry 克隆,任何 worker 都能并发注册新 fd,不需要回到 driver 线程。

2. 事件等待集中化:只有一个线程阻塞在 epoll_wait,避免多线程同时 poll 同一个 epoll fd 的惊群问题。

3. 唤醒路径短:driver 收到事件后直接操作 ScheduledIo 并调用 Waker::wake(),wake() 内部把任务推入 worker 队列,不需要跨线程消息传递。

代价是 ScheduledIo 需要处理并发访问(set_readiness 和 poll_readiness 可能同时发生),这通过原子操作和内部锁解决。

生产踩坑:is_shutdown 与 RUNTIME_SHUTTING_DOWN_ERROR

poll_ready 检查 ev.is_shutdown 📎 tokio/src/runtime/io/registration.rs:155-171,若为真返回 gone() 📎 tokio/src/runtime/io/registration.rs:265-267,即 RUNTIME_SHUTTING_DOWN_ERROR。

〔设计推断与架构权衡〕
这个检查的意义在于:runtime 关闭时,driver 的 shutdown 📎 tokio/src/runtime/io/driver.rs:174-182 会遍历所有注册并调用 io.shutdown(),把 is_shutdown 置位并唤醒所有等待者。如果不检查这个标志,任务可能在 runtime 已经停止调度后仍然尝试读 socket,导致未定义行为或挂起。生产环境中,如果你看到 RUNTIME_SHUTTING_DOWN_ERROR,通常意味着有任务在 runtime drop 之后仍在运行——检查是否有 spawn 的任务没有被正确 join。

另一个坑是 deregister_source 的 unpark 📎 tokio/src/runtime/io/driver.rs:328。如果 driver 正阻塞在 poll 里,且此时最后一个 Registration 被 drop,unpark 会唤醒 driver。但如果 driver 不在阻塞状态(比如正在处理其他事件),unpark 只是让下一次 turn 立即返回 📎 tokio/src/runtime/io/driver.rs:280-283。这个语义在 Handle::unpark 的文档注释里有说明。

设计思考:Reactor 的三个关键权衡

权衡一:Token 用指针而非索引。EXPOSE_IO.from_exposed_addr(token.0) 📎 tokio/src/runtime/io/driver.rs:220 把 mio::Token 直接当作 *const ScheduledIo 的地址。这避免了维护一个 Token → ScheduledIo 的映射表,查找是 O(1) 且无锁。代价是安全性依赖严格的生命周期管理:指针必须在注销且 driver 不再 poll 之后才能释放 📎 tokio/src/runtime/io/driver.rs:222-225。

权衡二:读写双 Waker 槽位。Registration 文档 📎 tokio/src/runtime/io/registration.rs:24-26 说「A registration instance represents two separate readiness streams」——读和写各有一个独立的 Waker 槽位。这允许同一个 socket 的读任务和写任务分别注册,互不干扰。但 poll_read_ready 的注释 📎 tokio/src/net/tcp/stream.rs:549-552 提醒:多次调用 poll_read_ready/poll_read/poll_peek 只有最后一次的 Waker 会被保留——读方向只有一个槽位。

权衡三:events_busy 的独立缓冲区。测试 📎 tokio/src/runtime/io/driver.rs:364-386 验证了这个行为:Driver::new(16, Some(2)) 创建 busy 容量为 2 的 driver,注册 5 个可读 source 后,非阻塞 turn 只取 2 个事件 📎 tokio/src/runtime/io/driver.rs:375-376,剩余 3 个留在内核队列,下次阻塞 turn 取到 📎 tokio/src/runtime/io/driver.rs:379-380。这防止了非阻塞 poll 一次性吞掉所有事件导致后续 poll 饥饿。

本章小结

本章追踪了 TcpStream::read 背后的完整 Reactor 链路:

  • 驱动层:Driver 独占 mio::Poll,turn 阻塞等待事件,用 EXPOSE_IO 把 Token 还原为 ScheduledIo 指针,调用 set_readiness + wake 触发 Waker。Handle 提供可跨线程的注册入口,unpark 用于打断阻塞。
  • 注册层:Registration 持有 Arc<ScheduledIo>,poll_ready 检查就绪位或存入 Waker,poll_io 用 WouldBlock 重试循环处理假阳性,try_io/async_io 分别服务同步和异步场景。
  • 状态层:ScheduledIo 是 fd 的状态槽,存储读写就绪位和双 Waker 槽位,是事件与任务之间的唯一桥梁。

本章思考与自测

Q1: 如果把 poll_io 里 WouldBlock 分支的 self.clear_readiness(ev) 删掉,在什么场景下会导致任务忙循环(busy-loop)?为什么?

参考解析:poll_io 的循环 📎 tokio/src/runtime/io/registration.rs:173-192 在 f() 返回 WouldBlock 时调用 clear_readiness(ev) 📎 tokio/src/runtime/io/registration.rs:187。ev 是 poll_ready 返回的 ReadyEvent,包含当前就绪位。clear_readiness 会把这些位从 ScheduledIo 里清掉。

如果不清理,下一次循环调用 poll_ready → poll_readiness 时,ScheduledIo 里仍然保留着旧的「可读」位,poll_readiness 会立即返回 Ready(因为就绪位非空),然后 f() 再次执行 read(),如果 socket 确实没数据,又返回 WouldBlock,循环继续。由于就绪位从未被清除,这个循环永远不会进入 Pending,任务会一直占用 CPU 轮询。

触发场景:多个任务共享同一个 socket 的读方向(虽然 Registration 文档 📎 tokio/src/runtime/io/registration.rs:28-33 说最多两个任务,但读方向只有一个槽位),或者 try_read 和 poll_read 混用。更常见的是:epoll 报告可读后,另一个线程抢先读走了数据,当前任务的 read() 返回 WouldBlock,此时必须清就绪位,否则会一直重试。

Q2: add_source 在 registry.register 失败时为什么要调用 registrations.remove?如果不调用会发生什么?

参考解析:add_source 📎 tokio/src/runtime/io/driver.rs:288-312 先 registrations.allocate 分配 ScheduledIo 📎 tokio/src/runtime/io/driver.rs:293,再 registry.register 向内核注册 📎 tokio/src/runtime/io/driver.rs:298。如果注册失败,ScheduledIo 已经分配但没有任何 fd 与之关联,如果不移除,它会永远留在 RegistrationSet 里。

注释 📎 tokio/src/runtime/io/driver.rs:296-297 明确说:「we should remove the scheduled_io from the registrations set if registering the source with the OS fails. Otherwise it will leak the scheduled_io.」——这是一个内存泄漏。

remove 调用 📎 tokio/src/runtime/io/driver.rs:300-303 用 unsafe 块包裹,因为 ScheduledIo 是 RegistrationSet 的一部分,移除操作需要保证没有其他引用。泄漏的后果:RegistrationSet 持续增长,Token 空间被浪费,最终可能导致 allocate 失败或内存耗尽。在高频创建/销毁连接的场景(如短连接服务器),如果注册失败率较高(比如 fd 耗尽),泄漏会加速资源枯竭。

Q3: deregister_source 中,为什么 unpark() 只在 registrations.deregister 返回 true 时调用?如果无条件调用会有什么问题?

参考解析:deregister_source 📎 tokio/src/runtime/io/driver.rs:315-334 的逻辑是:先 registry.deregister(source) 向内核注销 📎 tokio/src/runtime/io/driver.rs:322,然后 registrations.deregister 清理内部状态 📎 tokio/src/runtime/io/driver.rs:315-334,若返回 true 则 unpark() 📎 tokio/src/runtime/io/driver.rs:328。

registrations.deregister 返回 true 意味着这是最后一个引用,ScheduledIo 被真正移除。此时 driver 可能正阻塞在 poll 里等待这个 fd 的事件,但 fd 已经注销,内核不会再产生事件。unpark 通过 mio::Waker 往 epoll 塞一个 TOKEN_WAKEUP 事件 📎 tokio/src/runtime/io/driver.rs:280-283,让 poll 立即返回,driver 重新检查注册集合并可能退出阻塞。

如果无条件调用 unpark:每次注销一个非最后的引用都会唤醒 driver,造成不必要的唤醒。在大量连接共享同一个 ScheduledIo 的场景(比如 TcpStream 的 split 后读写两半),每次 drop 一个半都会唤醒 driver,增加 CPU 开销。更严重

本章我们拆解了 Reactor 如何把 epoll 事件翻译成 Waker 唤醒:从 TcpStream 的 poll_read_ready 出发,经过 Registration 的注册与查询,落到 ScheduledIo 的就绪位与 Waker 槽位,再由 Driver 在事件循环中根据 Token 定位并触发唤醒。关键设计包括:Token 即指针实现 O(1) 查找,读写双 Waker 槽位支持并发读写分离,events_busy 独立缓冲区防止事件饥饿,assume_ready 乐观猜测优化 accept 场景。至此,I/O 就绪通知的闭环已经完整。但异步运行时还需要处理另一类「就绪」——时间。下一章我们将剖析 tokio::time::sleep 与 timeout 的实现:定时器如何被插入时间轮、时间轮如何按到期时间分级、driver 如何计算下一次 park 的超时并触发到期任务。你会看到「时间也是一种 I/O 事件」这一统一抽象,以及 start_paused 与 test clock 如何让时间在测试中可控。

AI 赋能代码库精读 · 本地优先架构

读完了本章?为你自己的私有项目生成专属架构全景书

基于 Tauri 2 + Rust 本地原生引擎,100% 源码离线隐私安全,零代码上传云端。像阅读一本传世专著一样拆解你的复杂系统。

⚡ Tauri 2 · Rust 原生引擎 · 100% 离线私密安全 · 适配超百万行代码库

CHAPTER 06

第 6 章:时间轮与高精度定时器:Time Driver 的分级时间轮实现

官方源: tokio-rs/tokio · Commit @e800714a · 全书进度: 第 6 / 14 章

第 6 章:时间轮与高精度定时器:Time Driver 的分级时间轮实现

上一章我们追踪了 TcpStream::read 的完整链路,看到 ScheduledIo 如何把 epoll 的 fd 就绪事件翻译成 Waker 唤醒。但异步运行时还需要处理另一类「就绪」:一个 sleep(100ms) 的 Future,在 100ms 后必须被唤醒。这类事件不来自内核 fd,而来自「时间本身」。Tokio 的设计选择是把时间也当作一种 I/O 事件:Driver 结构体里只有一个字段 park: IoStack,它复用了 I/O driver 的 park/unpark 机制。当时间轮算出「下一次到期时刻」时,driver 就调用 park_timeout 让线程睡到那个时刻;被唤醒后再从时间轮里取出到期条目、触发它们的 Waker。这样,调度器只需要一个统一的 park 入口,就能同时等待「fd 就绪」和「定时器到期」两类事件。本章要回答三个问题:定时器如何被插入时间轮?时间轮如何按到期时间分级?driver 如何计算下一次 park 的超时并触发到期任务?

一、时间轮:六层 64 槽的哈希分级结构

直觉模型

想象一个机械钟表:秒针转一圈带动分针,分针转一圈带动时针。如果只有一根秒针,要表示「12 天后」就得数 100 万格;而分层之后,秒针只管 64 秒内的精度,分针管 64 分钟,时针管 64 小时——每一层只需 64 个槽位,就能覆盖到 2 年之后。

若没有分层,插入一个远期定时器要么需要 O(N) 遍历,要么需要巨大的数组。时间轮用「按到期时间分级」把插入和触发都压到近似 O(1)。

内存布局与字段

Wheel 的核心字段只有三个 📎 tokio/src/runtime/time/wheel/mod.rs:22-40:

rust
pub(crate) struct Wheel {
    elapsed: u64,                          // 自 wheel 创建以来经过的毫秒数
    levels: Box,      // 6 层,每层 64 槽
    pending: LinkedList,      // 已到期、待触发的条目
}

NUM_LEVELS = 6,BITS_PER_LEVEL = 6(即每层 64 槽)📎 tokio/src/runtime/time/wheel/mod.rs:45-47。MAX_DURATION = 1 << (6 * 6) = 1 << 36 毫秒,约 2 年 📎 tokio/src/runtime/time/wheel/mod.rs:50。

六层的粒度按文档注释是 📎 tokio/src/runtime/time/wheel/mod.rs:22-40:

层槽粒度覆盖范围
01 ms64 ms
164 ms~4 s
2~4 s~4 min
3~4 min~4 hr
4~4 hr~12 day
5~12 day~2 yr

pending 是一个侵入式链表(LinkedList<TimerShared>),存放已经从轮中取出、等待触发 Waker 的条目。注意它是 LinkedList 而非 Vec:条目本身内嵌在 TimerShared 里,插入/移除不需要分配。

场景驱动:插入一个 100ms 的 sleep

当 sleep(100ms) 首次被 poll 时,Sleep::poll_elapsed 会构造 Timer::new 并调用 init 📎 tokio/src/time/sleep.rs:436-440。init 最终调用 Handle::reregister,进而调用 Wheel::insert。

insert 的第一步是检查是否已过期 📎 tokio/src/runtime/time/wheel/mod.rs:90-98:

rust
let when = unsafe { item.sync_when() };
if when  usize {
    const SLOT_MASK: u64 = (1 = MAX_DURATION {
        return NUM_LEVELS - 1;
    }
    masked.ilog2() as usize / BITS_PER_LEVEL
}

这里用 elapsed ^ when 而非 when - elapsed,是一个精妙的技巧:XOR 的最高有效位反映了「两个时间戳从哪一位开始不同」,也就是「需要多粗的粒度才能区分它们」。| SLOT_MASK 把低 6 位强制置 1,避免 ilog2 落在同一槽内时算出过小的层。ilog2() / 6 把位宽映射到层号。如果 XOR 结果超过 MAX_DURATION(即超过 2 年),就强制塞进最高层——这就是「fudge the timer into the top level」。

对于 100ms 的 sleep,假设 elapsed 接近 0,when ≈ 100,elapsed ^ when ≈ 100,ilog2(100) = 6,6 / 6 = 1,所以落在第 1 层(64ms 粒度)。这意味着它会在第 1 层的某个槽里等待,直到时间推进到该槽的边界时才被下沉到第 0 层。

分级下沉:process_expiration

当 poll(now) 推进时间时,Wheel::poll 会循环调用 next_expiration 和 process_expiration 📎 tokio/src/runtime/time/wheel/mod.rs:142-166:

rust
pub(crate) fn poll(&mut self, now: u64) -> Option {
    loop {
        if let Some(handle) = self.pending.pop_back() {
            return Some(handle);
        }
        match self.next_expiration() {
            Some(ref expiration) if expiration.deadline  {
                self.process_expiration(expiration);
                self.set_elapsed(expiration.deadline);
            }
            _ => {
                self.set_elapsed(now);
                break;
            }
        }
    }
    self.pending.pop_back()
}

process_expiration 负责把某一层的到期条目「下沉」到下一层,或者(在第 0 层)标记为 pending 📎 tokio/src/runtime/time/wheel/mod.rs:218-251:

rust
let mut entries = self.take_entries(expiration);
while let Some(item) = entries.pop_back() {
    match unsafe { item.mark_pending(expiration.deadline) } {
        Ok(()) => {
            self.pending.push_front(item);   // 真正到期
        }
        Err(expiration_tick) => {
            let level = level_for(expiration.deadline, expiration_tick);
            unsafe { self.levels[level].add_entry(item); }  // 下沉到更低层
        }
    }
}

mark_pending 是关键:它检查条目的实际 deadline 是否已经到达。如果到达,返回 Ok(()),条目进入 pending 链表;如果还没到(只是所在槽的边界到了),返回 Err(expiration_tick),条目被重新插入到更细的层。

注意注释里强调的一点 📎 tokio/src/runtime/time/wheel/mod.rs:219-228:必须先把整个槽的条目全部取出再处理,因为某些条目可能被重新插入到同一个槽(当插入时间超过 MAX_DURATION 时会发生环绕)。如果边取边插,可能陷入无限循环。

下一到期时刻的计算

next_expiration 从低层到高层扫描,返回第一个非空的到期点 📎 tokio/src/runtime/time/wheel/mod.rs:169-191:

rust
fn next_expiration(&self) -> Option {
    if !self.pending.is_empty() {
        return Some(Expiration { level: 0, slot: 0, deadline: self.elapsed });
    }
    for (level_num, level) in self.levels.iter().enumerate() {
        if let Some(expiration) = level.next_expiration(self.elapsed) {
            debug_assert!(self.no_expirations_before(level_num + 1, expiration.deadline));
            return Some(expiration);
        }
    }
    None
}

如果 pending 非空,说明有已到期条目待触发,立即返回当前 elapsed 作为 deadline(这样 driver 会以 0 超时 park,马上回来处理)。否则逐层扫描,返回第一个有内容的槽的 deadline。debug_assert 验证了一个不变量:更高层不可能有比当前层更早的到期点。

mermaid
flowchart TD
    start["Wheel::poll(now)"] --> check_pending{"pending 非空?"}
    check_pending -->|是| pop["pop_back 返回 TimerHandle"]
    check_pending -->|否| next_exp{"next_expiration() 有到期点?"}
    next_exp -->|无| set_elapsed["set_elapsed(now) 后 break"]
    next_exp -->|有| cmp{"expiration.deadline |否| set_elapsed
    cmp -->|是| proc["process_expiration(expiration)"]
    proc --> take["take_entries 取出整槽"]
    take --> mark{"item.mark_pending()"}
    mark -->|Ok 已到期| push_pending["pending.push_front(item)"]
    mark -->|Err 未到期| reinsert["level_for 后 add_entry 下沉"]
    push_pending --> set_elapsed2["set_elapsed(expiration.deadline)"]
    reinsert --> set_elapsed2
    set_elapsed2 --> check_pending
    set_elapsed --> pop2["pending.pop_back() 返回"]

---

二、Driver 的 park 循环:把时间轮接到 I/O 栈上

直觉模型

时间轮本身不会「自己走」。它需要一个外部循环反复问它:「下一次到期是什么时候?」然后睡到那个时刻,醒来后再推进时间。这个循环就是 Driver::park_internal。它把「时间轮的下一次到期」翻译成一个 park_timeout 的时长,交给底层的 I/O 栈去睡。

若没有这个循环,定时器永远不会被触发——时间轮只是静态数据结构,需要有人「拨动」它。

数据结构:Driver 与 InnerState

Driver 只有一个字段 park: IoStack 📎 tokio/src/runtime/time/mod.rs:90-93。真正的状态在 Handle 里,通过 Inner 枚举区分传统实现和实验性实现 📎 tokio/src/runtime/time/mod.rs:95-127。传统实现的 InnerState 包含两个字段 📎 tokio/src/runtime/time/mod.rs:130-136:

rust
struct InnerState {
    next_wake: Option,   // 承诺的最早唤醒时刻
    wheel: wheel::Wheel,
}

next_wake 用 NonZeroU64 而非 Option<u64> 的嵌套,是为了利用 niche 优化——Option<NonZeroU64> 和 u64 同大小。它记录「driver 承诺在哪个 tick 之前会醒来」,用于 reregister 时判断是否需要 unpark。

is_shutdown 是独立的 AtomicBool,注释解释了为什么把它从 Mutex 里拆出来 📎 tokio/src/runtime/time/mod.rs:90-93:Handle 需要在不锁 mutex 的情况下检查 is_shutdown。这是一个典型的「读多写少」优化——shutdown 只发生一次,但检查可能频繁。

场景驱动:一次 park 的完整流程

park_internal 是核心 📎 tokio/src/runtime/time/mod.rs:213-256:

rust
fn park_internal(&mut self, rt_handle: &driver::Handle, limit: Option) {
    let handle = rt_handle.time();
    let mut lock = handle.inner.lock();
    assert!(!handle.is_shutdown());

    let next_wake = lock.wheel.next_expiration_time();
    lock.next_wake = next_wake.map(|t| NonZeroU64::new(t).unwrap_or_else(|| NonZeroU64::new(1).unwrap()));
    drop(lock);

    match next_wake {
        Some(when) => {
            let now = handle.time_source.now(rt_handle.clock());
            let mut duration = handle.time_source.tick_to_duration(when.saturating_sub(now));
            if duration > Duration::from_millis(0) {
                if let Some(limit) = limit {
                    duration = std::cmp::min(limit, duration);
                }
                self.park_thread_timeout(rt_handle, duration);
            } else {
                self.park.park_timeout(rt_handle, Duration::from_secs(0));
            }
        }
        None => {
            if let Some(duration) = limit {
                self.park_thread_timeout(rt_handle, duration);
            } else {
                self.park.park(rt_handle);
            }
        }
    }

    handle.process(rt_handle.clock());
}

分步解析:

1. 取锁、读下一次到期:lock.wheel.next_expiration_time() 返回 Option<u64>,即下一个到期 tick。同时把它写入 lock.next_wake,供 reregister 判断是否需要 unpark。

2. 释放锁:drop(lock) 必须在 park 之前,否则 park 期间其他线程无法插入定时器。

3. 计算 park 时长:when.saturating_sub(now) 得到剩余 tick 数,tick_to_duration 转成 Duration。注释指出这里实际上向上取整到 1ms 📎 tokio/src/runtime/time/mod.rs:228-230,避免微秒级 sleep 被 OS 当作零长度。

4. 处理 limit:如果调用方传了 limit(比如 park_timeout 的显式超时),取 min(limit, duration),保证不会睡过头。

5. 特殊情况:如果 duration == 0(已到期),用 park_timeout(0) 立即返回,不真正睡。

6. 无定时器时:如果 next_wake 为 None,有 limit 就 park_thread_timeout(limit),否则无限 park。

7. 唤醒后处理:handle.process(clock) 推进时间轮并触发到期条目。

process_at_time:触发到期条目

process 调用 process_at_time 📎 tokio/src/runtime/time/mod.rs:296-337:

rust
pub(self) fn process_at_time(&self, mut now: u64) {
    let mut waker_list = WakeList::new();
    let mut lock = self.inner.lock();

    if now ) {
    let waker = unsafe {
        let mut lock = self.inner.lock();
        if unsafe { entry.as_ref().might_be_registered() } {
            lock.wheel.remove(entry);
        }
        let entry = entry.as_ref().handle();
        if self.is_shutdown() {
            unsafe { entry.fire(Err(crate::time::error::Error::shutdown())) }
        } else {
            entry.set_expiration(new_tick);
            match unsafe { lock.wheel.insert(entry) } {
                Ok(when) => {
                    if lock.next_wake.is_none_or(|next_wake| when  unsafe {
                    entry.fire(Ok(()))
                },
            }
        }
    };
    if let Some(waker) = waker {
        waker.wake();
    }
}

关键逻辑:插入成功后,如果新到期时刻比 next_wake 更早,就调用 unpark.unpark() 唤醒 driver。这是因为 driver 可能正睡在一个更晚的时刻,需要被提前叫醒以重新计算 park 时长。

注意 unpark 是在持有锁时调用的,而 waker.wake() 是在释放锁后调用的。注释解释 📎 tokio/src/runtime/time/mod.rs:441:必须在调用 Waker 前释放锁以避免死锁。但 unpark 不同——它只是往 epoll 塞一个事件,不会回调用户代码,所以持锁调用是安全的。

mermaid
sequenceDiagram
    participant Sleep as Sleep::poll
    participant Handle as time::Handle
    participant Wheel as Wheel
    participant Driver as Driver::park_internal
    participant IoStack as IoStack

    Sleep->>Handle: reregister(unpark, new_tick, entry)
    Handle->>Handle: lock.inner.lock()
    Handle->>Wheel: wheel.remove(entry) [若已注册]
    Handle->>Wheel: wheel.insert(entry)
    Wheel-->>Handle: Ok(when)
    alt when >IoStack: unpark.unpark()
    end
    Handle->>Handle: drop(lock)
    Handle-->>Sleep: 返回 waker (若有)

    Note over Driver: 另一线程
    Driver->>Handle: lock.inner.lock()
    Driver->>Wheel: next_expiration_time()
    Wheel-->>Driver: Some(when)
    Driver->>Driver: drop(lock)
    Driver->>IoStack: park_timeout(duration)
    IoStack-->>Driver: 被 unpark 或超时
    Driver->>Handle: process(clock)
    Handle->>Wheel: poll(now)
    Wheel-->>Handle: TimerHandle
    Handle->>Sleep: waker.wake()

---

三、Sleep 与 Timeout:用户可见的 API 层

直觉模型

Sleep 是用户直接 .await 的 Future,Timeout 是包裹另一个 Future 的适配器。它们本身不管理时间轮,只是把「deadline」翻译成 tick,委托给 Timer 和 Handle。

Sleep 的内存布局

Sleep 用 pin_project! 宏定义 📎 tokio/src/time/sleep.rs:221-227:

rust
pub struct Sleep {
    deadline: Instant,
    driver: scheduler::Handle,
    inner: Inner,
    #[pin]
    timer: Option,
}

timer 是 Option<Timer> 且带 #[pin]:首次 poll 前是 None,首次 poll 时才创建 Timer 并注册。这种「惰性初始化」避免了在 sleep() 调用时就访问运行时——sleep() 可以在运行时外调用,只要在 .await 时才真正注册。

PinnedDrop 实现确保 drop 时取消定时器 📎 tokio/src/time/sleep.rs:230-235:

rust
impl PinnedDrop for Sleep {
    fn drop(this: Pin) {
        let this = this.project();
        if let Some(timer) = this.timer.as_pin_mut() {
            timer.cancel(this.driver);
        }
    }
}

poll_elapsed 的完整流程

poll_elapsed 是 Sleep 的核心 📎 tokio/src/time/sleep.rs:396-454:

rust
fn poll_elapsed(self: Pin, cx: &mut task::Context) -> Poll> {
    ready!(crate::trace::trace_leaf());
    let mut this = self.project();

    // coop 预算
    let coop = ready!(crate::task::coop::poll_proceed(cx));

    let handle = this.driver;
    let timer = match this.timer.as_mut().as_pin_mut() {
        Some(timer) => timer,
        None => {
            let time_source = handle.driver().time().time_source();
            let deadline = time_source.deadline_to_tick(*this.deadline);
            let timer = Timer::new(handle, deadline);
            this.timer.set(Some(timer));
            let mut timer = this.timer.as_pin_mut().unwrap();
            timer.as_mut().init(handle, deadline);
            timer
        }
    };

    let result = timer.poll_elapsed(cx, handle).map(move |r| {
        coop.made_progress();
        r
    });
    result
}

分步:

1. coop 预算检查:poll_proceed(cx) 消耗一次协作预算。如果预算耗尽,返回 Pending 并让出执行权。这是 Tokio 防止单个任务饿死其他任务的机制。

2. 惰性创建 Timer:如果 timer 是 None,把 deadline 转成 tick,创建 Timer 并调用 init 注册到时间轮。

3. 委托给 Timer::poll_elapsed:实际的到期检查由 Timer 完成。

4. 成功后标记进度:coop.made_progress() 表示这次 poll 有实际进展。

Timeout 的 poll:先 poll 值,再 poll 延迟

Timeout 的 poll 顺序很关键 📎 tokio/src/time/timeout.rs:210-224:

rust
fn poll(self: Pin, cx: &mut task::Context) -> Poll {
    let me = self.project();
    let had_budget_before = coop::has_budget_remaining();

    // 先 poll 被包裹的 future
    if let Poll::Ready(v) = me.value.poll(cx) {
        return Poll::Ready(Ok(v));
    }

    match me.delay.as_pin_mut() {
        Some(delay) => poll_delay(had_budget_before, delay, cx).map(Err),
        None => Poll::Pending,
    }
}

注释明确指出 📎 tokio/src/time/timeout.rs:24-26:future 先被 poll,然后才检查超时。所以如果 future 不 yield 就完成,它可能在超过 timeout 后仍返回 Ok。这是设计选择,不是 bug。

poll_delay 处理一个微妙的场景 📎 tokio/src/time/timeout.rs:229-251:

rust
fn poll_delay(had_budget_before: bool, delay: Pin, cx: &mut task::Context) -> Poll {
    let delay_poll = || match delay.poll(cx) {
        Poll::Ready(()) => Poll::Ready(Elapsed::new()),
        Poll::Pending => Poll::Pending,
    };

    let has_budget_now = coop::has_budget_remaining();

    if let (true, false) = (had_budget_before, has_budget_now) {
        // 如果预算是被底层 future 耗尽的,用无约束预算 poll delay
        coop::with_unconstrained(delay_poll)
    } else {
        delay_poll()
    }
}

逻辑:如果进入 poll 时还有预算,但 poll 完 value 后预算耗尽了,说明是 value 消耗了预算。此时如果用受限预算 poll delay,delay 可能立即返回 Pending,导致永远无法判断超时是否到达。所以用 with_unconstrained 临时解除预算限制。注释称之为「pathological cases」📎 tokio/src/time/timeout.rs:243-246。

timeout 的 deadline 溢出处理

timeout 函数用 checked_add 处理溢出 📎 tokio/src/time/timeout.rs:86-99:

rust
Timeout {
    value: future.into_future(),
    delay: match Instant::now().checked_add(duration) {
        Some(deadline) => Some(Sleep::new_timeout(deadline, trace::caller_location())),
        None => None,
    },
}

如果 Instant::now() + duration 溢出(duration 极大),delay 为 None,poll 时直接返回 Poll::Pending 📎 tokio/src/time/timeout.rs:222。这相当于「永不超时」,是合理的降级行为。

---

设计思考与生产踩坑

为什么用 XOR 而非减法计算层级? elapsed ^ when 的最高有效位直接反映「两个时间戳从哪一位开始不同」,这正是「需要多粗的粒度」的度量。减法 when - elapsed 在 elapsed 接近 when 时高位全为 0,ilog2 会算出过小的层。XOR 天然处理了环绕场景。

时间倒流保护的必要性 📎 tokio/src/runtime/time/mod.rs:301-309:Rust 保证 Instant 单调,但底层 OS 可能不保证。在 Windows 宿主上的 Linux VM 里,std 信任硬件时钟导致 Instant 倒退。Tokio 用 now = lock.wheel.elapsed() 钳制,避免 set_elapsed 的 assert 失败。

批量唤醒与死锁 📎 tokio/src/runtime/time/mod.rs:319:持有时间轮锁时调用 Waker 是危险的——Waker 可能触发任务重新 poll,进而调用 Sleep::reset,试图再次获取时间轮锁,造成死锁。WakeList 的批量机制在锁满时临时释放锁,是标准的「锁外回调」模式。

next_wake 的 niche 优化 📎 tokio/src/runtime/time/mod.rs:130-136:Option<NonZeroU64> 与 u64 同大小,因为 0 被用作 None 的 niche。但 tick 0 是合法值,所以代码用 NonZeroU64::new(t).unwrap_or_else(|| NonZeroU64::new(1).unwrap()) 把 0 映射到 1 📎 tokio/src/runtime/time/mod.rs:221。这是一个微妙的边界处理:tick 0 被当作 tick 1,最多导致 1ms 的额外唤醒。

process_expiration 的「先取后处理」 📎 tokio/src/runtime/time/wheel/mod.rs:219-228:必须先把整槽条目取出再处理,因为超过 MAX_DURATION 的条目会环绕并重新插入同一槽。如果边取边插,会无限循环。

Timeout 的 poll 顺序陷阱 📎 tokio/src/time/timeout.rs:24-26:future 先 poll,超时后检查。如果 future 是 CPU 密集且不 yield,它可能超过 timeout 仍返回 Ok。生产环境中不要依赖 timeout 来强制中断不合作的 future。

---

本章小结

本章拆解了 Tokio 时间驱动的三层结构:

1. 时间轮(Wheel):六层 64 槽的哈希分级结构,用 elapsed ^ when 的位宽决定条目层级,插入和触发近似 O(1)。pending 链表存放已到期条目,process_expiration 负责逐层下沉。

2. Driver(Driver::park_internal):把时间轮的 next_expiration_time 翻译成 park_timeout 时长,复用 I/O 栈的 park/unpark。process_at_time 在唤醒后推进时间轮、批量触发 Waker,并处理时间倒流和死锁防护。

3. 用户 API(Sleep / Timeout):Sleep 惰性创建 Timer 并注册,Timeout 先 poll value 再 poll delay,用 with_unconstrained 处理预算耗尽场景。

核心设计是「时间也是一种 I/O 事件」:driver 只有一个 park 入口,同时等待 fd 就绪和定时器到期。next_wake 记录承诺的唤醒时刻,reregister 在插入更早的定时器时 unpark 唤醒 driver 重新计算。

下一章我们将进入同步原语:Mutex、Semaphore 与通道如何实现异步等待。你会看到它们如何复用本章的 Waker 机制,以及「许可计数」与「等待队列」如何协作。

本章思考与自测

Q1: 如果把 Wheel::insert 中的 if when <= self.elapsed 改成 if when < self.elapsed(去掉等号),在什么场景下会导致定时器永远不被触发?

参考解析:when == self.elapsed 表示定时器的到期时刻恰好等于当前已推进的时间。原代码用 <= 把它判为 Elapsed,调用方立即触发 📎 tokio/src/runtime/time/wheel/mod.rs:96-98。如果改成 <,这个条目会被插入到 level_for(elapsed, when) 算出的层。由于 elapsed ^ when == 0,masked = 0 | SLOT_MASK = 63,ilog2(63) = 5,5 / 6 = 0,落在第 0 层。但第 0 层的 next_expiration 会返回一个 deadline >= elapsed 的槽,而 Wheel::poll 的条件是 expiration.deadline <= now。如果 now == elapsed,条件成立,process_expiration 会取出该条目,mark_pending(elapsed) 检查实际 deadline 是否到达——此时 when == elapsed,mark_pending 返回 Ok,条目进入 pending。所以实际上仍会被触发,但多绕了一圈。真正的风险在于:如果 elapsed 已经推进到 when 之后(when < elapsed),原代码返回 Elapsed 立即触发,改后则插入到一个已经过去的槽,next_expiration 可能返回 deadline < elapsed,set_elapsed 的 assert elapsed <= when 会失败 panic 📎 tokio/src/runtime/time/wheel/mod.rs:253-264。所以这个等号是防止 assert 失败的关键边界。

Q2: process_at_time 中 WakeList 满了之后为什么要 drop(lock) 再 wake_all() 再重新 lock?如果去掉这个 drop,在什么并发场景下会死锁?

参考解析:WakeList 收集 Waker,满了之后必须唤醒一批以腾出空间 📎 tokio/src/runtime/time/mod.rs:318-325。如果持有 self.inner.lock() 时调用 waker.wake(),被唤醒的任务可能立即在另一个线程(或同一线程的调度器)上运行,调用 Sleep::reset 或 Sleep::poll_elapsed,进而调用 Handle::reregister,而 reregister 的第一件事就是 self.inner.lock() 📎 tokio/src/runtime/time/mod.rs:405。由于 std::sync::Mutex 不可重入,同一线程会死锁;即使在不同线程,也会阻塞直到 process_at_time 释放锁,而 process_at_time 正等着 wake_all 返回,形成循环等待。注释明确说「To avoid deadlock, we must do this with the lock temporarily dropped」📎 tokio/src/runtime/time/mod.rs:319。drop 后重新 lock 时,时间轮状态可能已被其他线程修改(比如新定时器插入),所以 while let Some(entry) = lock.wheel.poll(now) 会继续从新状态取条目,这是安全的。

Q3: Timeout::poll 中 had_budget_before 和 has_budget_now 的组合判断 (true, false) 为什么只在「进入时有预算、poll 完 value 后没预算」时才用 with_unconstrained?如果反过来 (false, true) 会怎样?

参考解析:had_budget_before 在 poll value 之前记录 📎 tokio/src/time/timeout.rs:208-208,has_budget_now 在 poll value 之后记录 📎 tokio/src/time/timeout.rs:239。(true, false) 意味着预算是在 poll value 期间耗尽的,说明 value 是「预算消耗者」。此时如果用受限预算 poll delay,poll_proceed 会立即返回 Pending,delay 永远不会被真正检查,超时判断失效。所以用 with_unconstrained 临时解除限制 📎 tokio/src/time/timeout.rs:247。(false, true) 不可能发生——预算只能被消耗,不能被恢复(除非显式 with_unconstrained,但这里没有)。(false, false) 意味着进入时就没预算,此时 poll value 可能已经返回 Pending(因为 poll_proceed 失败),delay 也用受限预算 poll,两者都 pending,符合预期。(true, true) 是正常情况,预算充足,直接 poll delay。

至此,我们已经看清时间驱动如何复用 I/O driver 的 park/unpark 机制,让定时器与 fd 就绪共享同一个等待入口。时间轮的分级、到期计算与 Waker 触发,构成了异步运行时处理「时间就绪」的完整闭环。但异步等待不止于 I/O 与时间——当多个任务竞争同一把锁、或通过通道传递消息时,Waker 又该被存放到哪里?下一章我们将进入 tokio::sync 家族,看看 Mutex、Semaphore 与各类通道如何在「等待者队列 + Waker 唤醒」上做出不同取舍。

AI 赋能代码库精读 · 本地优先架构

读完了本章?为你自己的私有项目生成专属架构全景书

基于 Tauri 2 + Rust 本地原生引擎,100% 源码离线隐私安全,零代码上传云端。像阅读一本传世专著一样拆解你的复杂系统。

⚡ Tauri 2 · Rust 原生引擎 · 100% 离线私密安全 · 适配超百万行代码库

CHAPTER 07

第 7 章:同步原语深度剖析:tokio::sync (Mutex, RwLock, Notify, mpsc)

官方源: tokio-rs/tokio · Commit @e800714a · 全书进度: 第 7 / 14 章

第 7 章:同步原语深度剖析:tokio::sync (Mutex, RwLock, Notify, mpsc)

上一章揭示了时间如何被抽象为一种 I/O 事件,让定时器与 fd 就绪共享同一个 park/unpark 等待入口。然而,当多个任务竞争同一把锁或通过通道传递消息时,等待的对象不再是 fd 或时钟,而是另一个任务的状态变化。本章进入 tokio::sync 家族,探明一次 lock().await 或 recv().await 在阻塞时究竟把 Waker 存到了哪里,被唤醒时又如何被重新调度。

为什么异步 Mutex 不能复用 std 的实现

直觉模型:从「占着茅坑」到「让出座位」

std::sync::Mutex 的 lock() 在锁被占用时会阻塞当前线程——线程被操作系统挂起,直到锁释放。这在异步运行时里是灾难性的:一个 worker 线程可能同时驱动成百上千个任务,如果它因为等一把锁而阻塞,它承载的所有其他任务全部停摆。异步 Mutex 的核心诉求是:等锁时让出线程,把「我在等这把锁」这件事登记到一个队列里,然后返回 Pending,让执行器去跑别的任务。

Tokio 的 Mutex 没有自己实现等待队列,而是完全建立在信号量之上。

数据结构与内存布局

Mutex<T> 的字段极简:

📎 tokio/src/sync/mutex.rs:133-138

rust
pub struct Mutex {
    #[cfg(all(tokio_unstable, feature = "tracing"))]
    resource_span: tracing::Span,
    s: semaphore::Semaphore,
    c: UnsafeCell,
}

三个字段各司其职:s 是一个许可数为 1 的信号量,c 是 UnsafeCell<T> 包裹的受保护数据。注意这里的 semaphore 是 batch_semaphore 的别名 📎 tokio/src/sync/mutex.rs:3-3,也就是底层实现,而非 sync::Semaphore 那层公开封装。

MutexGuard<'a, T> 则只持有一个对 Mutex 的引用:

📎 tokio/src/sync/mutex.rs:151-157

rust
pub struct MutexGuard {
    #[cfg(all(tokio_unstable, feature = "tracing"))]
    resource_span: tracing::Span,
    lock: &'a Mutex,
}

这里有个关键设计:MutexGuard 不持有信号量许可对象,只持有 &Mutex。释放锁的动作发生在 Drop 里,直接调用 self.lock.s.release(1) 📎 tokio/src/sync/mutex.rs:959-961。这与 SemaphorePermit 持有 permits: usize 计数、在 Drop 时归还不同——Mutex 的许可数恒为 1,不需要计数。

Send/Sync 的边界值得单独看:

📎 tokio/src/sync/mutex.rs:258-259

rust
unsafe impl Send for Mutex where T: ?Sized + Send {}
unsafe impl Sync for Mutex where T: ?Sized + Send {}

Sync 只要求 T: Send 而非 T: Sync——这是合理的,因为互斥访问保证了同一时刻只有一个线程能触碰 T,跨线程传递 T 的所有权(Send)就够了,不需要 T 本身可被共享(Sync)。这正是 Mutex<T> 能把非 Sync 的 T 变成 Sync 的原因。

Step-by-Step:一次 lock().await 的完整旅程

代入场景:任务 A 调用 mutex.lock().await,此时锁空闲。

第一步,lock() 构造一个 async 块,内部先 self.acquire().await,成功后构造 MutexGuard 📎 tokio/src/sync/mutex.rs:434-443。

第二步,acquire() 直接委托给信号量:

📎 tokio/src/sync/mutex.rs:655-663

rust
async fn acquire(&self) {
    crate::trace::async_trace_leaf().await;
    self.s.acquire(1).await.unwrap_or_else(|_| {
        unreachable!()
    });
}

unwrap_or_else(|_| unreachable!()) 这行注释道出了设计约束:Mutex 从不显式 close 信号量,且独占持有它,所以 acquire 永远不会返回 Err。这是把「信号量关闭」这一错误路径在类型层面排除掉。

第三步,若锁被占用,s.acquire(1) 返回 Pending,当前任务的 Waker 被登记进信号量的等待队列。Waker 存在哪里? 答案在 batch_semaphore 的等待队列里(本章源码材料未展开该文件,但其角色是:每个等待者持有一个 Waker,按 FIFO 排队)。

第四步,持有锁的任务 B 释放锁时,MutexGuard::drop 调用 s.release(1) 📎 tokio/src/sync/mutex.rs:965-975,信号量把许可交给队首等待者并唤醒其 Waker,任务 A 被重新调度,acquire 返回 Ok,构造出 MutexGuard。

整个流程可以用下面的时序图刻画:

mermaid
sequenceDiagram
    participant TaskA as 任务 A
    participant Mutex as Mutex.s (batch_semaphore)
    participant TaskB as 任务 B (持锁者)
    participant Exec as Executor

    TaskA->>Mutex: acquire(1).await
    Mutex-->>TaskA: Pending (Waker 入队)
    TaskA->>Exec: 让出,调度其他任务
    Note over TaskB: 持有锁执行临界区
    TaskB->>Mutex: MutexGuard::drop -> release(1)
    Mutex->>TaskA: 唤醒队首 Waker
    Exec->>TaskA: 重新 poll
    TaskA->>Mutex: acquire(1) 重试
    Mutex-->>TaskA: Ok(()) 获得许可
    TaskA->>TaskA: 构造 MutexGuard

设计思考:FIFO 公平性与取消安全

文档明确声明 Tokio 的 Mutex 保证 FIFO 📎 tokio/src/sync/mutex.rs:20-22。这一公平性来自底层信号量的排队语义。公平的代价是:一次 lock 被取消(比如在 select! 中落败)会让你失去队列中的位置 📎 tokio/src/sync/mutex.rs:415-419。这不是 bug,而是 FIFO 队列的必然——取消意味着从队列中移除,重新 lock 就得重新排队。

另一个反直觉的设计是不投毒(no poisoning)。std::sync::Mutex 在持锁线程 panic 时会标记为 poisoned,后续 lock 返回 Err。Tokio 的 Mutex 不这么做:持锁者 panic 时锁会被正常释放 📎 tokio/src/sync/mutex.rs:122-125。文档警告,如果 panic 被捕获,受保护数据可能处于不一致状态。这是异步场景下的务实取舍——panic 在异步任务里通常意味着任务终止,投毒机制反而增加复杂度。

MutexGuard::map 系列方法值得一提。它允许把整个 MutexGuard<T> 降级为只保护某个子字段的 MappedMutexGuard<U>。实现上,它先用闭包算出子字段指针 data,再通过 skip_drop 把原 guard 拆解成不触发 Drop 的 MutexGuardInner,最后构造新的 guard 📎 tokio/src/sync/mutex.rs:869-883。skip_drop 用 ManuallyDrop + ptr::read 转移字段所有权,避免 Drop 被调用两次 📎 tokio/src/sync/mutex.rs:827-836。这是 Rust 里「转移所有权但不触发析构」的经典手法。

Semaphore:许可计数与等待队列如何实现背压

直觉模型:停车场的车位

信号量就像停车场:acquire 是开车进场,有空位就进,没空位就在门口排队;release 是开车离场,空出一个位子就通知队首的车进场。许可数就是车位总数,acquire_many(n) 就是一辆占 n 个车位的大车。

数据结构与内存布局

公开的 Semaphore 只是底层 batch_semaphore::Semaphore 的薄封装:

📎 tokio/src/sync/semaphore.rs:427-432

rust
pub struct Semaphore {
    ll_sem: ll::Semaphore,
    #[cfg(all(tokio_unstable, feature = "tracing"))]
    resource_span: tracing::Span,
}

SemaphorePermit<'a> 持有信号量引用和许可计数:

📎 tokio/src/sync/semaphore.rs:442-445

rust
pub struct SemaphorePermit {
    sem: &'a Semaphore,
    permits: usize,
}

permits 字段是理解 forget/merge/split 的关键。forget 把 permits 置零 📎 tokio/src/sync/semaphore.rs:1193-1195,这样 Drop 时归还 0 个许可——等价于「永久消耗」这些许可。split 从当前许可里切出 n 个给新 permit 📎 tokio/src/sync/semaphore.rs:1260-1271。merge 把另一个 permit 的计数合并进来,并断言两者来自同一信号量 📎 tokio/src/sync/semaphore.rs:1230-1240。

〔设计推断与架构权衡〕
MAX_PERMITS 是 usize::MAX >> 3 📎 tokio/src/sync/semaphore.rs:476-479。为什么右移 3 位? 底层 batch_semaphore 需要在高位比特里编码状态标志(如关闭标志),所以把可用许可数限制在低位,留出高位做标志位。这是把「计数 + 状态」压进单个 usize 的常见技巧。

Step-by-Step:acquire 与 release 的许可流转

场景:信号量初始 2 个许可,任务 A acquire(),任务 B acquire_many(2)。

acquire() 委托给 ll_sem.acquire(1),成功后构造 SemaphorePermit { permits: 1 } 📎 tokio/src/sync/semaphore.rs:614-631。acquire_many(2) 类似,但传 2 📎 tokio/src/sync/semaphore.rs:661-679。

若许可不足,ll_sem.acquire(n) 返回 Pending,Waker 入队。这里有个公平性细节:文档指出,如果队首是一个 acquire_many(5) 而当前只剩 3 个许可,即使后面有个 acquire(1) 能立刻满足,它也必须等——因为队首的大车占着队 📎 tokio/src/sync/semaphore.rs:19-24。这是严格 FIFO 的代价,避免了饥饿。

释放路径在 Drop:

📎 tokio/src/sync/semaphore.rs:1402-1404

rust
impl Drop for SemaphorePermit {
    fn drop(&mut self) {
        self.sem.add_permits(self.permits);
    }
}

add_permits 委托给 ll_sem.release(n) 📎 tokio/src/sync/semaphore.rs:568-570,底层把许可还给等待队列,唤醒能凑够许可的等待者。

内存序方面,文档给出了强保证:acquire、release、close 都是 AcqRel 操作,彼此全序,等价于单个原子变量上的 AcqRel 📎 tokio/src/sync/semaphore.rs:35-42。这意味着「先写数据再 release 许可」的写入,对「后 acquire 许可」的任务可见——信号量可以安全地在任务间传递数据。

设计思考:close 与背压

close() 让所有等待者收到 AcquireError,且后续 try_acquire 返回 Closed 📎 tokio/src/sync/semaphore.rs:1161-1163。这是优雅关闭的基础:当接收端不再需要数据时,close 信号量能让所有阻塞的发送者立刻失败返回,而不是永远等待。

背压的本质在 mpsc 里体现得最清楚。下一节会看到,mpsc 的容量控制就是用一个许可数等于 buffer 大小的信号量实现的。

通道家族:等待者队列与 Waker 唤醒的不同取舍

直觉模型:四种通道,四种等待策略

oneshot 是「一次性信封」——只能送一封信,发送方不等待(send 是同步的),接收方 await 等信。mpsc 是「有界传送带」——发送方在传送带满时等待,接收方在空时等待,容量由信号量控制。broadcast 和 watch 是「广播喇叭」——一个发送方,多个接收方,但两者对「落后」的处理截然不同。

本节源码材料聚焦 oneshot 和 mpsc::bounded,我们逐一拆解。

oneshot:用状态位编码的极简握手

oneshot 的 Inner 结构是理解其设计的核心:

📎 tokio/src/sync/oneshot.rs:386-409

rust
struct Inner {
    state: AtomicUsize,
    value: UnsafeCell>,
    tx_task: Task,
    rx_task: Task,
}

state 是一个 AtomicUsize,用位标志编码整个通道的状态。四个标志位定义在文件末尾:

📎 tokio/src/sync/oneshot.rs:1488-1505

rust
const RX_TASK_SET: usize = 0b00001;
const VALUE_SENT: usize = 0b00010;
const CLOSED: usize = 0b00100;
const TX_TASK_SET: usize = 0b01000;

value 是 UnsafeCell<Option<T>>,tx_task 和 rx_task 是 Task 类型,内部是 UnsafeCell<MaybeUninit<Waker>> 📎 tokio/src/sync/oneshot.rs:411-411。注意 MaybeUninit——Waker 可能未初始化,是否有效由 state 里的 RX_TASK_SET/TX_TASK_SET 位决定 📎 tokio/src/sync/oneshot.rs:396-399。

这个设计的精髓:VALUE_SENT 位不仅表示「值已发送」,还决定了 UnsafeCell 的访问权归属。注释写得非常明确 📎 tokio/src/sync/oneshot.rs:1491-1496:若 VALUE_SENT 置位,UnsafeCell 只能被接收方访问;若未置位,只能被发送方访问。这样就用一个原子位实现了无锁的所有权转移,避免了额外的锁。

send 的流程:

📎 tokio/src/sync/oneshot.rs:622-646

rust
pub fn send(mut self, t: T) -> Result {
    let inner = self.inner.take().unwrap();
    inner.value.with_mut(|ptr| unsafe {
        *ptr = Some(t);
    });
    if !inner.complete() {
        unsafe {
            return Err(inner.consume_value().unwrap());
        }
    }
    Ok(())
}

先把值写入 UnsafeCell(此时 VALUE_SENT 未置位,接收方不会访问),再调用 complete() 尝试置位 VALUE_SENT。complete() 是一个 CAS 循环:

📎 tokio/src/sync/oneshot.rs:1516-1549

rust
fn set_complete(cell: &AtomicUsize) -> State {
    let mut state = cell.load(Ordering::Relaxed);
    loop {
        if State(state).is_closed() {
            break;
        }
        match cell.compare_exchange_weak(
            state, state | VALUE_SENT, Ordering::AcqRel, Ordering::Acquire,
        ) {
            Ok(_) => break,
            Err(actual) => state = actual,
        }
    }
    State(state)
}

为什么用 CAS 而非简单的 fetch_or?注释解释得很清楚 📎 tokio/src/sync/oneshot.rs:1517-1529:如果通道已 CLOSED,就不能再置 VALUE_SENT。因为一旦置位,接收方会认为可以访问 UnsafeCell,而此时发送方正准备把值取回去(consume_value),两边同时访问就会数据竞争。所以 CAS 循环在发现 CLOSED 时提前 break,不置位。

complete() 返回后,如果成功置位且 RX_TASK_SET 已置位,就唤醒接收方:

📎 tokio/src/sync/oneshot.rs:1300-1315

rust
fn complete(&self) -> bool {
    let prev = State::set_complete(&self.state);
    if prev.is_closed() {
        return false;
    }
    if prev.is_rx_task_set() {
        unsafe {
            self.rx_task.with_task(Waker::wake_by_ref);
        }
    }
    true
}

接收方的 poll_recv 是状态机的核心:

📎 tokio/src/sync/oneshot.rs:1317-1384

它先加载状态,若 is_complete() 则直接 consume_value 返回;若 is_closed() 返回 Err;否则进入「登记 Waker」分支。登记时先检查 is_rx_task_set(),若已设置且 will_wake 判断是同一个 Waker 就不重复设置;若不同则先 unset 再 set。这里有个微妙的竞态处理:unset 之后如果发现 is_complete() 变真了,要把标志位重新 set 回去 📎 tokio/src/sync/oneshot.rs:1342-1344,否则 Waker 会在 Drop 时泄漏(因为 Drop 依赖标志位判断是否要 drop Waker)。

这个「unset 后重新 set」的模式在 poll_closed 里也出现 📎 tokio/src/sync/oneshot.rs:839-848,是 oneshot 处理并发唤醒的标准手法。

mpsc::bounded:信号量驱动的背压

mpsc 的容量控制完全交给信号量。channel 函数创建一个许可数等于 buffer 的信号量:

📎 tokio/src/sync/mpsc/bounded.rs:159-171

rust
pub fn channel(buffer: usize) -> (Sender, Receiver) {
    assert!(buffer > 0, "mpsc bounded channel requires buffer > 0");
    let semaphore = Semaphore {
        semaphore: semaphore::Semaphore::new(buffer),
        bound: buffer,
    };
    let (tx, rx) = chan::channel(semaphore);
    let tx = Sender::new(tx);
    let rx = Receiver::new(rx);
    (tx, rx)
}

Semaphore 是 mpsc 内部的包装,同时持有底层信号量和 bound(最大容量)📎 tokio/src/sync/mpsc/bounded.rs:176-179。bound 用于 max_capacity 查询,而 available_permits 给出当前容量 📎 tokio/src/sync/mpsc/bounded.rs:591-593。

发送路径 send 先 reserve 再 send:

📎 tokio/src/sync/mpsc/bounded.rs:816-824

rust
pub async fn send(&self, value: T) -> Result> {
    match self.reserve().await {
        Ok(permit) => {
            permit.send(value);
            Ok(())
        }
        Err(_) => Err(SendError(value)),
    }
}

reserve 内部调用 reserve_inner(1),后者先检查 n > max_capacity 直接返回错误,再 acquire(n) 📎 tokio/src/sync/mpsc/bounded.rs:1272-1311。这里有个精妙的 WakeReceiverOnDrop 守卫:

📎 tokio/src/sync/mpsc/bounded.rs:1286-1301

rust
struct WakeReceiverOnDrop {
    chan: &'a chan::Tx,
}
impl Drop for WakeReceiverOnDrop {
    fn drop(&mut self) {
        use chan::Semaphore;
        let semaphore = self.chan.semaphore();
        if semaphore.is_closed() && semaphore.is_idle() {
            self.chan.wake_rx();
        }
    }
}

注释解释了动机 📎 tokio/src/sync/mpsc/bounded.rs:1279-1285:如果 reserve 在拿到部分许可后被取消(比如 select! 落败),底层 Acquire 会在 Drop 时归还这些许可,但不会像 Permit 那样通知接收方。如果此时通道已关闭且空闲,接收方可能永远等不到「通道已关闭」的通知。这个守卫在 Drop 时补上这个唤醒。成功时用 mem::forget(guard) 取消守卫 📎 tokio/src/sync/mpsc/bounded.rs:1306-1306,因为成功路径由 Permit 接管通知职责。

Permit 的 Drop 也做同样的事:

📎 tokio/src/sync/mpsc/bounded.rs:1732-1745

rust
impl Drop for Permit {
    fn drop(&mut self) {
        use chan::Semaphore;
        let semaphore = self.chan.semaphore();
        semaphore.add_permit();
        if semaphore.is_closed() && semaphore.is_idle() {
            self.chan.wake_rx();
        }
    }
}

Permit::send 则用 mem::forget 跳过 Drop,避免归还许可 📎 tokio/src/sync/mpsc/bounded.rs:1721-1728。

接收路径 recv 用 poll_fn 包装 chan.recv(cx) 📎 tokio/src/sync/mpsc/bounded.rs:243-246。poll_recv 直接委托 📎 tokio/src/sync/mpsc/bounded.rs:650-652。真正的等待队列逻辑在 chan 模块(本章未展开),但可以推断:接收方 Waker 存在 chan::Rx 里,当发送方 send 时唤醒。

try_send 展示了非阻塞路径:

📎 tokio/src/sync/mpsc/bounded.rs:924-934

rust
pub fn try_send(&self, message: T) -> Result> {
    match self.chan.semaphore().semaphore.try_acquire(1) {
        Ok(()) => {}
        Err(TryAcquireError::Closed) => return Err(TrySendError::Closed(message)),
        Err(TryAcquireError::NoPermits) => return Err(TrySendError::Full(message)),
    }
    self.chan.send(message);
    Ok(())
}

try_acquire 的两种错误精确映射到 Closed 和 Full,区分了「通道关闭」和「缓冲区满」两种失败。

设计思考:取消安全与消息丢失

mpsc 文档反复强调取消安全 📎 tokio/src/sync/mpsc/bounded.rs:776-784:send 在 select! 中落败时,消息会被丢弃。要避免丢失,必须用 reserve 拿到 Permit 再 send——因为 Permit 已经预留了容量,send 是同步的、不会被打断。

recv 则是取消安全的 📎 tokio/src/sync/mpsc/bounded.rs:199-204:若 recv 在 select! 中落败,保证没有消息被消费。这是因为 recv 的 poll_recv 只在真正取到消息时才返回 Ready,Pending 时不动队列。

oneshot 的 Receiver 作为 Future 也是取消安全的 📎 tokio/src/sync/oneshot.rs:246-251。但要注意:oneshot 的 send 是同步的,所以不存在「send 被取消」的问题——要么发出去,要么 Err 返回原值。

设计思考与生产踩坑

坑一:用异步 Mutex 保护纯数据。 文档明确建议 📎 tokio/src/sync/mutex.rs:26-36:如果受保护的是纯数据(无 .await 需求),用 std::sync::Mutex 或 parking_lot 更快。异步 Mutex 的开销在于信号量的原子操作和可能的任务调度。只有当需要在持锁期间 .await(比如持锁访问数据库连接)时,才用异步 Mutex。

坑二:持锁跨 .await 导致死锁。 这是异步 Mutex 最危险的陷阱。如果任务 A 持锁后 .await 一个需要任务 B 完成的事件,而任务 B 又在等这把锁,就死锁了。std::sync::Mutex 的 guard 不是 Send(在可移动任务中),编译器会阻止跨 .await 持有;但异步 Mutex 的 guard 是 Send 📎 tokio/src/sync/mutex.rs:314-314,编译器不拦你,需要自己保证不形成循环等待。

坑三:reserve 后忘记 send。 Permit 的 Drop 会归还许可 📎 tokio/src/sync/mpsc/bounded.rs:1732-1745,所以不会泄漏容量。但如果通道已关闭且空闲,Drop 会唤醒接收方——这个唤醒是必要的,否则接收方可能永远等不到关闭通知。

坑四:oneshot 的 poll 可能虚假 Pending。 文档说明 📎 tokio/src/sync/oneshot.rs:236-242:即使消息已发送,poll 也可能返回 Pending。这不是 bug,而是并发竞态下的正常现象——调用者会被唤醒重试,消息不会丢失,只是延迟。

坑五:forget_permits 的语义。 forget_permits(n) 尝试减少 n 个许可,返回实际减少的数量 📎 tokio/src/sync/semaphore.rs:576-578。它不会阻塞,也不会唤醒等待者——只是单纯地「吞掉」许可。用于动态收缩信号量容量。

本章小结

本章揭示了 tokio::sync 的核心模式:所有异步等待原语都建立在「等待者队列 + Waker 唤醒」之上,而队列的具体实现因场景而异。

  • Mutex 复用许可数为 1 的信号量,MutexGuard 只持引用,Drop 时 release(1),FIFO 公平但不投毒。
  • Semaphore 是许可计数 + 等待队列,SemaphorePermit 用 permits 计数支持 forget/merge/split,MAX_PERMITS 右移 3 位为状态标志留位。
  • oneshot 用单个 AtomicUsize 的位标志编码状态,VALUE_SENT 位同时决定 UnsafeCell 的访问权归属,CAS 循环防止在 CLOSED 后置位。
  • mpsc::bounded 用许可数等于 buffer 的信号量实现背压,WakeReceiverOnDrop 守卫处理取消时的唤醒补偿。

本章思考与自测

Q: 如果把 set_complete 的 CAS 循环改成简单的 fetch_or(VALUE_SENT),在什么并发场景下会触发数据竞争?

参考解析:set_complete 用 CAS 循环而非 fetch_or 的原因在注释里写明 📎 tokio/src/sync/oneshot.rs:1517-1529:必须在置 VALUE_SENT 前检查 CLOSED。如果改成无条件 fetch_or,考虑这个时序:接收方先调用 close() 置 CLOSED 📎 tokio/src/sync/oneshot.rs:1569-1574,发送方随后 send 写入值并 fetch_or(VALUE_SENT)。此时 VALUE_SENT 和 CLOSED 同时置位,接收方的 poll_recv 看到 is_complete() 为真,会调用 consume_value 取走值 📎 tokio/src/sync/oneshot.rs:1325-1330;而发送方的 complete() 返回后,因为 prev.is_closed() 为真,会调用 consume_value 把值取回 📎 tokio/src/sync/oneshot.rs:1300-1315。两边同时访问 UnsafeCell,数据竞争。CAS 循环在发现 CLOSED 时提前 break,不置 VALUE_SENT,从而保证「关闭后发送方独占访问权」这一不变量。

Q: reserve_inner 里的 WakeReceiverOnDrop 守卫在成功路径上用 mem::forget 跳过,如果去掉这个 forget 会发生什么?

参考解析:守卫的 Drop 逻辑是「若信号量已关闭且空闲则唤醒接收方」📎 tokio/src/sync/mpsc/bounded.rs:1290-1298。成功路径上,acquire(n) 返回 Ok,调用者拿到许可并会构造 Permit,由 Permit 负责后续的通知职责。如果不去掉守卫,守卫在函数返回时 Drop,会额外检查一次「已关闭且空闲」——但此时许可已被 reserve_inner 的调用者持有,信号量并非空闲(is_idle 为假),所以实际上不会重复唤醒。但更关键的是语义清晰:成功路径的唤醒职责应完全由 Permit 承担,守卫只负责「取消/失败」路径的补偿。mem::forget 明确表达了「这条路径不需要守卫」的意图。如果去掉 forget 且恰好信号量处于「已关闭且空闲」的边界状态(比如 acquire 返回 Ok 但许可尚未被 Permit 接管),可能产生一次多余的唤醒——虽然不会导致错误,但会浪费一次调度。

Q: 若把 MutexGuard 改成持有信号量许可对象(像 SemaphorePermit 那样),会引入什么问题?

参考解析:当前 MutexGuard 只持有 &Mutex,Drop 时调用 self.lock.s.release(1) 📎 tokio/src/sync/mutex.rs:959-961。如果改成持有许可对象,会引入几个问题。其一,MutexGuard::map 系列方法需要把 guard 拆解成 MappedMutexGuard,只保护子字段 📎 tokio/src/sync/mutex.rs:869-883。当前设计下,MappedMutexGuard 只需持有 &Semaphore 和子字段指针 📎 tokio/src/sync/mutex.rs:190-199,Drop 时 self.s.release(1) 📎 tokio/src/sync/mutex.rs:1252-1262。如果 guard 持有许可对象,map 时就得转移许可对象的所有权,而 MappedMutexGuard 的字段布局会更复杂。其二,许可对象通常带 permits: usize 计数,对 Mutex 而言这个计数恒为 1,是冗余的。其三,MutexGuard 的 Send/Sync 边界已经通过 unsafe impl 精确控制 📎 tokio/src/sync/mutex.rs:260-263,持有许可对象会引入额外的 trait 约束。当前「只持引用 + 手动 release」的设计更轻量,也更容易支持 map。

至此,我们已经看清 tokio::sync 如何用「等待者队列 + Waker 唤醒」这一统一模式,支撑起 Mutex、Semaphore 与各类通道的异步等待。但并非所有阻塞都能被异步化——有些操作(如文件系统调用、CPU 密集计算)本质上会阻塞线程。下一章我们将进入 spawn_blocking 线程池与 block_on 的边界,看看 Tokio 如何在异步运行时与同步阻塞之间架起桥梁。

Waker 的存放位置因原语而异:Mutex/Semaphore 存在底层信号量的等待队列,oneshot 存在 Inner 的 tx_task/rx_task 字段,mpsc 存在 chan 模块的收发队列。但唤醒机制统一:状态变更时取出 Waker 调用 wake_by_ref,执行器重新调度任务。至此,异步原语内部的等待与唤醒已清晰可见。然而,并非所有代码都能异步化——下一章将探讨如何用 spawn_blocking 桥接阻塞操作,以及 block_on 如何在非异步上下文中驱动 Future。

AI 赋能代码库精读 · 本地优先架构

读完了本章?为你自己的私有项目生成专属架构全景书

基于 Tauri 2 + Rust 本地原生引擎,100% 源码离线隐私安全,零代码上传云端。像阅读一本传世专著一样拆解你的复杂系统。

⚡ Tauri 2 · Rust 原生引擎 · 100% 离线私密安全 · 适配超百万行代码库

CHAPTER 08

第 8 章:阻塞线程池与外部互操作:spawn_blocking 与任务隔离

官方源: tokio-rs/tokio · Commit @e800714a · 全书进度: 第 8 / 14 章

第 8 章:阻塞线程池与外部互操作:spawn_blocking 与任务隔离

上一章我们看到,异步 Mutex 和通道之所以能在等待时不占用线程,关键在于把 Waker 存进等待队列,等条件满足后再由唤醒者重新调度任务。但这一切的前提是任务能在 Pending 时主动让出线程。一旦代码调用 std::fs::read、libsqlite3 或纯 CPU 压缩循环,它就会霸占 worker 线程直到返回,期间该线程上的其他任务全部饿死。Tokio 的解法是把这类工作外包给独立的阻塞线程池,并用 block_on 在非异步上下文里驱动 Future。本章拆解这两条边界。

8.1 阻塞线程池的内存布局:Inner 与双实现队列

直觉模型:spawn_blocking 线程池就像餐厅的「外包帮工池」。前台服务员(worker 线程)只负责点单和传菜,遇到需要慢炖的菜就写一张工单丢进后厨的传菜窗口(队列),帮工(阻塞线程)从窗口取单。若没有这个池子,服务员就得亲自下厨,整个餐厅停摆。

核心结构。整个池子由 BlockingPool 持有,它只存两样东西:一个可克隆的 Spawner(投递入口)和一个 shutdown_rx(关闭信号接收端)📎 tokio/src/runtime/blocking/pool.rs:20-23。Spawner 内部是 Arc<Inner>,所有投递者共享同一份状态 📎 tokio/src/runtime/blocking/pool.rs:26-28。

Inner 是池子的全部状态,字段值得逐个看 📎 tokio/src/runtime/blocking/pool.rs:77-104:

  • inner_impl: InnerImpl:队列 + 通知 + 锁拓扑的实现,是一个枚举,有 Locked 和 Sharded 两个变体 📎 tokio/src/runtime/blocking/pool.rs:107-110。这是本章最关键的抽象——它把「单锁队列」和「分片队列」两种拓扑统一在一个接口下。
  • thread_cap: usize:线程数上限,即 max_blocking_threads。
  • scheduler_threads: usize:调度器 worker 线程数,用于在指标里扣除,使 num_blocking_threads 只统计阻塞线程 📎 tokio/src/runtime/blocking/pool.rs:455-460。
  • keep_alive: Duration:空闲线程存活时长,默认 KEEP_ALIVE = 10s 📎 tokio/src/runtime/blocking/pool.rs:231。
  • metrics: SpawnerMetrics:三个原子计数器——num_threads、num_idle_threads、queue_depth 📎 tokio/src/runtime/blocking/pool.rs:31-35。
〔设计推断与架构权衡〕
为什么用原子计数器而不是锁内字段? num_idle_threads 在 spawn_task 的热路径上被读取(判断是否需要唤醒空闲线程),如果它藏在 Mutex 里,每次投递都要先拿锁再读。把它做成 MetricAtomicUsize 后,投递路径可以在不持有队列锁的情况下先做一次快速判断。代价是这些计数与队列状态之间没有原子性保证,因此代码里用 num_notify 计数器来补偿——见下文。

线程管理状态。ThreadManagementState 被单独抽出来,供两种队列实现复用 📎 tokio/src/runtime/blocking/pool.rs:135-150:

  • shutdown: bool:关闭标志。
  • shutdown_tx: Option<shutdown::Sender>:每个 worker 线程持有一份克隆,全部 drop 后 shutdown_rx 收到通知。
  • last_exiting_thread: Option<JoinHandle<()>>:上一个超时退出的线程句柄。
  • worker_threads: HashMap<usize, JoinHandle<()>>:所有存活 worker 的句柄。
  • worker_thread_index: usize:单调递增的线程 ID 分配器。

last_exiting_thread 的设计动机在注释里写得很清楚:超时退出的线程会 join 上一个超时退出的线程,避免 Valgrind 误报 📎 tokio/src/runtime/blocking/pool.rs:135-150。worker_timed_out 正是这个链式 join 的实现——它移除自己的句柄,把旧的 last_exiting_thread 换出来返回给调用者去 join 📎 tokio/src/runtime/blocking/pool.rs:172-178。

任务封装。队列里存的是 Task,它包了一个 UnownedTask<BlockingSchedule> 和一个 Mandatory 标志 📎 tokio/src/runtime/blocking/pool.rs:187-191。Mandatory 决定关闭时这个任务是被丢弃还是强制执行:shutdown_or_run_if_mandatory 在 NonMandatory 时调 shutdown(),在 Mandatory 时调 run() 📎 tokio/src/runtime/blocking/pool.rs:223-228。这就是 spawn_blocking(非强制)与 spawn_mandatory_blocking(强制,供 fs 使用)的区别 📎 tokio/src/runtime/blocking/pool.rs:233-265。

单锁实现的内存布局。LockedImpl 是最原始的拓扑:一个 Mutex<LockedInner> 加一个 Condvar 📎 tokio/src/runtime/blocking/pool.rs:113-116。LockedInner 里是 VecDeque<Task>、num_notify: u32 和 thread_mgmt_state 📎 tokio/src/runtime/blocking/pool.rs:118-124。注意 num_notify 与 thread_mgmt_state 在同一个锁下,而 num_idle_threads 是锁外的原子量——这种「部分状态在锁内、部分在锁外」的混合布局,正是后面所有并发微妙性的根源。

8.2 投递路径:从 spawn_blocking 到线程唤醒

场景:异步任务里调用 tokio::task::spawn_blocking(move || heavy_compute(data)),此刻发生了什么?

第一步:装箱决策与任务构造。Spawner::spawn_blocking 先测量闭包大小 fn_size,然后根据 AutoBox::<F>::SHOULD_BOX 决定是否把闭包 Box 起来 📎 tokio/src/runtime/blocking/pool.rs:359-389。这是 Tokio 通用的「大 Future 自动装箱」策略:闭包过大时装箱,避免任务结构体膨胀。

进入 spawn_blocking_inner,先分配任务 ID,再用 blocking_task 把闭包包成一个 Future,最后用 task::unowned 构造出 UnownedTask 和 JoinHandle 📎 tokio/src/runtime/blocking/pool.rs:440-449。注意这里返回的是 (JoinHandle<R>, Result<(), SpawnError>) 二元组——句柄和投递结果分开返回。

第二步:投递结果的三种处理。回到 spawn_blocking,对 spawn_result 做匹配 📎 tokio/src/runtime/blocking/pool.rs:381-388:

  • Ok(()):正常,返回句柄。
  • Err(ShuttingDown):不 panic,仍然返回句柄。注释说明这是兼容性考虑——句柄永远不会 resolve,但调用方不会因为运行时正在关闭而崩溃。
  • Err(NoThreads(e)):OS 无法创建线程且池中无人接手,直接 panic。

第三步:入队与唤醒决策。spawn_task 把 on_no_idle 闭包传给 InnerImpl::spawn_task,由具体实现决定何时调用它 📎 tokio/src/runtime/blocking/pool.rs:462-506。看 LockedImpl::spawn_task 的临界区 📎 tokio/src/runtime/blocking/pool.rs:603-639:

rust
let mut locked = self.mutex.lock();

if locked.thread_mgmt_state.shutdown {
    task.task.shutdown();
    return Err(SpawnError::ShuttingDown);
}

locked.queue.push_back(task);
metrics.inc_queue_depth();

if metrics.num_idle_threads() == 0 {
    on_no_idle(&mut locked.thread_mgmt_state)?;
} else {
    metrics.dec_num_idle_threads();
    locked.num_notify += 1;
    self.condvar.notify_one();
}

这里有两个关键点。其一,关闭检查在入队之前,且即使任务是 Mandatory 也直接 shutdown()——注释解释:它在关闭开始之后才被调度,所以丢弃是合法的 📎 tokio/src/runtime/blocking/pool.rs:614-620。其二,唤醒决策依赖锁外的 num_idle_threads:若为 0,调 on_no_idle 尝试起新线程;否则递减空闲计数、递增 num_notify、notify_one。

num_notify 为什么必须存在? 因为 Condvar 可能产生虚假唤醒(spurious wakeup)。如果只用 notify_one 而不计数,一个虚假唤醒的线程会误以为有任务可取,结果发现队列为空又睡回去,而真正被唤醒的线程可能永远收不到通知。num_notify 把「合法唤醒」变成可计数的令牌:投递方 +1,被唤醒方在 num_notify != 0 时才认为唤醒合法并 -1 📎 tokio/src/runtime/blocking/pool.rs:674-684。

第四步:起新线程。on_no_idle 闭包在持有队列锁的情况下执行 📎 tokio/src/runtime/blocking/pool.rs:462-506。它先检查 num_threads == thread_cap,达到上限就直接返回 Ok(())——任务留在队列里等现有线程处理,这就是背压。否则克隆 shutdown_tx,调 spawn_thread 创建线程,成功后递增 num_threads、递增 worker_thread_index、把句柄插入 worker_threads。

spawn_thread 用 thread::Builder 设置线程名和栈大小,然后 spawn 一个闭包:进入运行时上下文 rt.enter(),调用 inner.run(id),最后 drop shutdown_tx 📎 tokio/src/runtime/blocking/pool.rs:508-528。

OS 线程创建失败的容错。spawn_thread 可能失败。代码对错误做了分类 📎 tokio/src/runtime/blocking/pool.rs:488-500:若是 WouldBlock(临时性错误,由 is_temporary_os_thread_error 判定 📎 tokio/src/runtime/blocking/pool.rs:750-752)且池中已有阻塞线程,则静默忽略——任务会被某个当前忙碌的线程最终取走。否则返回 SpawnError::NoThreads,最终导致 panic。

用一张控制流图总结投递路径的决策分支:

mermaid
flowchart TD
    call["Spawner::spawn_blocking(func)"] --> box{"AutoBox::SHOULD_BOX?"}
    box -->|是| boxed["Box::new(func)"]
    box -->|否| raw["func"]
    boxed --> inner["spawn_blocking_inner"]
    raw --> inner
    inner --> unowned["task::unowned -> Task + JoinHandle"]
    unowned --> spawn_task["InnerImpl::spawn_task"]
    spawn_task --> lock["LockedImpl: mutex.lock()"]
    lock --> shutting{"thread_mgmt_state.shutdown?"}
    shutting -->|是| discard["task.task.shutdown()"]
    discard --> err_sd["Err(ShuttingDown)"]
    shutting -->|否| push["queue.push_back(task)"]
    push --> idle{"num_idle_threads == 0?"}
    idle -->|是| on_no_idle["on_no_idle(thread_mgmt_state)"]
    on_no_idle --> cap{"num_threads == thread_cap?"}
    cap -->|是| backpressure["返回 Ok, 任务留队列"]
    cap -->|否| spawn_th["spawn_thread(shutdown_tx, rt, id)"]
    spawn_th --> th_ok{"spawn 成功?"}
    th_ok -->|是| reg["inc_num_threads, 注册 JoinHandle"]
    th_ok -->|否| tmp{"WouldBlock 且已有线程?"}
    tmp -->|是| ignore["忽略, 等忙碌线程取走"]
    tmp -->|否| err_nt["Err(NoThreads)"]
    idle -->|否| notify["dec_num_idle_threads, num_notify+=1, notify_one"]
    err_sd --> ret["返回 JoinHandle"]
    backpressure --> ret
    reg --> ret
    ignore --> ret
    err_nt --> panic_os["panic: OS can't spawn worker thread"]

8.3 worker 主循环:BUSY/IDLE 状态机与超时回收

直觉模型:每个阻塞线程就是一个「待命帮工」。有单时连续干活(BUSY),没单时打盹(IDLE),打盹超过 keep_alive 就下班(超时退出)。若没有超时回收,池子会永久保留峰值时创建的所有线程,浪费内存与内核调度开销。

主循环结构。LockedImpl::run_worker 是一个 'main 循环,内部交替处于 BUSY 和 IDLE 两个阶段 📎 tokio/src/runtime/blocking/pool.rs:642-735。注意:这里的 BUSY/IDLE 是循环内的阶段,不是显式枚举状态,所以下面用流程图而非状态图描述。

BUSY 阶段:内层 while let Some(task) = locked.queue.pop_front() 不断取任务 📎 tokio/src/runtime/blocking/pool.rs:655-661。取到后递减 queue_depth,drop 锁,执行 task.run(),再重新拿锁。drop 锁这一步至关重要——阻塞任务可能跑很久,绝不能持锁执行。

IDLE 阶段:队列空了,递增 num_idle_threads,设 is_counted_idle = true,然后进入等待循环 📎 tokio/src/runtime/blocking/pool.rs:663-696。核心是 condvar.wait_timeout(locked, keep_alive),返回后检查三件事:

1. num_notify != 0:合法唤醒。递减 num_notify,设 is_counted_idle = false(因为投递方已经递减过 num_idle_threads 了),break 回 BUSY 📎 tokio/src/runtime/blocking/pool.rs:674-684。

2. 未关闭且超时:调 worker_timed_out 拿到上一个退出线程的句柄,break 'main 退出循环 📎 tokio/src/runtime/blocking/pool.rs:689-693。

3. 否则是虚假唤醒,继续等待。

关闭时的队列排空。若 thread_mgmt_state.shutdown 为真,进入排空逻辑 📎 tokio/src/runtime/blocking/pool.rs:698-710:逐个弹出任务,drop 锁,调 task.shutdown_or_run_if_mandatory()——非强制任务被丢弃,强制任务照常执行。然后 break 退出主循环。

退出清理。线程退出前递减 num_threads 📎 tokio/src/runtime/blocking/pool.rs:714。若 is_counted_idle 为真,还要递减 num_idle_threads,并用 assert_ne!(prev_idle, 0) 断言没有下溢 📎 tokio/src/runtime/blocking/pool.rs:716-726。这个断言是调试期的护栏:一旦 num_idle_threads 记账出错,这里会立刻 panic 而不是让错误静默传播。

最后,若正在关闭且 num_threads == 0(最后一个线程),notify_one 唤醒可能在等待的关闭发起者 📎 tokio/src/runtime/blocking/pool.rs:728-730。返回 join_on_thread,由 Inner::run 在退出前 join 📎 tokio/src/runtime/blocking/pool.rs:755-771。

关闭握手。BlockingPool::shutdown 先调 begin_shutdown 拿到所有 worker 句柄 📎 tokio/src/runtime/blocking/pool.rs:310-312。LockedImpl::begin_shutdown 设置关闭标志、drop shutdown_tx、notify_all 唤醒所有等待线程 📎 tokio/src/runtime/blocking/pool.rs:740-745。然后 shutdown_rx.wait(timeout) 阻塞等待 📎 tokio/src/runtime/blocking/pool.rs:324。

shutdown::Receiver::wait 的实现很讲究 📎 tokio/src/runtime/blocking/shutdown.rs:37-70:先处理 timeout == 0 的快速路径直接返回 false;再调 try_enter_blocking_region() 进入阻塞区域,若失败且当前正在 panic 则返回 false,否则 panic 并给出「不能在异步上下文中 drop runtime」的提示 📎 tokio/src/runtime/blocking/shutdown.rs:44-57。最后根据 timeout 调 block_on_timeout 或 block_on 驱动那个 oneshot。

shutdown_tx 的机制是:每个 worker 线程持有一份 Arc<oneshot::Sender<()>> 的克隆 📎 tokio/src/runtime/blocking/shutdown.rs:12-14。所有线程退出后,所有克隆被 drop,Arc 计数归零,oneshot::Sender 被 drop,Receiver 收到通知。这就是「所有 Sender drop 后 Receiver 被唤醒」的经典模式。

mermaid
sequenceDiagram
    participant App as "应用线程 (drop Runtime)"
    participant Pool as "BlockingPool::shutdown"
    participant Locked as "LockedImpl"
    participant Worker as "阻塞 worker 线程"
    participant Rx as "shutdown::Receiver"

    App->>Pool: shutdown(timeout)
    Pool->>Locked: begin_shutdown()
    Locked->>Locked: thread_mgmt_state.begin_shutdown() 设 shutdown=true, shutdown_tx=None
    Locked->>Worker: condvar.notify_all()
    Locked-->>Pool: Some((last_exited_thread, workers))
    Pool->>Rx: wait(timeout)
    Worker->>Worker: 从 wait_timeout 醒来, 见 shutdown=true
    Worker->>Worker: 排空队列 shutdown_or_run_if_mandatory()
    Worker->>Worker: dec_num_threads, 退出 run_worker
    Worker->>Worker: drop(shutdown_tx) 克隆
    Worker-->>Rx: 最后一个 Sender drop, oneshot 完成
    Rx-->>Pool: 返回 true
    Pool->>Worker: join 所有 worker 句柄

8.4 block_on:在非异步上下文驱动 Future

直觉模型:block_on 是运行时的「正门」。它把当前线程变成临时的执行器,反复 poll 传入的 Future 直到完成。若没有它,main 函数就无法启动任何异步代码。

入口与装箱。Runtime::block_on 同样先测大小、按 SHOULD_BOX 决定是否 Box::pin,然后进 block_on_inner 📎 tokio/src/runtime/runtime.rs:343-350。block_on_inner 里有两段条件编译的 trace 包装(taskdump 和 tracing),然后 self.enter() 进入运行时上下文,最后按调度器类型分派 📎 tokio/src/runtime/runtime.rs:353-383:

rust
let _enter = self.enter();

match &self.scheduler {
    Scheduler::CurrentThread(exec) => exec.block_on(&self.handle.inner, future),
    Scheduler::MultiThread(exec) => exec.block_on(&self.handle.inner, future),
}

两种调度器的 block_on 语义不同,文档里说得很清楚 📎 tokio/src/runtime/runtime.rs:302-320:

  • 多线程调度器:Future 在 I/O 驱动和定时器上下文中运行,block_on 返回后已 spawn 的任务继续运行。
  • 当前线程调度器:block_on 可以被多个线程并发调用,第一个调用者取得 I/O 和定时器驱动的所有权,其他线程「钩入」它。第一个 block_on 完成后,其他线程可以「偷走」驱动。block_on 返回后已 spawn 的任务被挂起,再次调用 block_on 会恢复它们。

关键限制:不能在异步上下文中调用。文档明确 block_on 在异步执行上下文中调用会 panic 📎 tokio/src/runtime/runtime.rs:321-324。原因很直接:block_on 会阻塞当前线程直到 Future 完成,若当前线程本身是某个 worker 线程,就会阻塞整个执行器——这正是 spawn_blocking 要解决的问题,所以两者互斥。

关闭路径。Runtime::drop 按调度器类型分派 📎 tokio/src/runtime/runtime.rs:506-521:当前线程调度器需要先 try_set_current 进入上下文再 shutdown(保证任务在运行时上下文中被 drop);多线程调度器直接 shutdown(worker 线程本身已在上下文中)。shutdown_timeout 先关调度器再关阻塞池 📎 tokio/src/runtime/runtime.rs:457-461,shutdown_background 等价于 shutdown_timeout(Duration::from_nanos(0)) 📎 tokio/src/runtime/runtime.rs:494-496。

设计思考、错误恢复与生产踩坑

为什么 spawn_blocking 的 ShuttingDown 不 panic? 📎 tokio/src/runtime/blocking/pool.rs:383-384 注释说是兼容性考虑。spawn_blocking 返回 JoinHandle 而非 Result,若在关闭时 panic,会让「运行时正在关闭」这个可预期状态变成崩溃。返回一个永不 resolve 的句柄,调用方 await 时会一直挂起——但此时运行时已关闭,整个 block_on 也会退出,所以实际不会永久泄漏。

max_blocking_threads 的背压语义。默认值很大(512),因为 spawn_blocking 常用于文件 I/O。但文档警告:跑 CPU 密集任务时要用信号量限制并发,否则会创建大量线程 📎 tokio/src/task/blocking.rs:94-100。达到上限后任务在队列里排队,形成背压——但注意这个背压只作用于阻塞池,不会反压到异步调度器。

spawn_blocking 不可取消。文档明确:abort 对已开始运行的阻塞任务无效,任务会继续跑完 📎 tokio/src/task/blocking.rs:106-120。只有尚未开始的任务可能被 abort 阻止。关闭时运行时会等待所有已开始的阻塞任务,shutdown_timeout 超时后会泄漏这些线程。

num_idle_threads 的记账陷阱。is_counted_idle 标志的存在说明这个计数很容易出错。投递方在唤醒时递减 num_idle_threads,被唤醒方看到 num_notify != 0 后设 is_counted_idle = false,避免重复递减 📎 tokio/src/runtime/blocking/pool.rs:679-682。若这条路径有 bug,assert_ne!(prev_idle, 0) 会在退出时 panic 📎 tokio/src/runtime/blocking/pool.rs:722-725。生产环境若见到「num_idle_threads underflowed on thread exit」,说明池子的记账逻辑被破坏。

〔设计推断与架构权衡〕
last_exiting_thread 链式 join 的代价。超时退出的线程会 join 上一个超时退出的线程 📎 tokio/src/runtime/blocking/pool.rs:172-178。这形成一个 join 链:每个退出的线程都要等前一个真正结束。在高频创建/销毁阻塞线程的场景下,这条链可能变长,导致线程退出延迟累积。 这是为了避免 Valgrind 误报而做的权衡,正常生产环境影响有限,但在线程频繁超时的负载下值得关注。

InnerImpl 枚举抽象的意义。注释说明 Locked 变体的行为与重构前完全一致,而 Sharded 变体为未来的并发队列预留了对称的槽位 📎 tokio/src/runtime/blocking/pool.rs:537-539。spawn_task、run_worker、begin_shutdown 三个方法都通过枚举分派 📎 tokio/src/runtime/blocking/pool.rs:548-582。这种「枚举分派 + 每变体自持临界区」的设计,使得新增队列拓扑时不需要改动调用方。

本章小结

本章拆解了 Tokio 容纳同步代码的两条边界。spawn_blocking 把闭包投递到独立的阻塞线程池:Inner 持有队列、线程上限、存活时长与原子指标;LockedImpl 用单锁 + Condvar 实现队列,num_notify 计数器补偿虚假唤醒;worker 在 BUSY/IDLE 间循环,空闲超时后链式 join 退出;max_blocking_threads 达到上限后任务排队形成背压。block_on 则在非异步上下文驱动 Future,多线程与当前线程调度器语义不同,且严禁在异步上下文中调用。关闭路径通过 shutdown_tx 的 Arc 计数归零触发 oneshot,实现「所有 worker 退出后唤醒关闭发起者」的握手。

本章思考与自测

Q1: 若把 LockedImpl::spawn_task 中 if metrics.num_idle_threads() == 0 的判断改成恒为真(即每次都调 on_no_idle),在高并发投递场景下会发生什么?为什么?

参考解析:on_no_idle 会检查 num_threads == thread_cap,未达上限就创建新线程 📎 tokio/src/runtime/blocking/pool.rs:471-487。若判断恒为真,即使有空闲线程也会尝试起新线程,导致线程数迅速冲到 thread_cap。更严重的是,空闲线程不会被 notify_one 唤醒(因为走了 on_no_idle 分支而非 else 分支的 num_notify += 1; notify_one 📎 tokio/src/runtime/blocking/pool.rs:627-636),队列里的任务可能无人处理,直到某个新线程启动后才发现队列非空。这会造成「线程爆满但任务仍排队」的假死状态。原判断的意义正是:有空闲线程时优先唤醒它们,避免无谓的线程创建。

Q2: LockedImpl::run_worker 在 BUSY 阶段执行 task.run() 前会 drop(locked) 📎 tokio/src/runtime/blocking/pool.rs:657-658。如果去掉这个 drop,在什么场景下会触发死锁?

参考解析:task.run() 执行的是用户闭包,闭包内部完全可能再次调用 spawn_blocking 投递新任务。投递路径 LockedImpl::spawn_task 第一件事就是 self.mutex.lock() 📎 tokio/src/runtime/blocking/pool.rs:612。若 worker 持锁执行闭包,闭包内的投递就会尝试获取同一把锁,而 std::sync::Mutex 不可重入,直接死锁。此外,持锁执行长任务会阻塞所有其他投递者和 worker 的取任务操作,即使不死锁也会让整个池子串行化。drop(locked) 是必须的。

Q3: shutdown::Receiver::wait 在 try_enter_blocking_region() 失败且当前正在 panic 时返回 false,否则 panic 📎 tokio/src/runtime/blocking/shutdown.rs:44-57。为什么要在 panic 时特殊处理?如果去掉这个分支,在什么场景下会出问题?

参考解析:try_enter_blocking_region 失败意味着当前处于异步上下文,不允许阻塞。正常情况下应 panic 提示用户「不能在异步上下文中 drop runtime」。但如果当前线程已经在 panic(std::thread::panicking() 为真),再 panic 会导致双重 panic,Rust 默认行为是直接 abort 进程。场景:用户在异步任务里 drop 一个 Runtime,而该任务本身因为其他原因正在 panic,此时 drop 触发的 shutdown 会二次 panic。返回 false 让 shutdown 放弃等待,避免进程 abort,给用户保留看到原始 panic 信息的机会。这是「panic 安全」的典型处理。

阻塞线程池与 block_on 划定了异步运行时的能力边界:前者把无法让出线程的工作隔离到专用线程,后者让非异步入口也能驱动 Future。但这两条边界在代码里往往不是手写的——下一章我们将进入宏的世界,看看 #[tokio::main]、select! 和 join! 如何在编译期生成这些运行时代码。

AI 赋能代码库精读 · 本地优先架构

读完了本章?为你自己的私有项目生成专属架构全景书

基于 Tauri 2 + Rust 本地原生引擎,100% 源码离线隐私安全,零代码上传云端。像阅读一本传世专著一样拆解你的复杂系统。

⚡ Tauri 2 · Rust 原生引擎 · 100% 离线私密安全 · 适配超百万行代码库

CHAPTER 09

第 9 章:协作式调度预算:coop 机制如何防止异步任务饥饿

官方源: tokio-rs/tokio · Commit @e800714a · 全书进度: 第 9 / 14 章

第 9 章:协作式调度预算:coop 机制如何防止异步任务饥饿

上一章我们看到 block_on 与阻塞线程池如何划定异步运行时的能力边界,而用户几乎从不手写这些边界——他们写 #[tokio::main]、select!、join!,让宏在编译期把这些样板代码铺开。宏是 Tokio 给用户的第一层糖衣,也是编译期真正生成运行时代码的地方。本章聚焦 tokio-macros crate 与 tokio/src/macros/select.rs,拆解三条最常用的宏展开路径,重点回答一个问题:宏展开后,真实的调用链长什么样,以及为什么 select! 的取消安全语义必须单独警惕。

9.1 #[tokio::main]:把 async fn 改写成 Runtime::block_on

直觉模型:#[tokio::main] 就像一张「装修委托书」。你交出一间毛坯房(async fn main),它替你铺好水电(构建 Runtime)、装好门窗(enable_all),最后把你原本的家具(函数体)搬进去。若没有它,每个 main 都得手写 Builder::new_multi_thread().enable_all().build().unwrap().block_on(...),样板代码会淹没业务逻辑。

数据结构与内存布局

宏本身不产生运行时数据结构,但它解析出的配置被装进两个结构体。Configuration 是「解析期的可变累加器」,字段全是 Option,因为属性参数可能缺省、可能重复、可能非法 📎 tokio-macros/src/entry.rs:74-84。注意 worker_threads、start_paused、unhandled_panic 都带 Span——这是为了在报错时把错误定位到用户写的那一行,而不是宏内部 📎 tokio-macros/src/entry.rs:74-84。FinalConfig 则是「校验后的不可变结果」,flavor 不再是 Option,因为 build() 已经用 default_flavor 兜底 📎 tokio-macros/src/entry.rs:55-62。

RuntimeFlavor 只有三个变体:CurrentThread、Threaded、Local 📎 tokio-macros/src/entry.rs:10-14。from_str 里特意为历史遗留名字给出友好报错:single_thread 提示应叫 current_thread,basic_scheduler 提示已改名,threaded_scheduler 提示已改名 📎 tokio-macros/src/entry.rs:17-27。这是宏作为「用户第一接触面」的典型设计:错误信息即文档。

Step-by-Step 展开流程

代入场景:用户写下 #[tokio::main(flavor = "multi_thread", worker_threads = 4)] async fn main() { ... }。

第一步,main 入口先解析 item 为自定义的 ItemFn 📎 tokio-macros/src/entry.rs:577-580。这个 ItemFn 不是 syn::ItemFn,而是 Tokio 自己实现的解析器,原因写在注释里:它不想递归解析整条语句,只做「按 token tree 缓冲、遇到分号切分」的轻量解析 📎 tokio-macros/src/entry.rs:720-764。这避免了在宏里对函数体做完整 AST 构建的开销。

第二步,build_config 校验 async 关键字是否存在,缺失则报 "the async keyword is missing" 📎 tokio-macros/src/entry.rs:346-349。随后遍历属性参数,把 worker_threads、flavor、start_paused、crate、unhandled_panic、name 分派到对应 setter 📎 tokio-macros/src/entry.rs:369-399。注意 core_threads 被显式拒绝并提示已改名 📎 tokio-macros/src/entry.rs:379-382。

第三步,Configuration::build 做跨字段一致性校验。这里有三条关键约束:worker_threads 只允许 multi_thread 📎 tokio-macros/src/entry.rs:197-217;start_paused 只允许 current_thread/local 📎 tokio-macros/src/entry.rs:219-229;unhandled_panic 同样只允许 current_thread/local 📎 tokio-macros/src/entry.rs:231-241。若用户选了 multi_thread 但 rt-multi-thread feature 未开,报错信息会根据是否显式指定 flavor 而不同 📎 tokio-macros/src/entry.rs:209-216。

第四步,parse_knobs 生成代码。它先抹掉 asyncness 📎 tokio-macros/src/entry.rs:441,然后根据 flavor 选择 builder 起点:CurrentThread/Local 用 Builder::new_current_thread(),Threaded 用 Builder::new_multi_thread() 📎 tokio-macros/src/entry.rs:468-477。Local 特殊之处在于 build 调用是 build_local(Default::default()) 而非 build() 📎 tokio-macros/src/entry.rs:479-483。随后按需链式追加 .worker_threads(#v)、.start_paused(#v)、.unhandled_panic(...)、.name(#v) 📎 tokio-macros/src/entry.rs:485-497。

第五步,生成最终函数体。核心是 last_block:return #rt.enable_all().#build.expect("Failed building the Runtime").block_on(body) 📎 tokio-macros/src/entry.rs:509-522。注意那个显式 return,注释指向 tokio-rs/tokio#4636,是为了修复类型推断问题 📎 tokio-macros/src/entry.rs:508。

第六步,函数体被包成 async #body 并做类型检查。非 test 路径下,若返回类型不是 ! 且不含 impl Trait,会插入 if false { let _: &dyn Future<Output = #output_type> = &body; } 做编译期断言 📎 tokio-macros/src/entry.rs:551-571。test 路径则用 pin! 把 body 钉在栈上并转成 Pin<&mut dyn Future>,注释解释这是为了减少 block_on 泛型实例化的编译开销 📎 tokio-macros/src/entry.rs:526-548。

mermaid
flowchart TD
    entry["main(args, item)"] --> parse_item{"syn::parse2(item) 成功?"}
    parse_item -->|否| err_ret["token_stream_with_error 返回原始 item + 编译错误"]
    parse_item -->|是| check_main{"ident == main 且有参数?"}
    check_main -->|是| err_args["报错: main 不能接受参数"]
    check_main -->|否| parse_args["AttributeArgs::parse_terminated"]
    parse_args --> build_cfg["build_config 校验 async 与各字段"]
    build_cfg --> cfg_ok{"config 构建成功?"}
    cfg_ok -->|否| fallback["parse_knobs(DEFAULT_ERROR_CONFIG) + 错误"]
    cfg_ok -->|是| knobs["parse_knobs 生成 Builder 链 + block_on"]
    knobs --> out["输出同步 fn main"]

设计思考与生产踩坑

main 与 test 共享 parse_knobs,但默认 flavor 不同:test 默认 CurrentThread,main 默认 Threaded 📎 tokio-macros/src/entry.rs:91-94。这解释了为什么 #[tokio::test] 默认单线程——测试通常不需要多核,且单线程更容易复现。

一个容易被忽略的坑:宏展开后每次调用函数都会新建 Runtime。文档明确警告,若函数被频繁调用,应改用 Builder 复用 Runtime 📎 tokio-macros/src/lib.rs:31-35。把 #[tokio::main] 用在普通函数上是合法的,但每次调用都付一次 Runtime 构建成本。

另一个坑是 crate 重命名。当用户 use tokio as tokio1 时,宏内部默认生成的 tokio::runtime::Builder 会找不到路径,必须显式 crate = "tokio1" 📎 tokio-macros/src/lib.rs:239-264。parse_knobs 里 crate_path 的默认值是 Ident::new("tokio", ...) 📎 tokio-macros/src/entry.rs:456-462,这正是重命名场景报错的根源。

9.2 select!:多分支轮询、位掩码与随机公平性

直觉模型:select! 像一位「同时盯多个取餐窗口的服务员」。哪个窗口先出餐,他就端走哪份,其余窗口的排队作废。若没有它,用户得手写 poll_fn 把多个 Future 塞进一个元组逐个 poll,还要自己处理「某个分支就绪后其余分支该丢弃」的逻辑。

数据结构与内存布局

select! 展开后生成一个局部模块 __tokio_select_util,里面有一个枚举 Out 和一个类型别名 Mask 📎 tokio/src/macros/select.rs:615-619。Out 的变体名是 _0、_1……每个分支一个,外加一个 Disabled 表示所有分支都失效 📎 tokio-macros/src/select.rs:33-39。Mask 的底层类型按分支数动态选择:≤8 用 u8,≤16 用 u16,≤32 用 u32,≤64 用 u64,超过 64 直接 panic 📎 tokio-macros/src/select.rs:17-31。这个位掩码是 select! 的核心状态:第 i 位为 1 表示第 i 个分支已被禁用。

所有 Future 被存进一个元组 futures,每个元素先经 IntoFuture::into_future 转换 📎 tokio/src/macros/select.rs:654-656。注意这里先构造 futures_init 再逐个 into_future,注释解释这是为了利用临时生命周期延长 📎 tokio/src/macros/select.rs:641-646。随后 let mut futures = &mut futures; 把元组降级为可变引用,避免 poll_fn 闭包夺取所有权 📎 tokio/src/macros/select.rs:658-662。

Step-by-Step 轮询流程

代入场景:select! { v = stream1.next() => ..., v = stream2.next() => ..., else => break }。

第一步,宏入口规则匹配。若有 biased; 前缀,start=0 📎 tokio/src/macros/select.rs:801-803;否则 start 是一个随机表达式 thread_rng_n(BRANCHES) 📎 tokio/src/macros/select.rs:805-809。这就是文档所说的「默认随机挑选分支先检查」的公平性来源 📎 tokio/src/macros/select.rs:61-65。

第二步,归一化。tt-muncher 把每个分支规整成 (skip) pat = fut, if cond => handler, 形式,skip 是一串 _,长度等于该分支之前的 branch 数 📎 tokio/src/macros/select.rs:770-793。skip 既用于生成元组字段访问 futures_init.$($skip)*,也用于 count! 算出分支索引。

第三步,前置条件求值。对每个分支的 if $c,若为 false,则 disabled |= 1 << index 📎 tokio/src/macros/select.rs:631-636。注意:即使分支被禁用,其 $fut 表达式仍会被求值,只是不会被 poll 📎 tokio/src/macros/select.rs:39-41。

第四步,进入 poll_fn 闭包。先检查协作预算:ready!(poll_budget_available(cx)),预算耗尽直接返回 Pending 📎 tokio/src/macros/select.rs:664-667。这保证 select! 不会霸占 worker。

第五步,循环 for i in 0..BRANCHES,branch = (start + i) % BRANCHES 📎 tokio/src/macros/select.rs:680-685。对每个 branch:先查 disabled & mask == mask,已禁用则 continue 📎 tokio/src/macros/select.rs:694-699;否则从元组取出该 Future,用 Pin::new_unchecked 包一层(安全性依赖 Future 存于栈上且不被移动)📎 tokio/src/macros/select.rs:701-707;poll 之,Ready(out) 则先 disabled |= mask 再匹配模式 📎 tokio/src/macros/select.rs:710-730。

第六步,模式匹配。若 out 匹配 $bind,返回 Poll::Ready(Out::_i(out)) 📎 tokio/src/macros/select.rs:727-733;若不匹配,continue 继续轮询其他分支——这正是文档步骤 5 所说的「模式不匹配则禁用当前分支」📎 tokio/src/macros/select.rs:44-47。

第七步,循环结束。若 is_pending 为真返回 Pending,否则所有分支都失效,返回 Out::Disabled 📎 tokio/src/macros/select.rs:740-745。外层 match output 把 Out::_i 映射到对应 handler,Disabled 映射到 else 表达式 📎 tokio/src/macros/select.rs:749-755。

mermaid
flowchart TD
    start["poll_fn 闭包被调用"] --> budget{"poll_budget_available(cx)?"}
    budget -->|否| pending_budget["返回 Pending"]
    budget -->|是| init["is_pending = false; start = $start"]
    init --> loop{"i |否| check_pending{"is_pending?"}
    check_pending -->|是| pending["返回 Pending"]
    check_pending -->|否| disabled_out["返回 Out::Disabled"]
    loop -->|是| branch["branch = (start+i) % BRANCHES"]
    branch --> is_disabled{"disabled & mask == mask?"}
    is_disabled -->|是| next_i["i += 1"]
    is_disabled -->|否| poll_fut["Pin::new_unchecked(fut).poll(cx)"]
    poll_fut --> poll_res{"Poll 结果?"}
    poll_res -->|Pending| set_pending["is_pending = true; i += 1"]
    poll_res -->|Ready| disable["disabled |= mask"]
    disable --> pat_match{"out 匹配 $bind?"}
    pat_match -->|否| next_i
    pat_match -->|是| ready_out["返回 Out::_i(out)"]
    next_i --> loop
    set_pending --> loop

设计思考与生产踩坑

为什么用位掩码而不是 Vec<bool>? 位掩码是栈上单个整数,无堆分配,且 disabled |= mask 是单条指令。对于热路径上的 select!,这避免了每次迭代的堆访问。

为什么模式不匹配要禁用分支? 这是 select! 与「简单 race」的关键区别。考虑 Some(v) = stream.next() => ...,若 stream.next() 返回 None(流结束),模式不匹配,该分支被永久禁用,避免无限轮询一个已结束的流。文档示例正是靠这个语义收集两个流直到都结束 📎 tokio/src/macros/select.rs:198-223。

取消安全的真正含义:select! 一旦某分支就绪,其余分支的 Future 会被 drop。若被 drop 的 Future 已经消费了数据但尚未返回,数据就丢了。文档明确列出 read_exact、read_to_end、write_all 不取消安全 📎 tokio/src/macros/select.rs:119-124,而 Mutex::lock、Semaphore::acquire 因为排队公平性,取消会丢失队列位置 📎 tokio/src/macros/select.rs:126-133。判定方法:找 .await 点,若在 .await 处重启函数仍正确,则取消安全 📎 tokio/src/macros/select.rs:135-139。

if 前置条件的竞态陷阱:文档给了一个经典错误示例——用 if !sleep.is_elapsed() 守卫 sleep 分支,但 is_elapsed() 可能在 while 检查与 select! 之间变为 true,导致超时被漏掉 📎 tokio/src/macros/select.rs:336-376。正确写法是去掉 if,让 sleep 分支始终参与轮询,超时后 break 📎 tokio/src/macros/select.rs:378-405。

biased; 的代价:随机 RNG 有 CPU 成本,且某些场景需要确定的轮询顺序 📎 tokio/src/macros/select.rs:67-74。但 biased; 把公平性责任交给用户:若一个分支永远就绪,后面的分支会饿死 📎 tokio/src/macros/select.rs:75-81。

9.3 join! 与宏展开的工程约束

直觉模型:join! 像「同时等所有快递都到齐」。它不像 select! 那样谁先到就取消其余,而是把所有 Future 的 Ready 值聚合成一个元组。若没有它,用户得手写 poll_fn 维护每个 Future 的完成状态。

数据结构与内存布局

join! 的展开同样基于元组存 Future,但状态不是位掩码,而是一个「已完成值」的元组。每个 Future 完成后,其值被取出存入结果元组,对应槽位标记为已完成。与 select! 不同,join! 不会 drop 未完成的 Future——它必须等所有 Future 都完成才返回。

Step-by-Step 流程

join! 的轮询逻辑与 select! 共享「元组存 Future + poll_fn 驱动」的骨架,但语义相反:select! 是「任一就绪即返回」,join! 是「全部就绪才返回」。每轮 poll 遍历所有未完成的 Future,任一返回 Pending 则整体 Pending,全部 Ready 则聚合返回。

mermaid
flowchart LR
    subgraph input["输入"]
        f1["Future A"]
        f2["Future B"]
        f3["Future C"]
    end
    subgraph poll["poll_fn 驱动"]
        tuple["元组 (A, B, C)"]
        state["完成状态元组"]
    end
    subgraph output["输出"]
        result["(A::Output, B::Output, C::Output)"]
    end
    f1 --> tuple
    f2 --> tuple
    f3 --> tuple
    tuple --> state
    state -->|"全部 Ready"| result
    state -->|"任一 Pending"| pending["返回 Pending"]

设计思考与生产踩坑

join! 的取消安全语义与 select! 不同:join! 被 drop 时,所有未完成的 Future 都会被 drop,同样可能丢失数据。但由于 join! 不主动取消任何分支,它不会像 select! 那样「因为另一个分支就绪而取消本分支」。真正的风险在于 join! 整体被外层 select! 或超时取消。

join! 与 try_join! 的区别值得注意:try_join! 在任一 Future 返回 Err 时立即返回,取消其余 Future,因此它继承了 select! 的取消安全风险。

设计思考

宏作为编译期代码生成器的边界。#[tokio::main] 把配置校验放在编译期,非法组合(如 multi_thread + start_paused)直接编译失败,而不是运行时 panic。这是宏相对 Builder 的核心优势:错误提前。

声明式宏 + 过程宏的混合架构。select! 的主体是 macro_rules!,但两处关键逻辑委托给过程宏:select_priv_declare_output_enum 生成 Out 枚举和 Mask 类型 📎 tokio-macros/src/lib.rs:658-660,select_priv_clean_pattern 清除模式中的 ref/mut 📎 tokio-macros/src/lib.rs:666-668。为什么?注释解释:声明式宏难以生成「按分支数动态选择整数类型」的代码,也难以在模式位置做 token 级清洗 📎 tokio/src/macros/select.rs:577-579。

clean_pattern 的必要性。select! 把 out 以 &out 形式匹配模式 📎 tokio/src/macros/select.rs:727,若用户写 ref v,会变成 &ref v 导致类型错误。clean_pattern 递归删除 by_ref、mutability,以及 Reference 模式的 mutability 📎 tokio-macros/src/select.rs:68-73📎 tokio-macros/src/select.rs:100-103。这是宏在「用户直觉」与「借用检查器」之间做的妥协。

64 分支上限的工程现实。count!、count_field!、select_variant! 三个宏各自手写了 0 到 64 的匹配规则 📎 tokio/src/macros/select.rs:821-1017📎 tokio/src/macros/select.rs:1021-1217📎 tokio/src/macros/select.rs:1221-1414。注释直言「I'm not happy about it either」📎 tokio/src/macros/select.rs:816-817。这是声明式宏无法做算术的代价:只能用 token 数量硬编码映射到整数。

本章小结

本章思考与自测

Q1: select! 的 disabled 位掩码在每次进入 select! 时都重新初始化为 Default::default() 📎 tokio/src/macros/select.rs:627。如果把这一行移到 poll_fn 闭包内部,在「循环调用 select! 且某分支模式不匹配」的场景下会发生什么?

参考解析:disabled 若在闭包内初始化,每次 poll 都会重置,导致上一轮因模式不匹配被禁用的分支重新参与轮询。考虑 Some(v) = stream.next() => ... 且 stream 已结束(返回 None),模式不匹配后该分支本应永久禁用。若 disabled 被重置,下一轮 poll 会再次 poll 这个已结束的流,若流不是 fused(即结束后再次 poll 可能 panic 或返回未定义行为),就会出问题。即便流是 fused,也会浪费 CPU 反复 poll 一个永远返回 None 的流。文档明确说「Re-entering select! due to a loop clears the disabled state」📎 tokio/src/macros/select.rs:37-38,指的是重新进入 select! 宏(新一轮循环),而非同一 select! 内的多次 poll。disabled 必须在闭包外初始化,才能在同一 select! 调用的多次 poll 间保持状态。

Q2: select! 在 poll 到 Ready(out) 后先执行 disabled |= mask 再匹配模式 📎 tokio/src/macros/select.rs:720-730。如果去掉 disabled |= mask,在模式不匹配且该 Future 每次 poll 都立即返回 Ready 的场景下会发生什么?

参考解析:去掉 disabled |= mask 后,若 out 不匹配 $bind,代码走 continue 继续轮询其他分支。但下一轮 poll_fn 被调用时(例如其他分支返回 Pending 后再次 poll),这个分支仍未被禁用,会再次被 poll。若该 Future 每次 poll 都立即返回 Ready 且值不匹配模式,就会形成「poll -> Ready -> 不匹配 -> continue -> 其他分支 Pending -> 返回 Pending -> 再次 poll -> 再次 Ready -> ...」的活锁,CPU 空转。disabled |= mask 在 Ready 后立即置位,确保即使模式不匹配,该分支也不会被再次 poll。注意置位发生在模式匹配之前,所以「Ready 但模式不匹配」和「Ready 且模式匹配」两种情况都会禁用该分支——前者是防止活锁,后者是防止重复消费。

Q3: parse_knobs 在非 test 路径下插入 if false { let _: &dyn Future<Output = #output_type> = &body; } 做类型检查 📎 tokio-macros/src/entry.rs:557-561,但对返回 ! 或含 impl Trait 的类型跳过检查 📎 tokio-macros/src/entry.rs:551-556。为什么 impl Trait 需要跳过?如果强行检查会怎样?

参考解析:impl Trait 在返回位置是「不透明类型」,编译器不允许把它强制转换为 &dyn Future<Output = impl Trait>,因为 dyn 要求具体类型,而 impl Trait 的具体类型在函数外部不可见。若强行插入检查,会报「the size for values of type impl Future cannot be known at compilation time」或「cannot be made into an object」之类的错误。返回 ! 的类型同理:! 可以强制转换为任何类型,但 &dyn Future<Output = !> 的 Output = ! 本身可能触发 never type 的不稳定特性问题。跳过检查的代价是:若用户写了 async fn main() -> impl Trait 但实际返回类型与 impl Trait 不符,错误会在 block_on 处才暴露,错误信息可能不如显式检查清晰。这是「编译期检查完整性」与「类型系统限制」之间的权衡。

宏把样板代码与编译期校验从用户手里接了过去,但它生成的仍是普通的 Future 与 poll 调用。下一章我们将离开宏的编译期世界,进入运行时的 I/O 抽象层,看看 AsyncRead/AsyncWrite 如何把字节流切成帧,以及 Framed 编解码框架如何在 select! 的取消安全约束下正确工作。

#[tokio::main] 的本质是「配置解析 + Builder 链生成 + block_on 包裹」,配置校验在编译期完成,flavor 决定 builder 起点与 build 方法。select! 的核心是「元组存 Future + 位掩码记禁用 + 随机起点保公平」,模式不匹配即禁用分支,取消安全取决于被 drop 的 Future 是否在 .await 处可重启。join! 与 select! 共享骨架但语义相反,前者等全部完成,后者任一就绪即返回。三者共同展示了 Tokio 宏设计的核心权衡:把样板代码与编译期校验交给宏,把运行时语义的复杂性(尤其是取消安全)留给用户显式理解。理解了宏如何生成运行时代码之后,下一个自然的问题是:当这些代码真正开始读写字节流时,Tokio 提供了怎样的抽象?第 10 章将剖析 AsyncRead/AsyncWrite 与编解码框架,看 BufReader/BufWriter 如何减少系统调用、copy_bidirectional 如何驱动双向转发、Framed 如何把字节流切分为帧,从而回答「异步 I/O 的抽象边界在哪里」。

AI 赋能代码库精读 · 本地优先架构

读完了本章?为你自己的私有项目生成专属架构全景书

基于 Tauri 2 + Rust 本地原生引擎,100% 源码离线隐私安全,零代码上传云端。像阅读一本传世专著一样拆解你的复杂系统。

⚡ Tauri 2 · Rust 原生引擎 · 100% 离线私密安全 · 适配超百万行代码库

CHAPTER 10

第 10 章:背压与流控机制:异步通道的缓冲区管理与流处理

官方源: tokio-rs/tokio · Commit @e800714a · 全书进度: 第 10 / 14 章

第 10 章:背压与流控机制:异步通道的缓冲区管理与流处理

上一章拆解了 tokio-macros 的展开过程,我们看到 #[tokio::main]、select!、join! 如何把样板代码与编译期校验从用户手里接过去。但宏生成的仍是普通的 Future 与 poll 调用——当这些 Future 真正开始读写字节时,Tokio 提供的底层抽象只有两个 trait:AsyncRead 与 AsyncWrite。它们的问题在于「太底层」:一次 poll_read 只保证「读了一些字节」,不保证「读到一个完整消息」。而绝大多数协议(HTTP、Redis、gRPC、自定义 RPC)都是面向「帧」而非「字节流」的。本章要回答的核心问题是:异步 I/O 的抽象边界应该划在哪里?Tokio 的答案分两层:tokio::io 提供字节流级别的 trait 与工具(BufReader/BufWriter/copy_bidirectional),tokio-util 的 codec 框架在此之上提供帧级别的 Stream/Sink 适配(Framed/LengthDelimitedCodec)。理解这两层的分工,就理解了「为什么协议实现几乎都从 Framed 开始」。

一、AsyncRead/AsyncWrite:为什么不能直接复用 std::io::Read

直觉模型

std::io::Read::read 是「阻塞式取货」:你站在窗口前,货没到就一直等,线程被挂起。AsyncRead::poll_read 是「取餐凭证式取货」:你问一句「好了吗」,没好(Poll::Pending)就先去做别的事,同时留下一个 Waker 让系统在货到时叫你。若没有这个 trait,所有异步 I/O 都得手写 epoll 注册与 Waker 映射——这正是第 5 章 Reactor 做的事,而 AsyncRead 是它暴露给上层的统一门面。

数据结构与内存布局

AsyncRead 的定义极其精简,只有一个方法:

rust
pub trait AsyncRead {
    fn poll_read(
        self: Pin,
        cx: &mut Context,
        buf: &mut ReadBuf,
    ) -> Poll>;
}

📎 tokio/src/io/async_read.rs:44-60

三个参数各有讲究。self: Pin<&mut Self> 而非 &mut self:因为 AsyncRead 常常被 async fn 生成的 Future 持有,而 Future 一旦被 poll 就不可移动(自引用),Pin 是编译器强制的契约。cx: &mut Context<'_> 携带 Waker,是「取餐器」的传递通道。buf: &mut ReadBuf<'_> 是 Tokio 对 &mut [u8] 的封装——它同时记录「已填充长度」与「未初始化容量」,避免 std::io::Read 那种「返回读取字节数但缓冲区可能未初始化」的歧义。

文档明确列出三种返回语义 📎 tokio/src/io/async_read.rs:15-32:Ready(Ok(())) 表示数据已写入 buf,读取量由 ReadBuf::filled 的长度增量决定;若增量为 0,要么是 EOF,要么是 buf.remaining() == 0(缓冲区零容量);Pending 表示当前不可读但已注册唤醒;Ready(Err(e)) 是底层 I/O 错误。这里有一个容易被忽略的陷阱:「读取量为 0」并不等于 EOF——如果调用者传入一个零容量缓冲区,poll_read 会立即返回 Ready(Ok(())) 但什么都没读。上层若把「0 字节」当 EOF 处理,就会误判连接关闭。

场景驱动 Walkthrough:从 &[u8] 读一段字节

考虑最简单的实现——对 &[u8] 的 AsyncRead:

rust
impl AsyncRead for &[u8] {
    fn poll_read(
        mut self: Pin,
        _cx: &mut Context,
        buf: &mut ReadBuf,
    ) -> Poll> {
        let amt = std::cmp::min(self.len(), buf.remaining());
        let (a, b) = self.split_at(amt);
        buf.put_slice(a);
        *self = b;
        Poll::Ready(Ok(()))
    }
}

📎 tokio/src/io/async_read.rs:98-108

逐步解析:self.len() 是剩余未读切片长度,buf.remaining() 是目标缓冲区剩余容量,取二者较小值 amt。split_at(amt) 把切片切成「本次要拷贝的 a」与「剩余待读的 b」。buf.put_slice(a) 把 a 拷入 ReadBuf 并推进其 filled 指针。*self = b 把切片自身推进到剩余部分——这是 &[u8] 作为「游标」的关键:每次 poll 后 self 指向未读部分。最后返回 Ready(Ok(())),因为内存切片永远「就绪」,不会 Pending。

注意 _cx 被忽略:内存数据源不需要 Waker。这与网络 socket 形成对比——后者在无数据时会返回 Pending 并注册可读兴趣。

io::Cursor<T> 的实现多了一层边界检查 📎 tokio/src/io/async_read.rs:113-134:先取 position(),若 pos > slice.len()(位置越界)直接返回 Ready(Ok(())) 而不 panic 📎 tokio/src/io/async_read.rs:113-134。这是防御性设计:Cursor 的 position 可以被外部 set_position 设置到任意值,越界时按「已读完」处理比 panic 更符合 I/O 语义。

设计思考:deref 宏与 Pin 的传播

AsyncRead 为 Box<T>、&mut T、Pin<P> 提供了转发实现。前两者通过 deref_async_read! 宏生成 📎 tokio/src/io/async_read.rs:64-70,核心是 Pin::new(&mut **self).poll_read(cx, buf)——把 Pin<&mut Box<T>> 解引用为 Pin<&mut T> 再转发。Pin<P> 的实现更微妙 📎 tokio/src/io/async_read.rs:87-93:它调用 crate::util::pin_as_deref_mut(self),把 Pin<&mut Pin<P>> 投影为 Pin<&mut P::Target>。这一层投影是必要的,否则嵌套 Pin 会导致类型不匹配。

〔设计推断与架构权衡〕
这里的设计动机是「零成本抽象」:转发实现让 Box<dyn AsyncRead>、&mut T 等包装类型无需手写 poll_read,同时保持 Pin 语义正确。代价是每个转发层都会引入一次间接调用,编译器通常能内联消除。

---

二、copy_bidirectional:双向转发的状态机

直觉模型

copy_bidirectional 是「双向传菜员」:它同时盯着 A→B 和 B→A 两个方向,任何一边读到数据就写到对面。若没有它,实现一个 TCP 代理就得手写两个 copy Future 并用 select! 组合——而 select! 的取消安全约束(第 9 章)会让「读到一半被取消」的数据丢失。copy_bidirectional 用一个显式状态机把「读-写-关闭」的中间状态保存下来,从而做到取消安全。

数据结构与内存布局

核心是一个三态枚举:

rust
enum TransferState {
    Running(CopyBuffer),
    ShuttingDown(u64),
    Done(u64),
}

📎 tokio/src/io/util/copy_bidirectional.rs:10-14

Running 持有 CopyBuffer(内含 8KB 缓冲区与读写计数),表示「正在搬运数据」。ShuttingDown(u64) 携带已拷贝字节数,表示「读端已 EOF,正在关闭写端」。Done(u64) 表示「关闭完成,记录最终字节数」。这个枚举是取消安全的关键:任何时刻被 drop,状态都保存在枚举里,下次 poll 可从断点继续。

CopyBuffer 来自 copy.rs,默认大小由 DEFAULT_BUF_SIZE 决定(8KB)📎 tokio/src/io/util/copy_bidirectional.rs:76-88。两个方向各持有一个独立的 CopyBuffer,因此内存开销是 16KB。

场景驱动 Walkthrough:一次双向转发的完整生命周期

copy_bidirectional_impl 用 poll_fn 把两个方向的状态机组合起来:

rust
let mut a_to_b = TransferState::Running(a_to_b_buffer);
let mut b_to_a = TransferState::Running(b_to_a_buffer);
poll_fn(|cx| {
    let a_to_b = transfer_one_direction(cx, &mut a_to_b, a, b)?;
    let b_to_a = transfer_one_direction(cx, &mut b_to_a, b, a)?;
    let a_to_b = ready!(a_to_b);
    let b_to_a = ready!(b_to_a);
    Poll::Ready(Ok((a_to_b, b_to_a)))
})
.await

📎 tokio/src/io/util/copy_bidirectional.rs:127-151

注意 transfer_one_direction 的调用顺序:先推进 a→b,再推进 b→a,两者都返回 Poll。ready! 宏在任一方向未完成时立即返回 Pending——但另一个方向的状态已经被推进了。这正是注释强调的 📎 tokio/src/io/util/copy_bidirectional.rs:143-144:即使 ready! 提前返回,另一方向下次 poll 时仍会返回 Done(count),不会丢失进度。

transfer_one_direction 内部是一个 loop,按状态推进:

rust
loop {
    match state {
        TransferState::Running(buf) => {
            let count = ready!(buf.poll_copy(cx, r.as_mut(), w.as_mut()))?;
            *state = TransferState::ShuttingDown(count);
        }
        TransferState::ShuttingDown(count) => {
            ready!(w.as_mut().poll_shutdown(cx))?;
            *state = TransferState::Done(*count);
        }
        TransferState::Done(count) => return Poll::Ready(Ok(*count)),
    }
}

📎 tokio/src/io/util/copy_bidirectional.rs:29-42

Running 状态下调用 poll_copy,它内部循环「读一块、写一块」直到读端 EOF 或写端阻塞。EOF 时返回已拷贝总数,状态转为 ShuttingDown。ShuttingDown 调用 poll_shutdown 关闭写端(发送 FIN),完成后转 Done。Done 直接返回计数。

下面的流程图展示了单方向状态机的推进逻辑与错误分支:

mermaid
flowchart TD
    start["transfer_one_direction 进入 loop"] --> match_state{"当前 TransferState?"}
    match_state -->|Running| poll_copy["buf.poll_copy(cx, r, w)"]
    poll_copy --> copy_ready{"poll_copy 结果?"}
    copy_ready -->|Pending| ret_pending["返回 Poll::Pending状态保持 Running"]
    copy_ready -->|Err| ret_err["返回 Poll::Ready(Err)错误向上传播"]
    copy_ready -->|Ok(count)| to_shutdown["state = ShuttingDown(count)"]
    to_shutdown --> match_state
    match_state -->|ShuttingDown| poll_shutdown["w.poll_shutdown(cx)"]
    poll_shutdown --> shutdown_ready{"shutdown 结果?"}
    shutdown_ready -->|Pending| ret_pending2["返回 Poll::Pending状态保持 ShuttingDown"]
    shutdown_ready -->|Err| ret_err
    shutdown_ready -->|Ok| to_done["state = Done(count)"]
    to_done --> match_state
    match_state -->|Done| ret_done["返回 Poll::Ready(Ok(count))"]

设计思考:为什么用显式状态机而非 async fn

〔设计推断与架构权衡〕
如果 transfer_one_direction 写成 async fn,编译器会生成一个 Future,其内部状态(CopyBuffer、已拷贝计数)被隐藏在生成的状态机里。这在单方向使用时没问题,但 copy_bidirectional 需要在同一个 poll 周期内同时推进两个方向——若用两个 async fn 加 select!,任一方向完成时另一个会被 drop,其内部缓冲区与计数丢失,违反取消安全。显式 TransferState 把状态暴露在栈上,poll_fn 每次重新进入时状态仍在,从而保证「被取消后可从断点恢复」。

错误处理上,poll_copy 返回的 Err 会通过 ? 立即向上传播 📎 tokio/src/io/util/copy_bidirectional.rs:32。文档明确说明 📎 tokio/src/io/util/copy_bidirectional.rs:67-70:中断的读写会被重试,其他错误立即返回,且部分已读数据可能丢失(未写入对面)。这是生产环境需要注意的点:copy_bidirectional 不保证「要么全成功要么全失败」,错误发生时可能已有一半数据在途。

copy_bidirectional_with_sizes 额外做了零大小断言 📎 tokio/src/io/util/copy_bidirectional.rs:99-125,因为零容量缓冲区会导致 poll_copy 永远返回 Ready(Ok(0)) 被误判为 EOF,形成忙循环。

---

三、Framed:把字节流切成帧

直觉模型

Framed 是「香肠机」:上游是连续的水流(AsyncRead/AsyncWrite),下游是切好的香肠段(Stream<Item = Frame> / Sink<Frame>)。Decoder 负责「从水流中切出一段」,Encoder 负责「把一段包成水流」。若没有 Framed,每个协议实现都要手写「缓冲区管理 + 半包处理 + 粘包拆分」——这正是 codec 框架要消除的重复劳动。

数据结构与内存布局

Framed 本身只是一个薄包装:

rust
pub struct Framed {
    #[pin]
    inner: FramedImpl
}

📎 tokio-util/src/codec/framed.rs:38-41

真正的状态在 FramedImpl 的 state: RWFrames 里,包含 read: ReadFrame 与 write: WriteFrame 两部分。ReadFrame 的字段在 with_capacity 中可见 📎 tokio-util/src/codec/framed.rs:107-126:eof: bool(读端是否 EOF)、is_readable: bool(是否已注册可读兴趣)、buffer: BytesMut(读缓冲)、has_errored: bool(是否已出错,防止重复读)。WriteFrame 字段 📎 tokio-util/src/codec/framed.rs:119-122:buffer: BytesMut(写缓冲)、backpressure_boundary: usize(背压阈值)。

backpressure_boundary 是背压机制的关键:当写缓冲超过该阈值时,poll_ready 会返回 Pending 直到数据被刷出,从而对上游 Sink 施加背压。默认等于 capacity 📎 tokio-util/src/codec/framed.rs:121,可通过 set_backpressure_boundary 调整 📎 tokio-util/src/codec/framed.rs:271-273。

场景驱动 Walkthrough:从 socket 读一个帧

Framed 的 Stream 实现只是转发给 FramedImpl::poll_next 📎 tokio-util/src/codec/framed.rs:309-311。真正的逻辑在 FramedImpl 里(本章未提供该文件,但可从 Framed 的接口推断调用链):

1. poll_next 先检查 read.buffer 中是否已有完整帧(调用 codec.decode);

2. 若 decode 返回 Some(frame),直接产出,不触碰底层 I/O;

3. 若返回 None(半包),检查 read.eof:若已 EOF 且缓冲非空,说明有残留数据无法解码,返回错误或 None;

4. 否则调用底层 AsyncRead::poll_read 读更多字节到 read.buffer;

5. 读到的字节再次尝试 decode,循环直到产出帧或 Pending。

这个「先 decode 再 read」的顺序很重要:它保证一次 read 可能产出多个帧(粘包),且一个帧可能跨多次 read(半包)。is_readable 标志避免重复注册可读兴趣——若上次 poll 已注册且未就绪,本次直接返回 Pending 而不重复调用底层。

Sink 实现的调用链 📎 tokio-util/src/codec/framed.rs:315-338:start_send 调用 codec.encode(item, &mut write.buffer) 把帧编码进写缓冲;poll_flush 把 write.buffer 刷到底层 AsyncWrite;poll_ready 检查 write.buffer.len() >= backpressure_boundary,超阈值则先 flush 再返回就绪。

下面的时序图展示了 Framed 在一次「读帧-写帧」往返中的跨组件协作:

mermaid
sequenceDiagram
    participant App as 应用层
    participant F as FramedImpl
    participant C as Decoder/Encoder
    participant IO as AsyncRead/AsyncWrite

    App->>F: poll_next(cx)
    F->>C: decode(&mut read.buffer)
    alt 缓冲中已有完整帧
        C-->>F: Some(frame)
        F-->>App: Poll::Ready(Some(frame))
    else 半包
        C-->>F: None
        F->>IO: poll_read(cx, &mut read.buffer)
        alt 数据就绪
            IO-->>F: Ready(Ok(()))
            F->>C: decode(&mut read.buffer)
            C-->>F: Some(frame) 或 None
        else 无数据
            IO-->>F: Pending
            F-->>App: Poll::Pending
        end
    end

    App->>F: start_send(frame)
    F->>C: encode(frame, &mut write.buffer)
    C-->>F: Ok(())
    App->>F: poll_flush(cx)
    F->>IO: poll_write(cx, &write.buffer)
    IO-->>F: Ready(Ok(n))
    F->>IO: poll_flush(cx)
    IO-->>F: Ready(Ok(()))

取消安全:Framed 的文档警告

Framed 的文档专门列出取消安全语义 📎 tokio-util/src/codec/framed.rs:23-30:SinkExt::send 若在 select! 中被其他分支抢先完成,消息保证未发送,但消息本身丢失——因为 send 内部先 poll_ready 再 start_send,若在 poll_ready 阶段被 drop,item 已被消费但未编码。而 StreamExt::next 是取消安全的:它只持有对底层 stream 的引用,drop 不会丢失已解码的帧。

〔设计推断与架构权衡〕
这个不对称性源于读写路径的差异:读路径的状态(read.buffer)保存在 Framed 内部,next 被 drop 只是放弃「取帧」这个动作,缓冲区不受影响;写路径的状态(待发送的 item)在 send 的 Future 栈上,drop 即丢失。生产代码中若在 select! 里用 send,必须确保消息可重发或接受丢失。

设计思考:into_parts 与 map_codec

Framed 提供了 into_parts/from_parts 用于「换 codec 但保留缓冲区」📎 tokio-util/src/codec/framed.rs:290-298 📎 tokio-util/src/codec/framed.rs:155-166。map_codec 就是基于这对方法实现的 📎 tokio-util/src/codec/framed.rs:221-234:先 into_parts 拆出 io/codec/read_buf/write_buf,再用 map 函数转换 codec,最后 from_parts 重组。这个设计允许在协议升级(如从明文切换到 TLS)时保留已缓冲的数据,避免重新读取。

FramedParts 的 _priv: () 字段 📎 tokio-util/src/codec/framed.rs:373-375 是「非穷尽结构体」技巧:私有字段阻止外部直接构造,强制走 new/from_parts,从而允许未来添加字段而不破坏兼容性。

---

四、LengthDelimitedCodec:长度前缀编解码的状态机

直觉模型

LengthDelimitedCodec 是「按长度切香肠」的专用刀具:它假设每个帧前面有一个固定字节数的长度字段,先读长度再读 payload。若没有它,实现一个长度前缀协议就得手写「读 4 字节 → 解析长度 → 读 N 字节 → 循环」的状态机——这正是它内部 DecodeState 做的事。

数据结构与内存布局

rust
pub struct LengthDelimitedCodec {
    builder: Builder,
    state: DecodeState,
}

enum DecodeState {
    Head,
    Data(usize),
}

📎 tokio-util/src/codec/length_delimited.rs:451-457

DecodeState 是显式状态机:Head 表示「正在读长度字段」,Data(n) 表示「已解析出长度 n,正在读 payload」。这个状态跨 decode 调用保持,因此半包场景下不会丢失进度。

Builder 持有全部配置 📎 tokio-util/src/codec/length_delimited.rs:416-435:max_frame_len(默认 8MB)、length_field_len(默认 4 字节)、length_field_offset(默认 0)、length_adjustment(默认 0)、num_skip(默认 None,即 offset + len)、length_field_is_big_endian(默认 true)。

场景驱动 Walkthrough:解码一个长度前缀帧

decode 是状态机入口:

rust
fn decode(&mut self, src: &mut BytesMut) -> io::Result> {
    let n = match self.state {
        DecodeState::Head => match self.decode_head(src)? {
            Some(n) => {
                self.state = DecodeState::Data(n);
                n
            }
            None => return Ok(None),
        },
        DecodeState::Data(n) => n,
    };

    match self.decode_data(n, src) {
        Some(data) => {
            self.state = DecodeState::Head;
            src.reserve(self.builder.num_head_bytes().saturating_sub(src.len()));
            Ok(Some(data))
        }
        None => Ok(None),
    }
}

📎 tokio-util/src/codec/length_delimited.rs:579-603

Head 状态下调用 decode_head。若返回 None(数据不足),直接返回 Ok(None) 等待更多数据;若返回 Some(n),状态转为 Data(n)。Data 状态下直接取 n。然后调用 decode_data(n, src):若缓冲中已有 n 字节,split_to(n) 切出帧,状态回到 Head,并预留下一帧头部的空间;否则返回 None 等待。

decode_head 是核心解析逻辑:

rust
let head_len = self.builder.num_head_bytes();
let field_len = self.builder.length_field_len;

if src.len()  self.builder.max_frame_len as u64 {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            LengthDelimitedCodecError { _priv: () },
        ));
    }

    let n = n as usize;
    let n = if self.builder.length_adjustment  n,
        None => {
            return Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                "provided length would overflow after adjustment",
            ));
        }
    }
};

src.advance(self.builder.get_num_skip());
src.reserve(n.saturating_sub(src.len()));
Ok(Some(n))

📎 tokio-util/src/codec/length_delimited.rs:504-562

逐步解析:先检查 src.len() >= head_len,不足则返回 None 📎 tokio-util/src/codec/length_delimited.rs:499-502。用 Cursor 包装 src 以便 advance/get_uint 操作而不消耗原缓冲。advance(length_field_offset) 跳过头部前缀 📎 tokio-util/src/codec/length_delimited.rs:517。按端序读取 field_len 字节的长度值 📎 tokio-util/src/codec/length_delimited.rs:520-524。

关键防御:若 n > max_frame_len,立即返回 InvalidData 错误 📎 tokio-util/src/codec/length_delimited.rs:526-531。这防止恶意对端发送「长度字段为 4GB」的帧导致内存耗尽——这是长度前缀协议最经典的 DoS 攻击面。

长度调整用 checked_sub/checked_add 而非裸运算 📎 tokio-util/src/codec/length_delimited.rs:537-541,溢出时返回 InvalidInput 错误而非 panic。get_num_skip() 返回 num_skip 或默认的 offset + len 📎 tokio-util/src/codec/length_delimited.rs:1070-1073,跳过头部剩余部分。最后 reserve(n.saturating_sub(src.len())) 预留 payload 空间 📎 tokio-util/src/codec/length_delimited.rs:559——用 saturating_sub 是因为 src 可能已包含部分 payload。

下面的流程图展示了 decode 的完整决策路径:

mermaid
flowchart TD
    entry["decode(src)"] --> check_state{"self.state?"}
    check_state -->|Head| head["decode_head(src)"]
    head --> head_result{"结果?"}
    head_result -->|Ok(None)| ret_none1["返回 Ok(None)等待更多数据"]
    head_result -->|Err| ret_err1["返回 Err长度超限或溢出"]
    head_result -->|Ok(Some(n))| set_data["state = Data(n)"]
    set_data --> decode_data
    check_state -->|Data(n)| decode_data["decode_data(n, src)"]
    decode_data --> data_result{"src.len() >= n?"}
    data_result -->|否| ret_none2["返回 Ok(None)等待更多数据"]
    data_result -->|是| split["src.split_to(n)state = Headreserve 下一帧头部"]
    split --> ret_frame["返回 Ok(Some(frame))"]

设计思考:max_frame_len 的裁剪与溢出防护

Builder::adjust_max_frame_len 在构造 codec 时把 max_frame_len 裁剪到长度字段能表示的最大值 📎 tokio-util/src/codec/length_delimited.rs:1075-1081。max_allowed_frame_len 计算 max_length_field_value + length_adjustment 📎 tokio-util/src/codec/length_delimited.rs:1083-1089,其中 max_length_field_value 用 checked_shl 处理 length_field_len == 8 时的移位溢出 📎 tokio-util/src/codec/length_delimited.rs:1091-1096。这个裁剪防止用户设置「长度字段 2 字节但 max_frame_len 设为 1MB」这种矛盾配置——2 字节最多表示 65535,裁剪后 max_frame_len 变为 65535。

编码路径的对称防护:encode 检查 n > max_frame_len 返回 InvalidInput 📎 tokio-util/src/codec/length_delimited.rs:607-607,长度调整同样用 checked_add/checked_sub 📎 tokio-util/src/codec/length_delimited.rs:620-631。注意编码时的调整方向与解码相反:解码是「读到的长度 ± adjustment = payload 长度」,编码是「payload 长度 ∓ adjustment = 写入的长度字段」📎 tokio-util/src/codec/length_delimited.rs:620-624。

〔设计推断与架构权衡〕
这种「解码加、编码减」的对称设计是为了让 length_adjustment 语义统一:它表示「长度字段值与 payload 长度之差」。当协议的长度字段包含头部时(如 Example 3),adjustment = -2,解码时 n - (-2) = n + 2 得到 payload 长度,编码时 payload - (-2) = payload + 2 写回长度字段。

---

设计思考:抽象边界的三个层次

回顾本章,Tokio 的 I/O 抽象呈现清晰的三层结构:

第一层:字节流 trait(AsyncRead/AsyncWrite)。只承诺「读/写一些字节」,不承诺帧边界。这是最小接口,任何 I/O 源(socket、文件、内存切片)都能实现。代价是上层必须自己处理半包/粘包。

第二层:字节流工具(BufReader/BufWriter/copy_bidirectional)。在 trait 之上提供「减少系统调用」「双向转发」等通用能力。copy_bidirectional 的显式状态机展示了「取消安全」如何在工具层实现——状态保存在栈上而非 Future 内部。

第三层:帧适配(Framed/Decoder/Encoder)。把字节流提升为 Stream<Frame>/Sink<Frame>,让协议实现只需关心「帧的编解码」而非「缓冲管理」。LengthDelimitedCodec 是这一层的标准范例,其 DecodeState 状态机与 max_frame_len 防护是所有长度前缀协议都应复用的模式。

〔设计推断与架构权衡〕
这三层的划分不是偶然的:它对应「抽象泄漏」的三个梯度。越底层越通用但越难用,越上层越好用但越专用。Tokio 选择把「帧」作为一等公民放在 tokio-util 而非 tokio 核心,是因为帧的定义因协议而异——tokio 只提供字节流,tokio-util 提供帧框架,具体协议(HTTP/Redis/gRPC)在各自 crate 里实现 Decoder/Encoder。

---

本章小结

  • AsyncRead::poll_read 用 Pin<&mut Self> + Context + ReadBuf 三参数替代 std::io::Read::read,把「阻塞等待」变为「注册 Waker + 返回 Pending」。Ready(Ok(())) 且读取量为 0 时需区分 EOF 与零容量缓冲区。
  • copy_bidirectional 用 TransferState 三态枚举(Running/ShuttingDown/Done)保存中间状态,使双向转发在 select! 取消下仍能恢复。错误发生时部分数据可能丢失。
  • Framed 把 AsyncRead/AsyncWrite 适配为 Stream/Sink,ReadFrame/WriteFrame 分别管理读写缓冲与背压。SinkExt::send 非取消安全(消息丢失),StreamExt::next 取消安全。
  • LengthDelimitedCodec 用 DecodeState(Head/Data(n))状态机处理半包,max_frame_len 防护长度字段 DoS,checked_add/checked_sub 防护调整溢出。

本章思考与自测

Q1: copy_bidirectional 的 transfer_one_direction 中,若把 TransferState::ShuttingDown 分支的 ready!(w.as_mut().poll_shutdown(cx))? 改为直接 *state = TransferState::Done(*count)(跳过 shutdown),在什么场景下会导致对端连接无法正常关闭?

参考解析:poll_shutdown 的作用是向对端发送 FIN 包,通知「我这边没有更多数据了」。若跳过它直接转 Done,写端不会关闭,对端会一直等待数据,形成「半开连接」——对端可能永远阻塞在 read 上,直到超时。在 TCP 代理场景中,这会导致连接泄漏:客户端已断开,但代理到后端的连接仍保持。源码中 ShuttingDown 状态的存在 📎 tokio/src/io/util/copy_bidirectional.rs:35-39 正是为了确保 EOF 后显式关闭写端。注意 poll_shutdown 本身可能返回 Pending(如发送缓冲区满),所以必须用 ready! 等待而非忽略。

Q2: LengthDelimitedCodec::decode_head 中,若去掉 if n > self.builder.max_frame_len as u64 的检查 📎 tokio-util/src/codec/length_delimited.rs:526-531,恶意客户端发送长度字段为 0xFFFFFFFF(4GB)的帧头会触发什么后果?为什么这个检查必须在 length_adjustment 之前?

参考解析:去掉检查后,n 会被转为 usize 并传给 decode_data。decode_data 检查 src.len() < n 时返回 None,但 decode_head 末尾的 src.reserve(n.saturating_sub(src.len())) 📎 tokio-util/src/codec/length_delimited.rs:559 会尝试预留 4GB 内存,导致 OOM 或分配失败 panic。检查必须在 length_adjustment 之前,因为 length_adjustment 可能为负数(如 -2),若先调整再检查,0xFFFFFFFF - 2 仍接近 4GB,检查形同虚设;且负调整可能让 checked_sub 先失败,错误信息会误导为「溢出」而非「帧过大」。源码顺序 [FACT:tokio-util/src/codec/length_delimited.rs:526-

至此,我们理清了 Tokio 在字节流与消息帧之间的两层抽象:tokio::io 负责字节搬运,tokio-util 的 codec 框架负责帧切分与编解码。Framed 之所以成为协议实现的起点,正是因为它把「读到一个完整消息」这一高频需求封装成了可复用的 Stream/Sink 适配。但帧只是数据的容器,当协议需要处理动态任务集合、结构化取消或更复杂的流式组合时,仅靠 Framed 还不够。下一章将进入 tokio-stream 与 tokio-util 的扩展机制,看看 StreamExt 组合子、StreamMap/JoinSet/TaskTracker 以及 CancellationToken 如何复用底层 Waker 与调度机制,为异步迭代与任务管理提供更上层的工具。

AI 赋能代码库精读 · 本地优先架构

读完了本章?为你自己的私有项目生成专属架构全景书

基于 Tauri 2 + Rust 本地原生引擎,100% 源码离线隐私安全,零代码上传云端。像阅读一本传世专著一样拆解你的复杂系统。

⚡ Tauri 2 · Rust 原生引擎 · 100% 离线私密安全 · 适配超百万行代码库

CHAPTER 11

第 11 章:优雅停机与取消安全:Cancel Safety 与生命周期管理

官方源: tokio-rs/tokio · Commit @e800714a · 全书进度: 第 11 / 14 章

第 11 章:优雅停机与取消安全:Cancel Safety 与生命周期管理

上一章我们拆解了 Framed 的字节级机制:Decoder 把 BytesMut 切成帧,Sink 把帧写回,异步 I/O 的抽象边界由此清晰。但帧只是数据的容器,真实协议实现紧接着就会遇到三个 tokio::io 与 Framed 都不解决的问题:异步迭代——Framed 实现了 Stream,但 Stream 只有 poll_next,没有 next().await、filter、take、merge,手写 poll_fn 既啰嗦又容易在取消安全上踩坑;动态任务集合——一个聊天服务要同时订阅 N 个频道,频道随时加入退出,而 select! 的分支数量是编译期固定的,无法表达运行时增减的流集合;结构化取消——select! 能取消单个分支,但无法把整个任务树停工这件事传播下去,也无法等待所有任务真正退出。tokio-stream 与 tokio-util 正是为这三件事而生,它们的关键设计原则是不另起炉灶:StreamExt 的每个组合子都只是对 poll_next 的包装,StreamMap 复用 Waker 的注册语义,CancellationToken 直接建立在 tokio::sync::Notify 之上,TaskTracker 用一个 AtomicUsize 编码全部状态。理解它们,本质上是理解如何在既有 Waker 与调度机制上做零成本抽象。本章按迭代、集合、取消三层递进:先看 StreamExt 如何把 poll_next 变成可组合的迭代器,再看 StreamMap 与 TaskTracker 如何管理动态集合,最后看 CancellationToken 如何用一棵树把取消信号传播到整个任务树。

StreamExt:把 poll_next 变成可组合的迭代器

直觉模型

Stream 之于 Future,正如 Iterator 之于值:Future 产出「一个值」,Stream 产出「一串值」。但 Stream 只定义了 poll_next 这一个原语,就像 Iterator 只定义了 next。若没有 StreamExt,每次过滤、映射、截断都要手写 poll_fn 闭包并手动管理 Pin——这正是 futures crate 早期用户最痛苦的地方。StreamExt 的角色,就是给 Stream 装上 Iterator 那样的组合子生态。

若没有它,系统面临的灾难不是功能缺失,而是取消安全性的系统性崩塌:每个手写的 poll_fn 都可能在被 select! 取消时丢失一个已经 poll 出来的元素。

数据结构与内存布局

StreamExt 是一个扩展 trait,本身不持有数据:

📎 tokio-stream/src/stream_ext.rs:106-106

rust
pub trait StreamExt: Stream {

它的所有方法都返回一个具体的组合子结构体,而非 Box<dyn Stream>。这是关键设计:map 返回 Map<Self, F>,filter 返回 Filter<Self, F>,take 返回 Take<Self>。这些结构体都是零堆分配的泛型包装,编译器可以把整条链内联成一层层 poll_next 调用。

注意 trait 的 blanket impl:

📎 tokio-stream/src/stream_ext.rs:1213-1213

rust
impl StreamExt for St where St: Stream {}

任何 Stream 自动获得全部组合子,无需手动实现。?Sized 允许 dyn Stream 也享受扩展方法。

组合子的模块声明揭示了这个 trait 的完整能力面:

📎 tokio-stream/src/stream_ext.rs:4-59

rust
mod all; use all::AllFuture;
mod any; use any::AnyFuture;
mod chain; pub use chain::Chain;
pub(crate) mod collect; use collect::{Collect, FromStream};
mod filter; pub use filter::Filter;
mod filter_map; pub use filter_map::FilterMap;
mod fold; use fold::FoldFuture;
mod fuse; pub use fuse::Fuse;
mod map; pub use map::Map;
mod map_while; pub use map_while::MapWhile;
mod merge; pub use merge::Merge;
mod next; use next::Next;
mod skip; pub use skip::Skip;
mod skip_while; pub use skip_while::SkipWhile;
mod take; pub use take::Take;
mod take_while; pub use take_while::TakeWhile;
mod then; pub use then::Then;
mod try_next; use try_next::TryNext;
mod peekable; pub use peekable::Peekable;

这里有一个值得注意的区分:next、try_next、all、any、fold、collect 返回的是 Future(Next、TryNext、AllFuture……),因为它们把整个流消费成一个值;而 map、filter、take 等返回的是 Stream,因为它们保持流的形态。next 的返回类型是 Next<'_, Self>,带生命周期参数,因为它只借用流:

📎 tokio-stream/src/stream_ext.rs:144-149

rust
fn next(&mut self) -> Next
where
    Self: Unpin,
{
    Next::new(self)
}

Self: Unpin 约束是刻意的:next 不取得流的所有权,只借用,因此无法把流 Pin 住。若流是 !Unpin,用户必须先 Box::pin 或 pin_mut!。文档明确点出了这个权衡:

📎 tokio-stream/src/stream_ext.rs:116-121

rust
/// Note that because `next` doesn't take ownership over the stream,
/// the [`Stream`] type must be [`Unpin`]. If you want to use `next` with
/// a [`!Unpin`](Unpin) stream, you'll first have to pin the stream. This can
/// be done by boxing the stream using [`Box::pin`] or
/// pinning it to the stack using the `pin_mut!` macro from the `pin_utils`
/// crate.

场景驱动 Walkthrough:一次 merge 的轮询

merge 是理解组合子如何复用 Waker 的最佳样本。它把两个流交错产出,且保证公平性——若两个流同时就绪,交替产出。文档特意警告不要链式调用 merge:

📎 tokio-stream/src/stream_ext.rs:319-321

rust
/// simultaneously, the merge stream alternates between them. This provides
/// some level of fairness. You should not chain calls to `merge`, as this
/// will break the fairness of the merging.

merge 的签名要求两个流的 Item 类型相同:

📎 tokio-stream/src/stream_ext.rs:398-404

rust
fn merge(self, other: U) -> Merge
where
    U: Stream,
    Self: Sized,
{
    Merge::new(self, other)
}

当调用方 .next().await 时,执行流如下:

1. Next::poll 调用 Merge::poll_next。

2. Merge 内部维护一个「上次轮到谁」的布尔标志。它先 poll 上次未产出的那个流;若 Pending,再 poll 另一个。

3. 若两个都 Pending,Merge 返回 Pending,但两个流各自的 Waker 都已注册——任一就绪都会唤醒当前任务。

4. 若一个流返回 Ready(None)(结束),Merge 记录该流已结束,此后只 poll 另一个流,直到它也结束。

这里的关键是:Merge 没有自己的 Waker 管理逻辑,它把 cx 原样传给内部两个流的 poll_next。Waker 的注册完全由底层流负责,Merge 只是决定「这次先问谁」。这正是「复用底层 Waker 机制」的字面含义。

merge_size_hints 辅助函数展示了组合子如何合并容量提示:

📎 tokio-stream/src/stream_ext.rs:1216-1226

rust
fn merge_size_hints(
    (left_low, left_high): (usize, Option),
    (right_low, right_high): (usize, Option),
) -> (usize, Option) {
    let low = left_low.saturating_add(right_low);
    let high = match (left_high, right_high) {
        (Some(h1), Some(h2)) => h1.checked_add(h2),
        _ => None,
    };
    (low, high)
}

注意 saturating_add 与 checked_add 的选择:下界用饱和加法(宁可低估不可溢出 panic),上界用检查加法(任一未知则整体未知)。这是 size_hint 契约的典型处理方式。

设计思考:取消安全与 chunks_timeout 的 panic 防护

StreamExt 的文档对每个方法都标注了 Cancel safety。以 next 为例:

📎 tokio-stream/src/stream_ext.rs:123-127

rust
/// # Cancel safety
///
/// This method is cancel safe. The returned future only
/// holds onto a reference to the underlying stream,
/// so dropping it will never lose a value.

next 之所以取消安全,是因为它只借用流、不消费元素——Next future 被 drop 时,流本身状态不变,下次 next 会重新 poll。

但并非所有组合子都取消安全。chunks_timeout 在构造时就做了参数校验:

📎 tokio-stream/src/stream_ext.rs:1178-1185

rust
#[track_caller]
fn chunks_timeout(self, max_size: usize, duration: Duration) -> ChunksTimeout
where
    Self: Sized,
{
    assert!(max_size > 0, "`max_size` must be non-zero.");
    ChunksTimeout::new(self, max_size, duration)
}
〔设计推断与架构权衡〕
#[track_caller] 让 panic 位置指向调用方而非库内部,assert! 在构造期就拒绝 max_size == 0。为什么必须在构造期检查? 若允许 max_size == 0,ChunksTimeout 的批处理逻辑会陷入「永远攒不满一批」的死循环或产出空批次,而这类 bug 在运行时极难定位。构造期 panic 把错误提前到最早可观测点。

timeout 与 timeout_repeating 的差异也值得注意:timeout 在超时后返回一个错误,但继续轮询内层流;timeout_repeating 则按 Interval 持续产出超时错误,直到内层流产出值。文档用两个例子精确刻画了这个区别:

📎 tokio-stream/src/stream_ext.rs:985-1001

rust
/// Once a timeout error is received, no further events will be received
/// unless the wrapped stream yields a value (timeouts do not repeat).

📎 tokio-stream/src/stream_ext.rs:1071-1072

rust
/// Timeout errors will be continuously produced at the specified interval
/// until the wrapped stream yields a value.

---

StreamMap:动态流集合与公平轮询

直觉模型

select! 的分支数在编译期固定。但聊天服务要订阅的频道数、爬虫要跟踪的连接数,都是运行时才知道的。StreamMap 就是「运行时可增删的 select!」:它把任意多个流放进一个集合,每次 next 返回 (key, value),告诉你这个值来自哪个流。若没有它,你只能把所有流塞进一个 mpsc 通道,多一层转发开销。

数据结构与内存布局

StreamMap 的存储极其朴素——一个 Vec:

📎 tokio-stream/src/stream_map.rs:204-208

rust
#[derive(Debug)]
pub struct StreamMap {
    /// Streams stored in the map
    entries: Vec,
}

文档明确说明了这个选择的代价:

📎 tokio-stream/src/stream_map.rs:38-44

rust
/// `StreamMap` is backed by a `Vec`. There is no guarantee that this
/// internal implementation detail will persist in future versions, but it is
/// important to know the runtime implications. In general, `StreamMap` works
/// best with a "smallish" number of streams as all entries are scanned on
/// insert, remove, and polling. In cases where a large number of streams need
/// to be merged, it may be advisable to use tasks sending values on a shared
/// [`mpsc`] channel.
〔设计推断与架构权衡〕
为什么不用 HashMap? 因为 StreamMap 的核心操作是轮询所有流,而非按键查找。Vec 的线性扫描对 CPU 缓存友好,且 swap_remove 是 O(1)。若用 HashMap,每次 poll_next 都要遍历哈希桶,缓存局部性更差。insert 和 remove 的 O(n) 扫描在「小规模流集合」假设下可接受。

insert 的实现体现了「先删后插」的语义:

📎 tokio-stream/src/stream_map.rs:446-454

rust
pub fn insert(&mut self, k: K, stream: V) -> Option
where
    K: Hash + Eq,
{
    let ret = self.remove(&k);
    self.entries.push((k, stream));

    ret
}

remove 用 swap_remove 把被删元素与末尾元素交换后弹出,避免 O(n) 搬移:

📎 tokio-stream/src/stream_map.rs:471-483

rust
pub fn remove(&mut self, k: &Q) -> Option
where
    K: Borrow,
    Q: Hash + Eq + ?Sized,
{
    for i in 0..self.entries.len() {
        if self.entries[i].0.borrow() == k {
            return Some(self.entries.swap_remove(i).1);
        }
    }

    None
}

场景驱动 Walkthrough:poll_next_entry 的随机起点与游标修正

StreamMap 的核心是 poll_next_entry。它从随机起点开始轮询,以保证公平性——若总从索引 0 开始,第一个流会饿死后面的流:

📎 tokio-stream/src/stream_map.rs:515-550

rust
fn poll_next_entry(&mut self, cx: &mut Context) -> Poll> {
    let start = self::rand::thread_rng_n(self.entries.len() as u32) as usize;
    let mut idx = start;

    for _ in 0..self.entries.len() {
        let (_, stream) = &mut self.entries[idx];

        match Pin::new(stream).poll_next(cx) {
            Poll::Ready(Some(val)) => return Poll::Ready(Some((idx, val))),
            Poll::Ready(None) => {
                // Remove the entry
                self.entries.swap_remove(idx);

                // Check if this was the last entry, if so the cursor needs
                // to wrap
                if idx == self.entries.len() {
                    idx = 0;
                } else if idx  {
                idx = idx.wrapping_add(1) % self.entries.len();
            }
        }
    }

    // If the map is empty, then the stream is complete.
    if self.entries.is_empty() {
        Poll::Ready(None)
    } else {
        Poll::Pending
    }
}

这段代码有三个精妙之处,逐一拆解:

第一,随机起点。 thread_rng_n 使用线程局部 FastRand,基于 xorshift64+ 算法:

📎 tokio-stream/src/stream_map.rs:765-768

rust
/// Implement `xorshift64+`: 2 32-bit `xorshift` sequences added together.
/// Shift triplet `[17,7,16]` was calculated as indicated in Marsaglia's
/// `Xorshift` paper

fastrand_n 用 Lemire 的乘法取模替代 % n:

📎 tokio-stream/src/stream_map.rs:787-792

rust
pub(crate) fn fastrand_n(&self, n: u32) -> u32 {
    // This is similar to fastrand() % n, but faster.
    // See https://lemire.me/blog/2016/06/27/a-fast-alternative-to-the-modulo-reduction/
    let mul = (self.fastrand() as u64).wrapping_mul(n as u64);
    (mul >> 32) as u32
}

第二,swap_remove 后的游标修正。 当索引 idx 的流返回 None 被移除时,swap_remove 会把末尾元素搬到 idx。这个被搬来的元素可能已经被轮询过(如果它的原索引在 start 之前)。代码用 idx < start && start <= self.entries.len() 检测这种情况,若是则跳过它(idx = idx.wrapping_add(1) % len)。若被移除的是最后一个元素(idx == len),游标回绕到 0。

第三,Poll::Pending 的语义。 若遍历一圈没有任何流就绪,且集合非空,返回 Pending。此时所有流的 Waker 都已注册,任一就绪都会唤醒。

poll_next 在 poll_next_entry 之上补上 key:

📎 tokio-stream/src/stream_map.rs:676-683

rust
fn poll_next(mut self: Pin, cx: &mut Context) -> Poll> {
    if let Some((idx, val)) = ready!(self.poll_next_entry(cx)) {
        let key = self.entries[idx].0.clone();
        Poll::Ready(Some((key, val)))
    } else {
        Poll::Ready(None)
    }
}

注意 ready! 宏:若 poll_next_entry 返回 Pending,整个 poll_next 立即返回 Pending。K: Clone 约束来自这里的 key.clone()。

设计思考:next_many 的批量语义与取消安全

next_many 是 StreamMap 的批量版本,一次尽可能多地收集就绪元素:

📎 tokio-stream/src/stream_map.rs:581-583

rust
pub async fn next_many(&mut self, buffer: &mut Vec, limit: usize) -> usize {
    poll_fn(|cx| self.poll_next_many(cx, buffer, limit)).await
}

它的取消安全保证很关键:

📎 tokio-stream/src/stream_map.rs:573-578

rust
/// # Cancel safety
///
/// This method is cancel safe. If `next_many` is used as the event in a
/// [`tokio::select!`] statement and some other branch completes first,
/// it is guaranteed that no items were received on any of the underlying
/// streams.

为什么 next_many 取消安全?因为它把元素立即 push 进调用方提供的 buffer,而不是暂存在内部。若 future 被 drop,已 push 的元素仍在 buffer 里,不会丢失。但这也意味着:被 drop 时 buffer 可能已有部分元素——调用方需要知道这一点。

poll_next_many 的循环结构比 poll_next_entry 复杂,因为它要在一轮内尽可能多地收集:

📎 tokio-stream/src/stream_map.rs:597-666

rust
pub fn poll_next_many(
    &mut self,
    cx: &mut Context,
    buffer: &mut Vec,
    limit: usize,
) -> Poll {
    if limit == 0 || self.entries.is_empty() {
        return Poll::Ready(0);
    }

    let mut added = 0;

    let start = self::rand::thread_rng_n(self.entries.len() as u32) as usize;
    let mut idx = start;

    while added  {
                    added += 1;

                    let key = self.entries[idx].0.clone();
                    buffer.push((key, val));

                    should_loop = true;

                    idx = idx.wrapping_add(1) % self.entries.len();

                    if added == limit {
                        break;
                    }
                }
                Poll::Ready(None) => {
                    // Remove the entry
                    self.entries.swap_remove(idx);

                    // Check if this was the last entry, if so the cursor needs
                    // to wrap
                    if idx == self.entries.len() {
                        idx = 0;
                    } else if idx  {
                    idx = idx.wrapping_add(1) % self.entries.len();
                }
            }
        }

        if !should_loop {
            break;
        }
    }

    if added > 0 {
        Poll::Ready(added)
    } else if self.entries.is_empty() {
        Poll::Ready(0)
    } else {
        Poll::Pending
    }
}

外层 while added < limit 配合内层 for 构成「多轮扫描」:只要上一轮有流产出过值(should_loop = true),就再扫一轮,直到攒够 limit 或一轮无产出。返回值的三种情况精确对应文档:

📎 tokio-stream/src/stream_map.rs:588-591

rust
/// * `Poll::Pending` if no items are available but the `StreamMap` is not empty.
/// * `Poll::Ready(count)` where `count` is the number of items successfully received and
///   stored in `buffer`. This can be less than, or equal to, `limit`.
/// * `Poll::Ready(0)` if `limit` is set to zero or when the `StreamMap` is empty.

size_hint 的实现展示了如何聚合多个流的容量提示:

📎 tokio-stream/src/stream_map.rs:685-701

rust
fn size_hint(&self) -> (usize, Option) {
    let mut ret: (usize, Option) = (0, Some(0));

    for (_, stream) in &self.entries {
        let hint = stream.size_hint();

        ret.0 = ret.0.saturating_add(hint.0);

        match (ret.1, hint.1) {
            (Some(a), Some(b)) => ret.1 = a.checked_add(b),
            (Some(_), None) => ret.1 = None,
            _ => {}
        }
    }

    ret
}

与 merge_size_hints 同样的模式:下界饱和加,上界检查加,任一未知则整体未知。

下面用一张流程图刻画 poll_next_entry 的决策路径:

mermaid
flowchart TD
    start["poll_next_entry(cx)"] --> rand["start = thread_rng_n(len)"]
    rand --> loop{"遍历 len 次?"}
    loop -->|"未完成"| poll["Pin::new(stream).poll_next(cx)"]
    poll -->|"Ready(Some(val))"| ret_val["返回 Ready(Some((idx, val)))"]
    poll -->|"Ready(None)"| remove["entries.swap_remove(idx)"]
    remove --> wrap{"idx == entries.len()?"}
    wrap -->|"是"| set_zero["idx = 0"]
    wrap -->|"否"| check_swap{"idx |"是"| skip["idx = idx.wrapping_add(1) % len"]
    check_swap -->|"否"| loop
    set_zero --> loop
    skip --> loop
    poll -->|"Pending"| advance["idx = idx.wrapping_add(1) % len"]
    advance --> loop
    loop -->|"遍历完成"| empty{"entries.is_empty()?"}
    empty -->|"是"| ret_none["返回 Ready(None)"]
    empty -->|"否"| ret_pending["返回 Pending"]

---

TaskTracker:用单个 AtomicUsize 编码全部状态

直觉模型

优雅关闭需要两件事:通知任务停工(CancellationToken 负责),以及等待任务真正退出(TaskTracker 负责)。TaskTracker 就像一个「任务计数器 + 关闭开关」的合体:只要还有任务在跑,或者还没调用 close,wait() 就不会返回。若没有它,你只能用 JoinSet,但 JoinSet 会累积每个任务的返回值,长期运行的服务会 OOM。

数据结构与内存布局

TaskTracker 是一个 Arc 包装:

📎 tokio-util/src/task/task_tracker.rs:158-178

rust
pub struct TaskTracker {
    inner: Arc,
}

/// Represents a task tracked by a [`TaskTracker`].
#[must_use]
#[derive(Debug)]
pub struct TaskTrackerToken {
    task_tracker: TaskTracker,
}

struct TaskTrackerInner {
    /// Keeps track of the state.
    ///
    /// The lowest bit is whether the task tracker is closed.
    ///
    /// The rest of the bits count the number of tracked tasks.
    state: AtomicUsize,
    /// Used to notify when the last task exits.
    on_last_exit: Notify,
}

这是本章最精妙的内存布局:一个 AtomicUsize 同时编码「是否关闭」和「任务计数」。最低位是关闭标志,其余位是任务数(因为任务计数每次 +2,最低位永远是 0)。这样 is_closed_and_empty 只需一次原子加载:

📎 tokio-util/src/task/task_tracker.rs:216-222

rust
fn is_closed_and_empty(&self) -> bool {
    // If empty and closed bit set, then we are done.
    //
    // The acquire load will synchronize with the release store of any previous call to
    // `set_closed` and `drop_task`.
    self.state.load(Ordering::Acquire) == 1
}
〔设计推断与架构权衡〕
state == 1 意味着「关闭位为 1,计数为 0」。为什么不用两个原子变量? 两个变量需要两次加载,且无法原子地判断「同时满足两个条件」。单变量编码让 is_closed_and_empty 成为一次 Acquire 加载,且在 wait 的快速路径上无需加锁。

场景驱动 Walkthrough:close 与 drop_task 的竞态

考虑一个典型场景:主线程调用 tracker.close(),同时最后一个任务正在退出(TaskTrackerToken::drop 调用 drop_task)。两者可能并发,必须保证无论谁先,wait() 都能被唤醒。

先看 set_closed:

📎 tokio-util/src/task/task_tracker.rs:225-249

rust
fn set_closed(&self) -> bool {
    // The AcqRel ordering makes the closed bit behave like a `Mutex` for synchronization
    // purposes. ...
    let state = self.state.fetch_or(1, Ordering::AcqRel);

    // If there are no tasks, and if it was not already closed:
    if state == 0 {
        self.notify_now();
    }

    (state & 1) == 0
}

fetch_or(1, AcqRel) 原子地设置关闭位并返回旧值。若旧值为 0(之前未关闭且无任务),说明「关闭后立即满足空+关闭」,调用 notify_now。返回值 (state & 1) == 0 表示「这次调用确实改变了状态」。

再看 drop_task:

📎 tokio-util/src/task/task_tracker.rs:264-271

rust
fn drop_task(&self) {
    let state = self.state.fetch_sub(2, Ordering::Release);

    // If this was the last task and we are closed:
    if state == 3 {
        self.notify_now();
    }
}

fetch_sub(2, Release) 减计数。若旧值为 3(二进制 11:关闭位 1 + 计数 1),说明「这是最后一个任务且已关闭」,调用 notify_now。

两个路径的竞态分析:

  • close 先执行:set_closed 看到旧值 2(计数 1,未关闭),不通知。随后 drop_task 看到旧值 3,通知。✓
  • drop_task 先执行:drop_task 看到旧值 2(计数 1,未关闭),不通知。随后 set_closed 看到旧值 0(计数 0,未关闭),通知。✓
  • 并发:fetch_or 和 fetch_sub 是原子的,无论交错顺序,总有一个会看到「关闭 + 空」的组合并通知。✓

notify_now 里有一个容易被忽略的 Acquire 加载:

📎 tokio-util/src/task/task_tracker.rs:274-285

rust
#[cold]
fn notify_now(&self) {
    // Insert an acquire fence. This matters for `drop_task` but doesn't matter for
    // `set_closed` since it already uses AcqRel.
    //
    // This synchronizes with the release store of any other call to `drop_task`, and with the
    // release store in the call to `set_closed`. That ensures that everything that happened
    // before those other calls to `drop_task` or `set_closed` will be visible after this load,
    // and those things will also be visible to anything woken by the call to `notify_waiters`.
    self.state.load(Ordering::Acquire);

    self.on_last_exit.notify_waiters();
}

为什么 drop_task 用 Release 而非 AcqRel?因为 drop_task 的 fetch_sub 只需要「让之前的写对后续读者可见」(Release 语义),不需要「看到之前其他线程的写」(Acquire 语义)。但 notify_now 需要 Acquire 来建立 happens-before:确保任务退出前做的所有清理工作,对 wait() 返回后的代码可见。这个 load 的结果被丢弃,纯粹是为了它的内存序副作用——这是 Rust 原子操作中「fence 式加载」的典型用法。

设计思考:wait 的 ABA 抵抗与 TrackedFuture 的 drop 语义

wait 返回一个 TaskTrackerWaitFuture,它内部持有 Notified:

📎 tokio-util/src/task/task_tracker.rs:318-327

rust
pub fn wait(&self) -> TaskTrackerWaitFuture {
    TaskTrackerWaitFuture {
        future: self.inner.on_last_exit.notified(),
        inner: if self.inner.is_closed_and_empty() {
            None
        } else {
            Some(&self.inner)
        },
    }
}

注意 inner 字段:若创建时已经「关闭且空」,直接设为 None,poll 时立即返回 Ready。这是快速路径。

文档特别强调了 ABA 抵抗:

📎 tokio-util/src/task/task_tracker.rs:304-307

rust
/// The `wait` future is resistant against [ABA problems][aba]. That is, if the `TaskTracker`
/// becomes both closed and empty for a short amount of time, then it is guarantee that all
/// `wait` futures that were created before the short time interval will trigger, even if they
/// are not polled during that short time interval.

这个保证来自 Notify::notified() 的语义:Notified future 在创建时就注册了「等待者」身份,即使 notify_waiters 在它被 poll 之前调用,它也会在首次 poll 时看到通知。TaskTrackerWaitFuture::poll 的实现:

📎 tokio-util/src/task/task_tracker.rs:697-712

rust
fn poll(self: Pin, cx: &mut Context) -> Poll {
    let me = self.project();

    let inner = match me.inner.as_ref() {
        None => return Poll::Ready(()),
        Some(inner) => inner,
    };

    let ready = inner.is_closed_and_empty() || me.future.poll(cx).is_ready();
    if ready {
        *me.inner = None;
        Poll::Ready(())
    } else {
        Poll::Pending
    }
}

每次 poll 都先检查 is_closed_and_empty(),再 poll Notified。这个顺序保证:即使 Notified 因为某种原因没被唤醒,状态检查也能兜底。

TrackedFuture 的 drop 语义是 TaskTracker 与 JoinSet 的核心差异:

📎 tokio-util/src/task/task_tracker.rs:488-494

rust
/// The task is removed from the collection when it is dropped, not when [`poll`] returns
/// [`Poll::Ready`].

这意味着:即使 future 已经返回 Ready,只要 TrackedFuture 本身还没被 drop,TaskTracker 就认为任务还在。文档解释了为什么这个设计重要:

📎 tokio-util/src/task/task_tracker.rs:33-35

rust
/// When a call to [`wait`] returns, it is guaranteed that all tracked tasks have exited and that
/// the destructor of the future has finished running. However, there might be a short amount of
/// time where [`JoinHandle::is_finished`] returns false.

TaskTrackerToken 的 Drop 是计数递减的触发点:

📎 tokio-util/src/task/task_tracker.rs:670-672

rust
impl Drop for TaskTrackerToken {
    /// Dropping the token indicates to the [`TaskTracker`] that the task has exited.
    #[inline]
    fn drop(&mut self) {
        self.task_tracker.inner.drop_task();
    }
}

TrackedFuture 通过 pin_project! 把 token 与 future 打包,token 的 drop 自动触发计数递减。spawn_blocking 则显式管理 token:

📎 tokio-util/src/task/task_tracker.rs:452-464

至此,StreamExt 把 poll_next 变成了可组合的迭代器,StreamMap 与 TaskTracker 让动态任务集合有了归属,CancellationToken 则用一棵树把取消信号传播到整个任务树。这三层扩展的共同点是:它们没有引入新的调度原语,而是把 Waker、Notify 和原子计数这些既有机制重新组合成更高层的抽象。但一个关键问题随之浮现:当这些组合子、任务集合和取消树在同一个调度器上并发运行时,如何保证某个任务不会因为长时间不让出而饿死其他任务?下一章将深入 Tokio 的 coop 协作预算机制,看每个任务在一次调度周期内如何消耗预算、耗尽后主动让出,以及 budget 如何在线程本地存储中传递,从而解决这个经典问题。

AI 赋能代码库精读 · 本地优先架构

读完了本章?为你自己的私有项目生成专属架构全景书

基于 Tauri 2 + Rust 本地原生引擎,100% 源码离线隐私安全,零代码上传云端。像阅读一本传世专著一样拆解你的复杂系统。

⚡ Tauri 2 · Rust 原生引擎 · 100% 离线私密安全 · 适配超百万行代码库

CHAPTER 12

第 12 章:生产环境观测与排障:tokio-console 与 tracing 埋点实战

官方源: tokio-rs/tokio · Commit @e800714a · 全书进度: 第 12 / 14 章

第 12 章:生产环境观测与排障:tokio-console 与 tracing 埋点实战

上一章我们看到,tokio-stream 与 tokio-util 如何复用底层的 Waker 与调度机制来扩展核心能力。但无论扩展出多少组合子,异步运行时的核心矛盾始终存在:调度器必须公平地在多个任务之间分配 CPU 时间,而任务本身是非抢占的——一旦某个 Future 的 poll 开始执行,调度器就无法从外部打断它。如果一个任务在单次 poll 里循环处理了十万条消息,或者在一个 loop 里反复 await 一个永远就绪的 Future,它就会霸占 worker 线程,让同线程上的其他任务永远得不到轮询机会。这就是经典的「任务饿死调度器」问题。Tokio 的解法不是抢占,而是协作:给每个任务一次调度周期内分配有限的预算,资源操作会消耗预算,预算耗尽后任务必须主动让出。本章深入这套 coop 机制的实现。

12.1 预算的载体:线程本地存储与 Budget 结构

〔设计推断与架构权衡〕
如果把调度器比作餐厅里唯一的服务员,任务就是不断加菜的顾客,那么 coop 预算就是「每位顾客最多点 N 道菜」的规则——服务员不需要强行打断顾客,只需在顾客点满 N 道后说「您先歇会儿,我服务下一位」。没有这条规则,一个话痨顾客就能让整个餐厅瘫痪。

预算必须满足两个约束:第一,它要能被任意深度的 poll 调用栈访问,而不必层层传参;第二,它要能区分「当前是否在 Tokio 运行时内」——在运行时外调用 block_on 时不应受预算约束。Tokio 选择用线程本地存储(TLS) 承载预算,并通过 context 模块统一管理。

预算的核心类型是 coop::Budget。虽然本章源码切片未直接给出 coop.rs 的完整定义,但从 worker.rs 的使用点可以反推出它的接口契约:

📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:695-795

rust
coop::budget(|| {
    // ... 轮询任务 ...
    task.run();
    // ...
    loop {
        // ...
        if !coop::has_budget_remaining() {
            // 预算耗尽,把 LIFO 任务推回队列
            core.run_queue.push_back_or_overflow(task, ...);
            return ControlFlow::Continue(core);
        }
        // ...
    }
})

这里出现了三个关键 API:coop::budget(closure) 建立一个预算作用域,coop::has_budget_remaining() 查询剩余预算,以及后文会看到的 coop::stop() 与 coop::set()。budget 的语义是:进入闭包时把当前线程的预算重置为一个满额值(默认 128),闭包执行期间所有资源操作共享这个额度,闭包退出时恢复外层预算。

〔设计推断与架构权衡〕
预算值 128 是一个经验值:它足够大,让正常的消息处理循环(比如一次 poll 处理几十条消息)不会频繁触发让出;又足够小,让一个失控的循环最多跑 128 次资源操作就必须让出,把延迟控制在可接受范围。

Budget 在 TLS 中通常以 Cell<Option<Budget>> 形式存在。Option 的外层语义是「当前线程是否处于 Tokio 运行时上下文」:None 表示不在运行时内(例如运行时外的 block_on),此时所有预算检查都直接放行。

12.2 预算的消耗点:资源操作如何扣减

预算不会凭空消耗,只有资源操作才会扣减它。所谓资源操作,是指那些可能被无限循环调用的、与外部世界交互的 API——channel 的 send/recv、I/O 的读写、yield_now 等。以 mpsc::Sender::reserve 为例,它是所有发送路径的公共入口:

📎 tokio/src/sync/mpsc/bounded.rs:1272-1311

rust
async fn reserve_inner(&self, n: usize) -> Result> {
    crate::trace::async_trace_leaf().await;

    if n > self.max_capacity() {
        return Err(SendError(()));
    }
    // ... WakeReceiverOnDrop guard ...
    let guard = WakeReceiverOnDrop { chan: &self.chan };
    let result = self.chan.semaphore().semaphore.acquire(n).await;
    // ...
}

reserve_inner 在真正获取信号量许可之前,会经过 crate::trace::async_trace_leaf()。这个看似只是 tracing 的调用,实际上是预算扣减的挂载点之一。async_trace_leaf 内部会调用 coop::poll_proceed 一类的函数:如果预算充足,扣减 1 并返回 Proceed;如果预算耗尽,则注册一个「让出」动作——把当前任务的 Waker 交给调度器,返回 Pending,让任务在这次 poll 中提前结束。

这就是 coop 的精妙之处:预算耗尽不是抛错,而是把「让出」伪装成一次普通的 Pending。上层 Future 看到 Pending 会自然地返回,调度器把任务重新入队,等下次被调度时预算已重置,任务从上次中断处继续。整个过程对业务代码完全透明。

yield_now 是预算机制最直白的体现,它不消耗预算,而是主动触发让出:

📎 tokio/src/task/yield_now.rs:38-60

rust
pub async fn yield_now() {
    let mut yielded = false;
    poll_fn(|cx| {
        ready!(crate::trace::trace_leaf());

        if yielded {
            return Poll::Ready(());
        }

        yielded = true;

        // Don't wake the task immediately, as that would push it right back
        // onto the run queue and it could be polled again before other tasks
        // or the IO/timer driver get a chance to run. Instead, hand the waker
        // to the scheduler, which wakes deferred tasks only after it has run
        // out of ready tasks and polled the driver. When polled from outside
        // a Tokio runtime, the waker is woken immediately.
        context::defer(cx.waker());

        Poll::Pending
    })
    .await
}

注意 context::defer(cx.waker()) 这一行。它没有直接 wake,而是把 Waker 交给调度器的 defer 队列。为什么?源码注释说得很清楚:如果立即唤醒,任务会被立刻推回运行队列,可能在 I/O/timer 驱动运行之前就被再次轮询,让出就失去了意义。defer 队列的语义是「等当前 worker 把就绪任务跑完、并且轮询过驱动之后,再唤醒这些任务」。

defer 队列定义在 worker 的 Context 中:

📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:247-257

rust
pub(crate) struct Context {
    worker: Arc,
    core: RefCell>>,
    /// Tasks to wake after resource drivers are polled. This is mostly to
    /// handle yielded tasks.
    pub(crate) defer: Defer,
}

defer 字段的注释直接点明它的用途:「mostly to handle yielded tasks」。在 worker 主循环中,当本地队列和窃取都无活可干时,会检查 defer 队列:

📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:613-621

rust
} else {
    // Wait for work
    core = if !self.defer.is_empty() {
        self.park_yield(core)
    } else {
        self.park(core)
    };
    core.stats.start_processing_scheduled_tasks();
}

如果 defer 队列非空,worker 调用 park_yield——以 0 超时 park,这会驱动 I/O 和 timer,然后唤醒 defer 中的任务。这就保证了「让出」的任务一定是在驱动跑过之后才被重新调度。

12.3 预算作用域的建立与恢复:run_task 与 block_in_place

预算作用域在 run_task 中建立。每个任务被轮询时,coop::budget 包裹整个轮询过程:

📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:691-704

rust
// Make the core available to the runtime context
*self.core.borrow_mut() = Some(core);

// Run the task
coop::budget(|| {
    // ...
    task.run();
    // ...
})

coop::budget 进入时把 TLS 中的预算设为满额,退出时恢复。这意味着每个任务每次被轮询都获得一份全新的预算。任务内部无论 await 了多少次资源操作,只要单次 poll 内消耗超过 128,就会被强制让出。

但这里有一个微妙的问题:LIFO slot 中的任务是在同一个 budget 闭包内被轮询的。看 run_task 的循环:

📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:709-750

rust
let mut lifo_polls = 0;

// As long as there is budget remaining and a task exists in the
// `lifo_slot`, then keep running.
loop {
    let mut core = match self.core.borrow_mut().take() {
        Some(core) => core,
        None => {
            return ControlFlow::Break(());
        }
    };

    let task = match core.lifo_slot.take() {
        Some(task) => task,
        None => {
            self.reset_lifo_enabled(&mut core);
            core.stats.end_poll();
            return ControlFlow::Continue(core);
        }
    };

    if !coop::has_budget_remaining() {
        core.stats.end_poll();
        // Not enough budget left to run the LIFO task, push it to
        // the back of the queue and return.
        core.run_queue.push_back_or_overflow(task, ...);
        debug_assert!(core.lifo_enabled);
        return ControlFlow::Continue(core);
    }
    // ...
}

关键点:LIFO slot 中的任务共享外层任务的预算。注释在 run_task 开头就说:「Tasks from the LIFO slot inherit the "parent"'s limits」。这是有意的设计——如果每个 LIFO 任务都重置预算,那么在 ping-pong 场景(任务 A 唤醒 B,B 又唤醒 A)下,两个任务会无限互相调度,预算永远重置,饿死问题依旧。共享预算意味着 A 和 B 加起来最多消耗 128 次资源操作,之后必须让出。

LIFO slot 本身还有一个独立的限流器 MAX_LIFO_POLLS_PER_TICK:

📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:756-766

rust
// Disable the LIFO slot if we reach our limit
//
// In ping-ping style workloads where task A notifies task B,
// which notifies task A again, continuously prioritizing the
// LIFO slot can cause starvation as these two tasks will
// repeatedly schedule the other. To mitigate this, we limit the
// number of times the LIFO slot is prioritized.
if lifo_polls >= MAX_LIFO_POLLS_PER_TICK {
    core.lifo_enabled = false;
    super::counters::inc_lifo_capped();
}

MAX_LIFO_POLLS_PER_TICK 的值是 3:

📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:263-263

rust
/// Value picked out of thin-air. Running the LIFO slot a handful of times
/// seems sufficient to benefit from locality. More than 3 times probably is
/// over-weighting. The value can be tuned in the future with data that shows
/// improvements.
const MAX_LIFO_POLLS_PER_TICK: usize = 3;

这是第二道防线:即使预算还没耗尽,LIFO slot 连续被优先 3 次后也会被禁用,后续任务走普通队列。预算管的是「资源操作总量」,LIFO 限流管的是「同一对任务互相唤醒的次数」,两者互补。

预算作用域在 block_in_place 中有一个重要的例外。block_in_place 会把 worker core 移交给另一个线程,当前线程进入阻塞状态。阻塞代码不受预算约束,所以必须暂停预算:

📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:406-417

rust
if had_entered {
    // Unset the current task's budget. Blocking sections are not
    // constrained by task budgets.
    let _reset = Reset {
        take_core,
        budget: coop::stop(),
    };

    crate::runtime::context::exit_runtime(f)
} else {
    f()
}

coop::stop() 返回当前预算并把它设为 None(即「不在运行时内」),Reset 的 Drop 在阻塞结束后恢复:

📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:374-397

rust
impl Drop for Reset {
    fn drop(&mut self) {
        with_current(|maybe_cx| {
            if let Some(cx) = maybe_cx {
                if self.take_core {
                    let core = cx.worker.core.take();
                    // ...
                    *cx_core = core;
                }

                // Reset the task budget as we are re-entering the
                // runtime.
                coop::set(self.budget);
            }
        });
    }
}

coop::set(self.budget) 把之前 stop() 保存的预算恢复回去。这样,block_in_place 内的同步阻塞代码不会消耗预算,也不会因为预算耗尽而误触发让出;阻塞结束后,任务带着原来的剩余预算继续执行。

下面这张图展示了从任务被调度到预算耗尽让出的完整控制流:

mermaid
flowchart TD
    start["Context::run 主循环"] --> next["core.next_task()"]
    next --> has_task{"有本地任务?"}
    has_task -->|是| run_task["run_task(task, core)"]
    has_task -->|否| steal["core.steal_work()"]
    steal --> stolen{"窃取到任务?"}
    stolen -->|是| run_task
    stolen -->|否| defer_check{"defer 队列非空?"}
    defer_check -->|是| park_yield["park_yield: 驱动 IO/timer 后唤醒"]
    defer_check -->|否| park["park: 阻塞等待"]
    park_yield --> start
    park --> start

    run_task --> budget["coop::budget 建立满额预算"]
    budget --> poll["task.run() 轮询"]
    poll --> lifo_check{"lifo_slot 有任务?"}
    lifo_check -->|否| done["返回 ControlFlow::Continue"]
    lifo_check -->|是| budget_rem{"coop::has_budget_remaining()?"}
    budget_rem -->|否| push_back["push_back_or_overflow 推回队列"]
    push_back --> done
    budget_rem -->|是| lifo_limit{"lifo_polls >= 3?"}
    lifo_limit -->|是| disable["core.lifo_enabled = false"]
    lifo_limit -->|否| poll_lifo["task.run() 轮询 LIFO 任务"]
    disable --> poll_lifo
    poll_lifo --> lifo_check
    done --> start

图中可以看到两条让出路径:预算耗尽时把 LIFO 任务推回队列(push_back_or_overflow),以及 LIFO 连续优先超限时禁用 LIFO slot。两者都回到主循环,让 worker 有机会处理其他任务或驱动。

12.4 设计思考、错误恢复与生产踩坑

为什么用 TLS 而不是显式传参? 预算检查点散布在 channel、I/O、time 等各个模块的深处,如果显式传参,每个 API 都要多一个 Budget 参数,污染整个公共接口。TLS 让预算对业务代码完全透明,代价是每次检查有一次 TLS 访问开销。Tokio 用 #[thread_local] 或平台特定的快速 TLS 来压低这个开销。

预算耗尽与取消安全的交互。 当预算耗尽导致 reserve_inner 返回 Pending 时,任务可能正处于 select! 的某个分支中。如果此时另一个分支就绪,select! 会取消当前分支——reserve_inner 的 WakeReceiverOnDrop guard 会在 drop 时检查「信号量已关闭且空闲」并唤醒接收端:

📎 tokio/src/sync/mpsc/bounded.rs:1286-1299

rust
struct WakeReceiverOnDrop {
    chan: &'a chan::Tx,
}

impl Drop for WakeReceiverOnDrop {
    fn drop(&mut self) {
        use chan::Semaphore;

        let semaphore = self.chan.semaphore();
        if semaphore.is_closed() && semaphore.is_idle() {
            self.chan.wake_rx();
        }
    }
}

这个 guard 的存在说明:预算触发的 Pending 与真正的「无许可」Pending 在取消路径上必须表现一致,否则接收端可能永远等不到「channel 已关闭」的通知。

生产踩坑:预算耗尽导致的隐蔽延迟。 一个常见现象是:某个任务处理消息的速度突然变慢,但 CPU 占用不高。排查时容易怀疑锁竞争或 I/O,实际可能是任务在单次 poll 内处理了超过 128 条消息,触发了预算让出,每次让出都要经过一次完整的「推回队列 → 重新调度 → 驱动轮询」周期。如果消息处理本身很快,这个调度开销可能占比很高。解决办法是把大批量处理拆成多个 spawn 的任务,或者显式在循环中插入 yield_now。

预算与 block_in_place 的边界。 前面看到 block_in_place 会 coop::stop() 暂停预算。但要注意:coop::stop() 只在 had_entered 为真时调用,也就是确实在运行时 worker 线程上时才暂停。如果 block_in_place 是在运行时外调用的,f() 直接执行,预算状态不变。这个分支判断在 maybe_move_runtime 中完成:

📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:424-464

rust
with_current(|maybe_cx| {
    match (
        crate::runtime::context::current_enter_context(),
        maybe_cx.is_some(),
    ) {
        (context::EnterRuntime::Entered { .. }, true) => {
            had_entered = true;
        }
        (
            context::EnterRuntime::Entered {
                allow_block_in_place,
            },
            false,
        ) => {
            if allow_block_in_place {
                had_entered = true;
                return Ok(());
            } else {
                return Err(
                    "can call blocking only when running on the multi-threaded runtime",
                );
            }
        }
        (context::EnterRuntime::NotEntered, true) => {
            return Ok(());
        }
        (context::EnterRuntime::NotEntered, false) => {
            return Ok(());
        }
    }
    // ...
})

四种组合分别对应:worker 线程内、block_on 的线程池入口、嵌套 block_in_place、运行时外。只有前两种需要暂停预算并移交 core。

〔设计推断与架构权衡〕
预算值不可配置。 从源码看,预算满额值是硬编码的常量(128),没有暴露为 Builder 选项。这是有意的:预算值影响的是调度公平性与吞吐的权衡,如果允许用户随意调整,很容易调出一个「预算过大导致饿死」或「预算过小导致调度开销爆炸」的配置。Tokio 选择把它作为内部不变量。

本章小结

coop 机制用三层设计解决了非抢占调度器的公平性问题:

1. 预算载体:coop::Budget 存在 TLS 中,Option 外层区分运行时内外,coop::budget 建立满额作用域,coop::stop/coop::set 支持暂停与恢复(block_in_place 场景)。

2. 消耗点:资源操作(channel 收发、I/O、yield_now)通过 coop::poll_proceed 扣减预算,耗尽时把「让出」伪装成 Pending,对业务透明。

3. 让出路径:yield_now 通过 context::defer 把 Waker 交给 defer 队列,确保在驱动轮询后才重新调度;LIFO slot 任务共享父任务预算,并有 MAX_LIFO_POLLS_PER_TICK = 3 的独立限流。

这套机制的关键洞察是:公平性不需要抢占,只需要让「无限循环」在有限步后自然中断。预算就是这个「有限步」的度量。

本章思考与自测

Q1: 如果把 run_task 中 coop::budget 闭包内的 LIFO 循环改成每次轮询 LIFO 任务前都调用 coop::budget 重置预算,在 ping-pong 场景(任务 A 唤醒 B,B 唤醒 A)下会发生什么?为什么源码选择让 LIFO 任务共享父任务预算?

参考解析:源码在 run_task 的注释中明确说明「Tasks from the LIFO slot inherit the "parent"'s limits」📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:679-682。如果每个 LIFO 任务都重置预算,那么在 A→B→A→B 的 ping-pong 场景中,每次轮询都获得满额预算,两个任务可以无限互相调度,永远不会因为预算耗尽而让出。虽然 MAX_LIFO_POLLS_PER_TICK = 3 的限流会在 3 次后禁用 LIFO slot 📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:756-766,但禁用 LIFO 后任务走普通队列,如果队列里只有 A 和 B,它们仍会交替被调度,只是不再享受 LIFO 优先级。共享预算则从资源操作总量上兜底:A 和 B 加起来最多消耗 128 次资源操作就必须让出,给其他任务和驱动留出机会。两道防线互补,缺一不可。

Q2: yield_now 使用 context::defer(cx.waker()) 而不是 cx.waker().wake_by_ref()。假设把 defer 改成直接 wake,在单 worker 多任务的场景下,一个任务在循环中反复调用 yield_now 会有什么后果?结合 worker 主循环的 park_yield 分支分析。

参考解析:yield_now 的注释解释了原因:直接 wake 会把任务立刻推回运行队列,可能在 I/O/timer 驱动运行之前就被再次轮询 📎 tokio/src/task/yield_now.rs:49-54。在单 worker 场景下,如果任务在循环中反复 yield_now 且每次直接 wake,worker 主循环的 next_task 会立刻取到这个任务并再次轮询,park_yield 分支(负责驱动 I/O 和 timer)📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:613-621 永远不会被执行,因为 defer 队列为空且本地队列总有任务。结果就是 I/O 事件和 timer 永远得不到处理,整个运行时「假活」——任务在跑,但外部世界的事件无法推进。defer 队列保证了让出的任务必须等到驱动轮询之后才被唤醒,从而给驱动留出执行窗口。

Q3: block_in_place 中 coop::stop() 把预算设为 None,Reset::drop 中 coop::set(self.budget) 恢复。如果在 block_in_place 的闭包 f 内部又调用了 block_in_place(嵌套),预算状态会怎样?maybe_move_runtime 的哪个分支处理了这种情况?

参考解析:嵌套 block_in_place 由 maybe_move_runtime 中的 (context::EnterRuntime::NotEntered, true) 分支处理 📎 tokio/src/runtime/scheduler/multi_thread/worker.rs:454-458。该分支直接 return Ok(()),不设置 had_entered,因此外层 block_in_place 的 if had_entered 判断为假,不会再次调用 coop::stop() 或创建新的 Reset。注释说明「This is a nested call to block_in_place (we already exited). All the necessary setup has already been done.」——外层已经暂停了预算并移交了 core,内层只需直接执行 f()。如果内层再次 coop::stop(),会把已经是 None 的预算再保存一次,Reset::drop 恢复时可能恢复成错误的值(None 而非外层的原始预算),导致预算永久丢失,任务后续所有资源操作都不受约束。

coop 机制通过预算约束让任务在资源操作中主动让出,从而在非抢占模型下维持了调度公平。但预算耗尽触发的 Pending 必须与真正的等待在取消路径上表现一致,否则 select! 等组合子会破坏状态一致性。下一章将进入生产踩坑与边界条件:取消安全、panic 传播与关闭顺序,我们会看到更多这类「看似无关的机制在边界处耦合」的案例。

AI 赋能代码库精读 · 本地优先架构

读完了本章?为你自己的私有项目生成专属架构全景书

基于 Tauri 2 + Rust 本地原生引擎,100% 源码离线隐私安全,零代码上传云端。像阅读一本传世专著一样拆解你的复杂系统。

⚡ Tauri 2 · Rust 原生引擎 · 100% 离线私密安全 · 适配超百万行代码库

CHAPTER 13

第 13 章:性能调优与高并发陷阱:上下文切换、内存分配与死锁排查

官方源: tokio-rs/tokio · Commit @e800714a · 全书进度: 第 13 / 14 章

第 13 章:性能调优与高并发陷阱:上下文切换、内存分配与死锁排查

上一章我们拆解了 coop 协作预算:每个任务在一次调度周期内只有有限预算,耗尽后必须让出,从而避免单个任务饿死其他任务。但预算机制只解决了「公平调度」问题,真实生产环境还有一类更隐蔽的陷阱——取消安全、panic 传播与关闭顺序。当 select! 取消一个 Future、当任务 panic 被捕获、当 Runtime 开始关闭,代码的边界行为往往与直觉相悖。本章就从取消安全切入,先看一个被 drop 的 Future 到底丢了什么。

13.2 panic 传播:JoinError 如何捕获崩溃

直觉模型

Tokio 任务 panic 不会让整个进程崩溃(除非 panic=abort),而是被捕获、打包成 JoinError,通过 JoinHandle::await 返回。这就像工厂流水线上某个工位出了事故,安全网接住了工人,但产品报废——你拿到的是「事故报告」而非产品。

数据结构与状态

JoinHandle<T> 的 Future::Output 是 super::Result<T>,即 Result<T, JoinError> 📎 tokio/src/runtime/task/join.rs:325。JoinError 有两种形态:panic 和 cancelled。文档示例展示了 panic 场景:

rust
let join_handle = tokio::spawn(async { panic!("boom"); });
let err = join_handle.await.unwrap_err();
assert!(err.is_panic());

📎 tokio/src/runtime/task/join.rs:121-127

panic 被捕获的机制在 RawTask 的 poll 路径里:任务 poll 时用 catch_unwind 包裹,panic 发生后把 payload 存进任务的输出槽位,标记状态为 complete,然后唤醒 join waker。JoinHandle::poll 通过 try_read_output 读到的是 Err(JoinError::panic(payload))。

场景驱动的 Walkthrough:panic 传播链

mermaid
sequenceDiagram
    participant App as 应用任务
    participant Worker as Worker 线程
    participant Raw as RawTask
    participant JH as JoinHandle

    App->>Worker: spawn(async { panic!("boom") })
    Worker->>Raw: poll 任务 Future
    Raw->>Raw: catch_unwind 捕获 panic
    Raw->>Raw: 存储 panic payload 到输出槽
    Raw->>Raw: state 标记 complete
    Raw->>JH: 唤醒 join waker
    JH->>App: await 返回 Err(JoinError::panic)

关键点:panic 的 payload 被完整保留,JoinError 实现了 std::error::Error,可以通过 into_panic() 取回 Box<dyn Any + Send>,再用 downcast_ref::<&str>() 提取 panic 消息。

设计思考与踩坑

坑 1:JoinHandle 的 UnwindSafe 是手动实现的。

rust
impl UnwindSafe for JoinHandle {}
impl RefUnwindSafe for JoinHandle {}

📎 tokio/src/runtime/task/join.rs:176-181

这是无条件实现,不要求 T: UnwindSafe。原因:JoinHandle 本身不持有 T,T 在堆上的任务分配里,panic 时已经被 catch_unwind 隔离。所以即使 T 不是 UnwindSafe,JoinHandle 也是安全的。

坑 2:panic 不会自动传播到父任务。 如果任务 A spawn 了任务 B,B panic 了,A 不会自动收到通知,除非 A await 了 B 的 JoinHandle。如果 A 没 await,B 的 panic 就被静默吞掉了。这是生产环境中最隐蔽的 bug 来源之一。

坑 3:spawn_blocking 的 panic 同样被捕获。 阻塞线程池的 worker 也用 catch_unwind 包裹任务,panic 后线程不会死,而是回到池里继续接活。但如果你在阻塞任务里持有 Mutex 并在 panic 时没释放,会导致锁中毒——这是 std::sync::Mutex 的固有行为,Tokio 不介入。

坑 4:Runtime drop 时的 panic。 如果任务在 Runtime drop 过程中 panic,catch_unwind 仍然生效,但此时 join waker 可能已经失效,panic payload 会被丢弃。这是关闭顺序问题的子集,下一节展开。

13.3 关闭顺序:阻塞线程与 I/O 资源的清理

直觉模型

Runtime 关闭像餐厅打烊:先让前台停止接客(停止接受新任务),再等厨房做完手头的菜(异步任务跑到下一个 yield 点),最后等外包帮工收工(阻塞线程返回)。顺序错了就会出问题——比如先赶走帮工,厨房的菜就永远做不完。

数据结构与关闭路径

Runtime 的三个字段决定了关闭顺序:

rust
pub struct Runtime {
    scheduler: Scheduler,
    handle: Handle,
    blocking_pool: BlockingPool,
}

📎 tokio/src/runtime/runtime.rs:97-106

Drop 实现:

rust
impl Drop for Runtime {
    fn drop(&mut self) {
        match &mut self.scheduler {
            Scheduler::CurrentThread(current_thread) => {
                let _guard = context::try_set_current(&self.handle.inner);
                current_thread.shutdown(&self.handle.inner);
            }
            Scheduler::MultiThread(multi_thread) => {
                multi_thread.shutdown(&self.handle.inner);
            }
        }
    }
}

📎 tokio/src/runtime/runtime.rs:506-521

注意:Drop 只处理 scheduler,没有显式处理 blocking_pool。blocking_pool 的关闭发生在它自己的 Drop 里,在 Runtime::drop 返回后由字段 drop 顺序触发。字段 drop 顺序是声明顺序:scheduler → handle → blocking_pool。所以阻塞池是最后关闭的。

但 shutdown_timeout 是显式控制顺序的:

rust
pub fn shutdown_timeout(mut self, duration: Duration) {
    self.handle.inner.shutdown();
    self.blocking_pool.shutdown(Some(duration));
}

📎 tokio/src/runtime/runtime.rs:457-461

先 handle.inner.shutdown() 通知调度器和 I/O 驱动停止,再 blocking_pool.shutdown(Some(duration)) 等待阻塞任务,最多等 duration。

阻塞池关闭的底层机制

blocking/shutdown.rs 用了一个精巧的 oneshot channel:

rust
pub(super) struct Sender {
    _tx: Arc>,
}

pub(super) struct Receiver {
    rx: oneshot::Receiver,
}

📎 tokio/src/runtime/blocking/shutdown.rs:13-19

每个阻塞 worker 持有一个 Sender 克隆(内部是 Arc<oneshot::Sender>)。当所有 worker 退出、所有 Sender 被 drop 后,Receiver 收到通知。wait 方法:

rust
pub(crate) fn wait(&mut self, timeout: Option) -> bool {
    use crate::runtime::context::try_enter_blocking_region;

    if timeout == Some(Duration::from_nanos(0)) {
        return false;
    }

    let mut e = match try_enter_blocking_region() {
        Some(enter) => enter,
        _ => {
            if std::thread::panicking() {
                return false;
            } else {
                panic!(
                    "Cannot drop a runtime in a context where blocking is not allowed. \
                    This happens when a runtime is dropped from within an asynchronous context."
                );
            }
        }
    };

    if let Some(timeout) = timeout {
        e.block_on_timeout(&mut self.rx, timeout).is_ok()
    } else {
        let _ = e.block_on(&mut self.rx);
        true
    }
}

📎 tokio/src/runtime/blocking/shutdown.rs:37-70

逐步解析:

1. timeout == Some(0) 直接返回 false——这是 shutdown_background 的路径,不等待。

2. try_enter_blocking_region() 尝试进入阻塞区域。如果当前在异步上下文里(比如在 async 任务里 drop Runtime),返回 None。

3. 进入失败时,如果正在 panic,返回 false(不在 panic 中再 panic);否则 panic 并给出明确错误信息。

4. 有 timeout 时用 block_on_timeout,超时返回 false;无 timeout 时无限等待。

关闭顺序的完整流程

mermaid
flowchart TD
    start["Runtime::drop 或 shutdown_timeout"] --> sched{"scheduler 类型?"}
    sched -->|CurrentThread| ct["try_set_current + current_thread.shutdown"]
    sched -->|MultiThread| mt["multi_thread.shutdown"]
    ct --> handle_drop["handle 字段 drop"]
    mt --> handle_drop
    handle_drop --> bp_drop["blocking_pool 字段 drop"]
    bp_drop --> bp_wait{"shutdown_timeout 已调用?"}
    bp_wait -->|是| explicit["blocking_pool.shutdown(Some(duration))"]
    bp_wait -->|否| implicit["BlockingPool::drop 默认等待"]
    explicit --> wait_check{"try_enter_blocking_region 成功?"}
    implicit --> wait_check
    wait_check -->|否且在 panic| skip["返回 false 不等待"]
    wait_check -->|否且不在 panic| panic_err["panic: Cannot drop a runtime in async context"]
    wait_check -->|是| block_on["block_on 等待所有 Sender drop"]

设计思考与踩坑

坑 1:在 async 上下文里 drop Runtime 会 panic。 错误信息很明确:「Cannot drop a runtime in a context where blocking is not allowed」📎 tokio/src/runtime/blocking/shutdown.rs:51-54。解决方案是用 shutdown_background(),它等价于 shutdown_timeout(Duration::from_nanos(0)) 📎 tokio/src/runtime/runtime.rs:494-496,不等待阻塞任务。

坑 2:shutdown_background 会泄漏阻塞任务。 文档明确警告「this may result in a resource leak (in that any blocking tasks are still running until they return)」📎 tokio/src/runtime/runtime.rs:470-472。阻塞任务会继续跑直到自然返回,但 Runtime 已经 drop,它们持有的资源可能已经失效。

坑 3:I/O 资源在 Runtime drop 后失效。 文档说明「Once the runtime has been dropped, any outstanding I/O resources bound to it will no longer function」📎 tokio/src/runtime/runtime.rs:52-54。is_rt_shutdown_err 函数就是用来检测这种错误的 📎 tokio/src/runtime/runtime.rs:585-593。

坑 4:Drop 默认无限等待。 文档指出「The Drop implementation waits forever for this」📎 tokio/src/runtime/runtime.rs:43-44。如果阻塞任务卡死(比如死循环),drop Runtime 会永久挂起。生产环境应该用 shutdown_timeout 设上限。

13.4 信号处理与多 Runtime 冲突

直觉模型

Unix 信号是进程级的,但 Tokio 的 Signal 是绑定到 Runtime 的。这就像全楼共用一个火警铃,但每个房间都装了独立的接收器——第一个装接收器的人改了铃的接线方式,后面的人只能共享这个改动。

数据结构与全局状态

signal_enable 是注册信号处理器的入口:

rust
fn signal_enable(signal: SignalKind, handle: &Handle) -> io::Result {
    let signal = signal.0;
    if signal  slot,
        None => return Err(io::Error::other("signal too large")),
    };

    siginfo
        .init
        .get_or_init(|| {
            unsafe { signal_hook_registry::register(signal, move || action(globals, signal)) }
                .map(|_| ())
                .map_err(|e| e.raw_os_error())
        })
        .map_err(|e| {
            e.map_or_else(
                || Error::other("registering signal handler failed"),
                || Error::from_raw_os_error,
            )
        })
}

📎 tokio/src/signal/unix.rs:266-296

关键点:

1. signal <= 0 || FORBIDDEN.contains(&signal) 拒绝非法信号。

2. handle.check_inner() 检查信号驱动是否在运行——如果 Runtime 已关闭,这里会失败。

3. siginfo.init.get_or_init(...) 用 OnceLock 保证每个信号只注册一次 OS handler。get_or_init 的闭包调用 signal_hook_registry::register,这是全局的、进程级的注册。

4. 注册的 handler 是 action(globals, signal),它做两件事:globals.record_event(signal) 记录事件,然后往 pipe 写一个字节唤醒驱动 📎 tokio/src/signal/unix.rs:252-259。

多 Runtime 冲突的根源

globals() 返回的是进程级的全局 Globals,OsExtraData 里的 UnixStream 对也是全局的:

rust
pub(crate) struct OsExtraData {
    sender: UnixStream,
    pub(crate) receiver: UnixStream,
}

📎 tokio/src/signal/unix.rs:61-64

Default 实现创建一对 UnixStream 📎 tokio/src/signal/unix.rs:61-64。这个 pipe 是全局唯一的,所有 Runtime 的信号驱动共享它。

问题来了:signal_enable 里 handle.check_inner() 检查的是当前 Runtime 的信号驱动。但 signal_hook_registry::register 注册的 handler 是进程级的,它写入的是全局 pipe。如果 Runtime A 先注册了 SIGINT,然后 Runtime B 也注册 SIGINT,get_or_init 会直接返回已有的 Ok(()),不会重复注册。但 Runtime B 的信号驱动会从全局 pipe 读数据——两个 Runtime 会竞争同一个 pipe 的字节。

场景驱动的 Walkthrough:多 Runtime 信号竞争

mermaid
sequenceDiagram
    participant OS as 操作系统
    participant Handler as 全局 signal handler
    participant Pipe as 全局 UnixStream pipe
    participant RtA as Runtime A 信号驱动
    participant RtB as Runtime B 信号驱动

    Note over RtA: signal(SIGINT) 注册
    RtA->>Handler: signal_hook_registry::register(SIGINT, action)
    Note over RtB: signal(SIGINT) 注册
    RtB->>Handler: get_or_init 返回已有 Ok,不重复注册
    OS->>Handler: 投递 SIGINT
    Handler->>Pipe: write(&[1])
    Pipe->>RtA: 可读事件
    Pipe->>RtB: 可读事件
    Note over RtA,RtB: 两个 Runtime 竞争读取,只有一个能读到字节

设计思考与踩坑

坑 1:信号处理器永不卸载。 文档明确警告「Once a signal handler is registered with the process the underlying libc signal handler is never unregistered」📎 tokio/src/signal/unix.rs:379-380。即使 Signal 实例被 drop,后续信号仍会被 Tokio 捕获,默认行为不会恢复 📎 tokio/src/signal/unix.rs:338-340。

坑 2:信号会被合并。 文档说明「before poll is called, all signal notifications are coalesced into one item returned from poll」📎 tokio/src/signal/unix.rs:312-315。如果你收到 10 个 SIGINT 但只 poll 了一次,只会看到一个事件。这是 Unix 信号本身的特性(标准信号不排队),Tokio 没有额外合并。

坑 3:多 Runtime 下信号可能丢失。 由于全局 pipe 被多个 Runtime 竞争读取,一个 Runtime 可能读走字节而另一个永远等不到。生产环境应该只在一个 Runtime 里处理信号,或者用 signal_hook 自己管理。

坑 4:signal 函数 panic 条件。 文档说明「This function panics if there is no current reactor set, or if the rt feature flag is not enabled」📎 tokio/src/signal/unix.rs:398-405。在 Runtime 外调用 signal() 会 panic。

坑 5:recv() 的取消安全。 文档保证「This method is cancel safe. If you use it as a branch in tokio::select! and another branch completes first, then it is guaranteed that no signal is lost」📎 tokio/src/signal/unix.rs:423-427。这是因为信号事件存在全局的 EventInfo 里,recv() 只是读取,不消费底层状态。

设计思考

本章三个主题共享一个底层模式:状态的所有权决定了取消/关闭/信号的安全性。

  • JoinHandle 取消安全,因为输出在堆上,handle 只是引用。
  • Runtime 关闭顺序敏感,因为阻塞池和调度器共享 Handle,顺序错了会死锁或 panic。
  • 信号多 Runtime 冲突,因为 handler 和 pipe 是进程级全局状态,而 Signal 是 Runtime 级视图。

理解这个模式后,避坑清单可以归纳为三条原则:

1. 取消安全 = 状态在 Future 外部。 如果 Future 内部有缓冲区,drop 就会丢数据。JoinHandle、Signal::recv、tokio::sync::mpsc::Receiver::recv 都满足这个条件。

2. 关闭顺序 = 依赖方向的反序。 谁依赖谁,就先关被依赖者。调度器依赖 I/O 驱动,所以先关调度器;阻塞池独立,最后关。

3. 全局状态 = 多实例冲突。 任何进程级资源(信号 handler、pipe、文件描述符表)在多 Runtime 下都会冲突,要么限制单 Runtime,要么用外部同步。

本章小结

本章思考与自测

Q1:如果把 JoinHandle::poll 中的 coop::poll_proceed(cx) 去掉,在什么场景下会导致其他任务饿死?为什么 try_read_output 本身不消耗预算?

参考解析:coop::poll_proceed(cx) 在 📎 tokio/src/runtime/task/join.rs:325-325 处消耗协作预算。如果去掉,一个在循环里反复 select! 多个 JoinHandle 的任务可以在一次调度周期内无限轮询所有 handle,永不返回 Pending,从而饿死同 worker 上的其他任务。try_read_output 本身不消耗预算,因为它只是一次内存读取 + 可能的 waker 存储,不涉及 I/O 或锁竞争,开销极小。预算机制的设计意图是约束「可能长时间运行的操作」,而不是每次 poll 都收费。注意 coop.made_progress() 只在 ret.is_ready() 时调用 📎 tokio/src/runtime/task/join.rs:349-351,即只有真正拿到输出才归还预算——这是为了防止「轮询了但没结果」的操作累积消耗预算。

Q2:blocking/shutdown.rs 的 wait 方法中,如果 try_enter_blocking_region() 返回 None 且当前正在 panic,为什么选择返回 false 而不是继续等待?如果改成继续等待会发生什么?

参考解析:try_enter_blocking_region() 返回 None 表示当前在异步上下文里,不允许阻塞 📎 tokio/src/runtime/blocking/shutdown.rs:44-57。如果此时正在 panic,代码选择返回 false 不等待 📎 tokio/src/runtime/blocking/shutdown.rs:47-49。原因是:panic 展开过程中再次 panic 会导致进程 abort(double panic)。如果改成继续等待,就需要调用 block_on,而在异步上下文里 block_on 会 panic——在 panic 展开中 panic 会直接 abort 进程,丢失所有诊断信息。返回 false 让 drop 继续完成,panic 信息得以保留。这是一个「优雅降级」的设计:关闭不完整总比进程崩溃好。

Q3:假设你在 Runtime A 里创建了 Signal 监听 SIGTERM,然后把 Signal 移到 Runtime B 里 poll。signal_enable 里的 handle.check_inner() 检查的是哪个 Runtime?如果 Runtime A 先 drop 了,Runtime B 里的 Signal 还能收到信号吗?

参考解析:signal_enable 在 signal() 调用时执行,此时 handle 是 Runtime A 的 📎 tokio/src/signal/unix.rs:398-405。check_inner() 检查的是 Runtime A 的信号驱动 📎 tokio/src/signal/unix.rs:275。Signal 内部是 RxFuture,包装的是 watch::Receiver<()> 📎 tokio/src/signal/unix.rs:366-368,这个 receiver 注册在全局 Globals 的 EventInfo 上。如果 Runtime A drop 了,它的信号驱动停止从全局 pipe 读数据,但全局 handler 仍然会 record_event 并写 pipe。Runtime B 的信号驱动如果也在运行,会读到 pipe 数据并触发 EventInfo,从而唤醒 Signal 的 waker。所以 Runtime B 里的 Signal 可能还能收到信号,但取决于 Runtime B 是否有信号驱动在运行。如果 Runtime B 没有信号驱动(比如没启用 signal feature 或驱动已关闭),pipe 数据无人读取,Signal 永远等不到唤醒。这就是多 Runtime 信号处理的脆弱性。

章末过渡

取消安全、panic 传播、关闭顺序、信号冲突——这四个问题的共同根源是「状态所有权」在异步边界上的模糊。Tokio 通过把状态放在堆上、用引用计数管理生命周期、用 catch_unwind 隔离 panic、用全局 Globals 共享信号状态,给出了工程上可用的答案。但这些答案都有边界条件,生产环境必须显式处理。

下一章将进入架构权衡与未来演进:从 io_uring 到可插拔驱动。我们会看到 Tokio 如何在保持 API 稳定的前提下,为新一代 I/O 接口预留扩展空间,以及当前架构中哪些设计决策是历史包袱、哪些是前瞻布局。

至此,我们走完了 Tokio 生产环境中最容易踩坑的边界地带:取消安全依赖输出存储在堆上、try_read_output 的原子性;JoinHandle::drop 不取消任务,abort 才真正取消但对 spawn_blocking 无效;panic 被 catch_unwind 捕获后打包成 JoinError,不 await 就静默丢失;Runtime 关闭有严格顺序,在 async 上下文 drop 会 panic;信号 handler 是进程级全局状态,注册后永不卸载。这些规则背后是 Tokio 在正确性与性能之间的反复权衡。下一章我们将跳出具体机制,站在架构高度回顾这些权衡的由来,并展望 io_uring、驱动重构与自定义执行器接口将把 Tokio 带向何方。

AI 赋能代码库精读 · 本地优先架构

读完了本章?为你自己的私有项目生成专属架构全景书

基于 Tauri 2 + Rust 本地原生引擎,100% 源码离线隐私安全,零代码上传云端。像阅读一本传世专著一样拆解你的复杂系统。

⚡ Tauri 2 · Rust 原生引擎 · 100% 离线私密安全 · 适配超百万行代码库

CHAPTER 14

第 14 章:演进历程与架构沉思:Tokio 从微内核到工业级运行时的权衡

官方源: tokio-rs/tokio · Commit @e800714a · 全书进度: 第 14 / 14 章

第 14 章:演进历程与架构沉思:Tokio 从微内核到工业级运行时的权衡

上一章我们梳理了取消安全、panic 传播、关闭顺序与信号冲突这四类生产陷阱,它们看似分散,实则都指向同一个架构问题:状态所有权在异步边界上如何被清晰地划分。而划分所有权的方式,恰恰由运行时最底层的三个架构决策决定——任务如何被调度、I/O 事件如何被分发、并发正确性如何被验证。本章不再钻进某个具体函数的实现细节,而是站到架构高度,回顾 Tokio 在这些决策上的取舍,并沿着官方文档与源码中已经埋下的演进线索,看看 io_uring、驱动重构与自定义执行器接口会把 Tokio 带向何方。读完本章,你应该能回答一个实践问题:什么时候该扩展 Tokio,什么时候该绕开它。

一、三个历史权衡:为什么是现在这个样子

直觉模型

把 Tokio 想象成一家已经开了十年的餐厅。厨房的排班方式(work-stealing)、传菜员的独立编制(I/O 驱动与调度器分离)、以及后厨的卫生检查制度(loom 并发验证),都不是开业第一天就设计好的,而是在「客人变多、菜品变复杂」的过程中逐步演化出来的。理解这些演化,才能判断哪些设计是前瞻布局、哪些是历史包袱。

权衡一:work-stealing 而非全局队列

〔设计推断与架构权衡〕
全局队列的实现最简单:所有任务进一个 Mutex<VecDeque>,worker 线程抢锁取任务。但锁竞争会随核数增加而恶化,且缓存局部性差——任务在哪个核上被创建、在哪个核上被执行完全随机。

work-stealing 的取舍是:每个 worker 持有本地队列,spawn 时优先入本地队列(无锁、缓存友好),本地空了才去别的 worker 队列尾部窃取。代价是负载均衡有延迟,且窃取本身需要原子操作与内存屏障。Tokio 选择后者,是因为现代服务器动辄几十核,锁竞争的成本远高于偶发的窃取开销。

〔设计推断与架构权衡〕
这个决策的边界条件是:任务粒度不能太细。如果每个任务只做几微秒的工作,窃取与调度的开销占比就会失控。这也是为什么 Tokio 在 spawn_blocking 之外,还要求长任务主动 yield_now()——协作式调度本质上是在替 work-stealing 兜底。

权衡二:I/O 驱动独立于调度器

这是本章源码材料里最值得玩味的一处。看 tokio/src/runtime/io/mod.rs 的模块结构:

📎 tokio/src/runtime/io/mod.rs:5-22

rust
mod driver;
use driver::{Direction, Tick};
pub(crate) use driver::{Driver, Handle, ReadyEvent};

mod registration;
pub(crate) use registration::Registration;

mod registration_set;
use registration_set::RegistrationSet;

mod scheduled_io;
use scheduled_io::ScheduledIo;

mod metrics;
use metrics::IoDriverMetrics;

use crate::util::ptr_expose::PtrExposeDomain;
static EXPOSE_IO: PtrExposeDomain = PtrExposeDomain::new();

注意 driver、registration、scheduled_io 是三个独立模块,且对外只暴露 Driver、Handle、ReadyEvent、Registration 这几个类型。ScheduledIo 是 pub(crate) 的——它被 PtrExposeDomain 包裹,用于在 loom 测试下把裸指针暴露给并发检查。

〔设计推断与架构权衡〕
为什么 I/O 驱动不直接嵌进调度器?因为两者的生命周期与并发模型不同。调度器关心的是「哪个任务该跑」,I/O 驱动关心的是「哪个 fd 就绪了」。如果耦合,那么每次调度策略调整都要动 I/O 路径,反之亦然。更重要的是,block_on 单线程运行时也需要 I/O 驱动,但不需要 work-stealing 调度器——分离让两种运行时能复用同一套 I/O 实现。

权衡三:loom 做并发模型检验

tokio/src/loom/mod.rs 只有 14 行,却揭示了 Tokio 并发正确性的验证策略:

📎 tokio/src/loom/mod.rs:1-14

rust
//! This module abstracts over `loom` and `std::sync` depending on whether we
//! are running tests or not.

#![allow(unused)]

#[cfg(not(all(test, loom)))]
mod std;
#[cfg(not(all(test, loom)))]
pub(crate) use self::std::*;

#[cfg(all(test, loom))]
mod mocked;
#[cfg(all(test, loom))]
pub(crate) use self::mocked::*;

关键在 #[cfg(all(test, loom))] 这个条件:只有同时开启 test 和 loom 两个 cfg 时,才会用 mocked 模块替换 std。这意味着生产构建里根本没有 loom 的代码,零运行时开销。

〔设计推断与架构权衡〕
loom 的价值在于它能把「线程交错的所有可能顺序」穷举出来。像 ScheduledIo 里 AtomicUsize 的读改写、Waiters 链表的插入删除,这些在真实硬件上可能跑一百万次都不出错,但 loom 能在几秒内构造出触发竞态的交错。代价是测试运行慢、内存占用高,所以只能用于单元测试,不能进生产。

设计思考

这三个权衡有一个共同特征:它们都选择了「更复杂但更可扩展」的方案,并把复杂度限制在内部。work-stealing 的复杂度藏在调度器里,I/O 驱动的复杂度藏在 ScheduledIo 里,loom 的复杂度藏在 cfg 条件里。对外暴露的 API 始终是 spawn、TcpStream::read 这些简单接口。

〔设计推断与架构权衡〕
这也是判断「何时该扩展 Tokio」的第一条准则:如果你的需求能被现有 API 表达,就不要碰内部结构。一旦你开始依赖 pub(crate) 的类型或 tokio_unstable 的 cfg,就意味着你把自己绑在了 Tokio 的内部实现上,升级时会付出代价。

---

二、驱动重构:从「一个 waker 一个方向」到「任意兴趣集」

直觉模型

早期的 Tokio I/O 类型有个硬性限制:async fn read(&mut self) 需要 &mut self。这就像餐厅只有一个取餐窗口,同一时间只能有一个人排队——因为 waker 被存在 I/O 资源内部,而不是存在操作对应的 Future 里。tokio/docs/reactor-refactor.md 完整记录了这个限制的成因与重构方案。

旧架构的痛点

文档开篇就点明了问题:

📎 tokio/docs/reactor-refactor.md:16-20

rust
Currently, I/O types require `&mut self` for `async` functions. The reason for
this is the task's waker is stored in the I/O resource's internal state
(`ScheduledIo`) instead of in the future returned by the `async` function.
Because of this limitation, I/O types limit the number of wakers to one per
direction (a direction is either read-related events or write-related events).
〔设计推断与架构权衡〕
把 waker 存在资源内部,意味着「一个方向只能有一个等待者」。如果你同时想读和写同一个 TcpStream,就必须 split() 成两半,各自持有独立的 waker 槽。这就是 TcpStream::split() 存在的原因——它不是 API 设计偏好,而是内部数据结构的直接约束。

新架构:把 waker 移到 Future 里

重构的核心思路是「把 waker 从资源状态移到操作 Future 里」,从而支持每个操作注册多个 waker:

📎 tokio/docs/reactor-refactor.md:22-25

rust
Moving the waker from the internal I/O resource's state to the operation's
future enables multiple wakers to be registered per operation. The "intrusive
wake list" strategy used by `Notify` applies to this case, though there are some
concerns unique to the I/O driver.

新的 ScheduledIo 结构如下:

📎 tokio/docs/reactor-refactor.md:97-134

rust
#[derive(Debug)]
pub(crate) struct ScheduledIo {
    /// Resource's known state packed with other state that must be
    /// atomically updated.
    readiness: AtomicUsize,

    /// Tracks tasks waiting on the resource
    waiters: Mutex,
}

#[derive(Debug)]
struct Waiters {
    // List of intrusive waiters.
    list: LinkedList,

    /// Waiter used by `AsyncRead` implementations.
    reader: Option,

    /// Waiter used by `AsyncWrite` implementations.
    writer: Option,
}

// This struct is contained by the **future** returned by `readiness()`.
#[derive(Debug)]
struct Waiter {
    /// Intrusive linked-list pointers
    pointers: linked_list::Pointers,

    /// Waker for task waiting on I/O resource
    waiter: Option,

    /// Readiness events being waited on. This is
    /// the value passed to `readiness()`
    interest: mio::Ready,

    /// Should not be `Unpin`.
    _p: PhantomPinned,
}

这里有几个精妙的设计点值得展开:

第一,readiness 是 AtomicUsize,waiters 是 Mutex<Waiters>。 为什么不用一把锁保护两者?因为 readiness 的读操作极其频繁(每次 readiness() 调用都要检查),而写操作只在收到 mio 事件时发生。用原子变量让读路径无锁,是典型的读写分离优化。

第二,Waiter 是侵入式链表节点。 pointers: linked_list::Pointers<Waiter> 让 Waiter 本身成为链表的一部分,不需要额外分配节点。_p: PhantomPinned 明确标记它不可 Unpin——因为侵入式链表的节点地址一旦移动,链表就断了。

第三,reader 和 writer 两个 Option<Waker> 是给 AsyncRead/AsyncWrite 用的。 文档解释了原因:

📎 tokio/docs/reactor-refactor.md:210-213

rust
The `AsyncRead` and `AsyncWrite` traits use a "poll" based API. This means that
it is not possible to use an intrusive linked list to track the waker.
Additionally, there is no future associated with the operation which means it is
not possible to cancel interest in the readiness events.
〔设计推断与架构权衡〕
这是新旧两套机制的妥协共存:async fn 路径用侵入式链表(支持多等待者、可取消),poll 路径用固定槽位(不支持取消、但兼容 trait)。这种「两套机制并存」是渐进式重构的典型代价。

竞态条件与 tick 机制

重构中最棘手的问题是竞态。文档给了一个具体的死锁场景:

📎 tokio/docs/reactor-refactor.md:175-175

rust
If care is not taken, if between `mio_socket.read(buf)` returning and
`clear_readiness(event)` is called, a readiness event arrives, the `read()`
function could deadlock. This happens because the readiness event is received,
`clear_readiness()` unsets the readiness event, and on the next iteration,
`readiness().await` will block forever as a new readiness event is not received.

解决方案是引入 tick 机制,把 readiness 这个 AtomicUsize 拆成多个位段:

📎 tokio/docs/reactor-refactor.md:199-199

code
| shutdown | generation |  driver tick | readiness |
|----------+------------+--------------+-----------|
|   1 bit  |   7 bits   +    8 bits    +  16 bits  |
〔设计推断与架构权衡〕
这个位段布局是「用空间换正确性」的经典案例。tick 每次 mio::poll() 递增,ReadyEvent 携带读取时的 tick。clear_readiness() 只在 tick 匹配时才清除就绪状态——如果 tick 不匹配,说明期间有新事件到达,不能清除。这样就把「清除」和「新事件到达」的竞态消解在了一个原子读改写里。

下面这张流程图刻画了 readiness() 与 clear_readiness() 之间的决策路径:

mermaid
flowchart TD
    start["readiness(interest).await"] --> check_ready{"已知 readiness与 interest 有交集?"}
    check_ready -->|是| ret_event["返回 ReadyEvent携带当前 tick"]
    check_ready -->|否| wait["注册 Waiter 到ScheduledIo.waiters"]
    wait --> mio_poll["mio.poll() 收到事件tick 递增"]
    mio_poll --> notify["遍历 waitersinterest 匹配者唤醒"]
    notify --> ret_event
    ret_event --> do_read["mio_socket.read(buf)"]
    do_read --> read_ok{"read 结果?"}
    read_ok -->|Ok| done["返回 Ok(v)"]
    read_ok -->|WouldBlock| clear["clear_readiness(event)"]
    read_ok -->|其他 Err| err["返回 Err(e)"]
    clear --> tick_match{"event.tick ==当前 readiness.tick?"}
    tick_match -->|是| clear_ok["清除 readiness 位"]
    tick_match -->|否| skip["跳过清除保留新事件"]
    clear_ok --> start
    skip --> start

这张图的关键分支在 tick_match:如果 tick 不匹配,clear_readiness 必须放弃清除,否则会丢掉刚到达的事件,导致下一轮 readiness() 永久阻塞。

取消兴趣与内存泄漏

侵入式链表带来一个新问题:如果 readiness() 返回的 Future 被提前 drop,链表节点必须被摘除。文档明确警告:

📎 tokio/docs/reactor-refactor.md:144-148

rust
The future returned by `readiness()` uses an intrusive linked list to store the
waker with `ScheduledIo`. Because `readiness()` can be called concurrently, many
wakers may be stored simultaneously in the list. If the `readiness()` future is
dropped early, it is essential that the waker is removed from the list. This
prevents leaking memory.
〔设计推断与架构权衡〕
这正是上一章「取消安全」在 I/O 层的体现。readiness() 的 Future 必须在 Drop 实现里把自己从链表摘除,否则节点会永久留在 ScheduledIo 里,既泄漏内存,又会在下次事件到达时被错误唤醒。

设计思考与生产踩坑

为什么不用 Vec<Waker> 而用侵入式链表? 文档在讨论 &Resource 实现时给出了答案:

📎 tokio/docs/reactor-refactor.md:228-233

rust
It is only possible to implement `AsyncRead` and `AsyncWrite` for resource types
themselves and not for `&Resource`. Implementing the traits for `&Resource`
would permit concurrent operations to the resource. Because only a single waker
is stored per direction, any concurrent usage would result in deadlocks. An
alternate implementation would call for a `Vec` but this would result in
memory leaks.
〔设计推断与架构权衡〕
Vec<Waker> 的问题是:Future 被 drop 后,对应的 waker 留在 Vec 里无法定位删除,只能等下次事件到达时才发现「这个 waker 已经失效」。侵入式链表让节点地址就是 Future 内部字段的地址,drop 时能精确摘除。

生产踩坑点:TcpStream::by_ref() 返回的 TcpStreamRef 持有 read_waiter 和 write_waiter 两个节点:

📎 tokio/docs/reactor-refactor.md:238-244

rust
struct TcpStreamRef {
    stream: &'a TcpStream,

    // `Waiter` is the node in the intrusive waiter linked-list
    read_waiter: Waiter,
    write_waiter: Waiter,
}
〔设计推断与架构权衡〕
这意味着 TcpStreamRef 一旦被 drop,两个 waiter 节点同时失效。如果你在 select! 里用 by_ref() 的引用跨分支共享,要小心生命周期——TcpStreamRef 不能活得比 TcpStream 长,也不能在多个 select! 分支间被同时借用。

---

三、自定义执行器:TokioContext 与「绕开 Tokio」的边界

直觉模型

有时你不想用 Tokio 的调度器,只想借它的 I/O 和定时器。这就像你不想在餐厅堂食,只想用它的外卖窗口。examples/custom-executor.rs 展示了这种「混合模式」:用 futures::executor::ThreadPool 做调度,用 Tokio 做 I/O。

核心机制:TokioContext

整个例子的关键在 TokioContext 这个包装类型:

📎 examples/custom-executor.rs:51-54

rust
impl ThreadPool {
    fn spawn(&self, f: impl Future + Send + 'static) {
        let handle = self.rt.handle().clone();
        self.inner.spawn_ok(TokioContext::new(f, handle));
    }
}
〔设计推断与架构权衡〕
TokioContext::new(f, handle) 把 Future 和 Tokio 的 Handle 绑在一起。当外部执行器 poll 这个包装 Future 时,TokioContext 会先进入 Tokio 的运行时上下文(设置线程局部的 Handle),再 poll 内部的 f。这样 f 里调用 TcpListener::bind 时,就能找到 Tokio 的 I/O 驱动。

看整个例子的结构:

📎 examples/custom-executor.rs:38-48

rust
static EXECUTOR: Lazy = Lazy::new(|| {
    // Spawn tokio runtime on a single background thread
    // enabling IO and timers.
    let rt = tokio::runtime::Builder::new_multi_thread()
        .enable_all()
        .build()
        .unwrap();
    let inner = futures::executor::ThreadPool::builder().create().unwrap();

    ThreadPool { inner, rt }
});
〔设计推断与架构权衡〕
这里 Tokio 运行时被创建但没有被 block_on 驱动——它只是「存在」,提供 I/O 驱动和定时器。真正的任务调度由 futures::executor::ThreadPool 负责。这种模式下,Tokio 的 worker 线程实际上在空转(等待 I/O 事件),任务执行发生在 futures 的线程池里。

数据流:一次 TcpListener::bind 的跨执行器旅程

mermaid
sequenceDiagram
    participant App as "应用 (main)"
    participant FE as "futures::ThreadPool"
    participant TC as "TokioContext"
    participant TR as "tokio::Runtime (后台线程)"
    participant IO as "I/O 驱动 (mio)"

    App->>FE: spawn_ok(TokioContext::new(f, handle))
    FE->>TC: poll(cx)
    TC->>TC: enter(handle) 设置线程局部上下文
    TC->>TC: f.poll(cx) 执行 TcpListener::bind
    TC->>TR: 通过 Handle 访问 I/O 驱动
    TR->>IO: Registration::new 注册 fd
    IO-->>TR: 注册完成
    TR-->>TC: 返回 Pending 或 Ready
    TC-->>FE: 返回 poll 结果
    Note over FE,TR: I/O 就绪时,Tokio 驱动唤醒 wakerFE 重新调度该任务

这张时序图的关键在于:任务的 poll 发生在 futures 线程池,但 I/O 事件的等待发生在 Tokio 后台线程。两者通过 Handle 和 waker 连接。

设计思考:何时该绕开 Tokio

〔设计推断与架构权衡〕
这个例子的存在本身就是一个信号:Tokio 的架构允许「只用 I/O 驱动,不用调度器」。判断标准可以归纳为三条:

1. 如果你需要与已有的执行器生态集成(比如某些框架强制要求 futures::executor),用 TokioContext 是最小侵入的方案。

2. 如果你需要完全控制调度策略(比如实时系统要求确定性调度),Tokio 的 work-stealing 不满足需求,但它的 I/O 驱动仍然可用。

3. 如果你只是嫌 Tokio 的 API 复杂,那不该绕开——TokioContext 引入的跨执行器边界会带来新的调试难度,得不偿失。

生产踩坑点:TokioContext 模式下,Tokio 运行时的 block_on 从未被调用,意味着 Runtime::shutdown 的清理逻辑不会自动触发。你必须在程序退出前显式 drop Runtime,否则 I/O 驱动的后台线程可能不会优雅关闭。

与 io_uring 的关系

〔设计推断与架构权衡〕
tokio/src/runtime/io/mod.rs 顶部的 cfg 条件透露了 io_uring 的接入方式:

📎 tokio/src/runtime/io/mod.rs:1-4

rust
#![cfg_attr(
    not(all(feature = "rt", feature = "net", feature = "io-uring", tokio_unstable)),
    allow(dead_code)
)]

注意 feature = "io-uring" 和 tokio_unstable 同时出现。这意味着 io_uring 支持目前是实验性的,必须同时开启 unstable 特性才能编译。allow(dead_code) 则说明:当这些特性未开启时,模块里的部分代码不会被使用,编译器会警告——用 allow 压掉。

〔设计推断与架构权衡〕
io_uring 与 epoll 的根本区别在于:epoll 是「就绪通知」,io_uring 是「完成通知」。前者需要应用自己发起 read/write 系统调用,后者由内核直接完成 I/O 并返回结果。这对 Tokio 的 ScheduledIo 模型是巨大冲击——readiness() 的语义在 io_uring 下不再适用,需要一套全新的「提交-完成」抽象。这也是为什么 io_uring 支持迟迟停留在 unstable:它不是加一个后端那么简单,而是要重构整个 I/O 驱动的抽象层。

---

本章小结

本章从架构高度回顾了 Tokio 的三个核心权衡,并展望了三条演进路径:

历史权衡:

  • work-stealing 用调度复杂度换取多核扩展性,边界是任务粒度不能太细;
  • I/O 驱动独立于调度器,让 block_on 与多线程运行时复用同一套 I/O 实现;
  • loom 通过 cfg 条件在生产构建中完全消失,只在测试时穷举线程交错。

驱动重构(reactor-refactor.md):

  • 把 waker 从 ScheduledIo 内部移到操作 Future 里,用侵入式链表支持多等待者;
  • 用 AtomicUsize 的位段布局(shutdown/generation/tick/readiness)消解 clear_readiness 的竞态;
  • AsyncRead/AsyncWrite 因 poll 语义无法用侵入式链表,保留 reader/writer 固定槽位作为妥协。

未来演进:

  • io_uring 需要「提交-完成」新抽象,目前受 tokio_unstable 保护;
  • TokioContext 允许只用 I/O 驱动、不用调度器,但需手动管理 Runtime 生命周期;
  • 判断「扩展还是绕开」的准则:能用现有 API 表达就不碰内部结构。

本章思考与自测

Q1: 在 ScheduledIo 的 readiness 位段布局中,如果把 tick 字段从 8 位缩减到 4 位,在什么场景下会触发错误?请结合 clear_readiness 的 tick 匹配逻辑分析。

参考解析:tick 在每次 mio::poll() 时递增 📎 tokio/docs/reactor-refactor.md:185-185。clear_readiness 只在 event.tick == 当前 readiness.tick 时才清除就绪位 📎 tokio/docs/reactor-refactor.md:199-199。如果 tick 只有 4 位,那么每 16 次 poll 就会回绕。假设某个 ReadyEvent 携带 tick=15,在它被 clear_readiness 之前,mio 又 poll 了 1 次,tick 回绕到 0。此时 clear_readiness 发现 tick 不匹配(15 != 0),会错误地跳过清除——但实际上期间可能没有新事件到达,只是 tick 回绕了。这会导致就绪位被永久保留,后续 readiness() 立即返回但 read 仍然 WouldBlock,陷入忙循环。8 位 tick 在正常负载下足够(256 次 poll 内完成一次 read-clear 循环),但极端高并发下仍有回绕风险,这是位段布局的固有边界。

Q2: examples/custom-executor.rs 中,Tokio 运行时被创建但从未 block_on。如果此时调用 rt.shutdown_timeout(),会发生什么?为什么这个例子选择不调用?

参考解析:rt.shutdown_timeout() 会等待所有任务完成并关闭 I/O 驱动。但在这个例子里,任务实际运行在 futures::executor::ThreadPool 上 📎 examples/custom-executor.rs:51-54,Tokio 运行时里没有任务——它只提供 I/O 驱动。如果调用 shutdown_timeout,它会立即返回(因为没有任务),但 I/O 驱动的后台线程可能仍在运行。例子选择不调用,是因为 EXECUTOR 是 Lazy 静态变量,程序退出时由 Rust 的静态析构机制处理。真正的坑在于:如果 TokioContext 包装的 Future 还在运行,而 Runtime 被 drop,那么 Future 里的 I/O 操作会 panic(找不到运行时上下文)。生产环境必须确保所有 TokioContext Future 完成后才 drop Runtime。

Q3: 假设你要为 Tokio 添加一个基于 io_uring 的 I/O 后端。根据 reactor-refactor.md 中 readiness() 的语义,哪些部分可以直接复用,哪些必须重写?

参考解析:可以直接复用的是 Registration 的注册接口和 ScheduledIo 的 waiters 链表结构——它们管理的是「谁在等」,与底层是 epoll 还是 io_uring 无关。必须重写的是 readiness() 的语义:epoll 下它返回「fd 就绪」,io_uring 下没有「就绪」概念,只有「提交的 SQE 完成」。clear_readiness 的 tick 机制也需要重新设计——io_uring 的完成事件自带 user_data 标识,不需要 tick 来区分新旧事件。最根本的改动是:readiness() 返回的 Future 在 io_uring 下应该变成「提交 SQE 并等待 CQE」,这意味着 Waiter 结构需要携带 SQE 参数,而不仅仅是 interest。这也是为什么 io_uring 支持受 tokio_unstable 保护 📎 tokio/src/runtime/io/mod.rs:1-4——它不是替换后端,而是改变 I/O 驱动的抽象契约。

至此,我们完成了从具体陷阱到架构权衡的爬升。回顾全书,从 Future 的惰性求值到调度器的公平性,从取消安全到关闭顺序,再到本章的 io_uring 与可插拔驱动,所有讨论都围绕一个核心:在异步边界上清晰地划分状态所有权。Tokio 的架构并非一成不变,io_uring 的零拷贝 I/O、驱动层的解耦、自定义执行器接口的开放,都在推动它向更灵活、更高效的方向演进。当你合上这本书,希望留下的不是一堆 API 用法,而是一套判断力:知道何时该信任运行时,何时该介入底层,以及如何在生产环境中避开那些会咬人的组合。异步 Rust 的生态仍在快速生长,保持对源码与官方文档的追踪,比记住任何结论都更重要。

AI 赋能代码库精读 · 本地优先架构

读完了本章?为你自己的私有项目生成专属架构全景书

基于 Tauri 2 + Rust 本地原生引擎,100% 源码离线隐私安全,零代码上传云端。像阅读一本传世专著一样拆解你的复杂系统。

⚡ Tauri 2 · Rust 原生引擎 · 100% 离线私密安全 · 适配超百万行代码库

读懂任何复杂项目,你真正需要的是一本专著

本书由 AiReadCode 扫描官方开源仓库全自动编撰,结合真实不可变 Commit 节点与 FACT 药丸行号溯源,提供纯静态、零服务依赖的极致双栏交互式在线阅读体验。

在 GitHub 上 Star 本项目 ★ 浏览更多架构专著 →