Tokio：优雅关闭

本文概述如何在异步应用中正确实现关闭流程。

实现优雅关闭通常包含三个部分：

确定何时关闭。
通知程序的各个部分开始关闭。
等待其他部分完成关闭。

下文依次介绍这三部分。mini-redis提供了实际实现，尤其可参考src/server.rs与src/shutdown.rs。

确定何时关闭

具体条件取决于应用，但一个很常见的触发条件是收到操作系统信号，例如程序运行时在终端按Ctrl+C。Tokio提供tokio::signal::ctrl_c，它会等待，直到收到该信号。用法如下：

use tokio::signal;

#[tokio::main]
async fn main() {
    // ... spawn application as separate task ...

    match signal::ctrl_c().await {
        Ok(()) => {},
        Err(err) => {
            eprintln!("Unable to listen for shutdown signal: {}", err);
            // we also shut down in case of error
        },
    }

    // send shutdown signal to application and wait
}

如果存在多个关闭条件，可以用mpsc通道把关闭信号集中到一处，再用select!同时等待ctrl_c与通道。例如：

use tokio::signal;
use tokio::sync::mpsc;

#[tokio::main]
async fn main() {
    let (shutdown_send, mut shutdown_recv) = mpsc::unbounded_channel();

    // ... spawn application as separate task ...
    //
    // application uses shutdown_send in case a shutdown was issued from inside
    // the application

    tokio::select! {
        _ = signal::ctrl_c() => {},
        _ = shutdown_recv.recv() => {},
    }

    // send shutdown signal to application and wait
}

通知任务关闭

要通知一个或多个任务关闭，可以使用取消令牌CancellationToken。令牌通知任务响应取消请求并自行终止，方便实现优雅关闭。

要在多个任务间共享CancellationToken，必须克隆它，因为单一所有权规则要求每个值只有一个所有者。克隆得到的令牌与原令牌等效；一个被取消，另一个也会被取消。可以按需克隆任意数量，调用其中任意一个的cancel时，它们都会被取消。

在多个任务中使用CancellationToken的步骤如下：

创建新的CancellationToken。
对原令牌调用clone，创建可交给另一个任务使用的令牌。
把原令牌或克隆令牌传给需要响应取消请求的任务。
需要优雅关闭时，对任意一个令牌调用cancel。所有等待原令牌或克隆令牌取消请求的任务都会收到关闭通知。

以下代码展示这些步骤：

// Step 1: Create a new CancellationToken
let token = CancellationToken::new();

// Step 2: Clone the token for use in another task
let cloned_token = token.clone();

// Task 1 - Wait for token cancellation or a long time
let task1_handle = tokio::spawn(async move {
    tokio::select! {
        // Step 3: Using cloned token to listen to cancellation requests
        _ = cloned_token.cancelled() => {
            // The token was cancelled, task can shut down
        }
        _ = tokio::time::sleep(std::time::Duration::from_secs(9999)) => {
            // Long work has completed
        }
    }
});

// Task 2 - Cancel the original token after a small delay
tokio::spawn(async move {
    tokio::time::sleep(std::time::Duration::from_millis(10)).await;

    // Step 4: Cancel the original or cloned token to notify other tasks about shutting down gracefully
    token.cancel();
});

// Wait for tasks to complete
task1_handle.await.unwrap()

使用取消令牌时，不必在令牌取消后立即终止任务。可以先执行清理流程，例如把数据刷入文件或数据库，或通过连接发送关闭消息，然后再结束任务。

等待关闭完成

通知其他任务关闭后，还需要等待它们完成。一个简单办法是使用TaskTracker。它是一组任务的集合；其wait方法返回一个future，只有集合中的所有future都完成，并且跟踪器已经关闭时，该future才会完成。

下面的示例启动10个任务，然后用任务跟踪器等待它们结束。

use std::time::Duration;
use tokio::time::sleep;
use tokio_util::task::TaskTracker;

#[tokio::main]
async fn main() {
    let tracker = TaskTracker::new();

    for i in 0..10 {
        tracker.spawn(some_operation(i));
    }

    // Once we spawned everything, we close the tracker.
    tracker.close();

    // Wait for everything to finish.
    tracker.wait().await;

    println!("This is printed after all of the tasks.");
}

async fn some_operation(i: u64) {
    sleep(Duration::from_millis(100 * i)).await;
    println!("Task {} shutting down.", i);
}