diff --git a/CHANGELOG.md b/CHANGELOG.md index dff84bf..38a98c5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -44,6 +44,17 @@ says is recorded beside the finding, never written onto it. Both are reversible. ### Fixed +- A draining engine never finished shutting down (#192). `Worker.stop()` waits ten minutes by + default and both grace periods this project ships are two, so an engine with a long fetch in + flight was still waiting when SIGKILL arrived: the heartbeat, the API connection pool and the + metrics server were never stopped, and `engine_stopped` never appeared in the log. The wait is + now bounded by `ICEBERG_DRAIN_SECONDS` (default 90), which a deploy invariant holds below both + graces, and a task still running when the budget expires is named in an + `engine_drain_incomplete` warning before the engine goes. The drain policy — wait, then let the + lease hand the work to another engine rather than reporting a terminal failure — is written down + in [`docs/deployment.md`](./docs/deployment.md) § Draining an engine, along with why the + alternative was rejected. + - A scan could report clean coverage for a task that died halfway through a resumable source (#193). When a connector checkpoints, the engine's last submission carries only the *remainder* the progress batches did not — but the task-wide gap that says "this task never finished reading diff --git a/apps/engine/src/iceberg_engine/runner.py b/apps/engine/src/iceberg_engine/runner.py index 0f9840a..0c58682 100644 --- a/apps/engine/src/iceberg_engine/runner.py +++ b/apps/engine/src/iceberg_engine/runner.py @@ -575,11 +575,16 @@ def run_task( log.warning("scan_task_connector_failed", error=report.error) ENGINE_CONNECTOR_FAILURES.labels(source_type=lease.source_type).inc() except Interrupt as exc: - # Dramatiq's time limit and its shutdown both raise a `BaseException`, so - # neither handler here sees one and the task would end without reporting — - # leaving the API to wait out the lease and redeliver a task that will be - # interrupted at exactly the same point (#106). Report what the fetch got - # through, then let the interrupt go on killing the thread it was raised in. + # Dramatiq's time limit raises a `BaseException`, so neither handler here + # sees one and the task would end without reporting — leaving the API to + # wait out the lease and redeliver a task that will be interrupted at + # exactly the same point (#106). Report what the fetch got through, then + # let the interrupt go on killing the thread it was raised in. + # + # A drain does *not* arrive here: `notify_shutdown` is deliberately unset, + # so a shutdown never interrupts an actor and an abandoned task goes back + # through lease expiry instead of reporting a terminal failure. See + # `iceberg_engine.worker.drain` for why (#192). report.status = "failed" report.error = type(exc).__name__ _record_task_gap(report, lease, CoverageReason.TIMEOUT, connector) diff --git a/apps/engine/src/iceberg_engine/worker.py b/apps/engine/src/iceberg_engine/worker.py index 10f8325..bd619a2 100644 --- a/apps/engine/src/iceberg_engine/worker.py +++ b/apps/engine/src/iceberg_engine/worker.py @@ -251,7 +251,8 @@ def main(settings: EngineSettings | None = None) -> None: Docker stops containers with SIGTERM, so handling it is what makes ``docker compose down`` a clean shutdown rather than a ten-second wait - followed by SIGKILL. + followed by SIGKILL. What that shutdown does — and does not — promise the + tasks in flight is :func:`drain`. The consumer is started here rather than by the `dramatiq` CLI so that importing this module stays free of side effects. The CLI configures a broker @@ -282,7 +283,7 @@ def request_shutdown(signum: int, _frame: object) -> None: # Stop consuming before the metrics server goes: a message picked up during # shutdown still holds a lease, and finishing it is cheaper for the scan than # waiting out an expiry. - consumer.stop() + drain(consumer, resolved.drain_seconds) if heartbeat is not None: heartbeat.stop() # Only once both are stopped: until then, threads are still reporting through @@ -292,6 +293,43 @@ def request_shutdown(signum: int, _frame: object) -> None: logger.info("engine_stopped") +def drain(consumer: dramatiq.Worker, seconds: float) -> list[uuid.UUID]: + """Stop taking work and wait, bounded, for the tasks this engine already holds. + + Returns the tasks still running when the budget ran out, which are abandoned: + the worker threads are daemons, so they die with the process, and the API + reclaims each task when its lease lapses and hands it to another engine. A + connector that checkpoints resumes from its last flushed batch (#143); one + that does not re-reads its spec. **That latency is the deliberate price.** + + The alternative — interrupting the task so it reports what it has — was + considered and rejected (#192). An engine may only report `completed` or + `failed`, and a failed task is terminal: it is never reclaimed, it makes its + scan `partial`, and a partial scan may not auto-resolve findings (ADR 0009 + §4). So interrupting would trade one lease TTL of latency for a scan that + cannot close a secret somebody has already fixed — on every rolling deploy + that lands mid-scan. Waiting, then letting the lease do its job, keeps the + API's reclaim the single re-delivery authority (ADR 0009 §2). + + What the budget buys is the shutdown *after* it. Dramatiq's default is ten + minutes; every grace period this project ships is two, so the wait was still + running when SIGKILL arrived and the heartbeat, the client pool and the + metrics server were never stopped at all. + """ + consumer.stop(timeout=int(seconds * 1000)) + abandoned = TASKS.held() + if abandoned: + # Named, because these are the tasks a scan is about to sit on until their + # leases lapse — the one operator-visible cost of a drain that timed out. + logger.warning( + "engine_drain_incomplete", + drain_seconds=seconds, + tasks=[str(task_id) for task_id in abandoned], + detail="left to lease expiry and API reclaim", + ) + return abandoned + + def _start_heartbeat(settings: EngineSettings, client: EngineClient) -> Heartbeat | None: """Begin renewing leases, if this engine knows which engine it is. diff --git a/apps/engine/tests/test_worker.py b/apps/engine/tests/test_worker.py index bbf8429..8814659 100644 --- a/apps/engine/tests/test_worker.py +++ b/apps/engine/tests/test_worker.py @@ -1,7 +1,10 @@ import io import json import socket +import threading import urllib.request +import uuid +from time import monotonic import dramatiq import pytest @@ -13,14 +16,17 @@ from iceberg_core.logging import configure_logging from iceberg_core.tasks import SCAN_TASK_ACTOR, SCAN_TASK_QUEUE from iceberg_engine.worker import ( + TASKS, api_client, bootstrap, build_broker, close_api_client, + drain, register_connectors, run_scan_task, ) from pydantic import SecretStr +from structlog.testing import capture_logs def test_build_broker_defaults_to_stub() -> None: @@ -151,6 +157,53 @@ def test_shutdown_closes_the_shared_client() -> None: close_api_client() # a shutdown path may run twice +def test_a_drain_gives_up_on_a_long_task_rather_than_outliving_its_grace() -> None: + """The bug this budget exists for (#192). + + Dramatiq's `Worker.stop()` waits ten minutes by default. Every grace period + this project ships is two, so a fetch still running at SIGTERM held the wait + until SIGKILL — and the heartbeat, the API pool and the metrics server were + never stopped. The wait is now bounded, and the task it gives up on is named + so an operator can see which scan is about to sit on a lease. + """ + broker = build_broker() + running, release = threading.Event(), threading.Event() + task_id = uuid.uuid4() + + @dramatiq.actor(broker=broker, queue_name="drain-test") + def slow_task() -> None: + with TASKS.holding(task_id): + running.set() + release.wait(timeout=30) + + slow_task.send() + consumer = Worker(broker, queues={"drain-test"}, worker_timeout=50) + consumer.start() + try: + assert running.wait(timeout=5), "the fixture task never started" + started = monotonic() + with capture_logs() as events: + abandoned = drain(consumer, 0.25) + elapsed = monotonic() - started + finally: + release.set() + + assert abandoned == [task_id] + assert elapsed < 5, f"the drain waited {elapsed:.1f}s on a 0.25s budget" + assert [event for event in events if event["event"] == "engine_drain_incomplete"] + + +def test_a_drain_returns_as_soon_as_the_work_is_done() -> None: + """The budget is a ceiling, not a delay: the ordinary drain finds nothing in + flight and must not spend the grace period waiting to discover that.""" + consumer = Worker(build_broker(), queues={SCAN_TASK_QUEUE}, worker_timeout=50) + consumer.start() + started = monotonic() + + assert drain(consumer, 30.0) == [] + assert monotonic() - started < 5 + + def test_the_shipped_connectors_are_registered() -> None: """A lease names a source type and nothing else, so a type this image cannot resolve fails the task. Registration is explicit rather than by import-time diff --git a/deploy/compose/docker-compose.yml b/deploy/compose/docker-compose.yml index e164b51..f12be22 100644 --- a/deploy/compose/docker-compose.yml +++ b/deploy/compose/docker-compose.yml @@ -154,10 +154,12 @@ services: ICEBERG_ENGINE_TOKEN: ${ICEBERG_ENGINE_TOKEN:-} ICEBERG_WORKER_THREADS: ${ICEBERG_WORKER_THREADS:-4} # Matches terminationGracePeriodSeconds in the chart, and for the same - # reason: worker.py catches SIGTERM and finishes the task it has leased. - # Docker's default of 10s would SIGKILL a fetch in progress and leave the - # scan waiting out its 300s lease before anything reclaims it — so `make - # down` is only the clean shutdown the worker promises with this set. + # reason: worker.py catches SIGTERM and waits ICEBERG_DRAIN_SECONDS (90s) for + # the tasks it already holds. Docker's default of 10s would SIGKILL a fetch in + # progress and leave the scan waiting out its 300s lease before anything + # reclaims it — so `make down` is only the clean shutdown the worker promises + # with this set. It must stay *longer* than the drain budget, or the wait is + # killed before the heartbeat, the API pool and the metrics server are stopped. stop_grace_period: 120s depends_on: redis: diff --git a/deploy/helm/icebergsst/templates/engine-deployment.yaml b/deploy/helm/icebergsst/templates/engine-deployment.yaml index 0765daf..0dce151 100644 --- a/deploy/helm/icebergsst/templates/engine-deployment.yaml +++ b/deploy/helm/icebergsst/templates/engine-deployment.yaml @@ -55,9 +55,11 @@ spec: {{- end }} securityContext: {{- toYaml .Values.podSecurityContext | nindent 8 }} - # A draining engine finishes the task it leased rather than dropping it and - # waiting for the lease to lapse. Longer than the longest fetch, shorter - # than the lease. + # A draining engine waits ICEBERG_DRAIN_SECONDS (90s) for the tasks it + # already holds, so most finish rather than waiting out a lease. Longer than + # that budget — the shutdown after it only runs if the process is still + # alive — and shorter than the 300s lease, so a task this engine does + # abandon is reclaimed rather than held by a pod that is gone. terminationGracePeriodSeconds: 120 containers: - name: engine diff --git a/docs/deployment.md b/docs/deployment.md index 73abe87..f2906c3 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -176,6 +176,37 @@ JSONB, `postgresql_where` partial indexes, or an `ALTER` that SQLite only surviv table. A revision that is valid in one dialect and not the other fails there rather than in the pre-upgrade `Job`. +## Draining an engine + +A rolling deploy, `make down`, or a `kubectl delete pod` sends SIGTERM to an engine that may be +halfway through a fetch. What it does then is a policy, and the policy is: **wait, then hand the +work back to the lease** (#192). + +1. Dramatiq stops handing messages to worker threads. A thread that finishes its task exits rather + than taking another, and messages already pulled into the local queue go back to Redis. +2. The process waits up to `ICEBERG_DRAIN_SECONDS` (default 90) for the tasks it *already holds*. + Most finish inside it. +3. Anything still running when the budget expires is abandoned. The API reclaims it when its 300s + lease lapses and hands it to another engine, which — for Confluence, Jira and file shares — + resumes from the last checkpoint the old engine flushed rather than re-reading the scope. + The engine logs `engine_drain_incomplete` naming those tasks first. +4. The heartbeat stops, the API connection pool closes, the metrics server stops, and the process + logs `engine_stopped`. + +Step 4 is why the budget matters: it has to be **shorter than the termination grace period** +(`stop_grace_period` in compose, `terminationGracePeriodSeconds` in the chart — both 120s), or +SIGKILL lands mid-wait and none of it runs. A test holds the default against both. + +An abandoned task is not failed. An engine may only report a task `completed` or `failed`, and a +failed task is terminal — never reclaimed, and it makes its scan `partial`, which may not +auto-resolve findings (ADR 0009 §4). Interrupting in-flight work to report it would trade one +lease TTL of latency for a scan that cannot close a secret somebody has already fixed, on every +deploy that lands mid-scan. Waiting costs latency instead, and keeps lease expiry the single +re-delivery authority (ADR 0009 §2). + +Raise `ICEBERG_DRAIN_SECONDS` **and** both grace periods together if your fetches routinely run +longer than the budget; raising one without the other only moves where the work is lost. + ## Scaling model - Throughput scales by adding **engine** replicas — more Dramatiq consumers pulling scan tasks from Redis. diff --git a/packages/core/src/iceberg_core/config.py b/packages/core/src/iceberg_core/config.py index 8c70f1a..8b8ff72 100644 --- a/packages/core/src/iceberg_core/config.py +++ b/packages/core/src/iceberg_core/config.py @@ -280,6 +280,18 @@ class EngineSettings(CoreSettings): #: against load on the control plane; scale replicas as well (`make scale`). worker_threads: int = Field(default=4, ge=1, le=64) + #: How long a draining engine waits for the tasks it already holds before it + #: stops waiting and lets their leases lapse (#192). + #: + #: It has to fit *inside* the deployment's termination grace period, because + #: what follows the wait — stopping the heartbeat, closing the API pool, + #: stopping the metrics server — only happens if the process is still alive to + #: do it. Dramatiq's own default is ten minutes, five times either grace this + #: project ships, so the wait was still running when SIGKILL arrived and none + #: of the shutdown ran at all. `tests/test_deploy_invariants.py` holds this + #: default against both graces. + drain_seconds: float = Field(default=90.0, gt=0) + @lru_cache(maxsize=1) def get_core_settings() -> CoreSettings: diff --git a/tests/test_deploy_invariants.py b/tests/test_deploy_invariants.py index 4ee4fcb..61867b4 100644 --- a/tests/test_deploy_invariants.py +++ b/tests/test_deploy_invariants.py @@ -28,6 +28,7 @@ import pytest import yaml +from iceberg_core.config import EngineSettings REPO_ROOT = Path(__file__).resolve().parents[1] DOCKER_DIR = REPO_ROOT / "deploy" / "docker" @@ -403,20 +404,41 @@ def test_the_engine_service_can_be_scaled(compose: dict[str, Any]) -> None: def test_the_engine_service_is_given_time_to_finish_its_task(compose: dict[str, Any]) -> None: - """worker.py catches SIGTERM and finishes the task it leased. + """worker.py catches SIGTERM and waits for the task it leased. Docker's default 10s would SIGKILL a fetch in progress, and the scan then waits out its 300s lease before anything reclaims it — so the shutdown the worker documents is only clean with this set. Matches ``terminationGracePeriodSeconds`` in the chart, deliberately. """ - chart_grace = re.search( + assert _service(compose, "engine")["stop_grace_period"] == f"{_chart_grace_seconds()}s" + + +def test_a_draining_engine_finishes_shutting_down_before_it_is_killed( + compose: dict[str, Any], +) -> None: + """The drain budget has to fit inside every grace period that ships (#192). + + The wait is not the point of the wait: what follows it — stopping the + heartbeat, closing the API pool, stopping the metrics server — only happens if + the process is still alive to do it. A budget at or past the grace means + SIGKILL lands mid-wait and none of that runs, which is the state this project + shipped in, with dramatiq's ten-minute default against a two-minute grace. + """ + budget = EngineSettings.model_fields["drain_seconds"].default + compose_grace = float(_service(compose, "engine")["stop_grace_period"].removesuffix("s")) + + assert budget < compose_grace + assert budget < _chart_grace_seconds() + + +def _chart_grace_seconds() -> int: + grace = re.search( r"terminationGracePeriodSeconds:\s*(\d+)", (CHART_DIR / "templates" / "engine-deployment.yaml").read_text(), ) - - assert chart_grace is not None - assert _service(compose, "engine")["stop_grace_period"] == f"{chart_grace.group(1)}s" + assert grace is not None + return int(grace.group(1)) def test_the_engine_service_waits_for_a_token_that_only_the_api_can_mint(