rex/service_supervisor.py
Rama master · réplica de lectura del repositorio original.
"""
Service supervision and health management for Rex AI Assistant.
Provides automatic service restart on crash, exponential backoff retry logic,
and a simple health check endpoint.
"""
from __future__ import annotations
import json
import logging
import threading
import time
from collections.abc import Callable
from dataclasses import dataclass, field
from datetime import UTC, datetime
from http.server import BaseHTTPRequestHandler, HTTPServer
logger = logging.getLogger(__name__)
@dataclass
class ServiceHealthStatus:
"""Health status of a single service."""
name: str
is_running: bool
last_checked: datetime = field(default_factory=lambda: datetime.now(UTC))
last_error: str | None = None
restart_count: int = 0
uptime_seconds: float = 0.0
@dataclass
class ServiceMetrics:
"""Cumulative metrics for service supervision."""
total_restarts: int = 0
total_crashes: int = 0
total_uptime_seconds: float = 0.0
last_health_check: datetime = field(default_factory=lambda: datetime.now(UTC))
successful_checks: int = 0
failed_checks: int = 0
class ManagedService:
"""Represents a service managed by the supervisor."""
def __init__(
self,
name: str,
start_func: Callable[[], None],
stop_func: Callable[[], None],
health_check_func: Callable[[], bool] = None, # type: ignore[assignment]
metrics_func: Callable[[], dict] | None = None,
max_restart_attempts: int = 5,
backoff_multiplier: float = 1.5,
initial_backoff_seconds: float = 1.0,
):
self.name = name
self.start_func = start_func
self.stop_func = stop_func
self.health_check_func = health_check_func or (lambda: True) # type: ignore[truthy-function]
self.metrics_func = metrics_func or (lambda: {})
self.max_restart_attempts = max_restart_attempts
self.backoff_multiplier = backoff_multiplier
self.initial_backoff_seconds = initial_backoff_seconds
self.thread: threading.Thread | None = None
self.is_running = False
self.restart_count = 0
self.start_time: datetime | None = None
self.last_error: str | None = None
def status(self) -> ServiceHealthStatus:
"""Get current health status."""
uptime = 0.0
if self.start_time:
uptime = (datetime.now(UTC) - self.start_time).total_seconds()
return ServiceHealthStatus(
name=self.name,
is_running=self.is_running,
last_error=self.last_error,
restart_count=self.restart_count,
uptime_seconds=uptime,
)
class ServiceSupervisor:
"""
Supervises critical services with automatic restart and health checks.
Manages service lifecycle with exponential backoff retry logic and
provides a simple health check endpoint.
"""
def __init__(self, health_check_port: int = 8765):
self.services: dict[str, ManagedService] = {}
self.metrics = ServiceMetrics()
self.health_check_port = health_check_port
self.is_running = False
self._health_check_thread: threading.Thread | None = None
self._http_server_thread: threading.Thread | None = None
self._http_server: HTTPServer | None = None
self._lock = threading.RLock()
def register_service(
self,
name: str,
start_func: Callable[[], None],
stop_func: Callable[[], None],
health_check_func: Callable[[], bool] = None, # type: ignore[assignment]
metrics_func: Callable[[], dict] | None = None,
max_restart_attempts: int = 5,
backoff_multiplier: float = 1.5,
initial_backoff_seconds: float = 1.0,
) -> None:
"""Register a service to be supervised."""
with self._lock:
service = ManagedService(
name=name,
start_func=start_func,
stop_func=stop_func,
health_check_func=health_check_func,
metrics_func=metrics_func,
max_restart_attempts=max_restart_attempts,
backoff_multiplier=backoff_multiplier,
initial_backoff_seconds=initial_backoff_seconds,
)
self.services[name] = service
logger.info(f"Registered service: {name}")
def start(self) -> None:
"""Start all registered services."""
with self._lock:
if self.is_running:
logger.warning("Supervisor is already running")
return
logger.info("Starting service supervisor")
self.is_running = True
# Start all services
for service in self.services.values():
self._start_service(service)
# Start health check loop
self._health_check_thread = threading.Thread(
target=self._health_check_loop, daemon=True, name="rex-health-check"
)
self._health_check_thread.start()
# Start HTTP health endpoint
self._http_server_thread = threading.Thread(
target=self._run_health_server, daemon=True, name="rex-health-server"
)
self._http_server_thread.start()
logger.info(f"Health check server started on port {self.health_check_port}")
def stop(self) -> None:
"""Stop all services gracefully."""
with self._lock:
if not self.is_running:
logger.warning("Supervisor is not running")
return
logger.info("Stopping service supervisor")
self.is_running = False
# Stop all services in reverse order
for service in reversed(list(self.services.values())):
self._stop_service(service)
# Stop HTTP server
if self._http_server:
self._http_server.shutdown()
self._http_server = None
def _start_service(self, service: ManagedService) -> None:
"""Start a single service with error handling."""
try:
logger.info(f"Starting service: {service.name}")
service.start_time = datetime.now(UTC)
service.is_running = True
service.start_func()
logger.info(f"Service started: {service.name}")
except Exception as e:
logger.error(f"Failed to start service {service.name}: {e}")
service.is_running = False
service.last_error = str(e)
def _stop_service(self, service: ManagedService) -> None:
"""Stop a single service gracefully."""
try:
logger.info(f"Stopping service: {service.name}")
service.stop_func()
service.is_running = False
logger.info(f"Service stopped: {service.name}")
except Exception as e:
logger.error(f"Error stopping service {service.name}: {e}")
def _health_check_loop(self) -> None:
"""Continuously monitor service health."""
while self.is_running:
try:
self._perform_health_check()
time.sleep(5) # Check every 5 seconds
except Exception as e:
logger.error(f"Error in health check loop: {e}")
self.metrics.failed_checks += 1
time.sleep(5)
def _perform_health_check(self) -> None:
"""Check health of all services and restart if needed."""
with self._lock:
self.metrics.last_health_check = datetime.now(UTC)
all_healthy = True
for service in self.services.values():
try:
# Check if service thinks it's healthy
if service.is_running and not service.health_check_func():
logger.warning(f"Health check failed for {service.name}")
all_healthy = False
self._restart_service(service)
except Exception as e:
logger.error(f"Health check exception for {service.name}: {e}")
all_healthy = False
self._restart_service(service)
if all_healthy:
self.metrics.successful_checks += 1
else:
self.metrics.failed_checks += 1
def _restart_service(self, service: ManagedService, attempt: int = 0) -> None:
"""Restart a service with exponential backoff."""
if attempt >= service.max_restart_attempts:
logger.error(
f"Service {service.name} exceeded max restart attempts ({service.max_restart_attempts})"
)
service.is_running = False
return
# Calculate backoff
backoff_seconds = service.initial_backoff_seconds * (service.backoff_multiplier**attempt)
logger.info(
f"Restarting service {service.name} in {backoff_seconds:.1f}s (attempt {attempt + 1}/{service.max_restart_attempts})"
)
# Wait with backoff
time.sleep(backoff_seconds)
# Try to restart
try:
self._stop_service(service)
time.sleep(0.5)
self._start_service(service)
service.restart_count += 1
self.metrics.total_restarts += 1
logger.info(f"Service {service.name} restarted successfully")
except Exception as e:
logger.error(f"Failed to restart {service.name}: {e}")
service.last_error = str(e)
self.metrics.total_crashes += 1
# Retry with next backoff
self._restart_service(service, attempt + 1)
def get_service_status(self, service_name: str) -> ServiceHealthStatus | None:
"""Get health status of a specific service."""
with self._lock:
service = self.services.get(service_name)
if service:
return service.status()
return None
def get_all_status(self) -> dict[str, ServiceHealthStatus]:
"""Get health status of all services."""
with self._lock:
return {name: service.status() for name, service in self.services.items()}
def get_all_metrics(self) -> dict[str, dict]:
"""Get metrics payload from all services."""
with self._lock:
metrics: dict[str, dict] = {}
for name, service in self.services.items():
try:
metrics[name] = service.metrics_func()
except Exception as exc:
logger.warning("Failed to collect metrics for %s: %s", name, exc)
metrics[name] = {"error": str(exc)}
return metrics
def get_metrics(self) -> ServiceMetrics:
"""Get cumulative supervisor metrics."""
with self._lock:
return self.metrics
def _run_health_server(self) -> None:
"""Run HTTP health check server."""
try:
handler = self._create_health_handler()
self._http_server = HTTPServer(("127.0.0.1", self.health_check_port), handler)
logger.info(f"Health server listening on port {self.health_check_port}")
self._http_server.serve_forever()
except Exception as e:
logger.error(f"Failed to start health server: {e}")
def _create_health_handler(self):
"""Create HTTP request handler for health checks."""
supervisor = self
class HealthHandler(BaseHTTPRequestHandler):
def do_GET(self):
if self.path == "/health":
self.send_response(200)
self.send_header("Content-Type", "application/json")
self.end_headers()
status = supervisor.get_all_status()
all_healthy = all(s.is_running for s in status.values())
response = {
"status": "healthy" if all_healthy else "degraded",
"services": {
name: {
"running": s.is_running,
"restarts": s.restart_count,
"uptime_seconds": s.uptime_seconds,
"error": s.last_error,
}
for name, s in status.items()
},
"service_metrics": supervisor.get_all_metrics(),
"metrics": {
"total_restarts": supervisor.metrics.total_restarts,
"total_crashes": supervisor.metrics.total_crashes,
"successful_checks": supervisor.metrics.successful_checks,
"failed_checks": supervisor.metrics.failed_checks,
},
}
self.wfile.write(json.dumps(response).encode())
elif self.path == "/ready":
status = supervisor.get_all_status()
all_ready = all(s.is_running for s in status.values())
if all_ready:
self.send_response(200)
self.send_header("Content-Type", "text/plain")
self.end_headers()
self.wfile.write(b"Ready")
else:
self.send_response(503)
self.send_header("Content-Type", "text/plain")
self.end_headers()
self.wfile.write(b"Not ready")
elif self.path == "/metrics":
self.send_response(200)
self.send_header("Content-Type", "application/json")
self.end_headers()
payload = {
"supervisor": {
"total_restarts": supervisor.metrics.total_restarts,
"total_crashes": supervisor.metrics.total_crashes,
"successful_checks": supervisor.metrics.successful_checks,
"failed_checks": supervisor.metrics.failed_checks,
},
"services": {
name: {
"status": status.is_running,
"uptime_seconds": status.uptime_seconds,
"restarts": status.restart_count,
"last_error": status.last_error,
"metrics": supervisor.services[name].metrics_func(),
}
for name, status in supervisor.get_all_status().items()
},
}
self.wfile.write(json.dumps(payload).encode())
else:
self.send_response(404)
self.send_header("Content-Type", "text/plain")
self.end_headers()
self.wfile.write(b"Not found")
def log_message(self, format, *args):
# Suppress default logging
pass
return HealthHandler
__all__ = [
"ServiceSupervisor",
"ManagedService",
"ServiceHealthStatus",
"ServiceMetrics",
]