'commit'
This commit is contained in:
HuangHai
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()
|
||||
Reference in New Issue
Block a user