"""虎牙任务执行器：公共基础（由 huya_runner.py 按功能域拆分）。""" from __future__ import annotations import asyncio import copy import threading import time from datetime import UTC, datetime from typing import Any from loguru import logger from sqlalchemy.orm import Session from ..models import HuyaTask from .huya_service import cookie_value # 绑定跳转与充值下单共用（bind 与 recharge 两个域都引用了 sourceId） HUYA_RECHARGE_SOURCE_ID = "yellowcarlist" from typing import TYPE_CHECKING if TYPE_CHECKING: from .huya_runner import HuyaBatchRunner class HuyaBatchRunnerCore: """虎牙任务执行器公共基础：批次状态、日志、任务落库。""" def __init__( self, db: Session, batch_id: str, task_type: str, payload: dict | None = None, log_queue: asyncio.Queue | None = None, loop: asyncio.AbstractEventLoop | None = None, concurrency: int = 3, ): self.db = db self.batch_id = batch_id self.task_type = task_type self.payload = payload or {} self.log_queue = log_queue self.loop = loop self.concurrency = max(1, min(concurrency, 10)) self._stop = threading.Event() self._counter_lock = threading.Lock() self._started = 0 self._wss_sessions: list[Any] = [] def stop(self): self._stop.set() def _open_wss_session(self, uid: int, cookie: str, kind: str): """在当前任务线程创建并登记 WSS 会话，统一由 finally 回收。""" from core.huya.elite_session import HuyaEliteWssSession session = HuyaEliteWssSession( uid=uid, cookie=cookie, kind=kind, logger=lambda message: self._push_log("info", f"[{uid}] {message}"), ) self._wss_sessions.append(session) return session def _close_wss_sessions(self): for session in reversed(self._wss_sessions): try: session.close() except Exception: # noqa: BLE001 - cleanup must not mask task status logger.debug("[huya] WSS 会话清理异常") self._wss_sessions.clear() def _push_log(self, level: str, message: str): if level == "result": try: huya_batch_registry.mark_finished(self.batch_id) except NameError: pass if level != "result" and message: log_func = getattr(logger, level, logger.info) log_func(f"[huya] {message}") if self.log_queue and self.loop: asyncio.run_coroutine_threadsafe( self.log_queue.put({"level": level, "message": message}), self.loop, ) @staticmethod def _account_name(account_info: dict) -> str: return ( account_info.get("nickname") or account_info.get("username") or account_info.get("uid") or f"#{account_info.get('account_id')}" ) @staticmethod def _to_int(value) -> int: text = str(value or "").strip() return int(text) if text.isdigit() else 0 def _resolve_uid(self, account_info: dict) -> int: cookie = account_info.get("cookie") or "" return ( self._to_int(account_info.get("yyuid")) or self._to_int(account_info.get("uid")) or self._to_int(cookie_value(cookie, "yyuid")) or self._to_int(cookie_value(cookie, "udb_uid")) ) @staticmethod def _format_local_time(timestamp: int) -> str: if not timestamp: return "" return datetime.fromtimestamp(timestamp, UTC).strftime("%Y-%m-%d %H:%M:%S") @staticmethod def _parse_scheduled_time(value) -> datetime | None: text = str(value or "").strip() if not text: return None try: normalized = text.replace("Z", "+00:00") dt = datetime.fromisoformat(normalized) except ValueError: return None if dt.tzinfo is None: return dt.astimezone() return dt def _wait_until(self, when: datetime, uid: int) -> bool: target = when.timestamp() local_text = self._format_local_time(int(target)) self._push_log("info", f"[{uid}] 定时兑换等待到 {local_text}") while not self._stop.is_set(): remaining = target - time.time() if remaining <= 0: return True time.sleep(min(0.2, max(0.02, remaining))) return False def _mark_task( self, worker_db: Session, task: HuyaTask, status: str, message: str, result: dict | None = None, ): task.status = status task.message = message task.result = result task.finished_at = datetime.now(UTC) worker_db.commit() self._push_task_event(task) def _update_task_progress( self, worker_db: Session, task: HuyaTask, status: str, message: str, result: dict | None = None, ): task.status = status task.message = message if result is not None: task.result = result worker_db.commit() self._push_task_event(task) def _push_task_event(self, task: HuyaTask) -> None: """向批次 WS 推送任务状态；二维码图片留给详情接口按需读取。""" if not self.log_queue or not self.loop: return result = copy.deepcopy(task.result) if isinstance(task.result, dict) else None if result and isinstance(result.get("mini_qrcode_image"), str): result.pop("mini_qrcode_image", None) result["has_mini_qrcode"] = True event = { "level": "task", "message": "", "task": { "id": task.id, "batch_id": task.batch_id, "account_id": task.account_id, "task_type": task.task_type, "handbook_scope": getattr(task, "handbook_scope", "legacy") or "legacy", "status": task.status or "", "message": task.message or "", "result": result, "created_by": task.created_by, "created_at": task.created_at.isoformat() if task.created_at else None, "finished_at": task.finished_at.isoformat() if task.finished_at else None, }, } asyncio.run_coroutine_threadsafe(self.log_queue.put(event), self.loop) class HuyaBatchRegistry: """管理运行中的虎牙批次。""" def __init__(self): self._batches: dict[str, dict] = {} self._lock = threading.Lock() def _cleanup_locked(self, ttl_seconds: int = 300): now = time.time() expired = [ batch_id for batch_id, batch in self._batches.items() if batch.get("finished") and now - float(batch.get("finished_at") or now) > ttl_seconds ] for batch_id in expired: self._batches.pop(batch_id, None) def register( self, batch_id: str, log_queue: asyncio.Queue, loop: asyncio.AbstractEventLoop, runner: HuyaBatchRunner, ): with self._lock: self._cleanup_locked() self._batches[batch_id] = { "log_queue": log_queue, "loop": loop, "runner": runner, "finished": False, "finished_at": None, } def get(self, batch_id: str): with self._lock: self._cleanup_locked() return self._batches.get(batch_id) def active_ids(self) -> set[str]: with self._lock: self._cleanup_locked() return { batch_id for batch_id, batch in self._batches.items() if not batch.get("finished") } def mark_finished(self, batch_id: str): with self._lock: batch = self._batches.get(batch_id) if not batch: return batch["finished"] = True batch["finished_at"] = time.time() def pop(self, batch_id: str): with self._lock: return self._batches.pop(batch_id, None) huya_batch_registry = HuyaBatchRegistry()