Repository navigation
fix(worker): drain finalization on graceful shutdown (0.6) - #643
Conversation
Once a handler returned, the executor removed the attempt from the heartbeat set and finalized it in a detached task that Client::shutdown did not wait for. When the last handler returned as the drain ended, shutdown cancelled the completion batcher before that finalize reached it, and the direct-finalize fallback raced the process exit. The attempt stayed open with no heartbeat, so stale-heartbeat rescue ran the finished job again after heartbeat_staleness. Track finalizers in a TaskTracker and wait for them, within the drain timeout, before stopping the completion batcher, heartbeat and maintenance services. This is the shutdown part of #521 on main.
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. Please check the settings in the CodeRabbit UI or the
 |
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 7e6259e04d
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| // Graceful shutdown waits for this before it stops the completion | ||
| // batcher: the attempt has already left `in_flight`, so nothing | ||
| // heartbeats it, and a finalize lost at exit leaves it to rescue. | ||
| finalizers.spawn(async move { |
There was a problem hiding this comment.
When an outcome lifecycle hook is slow or hangs, this tracked future remains live through dispatch_lifecycle_event at line 572, so Client::shutdown waits until the entire drain deadline expires. This regresses the documented contract in docs/lifecycle-hooks.md:75-80 that hooks are best-effort and shutdown does not wait for them, turning user hook latency into shutdown latency. Track only the durable completion work and leave lifecycle dispatch detached.
Useful? React with 👍 / 👎.
hardbyte
added
the
full-ci
Problem
A graceful restart could run a finished job a second time. On a rolling deploy, a job whose handler returned while its runtime drained was occasionally rescued as a stale heartbeat about
heartbeat_stalenesslater and executed again on another runtime. The persona soak saw this after roughly one graceful restart in ten, on PostgreSQL 17 and 18.Once a handler returns, the executor removes the attempt from the heartbeat set and finalizes it in a detached
tokio::spawntask.Client::shutdowndrains the job tasks but not those finalizers. When the last handler returns just as the drain ends, shutdown cancels the completion batcher before that finalizer's request reaches it. The direct-finalize fallback then races the process exit: it fails on a closed pool or is dropped with the runtime. Nothing heartbeats the attempt any more, so it goes stale and rescue runs it again.A traced soak run shows the interleaving on one runtime (times in seconds):
Completed; its finalizer is spawned detachedCompleted; its finalizer is spawned detachedcompletion batcher stopped, falling back to direct finalizeshutdownreturns; the process closes its pool and exitsAll four rescues in that run have this shape. In each, a finalizer started and never finished on a runtime that had just shut down gracefully.
Change
Finalizers are spawned on a
TaskTracker.Client::shutdowncloses it and waits for it after draining the job tasks and before it cancels the completion batcher, heartbeat and maintenance. The wait shares the existing drain deadline. A finalize that outlives the timeout is left to rescue, as before.mainalready drains finalizers this way (#521). This change backports that shutdown behaviour without #521's reconnect retry.Tests
test_queue_storage_graceful_shutdown_finalizes_returned_attemptsruns a client on its own tokio runtime, so a process exit can be simulated. The test holds the attempt's receipt-completion advisory lock, which keeps the finalize in flight while shutdown drains. It completes the job through a follow-up spec, so the finalize runs in its own transaction. It then releases the lock, stops the runtime aftershutdownreturns, and starts a second runtime that rescues anything left open.runningafter the graceful exit, and the rescuing runtime runs the job a second time.completedunder the original lease and runs once.ci_timing.rscomes over frommainfor the staleness and wait bounds.Shard-ops persona soak, PostgreSQL 17, seed 5, 600 s, 40 graceful restarts per run:
cargo fmt,clippy -D warnings(workspace andawa-python), theawa-workerunit tests, and theexecutor_guard,enqueue_specandlifecycle_hooksuites pass on PostgreSQL 18.Independent review and validation
Includes the small #651 point-read fix. After a priority or queue move crossed a ready-ring wrap, retained tombstoned source rows could make
get_jobreturn old fields even after the batch succeeded. The available-candidate query now excludes full-identity tombstones (a 13-line SQL change), with priority and queue ring-wrap regressions. Both regressions fail before the fix and pass after it on PostgreSQL 17 and 18; existing aged-read and batch cases also pass. Required formatting, Clippy across all targets/features, and workspace build checks pass for the prepared source.Full CI passed all 29 jobs for
050e87bfdb68e1c3b83042806d44370e72e6d1f7.