trait経由で切り替え可能なロードバランサー設計と、4種のアルゴリズム比較

TCPコネクションの転送先を単一のバックエンドから複数のバックエンドへ広げると、「次のコネクションをどのバックエンドに送るか」という選択の問題が発生します。この選択ロジックをロードバランサーとして抽象化し、Round Robin・Random・Least Connections・Weightedの4種類を実装します。

ロードバランサーの抽象化

アルゴリズムを設定次第で差し替えられるようにするため、共通のトレイトを定義し、動的ディスパッチ（dyn）で扱います。

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

next_backendは次に転送すべきバックエンドを返すメソッドです。releaseはコネクション終了を通知するためのメソッドで、大半のアルゴリズムには不要なため空のデフォルト実装を持たせ、必要なアルゴリズムだけが上書きする設計にしています。

アルゴリズムの選択を実行時まで遅延させたいため、ジェネリクスではなく動的ディスパッチ（dyn）を採用しています。

// src/main.rs（抜粋）
let lb: Arc<dyn LoadBalancer> = Arc::new(Weighted::new(backends));

tokio::spawnで生成される各コネクションのタスクは、寿命の異なる複数の非同期タスクとして同じロードバランサーの状態を共有し続ける必要があるため、Arcで保持しています。dynとArcがそれぞれ何を解決しているかの詳細は、Rc/Arc/RefCell/MutexとSend/Syncの整理記事にまとめています。

Round Robin: ロックフリーなカウンタ

バックエンドを順番に一巡させるアルゴリズムです。状態は「次に選ぶインデックス」だけなので、AtomicUsizeのカウンタ1つで実現できます。

// src/load_balancer/round_robin.rs
pub struct RoundRobin {
    backends: Vec<SocketAddr>,
    counter: AtomicUsize,
}

impl LoadBalancer for RoundRobin {
    fn next_backend(&self) -> SocketAddr {
        let idx = self.counter.fetch_add(1, Ordering::Relaxed) % self.backends.len();
        self.backends[idx]
    }
}

fetch_addはCPUのアトミック命令を使い、「読む→1増やす→書く」を他スレッドに割り込まれず1操作として実行します。ロックを取らないため、Mutexより軽量です。

Random: 状態を持たない選択

ランダムに1つを選ぶだけなので、共有すべき状態自体が存在しません。

// src/load_balancer/random.rs
use rand::RngExt;

pub struct Random {
    backends: Vec<SocketAddr>,
}

impl LoadBalancer for Random {
    fn next_backend(&self) -> SocketAddr {
        let idx = rand::rng().random_range(0..self.backends.len());
        self.backends[idx]
    }
}

rand::rng()はスレッドローカルな乱数生成器を返すため、複数スレッドから同時に呼び出しても、生成器自体がスレッドごとに独立しており競合しません。

Least Connections: 選択と解放の対称性

処理中のコネクション数が最小のバックエンドを選ぶアルゴリズムです。バックエンドごとにAtomicUsizeのカウンタを持たせ、選択時に+1、コネクション終了時に-1します。

// src/load_balancer/least_connections.rs
impl LoadBalancer for LeastConnections {
    fn next_backend(&self) -> SocketAddr {
        let idx = self
            .active
            .iter()
            .enumerate()
            .min_by_key(|(_, count)| count.load(Ordering::Relaxed))
            .map(|(idx, _)| idx)
            .unwrap();
        self.active[idx].fetch_add(1, Ordering::Relaxed);
        self.backends[idx]
    }

    fn release(&self, backend_addr: SocketAddr) {
        if let Some(idx) = self.backends.iter().position(|&b| b == backend_addr) {
            self.active[idx].fetch_sub(1, Ordering::Relaxed);
        }
    }
}

+1と-1が対になっているため、呼び出し側（コネクションのハンドラ）は成功・失敗を問わず必ずreleaseを呼ぶ必要があります。転送処理を専用の関数に切り出し、その戻り値を見る前にreleaseを呼ぶ順序にすることで、この対称性を保証しています（tokioの`copy_bidirectional`によるバイト列転送を扱った記事で扱った転送処理を、この選択ロジックと組み合わせています）。

// src/proxy.rs（抜粋）
let backend_addr = lb.next_backend();
let result = forward(&mut inbound, backend_addr).await;
lb.release(backend_addr);
let (from_client, from_backend) = result?;

なお、複数バックエンドのカウントが同着（タイ）の場合、Iterator::min_by_keyは最初に見つかった要素を返す仕様です。処理中のコネクション数に実際の差がない限り、常に同じインデックスのバックエンドが優先的に選ばれるという偏りが生じます。これはLeast Connectionsというアルゴリズムの性質上、複数のコネクションが実際に重なって処理されている状況でなければ意味のある差別化ができないことの裏返しでもあります。

Weighted: 複数フィールドをまたぐのでMutexが要る

重みの比率に応じて振り分け頻度を変えるアルゴリズムです。ここではnginxなどで採用されているsmooth weighted round robinを実装します。各バックエンドに「現在の持ち点」を持たせ、選択のたびに全バックエンドの持ち点へ重みを加算し、最大のものを選んでその持ち点から重み合計を引く、という手順です。

// src/load_balancer/weighted.rs
pub struct Weighted {
    backends: Mutex<Vec<WeightedBackend>>,
    total_weight: i32,
}

impl LoadBalancer for Weighted {
    fn next_backend(&self) -> SocketAddr {
        let mut backends = self.backends.lock().unwrap();
        for b in backends.iter_mut() {
            b.current_weight += b.weight;
        }
        let selected = backends
            .iter_mut()
            .max_by_key(|b| b.current_weight)
            .unwrap();
        selected.current_weight -= self.total_weight;
        selected.addr
    }
}

これまでの3つのアルゴリズムはAtomicUsizeだけで実現できましたが、Weightedだけは「全バックエンドの持ち点を読んで加算し、最大値を選んで1つだけ更新する」という一連の操作を、他スレッドに割り込まれず1つの塊として実行する必要があります。個々のフィールドをそれぞれAtomicにしても、この一連の流れの途中に別スレッドが割り込むと不整合が起きるため、複数フィールドにまたがる操作にはMutexによる排他制御が必要です（単一の値の増減で済むAtomicと、複数フィールドにまたがる操作に要るMutexの違いは、Rc/Arc/RefCell/MutexとSend/Syncの整理記事で扱っています）。

重み5:1:1で設定した場合、選択順序はA, A, C, A, B, A, Aのように、重みの大きいバックエンドが連続しすぎないよう分散します（Iterator::max_by_keyは同着の場合に最後に見つかった要素を返す仕様のため、min_by_keyとは逆の挙動になる点に注意が必要です）。

モジュール分割とインラインテスト

4つのアルゴリズムを1ファイルにまとめると見通しが悪くなるため、load_balancerをモジュールとして切り出しました。

src/
  main.rs                          # 設定・起動ループ
  proxy.rs                         # 1コネクション分の転送処理
  load_balancer.rs                 # トレイト定義 + 各実装の再エクスポート
  load_balancer/
    round_robin.rs
    random.rs
    least_connections.rs
    weighted.rs

Rust 2018以降のモジュールシステムではmod.rsを使わず、「モジュール名と同じ.rsファイル」と「同じ名前のディレクトリ」の組み合わせで親子関係を表します。load_balancer.rs側でmod round_robin;のようにサブモジュールを宣言し、pub useで外部に再エクスポートすることで、main.rs側からはload_balancer::RoundRobinのようにフラットに参照できます。

各アルゴリズムのファイルには、そのロジックだけを検証する単体テストを添えています。next_backendはTCP接続や非同期処理を一切含まない純粋なロジックなので、tokioランタイムなしでcargo testだけで検証できます。

// src/load_balancer/round_robin.rs（抜粋）
#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn cycles_through_all_backends_in_order() {
        let backends: Vec<SocketAddr> = vec![
            "127.0.0.1:9001".parse().unwrap(),
            "127.0.0.1:9002".parse().unwrap(),
            "127.0.0.1:9003".parse().unwrap(),
        ];
        let lb = RoundRobin::new(backends.clone());

        let picked: Vec<SocketAddr> = (0..6).map(|_| lb.next_backend()).collect();

        assert_eq!(
            picked,
            vec![
                backends[0], backends[1], backends[2],
                backends[0], backends[1], backends[2],
            ]
        );
    }
}

選択ロジックと非同期I/Oの転送処理（forward/handle_connection）を別モジュールに分離したことで、前者は同期的な単体テストで、後者は実際のTCP接続を使った動作確認で、それぞれ適した方法で検証できる構成になっています。