Baseline: pr1 HYC下载站 v2.3 before security/functional fixes
This commit is contained in:
HYC Fixer
@@ -0,0 +1,677 @@
|
||||
#!/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
|
||||
Reference in New Issue
Block a user