'commit'
This commit is contained in:
HuangHai
421
DouYin/apiproxy/douyin/core/orchestrator.py
Normal file
421
DouYin/apiproxy/douyin/core/orchestrator.py
Normal file
@@ -0,0 +1,421 @@
|
||||
#!/usr/bin/env python
|
||||
# -*- coding: utf-8 -*-
|
||||
|
||||
"""
|
||||
下载编排器
|
||||
协调多种下载策略，实现智能降级和任务管理
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
import logging
|
||||
import uuid
|
||||
from typing import List, Dict, Any, Optional
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
|
||||
from apiproxy.douyin.strategies.base import (
|
||||
IDownloadStrategy,
|
||||
DownloadTask,
|
||||
DownloadResult,
|
||||
TaskType,
|
||||
TaskStatus
|
||||
)
|
||||
from apiproxy.douyin.strategies.api_strategy import EnhancedAPIStrategy
|
||||
from apiproxy.douyin.strategies.retry_strategy import RetryStrategy
|
||||
from .rate_limiter import AdaptiveRateLimiter, RateLimitConfig
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class OrchestratorConfig:
|
||||
"""编排器配置"""
|
||||
def __init__(
|
||||
self,
|
||||
max_concurrent: int = 5,
|
||||
enable_retry: bool = True,
|
||||
enable_rate_limit: bool = True,
|
||||
rate_limit_config: Optional[RateLimitConfig] = None,
|
||||
priority_queue: bool = True,
|
||||
save_progress: bool = True
|
||||
):
|
||||
self.max_concurrent = max_concurrent
|
||||
self.enable_retry = enable_retry
|
||||
self.enable_rate_limit = enable_rate_limit
|
||||
self.rate_limit_config = rate_limit_config or RateLimitConfig()
|
||||
self.priority_queue = priority_queue
|
||||
self.save_progress = save_progress
|
||||
|
||||
|
||||
class DownloadOrchestrator:
|
||||
"""下载任务编排器"""
|
||||
|
||||
def __init__(self, config: Optional[OrchestratorConfig] = None):
|
||||
"""
|
||||
初始化编排器
|
||||
|
||||
Args:
|
||||
config: 编排器配置
|
||||
"""
|
||||
self.config = config or OrchestratorConfig()
|
||||
self.strategies: List[IDownloadStrategy] = []
|
||||
self.rate_limiter = AdaptiveRateLimiter(self.config.rate_limit_config) if self.config.enable_rate_limit else None
|
||||
|
||||
# 任务队列
|
||||
self.pending_queue = asyncio.Queue()
|
||||
self.priority_tasks: List[DownloadTask] = []
|
||||
self.active_tasks: Dict[str, DownloadTask] = {}
|
||||
self.completed_tasks: List[DownloadTask] = []
|
||||
self.failed_tasks: List[DownloadTask] = []
|
||||
|
||||
# 工作线程
|
||||
self.workers: List[asyncio.Task] = []
|
||||
self.running = False
|
||||
|
||||
# 统计信息
|
||||
self.stats = {
|
||||
'total_tasks': 0,
|
||||
'completed_tasks': 0,
|
||||
'failed_tasks': 0,
|
||||
'retried_tasks': 0,
|
||||
'average_duration': 0.0,
|
||||
'success_rate': 0.0
|
||||
}
|
||||
|
||||
# 初始化默认策略
|
||||
self._init_default_strategies()
|
||||
|
||||
def _init_default_strategies(self):
|
||||
"""初始化默认策略"""
|
||||
# API策略
|
||||
api_strategy = EnhancedAPIStrategy()
|
||||
|
||||
# 如果启用重试，包装策略
|
||||
if self.config.enable_retry:
|
||||
api_strategy = RetryStrategy(api_strategy)
|
||||
|
||||
self.register_strategy(api_strategy)
|
||||
|
||||
def register_strategy(self, strategy: IDownloadStrategy):
|
||||
"""
|
||||
注册下载策略
|
||||
|
||||
Args:
|
||||
strategy: 下载策略实例
|
||||
"""
|
||||
self.strategies.append(strategy)
|
||||
# 按优先级排序
|
||||
self.strategies.sort(key=lambda s: s.get_priority(), reverse=True)
|
||||
logger.info(f"注册策略: {strategy.name} (优先级: {strategy.get_priority()})")
|
||||
|
||||
async def add_task(self, url: str, task_type: Optional[TaskType] = None, priority: int = 0) -> str:
|
||||
"""
|
||||
添加下载任务
|
||||
|
||||
Args:
|
||||
url: 下载URL
|
||||
task_type: 任务类型
|
||||
priority: 优先级（数值越大优先级越高）
|
||||
|
||||
Returns:
|
||||
任务ID
|
||||
"""
|
||||
# 自动识别任务类型
|
||||
if task_type is None:
|
||||
task_type = self._detect_task_type(url)
|
||||
|
||||
# 创建任务
|
||||
task = DownloadTask(
|
||||
task_id=str(uuid.uuid4()),
|
||||
url=url,
|
||||
task_type=task_type,
|
||||
priority=priority
|
||||
)
|
||||
|
||||
# 添加到队列
|
||||
if self.config.priority_queue and priority > 0:
|
||||
self.priority_tasks.append(task)
|
||||
self.priority_tasks.sort(key=lambda t: t.priority, reverse=True)
|
||||
else:
|
||||
await self.pending_queue.put(task)
|
||||
|
||||
self.stats['total_tasks'] += 1
|
||||
logger.info(f"添加任务: {task.task_id} ({task_type.value}) 优先级: {priority}")
|
||||
|
||||
return task.task_id
|
||||
|
||||
async def add_batch(self, urls: List[str], task_type: Optional[TaskType] = None) -> List[str]:
|
||||
"""
|
||||
批量添加任务
|
||||
|
||||
Args:
|
||||
urls: URL列表
|
||||
task_type: 任务类型
|
||||
|
||||
Returns:
|
||||
任务ID列表
|
||||
"""
|
||||
task_ids = []
|
||||
for i, url in enumerate(urls):
|
||||
# 批量任务使用递减优先级
|
||||
priority = len(urls) - i
|
||||
task_id = await self.add_task(url, task_type, priority)
|
||||
task_ids.append(task_id)
|
||||
|
||||
return task_ids
|
||||
|
||||
async def start(self):
|
||||
"""启动编排器"""
|
||||
if self.running:
|
||||
logger.warning("编排器已在运行")
|
||||
return
|
||||
|
||||
self.running = True
|
||||
logger.info(f"启动编排器，最大并发数: {self.config.max_concurrent}")
|
||||
|
||||
# 创建工作线程
|
||||
for i in range(self.config.max_concurrent):
|
||||
worker = asyncio.create_task(self._worker(i))
|
||||
self.workers.append(worker)
|
||||
|
||||
async def stop(self):
|
||||
"""停止编排器"""
|
||||
if not self.running:
|
||||
return
|
||||
|
||||
logger.info("停止编排器...")
|
||||
self.running = False
|
||||
|
||||
# 取消所有工作线程
|
||||
for worker in self.workers:
|
||||
worker.cancel()
|
||||
|
||||
# 等待工作线程结束
|
||||
await asyncio.gather(*self.workers, return_exceptions=True)
|
||||
self.workers.clear()
|
||||
|
||||
logger.info("编排器已停止")
|
||||
|
||||
async def wait_completion(self, timeout: Optional[float] = None):
|
||||
"""
|
||||
等待所有任务完成
|
||||
|
||||
Args:
|
||||
timeout: 超时时间（秒）
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
while self.running:
|
||||
# 检查是否所有任务都完成
|
||||
if (self.pending_queue.empty() and
|
||||
not self.priority_tasks and
|
||||
not self.active_tasks):
|
||||
logger.info("所有任务已完成")
|
||||
break
|
||||
|
||||
# 检查超时
|
||||
if timeout and (time.time() - start_time) > timeout:
|
||||
logger.warning(f"等待超时 ({timeout} 秒)")
|
||||
break
|
||||
|
||||
await asyncio.sleep(1)
|
||||
|
||||
# 计算统计信息
|
||||
self._calculate_stats()
|
||||
|
||||
async def _worker(self, worker_id: int):
|
||||
"""
|
||||
工作线程
|
||||
|
||||
Args:
|
||||
worker_id: 工作线程ID
|
||||
"""
|
||||
logger.info(f"工作线程 {worker_id} 启动")
|
||||
|
||||
while self.running:
|
||||
try:
|
||||
# 获取任务
|
||||
task = await self._get_next_task()
|
||||
if task is None:
|
||||
await asyncio.sleep(0.1)
|
||||
continue
|
||||
|
||||
# 标记为活动任务
|
||||
self.active_tasks[task.task_id] = task
|
||||
|
||||
# 限速控制
|
||||
if self.rate_limiter:
|
||||
await self.rate_limiter.acquire()
|
||||
|
||||
# 执行任务
|
||||
logger.info(f"工作线程 {worker_id} 开始处理任务: {task.task_id}")
|
||||
result = await self._execute_task(task)
|
||||
|
||||
# 移除活动任务
|
||||
del self.active_tasks[task.task_id]
|
||||
|
||||
# 处理结果
|
||||
if result.success:
|
||||
self.completed_tasks.append(task)
|
||||
self.stats['completed_tasks'] += 1
|
||||
logger.info(f"任务 {task.task_id} 完成")
|
||||
else:
|
||||
# 检查是否需要重试
|
||||
if task.increment_retry():
|
||||
logger.warning(f"任务 {task.task_id} 失败，准备重试 ({task.retry_count}/{task.max_retries})")
|
||||
await self.pending_queue.put(task)
|
||||
self.stats['retried_tasks'] += 1
|
||||
else:
|
||||
self.failed_tasks.append(task)
|
||||
self.stats['failed_tasks'] += 1
|
||||
logger.error(f"任务 {task.task_id} 最终失败: {result.error_message}")
|
||||
|
||||
# 保存进度
|
||||
if self.config.save_progress:
|
||||
await self._save_progress()
|
||||
|
||||
except asyncio.CancelledError:
|
||||
logger.info(f"工作线程 {worker_id} 被取消")
|
||||
break
|
||||
except Exception as e:
|
||||
logger.error(f"工作线程 {worker_id} 异常: {e}")
|
||||
await asyncio.sleep(1)
|
||||
|
||||
logger.info(f"工作线程 {worker_id} 结束")
|
||||
|
||||
async def _get_next_task(self) -> Optional[DownloadTask]:
|
||||
"""获取下一个任务"""
|
||||
# 优先处理高优先级任务
|
||||
if self.priority_tasks:
|
||||
return self.priority_tasks.pop(0)
|
||||
|
||||
# 从普通队列获取
|
||||
try:
|
||||
return await asyncio.wait_for(
|
||||
self.pending_queue.get(),
|
||||
timeout=0.1
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
return None
|
||||
|
||||
async def _execute_task(self, task: DownloadTask) -> DownloadResult:
|
||||
"""
|
||||
执行任务，尝试所有策略
|
||||
|
||||
Args:
|
||||
task: 下载任务
|
||||
|
||||
Returns:
|
||||
下载结果
|
||||
"""
|
||||
last_error = None
|
||||
|
||||
for strategy in self.strategies:
|
||||
try:
|
||||
# 检查策略是否能处理任务
|
||||
if not await strategy.can_handle(task):
|
||||
continue
|
||||
|
||||
logger.info(f"使用策略 {strategy.name} 处理任务 {task.task_id}")
|
||||
|
||||
# 执行下载
|
||||
result = await strategy.download(task)
|
||||
|
||||
if result.success:
|
||||
return result
|
||||
|
||||
last_error = result.error_message
|
||||
logger.warning(f"策略 {strategy.name} 失败: {last_error}")
|
||||
|
||||
except Exception as e:
|
||||
last_error = str(e)
|
||||
logger.error(f"策略 {strategy.name} 异常: {e}")
|
||||
|
||||
# 所有策略都失败
|
||||
return DownloadResult(
|
||||
success=False,
|
||||
task_id=task.task_id,
|
||||
error_message=f"所有策略都失败: {last_error}",
|
||||
retry_count=task.retry_count
|
||||
)
|
||||
|
||||
def _detect_task_type(self, url: str) -> TaskType:
|
||||
"""
|
||||
自动检测任务类型
|
||||
|
||||
Args:
|
||||
url: URL
|
||||
|
||||
Returns:
|
||||
任务类型
|
||||
"""
|
||||
url_lower = url.lower()
|
||||
|
||||
if '/user/' in url_lower:
|
||||
return TaskType.USER
|
||||
elif '/video/' in url_lower or '/note/' in url_lower:
|
||||
return TaskType.VIDEO
|
||||
elif '/music/' in url_lower:
|
||||
return TaskType.MUSIC
|
||||
elif '/mix/' in url_lower or '/collection/' in url_lower:
|
||||
return TaskType.MIX
|
||||
elif 'live.douyin.com' in url_lower:
|
||||
return TaskType.LIVE
|
||||
else:
|
||||
return TaskType.VIDEO # 默认为视频
|
||||
|
||||
def _calculate_stats(self):
|
||||
"""计算统计信息"""
|
||||
total = self.stats['total_tasks']
|
||||
if total > 0:
|
||||
self.stats['success_rate'] = self.stats['completed_tasks'] / total * 100
|
||||
|
||||
# 计算平均时长
|
||||
durations = []
|
||||
for task in self.completed_tasks:
|
||||
if hasattr(task, 'duration'):
|
||||
durations.append(task.duration)
|
||||
|
||||
if durations:
|
||||
self.stats['average_duration'] = sum(durations) / len(durations)
|
||||
|
||||
async def _save_progress(self):
|
||||
"""保存进度（可扩展为持久化到文件或数据库）"""
|
||||
# TODO: 实现进度保存逻辑
|
||||
pass
|
||||
|
||||
def get_stats(self) -> Dict[str, Any]:
|
||||
"""获取统计信息"""
|
||||
self._calculate_stats()
|
||||
return self.stats.copy()
|
||||
|
||||
def get_task_status(self, task_id: str) -> Optional[TaskStatus]:
|
||||
"""
|
||||
获取任务状态
|
||||
|
||||
Args:
|
||||
task_id: 任务ID
|
||||
|
||||
Returns:
|
||||
任务状态
|
||||
"""
|
||||
# 检查活动任务
|
||||
if task_id in self.active_tasks:
|
||||
return self.active_tasks[task_id].status
|
||||
|
||||
# 检查完成任务
|
||||
for task in self.completed_tasks:
|
||||
if task.task_id == task_id:
|
||||
return TaskStatus.COMPLETED
|
||||
|
||||
# 检查失败任务
|
||||
for task in self.failed_tasks:
|
||||
if task.task_id == task_id:
|
||||
return TaskStatus.FAILED
|
||||
|
||||
# 检查待处理任务
|
||||
for task in self.priority_tasks:
|
||||
if task.task_id == task_id:
|
||||
return TaskStatus.PENDING
|
||||
|
||||
return None
|
||||
503
DouYin/apiproxy/douyin/core/progress_tracker.py
Normal file
503
DouYin/apiproxy/douyin/core/progress_tracker.py
Normal file
@@ -0,0 +1,503 @@
|
||||
#!/usr/bin/env python
|
||||
# -*- coding: utf-8 -*-
|
||||
|
||||
"""
|
||||
实时进度跟踪系统
|
||||
支持WebSocket推送和进度监控
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import time
|
||||
import logging
|
||||
from typing import Dict, List, Optional, Callable, Any
|
||||
from dataclasses import dataclass, field, asdict
|
||||
from enum import Enum
|
||||
from datetime import datetime
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# 动态导入WebSocket支持
|
||||
try:
|
||||
import websockets
|
||||
from websockets.server import WebSocketServerProtocol
|
||||
WEBSOCKET_AVAILABLE = True
|
||||
except ImportError:
|
||||
WEBSOCKET_AVAILABLE = False
|
||||
logger.warning("websockets未安装，WebSocket功能不可用")
|
||||
|
||||
|
||||
class EventType(Enum):
|
||||
"""事件类型"""
|
||||
TASK_ADDED = "task_added"
|
||||
TASK_STARTED = "task_started"
|
||||
TASK_PROGRESS = "task_progress"
|
||||
TASK_COMPLETED = "task_completed"
|
||||
TASK_FAILED = "task_failed"
|
||||
TASK_RETRYING = "task_retrying"
|
||||
STATS_UPDATE = "stats_update"
|
||||
SPEED_UPDATE = "speed_update"
|
||||
ERROR = "error"
|
||||
INFO = "info"
|
||||
|
||||
|
||||
@dataclass
|
||||
class ProgressEvent:
|
||||
"""进度事件"""
|
||||
event_type: EventType
|
||||
task_id: Optional[str] = None
|
||||
data: Dict[str, Any] = field(default_factory=dict)
|
||||
timestamp: float = field(default_factory=time.time)
|
||||
|
||||
def to_dict(self) -> Dict:
|
||||
"""转换为字典"""
|
||||
return {
|
||||
'event_type': self.event_type.value,
|
||||
'task_id': self.task_id,
|
||||
'data': self.data,
|
||||
'timestamp': self.timestamp
|
||||
}
|
||||
|
||||
def to_json(self) -> str:
|
||||
"""转换为JSON"""
|
||||
return json.dumps(self.to_dict())
|
||||
|
||||
|
||||
@dataclass
|
||||
class TaskProgress:
|
||||
"""任务进度信息"""
|
||||
task_id: str
|
||||
url: str
|
||||
status: str
|
||||
progress: float = 0.0 # 0-100
|
||||
downloaded_bytes: int = 0
|
||||
total_bytes: int = 0
|
||||
speed: float = 0.0 # bytes/second
|
||||
eta: float = 0.0 # seconds
|
||||
start_time: float = field(default_factory=time.time)
|
||||
end_time: Optional[float] = None
|
||||
error_message: Optional[str] = None
|
||||
|
||||
def get_duration(self) -> float:
|
||||
"""获取耗时"""
|
||||
if self.end_time:
|
||||
return self.end_time - self.start_time
|
||||
return time.time() - self.start_time
|
||||
|
||||
def update_progress(self, downloaded: int, total: int):
|
||||
"""更新进度"""
|
||||
self.downloaded_bytes = downloaded
|
||||
self.total_bytes = total
|
||||
|
||||
if total > 0:
|
||||
self.progress = (downloaded / total) * 100
|
||||
|
||||
# 计算速度
|
||||
duration = self.get_duration()
|
||||
if duration > 0:
|
||||
self.speed = downloaded / duration
|
||||
|
||||
# 计算剩余时间
|
||||
if self.speed > 0 and total > downloaded:
|
||||
self.eta = (total - downloaded) / self.speed
|
||||
|
||||
def to_dict(self) -> Dict:
|
||||
"""转换为字典"""
|
||||
return {
|
||||
'task_id': self.task_id,
|
||||
'url': self.url,
|
||||
'status': self.status,
|
||||
'progress': round(self.progress, 2),
|
||||
'downloaded_bytes': self.downloaded_bytes,
|
||||
'total_bytes': self.total_bytes,
|
||||
'speed': round(self.speed, 2),
|
||||
'eta': round(self.eta, 2),
|
||||
'duration': round(self.get_duration(), 2),
|
||||
'error_message': self.error_message
|
||||
}
|
||||
|
||||
|
||||
class ProgressTracker:
|
||||
"""进度跟踪器"""
|
||||
|
||||
def __init__(self, enable_websocket: bool = True, ws_port: int = 8765):
|
||||
"""
|
||||
初始化进度跟踪器
|
||||
|
||||
Args:
|
||||
enable_websocket: 是否启用WebSocket
|
||||
ws_port: WebSocket端口
|
||||
"""
|
||||
self.enable_websocket = enable_websocket and WEBSOCKET_AVAILABLE
|
||||
self.ws_port = ws_port
|
||||
|
||||
# 任务进度
|
||||
self.tasks: Dict[str, TaskProgress] = {}
|
||||
|
||||
# 事件监听器
|
||||
self.listeners: List[Callable[[ProgressEvent], None]] = []
|
||||
|
||||
# WebSocket连接
|
||||
self.websocket_clients: List[WebSocketServerProtocol] = []
|
||||
self.websocket_server = None
|
||||
self.ws_task = None
|
||||
|
||||
# 统计信息
|
||||
self.stats = {
|
||||
'total_tasks': 0,
|
||||
'active_tasks': 0,
|
||||
'completed_tasks': 0,
|
||||
'failed_tasks': 0,
|
||||
'total_downloaded': 0,
|
||||
'average_speed': 0.0,
|
||||
'success_rate': 0.0
|
||||
}
|
||||
|
||||
# 速度历史（用于计算平均速度）
|
||||
self.speed_history = []
|
||||
self.max_speed_history = 100
|
||||
|
||||
def add_listener(self, listener: Callable[[ProgressEvent], None]):
|
||||
"""添加事件监听器"""
|
||||
self.listeners.append(listener)
|
||||
logger.debug(f"添加事件监听器: {listener}")
|
||||
|
||||
def remove_listener(self, listener: Callable[[ProgressEvent], None]):
|
||||
"""移除事件监听器"""
|
||||
if listener in self.listeners:
|
||||
self.listeners.remove(listener)
|
||||
|
||||
async def emit_event(self, event: ProgressEvent):
|
||||
"""触发事件"""
|
||||
# 通知监听器
|
||||
for listener in self.listeners:
|
||||
try:
|
||||
if asyncio.iscoroutinefunction(listener):
|
||||
await listener(event)
|
||||
else:
|
||||
listener(event)
|
||||
except Exception as e:
|
||||
logger.error(f"事件监听器执行失败: {e}")
|
||||
|
||||
# 推送到WebSocket客户端
|
||||
if self.websocket_clients:
|
||||
await self._broadcast_websocket(event.to_json())
|
||||
|
||||
async def add_task(self, task_id: str, url: str):
|
||||
"""添加任务"""
|
||||
self.tasks[task_id] = TaskProgress(
|
||||
task_id=task_id,
|
||||
url=url,
|
||||
status="pending"
|
||||
)
|
||||
|
||||
self.stats['total_tasks'] += 1
|
||||
|
||||
await self.emit_event(ProgressEvent(
|
||||
event_type=EventType.TASK_ADDED,
|
||||
task_id=task_id,
|
||||
data={'url': url}
|
||||
))
|
||||
|
||||
async def start_task(self, task_id: str):
|
||||
"""开始任务"""
|
||||
if task_id in self.tasks:
|
||||
self.tasks[task_id].status = "processing"
|
||||
self.tasks[task_id].start_time = time.time()
|
||||
|
||||
self.stats['active_tasks'] += 1
|
||||
|
||||
await self.emit_event(ProgressEvent(
|
||||
event_type=EventType.TASK_STARTED,
|
||||
task_id=task_id
|
||||
))
|
||||
|
||||
async def update_progress(
|
||||
self,
|
||||
task_id: str,
|
||||
downloaded: int,
|
||||
total: int,
|
||||
extra_data: Optional[Dict] = None
|
||||
):
|
||||
"""更新任务进度"""
|
||||
if task_id not in self.tasks:
|
||||
return
|
||||
|
||||
task = self.tasks[task_id]
|
||||
task.update_progress(downloaded, total)
|
||||
|
||||
# 更新速度历史
|
||||
if task.speed > 0:
|
||||
self.speed_history.append(task.speed)
|
||||
if len(self.speed_history) > self.max_speed_history:
|
||||
self.speed_history.pop(0)
|
||||
|
||||
# 计算平均速度
|
||||
self.stats['average_speed'] = sum(self.speed_history) / len(self.speed_history)
|
||||
|
||||
event_data = task.to_dict()
|
||||
if extra_data:
|
||||
event_data.update(extra_data)
|
||||
|
||||
await self.emit_event(ProgressEvent(
|
||||
event_type=EventType.TASK_PROGRESS,
|
||||
task_id=task_id,
|
||||
data=event_data
|
||||
))
|
||||
|
||||
async def complete_task(self, task_id: str, success: bool = True, error: Optional[str] = None):
|
||||
"""完成任务"""
|
||||
if task_id not in self.tasks:
|
||||
return
|
||||
|
||||
task = self.tasks[task_id]
|
||||
task.end_time = time.time()
|
||||
|
||||
if success:
|
||||
task.status = "completed"
|
||||
task.progress = 100.0
|
||||
self.stats['completed_tasks'] += 1
|
||||
event_type = EventType.TASK_COMPLETED
|
||||
else:
|
||||
task.status = "failed"
|
||||
task.error_message = error
|
||||
self.stats['failed_tasks'] += 1
|
||||
event_type = EventType.TASK_FAILED
|
||||
|
||||
self.stats['active_tasks'] = max(0, self.stats['active_tasks'] - 1)
|
||||
|
||||
# 更新成功率
|
||||
total_finished = self.stats['completed_tasks'] + self.stats['failed_tasks']
|
||||
if total_finished > 0:
|
||||
self.stats['success_rate'] = (self.stats['completed_tasks'] / total_finished) * 100
|
||||
|
||||
await self.emit_event(ProgressEvent(
|
||||
event_type=event_type,
|
||||
task_id=task_id,
|
||||
data=task.to_dict()
|
||||
))
|
||||
|
||||
async def retry_task(self, task_id: str, retry_count: int):
|
||||
"""重试任务"""
|
||||
if task_id in self.tasks:
|
||||
self.tasks[task_id].status = "retrying"
|
||||
|
||||
await self.emit_event(ProgressEvent(
|
||||
event_type=EventType.TASK_RETRYING,
|
||||
task_id=task_id,
|
||||
data={'retry_count': retry_count}
|
||||
))
|
||||
|
||||
async def update_stats(self):
|
||||
"""更新统计信息"""
|
||||
await self.emit_event(ProgressEvent(
|
||||
event_type=EventType.STATS_UPDATE,
|
||||
data=self.stats.copy()
|
||||
))
|
||||
|
||||
def get_task_progress(self, task_id: str) -> Optional[TaskProgress]:
|
||||
"""获取任务进度"""
|
||||
return self.tasks.get(task_id)
|
||||
|
||||
def get_active_tasks(self) -> List[TaskProgress]:
|
||||
"""获取活动任务"""
|
||||
return [
|
||||
task for task in self.tasks.values()
|
||||
if task.status in ["processing", "retrying"]
|
||||
]
|
||||
|
||||
def get_statistics(self) -> Dict[str, Any]:
|
||||
"""获取统计信息"""
|
||||
return self.stats.copy()
|
||||
|
||||
def clear_completed_tasks(self):
|
||||
"""清理已完成的任务"""
|
||||
completed_ids = [
|
||||
task_id for task_id, task in self.tasks.items()
|
||||
if task.status in ["completed", "failed"]
|
||||
]
|
||||
|
||||
for task_id in completed_ids:
|
||||
del self.tasks[task_id]
|
||||
|
||||
logger.info(f"清理了 {len(completed_ids)} 个已完成任务")
|
||||
|
||||
# WebSocket相关功能
|
||||
async def start_websocket_server(self):
|
||||
"""启动WebSocket服务器"""
|
||||
if not self.enable_websocket:
|
||||
logger.info("WebSocket功能未启用")
|
||||
return
|
||||
|
||||
try:
|
||||
self.websocket_server = await websockets.serve(
|
||||
self._handle_websocket,
|
||||
"localhost",
|
||||
self.ws_port
|
||||
)
|
||||
|
||||
logger.info(f"WebSocket服务器启动在 ws://localhost:{self.ws_port}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"启动WebSocket服务器失败: {e}")
|
||||
|
||||
async def stop_websocket_server(self):
|
||||
"""停止WebSocket服务器"""
|
||||
if self.websocket_server:
|
||||
self.websocket_server.close()
|
||||
await self.websocket_server.wait_closed()
|
||||
self.websocket_server = None
|
||||
logger.info("WebSocket服务器已停止")
|
||||
|
||||
async def _handle_websocket(self, websocket: WebSocketServerProtocol, path: str):
|
||||
"""处理WebSocket连接"""
|
||||
logger.info(f"新的WebSocket连接: {websocket.remote_address}")
|
||||
|
||||
# 添加客户端
|
||||
self.websocket_clients.append(websocket)
|
||||
|
||||
try:
|
||||
# 发送当前状态
|
||||
await websocket.send(json.dumps({
|
||||
'type': 'init',
|
||||
'data': {
|
||||
'tasks': {
|
||||
task_id: task.to_dict()
|
||||
for task_id, task in self.tasks.items()
|
||||
},
|
||||
'stats': self.stats
|
||||
}
|
||||
}))
|
||||
|
||||
# 保持连接
|
||||
async for message in websocket:
|
||||
try:
|
||||
data = json.loads(message)
|
||||
await self._handle_ws_message(websocket, data)
|
||||
except json.JSONDecodeError:
|
||||
logger.warning(f"无效的WebSocket消息: {message}")
|
||||
|
||||
except websockets.exceptions.ConnectionClosed:
|
||||
logger.info(f"WebSocket连接关闭: {websocket.remote_address}")
|
||||
finally:
|
||||
# 移除客户端
|
||||
if websocket in self.websocket_clients:
|
||||
self.websocket_clients.remove(websocket)
|
||||
|
||||
async def _handle_ws_message(self, websocket: WebSocketServerProtocol, data: Dict):
|
||||
"""处理WebSocket消息"""
|
||||
msg_type = data.get('type')
|
||||
|
||||
if msg_type == 'ping':
|
||||
await websocket.send(json.dumps({'type': 'pong'}))
|
||||
elif msg_type == 'get_stats':
|
||||
await websocket.send(json.dumps({
|
||||
'type': 'stats',
|
||||
'data': self.stats
|
||||
}))
|
||||
elif msg_type == 'get_tasks':
|
||||
await websocket.send(json.dumps({
|
||||
'type': 'tasks',
|
||||
'data': {
|
||||
task_id: task.to_dict()
|
||||
for task_id, task in self.tasks.items()
|
||||
}
|
||||
}))
|
||||
|
||||
async def _broadcast_websocket(self, message: str):
|
||||
"""广播WebSocket消息"""
|
||||
if not self.websocket_clients:
|
||||
return
|
||||
|
||||
# 并发发送给所有客户端
|
||||
disconnected = []
|
||||
|
||||
for client in self.websocket_clients:
|
||||
try:
|
||||
await client.send(message)
|
||||
except websockets.exceptions.ConnectionClosed:
|
||||
disconnected.append(client)
|
||||
|
||||
# 清理断开的连接
|
||||
for client in disconnected:
|
||||
self.websocket_clients.remove(client)
|
||||
|
||||
# 上下文管理器
|
||||
async def __aenter__(self):
|
||||
"""异步上下文管理器入口"""
|
||||
await self.start_websocket_server()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
||||
"""异步上下文管理器出口"""
|
||||
await self.stop_websocket_server()
|
||||
|
||||
|
||||
def create_console_listener(use_rich: bool = True):
|
||||
"""
|
||||
创建控制台进度监听器
|
||||
|
||||
Args:
|
||||
use_rich: 是否使用Rich库美化输出
|
||||
|
||||
Returns:
|
||||
监听器函数
|
||||
"""
|
||||
if use_rich:
|
||||
try:
|
||||
from rich.console import Console
|
||||
from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn
|
||||
from rich.table import Table
|
||||
|
||||
console = Console()
|
||||
progress = Progress(
|
||||
SpinnerColumn(),
|
||||
TextColumn("[progress.description]{task.description}"),
|
||||
BarColumn(),
|
||||
TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
|
||||
)
|
||||
|
||||
task_map = {}
|
||||
|
||||
def rich_listener(event: ProgressEvent):
|
||||
if event.event_type == EventType.TASK_ADDED:
|
||||
task_id = event.task_id
|
||||
if task_id not in task_map:
|
||||
task_map[task_id] = progress.add_task(
|
||||
f"[cyan]{event.data.get('url', 'Unknown')}",
|
||||
total=100
|
||||
)
|
||||
|
||||
elif event.event_type == EventType.TASK_PROGRESS:
|
||||
task_id = event.task_id
|
||||
if task_id in task_map:
|
||||
progress.update(
|
||||
task_map[task_id],
|
||||
completed=event.data.get('progress', 0)
|
||||
)
|
||||
|
||||
elif event.event_type == EventType.TASK_COMPLETED:
|
||||
console.print(f"[green]✓ 任务完成: {event.task_id}")
|
||||
|
||||
elif event.event_type == EventType.TASK_FAILED:
|
||||
console.print(f"[red]✗ 任务失败: {event.task_id} - {event.data.get('error_message', 'Unknown error')}")
|
||||
|
||||
return rich_listener
|
||||
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
# 默认简单监听器
|
||||
def simple_listener(event: ProgressEvent):
|
||||
timestamp = datetime.fromtimestamp(event.timestamp).strftime('%H:%M:%S')
|
||||
|
||||
if event.event_type == EventType.TASK_PROGRESS:
|
||||
progress = event.data.get('progress', 0)
|
||||
speed = event.data.get('speed', 0)
|
||||
print(f"[{timestamp}] 任务 {event.task_id}: {progress:.1f}% ({speed/1024/1024:.2f} MB/s)")
|
||||
elif event.event_type == EventType.TASK_COMPLETED:
|
||||
print(f"[{timestamp}] ✓ 任务完成: {event.task_id}")
|
||||
elif event.event_type == EventType.TASK_FAILED:
|
||||
print(f"[{timestamp}] ✗ 任务失败: {event.task_id}")
|
||||
|
||||
return simple_listener
|
||||
494
DouYin/apiproxy/douyin/core/queue_manager.py
Normal file
494
DouYin/apiproxy/douyin/core/queue_manager.py
Normal file
@@ -0,0 +1,494 @@
|
||||
#!/usr/bin/env python
|
||||
# -*- coding: utf-8 -*-
|
||||
|
||||
"""
|
||||
队列管理器
|
||||
支持任务持久化和断点恢复
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import sqlite3
|
||||
import time
|
||||
import logging
|
||||
import pickle
|
||||
from typing import List, Dict, Optional, Any
|
||||
from pathlib import Path
|
||||
from dataclasses import dataclass, asdict
|
||||
from enum import Enum
|
||||
|
||||
from apiproxy.douyin.strategies.base import DownloadTask, TaskStatus, TaskType
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class PersistentQueue:
|
||||
"""持久化队列管理器"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
db_path: str = "download_queue.db",
|
||||
max_size: int = 10000,
|
||||
checkpoint_interval: int = 60
|
||||
):
|
||||
"""
|
||||
初始化队列管理器
|
||||
|
||||
Args:
|
||||
db_path: 数据库文件路径
|
||||
max_size: 队列最大容量
|
||||
checkpoint_interval: 检查点保存间隔（秒）
|
||||
"""
|
||||
self.db_path = Path(db_path)
|
||||
self.max_size = max_size
|
||||
self.checkpoint_interval = checkpoint_interval
|
||||
|
||||
self.conn: Optional[sqlite3.Connection] = None
|
||||
self.queue = asyncio.Queue(maxsize=max_size)
|
||||
self._checkpoint_task = None
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
# 初始化数据库
|
||||
self._init_database()
|
||||
|
||||
# 恢复未完成的任务
|
||||
self._restore_tasks()
|
||||
|
||||
def _init_database(self):
|
||||
"""初始化数据库"""
|
||||
self.conn = sqlite3.connect(str(self.db_path), check_same_thread=False)
|
||||
cursor = self.conn.cursor()
|
||||
|
||||
# 创建任务表
|
||||
cursor.execute('''
|
||||
CREATE TABLE IF NOT EXISTS tasks (
|
||||
task_id TEXT PRIMARY KEY,
|
||||
url TEXT NOT NULL,
|
||||
task_type TEXT NOT NULL,
|
||||
priority INTEGER DEFAULT 0,
|
||||
status TEXT NOT NULL,
|
||||
retry_count INTEGER DEFAULT 0,
|
||||
max_retries INTEGER DEFAULT 3,
|
||||
metadata TEXT,
|
||||
created_at REAL NOT NULL,
|
||||
updated_at REAL NOT NULL,
|
||||
completed_at REAL,
|
||||
error_message TEXT,
|
||||
result TEXT
|
||||
)
|
||||
''')
|
||||
|
||||
# 创建索引
|
||||
cursor.execute('CREATE INDEX IF NOT EXISTS idx_status ON tasks(status)')
|
||||
cursor.execute('CREATE INDEX IF NOT EXISTS idx_priority ON tasks(priority DESC)')
|
||||
cursor.execute('CREATE INDEX IF NOT EXISTS idx_created_at ON tasks(created_at)')
|
||||
|
||||
# 创建进度表
|
||||
cursor.execute('''
|
||||
CREATE TABLE IF NOT EXISTS progress (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
timestamp REAL NOT NULL,
|
||||
total_tasks INTEGER,
|
||||
pending_tasks INTEGER,
|
||||
active_tasks INTEGER,
|
||||
completed_tasks INTEGER,
|
||||
failed_tasks INTEGER,
|
||||
success_rate REAL,
|
||||
average_duration REAL
|
||||
)
|
||||
''')
|
||||
|
||||
self.conn.commit()
|
||||
logger.info(f"数据库初始化完成: {self.db_path}")
|
||||
|
||||
def _restore_tasks(self):
|
||||
"""从数据库恢复未完成的任务"""
|
||||
cursor = self.conn.cursor()
|
||||
|
||||
# 将所有PROCESSING状态的任务重置为PENDING
|
||||
cursor.execute('''
|
||||
UPDATE tasks
|
||||
SET status = ?, updated_at = ?
|
||||
WHERE status = ?
|
||||
''', (TaskStatus.PENDING.value, time.time(), TaskStatus.PROCESSING.value))
|
||||
|
||||
# 获取所有待处理的任务
|
||||
cursor.execute('''
|
||||
SELECT task_id, url, task_type, priority, retry_count, max_retries, metadata, created_at
|
||||
FROM tasks
|
||||
WHERE status IN (?, ?)
|
||||
ORDER BY priority DESC, created_at ASC
|
||||
''', (TaskStatus.PENDING.value, TaskStatus.RETRYING.value))
|
||||
|
||||
restored_count = 0
|
||||
for row in cursor.fetchall():
|
||||
task = self._row_to_task(row)
|
||||
if task:
|
||||
try:
|
||||
self.queue.put_nowait(task)
|
||||
restored_count += 1
|
||||
except asyncio.QueueFull:
|
||||
break
|
||||
|
||||
self.conn.commit()
|
||||
|
||||
if restored_count > 0:
|
||||
logger.info(f"从数据库恢复了 {restored_count} 个未完成任务")
|
||||
|
||||
def _row_to_task(self, row: tuple) -> Optional[DownloadTask]:
|
||||
"""将数据库行转换为任务对象"""
|
||||
try:
|
||||
task_id, url, task_type, priority, retry_count, max_retries, metadata_str, created_at = row
|
||||
|
||||
metadata = {}
|
||||
if metadata_str:
|
||||
try:
|
||||
metadata = json.loads(metadata_str)
|
||||
except:
|
||||
pass
|
||||
|
||||
return DownloadTask(
|
||||
task_id=task_id,
|
||||
url=url,
|
||||
task_type=TaskType(task_type),
|
||||
priority=priority,
|
||||
retry_count=retry_count,
|
||||
max_retries=max_retries,
|
||||
metadata=metadata,
|
||||
created_at=created_at
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"转换任务失败: {e}")
|
||||
return None
|
||||
|
||||
async def add_task(self, task: DownloadTask) -> bool:
|
||||
"""
|
||||
添加任务到队列
|
||||
|
||||
Args:
|
||||
task: 下载任务
|
||||
|
||||
Returns:
|
||||
是否成功添加
|
||||
"""
|
||||
async with self._lock:
|
||||
try:
|
||||
# 保存到数据库
|
||||
cursor = self.conn.cursor()
|
||||
cursor.execute('''
|
||||
INSERT OR REPLACE INTO tasks (
|
||||
task_id, url, task_type, priority, status,
|
||||
retry_count, max_retries, metadata,
|
||||
created_at, updated_at
|
||||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
''', (
|
||||
task.task_id,
|
||||
task.url,
|
||||
task.task_type.value,
|
||||
task.priority,
|
||||
task.status.value,
|
||||
task.retry_count,
|
||||
task.max_retries,
|
||||
json.dumps(task.metadata),
|
||||
task.created_at,
|
||||
task.updated_at
|
||||
))
|
||||
self.conn.commit()
|
||||
|
||||
# 添加到内存队列
|
||||
await self.queue.put(task)
|
||||
|
||||
logger.debug(f"任务 {task.task_id} 已添加到队列")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"添加任务失败: {e}")
|
||||
return False
|
||||
|
||||
async def get_task(self, timeout: float = 1.0) -> Optional[DownloadTask]:
|
||||
"""
|
||||
从队列获取任务
|
||||
|
||||
Args:
|
||||
timeout: 超时时间
|
||||
|
||||
Returns:
|
||||
下载任务
|
||||
"""
|
||||
try:
|
||||
task = await asyncio.wait_for(self.queue.get(), timeout=timeout)
|
||||
|
||||
# 更新数据库状态
|
||||
await self.update_task_status(task.task_id, TaskStatus.PROCESSING)
|
||||
|
||||
return task
|
||||
|
||||
except asyncio.TimeoutError:
|
||||
return None
|
||||
|
||||
async def update_task_status(
|
||||
self,
|
||||
task_id: str,
|
||||
status: TaskStatus,
|
||||
error_message: Optional[str] = None,
|
||||
result: Optional[Dict] = None
|
||||
):
|
||||
"""
|
||||
更新任务状态
|
||||
|
||||
Args:
|
||||
task_id: 任务ID
|
||||
status: 新状态
|
||||
error_message: 错误信息
|
||||
result: 执行结果
|
||||
"""
|
||||
async with self._lock:
|
||||
cursor = self.conn.cursor()
|
||||
|
||||
update_fields = {
|
||||
'status': status.value,
|
||||
'updated_at': time.time()
|
||||
}
|
||||
|
||||
if error_message:
|
||||
update_fields['error_message'] = error_message
|
||||
|
||||
if result:
|
||||
update_fields['result'] = json.dumps(result)
|
||||
|
||||
if status == TaskStatus.COMPLETED:
|
||||
update_fields['completed_at'] = time.time()
|
||||
|
||||
# 构建UPDATE语句
|
||||
set_clause = ', '.join([f'{k} = ?' for k in update_fields.keys()])
|
||||
values = list(update_fields.values()) + [task_id]
|
||||
|
||||
cursor.execute(
|
||||
f'UPDATE tasks SET {set_clause} WHERE task_id = ?',
|
||||
values
|
||||
)
|
||||
self.conn.commit()
|
||||
|
||||
async def requeue_task(self, task: DownloadTask):
|
||||
"""
|
||||
将任务重新加入队列
|
||||
|
||||
Args:
|
||||
task: 下载任务
|
||||
"""
|
||||
task.retry_count += 1
|
||||
task.status = TaskStatus.RETRYING
|
||||
task.updated_at = time.time()
|
||||
|
||||
await self.add_task(task)
|
||||
logger.info(f"任务 {task.task_id} 重新加入队列 (重试 {task.retry_count}/{task.max_retries})")
|
||||
|
||||
def get_statistics(self) -> Dict[str, Any]:
|
||||
"""获取队列统计信息"""
|
||||
cursor = self.conn.cursor()
|
||||
|
||||
# 统计各状态任务数
|
||||
cursor.execute('''
|
||||
SELECT status, COUNT(*)
|
||||
FROM tasks
|
||||
GROUP BY status
|
||||
''')
|
||||
|
||||
status_counts = {}
|
||||
for status, count in cursor.fetchall():
|
||||
status_counts[status] = count
|
||||
|
||||
# 计算成功率
|
||||
total = sum(status_counts.values())
|
||||
completed = status_counts.get(TaskStatus.COMPLETED.value, 0)
|
||||
success_rate = (completed / total * 100) if total > 0 else 0
|
||||
|
||||
# 计算平均耗时
|
||||
cursor.execute('''
|
||||
SELECT AVG(completed_at - created_at)
|
||||
FROM tasks
|
||||
WHERE status = ? AND completed_at IS NOT NULL
|
||||
''', (TaskStatus.COMPLETED.value,))
|
||||
|
||||
avg_duration = cursor.fetchone()[0] or 0
|
||||
|
||||
stats = {
|
||||
'total_tasks': total,
|
||||
'pending_tasks': status_counts.get(TaskStatus.PENDING.value, 0),
|
||||
'processing_tasks': status_counts.get(TaskStatus.PROCESSING.value, 0),
|
||||
'completed_tasks': completed,
|
||||
'failed_tasks': status_counts.get(TaskStatus.FAILED.value, 0),
|
||||
'retrying_tasks': status_counts.get(TaskStatus.RETRYING.value, 0),
|
||||
'success_rate': success_rate,
|
||||
'average_duration': avg_duration,
|
||||
'queue_size': self.queue.qsize()
|
||||
}
|
||||
|
||||
return stats
|
||||
|
||||
async def save_progress(self):
|
||||
"""保存进度到数据库"""
|
||||
stats = self.get_statistics()
|
||||
|
||||
cursor = self.conn.cursor()
|
||||
cursor.execute('''
|
||||
INSERT INTO progress (
|
||||
timestamp, total_tasks, pending_tasks, active_tasks,
|
||||
completed_tasks, failed_tasks, success_rate, average_duration
|
||||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?)
|
||||
''', (
|
||||
time.time(),
|
||||
stats['total_tasks'],
|
||||
stats['pending_tasks'],
|
||||
stats['processing_tasks'],
|
||||
stats['completed_tasks'],
|
||||
stats['failed_tasks'],
|
||||
stats['success_rate'],
|
||||
stats['average_duration']
|
||||
))
|
||||
self.conn.commit()
|
||||
|
||||
logger.debug("进度已保存")
|
||||
|
||||
async def start_checkpoint(self):
|
||||
"""启动检查点保存任务"""
|
||||
if not self._checkpoint_task:
|
||||
self._checkpoint_task = asyncio.create_task(self._checkpoint_loop())
|
||||
logger.info("检查点保存任务已启动")
|
||||
|
||||
async def stop_checkpoint(self):
|
||||
"""停止检查点保存任务"""
|
||||
if self._checkpoint_task:
|
||||
self._checkpoint_task.cancel()
|
||||
try:
|
||||
await self._checkpoint_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
self._checkpoint_task = None
|
||||
logger.info("检查点保存任务已停止")
|
||||
|
||||
async def _checkpoint_loop(self):
|
||||
"""检查点保存循环"""
|
||||
while True:
|
||||
try:
|
||||
await asyncio.sleep(self.checkpoint_interval)
|
||||
await self.save_progress()
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
except Exception as e:
|
||||
logger.error(f"保存检查点失败: {e}")
|
||||
|
||||
def get_recent_progress(self, hours: int = 24) -> List[Dict]:
|
||||
"""
|
||||
获取最近的进度记录
|
||||
|
||||
Args:
|
||||
hours: 获取最近多少小时的记录
|
||||
|
||||
Returns:
|
||||
进度记录列表
|
||||
"""
|
||||
cursor = self.conn.cursor()
|
||||
since = time.time() - hours * 3600
|
||||
|
||||
cursor.execute('''
|
||||
SELECT timestamp, total_tasks, completed_tasks, failed_tasks, success_rate
|
||||
FROM progress
|
||||
WHERE timestamp > ?
|
||||
ORDER BY timestamp DESC
|
||||
LIMIT 100
|
||||
''', (since,))
|
||||
|
||||
records = []
|
||||
for row in cursor.fetchall():
|
||||
records.append({
|
||||
'timestamp': row[0],
|
||||
'total_tasks': row[1],
|
||||
'completed_tasks': row[2],
|
||||
'failed_tasks': row[3],
|
||||
'success_rate': row[4]
|
||||
})
|
||||
|
||||
return records
|
||||
|
||||
def cleanup_old_tasks(self, days: int = 7):
|
||||
"""
|
||||
清理旧任务记录
|
||||
|
||||
Args:
|
||||
days: 保留最近多少天的记录
|
||||
"""
|
||||
cursor = self.conn.cursor()
|
||||
cutoff = time.time() - days * 86400
|
||||
|
||||
cursor.execute('''
|
||||
DELETE FROM tasks
|
||||
WHERE status IN (?, ?) AND updated_at < ?
|
||||
''', (TaskStatus.COMPLETED.value, TaskStatus.FAILED.value, cutoff))
|
||||
|
||||
deleted = cursor.rowcount
|
||||
self.conn.commit()
|
||||
|
||||
if deleted > 0:
|
||||
logger.info(f"清理了 {deleted} 条旧任务记录")
|
||||
|
||||
def export_tasks(self, status: Optional[TaskStatus] = None) -> List[Dict]:
|
||||
"""
|
||||
导出任务列表
|
||||
|
||||
Args:
|
||||
status: 筛选状态
|
||||
|
||||
Returns:
|
||||
任务列表
|
||||
"""
|
||||
cursor = self.conn.cursor()
|
||||
|
||||
if status:
|
||||
cursor.execute('''
|
||||
SELECT * FROM tasks WHERE status = ?
|
||||
ORDER BY created_at DESC
|
||||
''', (status.value,))
|
||||
else:
|
||||
cursor.execute('''
|
||||
SELECT * FROM tasks
|
||||
ORDER BY created_at DESC
|
||||
''')
|
||||
|
||||
tasks = []
|
||||
columns = [desc[0] for desc in cursor.description]
|
||||
|
||||
for row in cursor.fetchall():
|
||||
task_dict = dict(zip(columns, row))
|
||||
# 解析JSON字段
|
||||
if task_dict.get('metadata'):
|
||||
try:
|
||||
task_dict['metadata'] = json.loads(task_dict['metadata'])
|
||||
except:
|
||||
pass
|
||||
if task_dict.get('result'):
|
||||
try:
|
||||
task_dict['result'] = json.loads(task_dict['result'])
|
||||
except:
|
||||
pass
|
||||
tasks.append(task_dict)
|
||||
|
||||
return tasks
|
||||
|
||||
def close(self):
|
||||
"""关闭数据库连接"""
|
||||
if self.conn:
|
||||
self.conn.close()
|
||||
self.conn = None
|
||||
logger.info("数据库连接已关闭")
|
||||
|
||||
async def __aenter__(self):
|
||||
"""异步上下文管理器入口"""
|
||||
await self.start_checkpoint()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
||||
"""异步上下文管理器出口"""
|
||||
await self.stop_checkpoint()
|
||||
await self.save_progress()
|
||||
self.close()
|
||||
308
DouYin/apiproxy/douyin/core/rate_limiter.py
Normal file
308
DouYin/apiproxy/douyin/core/rate_limiter.py
Normal file
@@ -0,0 +1,308 @@
|
||||
#!/usr/bin/env python
|
||||
# -*- coding: utf-8 -*-
|
||||
|
||||
"""
|
||||
智能限速器
|
||||
防止请求过快导致被封禁
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
import logging
|
||||
from collections import deque
|
||||
from typing import Optional, Dict, Any
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class RateLimitStrategy(Enum):
|
||||
"""限速策略枚举"""
|
||||
FIXED = "fixed" # 固定速率
|
||||
ADAPTIVE = "adaptive" # 自适应速率
|
||||
BURST = "burst" # 突发模式
|
||||
|
||||
|
||||
@dataclass
|
||||
class RateLimitConfig:
|
||||
"""限速配置"""
|
||||
max_per_second: int = 2
|
||||
max_per_minute: int = 30
|
||||
max_per_hour: int = 1000
|
||||
burst_size: int = 5
|
||||
strategy: RateLimitStrategy = RateLimitStrategy.ADAPTIVE
|
||||
cooldown_time: int = 60 # 触发限制后的冷却时间（秒）
|
||||
|
||||
|
||||
class AdaptiveRateLimiter:
|
||||
"""自适应限速器"""
|
||||
|
||||
def __init__(self, config: Optional[RateLimitConfig] = None):
|
||||
"""
|
||||
初始化限速器
|
||||
|
||||
Args:
|
||||
config: 限速配置
|
||||
"""
|
||||
self.config = config or RateLimitConfig()
|
||||
self.requests = deque()
|
||||
self.failures = deque()
|
||||
self.lock = asyncio.Lock()
|
||||
|
||||
# 当前限制值（可动态调整）
|
||||
self.current_max_per_second = self.config.max_per_second
|
||||
self.current_max_per_minute = self.config.max_per_minute
|
||||
self.current_max_per_hour = self.config.max_per_hour
|
||||
|
||||
# 统计信息
|
||||
self.stats = {
|
||||
'total_requests': 0,
|
||||
'blocked_requests': 0,
|
||||
'rate_adjustments': 0,
|
||||
'current_rate': self.current_max_per_second,
|
||||
'failure_rate': 0.0
|
||||
}
|
||||
|
||||
# 冷却状态
|
||||
self.cooldown_until = 0
|
||||
|
||||
async def acquire(self) -> bool:
|
||||
"""
|
||||
获取请求许可
|
||||
|
||||
Returns:
|
||||
是否获得许可
|
||||
"""
|
||||
async with self.lock:
|
||||
now = time.time()
|
||||
|
||||
# 检查是否在冷却期
|
||||
if self.cooldown_until > now:
|
||||
remaining = self.cooldown_until - now
|
||||
logger.warning(f"限速器处于冷却期，还需等待 {remaining:.1f} 秒")
|
||||
await asyncio.sleep(remaining)
|
||||
self.cooldown_until = 0
|
||||
|
||||
# 清理过期记录
|
||||
self._clean_old_records(now)
|
||||
|
||||
# 检查速率限制
|
||||
while not self._can_proceed(now):
|
||||
# 计算需要等待的时间
|
||||
wait_time = self._calculate_wait_time(now)
|
||||
if wait_time > 0:
|
||||
logger.debug(f"速率限制，等待 {wait_time:.2f} 秒")
|
||||
await asyncio.sleep(wait_time)
|
||||
now = time.time()
|
||||
self._clean_old_records(now)
|
||||
else:
|
||||
# 无法继续，记录被阻塞的请求
|
||||
self.stats['blocked_requests'] += 1
|
||||
return False
|
||||
|
||||
# 记录请求
|
||||
self.requests.append(now)
|
||||
self.stats['total_requests'] += 1
|
||||
|
||||
# 自适应调整
|
||||
if self.config.strategy == RateLimitStrategy.ADAPTIVE:
|
||||
self._adjust_rate()
|
||||
|
||||
return True
|
||||
|
||||
async def __aenter__(self):
|
||||
"""异步上下文管理器入口"""
|
||||
await self.acquire()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
||||
"""异步上下文管理器出口"""
|
||||
if exc_type is not None:
|
||||
# 发生异常，记录失败
|
||||
self.record_failure()
|
||||
|
||||
def record_failure(self):
|
||||
"""记录失败请求"""
|
||||
now = time.time()
|
||||
self.failures.append(now)
|
||||
|
||||
# 自适应调整
|
||||
if self.config.strategy == RateLimitStrategy.ADAPTIVE:
|
||||
self._handle_failure()
|
||||
|
||||
def _can_proceed(self, now: float) -> bool:
|
||||
"""检查是否可以继续请求"""
|
||||
# 检查每秒限制
|
||||
recent_second = [r for r in self.requests if now - r < 1]
|
||||
if len(recent_second) >= self.current_max_per_second:
|
||||
return False
|
||||
|
||||
# 检查每分钟限制
|
||||
recent_minute = [r for r in self.requests if now - r < 60]
|
||||
if len(recent_minute) >= self.current_max_per_minute:
|
||||
return False
|
||||
|
||||
# 检查每小时限制
|
||||
recent_hour = [r for r in self.requests if now - r < 3600]
|
||||
if len(recent_hour) >= self.current_max_per_hour:
|
||||
return False
|
||||
|
||||
# 突发模式检查
|
||||
if self.config.strategy == RateLimitStrategy.BURST:
|
||||
recent_burst = [r for r in self.requests if now - r < 0.1]
|
||||
if len(recent_burst) >= self.config.burst_size:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def _calculate_wait_time(self, now: float) -> float:
|
||||
"""计算需要等待的时间"""
|
||||
wait_times = []
|
||||
|
||||
# 计算每秒限制的等待时间
|
||||
recent_second = [r for r in self.requests if now - r < 1]
|
||||
if len(recent_second) >= self.current_max_per_second:
|
||||
oldest_in_second = min(recent_second)
|
||||
wait_times.append(1 - (now - oldest_in_second))
|
||||
|
||||
# 计算每分钟限制的等待时间
|
||||
recent_minute = [r for r in self.requests if now - r < 60]
|
||||
if len(recent_minute) >= self.current_max_per_minute:
|
||||
oldest_in_minute = min(recent_minute)
|
||||
wait_times.append(60 - (now - oldest_in_minute))
|
||||
|
||||
# 返回最小等待时间
|
||||
return min(wait_times) if wait_times else 0.1
|
||||
|
||||
def _clean_old_records(self, now: float):
|
||||
"""清理过期记录"""
|
||||
# 保留最近1小时的请求记录
|
||||
while self.requests and now - self.requests[0] > 3600:
|
||||
self.requests.popleft()
|
||||
|
||||
# 保留最近10分钟的失败记录
|
||||
while self.failures and now - self.failures[0] > 600:
|
||||
self.failures.popleft()
|
||||
|
||||
def _adjust_rate(self):
|
||||
"""自适应调整速率"""
|
||||
now = time.time()
|
||||
|
||||
# 计算失败率
|
||||
recent_failures = [f for f in self.failures if now - f < 60]
|
||||
recent_requests = [r for r in self.requests if now - r < 60]
|
||||
|
||||
if len(recent_requests) > 10:
|
||||
failure_rate = len(recent_failures) / len(recent_requests)
|
||||
self.stats['failure_rate'] = failure_rate
|
||||
|
||||
if failure_rate > 0.3:
|
||||
# 失败率过高，降低速率
|
||||
self._decrease_rate()
|
||||
elif failure_rate < 0.05 and len(recent_requests) > 20:
|
||||
# 失败率很低，尝试提高速率
|
||||
self._increase_rate()
|
||||
|
||||
def _handle_failure(self):
|
||||
"""处理失败，调整限速策略"""
|
||||
now = time.time()
|
||||
recent_failures = [f for f in self.failures if now - f < 10]
|
||||
|
||||
# 如果短时间内失败次数过多，触发冷却
|
||||
if len(recent_failures) >= 5:
|
||||
logger.warning(f"检测到频繁失败，进入冷却期 {self.config.cooldown_time} 秒")
|
||||
self.cooldown_until = now + self.config.cooldown_time
|
||||
self._decrease_rate()
|
||||
|
||||
def _decrease_rate(self):
|
||||
"""降低请求速率"""
|
||||
old_rate = self.current_max_per_second
|
||||
|
||||
self.current_max_per_second = max(1, int(self.current_max_per_second * 0.7))
|
||||
self.current_max_per_minute = max(10, int(self.current_max_per_minute * 0.7))
|
||||
self.current_max_per_hour = max(100, int(self.current_max_per_hour * 0.7))
|
||||
|
||||
if old_rate != self.current_max_per_second:
|
||||
self.stats['rate_adjustments'] += 1
|
||||
self.stats['current_rate'] = self.current_max_per_second
|
||||
logger.info(f"降低请求速率: {old_rate}/s -> {self.current_max_per_second}/s")
|
||||
|
||||
def _increase_rate(self):
|
||||
"""提高请求速率"""
|
||||
old_rate = self.current_max_per_second
|
||||
|
||||
# 不超过配置的最大值
|
||||
self.current_max_per_second = min(
|
||||
self.config.max_per_second,
|
||||
int(self.current_max_per_second * 1.2)
|
||||
)
|
||||
self.current_max_per_minute = min(
|
||||
self.config.max_per_minute,
|
||||
int(self.current_max_per_minute * 1.2)
|
||||
)
|
||||
self.current_max_per_hour = min(
|
||||
self.config.max_per_hour,
|
||||
int(self.current_max_per_hour * 1.2)
|
||||
)
|
||||
|
||||
if old_rate != self.current_max_per_second:
|
||||
self.stats['rate_adjustments'] += 1
|
||||
self.stats['current_rate'] = self.current_max_per_second
|
||||
logger.info(f"提高请求速率: {old_rate}/s -> {self.current_max_per_second}/s")
|
||||
|
||||
def get_stats(self) -> Dict[str, Any]:
|
||||
"""获取统计信息"""
|
||||
return self.stats.copy()
|
||||
|
||||
def reset_stats(self):
|
||||
"""重置统计信息"""
|
||||
self.stats = {
|
||||
'total_requests': 0,
|
||||
'blocked_requests': 0,
|
||||
'rate_adjustments': 0,
|
||||
'current_rate': self.current_max_per_second,
|
||||
'failure_rate': 0.0
|
||||
}
|
||||
|
||||
def set_cooldown(self, seconds: int):
|
||||
"""手动设置冷却时间"""
|
||||
self.cooldown_until = time.time() + seconds
|
||||
logger.info(f"手动设置冷却期 {seconds} 秒")
|
||||
|
||||
|
||||
class SimpleRateLimiter:
|
||||
"""简单限速器（固定速率）"""
|
||||
|
||||
def __init__(self, requests_per_second: float = 1.0):
|
||||
"""
|
||||
初始化简单限速器
|
||||
|
||||
Args:
|
||||
requests_per_second: 每秒允许的请求数
|
||||
"""
|
||||
self.requests_per_second = requests_per_second
|
||||
self.min_interval = 1.0 / requests_per_second
|
||||
self.last_request_time = 0
|
||||
self.lock = asyncio.Lock()
|
||||
|
||||
async def acquire(self):
|
||||
"""获取请求许可"""
|
||||
async with self.lock:
|
||||
now = time.time()
|
||||
time_since_last = now - self.last_request_time
|
||||
|
||||
if time_since_last < self.min_interval:
|
||||
sleep_time = self.min_interval - time_since_last
|
||||
await asyncio.sleep(sleep_time)
|
||||
|
||||
self.last_request_time = time.time()
|
||||
|
||||
async def __aenter__(self):
|
||||
"""异步上下文管理器入口"""
|
||||
await self.acquire()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
||||
"""异步上下文管理器出口"""
|
||||
pass
|
||||
Reference in New Issue
Block a user