From be43f2e546c70d71db04d050c5e9d9590eed5a1d Mon Sep 17 00:00:00 2001 From: omp Date: Fri, 7 Aug 2026 15:51:25 +1000 Subject: [PATCH] fix(worker): release operations when a task is cancelled or its terminal write fails Three paths let a task stop running while its async_operations row stays 'processing' owned by a live worker, permanently: 1. asyncio.CancelledError derives from BaseException, so it escapes every handler in _execute_task_inner; the done-callback then removes the task from _active_tasks and nothing in-process remembers the row. 2. _mark_failed in the 'except Exception' branch is itself a DB write; if it raises, the exception escapes with the row still claimed. 3. shutdown_graceful cancels remaining tasks after the drain timeout and returns immediately; the event loop may close before any cancelled task could reconcile its own row. recover_own_tasks only runs at startup, and no dead-worker logic applies because the worker is alive, so such rows are unreachable by any existing recovery path and clients polling them wait forever. Fix: a self-scoped release guarded on status='processing' AND worker_id=, applied at all three sites. Only the worker's own local knowledge is used, the same authority recover_own_tasks already exercises. The retry_count increment mirrors recover_own_tasks (#2675/#2834) so an interrupted task still counts against its retry budget. Fixes #3228 --- .../hindsight_api/worker/poller.py | 84 ++++++- .../tests/test_worker_task_release.py | 211 ++++++++++++++++++ 2 files changed, 294 insertions(+), 1 deletion(-) create mode 100644 hindsight-api-slim/tests/test_worker_task_release.py diff --git a/hindsight-api-slim/hindsight_api/worker/poller.py b/hindsight-api-slim/hindsight_api/worker/poller.py index a97ce0bebd..e184c7d65b 100644 --- a/hindsight-api-slim/hindsight_api/worker/poller.py +++ b/hindsight-api-slim/hindsight_api/worker/poller.py @@ -586,6 +586,65 @@ async def _mark_failed(self, operation_id: str, error_message: str, schema: str ) await self._maybe_update_parent_operation(operation_id, schema, conn) + async def _release_if_still_claimed(self, operation_id: str, schema: str | None) -> bool: + """Return an operation this worker still owns to 'pending'. + + Guarded on ``status = 'processing' AND worker_id = ``, so it is a + no-op the moment the task has written its own terminal state. That makes + it safe to call whenever a task stops running without reaching the + terminal-marking code: ``asyncio.CancelledError`` derives from + BaseException and so escapes every ``except Exception`` above, and + ``_mark_failed`` is itself a DB write that can raise. In both cases the + done-callback removes the task from ``_active_tasks``, so nothing + in-process remembers the row and ``recover_own_tasks`` — which runs only + at startup — never sees it either. + + The retry_count increment matches recover_own_tasks: an interrupted task + counts against the retry budget, so a task that reliably kills its + worker cannot be re-claimed forever (#2675 / #2834). + """ + table = fq_table("async_operations", schema) + async with self._backend.acquire() as conn: + result = await conn.execute( + f""" + UPDATE {table} + SET status = 'pending', worker_id = NULL, claimed_at = NULL, + retry_count = COALESCE(retry_count, 0) + 1, updated_at = now() + WHERE operation_id = $1 AND status = 'processing' AND worker_id = $2 + """, + operation_id, + self._worker_id, + ) + return bool(_updated_row_count(result)) + + async def _release_all_claimed(self) -> int: + """Return every operation still claimed by this worker to 'pending'. + + Runs across all configured schemas, like recover_own_tasks. + """ + total = 0 + for schema in await self._get_schemas(): + table = fq_table("async_operations", schema) + try: + async with self._backend.acquire() as conn: + result = await conn.execute( + f""" + UPDATE {table} + SET status = 'pending', worker_id = NULL, claimed_at = NULL, + retry_count = COALESCE(retry_count, 0) + 1, updated_at = now() + WHERE status = 'processing' AND worker_id = $1 + """, + self._worker_id, + ) + total += _updated_row_count(result) + except Exception: + schema_display = f'"{schema}"' if schema else str(schema) + logger.warning( + f"Worker {self._worker_id} could not release claimed tasks for schema {schema_display}", + exc_info=True, + ) + return total + async def _maybe_update_parent_operation(self, child_operation_id: str, schema: str | None, conn) -> None: """If this operation is a child of a batch_retain, update the parent status when all siblings are done. @@ -843,10 +902,23 @@ async def _execute_task_inner(self, task: ClaimedTask, holder: StageHolder | Non except RetryTaskAt as e: # Retry is not a terminal outcome — do not record a completion. await self._schedule_retry(task.operation_id, e.retry_at, str(e), task.schema) + except asyncio.CancelledError: + # Cancellation is not a terminal outcome: leave no metric, but the row + # must not stay 'processing'. Re-raise so the task still reports as + # cancelled to whoever cancelled it. + if await self._release_if_still_claimed(task.operation_id, task.schema): + logger.warning(f"Task {task.operation_id} was cancelled; returned operation to 'pending' for re-claim") + raise except Exception as e: logger.error(f"Task {task.operation_id} failed: {e}") traceback.print_exc() - await self._mark_failed(task.operation_id, str(e), task.schema) + try: + await self._mark_failed(task.operation_id, str(e), task.schema) + except Exception: + # Marking failed is itself a DB write. If it cannot land, the row + # would otherwise stay 'processing' with no owner running it. + logger.exception(f"Could not mark task {task.operation_id} failed; releasing it for re-claim") + await self._release_if_still_claimed(task.operation_id, task.schema) terminal_success = False # Record the metric outside the executor's exception scope so a metrics @@ -1272,6 +1344,16 @@ async def shutdown_graceful(self, timeout: float = 30.0): if not info.bg_task.done(): info.bg_task.cancel() + # Cancelled tasks cannot be trusted to land their own release — the loop + # may close first — so reclaim everything this worker still owns in one + # statement. Same shape as recover_own_tasks, run at shutdown instead of + # only at startup, so a restart no longer strands in-flight work. + released = await self._release_all_claimed() + if released: + logger.warning( + f"Worker {self._worker_id} returned {released} in-flight operations to 'pending' on shutdown" + ) + async def _log_progress_if_due(self): """Log progress stats every PROGRESS_LOG_INTERVAL seconds. diff --git a/hindsight-api-slim/tests/test_worker_task_release.py b/hindsight-api-slim/tests/test_worker_task_release.py new file mode 100644 index 0000000000..bac9c12ab7 --- /dev/null +++ b/hindsight-api-slim/tests/test_worker_task_release.py @@ -0,0 +1,211 @@ +""" +Tests for the worker releasing operations it still owns when a task stops +running without reaching its terminal-marking code. + +Three paths previously stranded rows in status='processing' forever: +- asyncio.CancelledError escaping _execute_task_inner (BaseException, so no + `except Exception` catches it), +- _mark_failed itself raising (it is a DB write), +- shutdown_graceful cancelling in-flight tasks and returning without any + DB reconciliation. + +See issue #3228. +""" + +import asyncio +import contextlib +import json +import uuid + +import pytest +import pytest_asyncio + + +async def _ensure_bank(pool, bank_id: str) -> None: + """Upsert a minimal bank row so FK on async_operations passes.""" + await pool.execute( + "INSERT INTO banks (bank_id, name) VALUES ($1, $2) ON CONFLICT DO NOTHING", + bank_id, + bank_id, + ) + + +# Use loadgroup to ensure these tests run in the same worker +# since they share database state +pytestmark = pytest.mark.xdist_group("worker_tests") + + +@pytest_asyncio.fixture +async def backend(pg0_db_url): + """Create a DatabaseBackend for worker tests.""" + from hindsight_api.engine.db import create_database_backend + from hindsight_api.pg0 import resolve_database_url + + resolved_url = await resolve_database_url(pg0_db_url) + + b = create_database_backend("postgresql") + await b.initialize(resolved_url, min_size=2, max_size=10, command_timeout=30) + yield b + await b.shutdown() + + +@pytest_asyncio.fixture +async def pool(backend): + """Expose the raw asyncpg pool from the backend for direct DB access in tests.""" + yield backend.get_pool() + + +@pytest_asyncio.fixture +async def clean_operations(pool): + """Clean up async_operations table before and after tests.""" + await pool.execute("DELETE FROM async_operations WHERE status = 'pending'") + yield + await pool.execute("DELETE FROM async_operations WHERE bank_id LIKE 'test-worker-%'") + + +async def _insert_pending(pool, bank_id: str) -> uuid.UUID: + """Insert one claimable pending operation, returning its id.""" + op_id = uuid.uuid4() + payload = json.dumps({"type": "test_task", "bank_id": bank_id, "operation_id": str(op_id)}) + await pool.execute( + """ + INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload) + VALUES ($1, $2, 'test', 'pending', $3::jsonb) + """, + op_id, + bank_id, + payload, + ) + return op_id + + +async def _fetch_row(pool, op_id): + return await pool.fetchrow( + "SELECT status, worker_id, claimed_at, retry_count FROM async_operations WHERE operation_id = $1", + op_id, + ) + + +async def _wait_for_status(pool, op_id, status: str, timeout: float = 2.0): + """Poll until the row reaches `status` or timeout; return the final row.""" + deadline = asyncio.get_event_loop().time() + timeout + row = await _fetch_row(pool, op_id) + while row["status"] != status and asyncio.get_event_loop().time() < deadline: + await asyncio.sleep(0.05) + row = await _fetch_row(pool, op_id) + return row + + +class TestTaskRelease: + @pytest.mark.asyncio + async def test_cancelled_task_returns_operation_to_pending(self, pool, backend, clean_operations): + """A cancelled in-flight task must not leave its row 'processing'.""" + from hindsight_api.worker import WorkerPoller + + bank_id = f"test-worker-{uuid.uuid4().hex[:8]}" + await _ensure_bank(pool, bank_id) + op_id = await _insert_pending(pool, bank_id) + + started = asyncio.Event() + block = asyncio.Event() + + async def blocking_executor(task_dict): + started.set() + await block.wait() + + poller = WorkerPoller( + backend=backend, + worker_id="test-release-worker", + executor=blocking_executor, + ) + + claimed = await poller.claim_batch() + ours = [t for t in claimed if t.operation_id == str(op_id)] + assert len(ours) == 1, "test operation should be claimed" + + row = await _fetch_row(pool, op_id) + assert row["status"] == "processing" + assert row["worker_id"] == "test-release-worker" + + await poller.execute_task(ours[0]) + await asyncio.wait_for(started.wait(), timeout=5) + + bg_task = poller._active_tasks[str(op_id)].bg_task + bg_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await bg_task + + row = await _wait_for_status(pool, op_id, "pending") + assert row["status"] == "pending", "cancelled task must release its row" + assert row["worker_id"] is None + assert row["claimed_at"] is None + assert row["retry_count"] == 1, "an interrupted run counts against the retry budget" + + @pytest.mark.asyncio + async def test_failed_terminal_write_returns_operation_to_pending(self, pool, backend, clean_operations): + """If _mark_failed itself raises, the row must be released, not stranded.""" + from hindsight_api.worker import WorkerPoller + + bank_id = f"test-worker-{uuid.uuid4().hex[:8]}" + await _ensure_bank(pool, bank_id) + op_id = await _insert_pending(pool, bank_id) + + async def failing_executor(task_dict): + raise RuntimeError("executor blew up") + + poller = WorkerPoller( + backend=backend, + worker_id="test-release-worker", + executor=failing_executor, + ) + + async def broken_mark_failed(operation_id, error_message, schema): + raise RuntimeError("pool exhausted") + + poller._mark_failed = broken_mark_failed + + claimed = await poller.claim_batch() + ours = [t for t in claimed if t.operation_id == str(op_id)] + assert len(ours) == 1 + + await poller.execute_task(ours[0]) + + row = await _wait_for_status(pool, op_id, "pending") + assert row["status"] == "pending", "a failed terminal write must not strand the row" + assert row["worker_id"] is None + assert row["retry_count"] == 1 + + @pytest.mark.asyncio + async def test_shutdown_graceful_releases_inflight_operations(self, pool, backend, clean_operations): + """shutdown_graceful must not strand rows it cancelled.""" + from hindsight_api.worker import WorkerPoller + + bank_id = f"test-worker-{uuid.uuid4().hex[:8]}" + await _ensure_bank(pool, bank_id) + op_id = await _insert_pending(pool, bank_id) + + started = asyncio.Event() + block = asyncio.Event() + + async def blocking_executor(task_dict): + started.set() + await block.wait() + + poller = WorkerPoller( + backend=backend, + worker_id="test-release-worker", + executor=blocking_executor, + ) + + claimed = await poller.claim_batch() + ours = [t for t in claimed if t.operation_id == str(op_id)] + assert len(ours) == 1 + + await poller.execute_task(ours[0]) + await asyncio.wait_for(started.wait(), timeout=5) + + await poller.shutdown_graceful(timeout=0.1) + + row = await _wait_for_status(pool, op_id, "pending") + assert row["status"] == "pending", "shutdown must return in-flight rows to pending" + assert row["worker_id"] is None