atlas-iac/services/hermes/scripts/execution_pool_server.py

#!/usr/bin/env python3
"""Versioned HTTP and maintenance loop for the execution-pool coordinator."""

from __future__ import annotations

import argparse
import os
import sqlite3
import sys
import threading
import time
from collections.abc import Iterable
from http.server import BaseHTTPRequestHandler
from pathlib import Path
from typing import Any

from execution_pool_protocol import (
    MAX_WIRE_BYTES,
    PROTOCOL_VERSION,
    BoundedHTTPServer,
    ProtocolError,
    canonical_json,
    parse_wire,
    read_key,
)
from execution_pool_store import PoolStore


DATA_ROOT = Path(os.environ.get("HERMES_HOME", "/opt/data"))
STATE_ROOT = DATA_ROOT / "execution-pool"
KEY_PATH = Path(
    os.environ.get(
        "HERMES_EXECUTION_POOL_KEY_FILE", "/runtime-access/execution-pool-key"
    )
)
PORT = int(os.environ.get("HERMES_EXECUTION_POOL_PORT", "9007"))
RETENTION_SECONDS = int(
    os.environ.get("HERMES_EXECUTION_POOL_RETENTION_SECONDS", "1209600")
)


def handler_factory(coordinator: Any) -> type[BaseHTTPRequestHandler]:
    class Handler(BaseHTTPRequestHandler):
        server_version = f"hermes-execution-pool/{PROTOCOL_VERSION}"

        def _reply(self, status: int, value: dict[str, Any]) -> None:
            body = canonical_json(value)
            self.send_response(status)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(body)))
            self.send_header("Cache-Control", "no-store")
            self.end_headers()
            self.wfile.write(body)

        def do_GET(self) -> None:  # noqa: N802
            if self.path != "/ready":
                self._reply(404, {"error": "not found"})
                return
            try:
                coordinator.store.available_ordinals()
                self._reply(200, {"ready": True, "version": PROTOCOL_VERSION})
            except (OSError, sqlite3.Error):
                self._reply(503, {"ready": False, "version": PROTOCOL_VERSION})

        def do_POST(self) -> None:  # noqa: N802
            try:
                length = int(self.headers.get("Content-Length", "0"))
                if length <= 0 or length > MAX_WIRE_BYTES:
                    raise ProtocolError("invalid content length")
                envelope = parse_wire(self.rfile.read(length))
                routes = {
                    "/v1/poll": coordinator.poll,
                    "/v1/heartbeat": coordinator.heartbeat,
                    "/v1/result": coordinator.result,
                }
                if self.path not in routes:
                    self._reply(404, {"error": "not found"})
                    return
                self._reply(200, routes[self.path](envelope))
            except ProtocolError as error:
                self._reply(409, {"error": str(error)})
            except Exception as error:
                print(
                    f"pool request failed: {type(error).__name__}: {error}",
                    file=sys.stderr,
                    flush=True,
                )
                self._reply(503, {"error": "coordinator unavailable"})

        def log_message(self, _format: str, *_arguments: Any) -> None:
            return

    return Handler


def maintenance_cycle(operations: Iterable[Any]) -> int:
    """Run one maintenance pass; a failing operation never aborts the others.

    Startup uses this too. A durable row that one pass cannot process must not
    stop the coordinator from binding its port, because the store lives on a PVC
    and would otherwise turn a single bad row into a permanent crash loop.
    """
    deferred = 0
    for operation in operations:
        try:
            operation()
        except Exception as error:  # noqa: BLE001 - the loop is the recovery mechanism
            deferred += 1
            print(
                f"pool maintenance deferred: {type(error).__name__}: {error}",
                file=sys.stderr,
                flush=True,
            )
    return deferred


def run(coordinator_type: type[Any]) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--once", action="store_true")
    args = parser.parse_args()
    store = PoolStore(STATE_ROOT / "assignments.db")
    coordinator = coordinator_type(read_key(KEY_PATH), store)
    operations = (
        coordinator.expire_leases,
        coordinator.recover_results,
        coordinator.reconcile,
        coordinator.dispatch,
        lambda: store.garbage_collect(RETENTION_SECONDS),
    )
    maintenance_cycle(operations)
    if args.once:
        return 0
    server = BoundedHTTPServer(
        ("0.0.0.0", PORT), handler_factory(coordinator), max_workers=8
    )
    threading.Thread(target=server.serve_forever, daemon=True).start()
    while True:
        maintenance_cycle(operations)
        time.sleep(5)