366 lines
13 KiB
Python
366 lines
13 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
from __future__ import annotations
|
||
|
||
import heapq
|
||
import threading
|
||
import time
|
||
from concurrent.futures import ThreadPoolExecutor, wait
|
||
from dataclasses import dataclass
|
||
|
||
import database
|
||
from app_logger import get_logger
|
||
from services.state import safe_get_account, safe_get_task, safe_remove_task, safe_set_task
|
||
from services.task_batches import _batch_task_record_result, _get_batch_id_from_source
|
||
|
||
logger = get_logger("app")
|
||
|
||
# VIP优先级队列（仅用于可视化/调试）
|
||
vip_task_queue = [] # VIP用户任务队列
|
||
normal_task_queue = [] # 普通用户任务队列
|
||
task_queue_lock = threading.Lock()
|
||
|
||
@dataclass
|
||
class _TaskRequest:
|
||
user_id: int
|
||
account_id: str
|
||
browse_type: str
|
||
enable_screenshot: bool
|
||
source: str
|
||
retry_count: int
|
||
submitted_at: float
|
||
is_vip: bool
|
||
seq: int
|
||
canceled: bool = False
|
||
done_callback: object = None
|
||
|
||
|
||
class TaskScheduler:
|
||
"""全局任务调度器：队列排队，不为每个任务单独创建线程。"""
|
||
|
||
def __init__(self, max_global: int, max_per_user: int, max_queue_size: int = 1000, run_task_fn=None):
|
||
self.max_global = max(1, int(max_global))
|
||
self.max_per_user = max(1, int(max_per_user))
|
||
self.max_queue_size = max(1, int(max_queue_size))
|
||
|
||
self._cond = threading.Condition()
|
||
self._pending = [] # heap: (priority, submitted_at, seq, task)
|
||
self._pending_by_account = {} # {account_id: task}
|
||
self._seq = 0
|
||
self._known_account_ids = set()
|
||
|
||
self._running_global = 0
|
||
self._running_by_user = {} # {user_id: running_count}
|
||
|
||
self._executor_max_workers = self.max_global
|
||
self._executor = ThreadPoolExecutor(max_workers=self._executor_max_workers, thread_name_prefix="TaskWorker")
|
||
|
||
self._futures_lock = threading.Lock()
|
||
self._active_futures = set()
|
||
|
||
self._running = True
|
||
self._run_task_fn = run_task_fn
|
||
self._dispatcher_thread = threading.Thread(target=self._dispatch_loop, daemon=True, name="TaskDispatcher")
|
||
self._dispatcher_thread.start()
|
||
|
||
def _track_future(self, future) -> None:
|
||
with self._futures_lock:
|
||
self._active_futures.add(future)
|
||
try:
|
||
future.add_done_callback(self._untrack_future)
|
||
except Exception:
|
||
pass
|
||
|
||
def _untrack_future(self, future) -> None:
|
||
with self._futures_lock:
|
||
self._active_futures.discard(future)
|
||
|
||
def shutdown(self, timeout: float = 5.0):
|
||
"""停止调度器（用于进程退出清理）"""
|
||
with self._cond:
|
||
self._running = False
|
||
self._cond.notify_all()
|
||
|
||
try:
|
||
self._dispatcher_thread.join(timeout=timeout)
|
||
except Exception:
|
||
pass
|
||
|
||
# 等待已提交的任务收尾（最多等待 timeout 秒），避免遗留 active_task 干扰后续调度/测试
|
||
try:
|
||
deadline = time.time() + max(0.0, float(timeout or 0))
|
||
while True:
|
||
with self._futures_lock:
|
||
pending = [f for f in self._active_futures if not f.done()]
|
||
if not pending:
|
||
break
|
||
remaining = deadline - time.time()
|
||
if remaining <= 0:
|
||
break
|
||
wait(pending, timeout=remaining)
|
||
except Exception:
|
||
pass
|
||
|
||
try:
|
||
self._executor.shutdown(wait=False)
|
||
except Exception:
|
||
pass
|
||
|
||
# 最后兜底：清理本调度器提交过的 active_task，避免测试/重启时被“任务已在运行中”误拦截
|
||
try:
|
||
with self._cond:
|
||
known_ids = set(self._known_account_ids) | set(self._pending_by_account.keys())
|
||
self._pending.clear()
|
||
self._pending_by_account.clear()
|
||
self._cond.notify_all()
|
||
for account_id in known_ids:
|
||
safe_remove_task(account_id)
|
||
except Exception:
|
||
pass
|
||
|
||
def update_limits(self, max_global: int = None, max_per_user: int = None, max_queue_size: int = None):
|
||
"""动态更新并发/队列上限（不影响已在运行的任务）"""
|
||
with self._cond:
|
||
if max_per_user is not None:
|
||
self.max_per_user = max(1, int(max_per_user))
|
||
if max_queue_size is not None:
|
||
self.max_queue_size = max(1, int(max_queue_size))
|
||
|
||
if max_global is not None:
|
||
new_max_global = max(1, int(max_global))
|
||
self.max_global = new_max_global
|
||
if new_max_global > self._executor_max_workers:
|
||
# 立即关闭旧线程池，防止资源泄漏
|
||
old_executor = self._executor
|
||
self._executor_max_workers = new_max_global
|
||
self._executor = ThreadPoolExecutor(
|
||
max_workers=self._executor_max_workers, thread_name_prefix="TaskWorker"
|
||
)
|
||
# 立即关闭旧线程池
|
||
try:
|
||
old_executor.shutdown(wait=False)
|
||
logger.info(f"线程池已扩容：{old_executor._max_workers} -> {self._executor_max_workers}")
|
||
except Exception as e:
|
||
logger.warning(f"关闭旧线程池失败: {e}")
|
||
|
||
self._cond.notify_all()
|
||
|
||
def get_queue_state_snapshot(self) -> dict:
|
||
"""获取调度器队列/运行状态快照（用于前端展示/监控）。"""
|
||
with self._cond:
|
||
pending_tasks = [t for t in self._pending_by_account.values() if t and not t.canceled]
|
||
pending_tasks.sort(key=lambda t: (0 if t.is_vip else 1, t.submitted_at, t.seq))
|
||
|
||
positions = {}
|
||
for idx, t in enumerate(pending_tasks):
|
||
positions[t.account_id] = {"queue_position": idx + 1, "queue_ahead": idx, "is_vip": bool(t.is_vip)}
|
||
|
||
return {
|
||
"pending_total": len(pending_tasks),
|
||
"running_total": int(self._running_global),
|
||
"running_by_user": dict(self._running_by_user),
|
||
"positions": positions,
|
||
}
|
||
|
||
def submit_task(
|
||
self,
|
||
user_id: int,
|
||
account_id: str,
|
||
browse_type: str,
|
||
enable_screenshot: bool = True,
|
||
source: str = "manual",
|
||
retry_count: int = 0,
|
||
is_vip: bool = None,
|
||
done_callback=None,
|
||
):
|
||
"""提交任务进入队列（返回: (ok, message)）"""
|
||
if not user_id or not account_id:
|
||
return False, "参数错误"
|
||
|
||
submitted_at = time.time()
|
||
if is_vip is None:
|
||
try:
|
||
is_vip = bool(database.is_user_vip(user_id))
|
||
except Exception:
|
||
is_vip = False
|
||
else:
|
||
is_vip = bool(is_vip)
|
||
|
||
with self._cond:
|
||
if not self._running:
|
||
return False, "调度器未运行"
|
||
if len(self._pending_by_account) >= self.max_queue_size:
|
||
return False, "任务队列已满，请稍后再试"
|
||
if account_id in self._pending_by_account:
|
||
return False, "任务已在队列中"
|
||
if safe_get_task(account_id) is not None:
|
||
return False, "任务已在运行中"
|
||
|
||
self._seq += 1
|
||
task = _TaskRequest(
|
||
user_id=user_id,
|
||
account_id=account_id,
|
||
browse_type=browse_type,
|
||
enable_screenshot=bool(enable_screenshot),
|
||
source=source,
|
||
retry_count=int(retry_count or 0),
|
||
submitted_at=submitted_at,
|
||
is_vip=is_vip,
|
||
seq=self._seq,
|
||
done_callback=done_callback,
|
||
)
|
||
self._pending_by_account[account_id] = task
|
||
self._known_account_ids.add(account_id)
|
||
priority = 0 if is_vip else 1
|
||
heapq.heappush(self._pending, (priority, task.submitted_at, task.seq, task))
|
||
self._cond.notify_all()
|
||
|
||
# 用于可视化/调试：记录队列
|
||
with task_queue_lock:
|
||
if is_vip:
|
||
vip_task_queue.append(account_id)
|
||
else:
|
||
normal_task_queue.append(account_id)
|
||
|
||
return True, "已加入队列"
|
||
|
||
def cancel_pending_task(self, user_id: int, account_id: str) -> bool:
|
||
"""取消尚未开始的排队任务（已运行的任务由 should_stop 控制）"""
|
||
canceled_task = None
|
||
with self._cond:
|
||
task = self._pending_by_account.pop(account_id, None)
|
||
if not task:
|
||
return False
|
||
task.canceled = True
|
||
canceled_task = task
|
||
self._cond.notify_all()
|
||
|
||
# 从可视化队列移除
|
||
with task_queue_lock:
|
||
if account_id in vip_task_queue:
|
||
vip_task_queue.remove(account_id)
|
||
if account_id in normal_task_queue:
|
||
normal_task_queue.remove(account_id)
|
||
|
||
# 批次任务：取消也要推进完成计数，避免批次缓存常驻
|
||
try:
|
||
batch_id = _get_batch_id_from_source(canceled_task.source)
|
||
if batch_id:
|
||
acc = safe_get_account(user_id, account_id)
|
||
if acc:
|
||
account_name = acc.remark if acc.remark else acc.username
|
||
else:
|
||
account_name = account_id
|
||
_batch_task_record_result(
|
||
batch_id=batch_id,
|
||
account_name=account_name,
|
||
screenshot_path=None,
|
||
total_items=0,
|
||
total_attachments=0,
|
||
)
|
||
except Exception:
|
||
pass
|
||
|
||
return True
|
||
|
||
def _dispatch_loop(self):
|
||
while True:
|
||
task = None
|
||
with self._cond:
|
||
if not self._running:
|
||
return
|
||
|
||
if not self._pending or self._running_global >= self.max_global:
|
||
self._cond.wait(timeout=0.5)
|
||
continue
|
||
|
||
task = self._pop_next_runnable_locked()
|
||
if task is None:
|
||
self._cond.wait(timeout=0.5)
|
||
continue
|
||
|
||
self._running_global += 1
|
||
self._running_by_user[task.user_id] = self._running_by_user.get(task.user_id, 0) + 1
|
||
|
||
# 从队列移除（可视化）
|
||
with task_queue_lock:
|
||
if task.account_id in vip_task_queue:
|
||
vip_task_queue.remove(task.account_id)
|
||
if task.account_id in normal_task_queue:
|
||
normal_task_queue.remove(task.account_id)
|
||
|
||
try:
|
||
future = self._executor.submit(self._run_task_wrapper, task)
|
||
self._track_future(future)
|
||
safe_set_task(task.account_id, future)
|
||
except Exception:
|
||
with self._cond:
|
||
self._running_global = max(0, self._running_global - 1)
|
||
# 使用默认值 0 与增加时保持一致
|
||
self._running_by_user[task.user_id] = max(0, self._running_by_user.get(task.user_id, 0) - 1)
|
||
if self._running_by_user.get(task.user_id) == 0:
|
||
self._running_by_user.pop(task.user_id, None)
|
||
self._cond.notify_all()
|
||
|
||
def _pop_next_runnable_locked(self):
|
||
"""在锁内从优先队列取出“可运行”的任务，避免VIP任务占位阻塞普通任务。"""
|
||
if not self._pending:
|
||
return None
|
||
|
||
skipped = []
|
||
selected = None
|
||
|
||
while self._pending:
|
||
_, _, _, task = heapq.heappop(self._pending)
|
||
|
||
if task.canceled:
|
||
continue
|
||
if self._pending_by_account.get(task.account_id) is not task:
|
||
continue
|
||
|
||
running_for_user = self._running_by_user.get(task.user_id, 0)
|
||
if running_for_user >= self.max_per_user:
|
||
skipped.append(task)
|
||
continue
|
||
|
||
selected = task
|
||
break
|
||
|
||
for t in skipped:
|
||
priority = 0 if t.is_vip else 1
|
||
heapq.heappush(self._pending, (priority, t.submitted_at, t.seq, t))
|
||
|
||
if selected is None:
|
||
return None
|
||
|
||
self._pending_by_account.pop(selected.account_id, None)
|
||
return selected
|
||
|
||
def _run_task_wrapper(self, task: _TaskRequest):
|
||
try:
|
||
if callable(self._run_task_fn):
|
||
self._run_task_fn(
|
||
user_id=task.user_id,
|
||
account_id=task.account_id,
|
||
browse_type=task.browse_type,
|
||
enable_screenshot=task.enable_screenshot,
|
||
source=task.source,
|
||
retry_count=task.retry_count,
|
||
)
|
||
finally:
|
||
try:
|
||
if callable(task.done_callback):
|
||
task.done_callback()
|
||
except Exception:
|
||
pass
|
||
safe_remove_task(task.account_id)
|
||
with self._cond:
|
||
self._running_global = max(0, self._running_global - 1)
|
||
# 使用默认值 0 与增加时保持一致
|
||
self._running_by_user[task.user_id] = max(0, self._running_by_user.get(task.user_id, 0) - 1)
|
||
if self._running_by_user.get(task.user_id) == 0:
|
||
self._running_by_user.pop(task.user_id, None)
|
||
self._cond.notify_all()
|
||
|
||
|