Shutdown waited on asyncio.gather(*active_tasks) with no timeout, and released
work-unit claims only via the finally in initialize(). Both assume the process
gets as long as it needs.
It does not. A platform can terminate the process on a budget shorter than a
single work unit takes -- a deriver LLM call alone can run for tens of seconds
-- and when that budget expires the process is killed outright. The unbounded
wait is then still in progress, initialize() never returns, its finally never
runs, and the claims are left held. The work is not lost, but nothing may touch
it until cleanup_stale_work_units notices, which takes
STALE_SESSION_TIMEOUT_MINUTES plus a jittered poll interval.
So shutdown is a race to hand the claims back, not an opportunity to finish the
work. Two changes:
Drain is bounded by SHUTDOWN_DRAIN_TIMEOUT_SECONDS (default 8). Work that fits
in the budget still completes; work that does not is cancelled rather than
waited on, since the work is reclaimable and the remaining budget is not.
cleanup() is called at the end of shutdown() rather than relying on the finally
in initialize(). Claims are given up while the process is still alive, so
another worker repicks the work on its next poll instead of waiting out the
stale sweep. cleanup() clears worker_ownership, so the existing call in the
finally becomes a no-op and the ordering is idempotent.
The drain deliberately runs before the release rather than after. Releasing
first would leave a window in which a task that completes during shutdown has
already had its claim taken by another worker, and the same queue items get
derived twice. Cancelling first and releasing after keeps that window closed,
and both steps fit inside a short budget.
Cancelling also keeps the error reporting honest without special-casing it.
asyncio.CancelledError derives from BaseException rather than Exception, so an
abandoned task passes straight through the `except Exception` handlers that
report to Sentry. Shutdown stops generating spurious errors on its own, while a
task that genuinely fails during the drain, or a cleanup() that cannot reach the
database, still reports as it should.
Tests cover a task that outlives the budget being abandoned with cleanup still
awaited, and a task that fits in the budget still finishing.
26be90f60d