scirs2_core/distributed/cluster/
registry.rs1use crate::error::CoreResult;
7use std::collections::HashMap;
8
9use super::types::{NodeInfo, NodeStatus};
10
11#[derive(Debug)]
13pub struct NodeRegistry {
14 nodes: HashMap<String, NodeInfo>,
15 node_status: HashMap<String, NodeStatus>,
16}
17
18impl Default for NodeRegistry {
19 fn default() -> Self {
20 Self::new()
21 }
22}
23
24impl NodeRegistry {
25 pub fn new() -> Self {
27 Self {
28 nodes: HashMap::new(),
29 node_status: HashMap::new(),
30 }
31 }
32
33 pub fn register_node(&mut self, nodeinfo: NodeInfo) -> CoreResult<bool> {
36 let is_new = !self.nodes.contains_key(&nodeinfo.id);
37 self.nodes.insert(nodeinfo.id.clone(), nodeinfo.clone());
38 self.node_status
39 .insert(nodeinfo.id.clone(), nodeinfo.status);
40 Ok(is_new)
41 }
42
43 pub fn get_all_nodes(&self) -> Vec<NodeInfo> {
45 self.nodes.values().cloned().collect()
46 }
47
48 pub fn get_healthy_nodes(&self) -> Vec<NodeInfo> {
50 self.nodes
51 .values()
52 .filter(|node| self.node_status.get(&node.id) == Some(&NodeStatus::Healthy))
53 .cloned()
54 .collect()
55 }
56
57 pub fn get_nodes_by_status(&self, status: NodeStatus) -> Vec<NodeInfo> {
59 self.nodes
60 .values()
61 .filter(|node| self.node_status.get(&node.id) == Some(&status))
62 .cloned()
63 .collect()
64 }
65
66 pub fn get_node(&self, nodeid: &str) -> Option<&NodeInfo> {
68 self.nodes.get(nodeid)
69 }
70
71 pub fn get_node_status(&self, nodeid: &str) -> Option<NodeStatus> {
73 self.node_status.get(nodeid).copied()
74 }
75
76 pub fn update_node_status(&mut self, nodeid: &str, status: NodeStatus) -> CoreResult<()> {
78 if let Some(node) = self.nodes.get_mut(nodeid) {
79 node.status = status;
80 self.node_status.insert(nodeid.to_string(), status);
81 }
82 Ok(())
83 }
84
85 pub fn remove_node(&mut self, nodeid: &str) -> Option<NodeInfo> {
87 self.node_status.remove(nodeid);
88 self.nodes.remove(nodeid)
89 }
90
91 pub fn node_count(&self) -> usize {
93 self.nodes.len()
94 }
95
96 pub fn healthy_node_count(&self) -> usize {
98 self.get_healthy_nodes().len()
99 }
100
101 pub fn contains_node(&self, nodeid: &str) -> bool {
103 self.nodes.contains_key(nodeid)
104 }
105
106 pub fn clear(&mut self) {
108 self.nodes.clear();
109 self.node_status.clear();
110 }
111
112 pub fn get_stale_nodes(&self, max_age: std::time::Duration) -> Vec<NodeInfo> {
114 let now = std::time::Instant::now();
115 self.nodes
116 .values()
117 .filter(|node| now.duration_since(node.last_seen) > max_age)
118 .cloned()
119 .collect()
120 }
121
122 pub fn update_node_last_seen(&mut self, nodeid: &str) {
124 if let Some(node) = self.nodes.get_mut(nodeid) {
125 node.last_seen = std::time::Instant::now();
126 }
127 }
128}