Don't hang forever when cancelling a worker whose database connection is unresponsive

[tonycpsu]

Closes #1513

The problem

run_worker_async() can never be cancelled if the worker's database connection stops responding. The worker keeps running, processes nothing, logs nothing, and the only way out is to kill the process.

This is not theoretical: it caused a 7-hour silent stall in production, and the log evidence is unambiguous. A provider-side event on the managed-PostgreSQL connection path reset the server side of the worker's established connections — Postgres itself never restarted (its checkpoints run uninterrupted through the incident) but logged could not receive data from client: Connection reset by peer at the event moment — while the worker's side of those sockets stayed open. From that instant the worker stopped fetching jobs entirely; run_worker_async() neither returned nor raised, and not a single procrastinate ERROR was logged for 7 hours. The process was otherwise alive the whole time — it kept serving HTTP and even kept deferring jobs successfully on fresh connections — it just never consumed the queue again, until the process was restarted. TCP keepalives and tcp_user_timeout cannot catch this: the middlebox answers the probes, and nothing sent is left unacknowledged.

Why it hangs

Worker.run() shields the run loop task and relies on the stop event for a graceful shutdown:

try:
    await asyncio.shield(loop_task)
except asyncio.CancelledError:
    self.stop()
    await loop_task   # <-- hangs forever
    raise

stop() only sets _stop_event, and _fetch_and_process_jobs() only checks it between iterations. A fetch that never returns therefore never observes the stop event. Because the loop task is shielded, nothing ever cancels it either — so await loop_task waits for a task that cannot finish, and the shutdown hangs forever.

Observed against main, with the worker parked exactly there:

run_worker_async: ['app.py:348 run_worker_async', 'worker.py:468 run']
worker loop:      ['worker.py:701 _run_loop', 'worker.py:399 _fetch_and_process_jobs',
                   'manager.py:212 fetch_job', 'psycopg_connector.py:231 execute_query_one_async',
                   'contextlib.py:210 __aenter__']
update_heartbeats: [... 'manager.py:1018 update_heartbeat', ... 'contextlib.py:210 __aenter__']

The change

Adds an opt-in shutdown_timeout. When run() is cancelled, the worker waits up to shutdown_timeout for the run loop to stop; if it hasn't, it cancels it; if it does not react to the cancellation either, it is abandoned and the cancellation is propagated regardless, so cancelling a worker always terminates.

The same bound is applied to unregister_worker() in _shutdown(): verified by stack traces, the run loop otherwise re-enters an unbounded database round-trip through a pool that may still hold dead connections, after surviving the fetch cancellation. Unregistering is best effort by design — a worker that could not unregister is eventually pruned through its stale heartbeat (prune_stalled_workers) — so on timeout it is cancelled and logged, not waited for. This is also what lets an abandoned run loop terminate on its own shortly after (observed: ~10s) instead of leaking for the life of the process, once the driver honors cancellation.

Default is None, which is exactly today's behavior (asyncio.wait(..., timeout=None) waits just like await loop_task did), so nothing changes for existing users unless they opt in.

I deliberately did not reuse shutdown_graceful_timeout for this. It bounds how long the worker waits for jobs, and it is consumed inside _shutdown(), which runs inside the run loop task. Using the same budget for the outer wait cancels the run loop while it is still legitimately aborting jobs — this breaks test_stopping_worker_aborts_job_after_timeout[cancel], which is how I found it. The two timeouts have to be separate; shutdown_timeout should be greater than shutdown_graceful_timeout.

What this does NOT fix (deliberate)

• The root cause is in psycopg, not here. I localized it: AsyncConnection.wait() (connection_async.py:520-529) catches CancelledError, then re-enters waiting.wait_async() on the same dead socket with no deadline, so the pending CancelledError is never re-raised. A task stays cancelling=1, cancelled=False forever. Verified with raw psycopg, no procrastinate and no pool involved, on psycopg 3.3.4 with both the binary (libpq 18) and pure-python (libpq 17) implementations. This is why the "did not react to cancellation" path exists at all; once psycopg bounds that wait, the cancellation lands promptly and the loop task unwinds cleanly. Reported upstream with a standalone reproduction: Async: task cancellation is never delivered when the server stops responding on an established connection psycopg/psycopg#1371.

• Abandoning the run loop leaks it. The abandoned task, its wedged connection, and the heartbeat side task stay wedged for the life of the process. This is a real cost, accepted deliberately: procrastinate cannot force an unresponsive round-trip to unwind, and hanging the caller forever is worse. It is bounded and logged rather than silent. (Concretely, the leaked backends hold idle in transaction locks server-side until the connection actually dies.)

This abandon path is load-bearing, not a temporary workaround. I verified that a wedged pool.connection() cannot be cancelled even with the psycopg core bug fixed, because there is a second, independent swallow in psycopg_pool: pool connections are not autocommit, so check_connection() runs try: await conn.execute("") / finally: await conn.set_autocommit(False). The finally performs I/O on a connection that is by then dead, so it raises, and an exception raised in a finally replaces the in-flight CancelledError. The pool's _getconn_with_check_loop then reads that as "this connection failed its check", discards it, and retries the next one — so the cancellation is lost and the caller eventually gets a PoolTimeout rather than CancelledError. Measured: a cancelled pool.connection() was still not cancelled 40s later. Until that is fixed too, procrastinate cannot rely on cancelling the loop task, and this timeout is the only thing that guarantees shutdown terminates.

• This does not fix Worker crashes on database connection loss during fetch loop #1523. That is a different failure mode — a ConnectorException crashing out of the fetch loop when the outage exceeds the pool timeout, i.e. a crash, not a hang. This PR only addresses shutdown.

• It does not make the worker recover from a wedged connection. The fetch loop still stalls silently (no error, no fetch) until something cancels the worker. Bounding the fetch itself is impossible today: I verified that pool.connection(timeout=N), asyncio.wait_for() and asyncio.timeout() all fail to fire on a half-open connection, because they all depend on the cancellation that psycopg swallows. Setting check=None on the pool doesn't help either — the hang just moves from check_connection to the query. Detecting the stall is left to the caller for now.

Testing

• New regression test test_cancelling_run_does_not_hang_on_stuck_loop_task. It models a fetch that never returns and does not honor the first cancellation, which is the essential property — a plainly cancellable sleep passes even without the fix and would give false confidence. The test fails on main (asserting run() never returned) and passes with the fix. It uses asyncio.wait, never wait_for, because wait_for cancels only once and would itself hang on the swallowed cancellation rather than failing the test.
• Verified end-to-end against a real half-open TCP connection (a proxy that strands established sockets while new connections keep working, simulating a failover behind a load balancer): cancelling the worker went from hanging forever to terminating, emitting Worker loop did not stop within shutdown_timeout. Cancelling it and Worker loop did not react to cancellation. Abandoning it.

---

Disclosure: this change was written with the assistance of an LLM (Claude). The bug was reproduced, the root cause localized, and the fix verified end-to-end against a real half-open connection by a human-reviewed process; the reasoning and evidence are above.

🤖 Generated with Claude Code

Summary by CodeRabbit

• New Features
  • Added an optional shutdown_timeout setting to bound how long workers wait to stop during cancellation and cleanup.
• Documentation
  • Updated shutdown guidance to cover the “stuck run loop” scenario and include examples showing shutdown_graceful_timeout vs shutdown_timeout.
• Bug Fixes
  • Prevented worker shutdown from hanging indefinitely when shutdown involves unresponsive database calls or worker unregistration.
• Tests
  • Added regression coverage for bounded shutdown behavior and for warnings when shutdown_timeout is set incorrectly.

[coderabbitai]

📝 Walkthrough

Walkthrough

Adds optional shutdown_timeout configuration, bounds worker loop cancellation and unregistration waits, logs timeout outcomes, consumes late task failures, and documents and tests shutdown behavior with unresponsive database operations.

Changes

Worker shutdown timeout

Layer / File(s)	Summary
Shutdown timeout configuration and documentation procrastinate/app.py, procrastinate/worker.py, docs/howto/advanced/shutdown.md, tests/unit/test_worker.py	Adds and documents the optional shutdown_timeout setting, validates its relationship to shutdown_graceful_timeout, and tests the configuration warning.
Bounded run-loop cancellation procrastinate/worker.py, tests/unit/test_worker.py	Bounds waiting for a cancelled worker loop, cancels the loop after timeout, consumes late task failures, and tests shutdown with a stuck job fetch.
Bounded worker unregistration procrastinate/worker.py, tests/unit/test_worker.py	Bounds worker unregistration, cancels timed-out tasks, consumes late failures, logs warnings, and tests shutdown with a stuck unregister call.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Possibly related issues

• Worker hangs on shutdown when DB connection is broken #1513: Addresses worker shutdown hangs when cancellation is blocked by stuck database-related operations.
• One-shot worker hangs forever at shutdown if a side task swallows its cancellation (periodic deferrer + connection-pool cancellation race) #1591: Addresses indefinitely hanging worker shutdown by bounding waits for stuck or cancellation-resistant tasks.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly matches the main change: adding bounded worker cancellation for an unresponsive database connection.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (1)

procrastinate/worker.py (1)

34-104: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Consider warning when shutdown_timeout <= shutdown_graceful_timeout.

Docs state shutdown_timeout "should be greater than shutdown_graceful_timeout" since the run loop waits for jobs to finish before exiting, but nothing validates this at construction time. A misconfigured pair (e.g. equal or shutdown_timeout smaller) silently aborts jobs before the graceful window elapses.

♻️ Example validation

         self.shutdown_graceful_timeout = shutdown_graceful_timeout
         self.shutdown_timeout = shutdown_timeout
+        if (
+            shutdown_timeout is not None
+            and shutdown_graceful_timeout is not None
+            and shutdown_timeout <= shutdown_graceful_timeout
+        ):
+            self.logger.warning(
+                "shutdown_timeout should be greater than shutdown_graceful_timeout"
+            )

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@procrastinate/worker.py` around lines 34 - 104, In Worker.__init__, validate
the configured shutdown_graceful_timeout and shutdown_timeout after assigning
them: when both are set and shutdown_timeout is less than or equal to
shutdown_graceful_timeout, emit a warning using the worker logger. Preserve
valid configurations and None values without warning, and do not change shutdown
behavior.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@procrastinate/worker.py`:
- Around line 467-508: The cancellation branch in the worker shutdown flow must
consume a completed loop_task after the second bounded wait, just as
unregister_task does. After the wait following loop_task.cancel(), check whether
the task is done and await it to retrieve any late exception, while preserving
the existing abandonment warning and non-blocking behavior when it remains
pending.

---

Nitpick comments:
In `@procrastinate/worker.py`:
- Around line 34-104: In Worker.__init__, validate the configured
shutdown_graceful_timeout and shutdown_timeout after assigning them: when both
are set and shutdown_timeout is less than or equal to shutdown_graceful_timeout,
emit a warning using the worker logger. Preserve valid configurations and None
values without warning, and do not change shutdown behavior.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 8d6dc9cd-fc69-4199-bdc6-8329e4d47fc1

📥 Commits

Reviewing files that changed from the base of the PR and between 7feebed and 316ca75.

📒 Files selected for processing (4)

• docs/howto/advanced/shutdown.md
• procrastinate/app.py
• procrastinate/worker.py
• tests/unit/test_worker.py

[tonycpsu]

For what it's worth, I was hoping to submit issues for psycopg as well but I think the author over there is a bit more hostile to AI-assisted fixes, which is why I'm starting here.

[tonycpsu]

Addressed both review items in 60dc4f7: the abandoned loop task's outcome is now consumed via a done callback (same pattern as the unregister task), and the worker warns at construction when shutdown_timeout <= shutdown_graceful_timeout, with a unit test for the warning.