From fed7f56ad066a82324c8143b2c80f72d96df1fea Mon Sep 17 00:00:00 2001 From: oliver Date: Thu, 28 May 2026 10:22:11 +0800 Subject: [PATCH] fix: collection first-run pending and module tab cross-nav Defer collection scheduling to BackgroundTasks after commit; claim pending runs atomically with retries. Only handle focus-module broadcasts for the matching moduleId so opening one module does not rewrite other tabs. Co-authored-by: Cursor --- netx_api/collection_router.py | 48 ++++++++++++--- netx_api/collection_service.py | 37 ++++++++---- netx_api/ne_collect_runner.py | 106 +++++++++++++++++++++++++++------ web/src/utils/moduleWindows.ts | 19 +++--- web/src/utils/tabChannel.ts | 2 + 5 files changed, 169 insertions(+), 43 deletions(-) diff --git a/netx_api/collection_router.py b/netx_api/collection_router.py index 2c297f8..e67cdbd 100644 --- a/netx_api/collection_router.py +++ b/netx_api/collection_router.py @@ -1,6 +1,6 @@ from __future__ import annotations -from fastapi import APIRouter, Depends, Query +from fastapi import APIRouter, BackgroundTasks, Depends, Query from fastapi.responses import FileResponse, Response from sqlalchemy.orm import Session @@ -20,6 +20,7 @@ from .collection_service import ( from .collection_schemas import CollectionJobCreate from .db import get_db from .models import NeCollectionRun +from .ne_collect_runner import dispatch_collection_runs router = APIRouter(prefix="/v1/ne-collections", tags=["ne-collections"]) @@ -34,8 +35,19 @@ def api_eligible_ne( @router.post("") -def api_create_collection(body: CollectionJobCreate, db: Session = Depends(get_db)): - return create_and_start_collection(db, body).model_dump() +def api_create_collection( + body: CollectionJobCreate, + background_tasks: BackgroundTasks, + db: Session = Depends(get_db), +): + out, payload = create_and_start_collection(db, body) + background_tasks.add_task( + dispatch_collection_runs, + payload["job_id"], + payload["run_ids"], + payload["commands"], + ) + return out.model_dump() @router.get("") @@ -92,13 +104,35 @@ def api_pause_collection(job_id: str, db: Session = Depends(get_db)): @router.post("/{job_id}/restart") -def api_restart_collection(job_id: str, db: Session = Depends(get_db)): - return restart_collection_job(db, job_id).model_dump() +def api_restart_collection( + job_id: str, + background_tasks: BackgroundTasks, + db: Session = Depends(get_db), +): + out, payload = restart_collection_job(db, job_id) + background_tasks.add_task( + dispatch_collection_runs, + payload["job_id"], + payload["run_ids"], + payload["commands"], + ) + return out.model_dump() @router.post("/{job_id}/retry-failed") -def api_retry_failed_collection(job_id: str, db: Session = Depends(get_db)): - return retry_failed_collection_job(db, job_id).model_dump() +def api_retry_failed_collection( + job_id: str, + background_tasks: BackgroundTasks, + db: Session = Depends(get_db), +): + out, payload = retry_failed_collection_job(db, job_id) + background_tasks.add_task( + dispatch_collection_runs, + payload["job_id"], + payload["run_ids"], + payload["commands"], + ) + return out.model_dump() @router.delete("/{job_id}") diff --git a/netx_api/collection_service.py b/netx_api/collection_service.py index ec46a53..9e41382 100644 --- a/netx_api/collection_service.py +++ b/netx_api/collection_service.py @@ -7,7 +7,7 @@ import shutil import zipfile from datetime import datetime from pathlib import Path -from typing import Any +from typing import Any, TypedDict from fastapi import HTTPException from sqlalchemy import func, or_ @@ -21,12 +21,17 @@ from .collection_job_state import ( _sync_job_counts, ) from .collection_schemas import CollectionJobCreate, CollectionJobOut, CollectionRunOut -from .ne_collect_runner import schedule_collection_runs from .ne_collection_paths import clear_run_output_files, collection_data_root _log = logging.getLogger("netx.collection") +class CollectionSchedulePayload(TypedDict): + job_id: str + run_ids: list[str] + commands: list[str] + + def _now() -> datetime: return datetime.now() @@ -130,7 +135,9 @@ def list_eligible_ne(db: Session, *, page: int = 1, page_size: int = 200) -> dic return {"total": total, "page": page, "page_size": page_size, "items": items} -def create_and_start_collection(db: Session, body: CollectionJobCreate) -> CollectionJobOut: +def create_and_start_collection( + db: Session, body: CollectionJobCreate +) -> tuple[CollectionJobOut, CollectionSchedulePayload]: commands = _parse_commands(body.commands) if not commands: raise HTTPException(status_code=400, detail="commands_empty") @@ -184,8 +191,12 @@ def create_and_start_collection(db: Session, body: CollectionJobCreate) -> Colle db.commit() db.refresh(job) - schedule_collection_runs(str(job.id), run_ids, commands) - return job_to_out(job, output_count=0) + payload: CollectionSchedulePayload = { + "job_id": str(job.id), + "run_ids": run_ids, + "commands": commands, + } + return job_to_out(job, output_count=0), payload def list_collection_jobs(db: Session, *, page: int = 1, page_size: int = 20) -> dict[str, Any]: @@ -319,7 +330,7 @@ def _start_job_retry( commands: list[str], *, reset_all_counts: bool, -) -> CollectionJobOut: +) -> tuple[CollectionJobOut, CollectionSchedulePayload]: if not retry_ids: raise HTTPException(status_code=400, detail="collection_nothing_to_retry") now = _now() @@ -336,11 +347,15 @@ def _start_job_retry( job.last_run_at = now db.commit() db.refresh(job) - schedule_collection_runs(job_id, retry_ids, commands) - return job_to_out(job, output_count=_output_count_for_job(db, job_id)) + payload: CollectionSchedulePayload = { + "job_id": job_id, + "run_ids": retry_ids, + "commands": commands, + } + return job_to_out(job, output_count=_output_count_for_job(db, job_id)), payload -def restart_collection_job(db: Session, job_id: str) -> CollectionJobOut: +def restart_collection_job(db: Session, job_id: str) -> tuple[CollectionJobOut, CollectionSchedulePayload]: job = db.get(NeCollectionJob, job_id) if not job: raise HTTPException(status_code=404, detail="collection_job_not_found") @@ -361,7 +376,9 @@ def restart_collection_job(db: Session, job_id: str) -> CollectionJobOut: return _start_job_retry(db, job, job_id, retry_ids, commands, reset_all_counts=True) -def retry_failed_collection_job(db: Session, job_id: str) -> CollectionJobOut: +def retry_failed_collection_job( + db: Session, job_id: str +) -> tuple[CollectionJobOut, CollectionSchedulePayload]: job = db.get(NeCollectionJob, job_id) if not job: raise HTTPException(status_code=404, detail="collection_job_not_found") diff --git a/netx_api/ne_collect_runner.py b/netx_api/ne_collect_runner.py index 172f243..4e1761d 100644 --- a/netx_api/ne_collect_runner.py +++ b/netx_api/ne_collect_runner.py @@ -2,6 +2,7 @@ from __future__ import annotations import logging import re +import time import traceback from concurrent.futures import ThreadPoolExecutor, TimeoutError as FuturesTimeout from datetime import datetime @@ -97,31 +98,70 @@ def _update_run(run_id: str, **fields: Any) -> None: db.close() -def _collection_aborted(job_id: str, run_id: str) -> bool: +def _job_is_paused(job_id: str) -> bool: db = SessionLocal() try: job = db.get(NeCollectionJob, job_id) - run = db.get(NeCollectionRun, run_id) - if not job or not run: - return True - if str(job.status or "") == "paused": - if str(run.status or "") == "pending": - run.status = "cancelled" - run.message = "paused" - run.ended_at = datetime.now() - db.commit() - return True - if str(run.status or "") in ("cancelled", "success", "fail"): - return True - return False + return not job or str(job.status or "") == "paused" finally: db.close() -def _run_single(job_id: str, run_id: str, commands: list[str]) -> None: - if _collection_aborted(job_id, run_id): +def _claim_run(job_id: str, run_id: str) -> bool: + """Atomically move a pending run to running; retry while DB rows become visible.""" + for attempt in range(10): db = SessionLocal() try: + run = db.get(NeCollectionRun, run_id) + job = db.get(NeCollectionJob, job_id) + if not run or not job: + time.sleep(0.05 * (attempt + 1)) + continue + job_status = str(job.status or "") + run_status = str(run.status or "") + if job_status == "paused": + if run_status == "pending": + run.status = "cancelled" + run.message = "paused" + run.ended_at = datetime.now() + db.commit() + return False + if job_status != "running": + return False + if run_status == "running": + return True + if run_status != "pending": + return False + run.status = "running" + run.message = "collecting" + run.started_at = datetime.now() + db.commit() + return True + finally: + db.close() + time.sleep(0.05 * (attempt + 1)) + return False + + +def _run_single(job_id: str, run_id: str, commands: list[str]) -> None: + if not _claim_run(job_id, run_id): + db = SessionLocal() + try: + run = db.get(NeCollectionRun, run_id) + st = str(run.status or "") if run else "" + if st in ("cancelled", "success", "fail"): + sync_job_progress(db, job_id) + finalize_collection_job(db, job_id) + return + if st == "pending": + _log.warning("collection claim failed job=%s run=%s", job_id, run_id) + _update_run( + run_id, + status="fail", + message="collection_claim_failed", + ended_at=datetime.now(), + ) + sync_job_progress(db, job_id) finalize_collection_job(db, job_id) finally: db.close() @@ -135,8 +175,9 @@ def _run_single(job_id: str, run_id: str, commands: list[str]) -> None: if not ne: _update_run(run_id, status="fail", message="managed_ne_not_found", ended_at=datetime.now()) return - started_at = datetime.now() - _update_run(run_id, status="running", message="collecting", started_at=started_at) + if _job_is_paused(job_id): + _update_run(run_id, status="cancelled", message="paused", ended_at=datetime.now()) + return try: creds = get_device_credentials(ne) output = _collect_with_timeout(creds, commands) @@ -174,10 +215,37 @@ def _run_single(job_id: str, run_id: str, commands: list[str]) -> None: db2.close() +def _run_collect_safe(job_id: str, run_id: str, commands: list[str]) -> None: + try: + _run_single(job_id, run_id, commands) + except Exception: + _log.exception("collection task crashed job=%s run=%s", job_id, run_id) + _update_run( + run_id, + status="fail", + message="collection_worker_crashed", + ended_at=datetime.now(), + ) + db = SessionLocal() + try: + sync_job_progress(db, job_id) + finalize_collection_job(db, job_id) + finally: + db.close() + + def schedule_collection_runs(job_id: str, run_ids: list[str], commands: list[str]) -> int: pool = _executor_pool() + cmd_list = list(commands) submitted = 0 for run_id in run_ids: - pool.submit(_run_single, job_id, run_id, list(commands)) + pool.submit(_run_collect_safe, job_id, str(run_id), cmd_list) submitted += 1 + if submitted: + _log.info("scheduled collection job=%s runs=%s", job_id, submitted) return submitted + + +def dispatch_collection_runs(job_id: str, run_ids: list[str], commands: list[str]) -> int: + """Entry point for FastAPI BackgroundTasks after the request transaction commits.""" + return schedule_collection_runs(job_id, run_ids, commands) diff --git a/web/src/utils/moduleWindows.ts b/web/src/utils/moduleWindows.ts index 7183dd4..717770e 100644 --- a/web/src/utils/moduleWindows.ts +++ b/web/src/utils/moduleWindows.ts @@ -36,13 +36,18 @@ export function registerModuleWindow(moduleId: string): () => void { /* ignore */ } - return registerBroadcastListener(MODULE_CHANNEL, moduleAckKey(moduleId), (data) => { - const targetPath = String(data.path || "").trim(); - if (targetPath && window.location.pathname !== targetPath) { - window.location.assign(targetPath); - } - window.focus(); - }); + return registerBroadcastListener( + MODULE_CHANNEL, + moduleAckKey(moduleId), + (data) => { + const targetPath = String(data.path || "").trim(); + if (targetPath && window.location.pathname !== targetPath) { + window.location.assign(targetPath); + } + window.focus(); + }, + (data) => data.moduleId === moduleId, + ); } export function openOrFocusModule({ moduleId, path }: ModuleWindowSpec): void { diff --git a/web/src/utils/tabChannel.ts b/web/src/utils/tabChannel.ts index 484521b..9fc1f5c 100644 --- a/web/src/utils/tabChannel.ts +++ b/web/src/utils/tabChannel.ts @@ -10,6 +10,7 @@ export function registerBroadcastListener( channelName: string, ackStorageKey: string, onMessage: (data: T) => void, + shouldHandle?: (data: T) => boolean, ): () => void { if (typeof BroadcastChannel === "undefined") { return () => undefined; @@ -19,6 +20,7 @@ export function registerBroadcastListener( channel.onmessage = (ev: MessageEvent) => { const data = ev.data; if (!data?.requestId) return; + if (shouldHandle && !shouldHandle(data)) return; try { sessionStorage.setItem(ackStorageKey, data.requestId); } catch {