mocra_core/common/
registry.rs

1use crate::cacheable::{CacheAble, CacheService};
2use serde::{Deserialize, Serialize};
3use std::sync::Arc;
4use std::time::Duration;
5
6/// Information about a registered node in the distributed system.
7#[derive(Debug, Clone, Serialize, Deserialize)]
8pub struct NodeInfo {
9    /// Unique identifier for the node
10    pub id: String,
11    /// IP address of the node
12    pub ip: String,
13    /// Hostname of the node
14    pub hostname: String,
15    /// Timestamp of the last heartbeat
16    pub last_heartbeat: u64,
17    /// Software version of the node
18    pub version: String,
19}
20
21impl CacheAble for NodeInfo {
22    fn field() -> impl AsRef<str> {
23        "nodes".to_string()
24    }
25}
26
27/// Registry for managing active nodes in the cluster.
28///
29/// Handles registration, heartbeats, and discovery of other nodes via the cache service.
30pub struct NodeRegistry {
31    cache: Arc<CacheService>,
32    node_id: String,
33    ttl: Duration,
34}
35
36impl NodeRegistry {
37    /// Creates a new NodeRegistry instance.
38    pub fn new(cache: Arc<CacheService>, node_id: String, ttl: Duration) -> Self {
39        Self {
40            cache,
41            node_id,
42            ttl,
43        }
44    }
45
46    /// Sends a heartbeat to the registry to indicate this node is alive.
47    ///
48    /// Updates the node information in the cache service with a TTL.
49    pub async fn heartbeat(
50        &self,
51        ip: &str,
52        hostname: &str,
53        version: &str,
54    ) -> Result<(), crate::errors::Error> {
55        let now = std::time::SystemTime::now()
56            .duration_since(std::time::UNIX_EPOCH)
57            .unwrap()
58            .as_secs();
59        let info = NodeInfo {
60            id: self.node_id.clone(),
61            ip: ip.to_string(),
62            hostname: hostname.to_string(),
63            last_heartbeat: now,
64            version: version.to_string(),
65        };
66
67        // 1. Save detailed info with TTL (Key retrieval)
68        info.send_with_ttl(&self.node_id, &self.cache, self.ttl)
69            .await
70            .map_err(crate::errors::Error::from)?;
71
72        // 2. Update ZSET index for efficient querying (Range retrieval)
73        // Key: namespace:registry:nodes_index
74        let index_key = format!("{}:registry:nodes_index", self.cache.namespace());
75        let _ = self
76            .cache
77            .zadd(&index_key, now as f64, self.node_id.as_bytes())
78            .await
79            .map_err(crate::errors::Error::from)?;
80
81        Ok(())
82    }
83
84    /// Retrieves a list of all currently active nodes in the cluster.
85    pub async fn get_active_nodes(&self) -> Result<Vec<NodeInfo>, crate::errors::Error> {
86        let now = std::time::SystemTime::now()
87            .duration_since(std::time::UNIX_EPOCH)
88            .unwrap()
89            .as_secs();
90        let cutoff = now.saturating_sub(self.ttl.as_secs());
91        let index_key = format!("{}:registry:nodes_index", self.cache.namespace());
92
93        // 1. Lazy Cleanup: Remove expired nodes from index
94        // We do this on read or via a separate maintenance task. Doing on read is fine for registry.
95        let _ = self
96            .cache
97            .zremrangebyscore(&index_key, f64::NEG_INFINITY, cutoff as f64)
98            .await;
99
100        // 2. Get active Node IDs
101        let active_ids_bytes = self
102            .cache
103            .zrangebyscore(&index_key, cutoff as f64, f64::INFINITY)
104            .await
105            .map_err(crate::errors::Error::from)?;
106
107        if active_ids_bytes.is_empty() {
108            return Ok(Vec::new());
109        }
110
111        let mut keys = Vec::with_capacity(active_ids_bytes.len());
112        for id_bytes in active_ids_bytes {
113            let id = String::from_utf8(id_bytes).unwrap_or_default();
114            if !id.is_empty() {
115                // Reconstruct the key used by NodeInfo::send_with_ttl
116                // Format: {namespace}:nodes:{id}
117                keys.push(format!("{}:nodes:{}", self.cache.namespace(), id));
118            }
119        }
120
121        if keys.is_empty() {
122            return Ok(Vec::new());
123        }
124
125        let key_refs: Vec<&str> = keys.iter().map(|s| s.as_str()).collect();
126        let results = self
127            .cache
128            .mget(&key_refs)
129            .await
130            .map_err(crate::errors::Error::from)?;
131
132        let mut nodes = Vec::with_capacity(results.len());
133        for bytes in results.into_iter().flatten() {
134            if let Ok(node) = serde_json::from_slice::<NodeInfo>(&bytes) {
135                nodes.push(node);
136            }
137        }
138        Ok(nodes)
139    }
140
141    pub async fn deregister(&self) -> Result<(), crate::errors::Error> {
142        NodeInfo::delete(&self.node_id, &self.cache)
143            .await
144            .map_err(crate::errors::Error::from)
145    }
146}
mocra_core/common/registry.rs

mocra_core/common/
registry.rs
