ri/service_mesh/
health_check.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
18//! # Health Check Module
19//! 
20//! This module provides health checking functionality for the Ri service mesh. It allows
21//! monitoring the health of services using various protocols and provides comprehensive
22//! health status information.
23//! 
24//! ## Key Components
25//! 
26//! - **RiHealthCheckConfig**: Configuration for health checks
27//! - **RiHealthCheckResult**: Result of a health check
28//! - **RiHealthCheckType**: Supported health check types
29//! - **RiHealthCheckProvider**: Trait for implementing health check providers
30//! - **RiHttpHealthCheckProvider**: HTTP health check implementation
31//! - **RiTcpHealthCheckProvider**: TCP health check implementation
32//! - **RiHealthChecker**: Main health checking service
33//! - **RiHealthStatus**: Health status enum
34//! - **RiHealthSummary**: Summary of health check results
35//! 
36//! ## Design Principles
37//! 
38//! 1. **Protocol Agnostic**: Supports multiple health check protocols (HTTP, TCP, gRPC, custom)
39//! 2. **Async-First**: All health check operations are asynchronous
40//! 3. **Extensible**: Easy to implement new health check providers
41//! 4. **Configurable**: Highly configurable health check parameters
42//! 5. **Real-time Monitoring**: Background tasks for continuous health monitoring
43//! 6. **Comprehensive Results**: Detailed health check results with response times and error messages
44//! 7. **Health Summary**: Aggregated health status with success rates and average response times
45//! 8. **Thread-safe**: Uses Arc and RwLock for safe concurrent access
46//! 9. **Graceful Shutdown**: Proper cleanup of background tasks
47//! 10. **Error Handling**: Comprehensive error handling with RiResult
48//! 
49//! ## Usage
50//! 
51//! ```rust
52//! use ri::prelude::*;
53//! use std::time::Duration;
54//! 
55//! async fn example() -> RiResult<()> {
56//!     // Create a health checker with 30-second intervals
57//!     let health_checker = RiHealthChecker::new(Duration::from_secs(30));
58//!     
59//!     // Register a health check for a service
60//!     let config = RiHealthCheckConfig {
61//!         endpoint: "/health".to_string(),
62//!         method: "GET".to_string(),
63//!         timeout: Duration::from_secs(5),
64//!         expected_status_code: 200,
65//!         expected_response_body: None,
66//!         headers: FxHashMap::default(),
67//!     };
68//!     
69//!     health_checker.register_health_check(
70//!         "example-service",
71//!         "http://localhost:8080",
72//!         RiHealthCheckType::Http,
73//!         config
74//!     ).await?;
75//!     
76//!     // Start background health checks
77//!     health_checker.start_health_check("example-service", "http://localhost:8080").await?;
78//!     
79//!     // Get health summary
80//!     let summary = health_checker.get_service_health_summary("example-service").await?;
81//!     println!("Service health: {:?}", summary.overall_status);
82//!     println!("Success rate: {:.2}%", summary.success_rate);
83//!     
84//!     Ok(())
85//! }
86//! ```
87
88use async_trait::async_trait;
89use serde::{Deserialize, Serialize};
90use std::collections::HashMap as FxHashMap;
91use std::sync::Arc;
92use std::time::{Duration, SystemTime};
93use tokio::sync::RwLock;
94use tokio::task::JoinHandle;
95
96#[cfg(feature = "pyo3")]
97use pyo3::PyResult;
98#[cfg(feature = "service_mesh")]
99use hyper;
100
101use crate::core::{RiResult, RiError};
102use crate::observability::{RiTracer, RiSpanKind, RiSpanStatus};
103
104/// Configuration for health checks.
105///
106/// This struct defines the parameters for performing health checks, including
107/// endpoint, HTTP method, timeout, expected status code, and custom headers.
108#[derive(Debug, Clone, Serialize, Deserialize)]
109pub struct RiHealthCheckConfig {
110    /// Health check endpoint path
111    pub endpoint: String,
112    /// HTTP method to use for health checks
113    pub method: String,
114    /// Timeout for health check requests
115    pub timeout: Duration,
116    /// Expected HTTP status code for a healthy service
117    pub expected_status_code: u16,
118    /// Optional expected response body for validation
119    pub expected_response_body: Option<String>,
120    /// Custom headers to include in health check requests
121    pub headers: FxHashMap<String, String>,
122}
123
124impl Default for RiHealthCheckConfig {
125    /// Creates a default health check configuration.
126    ///
127    /// # Returns
128    ///
129    /// A `RiHealthCheckConfig` instance with default values
130    fn default() -> Self {
131        Self {
132            endpoint: "/health".to_string(),
133            method: "GET".to_string(),
134            timeout: Duration::from_secs(5),
135            expected_status_code: 200,
136            expected_response_body: None,
137            headers: FxHashMap::default(),
138        }
139    }
140}
141
142/// Result of a health check operation.
143///
144/// This struct contains detailed information about the result of a health check,
145/// including whether the service is healthy, response time, and error messages if any.
146#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
147#[derive(Debug, Clone)]
148pub struct RiHealthCheckResult {
149    /// Name of the service being checked
150    pub service_name: String,
151    /// Endpoint used for the health check
152    pub endpoint: String,
153    /// Whether the service is considered healthy
154    pub is_healthy: bool,
155    /// HTTP status code received (if applicable)
156    pub status_code: Option<u16>,
157    /// Time taken to perform the health check
158    pub response_time: Duration,
159    /// Error message if the health check failed
160    pub error_message: Option<String>,
161    /// Timestamp when the health check was performed
162    pub timestamp: SystemTime,
163}
164
165#[cfg(feature = "pyo3")]
166#[pyo3::prelude::pymethods]
167impl RiHealthCheckResult {
168    fn get_service_name(&self) -> String {
169        self.service_name.clone()
170    }
171    
172    fn get_endpoint(&self) -> String {
173        self.endpoint.clone()
174    }
175    
176    fn get_is_healthy(&self) -> bool {
177        self.is_healthy
178    }
179    
180    fn get_status_code(&self) -> Option<u16> {
181        self.status_code
182    }
183    
184    fn get_response_time_ms(&self) -> u64 {
185        self.response_time.as_millis() as u64
186    }
187    
188    fn get_error_message(&self) -> Option<String> {
189        self.error_message.clone()
190    }
191}
192
193/// Types of health checks supported.
194///
195/// This enum defines the different protocols that can be used for health checking.
196#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
197#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
198pub enum RiHealthCheckType {
199    /// HTTP health check
200    Http,
201    /// TCP health check
202    Tcp,
203    /// gRPC health check
204    Grpc,
205    /// Custom health check implementation
206    Custom,
207}
208
209/// Trait for implementing health check providers.
210///
211/// This trait defines the interface for health check providers, allowing for
212/// different health check implementations based on protocol.
213#[async_trait]
214pub trait RiHealthCheckProvider: Send + Sync {
215    /// Performs a health check on the specified endpoint.
216    ///
217    /// # Parameters
218    ///
219    /// - `endpoint`: The endpoint to check
220    /// - `config`: Health check configuration
221    ///
222    /// # Returns
223    ///
224    /// A `RiResult<RiHealthCheckResult>` containing the health check result
225    async fn check_health(&self, endpoint: &str, config: &RiHealthCheckConfig) -> RiResult<RiHealthCheckResult>;
226}
227
228/// HTTP health check provider.
229///
230/// This struct implements the `RiHealthCheckProvider` trait for HTTP health checks.
231pub struct RiHttpHealthCheckProvider;
232
233#[async_trait]
234impl RiHealthCheckProvider for RiHttpHealthCheckProvider {
235    /// Performs an HTTP health check on the specified endpoint.
236    ///
237    /// # Parameters
238    ///
239    /// - `endpoint`: The HTTP endpoint to check
240    /// - `config`: Health check configuration
241    ///
242    /// # Returns
243    ///
244    /// A `RiResult<RiHealthCheckResult>` containing the health check result
245    #[cfg(feature = "service_mesh")]
246    async fn check_health(&self, endpoint: &str, _config: &RiHealthCheckConfig) -> RiResult<RiHealthCheckResult> {
247        let start_time = SystemTime::now();
248        
249        let client = hyper::Client::new();
250
251        let uri: hyper::Uri = endpoint.parse()
252            .map_err(|e| RiError::ServiceMesh(format!("Invalid URI: {e}")))?;
253
254        let req = hyper::Request::builder()
255            .method(_config.method.as_str())
256            .uri(uri)
257            .body(hyper::Body::empty())
258            .map_err(|e| RiError::ServiceMesh(format!("Failed to build request: {e}")))?;
259
260        match client.request(req).await {
261            Ok(response) => {
262                let status_code = response.status().as_u16();
263                let is_healthy = status_code == _config.expected_status_code;
264                let response_time = SystemTime::now().duration_since(start_time)
265                    .unwrap_or(Duration::from_secs(0));
266
267                let error_message = if !is_healthy {
268                    Some(format!("Expected status code {}, got {}", _config.expected_status_code, status_code))
269                } else {
270                    None
271                };
272
273                Ok(RiHealthCheckResult {
274                    service_name: "unknown".to_string(),
275                    endpoint: endpoint.to_string(),
276                    is_healthy,
277                    status_code: Some(status_code),
278                    response_time,
279                    error_message,
280                    timestamp: SystemTime::now(),
281                })
282            }
283            Err(e) => {
284                let response_time = SystemTime::now().duration_since(start_time)
285                    .unwrap_or(Duration::from_secs(0));
286
287                Ok(RiHealthCheckResult {
288                    service_name: "unknown".to_string(),
289                    endpoint: endpoint.to_string(),
290                    is_healthy: false,
291                    status_code: None,
292                    response_time,
293                    error_message: Some(e.to_string()),
294                    timestamp: SystemTime::now(),
295                })
296            }
297        }
298    }
299    
300    #[cfg(not(feature = "service_mesh"))]
301    async fn check_health(&self, endpoint: &str, _config: &RiHealthCheckConfig) -> RiResult<RiHealthCheckResult> {
302        // If service_mesh feature is not enabled, assume all endpoints are healthy
303        Ok(RiHealthCheckResult {
304            service_name: "unknown".to_string(),
305            endpoint: endpoint.to_string(),
306            is_healthy: true,
307            status_code: Some(_config.expected_status_code),
308            response_time: Duration::from_secs(0),
309            error_message: None,
310            timestamp: SystemTime::now(),
311        })
312    }
313}
314
315/// TCP health check provider.
316///
317/// This struct implements the `RiHealthCheckProvider` trait for TCP health checks.
318pub struct RiTcpHealthCheckProvider;
319
320#[async_trait]
321impl RiHealthCheckProvider for RiTcpHealthCheckProvider {
322    /// Performs a TCP health check on the specified endpoint.
323    ///
324    /// # Parameters
325    ///
326    /// - `endpoint`: The TCP endpoint to check (format: "host:port")
327    /// - `config`: Health check configuration
328    ///
329    /// # Returns
330    ///
331    /// A `RiResult<RiHealthCheckResult>` containing the health check result
332    async fn check_health(&self, endpoint: &str, _config: &RiHealthCheckConfig) -> RiResult<RiHealthCheckResult> {
333        let start_time = SystemTime::now();
334        
335        match tokio::net::TcpStream::connect(endpoint).await {
336            Ok(_) => {
337                let response_time = SystemTime::now().duration_since(start_time)
338                    .unwrap_or(Duration::from_secs(0));
339
340                Ok(RiHealthCheckResult {
341                    service_name: "unknown".to_string(),
342                    endpoint: endpoint.to_string(),
343                    is_healthy: true,
344                    status_code: None,
345                    response_time,
346                    error_message: None,
347                    timestamp: SystemTime::now(),
348                })
349            }
350            Err(e) => {
351                let response_time = SystemTime::now().duration_since(start_time)
352                    .unwrap_or(Duration::from_secs(0));
353
354                Ok(RiHealthCheckResult {
355                    service_name: "unknown".to_string(),
356                    endpoint: endpoint.to_string(),
357                    is_healthy: false,
358                    status_code: None,
359                    response_time,
360                    error_message: Some(e.to_string()),
361                    timestamp: SystemTime::now(),
362                })
363            }
364        }
365    }
366}
367
368/// gRPC health check provider.
369///
370/// This struct implements the `RiHealthCheckProvider` trait for gRPC health checks.
371pub struct RiGrpcHealthCheckProvider;
372
373#[async_trait]
374impl RiHealthCheckProvider for RiGrpcHealthCheckProvider {
375    /// Performs a gRPC health check on the specified endpoint.
376    ///
377    /// # Parameters
378    ///
379    /// - `endpoint`: The gRPC endpoint to check (format: "host:port")
380    /// - `config`: Health check configuration
381    ///
382    /// # Returns
383    ///
384    /// A `RiResult<RiHealthCheckResult>` containing the health check result
385    async fn check_health(&self, endpoint: &str, _config: &RiHealthCheckConfig) -> RiResult<RiHealthCheckResult> {
386        let start_time = SystemTime::now();
387        
388        // Simple gRPC health check implementation using TCP connection
389        // In a full implementation, this would use the gRPC health check service
390        match tokio::net::TcpStream::connect(endpoint).await {
391            Ok(_) => {
392                let response_time = SystemTime::now().duration_since(start_time)
393                    .unwrap_or(Duration::from_secs(0));
394
395                Ok(RiHealthCheckResult {
396                    service_name: "unknown".to_string(),
397                    endpoint: endpoint.to_string(),
398                    is_healthy: true,
399                    status_code: None,
400                    response_time,
401                    error_message: None,
402                    timestamp: SystemTime::now(),
403                })
404            }
405            Err(e) => {
406                let response_time = SystemTime::now().duration_since(start_time)
407                    .unwrap_or(Duration::from_secs(0));
408
409                Ok(RiHealthCheckResult {
410                    service_name: "unknown".to_string(),
411                    endpoint: endpoint.to_string(),
412                    is_healthy: false,
413                    status_code: None,
414                    response_time,
415                    error_message: Some(e.to_string()),
416                    timestamp: SystemTime::now(),
417                })
418            }
419        }
420    }
421}
422
423/// Main health checker service.
424///
425/// This struct provides the core functionality for managing health checks, including
426/// registering health checks, starting background monitoring, and retrieving health status.
427#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
428pub struct RiHealthChecker {
429    check_interval: Duration,
430    providers: Arc<RwLock<FxHashMap<RiHealthCheckType, Box<dyn RiHealthCheckProvider>>>>,
431    check_results: Arc<RwLock<FxHashMap<String, Vec<RiHealthCheckResult>>>>,
432    background_tasks: Arc<RwLock<Vec<JoinHandle<()>>>>,
433    tracer: Option<Arc<RiTracer>>,
434}
435
436impl RiHealthChecker {
437    pub fn new(check_interval: Duration) -> Self {
438        let mut providers: FxHashMap<RiHealthCheckType, Box<dyn RiHealthCheckProvider>> = FxHashMap::default();
439        providers.insert(RiHealthCheckType::Http, Box::new(RiHttpHealthCheckProvider));
440        providers.insert(RiHealthCheckType::Tcp, Box::new(RiTcpHealthCheckProvider));
441        providers.insert(RiHealthCheckType::Grpc, Box::new(RiGrpcHealthCheckProvider));
442
443        Self {
444            check_interval,
445            providers: Arc::new(RwLock::new(providers)),
446            check_results: Arc::new(RwLock::new(FxHashMap::default())),
447            background_tasks: Arc::new(RwLock::new(Vec::new())),
448            tracer: None,
449        }
450    }
451    
452    pub fn with_tracer(mut self, tracer: Arc<RiTracer>) -> Self {
453        self.tracer = Some(tracer);
454        self
455    }
456    
457    pub fn set_tracer(&mut self, tracer: Arc<RiTracer>) {
458        self.tracer = Some(tracer);
459    }
460    
461    /// Validates an endpoint URL to prevent SSRF attacks.
462    ///
463    /// # Security
464    ///
465    /// This method validates:
466    /// 1. URL scheme must be HTTP or HTTPS
467    /// 2. URL must be well-formed
468    /// 3. Warns if URL points to private IP address
469    /// 4. Blocks file://, gopher://, and other dangerous schemes
470    fn validate_endpoint_url(endpoint: &str) -> RiResult<()> {
471        // Check URL length
472        if endpoint.is_empty() || endpoint.len() > 2048 {
473            return Err(RiError::ServiceMesh(
474                "Endpoint URL must be 1-2048 characters".to_string()
475            ));
476        }
477
478        // Parse and validate URL
479        let parsed_url = url::Url::parse(endpoint)
480            .map_err(|e| RiError::ServiceMesh(format!("Invalid endpoint URL: {}", e)))?;
481
482        // Only allow HTTP and HTTPS schemes
483        let scheme = parsed_url.scheme();
484        if scheme != "http" && scheme != "https" {
485            log::warn!(
486                "[Ri.HealthCheck] Blocked non-HTTP(S) endpoint: scheme={} url={}",
487                scheme, endpoint
488            );
489            return Err(RiError::ServiceMesh(
490                format!("Invalid URL scheme '{}'. Only HTTP and HTTPS are allowed for health checks.", scheme)
491            ));
492        }
493
494        // Check for private IP addresses (informational warning)
495        if let Some(host) = parsed_url.host_str() {
496            // Check for localhost
497            if host == "localhost" || host == "127.0.0.1" || host == "::1" {
498                log::warn!(
499                    "[Ri.HealthCheck] Health check endpoint points to localhost: {}",
500                    endpoint
501                );
502            }
503            
504            // Check for private IP ranges
505            if let Ok(ip) = host.parse::<std::net::IpAddr>() {
506                let is_private_or_link_local = match ip {
507                    std::net::IpAddr::V4(ipv4) => ipv4.is_private() || ipv4.is_link_local(),
508                    std::net::IpAddr::V6(ipv6) => ipv6.is_unicast_link_local(),
509                };
510                if ip.is_loopback() || is_private_or_link_local {
511                    log::warn!(
512                        "[Ri.HealthCheck] Health check endpoint points to private IP {}: {}",
513                        ip, endpoint
514                    );
515                }
516            }
517        }
518
519        // Check for credentials in URL
520        if parsed_url.username() != "" || parsed_url.password().is_some() {
521            log::warn!(
522                "[Ri.HealthCheck] Health check endpoint URL contains credentials: {}",
523                endpoint.split('@').last().unwrap_or(endpoint)
524            );
525        }
526
527        Ok(())
528    }
529
530
531    /// Registers a health check for a service.
532    ///
533    /// This method registers a health check for a service and performs an immediate check.
534    ///
535    /// # Parameters
536    ///
537    /// - `service_name`: Name of the service to check
538    /// - `endpoint`: Endpoint URL for health checks
539    /// - `check_type`: Type of health check to perform
540    /// - `config`: Health check configuration
541    ///
542    /// # Returns
543    ///
544    /// A `RiResult<()>` indicating success or failure
545    ///
546    /// # Security
547    ///
548    /// This method validates the endpoint URL to prevent SSRF attacks:
549    /// - Only HTTP and HTTPS schemes are allowed
550    /// - Private IP addresses are logged as warnings
551    /// - File://, gopher://, and other schemes are blocked
552    pub async fn register_health_check(
553        &self,
554        service_name: &str,
555        endpoint: &str,
556        check_type: RiHealthCheckType,
557        config: RiHealthCheckConfig,
558    ) -> RiResult<()> {
559        // Security: Validate endpoint URL to prevent SSRF
560        Self::validate_endpoint_url(endpoint)?;
561        
562        let span_id = if let Some(tracer) = &self.tracer {
563            let span_id = tracer.start_span_from_context(
564                format!("health_check:{}", service_name),
565                RiSpanKind::Internal,
566            );
567            if let Some(ref sid) = span_id {
568                let _ = tracer.span_mut(sid, |span| {
569                    span.set_attribute("service_name".to_string(), service_name.to_string());
570                    span.set_attribute("endpoint".to_string(), endpoint.to_string());
571                    span.set_attribute("check_type".to_string(), format!("{:?}", check_type));
572                });
573            }
574            span_id
575        } else {
576            None
577        };
578
579        let result = self.register_health_check_internal(service_name, endpoint, check_type, config).await;
580
581        if let (Some(tracer), Some(sid)) = (&self.tracer, span_id) {
582            let status = match &result {
583                Ok(_) => RiSpanStatus::Ok,
584                Err(e) => RiSpanStatus::Error(e.to_string()),
585            };
586            let _ = tracer.end_span(&sid, status);
587        }
588
589        result
590    }
591    
592    async fn register_health_check_internal(
593        &self,
594        service_name: &str,
595        endpoint: &str,
596        check_type: RiHealthCheckType,
597        config: RiHealthCheckConfig,
598    ) -> RiResult<()> {
599        let providers = self.providers.read().await;
600        let provider = providers.get(&check_type)
601            .ok_or_else(|| RiError::ServiceMesh(format!("Health check provider for {check_type:?} not found")))?;
602
603        let result = provider.check_health(endpoint, &config).await?;
604        
605        let mut check_results = self.check_results.write().await;
606        let service_results = check_results.entry(service_name.to_string())
607            .or_insert_with(Vec::new);
608        service_results.push(result);
609
610        Ok(())
611    }
612
613    /// Starts background health checks for a service.
614    /// 
615    /// This method creates a background task that periodically checks the health of a service.
616    /// 
617    /// # Parameters
618    /// 
619    /// - `service_name`: Name of the service to check
620    /// - `endpoint`: Endpoint URL for health checks
621    /// 
622    /// # Returns
623    /// 
624    /// A `RiResult<()>` indicating success or failure
625    pub async fn start_health_check(&self, service_name: &str, endpoint: &str) -> RiResult<()> {
626        let mut tasks = self.background_tasks.write().await;
627        
628        let service_name_clone = service_name.to_string();
629        let endpoint_clone = endpoint.to_string();
630        let check_interval = self.check_interval;
631        let providers = Arc::clone(&self.providers);
632        let check_results = Arc::clone(&self.check_results);
633
634        // Determine health check type based on endpoint URL scheme
635        let check_type = if endpoint.starts_with("grpc://") || endpoint.starts_with("grpcs://") {
636            RiHealthCheckType::Grpc
637        } else if endpoint.starts_with("http://") || endpoint.starts_with("https://") {
638            RiHealthCheckType::Http
639        } else {
640            // Assume TCP for other protocols
641            RiHealthCheckType::Tcp
642        };
643
644        let task = tokio::spawn(async move {
645            let mut interval = tokio::time::interval(check_interval);
646            let config = RiHealthCheckConfig::default();
647            
648            loop {
649                interval.tick().await;
650                
651                let providers_guard = providers.read().await;
652                if let Some(provider) = providers_guard.get(&check_type) {
653                    match provider.check_health(&endpoint_clone, &config).await {
654                        Ok(result) => {
655                            let mut results = check_results.write().await;
656                            let service_results = results.entry(service_name_clone.clone())
657                                .or_insert_with(Vec::new);
658                            
659                            // Add new result to the end
660                            service_results.push(result);
661                            
662                            // Keep only the most recent 100 results per service to avoid memory issues
663                            if service_results.len() > 100 {
664                                service_results.drain(0..service_results.len() - 100);
665                            }
666                        }
667                        Err(e) => {
668                            log::warn!("Health check failed for {endpoint_clone}: {e}");
669                        }
670                    }
671                }
672            }
673        });
674
675        tasks.push(task);
676        Ok(())
677    }
678    
679    /// Stops health checks for a specific service endpoint.
680    /// 
681    /// This method clears the health check results for the specified service.
682    /// The background task will continue running but will no longer record results.
683    /// 
684    /// # Parameters
685    /// 
686    /// - `service_name`: Name of the service
687    /// - `endpoint`: Endpoint URL
688    /// 
689    /// # Returns
690    /// 
691    /// A `RiResult<()>` indicating success or failure
692    pub async fn stop_health_check(&self, service_name: &str, _endpoint: &str) -> RiResult<()> {
693        let mut results = self.check_results.write().await;
694        results.remove(service_name);
695        Ok(())
696    }
697    
698    /// Starts background health checks for a service with a specific health check type.
699    /// 
700    /// This method creates a background task that periodically checks the health of a service
701    /// using the specified health check type.
702    /// 
703    /// # Parameters
704    /// 
705    /// - `service_name`: Name of the service to check
706    /// - `endpoint`: Endpoint URL for health checks
707    /// - `check_type`: Type of health check to perform
708    /// 
709    /// # Returns
710    /// 
711    /// A `RiResult<()>` indicating success or failure
712    pub async fn start_health_check_with_type(
713        &self, 
714        service_name: &str, 
715        endpoint: &str,
716        check_type: RiHealthCheckType
717    ) -> RiResult<()> {
718        let mut tasks = self.background_tasks.write().await;
719        
720        let service_name_clone = service_name.to_string();
721        let endpoint_clone = endpoint.to_string();
722        let check_interval = self.check_interval;
723        let providers = Arc::clone(&self.providers);
724        let check_results = Arc::clone(&self.check_results);
725        let check_type_clone = check_type;
726
727        let task = tokio::spawn(async move {
728            let mut interval = tokio::time::interval(check_interval);
729            let config = RiHealthCheckConfig::default();
730            
731            loop {
732                interval.tick().await;
733                
734                let providers_guard = providers.read().await;
735                if let Some(provider) = providers_guard.get(&check_type_clone) {
736                    match provider.check_health(&endpoint_clone, &config).await {
737                        Ok(result) => {
738                            let mut results = check_results.write().await;
739                            let service_results = results.entry(service_name_clone.clone())
740                                .or_insert_with(Vec::new);
741                            
742                            // Add new result to the end
743                            service_results.push(result);
744                            
745                            // Keep only the most recent 100 results per service to avoid memory issues
746                            if service_results.len() > 100 {
747                                service_results.drain(0..service_results.len() - 100);
748                            }
749                        }
750                        Err(e) => {
751                            log::warn!("Health check failed for {endpoint_clone}: {e}");
752                        }
753                    }
754                }
755            }
756        });
757
758        tasks.push(task);
759        Ok(())
760    }
761
762    /// Gets the health check results for a service.
763    ///
764    /// # Parameters
765    ///
766    /// - `service_name`: Name of the service to get results for
767    ///
768    /// # Returns
769    ///
770    /// A `RiResult<Vec<RiHealthCheckResult>>` containing the health check results
771    pub async fn get_health_status(&self, service_name: &str) -> RiResult<Vec<RiHealthCheckResult>> {
772        let check_results = self.check_results.read().await;
773        let results = check_results.get(service_name)
774            .cloned()
775            .unwrap_or_default();
776
777        Ok(results)
778    }
779    
780    /// Gets the latest health check result for a service.
781    ///
782    /// # Parameters
783    ///
784    /// - `service_name`: Name of the service to get the latest result for
785    ///
786    /// # Returns
787    ///
788    /// A `RiResult<Option<RiHealthCheckResult>>` containing the latest health check result if available
789    pub async fn get_latest_health_status(&self, service_name: &str) -> RiResult<Option<RiHealthCheckResult>> {
790        let check_results = self.check_results.read().await;
791        let latest_result = check_results.get(service_name)
792            .and_then(|results| results.last().cloned());
793
794        Ok(latest_result)
795    }
796    
797    /// Gets the health check results for a service within a specified time window.
798    ///
799    /// # Parameters
800    ///
801    /// - `service_name`: Name of the service to get results for
802    /// - `time_window`: Time window to filter results by
803    ///
804    /// # Returns
805    ///
806    /// A `RiResult<Vec<RiHealthCheckResult>>` containing the filtered health check results
807    pub async fn get_health_status_within(&self, service_name: &str, time_window: Duration) -> RiResult<Vec<RiHealthCheckResult>> {
808        let check_results = self.check_results.read().await;
809        let now = SystemTime::now();
810        
811        let results = check_results.get(service_name)
812            .map(|results| {
813                results.iter()
814                    .filter(|r| {
815                        if let Ok(elapsed) = now.duration_since(r.timestamp) {
816                            elapsed <= time_window
817                        } else {
818                            false
819                        }
820                    })
821                    .cloned()
822                    .collect()
823            })
824            .unwrap_or_default();
825
826        Ok(results)
827    }
828
829    /// Gets a health summary for a service.
830    ///
831    /// This method aggregates health check results to provide a summary of the service's health,
832    /// including success rate, average response time, and overall status.
833    ///
834    /// # Parameters
835    ///
836    /// - `service_name`: Name of the service to get a summary for
837    ///
838    /// # Returns
839    ///
840    /// A `RiResult<RiHealthSummary>` containing the health summary
841    pub async fn get_service_health_summary(&self, service_name: &str) -> RiResult<RiHealthSummary> {
842        let results = self.get_health_status(service_name).await?;
843        
844        if results.is_empty() {
845            return Ok(RiHealthSummary {
846                service_name: service_name.to_string(),
847                total_checks: 0,
848                healthy_checks: 0,
849                unhealthy_checks: 0,
850                success_rate: 0.0,
851                average_response_time: Duration::from_secs(0),
852                last_check_time: None,
853                overall_status: RiHealthStatus::Unknown,
854            });
855        }
856
857        let total_checks = results.len();
858        let healthy_checks = results.iter().filter(|r| r.is_healthy).count();
859        let unhealthy_checks = total_checks - healthy_checks;
860        let success_rate = (healthy_checks as f64) / (total_checks as f64) * 100.0;
861
862        let total_response_time: Duration = results.iter()
863            .map(|r| r.response_time)
864            .sum();
865        let average_response_time = total_response_time / total_checks as u32;
866
867        let last_check_time = results.last().map(|r| r.timestamp);
868
869        let overall_status = if success_rate >= 80.0 {
870            RiHealthStatus::Healthy
871        } else if success_rate >= 50.0 {
872            RiHealthStatus::Degraded
873        } else {
874            RiHealthStatus::Unhealthy
875        };
876
877        Ok(RiHealthSummary {
878            service_name: service_name.to_string(),
879            total_checks,
880            healthy_checks,
881            unhealthy_checks,
882            success_rate,
883            average_response_time,
884            last_check_time,
885            overall_status,
886        })
887    }
888
889    /// Starts background health check tasks.
890    ///
891    /// This method initializes and starts all background health monitoring tasks,
892    /// including periodic health checks for registered services and cleanup tasks.
893    ///
894    /// # Returns
895    ///
896    /// A `RiResult<()>` indicating success or failure
897    pub async fn start_background_tasks(&self) -> RiResult<()> {
898        // Start periodic cleanup task to remove old health check results
899        let check_results = Arc::clone(&self.check_results);
900        let cleanup_interval = self.check_interval * 10; // Cleanup every 10 check intervals
901        
902        let cleanup_task = tokio::spawn(async move {
903            let mut interval = tokio::time::interval(cleanup_interval);
904            
905            loop {
906                interval.tick().await;
907                
908                let mut results = check_results.write().await;
909                let now = SystemTime::now();
910                let max_age = Duration::from_secs(3600); // Keep results for 1 hour
911                
912                // Remove health check results older than max_age
913                for service_results in results.values_mut() {
914                    service_results.retain(|result| {
915                        now.duration_since(result.timestamp)
916                            .map(|age| age < max_age)
917                            .unwrap_or(false)
918                    });
919                }
920                
921                // Remove services with no recent results
922                results.retain(|_, results| !results.is_empty());
923            }
924        });
925        
926        // Store cleanup task
927        let mut tasks = self.background_tasks.write().await;
928        tasks.push(cleanup_task);
929        
930        log::info!("Background health check tasks started successfully");
931        Ok(())
932    }
933
934    /// Stops all background health check tasks.
935    ///
936    /// This method aborts all running background health check tasks and cleans up resources.
937    ///
938    /// # Returns
939    ///
940    /// A `RiResult<()>` indicating success or failure
941    pub async fn stop_background_tasks(&self) -> RiResult<()> {
942        let mut tasks = self.background_tasks.write().await;
943        for task in tasks.drain(..) {
944            task.abort();
945        }
946        Ok(())
947    }
948
949    /// Performs a health check on the health checker itself.
950    ///
951    /// # Returns
952    ///
953    /// A `RiResult<bool>` indicating whether the health checker is healthy
954    pub async fn health_check(&self) -> RiResult<bool> {
955        Ok(true)
956    }
957}
958
959#[cfg(feature = "pyo3")]
960/// Python bindings for RiHealthChecker
961#[pyo3::prelude::pymethods]
962impl RiHealthChecker {
963    #[new]
964    fn py_new(check_interval: u64) -> PyResult<Self> {
965        Ok(Self::new(Duration::from_secs(check_interval)))
966    }
967    
968    /// Get service health summary from Python
969    #[pyo3(name = "get_service_health_summary")]
970    fn get_service_health_summary_impl(&self, service_name: String) -> PyResult<RiHealthSummary> {
971        let rt = tokio::runtime::Runtime::new().map_err(|e| {
972            pyo3::exceptions::PyRuntimeError::new_err(format!("Failed to create runtime: {}", e))
973        })?;
974        
975        rt.block_on(async {
976            self.get_service_health_summary(&service_name)
977                .await
978                .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("Failed to get health summary: {e}")))
979        })
980    }
981    
982    /// Start health check from Python
983    #[pyo3(name = "start_health_check")]
984    fn start_health_check_impl(&self, service_name: String, endpoint: String) -> PyResult<()> {
985        let rt = tokio::runtime::Runtime::new().map_err(|e| {
986            pyo3::exceptions::PyRuntimeError::new_err(format!("Failed to create runtime: {}", e))
987        })?;
988        
989        rt.block_on(async {
990            self.start_health_check(&service_name, &endpoint)
991                .await
992                .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("Failed to start health check: {e}")))
993        })
994    }
995    
996    /// Stop health check from Python
997    #[pyo3(name = "stop_health_check")]
998    fn stop_health_check_impl(&self, service_name: String, endpoint: String) -> PyResult<()> {
999        let rt = tokio::runtime::Runtime::new().map_err(|e| {
1000            pyo3::exceptions::PyRuntimeError::new_err(format!("Failed to create runtime: {}", e))
1001        })?;
1002        
1003        rt.block_on(async {
1004            self.stop_health_check(&service_name, &endpoint)
1005                .await
1006                .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("Failed to stop health check: {e}")))
1007        })
1008    }
1009    
1010    /// Get health status from Python
1011    #[pyo3(name = "get_health_status")]
1012    fn get_health_status_impl(&self, service_name: String) -> PyResult<Vec<RiHealthCheckResult>> {
1013        let rt = tokio::runtime::Runtime::new().map_err(|e| {
1014            pyo3::exceptions::PyRuntimeError::new_err(format!("Failed to create runtime: {}", e))
1015        })?;
1016        
1017        rt.block_on(async {
1018            self.get_health_status(&service_name)
1019                .await
1020                .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("Failed to get health status: {e}")))
1021        })
1022    }
1023}
1024
1025/// Health status enum.
1026///
1027/// This enum represents the overall health status of a service.
1028#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
1029#[derive(Debug, Clone)]
1030pub enum RiHealthStatus {
1031    /// Service is healthy
1032    Healthy,
1033    /// Service is degraded but still functional
1034    Degraded,
1035    /// Service is unhealthy
1036    Unhealthy,
1037    /// Health status is unknown
1038    Unknown,
1039}
1040
1041/// Summary of health check results.
1042///
1043/// This struct provides an aggregated view of a service's health, including
1044/// total checks, success rate, average response time, and overall status.
1045#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
1046#[derive(Debug, Clone)]
1047pub struct RiHealthSummary {
1048    /// Name of the service
1049    pub service_name: String,
1050    /// Total number of health checks performed
1051    pub total_checks: usize,
1052    /// Number of successful health checks
1053    pub healthy_checks: usize,
1054    /// Number of failed health checks
1055    pub unhealthy_checks: usize,
1056    /// Success rate percentage (0.0 to 100.0)
1057    pub success_rate: f64,
1058    /// Average response time for health checks
1059    pub average_response_time: Duration,
1060    /// Timestamp of the last health check
1061    pub last_check_time: Option<SystemTime>,
1062    /// Overall health status
1063    pub overall_status: RiHealthStatus,
1064}
1065
1066#[cfg(feature = "pyo3")]
1067#[pyo3::prelude::pymethods]
1068impl RiHealthSummary {
1069    fn get_service_name(&self) -> String {
1070        self.service_name.clone()
1071    }
1072    
1073    fn get_total_checks(&self) -> usize {
1074        self.total_checks
1075    }
1076    
1077    fn get_healthy_checks(&self) -> usize {
1078        self.healthy_checks
1079    }
1080    
1081    fn get_unhealthy_checks(&self) -> usize {
1082        self.unhealthy_checks
1083    }
1084    
1085    fn get_success_rate(&self) -> f64 {
1086        self.success_rate
1087    }
1088    
1089    fn get_average_response_time_ms(&self) -> u64 {
1090        self.average_response_time.as_millis() as u64
1091    }
1092    
1093    fn get_overall_status(&self) -> String {
1094        match self.overall_status {
1095            RiHealthStatus::Healthy => "Healthy".to_string(),
1096            RiHealthStatus::Degraded => "Degraded".to_string(),
1097            RiHealthStatus::Unhealthy => "Unhealthy".to_string(),
1098            RiHealthStatus::Unknown => "Unknown".to_string(),
1099        }
1100    }
1101}
ri/service_mesh/health_check.rs

ri/service_mesh/
health_check.rs
