rs-zero 0.2.6

Rust-first microservice framework inspired by go-zero engineering practices
Documentation
use std::collections::BTreeMap;

use async_trait::async_trait;

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

/// Discovery adapter backed by a static service map.
#[derive(Debug, Clone, Default)]
pub struct StaticDiscovery {
    services: BTreeMap<String, Vec<ServiceInstance>>,
}

impl StaticDiscovery {
    /// Creates an empty static discovery map.
    pub fn new() -> Self {
        Self::default()
    }

    /// Adds an instance and returns the updated map.
    pub fn with_instance(mut self, instance: ServiceInstance) -> Self {
        self.services
            .entry(instance.service.clone())
            .or_default()
            .push(instance);
        self
    }
}

#[async_trait]
impl Discovery for StaticDiscovery {
    async fn discover(&self, service: &str) -> DiscoveryResult<Vec<ServiceInstance>> {
        let instances = self
            .services
            .get(service)
            .map(|items| {
                items
                    .iter()
                    .filter(|instance| instance.healthy)
                    .cloned()
                    .collect::<Vec<_>>()
            })
            .unwrap_or_default();
        if instances.is_empty() {
            Err(DiscoveryError::NoInstances {
                service: service.to_string(),
            })
        } else {
            Ok(instances)
        }
    }
}

#[cfg(test)]
mod tests {
    use super::{Discovery, StaticDiscovery};
    use crate::discovery::{InstanceEndpoint, ServiceInstance};

    #[tokio::test]
    async fn static_discovery_filters_unhealthy_instances() {
        let endpoint = InstanceEndpoint::new("127.0.0.1", 8080).expect("endpoint");
        let discovery = StaticDiscovery::new()
            .with_instance(ServiceInstance::new("api", "api-1", endpoint.clone()))
            .with_instance(ServiceInstance::new("api", "api-2", endpoint).with_health(false));

        let instances = discovery.discover("api").await.expect("instances");
        assert_eq!(instances.len(), 1);
        assert_eq!(instances[0].id, "api-1");
    }
}