scirs2_core/distributed/cluster/
registry.rs

1//! Node registry for tracking cluster members
2//!
3//! This module provides the NodeRegistry for managing cluster node
4//! information, status tracking, and node health monitoring.
5
6use crate::error::CoreResult;
7use std::collections::HashMap;
8
9use super::types::{NodeInfo, NodeStatus};
10
11/// Node registry for tracking cluster members
12#[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    /// Create a new node registry
26    pub fn new() -> Self {
27        Self {
28            nodes: HashMap::new(),
29            node_status: HashMap::new(),
30        }
31    }
32
33    /// Register a new node in the cluster
34    /// Returns true if the node is new, false if it was already registered
35    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    /// Get all registered nodes
44    pub fn get_all_nodes(&self) -> Vec<NodeInfo> {
45        self.nodes.values().cloned().collect()
46    }
47
48    /// Get only healthy nodes
49    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    /// Get nodes with a specific status
58    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    /// Get a specific node by ID
67    pub fn get_node(&self, nodeid: &str) -> Option<&NodeInfo> {
68        self.nodes.get(nodeid)
69    }
70
71    /// Get the status of a specific node
72    pub fn get_node_status(&self, nodeid: &str) -> Option<NodeStatus> {
73        self.node_status.get(nodeid).copied()
74    }
75
76    /// Update the status of a specific node
77    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    /// Remove a node from the registry
86    pub fn remove_node(&mut self, nodeid: &str) -> Option<NodeInfo> {
87        self.node_status.remove(nodeid);
88        self.nodes.remove(nodeid)
89    }
90
91    /// Get the total number of registered nodes
92    pub fn node_count(&self) -> usize {
93        self.nodes.len()
94    }
95
96    /// Get the number of healthy nodes
97    pub fn healthy_node_count(&self) -> usize {
98        self.get_healthy_nodes().len()
99    }
100
101    /// Check if a node is registered
102    pub fn contains_node(&self, nodeid: &str) -> bool {
103        self.nodes.contains_key(nodeid)
104    }
105
106    /// Clear all nodes from the registry
107    pub fn clear(&mut self) {
108        self.nodes.clear();
109        self.node_status.clear();
110    }
111
112    /// Get nodes that haven't been seen recently
113    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    /// Update a node's last seen timestamp
123    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}