first

domain-api/app/services/juming_task_service.py

@@ -0,0 +1,206 @@
+from __future__ import annotations
+
+import threading
+from datetime import datetime
+from uuid import uuid4
+
+from app.core.files import load_juming_records, save_juming_records
+from app.services.juming_service import TaskStoppedError, crawl_juming
+
+
+_JUMING_TASK_LOCK = threading.Lock()
+_MAX_LOG_LINES = 400
+
+
+def _now() -> str:
+    return datetime.now().isoformat(sep=" ", timespec="seconds")
+
+
+def list_juming_tasks() -> list[dict]:
+    return load_juming_records()
+
+
+def _save_tasks(tasks: list[dict]) -> None:
+    save_juming_records(tasks)
+
+
+def _append_log(task_id: str, message: str) -> None:
+    with _JUMING_TASK_LOCK:
+        tasks = load_juming_records()
+        target = next((item for item in tasks if item["task_id"] == task_id), None)
+        if not target:
+            return
+        logs = list(target.get("logs") or [])
+        logs.append(f"[{_now()}] {message}")
+        target["logs"] = logs[-_MAX_LOG_LINES:]
+        target["updated_at"] = _now()
+        _save_tasks(tasks)
+
+
+def _update_task(task_id: str, **patch: object) -> dict | None:
+    with _JUMING_TASK_LOCK:
+        tasks = load_juming_records()
+        target = next((item for item in tasks if item["task_id"] == task_id), None)
+        if not target:
+            return None
+        target.update(patch)
+        target["updated_at"] = _now()
+        _save_tasks(tasks)
+        return dict(target)
+
+
+def _is_stop_requested(task_id: str) -> bool:
+    tasks = load_juming_records()
+    target = next((item for item in tasks if item["task_id"] == task_id), None)
+    return bool(target and target.get("cancel_requested"))
+
+
+def _set_phase(task_id: str, phase: str, message: str | None = None) -> None:
+    phase_labels = {
+        "queued": "排队中",
+        "starting": "启动中",
+        "fetching": "抓取中",
+        "importing": "入库中",
+        "completed": "已完成",
+        "failed": "失败",
+        "stopping": "停止中",
+        "stopped": "已停止",
+    }
+    patch: dict[str, object] = {
+        "phase": phase,
+        "phase_label": phase_labels.get(phase, phase),
+    }
+    if message:
+        patch["message"] = message
+    _update_task(task_id, **patch)
+
+
+def _log_and_track_phase(task_id: str, message: str) -> None:
+    if "开始采集" in message or "正在抓取" in message:
+        _set_phase(task_id, "fetching", message)
+    elif "开始入库处理" in message or "入库完成" in message or "已存在域名" in message:
+        _set_phase(task_id, "importing", message)
+    _append_log(task_id, message)
+
+
+def _run_juming_task(task_id: str, payload: dict) -> None:
+    _update_task(task_id, status="running", started_at=_now(), message="聚名采集任务开始执行", cancel_requested=False)
+    _set_phase(task_id, "starting", "聚名采集任务开始执行")
+    _append_log(task_id, "任务已启动，正在准备读取聚名登录态")
+    try:
+        result = crawl_juming(
+            payload,
+            log=lambda message: _log_and_track_phase(task_id, str(message)),
+            should_stop=lambda: _is_stop_requested(task_id),
+        )
+        stats = result.get("stats", {})
+        _update_task(
+            task_id,
+            status="completed",
+            completed_at=_now(),
+            result=result,
+            cancel_requested=False,
+            message=(
+                f"采集完成：抓取 {result.get('domains_found', 0)} 个域名，"
+                f"新增 {stats.get('added', 0)}，已存在 {stats.get('exists', 0)}，无效 {stats.get('invalid', 0)}"
+            ),
+        )
+        _set_phase(task_id, "completed")
+        _append_log(task_id, "任务执行完成")
+    except TaskStoppedError as exc:
+        _update_task(
+            task_id,
+            status="stopped",
+            completed_at=_now(),
+            cancel_requested=False,
+            message=str(exc),
+        )
+        _set_phase(task_id, "stopped")
+        _append_log(task_id, "任务已按请求停止")
+    except Exception as exc:
+        _update_task(
+            task_id,
+            status="failed",
+            completed_at=_now(),
+            cancel_requested=False,
+            message=f"采集失败：{exc}",
+        )
+        _set_phase(task_id, "failed")
+        _append_log(task_id, f"任务执行失败：{exc}")
+
+
+def create_juming_task(payload: dict) -> dict:
+    task_id = uuid4().hex
+    mode = str(payload.get("mode") or "delete_list").strip() or "delete_list"
+    record = {
+        "task_id": task_id,
+        "mode": mode,
+        "payload": dict(payload or {}),
+        "status": "queued",
+        "phase": "queued",
+        "phase_label": "排队中",
+        "cancel_requested": False,
+        "message": "采集任务已创建，等待后台执行",
+        "created_at": _now(),
+        "updated_at": _now(),
+        "started_at": "",
+        "completed_at": "",
+        "result": None,
+        "logs": [f"[{_now()}] 已创建聚名采集任务，等待后台执行"],
+    }
+
+    with _JUMING_TASK_LOCK:
+        tasks = load_juming_records()
+        tasks.insert(0, record)
+        _save_tasks(tasks)
+
+    worker = threading.Thread(target=_run_juming_task, args=(task_id, dict(payload or {})), daemon=True)
+    worker.start()
+    return record
+
+
+def retry_juming_task(task_id: str) -> dict:
+    with _JUMING_TASK_LOCK:
+        tasks = load_juming_records()
+        target = next((item for item in tasks if item["task_id"] == task_id), None)
+        if not target:
+            raise ValueError("聚名采集任务不存在")
+        if target.get("status") == "running":
+            raise ValueError("聚名采集任务正在运行，不能重复执行")
+        target["status"] = "queued"
+        target["phase"] = "queued"
+        target["phase_label"] = "排队中"
+        target["cancel_requested"] = False
+        target["message"] = "采集任务已重新加入队列"
+        target["started_at"] = ""
+        target["completed_at"] = ""
+        target["updated_at"] = _now()
+        target["result"] = None
+        target["logs"] = [f"[{_now()}] 已重新加入队列，等待后台执行"]
+        payload = dict(target.get("payload") or {})
+        record = dict(target)
+        _save_tasks(tasks)
+
+    worker = threading.Thread(target=_run_juming_task, args=(task_id, payload), daemon=True)
+    worker.start()
+    return record
+
+
+def request_stop_juming_task(task_id: str) -> dict:
+    with _JUMING_TASK_LOCK:
+        tasks = load_juming_records()
+        target = next((item for item in tasks if item["task_id"] == task_id), None)
+        if not target:
+            raise ValueError("聚名采集任务不存在")
+        if target.get("status") in {"completed", "failed", "stopped"}:
+            raise ValueError("当前任务已结束，无需停止")
+        target["cancel_requested"] = True
+        target["phase"] = "stopping"
+        target["phase_label"] = "停止中"
+        target["message"] = "已发送停止请求，等待当前步骤安全退出"
+        target["updated_at"] = _now()
+        logs = list(target.get("logs") or [])
+        logs.append(f"[{_now()}] 已收到停止请求，等待当前步骤安全退出")
+        target["logs"] = logs[-_MAX_LOG_LINES:]
+        _save_tasks(tasks)
+        return dict(target)