"""登录服务：复用 core/ 核心模块，在线程池中并发执行登录并推送日志。""" import asyncio import threading import time import uuid from types import SimpleNamespace from concurrent.futures import ThreadPoolExecutor, as_completed from datetime import datetime, timezone from typing import Optional from sqlalchemy.orm import Session from core.douyu import DouyuLogin from core.douyu.proxy import resolve_working_proxy, get_proxy_manager from ..models import Account as AccountModel, LoginTask, ProxyConfig as ProxyConfigModel class LoginBatchRunner: """批量登录执行器，在线程中运行，通过 ThreadPoolExecutor 并发登录多个账号。""" # 代理池耗尽时等待恢复的最大秒数 _PROXY_WAIT_MAX = 60 def __init__( self, db: Session, account_ids: list[int], created_by: int, creator_permissions: list[str], max_geetest_retries: int = 5, max_proxy_retries: int = 0, max_login_retries: int = 3, max_total_time: float = 300, proxy_config: Optional[ProxyConfigModel] = None, log_queue: Optional[asyncio.Queue] = None, loop: Optional[asyncio.AbstractEventLoop] = None, concurrency: int = 3, ): self.db = db self.account_ids = account_ids self.created_by = created_by self.creator_permissions = creator_permissions self.max_geetest_retries = max_geetest_retries self.max_proxy_retries = max_proxy_retries self.max_login_retries = max_login_retries self.max_total_time = max_total_time self.proxy_config = proxy_config self.log_queue = log_queue self.loop = loop self.batch_id = uuid.uuid4().hex[:12] self.concurrency = max(1, min(concurrency, 10)) # 限制 1-10 self._stop = threading.Event() self._counter_lock = threading.Lock() self._completed = 0 # 共享代理管理器（带锁，避免并发白名单限流；极验失败时可刷新代理） self._shared_proxy_manager = None if proxy_config and proxy_config.enabled and proxy_config.api_url: wl_uid = proxy_config.whitelist_uid or "" if proxy_config.whitelist_enabled else "" wl_ukey = proxy_config.whitelist_ukey or "" if proxy_config.whitelist_enabled else "" self._shared_proxy_manager = get_proxy_manager( proxy_config.api_url, whitelist_uid=wl_uid, whitelist_ukey=wl_ukey, ) def stop(self): self._stop.set() def _sleep_or_stop(self, seconds: float) -> bool: """可中断等待；返回 True 表示收到停止信号。""" deadline = time.monotonic() + seconds while time.monotonic() < deadline: if self._stop.is_set(): return True time.sleep(min(0.2, deadline - time.monotonic())) return self._stop.is_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, ) def _resolve_static_proxy(self) -> tuple[Optional[dict], str]: """ 解析静态代理配置（仅处理无代理和静态代理场景）。 API代理由 DouyuLogin 通过 proxy_manager 内部管理， 不在此处预先获取——登录过程中的代理切换（极验失败、整体重试） 都在 DouyuLogin 内部自治完成。 """ if not self.proxy_config or not self.proxy_config.enabled: return None, '' # 静态代理 if self.proxy_config.http or self.proxy_config.https: proxy_url = self.proxy_config.http or self.proxy_config.https return {'http': proxy_url, 'https': proxy_url}, f'使用静态代理: {proxy_url}' # API代理：不在此处获取，由 DouyuLogin 通过 proxy_manager 内部管理 return None, '' def _execute_one(self, task_id: int, acc_info: dict, total: int): """在独立线程中执行单个账号登录，使用独立的 DB 会话和代理。""" if self._stop.is_set(): self._push_log("warning", f"任务已停止，跳过: {acc_info['username']}") return worker_db = SessionLocal() try: task = worker_db.query(LoginTask).filter(LoginTask.id == task_id).first() if not task: return task.status = "running" worker_db.commit() with self._counter_lock: self._completed += 1 current = self._completed self._push_log("info", f"[{current}/{total}] 开始登录: {acc_info['username']}") # 解析代理配置 proxy_dict, proxy_msg = self._resolve_static_proxy() if proxy_msg: self._push_log("info", f"[{current}] {proxy_msg}") # API代理模式下，验证代理池是否有可用代理；池空时等待恢复（最多 _PROXY_WAIT_MAX 秒） # 注意：不预取代理传给 DouyuLogin，让 DouyuLogin 通过 proxy_manager 内部自治管理 is_api_proxy = ( self.proxy_config and self.proxy_config.enabled and not (self.proxy_config.http or self.proxy_config.https) and self._shared_proxy_manager ) if is_api_proxy: # 先验证代理池是否有可用代理（获取后立即归还，不占用） test_proxy = self._shared_proxy_manager.get_proxy() if test_proxy: self._shared_proxy_manager.release_proxy(test_proxy) self._push_log("info", f"[{current}] 代理池可用，由 DouyuLogin 内部管理代理获取与切换") else: # 代理池暂时耗尽，等待冷却代理恢复或新代理入池 self._push_log("warning", f"[{current}] 代理池暂时耗尽，等待恢复...") pool_available = False for wait_sec in range(0, self._PROXY_WAIT_MAX, 10): if self._stop.is_set(): task.status = "error" task.message = "任务已停止" task.finished_at = datetime.now(timezone.utc) worker_db.commit() return if self._sleep_or_stop(10): task.status = "error" task.message = "任务已停止" task.finished_at = datetime.now(timezone.utc) worker_db.commit() return test_proxy = self._shared_proxy_manager.get_proxy() if test_proxy: self._shared_proxy_manager.release_proxy(test_proxy) pool_available = True self._push_log("info", f"[{current}] 代理池恢复，由 DouyuLogin 内部管理代理") break remaining = self._PROXY_WAIT_MAX - wait_sec - 10 self._push_log("info", f"[{current}] 代理池仍为空，继续等待... (剩余 {remaining}s)") if not pool_available: task.status = "error" task.message = f"代理池耗尽，等待 {self._PROXY_WAIT_MAX}s 后仍无可用代理" task.finished_at = datetime.now(timezone.utc) worker_db.commit() self._push_log("error", f"[{current}] {acc_info['username']} 代理池耗尽，放弃") return # 静态代理启用但配置为空 → 不可用 if self.proxy_config and self.proxy_config.enabled and not (self.proxy_config.http or self.proxy_config.https) and not self._shared_proxy_manager and not proxy_dict: task.status = "error" task.message = "代理不可用: 未配置代理" task.finished_at = datetime.now(timezone.utc) worker_db.commit() self._push_log("error", f"[{current}] {acc_info['username']} 代理不可用") return try: account = SimpleNamespace( username=acc_info["username"], password=acc_info["password"], email=acc_info["email"], email_password=acc_info["email_password"], email_imap_server=acc_info["email_imap_server"] or "", email_imap_port=acc_info["email_imap_port"] or 993, email_imap_ssl=acc_info["email_imap_ssl"], ) loginer = DouyuLogin( account, proxy=proxy_dict, max_geetest_retries=self.max_geetest_retries, max_proxy_retries=self.max_proxy_retries, max_login_retries=self.max_login_retries, max_total_time=self.max_total_time, proxy_manager=self._shared_proxy_manager, stop_event=self._stop, ) result = loginer.login() if result.success: task.status = "success" task.cookie = result.cookie task.message = result.message or "登录成功" self._push_log("success", f"[{current}] {acc_info['username']} {task.message}") else: task.status = "failed" task.message = result.message self._push_log("error", f"[{current}] {acc_info['username']} 登录失败: {result.message}") except Exception as e: task.status = "error" task.message = str(e) self._push_log("error", f"[{current}] {acc_info['username']} 登录异常: {e}") task.finished_at = datetime.now(timezone.utc) worker_db.commit() finally: worker_db.close() def run(self): """在线程中执行批量登录。""" batch_id = self.batch_id concurrency = self.concurrency self._push_log("info", f"批量登录任务 {batch_id} 开始，共 {len(self.account_ids)} 个账号，并发数: {concurrency}") try: # 创建或复用任务记录（顺序执行，线程安全） task_infos: list[dict] = [] # {task_id, acc_info} for aid in self.account_ids: acc = self.db.query(AccountModel).filter(AccountModel.id == aid).first() if not acc: continue # 权限检查：客服只能跑分配给自己的 if "login:view_all" not in self.creator_permissions: if acc.assigned_to != self.created_by: self._push_log("warning", f"跳过无权账号: {acc.username}") continue # 复用该账号最近一条失败任务记录，避免重复产生多条 existing_task = ( self.db.query(LoginTask) .filter(LoginTask.account_id == aid, LoginTask.status.in_(["failed", "error"])) .order_by(LoginTask.id.desc()) .first() ) if existing_task: existing_task.batch_id = batch_id existing_task.status = "pending" existing_task.cookie = "" existing_task.message = "" existing_task.finished_at = None task = existing_task else: task = LoginTask( batch_id=batch_id, account_id=aid, status="pending", created_by=self.created_by, ) self.db.add(task) self.db.flush() # 获取 task.id task_infos.append({ "task_id": task.id, "acc_info": { "username": acc.username, "password": acc.password, "email": acc.email, "email_password": acc.email_password, "email_imap_server": acc.email_imap_server or "", "email_imap_port": acc.email_imap_port or 993, "email_imap_ssl": acc.email_imap_ssl if acc.email_imap_ssl is not None else True, }, }) self.db.commit() total = len(task_infos) if total == 0: self._push_log("warning", "没有可执行的账号") self._push_log("result", "") return # 并发执行登录，每个账号独立获取代理 with ThreadPoolExecutor(max_workers=concurrency) as executor: futures = [] for item in task_infos: if self._stop.is_set(): self._push_log("warning", "任务已停止，跳过剩余账号") break future = executor.submit( self._execute_one, item["task_id"], item["acc_info"], total, ) futures.append(future) # 等待所有任务完成 for future in as_completed(futures): try: future.result() except Exception as e: self._push_log("error", f"Worker 异常: {e}") self._push_log("info", f"批量登录任务 {batch_id} 完成") self._push_log("result", "") finally: # 确保 DB Session 被关闭，避免连接泄漏 self.db.close() class BatchRegistry: """管理运行中的登录批次状态。""" def __init__(self): self._batches: dict[str, dict] = {} self._lock = threading.Lock() def register(self, batch_id: str, log_queue: asyncio.Queue, loop: asyncio.AbstractEventLoop, runner: LoginBatchRunner): with self._lock: self._batches[batch_id] = { "log_queue": log_queue, "loop": loop, "runner": runner, } def get(self, batch_id: str): with self._lock: return self._batches.get(batch_id) def pop(self, batch_id: str): with self._lock: return self._batches.pop(batch_id, None) # 模块级单例 batch_registry = BatchRegistry() # 在模块末尾导入 SessionLocal（避免循环导入） from ..database import SessionLocal