Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 5 additions & 1 deletion netengine/api/routes.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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,
Expand Down Expand Up @@ -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),
}

Expand Down
22 changes: 20 additions & 2 deletions netengine/cli/doctor.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@

from __future__ import annotations

import asyncio
import json
from dataclasses import asdict
from pathlib import Path
Expand All @@ -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,
Expand All @@ -26,6 +28,7 @@
run_preflight,
)
from netengine.spec.loader import SpecLoadError, load_spec
from netengine.workers.registry import expected_worker_statuses

__all__ = [
"DoctorCheckResult",
Expand Down Expand Up @@ -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")
80 changes: 69 additions & 11 deletions netengine/core/consumer_supervisor.py
Original file line number Diff line number Diff line change
@@ -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

Expand All @@ -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."""
Expand All @@ -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")
Expand All @@ -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()}
6 changes: 5 additions & 1 deletion netengine/handlers/phase_pki.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
4 changes: 2 additions & 2 deletions netengine/monitoring/service.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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:
Expand Down
21 changes: 13 additions & 8 deletions netengine/phases/phase_ands.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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,
Expand Down
2 changes: 1 addition & 1 deletion netengine/phases/phase_registries.py
Original file line number Diff line number Diff line change
Expand Up @@ -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),
)

Expand Down
54 changes: 42 additions & 12 deletions netengine/phases/phase_services.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,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.events.schema import EventEnvelope
from netengine.handlers._base import BasePhaseHandler
from netengine.handlers.context import PhaseContext
Expand Down Expand Up @@ -112,21 +112,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),
)
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)
# 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)
Expand Down
Loading