a3s_box_runtime/scale/
registry.rs

1//! Instance registry for multi-node deployments.
2
3use std::collections::HashMap;
4
5use chrono::{DateTime, Utc};
6use serde::{Deserialize, Serialize};
7
8/// Registry for instance self-registration in standalone multi-node deployments.
9///
10/// Each Box host registers its instances with the registry so Gateway
11/// can discover endpoints for traffic routing.
12pub struct InstanceRegistry {
13    /// Registered instances: instance_id → registration
14    pub(super) entries: HashMap<String, RegistryEntry>,
15}
16
17/// A registered instance entry.
18#[derive(Debug, Clone, Serialize, Deserialize)]
19pub(super) struct RegistryEntry {
20    pub(super) instance_id: String,
21    pub(super) service: String,
22    pub(super) endpoint: String,
23    pub(super) host_id: String,
24    pub(super) metadata: HashMap<String, String>,
25    pub(super) registered_at: DateTime<Utc>,
26    pub(super) last_heartbeat: DateTime<Utc>,
27}
28
29impl Default for InstanceRegistry {
30    fn default() -> Self {
31        Self::new()
32    }
33}
34
35impl InstanceRegistry {
36    /// Create a new empty registry.
37    pub fn new() -> Self {
38        Self {
39            entries: HashMap::new(),
40        }
41    }
42
43    /// Register an instance.
44    pub fn register(
45        &mut self,
46        instance_id: &str,
47        service: &str,
48        endpoint: &str,
49        host_id: &str,
50        metadata: HashMap<String, String>,
51    ) {
52        let now = Utc::now();
53        self.entries.insert(
54            instance_id.to_string(),
55            RegistryEntry {
56                instance_id: instance_id.to_string(),
57                service: service.to_string(),
58                endpoint: endpoint.to_string(),
59                host_id: host_id.to_string(),
60                metadata,
61                registered_at: now,
62                last_heartbeat: now,
63            },
64        );
65    }
66
67    /// Deregister an instance.
68    pub fn deregister(&mut self, instance_id: &str) -> bool {
69        self.entries.remove(instance_id).is_some()
70    }
71
72    /// Record a heartbeat for an instance.
73    pub fn heartbeat(&mut self, instance_id: &str) -> bool {
74        if let Some(entry) = self.entries.get_mut(instance_id) {
75            entry.last_heartbeat = Utc::now();
76            true
77        } else {
78            false
79        }
80    }
81
82    /// Get all endpoints for a service (for load balancing).
83    pub fn endpoints(&self, service: &str) -> Vec<String> {
84        self.entries
85            .values()
86            .filter(|e| e.service == service)
87            .map(|e| e.endpoint.clone())
88            .collect()
89    }
90
91    /// Get all instances for a service.
92    pub fn instances_for_service(&self, service: &str) -> Vec<&str> {
93        self.entries
94            .values()
95            .filter(|e| e.service == service)
96            .map(|e| e.instance_id.as_str())
97            .collect()
98    }
99
100    /// Get all instances on a specific host.
101    pub fn instances_on_host(&self, host_id: &str) -> Vec<&str> {
102        self.entries
103            .values()
104            .filter(|e| e.host_id == host_id)
105            .map(|e| e.instance_id.as_str())
106            .collect()
107    }
108
109    /// Remove stale entries that haven't sent a heartbeat within the given duration.
110    pub fn evict_stale(&mut self, max_age: chrono::Duration) -> Vec<String> {
111        let cutoff = Utc::now() - max_age;
112        let stale: Vec<String> = self
113            .entries
114            .iter()
115            .filter(|(_, e)| e.last_heartbeat < cutoff)
116            .map(|(id, _)| id.clone())
117            .collect();
118
119        for id in &stale {
120            self.entries.remove(id);
121        }
122        stale
123    }
124
125    /// Total number of registered instances.
126    pub fn len(&self) -> usize {
127        self.entries.len()
128    }
129
130    /// Whether the registry is empty.
131    pub fn is_empty(&self) -> bool {
132        self.entries.is_empty()
133    }
134
135    /// List all unique services in the registry.
136    pub fn services(&self) -> Vec<String> {
137        let mut svcs: Vec<String> = self.entries.values().map(|e| e.service.clone()).collect();
138        svcs.sort();
139        svcs.dedup();
140        svcs
141    }
142}
a3s_box_runtime/scale/registry.rs

a3s_box_runtime/scale/
registry.rs
