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 _ACTIVE_TASK_IDS: set[str] = set() def _now() -> str: return datetime.now().isoformat(sep=" ", timespec="seconds") def list_juming_tasks() -> list[dict]: _cleanup_orphaned_tasks() 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 _cleanup_orphaned_tasks() -> None: with _JUMING_TASK_LOCK: tasks = load_juming_records() changed = False for item in tasks: status = str(item.get("status") or "").strip() task_id = str(item.get("task_id") or "").strip() if status == "running" and task_id and task_id not in _ACTIVE_TASK_IDS: item["status"] = "failed" item["phase"] = "failed" item["phase_label"] = "失败" item["cancel_requested"] = False item["message"] = "任务因服务重启或进程中断而终止，请重新执行" item["updated_at"] = _now() logs = list(item.get("logs") or []) logs.append(f"[{_now()}] 任务因服务重启或进程中断而终止，请重新执行") item["logs"] = logs[-_MAX_LOG_LINES:] changed = True if changed: _save_tasks(tasks) def _ensure_no_active_task() -> None: _cleanup_orphaned_tasks() tasks = load_juming_records() active = next( ( item for item in tasks if str(item.get("status") or "").strip() == "running" ), None, ) if active: raise ValueError(f"已有聚名采集任务正在运行：{active.get('task_id')}") 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 or "已存在域名" in message ): _set_phase(task_id, "importing", message) _append_log(task_id, message) def _run_juming_task(task_id: str, payload: dict) -> None: _ACTIVE_TASK_IDS.add(task_id) _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}") finally: _ACTIVE_TASK_IDS.discard(task_id) def create_juming_task(payload: dict) -> dict: _ensure_no_active_task() 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) _ACTIVE_TASK_IDS.add(task_id) 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: _ensure_no_active_task() 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)