svcforge/services/worker/handlers.py

"""Task handlers.

Every handler here obeys one rule: running it twice must equal running it once.

That is not a nicety. A worker can be SIGKILLed after helm has installed the release but
before the DB row says so; the lease expires; another worker claims the same task and runs
this function again. If the handler is not idempotent, the tenant gets two Elasticsearches
and you get a bill. Idempotency is what makes the crash safe, and it is bought in two
places: a deterministic `release_name`, and adapters that state desired state
(`helm upgrade --install`) instead of issuing imperative commands.
"""

from __future__ import annotations

from collections.abc import Awaitable, Callable
from typing import Any

from services.worker.deps import WorkerDeps
from svcforge_core.domain.models import CatalogEntry, Instance, Task, TaskKind
from svcforge_core.domain.states import InstanceState


class HandlerError(RuntimeError):
    """A task failed in a way worth retrying. The message lands in tasks.last_error."""


async def _load_instance(task: Task, deps: WorkerDeps) -> Instance:
    async with deps.pool.connection() as conn, conn.cursor() as cur:
        await cur.execute(
            """select id, team, service_type, size, state, namespace, release_name,
                      chart_version, endpoint, error, expires_at, created_at, updated_at
                 from instances where id = %s""",
            (task.instance_id,),
        )
        row = await cur.fetchone()
    if row is None:
        raise HandlerError(f"instance {task.instance_id} vanished")
    return Instance.model_validate(row)


def _values_for(inst: Instance, entry: CatalogEntry) -> dict[str, Any]:
    """Turn a catalog size into helm values."""
    size = entry.sizes.get(inst.size)
    if size is None:
        raise HandlerError(f"size {inst.size!r} not in catalog for {inst.service_type!r}")
    return {"replicaCount": size.replicas, "resources": size.resources}


async def handle_provision(task: Task, deps: WorkerDeps) -> None:
    """Install the release and mark the instance ready. Idempotent."""
    inst = await _load_instance(task, deps)

    if inst.state is InstanceState.READY:
        # A previous attempt already finished; the crash was after the work, before the
        # bookkeeping. Nothing to do — and re-installing would be the bug.
        return

    entry = deps.catalog.get(inst.service_type)
    if entry is None:
        raise HandlerError(f"unknown service_type {inst.service_type!r}")

    # Best-effort CAS. It returning False means someone else moved the row; the helm call
    # below is idempotent either way, so this is bookkeeping, not a lock.
    await deps.instances.update_state(inst.id, InstanceState.REQUESTED, InstanceState.PROVISIONING)

    await deps.provisioner.install(
        release=inst.release_name,
        ns=inst.namespace,
        entry=entry,
        values=_values_for(inst, entry),
    )

    endpoint = f"http://{inst.release_name}.{inst.namespace}.svc.cluster.local"
    ok = await deps.instances.update_state(
        inst.id, InstanceState.PROVISIONING, InstanceState.READY, endpoint=endpoint
    )
    if ok:
        await deps.notifier.send(
            "instance.ready",
            f"instance {inst.id} is ready at {endpoint}",
            {"instance_id": str(inst.id), "team": inst.team, "service_type": inst.service_type},
        )


async def handle_deprovision(task: Task, deps: WorkerDeps) -> None:
    """Remove the release and mark the instance deleted. Idempotent."""
    inst = await _load_instance(task, deps)

    if inst.state is InstanceState.DELETED:
        return

    # `helm uninstall` of an already-gone release is not an error to us: the adapter
    # swallows not-found, because the desired state — no release — is already true.
    await deps.provisioner.uninstall(release=inst.release_name, ns=inst.namespace)
    await deps.instances.update_state(inst.id, InstanceState.DELETING, InstanceState.DELETED)


async def handle_upgrade(task: Task, deps: WorkerDeps) -> None:
    """Upgrade the release to the catalog's pinned version, then record it.

    `instances.chart_version` is written only AFTER helm reports success. That column is
    what the day-2 work-list query compares against, so writing it optimistically would
    make the fleet look upgraded while it isn't.
    """
    inst = await _load_instance(task, deps)
    entry = deps.catalog.get(inst.service_type)
    if entry is None:
        raise HandlerError(f"unknown service_type {inst.service_type!r}")

    if inst.chart_version == entry.chart_version:
        return  # already there

    await deps.provisioner.install(
        release=inst.release_name,
        ns=inst.namespace,
        entry=entry,
        values=_values_for(inst, entry),
    )

    async with deps.pool.connection() as conn, conn.cursor() as cur:
        await cur.execute(
            "update instances set chart_version = %s, updated_at = now() where id = %s",
            (entry.chart_version, inst.id),
        )


async def handle_verify(task: Task, deps: WorkerDeps) -> None:
    """Post-upgrade health probe. On failure, halt the whole rollout for this service type.

    One column decides whether the fleet keeps rolling. The work-list query returns nothing
    while `rollout_state='halted'`, so a bad chart stops after the first tenant instead of
    after all of them. You clear it with SQL, deliberately: an automatic un-halt would just
    resume breaking things.
    """
    inst = await _load_instance(task, deps)
    releases = {r.name for r in await deps.provisioner.list_releases()}

    if inst.release_name in releases:
        return

    async with deps.pool.connection() as conn, conn.cursor() as cur:
        await cur.execute(
            """insert into catalog_versions (service_type, rollout_state)
                    values (%s, 'halted')
               on conflict (service_type) do update set rollout_state = 'halted'""",
            (inst.service_type,),
        )
    await deps.notifier.send(
        "rollout.halted",
        f"rollout halted for {inst.service_type}: {inst.release_name} failed verify",
        {"instance_id": str(inst.id), "team": inst.team, "service_type": inst.service_type},
    )
    raise HandlerError(f"verify failed for {inst.release_name}; rollout halted")


HANDLERS: dict[TaskKind, Callable[[Task, WorkerDeps], Awaitable[None]]] = {
    TaskKind.PROVISION: handle_provision,
    TaskKind.DEPROVISION: handle_deprovision,
    TaskKind.UPGRADE: handle_upgrade,
    TaskKind.VERIFY: handle_verify,
}