Tokio 异步服务优雅停机：取消、排空与任务收尾

线上进程收到 SIGTERM 时，如果立即退出，正在处理的请求或队列任务可能被中断；如果永远等待所有任务完成，停机又可能卡住。优雅停机需要明确顺序：停止接收新工作、通知任务收尾、等待有限时间，最后再强制退出。

下面用 Tokio 有界 mpsc 队列演示一个 worker 的生命周期。它适用于理解控制流，不是持久队列：进程被强杀时，内存里的消息仍会丢失。

关闭流程

sequenceDiagram
    participant OS as SIGTERM / Ctrl-C
    participant Main as Main task
    participant W as Worker
    participant Q as Bounded channel
    OS->>Main: signal received
    Main->>W: CancellationToken.cancel()
    W->>Q: close receiver
    W->>Q: drain buffered jobs
    W-->>Main: task completed
    Main->>Main: wait until deadline
    Main->>W: abort only after timeout

Tokio 的官方指南把停机拆成三个阶段：发现停机信号、通知程序各部分、等待它们完成。CancellationToken 适合把同一个关闭信号传给多个任务；JoinHandle 让主任务能观察子任务是否已结束。

Worker：停止接收后排空缓冲区

use tokio::sync::mpsc;
use tokio_util::sync::CancellationToken;

async fn process(job: String) {
    // 实际服务在这里调用领域逻辑，并在成功后确认任务。
    println!("processing {job}");
}

async fn worker(mut rx: mpsc::Receiver<String>, shutdown: CancellationToken) {
    loop {
        tokio::select! {
            biased;
            // 先检查关闭信号，队列持续有消息时也优先进入收尾。
            _ = shutdown.cancelled() => {
                rx.close(); // 拒绝后续 send，但允许读取已缓冲消息
                while let Some(job) = rx.recv().await {
                    process(job).await;
                }
                break;
            }
            job = rx.recv() => match job {
                Some(job) => process(job).await,
                None => break,
            }
        }
    }
}

这里在收到关闭信号后，先 close() 再把已缓冲消息处理完。Tokio 的 Receiver::recv 是 cancellation-safe 的：如果 select! 的其他分支先完成，不会悄悄取走一条消息。

有界 channel 的容量是背压的一部分。生产者执行 send().await 时，队列满了会等待；不要为了追求吞吐量改成无界队列，否则突发流量可能把进程内存耗尽。

主任务：设置等待上限

以下示例可放入 Tokio 二进制项目。Cargo 依赖需要启用 Tokio 的 macros、rt-multi-thread、signal、sync、time 功能，以及 tokio-util 的 rt 功能。

use std::time::Duration;
use tokio::{signal, sync::mpsc, time::timeout};
use tokio_util::sync::CancellationToken;

async fn shutdown_signal() -> std::io::Result<()> {
    #[cfg(unix)]
    {
        let mut terminate = signal::unix::signal(signal::unix::SignalKind::terminate())?;
        tokio::select! {
            result = signal::ctrl_c() => result,
            _ = terminate.recv() => Ok(()),
        }
    }
    #[cfg(not(unix))]
    {
        signal::ctrl_c().await
    }
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let (_tx, rx) = mpsc::channel::<String>(128);
    let shutdown = CancellationToken::new();
    let worker_shutdown = shutdown.clone();
    let mut worker_task = tokio::spawn(worker(rx, worker_shutdown));

    tokio::select! {
        result = shutdown_signal() => result?,
        result = &mut worker_task => {
            result?;
            return Ok(());
        }
    }

    shutdown.cancel();
    match timeout(Duration::from_secs(20), &mut worker_task).await {
        Ok(result) => result?,
        Err(_) => worker_task.abort(),
    }
    Ok(())
}

在真实服务中，_tx 会由请求处理器或任务生产者持有。收到关闭信号后，还应先让 HTTP server 停止接受新连接，再通知后台 worker。超时后的 abort() 是兜底动作；需要记录哪些任务未完成，便于告警与恢复。

常见踩坑

只发取消信号，不等待任务。 子任务可能还没完成 flush、ack 或状态回写，主程序就退出了。
把取消当成事务回滚。 任务在外部服务已产生副作用后被取消，不能自动撤销；依赖幂等键或补偿操作。
没有关闭队列接收端。 关闭后仍允许生产者继续入队，排空过程就没有清晰终点。
对所有任务一刀切立即 abort。 网络读写、锁获取和自定义 future 的取消安全性不同；检查每个 .await 被中断后的状态。
把内存 channel 当作可靠消息系统。 进程崩溃或强制超时会丢失未持久化的消息。需要跨重启恢复时，使用持久队列或数据库任务表，并在处理成功后 ack。
忽略运行环境的信号约定。 容器编排通常发送 SIGTERM 并有固定宽限期，服务内超时必须短于平台的最终强杀时间。

实践验证

先让 process 执行一个可观测的延迟工作，启动服务后发送 Ctrl-C，确认新任务停止进入、队列中的任务完成、主任务在截止时间内退出。再模拟一个永不结束的任务，确认超时会触发告警并走强制收尾路径。不要仅凭本地 Ctrl-C 测试就宣称生产停机可靠。

关闭流程

Worker：停止接收后排空缓冲区

主任务：设置等待上限

常见踩坑

实践验证

参考资料

相关文章