From 96eb92710bba8b1f0239433c666e8c2ed6a240e7 Mon Sep 17 00:00:00 2001 From: Curry Date: Sun, 19 Jul 2026 01:03:45 +0800 Subject: [PATCH] fix(executor): retain task refs + global pipeline semaphore + off-loop broker dispatch [C5,C6,C9] MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - C5: LocalExecutor.dispatch_collection/dispatch_scheduled_collection now hold a strong reference to their asyncio.Task (in new _collection_tasks / _scheduled_collection_tasks dicts, popped on completion) instead of discarding the return of create_task, which asyncio may otherwise GC mid-flight — same fix already in place for dispatch_acquisition. dispatch_scheduled_collection also gains dispatch_acquisition's skip-if-inflight check, since schedule_id (unlike task_id) is not unique per tick. - C6: local.py gains a process-wide asyncio.Semaphore bounding concurrently RUNNING pipeline executions, independent of the per-domain cap in pipeline/domain_limiter.py. The limit is configurable via the new Settings.local_max_concurrent_pipelines field (env LOCAL_MAX_CONCURRENT_PIPELINES, default 8). The semaphore wraps pipeline execution (inside _run_collection/_run_scheduled), not dispatch, so dispatch_collection/dispatch_scheduled_collection stay fast and non-blocking; a task queued behind the semaphore is already referenced per the C5 fix, so it can't be GC'd while waiting. - C9: CeleryExecutor.dispatch_collection/dispatch_scheduled_collection wrap their apply_async calls in asyncio.to_thread so the synchronous broker round-trip no longer runs inline on the event loop. Tests extended in tests/unit/test_executor.py for all three findings. --- backend/config.py | 9 ++ backend/executor/celery_exec.py | 17 +- backend/executor/local.py | 81 +++++++++- tests/unit/test_executor.py | 269 ++++++++++++++++++++++++++++++++ 4 files changed, 366 insertions(+), 10 deletions(-) diff --git a/backend/config.py b/backend/config.py index 40bef82e..a558e45c 100644 --- a/backend/config.py +++ b/backend/config.py @@ -24,6 +24,15 @@ class Settings(BaseSettings): # Task execution mode: "local" (in-process asyncio) or "celery" (distributed) task_executor: Literal["local", "celery"] = "local" + # AUDIT C6: process-wide cap on concurrently-RUNNING pipeline executions in + # the local (in-process asyncio) executor — independent of the per-domain + # cap in pipeline/domain_limiter.py. Bounds how many schedules firing on + # the same tick plus manual/webhook triggers can drive Chrome/opencli + # subprocesses at once on the one event loop. Only meaningful when + # task_executor="local" (celery fans out across worker processes + # instead). Env: LOCAL_MAX_CONCURRENT_PIPELINES. + local_max_concurrent_pipelines: int = 8 + # Collection orchestrator: # admin — API内置 scheduler.py / Celery Beat 驱动定时采集（默认） # iii — III engine + schedule-bootstrap 驱动 cron；API 仅保留 UI/手动任务 diff --git a/backend/executor/celery_exec.py b/backend/executor/celery_exec.py index d86e5e34..d9ad4448 100644 --- a/backend/executor/celery_exec.py +++ b/backend/executor/celery_exec.py @@ -1,5 +1,6 @@ """Celery-based distributed executor.""" +import asyncio import logging from backend.executor.base import AbstractExecutor @@ -29,8 +30,14 @@ async def cancel_acquisition(self, execution_id: str) -> None: async def dispatch_collection(self, task_id: str, parameters: dict) -> dict: from backend.worker.tasks import run_collection - result = run_collection.apply_async( - kwargs={"task_id": task_id, "parameters": parameters} + + # AUDIT C9: apply_async is a synchronous broker round-trip (network + # I/O to Redis/RabbitMQ). Called directly here it would block this + # event loop for every dispatch; to_thread moves the round-trip off + # the loop. + result = await asyncio.to_thread( + run_collection.apply_async, + kwargs={"task_id": task_id, "parameters": parameters}, ) return {"task_id": task_id, "celery_task_id": result.id} @@ -38,6 +45,8 @@ async def dispatch_scheduled_collection( self, schedule_id: str, source_id: str, parameters: dict ) -> None: from backend.worker.tasks import run_scheduled_collection - run_scheduled_collection.apply_async( - kwargs={"schedule_id": schedule_id, "source_id": source_id, "parameters": parameters} + + await asyncio.to_thread( + run_scheduled_collection.apply_async, + kwargs={"schedule_id": schedule_id, "source_id": source_id, "parameters": parameters}, ) diff --git a/backend/executor/local.py b/backend/executor/local.py index 6585afe2..9b9ec9df 100644 --- a/backend/executor/local.py +++ b/backend/executor/local.py @@ -3,10 +3,29 @@ import asyncio import logging +from backend.config import get_settings from backend.executor.base import AbstractExecutor logger = logging.getLogger(__name__) +# AUDIT C6: process-wide cap on concurrently-RUNNING pipeline executions +# (manual/webhook dispatch_collection + every dispatch_scheduled_collection), +# independent of the per-domain cap in pipeline/domain_limiter.py. Keyed by +# event loop id, mirroring domain_limiter._semaphore, so a semaphore is never +# reused across event loops (production runs one loop and shares correctly; +# tests get a fresh loop each and never touch a stale entry). +_pipeline_semaphores: dict[int, asyncio.Semaphore] = {} + + +def _pipeline_semaphore() -> asyncio.Semaphore: + loop = asyncio.get_running_loop() + key = id(loop) + sem = _pipeline_semaphores.get(key) + if sem is None: + sem = asyncio.Semaphore(get_settings().local_max_concurrent_pipelines) + _pipeline_semaphores[key] = sem + return sem + def _log_task_exception(task: asyncio.Task) -> None: """Log any unhandled exception from a background asyncio task.""" @@ -21,6 +40,12 @@ def _log_task_exception(task: asyncio.Task) -> None: class LocalExecutor(AbstractExecutor): def __init__(self) -> None: self._acquisition_tasks: dict[str, asyncio.Task[None]] = {} + # AUDIT C5: distinct dicts from _acquisition_tasks — task_id (manual/ + # webhook dispatch) and schedule_id (scheduled dispatch) are different + # id spaces, so each gets its own strong-reference table instead of + # sharing one dict with two key schemes. + self._collection_tasks: dict[str, asyncio.Task[dict]] = {} + self._scheduled_collection_tasks: dict[str, asyncio.Task[dict]] = {} async def dispatch_acquisition(self, execution_id: str) -> None: from backend.acquisition.runner import run_acquisition_execution @@ -47,15 +72,59 @@ async def cancel_acquisition(self, execution_id: str) -> None: except asyncio.CancelledError: pass - async def dispatch_collection(self, task_id: str, parameters: dict) -> dict: + async def _run_collection(self, task_id: str, parameters: dict) -> dict: from backend.pipeline.runner import run_collection_pipeline - t = asyncio.create_task(run_collection_pipeline(task_id, parameters)) - t.add_done_callback(_log_task_exception) + + # AUDIT C6: the semaphore wraps execution, not dispatch — the task + # blocks on it here, inside itself, so dispatch_collection stays fast + # and non-blocking. + async with _pipeline_semaphore(): + return await run_collection_pipeline(task_id, parameters) + + async def _run_scheduled( + self, schedule_id: str, source_id: str, parameters: dict + ) -> dict: + from backend.pipeline.runner import run_scheduled_pipeline + + async with _pipeline_semaphore(): + return await run_scheduled_pipeline(schedule_id, source_id, parameters) + + async def dispatch_collection(self, task_id: str, parameters: dict) -> dict: + # AUDIT C5: hold a strong reference exactly like dispatch_acquisition + # does. task_id is unique per call (each manual/webhook trigger has + # its own CollectionTask row), so unlike dispatch_scheduled_collection + # below there is no skip-if-inflight check — just track-and-pop. The + # task is stored before it can even attempt to acquire the semaphore, + # so it stays referenced (and un-GC'd) for the whole time it's + # blocked waiting on _pipeline_semaphore(), not just while running. + task = asyncio.create_task(self._run_collection(task_id, parameters)) + task.add_done_callback(_log_task_exception) + self._collection_tasks[task_id] = task + task.add_done_callback( + lambda completed: self._collection_tasks.pop(task_id, None) + if self._collection_tasks.get(task_id) is completed + else None + ) return {"task_id": task_id} async def dispatch_scheduled_collection( self, schedule_id: str, source_id: str, parameters: dict ) -> None: - from backend.pipeline.runner import run_scheduled_pipeline - t = asyncio.create_task(run_scheduled_pipeline(schedule_id, source_id, parameters)) - t.add_done_callback(_log_task_exception) + # AUDIT C5: schedule_id has no per-run unique id at dispatch time (a + # schedule fires the same schedule_id on every tick), so this mirrors + # dispatch_acquisition's skip-if-inflight precedent — a schedule whose + # previous run hasn't finished yet is skipped rather than allowed to + # overlap itself. + current = self._scheduled_collection_tasks.get(schedule_id) + if current is not None and not current.done(): + return + task = asyncio.create_task( + self._run_scheduled(schedule_id, source_id, parameters) + ) + task.add_done_callback(_log_task_exception) + self._scheduled_collection_tasks[schedule_id] = task + task.add_done_callback( + lambda completed: self._scheduled_collection_tasks.pop(schedule_id, None) + if self._scheduled_collection_tasks.get(schedule_id) is completed + else None + ) diff --git a/tests/unit/test_executor.py b/tests/unit/test_executor.py index 6afc78ca..1131d8e4 100644 --- a/tests/unit/test_executor.py +++ b/tests/unit/test_executor.py @@ -1,10 +1,26 @@ """Unit tests for backend executor modules.""" import asyncio +import threading from unittest.mock import AsyncMock, MagicMock, patch import pytest +from backend.executor import local as local_executor_module + + +@pytest.fixture(autouse=True) +def _clear_pipeline_semaphore_registry(): + """AUDIT C6's semaphore registry is a module-level dict keyed by event + loop id (mirrors backend.pipeline.domain_limiter._semaphores). Clear it + around every test in this file so a stale entry from a prior test's + (closed) event loop can never be reused just because CPython recycled + its id() for a new loop object.""" + local_executor_module._pipeline_semaphores.clear() + yield + local_executor_module._pipeline_semaphores.clear() + + # ── LocalExecutor ────────────────────────────────────────────────────────────── @pytest.mark.asyncio @@ -74,6 +90,157 @@ async def test_local_executor_dispatch_scheduled_collection(): await asyncio.sleep(0) +# ── LocalExecutor: AUDIT C5 (task reference retention) ───────────────────────── + +@pytest.mark.asyncio +async def test_local_executor_collection_task_retained_while_inflight_and_popped_after(): + """dispatch_collection's task is held in a strong-reference dict for as + long as the pipeline is running, and popped once it completes — the fix + for C5 (an unreferenced create_task can be GC'd mid-flight).""" + from backend.executor.local import LocalExecutor + + gate = asyncio.Event() + + async def fake_pipeline(task_id, parameters): + await gate.wait() + return {} + + executor = LocalExecutor() + with patch("backend.pipeline.runner.run_collection_pipeline", new=fake_pipeline): + result = await executor.dispatch_collection("task-abc", {}) + await asyncio.sleep(0) + + assert result == {"task_id": "task-abc"} + assert "task-abc" in executor._collection_tasks + task = executor._collection_tasks["task-abc"] + assert isinstance(task, asyncio.Task) + assert not task.done() + + gate.set() + await task + await asyncio.sleep(0) # let the pop done-callback run + + assert "task-abc" not in executor._collection_tasks + + +@pytest.mark.asyncio +async def test_local_executor_scheduled_collection_skips_when_prior_run_inflight(): + """A schedule whose previous run hasn't finished yet is skipped rather + than dispatched again — mirrors dispatch_acquisition's existing + skip-if-inflight precedent (AUDIT C5).""" + from backend.executor.local import LocalExecutor + + gate = asyncio.Event() + calls: list[tuple] = [] + + async def fake_scheduled_pipeline(schedule_id, source_id, parameters): + calls.append((schedule_id, source_id, parameters)) + await gate.wait() + return {} + + executor = LocalExecutor() + with patch("backend.pipeline.runner.run_scheduled_pipeline", new=fake_scheduled_pipeline): + await executor.dispatch_scheduled_collection("sched-1", "src-1", {"a": 1}) + await asyncio.sleep(0) + assert "sched-1" in executor._scheduled_collection_tasks + first_task = executor._scheduled_collection_tasks["sched-1"] + + # Second tick while the first run is still in-flight must be a no-op. + await executor.dispatch_scheduled_collection("sched-1", "src-1", {"a": 2}) + await asyncio.sleep(0) + assert executor._scheduled_collection_tasks["sched-1"] is first_task + + gate.set() + await first_task + await asyncio.sleep(0) + + assert len(calls) == 1 + assert calls[0] == ("sched-1", "src-1", {"a": 1}) + assert "sched-1" not in executor._scheduled_collection_tasks + + +@pytest.mark.asyncio +async def test_local_executor_scheduled_collection_dispatches_again_once_prior_run_done(): + """Skip-if-inflight is not a permanent lockout: once a run finishes, the + next tick for the same schedule_id dispatches a fresh task.""" + from backend.executor.local import LocalExecutor + + calls: list[str] = [] + + async def fake_scheduled_pipeline(schedule_id, source_id, parameters): + calls.append(schedule_id) + return {} + + executor = LocalExecutor() + with patch("backend.pipeline.runner.run_scheduled_pipeline", new=fake_scheduled_pipeline): + await executor.dispatch_scheduled_collection("sched-1", "src-1", {}) + await asyncio.sleep(0) + await executor.dispatch_scheduled_collection("sched-1", "src-1", {}) + await asyncio.sleep(0) + + assert calls == ["sched-1", "sched-1"] + + +# ── LocalExecutor: AUDIT C6 (global pipeline concurrency semaphore) ──────────── + +@pytest.mark.asyncio +async def test_local_executor_pipeline_semaphore_bounds_concurrency(monkeypatch): + """AUDIT C6: no more than settings.local_max_concurrent_pipelines + collection runs execute at once, even when more are dispatched than + that (mirrors backend/pipeline/domain_limiter's own peak-concurrency + test style: real tiny hold + asyncio.gather + peak check). + + AUDIT C5: every dispatched task — including the ones that will have to + wait on the semaphore — is referenced in _collection_tasks the instant + it's created, well before any of them finish.""" + from backend.config import get_settings + from backend.executor.local import LocalExecutor + + monkeypatch.setenv("LOCAL_MAX_CONCURRENT_PIPELINES", "2") + get_settings.cache_clear() + try: + state = {"cur": 0, "peak": 0} + + async def fake_pipeline(task_id, parameters): + state["cur"] += 1 + state["peak"] = max(state["peak"], state["cur"]) + await asyncio.sleep(0.02) + state["cur"] -= 1 + return {} + + executor = LocalExecutor() + with patch("backend.pipeline.runner.run_collection_pipeline", new=fake_pipeline): + for i in range(6): + await executor.dispatch_collection(f"task-{i}", {}) + + # All 6 are referenced immediately, before any of them have run + # far enough to finish — none can be silently GC'd (C5), whether + # they're running or still queued behind the semaphore. + assert len(executor._collection_tasks) == 6 + + await asyncio.gather(*list(executor._collection_tasks.values())) + + assert state["peak"] == 2 + finally: + get_settings.cache_clear() + + +# ── Settings: local_max_concurrent_pipelines (AUDIT C6) ──────────────────────── + +def test_settings_local_max_concurrent_pipelines_defaults_to_eight(): + from backend.config import Settings + + assert Settings().local_max_concurrent_pipelines == 8 + + +def test_settings_local_max_concurrent_pipelines_honors_env_override(monkeypatch): + from backend.config import Settings + + monkeypatch.setenv("LOCAL_MAX_CONCURRENT_PIPELINES", "3") + + assert Settings().local_max_concurrent_pipelines == 3 + + def test_log_task_exception_logs_error(): """_log_task_exception logs when task has an unhandled exception.""" from backend.executor.local import _log_task_exception @@ -197,3 +364,105 @@ async def test_celery_executor_dispatch_scheduled_collection(): mock_task.apply_async.assert_called_once_with( kwargs={"schedule_id": "sched-1", "source_id": "src-1", "parameters": {"k": "v"}} ) + + +# ── CeleryExecutor: AUDIT C9 (apply_async off the event loop) ────────────────── + +@pytest.mark.asyncio +async def test_celery_executor_dispatch_collection_offloads_apply_async_to_thread(): + """apply_async is invoked through asyncio.to_thread rather than directly, + so the broker round-trip never runs inline on the event loop (C9).""" + from backend.executor.celery_exec import CeleryExecutor + + executor = CeleryExecutor() + mock_result = MagicMock() + mock_result.id = "celery-task-abc" + mock_task = MagicMock() + mock_task.apply_async = MagicMock(return_value=mock_result) + + with ( + patch("backend.worker.tasks.run_collection", mock_task), + patch( + "backend.executor.celery_exec.asyncio.to_thread", + new=AsyncMock(side_effect=lambda fn, **kw: fn(**kw)), + ) as to_thread_mock, + ): + result = await executor.dispatch_collection("task-999", {"p": 1}) + + to_thread_mock.assert_called_once_with( + mock_task.apply_async, kwargs={"task_id": "task-999", "parameters": {"p": 1}} + ) + assert result == {"task_id": "task-999", "celery_task_id": "celery-task-abc"} + + +@pytest.mark.asyncio +async def test_celery_executor_dispatch_scheduled_collection_offloads_apply_async_to_thread(): + """Same off-loop guarantee for the scheduled-dispatch path (C9).""" + from backend.executor.celery_exec import CeleryExecutor + + executor = CeleryExecutor() + mock_task = MagicMock() + mock_task.apply_async = MagicMock(return_value=MagicMock()) + + with ( + patch("backend.worker.tasks.run_scheduled_collection", mock_task), + patch( + "backend.executor.celery_exec.asyncio.to_thread", + new=AsyncMock(side_effect=lambda fn, **kw: fn(**kw)), + ) as to_thread_mock, + ): + await executor.dispatch_scheduled_collection("sched-1", "src-1", {"k": "v"}) + + to_thread_mock.assert_called_once_with( + mock_task.apply_async, + kwargs={"schedule_id": "sched-1", "source_id": "src-1", "parameters": {"k": "v"}}, + ) + + +@pytest.mark.asyncio +async def test_celery_executor_dispatch_collection_does_not_block_event_loop(): + """AUDIT C9, end-to-end with the real asyncio.to_thread: while + apply_async is 'in flight' (blocking its worker thread), other + coroutines scheduled on this same loop keep making progress instead of + stalling behind the broker round-trip.""" + from backend.executor.celery_exec import CeleryExecutor + + executor = CeleryExecutor() + release = threading.Event() + + def slow_apply_async(**kwargs): + # Stands in for a slow/blocking broker round-trip. A thread-safe + # (not asyncio) primitive because this runs in a to_thread worker + # thread, not on the event loop. + if not release.wait(timeout=2): + raise AssertionError("test never released the blocking call") + result = MagicMock() + result.id = "celery-task-abc" + return result + + mock_task = MagicMock() + mock_task.apply_async = slow_apply_async + + probe_ticks = 0 + + async def probe(): + nonlocal probe_ticks + for _ in range(20): + await asyncio.sleep(0.01) + probe_ticks += 1 + + with patch("backend.worker.tasks.run_collection", mock_task): + dispatch_task = asyncio.create_task(executor.dispatch_collection("task-999", {})) + probe_task = asyncio.create_task(probe()) + + # Give dispatch_collection a moment to reach and block inside + # to_thread's worker thread. + await asyncio.sleep(0.05) + ticks_while_blocked = probe_ticks + release.set() + + result = await dispatch_task + await probe_task + + assert ticks_while_blocked > 0, "event loop made no progress while apply_async was in flight" + assert result == {"task_id": "task-999", "celery_task_id": "celery-task-abc"}