use serde::{Deserialize, Serialize}; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; use std::time::Instant; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::TcpListener; type Result = std::result::Result>; #[derive(Debug, Clone, Serialize, Deserialize)] pub struct HealthStatus { pub status: String, pub version: String, pub uptime_seconds: u64, pub connected: bool, } /// Start a simple health check HTTP server using raw TCP. /// /// The health check returns: /// - 200 OK with status "healthy" when WebSocket is connected /// - 503 Service Unavailable with status "unhealthy" when WebSocket is disconnected /// /// This allows Kubernetes readiness probes to remove the pod from service /// when the WebSocket connection is broken. pub async fn start_health_server(port: u16, connected: Arc) -> Result<()> { let start_time = Arc::new(Instant::now()); let addr = format!("0.0.0.0:{}", port); crate::log_info!("Starting health endpoint on {}", addr); let listener = TcpListener::bind(&addr).await?; loop { match listener.accept().await { Ok((mut socket, _)) => { let start_time = Arc::clone(&start_time); let connected = Arc::clone(&connected); tokio::spawn(async move { let mut buffer = [0u8; 1024]; // Read the HTTP request if let Ok(n) = socket.read(&mut buffer).await { if n > 0 { let request = String::from_utf8_lossy(&buffer[..n]); // Check if this is a GET request to /health if request.starts_with("GET /health") { let is_connected = connected.load(Ordering::Relaxed); let uptime = start_time.elapsed().as_secs(); let status = HealthStatus { status: if is_connected { "healthy".to_string() } else { "unhealthy".to_string() }, version: env!("CARGO_PKG_VERSION").to_string(), uptime_seconds: uptime, connected: is_connected, }; let json = serde_json::to_string(&status).unwrap_or_else(|_| { r#"{"status":"error","message":"Failed to serialize health status"}"# .to_string() }); // Return 200 if connected, 503 if not let http_status = if is_connected { "200 OK" } else { "503 Service Unavailable" }; let response = format!( "HTTP/1.1 {}\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n{}", http_status, json.len(), json ); let _ = socket.write_all(response.as_bytes()).await; } else { // 404 for other paths let response = "HTTP/1.1 404 Not Found\r\nContent-Length: 9\r\n\r\nNot found"; let _ = socket.write_all(response.as_bytes()).await; } } } let _ = socket.shutdown().await; }); } Err(e) => { crate::log_warn!("Failed to accept health check connection: {}", e); } } } } #[cfg(test)] mod tests { use super::*; #[test] fn test_health_status_serialization_healthy() { let status = HealthStatus { status: "healthy".to_string(), version: "0.1.0".to_string(), uptime_seconds: 42, connected: true, }; let json = serde_json::to_string(&status).unwrap(); assert!(json.contains(r#""status":"healthy""#)); assert!(json.contains(r#""version":"0.1.0""#)); assert!(json.contains(r#""uptime_seconds":42"#)); assert!(json.contains(r#""connected":true"#)); } #[test] fn test_health_status_serialization_unhealthy() { let status = HealthStatus { status: "unhealthy".to_string(), version: "0.1.0".to_string(), uptime_seconds: 100, connected: false, }; let json = serde_json::to_string(&status).unwrap(); assert!(json.contains(r#""status":"unhealthy""#)); assert!(json.contains(r#""connected":false"#)); } #[test] fn test_health_status_deserialization() { let json = r#"{"status":"healthy","version":"0.1.0","uptime_seconds":42,"connected":true}"#; let status: HealthStatus = serde_json::from_str(json).unwrap(); assert_eq!(status.status, "healthy"); assert_eq!(status.version, "0.1.0"); assert_eq!(status.uptime_seconds, 42); assert!(status.connected); } #[test] fn test_health_status_clone() { let status = HealthStatus { status: "healthy".to_string(), version: "0.1.0".to_string(), uptime_seconds: 42, connected: true, }; let cloned = status.clone(); assert_eq!(status.status, cloned.status); assert_eq!(status.version, cloned.version); assert_eq!(status.uptime_seconds, cloned.uptime_seconds); assert_eq!(status.connected, cloned.connected); } // Note: start_health_server is tested manually/via integration tests // Unit testing it requires complex async server mocking which is not practical }