Files

427 lines
15 KiB
Python

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
镜像源健康检查模块
自动检测上游镜像源的可用性,支持故障切换
"""
import os
import sys
import json
import time
import threading
import logging
import urllib.request
import urllib.error
from datetime import datetime, timedelta
from typing import Dict, List, Optional, Callable
from enum import Enum
from dataclasses import dataclass, field
logger = logging.getLogger(__name__)
class HealthStatus(Enum):
"""健康状态"""
UNKNOWN = "unknown"
HEALTHY = "healthy"
DEGRADED = "degraded"
UNHEALTHY = "unhealthy"
@dataclass
class HealthCheckResult:
"""健康检查结果"""
source_name: str
status: HealthStatus
response_time: float # 毫秒
http_status: Optional[int] = None
error_message: Optional[str] = None
last_check: Optional[datetime] = None
consecutive_failures: int = 0
total_checks: int = 0
success_rate: float = 100.0
details: Dict = field(default_factory=dict)
class HealthChecker:
"""健康检查器"""
def __init__(self, config: dict = None):
"""
初始化健康检查器
Args:
config: 健康检查配置
"""
self.config = config or {}
self.default_timeout = self.config.get('timeout', 10) # 秒
self.default_interval = self.config.get('interval', 60) # 秒
self.max_retries = self.config.get('max_retries', 3)
self.failure_threshold = self.config.get('failure_threshold', 3) # 连续失败次数阈值
# 检查结果存储
self._results: Dict[str, HealthCheckResult] = {}
self._lock = threading.Lock()
# 状态变化回调
self._on_status_change: Optional[Callable] = None
def set_status_change_callback(self, callback: Callable):
"""设置状态变化回调"""
self._on_status_change = callback
def check_source(self, source_name: str, source_config: dict) -> HealthCheckResult:
"""
检查单个镜像源的健康状态
Args:
source_name: 镜像源名称
source_config: 镜像源配置
Returns:
HealthCheckResult: 检查结果
"""
url = source_config.get('url', '')
if not url:
return HealthCheckResult(
source_name=source_name,
status=HealthStatus.UNHEALTHY,
response_time=0,
error_message="No URL configured",
last_check=datetime.now()
)
start_time = time.time()
http_status = None
error_message = None
details = {}
try:
# 发送 HEAD 请求(更轻量)
req = urllib.request.Request(
url.rstrip('/') + '/',
method='HEAD',
headers={
'User-Agent': 'HYC-Mirror-HealthCheck/1.0',
'Accept': '*/*'
}
)
with urllib.request.urlopen(req, timeout=self.default_timeout) as response:
http_status = response.status
details['headers'] = dict(response.headers)
except urllib.error.HTTPError as e:
http_status = e.code
error_message = f"HTTP {e.code}"
except urllib.error.URLError as e:
error_message = f"Connection error: {str(e.reason)}"
except Exception as e:
error_message = str(e)
response_time = (time.time() - start_time) * 1000 # 转换为毫秒
# 判断健康状态
if error_message:
if http_status and 400 <= http_status < 500:
status = HealthStatus.DEGRADED # 客户端错误,可能暂时
else:
status = HealthStatus.UNHEALTHY
elif http_status and 200 <= http_status < 400:
status = HealthStatus.HEALTHY
elif http_status:
status = HealthStatus.DEGRADED
else:
status = HealthStatus.UNHEALTHY
# 计算统计数据
with self._lock:
if source_name not in self._results:
self._results[source_name] = HealthCheckResult(
source_name=source_name,
status=status,
response_time=response_time,
last_check=datetime.now(),
consecutive_failures=0,
total_checks=1
)
else:
old_result = self._results[source_name]
consecutive_failures = old_result.consecutive_failures + (1 if status == HealthStatus.UNHEALTHY else 0)
total_checks = old_result.total_checks + 1
success_rate = ((total_checks - consecutive_failures) / total_checks) * 100
self._results[source_name] = HealthCheckResult(
source_name=source_name,
status=status,
response_time=response_time,
http_status=http_status,
error_message=error_message,
last_check=datetime.now(),
consecutive_failures=consecutive_failures,
total_checks=total_checks,
success_rate=success_rate,
details=details
)
# 检查状态变化,触发回调
if self._on_status_change and source_name in self._results:
old_status = self._results[source_name].status
if old_status != status:
try:
self._on_status_change(source_name, old_status, status)
except Exception as e:
logger.error(f"状态变化回调执行失败: {e}")
return self._results[source_name]
def get_all_results(self) -> List[HealthCheckResult]:
"""获取所有检查结果"""
with self._lock:
return list(self._results.values())
def get_result(self, source_name: str) -> Optional[HealthCheckResult]:
"""获取指定源的结果"""
with self._lock:
return self._results.get(source_name)
def is_healthy(self, source_name: str) -> bool:
"""检查源是否健康"""
result = self.get_result(source_name)
if result is None:
return True # 未检查过的默认健康
return result.status == HealthStatus.HEALTHY
def get_unhealthy_sources(self) -> List[str]:
"""获取不健康的源列表"""
with self._lock:
return [name for name, result in self._results.items()
if result.status == HealthStatus.UNHEALTHY]
def get_stats(self) -> dict:
"""获取健康检查统计"""
with self._lock:
total = len(self._results)
healthy = sum(1 for r in self._results.values() if r.status == HealthStatus.HEALTHY)
degraded = sum(1 for r in self._results.values() if r.status == HealthStatus.DEGRADED)
unhealthy = sum(1 for r in self._results.values() if r.status == HealthStatus.UNHEALTHY)
avg_response_time = 0
if total > 0:
avg_response_time = sum(r.response_time for r in self._results.values()) / total
return {
'total_sources': total,
'healthy': healthy,
'degraded': degraded,
'unhealthy': unhealthy,
'avg_response_time_ms': round(avg_response_time, 2),
'timestamp': datetime.now().isoformat()
}
class MirrorFailoverManager:
"""镜像源故障切换管理器"""
def __init__(self, config: dict = None):
"""
初始化故障切换管理器
Args:
config: 配置,包含镜像源列表
"""
self.config = config or {}
self.mirrors: Dict[str, Dict] = self.config.get('mirrors', {})
# 启用故障切换
self.failover_enabled = self.config.get('failover_enabled', True)
self.failover_threshold = self.config.get('failover_threshold', 3) # 连续失败次数
# 健康检查器
self.health_checker = HealthChecker(self.config.get('health_check', {}))
# 当前活跃源
self._active_source: Dict[str, str] = {} # mirror_type -> source_name
self._source_priority: Dict[str, List[str]] = {} # mirror_type -> [优先列表]
# 故障切换历史
self._failover_history: List[dict] = []
# 回调
self._on_failover: Optional[Callable] = None
def set_failover_callback(self, callback: Callable):
"""设置故障切换回调"""
self._on_failover = callback
def initialize(self):
"""初始化,确定各镜像类型的首选源"""
for mirror_type, mirror_config in self.mirrors.items():
if not isinstance(mirror_config, dict):
continue
sources = mirror_config.get('sources', [])
if sources:
# 使用配置的优先列表
self._source_priority[mirror_type] = sources
else:
# 使用内置的默认优先列表
self._source_priority[mirror_type] = self._get_default_priority(mirror_type)
# 选择首选源
if self._source_priority[mirror_type]:
self._active_source[mirror_type] = self._source_priority[mirror_type][0]
def _get_default_priority(self, mirror_type: str) -> List[str]:
"""获取默认的镜像源优先列表"""
priorities = {
'docker': ['docker.io', 'docker.mirrors.aliyun.com', 'dockerhub.azk8s.cn'],
'apt': ['archive.ubuntu.com', 'mirrors.aliyun.com', 'security.ubuntu.com'],
'yum': ['mirror.centos.org', 'mirrors.aliyun.com'],
'pypi': ['pypi.org', 'pypi.mirrors.aliyun.com'],
'npm': ['registry.npmjs.org', 'registry.npmmirror.com'],
'go': ['proxy.golang.org', 'goproxy.cn']
}
return priorities.get(mirror_type, [])
def check_all(self) -> Dict[str, HealthCheckResult]:
"""检查所有镜像源"""
results = {}
for mirror_type, mirror_config in self.mirrors.items():
if not isinstance(mirror_config, dict):
continue
sources = mirror_config.get('sources', [])
for source_name in sources:
if source_name not in results:
result = self.health_checker.check_source(source_name, {'url': self._get_source_url(mirror_type, source_name)})
results[source_name] = result
return results
def _get_source_url(self, mirror_type: str, source_name: str) -> str:
"""获取源 URL"""
# 从配置中获取
sources_config = self.mirrors.get(mirror_type, {}).get('sources_config', {})
if source_name in sources_config:
return sources_config[source_name].get('url', '')
# 从 URL 模板生成
url_template = self.mirrors.get(mirror_type, {}).get('url_template', '')
if url_template and '{mirror}' in url_template:
return url_template.replace('{mirror}', source_name)
return ''
def get_active_source(self, mirror_type: str) -> Optional[str]:
"""获取当前活跃的镜像源"""
return self._active_source.get(mirror_type)
def get_source_for_request(self, mirror_type: str, original_url: str) -> str:
"""
根据故障切换策略获取请求的源 URL
Args:
mirror_type: 镜像类型
original_url: 原始 URL
Returns:
str: 实际请求的 URL
"""
if not self.failover_enabled:
return original_url
active_source = self._active_source.get(mirror_type)
if not active_source:
return original_url
source_url = self._get_source_url(mirror_type, active_source)
if not source_url:
return original_url
# 替换 URL 中的主机部分
try:
from urllib.parse import urlparse
parsed = urlparse(original_url)
# 构建新 URL
new_url = f"{parsed.scheme}://{source_url}{parsed.path}"
if parsed.query:
new_url += f"?{parsed.query}"
return new_url
except Exception:
return original_url
def perform_failover(self, mirror_type: str) -> bool:
"""
对指定镜像类型执行故障切换
Args:
mirror_type: 镜像类型
Returns:
bool: 是否成功切换
"""
priority_list = self._source_priority.get(mirror_type, [])
if not priority_list:
return False
current_index = 0
if mirror_type in self._active_source:
try:
current_index = priority_list.index(self._active_source[mirror_type])
except ValueError:
pass
# 查找下一个健康的源
old_source = self._active_source.get(mirror_type)
for i in range(current_index + 1, len(priority_list)):
source_name = priority_list[i]
result = self.health_checker.get_result(source_name)
if result and result.status == HealthStatus.HEALTHY:
self._active_source[mirror_type] = source_name
# 记录故障切换
failover_record = {
'timestamp': datetime.now().isoformat(),
'mirror_type': mirror_type,
'old_source': old_source,
'new_source': source_name,
'reason': 'Health check failed'
}
self._failover_history.append(failover_record)
# 保持历史记录在合理范围内
if len(self._failover_history) > 100:
self._failover_history = self._failover_history[-50:]
# 触发回调
if self._on_failover:
try:
self._on_failover(mirror_type, old_source, source_name)
except Exception as e:
logger.error(f"故障切换回调执行失败: {e}")
return True
return False
def get_failover_history(self, limit: int = 10) -> List[dict]:
"""获取故障切换历史"""
return self._failover_history[-limit:]
def get_health_summary(self) -> dict:
"""获取健康状态摘要"""
return {
'failover_enabled': self.failover_enabled,
'health': self.health_checker.get_stats(),
'active_sources': self._active_source.copy(),
'failover_history_count': len(self._failover_history)
}