488 lines
20 KiB
Python
488 lines
20 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
任务处理器 - 后台并发处理队列中的任务
|
||
支持动态调整并发数，通过 SocketIO 实时推送进度和日志
|
||
"""
|
||
import threading
|
||
import time
|
||
import logging
|
||
import traceback
|
||
import psutil
|
||
from task_queue import get_task_queue, TaskStatus
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# 全局变量，用于存储 socketio 实例
|
||
_socketio_instance = None
|
||
|
||
def set_socketio(socketio):
|
||
"""设置 SocketIO 实例"""
|
||
global _socketio_instance
|
||
_socketio_instance = socketio
|
||
|
||
def emit_log(task_id, message, level='info'):
|
||
"""保存日志到数据库"""
|
||
from datetime import datetime
|
||
timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
|
||
|
||
# 保存到数据库
|
||
try:
|
||
from database import get_database
|
||
db = get_database()
|
||
db.add_task_log(task_id, message, level, timestamp)
|
||
except Exception as e:
|
||
logger.error(f"保存日志到数据库失败: {e}")
|
||
|
||
logger.info(f"[{task_id}] {message}")
|
||
|
||
def emit_progress(task_id, progress, current_step='', **kwargs):
|
||
"""更新任务进度"""
|
||
logger.info(f"[{task_id}] 进度: {progress}% - {current_step}")
|
||
|
||
|
||
class TaskWorker:
|
||
"""任务处理工作线程池（支持动态并发）"""
|
||
|
||
def __init__(self, min_workers=1, max_workers=3):
|
||
self.queue = get_task_queue()
|
||
self.running = False
|
||
self.min_workers = min_workers # 最小并发数
|
||
self.max_workers = max_workers # 最大并发数
|
||
self.current_workers = min_workers # 当前并发数
|
||
self.worker_threads = [] # 工作线程列表
|
||
self.processing_tasks = set() # 正在处理的任务ID集合
|
||
self.lock = threading.Lock()
|
||
|
||
def start(self):
|
||
"""启动工作线程池"""
|
||
if self.running:
|
||
logger.warning("工作线程池已经在运行")
|
||
return
|
||
|
||
self.running = True
|
||
|
||
# 启动初始工作线程
|
||
for i in range(self.min_workers):
|
||
self._start_worker(i)
|
||
|
||
# 启动动态调整线程
|
||
self.adjust_thread = threading.Thread(target=self._adjust_workers, daemon=True)
|
||
self.adjust_thread.start()
|
||
|
||
logger.info(f"任务处理器已启动（初始并发数: {self.min_workers}，最大并发数: {self.max_workers}）")
|
||
|
||
def _start_worker(self, worker_id):
|
||
"""启动一个工作线程"""
|
||
thread = threading.Thread(target=self._work_loop, args=(worker_id,), daemon=True)
|
||
thread.start()
|
||
with self.lock:
|
||
self.worker_threads.append(thread)
|
||
logger.info(f"工作线程 #{worker_id} 已启动")
|
||
|
||
def stop(self):
|
||
"""停止工作线程池"""
|
||
self.running = False
|
||
for thread in self.worker_threads:
|
||
if thread and thread.is_alive():
|
||
thread.join(timeout=5)
|
||
logger.info("任务处理器已停止")
|
||
|
||
def _work_loop(self, worker_id):
|
||
"""工作循环（单个线程）"""
|
||
logger.info(f"工作线程 #{worker_id} 进入循环")
|
||
|
||
while self.running:
|
||
try:
|
||
# 获取待处理任务
|
||
task = self.queue.get_pending_task()
|
||
|
||
if task:
|
||
task_id = task["task_id"]
|
||
|
||
# 检查是否已经有其他线程在处理这个任务
|
||
with self.lock:
|
||
if task_id in self.processing_tasks:
|
||
continue
|
||
self.processing_tasks.add(task_id)
|
||
|
||
try:
|
||
logger.info(f"工作线程 #{worker_id} 开始处理任务: {task_id}")
|
||
self._process_task(task, worker_id)
|
||
finally:
|
||
# 处理完成后从集合中移除
|
||
with self.lock:
|
||
self.processing_tasks.discard(task_id)
|
||
else:
|
||
# 没有任务，休息一会
|
||
time.sleep(2)
|
||
|
||
except Exception as e:
|
||
logger.error(f"工作线程 #{worker_id} 错误: {e}")
|
||
logger.error(traceback.format_exc())
|
||
time.sleep(5)
|
||
|
||
logger.info(f"工作线程 #{worker_id} 退出循环")
|
||
|
||
def _adjust_workers(self):
|
||
"""动态调整工作线程数量"""
|
||
logger.info("动态调整线程已启动")
|
||
|
||
while self.running:
|
||
try:
|
||
time.sleep(10) # 每10秒检查一次
|
||
|
||
# 获取系统资源信息
|
||
cpu_percent = psutil.cpu_percent(interval=1)
|
||
memory_percent = psutil.virtual_memory().percent
|
||
|
||
# 获取队列信息
|
||
pending_count = len([t for t in self.queue.get_all_tasks() if t.get('status') == 'pending'])
|
||
processing_count = len(self.processing_tasks)
|
||
|
||
# 决策逻辑
|
||
target_workers = self._calculate_target_workers(
|
||
pending_count,
|
||
processing_count,
|
||
cpu_percent,
|
||
memory_percent
|
||
)
|
||
|
||
# 调整线程数
|
||
if target_workers > self.current_workers:
|
||
# 增加线程
|
||
for i in range(self.current_workers, target_workers):
|
||
self._start_worker(i)
|
||
logger.info(f"增加工作线程: {self.current_workers} -> {target_workers}")
|
||
self.current_workers = target_workers
|
||
elif target_workers < self.current_workers:
|
||
# 减少线程（自然退出，不强制终止）
|
||
logger.info(f"准备减少工作线程: {self.current_workers} -> {target_workers}")
|
||
self.current_workers = target_workers
|
||
|
||
except Exception as e:
|
||
logger.error(f"调整线程数错误: {e}")
|
||
logger.error(traceback.format_exc())
|
||
time.sleep(30)
|
||
|
||
def _calculate_target_workers(self, pending_count, processing_count, cpu_percent, memory_percent):
|
||
"""计算目标线程数"""
|
||
# 基本逻辑：
|
||
# 1. 如果没有待处理任务，保持最小线程数
|
||
# 2. 如果有很多待处理任务，且系统资源充足，增加线程
|
||
# 3. 如果系统资源紧张，减少线程
|
||
|
||
if pending_count == 0:
|
||
return self.min_workers
|
||
|
||
# 系统资源紧张（CPU>80% 或 内存>85%）
|
||
if cpu_percent > 80 or memory_percent > 85:
|
||
logger.warning(f"系统资源紧张 (CPU: {cpu_percent}%, 内存: {memory_percent}%)")
|
||
return max(self.min_workers, self.current_workers - 1)
|
||
|
||
# 系统资源充足
|
||
if cpu_percent < 50 and memory_percent < 70:
|
||
# 根据待处理任务数决定线程数
|
||
if pending_count >= 3:
|
||
return min(self.max_workers, self.current_workers + 1)
|
||
elif pending_count >= 1:
|
||
return min(self.max_workers, max(2, self.current_workers))
|
||
|
||
# 默认保持当前线程数
|
||
return self.current_workers
|
||
|
||
def _process_task(self, task, worker_id):
|
||
"""处理单个任务"""
|
||
task_id = task["task_id"]
|
||
logger.info(f"工作线程 #{worker_id} 开始处理任务: {task_id}")
|
||
emit_log(task_id, f"任务开始处理 (Worker #{worker_id})")
|
||
|
||
# 获取当前重试次数
|
||
retry_count = task.get("retry_count", 0)
|
||
|
||
try:
|
||
# 更新状态为处理中
|
||
self.queue.update_task_status(
|
||
task_id,
|
||
TaskStatus.PROCESSING,
|
||
current_step="准备处理"
|
||
)
|
||
emit_progress(task_id, 5, "准备处理")
|
||
|
||
# 导入必要的模块
|
||
from app import BaijiahaoScraper
|
||
import pandas as pd
|
||
import os
|
||
|
||
# 步骤1: 解析URL获取UK
|
||
self.queue.update_task_progress(task_id, 10, "解析URL获取UK")
|
||
emit_progress(task_id, 10, "解析URL获取UK")
|
||
emit_log(task_id, f"URL: {task['url']}")
|
||
|
||
url = task["url"]
|
||
use_proxy = task.get("use_proxy", False)
|
||
proxy_api_url = task.get("proxy_api_url")
|
||
articles_only = task.get("articles_only", True) # 获取是否仅爬取文章
|
||
|
||
if use_proxy:
|
||
emit_log(task_id, "已启用代理IP池", "info")
|
||
|
||
if articles_only:
|
||
emit_log(task_id, "已启用文章过滤（跳过视频内容）", "info")
|
||
|
||
# 提取app_id
|
||
import re
|
||
app_id_match = re.search(r'app_id=([^&\s]+)', url)
|
||
if not app_id_match:
|
||
raise Exception("无法从 URL 中提取 app_id")
|
||
|
||
app_id = app_id_match.group(1)
|
||
emit_log(task_id, f"解析到 app_id: {app_id}")
|
||
|
||
# 获取UK
|
||
emit_log(task_id, "正在获取用户 UK...")
|
||
try:
|
||
uk, cookies = BaijiahaoScraper.get_uk_from_app_id(
|
||
app_id,
|
||
use_proxy=use_proxy,
|
||
proxy_api_url=proxy_api_url
|
||
)
|
||
emit_log(task_id, f"成功获取 UK: {uk[:20]}...")
|
||
except Exception as uk_error:
|
||
emit_log(task_id, f"获取UK失败: {str(uk_error)}", "error")
|
||
raise
|
||
|
||
# 步骤2: 初始化爬虫
|
||
self.queue.update_task_progress(task_id, 20, "初始化爬虫")
|
||
emit_progress(task_id, 20, "初始化爬虫")
|
||
emit_log(task_id, "初始化爬虫实例...")
|
||
|
||
scraper = BaijiahaoScraper(
|
||
uk=uk,
|
||
cookies=cookies,
|
||
use_proxy=use_proxy,
|
||
proxy_api_url=proxy_api_url
|
||
)
|
||
|
||
# 步骤3: 获取文章列表
|
||
months = task.get("months", 6)
|
||
|
||
# 检查是否有断点续传数据
|
||
last_page = task.get("last_page", 0)
|
||
last_ctime = task.get("last_ctime")
|
||
start_page = 1
|
||
start_ctime = None
|
||
|
||
if last_page > 0 and last_ctime:
|
||
# 断点续传
|
||
start_page = last_page
|
||
start_ctime = last_ctime
|
||
emit_log(task_id, f"🔄 检测到断点数据，从第{start_page}页继续爬取", "info")
|
||
|
||
# 检查缓存中是否有数据
|
||
from database import get_database
|
||
db = get_database()
|
||
cached_count = db.get_cached_article_count(task_id)
|
||
if cached_count > 0:
|
||
emit_log(task_id, f"💾 已缓存 {cached_count} 篇文章，将继续爬取...", "info")
|
||
else:
|
||
# 新任务，清除之前的缓存（如果有）
|
||
from database import get_database
|
||
db = get_database()
|
||
db.clear_article_cache(task_id)
|
||
|
||
self.queue.update_task_progress(task_id, 30, f"获取文章列表（近{months}个月）")
|
||
emit_progress(task_id, 30, f"获取文章列表（近{months}个月）")
|
||
emit_log(task_id, f"开始获取近 {months} 个月的文章...")
|
||
emit_log(task_id, "提示：抓取过程较慢（8-12秒/页），请耐心等待...", "info")
|
||
emit_log(task_id, "系统正在使用代理IP池抓取数据，过程中会自动切换IP应对反爬...", "info")
|
||
|
||
# 定义保存回调函数：每页数据立即保存到数据库
|
||
from database import get_database
|
||
db = get_database()
|
||
|
||
def save_page_to_db(page, articles, ctime):
|
||
"""保存每页数据到数据库缓存"""
|
||
if articles:
|
||
db.save_articles_batch(task_id, articles, page)
|
||
emit_log(task_id, f"💾 第{page}页数据已保存，{len(articles)}篇文章", "success")
|
||
|
||
# 更新任务的断点信息
|
||
self.queue.update_task_status(
|
||
task_id,
|
||
TaskStatus.PROCESSING,
|
||
last_page=page,
|
||
last_ctime=ctime
|
||
)
|
||
|
||
# 更新进度（粗略估计，30-80%区间）
|
||
total_cached = db.get_cached_article_count(task_id)
|
||
progress = min(30 + int(page * 2), 80) # 每页增加2%，最多80%
|
||
self.queue.update_task_progress(
|
||
task_id,
|
||
progress,
|
||
f"正在抓取第{page}页...",
|
||
processed_articles=total_cached
|
||
)
|
||
emit_progress(task_id, progress, f"正在抓取第{page}页...", processed_articles=total_cached)
|
||
|
||
# 调用 get_articles，传入回调函数和断点参数
|
||
result = scraper.get_articles(
|
||
months=months,
|
||
app_id=app_id,
|
||
articles_only=articles_only,
|
||
task_id=task_id,
|
||
on_page_fetched=save_page_to_db,
|
||
start_page=start_page,
|
||
start_ctime=start_ctime
|
||
)
|
||
|
||
if not result or not result.get('completed'):
|
||
# 未完成，保留断点信息以便续传
|
||
raise Exception(f"抓取未完成，已保存到第{result.get('last_page', start_page)}页")
|
||
|
||
# 从数据库读取全部缓存文章
|
||
articles = db.get_cached_articles(task_id)
|
||
|
||
if not articles:
|
||
raise Exception("未获取到文章数据")
|
||
|
||
# 更新总文章数
|
||
total = len(articles)
|
||
self.queue.update_task_status(
|
||
task_id,
|
||
TaskStatus.PROCESSING,
|
||
total_articles=total
|
||
)
|
||
emit_log(task_id, f"成功获取 {total} 篇文章", "success")
|
||
|
||
# 步骤4: 生成Excel（直接使用数据库中的数据）
|
||
self.queue.update_task_progress(task_id, 90, "生成Excel文件")
|
||
emit_progress(task_id, 90, "生成Excel文件")
|
||
emit_log(task_id, "正在生成 Excel 文件...")
|
||
|
||
df = pd.DataFrame(articles)
|
||
|
||
# 生成文件名
|
||
timestamp = time.strftime("%Y%m%d_%H%M%S")
|
||
filename = f"百家号文章_{app_id}_{timestamp}.xlsx"
|
||
result_file = os.path.join(self.queue.results_dir, filename)
|
||
|
||
# 保存Excel
|
||
with pd.ExcelWriter(result_file, engine='openpyxl') as writer:
|
||
df.to_excel(writer, index=False, sheet_name='文章列表')
|
||
|
||
# 调整列宽
|
||
worksheet = writer.sheets['文章列表']
|
||
worksheet.column_dimensions['A'].width = 80 # 标题列
|
||
worksheet.column_dimensions['B'].width = 20 # 时间列
|
||
|
||
emit_log(task_id, f"Excel 文件已生成: {filename}")
|
||
|
||
# 清除缓存数据（任务已完成）
|
||
db.clear_article_cache(task_id)
|
||
emit_log(task_id, "🗑️ 已清除缓存数据", "info")
|
||
|
||
# 步骤5: 完成
|
||
self.queue.update_task_status(
|
||
task_id,
|
||
TaskStatus.COMPLETED,
|
||
progress=100,
|
||
current_step="处理完成",
|
||
result_file=filename,
|
||
processed_articles=total
|
||
)
|
||
emit_progress(task_id, 100, "处理完成", result_file=filename)
|
||
emit_log(task_id, f"✅ 任务完成！导出 {total} 篇文章", "success")
|
||
|
||
logger.info(f"工作线程 #{worker_id} 任务完成: {task_id}, 导出 {total} 篇文章")
|
||
|
||
except Exception as e:
|
||
error_msg = str(e)
|
||
logger.error(f"工作线程 #{worker_id} 任务失败: {task_id}, 错误: {error_msg}")
|
||
|
||
# 记录详细错误堆栈
|
||
error_traceback = traceback.format_exc()
|
||
logger.error(error_traceback)
|
||
|
||
# 将错误堆栈也推送到前端（分行推送）
|
||
emit_log(task_id, f"❌ 任务失败: {error_msg}", "error")
|
||
|
||
# 推送错误详情（每行作为独立日志）
|
||
for line in error_traceback.split('\n'):
|
||
if line.strip():
|
||
emit_log(task_id, line, "error")
|
||
|
||
# 判断是否需要重试或暂停
|
||
retry_count += 1
|
||
|
||
# 检查是否有缓存数据（如果有，说明部分成功）
|
||
from database import get_database
|
||
db = get_database()
|
||
cached_count = db.get_cached_article_count(task_id)
|
||
|
||
# 如果已经有缓存数据，说明部分成功，增加重试次数
|
||
max_retries = 10 if cached_count > 0 else 3 # 有缓存时允许10次重试
|
||
|
||
# 如果连续失败超过上限，暂停任务10分钟
|
||
if retry_count >= max_retries:
|
||
from datetime import datetime
|
||
paused_at = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
|
||
|
||
if cached_count > 0:
|
||
emit_log(task_id, f"⚠️ 连续失败{retry_count}次，已缓存{cached_count}篇文章，10分钟后继续尝试", "warning")
|
||
else:
|
||
emit_log(task_id, f"⚠️ 连续失败{retry_count}次，任务将暂停10分钟后自动重试", "warning")
|
||
|
||
self.queue.update_task_status(
|
||
task_id,
|
||
TaskStatus.PAUSED,
|
||
error=error_msg,
|
||
last_error=error_msg,
|
||
retry_count=retry_count,
|
||
paused_at=paused_at,
|
||
current_step=f"暂停中（10分钟后重试） - 错误: {error_msg}"
|
||
)
|
||
emit_progress(task_id, 0, f"暂停中（10分钟后重试）")
|
||
|
||
logger.warning(f"任务 {task_id} 已暂停，将在 {paused_at} 后10分钟恢复，已缓存{cached_count}篇")
|
||
else:
|
||
# 重试次数未达到上限，标记为待处理状态，等待下次重试
|
||
if cached_count > 0:
|
||
emit_log(task_id, f"⚠️ 任务失败，将进行第{retry_count + 1}次重试（已缓存{cached_count}篇）", "warning")
|
||
else:
|
||
emit_log(task_id, f"⚠️ 任务失败，将进行第{retry_count + 1}次重试", "warning")
|
||
|
||
self.queue.update_task_status(
|
||
task_id,
|
||
TaskStatus.PENDING,
|
||
error=error_msg,
|
||
last_error=error_msg,
|
||
retry_count=retry_count,
|
||
current_step=f"等待重试 (已失败{retry_count}次) - {error_msg}"
|
||
)
|
||
emit_progress(task_id, 0, f"等待重试")
|
||
|
||
|
||
# 全局工作线程
|
||
_worker = None
|
||
|
||
|
||
def get_task_worker():
|
||
"""获取全局任务处理器实例"""
|
||
global _worker
|
||
if _worker is None:
|
||
_worker = TaskWorker()
|
||
return _worker
|
||
|
||
|
||
def start_task_worker():
|
||
"""启动任务处理器"""
|
||
worker = get_task_worker()
|
||
worker.start()
|
||
|
||
|
||
def stop_task_worker():
|
||
"""停止任务处理器"""
|
||
worker = get_task_worker()
|
||
worker.stop()
|