TCP接続によるアクティブヘルスチェックと、ロードバランサーからの動的除外

複数バックエンドへの振り分けだけでは、停止しているバックエンドにも構わず転送してしまいます。バックエンドの生死を定期的なTCP接続確認で判定し、死んでいるバックエンドをロードバランサーの選択候補から動的に除外する仕組みを実装します。

生死フラグを持つ共有状態

バックエンドのアドレス一覧と、各バックエンドの生死フラグをセットで持つBackendsを導入します。ロードバランサーとヘルスチェックタスクの両方から参照される状態なので、Arcで共有します。

// src/health.rs
pub struct Backends {
    addrs: Vec<SocketAddr>,
    healthy: Vec<AtomicBool>,
}

impl Backends {
    pub fn new(addrs: Vec<SocketAddr>) -> Self {
        let healthy = addrs.iter().map(|_| AtomicBool::new(true)).collect();
        Self { addrs, healthy }
    }

    pub fn addrs(&self) -> &[SocketAddr] {
        &self.addrs
    }

    pub fn is_healthy(&self, idx: usize) -> bool {
        self.healthy[idx].load(Ordering::Relaxed)
    }

    fn set_healthy(&self, idx: usize, healthy: bool) {
        self.healthy[idx].store(healthy, Ordering::Relaxed);
    }
}

初期状態は全バックエンドを健全(true)としています。起動直後にヘルスチェックが1周するまでの間は、実際には死んでいるバックエンドにも転送を試みる可能性がありますが、これは後述するtokio::time::intervalが起動直後に即座に1回目のtickを発行する仕様で実質的にカバーされます。

並行・タイムアウト付きの定期チェック

// src/health.rs(続き)
pub async fn run_health_checks(backends: Arc<Backends>, interval: Duration, timeout: Duration) {
    let mut ticker = tokio::time::interval(interval);
    loop {
        ticker.tick().await;
        for idx in 0..backends.addrs().len() {
            let backends = Arc::clone(&backends);
            tokio::spawn(async move {
                let addr = backends.addrs()[idx];
                let healthy = tokio::time::timeout(timeout, TcpStream::connect(addr))
                    .await
                    .is_ok_and(|connect_result| connect_result.is_ok());
                backends.set_healthy(idx, healthy);
            });
        }
    }
}

各バックエンドへの接続確認を個別のtokio::spawnに分けているのは、1台の応答が遅い(タイムアウト待ちになる)ことが、他のバックエンドの判定を遅延させないようにするためです。全バックエンドをまとめて1つのタスクで順番にチェックすると、最悪の場合「バックエンド数 × タイムアウト時間」だけ次のチェック開始が遅れます。

判定自体はtokio::time::timeoutでTCP接続の試行を包み、指定時間内に接続が確立できたかどうかだけを見ます。L4の範囲での生死判定なので、アプリケーション層の応答内容までは見ません。

「選べないかもしれない」をtraitに反映する

ロードバランサーが選択できるバックエンドが1台もない(全滅している)状況が起こり得るようになったため、LoadBalancerトレイトの戻り値をOption<SocketAddr>に変更しました。

// src/load_balancer.rs
pub trait LoadBalancer: Send + Sync {
    fn next_backend(&self) -> Option<SocketAddr>;
    fn release(&self, _backend_addr: SocketAddr) {}
}

各アルゴリズムの実装は、まず生きているインデックスだけに絞り込んでから、これまでの選択ロジックを適用する形に変わります。Round Robinを例にすると次のようになります。

// src/load_balancer/round_robin.rs(抜粋)
fn next_backend(&self) -> Option<SocketAddr> {
    let healthy_indices: Vec<usize> = (0..self.backends.addrs().len())
        .filter(|&idx| self.backends.is_healthy(idx))
        .collect();
    if healthy_indices.is_empty() {
        return None;
    }
    let pick = self.counter.fetch_add(1, Ordering::Relaxed) % healthy_indices.len();
    Some(self.backends.addrs()[healthy_indices[pick]])
}

候補数(healthy_indices.len())を分母にしているため、バックエンドが離脱・復帰するたびに巡回の周期そのものが変化します。これはRound Robinというアルゴリズムの性質上自然な挙動で、常に「その時点で生きている集合」を一巡する動きになります。

Weightedのように複数バックエンドの状態をまたいで計算するアルゴリズムでは、重みの合計値も生きているバックエンドだけで毎回計算し直します。死んでいるバックエンドには持ち点を積み増さないため、復帰した瞬間に停止期間中の分まで一気に選ばれる、という偏りは発生しません。

呼び出し側: let-elseで早期リターンする

next_backendNoneを返す場合、そのコネクションは転送先がないのでエラーとして扱います。

// src/proxy.rs(抜粋)
let Some(backend_addr) = lb.next_backend() else {
    return Err(std::io::Error::other("no healthy backend available"));
};

let-else構文は、右辺のパターンマッチが失敗した場合にelseブロック側で必ず関数を抜ける(return/continue/break/panicのいずれか)ことをコンパイラが要求します。matchで同じことを書くよりネストが増えず、「正常系だけがこの先に続く」という読み方ができます。

動作確認

3台のバックエンドを起動した状態でRound Robinの巡回を確認したあと、1台を停止し、ヘルスチェックの間隔(数秒)を空けてから振り分け先を確認すると、停止したバックエンドが巡回から外れ、残りの2台だけで循環します。停止したバックエンドを再起動し、再度間隔を空けて確認すると、巡回に復帰します。死活の変化がロードバランサーの選択に即座に反映される、という一連の流れが確認できました。