Skip to main content

rnacos/naming/
service_index.rs

1use super::model::ServiceKey;
2use crate::common::model::privilege::NamespacePrivilegeGroup;
3use crate::common::string_utils::StringUtils;
4use crate::namespace::model::{NamespaceActorReq, WeakNamespaceFromType, WeakNamespaceParam};
5use crate::namespace::NamespaceActor;
6use actix::Addr;
7use serde::{Deserialize, Serialize};
8use std::{
9    collections::{BTreeMap, BTreeSet},
10    sync::Arc,
11};
12
13#[derive(Debug, Clone, Default, Serialize, Deserialize)]
14pub struct ServiceQueryParam {
15    pub namespace_id: Option<Arc<String>>,
16    pub group: Option<Arc<String>>,
17    pub service: Option<Arc<String>>,
18    pub like_group: Option<String>,
19    pub like_service: Option<String>,
20    pub namespace_privilege: NamespacePrivilegeGroup,
21    pub offset: usize,
22    pub limit: usize,
23}
24
25impl ServiceQueryParam {
26    pub fn match_namespace_id(&self, g: &Arc<String>) -> bool {
27        if let Some(namespace_id) = &self.namespace_id {
28            namespace_id.is_empty() || StringUtils::eq(g, namespace_id)
29        } else {
30            true
31        }
32    }
33
34    pub fn match_group(&self, g: &Arc<String>) -> bool {
35        if let Some(group) = &self.group {
36            group.is_empty() || StringUtils::eq(g, group)
37        } else if let Some(like_group) = &self.like_group {
38            like_group.is_empty() || StringUtils::like(g, like_group).is_some()
39        } else {
40            true
41        }
42    }
43    pub fn match_service(&self, s: &Arc<String>) -> bool {
44        if let Some(service) = &self.service {
45            service.is_empty() || StringUtils::eq(s, service)
46        } else if let Some(like_service) = &self.like_service {
47            like_service.is_empty() || StringUtils::like(s, like_service).is_some()
48        } else {
49            true
50        }
51    }
52}
53
54#[derive(Debug, Clone, Default)]
55pub struct ServiceIndex {
56    //pub namespace_id:Arc<String>,
57    pub(crate) group_service: BTreeMap<Arc<String>, BTreeSet<Arc<String>>>,
58    pub(crate) service_size: usize,
59}
60
61impl ServiceIndex {
62    pub(crate) fn new() -> Self {
63        Default::default()
64    }
65
66    pub(crate) fn insert_service(&mut self, group: Arc<String>, service: Arc<String>) -> bool {
67        if let Some(set) = self.group_service.get_mut(&group) {
68            if set.insert(service) {
69                self.service_size += 1;
70                return true;
71            }
72        } else {
73            let mut set = BTreeSet::new();
74            set.insert(service);
75            self.group_service.insert(group, set);
76            self.service_size += 1;
77            return true;
78        }
79        false
80    }
81
82    pub(crate) fn remove_service(
83        &mut self,
84        group: &Arc<String>,
85        service: &Arc<String>,
86    ) -> (bool, usize) {
87        let b = if let Some(set) = self.group_service.get_mut(group) {
88            let b = set.remove(service);
89            if b {
90                self.service_size -= 1;
91                if set.is_empty() {
92                    self.group_service.remove(group);
93                }
94            }
95            b
96        } else {
97            false
98        };
99        (b, self.group_service.len())
100    }
101
102    /*
103    pub(crate) fn query_service_list(&self,namespace_id:&Arc<String>,group_key:&Arc<String>,service_key:&Arc<String>)
104        -> Vec<ServiceKey> {
105        let mut rlist=vec![];
106        for (g,set) in &self.group_service{
107            if StringUtils::is_empty(group_key) || StringUtils::like(g, &group_key).is_some() {
108                for s in set {
109                    if StringUtils::is_empty(service_key) || StringUtils::like(s, &service_key).is_some() {
110                        let service_key = ServiceKey::new_by_arc(namespace_id.clone(),g.clone(),s.clone());
111                        rlist.push(service_key);
112                    }
113                }
114            }
115        }
116        rlist
117    }
118    */
119
120    pub(crate) fn query_service_page(
121        &self,
122        namespace_id: &Arc<String>,
123        limit: usize,
124        param: &ServiceQueryParam,
125    ) -> (usize, Vec<ServiceKey>) {
126        let mut rlist = vec![];
127        let end_index = param.offset + limit;
128        let mut index = 0;
129        for (g, set) in &self.group_service {
130            if param.match_group(g) {
131                for s in set {
132                    if param.match_service(s) {
133                        if index >= param.offset && index < end_index {
134                            let service_key =
135                                ServiceKey::new_by_arc(namespace_id.clone(), g.clone(), s.clone());
136                            rlist.push(service_key);
137                        }
138                        index += 1;
139                    }
140                }
141            }
142        }
143        (index, rlist)
144    }
145
146    pub fn get_service_count(&self) -> usize {
147        let mut sum = 0;
148        for set in self.group_service.values() {
149            sum += set.len();
150        }
151        sum
152    }
153}
154
155#[derive(Debug, Clone, Default)]
156pub struct NamespaceIndex {
157    pub namespace_group: BTreeMap<Arc<String>, ServiceIndex>,
158    pub service_size: usize,
159    pub(crate) namespace_actor: Option<Addr<NamespaceActor>>,
160}
161
162impl NamespaceIndex {
163    pub fn new() -> Self {
164        Default::default()
165    }
166
167    pub fn insert_service(&mut self, key: ServiceKey) -> bool {
168        self.do_insert_service(key.namespace_id, key.group_name, key.service_name)
169    }
170
171    pub fn remove_service(&mut self, key: &ServiceKey) -> bool {
172        self.do_remove_service(&key.namespace_id, &key.group_name, &key.service_name)
173    }
174
175    pub(crate) fn do_insert_service(
176        &mut self,
177        namespace_id: Arc<String>,
178        group: Arc<String>,
179        service: Arc<String>,
180    ) -> bool {
181        let mut result = false;
182        if let Some(service_index) = self.namespace_group.get_mut(&namespace_id) {
183            if service_index.insert_service(group, service) {
184                self.service_size += 1;
185                result = true;
186            }
187        } else {
188            let mut service_index = ServiceIndex::new();
189            if service_index.insert_service(group, service) {
190                self.service_size += 1;
191                result = true;
192            }
193            self.notify_namespace_change(
194                WeakNamespaceParam {
195                    namespace_id: namespace_id.clone(),
196                    from_type: WeakNamespaceFromType::Naming,
197                },
198                false,
199            );
200            self.namespace_group.insert(namespace_id, service_index);
201        }
202        result
203    }
204
205    pub(crate) fn do_remove_service(
206        &mut self,
207        namespace_id: &Arc<String>,
208        group: &Arc<String>,
209        service: &Arc<String>,
210    ) -> bool {
211        let mut result = false;
212        if let Some(service_index) = self.namespace_group.get_mut(namespace_id) {
213            let (b, group_size) = service_index.remove_service(group, service);
214            if b {
215                self.service_size -= 1;
216                result = true;
217            }
218            if group_size == 0 {
219                self.notify_namespace_change(
220                    WeakNamespaceParam {
221                        namespace_id: namespace_id.clone(),
222                        from_type: WeakNamespaceFromType::Naming,
223                    },
224                    true,
225                );
226                self.namespace_group.remove(namespace_id);
227            }
228        }
229        result
230    }
231
232    pub fn query_service_page(&self, param: &ServiceQueryParam) -> (usize, Vec<ServiceKey>) {
233        let mut rlist = vec![];
234        let mut size = 0;
235        let mut limit = param.limit;
236        if let Some(namespace_id) = &param.namespace_id {
237            if param.namespace_privilege.check_permission(namespace_id) {
238                if let Some(index) = self.namespace_group.get(namespace_id) {
239                    return index.query_service_page(namespace_id, limit, param);
240                }
241            }
242        } else {
243            for (namespace_id, service_index) in &self.namespace_group {
244                if param.namespace_privilege.check_permission(namespace_id) {
245                    let (sub_size, mut sub_list) =
246                        service_index.query_service_page(namespace_id, limit, param);
247                    size += sub_size;
248                    limit -= sub_list.len();
249                    rlist.append(&mut sub_list);
250                }
251            }
252        }
253        (size, rlist)
254    }
255
256    fn notify_namespace_change(&self, param: WeakNamespaceParam, is_remove: bool) {
257        if let Some(act) = &self.namespace_actor {
258            if is_remove {
259                act.do_send(NamespaceActorReq::RemoveWeak(param));
260            } else {
261                act.do_send(NamespaceActorReq::SetWeak(param));
262            }
263        }
264    }
265
266    pub fn get_tenant_count(&self) -> usize {
267        self.namespace_group.len()
268    }
269
270    pub fn get_service_count(&self) -> (usize, usize) {
271        let mut group_sum = 0;
272        let mut sum = 0;
273        for service in self.namespace_group.values() {
274            group_sum += service.group_service.len();
275            sum += service.get_service_count();
276        }
277        (group_sum, sum)
278    }
279}
280
281#[test]
282fn add_service() {
283    let mut index = NamespaceIndex::new();
284    let key1 = ServiceKey::new("1", "1", "1");
285    let key2 = ServiceKey::new("1", "1", "2");
286    let key3 = ServiceKey::new("1", "2", "1");
287    index.insert_service(key1.clone());
288    assert!(index.service_size == 1);
289    index.insert_service(key2.clone());
290    assert!(index.service_size == 2);
291    index.insert_service(key3.clone());
292    assert!(index.service_size == 3);
293    assert!(index.namespace_group.len() == 1);
294    //remove
295    index.remove_service(&key1);
296    assert!(index.service_size == 2);
297    index.remove_service(&key1);
298    assert!(index.service_size == 2);
299    index.remove_service(&key2);
300    assert!(index.service_size == 1);
301    index.remove_service(&key3);
302    index.remove_service(&key3);
303    assert!(index.service_size == 0);
304    assert!(index.namespace_group.is_empty());
305}
306
307#[test]
308fn query_service() {
309    let mut index = NamespaceIndex::new();
310    let key1 = ServiceKey::new("1", "1", "1");
311    let key2 = ServiceKey::new("1", "1", "2");
312    let key3 = ServiceKey::new("1", "2", "1");
313    let key4 = ServiceKey::new("1", "2", "2");
314    let key5 = ServiceKey::new("2", "1", "1");
315    index.insert_service(key1.clone());
316    index.insert_service(key2.clone());
317    index.insert_service(key3.clone());
318    index.insert_service(key4.clone());
319    index.insert_service(key5.clone());
320
321    let mut param = ServiceQueryParam {
322        namespace_id: None,
323        limit: 0xffff_ffff,
324        ..ServiceQueryParam::default()
325    };
326    let (size, list) = index.query_service_page(&param);
327    assert!(size == 5);
328    assert!(list.len() == 5);
329
330    param.limit = 2;
331    let (size, list) = index.query_service_page(&param);
332    assert!(size == 5);
333    assert!(list.len() == 2);
334
335    index.remove_service(&key1);
336    index.remove_service(&key2);
337    index.remove_service(&key3);
338    index.remove_service(&key4);
339    index.remove_service(&key5);
340
341    param.limit = 2;
342    let (size, list) = index.query_service_page(&param);
343    assert!(size == 0);
344    assert!(list.is_empty());
345}