678 lines
22 KiB
Python
678 lines
22 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
|
|
"""
|
|
缓存预热模块
|
|
预热常用镜像缓存提高命中率
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import json
|
|
import time
|
|
import threading
|
|
import logging
|
|
import requests
|
|
from datetime import datetime
|
|
from typing import Dict, List, Optional, Callable
|
|
from dataclasses import dataclass, field
|
|
from enum import Enum
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class PrewarmPriority(Enum):
|
|
"""预热优先级"""
|
|
CRITICAL = "critical" # 关键，必须预热
|
|
HIGH = "high" # 高优先级
|
|
MEDIUM = "medium" # 中优先级
|
|
LOW = "low" # 低优先级
|
|
|
|
|
|
class ItemStatus(Enum):
|
|
"""项目状态"""
|
|
PENDING = "pending"
|
|
IN_PROGRESS = "in_progress"
|
|
SUCCESS = "success"
|
|
FAILED = "failed"
|
|
SKIPPED = "skipped"
|
|
|
|
|
|
@dataclass
|
|
class PrewarmItem:
|
|
"""预热项目"""
|
|
id: str
|
|
mirror_type: str
|
|
item_name: str
|
|
url: str
|
|
priority: str
|
|
status: str = field(default=ItemStatus.PENDING.value)
|
|
attempts: int = 0
|
|
max_attempts: int = 3
|
|
response_time_ms: float = 0
|
|
error_message: str = None
|
|
prewarmed_at: float = None
|
|
size_bytes: int = 0
|
|
|
|
def to_dict(self) -> Dict:
|
|
return {
|
|
'id': self.id,
|
|
'mirror_type': self.mirror_type,
|
|
'item_name': self.item_name,
|
|
'url': self.url,
|
|
'priority': self.priority,
|
|
'status': self.status,
|
|
'attempts': self.attempts,
|
|
'response_time_ms': self.response_time_ms,
|
|
'error_message': self.error_message,
|
|
'prewarmed_at': self.prewarmed_at,
|
|
'size_bytes': self.size_bytes
|
|
}
|
|
|
|
|
|
@dataclass
|
|
class PrewarmTarget:
|
|
"""预热目标配置"""
|
|
mirror_type: str
|
|
priority: str
|
|
limit: int
|
|
items: List[str] = field(default_factory=list) # 指定的预热项目列表
|
|
tags: List[str] = field(default_factory=list) # 按标签筛选
|
|
|
|
|
|
class CachePrewarmer:
|
|
"""缓存预热管理器"""
|
|
|
|
def __init__(self, config: Dict = None):
|
|
"""
|
|
初始化缓存预热管理器
|
|
|
|
Args:
|
|
config: 预热配置
|
|
"""
|
|
self.config = config or {}
|
|
self.enabled = self.config.get('enabled', False)
|
|
|
|
# 预热目标
|
|
self.targets = self._parse_targets(self.config.get('targets', []))
|
|
|
|
# 调度配置
|
|
self.schedule = self.config.get('schedule', '0 3 * * *') # 默认每天凌晨3点
|
|
self.batch_size = self.config.get('batch_size', 10)
|
|
|
|
# 状态管理
|
|
self._items: Dict[str, PrewarmItem] = {}
|
|
self._items_lock = threading.Lock()
|
|
self._is_running = False
|
|
self._run_lock = threading.Lock()
|
|
|
|
# 历史记录
|
|
self._history: List[Dict] = []
|
|
self._history_lock = threading.Lock()
|
|
|
|
# 回调函数
|
|
self._on_start: Optional[Callable] = None
|
|
self._on_complete: Optional[Callable] = None
|
|
self._on_item_complete: Optional[Callable] = None
|
|
self._on_error: Optional[Callable] = None
|
|
|
|
# HTTP 会话
|
|
self._session = requests.Session()
|
|
self._session.headers.update({
|
|
'User-Agent': 'HYC-Mirror-Prewarmer/1.0'
|
|
})
|
|
|
|
# 超时设置
|
|
self._request_timeout = self.config.get('request_timeout', 30)
|
|
|
|
# 常用镜像包列表
|
|
self._popular_items = self._load_popular_items()
|
|
|
|
def _parse_targets(self, targets_config: List[Dict]) -> List[PrewarmTarget]:
|
|
"""解析预热目标配置"""
|
|
targets = []
|
|
for t in targets_config:
|
|
targets.append(PrewarmTarget(
|
|
mirror_type=t.get('mirror_type', 'docker'),
|
|
priority=t.get('priority', 'medium'),
|
|
limit=t.get('limit', 100),
|
|
items=t.get('items', []),
|
|
tags=t.get('tags', [])
|
|
))
|
|
return targets
|
|
|
|
def _load_popular_items(self) -> Dict[str, List[str]]:
|
|
"""加载常用镜像包列表"""
|
|
return {
|
|
'docker': [
|
|
'library/alpine:latest',
|
|
'library/ubuntu:latest',
|
|
'library/debian:latest',
|
|
'library/centos:latest',
|
|
'library/nginx:latest',
|
|
'library/python:3.9',
|
|
'library/python:3.10',
|
|
'library/node:18',
|
|
'library/node:20',
|
|
'library/go:1.20',
|
|
'library/redis:alpine',
|
|
'library/mysql:8',
|
|
'library/postgres:15',
|
|
],
|
|
'pip': [
|
|
'requests',
|
|
'numpy',
|
|
'pandas',
|
|
'flask',
|
|
'django',
|
|
'scipy',
|
|
'scikit-learn',
|
|
'torch',
|
|
'tensorflow',
|
|
'celery',
|
|
'pytest',
|
|
'black',
|
|
],
|
|
'npm': [
|
|
'react',
|
|
'vue',
|
|
'angular',
|
|
'lodash',
|
|
'express',
|
|
'axios',
|
|
'typescript',
|
|
'webpack',
|
|
'vite',
|
|
'eslint',
|
|
],
|
|
'apt': [
|
|
'ubuntu-desktop',
|
|
'ubuntu-standard',
|
|
'nginx',
|
|
'python3-pip',
|
|
'nodejs',
|
|
'docker.io',
|
|
],
|
|
'yum': [
|
|
'epel-release',
|
|
'nginx',
|
|
'docker-ce',
|
|
'python3-pip',
|
|
'nodejs',
|
|
],
|
|
'go': [
|
|
'golang.org/x/tools',
|
|
'github.com/gin-gonic/gin',
|
|
'github.com/beego/beego',
|
|
'github.com/gorilla/mux',
|
|
],
|
|
}
|
|
|
|
def _get_base_url(self, mirror_type: str) -> str:
|
|
"""获取镜像源基础 URL"""
|
|
mirrors = self.config.get('mirrors', {})
|
|
|
|
if isinstance(mirrors, dict):
|
|
mirror_config = mirrors.get(mirror_type, {})
|
|
if isinstance(mirror_config, dict):
|
|
sources = mirror_config.get('sources', [])
|
|
if sources:
|
|
source_config = mirror_config.get('sources_config', {}).get(sources[0], {})
|
|
return source_config.get('url', '')
|
|
|
|
return ''
|
|
|
|
def _generate_url(self, mirror_type: str, item_name: str) -> str:
|
|
"""生成预热 URL"""
|
|
base_url = self._get_base_url(mirror_type)
|
|
|
|
if not base_url:
|
|
return ''
|
|
|
|
if mirror_type == 'docker':
|
|
return f"{base_url}/v2/{item_name}/manifests/latest"
|
|
elif mirror_type == 'pip':
|
|
return f"{base_url}/simple/{item_name}/"
|
|
elif mirror_type == 'npm':
|
|
return f"{base_url}/{item_name}"
|
|
elif mirror_type == 'apt':
|
|
return f"{base_url}/dists/{item_name}/InRelease"
|
|
elif mirror_type == 'yum':
|
|
return f"{base_url}/repodata/repomd.xml"
|
|
elif mirror_type == 'go':
|
|
return f"{base_url}/{item_name}?go-get=1"
|
|
else:
|
|
return f"{base_url}/{item_name}"
|
|
|
|
def _create_item(
|
|
self,
|
|
mirror_type: str,
|
|
item_name: str,
|
|
priority: str = 'medium'
|
|
) -> PrewarmItem:
|
|
"""创建预热项目"""
|
|
url = self._generate_url(mirror_type, item_name)
|
|
|
|
return PrewarmItem(
|
|
id=f"{mirror_type}_{item_name}_{int(time.time())}",
|
|
mirror_type=mirror_type,
|
|
item_name=item_name,
|
|
url=url,
|
|
priority=priority
|
|
)
|
|
|
|
def add_item(self, item: PrewarmItem):
|
|
"""添加预热项目"""
|
|
with self._items_lock:
|
|
self._items[item.id] = item
|
|
|
|
def add_items_batch(self, mirror_type: str, items: List[str], priority: str = 'medium'):
|
|
"""批量添加预热项目"""
|
|
for item_name in items:
|
|
item = self._create_item(mirror_type, item_name, priority)
|
|
self.add_item(item)
|
|
|
|
def _prewarm_item(self, item: PrewarmItem) -> PrewarmItem:
|
|
"""
|
|
预热单个项目
|
|
|
|
Args:
|
|
item: 预热项目
|
|
|
|
Returns:
|
|
更新后的项目
|
|
"""
|
|
item.attempts += 1
|
|
item.status = ItemStatus.IN_PROGRESS.value
|
|
|
|
try:
|
|
start_time = time.time()
|
|
|
|
response = self._session.get(
|
|
item.url,
|
|
timeout=self._request_timeout,
|
|
allow_redirects=True
|
|
)
|
|
|
|
response.raise_for_status()
|
|
|
|
item.response_time_ms = round((time.time() - start_time) * 1000, 2)
|
|
item.size_bytes = len(response.content)
|
|
item.status = ItemStatus.SUCCESS.value
|
|
item.prewarmed_at = time.time()
|
|
|
|
logger.debug(f"Prewarmed {item.mirror_type}/{item.item_name}: {item.response_time_ms}ms")
|
|
|
|
except requests.exceptions.Timeout:
|
|
item.error_message = f"Timeout after {self._request_timeout}s"
|
|
if item.attempts < item.max_attempts:
|
|
item.status = ItemStatus.PENDING.value
|
|
else:
|
|
item.status = ItemStatus.FAILED.value
|
|
|
|
except requests.exceptions.HTTPError as e:
|
|
item.error_message = f"HTTP Error: {e.response.status_code}"
|
|
if item.attempts < item.max_attempts:
|
|
item.status = ItemStatus.PENDING.value
|
|
else:
|
|
item.status = ItemStatus.FAILED.value
|
|
|
|
except Exception as e:
|
|
item.error_message = str(e)
|
|
item.status = ItemStatus.FAILED.value
|
|
|
|
return item
|
|
|
|
def run(self, targets: List[PrewarmTarget] = None) -> Dict:
|
|
"""
|
|
执行预热
|
|
|
|
Args:
|
|
targets: 预热目标列表，None 表示使用配置的默认目标
|
|
|
|
Returns:
|
|
预热结果
|
|
"""
|
|
with self._run_lock:
|
|
if self._is_running:
|
|
return {
|
|
'success': False,
|
|
'error': 'Prewarm already running'
|
|
}
|
|
|
|
self._is_running = True
|
|
|
|
start_time = time.time()
|
|
result = {
|
|
'success': True,
|
|
'total_items': 0,
|
|
'success_count': 0,
|
|
'failed_count': 0,
|
|
'skipped_count': 0,
|
|
'elapsed_seconds': 0,
|
|
'errors': []
|
|
}
|
|
|
|
try:
|
|
# 通知开始
|
|
if self._on_start:
|
|
try:
|
|
self._on_start()
|
|
except Exception as e:
|
|
logger.error(f"Prewarm start callback failed: {e}")
|
|
|
|
# 确定要预热的目标
|
|
if targets is None:
|
|
targets = self.targets
|
|
|
|
# 如果没有指定目标，使用流行项目
|
|
if not targets:
|
|
for mirror_type, items in self._popular_items.items():
|
|
targets.append(PrewarmTarget(
|
|
mirror_type=mirror_type,
|
|
priority='medium',
|
|
limit=len(items),
|
|
items=items
|
|
))
|
|
|
|
# 添加项目到队列
|
|
total_added = 0
|
|
for target in targets:
|
|
if target.items:
|
|
# 使用指定的预热项目
|
|
for item_name in target.items[:target.limit]:
|
|
item = self._create_item(target.mirror_type, item_name, target.priority)
|
|
self.add_item(item)
|
|
total_added += 1
|
|
else:
|
|
# 使用流行项目列表
|
|
popular = self._popular_items.get(target.mirror_type, [])
|
|
for item_name in popular[:target.limit]:
|
|
item = self._create_item(target.mirror_type, item_name, target.priority)
|
|
self.add_item(item)
|
|
total_added += 1
|
|
|
|
result['total_items'] = total_added
|
|
|
|
# 按优先级排序
|
|
priority_order = {
|
|
'critical': 0,
|
|
'high': 1,
|
|
'medium': 2,
|
|
'low': 3
|
|
}
|
|
|
|
with self._items_lock:
|
|
sorted_items = sorted(
|
|
self._items.values(),
|
|
key=lambda x: (priority_order.get(x.priority, 99), x.id)
|
|
)
|
|
|
|
# 分批执行
|
|
completed = 0
|
|
with ThreadPoolExecutor(max_workers=self.batch_size) as executor:
|
|
futures = {
|
|
executor.submit(self._prewarm_item, item): item
|
|
for item in sorted_items
|
|
}
|
|
|
|
for future in as_completed(futures):
|
|
item = futures[future]
|
|
try:
|
|
updated_item = future.result()
|
|
|
|
# 更新项目状态
|
|
with self._items_lock:
|
|
self._items[item.id] = updated_item
|
|
|
|
# 统计
|
|
if updated_item.status == ItemStatus.SUCCESS.value:
|
|
result['success_count'] += 1
|
|
elif updated_item.status == ItemStatus.FAILED.value:
|
|
result['failed_count'] += 1
|
|
result['errors'].append({
|
|
'item': updated_item.item_name,
|
|
'error': updated_item.error_message
|
|
})
|
|
else:
|
|
result['skipped_count'] += 1
|
|
|
|
completed += 1
|
|
|
|
# 回调
|
|
if self._on_item_complete:
|
|
try:
|
|
self._on_item_complete(updated_item)
|
|
except Exception as e:
|
|
logger.error(f"Item complete callback failed: {e}")
|
|
|
|
except Exception as e:
|
|
logger.error(f"Prewarm execution error: {e}")
|
|
result['failed_count'] += 1
|
|
result['errors'].append({
|
|
'item': item.item_name,
|
|
'error': str(e)
|
|
})
|
|
|
|
result['elapsed_seconds'] = round(time.time() - start_time, 2)
|
|
|
|
# 记录历史
|
|
self._add_to_history(result)
|
|
|
|
# 通知完成
|
|
if self._on_complete:
|
|
try:
|
|
self._on_complete(result)
|
|
except Exception as e:
|
|
logger.error(f"Prewarm complete callback failed: {e}")
|
|
|
|
except Exception as e:
|
|
logger.error(f"Prewarm failed: {e}")
|
|
result['success'] = False
|
|
result['errors'].append({'error': str(e)})
|
|
|
|
finally:
|
|
self._is_running = False
|
|
|
|
return result
|
|
|
|
def _add_to_history(self, result: Dict):
|
|
"""添加历史记录"""
|
|
record = {
|
|
'timestamp': datetime.now().isoformat(),
|
|
'success': result.get('success', False),
|
|
'total_items': result.get('total_items', 0),
|
|
'success_count': result.get('success_count', 0),
|
|
'failed_count': result.get('failed_count', 0),
|
|
'elapsed_seconds': result.get('elapsed_seconds', 0)
|
|
}
|
|
|
|
with self._history_lock:
|
|
self._history.append(record)
|
|
# 只保留最近 50 条记录
|
|
self._history = self._history[-50:]
|
|
|
|
def get_status(self) -> Dict:
|
|
"""获取预热状态"""
|
|
with self._items_lock:
|
|
items = list(self._items.values())
|
|
|
|
total = len(items)
|
|
success = sum(1 for i in items if i.status == ItemStatus.SUCCESS.value)
|
|
failed = sum(1 for i in items if i.status == ItemStatus.FAILED.value)
|
|
in_progress = sum(1 for i in items if i.status == ItemStatus.IN_PROGRESS.value)
|
|
pending = sum(1 for i in items if i.status == ItemStatus.PENDING.value)
|
|
|
|
return {
|
|
'enabled': self.enabled,
|
|
'is_running': self._is_running,
|
|
'total_items': total,
|
|
'success_count': success,
|
|
'failed_count': failed,
|
|
'in_progress_count': in_progress,
|
|
'pending_count': pending,
|
|
'success_rate': (success / total * 100) if total > 0 else 0,
|
|
'targets_count': len(self.targets)
|
|
}
|
|
|
|
def get_items(
|
|
self,
|
|
status: str = None,
|
|
mirror_type: str = None,
|
|
limit: int = 50
|
|
) -> List[Dict]:
|
|
"""获取预热项目列表"""
|
|
with self._items_lock:
|
|
items = [i.to_dict() for i in self._items.values()]
|
|
|
|
# 过滤
|
|
if status:
|
|
items = [i for i in items if i['status'] == status]
|
|
if mirror_type:
|
|
items = [i for i in items if i['mirror_type'] == mirror_type]
|
|
|
|
return items[-limit:]
|
|
|
|
def get_history(self, limit: int = 20) -> List[Dict]:
|
|
"""获取预热历史"""
|
|
with self._history_lock:
|
|
return list(self._history[-limit:])
|
|
|
|
def get_stats(self) -> Dict:
|
|
"""获取统计信息"""
|
|
status = self.get_status()
|
|
history = self.get_history(10)
|
|
|
|
avg_duration = 0
|
|
if history:
|
|
durations = [h['elapsed_seconds'] for h in history if 'elapsed_seconds' in h]
|
|
if durations:
|
|
avg_duration = sum(durations) / len(durations)
|
|
|
|
return {
|
|
'enabled': self.enabled,
|
|
'is_running': self._is_running,
|
|
'total_prewarmed': status['success_count'],
|
|
'total_failed': status['failed_count'],
|
|
'success_rate': status['success_rate'],
|
|
'avg_duration_seconds': round(avg_duration, 2),
|
|
'recent_runs': len(history),
|
|
'targets': [
|
|
{
|
|
'mirror_type': t.mirror_type,
|
|
'priority': t.priority,
|
|
'limit': t.limit
|
|
}
|
|
for t in self.targets
|
|
]
|
|
}
|
|
|
|
def clear_items(self, status: str = None):
|
|
"""清除预热项目"""
|
|
with self._items_lock:
|
|
if status:
|
|
self._items = {
|
|
k: v for k, v in self._items.items()
|
|
if v.status != status
|
|
}
|
|
else:
|
|
self._items = {}
|
|
|
|
def set_start_callback(self, callback: Callable):
|
|
"""设置开始回调"""
|
|
self._on_start = callback
|
|
|
|
def set_complete_callback(self, callback: Callable):
|
|
"""设置完成回调"""
|
|
self._on_complete = callback
|
|
|
|
def set_item_complete_callback(self, callback: Callable):
|
|
"""设置项目完成回调"""
|
|
self._on_item_complete = callback
|
|
|
|
def set_error_callback(self, callback: Callable):
|
|
"""设置错误回调"""
|
|
self._on_error = callback
|
|
|
|
def get_popular_items(self, mirror_type: str) -> List[str]:
|
|
"""获取指定镜像类型的流行项目列表"""
|
|
return self._popular_items.get(mirror_type, [])
|
|
|
|
def add_popular_items_to_queue(
|
|
self,
|
|
mirror_type: str,
|
|
limit: int = None,
|
|
priority: str = 'medium'
|
|
):
|
|
"""
|
|
添加流行项目到预热队列
|
|
|
|
Args:
|
|
mirror_type: 镜像类型
|
|
limit: 数量限制
|
|
priority: 优先级
|
|
"""
|
|
popular = self._popular_items.get(mirror_type, [])
|
|
if limit:
|
|
popular = popular[:limit]
|
|
|
|
self.add_items_batch(mirror_type, popular, priority)
|
|
|
|
|
|
class PrewarmScheduler:
|
|
"""预热调度器"""
|
|
|
|
def __init__(self, config: Dict = None):
|
|
self.config = config or {}
|
|
self._scheduler = None
|
|
self._running = False
|
|
|
|
def start(self, prewarmer: CachePrewarmer):
|
|
"""启动调度器"""
|
|
if self._running:
|
|
return
|
|
|
|
schedule = self.config.get('schedule', '0 3 * * *')
|
|
logger.info(f"Starting prewarm scheduler with schedule: {schedule}")
|
|
|
|
# 使用简单的时间间隔检查
|
|
self._running = True
|
|
self._scheduler_thread = threading.Thread(
|
|
target=self._run_scheduler,
|
|
args=(prewarmer,),
|
|
daemon=True
|
|
)
|
|
self._scheduler_thread.start()
|
|
|
|
def _run_scheduler(self, prewarmer: CachePrewarmer):
|
|
"""运行调度器"""
|
|
import croniter
|
|
|
|
schedule = self.config.get('schedule', '0 3 * * *')
|
|
|
|
try:
|
|
cron = croniter.croniter(schedule, datetime.now())
|
|
next_run = cron.get_next(datetime)
|
|
except Exception as e:
|
|
logger.error(f"Invalid cron schedule: {e}")
|
|
return
|
|
|
|
while self._running:
|
|
now = datetime.now()
|
|
if now >= next_run:
|
|
logger.info("Running scheduled prewarm")
|
|
try:
|
|
prewarmer.run()
|
|
except Exception as e:
|
|
logger.error(f"Scheduled prewarm failed: {e}")
|
|
cron = croniter.croniter(schedule, datetime.now())
|
|
next_run = cron.get_next(datetime)
|
|
|
|
# 检查间隔
|
|
time.sleep(60)
|
|
|
|
def stop(self):
|
|
"""停止调度器"""
|
|
self._running = False
|