From bd8b09876beede323f50d4e39c52ed500dd45079 Mon Sep 17 00:00:00 2001 From: Conrad Date: Mon, 10 Aug 2026 15:32:07 -0400 Subject: [PATCH 1/7] build: Track wool 0.14 for the worker idle RPC wool 0.14 adds an idle RPC to every worker and a WorkerConnection.idle() client for it, which is the primitive the idle-based worker shutdown needs. The LocalWorker construction surface and every other wool import cfdb uses are unchanged from 0.13.0. The specifier is compatible-release on 0.14 rather than exact: patch releases are picked up, 0.15 is not, since a minor bump may move the wire protocol the dispatch channel rides on. Because wool admits a worker only when the proxy's version is at most the worker's, any lock refresh that moves the resolved version is a fleet-rejection event rather than a routine dependency update, so refresh the lock deliberately and drain the fleet when it moves. --- pyproject.toml | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 0a7f58c..d6ce2bb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,7 +22,14 @@ dependencies = [ "requests", "strawberry-graphql", "uvicorn", - "wool~=0.13.0", + # Compatible-release on 0.14 now that the final has shipped: patch + # releases are picked up, 0.15 is not, since a minor bump may move + # the wire protocol the dispatch channel depends on. Note that even + # a patch move is a fleet-rejection event — wool admits a worker + # only when the proxy's version is at most the worker's — so + # refreshing the lock deliberately, and drain the fleet when it + # moves (see README's version-bump deploy caveat). + "wool~=0.14.0", ] description = "A Python utility for parsing and normalizing various DCC datapackages." dynamic = ["version"] From 56344fdd210739b2d1068b3e42df2a3a61ecf26a Mon Sep 17 00:00:00 2001 From: Conrad Date: Mon, 10 Aug 2026 15:32:15 -0400 Subject: [PATCH 2/7] feat: Shut workers down on continuous idle, draining before exit The static max-lifetime ceiling was the only self-termination path because wool exposed no per-job activity signal: the worker could not tell idle from busy, so drained-to-idle workers billed Fargate for hours and the ceiling could preempt a worker mid-task. The serve loop now dials its own worker over loopback (the same channel-back-to-own-subprocess pattern wool uses for graceful stop, so mTLS credentials and the identity SAN verify unchanged) and polls the idle RPC on a 15 s cadence, exiting once continuous idle crosses CFDB_WORKER_IDLE_TIMEOUT_SECONDS (default 600 s, 0 disables). A busy worker reports zero idle, so the idle exit never fires while a task is running. Both self-termination exits now stop the worker with a drain grace rather than wool's default immediate cancel. An idle reading is a snapshot, so it cannot rule out a dispatch accepted between the final poll and the teardown, and a max-lifetime expiry can land mid-job on a worker running jobs back to back; in either case the task has already been marked running on the API side, where a graceless cancel is finalized as a terminal failure instead of being re-queued. The grace (CFDB_WORKER_MAX_LIFETIME_GRACE_SECONDS, default 6 h) returns instantly when the docket is empty and is sized above the API's 4 h per-job duration cap, so the only work it can cancel is work already past that bound. That in turn lets the ceiling rise to 12 h without reintroducing mid-task preemption, at a worst-case worker uptime of lifetime plus grace. The signal paths keep the immediate cancel, since ECS bounds SIGTERM with SIGKILL anyway. Poll failures are handled without either killing a healthy worker or flooding the logs: IdleUnavailable disables polling for the process lifetime, an isolated failure retries on the next cadence, and after CFDB_WORKER_IDLE_POLL_FAILURE_LIMIT consecutive failures the worker escalates once to ERROR and stops polling. Those messages now name the failing exception type and describe the remaining bound honestly, rather than claiming a max-lifetime backstop that is disabled. --- src/cfdb/workflows/worker_main.py | 317 +++++++++++++++++++++++++++--- 1 file changed, 285 insertions(+), 32 deletions(-) diff --git a/src/cfdb/workflows/worker_main.py b/src/cfdb/workflows/worker_main.py index 471333a..96725d0 100644 --- a/src/cfdb/workflows/worker_main.py +++ b/src/cfdb/workflows/worker_main.py @@ -3,8 +3,11 @@ This module is the ``CMD`` for the worker container image. It boots a ``wool.LocalWorker`` on a known port, exposes a tiny HTTP health endpoint the ECS health check probes, handles SIGTERM cleanly so ``ecs.stop_task`` -cycles drain in-flight work, and self-terminates after a configurable -maximum lifetime so workers don't accumulate when the dispatch rate falls. +cycles drain in-flight work, and self-terminates once it has been +continuously idle beyond a configurable threshold so workers don't +accumulate when the dispatch rate falls — with a maximum-lifetime +ceiling retained as a backstop for workers whose idle reporting is +wedged or whose job is stuck. ECS owns the worker *lifecycle* — registration, IP, status, health — and ``EcsDiscovery`` polls it for all of that. But two fields of the wool @@ -27,10 +30,24 @@ * ``CFDB_WORKER_GRPC_PORT`` — gRPC port wool binds (default 50051). * ``CFDB_WORKER_HEALTH_PORT`` — HTTP ``/health`` port the ECS ``healthCheck`` probes (default 8080). +* ``CFDB_WORKER_IDLE_TIMEOUT_SECONDS`` — continuous idle beyond + which the worker exits; 0 disables. The primary reaper: a busy + worker reports zero idle, so the idle exit never fires mid-task, + and a dispatch racing the teardown drains under the + self-termination grace below. +* ``CFDB_WORKER_IDLE_POLL_INTERVAL_SECONDS`` — cadence of the idle + poll (default 15). Values below the loop's 1 s wakeup are + effectively floored at 1 s. +* ``CFDB_WORKER_IDLE_POLL_FAILURE_LIMIT`` — consecutive idle-poll + failures before the worker escalates once to ERROR and disables + idle shutdown (default 20). * ``CFDB_WORKER_MAX_LIFETIME_SECONDS`` — hard ceiling on worker - uptime; 0 disables. One hour above - :data:`cfdb.workflows.WORKFLOW_DURATION_CAP_S` so a worker started - shortly before a long sort can still outlive the job. + uptime; 0 disables. The backstop behind the idle timeout for a + worker whose idle reporting is wedged or whose job is stuck. +* ``CFDB_WORKER_MAX_LIFETIME_GRACE_SECONDS`` — how long a + self-terminating exit (idle or max-lifetime) waits for in-flight + tasks to drain before cancelling them; worst-case worker uptime is + therefore lifetime + grace. * ``CFDB_WORKER_DRAIN_GRACE_SECONDS`` — how long ``/health`` returns 503 after SIGTERM before tearing down the gRPC port. A second SIGTERM short-circuits. @@ -75,16 +92,69 @@ #: ``healthCheck`` can ``curl`` it without speaking gRPC. DEFAULT_HEALTH_PORT = 8080 -#: Default maximum wall-clock lifetime of a worker process. Wool exposes -#: no per-job activity hook today, so the worker can't tell idle from -#: busy; this is a hard ceiling — ECS replaces the task after this long. -#: Sized one hour above :data:`cfdb.workflows.WORKFLOW_DURATION_CAP_S` -#: (default 4 h) so a worker started shortly before a long sort can -#: still outlive the job. Note: max-lifetime expiry exits without the -#: drain-grace window the SIGTERM path provides — operators that want -#: a cleaner handoff should rely on ECS rolling tasks via service -#: updates rather than waiting for max-lifetime to fire. -DEFAULT_MAX_LIFETIME_SECONDS = 5 * 60 * 60 +#: Default continuous-idle threshold beyond which the worker exits. +#: This is the primary reaper: wool's ``idle`` RPC reports seconds since +#: the worker's in-flight task set last became empty (zero while any +#: task runs), so the idle exit never fires while a task is running, +#: and a dispatch accepted in the teardown window drains under +#: :data:`DEFAULT_MAX_LIFETIME_GRACE_SECONDS` rather than being +#: cancelled. Ten minutes rides out ordinary dispatch gaps between +#: queued jobs while reclaiming a drained-to-idle Fargate task ~70× +#: sooner than the max-lifetime ceiling would. ``0`` disables idle +#: shutdown and restores the pure max-lifetime behavior. +DEFAULT_IDLE_TIMEOUT_SECONDS = 600.0 + +#: Default cadence at which the serve loop polls its own worker's +#: ``idle`` RPC. The RPC returns the accumulated continuous idle +#: duration, so the cadence only bounds threshold overshoot — the +#: worker exits within one poll interval of crossing the timeout — +#: and polls deliberately do not disturb the measurement. The loop +#: wakes once per :data:`_STOP_POLL_INTERVAL_SECONDS`, so values +#: below 1 s are effectively floored at 1 s. +DEFAULT_IDLE_POLL_INTERVAL_SECONDS = 15.0 + +#: Per-poll gRPC deadline for the ``idle`` RPC. The dial is loopback, +#: so a slow answer means the worker subprocess is broken rather than +#: busy; a poll that times out is logged and retried on the next +#: cadence, and a worker that never answers is bounded by the +#: max-lifetime backstop. +_IDLE_POLL_RPC_TIMEOUT_SECONDS = 5.0 + +#: Consecutive idle-poll failures tolerated before the loop escalates +#: once to ERROR and disables further polling. Twenty at the 15 s +#: cadence is ~5 minutes of sustained failure — far beyond any +#: transient blip — after which continuing to warn every cadence adds +#: noise without information. The max-lifetime backstop still bounds +#: the worker once polling is disabled. +DEFAULT_IDLE_POLL_FAILURE_LIMIT = 20 + +#: Default maximum wall-clock lifetime of a worker process — the +#: backstop behind idle-based shutdown, not the primary reaper. It +#: bounds the two cases the idle timeout cannot: a worker whose +#: ``idle`` RPC is wedged or unimplemented (so idle polling yields +#: nothing), and a stuck job that holds the in-flight set non-empty +#: past any reasonable runtime. Because expiry now drains in-flight +#: work for up to :data:`DEFAULT_MAX_LIFETIME_GRACE_SECONDS` rather +#: than cancelling it, the ceiling can sit well above any healthy +#: job's runtime: a worker that stays busy back-to-back is reaped at +#: the first expiry whose drain completes, not mid-task. +DEFAULT_MAX_LIFETIME_SECONDS = 12 * 60 * 60 + +#: How long a self-terminating exit — idle timeout or max-lifetime — +#: waits for in-flight tasks to drain before cancelling them +#: (``wool.Worker.stop(grace=...)``). On a genuinely idle exit the +#: docket is empty and the drain returns instantly, so the grace +#: costs nothing in the common case; it exists for the two cases +#: where work is in flight at stop time: a dispatch accepted between +#: the final idle poll and the teardown, and a max-lifetime expiry on +#: a busy worker. Sized above the API's 4 h +#: :data:`cfdb.workflows.WORKFLOW_DURATION_CAP_S` so a healthy job +#: always finishes inside it — the only work ever cancelled is work +#: already past the API's own viability bound. Worst-case worker +#: uptime is therefore lifetime + grace (18 h at the defaults). The +#: signal path does not use this grace: ECS bounds SIGTERM with +#: SIGKILL, so a long drain there is unreachable anyway. +DEFAULT_MAX_LIFETIME_GRACE_SECONDS = 6 * 60 * 60 #: How long to keep returning ``503`` on ``/health`` after SIGTERM, #: giving ECS a chance to observe ``unhealthy`` and drain at the load @@ -328,11 +398,33 @@ async def _publish_worker_metadata( return +def _lifetime_bound_description(max_lifetime_seconds: float) -> str: + """Describe what still bounds the worker once idle polling stops. + + The disable-polling log lines close with this so they stay honest + when the max-lifetime backstop is itself disabled (``0``): claiming + a "backstop (0s)" would assert a bound that does not exist. + """ + if max_lifetime_seconds > 0: + return ( + f"the max-lifetime backstop ({max_lifetime_seconds:.0f}s) " + "still bounds this worker" + ) + return ( + "no max-lifetime backstop is configured — this worker's uptime " + "is now unbounded" + ) + + async def serve( *, worker_port: int = DEFAULT_WORKER_PORT, health_port: int = DEFAULT_HEALTH_PORT, + idle_timeout_seconds: float = DEFAULT_IDLE_TIMEOUT_SECONDS, + idle_poll_interval_seconds: float = DEFAULT_IDLE_POLL_INTERVAL_SECONDS, + idle_poll_failure_limit: int = DEFAULT_IDLE_POLL_FAILURE_LIMIT, max_lifetime_seconds: float = DEFAULT_MAX_LIFETIME_SECONDS, + max_lifetime_grace_seconds: float = DEFAULT_MAX_LIFETIME_GRACE_SECONDS, drain_grace_seconds: float = DEFAULT_DRAIN_GRACE_SECONDS, tls_ca: Optional[str] = None, tls_cert: Optional[str] = None, @@ -340,9 +432,24 @@ async def serve( publish_attempts: int = DEFAULT_PUBLISH_ATTEMPTS, publish_backoff_seconds: float = DEFAULT_PUBLISH_BACKOFF_SECONDS, ) -> int: - """Run the worker until SIGTERM or maximum lifetime elapses. - - Returns ``0`` on clean shutdown (SIGTERM, SIGINT, or max-lifetime). + """Run the worker until SIGTERM, idle timeout, or maximum lifetime. + + The primary self-termination path is idle-based: on a + ``idle_poll_interval_seconds`` cadence the loop asks its own worker + (via wool's ``idle`` RPC, over loopback) how long it has been + continuously idle, and exits once that crosses + ``idle_timeout_seconds``. A busy worker reports zero idle, so the + idle exit never fires while a task is running; + ``idle_timeout_seconds=0`` disables it. ``max_lifetime_seconds`` + remains as the backstop for a worker whose idle reporting is wedged + or whose job is stuck. Both self-termination exits stop the worker + with ``grace=max_lifetime_grace_seconds``, draining any in-flight + task — a dispatch that raced the idle teardown, or the job a + max-lifetime expiry interrupted — before cancelling; the signal + paths keep wool's immediate cancel. + + Returns ``0`` on clean shutdown (SIGTERM, SIGINT, idle timeout, or + max-lifetime). Bind failures and other early-startup errors raise out — ``main`` propagates them and the process exits with a Python traceback, which surfaces the cause in container logs more clearly than a @@ -415,8 +522,22 @@ def _signal_handler_threaded(*_: object) -> None: # too_many_pings; see cfdb.workflows.grpc_options. options=worker_grpc_options(), ) + idle_conn: Optional["wool.WorkerConnection"] = None + # Grace passed to ``worker.stop()`` in the teardown. ``None`` is + # wool's immediate-cancel; the self-termination exits below replace + # it with ``max_lifetime_grace_seconds`` so their teardown drains + # in-flight work, while the signal paths keep the immediate cancel. + stop_grace: Optional[float] = None await worker.start() try: + if idle_timeout_seconds > 0: + # Dial our own worker over loopback to poll its idle RPC — + # the same channel-back-to-own-subprocess pattern wool uses + # for graceful stop, so the shared credentials (and their + # identity SAN) verify the same way they do on that channel. + idle_conn = wool.WorkerConnection( + f"127.0.0.1:{worker_port}", credentials=credentials + ) # Publish before entering the serve loop. The worker is already # accepting gRPC by now, but until its tags land EcsDiscovery # deliberately will not advertise it, so there is no window in @@ -437,25 +558,39 @@ def _signal_handler_threaded(*_: object) -> None: "enabled" if credentials is not None else "disabled", WORKER_MAX_CONCURRENT_TASKS if backpressure is not None else "unbounded", ) + next_idle_poll = loop.time() + idle_poll_failures = 0 while True: - # Check the self-termination path first. Max-lifetime is a - # local hard ceiling: the gap between ``stop_event.set()`` - # and ``worker.stop()`` is microseconds — no health probe - # will actually fire during it, so this is a defense-in- - # depth flip rather than a real drain window. Operators - # that need a true drain handoff on max-lifetime should - # roll tasks via ECS service updates instead of relying - # on the self-timeout. Flipping /health to 503 still costs - # nothing and keeps the in-process ordering consistent - # with the signal path's drain semantics. + # Check the self-termination paths first. Both skip the + # signal path's /health drain window (the gap between + # ``stop_event.set()`` and ``worker.stop()`` is microseconds + # — no health probe will actually fire during it) but set + # ``stop_grace`` so the teardown *drains* in-flight work + # instead of cancelling it. That grace is what makes these + # exits safe: an idle snapshot cannot rule out a dispatch + # accepted between the final poll and the stop, and a + # max-lifetime expiry can land mid-job on a busy worker — + # in either case the accepted task has already been marked + # running on the API side, where a graceless cancel would + # finalize the job as terminally failed rather than + # re-queue it. With the grace, in-flight work runs to + # completion (instantly when the docket is truly empty) + # and only the drain's own timeout — sized above the API's + # per-job duration cap — ever cancels anything. The signal + # paths leave ``stop_grace`` at wool's immediate cancel: + # ECS bounds SIGTERM with SIGKILL, so a long drain there + # could never complete anyway. if ( max_lifetime_seconds > 0 and (loop.time() - started_at) >= max_lifetime_seconds ): logger.info( - "Max lifetime (%.0fs) reached — exiting", + "Max lifetime (%.0fs) reached — draining in-flight " + "work for up to %.0fs, then exiting", max_lifetime_seconds, + max_lifetime_grace_seconds, ) + stop_grace = max_lifetime_grace_seconds stop_event.set() break if stop_event.is_set(): @@ -476,6 +611,62 @@ def _signal_handler_threaded(*_: object) -> None: except asyncio.TimeoutError: pass break + if idle_conn is not None and loop.time() >= next_idle_poll: + next_idle_poll = loop.time() + idle_poll_interval_seconds + try: + idle = await idle_conn.idle( + timeout=_IDLE_POLL_RPC_TIMEOUT_SECONDS + ) + except wool.IdleUnavailable: + # Structurally unreachable when the worker is this + # same wool install, but a skew scenario must not + # crash-loop the poll: fall back to the + # max-lifetime backstop. The connection stays + # referenced so the teardown below still closes it. + logger.warning( + "Worker does not implement idle reporting — " + "disabling idle shutdown; %s", + _lifetime_bound_description(max_lifetime_seconds), + ) + next_idle_poll = float("inf") + except Exception as exc: + # Transient or unexpected poll failure. A busy + # worker must never die to a flaky poll, so retry + # on the next cadence — but sustained failure is + # not a blip (a permanent TLS misconfiguration or + # a dead subprocess looks exactly like this), so + # after enough consecutive misses escalate once to + # ERROR and stop polling: one actionable CloudWatch + # signal instead of hours of identical warnings. + idle_poll_failures += 1 + if idle_poll_failures >= idle_poll_failure_limit: + logger.error( + "Idle poll failed %d consecutive times " + "(last: %s: %s) — disabling idle shutdown; %s", + idle_poll_failures, + type(exc).__name__, + exc, + _lifetime_bound_description(max_lifetime_seconds), + ) + next_idle_poll = float("inf") + else: + logger.warning( + "Idle poll failed (%s: %s), retrying in %.0fs", + type(exc).__name__, + exc, + idle_poll_interval_seconds, + ) + else: + idle_poll_failures = 0 + if idle >= idle_timeout_seconds: + logger.info( + "Idle for %.0fs (threshold %.0fs) — exiting", + idle, + idle_timeout_seconds, + ) + stop_grace = max_lifetime_grace_seconds + stop_event.set() + break try: await asyncio.wait_for( stop_event.wait(), timeout=_STOP_POLL_INTERVAL_SECONDS @@ -484,8 +675,13 @@ def _signal_handler_threaded(*_: object) -> None: continue return 0 finally: + if idle_conn is not None: + try: + await idle_conn.close() + except Exception: + logger.exception("idle connection close failed during shutdown") try: - await worker.stop() + await worker.stop(grace=stop_grace) except Exception: logger.exception("worker.stop() failed during shutdown") await _shutdown_health_server(health_runner) @@ -560,13 +756,61 @@ async def _shutdown_health_server(runner: Optional["web.AppRunner"]) -> None: show_default=True, help="HTTP port the ECS health-check endpoint binds.", ) +@click.option( + "--idle-timeout-seconds", + type=click.FloatRange(min=0), + envvar="CFDB_WORKER_IDLE_TIMEOUT_SECONDS", + default=DEFAULT_IDLE_TIMEOUT_SECONDS, + show_default=True, + help=( + "Continuous idle seconds beyond which the worker exits; " + "0 disables idle shutdown." + ), +) +@click.option( + "--idle-poll-interval-seconds", + type=click.FloatRange(min=0), + envvar="CFDB_WORKER_IDLE_POLL_INTERVAL_SECONDS", + default=DEFAULT_IDLE_POLL_INTERVAL_SECONDS, + show_default=True, + help=( + "Cadence of the idle poll in seconds. Values below the loop's " + "1 s wakeup are effectively floored at 1 s." + ), +) +@click.option( + "--idle-poll-failure-limit", + type=click.IntRange(min=1), + envvar="CFDB_WORKER_IDLE_POLL_FAILURE_LIMIT", + default=DEFAULT_IDLE_POLL_FAILURE_LIMIT, + show_default=True, + help=( + "Consecutive idle-poll failures before escalating once to " + "ERROR and disabling idle shutdown for this worker." + ), +) @click.option( "--max-lifetime-seconds", type=click.FloatRange(min=0), envvar="CFDB_WORKER_MAX_LIFETIME_SECONDS", default=DEFAULT_MAX_LIFETIME_SECONDS, show_default=True, - help="Hard ceiling on worker uptime in seconds; 0 disables.", + help=( + "Hard ceiling on worker uptime in seconds; 0 disables. The " + "backstop behind the idle timeout, not the primary reaper." + ), +) +@click.option( + "--max-lifetime-grace-seconds", + type=click.FloatRange(min=0), + envvar="CFDB_WORKER_MAX_LIFETIME_GRACE_SECONDS", + default=DEFAULT_MAX_LIFETIME_GRACE_SECONDS, + show_default=True, + help=( + "Seconds a self-terminating exit (idle or max-lifetime) drains " + "in-flight work before cancelling it. Worst-case uptime is " + "max lifetime plus this grace." + ), ) @click.option( "--drain-grace-seconds", @@ -619,7 +863,11 @@ async def _shutdown_health_server(runner: Optional["web.AppRunner"]) -> None: def main( worker_port: int, health_port: int, + idle_timeout_seconds: float, + idle_poll_interval_seconds: float, + idle_poll_failure_limit: int, max_lifetime_seconds: float, + max_lifetime_grace_seconds: float, drain_grace_seconds: float, tls_ca: Optional[str], tls_cert: Optional[str], @@ -630,7 +878,8 @@ def main( """ECS Fargate worker entrypoint — invoked by the container CMD. Boots a wool gRPC worker, exposes /health for ECS to probe, and - self-terminates after the max-lifetime ceiling. SIGTERM begins a + self-terminates once continuously idle beyond the idle timeout + (with the max-lifetime ceiling as a backstop). SIGTERM begins a drain grace window during which /health returns 503 so the load balancer can drop the worker before the gRPC port closes. """ @@ -640,7 +889,11 @@ def main( serve( worker_port=worker_port, health_port=health_port, + idle_timeout_seconds=idle_timeout_seconds, + idle_poll_interval_seconds=idle_poll_interval_seconds, + idle_poll_failure_limit=idle_poll_failure_limit, max_lifetime_seconds=max_lifetime_seconds, + max_lifetime_grace_seconds=max_lifetime_grace_seconds, drain_grace_seconds=drain_grace_seconds, tls_ca=tls_ca, tls_cert=tls_cert, From 4bc8735d9a15a572871a529be2a35350269cda29 Mon Sep 17 00:00:00 2001 From: Conrad Date: Mon, 10 Aug 2026 15:32:21 -0400 Subject: [PATCH 3/7] test: Cover idle shutdown in the worker entrypoint MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Extends the CLI tests with the new idle and lifetime knobs and adds a suite exercising the serve loop's idle path: exit on crossing the threshold, immunity while busy, idle accumulating below the threshold, the 0-disables sentinel, the IdleUnavailable fallback to max-lifetime, transient and non-transient poll-failure retry, escalation after consecutive failures, the loopback dial with the worker's own credentials, and connection teardown including a failing close. The exit-path assertions name the path taken rather than only the outcome, since returning 0 with the worker stopped is equally true of the max-lifetime backstop — a broken idle comparison would otherwise still pass, just slower. Both self-termination tests also pin the drain grace passed to stop, which is what keeps a task racing the teardown from being cancelled into a terminal job failure. The health server is no longer patched out: these tests already bind port 0, so the real one comes up on an ephemeral port and the helper stops reaching into module privates. --- tests/test_workflows/test_worker_main.py | 492 ++++++++++++++++++++++- 1 file changed, 491 insertions(+), 1 deletion(-) diff --git a/tests/test_workflows/test_worker_main.py b/tests/test_workflows/test_worker_main.py index fecb527..3211def 100644 --- a/tests/test_workflows/test_worker_main.py +++ b/tests/test_workflows/test_worker_main.py @@ -5,8 +5,10 @@ import logging from unittest.mock import patch +import grpc import pytest import pytest_asyncio +import wool from botocore.exceptions import ClientError from click.testing import CliRunner @@ -51,7 +53,11 @@ def test_main_uses_documented_defaults_when_no_args_or_env(self, monkeypatch): for var in ( "CFDB_WORKER_GRPC_PORT", "CFDB_WORKER_HEALTH_PORT", + "CFDB_WORKER_IDLE_TIMEOUT_SECONDS", + "CFDB_WORKER_IDLE_POLL_INTERVAL_SECONDS", + "CFDB_WORKER_IDLE_POLL_FAILURE_LIMIT", "CFDB_WORKER_MAX_LIFETIME_SECONDS", + "CFDB_WORKER_MAX_LIFETIME_GRACE_SECONDS", "CFDB_WORKER_DRAIN_GRACE_SECONDS", "CFDB_WORKER_TLS_CA", "CFDB_WORKER_TLS_CERT", @@ -66,7 +72,20 @@ def test_main_uses_documented_defaults_when_no_args_or_env(self, monkeypatch): assert exit_code == 0 assert captured["worker_port"] == worker_main.DEFAULT_WORKER_PORT assert captured["health_port"] == worker_main.DEFAULT_HEALTH_PORT + assert captured["idle_timeout_seconds"] == worker_main.DEFAULT_IDLE_TIMEOUT_SECONDS + assert ( + captured["idle_poll_interval_seconds"] + == worker_main.DEFAULT_IDLE_POLL_INTERVAL_SECONDS + ) + assert ( + captured["idle_poll_failure_limit"] + == worker_main.DEFAULT_IDLE_POLL_FAILURE_LIMIT + ) assert captured["max_lifetime_seconds"] == worker_main.DEFAULT_MAX_LIFETIME_SECONDS + assert ( + captured["max_lifetime_grace_seconds"] + == worker_main.DEFAULT_MAX_LIFETIME_GRACE_SECONDS + ) assert captured["drain_grace_seconds"] == worker_main.DEFAULT_DRAIN_GRACE_SECONDS assert captured["tls_ca"] is None assert captured["tls_cert"] is None @@ -85,7 +104,11 @@ def test_main_with_env_overrides(self, monkeypatch): # Arrange monkeypatch.setenv("CFDB_WORKER_GRPC_PORT", "60001") monkeypatch.setenv("CFDB_WORKER_HEALTH_PORT", "9001") + monkeypatch.setenv("CFDB_WORKER_IDLE_TIMEOUT_SECONDS", "60") + monkeypatch.setenv("CFDB_WORKER_IDLE_POLL_INTERVAL_SECONDS", "5") + monkeypatch.setenv("CFDB_WORKER_IDLE_POLL_FAILURE_LIMIT", "7") monkeypatch.setenv("CFDB_WORKER_MAX_LIFETIME_SECONDS", "1800") + monkeypatch.setenv("CFDB_WORKER_MAX_LIFETIME_GRACE_SECONDS", "900") monkeypatch.setenv("CFDB_WORKER_DRAIN_GRACE_SECONDS", "10") # Act @@ -95,7 +118,11 @@ def test_main_with_env_overrides(self, monkeypatch): assert exit_code == 0 assert captured["worker_port"] == 60001 assert captured["health_port"] == 9001 + assert captured["idle_timeout_seconds"] == 60.0 + assert captured["idle_poll_interval_seconds"] == 5.0 + assert captured["idle_poll_failure_limit"] == 7 assert captured["max_lifetime_seconds"] == 1800.0 + assert captured["max_lifetime_grace_seconds"] == 900.0 assert captured["drain_grace_seconds"] == 10.0 def test_main_cli_flags_override_env_vars(self, monkeypatch): @@ -111,13 +138,26 @@ def test_main_cli_flags_override_env_vars(self, monkeypatch): """ # Arrange monkeypatch.setenv("CFDB_WORKER_GRPC_PORT", "60001") + monkeypatch.setenv("CFDB_WORKER_IDLE_TIMEOUT_SECONDS", "60") + monkeypatch.setenv("CFDB_WORKER_MAX_LIFETIME_GRACE_SECONDS", "7200") # Act - exit_code, captured = _invoke(["--worker-port", "55555"]) + exit_code, captured = _invoke( + [ + "--worker-port", + "55555", + "--idle-timeout-seconds", + "30", + "--max-lifetime-grace-seconds", + "3600", + ] + ) # Assert assert exit_code == 0 assert captured["worker_port"] == 55555 + assert captured["idle_timeout_seconds"] == 30.0 + assert captured["max_lifetime_grace_seconds"] == 3600.0 def test_main_rejects_out_of_range_worker_port(self, monkeypatch): """Test that a port outside [1, 65535] is rejected at parse time. @@ -271,6 +311,456 @@ async def test_serve_should_disable_backpressure_when_threshold_zero( assert local_worker.call_args.kwargs["backpressure"] is None +def _arrange_idle_serve(mocker, monkeypatch, *, idle_effect): + """Patch ``serve``'s collaborators for a full run-loop pass. + + Unlike ``_arrange_serve`` the fake worker starts cleanly, so ``serve`` + enters its run loop and exercises the idle-poll path for real; the + health server is left unpatched — tests pass ``health_port=0`` so it + binds an ephemeral port for the run's duration. ``idle_effect`` + becomes the fake connection's ``idle`` side effect. Returns + ``(worker, connection_cls, connection)``. + """ + monkeypatch.delenv("ECS_CONTAINER_METADATA_URI_V4", raising=False) + mocker.patch.object(worker_main, "build_worker_credentials", return_value=None) + worker_instance = mocker.Mock() + worker_instance.start = mocker.AsyncMock() + worker_instance.stop = mocker.AsyncMock() + mocker.patch.object( + worker_main.wool, "LocalWorker", return_value=worker_instance + ) + connection = mocker.Mock() + connection.idle = mocker.AsyncMock(side_effect=idle_effect) + connection.close = mocker.AsyncMock() + connection_cls = mocker.patch.object( + worker_main.wool, "WorkerConnection", return_value=connection + ) + return worker_instance, connection_cls, connection + + +class TestServeIdleShutdown: + @pytest.mark.asyncio + async def test_serve_should_exit_when_idle_exceeds_timeout( + self, mocker, monkeypatch, caplog + ): + """Test that crossing the idle threshold shuts the worker down. + + Given: + A worker whose idle RPC reports more continuous idle time + than the configured idle timeout. + When: + ``serve`` is run. + Then: + It should exit through the idle path on the first poll — + with a bounded per-poll RPC deadline, and stopping the + worker with the drain grace so a dispatch racing the + teardown completes instead of being cancelled. + """ + # Arrange + worker, _, connection = _arrange_idle_serve( + mocker, monkeypatch, idle_effect=[10.0] + ) + + # Act + with caplog.at_level(logging.INFO, logger=worker_main.__name__): + result = await worker_main.serve( + worker_port=0, + health_port=0, + idle_timeout_seconds=5.0, + idle_poll_interval_seconds=0.01, + max_lifetime_seconds=60.0, + max_lifetime_grace_seconds=7.5, + ) + + # Assert + assert result == 0 + assert any("Idle for" in record.message for record in caplog.records) + assert connection.idle.await_count == 1 + poll_timeout = connection.idle.await_args.kwargs["timeout"] + assert 0 < poll_timeout < float("inf") + worker.stop.assert_awaited_once_with(grace=7.5) + + @pytest.mark.asyncio + async def test_serve_should_keep_serving_when_worker_busy( + self, mocker, monkeypatch, caplog + ): + """Test that a busy worker is never reaped by the idle path. + + Given: + A worker whose idle RPC always reports zero (work in + flight) and a short max lifetime. + When: + ``serve`` is run. + Then: + It should poll idle at least once without exiting on it and + terminate via the max-lifetime backstop instead — stopping + the worker with the drain grace so the in-flight job the + expiry interrupted completes rather than being cancelled. + """ + # Arrange + worker, _, connection = _arrange_idle_serve( + mocker, monkeypatch, idle_effect=lambda **_: 0.0 + ) + + # Act + with caplog.at_level(logging.INFO, logger=worker_main.__name__): + result = await worker_main.serve( + worker_port=0, + health_port=0, + idle_timeout_seconds=5.0, + idle_poll_interval_seconds=0.01, + max_lifetime_seconds=0.5, + max_lifetime_grace_seconds=7.5, + ) + + # Assert + assert result == 0 + assert connection.idle.await_count >= 1 + assert any("Max lifetime" in record.message for record in caplog.records) + worker.stop.assert_awaited_once_with(grace=7.5) + + @pytest.mark.asyncio + async def test_serve_should_not_dial_idle_connection_when_timeout_zero( + self, mocker, monkeypatch + ): + """Test that the disable sentinel skips the idle connection. + + Given: + ``idle_timeout_seconds`` set to 0 and a short max lifetime. + When: + ``serve`` is run. + Then: + It should never construct a ``WorkerConnection``, restoring + the pure max-lifetime behavior. + """ + # Arrange + _, connection_cls, _ = _arrange_idle_serve( + mocker, monkeypatch, idle_effect=[0.0] + ) + + # Act + result = await worker_main.serve( + worker_port=0, + health_port=0, + idle_timeout_seconds=0.0, + max_lifetime_seconds=1e-6, + ) + + # Assert + assert result == 0 + connection_cls.assert_not_called() + + @pytest.mark.asyncio + async def test_serve_should_fall_back_to_max_lifetime_when_idle_rpc_unimplemented( + self, mocker, monkeypatch + ): + """Test that a worker without the idle RPC disables idle polling. + + Given: + A worker whose idle RPC raises ``IdleUnavailable`` and a + short max lifetime. + When: + ``serve`` is run. + Then: + It should poll exactly once, keep serving, and exit via the + max-lifetime backstop — a version-skew scenario must not + crash-loop the poll. + """ + # Arrange + _, _, connection = _arrange_idle_serve( + mocker, + monkeypatch, + idle_effect=wool.IdleUnavailable("no idle rpc"), + ) + + # Act + result = await worker_main.serve( + worker_port=0, + health_port=0, + idle_timeout_seconds=5.0, + idle_poll_interval_seconds=0.01, + max_lifetime_seconds=0.5, + ) + + # Assert + assert result == 0 + assert connection.idle.await_count == 1 + connection.close.assert_awaited_once() + + @pytest.mark.asyncio + async def test_serve_should_keep_polling_when_idle_poll_fails_transiently( + self, mocker, monkeypatch, caplog + ): + """Test that a flaky idle poll is retried rather than fatal. + + Given: + A worker whose idle RPC fails transiently once and then + reports idle time beyond the threshold. + When: + ``serve`` is run. + Then: + It should survive the failed poll and exit via the idle + path on the next cadence, so a flaky poll never kills a + worker. + """ + # Arrange + _, _, connection = _arrange_idle_serve( + mocker, + monkeypatch, + idle_effect=[ + wool.TransientRpcError(grpc.StatusCode.UNAVAILABLE, "poll failed"), + 10.0, + ], + ) + + # Act + with caplog.at_level(logging.INFO, logger=worker_main.__name__): + result = await worker_main.serve( + worker_port=0, + health_port=0, + idle_timeout_seconds=5.0, + idle_poll_interval_seconds=0.01, + max_lifetime_seconds=60.0, + ) + + # Assert + assert result == 0 + assert connection.idle.await_count == 2 + assert any("Idle for" in record.message for record in caplog.records) + + @pytest.mark.asyncio + async def test_serve_should_dial_loopback_with_the_worker_credentials( + self, mocker, monkeypatch, caplog + ): + """Test that the idle connection targets the worker's own port. + + Given: + A credentials builder returning a sentinel object and a + non-default worker port. + When: + ``serve`` is run until the idle path exits it. + Then: + It should construct the ``WorkerConnection`` against + loopback at the worker port with those same credentials, so + the idle poll verifies mTLS the same way wool's own + drain channel does. + """ + # Arrange + credentials = object() + _, connection_cls, _ = _arrange_idle_serve( + mocker, monkeypatch, idle_effect=[10.0] + ) + mocker.patch.object( + worker_main, "build_worker_credentials", return_value=credentials + ) + + # Act + with caplog.at_level(logging.INFO, logger=worker_main.__name__): + await worker_main.serve( + worker_port=50055, + health_port=0, + idle_timeout_seconds=5.0, + idle_poll_interval_seconds=0.01, + max_lifetime_seconds=60.0, + ) + + # Assert + connection_cls.assert_called_once_with( + "127.0.0.1:50055", credentials=credentials + ) + assert any("Idle for" in record.message for record in caplog.records) + + @pytest.mark.asyncio + async def test_serve_should_close_idle_connection_on_shutdown( + self, mocker, monkeypatch, caplog + ): + """Test that shutdown releases the idle connection's resources. + + Given: + A worker that exits via the idle path. + When: + ``serve`` returns. + Then: + It should close the ``WorkerConnection`` so pooled channels + are released alongside the worker's own teardown. + """ + # Arrange + _, _, connection = _arrange_idle_serve( + mocker, monkeypatch, idle_effect=[10.0] + ) + + # Act + with caplog.at_level(logging.INFO, logger=worker_main.__name__): + await worker_main.serve( + worker_port=0, + health_port=0, + idle_timeout_seconds=5.0, + idle_poll_interval_seconds=0.01, + max_lifetime_seconds=60.0, + ) + + # Assert + connection.close.assert_awaited_once() + assert any("Idle for" in record.message for record in caplog.records) + + @pytest.mark.asyncio + async def test_serve_should_disable_idle_polling_when_poll_fails_persistently( + self, mocker, monkeypatch, caplog + ): + """Test that sustained poll failure escalates once and stops polling. + + Given: + A worker whose idle RPC fails on every poll and a failure + limit of two. + When: + ``serve`` is run. + Then: + It should emit exactly one ERROR naming the max-lifetime + bound, poll no further, and exit via the backstop — one + actionable signal instead of hours of identical warnings. + """ + # Arrange + _, _, connection = _arrange_idle_serve( + mocker, monkeypatch, idle_effect=RuntimeError("subprocess dead") + ) + + # Act + with caplog.at_level(logging.INFO, logger=worker_main.__name__): + result = await worker_main.serve( + worker_port=0, + health_port=0, + idle_timeout_seconds=5.0, + idle_poll_interval_seconds=0.01, + idle_poll_failure_limit=2, + max_lifetime_seconds=2.5, + ) + + # Assert + assert result == 0 + errors = [r for r in caplog.records if r.levelno == logging.ERROR] + assert len(errors) == 1 + assert "disabling idle shutdown" in errors[0].getMessage() + assert "max-lifetime backstop" in errors[0].getMessage() + assert connection.idle.await_count == 2 + + @pytest.mark.asyncio + async def test_serve_should_keep_polling_when_idle_poll_fails_nontransiently( + self, mocker, monkeypatch, caplog + ): + """Test that isolated failures below the limit never escalate. + + Given: + A worker whose idle RPC fails non-transiently, succeeds + (resetting the consecutive-failure count), fails again, and + then reports idle beyond the threshold — with a failure + limit of two. + When: + ``serve`` is run. + Then: + It should retry through both isolated failures without + escalating to ERROR and exit via the idle path, so only + *consecutive* failures count toward the limit. + """ + # Arrange + _, _, connection = _arrange_idle_serve( + mocker, + monkeypatch, + idle_effect=[ + RuntimeError("blip"), + 0.0, + RuntimeError("blip"), + 10.0, + ], + ) + + # Act + with caplog.at_level(logging.INFO, logger=worker_main.__name__): + result = await worker_main.serve( + worker_port=0, + health_port=0, + idle_timeout_seconds=5.0, + idle_poll_interval_seconds=0.01, + idle_poll_failure_limit=2, + max_lifetime_seconds=60.0, + ) + + # Assert + assert result == 0 + assert connection.idle.await_count == 4 + assert not [r for r in caplog.records if r.levelno == logging.ERROR] + assert any("Idle for" in record.message for record in caplog.records) + + @pytest.mark.asyncio + async def test_serve_should_keep_serving_when_idle_below_threshold( + self, mocker, monkeypatch, caplog + ): + """Test that partial idle accumulation does not trigger the exit. + + Given: + A worker whose idle RPC reports idle time strictly between + zero and the threshold, and a short max lifetime. + When: + ``serve`` is run. + Then: + It should keep serving through those polls and exit via the + max-lifetime backstop, so a worker inside an ordinary + dispatch gap is not reaped early. + """ + # Arrange + _, _, connection = _arrange_idle_serve( + mocker, monkeypatch, idle_effect=lambda **_: 3.0 + ) + + # Act + with caplog.at_level(logging.INFO, logger=worker_main.__name__): + result = await worker_main.serve( + worker_port=0, + health_port=0, + idle_timeout_seconds=5.0, + idle_poll_interval_seconds=0.01, + max_lifetime_seconds=0.5, + ) + + # Assert + assert result == 0 + assert connection.idle.await_count >= 1 + assert any("Max lifetime" in record.message for record in caplog.records) + + @pytest.mark.asyncio + async def test_serve_should_stop_worker_when_idle_connection_close_fails( + self, mocker, monkeypatch + ): + """Test that a close failure does not skip the worker teardown. + + Given: + A worker exiting via the idle path whose ``WorkerConnection`` + raises on ``close``. + When: + ``serve`` runs to completion. + Then: + It should still return 0 and stop the worker, so a teardown + hiccup on the poll channel never leaks the worker itself. + """ + # Arrange + worker, _, connection = _arrange_idle_serve( + mocker, monkeypatch, idle_effect=[10.0] + ) + connection.close = mocker.AsyncMock(side_effect=RuntimeError("close failed")) + + # Act + result = await worker_main.serve( + worker_port=0, + health_port=0, + idle_timeout_seconds=5.0, + idle_poll_interval_seconds=0.01, + max_lifetime_seconds=60.0, + ) + + # Assert + assert result == 0 + worker.stop.assert_awaited_once() + + #: Task ARN the fake ECS metadata endpoint reports. _TASK_ARN = "arn:aws:ecs:us-east-2:605134458779:task/cfdb-cluster/abc123" From 89b78c0593fe6fb51ce621e552035773c6e11099 Mon Sep 17 00:00:00 2001 From: Conrad Date: Mon, 10 Aug 2026 15:32:27 -0400 Subject: [PATCH 4/7] feat: Wire the worker reaping knobs through the workers stack Adds a WorkerIdleTimeoutSeconds parameter (default 600) feeding CFDB_WORKER_IDLE_TIMEOUT_SECONDS in the worker container, and a WorkerMaxLifetimeGraceSeconds parameter (default 21600) feeding CFDB_WORKER_MAX_LIFETIME_GRACE_SECONDS, the drain a self-terminating exit grants in-flight work before cancelling it. The max-lifetime ceiling rises to 43200 to match the code default, which the grace makes safe. The idle-timeout description also records the tuning constraint that is otherwise only discoverable by hitting it: an idle timeout at or below the scheduler's retry cadence plus worker cold-start time lets overflow-spawned workers reap themselves before the next retry tick reaches them, and the jobs ride the queue to the capacity deadline. Landing the parameters requires a one-time privileged cloudformation deploy per environment, as with every workers-stack change. --- cloudformation/workers.yml | 38 +++++++++++++++++++++++++++++++++++--- 1 file changed, 35 insertions(+), 3 deletions(-) diff --git a/cloudformation/workers.yml b/cloudformation/workers.yml index c8c812f..59c1884 100644 --- a/cloudformation/workers.yml +++ b/cloudformation/workers.yml @@ -52,14 +52,42 @@ Parameters: Days after which cached artifacts expire. Cache keys are content-addressed by upstream md5, so expiry is safe — a missing artifact is simply re-materialized on the next request. + WorkerIdleTimeoutSeconds: + Type: String + Default: "600" + AllowedPattern: "^[0-9]+$" + Description: > + Continuous idle seconds beyond which a worker exits + (CFDB_WORKER_IDLE_TIMEOUT_SECONDS); 0 disables idle shutdown. + The primary reaper — a busy worker reports zero idle, so the + idle exit never fires mid-task, and a dispatch racing the + teardown drains under WorkerMaxLifetimeGraceSeconds. Keep this + comfortably above CFDB_WORKFLOW_RETRY_INTERVAL_S (default 120 s) + plus worker cold-start time (at least 2-3x), or overflow-spawned + workers idle-reap before the scheduler's next retry tick reaches + them and jobs ride the queue to the capacity deadline. WorkerMaxLifetimeSeconds: Type: String - Default: "18000" + Default: "43200" AllowedPattern: "^[0-9]+$" Description: > Hard ceiling on worker uptime (CFDB_WORKER_MAX_LIFETIME_SECONDS); - one hour above the 4 h workflow duration cap so a worker started - shortly before a long sort can still outlive the job. + the backstop behind the idle timeout for a worker whose idle + reporting is wedged or whose job is stuck. Expiry drains + in-flight work for up to WorkerMaxLifetimeGraceSeconds rather + than cancelling it, so worst-case worker uptime is this ceiling + plus that grace (18 h at the defaults). + WorkerMaxLifetimeGraceSeconds: + Type: String + Default: "21600" + AllowedPattern: "^[0-9]+$" + Description: > + Seconds a self-terminating exit — idle timeout or max lifetime — + drains in-flight work before cancelling it + (CFDB_WORKER_MAX_LIFETIME_GRACE_SECONDS). Instant when the worker + is genuinely idle; sized above the 4 h workflow duration cap so + the only work ever cancelled is work already past the API's own + viability bound. WorkerDrainGraceSeconds: Type: String Default: "120" @@ -360,8 +388,12 @@ Resources: Value: "50051" - Name: CFDB_WORKER_HEALTH_PORT Value: "8080" + - Name: CFDB_WORKER_IDLE_TIMEOUT_SECONDS + Value: !Ref WorkerIdleTimeoutSeconds - Name: CFDB_WORKER_MAX_LIFETIME_SECONDS Value: !Ref WorkerMaxLifetimeSeconds + - Name: CFDB_WORKER_MAX_LIFETIME_GRACE_SECONDS + Value: !Ref WorkerMaxLifetimeGraceSeconds - Name: CFDB_WORKER_DRAIN_GRACE_SECONDS Value: !Ref WorkerDrainGraceSeconds # The worker tags its own task at startup (worker-metadata From 23b6a7d3045ca12793c3c0b6027e95dc0df5a89c Mon Sep 17 00:00:00 2001 From: Conrad Date: Mon, 10 Aug 2026 15:32:31 -0400 Subject: [PATCH 5/7] test: Pin the workers-stack worker-reaping wiring and defaults The template must both wire each reaping knob from its parameter and keep the parameter default equal to the code default; a miss on either silently deploys a fleet whose reaping cadence disagrees with what the code documents. Covers the idle timeout, the max-lifetime ceiling, and the self-termination drain grace. --- tests/test_cloudformation.py | 123 +++++++++++++++++++++++++++++++++++ 1 file changed, 123 insertions(+) diff --git a/tests/test_cloudformation.py b/tests/test_cloudformation.py index 3c7682d..5fd8e20 100644 --- a/tests/test_cloudformation.py +++ b/tests/test_cloudformation.py @@ -325,6 +325,129 @@ def test_worker_container_should_declare_aws_region(): assert plain_env.get("AWS_REGION") == {"Ref": "AWS::Region"} +def test_worker_idle_timeout_should_be_wired_and_agree_with_the_application_default(): + """Test that the idle-timeout knob reaches the worker and matches the code. + + Given: + The workers template's WorkerIdleTimeoutSeconds parameter and + container definition. The worker's primary reaper is its idle + timeout; a template that failed to wire the env var would + silently leave every deployed worker on the code default, and a + drifted template default would make the deployed fleet disagree + with what the code documents. + When: + The container's Environment entries and the parameter default + are read. + Then: + ``CFDB_WORKER_IDLE_TIMEOUT_SECONDS`` should reference the + parameter, and the parameter default should equal + ``worker_main.DEFAULT_IDLE_TIMEOUT_SECONDS``. + """ + # Arrange + from cfdb.workflows.worker_main import DEFAULT_IDLE_TIMEOUT_SECONDS + + template = _load_template("workers.yml") + container = template["Resources"]["WorkerTaskDefinition"]["Properties"][ + "ContainerDefinitions" + ][0] + + # Act + plain_env = { + entry["Name"]: entry.get("Value") + for entry in container["Environment"] + if isinstance(entry, dict) and "Name" in entry + } + default = template["Parameters"]["WorkerIdleTimeoutSeconds"]["Default"] + + # Assert + assert plain_env.get("CFDB_WORKER_IDLE_TIMEOUT_SECONDS") == { + "Ref": "WorkerIdleTimeoutSeconds" + } + assert float(default) == DEFAULT_IDLE_TIMEOUT_SECONDS + + +def test_worker_max_lifetime_should_be_wired_and_agree_with_the_application_default(): + """Test that the max-lifetime knob reaches the worker and matches the code. + + Given: + The workers template's WorkerMaxLifetimeSeconds parameter and + container definition. The ceiling is the backstop behind idle + shutdown; a template that failed to wire the env var would + silently leave every deployed worker on the code default, and a + drifted template default would make the deployed fleet disagree + with what the code documents. + When: + The container's Environment entries and the parameter default + are read. + Then: + ``CFDB_WORKER_MAX_LIFETIME_SECONDS`` should reference the + parameter, and the parameter default should equal + ``worker_main.DEFAULT_MAX_LIFETIME_SECONDS``. + """ + # Arrange + from cfdb.workflows.worker_main import DEFAULT_MAX_LIFETIME_SECONDS + + template = _load_template("workers.yml") + container = template["Resources"]["WorkerTaskDefinition"]["Properties"][ + "ContainerDefinitions" + ][0] + + # Act + plain_env = { + entry["Name"]: entry.get("Value") + for entry in container["Environment"] + if isinstance(entry, dict) and "Name" in entry + } + default = template["Parameters"]["WorkerMaxLifetimeSeconds"]["Default"] + + # Assert + assert plain_env.get("CFDB_WORKER_MAX_LIFETIME_SECONDS") == { + "Ref": "WorkerMaxLifetimeSeconds" + } + assert float(default) == DEFAULT_MAX_LIFETIME_SECONDS + + +def test_worker_max_lifetime_grace_should_be_wired_and_agree_with_the_application_default(): + """Test that the self-termination grace reaches the worker and matches the code. + + Given: + The workers template's WorkerMaxLifetimeGraceSeconds parameter + and container definition. The grace is what makes the idle and + max-lifetime exits drain in-flight work instead of cancelling + it; a template that failed to wire the env var or drifted from + the code default would silently change how much drain a deployed + worker actually grants. + When: + The container's Environment entries and the parameter default + are read. + Then: + ``CFDB_WORKER_MAX_LIFETIME_GRACE_SECONDS`` should reference the + parameter, and the parameter default should equal + ``worker_main.DEFAULT_MAX_LIFETIME_GRACE_SECONDS``. + """ + # Arrange + from cfdb.workflows.worker_main import DEFAULT_MAX_LIFETIME_GRACE_SECONDS + + template = _load_template("workers.yml") + container = template["Resources"]["WorkerTaskDefinition"]["Properties"][ + "ContainerDefinitions" + ][0] + + # Act + plain_env = { + entry["Name"]: entry.get("Value") + for entry in container["Environment"] + if isinstance(entry, dict) and "Name" in entry + } + default = template["Parameters"]["WorkerMaxLifetimeGraceSeconds"]["Default"] + + # Assert + assert plain_env.get("CFDB_WORKER_MAX_LIFETIME_GRACE_SECONDS") == { + "Ref": "WorkerMaxLifetimeGraceSeconds" + } + assert float(default) == DEFAULT_MAX_LIFETIME_GRACE_SECONDS + + def test_worker_tls_identity_should_only_render_when_mtls_is_enabled(): """Test that the worker identity env var is gated on mTLS. From 63f913aa000b95294c5d3e5e89d9846df232468a Mon Sep 17 00:00:00 2001 From: Conrad Date: Mon, 10 Aug 2026 15:32:38 -0400 Subject: [PATCH 6/7] docs: Reframe worker reaping around the idle timeout MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The leaky-bucket and version-bump-wedge sections and the provisioner's orphan-worker note all described max-lifetime as the only reaper; idle shutdown is now the primary one, shrinking the wedge window from hours to minutes once the fleet runs an idle-aware image. The wedge section also notes that a bump from a pre-idle fleet still needs the manual drain, since old workers cannot self-reap early, and that the eventual refresh of the pre-release wool pin is itself such a bump. The worker-knobs paragraph states what the reaping guarantee actually is — a busy worker reports zero idle, and a task that races the teardown drains rather than dying — along with the resulting worst-case worker uptime and the sizing constraint on the idle timeout relative to the scheduler's retry cadence. The mTLS identity section names the idle poll as a third consumer of the identity contract, since a worker now dials its own subprocess over loopback: under the address-verification opt-out the worker leaf needs a loopback SAN, or every poll fails its handshake and idle shutdown quietly degrades to the backstop. --- README.md | 8 ++++---- src/cfdb/workflows/provisioner.py | 10 ++++++---- 2 files changed, 10 insertions(+), 8 deletions(-) diff --git a/README.md b/README.md index 514a48c..50e6b3b 100644 --- a/README.md +++ b/README.md @@ -183,7 +183,7 @@ As with dev, confirm after the first promotion that the running prod tasks refer A few operational caveats of the moving-tag pattern: - **Version skew across a deploy.** The API service rolls immediately, but the worker fleet only advances when the next `EcsProvisioner` `RunTask` pulls the freshly-pushed `:dev`. Between the API roll and the next worker launch, a newly-rolled API can dispatch to in-flight workers still running the prior `:dev` image — so a deploy has a transient window where the API and worker code can be one commit apart. Keep the API↔worker dispatch contract backward-compatible across adjacent commits. -- **A wool version bump is a deliberate exception to that.** wool admits a worker only when the proxy's version is `<=` the worker's within the same major (`is_version_compatible`), applied as a discovery filter, and `wool.protocol.__version__` is just the installed package version. Since the pipeline rolls the API first, a bump means the new API rejects **every** in-flight worker until the fleet turns over. It self-heals — jobs stay `pending` and the durable scheduler drains them once fresh workers spawn — but two second-order effects are worth knowing. Rejected-but-running workers still count toward `ECS_MAX_WORKERS` in the provisioner's pre-`RunTask` census, and nothing reaps them before `CFDB_WORKER_MAX_LIFETIME_SECONDS` (5 h, longer than the 4 h dispatch deadline), so a bump landing on a near-capped fleet can wedge spawning long enough to fail jobs `capacity:`. And the symptom is `NoWorkersAvailable`, indistinguishable from the TLS failures above. To make it a non-event, drain the fleet as part of the deploy that carries the bump — `aws ecs list-tasks --cluster --family ` then `stop-task` on each, or simply confirm it is empty before promoting. The rollback direction is safe: an older API against newer workers passes the gate. +- **A wool version bump is a deliberate exception to that.** wool admits a worker only when the proxy's version is `<=` the worker's within the same major (`is_version_compatible`), applied as a discovery filter, and `wool.protocol.__version__` is just the installed package version. Since the pipeline rolls the API first, a bump means the new API rejects **every** in-flight worker until the fleet turns over. It self-heals — jobs stay `pending` and the durable scheduler drains them once fresh workers spawn — but two second-order effects are worth knowing. Rejected-but-running workers still count toward `ECS_MAX_WORKERS` in the provisioner's pre-`RunTask` census. Since they receive no dispatches, idle shutdown reaps them after `CFDB_WORKER_IDLE_TIMEOUT_SECONDS` (default 10 min), so the wedge window is minutes rather than the hours the max-lifetime ceiling would allow — but only when the running fleet is already on an idle-aware image (wool ≥ 0.14): a bump *from* an older fleet leaves the old workers unable to self-reap early, and a near-capped fleet can wedge spawning long enough to fail jobs `capacity:`. And the symptom is `NoWorkersAvailable`, indistinguishable from the TLS failures above. To make it a non-event, drain the fleet as part of the deploy that carries the bump — `aws ecs list-tasks --cluster --family ` then `stop-task` on each, or simply confirm it is empty before promoting. Note that a bump does not require a wool code change: cfdb tracks `wool~=0.14.0`, so any lock refresh that moves the resolved version — a patch release, not just a deliberate upgrade — is itself a bump requiring this same procedure. The rollback direction is safe: an older API against newer workers passes the gate. - **Single-environment moving tag.** `MOVING_TAG` is hard-coded to `dev` in the workflow, so this pipeline targets exactly one environment. A second environment (e.g. `prod`) would need its own moving tag, repo variables, and task-def wiring — not yet parameterized. - **Rollout wait ceiling.** `aws ecs wait services-stable` polls for up to ~10 minutes (40 attempts × 15 s) before timing out. A genuinely slow or wedged rollout will fail the workflow at that ceiling even though the `update-service` call itself succeeded; the deploy may still converge afterward, or the circuit breaker (below) may roll it back. - **Stale GitHub secrets.** The old `BACKEND_STACK_NAME` and `WORKERS_STACK_NAME` GitHub secrets are no longer used by this workflow (it no longer runs `cloudformation deploy`) and can be deleted. @@ -842,7 +842,7 @@ Cache keys are content-addressed using each file's upstream `md5`, so a byte cha **Bounded concurrency, durable queuing, and admission control.** Dispatch is bounded on three cooperating layers so an unauthenticated burst on `/data` and `/index` can't oversubscribe the worker fleet or queue unbounded work: - **Per-worker backpressure** — each worker accepts at most `CFDB_WORKER_MAX_CONCURRENT_TASKS` tasks at once (default `1`), serializing the subprocess pipelines on a 1-vCPU worker. A worker at capacity rejects the dispatch and the API's priority load balancer rotates to the next worker. -- **Priority (leaky-bucket) load balancing** — the API offers each task to discovered workers in a stable order, so load concentrates on the lowest-ordered workers and over-provisioned workers drain to idle and self-reap (via `CFDB_WORKER_MAX_LIFETIME_SECONDS`) instead of every worker carrying a thin perpetual slice. +- **Priority (leaky-bucket) load balancing** — the API offers each task to discovered workers in a stable order, so load concentrates on the lowest-ordered workers and over-provisioned workers drain to idle and self-reap (on the ECS profile, via `CFDB_WORKER_IDLE_TIMEOUT_SECONDS`, with `CFDB_WORKER_MAX_LIFETIME_SECONDS` as the backstop) instead of every worker carrying a thin perpetual slice. - **Durable queue + retry-to-deadline** — when no worker has capacity, the job is **not** failed and does **not** block the request: it stays `pending` and a durable, Mongo-backed scheduler re-attempts dispatch every `CFDB_WORKFLOW_RETRY_INTERVAL_S` (plus jitter) until a worker frees up or the `CFDB_WORKFLOW_DISPATCH_DEADLINE_S` deadline elapses (then it is failed with a `capacity:`-prefixed error). Because the queue lives in Mongo, an API restart resumes it. On every scheduler tick (including the first, on boot) an orphan-recovery sweep re-queues jobs a crash left mid-flight — a `running` job whose API consumer died, or a fresh `pending` claim that never rescheduled — once they pass the stale threshold (`CFDB_WORKFLOW_STALE_THRESHOLD_S`), so recovery is autonomous and does not wait for a client to re-request the file. Recovery shares the same deadline clock as a fresh job: the re-queue preserves the original submission time, so an orphan older than `CFDB_WORKFLOW_DISPATCH_DEADLINE_S` is failed `capacity:` on its first recovery attempt rather than resumed (its committed cache artifacts survive for a later fresh `GET` to reuse) — recovery is best-effort, not unbounded. On the ECS profile, an overflow also requests one bounded worker spawn (the leaky bucket overflowing), inverting the old unconditional per-request spawn. - **Admission ceiling** — once `CFDB_WORKFLOW_MAX_ACTIVE` workflows are active (`pending` + `running`), further preprocessing requests are shed with `429 Retry-After` rather than queued, so the backlog itself is bounded. The check runs before the per-file mutex, so at the ceiling even a re-`GET` for a file whose workflow is already in flight is shed with `429` (rather than attaching to the in-flight job) and the client retries — the deliberate trade for shedding before an unbounded admission race. The readiness `/status` probes never dispatch and so never `429`. @@ -887,7 +887,7 @@ When the API runs on ECS Fargate (or LocalStack-backed dev that mirrors prod end - `ECS_WORKER_SECURITY_GROUPS` — comma-separated awsvpc security group IDs. Optional — when empty, ECS applies the VPC default SG. - `ECS_WORKER_ASSIGN_PUBLIC_IP` — `ENABLED` or `DISABLED` (default `DISABLED`). Production should leave this disabled and reach AWS via VPC endpoints; LocalStack accepts either value. -The worker container's `CMD` is `python -m cfdb.workflows.worker_main`. Worker-side knobs (gRPC port, health port, max lifetime, drain grace) are documented under `--help` on that command; their env vars are `CFDB_WORKER_GRPC_PORT`, `CFDB_WORKER_HEALTH_PORT`, `CFDB_WORKER_MAX_LIFETIME_SECONDS`, and `CFDB_WORKER_DRAIN_GRACE_SECONDS`. The worker task definition MUST declare a `healthCheck` against the gRPC port; without one ECS reports `healthStatus: UNKNOWN` indefinitely and the worker is never advertised to discovery. +The worker container's `CMD` is `python -m cfdb.workflows.worker_main`. Worker-side knobs (gRPC port, health port, idle timeout, idle poll cadence, max lifetime, self-termination grace, drain grace) are documented under `--help` on that command; their env vars are `CFDB_WORKER_GRPC_PORT`, `CFDB_WORKER_HEALTH_PORT`, `CFDB_WORKER_IDLE_TIMEOUT_SECONDS`, `CFDB_WORKER_IDLE_POLL_INTERVAL_SECONDS`, `CFDB_WORKER_IDLE_POLL_FAILURE_LIMIT`, `CFDB_WORKER_MAX_LIFETIME_SECONDS`, `CFDB_WORKER_MAX_LIFETIME_GRACE_SECONDS`, and `CFDB_WORKER_DRAIN_GRACE_SECONDS`. A worker self-terminates once continuously idle beyond `CFDB_WORKER_IDLE_TIMEOUT_SECONDS` (default 600 s; 0 disables) — measured via wool's `idle` RPC, which reports zero while any task runs, so the idle exit never fires while a task is running — with `CFDB_WORKER_MAX_LIFETIME_SECONDS` (default 12 h) retained as the backstop for a worker whose idle reporting is wedged or whose job is stuck. Both self-termination exits stop the worker with a drain grace (`CFDB_WORKER_MAX_LIFETIME_GRACE_SECONDS`, default 6 h): a dispatch accepted in the idle teardown window, or the job a max-lifetime expiry lands on, runs to completion instead of being cancelled into a terminal job failure, so worst-case worker uptime is lifetime + grace (18 h at the defaults) and the grace — sized above the API's 4 h `CFDB_WORKFLOW_DURATION_CAP_S` — only ever cancels work already past the API's own viability bound. Keep the idle timeout comfortably above `CFDB_WORKFLOW_RETRY_INTERVAL_S` (default 120 s) plus worker cold-start time (≥ 2–3×): wool's idle clock starts at worker startup, so an aggressive timeout lets overflow-spawned workers idle-reap before the scheduler's next retry tick ever reaches them, and jobs ride the queue to the `capacity:` deadline while everything looks green. The worker task definition MUST declare a `healthCheck` against the gRPC port; without one ECS reports `healthStatus: UNKNOWN` indefinitely and the worker is never advertised to discovery. **Workers publish their own metadata.** ECS reports a task's address and health, but not what is running inside the container — and two fields of wool's `WorkerMetadata` are knowable only to the worker: the wool protocol version it runs, and whether it configured TLS. wool gates worker admission on both, so a value the API invented for them would be a value that silently rejects the entire fleet. After starting, each worker therefore tags its own ECS task (`wool.version`, `wool.secure`) with what wool authored for it, and `EcsDiscovery` reads those tags back via `DescribeTasks … --include TAGS`. ECS supplies liveness; the worker supplies identity. @@ -912,7 +912,7 @@ By default the API↔worker gRPC dispatch channel is plaintext, gated only by ne The configuration is gating-by-presence: when all three of `CFDB_WORKER_TLS_CA`, `CFDB_WORKER_TLS_CERT`, and `CFDB_WORKER_TLS_KEY` are unset the plaintext path is used unchanged (local PoC dev needs no certs); when all three are set mTLS is enforced. A *partial* configuration (some set, some not) fails fast at startup rather than silently degrading to plaintext. -**Identity.** TLS normally verifies a server's certificate against the address the client dialed, which is a problem here: workers answer wherever they happen to come up. `EcsDiscovery` reaches each Fargate worker at the awsvpc IP assigned at launch, and a containerized local worker answers on a bridge IP — neither is knowable when the certificate is minted. `CFDB_WORKER_TLS_IDENTITY` (default `cfdb-worker`) points the API at a fixed logical name instead, so the worker leaf carries one stable SAN rather than an enumeration of every address it might be reached at. Chain and SAN verification both still happen; only the name being matched changes. It is a client-side setting — but "client" is a property of a connection, not of a process: the API is the client on the dispatch channel, and each worker is the client on the one channel wool opens back to its own subprocess to drain it, verifying the same SAN. Both sides therefore read the variable and MUST agree on it; a worker left at a different value drains by force-reap instead of gracefully, losing in-flight work with no TLS error anywhere. Setting it to the empty string restores address verification, which is only workable when workers are reached at a fixed, certified address. +**Identity.** TLS normally verifies a server's certificate against the address the client dialed, which is a problem here: workers answer wherever they happen to come up. `EcsDiscovery` reaches each Fargate worker at the awsvpc IP assigned at launch, and a containerized local worker answers on a bridge IP — neither is knowable when the certificate is minted. `CFDB_WORKER_TLS_IDENTITY` (default `cfdb-worker`) points the API at a fixed logical name instead, so the worker leaf carries one stable SAN rather than an enumeration of every address it might be reached at. Chain and SAN verification both still happen; only the name being matched changes. It is a client-side setting — but "client" is a property of a connection, not of a process: the API is the client on the dispatch channel, and each worker is the client on the one channel wool opens back to its own subprocess to drain it, verifying the same SAN. Both sides therefore read the variable and MUST agree on it; a worker left at a different value drains by force-reap instead of gracefully, losing in-flight work with no TLS error anywhere. The idle-shutdown poll is a third consumer of the same contract: each worker dials its own subprocess at `127.0.0.1` to read wool's `idle` RPC, verifying the same identity SAN. Setting it to the empty string restores address verification, which is only workable when workers are reached at a fixed, certified address — and for the idle poll that means the worker leaf must carry a loopback SAN (the generator's default leaf includes `DNS:localhost` and `IP:127.0.0.1`), or every poll fails its handshake and idle shutdown silently degrades to the max-lifetime backstop with only repeated poll warnings as the symptom. Generate a local CA and the worker + API leaf certs with: diff --git a/src/cfdb/workflows/provisioner.py b/src/cfdb/workflows/provisioner.py index f5e1300..ec1f7c6 100644 --- a/src/cfdb/workflows/provisioner.py +++ b/src/cfdb/workflows/provisioner.py @@ -467,10 +467,12 @@ async def _submit_run_task(self, kwargs: dict[str, Any]) -> dict[str, Any]: The future is recorded in ``_pending`` so ``aclose`` can drain threads mid-flight. A done-callback logs any ARNs the boto thread produced *after* the awaiting coroutine was cancelled — - these are orphan workers that no caller can claim, and their - only safety net is the worker's own ``CFDB_WORKER_MAX_LIFETIME`` - ceiling. Surfacing them at WARNING gives operators a chance - to reap manually before the ceiling fires. + these are orphan workers that no caller can claim. Left unused, + they self-reap once continuously idle beyond + ``CFDB_WORKER_IDLE_TIMEOUT_SECONDS`` (minutes), with the + ``CFDB_WORKER_MAX_LIFETIME_SECONDS`` ceiling as the outer + bound. Surfacing them at WARNING still gives operators a + chance to reap manually first. """ slot = _SubmittedRunTask() future = self._executor.submit(self._client.run_task, **kwargs) From 4390f6c9eec61bd622373e613084ab464069751b Mon Sep 17 00:00:00 2001 From: Conrad Date: Wed, 12 Aug 2026 13:35:14 -0400 Subject: [PATCH 7/7] test: Pin wool's stop-grace contract against a real worker Both worker self-termination paths stop with a drain grace, and that grace is the only reason they are safe. An idle reading is a snapshot, so it cannot rule out a dispatch accepted between the final poll and the teardown, and a max-lifetime expiry can land mid-job outright. In either case the API has already marked that task running, where a mid-stream cancel is finalized as a terminal job failure rather than re-queued. Every unit test on that path mocks the worker, so the suite asserts only that a grace is passed; nothing verifies wool honors it. A wool change that cancelled regardless would leave the suite green while restoring the defect the grace exists to prevent. The two tests are a controlled pair over one real worker and one real dispatch, identical but for the grace and with opposite outcomes. The graceless case is what makes the drain meaningful: without it, a passing drain could just as well be a task that finished on its own. Both bound the wait, so a task that hangs fails the test rather than passing as something-went-wrong. The routine lives in the shared routines module because cloudpickle must resolve it by reference across the worker boundary, and the pool is bound to a reserved ephemeral port because dispatch needs a session while the stop has to target one known worker. --- tests/integration/routines.py | 12 ++ tests/integration/test_worker_stop_grace.py | 135 ++++++++++++++++++++ 2 files changed, 147 insertions(+) create mode 100644 tests/integration/test_worker_stop_grace.py diff --git a/tests/integration/routines.py b/tests/integration/routines.py index 5f152f4..f639113 100644 --- a/tests/integration/routines.py +++ b/tests/integration/routines.py @@ -115,6 +115,18 @@ async def echo(value: str) -> str: return value +@wool.routine +async def sleep_then_return(seconds: float, value: str) -> str: + """Occupy the worker for ``seconds``, then return ``value``. + + Gives the stop-grace tests a task that is reliably still in flight + when the stop lands, so the question they ask — does the drain let + it finish, or does it die mid-run — has a deterministic answer. + """ + await asyncio.sleep(seconds) + return value + + def stub_file_meta() -> dict[str, Any]: """Return a minimal BAM file_meta accepted by ``StubProcessor``.""" return { diff --git a/tests/integration/test_worker_stop_grace.py b/tests/integration/test_worker_stop_grace.py new file mode 100644 index 0000000..78a30cd --- /dev/null +++ b/tests/integration/test_worker_stop_grace.py @@ -0,0 +1,135 @@ +"""Integration test for the stop-grace contract the idle exit relies on. + +``worker_main.serve`` stops the worker with a grace on both of its +self-termination paths, and that grace is the whole reason those exits +are safe: an idle reading is a snapshot, so it cannot rule out a +dispatch accepted between the final poll and the teardown, and a +max-lifetime expiry can land mid-job outright. In either case the API +has already marked that task running, where a mid-stream cancel is +finalized as a terminal job failure rather than re-queued. + +Every unit test on that path mocks ``wool.LocalWorker``, so the suite +asserts only that cfdb *passes* a grace — nothing verifies wool honors +it. A wool change that made ``stop`` cancel regardless would leave the +whole suite green while silently restoring the defect the grace exists +to prevent. These tests start a real worker, dispatch a real task, and +stop it mid-flight. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import functools +import socket + +import pytest +import wool + +from tests.integration.routines import sleep_then_return + + +pytestmark = pytest.mark.integration + +#: How long the dispatched task occupies the worker. Long enough that +#: the stop below reliably lands while it is still running, short +#: enough not to drag the suite. +TASK_SECONDS = 4.0 + +#: Grace allowed for the drain — comfortably longer than the task, so a +#: completed task means the drain waited rather than that the drain +#: happened to outlast a tiny task. +DRAIN_GRACE_SECONDS = 30.0 + + +def _free_port() -> int: + """Reserve an ephemeral port and return it. + + The worker has to answer at an address the test can dial, so the + port cannot be left to ``0``; binding and releasing avoids the + collisions a hard-coded port would cause on a busy machine. + """ + with socket.socket() as probe: + probe.bind(("127.0.0.1", 0)) + return probe.getsockname()[1] + + +@contextlib.asynccontextmanager +async def _worker_on_known_port(): + """Yield ``(connection, port)`` for a pool-managed worker. + + Dispatch needs a pool session, but the stop has to target one known + worker, so the pool is given a factory bound to a reserved port and + the test dials that port directly. + """ + port = _free_port() + factory = functools.partial(wool.LocalWorker, host="127.0.0.1", port=port) + async with wool.WorkerPool( + spawn=1, worker=factory, discovery=wool.LocalDiscovery(), lazy=False + ): + connection = wool.WorkerConnection(f"127.0.0.1:{port}") + try: + yield connection, port + finally: + await connection.close() + + +class TestWorkerStopGrace: + @pytest.mark.asyncio + async def test_stop_should_drain_an_in_flight_task_when_given_a_grace(self): + """Test that a graced stop lets a running task finish. + + Given: + A real worker running a dispatched task that outlasts the + moment the stop is issued. + When: + The worker is stopped with a grace longer than the task's + remaining runtime. + Then: + It should let the task run to completion and return its + result, pinning the contract that makes cfdb's idle and + max-lifetime exits safe rather than merely narrow. + """ + async with _worker_on_known_port() as (connection, _): + # Arrange — put a task in flight, then let it get underway + task = asyncio.create_task( + sleep_then_return(TASK_SECONDS, "drained") + ) + await asyncio.sleep(TASK_SECONDS / 4) + assert not task.done() + + # Act — stop the worker out from under the running task + await connection.stop(grace=DRAIN_GRACE_SECONDS) + + # Assert + assert await task == "drained" + + @pytest.mark.asyncio + async def test_stop_should_cancel_an_in_flight_task_when_given_no_grace(self): + """Test that a graceless stop kills a running task. + + Given: + A real worker running a dispatched task that outlasts the + moment the stop is issued. + When: + The worker is stopped without a grace — wool's default. + Then: + It should cancel the task rather than return its value, so + the drain asserted above is demonstrably the grace's doing + and not an artifact of the task finishing on its own. + """ + async with _worker_on_known_port() as (connection, _): + # Arrange — put a task in flight, then let it get underway + task = asyncio.create_task( + sleep_then_return(TASK_SECONDS, "drained") + ) + await asyncio.sleep(TASK_SECONDS / 4) + assert not task.done() + + # Act — stop with wool's default: no grace + await connection.stop() + + # Assert — a hung task raises TimeoutError here instead, so + # the cancellation is pinned rather than merely "not done" + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=TASK_SECONDS)