rs-zero 0.2.7

Rust-first microservice framework inspired by go-zero engineering practices
Documentation
use std::sync::atomic::{AtomicUsize, Ordering};

use crate::discovery::{DiscoveryError, DiscoveryResult, ServiceInstance};

/// Selects one service instance from a discovered set.
pub trait InstanceSelector: Send + Sync {
    /// Returns one instance, or an error when no healthy instance exists.
    fn select(
        &self,
        service: &str,
        instances: &[ServiceInstance],
    ) -> DiscoveryResult<ServiceInstance>;
}

/// Round-robin selector that expands entries by their configured weight.
#[derive(Debug, Default)]
pub struct RoundRobinSelector {
    cursor: AtomicUsize,
}

impl RoundRobinSelector {
    /// Creates a selector starting at index zero.
    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");
    }
}