Rust Web：Axum 优雅停机（Graceful Shutdown）

Table of contents

为什么需要优雅停机

暴力停机（kill -9 / 直接断电）:
  正在处理的请求 → 连接直接断开，客户端收到错误
  正在写数据库的事务 → 可能中断，数据不一致
  正在刷盘的日志 → 丢失
  客户端重试 → 可能造成重复操作

优雅停机:
  1. 收到停机信号
  2. 停止接受新连接
  3. 等待正在处理的请求完成
  4. 清理资源（关闭数据库连接池、刷日志、导出最后的 metrics）
  5. 退出

基本实现

use axum::{routing::get, Router};
use tokio::signal;

#[tokio::main]
async fn main() {
    let app = Router::new()
        .route("/health", get(|| async { "OK" }));

    let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await.unwrap();
    println!("listening on :3000");

    axum::serve(listener, app)
        .with_graceful_shutdown(shutdown_signal())  // 关键：注册停机信号
        .await
        .unwrap();

    println!("server shut down");
}

async fn shutdown_signal() {
    // Ctrl+C（SIGINT）
    let ctrl_c = async {
        signal::ctrl_c()
            .await
            .expect("failed to install Ctrl+C handler");
    };

    // SIGTERM（Docker / k8s 停止容器时发送的信号）
    #[cfg(unix)]
    let terminate = async {
        signal::unix::signal(signal::unix::SignalKind::terminate())
            .expect("failed to install SIGTERM handler")
            .recv()
            .await;
    };

    #[cfg(not(unix))]
    let terminate = std::future::pending::<()>();

    tokio::select! {
        _ = ctrl_c => println!("received Ctrl+C"),
        _ = terminate => println!("received SIGTERM"),
    }
}

with_graceful_shutdown 接收一个 Future，这个 Future 完成时 axum 开始停机：

1. 停止 accept 新连接
2. 等所有正在处理的请求完成
3. axum::serve(...).await 返回，继续执行后面的清理代码

加上资源清理

#[tokio::main]
async fn main() {
    // 初始化资源
    let pool = PgPool::connect("postgres://...").await.unwrap();
    let redis = redis::Client::open("redis://...").unwrap();

    let state = AppState {
        db: pool.clone(),
        redis: redis.clone(),
    };

    let app = Router::new()
        .route("/health", get(health))
        .route("/players", get(list_players))
        .with_state(state);

    let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await.unwrap();
    println!("listening on :3000");

    // 启动服务（收到信号后这里会返回）
    axum::serve(listener, app)
        .with_graceful_shutdown(shutdown_signal())
        .await
        .unwrap();

    // ===== 到这里说明信号已收到，所有请求已处理完毕 =====

    // 清理资源
    println!("closing database connections...");
    pool.close().await;

    println!("flushing traces...");
    // opentelemetry_sdk::trace::TracerProvider::shutdown()

    println!("server shut down cleanly");
}

加上超时保护

有些请求可能很慢或卡住了，不能无限等待：

use std::time::Duration;
use tokio::time::timeout;

#[tokio::main]
async fn main() {
    let app = build_app();
    let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await.unwrap();

    let server = axum::serve(listener, app)
        .with_graceful_shutdown(shutdown_signal());

    // 给优雅停机一个最大等待时间
    match timeout(Duration::from_secs(30), server).await {
        Ok(_) => println!("all requests completed, shutting down"),
        Err(_) => println!("shutdown timed out after 30s, forcing exit"),
    }

    cleanup().await;
}

在 Docker / Kubernetes 中的行为

Docker:
  docker stop <container>
    → Docker 发 SIGTERM 给容器 PID 1
    → 等待 10 秒（可配置 --stop-timeout）
    → 如果还没退出，发 SIGKILL 强杀

  所以你的优雅停机必须在 Docker 的 stop-timeout 内完成

Kubernetes:
  Pod 终止流程:
    1. K8s 发 SIGTERM
    2. 从 Service 负载均衡中摘除（不再转发新请求）
    3. 等待 terminationGracePeriodSeconds（默认 30 秒）
    4. 如果还没退出，发 SIGKILL

  Dockerfile:
    # 确保你的程序是 PID 1，能直接收到 SIGTERM
    CMD ["mini_tarkov_server"]
    # 不要用 shell 包一层：CMD ["sh", "-c", "mini_tarkov_server"]
    # 否则 SIGTERM 发给 sh 而不是你的程序

完整示例

use axum::{routing::get, Router, extract::State};
use sqlx::PgPool;
use std::time::Duration;
use tokio::{signal, time::timeout};
use tracing::info;

#[derive(Clone)]
struct AppState {
    db: PgPool,
}

async fn health() -> &'static str { "OK" }

async fn list_players(State(state): State<AppState>) -> String {
    // 模拟慢查询
    tokio::time::sleep(Duration::from_secs(2)).await;
    "[]".to_string()
}

async fn shutdown_signal() {
    let ctrl_c = async {
        signal::ctrl_c().await.expect("failed to install Ctrl+C handler");
    };

    #[cfg(unix)]
    let terminate = async {
        signal::unix::signal(signal::unix::SignalKind::terminate())
            .expect("failed to install SIGTERM handler")
            .recv()
            .await;
    };

    #[cfg(not(unix))]
    let terminate = std::future::pending::<()>();

    tokio::select! {
        _ = ctrl_c => {},
        _ = terminate => {},
    }

    info!("shutdown signal received, waiting for requests to complete...");
}

#[tokio::main]
async fn main() {
    tracing_subscriber::fmt::init();

    let pool = PgPool::connect("postgres://tarkov:tarkov123@localhost/mini_tarkov")
        .await
        .expect("failed to connect to database");

    let state = AppState { db: pool.clone() };

    let app = Router::new()
        .route("/health", get(health))
        .route("/players", get(list_players))
        .with_state(state);

    let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await.unwrap();
    info!("listening on :3000");

    let server = axum::serve(listener, app)
        .with_graceful_shutdown(shutdown_signal());

    match timeout(Duration::from_secs(30), server).await {
        Ok(_) => info!("all requests completed"),
        Err(_) => info!("shutdown timed out, forcing exit"),
    }

    info!("closing database pool...");
    pool.close().await;

    info!("server shut down cleanly");
}

停机流程图

                  运行中
                    │
       收到 SIGTERM / Ctrl+C
                    │
                    ▼
          shutdown_signal() 完成
                    │
                    ▼
    ┌───────────────────────────────┐
    │  axum 停止 accept 新连接       │
    │  正在处理的请求继续执行          │
    └───────────┬───────────────────┘
                │
        ┌───────┴────────┐
        │                │
   所有请求完成      30 秒超时
        │                │
        └───────┬────────┘
                │
                ▼
    ┌────────────────────────┐
    │  清理资源               │
    │  - pool.close()        │
    │  - tracer.shutdown()   │
    │  - flush logs          │
    └───────────┬────────────┘
                │
                ▼
           进程退出