feat(worker): until-idle drain mode (#53)

[rcbevans]

Summary

Adds --until-idle mode to taskq worker and worker_main(until_idle=True) for batch-processing use cases. The worker drains all jobs in its subscribed queues through the existing four-phase graceful shutdown path (DRAINING → CANCELLING → FORCING → ABANDONING) and exits with a status code: 0 if every job succeeded, 2 if any failed, 3 if a max-runtime cap was hit. No new shutdown plumbing — the drain monitor triggers the same orchestrate_shutdown that SIGTERM uses.

Issue addressed

Closes #53 — requested a worker option to exit when queues are drained, with non-zero exit codes for failures and timeout.

Implementation

Drain monitor loop (src/taskq/worker/drain.py — new): drain_monitor_loop() is spawned as a sibling coroutine (with may_return=True) only when until_idle=True. It polls backend.count_active_jobs(queues) and deps.active_jobs.count() every idle_poll_interval. When both stay zero for idle_settle_window, it triggers orchestrate_shutdown via _trigger_drain_shutdown(), which creates the orchestration task and appends it to orchestrator_holder — exactly as the SIGTERM handler does. It does not set shutdown_event directly; orchestrate_shutdown's finally block does that at the correct point. If max_runtime is exceeded before drain completes, it triggers shutdown with exit code 3.

Double-orchestration guard (H2): _trigger_drain_shutdown and _on_shutdown_signal in shutdown.py both check orchestrator_holder and deps.shutdown_phase before creating a second orchestration task. First trigger wins — if SIGTERM fires while the drain monitor is mid-orchestration, the signal handler returns early, and vice versa.

drain_failures counter (src/taskq/worker/deps.py): New int field on WorkerDeps, defaulting to 0. Incremented by di_consumer_loop in run.py when dispatch_one_job returns "failed", or when dispatch_one_job raises an exception. The drain monitor reads this counter to pick exit code 0 vs 2. "scheduled" (snooze/retry) does not increment — a retried job is not a drain failure.

dispatch_one_job return type (src/taskq/worker/dispatch.py): Changed from -> None to -> AttemptOutcome. The outcome variable was already computed for metrics; return outcome was added after the try/except/finally block (not in finally, to preserve CancelledError propagation). di_consumer_loop captures the return value.

count_active_jobs backend method: Added to the Backend protocol, PostgresBackend, InMemoryBackend, and FakeBackend. The SQL query counts jobs with status in ('pending', 'scheduled', 'running') filtered by queue. The InMemoryBackend implementation reuses ACTIVE_STATUSES from statemachine.py. Protocol version stays at 3 (additive, loud-failure on missing method).

WorkerSettings fields (src/taskq/settings.py): Three new fields — idle_settle_window (default 2.0s, ge=0), idle_poll_interval (default 1.0s, ge=0.1), idle_max_runtime (default None, gt=0). All overridable via env vars (TASKQ_IDLE_SETTLE_WINDOW, TASKQ_IDLE_POLL_INTERVAL, TASKQ_IDLE_MAX_RUNTIME) and CLI flags.

worker_main / _main wiring (src/taskq/worker/_bootstrap.py): Both functions accept until_idle, idle_settle_window, idle_poll_interval, max_runtime params. When until_idle=True, settings overrides are resolved, a cron-incompatibility warning is logged if cron schedules are present, and drain_monitor_loop is spawned inside the TaskGroup.

CLI (src/taskq/cli.py): Added --until-idle, --idle-settle-window, --idle-poll-interval, --max-runtime options to taskq worker, passed through to worker_main().

Test coverage

File	Type	Coverage
tests/test_count_active_jobs.py	Unit + integration	8 unit tests (InMemoryBackend: empty queues, no jobs, pending, running, scheduled, terminal excluded, multi-queue, queue subset) + 1 PostgresBackend integration test
tests/test_worker_settings_drain.py	Unit	12 tests: defaults, env overrides, validation (negative settle window, poll below 0.1, zero max_runtime), float type checks
tests/test_worker_drain.py	Unit	WorkerDeps.drain_failures default; dispatch_one_job return type; di_consumer_loop drain_failures increment on failure/exception/no-increment on success/scheduled; drain monitor: triggers on idle, exit code 2 on failures, exit code 3 on timeout, settle reset on new jobs, no trigger when active jobs, skips when orchestration already active, continues after transient count error
tests/test_worker_main.py	Unit	Drain monitor spawned with until_idle, exit code 2 on failures, exit code 3 on timeout, no drain monitor without until_idle, watchdog disarm ordering
tests/test_cli_worker.py	Unit	--until-idle flag passed to worker_main, default false, --idle-settle-window passed
tests/test_backend_protocol.py	Unit	Updated TestMethodCount: 36 → 37 members, added count_active_jobs to expected names
tests/e2e/test_until_idle.py	E2E	4 tests with real Postgres: drain and exit 0, exit 2 on failures, exit 3 on timeout, scheduled jobs drain before exit

Documentation

• docs/guides/workers.md — new "Until-idle mode" section with CLI and Python usage, exit codes, and cron incompatibility note
• docs/guides/cli.md — --until-idle and related flags in the worker options table; exit codes 2 and 3 documented
• docs/architecture.md — DrainMonitor box added to architecture diagram; count_active_jobs added to Backend protocol listing
• docs/specs/2026-07-29-worker-drain.md — full design spec

---

Closes #53