Await-Tree: Async Rust 可观测性的灵丹妙药

作者：赵梓淇 (Bugen Zhao) RisingWave Labs 内核开发工程师

背景

Async Rust 通过一套简明的无栈协程抽象，为开发者提供了灵活且高效的异步编程能力；但其多变的调度和执行模型，也使得并发编程问题的调试变得格外棘手。在本文中，我们将介绍 Await-Tree 这一 Async Rust 调试工具，其基于 RisingWave 分布式流式数据库对于 Async Rust 的深入实践，允许开发者实时地将 Async Task 的执行状态以树状结构导出，分析 Task 内部的异步调用链及 Task 间的依赖阻塞关系，以极低的运行时代价大幅提升系统的可观测性和可调试性。

本篇内容主要分为三个部分：

1. 回顾 Async Rust 的设计与痛点
2. 介绍 Await-Tree 的设计原理与实现
3. 展示 Await-Tree 在 RisingWave 中的真实应用案例

Async Rust 的设计与痛点

Async Rust 回顾

Async Rust 是 Rust 编程语言与异步编程的结合。作为一种基于事件的编程范式，它首先继承了异步编程语言的共同优势，如：

• 引入 async/await 关键字，使得开发者能够通过同步编程的思维来编写和理解异步代码，进而充分利用异步接口以实现高性能 I/O；
• 高性能的用户态协作式调度，使得开发者能够在应用程序中以极低的性能开销创建大量协程，进而轻松地实现高并发。

此外，Async Rust 也继承了 Rust 编程语言自身独特的高性能、安全性等特点，使其在一众异步编程语言中脱颖而出，成为包括数据库等下一代分布式系统的首选编程语言。在这一方面，Async Rust 的独特优势包括：

• Ownership 和 Lifetime 机制、Send/Sync/Pin 等语言设计，使得开发者可以通过编译期的类型系统约束来保证并发编程的安全性和健壮性；
• 在此基础上采用了无栈协程的设计，避免了为每个协程分配独立的堆栈并在调度时切换，进一步降低了异步编程的内存占用和性能开销。

Async Rust 使用 Future 这一结构作为其无栈协程的基本抽象。从直观上来看，Future 可以理解为一个以 Enum 表示的状态机，Runtime 通过从外部不断调用 poll 函数来进行状态切换。以该图为例，一个 Future 被 Runtime 调度运行的过程可以简化地描述为：

1. 状态机初始状态为 Init。在 Runtime 第一次调用 poll 时，Future 会执行一段同步代码，直至进入下一个状态，如图中的 WaitForXx 等，表示该 Future 被网络、磁盘等 I/O 或锁、信号量等资源阻塞，而无法继续执行。此时，Future 会将 Runtime 提供的 Waker 保存注册为资源可用时的 Callback 函数，然后从 poll 函数中返回 Pending。
2. 当 Runtime 得知该 Future 处于 Pending 状态时，将其移动至等待队列中不再调度，转而切换至下一个任务进行调度。
3. 一段时间后，资源变为可用状态并调用 Callback 函数唤醒该 Future，使得 Runtime 重新将该 Future 放回调度队列中，等待下一次调度执行。
4. Future 被 Runtime 调度执行后，其 poll 函数再次被调用。Future 将从状态机中恢复其执行状态，并重复步骤 1~3 中描述的过程，直至 Future 进入 Done 状态，返回 Ready 代表其执行完毕。

在 Async Rust 中，Runtime 提供的底层 Future 常常需要通过手动实现 poll 函数来定义其行为。幸运的是，Async Rust 为我们提供了 async/await 关键字，使得开发者可以直接采用编写同步代码的思想来嵌套组合这些 Future，由编译器为每个 Async 函数生成匿名的 Future 类型。这些单线程执行的 Future 最终形成了 Task 的概念，即为作为 Runtime 的基本调度单元。

Async Rust 观测与调试的痛点

世上没有免费的午餐。Async Rust 通过一系列设计带来了优秀的性能、安全性与灵活性的同时，也必然有所牺牲，其中最为突出的便是其在观测与调试上带来的不便。

• Future 灵活的可组合性

  在 Async Rust 中，开发者们除了通过像编写同步代码一样 await 调用异步函数之外，还可以通过一系列第三方库提供的 join、select 等函数使得多个 Future 在同一 Task 内并发执行，甚至实现更加复杂的控制流。值得指出的是，得益于 Async Rust 中 Future 的灵活设计，这些控制流并不需要 Async Rust 提供任何额外的语言特性：它们仅仅是通过手动实现 poll函数，进而定制其中的执行和调度逻辑。

  然而，对于这些协程内部的并发执行，从 CPU 的角度来看，其仍然是单线程地在不同 Future 之间交错执行而已，无法感知开发者在业务或逻辑上所理解的"同时执行"。该特性使得所有基于线程的观测与调试工具都不能直观地反映当前某一 Task 的执行状态。例如，在 Async 上下文中打印 Backtrace，我们只能看到当前被调度执行的 Future 的调用栈，而无从得知 Task 内其它所有"正在"并发执行的 Future 的状态。出现这一现象的本质原因在于：当开发者在异步编程时，其执行在逻辑上可能形成一个"调用树"，而不仅仅是线程视角下的"调用栈"。

• 用户态调度的无栈协程

  正如我们前文所提到的，有栈协程实现需要为每个协程分配一段独立的堆栈内存，并在调度执行时进行切换；相反，无栈协程将每个协程的状态维护在状态机中，只有在执行时才通过 poll 函数调用将状态恢复到当前执行线程的堆栈中，而在无法继续执行时会重新将堆栈中的状态保存到状态机的结构中。一般而言，无栈协程的这一特点使其在内存管理上具有一定挑战，如栈上局部变量的引用维护等，Rust 通过 Pin 的概念从类型系统上解决了这一问题。

  尽管内存安全性得到了保证，无栈协程"不执行时不存在堆栈空间"的特点，使得我们无法通过基于线程的观测与调试工具来还原一个 Async Task 在阻塞发生前的执行状态，进而分析其被阻塞的原因，调试死锁等常见并发编程问题。在采用同步编程时，我们可以轻易地通过 Attach Debugger 等方式查看一个线程 Park 前的调用栈；而对于 Rust 异步编程，同样的办法我们只能得到看到 Runtime 的 Worker Thread 目前因为没有可执行的 Task 而同步阻塞在其任务队列上：至于为何 Task 异步阻塞，我们则无从得知。

不难看出，解决上述 Async Rust 痛点的核心问题在于：我们需要一种原生为 Task 设计的观测和调试工具，其能够隐藏"Task 在线程上物理执行"的实现细节，从 Task 角度出发来反映其逻辑上的执行状态。例如，由于并发执行的存在，Task 的执行状态将会用一棵树（而非栈）来表示；当一个 Task 即将被阻塞时，由于其逻辑上仍处于"尚未完成"的执行状态，其调用树也应被保留以供后续观测与调试。

Await-Tree 的设计原理与实现

基于上述思想，我们推出了 Await-Tree 这一 Async Rust 的调试工具。顾名思义，Await-Tree 能够在运行时追踪关键 Future （即 await 点）的整个生命周期中的控制流，将 Task 的逻辑执行状态实时地维护为一棵树，允许在任意时刻查询或导出，进而调试 Async Rust 在执行中遇到的死锁等问题。

我们以一个简单的 Async Rust 程序为例，介绍 Await-Tree 的使用方法与效果：

在这段代码中，我们涉及了一些简单的 Async 函数嵌套以及 Join 的并发执行。与平常不同的是，我们在每一个需要追踪的关键 Future 后额外添加了 .instrument_await 并为其指定了一个名字：该名字既可以是静态的字符串常量，也可以带有额外的运行时信息。

在执行 foo 函数并等待 1 秒后，该 Task 的所有分支都应处于 sleep 等待唤醒的阻塞状态中，整个 Task 的栈空间应该暂存于 Runtime 的某个状态机中。然而，我们可以通过 Await-Tree 以树状结构还原其逻辑上的执行状态，并得知每个 Future 已经持续的时间。

继续执行 2 秒后，由于部分 sleep 已经被唤醒，Task 进入了新的执行状态，此时导出的 Await-Tree 则同样能够反映该更新后的执行状态，Future 的持续时间也随之更新。和 2 秒前的树不同的是，这棵树中额外存在一个 current 指针，代表 CPU 正在执行的位置：current 指针的存在代表该 Task 目前正在执行某个 Future 的 poll 函数，而没有被阻塞。

Await-Tree 的设计细节

在直观地了解了 Await-Tree 的用法后，我们不妨进一步深入研究其设计细节。Await-Tree 的结构是随着 Future 的 poll 函数的控制流而维护的，为了让 Await-Tree 能够正确地反映一个 Async Task 的执行状态，其中最关键的便是需要充分理解 Future 整个生命周期中可能存在的控制流。

我们以一个稍微复杂、但在实际生产中十分常见的 Async Rust 程序为例，介绍 Await-Tree 是如何响应 Future 的控制流并进行维护的。在这个例子中，我们首先为 Async 函数 query 的执行设定一个超时，并在超时后打印一条告警信息，最后继续等待 query 执行完毕。

• Future 的构造
  设想我们由 serve 函数调用到 handle 函数。在刚刚进入 handle 时，我们首先构造了 query、select、timeout 这三个 Future。一般而言，Future 的执行是 lazy 的，因此构造 Future 并不会对 Await-Tree 产生影响：current 指针依然指向 handle。

• Future 第一次 poll
  接下来，由于我们在 select 上调用了 await，即第一次在 select 上调用了 poll，Await-Tree 会创建一个新的节点挂载为 current 的子节点，并移动 current 指针的位置。在 select``poll 函数的实现逻辑中，它会分别尝试调用两个 Future 的 poll 函数。类似地，Await-Tree 会为它们创建一个新的节点，挂载并移动 current 指针。

• Future 返回 Pending
  在执行了 query 中的一些同步代码后，我们假设它阻塞在了某个网络 I/O 上，进而其 poll 函数返回 Pending。虽然我们在物理上退出了 poll 函数，但该 Future 在逻辑上依然处于正在执行的过程中。因此，我们保留其节点，只是将 current 的位置移动到父节点 select 上。对于 select 的另一边 timeout 则重复同样的逻辑。

• Task 整体 Pending
  select 的 poll 函数返回 Pending，使得它的 Caller 也开始递归地返回 Pending，最终导致整个 Task 返回 Pending。此时由于整个 Task 处于被阻塞的状态，current 指针将不再指向任何一个节点；然而，Task 被挂起前 Await-Tree 本身的树结构依然保留，并能被其它线程所获取到。

• Future 再次 poll
  一段时间后，timeout 首先被唤醒，这使得 Runtime 开始重新调度执行该 Task，进而递归地调用各个 Future 的 poll 函数。由于各个 Future 已经处于运行状态，其 Await-Tree 上的节点已经存在，因此整个过程只需移动 current 指针。

• Future 返回 Ready
  timeout 被再次 poll 后立刻返回 Ready。由于 select 的语义和实现逻辑，它也会因此返回 Ready。因为 timeout 已经执行完毕，因此其在 Await-Tree 上对应的节点可以被删除；而对于 query，由于它还没有执行完毕，且我们向 select 传入的只是该 Future 的引用，select 的返回并不会导致 query 的取消。此时，由于父节点已删除，query 将被从 Await-Tree 上卸载并不再于其它节点连接。

• Future 调用关系重建
  在打印了告警日志后，我们在 handle 函数中继续等待 query 的完成。此时，Await-Tree 中的 query 节点发现其父节点发生了变化，将重建和 handle 的父子调用关系。

• Future 返回 Ready
  在一段时间后，query 执行完毕返回 Ready，使得 handle 递归地返回 Ready，Await-Tree 也恢复到了其初始状态。
  此外，Future 还有 Cancel 这一特殊的控制流，其在 Await-Tree 上的行为和 Ready 相似，只是无需操作 current 指针，这里就不再赘述。

Await-Tree 的实现

在深入理解了 Await-Tree 应该如何针对 Future 的不同控制流进行维护后，我们便可以进一步讨论 Await-Tree 的编程实现。

考虑到一棵 Await-Tree 表示了一个 Task 的执行状态，其物理执行是单线程的，因此对其的维护操作也无需引入任何线程间的竞争；此外，为了简化接口设计，我们需要将 Await-Tree 的数据结构存放在一个全局可访问的上下文中。据此，我们选择将 Await-Tree 维护在 Task-Local Storage 里，可以理解为 Thread-Local Storage 的协程版本。

在 Rust 中实现一个链式数据结构是比较困难的，还可能由于 Unsafe Rust 的误用导致内存安全问题。因此，我们使用了基于 Arena 的 Index Tree 来实现 Await-Tree，每个 Future 只需维护一个 Arena 中的 ID 即可访问 Await-Tree 对应的节点。这样的设计使得 Await-Tree 中不含任何 Unsafe 代码。

在接口设计上，我们采用了与 futures::FutureExt 相同的 Future Adaptor 的设计，通过为所有的 Future 实现 InstrumentAwait Trait，开发者可以在 Future 上直接调用 .instrument_await 以构造一个包装后的 Future，从而使得该 Future 能够被 Await-Tree 追踪维护。

在 Instrumented 的 poll 逻辑中，我们按照上节中介绍的设计细节实现并维护了一个 Await-Tree 节点的状态机。值得指出的是，即使需要追踪的 Future 是通过手动定制 poll 函数的逻辑实现的（如常见的 Join/Select 等），只要其行为是"结构化"的，Await-Tree 几乎都能通过其严谨而完备的实现正确追踪其执行状态，而无需做任何特殊的感知或处理，这也使得 Await-Tree 能够自然地与各种 Future 或 Runtime 库兼容。

Await-Tree 在 RisingWave 中的应用案例

Await-Tree 是基于 RisingWave 对 Async Rust 的深入实践后设计提出的。作为一款基于 SQL 的新一代云原生流式数据库，RisingWave 通过分布式的流计算任务来对系统中的 Materialized View 进行实时增量维护，并借助 S3 上的自研共享存储引擎来管理流计算状态和表数据。相比 OLAP 系统而言，RisingWave 的计算任务需要长期运行，对于稳定性的要求极高；此外，RisingWave 的流计算算子逻辑更加复杂，计算与状态存储请求穿插进行，对于 Async 的依赖性很强。为此，RisingWave 长期在生产环境中部署启用 Await-Tree，以极低的运行时开销大幅提高了 Async Rust 的可观测性和可调试性，帮助我们解决了数个棘手的 Async 死锁问题。

Backtrace 的补充

Await-Tree 原生面向 Async Rust 中的 Task 设计，可以提供 Task 逻辑上的调用树，因此可以直接用来做为 Thread Backtrace 的补充。例如，RisingWave 在 Panic 发生时会导出当前 Task 的 Await-Tree 并打印到日志中，通过运行时附加的动态信息，可以更加清晰地了解发生问题的 Executor 及其所属的 Materialized View，甚至包括同 Task 内的其它并发 Future，帮助开发者更加轻松地定位问题所在。

调试 Async Stuck

RisingWave 采用了类似 MPP 的执行引擎设计，数据处理根据分布的不同被切分为 Fragment，每个 Fragment 由多个独立执行的 Actor 在各自的 Async Task 上执行；Actor 之间通过本地 Channel 或远程 gRPC 进行通信，而 Actor 内部由多个 Executor 嵌套组合、同时存在并发执行的场景。由于引擎模型和计算逻辑较为复杂，RisingWave 在早期开发中常常遇到 Async Stuck 的死锁问题。Await-Tree 的引入大幅简化了此类问题的调试过程。

案例 1：Future Detach

RisingWave 的状态存储引擎的后端位于 S3 对象存储上，为了提升读写性能，我们自研了一套 Tiered Cache 层以更多地将 S3 上的存储 Block 缓存在内存和本地硬盘上。当对某一 Block 的请求未能命中 Cache 需要进行网络请求时，一个很常见的优化便是将后续同一 Block 的请求拦截，使得它们只需等待第一次的网络请求返回即可，该行为也被称为 Single Flight。

在 RisingWave 引入该优化后，我们常常在测试中遇到某个 Streaming Job 卡住而无法继续执行的问题。通过打印所有 Actor 的 Await-Tree，我们得到了如下的执行状态，发现其中 Block Cache 的网络请求从整个 Await-Tree 上卸载了（Detached），这说明 Actor 即使被网络请求的完成唤醒，也没能重新调用 get 对应 Future 的 poll 函数，进而无法通知当前执行路径上的 wait Future 完成，导致死锁的发生。

通过分析代码发现，我们的 Join 算子通过 select 两个上游 Executor 作为输入，从两端持续获取需要处理的数据，这便导致了潜在的问题：当某一端输入准备好下一批需要处理的后，另一端输入可能处于任意状态，例如正在等待某个 Block 的网络请求。然而，此时 Join 算子会获得该 Task 的执行权，进而开始对这批数据进行处理。如果在这个过程中，Join 算子也需要访问状态存储，并且恰好访问到了和此前另一端正在访问的相同的 Block，它就会根据 Single Flight 优化选择等待之前的请求完成，进而暂时性地导致当前 Actor 处于 Pending 状态。

一段时间后，此前的 get 请求完成，Actor 被唤醒，其 poll 函数被重新调用。然而不幸的是，由于 Join 算子目前独占了此 Actor 的执行权，本次唤醒并不能正确唤醒调用 get 的 poll函数，因此也就没能通知其它 wait Future 完成。在 Join 的 wait 看来，本次 Task 唤醒只是一次 Suspicous Wakeup，因此依然返回 Pending，但从此导致当前 Actor 永远失去了下一次被唤醒调度的机会。

定位到问题所在后，修复它就变得相当容易：出现此 Async Stuck 的根本原因在于，Block Cache 的 Single Flight 实现假设了所有的请求在调度执行上都拥有平等的地位，而 Actor 内部的 Executor 嵌套调用违背了这一原则。因此，我们只需将 Block Cache 的请求和通知唤醒过程通过 Spawn 一个新的 Task 来完成，使其永远能够被独立调度，该问题就迎刃而解了。

案例 2：环形资源依赖

除了单个 Task 内部的 Stuck 问题之外，得益于允许嵌入运行时信息，多个 Task 的 Await-Tree 也可以被轻松地关联起来，进而调试多个 Task 形成的依赖死锁等问题。

流计算场景的负载多变，为此，RisingWave 拥有独特的线上扩缩容能力，使得计算任务可以同时具有可控的延迟和极高的资源利用率。然而，在面对突发的负载提升时，RisingWave 会首先通过 Actor 之间的反压（Back Pressure）机制保证计算任务能够稳定运行。

RisingWave 的 Actor 之间通过本地 Channel 或远程 gRPC Streaming 通信。随着 Actor 并行程度的提升，为了避免远程通信占用过多端口 fd 等系统资源，我们引入了复用 gRPC 连接的优化。然而在引入该优化后，我们在极端测试中发现：当 Actor 并行度达到单进程 100+ 后，高负载下很容易出现整个 Streaming Job 卡住而无法继续执行的问题。基于此前解决 Async Stuck 的经验，我们选择直接导出所有的 Actor 的 Await-Tree，并发现了这样的现象：

在导出的 Await-Tree 之间，我们发现不同 Actor 上报的阻塞原因之间形成了一个环：下游 Actor 认为上游 Actor 没有产出数据，导致无法继续处理；而上游 Actor 又认为下游 Actor 没能及时消费数据导致了反压，进而无法继续处理------这对于 DAG 执行图而言是不可能出现的情况。通过对比观察，我们将问题锁定在了 Actor 远程通信的 gRPC Streaming 上；在进一步研究 gRPC 库 tonic 的内部实现后，我们发现它对于复用的 gRPC 连接设置有总 Window Size，由不同的 gRPC Stream 共享：一旦其中某个 Stream 占用的 Window Size 过大，可能导致其它随机 Stream 被反压，这无疑违背了 DAG 执行图中反压按照拓扑排序传播的原则。

不难看出，该 Async Stuck 问题的原因和案例 1 如出一辙：即都在本身具有依赖关系的资源上引入了错误的平等假设并形成了反向的依赖，进而导致了环形资源依赖，即死锁的发生。在定位到问题所在后，我们选择取消了 gRPC 在连接层面上的拥塞控制，但保留每个 gRPC Stream 独立进行拥塞控制的能力，从而在保留 gRPC 反压机制的同时，避免不同 Actor 之间的通信互相干扰。

总结

在本文中，我们介绍了 Await-Tree 这一 Async Rust 可观测性的"灵丹妙药"：Await-Tree 是一个原生为 Async Rust 设计的 Backtrace 工具，允许开发者实时观测各个 Async Task 的执行状态，并直观分析不同 Future 或 Task 间的依赖阻塞关系。在 RisingWave 中，Await-Tree 已长期于生产环境启用，并数次帮助开发者们解决了棘手的 Async Stuck 问题。

Await-Tree 已于 GitHub 开源：github.com/risingwavel...，欢迎各位读者将 Await-Tree 集成在自己的 Async Rust 系统中。

关于 RisingWave

RisingWave是一款分布式 SQL 流处理数据库，旨在帮助用户降低实时应用的的开发成本。作为专为云上分布式流处理而设计的系统，RisingWave 为用户提供了与 PostgreSQL 类似的使用体验，并且具备比 Flink 高出 10 倍的性能以及更低的成本。了解更多：

GitHub: risingwave.com/github

官网: risingwave.com

**微信公众号：**RisingWave中文开源社区