rskit_discovery/
memory.rs

1use std::collections::HashMap;
2use std::sync::Arc;
3
4use async_trait::async_trait;
5use tokio::sync::RwLock;
6
7use rskit_errors::{AppError, AppResult};
8
9use crate::instance::ServiceInstance;
10use crate::traits::{Discovery, Registry};
11
12/// In-memory service registry useful for testing and local development.
13pub struct InMemoryDiscovery {
14    instances: Arc<RwLock<HashMap<String, Vec<ServiceInstance>>>>,
15}
16
17impl InMemoryDiscovery {
18    /// Create an empty discovery registry.
19    pub fn new() -> Self {
20        Self {
21            instances: Arc::new(RwLock::new(HashMap::new())),
22        }
23    }
24
25    /// Add an instance under the given service name.
26    pub async fn add(&self, service: &str, instance: ServiceInstance) {
27        let mut map = self.instances.write().await;
28        map.entry(service.to_owned())
29            .or_insert_with(Vec::new)
30            .push(instance);
31    }
32
33    /// Remove an instance by id from the given service.
34    pub async fn remove(&self, service: &str, id: &str) {
35        let mut map = self.instances.write().await;
36        if let Some(instances) = map.get_mut(service) {
37            instances.retain(|i| i.id != id);
38        }
39    }
40}
41
42impl Default for InMemoryDiscovery {
43    fn default() -> Self {
44        Self::new()
45    }
46}
47
48#[async_trait]
49impl Discovery for InMemoryDiscovery {
50    async fn resolve(&self, service: &str) -> AppResult<Vec<ServiceInstance>> {
51        let map = self.instances.read().await;
52        Ok(map.get(service).cloned().unwrap_or_default())
53    }
54}
55
56#[async_trait]
57impl Registry for InMemoryDiscovery {
58    async fn register(&self, instance: &ServiceInstance) -> AppResult<()> {
59        self.add(&instance.name, instance.clone()).await;
60        Ok(())
61    }
62
63    async fn deregister(&self, id: &str) -> AppResult<()> {
64        let mut map = self.instances.write().await;
65        let mut found = false;
66        for instances in map.values_mut() {
67            let before = instances.len();
68            instances.retain(|i| i.id != id);
69            if instances.len() != before {
70                found = true;
71            }
72        }
73        if found {
74            Ok(())
75        } else {
76            Err(AppError::not_found("instance", Some(id)))
77        }
78    }
79}
80
81#[cfg(test)]
82mod tests {
83    use super::*;
84    use crate::traits::Registry;
85
86    fn instance(id: &str, name: &str) -> ServiceInstance {
87        ServiceInstance {
88            id: id.into(),
89            name: name.into(),
90            address: "127.0.0.1".into(),
91            port: 8080,
92            healthy: true,
93            weight: 1,
94            tags: vec![],
95            metadata: HashMap::new(),
96        }
97    }
98
99    #[tokio::test]
100    async fn register_and_resolve() {
101        let disco = InMemoryDiscovery::new();
102        let inst = instance("a", "svc");
103        disco.register(&inst).await.unwrap();
104
105        let resolved = disco.resolve("svc").await.unwrap();
106        assert_eq!(resolved.len(), 1);
107        assert_eq!(resolved[0].id, "a");
108    }
109
110    #[tokio::test]
111    async fn deregister_removes() {
112        let disco = InMemoryDiscovery::new();
113        disco.register(&instance("a", "svc")).await.unwrap();
114        disco.register(&instance("b", "svc")).await.unwrap();
115
116        disco.deregister("a").await.unwrap();
117        let resolved = disco.resolve("svc").await.unwrap();
118        assert_eq!(resolved.len(), 1);
119        assert_eq!(resolved[0].id, "b");
120    }
121
122    #[tokio::test]
123    async fn deregister_not_found() {
124        let disco = InMemoryDiscovery::new();
125        assert!(disco.deregister("nope").await.is_err());
126    }
127
128    #[tokio::test]
129    async fn resolve_empty() {
130        let disco = InMemoryDiscovery::new();
131        let resolved = disco.resolve("unknown").await.unwrap();
132        assert!(resolved.is_empty());
133    }
134}
rskit_discovery/memory.rs

rskit_discovery/
memory.rs
