494 lines
15 KiB
Python
494 lines
15 KiB
Python
|
#!/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()
|