增加重试

2026-06-22 15:45:48 +08:00
parent 2443d23316
commit 9fba12aad0
8 changed files with 199 additions and 60 deletions
@@ -1,8 +1,9 @@
-"""登录服务：复用 core/ 核心模块，在线程池中执行登录并推送日志。"""
+"""登录服务：复用 core/ 核心模块，在线程池中并发执行登录并推送日志。"""

 import asyncio
 import threading
 import uuid
+from concurrent.futures import ThreadPoolExecutor, as_completed
 from datetime import datetime
 from typing import Optional

@@ -16,7 +17,7 @@ from ..permissions import has_permission


 class LoginBatchRunner:
-    """批量登录执行器，在线程中运行，通过 asyncio.Queue 推送日志。"""
+    """批量登录执行器，在线程中运行，通过 ThreadPoolExecutor 并发登录多个账号。"""

    def __init__(
        self,
@@ -28,6 +29,7 @@ class LoginBatchRunner:
        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
@@ -38,7 +40,10 @@ class LoginBatchRunner:
        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

    def stop(self):
        self._stop.set()
@@ -86,13 +91,73 @@ class LoginBatchRunner:

        return None, ''

+    def _execute_one(self, task_id: int, acc_info: dict, proxy_dict: Optional[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']}")
+
+            try:
+                account = Account(
+                    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,
+                )
+
+                loginer = DouyuLogin(
+                    account,
+                    proxy=proxy_dict,
+                    max_geetest_retries=self.max_geetest_retries,
+                )
+                result = loginer.login()
+
+                if result.success:
+                    task.status = "success"
+                    task.cookie = result.cookie
+                    task.message = "登录成功"
+                    self._push_log("success", f"[{current}] {acc_info['username']} 登录成功")
+                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.utcnow()
+            worker_db.commit()
+
+        finally:
+            worker_db.close()
+
    def run(self):
        """在线程中执行批量登录。"""
        batch_id = self.batch_id
-        self._push_log("info", f"批量登录任务 {batch_id} 开始，共 {len(self.account_ids)} 个账号")
+        concurrency = self.concurrency
+        self._push_log("info", f"批量登录任务 {batch_id} 开始，共 {len(self.account_ids)} 个账号，并发数: {concurrency}")

-        # 创建任务记录
-        tasks = []
+        # 创建任务记录（顺序执行，线程安全）
+        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:
@@ -110,9 +175,26 @@ class LoginBatchRunner:
                created_by=self.created_by,
            )
            self.db.add(task)
-            tasks.append((task, acc))
+            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,
+                },
+            })

        self.db.commit()
+        total = len(task_infos)
+        if total == 0:
+            self._push_log("warning", "没有可执行的账号")
+            self._push_log("result", "")
+            return

        # 代理预检
        proxy_dict, proxy_msg = self._resolve_proxy()
@@ -122,57 +204,46 @@ class LoginBatchRunner:
        # 如果启用了代理但预检失败，终止任务
        if self.proxy_config and self.proxy_config.enabled and not proxy_dict:
            self._push_log("error", f"代理不可用，任务终止: {proxy_msg}")
-            for task, _ in tasks:
-                task.status = "error"
-                task.message = f"代理不可用: {proxy_msg}"
-                task.finished_at = datetime.utcnow()
-            self.db.commit()
+            for item in task_infos:
+                worker_db = SessionLocal()
+                try:
+                    task = worker_db.query(LoginTask).filter(LoginTask.id == item["task_id"]).first()
+                    if task:
+                        task.status = "error"
+                        task.message = f"代理不可用: {proxy_msg}"
+                        task.finished_at = datetime.utcnow()
+                        worker_db.commit()
+                finally:
+                    worker_db.close()
+            self._push_log("result", "")
            return

-        for i, (task, acc) in enumerate(tasks):
-            if self._stop.is_set():
-                self._push_log("warning", "任务已停止")
-                break
-
-            task.status = "running"
-            self.db.commit()
-
-            self._push_log("info", f"[{i+1}/{len(tasks)}] 开始登录: {acc.username}")
-
-            try:
-                account = Account(
-                    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,
+        # 并发执行登录
+        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"],
+                    proxy_dict,
+                    total,
                )
+                futures.append(future)

-                loginer = DouyuLogin(
-                    account,
-                    proxy=proxy_dict,
-                    max_geetest_retries=self.max_geetest_retries,
-                )
-                result = loginer.login()
-
-                if result.success:
-                    task.status = "success"
-                    task.cookie = result.cookie
-                    task.message = "登录成功"
-                    self._push_log("success", f"[{i+1}] {acc.username} 登录成功")
-                else:
-                    task.status = "failed"
-                    task.message = result.message
-                    self._push_log("error", f"[{i+1}] {acc.username} 登录失败: {result.message}")
-
-            except Exception as e:
-                task.status = "error"
-                task.message = str(e)
-                self._push_log("error", f"[{i+1}] {acc.username} 登录异常: {e}")
-
-            task.finished_at = datetime.utcnow()
-            self.db.commit()
+            # 等待所有任务完成
+            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", "")
+
+
+# 在模块末尾导入 SessionLocal（避免循环导入）
+from ..database import SessionLocal