use std::sync::atomic::{AtomicUsize, Ordering};
use crate::discovery::{DiscoveryError, DiscoveryResult, ServiceInstance};
pub trait InstanceSelector: Send + Sync {
fn select(
&self,
service: &str,
instances: &[ServiceInstance],
) -> DiscoveryResult<ServiceInstance>;
}
#[derive(Debug, Default)]
pub struct RoundRobinSelector {
cursor: AtomicUsize,
}
impl RoundRobinSelector {
pub fn new() -> Self {
Self::default()
}
}
impl InstanceSelector for RoundRobinSelector {
fn select(
&self,
service: &str,
instances: &[ServiceInstance],
) -> DiscoveryResult<ServiceInstance> {
let weighted = instances
.iter()
.filter(|instance| instance.healthy)
.flat_map(|instance| std::iter::repeat_n(instance, instance.weight.max(1) as usize))
.collect::<Vec<_>>();
if weighted.is_empty() {
return Err(DiscoveryError::NoInstances {
service: service.to_string(),
});
}
let index = self.cursor.fetch_add(1, Ordering::Relaxed) % weighted.len();
Ok(weighted[index].clone())
}
}
#[cfg(test)]
mod tests {
use super::{InstanceSelector, RoundRobinSelector};
use crate::discovery::{InstanceEndpoint, ServiceInstance};
#[test]
fn round_robin_selects_stable_sequence() {
let endpoint = InstanceEndpoint::new("127.0.0.1", 8080).expect("endpoint");
let instances = vec![
ServiceInstance::new("api", "a", endpoint.clone()),
ServiceInstance::new("api", "b", endpoint),
];
let selector = RoundRobinSelector::new();
assert_eq!(selector.select("api", &instances).expect("first").id, "a");
assert_eq!(selector.select("api", &instances).expect("second").id, "b");
assert_eq!(selector.select("api", &instances).expect("third").id, "a");
}
}