Pick Backends with Round-Robin Selection

use std::sync::atomic::{AtomicUsize, Ordering};

#[derive(Debug, Clone, PartialEq, Eq)]
struct Backend {
    name: &'static str,
    healthy: bool,
}

struct LoadBalancer {
    backends: Vec<Backend>,
    cursor: AtomicUsize,
}

impl LoadBalancer {
    fn new(backends: Vec<Backend>) -> Self {
        Self {
            backends,
            cursor: AtomicUsize::new(0),
        }
    }

    fn next(&self) -> Option<&Backend> {
        let count = self.backends.len();
        if count == 0 {
            return None;
        }
        for _ in 0..count {
            let index = self.cursor.fetch_add(1, Ordering::Relaxed) % count;
            let backend = &self.backends[index];
            if backend.healthy {
                return Some(backend);
            }
        }
        None
    }

    fn set_healthy(&mut self, name: &str, healthy: bool) {
        if let Some(backend) = self.backends.iter_mut().find(|item| item.name == name) {
            backend.healthy = healthy;
        }
    }
}

fn main() {
    let mut balancer = LoadBalancer::new(vec![
        Backend { name: "a", healthy: true },
        Backend { name: "b", healthy: false },
        Backend { name: "c", healthy: true },
    ]);

    let picks: Vec<&str> = (0..4)
        .filter_map(|_| balancer.next().map(|backend| backend.name))
        .collect();
    println!("{picks:?}");
    assert_eq!(picks, vec!["a", "c", "a", "c"]);

    balancer.set_healthy("b", true);
    let names: Vec<&str> = (0..3)
        .filter_map(|_| balancer.next().map(|backend| backend.name))
        .collect();
    assert_eq!(names, vec!["a", "b", "c"]);
    assert_eq!(LoadBalancer::new(Vec::new()).next(), None);
}