batata_client/naming/
service_info.rs1
2use dashmap::DashMap;
3
4use crate::api::naming::{Instance, Service};
5use crate::common::build_service_key;
6
7pub struct ServiceInfoCache {
9 services: DashMap<String, Service>,
11}
12
13impl ServiceInfoCache {
14 pub fn new() -> Self {
15 Self {
16 services: DashMap::new(),
17 }
18 }
19
20 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 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 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 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 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 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 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 pub fn keys(&self) -> Vec<String> {
93 self.services.iter().map(|r| r.key().clone()).collect()
94 }
95
96 pub fn clear(&self) {
98 self.services.clear();
99 }
100
101 pub fn len(&self) -> usize {
103 self.services.len()
104 }
105
106 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}