Bound the engine drain so the shutdown after it actually runs (#192)

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

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)

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.
 

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

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:

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

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.

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:

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(