rs_zero/discovery/
selector.rs

1use std::sync::atomic::{AtomicUsize, Ordering};
2
3use crate::discovery::{DiscoveryError, DiscoveryResult, ServiceInstance};
4
5/// Selects one service instance from a discovered set.
6pub trait InstanceSelector: Send + Sync {
7    /// Returns one instance, or an error when no healthy instance exists.
8    fn select(
9        &self,
10        service: &str,
11        instances: &[ServiceInstance],
12    ) -> DiscoveryResult<ServiceInstance>;
13}
14
15/// Round-robin selector that expands entries by their configured weight.
16#[derive(Debug, Default)]
17pub struct RoundRobinSelector {
18    cursor: AtomicUsize,
19}
20
21impl RoundRobinSelector {
22    /// Creates a selector starting at index zero.
23    pub fn new() -> Self {
24        Self::default()
25    }
26}
27
28impl InstanceSelector for RoundRobinSelector {
29    fn select(
30        &self,
31        service: &str,
32        instances: &[ServiceInstance],
33    ) -> DiscoveryResult<ServiceInstance> {
34        let weighted = instances
35            .iter()
36            .filter(|instance| instance.healthy)
37            .flat_map(|instance| std::iter::repeat_n(instance, instance.weight.max(1) as usize))
38            .collect::<Vec<_>>();
39        if weighted.is_empty() {
40            return Err(DiscoveryError::NoInstances {
41                service: service.to_string(),
42            });
43        }
44        let index = self.cursor.fetch_add(1, Ordering::Relaxed) % weighted.len();
45        Ok(weighted[index].clone())
46    }
47}
48
49#[cfg(test)]
50mod tests {
51    use super::{InstanceSelector, RoundRobinSelector};
52    use crate::discovery::{InstanceEndpoint, ServiceInstance};
53
54    #[test]
55    fn round_robin_selects_stable_sequence() {
56        let endpoint = InstanceEndpoint::new("127.0.0.1", 8080).expect("endpoint");
57        let instances = vec![
58            ServiceInstance::new("api", "a", endpoint.clone()),
59            ServiceInstance::new("api", "b", endpoint),
60        ];
61        let selector = RoundRobinSelector::new();
62        assert_eq!(selector.select("api", &instances).expect("first").id, "a");
63        assert_eq!(selector.select("api", &instances).expect("second").id, "b");
64        assert_eq!(selector.select("api", &instances).expect("third").id, "a");
65    }
66}
rs_zero/discovery/selector.rs

rs_zero/discovery/
selector.rs
