Day 044 — watchチャンネルでgraceful shutdownを実装する。select!の片方の腕が「二度と来ない」と分かった後も待ち続けるとビジーループになる

📚 背景知識（読んでから問題へ）

Day 042 で Semaphore による最大並行度の制御を、Day 043 で mpsc による進捗通知を実装しました。今日はこのテーマの最後のピース——外部からのシャットダウン信号でバッチ処理を安全に打ち切る——に取り組みます。

本番運用では「全タスクが終わるまで待つ」だけでは不十分な場面があります。Kubernetes は Pod を止める前に SIGTERM を送り、猶予期間内にプロセスが自発的に終了することを期待します。管理画面の「バッチを今すぐ中断」ボタンも同じ要求です。total_timeout による強制打ち切り（Day 041・042）は「時間切れ」を検知する仕組みでしたが、「誰かが明示的に止めろと言った」という外部イベントを検知する仕組みはまだありません。

これを実現する自然な道具が tokio::sync::watch です。watch チャンネルは「最新の1つの値を、複数の受信者にブロードキャストする」ためのチャンネルで、「シャットダウン中かどうか（bool）」のような『現在の状態』を表現するのに向いています（mpsc が『次々に届く個別のメッセージの列』に向いているのとは設計思想が異なります）。Receiver::changed() は「値が（前回観測してから）変わるまで待つ」Future を返し、これと set.join_next()（次のタスク完了を待つFuture）を tokio::select! で競争させることで、「次のタスクが終わるのを待ちながら、同時にシャットダウン信号にも耳を傾ける」というループが書けます。

📝 問題

Day 042 の run_batch_bounded を土台に、外部からのシャットダウン信号を受け取って早期に打ち切れる run_batch_cancellable を実装してください。

• シグネチャ: async fn run_batch_cancellable(inputs: Vec<u32>, max_concurrency: usize, per_task_timeout: Duration, total_timeout: Duration, shutdown: tokio::sync::watch::Receiver<bool>) -> BatchReport
• fetch 関数は Day 042 と同一のものを使用してください
• BatchReport に aborted_by_shutdown: usize フィールドを追加し、total_timeout による打ち切り（aborted_by_total_timeout）と区別して集計してください
• 集計ループは tokio::select! で「次のタスク完了」と「シャットダウン信号の変化」を競争させる形にしてください。シャットダウン信号（*shutdown.borrow() == true）を検知したら set.abort_all() で残りのタスクを打ち切り、aborted_by_shutdown に計上してください
• 罠を踏まないこと: watch チャンネルの送信側（Sender）が全てdropされた場合、Receiver::changed() は以後ずっと即座に Err を返すようになります。これを何も考えずに select! のループに書くと、送信側が二度と true を送ってこないと分かった後も、その分岐が「即座に完了するFuture」として毎回選ばれ続け、もう一方の set.join_next() の分岐が実行される機会を奪う（＝タスクの完了に気づけずバッチが進まなくなる）危険があります。この罠を回避する実装にしてください
• 以下の3つをテストで検証してください
  1. シャットダウン信号が一度も送られない場合、通常通り全タスクが完了すること
  2. シャットダウン信号が処理の途中で送られた場合、早期に打ち切られ aborted_by_shutdown が計上されること
  3. シャットダウン信号の送信側が「一度も送らずにdrop」された場合でも、上記の罠でハングせず正常に完了すること

📋 コピーasync fn fetch(input: u32) -> Result {
    tokio::time::sleep(Duration::from_millis(((input % 5) * 40) as u64)).await;
    if input % 7 == 0 {
        return Err(format!("input {input} is not allowed (divisible by 7)"));
    }
    if input % 11 == 0 {
        panic!("input {input} triggered an unexpected panic");
    }
    Ok(input * input)
}

🔍 ヒント（段階的開示）

📋 コピーlet mut shutdown_active = true;
let mut shutdown_triggered = false;

loop {
    tokio::select! {
        biased;
        changed = shutdown.changed(), if shutdown_active => {
            match changed {
                Ok(()) if *shutdown.borrow() => {
                    shutdown_triggered = true;
                    break;
                }
                Ok(()) => {} // false→false等、無関係な変化なので無視して継続
                Err(_) => {
                    shutdown_active = false; // 送信側消滅: この腕はもう待たない
                }
            }
        }
        maybe_result = set.join_next() => {
            match maybe_result {
                Some(join_result) => { /* Day 042と同じ集計処理 */ }
                None => break, // 全タスク完了
            }
        }
    }
}

✅ 模範解答

📋 コピー// Cargo.toml に以下を追加して検証すること:
// [dependencies]
// tokio = { version = "1", features = ["full"] }

use std::sync::Arc;
use std::time::Duration;
use tokio::sync::{watch, Semaphore};
use tokio::task::JoinSet;

async fn fetch(input: u32) -> Result {
    tokio::time::sleep(Duration::from_millis(((input % 5) * 40) as u64)).await;
    if input % 7 == 0 {
        return Err(format!("input {input} is not allowed (divisible by 7)"));
    }
    if input % 11 == 0 {
        panic!("input {input} triggered an unexpected panic");
    }
    Ok(input * input)
}

#[derive(Debug, Default, PartialEq, Eq)]
struct BatchReport {
    succeeded: Vec,
    application_errors: usize,
    panicked: usize,
    timed_out: usize,
    aborted_by_total_timeout: usize,
    aborted_by_shutdown: usize,
}

/// `Semaphore`で並行度を制限しつつ、`watch`チャンネル経由の外部シャットダウン
/// 信号と`total_timeout`のどちらでも早期打ち切りができるバッチ処理。
///
/// 設計判断: `shutdown_active`フラグで「送信側が消滅した後のshutdown.changed()腕」
/// をselect!の前置条件(if)で無効化する。無効化しないと、送信側drop後は
/// changed()が毎回即座にErrを返す「常に準備完了のFuture」になり、biased select!
/// が他の腕(join_next())を一切ポーリングしなくなるビジーループに陥る。
async fn run_batch_cancellable(
    inputs: Vec,
    max_concurrency: usize,
    per_task_timeout: Duration,
    total_timeout: Duration,
    mut shutdown: watch::Receiver,
) -> BatchReport {
    let semaphore = Arc::new(Semaphore::new(max_concurrency));
    let mut set: JoinSet, tokio::time::error::Elapsed>> =
        JoinSet::new();

    for input in inputs {
        let semaphore = Arc::clone(&semaphore);
        set.spawn(async move {
            let _permit = semaphore
                .acquire_owned()
                .await
                .expect("semaphore is never closed");
            tokio::time::timeout(per_task_timeout, fetch(input)).await
        });
    }

    let mut report = BatchReport::default();
    let mut shutdown_triggered = false;
    let mut shutdown_active = true;

    let timed_out_overall = tokio::time::timeout(total_timeout, async {
        loop {
            tokio::select! {
                biased;
                changed = shutdown.changed(), if shutdown_active => {
                    match changed {
                        Ok(()) if *shutdown.borrow() => {
                            shutdown_triggered = true;
                            break;
                        }
                        Ok(()) => {
                            // false→falseのような無関係な変化。監視を継続する。
                        }
                        Err(_) => {
                            // 送信側(Sender)が全てdropされた = 今後二度とtrueは
                            // 送られてこない。この腕を無効化しないと、biased
                            // selectは毎回ここをErrで即座に選び続け、
                            // join_next()の腕が永遠にポーリングされなくなる。
                            shutdown_active = false;
                        }
                    }
                }
                maybe_result = set.join_next() => {
                    match maybe_result {
                        Some(join_result) => match join_result {
                            Ok(Ok(Ok(value))) => report.succeeded.push(value),
                            Ok(Ok(Err(_app_err))) => report.application_errors += 1,
                            Ok(Err(_elapsed)) => report.timed_out += 1,
                            Err(join_err) if join_err.is_panic() => report.panicked += 1,
                            Err(_cancelled) => {}
                        },
                        None => break, // 全タスク完了
                    }
                }
            }
        }
    })
    .await
    .is_err();

    if timed_out_overall || shutdown_triggered {
        set.abort_all();
        while let Some(join_result) = set.join_next().await {
            match join_result {
                Ok(Ok(Ok(value))) => report.succeeded.push(value),
                Ok(Ok(Err(_app_err))) => report.application_errors += 1,
                Ok(Err(_elapsed)) => report.timed_out += 1,
                Err(join_err) if join_err.is_panic() => report.panicked += 1,
                Err(_cancelled) => {
                    if shutdown_triggered {
                        report.aborted_by_shutdown += 1;
                    } else {
                        report.aborted_by_total_timeout += 1;
                    }
                }
            }
        }
    }

    report
}

#[tokio::main]
async fn main() {
    let (_tx, rx) = watch::channel(false);
    let report = run_batch_cancellable(
        (1..=20).collect(),
        4,
        Duration::from_secs(2),
        Duration::from_secs(5),
        rx,
    )
    .await;
    println!("{report:#?}");
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::time::Instant;

    #[tokio::test]
    async fn completes_normally_when_shutdown_never_fires() {
        // 送信側を握ったまま一度もtrueを送らない = shutdown.changed()は
        // ずっとPendingのまま、join_next()側だけで正常に完了するはず。
        let (_tx, rx) = watch::channel(false);
        let inputs = vec![1, 2, 3, 7, 11, 14];

        let report =
            run_batch_cancellable(inputs, 3, Duration::from_secs(2), Duration::from_secs(5), rx)
                .await;

        assert_eq!(report.succeeded.len(), 3); // 1, 2, 3
        assert_eq!(report.application_errors, 2); // 7, 14
        assert_eq!(report.panicked, 1); // 11
        assert_eq!(report.aborted_by_shutdown, 0);
        assert_eq!(report.aborted_by_total_timeout, 0);
    }

    #[tokio::test]
    async fn shutdown_signal_triggers_early_abort() {
        // 全て160ms sleepする10件(7・11の倍数は避けて時間測定のノイズを消す)。
        // 同時実行数2なので、無制限なら800ms前後で全件完了する。
        let inputs: Vec = vec![4, 9, 19, 24, 29, 34, 39, 54, 59, 64];
        let (tx, rx) = watch::channel(false);

        let batch_fut =
            run_batch_cancellable(inputs, 2, Duration::from_secs(5), Duration::from_secs(5), rx);
        let trigger_fut = async {
            tokio::time::sleep(Duration::from_millis(100)).await;
            let _ = tx.send(true);
        };

        let start = Instant::now();
        let (report, ()) = tokio::join!(batch_fut, trigger_fut);
        let elapsed = start.elapsed();

        assert!(
            elapsed < Duration::from_millis(600),
            "shutdown信号後も打ち切られず走り続けている可能性: elapsed = {elapsed:?}"
        );
        assert!(
            report.aborted_by_shutdown > 0,
            "shutdown信号によるabortが計上されていない"
        );
        assert_eq!(report.aborted_by_total_timeout, 0);
    }

    #[tokio::test]
    async fn shutdown_sender_dropped_without_signal_does_not_hang() {
        // 送信側を一度もtrueを送らないままdropする。shutdown_activeガードが
        // なければ、changed()が毎回即座にErrを返し続けるビジーループに陥り、
        // join_next()の腕が進まず外側のtimeout(2秒)がErrになるはず。
        let (tx, rx) = watch::channel(false);
        drop(tx);
        let inputs = vec![1, 2, 3];

        let result = tokio::time::timeout(
            Duration::from_secs(2),
            run_batch_cancellable(inputs, 2, Duration::from_secs(5), Duration::from_secs(5), rx),
        )
        .await;

        let report =
            result.expect("shutdown_activeガードがないとここでハング/タイムアウトするはず");
        assert_eq!(report.succeeded.len(), 3);
    }
}

🪜 Step-by-Step 解説

💡 設計思想・なぜこう書くのか

tokio::select!に複数の腕を並べるだけでは、「今後絶対に起こらないイベント」を区別できません。Rustの型システムはResultで成功/失敗を表現させますが、「このFutureはもう意味のある結果を返さない」という時間的な性質までは型では表現しきれず、ifガードで明示的にプログラマが管理する必要があります。これは所有権システムがメモリ安全性を静的に保証してくれるのとは対照的に、並行処理のプロトコル設計は依然として書き手の責任範囲であることを改めて示す例です（Day 043で見た「Senderのクローンを持ちすぎる」問題と同じ種類の教訓です）。

🛑 コンパイルエラーが出た場合

shutdown: watch::Receiver<bool>を関数引数で受け取る際にmutを付け忘れると、以下のようなエラーになります。

📋 コピー// アンチパターン: mutを付け忘れる
async fn run_batch_cancellable(
    inputs: Vec,
    max_concurrency: usize,
    per_task_timeout: Duration,
    total_timeout: Duration,
    shutdown: watch::Receiver, // mut shutdown ではない
) -> BatchReport {
    // ...
    tokio::select! {
        changed = shutdown.changed() => { /* ... */ } // ここでエラー
        // ...
    }
}

🌐 他言語との比較

Go の context.Done() も「クローズされたチャンネルは即座に読める」という似た性質を持ちますが、Goのselectは「一度クローズされたらそのままDone()分岐を選び続けて終了処理へ向かう」という設計を前提にしている場面がほとんどで、今回のような「無効化しないとハング側の分岐が飢餓する」問題は起きにくい書き方が主流です。Rustではselect!をループで回し続ける設計自体を多用するため、「もう起きないイベントを待ち続けない」という一手間が顕在化しやすいと言えます。

🏆 実務での使いどころ

• Kubernetes Podのgraceful shutdown: SIGTERMハンドラでwatch::Sender::send(true)し、進行中のバッチ処理に「猶予期間内に安全に終われ」と伝える
• 管理画面からの処理中断: 「このジョブを中断」ボタンのAPIコールをwatchチャンネルの送信に変換し、実行中のワーカーに伝播させる
• 段階的な機能フラグの反映: watchは「シャットダウンbool」に限らず「現在の設定値」全般のブロードキャストに使える（例: レート制限の上限値をホットリロードする）

⚠️ よくある誤解・ミス

🚀 次のステップ

• 発展: shutdownが発火した際、即座にabort_all()するのではなく「猶予期間（grace period）」を設け、その間は新規タスクのspawnを止めつつ実行中のタスクの自然な完了を待ち、猶予期間を過ぎたタスクだけabortする設計に変更する
• 次回予告: async trait（async-trait / RPITIT） — 非同期メソッドを持つトレイトの設計

🎯 自己評価

自分の回答

気づき・メモ