ri/service_mesh/
service_discovery.rs

1//! Copyright © 2025-2026 Wenze Wei. All Rights Reserved.
2//!
3//! This file is part of Ri.
4//! The Ri project belongs to the Dunimd Team.
5//!
6//! Licensed under the Apache License, Version 2.0 (the "License");
7//! You may not use this file except in compliance with the License.
8//! You may obtain a copy of the License at
9//!
10//!     http://www.apache.org/licenses/LICENSE-2.0
11//!
12//! Unless required by applicable law or agreed to in writing, software
13//! distributed under the License is distributed on an "AS IS" BASIS,
14//! WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15//! See the License for the specific language governing permissions and
16//! limitations under the License.
17
18use serde::{Deserialize, Serialize};
19use std::collections::HashMap as FxHashMap;
20use std::sync::Arc;
21use std::time::{Duration, SystemTime};
22use tokio::sync::RwLock;
23use tokio::task::JoinHandle;
24
25#[cfg(feature = "etcd")]
26use etcd_client::{Client, PutOptions};
27#[cfg(feature = "etcd")]
28use tokio::sync::Mutex;
29
30use crate::core::{RiResult, RiError};
31
32#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
33#[derive(Debug, Clone, Serialize, Deserialize)]
34pub struct RiServiceInstance {
35    pub id: String,
36    pub service_name: String,
37    pub host: String,
38    pub port: u16,
39    pub metadata: FxHashMap<String, String>,
40    pub registered_at: SystemTime,
41    pub last_heartbeat: SystemTime,
42    pub status: RiServiceStatus,
43}
44
45#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
46#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
47pub enum RiServiceStatus {
48    Starting,
49    Running,
50    Stopping,
51    Stopped,
52    Unhealthy,
53}
54
55#[cfg(feature = "pyo3")]
56#[pyo3::prelude::pymethods]
57impl RiServiceInstance {
58    #[new]
59    fn py_new(
60        id: String,
61        service_name: String,
62        host: String,
63        port: u16,
64    ) -> Self {
65        Self {
66            id,
67            service_name,
68            host,
69            port,
70            metadata: FxHashMap::default(),
71            registered_at: SystemTime::now(),
72            last_heartbeat: SystemTime::now(),
73            status: RiServiceStatus::Starting,
74        }
75    }
76
77    #[getter]
78    fn id(&self) -> &str {
79        &self.id
80    }
81
82    #[getter]
83    fn service_name(&self) -> &str {
84        &self.service_name
85    }
86
87    #[getter]
88    fn host(&self) -> &str {
89        &self.host
90    }
91
92    #[getter]
93    fn port(&self) -> u16 {
94        self.port
95    }
96
97    #[getter]
98    fn status(&self) -> RiServiceStatus {
99        self.status.clone()
100    }
101}
102
103#[derive(Clone)]
104pub struct RiServiceRegistry {
105    services: Arc<RwLock<FxHashMap<String, Vec<RiServiceInstance>>>>,
106    instance_index: Arc<RwLock<FxHashMap<String, RiServiceInstance>>>,
107    #[cfg(feature = "etcd")]
108    etcd_client: Option<Arc<Mutex<Client>>>,
109    _etcd_prefix: String,
110}
111
112impl std::fmt::Debug for RiServiceRegistry {
113    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
114        let mut debug_struct = f.debug_struct("RiServiceRegistry");
115        debug_struct.field("services", &self.services)
116            .field("instance_index", &self.instance_index);
117        
118        #[cfg(feature = "etcd")]
119        debug_struct.field("etcd_client", &self.etcd_client.is_some());
120        
121        debug_struct.field("_etcd_prefix", &self._etcd_prefix)
122            .finish()
123    }
124}
125
126impl Default for RiServiceRegistry {
127    fn default() -> Self {
128        Self::new(None, "/dms/services".to_string())
129    }
130}
131
132impl RiServiceRegistry {
133    #[cfg(feature = "etcd")]
134    pub fn new(etcd_client: Option<Client>, etcd_prefix: String) -> Self {
135        Self {
136            services: Arc::new(RwLock::new(FxHashMap::default())),
137            instance_index: Arc::new(RwLock::new(FxHashMap::default())),
138            etcd_client: etcd_client.map(|c| Arc::new(Mutex::new(c))),
139            _etcd_prefix: etcd_prefix,
140        }
141    }
142    
143    #[cfg(not(feature = "etcd"))]
144    pub fn new(_etcd_client: Option<()>, etcd_prefix: String) -> Self {
145        Self {
146            services: Arc::new(RwLock::new(FxHashMap::default())),
147            instance_index: Arc::new(RwLock::new(FxHashMap::default())),
148            _etcd_prefix: etcd_prefix,
149        }
150    }
151
152    pub async fn register_service(&self, instance: RiServiceInstance) -> RiResult<()> {
153        // Security: Validate service instance to prevent spoofing attacks
154        self.validate_service_instance(&instance)?;
155        
156        // Update in-memory registry
157        let mut services = self.services.write().await;
158        let mut instance_index = self.instance_index.write().await;
159
160        services.entry(instance.service_name.clone())
161            .or_insert_with(Vec::new)
162            .push(instance.clone());
163
164        instance_index.insert(instance.id.clone(), instance.clone());
165        
166        // Persist to etcd if client is available
167        #[cfg(feature = "etcd")]
168        if let Some(client) = &self.etcd_client {
169            let key = format!("{}/{}/{}", self._etcd_prefix, instance.service_name, instance.id);
170            let value = serde_json::to_string(&instance)?;
171            
172            // Set with TTL of 5 minutes (300 seconds)
173            let mut client_guard = client.lock().await;
174            client_guard.put(key.as_bytes(), value.as_bytes(), Some(PutOptions::new().with_lease(300)))
175                .await
176                .map_err(|e| RiError::ServiceMesh(format!("Failed to register service in etcd: {}", e)))?;
177            drop(client_guard);
178        }
179
180        Ok(())
181    }
182
183    /// Validates a service instance to prevent spoofing and injection attacks.
184    ///
185    /// # Security
186    ///
187    /// This method validates:
188    /// 1. Service name format (alphanumeric, dash, underscore only)
189    /// 2. Instance ID format (alphanumeric, dash, underscore only)
190    /// 3. Host format (valid IP address or hostname)
191    /// 4. Port range (1-65535)
192    /// 5. Metadata key/value format
193    fn validate_service_instance(&self, instance: &RiServiceInstance) -> RiResult<()> {
194        // Validate service name
195        if instance.service_name.is_empty() || instance.service_name.len() > 128 {
196            return Err(RiError::ServiceMesh(
197                "Service name must be 1-128 characters".to_string()
198            ));
199        }
200        
201        for c in instance.service_name.chars() {
202            if !c.is_ascii_alphanumeric() && c != '-' && c != '_' {
203                return Err(RiError::ServiceMesh(
204                    "Service name can only contain alphanumeric characters, dash, and underscore".to_string()
205                ));
206            }
207        }
208
209        // Validate instance ID
210        if instance.id.is_empty() || instance.id.len() > 128 {
211            return Err(RiError::ServiceMesh(
212                "Instance ID must be 1-128 characters".to_string()
213            ));
214        }
215        
216        for c in instance.id.chars() {
217            if !c.is_ascii_alphanumeric() && c != '-' && c != '_' {
218                return Err(RiError::ServiceMesh(
219                    "Instance ID can only contain alphanumeric characters, dash, and underscore".to_string()
220                ));
221            }
222        }
223
224        // Validate host
225        if instance.host.is_empty() || instance.host.len() > 256 {
226            return Err(RiError::ServiceMesh(
227                "Host must be 1-256 characters".to_string()
228            ));
229        }
230        
231        // Check for invalid characters in host
232        for c in instance.host.chars() {
233            if !c.is_ascii_alphanumeric() && c != '.' && c != '-' && c != ':' && c != '_' {
234                return Err(RiError::ServiceMesh(
235                    format!("Invalid character in host: '{}'", c)
236                ));
237            }
238        }
239
240        // Validate port
241        if instance.port == 0 {
242            return Err(RiError::ServiceMesh(
243                "Port cannot be 0".to_string()
244            ));
245        }
246
247        // Validate metadata
248        for (key, value) in &instance.metadata {
249            if key.len() > 128 {
250                return Err(RiError::ServiceMesh(
251                    format!("Metadata key too long: {}", key)
252                ));
253            }
254            
255            if value.len() > 1024 {
256                return Err(RiError::ServiceMesh(
257                    format!("Metadata value too long for key: {}", key)
258                ));
259            }
260            
261            // Check for control characters in metadata
262            if key.chars().any(|c| c.is_control()) {
263                return Err(RiError::ServiceMesh(
264                    "Metadata key contains control characters".to_string()
265                ));
266            }
267            
268            if value.chars().any(|c| c.is_control()) {
269                return Err(RiError::ServiceMesh(
270                    format!("Metadata value for key '{}' contains control characters", key)
271                ));
272            }
273        }
274
275        Ok(())
276    }
277
278    pub async fn deregister_service(&self, instance_id: &str) -> RiResult<()> {
279        let mut instance_index = self.instance_index.write().await;
280        
281        if let Some(instance) = instance_index.remove(instance_id) {
282            // Update in-memory registry
283            let mut services = self.services.write().await;
284            if let Some(instances) = services.get_mut(&instance.service_name) {
285                instances.retain(|inst| inst.id != instance_id);
286                
287                if instances.is_empty() {
288                    services.remove(&instance.service_name);
289                }
290            }
291            
292            // Remove from etcd if client is available
293            #[cfg(feature = "etcd")]
294            if let Some(client) = &self.etcd_client {
295                let key = format!("{}/{}/{}", self._etcd_prefix, instance.service_name, instance_id);
296                let mut client_guard = client.lock().await;
297                client_guard.delete(key.as_bytes(), None)
298                    .await
299                    .map_err(|e| RiError::ServiceMesh(format!("Failed to deregister service in etcd: {}", e)))?;
300                drop(client_guard);
301            }
302        }
303
304        Ok(())
305    }
306
307    pub async fn get_service_instances(&self, service_name: &str) -> RiResult<Vec<RiServiceInstance>> {
308        let services = self.services.read().await;
309        let instances = services.get(service_name)
310            .cloned()
311            .unwrap_or_default();
312
313        Ok(instances)
314    }
315
316    pub async fn get_all_services(&self) -> RiResult<Vec<String>> {
317        let services = self.services.read().await;
318        let service_names: Vec<String> = services.keys().cloned().collect();
319        Ok(service_names)
320    }
321
322    pub async fn update_heartbeat(&self, instance_id: &str) -> RiResult<()> {
323        let mut instance_index = self.instance_index.write().await;
324        
325        if let Some(instance) = instance_index.get_mut(instance_id) {
326            instance.last_heartbeat = SystemTime::now();
327            
328            // Update in etcd if client is available
329            #[cfg(feature = "etcd")]
330            if let Some(client) = &self.etcd_client {
331                let key = format!("{}/{}/{}", self._etcd_prefix, instance.service_name, instance_id);
332                let value = serde_json::to_string(instance)?;
333                let mut client_guard = client.lock().await;
334                client_guard.put(key.as_bytes(), value.as_bytes(), None)
335                    .await
336                    .map_err(|e| RiError::ServiceMesh(format!("Failed to update service heartbeat in etcd: {}", e)))?;
337                drop(client_guard);
338            }
339        }
340
341        Ok(())
342    }
343    
344    pub async fn update_instance_status(&self, instance_id: &str, status: RiServiceStatus) -> RiResult<()> {
345        let mut instance_index = self.instance_index.write().await;
346        
347        if let Some(instance) = instance_index.get_mut(instance_id) {
348            instance.status = status;
349            instance.last_heartbeat = SystemTime::now();
350            
351            // Update in etcd if client is available
352            #[cfg(feature = "etcd")]
353            if let Some(client) = &self.etcd_client {
354                let key = format!("{}/{}/{}", self._etcd_prefix, instance.service_name, instance_id);
355                let value = serde_json::to_string(instance)?;
356                let mut client_guard = client.lock().await;
357                client_guard.put(key.as_bytes(), value.as_bytes(), None)
358                    .await
359                    .map_err(|e| RiError::ServiceMesh(format!("Failed to update service status in etcd: {}", e)))?;
360                drop(client_guard);
361            }
362        }
363
364        Ok(())
365    }
366
367    pub async fn get_healthy_instances(&self, service_name: &str) -> RiResult<Vec<RiServiceInstance>> {
368        let instances = self.get_service_instances(service_name).await?;
369        let healthy_instances: Vec<RiServiceInstance> = instances
370            .into_iter()
371            .filter(|inst| inst.status == RiServiceStatus::Running)
372            .collect();
373
374        Ok(healthy_instances)
375    }
376
377    pub async fn cleanup_expired_instances(&self, expiration_duration: Duration) -> RiResult<()> {
378        let now = SystemTime::now();
379        let mut expired_instances = Vec::with_capacity(4);
380
381        {
382            let instance_index = self.instance_index.read().await;
383            for (id, instance) in instance_index.iter() {
384                if let Ok(elapsed) = now.duration_since(instance.last_heartbeat) {
385                    if elapsed > expiration_duration {
386                        expired_instances.push(id.clone());
387                    }
388                }
389            }
390        }
391
392        for instance_id in expired_instances {
393            self.deregister_service(&instance_id).await?;
394        }
395
396        Ok(())
397    }
398    
399    /// Sync registry from etcd
400    #[cfg(feature = "etcd")]
401    pub async fn sync_from_etcd(&self) -> RiResult<()> {
402        if let Some(client) = &self.etcd_client {
403            // List all services from etcd
404            let prefix = format!("{}/", self._etcd_prefix);
405            let mut client_guard = client.lock().await;
406            let response = client_guard.get(prefix.as_bytes(), Some(etcd_client::GetOptions::new().with_prefix()))
407                .await
408                .map_err(|e| RiError::ServiceMesh(format!("Failed to sync from etcd: {}", e)))?;
409            drop(client_guard);
410            
411            // Clear current in-memory registry
412            let mut services = self.services.write().await;
413            let mut instance_index = self.instance_index.write().await;
414            services.clear();
415            instance_index.clear();
416            
417            // Reconstruct from etcd data
418            for kv in response.kvs() {
419                let instance: RiServiceInstance = serde_json::from_slice(kv.value())?;
420                
421                services.entry(instance.service_name.clone())
422                    .or_insert_with(Vec::new)
423                    .push(instance.clone());
424                
425                instance_index.insert(instance.id.clone(), instance);
426            }
427        }
428        
429        Ok(())
430    }
431    
432    /// Start etcd watcher to sync changes in real-time
433    #[cfg(feature = "etcd")]
434    pub async fn start_etcd_watcher(&self) -> RiResult<JoinHandle<()>> {
435        if let Some(client) = &self.etcd_client {
436            let client = client.clone();
437            let prefix = self._etcd_prefix.clone();
438            let registry = self.clone();
439            
440            let handle = tokio::spawn(async move {
441                let prefix = format!("{}/", prefix);
442                
443                loop {
444                    let mut client_guard = client.lock().await;
445                    let (_watcher, mut stream) = client_guard.watch(prefix.as_bytes(), Some(etcd_client::WatchOptions::new().with_prefix()))
446                        .await
447                        .expect("Failed to start etcd watcher");
448                    drop(client_guard); // Release the lock before awaiting
449                    
450                    while let Ok(Some(watch_response)) = stream.message().await {
451                        for event in watch_response.events() {
452                            match event.event_type() {
453                                etcd_client::EventType::Put => {
454                                    // Update or add instance
455                                    if let Some(kv) = event.kv() {
456                                        if let Ok(instance) = serde_json::from_slice(kv.value()) {
457                                            let _ = registry.register_service(instance).await;
458                                        }
459                                    }
460                                },
461                                etcd_client::EventType::Delete => {
462                                    // Delete instance - we'd need to parse the key to get instance ID
463                                    // This is simplified for now
464                                },
465                            }
466                        }
467                    }
468                }
469            });
470            
471            Ok(handle)
472        } else {
473            Err(RiError::ServiceMesh("No etcd client available".to_string()))
474        }
475    }
476}
477
478#[derive(Debug, Clone, Serialize, Deserialize)]
479pub struct RiEtcdConfig {
480    pub endpoints: Vec<String>,
481    pub username: Option<String>,
482    pub password: Option<String>,
483    pub prefix: String,
484}
485
486impl Default for RiEtcdConfig {
487    fn default() -> Self {
488        Self {
489            endpoints: vec!["http://localhost:2379".to_string()],
490            username: None,
491            password: None,
492            prefix: "/dms/services".to_string(),
493        }
494    }
495}
496
497#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
498#[derive(Debug, Clone)]
499pub struct RiServiceDiscovery {
500    enabled: bool,
501    registry: Arc<RiServiceRegistry>,
502    background_tasks: Arc<RwLock<Vec<JoinHandle<()>>>>,
503    cleanup_interval: Duration,
504    _etcd_config: Option<RiEtcdConfig>,
505}
506
507impl RiServiceDiscovery {
508    pub fn new(enabled: bool) -> Self {
509        Self {
510            enabled,
511            registry: Arc::new(RiServiceRegistry::new(None, "/dms/services".to_string())),
512            background_tasks: Arc::new(RwLock::new(Vec::new())),
513            cleanup_interval: Duration::from_secs(60),
514            _etcd_config: None,
515        }
516    }
517    
518    #[cfg(feature = "etcd")]
519    pub async fn new_with_etcd(enabled: bool, etcd_config: RiEtcdConfig) -> RiResult<Self> {
520        // Create etcd client
521        let client = Client::connect(etcd_config.endpoints.clone(), None)
522            .await
523            .map_err(|e| RiError::ServiceMesh(format!("Failed to connect to etcd: {}", e)))?;
524        
525        let registry = Arc::new(RiServiceRegistry::new(Some(client), etcd_config.prefix.clone()));
526        
527        let discovery = Self {
528            enabled,
529            registry,
530            background_tasks: Arc::new(RwLock::new(Vec::new())),
531            cleanup_interval: Duration::from_secs(60),
532            _etcd_config: Some(etcd_config),
533        };
534        
535        // Sync from etcd on startup
536        discovery.registry.sync_from_etcd().await?;
537        
538        Ok(discovery)
539    }
540
541    pub async fn register_service(
542        &self,
543        service_name: &str,
544        host: &str,
545        port: u16,
546        metadata: FxHashMap<String, String>,
547    ) -> RiResult<String> {
548        if !self.enabled {
549            return Err(RiError::ServiceMesh("Service discovery is disabled".to_string()));
550        }
551
552        let instance_id = format!("{service_name}:{host}:{port}");
553        let instance = RiServiceInstance {
554            id: instance_id.clone(),
555            service_name: service_name.to_string(),
556            host: host.to_string(),
557            port,
558            metadata,
559            registered_at: SystemTime::now(),
560            last_heartbeat: SystemTime::now(),
561            status: RiServiceStatus::Starting,
562        };
563
564        self.registry.register_service(instance).await?;
565        Ok(instance_id)
566    }
567
568    pub async fn deregister_service(&self, instance_id: &str) -> RiResult<()> {
569        if !self.enabled {
570            return Err(RiError::ServiceMesh("Service discovery is disabled".to_string()));
571        }
572
573        self.registry.deregister_service(instance_id).await
574    }
575
576    pub async fn discover_service(&self, service_name: &str) -> RiResult<Vec<RiServiceInstance>> {
577        if !self.enabled {
578            return Err(RiError::ServiceMesh("Service discovery is disabled".to_string()));
579        }
580
581        self.registry.get_healthy_instances(service_name).await
582    }
583
584    pub async fn update_heartbeat(&self, instance_id: &str) -> RiResult<()> {
585        if !self.enabled {
586            return Err(RiError::ServiceMesh("Service discovery is disabled".to_string()));
587        }
588
589        self.registry.update_heartbeat(instance_id).await
590    }
591    
592    pub async fn set_service_status(&self, instance_id: &str, status: RiServiceStatus) -> RiResult<()> {
593        if !self.enabled {
594            return Err(RiError::ServiceMesh("Service discovery is disabled".to_string()));
595        }
596
597        self.registry.update_instance_status(instance_id, status).await
598    }
599
600    pub async fn get_service_instances(&self, service_name: &str) -> RiResult<Vec<RiServiceInstance>> {
601        if !self.enabled {
602            return Err(RiError::ServiceMesh("Service discovery is disabled".to_string()));
603        }
604
605        self.registry.get_service_instances(service_name).await
606    }
607
608    pub async fn get_all_services(&self) -> RiResult<Vec<String>> {
609        if !self.enabled {
610            return Err(RiError::ServiceMesh("Service discovery is disabled".to_string()));
611        }
612
613        self.registry.get_all_services().await
614    }
615
616    pub async fn start_background_tasks(&self) -> RiResult<()> {
617        if !self.enabled {
618            return Ok(());
619        }
620
621        let registry_clone = Arc::clone(&self.registry);
622
623        let cleanup_interval = self.cleanup_interval;
624
625        // Start cleanup task
626        let cleanup_task = tokio::spawn(async move {
627            let mut interval = tokio::time::interval(cleanup_interval);
628            loop {
629                interval.tick().await;
630                if let Err(e) = registry_clone.cleanup_expired_instances(Duration::from_secs(300)).await {
631                    log::warn!("Failed to cleanup expired instances: {e}");
632                }
633            }
634        });
635
636        let mut tasks = self.background_tasks.write().await;
637        tasks.push(cleanup_task);
638        
639        // Start etcd watcher if etcd is configured
640        #[cfg(feature = "etcd")]
641        if self._etcd_config.is_some() {
642            let watcher_task = self.registry.start_etcd_watcher().await?;
643            tasks.push(watcher_task);
644            
645            // Start periodic sync from etcd (every 30 seconds)
646            let registry_clone = Arc::clone(&self.registry);
647            let sync_task = tokio::spawn(async move {
648                let mut interval = tokio::time::interval(Duration::from_secs(30));
649                loop {
650                    interval.tick().await;
651                    if let Err(e) = registry_clone.sync_from_etcd().await {
652                        log::warn!("Failed to sync from etcd: {e}");
653                    }
654                }
655            });
656            tasks.push(sync_task);
657        }
658
659        Ok(())
660    }
661
662    pub async fn stop_background_tasks(&self) -> RiResult<()> {
663        let mut tasks = self.background_tasks.write().await;
664        for task in tasks.drain(..) {
665            task.abort();
666        }
667        Ok(())
668    }
669
670    pub async fn health_check(&self) -> RiResult<bool> {
671        Ok(self.enabled)
672    }
673}
ri/service_mesh/service_discovery.rs

ri/service_mesh/
service_discovery.rs
