实现虎牙积分查询执行器
This commit is contained in:
yml2213
+62
-11
@@ -1,11 +1,13 @@
|
||||
"""虎牙基础管理路由"""
|
||||
|
||||
import asyncio
|
||||
import threading
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query, WebSocket
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query, WebSocket, WebSocketDisconnect
|
||||
from sqlalchemy.orm import Session, joinedload
|
||||
|
||||
from ..database import get_db
|
||||
from ..database import SessionLocal, get_db
|
||||
from ..deps import authenticate_websocket, require_permission
|
||||
from ..models import HuyaAccount, HuyaConfig, HuyaGoodsSnapshot, HuyaTask, User
|
||||
from ..schemas import (
|
||||
@@ -23,6 +25,7 @@ from ..services.huya_service import (
|
||||
ensure_huya_config,
|
||||
import_huya_cookies,
|
||||
)
|
||||
from ..services.huya_runner import HuyaBatchRunner, huya_batch_registry
|
||||
|
||||
|
||||
router = APIRouter(prefix="/api/huya", tags=["虎牙"])
|
||||
@@ -197,12 +200,12 @@ def list_goods(
|
||||
|
||||
|
||||
@router.post("/tasks/batch")
|
||||
def create_task_batch(
|
||||
async def create_task_batch(
|
||||
req: HuyaTaskBatchRequest,
|
||||
db: Session = Depends(get_db),
|
||||
current: User = Depends(require_permission("huya:task")),
|
||||
):
|
||||
"""创建虎牙任务记录，真实执行器后续接入。"""
|
||||
"""创建虎牙任务记录并启动后台执行器。"""
|
||||
if not req.account_ids:
|
||||
raise HTTPException(status_code=400, detail="请选择虎牙账号")
|
||||
try:
|
||||
@@ -217,6 +220,24 @@ def create_task_batch(
|
||||
raise HTTPException(status_code=400, detail=str(exc)) from exc
|
||||
if count == 0:
|
||||
raise HTTPException(status_code=400, detail="没有有效的虎牙账号")
|
||||
|
||||
log_queue = asyncio.Queue()
|
||||
loop = asyncio.get_running_loop()
|
||||
thread_db = SessionLocal()
|
||||
runner = HuyaBatchRunner(
|
||||
db=thread_db,
|
||||
batch_id=batch_id,
|
||||
task_type=req.task_type,
|
||||
payload=req.payload,
|
||||
log_queue=log_queue,
|
||||
loop=loop,
|
||||
concurrency=req.concurrency,
|
||||
)
|
||||
huya_batch_registry.register(batch_id, log_queue, loop, runner)
|
||||
|
||||
thread = threading.Thread(target=runner.run, daemon=True)
|
||||
thread.start()
|
||||
|
||||
return {"batch_id": batch_id, "count": count, "success": True}
|
||||
|
||||
|
||||
@@ -234,17 +255,47 @@ def list_tasks(
|
||||
return [_task_out(task) for task in tasks]
|
||||
|
||||
|
||||
@router.post("/stop/{batch_id}")
|
||||
def stop_batch(
|
||||
batch_id: str,
|
||||
current: User = Depends(require_permission("huya:task")),
|
||||
):
|
||||
"""停止正在运行的虎牙批次。"""
|
||||
batch = huya_batch_registry.get(batch_id)
|
||||
if batch:
|
||||
batch["runner"].stop()
|
||||
return {"message": "已发送停止信号", "success": True}
|
||||
raise HTTPException(status_code=404, detail="批次不存在或已结束")
|
||||
|
||||
|
||||
@router.websocket("/ws/{batch_id}")
|
||||
async def ws_huya_logs(websocket: WebSocket, batch_id: str):
|
||||
"""虎牙实时日志占位通道。"""
|
||||
"""虎牙实时日志推送通道。"""
|
||||
user = authenticate_websocket(websocket)
|
||||
if not user:
|
||||
await websocket.close(code=1008, reason="未授权")
|
||||
return
|
||||
await websocket.accept()
|
||||
await websocket.send_json({
|
||||
"level": "warning",
|
||||
"message": f"虎牙批次 {batch_id} 已创建，真实执行器尚未接入",
|
||||
})
|
||||
await websocket.send_json({"level": "result", "message": ""})
|
||||
await websocket.close()
|
||||
|
||||
batch = huya_batch_registry.get(batch_id)
|
||||
if not batch:
|
||||
await websocket.send_json({"level": "error", "message": "批次不存在或已结束"})
|
||||
await websocket.close()
|
||||
return
|
||||
|
||||
log_queue: asyncio.Queue = batch["log_queue"]
|
||||
|
||||
try:
|
||||
while True:
|
||||
try:
|
||||
msg = await asyncio.wait_for(log_queue.get(), timeout=30)
|
||||
await websocket.send_json(msg)
|
||||
if msg.get("level") == "result":
|
||||
await asyncio.sleep(0.1)
|
||||
break
|
||||
except asyncio.TimeoutError:
|
||||
await websocket.send_json({"level": "heartbeat", "message": ""})
|
||||
except WebSocketDisconnect:
|
||||
pass
|
||||
finally:
|
||||
huya_batch_registry.pop(batch_id)
|
||||
|
||||
@@ -0,0 +1,290 @@
|
||||
"""虎牙任务批次执行器。"""
|
||||
|
||||
import asyncio
|
||||
import threading
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from datetime import datetime, timezone
|
||||
from typing import Optional
|
||||
|
||||
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 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: None)
|
||||
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()
|
||||
@@ -140,7 +140,7 @@ def create_planned_tasks(
|
||||
created_by: int,
|
||||
payload: dict | None = None,
|
||||
) -> tuple[str, int]:
|
||||
"""创建虎牙任务记录，真实执行器后续接入。"""
|
||||
"""创建虎牙任务记录，等待后台执行器消费。"""
|
||||
if task_type not in SUPPORTED_TASK_TYPES:
|
||||
raise ValueError("不支持的任务类型")
|
||||
|
||||
@@ -153,7 +153,7 @@ def create_planned_tasks(
|
||||
account_id=account.id,
|
||||
task_type=task_type,
|
||||
status="planned",
|
||||
message="任务已创建，等待虎牙执行器接入",
|
||||
message="任务已创建，等待执行",
|
||||
result={"payload": payload} if payload else None,
|
||||
created_by=created_by,
|
||||
))
|
||||
|
||||
Reference in New Issue
Block a user