diff --git a/src/backend/base/langflow/services/background_execution/executor.py b/src/backend/base/langflow/services/background_execution/executor.py index 8c4fb99eb6..fa485a8d79 100644 --- a/src/backend/base/langflow/services/background_execution/executor.py +++ b/src/backend/base/langflow/services/background_execution/executor.py @@ -43,15 +43,36 @@ class InProcessExecutor: async def stop(self) -> None: self._closed = True + # Cancel the in-flight JOB tasks directly first, then await them. A job + # that catches its cancel (a user-stop reconciles to CANCELLED) would + # otherwise leave the worker's ``await task`` in cancellation limbo, so we + # do not rely on cancellation propagating through the worker. Awaiting the + # job tasks here lets each one's shielded terminal reconcile finish BEFORE + # teardown — without it a reconcile write races a closing DB engine and a + # "Task was destroyed but it is pending" warning is logged. + # return_exceptions swallows the CancelledError each task raises so one + # cancel cannot mask the others. + # Cancel the in-flight JOB tasks directly first, then await them. A job + # that catches its cancel (a user-stop reconciles to CANCELLED) would + # otherwise leave the worker's ``await task`` in cancellation limbo, so we + # do not rely on cancellation propagating through the worker. Awaiting the + # job tasks here lets each one's shielded terminal reconcile finish BEFORE + # teardown — without it a reconcile write races a closing DB engine and a + # "Task was destroyed but it is pending" warning is logged. + # return_exceptions swallows the CancelledError each task raises so one + # cancel cannot mask the others. + in_flight = list(self._in_flight.values()) + for task in in_flight: + task.cancel() + if in_flight: + await asyncio.gather(*in_flight, return_exceptions=True) + self._in_flight.clear() for task in self._workers: task.cancel() for task in self._workers: with contextlib.suppress(asyncio.CancelledError): await task self._workers.clear() - for task in list(self._in_flight.values()): - task.cancel() - self._in_flight.clear() async def submit(self, key: str, coro_factory: CoroFactory) -> None: """Enqueue a job. A free worker picks it up and invokes the factory.""" diff --git a/src/backend/tests/unit/background_execution/test_executor_hard_proof.py b/src/backend/tests/unit/background_execution/test_executor_hard_proof.py index ed1382452f..043ec3a480 100644 --- a/src/backend/tests/unit/background_execution/test_executor_hard_proof.py +++ b/src/backend/tests/unit/background_execution/test_executor_hard_proof.py @@ -93,6 +93,49 @@ async def test_pool_survives_repeated_job_cancellation(monkeypatch): await executor.stop() +async def test_stop_awaits_in_flight_tasks(): + """``stop()`` leaves every in-flight task ``done()`` when it returns. + + A real job cancelled mid-flight reconciles its terminal status inside a + shielded ``finally`` (a DB write). ``stop()`` cancels the in-flight tasks and + then ``gather``s them so that work finishes before teardown, rather than + racing a closing DB engine and leaving a "Task was destroyed but it is + pending" warning. We capture the in-flight task handle and assert it (and its + multi-turn shielded finalizer) is done the instant ``stop()`` returns — the + awaited-teardown contract the gather guarantees. + """ + running = asyncio.Event() + finalized = asyncio.Event() + + async def _job() -> None: + try: + running.set() + await asyncio.sleep(3600) # block until cancelled + finally: + # The runner's terminal reconcile is a shielded await (a DB write) + # that spans several event-loop turns. Model it so the task only + # finishes after a real awaited finalizer, not the instant the + # cancel is delivered. + for _ in range(10): + await asyncio.shield(asyncio.sleep(0.02)) + finalized.set() + + executor = InProcessExecutor(max_concurrency=1) + await executor.start() + await executor.submit("job", _job) + await asyncio.wait_for(running.wait(), timeout=2.0) + + # Capture the in-flight task before stop() clears the registry. + in_flight = executor._in_flight["job"] + + await executor.stop() + + # If stop() awaited the in-flight task, it (and its shielded finalizer) are + # done the moment stop() returns. A non-awaiting stop() leaves it pending. + assert in_flight.done(), "stop() returned without awaiting the in-flight task" + assert finalized.is_set(), "shielded finalizer did not complete before stop() returned" + + async def test_pool_survives_job_exception(): """A job raising a non-CancelledError must not kill the worker.""" executor = InProcessExecutor(max_concurrency=1) diff --git a/src/backend/tests/unit/background_execution/test_facade_hard_proof.py b/src/backend/tests/unit/background_execution/test_facade_hard_proof.py index 1c532664de..90601c55d5 100644 --- a/src/backend/tests/unit/background_execution/test_facade_hard_proof.py +++ b/src/backend/tests/unit/background_execution/test_facade_hard_proof.py @@ -381,6 +381,59 @@ async def test_events_replay_frames_are_sse_framed(hard_proof_job_service): assert f"id: {seq}".encode() in frame, f"replayed frame missing id: {frame!r}" +async def test_executor_stop_applies_terminal_reconcile(hard_proof_job_service): + """``executor.stop()`` lets an in-flight stopped job's reconcile land. + + Drives the REAL runner on the executor. The job blocks mid-run; a STOP signal + is written, then ``executor.stop()`` tears the pool down. stop() cancels and + gathers the in-flight task so the runner's shielded terminal reconcile applies + BEFORE it returns, and the durable row reads CANCELLED. The job swallows its + cancellation (user-stop path): stop() cancels the job task directly and + gathers it rather than waiting on the worker's absorbed-cancel ``await``, + which leaves the worker in cancellation limbo and stalls teardown. The + ``wait_for`` budget plus pytest-timeout turn that stall into a failure. Real + SQLite and Postgres. + """ + import asyncio + + from langflow.services.background_execution.executor import InProcessExecutor + + job_service = hard_proof_job_service + job_id = uuid4() + await job_service.create_job(job_id=job_id, flow_id=uuid4(), user_id=uuid4()) + + started = asyncio.Event() + release = asyncio.Event() + + async def blocking_source(**_kwargs) -> AsyncIterator[tuple[bytes, str]]: + yield _frame("build_start", {}) + started.set() + await release.wait() # block until cancelled + yield _frame("end", {}) # unreachable + + bus = InMemoryLiveBus() + runner = _runner(job_service, bus, job_id, blocking_source) + executor = InProcessExecutor(max_concurrency=1) + await executor.start() + + async def _coro() -> None: + await runner.run(job_id=job_id, source_kwargs={}) + + await executor.submit(str(job_id), _coro) + await asyncio.wait_for(started.wait(), timeout=5.0) + + await job_service.write_signal(job_id, SignalType.STOP) + # The fixed stop() (cancel+gather job tasks first) returns in well under a + # second. The original ordering (await workers first) leaves the worker in + # cancellation limbo behind the job's swallowed cancel and only unblocks once + # the reconcile's DB writes drain through lock backoff (many seconds), so a + # tight budget here reliably distinguishes the two. + await asyncio.wait_for(executor.stop(), timeout=5.0) + + job = await job_service.get_job_by_job_id(job_id) + assert job.status == JobStatus.CANCELLED + + async def test_stop_signal_is_marked_consumed(hard_proof_job_service): """When the runner acts on a STOP, the signal row is stamped ``consumed_at``.