From aecbb58181222985e6c16a3fb253f0b2265bebf4 Mon Sep 17 00:00:00 2001 From: aidankhogg Date: Tue, 30 Jun 2026 00:40:05 +0100 Subject: [PATCH] Expand background worker supervision --- netengine/api/routes.py | 6 +- netengine/cli/doctor.py | 22 +++++- netengine/core/consumer_supervisor.py | 80 ++++++++++++++++--- netengine/handlers/phase_pki.py | 6 +- netengine/monitoring/service.py | 4 +- netengine/phases/phase_ands.py | 21 +++-- netengine/phases/phase_registries.py | 2 +- netengine/phases/phase_services.py | 56 ++++++++++---- netengine/workers/dlq_worker.py | 81 ++++++++++++++++++++ netengine/workers/registry.py | 52 +++++++++++++ tests/test_background_workers.py | 106 ++++++++++++++++++++++++++ 11 files changed, 397 insertions(+), 39 deletions(-) create mode 100644 netengine/workers/dlq_worker.py create mode 100644 netengine/workers/registry.py create mode 100644 tests/test_background_workers.py diff --git a/netengine/api/routes.py b/netengine/api/routes.py index 109b434..f0a74c4 100644 --- a/netengine/api/routes.py +++ b/netengine/api/routes.py @@ -12,7 +12,7 @@ from typing import Any import yaml -from fastapi import APIRouter, Depends, HTTPException, Request +from fastapi import APIRouter, Depends, HTTPException from pydantic import BaseModel, Field from netengine.api.auth import _extract_roles, require_admin, require_auth @@ -22,6 +22,7 @@ from netengine.events.queues import PRIMARY_QUEUES, Queue, dlq_for from netengine.logs import get_logger from netengine.phase_labels import PHASE_LABELS +from netengine.workers.registry import expected_worker_statuses from netengine.security.redaction import ( _contains_private_pem, _is_secret_field, @@ -110,10 +111,13 @@ async def health() -> dict[str, Any]: ), } + workers = expected_worker_statuses(state, pgmq_enabled=events_ok) + return { "status": overall, "phases": phases, "events": events, + "workers": workers, "last_error_present": bool(state.last_error), } diff --git a/netengine/cli/doctor.py b/netengine/cli/doctor.py index edd5a46..096794c 100644 --- a/netengine/cli/doctor.py +++ b/netengine/cli/doctor.py @@ -7,6 +7,7 @@ from __future__ import annotations +import asyncio import json from dataclasses import asdict from pathlib import Path @@ -15,7 +16,8 @@ import click from netengine.cli.env import db_url_from_env -from netengine.core.state import get_state_file +from netengine.core.db_client import pgmq_available +from netengine.core.state import RuntimeState, get_state_file from netengine.diagnostic import preflight as _preflight from netengine.diagnostic.preflight import ( DoctorCheckResult, @@ -26,6 +28,7 @@ run_preflight, ) from netengine.spec.loader import SpecLoadError, load_spec +from netengine.workers.registry import expected_worker_statuses __all__ = [ "DoctorCheckResult", @@ -140,10 +143,25 @@ def doctor( results = run_checks( db_url, state_file, skip_db=skip_db, spec_subnets=spec_subnets ) + state = RuntimeState.load() + try: + pgmq_ok = False if skip_db else asyncio.run(pgmq_available())[0] + except Exception: + pgmq_ok = False + workers = expected_worker_statuses(state, pgmq_enabled=pgmq_ok) + if as_json: - click.echo(json.dumps([asdict(r) for r in results], indent=2)) + click.echo( + json.dumps({"checks": [asdict(r) for r in results], "workers": workers}, indent=2) + ) else: click.echo("NetEngine doctor preflight report") _print_report(results) + if workers: + click.echo("\nBackground workers") + for worker in workers.values(): + detail = worker.get("disabled_reason") or worker.get("last_error") or "" + suffix = f" — {detail}" if detail else "" + click.echo(f" {worker['state']}: {worker['name']}{suffix}") if any(r.status == DoctorStatus.FAIL and r.required for r in results): raise click.ClickException("required doctor checks failed") diff --git a/netengine/core/consumer_supervisor.py b/netengine/core/consumer_supervisor.py index f873896..ed057a8 100644 --- a/netengine/core/consumer_supervisor.py +++ b/netengine/core/consumer_supervisor.py @@ -1,5 +1,7 @@ import asyncio -from typing import Any, Callable, Coroutine, Dict +from dataclasses import asdict, dataclass +from datetime import UTC, datetime +from typing import Any, Callable, Coroutine, Dict, Literal from netengine.logs import get_logger @@ -9,16 +11,51 @@ _BACKOFF_MAX = 60 +WorkerState = Literal["registered", "running", "failed", "disabled", "stopped", "completed"] + + +@dataclass +class WorkerStatus: + """Structured status for a supervised background worker.""" + + name: str + state: WorkerState + restarts: int = 0 + last_error: str | None = None + disabled_reason: str | None = None + started_at: str | None = None + stopped_at: str | None = None + last_crashed_at: str | None = None + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + class ConsumerSupervisor: """Manages long-running consumer tasks with automatic restart on failure.""" def __init__(self) -> None: self.tasks: Dict[str, asyncio.Task[Any]] = {} self.consumers: Dict[str, Callable[[], Coroutine[Any, Any, None]]] = {} + self._statuses: Dict[str, WorkerStatus] = {} + + @staticmethod + def _now() -> str: + return datetime.now(UTC).isoformat() - def register(self, name: str, consumer_coro: Callable[[], Coroutine[Any, Any, None]]) -> None: - """Register a consumer coroutine function.""" - self.consumers[name] = consumer_coro + def register(self, name: str | Any, consumer_coro: Callable[[], Coroutine[Any, Any, None]]) -> None: + """Register a consumer coroutine function with a stable operator-visible name.""" + stable_name = str(name) + self.consumers[stable_name] = consumer_coro + self._statuses[stable_name] = WorkerStatus(name=stable_name, state="registered") + + def register_disabled(self, name: str | Any, reason: str) -> None: + """Record a disabled worker so operators can see missing dependencies.""" + stable_name = str(name) + self.consumers.pop(stable_name, None) + self._statuses[stable_name] = WorkerStatus( + name=stable_name, state="disabled", disabled_reason=reason + ) async def start_all(self) -> None: """Start all registered consumers.""" @@ -36,15 +73,31 @@ async def supervised_consumer() -> None: try: logger.info(f"Starting consumer: {name}") await consumer_func() + status = self._statuses.setdefault(name, WorkerStatus(name=name, state="completed")) + status.state = "completed" + status.stopped_at = self._now() delay = _BACKOFF_BASE # reset on clean exit except asyncio.CancelledError: + status = self._statuses.setdefault(name, WorkerStatus(name=name, state="stopped")) + status.state = "stopped" + status.stopped_at = self._now() logger.info(f"Consumer {name} cancelled") break except Exception as e: + status = self._statuses.setdefault(name, WorkerStatus(name=name, state="failed")) + status.state = "failed" + status.last_error = str(e) + status.last_crashed_at = self._now() + status.restarts += 1 logger.error(f"Consumer {name} crashed: {e}. Restarting in {delay}s...") await asyncio.sleep(delay) + status.state = "running" delay = min(delay * 2, _BACKOFF_MAX) + if name in self._statuses and self._statuses[name].state == "disabled": + logger.info(f"Consumer {name} not started because it is disabled") + return + self._statuses[name] = WorkerStatus(name=name, state="running", started_at=self._now()) task: asyncio.Task[None] = asyncio.create_task(supervised_consumer()) self.tasks[name] = task logger.info(f"Consumer {name} started") @@ -60,11 +113,16 @@ async def stop_all(self) -> None: pass def get_status(self) -> Dict[str, str]: - """Get status of all consumers.""" - status: Dict[str, str] = {} + """Get legacy status of all consumers.""" + return {name: status.state for name, status in self._statuses.items()} + + def get_structured_status(self) -> dict[str, dict[str, Any]]: + """Get structured operator/doctor status for all registered workers.""" for name, task in self.tasks.items(): - if task.done(): - status[name] = "crashed" if task.exception() else "completed" - else: - status[name] = "running" - return status + status = self._statuses.setdefault(name, WorkerStatus(name=name, state="registered")) + if task.done() and status.state == "running": + exc = task.exception() + status.state = "failed" if exc else "completed" + status.last_error = str(exc) if exc else None + status.stopped_at = self._now() + return {name: status.to_dict() for name, status in self._statuses.items()} diff --git a/netengine/handlers/phase_pki.py b/netengine/handlers/phase_pki.py index 59b235e..8838141 100644 --- a/netengine/handlers/phase_pki.py +++ b/netengine/handlers/phase_pki.py @@ -203,10 +203,14 @@ def _register_rotation_worker(self, context: PhaseContext, pki: PKIHandler, spec if context.pgmq_client: rotation_worker = PKICertRotationWorker(pki, context.pgmq_client, rotation_configs) - context.consumer_supervisor.register("pki_cert_rotation", rotation_worker.run) + context.consumer_supervisor.register("pki.cert_rotation", rotation_worker.run) context.logger.info( f"PKI certificate rotation worker registered ({len(rotation_configs)} cert types)" ) + else: + context.consumer_supervisor.register_disabled( + "pki.cert_rotation", "pgmq unavailable; certificate rotation events disabled" + ) async def _prepare_app_cert_rotation(self, cn: str, cert_metadata: dict) -> None: """Called before rotating app cert - prepare for transition. diff --git a/netengine/monitoring/service.py b/netengine/monitoring/service.py index 3fd0213..0dae410 100644 --- a/netengine/monitoring/service.py +++ b/netengine/monitoring/service.py @@ -22,7 +22,7 @@ class MonitoringService: monitoring. """ - def __init__(self, spec: NetEngineSpec, interval_seconds: float = 60.0) -> None: + def __init__(self, spec: NetEngineSpec, interval_seconds: float = 60.0, pgmq: PGMQClient | None = None) -> None: """Initialize monitoring service. Args: @@ -32,7 +32,7 @@ def __init__(self, spec: NetEngineSpec, interval_seconds: float = 60.0) -> None: self._spec = spec self._interval_seconds = interval_seconds self._runner = build_runner(spec) - self._pgmq = PGMQClient() + self._pgmq = pgmq or PGMQClient() self._running = False async def start(self) -> None: diff --git a/netengine/phases/phase_ands.py b/netengine/phases/phase_ands.py index 0e5b017..f8d43f2 100644 --- a/netengine/phases/phase_ands.py +++ b/netengine/phases/phase_ands.py @@ -13,7 +13,7 @@ import ipaddress import json from datetime import UTC, datetime -from typing import Any, Optional +from typing import Any from netengine.events.emitter import emit_event from netengine.events.queues import Queue @@ -141,11 +141,17 @@ async def execute(self, context: PhaseContext) -> None: # Register org-admission consumer through ConsumerSupervisor so # crashes are visible and the task is gracefully shut down. - if context.pgmq_client is not None: - context.consumer_supervisor.register( # type: ignore[union-attr] - "org_admission_events", - lambda: self._consume_org_admission_events(context, docker, gateway, ands_spec), - ) + if context.consumer_supervisor is not None: + if context.pgmq_client is not None: + context.consumer_supervisor.register( + "ands.org_admission", + lambda: self._consume_org_admission_events(context, docker, gateway, ands_spec), + ) + else: + context.consumer_supervisor.register_disabled( + "ands.org_admission", + "pgmq unavailable; new organization AND provisioning disabled", + ) except Exception as e: context.runtime_state.last_error = str(e) @@ -273,8 +279,7 @@ async def _provision_and( # 6. Register DNS zone dns_suffix = and_instance.dns_suffix or f"{and_instance.org}.internal" try: - from netengine.handlers.dns import DNSHandler - + dns = DNSHandler() await dns.add_zone_record( context=context, diff --git a/netengine/phases/phase_registries.py b/netengine/phases/phase_registries.py index 7ab376a..fd34573 100644 --- a/netengine/phases/phase_registries.py +++ b/netengine/phases/phase_registries.py @@ -60,7 +60,7 @@ async def execute(self, context: PhaseContext) -> None: # 5. Wire pgmq consumer for DNS updates through supervisor context.consumer_supervisor.register( # type: ignore[union-attr] - Queue.DNS_UPDATES, + "dns_updates", lambda: self._consume_dns_updates(context), ) diff --git a/netengine/phases/phase_services.py b/netengine/phases/phase_services.py index 74eab37..a6d2687 100644 --- a/netengine/phases/phase_services.py +++ b/netengine/phases/phase_services.py @@ -13,7 +13,7 @@ from typing import Any from netengine.events.emitter import emit_event -from netengine.events.queues import Queue +from netengine.events.queues import Queue, dlq_for from netengine.handlers._base import BasePhaseHandler from netengine.handlers.context import PhaseContext from netengine.handlers.dns import DNSHandler @@ -107,21 +107,51 @@ async def execute(self, context: PhaseContext) -> None: payload={"services": list(services_output.keys())}, ) - # Register background consumers + # Register background consumers with stable, operator-visible names. if context.consumer_supervisor is not None: - context.consumer_supervisor.register( - "org_admission_events", - lambda: self._consume_org_admission_events(context, docker, dns), - ) - - # Register monitoring service (always-running health checks) + if context.pgmq_client is not None: + context.consumer_supervisor.register( + "services.org_admission", + lambda: self._consume_org_admission_events(context, docker, dns), + ) + + from netengine.workers.dlq_worker import DLQReplayWorker + + services_dlq = DLQReplayWorker( + context.pgmq_client, + Queue.SERVICES_ADMISSIONS, + dlq_for(Queue.SERVICES_ADMISSIONS), + ) + context.consumer_supervisor.register( + "dlq.services_admissions", services_dlq.run + ) + else: + context.consumer_supervisor.register_disabled( + "services.org_admission", + "pgmq unavailable; per-org service provisioning disabled", + ) + context.consumer_supervisor.register_disabled( + "dlq.services_admissions", + "pgmq unavailable; DLQ replay disabled", + ) + + # Register monitoring service (always-running health checks). Event + # publication is disabled visibly when pgmq is unavailable. from netengine.monitoring import MonitoringService - monitoring_service = MonitoringService(spec, interval_seconds=60.0) - context.consumer_supervisor.register( - "monitoring_service", - monitoring_service.start, - ) + if context.pgmq_client is not None: + monitoring_service = MonitoringService( + spec, interval_seconds=60.0, pgmq=context.pgmq_client + ) + context.consumer_supervisor.register( + "monitoring.world_health", + monitoring_service.start, + ) + else: + context.consumer_supervisor.register_disabled( + "monitoring.world_health", + "pgmq unavailable; world health events disabled", + ) except Exception as e: runtime_state.last_error = str(e) diff --git a/netengine/workers/dlq_worker.py b/netengine/workers/dlq_worker.py new file mode 100644 index 0000000..2fc4c4d --- /dev/null +++ b/netengine/workers/dlq_worker.py @@ -0,0 +1,81 @@ +"""Dead-letter queue watch and guarded replay workers.""" + +import asyncio +import json +from typing import Any + +from netengine.core.pgmq_client import PGMQClient +from netengine.events.queues import Queue +from netengine.events.schema import EventEnvelope +from netengine.logs import get_logger + +logger = get_logger(__name__) + + +class DLQReplayWorker: + """Watch a DLQ and replay only messages explicitly marked resolved. + + Operators or remediation jobs must set ``payload.dlq_resolved`` to true before + this worker will move an event back to its primary queue. This avoids hot-loop + replay of unresolved poison messages. + """ + + def __init__( + self, + pgmq: PGMQClient, + primary_queue: Queue, + dlq_queue: Queue, + *, + poll_interval_seconds: float = 5.0, + ) -> None: + self._pgmq = pgmq + self._primary_queue = primary_queue + self._dlq_queue = dlq_queue + self._poll_interval_seconds = poll_interval_seconds + + async def run(self) -> None: + logger.info( + "Starting DLQ replay worker for %s -> %s", self._dlq_queue, self._primary_queue + ) + while True: + replayed = await self.replay_once() + if not replayed: + await asyncio.sleep(self._poll_interval_seconds) + + async def replay_once(self) -> bool: + """Attempt one guarded replay. Returns True when a message was replayed.""" + msg = await self._pgmq.receive(self._dlq_queue) + if not msg: + return False + + envelope = self._decode(msg) + if not envelope.payload.get("dlq_resolved"): + logger.warning( + "Leaving unresolved DLQ message %s on %s; payload.dlq_resolved is not true", + msg["msg_id"], + self._dlq_queue, + ) + return False + + replay_payload: dict[str, Any] = dict(envelope.payload) + replay_payload.pop("dlq_reason", None) + replay_payload.pop("dlq_resolved", None) + replay = EventEnvelope( + event_id=envelope.event_id, + correlation_id=envelope.correlation_id, + event_type=envelope.event_type, + emitted_by=envelope.emitted_by, + emitted_at=envelope.emitted_at, + payload=replay_payload, + parent_event_id=envelope.parent_event_id, + retry_count=0, + ) + await self._pgmq.send(self._primary_queue, replay) + await self._pgmq.delete(self._dlq_queue, msg["msg_id"]) + logger.info("Replayed resolved DLQ message %s to %s", msg["msg_id"], self._primary_queue) + return True + + @staticmethod + def _decode(msg: dict[str, Any]) -> EventEnvelope: + payload = json.loads(msg["message"]) + return EventEnvelope(**payload) diff --git a/netengine/workers/registry.py b/netengine/workers/registry.py new file mode 100644 index 0000000..72f215a --- /dev/null +++ b/netengine/workers/registry.py @@ -0,0 +1,52 @@ +"""Background worker inventory and status helpers for operator surfaces.""" + +from typing import Any + +from netengine.core.state import RuntimeState + +PGMQ_WORKERS: tuple[tuple[str, str], ...] = ( + ("pki.cert_rotation", "PKI certificate rotation"), + ("ands.org_admission", "AND organization admission provisioning"), + ("services.org_admission", "World-services organization admission provisioning"), + ("monitoring.world_health", "Monitoring world-health publisher"), + ("drift_detection", "Phase drift detection"), + ("dlq.services_admissions", "Services admissions DLQ replay"), +) +ALWAYS_WORKERS: tuple[tuple[str, str], ...] = ( + ("whois_server", "WHOIS registry server"), + ("dns_updates", "Registry DNS update consumer"), +) + + +def expected_worker_statuses( + state: RuntimeState, *, pgmq_enabled: bool, live: dict[str, dict[str, Any]] | None = None +) -> dict[str, dict[str, Any]]: + """Build structured worker status for API/doctor without requiring in-process tasks.""" + live = live or {} + statuses: dict[str, dict[str, Any]] = {} + + def add(name: str, description: str, enabled: bool, reason: str | None = None) -> None: + if name in live: + statuses[name] = {**live[name], "description": description} + return + statuses[name] = { + "name": name, + "description": description, + "state": "registered" if enabled else "disabled", + "disabled_reason": reason, + "last_error": None, + "restarts": 0, + } + + if state.phase_completed.get("5"): + for name, desc in ALWAYS_WORKERS: + add(name, desc, True) + if state.phase_completed.get("3"): + add(PGMQ_WORKERS[0][0], PGMQ_WORKERS[0][1], pgmq_enabled, "pgmq unavailable") + if state.phase_completed.get("7"): + add(PGMQ_WORKERS[1][0], PGMQ_WORKERS[1][1], pgmq_enabled, "pgmq unavailable") + if state.phase_completed.get("8"): + for name, desc in PGMQ_WORKERS[2:4] + PGMQ_WORKERS[5:]: + add(name, desc, pgmq_enabled, "pgmq unavailable") + add(PGMQ_WORKERS[4][0], PGMQ_WORKERS[4][1], pgmq_enabled, "pgmq unavailable") + return statuses diff --git a/tests/test_background_workers.py b/tests/test_background_workers.py new file mode 100644 index 0000000..f5651cd --- /dev/null +++ b/tests/test_background_workers.py @@ -0,0 +1,106 @@ +import asyncio +import json +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from netengine.core.consumer_supervisor import ConsumerSupervisor +from netengine.events.queues import Queue, dlq_for +from netengine.events.schema import EventEnvelope +from netengine.phases.phase_services import ServicesPhaseHandler +from netengine.workers.dlq_worker import DLQReplayWorker + + +@pytest.mark.asyncio +async def test_consumer_supervisor_reports_crash_and_graceful_shutdown(monkeypatch): + supervisor = ConsumerSupervisor() + first_run = asyncio.Event() + + async def crashing_worker(): + first_run.set() + raise RuntimeError("boom") + + async def long_worker(): + await asyncio.Event().wait() + + supervisor.register("worker.crashes", crashing_worker) + supervisor.register("worker.long", long_worker) + + await supervisor.start_all() + await asyncio.wait_for(first_run.wait(), timeout=1) + await asyncio.sleep(0) + + crashed = supervisor.get_structured_status()["worker.crashes"] + assert crashed["state"] == "failed" + assert crashed["last_error"] == "boom" + assert crashed["restarts"] == 1 + + await supervisor.stop_all() + stopped = supervisor.get_structured_status()["worker.long"] + assert stopped["state"] == "stopped" + + +@pytest.mark.asyncio +async def test_services_workers_register_disabled_when_pgmq_unavailable(minimal_spec): + context = MagicMock() + context.spec = minimal_spec + context.runtime_state = MagicMock() + context.runtime_state.substrate_output = {"ok": True} + context.runtime_state.dns_output = {"ok": True} + context.runtime_state.pki_output = {"ok": True} + context.runtime_state.identity_platform_output = {"ok": True} + context.runtime_state.world_registry_output = {"ok": True} + context.runtime_state.identity_inworld_output = {"ok": True} + context.runtime_state.ands_output = {"ok": True} + context.consumer_supervisor = ConsumerSupervisor() + context.pgmq_client = None + context.logger = MagicMock() + + with patch("netengine.phases.phase_services.DockerHandler"), patch( + "netengine.phases.phase_services.DNSHandler" + ), patch("netengine.phases.phase_services.PKIHandler"), patch( + "netengine.phases.phase_services.MailHandler" + ) as mail_cls, patch("netengine.phases.phase_services.StorageHandler") as storage_cls: + mail_cls.return_value.deploy_postfix = AsyncMock(return_value={"container_id": "mail"}) + storage_cls.return_value.deploy_minio = AsyncMock(return_value={"container_id": "minio"}) + await ServicesPhaseHandler().execute(context) + + status = context.consumer_supervisor.get_structured_status() + assert status["services.org_admission"]["state"] == "disabled" + assert status["dlq.services_admissions"]["state"] == "disabled" + assert status["monitoring.world_health"]["state"] == "disabled" + + +@pytest.mark.asyncio +async def test_dlq_replay_requires_resolved_marker(): + pgmq = AsyncMock() + unresolved = EventEnvelope.create( + event_type="services.test", emitted_by="test", payload={"dlq_reason": "still broken"} + ) + pgmq.receive.return_value = {"msg_id": 7, "message": json.dumps(unresolved.to_dict())} + + worker = DLQReplayWorker(pgmq, Queue.SERVICES_ADMISSIONS, dlq_for(Queue.SERVICES_ADMISSIONS)) + + assert await worker.replay_once() is False + pgmq.send.assert_not_called() + pgmq.delete.assert_not_called() + + +@pytest.mark.asyncio +async def test_dlq_replay_moves_resolved_message_back_to_primary(): + pgmq = AsyncMock() + resolved = EventEnvelope.create( + event_type="services.test", + emitted_by="test", + payload={"dlq_reason": "fixed", "dlq_resolved": True, "org": "acme"}, + ) + pgmq.receive.return_value = {"msg_id": 8, "message": json.dumps(resolved.to_dict())} + + worker = DLQReplayWorker(pgmq, Queue.SERVICES_ADMISSIONS, dlq_for(Queue.SERVICES_ADMISSIONS)) + + assert await worker.replay_once() is True + sent_queue, sent_event = pgmq.send.call_args.args + assert sent_queue == Queue.SERVICES_ADMISSIONS + assert sent_event.retry_count == 0 + assert sent_event.payload == {"org": "acme"} + pgmq.delete.assert_awaited_once_with(dlq_for(Queue.SERVICES_ADMISSIONS), 8)