Skip to main content

rs_zero/discovery/
static_discovery.rs

1use std::collections::BTreeMap;
2
3use async_trait::async_trait;
4
5use crate::discovery::{Discovery, DiscoveryError, DiscoveryResult, ServiceInstance};
6
7/// Discovery adapter backed by a static service map.
8#[derive(Debug, Clone, Default)]
9pub struct StaticDiscovery {
10    services: BTreeMap<String, Vec<ServiceInstance>>,
11}
12
13impl StaticDiscovery {
14    /// Creates an empty static discovery map.
15    pub fn new() -> Self {
16        Self::default()
17    }
18
19    /// Adds an instance and returns the updated map.
20    pub fn with_instance(mut self, instance: ServiceInstance) -> Self {
21        self.services
22            .entry(instance.service.clone())
23            .or_default()
24            .push(instance);
25        self
26    }
27}
28
29#[async_trait]
30impl Discovery for StaticDiscovery {
31    async fn discover(&self, service: &str) -> DiscoveryResult<Vec<ServiceInstance>> {
32        let instances = self
33            .services
34            .get(service)
35            .map(|items| {
36                items
37                    .iter()
38                    .filter(|instance| instance.healthy)
39                    .cloned()
40                    .collect::<Vec<_>>()
41            })
42            .unwrap_or_default();
43        if instances.is_empty() {
44            Err(DiscoveryError::NoInstances {
45                service: service.to_string(),
46            })
47        } else {
48            Ok(instances)
49        }
50    }
51}
52
53#[cfg(test)]
54mod tests {
55    use super::{Discovery, StaticDiscovery};
56    use crate::discovery::{InstanceEndpoint, ServiceInstance};
57
58    #[tokio::test]
59    async fn static_discovery_filters_unhealthy_instances() {
60        let endpoint = InstanceEndpoint::new("127.0.0.1", 8080).expect("endpoint");
61        let discovery = StaticDiscovery::new()
62            .with_instance(ServiceInstance::new("api", "api-1", endpoint.clone()))
63            .with_instance(ServiceInstance::new("api", "api-2", endpoint).with_health(false));
64
65        let instances = discovery.discover("api").await.expect("instances");
66        assert_eq!(instances.len(), 1);
67        assert_eq!(instances[0].id, "api-1");
68    }
69}