Tokio 在同步上下文中执行异步代码-526互联

从 `spawn` 说起

Tokio 库中有两个同名的量, 它们都叫 spawn, 但是却有着显著的区别:

其中一个是 tokio::runtime::Runtime 结构体的方法 (method), 另一个是 tokio::task 模块的一个函数, 同时也是你使用 tokio::spawn 时直接使用的那个. 从这个特征来看, 两者使用的方法是截然不同的, 但背后的原理却有相似之处.

显式绑定运行时的 `spawn` 方法

首先看 Runtime 结构体的方法 Runtime::spawn. 这个方法的目的是将一个实现了 Send trait 的 Future 对象送入给定的运行时中 (常常是一个线程池), 然后由对应的运行时执行器反复 poll 这个 Future 直到它执行完毕. 被送到执行器的 Future 会被立刻开始执行, 直到执行结束, 或者首次遇到 .await. 如下面的例子:

use std::time::Duration;

fn main() {
    let rt = tokio::runtime::Runtime::new().unwrap();

    rt.spawn(async move {
        println!("Printing in a future (L1).");
        tokio::time::sleep(Duration::from_secs(1)).await;
        println!("Printing in a future (L2).");
    });

    println!("Runtime terminated.");
}

我们先通过 Runtime::new() 创建了一个新的运行时, 随后 rt.spawn 会使用传入的 Future 创建新的异步任务并开始执行. 但是, 取决于操作系统对于多个线程的调度方式不同, 有可能在程序输出 Printing in a future (L1). 之前就已经结束了, 也有可能分别输出了:

Printing in a future (L1).
Runtime terminated.

但是一定不会输出的是 Printing in a future (L2)., 因为当这个 Future 执行遇到 .await 时, 它的执行就被暂停了. 在等待这 1 秒钟的过程中, 程序一定已经结束执行, 因此第二个语句不会被执行.

绑定当前运行时的 `spawn` 函数

从 tokio::task::spawn 或者 tokio::spawn 函数的调用逻辑来看, 当我们调用它的时候, 会执行:

tokio::runtime::Handle::current();

这个方法被用来获取当前上下文中运行时的 handle, 从而向对应的 scheduler 中添加新的 Task. 听起来非常简单, 我们把上例中的运行时构造步骤进行更改, 并将 rt.spawn 替换为函数调用 tokio::spawn:

use std::time::Duration;

fn main() {
    let rt = tokio::runtime::Builder::new_current_thread().enable_all().build().unwrap();

    etokio::spawn(async move {
        println!("Printing in a future (L1).");
        tokio::time::sleep(Duration::from_secs(1)).await;
        println!("Printing in a future (L2).");
    });

    println!("Runtime terminated.");
}

执行这段代码, 出现了这样的运行时错误:

thread 'main' panicked at 'there is no reactor running, must be called from the context of a Tokio 1.x runtime', src/main.rs:8:5
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace

查看堆栈跟踪信息, 可以发现问题就出在调用 tokio::runtime::scheduler::Handle::current 的过程中. 这段错误信息也说明了 Tokio 在尝试将 Future 添加到某个 scheduler 时出现了上下文中找不到带有 Tokio 运行时的问题. 所以, 我们需要一种方法将当前的线程和新创建的 Tokio 关联起来. Tokio 为每一个 Runtime 实例都提供了一个 enter 方法. 这个方法内部会尝试将本线程的静态 Tokio Context 实例关联到本当前 Runtime 之上:

let _enter_guard = rt.enter();

需要注意的是, 我们不能直接使用 _ 作为变量名. 使用 _ 相当于 rt.enter() 的返回值直接丢弃了, 等价于这个用法:

let _enter_guard = rt.enter();
drop(_enter_guard);

因此, 后续代码执行时仍然无法通过 Handle::current 捕获当前的运行时 handle, 故同样的错误还会出现.

再聊聊 `block_on`

刚才我们聊到了 Tokio 中 Runtime 的 spawn 方法可以用来在绑定的运行时 scheduler 上产生新的 Future, 并立即开始执行直到阻塞或结束. 但其实 Runtime 上还有另外一个方法也非常重要, 那就是 block_on.

block_on 方法有个重要的特征, 就是当它开始执行时, 一定会阻塞当前的线程, 直到执行完毕后退出. 我们回顾刚才的例子:

use std::time::Duration;

fn main() {
    let rt = tokio::runtime::Runtime::new().unwrap();
    let _enter_guard = rt.enter();

    tokio::spawn(async move {
        println!("Printing in a future (L1).");
        tokio::time::sleep(Duration::from_secs(1)).await;
        println!("Printing in a future (L2).");
    });

    println!("Runtime terminated.");
}

如果想让 Future 中两个输出语句 (L1 和 L2) 都能够输出, 显然程序不可以在遇到 .await 后结束执行. 由于我们使用了多线程的运行时, 最显而易见的方法就是阻塞主线程, 避免它意外退出. 最简单直接的方法就是使用 rt.block_on 方法:

let rt = tokio::runtime::Runtime::new().unwrap();

rt.block_on(async move {
	println!("Printing in a future (L1).");
	tokio::time::sleep(Duration::from_secs(1)).await;
	println!("Printing in a future (L2).");
});

println!("Runtime terminated.");

执行上述程序, 我们得到了正确的输出结果:

Printing in a future (L1).
Printing in a future (L2).
Runtime terminated.

但是怎样证明 rt.block_on 方法阻塞了当前线程的同时, 还有其他的线程在执行任务呢?

use std::time::Duration;

fn main() {
    println!("Main thread: {:?}", std::thread::current().id());

    let rt = tokio::runtime::Runtime::new().unwrap();

    rt.spawn(async move {
        println!(
            "[{:?}] Printing in a future (L0).",
            std::thread::current().id()
        );
    });

    rt.block_on(async move {
        println!(
            "[{:?}] Printing in a future (L1).",
            std::thread::current().id()
        );
        tokio::time::sleep(Duration::from_secs(1)).await;
        println!(
            "[{:?}] Printing in a future (L2).",
            std::thread::current().id()
        );
    });

    println!("Runtime terminated.");
}

我们在执行 rt.block_on 来阻塞当前线程之前, 先调用 rt.spawn 方法新增了一个异步任务. 其次, 在主线程, block_on 的 Future 和 spawn 的 Future 里面, 我们分别输出当前的线程号 (ThreadId). 执行结果如下:

Main thread: ThreadId(1)
[ThreadId(1)] Printing in a future (L1).
[ThreadId(7)] Printing in a future (L0).
[ThreadId(1)] Printing in a future (L2).
Runtime terminated.

可以看到, block_on 的线程 (1) 就是主线程 (1), 而 spawn 的 Future 在完全不同的另一个线程 (7) 中执行. 反复执行多次, 可以看到主线程始终不变 (1), 但 spawn 线程却每次都发生变化. 这个实验同时说明了以下几个问题:

Runtime::new() 方法会创建多线程 scheduler 来执行代码;
Runtime::block_on() 方法会阻塞所在的线程执行;
Runtime::spawn() 方法会在以类似线程池的方式来创建异步任务, 然后由不同的线程来获取执行.

除了 Runtime 上的 block_on 方法, 还有另外两种: Handle::block_on 和 tokio::task::LocalSet::block_on. 它们的区别又是什么呢?

`Handle::block_on` 方法

Handle::block_on 与 Runtime::block_on 的差异与当前的运行时类型密切相关. Runtime 的类型包括两种: 当前线程执行 或 多线程执行. 其中, 多线程执行方式下, Handle::block_on 与 Runtime::block_on 没有差别, 但是在当前线程执行时, 有着巨大的差异, 并且直接决定了我们的程序能否正常执行.

当前线程执行方式下, Future 中涉及 IO 相关或是计时器相关的阻塞都不会在 Handle::block_on 中执行, 它会一直阻塞永不退出. 相比之下, Runtime::block_on 则没有任何限制, 可以正确阻塞线程, 也可以完成 IO 和计时器相关的阻塞操作. 因此, 如果可能, 尽量避免使用 Handle::block_on. 比如下面的这个例子就很好说明了这个问题:

use std::time::Duration;

macro_rules! tid {
    () => {
        std::thread::current().id()
    };
}

fn main() {
    println!("Main thread: {:?}", std::thread::current().id());

    let rt = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .unwrap();
    let rt = tokio::runtime::Runtime::new().unwrap();

    rt.spawn(async move {
        println!("[{:?}] From future (1).", tid!());
        tokio::time::sleep(Duration::from_secs(1)).await;
        println!("[{:?}] From future (1), after sleep 1 second.", tid!());
    });

    // rt.block_on(async move {
    //     println!("[{:?}] From future (2)", tid!());
    //     tokio::time::sleep(Duration::from_secs(1)).await;
    //     println!("[{:?}] From future (2), after sleep 1 second.", tid!());
    // });
    
    rt.handle().block_on(async move {
        println!("[{:?}] From future (2)", tid!());
        tokio::time::sleep(Duration::from_secs(1)).await;
        println!("[{:?}] From future (2), after sleep 1 second.", tid!());
    });

    println!("Runtime terminated.");
}

通过分别注释 rt.handle().block_on 和 rt.block_on 的部分, 可以明确观察到上述现象.

`LocalSet::block_on` 方法

聊到这个方法, 就不能不谈多线程方式下的运行时了. 早前我们探究 spawn 时已经证实, Tokio 确实会采用线程池来执行这些任务. 但是有的情况下, 我们的任务没有实现 Send trait, 也就是不允许安全地在不同线程之间转移. 所以为了解决这个问题, 就需要使用 LocalSet::block_on 方法了, 参见文档.