212 lines
6.2 KiB
Python
212 lines
6.2 KiB
Python
"""虎牙任务执行器：公共基础（由 huya_runner.py 按功能域拆分）。"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import threading
|
||
import time
|
||
from datetime import datetime, timezone
|
||
from typing import Optional
|
||
|
||
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: 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":
|
||
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).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: Optional[dict] = None,
|
||
):
|
||
task.status = status
|
||
task.message = message
|
||
task.result = result
|
||
task.finished_at = datetime.now(timezone.utc)
|
||
worker_db.commit()
|
||
|
||
def _update_task_progress(
|
||
self,
|
||
worker_db: Session,
|
||
task: HuyaTask,
|
||
status: str,
|
||
message: str,
|
||
result: Optional[dict] = None,
|
||
):
|
||
task.status = status
|
||
task.message = message
|
||
if result is not None:
|
||
task.result = result
|
||
worker_db.commit()
|
||
|
||
|
||
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()
|