dory-processor-sdk 0.0.1__py3-none-any.whl

dory/health/server.py

@@ -0,0 +1,635 @@
+"""
+Health and metrics HTTP server.
+
+Provides endpoints for:
+- /health - Liveness probe
+- /ready - Readiness probe (matches Kubernetes convention)
+- /metrics - Prometheus metrics
+- /state - State transfer (GET/POST) for pod migration (authenticated, with timeout)
+- /prestop - PreStop hook handler for graceful shutdown
+"""
+
+import json
+import logging
+import os
+import secrets
+import time
+from concurrent.futures import ThreadPoolExecutor, TimeoutError as FuturesTimeoutError
+from typing import TYPE_CHECKING, Callable, Awaitable
+
+from aiohttp import web
+
+import dory
+from dory.health.probes import LivenessProbe, ReadinessProbe
+from dory.utils.errors import DoryHealthError
+from dory.migration.transfer import (
+    TransferConfig,
+    validate_state_size,
+    StateSizeExceeded,
+    StateTransferTimeout,
+    log_transfer_summary,
+    TransferMetrics,
+    ORCHESTRATOR_STATE_TIMEOUT_SEC,
+    ORCHESTRATOR_MAX_STATE_SIZE,
+)
+
+if TYPE_CHECKING:
+    from dory.metrics.collector import MetricsCollector
+
+logger = logging.getLogger(__name__)
+
+# SDK detection header — must match orchestrator's sdkcheck.SDKVersionHeader
+SDK_VERSION_HEADER = "X-Dory-SDK-Version"
+
+# Type aliases for callbacks
+StateGetter = Callable[[], dict]
+StateRestorer = Callable[[dict], Awaitable[None]]
+PreStopHandler = Callable[[], Awaitable[None]]
+
+
+class HealthServer:
+    """
+    HTTP server for health probes, metrics, and state transfer.
+
+    Runs on a separate port from the main application.
+    Provides endpoints required by Dory Orchestrator for:
+    - Health probes (liveness/readiness)
+    - Prometheus metrics
+    - State transfer during pod migration
+    - PreStop hook for graceful shutdown
+    """
+
+    # Fixed endpoint paths (not configurable)
+    HEALTH_PATH = "/health"
+    READY_PATH = "/ready"
+    METRICS_PATH = "/metrics"
+
+    def __init__(
+        self,
+        port: int = 8080,
+        metrics_collector: "MetricsCollector | None" = None,
+        state_getter: StateGetter | None = None,
+        state_restorer: StateRestorer | None = None,
+        prestop_handler: PreStopHandler | None = None,
+        state_token: str | None = None,
+        transfer_config: TransferConfig | None = None,
+    ):
+        """
+        Initialize health server.
+
+        Args:
+            port: Port to listen on
+            metrics_collector: Optional metrics collector for /metrics endpoint
+            state_getter: Callback to get processor state for /state GET
+            state_restorer: Callback to restore processor state for /state POST
+            prestop_handler: Callback for /prestop PreStop hook
+            state_token: Authentication token for /state endpoints. If not provided,
+                reads from DORY_STATE_TOKEN environment variable. If neither is set,
+                state endpoints are unauthenticated (not recommended in production).
+            transfer_config: Configuration for state transfer timeouts and size limits.
+                If not provided, uses defaults aligned with Orchestrator (25s timeout, 8MB max).
+        """
+        self._port = port
+        self._metrics_collector = metrics_collector
+        self._state_getter = state_getter
+        self._state_restorer = state_restorer
+        self._prestop_handler = prestop_handler
+
+        # State endpoint authentication token (matches Orchestrator's DORY_STATE_TOKEN)
+        self._state_token = state_token or os.environ.get("DORY_STATE_TOKEN")
+        if not self._state_token:
+            logger.warning(
+                "DORY_STATE_TOKEN not configured - state endpoints are unauthenticated. "
+                "Set DORY_STATE_TOKEN environment variable for production deployments."
+            )
+
+        # State transfer configuration (aligned with Orchestrator limits)
+        self._transfer_config = transfer_config or TransferConfig()
+        logger.info(
+            f"State transfer configured: capture_timeout={self._transfer_config.capture_timeout_sec}s "
+            f"(Orchestrator: {ORCHESTRATOR_STATE_TIMEOUT_SEC}s), "
+            f"max_size={self._transfer_config.max_size_bytes:,}B "
+            f"(Orchestrator: {ORCHESTRATOR_MAX_STATE_SIZE:,}B)"
+        )
+
+        # Thread pool for running synchronous state getter with timeout
+        self._executor = ThreadPoolExecutor(max_workers=2, thread_name_prefix="state-capture")
+
+        self._liveness = LivenessProbe()
+        self._readiness = ReadinessProbe()
+
+        self._app: web.Application | None = None
+        self._runner: web.AppRunner | None = None
+        self._site: web.TCPSite | None = None
+
+    @property
+    def liveness_probe(self) -> LivenessProbe:
+        """Get liveness probe for adding custom checks."""
+        return self._liveness
+
+    @property
+    def readiness_probe(self) -> ReadinessProbe:
+        """Get readiness probe for adding custom checks."""
+        return self._readiness
+
+    def mark_ready(self) -> None:
+        """Mark the application as ready to receive traffic."""
+        self._readiness.mark_ready()
+
+    def mark_not_ready(self) -> None:
+        """Mark the application as not ready."""
+        self._readiness.mark_not_ready()
+
+    def set_state_getter(self, getter: StateGetter) -> None:
+        """Set the callback for getting processor state."""
+        self._state_getter = getter
+
+    def set_state_restorer(self, restorer: StateRestorer) -> None:
+        """Set the callback for restoring processor state."""
+        self._state_restorer = restorer
+
+    def set_prestop_handler(self, handler: PreStopHandler) -> None:
+        """Set the callback for PreStop hook."""
+        self._prestop_handler = handler
+
+    def set_state_token(self, token: str) -> None:
+        """Set the authentication token for state endpoints."""
+        self._state_token = token
+
+    @property
+    def port(self) -> int:
+        """Get the actual port the server is running on."""
+        return self._port
+
+    def _find_available_port(self, start: int = 8080, end: int = 9000) -> int:
+        """
+        Find an available port in the given range.
+
+        Args:
+            start: Start of port range
+            end: End of port range
+
+        Returns:
+            Available port number
+
+        Raises:
+            DoryHealthError: If no available port found
+        """
+        import socket
+
+        for port in range(start, end):
+            try:
+                with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+                    s.bind(("0.0.0.0", port))
+                    logger.debug(f"Found available port: {port}")
+                    return port
+            except OSError:
+                continue
+
+        raise DoryHealthError(f"No available port found in range {start}-{end}")
+
+    def _verify_state_auth(self, request: web.Request) -> tuple[bool, str]:
+        """
+        Verify authentication for state endpoints.
+
+        Uses Bearer token authentication matching Orchestrator's implementation.
+        Token comparison uses constant-time comparison to prevent timing attacks.
+
+        Args:
+            request: The incoming HTTP request
+
+        Returns:
+            Tuple of (is_authenticated, error_message)
+        """
+        # If no token configured, allow unauthenticated access (with warning logged at startup)
+        if not self._state_token:
+            return True, ""
+
+        auth_header = request.headers.get("Authorization", "")
+
+        if not auth_header:
+            return False, "Missing Authorization header"
+
+        # Expected format: "Bearer <token>"
+        if not auth_header.startswith("Bearer "):
+            return False, "Invalid Authorization header format (expected 'Bearer <token>')"
+
+        provided_token = auth_header[7:]  # Strip "Bearer " prefix
+
+        # Use constant-time comparison to prevent timing attacks
+        if not secrets.compare_digest(provided_token, self._state_token):
+            return False, "Invalid authentication token"
+
+        return True, ""
+
+    async def start(self) -> None:
+        """
+        Start the health server.
+
+        If port is 0, automatically finds an available port.
+
+        Raises:
+            DoryHealthError: If server fails to start
+        """
+        try:
+            self._app = web.Application()
+            self._setup_routes()
+
+            self._runner = web.AppRunner(self._app)
+            await self._runner.setup()
+
+            # If port is 0, find an available port (development mode)
+            port = self._port
+            if port == 0:
+                # Warn if running in Kubernetes with auto-port selection
+                if os.environ.get("KUBERNETES_SERVICE_HOST") or os.environ.get("DORY_POD_NAME"):
+                    logger.warning(
+                        "Auto-port selection (port=0) is not recommended in Kubernetes. "
+                        "K8s probes require a fixed port. Set DORY_HEALTH_PORT or use production preset."
+                    )
+                port = self._find_available_port()
+                self._port = port
+
+            self._site = web.TCPSite(
+                self._runner,
+                host="0.0.0.0",
+                port=port,
+            )
+            await self._site.start()
+
+            logger.info(f"Health server started on port {self._port}")
+
+        except Exception as e:
+            raise DoryHealthError(f"Failed to start health server: {e}", cause=e)
+
+    async def stop(self) -> None:
+        """Stop the health server."""
+        if self._runner:
+            await self._runner.cleanup()
+
+        # Shutdown executor
+        self._executor.shutdown(wait=False)
+
+        logger.info("Health server stopped")
+
+    def _setup_routes(self) -> None:
+        """Configure HTTP routes."""
+        self._app.router.add_get(self.HEALTH_PATH, self._handle_health)
+        self._app.router.add_get(self.READY_PATH, self._handle_ready)
+        self._app.router.add_get(self.METRICS_PATH, self._handle_metrics)
+
+        # State transfer endpoints (required by Dory Orchestrator)
+        self._app.router.add_get("/state", self._handle_state_get)
+        self._app.router.add_post("/state", self._handle_state_post)
+
+        # PreStop hook endpoint (required by Dory Orchestrator)
+        self._app.router.add_get("/prestop", self._handle_prestop)
+
+        # Root endpoint for basic info
+        self._app.router.add_get("/", self._handle_root)
+
+    async def _handle_root(self, request: web.Request) -> web.Response:
+        """Handle root endpoint."""
+        return web.json_response(
+            {
+                "service": "dory-processor",
+                "sdk_version": dory.__version__,
+                "endpoints": [
+                    self.HEALTH_PATH,
+                    self.READY_PATH,
+                    self.METRICS_PATH,
+                    "/state",
+                    "/prestop",
+                ],
+            },
+            headers={SDK_VERSION_HEADER: dory.__version__},
+        )
+
+    async def _handle_health(self, request: web.Request) -> web.Response:
+        """
+        Handle liveness probe.
+
+        Returns 200 if alive, 503 if unhealthy.
+        Includes X-Dory-SDK-Version header for orchestrator SDK detection.
+        """
+        result = await self._liveness.check()
+
+        status = 200 if result.healthy else 503
+        return web.json_response(
+            result.to_dict(),
+            status=status,
+            headers={SDK_VERSION_HEADER: dory.__version__},
+        )
+
+    async def _handle_ready(self, request: web.Request) -> web.Response:
+        """
+        Handle readiness probe.
+
+        Returns 200 if ready, 503 if not ready.
+        """
+        result = await self._readiness.check()
+
+        status = 200 if result.healthy else 503
+        return web.json_response(result.to_dict(), status=status)
+
+    async def _handle_metrics(self, request: web.Request) -> web.Response:
+        """
+        Handle Prometheus metrics endpoint.
+
+        Returns metrics in Prometheus text format.
+        """
+        if self._metrics_collector is None:
+            return web.Response(
+                text="# No metrics collector configured\n",
+                content_type="text/plain",
+            )
+
+        try:
+            metrics_text = self._metrics_collector.export_prometheus()
+            return web.Response(
+                text=metrics_text,
+                content_type="text/plain; version=0.0.4",
+                charset="utf-8",
+            )
+        except Exception as e:
+            logger.error(f"Error exporting metrics: {e}")
+            return web.Response(
+                text=f"# Error exporting metrics: {e}\n",
+                content_type="text/plain",
+                status=500,
+            )
+
+    async def _handle_state_get(self, request: web.Request) -> web.Response:
+        """
+        Handle GET /state for state capture during migration.
+
+        Called by Dory Orchestrator to capture state from old pod
+        before transferring to new pod.
+
+        Requires authentication via Bearer token (DORY_STATE_TOKEN).
+        Enforces timeout and size limits to prevent Orchestrator timeouts.
+
+        Returns:
+            JSON response with processor state
+        """
+        # Verify authentication
+        is_authenticated, auth_error = self._verify_state_auth(request)
+        if not is_authenticated:
+            logger.warning(f"State GET authentication failed: {auth_error}")
+            return web.json_response(
+                {"error": auth_error},
+                status=401,
+            )
+
+        if self._state_getter is None:
+            logger.warning("State getter not configured, returning empty state")
+            return web.json_response({
+                "error": "state_getter not configured",
+                "data": {},
+            }, status=503)
+
+        start_time = time.monotonic()
+        metrics: TransferMetrics | None = None
+
+        try:
+            # Run state getter with timeout in thread pool
+            # (state_getter is synchronous, but we need timeout control)
+            loop = __import__("asyncio").get_event_loop()
+            state = await loop.run_in_executor(
+                self._executor,
+                self._state_getter,
+            )
+
+            # Check if we're already over timeout (executor doesn't enforce timeout)
+            capture_duration = time.monotonic() - start_time
+            if capture_duration > self._transfer_config.capture_timeout_sec:
+                metrics = TransferMetrics(
+                    duration_sec=capture_duration,
+                    size_bytes=0,
+                    size_ratio=0,
+                    timed_out=True,
+                    size_exceeded=False,
+                )
+                log_transfer_summary("capture", metrics, self._transfer_config)
+                return web.json_response(
+                    {
+                        "error": f"State capture timed out after {capture_duration:.1f}s "
+                        f"(limit: {self._transfer_config.capture_timeout_sec}s)",
+                        "timeout": True,
+                    },
+                    status=504,  # Gateway Timeout
+                )
+
+            # Serialize and validate size
+            state_json = json.dumps(state)
+            size_bytes = len(state_json.encode("utf-8"))
+
+            # Validate size
+            try:
+                validate_state_size(
+                    state_json,
+                    max_size=self._transfer_config.max_size_bytes,
+                    warn_threshold=self._transfer_config.size_warn_threshold,
+                )
+            except StateSizeExceeded as e:
+                metrics = TransferMetrics(
+                    duration_sec=time.monotonic() - start_time,
+                    size_bytes=size_bytes,
+                    size_ratio=size_bytes / self._transfer_config.max_size_bytes,
+                    timed_out=False,
+                    size_exceeded=True,
+                )
+                log_transfer_summary("capture", metrics, self._transfer_config)
+                return web.json_response(
+                    {
+                        "error": str(e),
+                        "size_exceeded": True,
+                        "size_bytes": size_bytes,
+                        "max_bytes": self._transfer_config.max_size_bytes,
+                    },
+                    status=413,  # Payload Too Large
+                )
+
+            # Success - log metrics
+            metrics = TransferMetrics(
+                duration_sec=time.monotonic() - start_time,
+                size_bytes=size_bytes,
+                size_ratio=size_bytes / self._transfer_config.max_size_bytes,
+                timed_out=False,
+                size_exceeded=False,
+            )
+            log_transfer_summary("capture", metrics, self._transfer_config)
+
+            logger.info(
+                "State captured for transfer",
+                extra={
+                    "state_keys": list(state.keys()),
+                    "size_bytes": size_bytes,
+                    "duration_sec": metrics.duration_sec,
+                },
+            )
+            return web.json_response(state)
+
+        except FuturesTimeoutError:
+            metrics = TransferMetrics(
+                duration_sec=time.monotonic() - start_time,
+                size_bytes=0,
+                size_ratio=0,
+                timed_out=True,
+                size_exceeded=False,
+            )
+            log_transfer_summary("capture", metrics, self._transfer_config)
+            return web.json_response(
+                {
+                    "error": f"State capture timed out after {self._transfer_config.capture_timeout_sec}s",
+                    "timeout": True,
+                },
+                status=504,
+            )
+
+        except Exception as e:
+            logger.error(f"Failed to capture state: {e}")
+            return web.json_response(
+                {"error": f"Failed to capture state: {e}"},
+                status=500,
+            )
+
+    async def _handle_state_post(self, request: web.Request) -> web.Response:
+        """
+        Handle POST /state for state restoration during migration.
+
+        Called by Dory Orchestrator to restore state to new pod
+        after capturing from old pod.
+
+        Requires authentication via Bearer token (DORY_STATE_TOKEN).
+        Enforces timeout to prevent Orchestrator timeouts.
+
+        Returns:
+            JSON response confirming state restoration
+        """
+        import asyncio
+
+        # Verify authentication
+        is_authenticated, auth_error = self._verify_state_auth(request)
+        if not is_authenticated:
+            logger.warning(f"State POST authentication failed: {auth_error}")
+            return web.json_response(
+                {"error": auth_error},
+                status=401,
+            )
+
+        if self._state_restorer is None:
+            logger.warning("State restorer not configured")
+            return web.json_response({
+                "error": "state_restorer not configured",
+            }, status=503)
+
+        start_time = time.monotonic()
+
+        try:
+            # Read request body
+            body = await request.read()
+            size_bytes = len(body)
+
+            # Validate incoming state size
+            if size_bytes > self._transfer_config.max_size_bytes:
+                logger.error(
+                    f"Incoming state size ({size_bytes:,}B) exceeds maximum "
+                    f"({self._transfer_config.max_size_bytes:,}B)"
+                )
+                return web.json_response(
+                    {
+                        "error": f"State size ({size_bytes:,} bytes) exceeds maximum "
+                        f"({self._transfer_config.max_size_bytes:,} bytes)",
+                        "size_exceeded": True,
+                    },
+                    status=413,
+                )
+
+            state = json.loads(body.decode("utf-8"))
+            logger.info(
+                "Restoring state from transfer",
+                extra={"state_keys": list(state.keys()), "size_bytes": size_bytes},
+            )
+
+            # Execute restore with timeout
+            try:
+                await asyncio.wait_for(
+                    self._state_restorer(state),
+                    timeout=self._transfer_config.restore_timeout_sec,
+                )
+            except asyncio.TimeoutError:
+                duration = time.monotonic() - start_time
+                logger.error(
+                    f"State restore timed out after {duration:.1f}s "
+                    f"(limit: {self._transfer_config.restore_timeout_sec}s)"
+                )
+                return web.json_response(
+                    {
+                        "error": f"State restore timed out after {self._transfer_config.restore_timeout_sec}s",
+                        "timeout": True,
+                    },
+                    status=504,
+                )
+
+            duration = time.monotonic() - start_time
+            logger.info(
+                f"State restored successfully in {duration:.2f}s",
+                extra={"size_bytes": size_bytes, "duration_sec": duration},
+            )
+
+            # Warn if restore took significant time
+            if duration > self._transfer_config.restore_timeout_sec * 0.5:
+                logger.warning(
+                    f"State restore took {duration:.1f}s "
+                    f"({duration/self._transfer_config.restore_timeout_sec:.0%} of "
+                    f"{self._transfer_config.restore_timeout_sec}s timeout)"
+                )
+
+            return web.json_response({
+                "status": "ok",
+                "message": "State restored",
+                "duration_sec": round(duration, 3),
+                "size_bytes": size_bytes,
+            })
+
+        except json.JSONDecodeError as e:
+            logger.error(f"Invalid JSON in state restore request: {e}")
+            return web.json_response(
+                {"error": f"Invalid JSON: {e}"},
+                status=400,
+            )
+        except Exception as e:
+            logger.error(f"Failed to restore state: {e}")
+            return web.json_response(
+                {"error": f"Failed to restore state: {e}"},
+                status=500,
+            )
+
+    async def _handle_prestop(self, request: web.Request) -> web.Response:
+        """
+        Handle GET /prestop for PreStop hook.
+
+        Called by Kubernetes PreStop hook before pod termination.
+        Allows the application to prepare for graceful shutdown.
+
+        Returns:
+            JSON response confirming prestop handling
+        """
+        logger.info("PreStop hook invoked - preparing for shutdown")
+
+        # Mark as not ready to stop receiving new traffic
+        self._readiness.mark_not_ready()
+
+        if self._prestop_handler:
+            try:
+                await self._prestop_handler()
+                logger.info("PreStop handler completed")
+            except Exception as e:
+                logger.error(f"PreStop handler error: {e}")
+                # Continue anyway - don't block shutdown
+
+        return web.json_response({
+            "status": "ok",
+            "message": "PreStop hook processed, ready for termination",
+        })