fix(worker): release operations when a task is cancelled or its terminal write fails

[russellbrenner]

Summary

When a worker's in-flight task stops running without reaching its terminal-marking code, its async_operations row stays status='processing' owned by a live worker, forever. The worker has already forgotten it (_cleanup_task removed it from _active_tasks and freed the slot), recover_own_tasks only runs at startup, and no dead-worker handling applies because the worker is alive. Clients polling the operation wait indefinitely.

Observed in production (0.8.6-slim, k8s, single worker): a pod running 39h without restart held two rows processing for 38h — retry_count=0, absent from the worker's own [WORKER_TASK] in-flight logging, zero log mentions in 20 minutes, not counted in slots=9/12. Full evidence in #3228.

The three paths

1. asyncio.CancelledError escapes _execute_task_inner. The handler chain catches _WallTimeoutExceeded, DeferOperation, RetryTaskAt, and Exception. CancelledError derives from BaseException, so it escapes all four; there is no finally.
2. _mark_failed is itself a DB write. If it raises inside the except Exception branch (pool exhaustion, connection reset, statement timeout), the row is stranded. Likely the cause of the observed rows, since no shutdown occurred on that pod.
3. shutdown_graceful cancels and returns. After the drain timeout it calls bg_task.cancel() on each in-flight task and returns immediately; the event loop may close before any cancelled task could reconcile its own row. Every rollout with in-flight work over the grace period strands it.

The fix

A self-scoped release, applied at all three sites:

UPDATE async_operations
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  -- self

• _execute_task_inner: new except asyncio.CancelledError: branch that releases and re-raises (no metric — cancellation is not a terminal outcome).
• except Exception: _mark_failed wrapped so its own failure releases instead of stranding.
• shutdown_graceful: after the cancel loop, a bulk release of everything still claimed by this worker id across all schemas — the same shape as recover_own_tasks, run at shutdown instead of only at the next startup.

Scope note: this deliberately does not touch other workers' rows. I read the history on #992/#1441/#1510 and the line drawn there — no cross-worker coordination, no liveness inference. The guard worker_id = <self> means every write here uses only the worker's own local knowledge, the same authority recover_own_tasks already exercises at startup; it is a no-op the moment a task lands its own terminal state, so there is no duplicate-work race. It also answers "we need to figure out why workers are stuck" (#992) — these three paths are why.

The retry_count increment mirrors recover_own_tasks (#2675/#2834): a task that reliably kills its worker still exhausts its retry budget rather than being re-claimed forever.

Complementary to (not overlapping) #3092, which makes the drain timeout configurable; this PR makes the post-timeout cancellation stop stranding rows regardless of the timeout's value.

Tests

tests/test_worker_task_release.py, three tests against real Postgres (pg0), all failing on main and passing with the fix:

• test_cancelled_task_returns_operation_to_pending — claim, block the executor, cancel the tracked bg_task; row must end pending/worker_id NULL/retry_count 1.
• test_failed_terminal_write_returns_operation_to_pending — executor raises, _mark_failed monkeypatched to raise; row must end pending, not stranded processing.
• test_shutdown_graceful_releases_inflight_operations — shutdown_graceful(timeout=0.1) with a blocked task in flight; row must end pending.

$ uv run pytest -n0 tests/test_worker_task_release.py
3 passed
# with the poller.py change stashed:
3 failed  (rows stay 'processing')
$ uv run pytest -n0 tests/test_worker.py
102 passed

Fixes #3228

[Copilot]

Pull request overview

This PR addresses a production issue where async_operations rows can remain stuck in status='processing' and owned by a live worker when an in-flight task exits early (cancellation) or its terminal DB write fails, causing clients polling the operation to wait indefinitely.

Changes:

• Add guarded “self-release” helpers in WorkerPoller to return still-claimed operations back to pending with retry_count incremented.
• Ensure asyncio.CancelledError and terminal-write failures release stranded operations rather than leaving them stuck in processing.
• Add regression tests covering cancellation, terminal-write failure, and graceful shutdown cleanup behavior.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
hindsight-api-slim/hindsight_api/worker/poller.py	Adds self-scoped release logic and wires it into cancellation, terminal-write failure handling, and graceful shutdown.
hindsight-api-slim/tests/test_worker_task_release.py	Adds regression tests validating operations are returned to pending in the three previously-stranding paths.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.