oximedia_distributed/
node_registry.rs

1//! Node registry for cluster membership management.
2//!
3//! Tracks all known nodes in the distributed cluster, their capabilities,
4//! roles, and current status.  Provides efficient lookup, filtering by role,
5//! and node lifecycle management.
6
7#![allow(dead_code)]
8
9// ---------------------------------------------------------------------------
10// NodeRole
11// ---------------------------------------------------------------------------
12
13/// The role a node plays in the cluster.
14#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
15pub enum NodeRole {
16    /// Central coordinator node.
17    Coordinator,
18    /// General-purpose compute/encoding worker.
19    Worker,
20    /// Storage node (not involved in compute).
21    Storage,
22    /// Gateway / load-balancer node.
23    Gateway,
24}
25
26impl NodeRole {
27    /// Returns the human-readable name of this role.
28    #[must_use]
29    pub fn as_str(self) -> &'static str {
30        match self {
31            Self::Coordinator => "coordinator",
32            Self::Worker => "worker",
33            Self::Storage => "storage",
34            Self::Gateway => "gateway",
35        }
36    }
37}
38
39// ---------------------------------------------------------------------------
40// NodeStatus
41// ---------------------------------------------------------------------------
42
43/// Operational status of a registered node.
44#[derive(Debug, Clone, Copy, PartialEq, Eq)]
45pub enum NodeStatus {
46    /// Node is healthy and available.
47    Healthy,
48    /// Node is available but operating in a degraded state.
49    Degraded,
50    /// Node is temporarily suspended (maintenance mode).
51    Suspended,
52    /// Node has been removed from the active cluster.
53    Removed,
54}
55
56impl NodeStatus {
57    /// Returns `true` if the node can accept work.
58    #[must_use]
59    pub fn is_available(self) -> bool {
60        matches!(self, Self::Healthy | Self::Degraded)
61    }
62}
63
64// ---------------------------------------------------------------------------
65// NodeInfo
66// ---------------------------------------------------------------------------
67
68/// Information about a registered cluster node.
69#[derive(Debug, Clone)]
70pub struct NodeInfo {
71    /// Unique node identifier.
72    pub id: String,
73    /// Network address (host:port).
74    pub address: String,
75    /// Assigned role in the cluster.
76    pub role: NodeRole,
77    /// Current operational status.
78    pub status: NodeStatus,
79    /// Number of CPU cores.
80    pub cpu_cores: u32,
81    /// Total RAM in megabytes.
82    pub memory_mb: u32,
83    /// Whether this node has GPU acceleration.
84    pub has_gpu: bool,
85    /// Unix epoch ms when the node was registered.
86    pub registered_at_ms: u64,
87    /// Unix epoch ms of the last status update.
88    pub last_seen_ms: u64,
89    /// Arbitrary tags (e.g., "av1", "region:eu-west-1").
90    pub tags: Vec<String>,
91}
92
93impl NodeInfo {
94    /// Create a new healthy node registration.
95    #[must_use]
96    pub fn new(
97        id: impl Into<String>,
98        address: impl Into<String>,
99        role: NodeRole,
100        cpu_cores: u32,
101        memory_mb: u32,
102        has_gpu: bool,
103        now_ms: u64,
104    ) -> Self {
105        Self {
106            id: id.into(),
107            address: address.into(),
108            role,
109            status: NodeStatus::Healthy,
110            cpu_cores,
111            memory_mb,
112            has_gpu,
113            registered_at_ms: now_ms,
114            last_seen_ms: now_ms,
115            tags: Vec::new(),
116        }
117    }
118
119    /// Add a tag to this node.
120    pub fn with_tag(mut self, tag: impl Into<String>) -> Self {
121        self.tags.push(tag.into());
122        self
123    }
124
125    /// Update the last-seen timestamp and optionally the status.
126    pub fn touch(&mut self, now_ms: u64, status: NodeStatus) {
127        self.last_seen_ms = now_ms;
128        self.status = status;
129    }
130
131    /// Returns `true` if the node can accept work.
132    #[must_use]
133    pub fn is_available(&self) -> bool {
134        self.status.is_available()
135    }
136
137    /// Returns `true` if the node has a specific tag.
138    #[must_use]
139    pub fn has_tag(&self, tag: &str) -> bool {
140        self.tags.iter().any(|t| t == tag)
141    }
142
143    /// Age of the registration in milliseconds.
144    #[must_use]
145    pub fn age_ms(&self, now_ms: u64) -> u64 {
146        now_ms.saturating_sub(self.registered_at_ms)
147    }
148}
149
150// ---------------------------------------------------------------------------
151// NodeRegistry
152// ---------------------------------------------------------------------------
153
154/// Central registry of all cluster nodes.
155///
156/// Supports fast lookup by ID and filtered queries by role/status/tag.
157#[derive(Debug, Default)]
158pub struct NodeRegistry {
159    nodes: Vec<NodeInfo>,
160}
161
162impl NodeRegistry {
163    /// Create an empty registry.
164    #[must_use]
165    pub fn new() -> Self {
166        Self::default()
167    }
168
169    /// Register a new node.  If a node with the same ID already exists, it is
170    /// replaced.
171    pub fn register(&mut self, node: NodeInfo) {
172        if let Some(existing) = self.nodes.iter_mut().find(|n| n.id == node.id) {
173            *existing = node;
174        } else {
175            self.nodes.push(node);
176        }
177    }
178
179    /// Remove a node by ID, returning it if found.
180    pub fn deregister(&mut self, id: &str) -> Option<NodeInfo> {
181        if let Some(pos) = self.nodes.iter().position(|n| n.id == id) {
182            Some(self.nodes.remove(pos))
183        } else {
184            None
185        }
186    }
187
188    /// Look up a node by ID.
189    #[must_use]
190    pub fn get(&self, id: &str) -> Option<&NodeInfo> {
191        self.nodes.iter().find(|n| n.id == id)
192    }
193
194    /// Mutable access to a node by ID.
195    #[must_use]
196    pub fn get_mut(&mut self, id: &str) -> Option<&mut NodeInfo> {
197        self.nodes.iter_mut().find(|n| n.id == id)
198    }
199
200    /// Total number of registered nodes.
201    #[must_use]
202    pub fn len(&self) -> usize {
203        self.nodes.len()
204    }
205
206    /// Returns `true` if the registry has no nodes.
207    #[must_use]
208    pub fn is_empty(&self) -> bool {
209        self.nodes.is_empty()
210    }
211
212    /// All nodes with a given role.
213    #[must_use]
214    pub fn by_role(&self, role: NodeRole) -> Vec<&NodeInfo> {
215        self.nodes.iter().filter(|n| n.role == role).collect()
216    }
217
218    /// All nodes with a given status.
219    #[must_use]
220    pub fn by_status(&self, status: NodeStatus) -> Vec<&NodeInfo> {
221        self.nodes.iter().filter(|n| n.status == status).collect()
222    }
223
224    /// All available (Healthy or Degraded) worker nodes.
225    #[must_use]
226    pub fn available_workers(&self) -> Vec<&NodeInfo> {
227        self.nodes
228            .iter()
229            .filter(|n| n.role == NodeRole::Worker && n.is_available())
230            .collect()
231    }
232
233    /// All nodes that have a specific tag.
234    #[must_use]
235    pub fn by_tag(&self, tag: &str) -> Vec<&NodeInfo> {
236        self.nodes.iter().filter(|n| n.has_tag(tag)).collect()
237    }
238
239    /// Mark a node as removed (soft delete).
240    pub fn remove_node(&mut self, id: &str, now_ms: u64) {
241        if let Some(n) = self.get_mut(id) {
242            n.touch(now_ms, NodeStatus::Removed);
243        }
244    }
245
246    /// Evict nodes that have not been seen within `ttl_ms` of `now_ms`.
247    ///
248    /// Returns the IDs of evicted nodes.
249    pub fn evict_stale(&mut self, now_ms: u64, ttl_ms: u64) -> Vec<String> {
250        let cutoff = now_ms.saturating_sub(ttl_ms);
251        let mut evicted = Vec::new();
252        self.nodes.retain(|n| {
253            if n.last_seen_ms < cutoff {
254                evicted.push(n.id.clone());
255                false
256            } else {
257                true
258            }
259        });
260        evicted
261    }
262
263    /// Summary counts: `(total, healthy, degraded, suspended, removed)`.
264    #[must_use]
265    pub fn status_summary(&self) -> (usize, usize, usize, usize, usize) {
266        let total = self.nodes.len();
267        let healthy = self.by_status(NodeStatus::Healthy).len();
268        let degraded = self.by_status(NodeStatus::Degraded).len();
269        let suspended = self.by_status(NodeStatus::Suspended).len();
270        let removed = self.by_status(NodeStatus::Removed).len();
271        (total, healthy, degraded, suspended, removed)
272    }
273}
274
275// ---------------------------------------------------------------------------
276// Tests
277// ---------------------------------------------------------------------------
278
279#[cfg(test)]
280mod tests {
281    use super::*;
282
283    fn worker(id: &str, now_ms: u64) -> NodeInfo {
284        NodeInfo::new(
285            id,
286            format!("10.0.0.1:5000"),
287            NodeRole::Worker,
288            8,
289            16_384,
290            false,
291            now_ms,
292        )
293    }
294
295    fn gpu_worker(id: &str, now_ms: u64) -> NodeInfo {
296        NodeInfo::new(
297            id,
298            format!("10.0.0.2:5000"),
299            NodeRole::Worker,
300            16,
301            32_768,
302            true,
303            now_ms,
304        )
305        .with_tag("gpu")
306    }
307
308    fn coordinator(id: &str, now_ms: u64) -> NodeInfo {
309        NodeInfo::new(
310            id,
311            format!("10.0.0.3:5000"),
312            NodeRole::Coordinator,
313            4,
314            8_192,
315            false,
316            now_ms,
317        )
318    }
319
320    // ── NodeRole ─────────────────────────────────────────────────────────
321
322    #[test]
323    fn test_node_role_as_str() {
324        assert_eq!(NodeRole::Worker.as_str(), "worker");
325        assert_eq!(NodeRole::Coordinator.as_str(), "coordinator");
326        assert_eq!(NodeRole::Storage.as_str(), "storage");
327        assert_eq!(NodeRole::Gateway.as_str(), "gateway");
328    }
329
330    // ── NodeStatus ───────────────────────────────────────────────────────
331
332    #[test]
333    fn test_healthy_is_available() {
334        assert!(NodeStatus::Healthy.is_available());
335    }
336
337    #[test]
338    fn test_degraded_is_available() {
339        assert!(NodeStatus::Degraded.is_available());
340    }
341
342    #[test]
343    fn test_suspended_not_available() {
344        assert!(!NodeStatus::Suspended.is_available());
345    }
346
347    #[test]
348    fn test_removed_not_available() {
349        assert!(!NodeStatus::Removed.is_available());
350    }
351
352    // ── NodeInfo ─────────────────────────────────────────────────────────
353
354    #[test]
355    fn test_node_info_initial_status_healthy() {
356        let n = worker("n0", 1000);
357        assert_eq!(n.status, NodeStatus::Healthy);
358    }
359
360    #[test]
361    fn test_node_info_has_tag() {
362        let n = gpu_worker("n0", 1000);
363        assert!(n.has_tag("gpu"));
364        assert!(!n.has_tag("av1"));
365    }
366
367    #[test]
368    fn test_node_info_age_ms() {
369        let n = worker("n0", 1000);
370        assert_eq!(n.age_ms(3000), 2000);
371    }
372
373    #[test]
374    fn test_node_info_touch_updates_status() {
375        let mut n = worker("n0", 1000);
376        n.touch(2000, NodeStatus::Degraded);
377        assert_eq!(n.status, NodeStatus::Degraded);
378        assert_eq!(n.last_seen_ms, 2000);
379    }
380
381    // ── NodeRegistry ─────────────────────────────────────────────────────
382
383    #[test]
384    fn test_registry_empty() {
385        let reg = NodeRegistry::new();
386        assert!(reg.is_empty());
387        assert_eq!(reg.len(), 0);
388    }
389
390    #[test]
391    fn test_registry_register_and_get() {
392        let mut reg = NodeRegistry::new();
393        reg.register(worker("n0", 1000));
394        let n = reg.get("n0").expect("get should return a value");
395        assert_eq!(n.id, "n0");
396    }
397
398    #[test]
399    fn test_registry_register_replaces_existing() {
400        let mut reg = NodeRegistry::new();
401        reg.register(worker("n0", 1000));
402        reg.register(gpu_worker("n0", 2000)); // same ID → replace
403        assert_eq!(reg.len(), 1);
404        assert!(reg.get("n0").expect("get should return a value").has_gpu);
405    }
406
407    #[test]
408    fn test_registry_deregister() {
409        let mut reg = NodeRegistry::new();
410        reg.register(worker("n0", 1000));
411        let removed = reg.deregister("n0");
412        assert!(removed.is_some());
413        assert!(reg.is_empty());
414    }
415
416    #[test]
417    fn test_registry_by_role() {
418        let mut reg = NodeRegistry::new();
419        reg.register(worker("w0", 1000));
420        reg.register(worker("w1", 1000));
421        reg.register(coordinator("c0", 1000));
422        assert_eq!(reg.by_role(NodeRole::Worker).len(), 2);
423        assert_eq!(reg.by_role(NodeRole::Coordinator).len(), 1);
424    }
425
426    #[test]
427    fn test_registry_available_workers() {
428        let mut reg = NodeRegistry::new();
429        let mut w1 = worker("w0", 1000);
430        w1.status = NodeStatus::Suspended;
431        reg.register(w1);
432        reg.register(worker("w1", 1000));
433        assert_eq!(reg.available_workers().len(), 1);
434    }
435
436    #[test]
437    fn test_registry_by_tag() {
438        let mut reg = NodeRegistry::new();
439        reg.register(gpu_worker("g0", 1000));
440        reg.register(worker("w0", 1000));
441        assert_eq!(reg.by_tag("gpu").len(), 1);
442    }
443
444    #[test]
445    fn test_registry_remove_node_soft_delete() {
446        let mut reg = NodeRegistry::new();
447        reg.register(worker("w0", 1000));
448        reg.remove_node("w0", 2000);
449        assert_eq!(
450            reg.get("w0").expect("get should return a value").status,
451            NodeStatus::Removed
452        );
453        assert!(!reg
454            .get("w0")
455            .expect("get should return a value")
456            .is_available());
457    }
458
459    #[test]
460    fn test_registry_evict_stale() {
461        let mut reg = NodeRegistry::new();
462        reg.register(worker("old", 100));
463        reg.register(worker("fresh", 5000));
464        let evicted = reg.evict_stale(6000, 1000); // cutoff = 5000 → old removed
465        assert!(evicted.contains(&"old".to_string()));
466        assert_eq!(reg.len(), 1);
467    }
468
469    #[test]
470    fn test_registry_status_summary() {
471        let mut reg = NodeRegistry::new();
472        reg.register(worker("w0", 1000));
473        reg.register(worker("w1", 1000));
474        let (total, healthy, _, _, _) = reg.status_summary();
475        assert_eq!(total, 2);
476        assert_eq!(healthy, 2);
477    }
478}
oximedia_distributed/node_registry.rs

oximedia_distributed/
node_registry.rs
