batata_client/naming/
service_info.rs

1
2use dashmap::DashMap;
3
4use crate::api::naming::{Instance, Service};
5use crate::common::build_service_key;
6
7/// Service information cache
8pub struct ServiceInfoCache {
9    /// Cached services: key -> Service
10    services: DashMap<String, Service>,
11}
12
13impl ServiceInfoCache {
14    pub fn new() -> Self {
15        Self {
16            services: DashMap::new(),
17        }
18    }
19
20    /// Get service from cache
21    pub fn get(&self, namespace: &str, group_name: &str, service_name: &str) -> Option<Service> {
22        let key = build_service_key(service_name, group_name, namespace);
23        self.services.get(&key).map(|r| r.value().clone())
24    }
25
26    /// Put service into cache
27    pub fn put(&self, namespace: &str, service: Service) {
28        let key = build_service_key(&service.name, &service.group_name, namespace);
29        self.services.insert(key, service);
30    }
31
32    /// Remove service from cache
33    pub fn remove(
34        &self,
35        namespace: &str,
36        group_name: &str,
37        service_name: &str,
38    ) -> Option<Service> {
39        let key = build_service_key(service_name, group_name, namespace);
40        self.services.remove(&key).map(|(_, v)| v)
41    }
42
43    /// Check if service exists in cache
44    pub fn contains(&self, namespace: &str, group_name: &str, service_name: &str) -> bool {
45        let key = build_service_key(service_name, group_name, namespace);
46        self.services.contains_key(&key)
47    }
48
49    /// Get all healthy instances for a service
50    pub fn get_healthy_instances(
51        &self,
52        namespace: &str,
53        group_name: &str,
54        service_name: &str,
55    ) -> Vec<Instance> {
56        self.get(namespace, group_name, service_name)
57            .map(|s| s.hosts.into_iter().filter(|i| i.healthy && i.enabled).collect())
58            .unwrap_or_default()
59    }
60
61    /// Get all instances for a service
62    pub fn get_all_instances(
63        &self,
64        namespace: &str,
65        group_name: &str,
66        service_name: &str,
67    ) -> Vec<Instance> {
68        self.get(namespace, group_name, service_name)
69            .map(|s| s.hosts)
70            .unwrap_or_default()
71    }
72
73    /// Update service instances
74    pub fn update_instances(
75        &self,
76        namespace: &str,
77        group_name: &str,
78        service_name: &str,
79        instances: Vec<Instance>,
80    ) {
81        let key = build_service_key(service_name, group_name, namespace);
82        if let Some(mut service) = self.services.get_mut(&key) {
83            service.hosts = instances;
84        } else {
85            let mut service = Service::new(service_name, group_name);
86            service.hosts = instances;
87            self.services.insert(key, service);
88        }
89    }
90
91    /// Get all cached service keys
92    pub fn keys(&self) -> Vec<String> {
93        self.services.iter().map(|r| r.key().clone()).collect()
94    }
95
96    /// Clear all cached services
97    pub fn clear(&self) {
98        self.services.clear();
99    }
100
101    /// Get cache size
102    pub fn len(&self) -> usize {
103        self.services.len()
104    }
105
106    /// Check if cache is empty
107    pub fn is_empty(&self) -> bool {
108        self.services.is_empty()
109    }
110}
111
112impl Default for ServiceInfoCache {
113    fn default() -> Self {
114        Self::new()
115    }
116}
117
118#[cfg(test)]
119mod tests {
120    use super::*;
121
122    #[test]
123    fn test_service_info_cache() {
124        let cache = ServiceInfoCache::new();
125
126        let mut service = Service::new("test-service", "DEFAULT_GROUP");
127        service.hosts.push(Instance::new("127.0.0.1", 8080));
128
129        cache.put("public", service);
130
131        assert!(cache.contains("public", "DEFAULT_GROUP", "test-service"));
132
133        let instances = cache.get_all_instances("public", "DEFAULT_GROUP", "test-service");
134        assert_eq!(instances.len(), 1);
135        assert_eq!(instances[0].ip, "127.0.0.1");
136    }
137}