svcforge: reference implementation

Complete working build of the system learn-python/ teaches.
164 tests, mypy --strict clean, domain coverage 99%.

services/worker/main.py

@@ -0,0 +1,175 @@
+"""The claim loop.
+
+Poll every 5 seconds. Claim while a semaphore slot is free. Run the handler. Report.
+That is the whole design, and the restraint is the point: LISTEN/NOTIFY would shave the
+latency, is fire-and-forget so it can never replace the poll anyway, is strictly extra
+code, and does not exist on pgbouncer's transaction pooler. The poll is not a placeholder
+for something better.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import signal
+import time
+
+from opentelemetry import trace
+
+from services.worker.deps import WorkerDeps
+from services.worker.handlers import HANDLERS
+from svcforge_core import obs
+from svcforge_core.adapters.clock import SystemClock
+from svcforge_core.adapters.helm import HelmProvisioner
+from svcforge_core.adapters.notify import LogNotifier
+from svcforge_core.domain.catalog import load_catalog
+from svcforge_core.domain.models import Task
+from svcforge_core.repo.db import make_pool
+from svcforge_core.repo.instances import InstanceRepo
+from svcforge_core.repo.tasks import TaskRepo
+from svcforge_core.settings import Settings, load_settings
+
+log = obs.get_logger("svcforge.worker")
+
+
+async def _sleep_or_stop(stop: asyncio.Event, seconds: float) -> None:
+    """Sleep, but wake immediately on shutdown.
+
+    `await asyncio.sleep(5)` would make every SIGTERM cost up to five seconds of
+    Kubernetes waiting on terminationGracePeriod for no reason.
+    """
+    with contextlib.suppress(TimeoutError):
+        await asyncio.wait_for(stop.wait(), timeout=seconds)
+
+
+async def _run_one(deps: WorkerDeps, task: Task, sem: asyncio.Semaphore) -> None:
+    """Run one task to a terminal report. Never lets an exception escape the TaskGroup."""
+    try:
+        # Every log line from here carries instance_id/task_id/team. Bound once, at claim,
+        # rather than passed down: the alternative is threading three arguments through
+        # every function that might log, and the first one anyone forgets is the one you
+        # need at 3am.
+        obs.bind_task_context(task.instance_id, task.id, team=task.team or "unknown")
+        log.info("task claimed", kind=task.kind.value, attempt=task.attempts)
+        obs.TASKS_CLAIMED.labels(kind=task.kind.value).inc()
+
+        handler = HANDLERS.get(task.kind)
+        if handler is None:
+            await deps.tasks.fail(task.id, f"no handler for {task.kind}", max_attempts=1)
+            return
+
+        # Re-parent to the span that enqueued this task. Without the stored traceparent
+        # the worker's span starts a brand-new trace, and the POST that caused the work
+        # is in a different trace to the helm call that did it.
+        ctx = obs.context_from_traceparent(task.traceparent)
+        started = time.monotonic()
+        with obs.tracer().start_as_current_span(
+            f"task.{task.kind.value}",
+            context=ctx,
+            kind=trace.SpanKind.CONSUMER,
+        ) as span:
+            span.set_attribute("task.id", task.id)
+            span.set_attribute("task.kind", task.kind.value)
+            span.set_attribute("instance.id", str(task.instance_id))
+            try:
+                await handler(task, deps)
+            except asyncio.CancelledError:
+                # CancelledError inherits from BaseException, so `except Exception` below
+                # would never see it. Catch it only to release the claim, then ALWAYS
+                # re-raise: swallowing it breaks cancellation for everyone above us.
+                await deps.tasks.fail(task.id, "cancelled", max_attempts=deps.settings.max_attempts)
+                raise
+            except Exception as exc:
+                log.exception("task failed", kind=task.kind.value, error=str(exc))
+                span.record_exception(exc)
+                span.set_status(trace.Status(trace.StatusCode.ERROR, str(exc)))
+                obs.TASKS_FAILED.labels(kind=task.kind.value).inc()
+                await deps.tasks.fail(task.id, str(exc), max_attempts=deps.settings.max_attempts)
+            else:
+                obs.PROVISION_TIME.observe(time.monotonic() - started)
+                await deps.tasks.complete(task.id)
+    finally:
+        sem.release()
+
+
+async def run_worker(deps: WorkerDeps, stop: asyncio.Event) -> None:
+    """Claim and run until told to stop, then drain what is in flight.
+
+    Draining is what makes a rolling deploy invisible. Exiting the `async with` block
+    awaits every in-flight handler, so a pod that is being replaced finishes the provision
+    it already started instead of abandoning it half-done for the lease to clean up
+    five minutes later.
+    """
+    sem = asyncio.Semaphore(deps.settings.worker_concurrency)
+    worker_id = deps.settings.worker_id
+
+    async with asyncio.TaskGroup() as tg:
+        while not stop.is_set():
+            await sem.acquire()
+            if stop.is_set():
+                sem.release()
+                break
+
+            try:
+                task = await deps.tasks.claim(worker_id)
+            except Exception:
+                # A DB blip must not kill the worker; back off and try again.
+                log.exception("claim failed")
+                sem.release()
+                await _sleep_or_stop(stop, deps.settings.poll_interval_s)
+                continue
+
+            if task is None:
+                sem.release()
+                await _sleep_or_stop(stop, deps.settings.poll_interval_s)
+                continue
+
+            tg.create_task(_run_one(deps, task, sem))
+    # TaskGroup.__aexit__ awaited the in-flight handlers. Now it is safe to exit 0.
+
+
+async def _amain() -> None:
+    settings: Settings = load_settings()
+
+    # Before anything else: nothing logged above this line is structured, and the metrics
+    # the SvcforgeTaskFailed / SvcforgeProvisionSlow alerts query do not exist until the
+    # registry is up.
+    obs.setup("svcforge-worker", settings)
+    obs.start_metrics_server(settings.metrics_port)
+
+    pool = make_pool(settings.pg_dsn.unicode_string(), settings.pool_min_size, settings.pool_max_size)
+    await pool.open(wait=True)
+
+    deps = WorkerDeps(
+        pool=pool,
+        instances=InstanceRepo(pool),
+        tasks=TaskRepo(pool),
+        provisioner=HelmProvisioner(helm_bin=settings.helm_bin, timeout_s=int(settings.helm_timeout_s)),
+        notifier=LogNotifier(),
+        clock=SystemClock(),
+        catalog=load_catalog(settings.catalog_path),
+        settings=settings,
+    )
+
+    stop = asyncio.Event()
+    loop = asyncio.get_running_loop()
+    for sig in (signal.SIGTERM, signal.SIGINT):
+        # add_signal_handler, NOT signal.signal. signal.signal runs the handler at an
+        # arbitrary bytecode boundary on whatever thread the C-level handler lands on,
+        # and the event loop will not notice until its next timer fires. This one is
+        # loop-safe: the callback runs as a normal loop callback.
+        loop.add_signal_handler(sig, stop.set)
+
+    try:
+        await run_worker(deps, stop)
+    finally:
+        await pool.close()
+
+
+def main() -> None:
+    """One asyncio.run, at the top, never nested."""
+    asyncio.run(_amain())
+
+
+if __name__ == "__main__":
+    main()