SIGTERM/Ctrl+Cによるグレースフルシャットダウンと、2回目のシグナルが効かなくなる罠

プロセスをいきなり終了させると、転送中のコネクションが強制的に切断されます。SIGTERM（あるいはローカルでのCtrl+C）を受け取ったら、新規コネクションの受け付けだけを止め、すでに転送中のコネクションは自然に終わるまで待つ、グレースフルシャットダウンを実装します。

JoinSetでコネクションを追跡する

これまで各コネクションはtokio::spawnで個別にスポーンし、そのハンドルは捨てていました。シャットダウン時に「今どれだけのコネクションが転送中か」「それらが全部終わるのを待つ」を行うには、スポーンしたタスクをまとめて追跡する必要があります。tokio::task::JoinSetに置き換えます。

// src/main.rs（抜粋）
let mut tasks = JoinSet::new();

// tokio::spawn(async move { ... }) だった箇所を
tasks.spawn(async move {
    if let Err(e) = handle_connection(inbound, peer_addr, lb, backends, pools, FAILURE_THRESHOLD).await {
        eprintln!("connection error ({peer_addr}): {e}");
    }
});

acceptとシグナル待ちをtokio::select!で両立する

acceptループは、これまでlistener.accept().awaitだけを待っていました。ここにシグナル待ちを追加します。tokio::select!は複数のFutureを同時に待ち、最初に完了したものの分岐だけを実行するマクロです。

// src/main.rs（抜粋）
let mut sigterm = signal(SignalKind::terminate())?;
let mut sigint = signal(SignalKind::interrupt())?;

loop {
    tokio::select! {
        accepted = listener.accept() => {
            let (inbound, peer_addr) = accepted?;
            // ... tasks.spawn(...) で処理
        }
        _ = sigterm.recv() => {
            println!("received SIGTERM, no longer accepting new connections");
            break;
        }
        _ = sigint.recv() => {
            println!("received Ctrl+C, no longer accepting new connections");
            break;
        }
    }
}

シグナルを受けたらbreakでループを抜け、以降はacceptを呼びません。SIGTERM（オーケストレータが送る想定）とCtrl+C（ローカル検証用）の両方を同じ扱いにしています。

2回目のシグナルが誰にも届かなくなる罠

ループを抜けたあと、JoinSet::join_next()がNoneを返すまで回せば全コネクションの完了を待てます。ただし、drain中もsigterm/sigintの.recv()を同じtokio::select!で待ち受け続ける必要があります。ここを外すと、drain中に送った2回目のSIGTERMやCtrl+Cはどこにも届かず、プロセスは一切反応しなくなります。

// src/main.rs（抜粋）
println!("draining {} in-flight connection(s)", tasks.len());
let drain = async {
    while tasks.join_next().await.is_some() {}
};
tokio::select! {
    _ = drain => {
        println!("all connections drained, shutdown complete");
    }
    _ = tokio::time::sleep(SHUTDOWN_DRAIN_TIMEOUT) => {
        eprintln!("drain timed out after {SHUTDOWN_DRAIN_TIMEOUT:?}, exiting with {} connection(s) still in flight", tasks.len());
    }
    _ = sigterm.recv() => {
        eprintln!("received a second SIGTERM during drain, forcing immediate shutdown");
    }
    _ = sigint.recv() => {
        eprintln!("received a second Ctrl+C during drain, forcing immediate shutdown");
    }
}

sigtermとsigintは、tokio::signal::unix::signalが返すSignal型のハンドルで、.recv()を何度でも呼び直せるように設計されています（tokio::signal::ctrl_c()のような一度きりのFutureではありません）。acceptループとdrainフェーズの両方で同じ変数をそのまま使い回せるのはこのためです。

備考: tokio::signal::unix::signal（やtokio::signal::ctrl_c()）を一度でも呼ぶと、OSのデフォルトのシグナル挙動（Ctrl+Cで即座にプロセス終了）が、tokio独自のハンドラにプロセス全体・恒久的に置き換わります。この置き換えは元に戻らないため、届いたシグナルをどこかの.recv()が確実に拾える状態にしておかないと、シグナルは行き場を失ってそのまま捨てられます。drainフェーズでsigterm/sigintを待ち受けから外していると、2回目のCtrl+Cが「本当に何も起こさない」状態になるのはこのためです。

drain・タイムアウト・2回目のシグナルのどれが最初に完了しても、tasks（JoinSet）は関数の終わりでスコープを抜けて破棄されます。JoinSetはドロップされる際に残っているタスクをすべて中断するため、2回目のシグナルを受けた場合も、タイムアウトした場合も、転送中のコネクションを打ち切って確実にプロセスを終了できます。