194 lines
7.4 KiB
Python
194 lines
7.4 KiB
Python
"""
|
|
Service registry for managing microservice discovery and health monitoring
|
|
"""
|
|
|
|
import aiohttp
|
|
import asyncio
|
|
from datetime import datetime
|
|
from typing import Dict, List, Optional
|
|
import logging
|
|
|
|
from models import ServiceConfig, ServiceHealth
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class ServiceRegistry:
|
|
"""Service registry for microservice management"""
|
|
|
|
def __init__(self):
|
|
self.services: Dict[str, ServiceConfig] = {}
|
|
self.service_health: Dict[str, ServiceHealth] = {}
|
|
self.session: Optional[aiohttp.ClientSession] = None
|
|
|
|
async def initialize(self):
|
|
"""Initialize the service registry"""
|
|
self.session = aiohttp.ClientSession(
|
|
timeout=aiohttp.ClientTimeout(total=10)
|
|
)
|
|
logger.info("Service registry initialized")
|
|
|
|
async def close(self):
|
|
"""Close the service registry"""
|
|
if self.session:
|
|
await self.session.close()
|
|
logger.info("Service registry closed")
|
|
|
|
async def register_services(self, services: Dict[str, ServiceConfig]):
|
|
"""Register multiple services"""
|
|
self.services.update(services)
|
|
|
|
# Initialize health status for all services
|
|
for service_name, config in services.items():
|
|
self.service_health[service_name] = ServiceHealth(
|
|
service=service_name,
|
|
status="unknown",
|
|
last_check=datetime.utcnow()
|
|
)
|
|
|
|
logger.info(f"Registered {len(services)} services")
|
|
|
|
# Perform initial health check
|
|
await self.update_all_service_health()
|
|
|
|
async def register_service(self, service_config: ServiceConfig):
|
|
"""Register a single service"""
|
|
self.services[service_config.name] = service_config
|
|
self.service_health[service_config.name] = ServiceHealth(
|
|
service=service_config.name,
|
|
status="unknown",
|
|
last_check=datetime.utcnow()
|
|
)
|
|
|
|
logger.info(f"Registered service: {service_config.name}")
|
|
|
|
# Check health of the newly registered service
|
|
await self.check_service_health(service_config.name)
|
|
|
|
async def unregister_service(self, service_name: str):
|
|
"""Unregister a service"""
|
|
self.services.pop(service_name, None)
|
|
self.service_health.pop(service_name, None)
|
|
logger.info(f"Unregistered service: {service_name}")
|
|
|
|
async def check_service_health(self, service_name: str) -> ServiceHealth:
|
|
"""Check health of a specific service"""
|
|
service_config = self.services.get(service_name)
|
|
if not service_config:
|
|
logger.error(f"Service {service_name} not found in registry")
|
|
return ServiceHealth(
|
|
service=service_name,
|
|
status="unknown",
|
|
last_check=datetime.utcnow(),
|
|
error_message="Service not registered"
|
|
)
|
|
|
|
start_time = datetime.utcnow()
|
|
|
|
try:
|
|
health_url = f"{service_config.base_url}{service_config.health_endpoint}"
|
|
|
|
async with self.session.get(health_url) as response:
|
|
end_time = datetime.utcnow()
|
|
response_time = (end_time - start_time).total_seconds() * 1000
|
|
|
|
if response.status == 200:
|
|
health_data = await response.json()
|
|
status = "healthy" if health_data.get("status") in ["healthy", "ok"] else "unhealthy"
|
|
|
|
health = ServiceHealth(
|
|
service=service_name,
|
|
status=status,
|
|
response_time_ms=response_time,
|
|
last_check=end_time
|
|
)
|
|
else:
|
|
health = ServiceHealth(
|
|
service=service_name,
|
|
status="unhealthy",
|
|
response_time_ms=response_time,
|
|
last_check=end_time,
|
|
error_message=f"HTTP {response.status}"
|
|
)
|
|
|
|
except aiohttp.ClientError as e:
|
|
health = ServiceHealth(
|
|
service=service_name,
|
|
status="unhealthy",
|
|
last_check=datetime.utcnow(),
|
|
error_message=f"Connection error: {str(e)}"
|
|
)
|
|
except Exception as e:
|
|
health = ServiceHealth(
|
|
service=service_name,
|
|
status="unhealthy",
|
|
last_check=datetime.utcnow(),
|
|
error_message=f"Health check failed: {str(e)}"
|
|
)
|
|
|
|
# Update health status
|
|
self.service_health[service_name] = health
|
|
|
|
# Log health status changes
|
|
if health.status != "healthy":
|
|
logger.warning(f"Service {service_name} health check failed: {health.error_message}")
|
|
|
|
return health
|
|
|
|
async def update_all_service_health(self):
|
|
"""Update health status for all registered services"""
|
|
health_checks = [
|
|
self.check_service_health(service_name)
|
|
for service_name in self.services.keys()
|
|
]
|
|
|
|
if health_checks:
|
|
await asyncio.gather(*health_checks, return_exceptions=True)
|
|
|
|
# Log summary
|
|
healthy_count = sum(1 for h in self.service_health.values() if h.status == "healthy")
|
|
total_count = len(self.services)
|
|
logger.info(f"Health check complete: {healthy_count}/{total_count} services healthy")
|
|
|
|
async def get_service_health(self, service_name: str) -> Optional[ServiceHealth]:
|
|
"""Get health status of a specific service"""
|
|
return self.service_health.get(service_name)
|
|
|
|
async def get_all_service_health(self) -> Dict[str, Dict]:
|
|
"""Get health status of all services"""
|
|
health_dict = {}
|
|
for service_name, health in self.service_health.items():
|
|
health_dict[service_name] = {
|
|
"status": health.status,
|
|
"response_time_ms": health.response_time_ms,
|
|
"last_check": health.last_check.isoformat(),
|
|
"error_message": health.error_message
|
|
}
|
|
return health_dict
|
|
|
|
async def is_service_healthy(self, service_name: str) -> bool:
|
|
"""Check if a service is healthy"""
|
|
health = self.service_health.get(service_name)
|
|
return health is not None and health.status == "healthy"
|
|
|
|
async def get_healthy_services(self) -> List[str]:
|
|
"""Get list of healthy service names"""
|
|
return [
|
|
service_name
|
|
for service_name, health in self.service_health.items()
|
|
if health.status == "healthy"
|
|
]
|
|
|
|
def get_service_config(self, service_name: str) -> Optional[ServiceConfig]:
|
|
"""Get configuration for a specific service"""
|
|
return self.services.get(service_name)
|
|
|
|
def get_all_services(self) -> Dict[str, ServiceConfig]:
|
|
"""Get all registered services"""
|
|
return self.services.copy()
|
|
|
|
async def get_service_url(self, service_name: str) -> Optional[str]:
|
|
"""Get base URL for a healthy service"""
|
|
if await self.is_service_healthy(service_name):
|
|
service_config = self.services.get(service_name)
|
|
return service_config.base_url if service_config else None
|
|
return None |