介绍

相比其他语言，Rust 的异步很有趣、很强大、很独特。而 Tokio 是 Rust 社区的一个非常流行的异步 runtime 实现，非常值得学习。本文会讲解 Tokio 的内部源码实现，以及一些精彩设计。

感谢 @aquarhead, @tyrchen 和 @ZhangHanDong 的建议。

Rust async 简介

在讲 Tokio 之前，不得不先讲一下 Rust 的异步编程，因为和很多语言不太一样。

对于 Rust Async 熟悉的可以跳过这章

对于 Erlang/Go/Nodejs 等语言，异步 runtime 都是内置于语言本身，开箱即用。但 Rust 作为一门系统级语言，并不想局限于一种实现，于是另辟蹊径，提供了 Future/async/await 等基本功能，实现了类似于 Nodejs Promise 的 task ，但把调度和运行 future 交给第三方实现，如 Tokio、async-std。例如下图，左边是 Rust 语言提供的基本功能，驱动了 task 的执行，而右边是第三方 runtime 需要实现的调度机制。

https://excalidraw.com/#json=5287000177377280,_EjA-elJg02sgC71T8uXLQ

Future, async, await

Future 是 Rust 的一个 trait（类似于 interface），表示一个异步任务（task），是"零成本"(zero-cost abstraction)的轻量级线程（类似 promise），会被交给 runtime 调度和执行。Future trait 需要实现 poll 这个方法，在 poll 中判断这个任务是否执行完，如果执行完（比如 IO 数据准备好），就返回 Ready，否则返回 Pending。

我们的代码不会直接调用 poll，而是通过 Rust 的关键字 .await 来执行这个 future，await 会被 Rust 在编译时生成代码来调用 poll，如果返回 Pending 则被 runtime 挂起（比如重新放到任务队列中）。当有 event 产生时，挂起的 future 会被唤醒，Rust 会再次调用 future 的 poll，如果此时返回 Ready 就执行完成。

除了直接实现 Future trait 以外，还可以通过 async 把一个 function 或者一个代码 block 转变为一个 Future。在 async 中可以调用其他 future 的 .await 来等待子 future 变成 Ready 状态。

struct HelloFuture { ready: bool, waker: ... }

impl Future for HelloFuture {          // 1. custom Future(leaf)
    fn poll(self: Self, ctx: &mut Context<'_>) -> Poll<()> {
        if self.ready {
            Poll::Ready(())
        } else {
            // store waker in ctx somewhere
            Poll::Pending
        }
    }
}

async fn hello_world() {                 // 2. generated Future by async
    println!("before await");
    HelloFuture { ready: false }.await;  // 3. HelloFuture is pending, then park
    println!("Hello, world!");
}

fn main() {
    let task = hello_world(); // 4. task is a generated Future
    // ... reactor code is ignored here, which will wake futures
    runtime::spawn(task);     // 5. task is run by runtime
}

例如上边这段代码，HelloFuture 是一个通过实现 Future trait 实现的 Future， hello_world 是 async 函数变成的 Future，并被传入 runtime 来执行（这里 runtime::spawn 只是示例）。hello_world 中又调用了 HelloFuture.await，因为 ready 是 false，所以 hello_world 会被挂起，直到 HelloFuture 被唤醒。

上述的这种嵌套的 futures 可以组成一个 Future 树，一般叶子节点都是由 runtime（如 Tokio）自己通过实现 Future trait 来实现的，如 io、tcp、time 等操作。非叶子节点则由库代码或用户通过 async 调用不同的子 future 实现的。root future 会被提交给 runtime 来执行，runtime 通过调度器来调用 root future，然后 root 再一级级往下调用 poll。

https://excalidraw.com/#json=6697978081312768,u9VYjoGonibMqcPX8VWeGg

生成状态机和 stackless

Future 会被 Rust 编译为一个状态机的代码，当执行到子 Future 的 await 的时候，会进入下一个状态，所以下次执行时可以从 await 的地方继续执行。因此 Rust 不需要预先为 future 分配独立的栈（stackless），是 zero-cost abstraction。但也因为如此，future 只能在 await 的地方调度走，是 cooperation scheduling（协同调度），而且很难做抢占式调度，这点和 stackful 的 Go/Erlang 不一样。状态机的示意伪代码如下：


#![allow(unused)]
fn main() {
use std::{future::Future, task::Poll};

enum HelloWorldState {
    Start,
    Await1(HelloFuture),
    Done,
}

impl Future for HelloWorldState {
    type Output = ();
    fn poll(&mut self: HelloWorldState, ctx: &mut Context<'_>) -> Poll<Output> {
        match self {
            HelloWorldState::Start => {
                println!("before await"); // code before await

                let hello = HelloFuture { ready: false };

                *self = HelloWorldState::Await1(hello);
                self.poll();              // re-poll after first state change
            },
            HelloWorldState::Await1(hello) => {
                match hello.poll(ctx) {      // await by poll
                    Poll::Pending => {
                        Poll::Pending
                    },
                    Poll::Ready(output) => {
                        println!("Hello, world!"); // code after await
                        *self = HelloWorldState::Done;
                        let output = ();
                        Poll::Ready(output)
                    }
                }
            },
            HelloWorldState::Done => {
                panic!("can't go here")
            }
        }
    }
}
}

在这段示意代码中， async fn hello_world() 被变成了一个 enum 的状态和它的 poll 方法，初始状态为 Start，第一次执行 poll 时会执行 .await 之前的代码，并改变当前状态为 Await1 。下次再被 poll 时，因为状态是 Await1，会进入第二个分支并执行 hello.poll()，如果 hello 还没完成，会返回 Pending，否则会执行 .await 之后的代码。

可以看到，Rust 主要提供了这些基础的工具和代码生成，而如何管理 OS 线程、如何调度任务、如何 poll events、如何唤醒 pending 的 tasks 等等，都需要 runtime 自己实现，可以实现为一个类似于 Erlang/Go 的 N:M 模型，也可以实现 Nodejs 这样单线程事件驱动的模型。

需要注意的是，如果使用了 async，很多标准库中的同步阻塞的库就不应该直接使用，否则会阻塞其他 task 的调度和运行。比如执行 std::println! 或者执行纯 CPU 计算的时候，因为没有 .await 这样的 yield point，所以无法被调度走，当前 OS 线程必须等到这个操作完成，执行到下一个 yield point 后，才能执行其他 task，这段时间内这个 OS 线程中的任务都被阻塞了。

一般 runtime 会封装好一些常用的模块，如 IO/TCP/timer 等等，写代码时注意使用 runtime 封装的而不用标准库的就行。像 Tokio 这样的 runtime 还提供了专门用来运行这类阻塞操作的专用线程池（类似于 Erlang 的 dirty scheduler），从而不影响 reactor 或者其他 task 的执行。

参考

Tokio 概览

初识 Tokio

Tokio 是一个 Rust 异步运行时库，底层基于 epoll/kqueue 这样的跨平台多路复用 IO 以及 event loop，目前正在支持 io_uring。它的 scheduler 和 Erlang/Go 实现的 N:M threads 类似，线程会执行 Task，可以充分利用多核。Task 是 Rust 基于 Future 抽象出的一种绿色线程，因为不需要预先分配多余的栈内存，可以创建大量 task，很适合做 IO 密集型应用。

虽然如 1.1 所说，Rust 本身不提供异步运行时，但因为 Rust 强大的 macro，我们可以非常方便地使用如 Tokio 这类第三方的 runtime。

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    let listener = TcpListener::bind("127.0.0.1:8080").await?;  // listen

    loop {
        let (mut socket, _) = listener.accept().await?; // async wait for incoming tcp socket

        tokio::spawn(async move {                       // create async task and let Tokio process it
            let mut buf = vec![0; 1024];

            loop {                                      // read and write data back until EOF
                let n = socket.read(&mut buf).await?;   // async wait for incoming data

                if n == 0 { return; }

                socket.write_all(&buf[0..n]).await?;    // async wait socket is ready to write and write data
            }
        });
    }
}

如上边这段代码，在 Tokio 中可以很容易地写一个常见的 TCP server，主线程 listen 端口，并在循环中接受连接，每一个连接的处理都在一个 Future 中完成，当等待 IO 时，这个 future 会让出 CPU 给其他 future，于是我们就有了一个高性能、高并发的 TCP server。之后代码解读也会以这段代码为示例。

架构概览

这里的 "magic" 就在 #[tokio::main] 这个宏，它会把代码预处理成这样：

fn main() {
    tokio::runtime::Builder::new_multi_thread().enable_all()
        .build().unwrap()
        .block_on(async {
            // async main
        })
}

程序启动后，在 build 中会初始化各种需要的数据、IO 资源，以及启动 worker 线程，然后在主线程中运行 async 代码 block，也就是我们自己写的那个 async main。

tokio overview link

上图大致描绘了以 echo 这个 example 为例的一个 Tokio runtime 的大致架构，对于理解 Tokio 很有帮助，之后的讲解也会再次提及。

图右的 worker 线程数量一般和核数相同，会执行 tokio::spawn 提交的 futures，当没有可执行的 task 时，会通过 epoll/kqueue 来 poll events，这部分工作由 reactor 负责。当被 events 唤醒后，会继续尝试执行 tasks，并这样循环下去。

在图的左边，runtime 在主线程 block_on 中会 poll 我们的 main 函数 future，执行到 listener.accept().await? 会返回 Pending，于是主线程会被挂起（park），在这里是等待信号量并休眠。接下来是三个事件：

（wake1）当收到 TCP 连接时，worker thread 会在 poll events 中拿到 events，并发送信号量给主线程。主线程会从 park 中被唤醒，然后执行 tokio::spawn。worker 线程会继续循环，也就是 poll events。
（wake2）主线程在 tokio::spawn 中会先把 TCP 连接的 future 放到 run queue 中，然后唤醒 worker thread，再回到等待 TCP accept 中。worker 线程被唤醒后会从 run queue 中取出 task 并执行，也就是 let mut buf = vec![0; 1024]; 开始的那段代码。

当 worker 线程执行 socket.read(&mut buf).await? 时，因为还数据还没有准备好，不能 read，会返回 Pending，再执行其他 tasks 或者等到 IO events。
（wake3）当 OS 收到 TCP 数据时，worker 线程会收到 events，并把之前未执行完的 task 放到 run queue 中，然后从 run queue 中取出并执行，这里会调用 syscall read 读取收到的数据，最后把数据写回 client。随后如果客户端关闭了连接，这个 task 就执行结束，worker 线程就会执行其他 task 或者等待 events。

这里有两个值得注意的地方。第一， tasks run queue 是有多个的，包括每个 worker 自己的 queue 和 global queue，worker 会优先从自己 queue 中取 task，在 3.2 中会详细讲解。另一个是，多个 worker 线程会并发地执行 tasks，但只有一个 worker 线程会作为 reactor 来 poll events，新的 events 可能是由 reactor 自己，也可能是由其他 worker 线程来执行对应的 future。

代码目录和结构

从 Cargo.toml 可以看出 Tokio 是一个包含了多个子 package 的 workspace，主要包括 tokio, tokio-macros, tokio-stream, tokio-util 和测试、examples 等其他代码。tokio-stream 是 Stream 的实现，tokio-util 是给 Tokio 的使用者用的，我们先暂时不管，所以主要的代码在 tokio 和 tokio-macros 两个子 package，我们看下它们的代码量：


#![allow(unused)]
fn main() {
-------------------------------------------------------------------------------
Language                     files          blank        comment           code
-------------------------------------------------------------------------------
tokio/src
Rust                           256           6332          23616          26233

tokio-macros/src
Rust                             3             50            266            410
}

代码主要集中在 tokio 中，另外还有差不多数量的 comment，可见文档和注释写的非常多。一共有2万多行代码，不算很多，也包含了不少测试代码，还有部分代码是用来实现一些标准库对应的异步版本，所以实际上我们一开始需要关注的代码量没有很多。

另外，还有一些主要的依赖：

bytes: 处理 bytes 的工具
mio: 封装了跨平台的 IO 操作，比如 epoll, kqueue 等
parking_lot: 实现了很多同步原语，如锁、信号量

我们看一下 tokio 的子模块：

.
# core
├── lib.rs        // library file
├── blocking.rs   // 提供 blocking 操作的封装
├── coop.rs       // 帮助实现更好的协同式调度
├── future        // future 操作的一些封装
├── park          // 类似于 std::thread::park，但更加通用
├── runtime       // Tokio runtime 的核心，包括 event loop，任务管理、调度，线程池等等
├── sync          // 让不同 task 用来进行同步的工具，如 channel 和 Mutex
├── task          // 上文介绍的 task

# async std in Tokio
├── io            // IO 操作的封装，相当于异步的 std::io，也是构建 net、fs 等子模块的基础
├── net           // TCP/UDP/Unix 的封装，类似于 std::net
├── fs            // 异步的 std::fs
├── process       // 异步进程管理，比如可以异步地运行一个子进程，类似于 std::process
├── signal        // 异步的信号处理，如 ctrl-c
├── time          // 时间相关的模块，如 Sleep

# utils
├── loom          // 统一了 std 和 loom(github.com/tokio-rs/loom) 的接口来方便测试
├── macros        // 一些公用的 macro，主要是声明宏。而 tokio-macro 主要是过程宏
└── util          // tokio 内部代码通用的工具模块

一开始我们主要关注 core 部分的模块、echo example 涉及的 io 和 net 模块，以及代码涉及到的 macros 和 util 模块。

另外，Tokio 里用到一些 feature flag 来允许定制化一些功能，如 rt-multi-thread 会开启多线程调度器， full 会开启几乎所有 feature，实际当中也可以根据需要关闭一些。

Tokio runtime 源码解读

在这部分，我们以 echo example 为例详细看一下 Tokio 的源码，会涉及到 Tokio runtime 的主要逻辑，熟悉了之后对 Tokio 乃至 Rust Async 就比较清楚了，再去看其他代码也会容易许多。

代码基于 tag tokio-1.5.0，且为了方便起见关闭了 time、process 和 signal 三个 feature，后边在 3.3 再开启。

Tokio runtime 启动

这是 echo example 的代码：

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    let listener = TcpListener::bind("127.0.0.1:8080").await?;  // listen

    loop {
        let (mut socket, _) = listener.accept().await?; // async wait for incoming tcp socket

        tokio::spawn(async move {                       // create async task and let Tokio process it
            let mut buf = vec![0; 1024];

            loop {                                      // read and write data back until EOF
                let n = socket.read(&mut buf).await?;   // async wait for incoming data

                if n == 0 { return; }

                socket.write_all(&buf[0..n]).await?;    // async wait socket is ready to write and write data
            }
        });
    }
}

和普通的同步代码不同，Tokio 需要我们写一个 async 的 main 函数，它主要是靠 #[tokio::main] 宏来生成代码，文档已经写得很清楚了，这里就不再赘述，只要知道它会被编译成下边这样就行。其实我们也可以根据需要在自己的 main 函数中，调用 API 来完成 runtime 初始化，而不通过 Tokio 的默认 macro。

fn main() {
    tokio::runtime::Builder::new_multi_thread()
        .enable_all()
        .build()
        .unwrap()
        .block_on(async {
            let listener = TcpListener::bind("127.0.0.1:8080").await?;
            // ...
        })
}

Runtime 初始化

这是 build() 方法（做了简化，以后的代码示例也会做适当的简化）：


#![allow(unused)]
fn main() {
let (driver, resources) = driver::Driver::new(self.get_cfg())?;

let (scheduler, launch) = ThreadPool::new(core_threads, Parker::new(driver));
let spawner = Spawner::ThreadPool(scheduler.spawner().clone());

// Create the blocking pool
let blocking_pool = blocking::create_blocking_pool(self, self.max_blocking_threads + core_threads);
let blocking_spawner = blocking_pool.spawner().clone();

// Create the runtime handle
let handle = Handle {
    spawner,
    io_handle: resources.io_handle,
    blocking_spawner,
};

// Spawn the thread pool workers
let _enter = crate::runtime::context::enter(handle.clone());
launch.launch();

Ok(Runtime {
    kind: Kind::ThreadPool(scheduler),
    handle,
    blocking_pool,
})
}

Tokio Internals