Keep heartbeat alive during graceful shutdown

[sebdiem]

A draining worker is still alive and processing its in-flight jobs, but _shutdown cancelled every side task (including the heartbeat) before waiting for those jobs to drain. The frozen heartbeat made the worker look dead to stalled-worker recovery, which could then reclaim jobs that were still being processed.

Cancel the other side tasks immediately (a draining worker must not fetch new jobs, defer periodics or listen for notifications), but keep the heartbeat beating until in-flight jobs have drained or been aborted, and only then cancel it before unregistering the worker.

This is kind of fixable on library's users side with something like this but it looks like it would be worth it to address this in the library itself.

HEARTBEAT_SIDE_TASK_NAME = "update_heartbeats"

class DrainHeartbeatWorker(Worker):
    """Worker that keeps beating while it gracefully drains in-flight jobs."""

    async def _shutdown(self, side_tasks: list) -> None:
        heartbeat_tasks = [
            t for t in side_tasks if t.get_name() == HEARTBEAT_SIDE_TASK_NAME
        ]
        other_tasks = [
            t for t in side_tasks if t.get_name() != HEARTBEAT_SIDE_TASK_NAME
        ]

        await utils.cancel_and_capture_errors(other_tasks)
        try:
            await super()._shutdown(side_tasks=[])
        finally:
            await utils.cancel_and_capture_errors(heartbeat_tasks)

Successful PR Checklist:

• Tests
  • (not applicable?)
• Documentation
  • (not applicable?)

PR label(s):

• PR type: breaking 💥 Contains breaking changes
• PR type: feature ⭐️ Contains new features
• PR type: bugfix 🕵️ Contains bug fix
• PR type: miscellaneous 👾 Contains misc changes
• PR type: dependencies 🤖 Contains only dependencies updates
• PR type: documentation 📚 Contains documentation updates

Summary by CodeRabbit

• Bug Fixes

  • Improved worker graceful shutdown sequence to maintain heartbeat signals while draining in-flight jobs, ensuring proper coordination and reliability before full unregistration.

• Tests

  • Added test coverage validating heartbeat behavior during graceful shutdown.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 0bc57a6e-3345-4d81-9c8a-18cc7c53f2ec

📥 Commits

Reviewing files that changed from the base of the PR and between 1433612 and 0013bff.

📒 Files selected for processing (2)

• procrastinate/worker.py
• tests/unit/test_worker.py

---

📝 Walkthrough

Walkthrough

The worker's shutdown sequence is modified so the heartbeat task runs independently from other side tasks. _start_side_tasks now returns a (heartbeat_task, side_tasks) tuple; _shutdown accepts heartbeat_task separately and cancels it only after in-flight jobs drain, before unregistering the worker. A new test verifies the heartbeat advances during the drain window.

Changes

Heartbeat lifecycle during graceful shutdown

Layer / File(s)	Summary
Heartbeat task separation and shutdown sequencing procrastinate/worker.py	_start_side_tasks creates heartbeat_task independently and returns (heartbeat_task, side_tasks). _run_loop unpacks the pair and adds the heartbeat to the side-task failure monitor. _shutdown gains a heartbeat_task parameter, keeps the heartbeat alive through the drain phase, then cancels it via cancel_and_capture_errors before unregistering the worker.
Test: heartbeat advances during drain tests/unit/test_worker.py	New test test_heartbeat_continues_during_graceful_shutdown starts a worker with update_heartbeat_interval, queues a long-running job, calls worker.stop() mid-drain, asserts heartbeat timestamps increase, then verifies worker.worker_id is None and the connector registry is empty after shutdown.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Poem

🐇 Hop hop, the heartbeat must not sleep,
While long-running jobs their promises keep!
We drain the queue with a steady beat,
Then cancel the pulse — the shutdown's complete.
A tidy registry, workers all gone,
The rabbit ensures the heartbeat beats on! 🕐

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 45.45% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title directly and clearly summarizes the main change: keeping the heartbeat alive during graceful shutdown, which is the core purpose of all modifications across both files.
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.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ 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 and usage tips.}

[owenh000]

I opened issue #1598 for a closely related problem. It seems to me that the poll_jobs_to_abort side task should also be kept running, like the heartbeat. Otherwise it is impossible to abort a running job while its worker is shutting down.