"""虎牙任务批次执行器。""" import asyncio import threading from concurrent.futures import ThreadPoolExecutor, as_completed from datetime import datetime, timezone from typing import Optional from loguru import logger from sqlalchemy.orm import Session, joinedload from core.huya import HuyaHttpClient from ..database import SessionLocal from ..models import HuyaAccount, HuyaConfig, HuyaTask from .huya_service import cookie_value class HuyaBatchRunner: """批量执行虎牙任务，通过队列推送实时日志。""" def __init__( self, db: Session, batch_id: str, task_type: str, payload: Optional[dict] = None, log_queue: Optional[asyncio.Queue] = None, loop: Optional[asyncio.AbstractEventLoop] = 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 def stop(self): self._stop.set() def _push_log(self, level: str, message: str): 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")) ) def _mark_task( self, worker_db: Session, task: HuyaTask, status: str, message: str, result: Optional[dict] = None, ): task.status = status task.message = message task.result = result task.finished_at = datetime.now(timezone.utc) worker_db.commit() def _execute_query_points( self, worker_db: Session, task: HuyaTask, account: HuyaAccount, account_info: dict, config_info: dict, ): sid = str(self.payload.get("sid") or config_info.get("sid") or "").strip() if not sid: self._mark_task(worker_db, task, "failed", "请先配置虎牙活动 SID") return sid_int = self._to_int(sid) if not sid_int: self._mark_task(worker_db, task, "failed", f"虎牙活动 SID 无效: {sid}") return uid = self._resolve_uid(account_info) if not uid: self._mark_task(worker_db, task, "failed", "无法从账号或 Cookie 解析 yyuid") return cookie = account_info.get("cookie") or "" if not cookie: self._mark_task(worker_db, task, "failed", "账号 Cookie 为空") return client = HuyaHttpClient(logger=lambda msg: self._push_log("info", f"[{uid}] {msg}")) response = client.query_user_score(uid=uid, cookie=cookie, sid=sid_int) if response is None: self._mark_task(worker_db, task, "error", "虎牙积分接口无响应") return result = response.to_dict() result["sid"] = sid_int if response.status != 200: self._mark_task( worker_db, task, "failed", response.msg or f"虎牙积分查询失败: {response.status}", result, ) return points = response.available_score account.points = points account.status = "points_queried" account.updated_at = datetime.now(timezone.utc) self._mark_task(worker_db, task, "success", f"积分: {points}", result) def _execute_one(self, task_id: int, account_info: dict, config_info: dict, total: int): worker_db = SessionLocal() try: task = worker_db.query(HuyaTask).filter(HuyaTask.id == task_id).first() account = worker_db.query(HuyaAccount).filter(HuyaAccount.id == account_info["account_id"]).first() if not task or not account: return if self._stop.is_set(): self._mark_task(worker_db, task, "failed", "任务已停止") return task.status = "running" task.message = "执行中" task.finished_at = None worker_db.commit() with self._counter_lock: self._started += 1 current = self._started name = self._account_name(account_info) self._push_log("info", f"[{current}/{total}] 开始虎牙任务: {name}") if self.task_type != "query_points": self._mark_task(worker_db, task, "failed", "该虎牙任务执行器暂未实现") self._push_log("warning", f"[{current}] {name} 暂未实现: {self.task_type}") return try: self._execute_query_points(worker_db, task, account, account_info, config_info) worker_db.refresh(task) if task.status == "success": self._push_log("success", f"[{current}] {name} {task.message}") else: self._push_log("error", f"[{current}] {name} {task.message}") except Exception as exc: self._mark_task(worker_db, task, "error", f"执行异常: {exc}") self._push_log("error", f"[{current}] {name} 执行异常: {exc}") finally: worker_db.close() def run(self): """在线程中执行虎牙批次任务。""" self._push_log( "info", f"虎牙批次 {self.batch_id} 开始，共执行 {self.task_type}，并发数: {self.concurrency}", ) try: config = self.db.query(HuyaConfig).first() config_info = { "sid": config.sid if config else "", "outer_act_id": config.outer_act_id if config else "", "bind_act_id": config.bind_act_id if config else "", "pay_channel": config.pay_channel if config else "", } tasks = ( self.db.query(HuyaTask) .options(joinedload(HuyaTask.account)) .filter(HuyaTask.batch_id == self.batch_id) .order_by(HuyaTask.id.asc()) .all() ) task_infos = [] for task in tasks: account = task.account if not account: task.status = "error" task.message = "账号不存在" task.finished_at = datetime.now(timezone.utc) continue task.status = "pending" task.message = "等待执行" task.finished_at = None task_infos.append({ "task_id": task.id, "account_info": { "account_id": account.id, "uid": account.uid or "", "yyuid": account.yyuid or "", "username": account.username or "", "nickname": account.nickname or "", "cookie": account.cookie or "", }, }) self.db.commit() total = len(task_infos) if total == 0: self._push_log("warning", "没有可执行的虎牙任务") self._push_log("result", "") return with ThreadPoolExecutor(max_workers=self.concurrency) as executor: futures = [] for item in task_infos: if self._stop.is_set(): self._push_log("warning", "任务已停止，跳过剩余账号") break futures.append(executor.submit( self._execute_one, item["task_id"], item["account_info"], config_info, total, )) for future in as_completed(futures): try: future.result() except Exception as exc: self._push_log("error", f"虎牙 Worker 异常: {exc}") self._push_log("info", f"虎牙批次 {self.batch_id} 完成") self._push_log("result", "") except Exception as exc: self._push_log("error", f"虎牙批次执行异常: {exc}") self._push_log("result", "") finally: self.db.close() class HuyaBatchRegistry: """管理运行中的虎牙批次。""" def __init__(self): self._batches: dict[str, dict] = {} def register( self, batch_id: str, log_queue: asyncio.Queue, loop: asyncio.AbstractEventLoop, runner: HuyaBatchRunner, ): self._batches[batch_id] = { "log_queue": log_queue, "loop": loop, "runner": runner, } def get(self, batch_id: str): return self._batches.get(batch_id) def pop(self, batch_id: str): return self._batches.pop(batch_id, None) huya_batch_registry = HuyaBatchRegistry()