dory-sdk 2.1.0__py3-none-any.whl → 2.1.4__py3-none-any.whl

dory/edge/heartbeat.py

@@ -0,0 +1,598 @@
+"""Edge heartbeat service for Orchestrator integration.
+
+Provides connectivity monitoring and health reporting for edge nodes
+with intermittent connectivity to the cloud/Orchestrator.
+"""
+
+import asyncio
+import logging
+import os
+import time
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Callable, Awaitable
+
+import aiohttp
+
+from dory.edge.fencing import FencingToken
+from dory.edge.role import ProcessorRole
+
+logger = logging.getLogger(__name__)
+
+
+class ConnectivityStatus(Enum):
+    """Connectivity status to Orchestrator/cloud."""
+
+    CONNECTED = "connected"
+    DEGRADED = "degraded"  # High latency or packet loss
+    DISCONNECTED = "disconnected"
+    UNKNOWN = "unknown"
+
+
+@dataclass
+class HeartbeatConfig:
+    """Configuration for edge heartbeat service."""
+
+    # Orchestrator endpoint for heartbeat reporting
+    orchestrator_url: str | None = None
+
+    # Heartbeat interval in seconds
+    interval_sec: float = 10.0
+
+    # Timeout for heartbeat requests
+    timeout_sec: float = 5.0
+
+    # Number of missed heartbeats before considered disconnected
+    missed_threshold: int = 3
+
+    # Number of consecutive successes to transition from degraded to connected
+    recovery_threshold: int = 2
+
+    # Latency threshold for degraded status (ms)
+    latency_threshold_ms: float = 500.0
+
+    # Enable automatic role demotion on disconnect
+    auto_demote_on_disconnect: bool = True
+
+    # Grace period before demoting (seconds)
+    demote_grace_period_sec: float = 30.0
+
+    def __post_init__(self):
+        """Load from environment if not provided."""
+        if not self.orchestrator_url:
+            self.orchestrator_url = os.environ.get(
+                "DORY_ORCHESTRATOR_URL",
+                "http://dory-orchestrator:8080",
+            )
+
+
+@dataclass
+class HeartbeatPayload:
+    """Payload sent to Orchestrator in heartbeat."""
+
+    processor_id: str
+    node_id: str
+    role: str
+    epoch: int | None
+    timestamp: float = field(default_factory=time.time)
+    status: str = "healthy"
+    last_state_sync: float | None = None
+    state_size_bytes: int | None = None
+    uptime_sec: float = 0.0
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary for JSON serialization."""
+        return {
+            "processor_id": self.processor_id,
+            "node_id": self.node_id,
+            "role": self.role,
+            "epoch": self.epoch,
+            "timestamp": self.timestamp,
+            "status": self.status,
+            "last_state_sync": self.last_state_sync,
+            "state_size_bytes": self.state_size_bytes,
+            "uptime_sec": self.uptime_sec,
+        }
+
+
+@dataclass
+class HeartbeatResponse:
+    """Response from Orchestrator heartbeat endpoint."""
+
+    acknowledged: bool
+    orchestrator_time: float
+    directive: str | None = None  # "continue", "demote", "promote", "shutdown"
+    message: str | None = None
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "HeartbeatResponse":
+        """Create from dictionary."""
+        return cls(
+            acknowledged=data.get("acknowledged", False),
+            orchestrator_time=data.get("orchestrator_time", time.time()),
+            directive=data.get("directive"),
+            message=data.get("message"),
+        )
+
+
+@dataclass
+class ConnectivityMetrics:
+    """Metrics tracking connectivity health."""
+
+    total_heartbeats: int = 0
+    successful_heartbeats: int = 0
+    failed_heartbeats: int = 0
+    consecutive_failures: int = 0
+    consecutive_successes: int = 0
+    last_success_time: float | None = None
+    last_failure_time: float | None = None
+    last_latency_ms: float | None = None
+    avg_latency_ms: float = 0.0
+    max_latency_ms: float = 0.0
+
+    def record_success(self, latency_ms: float) -> None:
+        """Record a successful heartbeat."""
+        self.total_heartbeats += 1
+        self.successful_heartbeats += 1
+        self.consecutive_successes += 1
+        self.consecutive_failures = 0
+        self.last_success_time = time.time()
+        self.last_latency_ms = latency_ms
+
+        # Update latency stats
+        if self.successful_heartbeats == 1:
+            self.avg_latency_ms = latency_ms
+        else:
+            # Exponential moving average
+            self.avg_latency_ms = 0.9 * self.avg_latency_ms + 0.1 * latency_ms
+
+        self.max_latency_ms = max(self.max_latency_ms, latency_ms)
+
+    def record_failure(self) -> None:
+        """Record a failed heartbeat."""
+        self.total_heartbeats += 1
+        self.failed_heartbeats += 1
+        self.consecutive_failures += 1
+        self.consecutive_successes = 0
+        self.last_failure_time = time.time()
+
+    def get_success_rate(self) -> float:
+        """Get heartbeat success rate."""
+        if self.total_heartbeats == 0:
+            return 0.0
+        return self.successful_heartbeats / self.total_heartbeats
+
+
+# Type alias for connectivity change callback
+ConnectivityChangeCallback = Callable[[ConnectivityStatus, ConnectivityStatus], Awaitable[None]]
+
+
+class HeartbeatManager:
+    """Manager for edge heartbeat reporting and connectivity monitoring.
+
+    Periodically sends heartbeats to the Orchestrator and tracks
+    connectivity status. Can automatically trigger role transitions
+    when connectivity is lost.
+
+    Usage:
+        config = HeartbeatConfig(
+            orchestrator_url="http://orchestrator:8080",
+            interval_sec=10.0,
+        )
+        manager = HeartbeatManager(config)
+
+        # Set processor info
+        manager.set_processor_info(
+            processor_id="my-processor",
+            node_id="edge-node-1",
+        )
+
+        # Start heartbeat loop
+        await manager.start()
+
+        # Check connectivity
+        if manager.is_connected():
+            # Safe to sync state
+            pass
+
+        # Stop on shutdown
+        await manager.stop()
+    """
+
+    def __init__(self, config: HeartbeatConfig | None = None):
+        """Initialize heartbeat manager.
+
+        Args:
+            config: Heartbeat configuration
+        """
+        self._config = config or HeartbeatConfig()
+        self._processor_id: str | None = None
+        self._node_id: str | None = None
+        self._role: ProcessorRole = ProcessorRole.INITIALIZING
+        self._fencing_token: FencingToken | None = None
+
+        self._status = ConnectivityStatus.UNKNOWN
+        self._metrics = ConnectivityMetrics()
+        self._callbacks: list[ConnectivityChangeCallback] = []
+
+        self._session: aiohttp.ClientSession | None = None
+        self._heartbeat_task: asyncio.Task | None = None
+        self._running = False
+        self._start_time: float = 0
+        self._last_state_sync: float | None = None
+        self._state_size_bytes: int | None = None
+
+        # Track disconnect time for grace period
+        self._disconnect_time: float | None = None
+
+    @property
+    def status(self) -> ConnectivityStatus:
+        """Current connectivity status."""
+        return self._status
+
+    @property
+    def metrics(self) -> ConnectivityMetrics:
+        """Connectivity metrics."""
+        return self._metrics
+
+    def is_connected(self) -> bool:
+        """Check if connected to Orchestrator."""
+        return self._status == ConnectivityStatus.CONNECTED
+
+    def is_disconnected(self) -> bool:
+        """Check if disconnected from Orchestrator."""
+        return self._status == ConnectivityStatus.DISCONNECTED
+
+    def set_processor_info(
+        self,
+        processor_id: str,
+        node_id: str,
+        role: ProcessorRole | None = None,
+        fencing_token: FencingToken | None = None,
+    ) -> None:
+        """Set processor information for heartbeat payload.
+
+        Args:
+            processor_id: Processor identifier
+            node_id: Edge node identifier
+            role: Current processor role
+            fencing_token: Current fencing token
+        """
+        self._processor_id = processor_id
+        self._node_id = node_id
+        if role:
+            self._role = role
+        if fencing_token:
+            self._fencing_token = fencing_token
+
+    def update_role(self, role: ProcessorRole) -> None:
+        """Update current processor role."""
+        self._role = role
+
+    def update_fencing_token(self, token: FencingToken | None) -> None:
+        """Update current fencing token."""
+        self._fencing_token = token
+
+    def record_state_sync(self, size_bytes: int | None = None) -> None:
+        """Record a successful state sync."""
+        self._last_state_sync = time.time()
+        if size_bytes is not None:
+            self._state_size_bytes = size_bytes
+
+    def add_connectivity_callback(self, callback: ConnectivityChangeCallback) -> None:
+        """Register callback for connectivity changes.
+
+        Args:
+            callback: Async function called on status changes
+        """
+        self._callbacks.append(callback)
+
+    def remove_connectivity_callback(self, callback: ConnectivityChangeCallback) -> None:
+        """Remove connectivity callback."""
+        if callback in self._callbacks:
+            self._callbacks.remove(callback)
+
+    async def start(self) -> None:
+        """Start heartbeat service."""
+        if self._running:
+            return
+
+        self._running = True
+        self._start_time = time.time()
+
+        # Create HTTP session
+        timeout = aiohttp.ClientTimeout(total=self._config.timeout_sec)
+        self._session = aiohttp.ClientSession(timeout=timeout)
+
+        # Start heartbeat loop
+        self._heartbeat_task = asyncio.create_task(self._heartbeat_loop())
+
+        logger.info(
+            f"Started heartbeat service: interval={self._config.interval_sec}s, "
+            f"orchestrator={self._config.orchestrator_url}"
+        )
+
+    async def stop(self) -> None:
+        """Stop heartbeat service."""
+        self._running = False
+
+        if self._heartbeat_task:
+            self._heartbeat_task.cancel()
+            try:
+                await self._heartbeat_task
+            except asyncio.CancelledError:
+                pass
+            self._heartbeat_task = None
+
+        if self._session:
+            await self._session.close()
+            self._session = None
+
+        logger.info("Stopped heartbeat service")
+
+    async def send_heartbeat(self) -> HeartbeatResponse | None:
+        """Send a single heartbeat to the Orchestrator.
+
+        Returns:
+            HeartbeatResponse if successful, None if failed
+        """
+        if not self._session or not self._processor_id or not self._node_id:
+            return None
+
+        payload = HeartbeatPayload(
+            processor_id=self._processor_id,
+            node_id=self._node_id,
+            role=self._role.value,
+            epoch=self._fencing_token.epoch if self._fencing_token else None,
+            last_state_sync=self._last_state_sync,
+            state_size_bytes=self._state_size_bytes,
+            uptime_sec=time.time() - self._start_time,
+        )
+
+        url = f"{self._config.orchestrator_url}/api/v1/edge/heartbeat"
+
+        start_time = time.monotonic()
+        try:
+            async with self._session.post(
+                url,
+                json=payload.to_dict(),
+                headers={"Content-Type": "application/json"},
+            ) as resp:
+                latency_ms = (time.monotonic() - start_time) * 1000
+
+                if resp.status == 200:
+                    data = await resp.json()
+                    response = HeartbeatResponse.from_dict(data)
+
+                    self._metrics.record_success(latency_ms)
+                    await self._update_status(latency_ms)
+
+                    logger.debug(
+                        f"Heartbeat succeeded: latency={latency_ms:.1f}ms, "
+                        f"directive={response.directive}"
+                    )
+
+                    return response
+                else:
+                    logger.warning(
+                        f"Heartbeat returned status {resp.status}"
+                    )
+                    self._metrics.record_failure()
+                    await self._update_status(None)
+                    return None
+
+        except asyncio.TimeoutError:
+            logger.warning(
+                f"Heartbeat timeout after {self._config.timeout_sec}s"
+            )
+            self._metrics.record_failure()
+            await self._update_status(None)
+            return None
+
+        except aiohttp.ClientError as e:
+            logger.warning(f"Heartbeat failed: {e}")
+            self._metrics.record_failure()
+            await self._update_status(None)
+            return None
+
+    async def _heartbeat_loop(self) -> None:
+        """Background heartbeat loop."""
+        while self._running:
+            try:
+                response = await self.send_heartbeat()
+
+                # Handle directives from Orchestrator
+                if response and response.directive:
+                    await self._handle_directive(response.directive)
+
+                await asyncio.sleep(self._config.interval_sec)
+
+            except asyncio.CancelledError:
+                break
+            except Exception as e:
+                logger.error(f"Heartbeat loop error: {e}")
+                await asyncio.sleep(self._config.interval_sec)
+
+    async def _update_status(self, latency_ms: float | None) -> None:
+        """Update connectivity status based on metrics."""
+        old_status = self._status
+
+        if latency_ms is not None:
+            # Successful heartbeat
+            self._disconnect_time = None
+
+            if latency_ms > self._config.latency_threshold_ms:
+                # High latency - degraded
+                new_status = ConnectivityStatus.DEGRADED
+            elif (
+                self._status == ConnectivityStatus.DEGRADED
+                and self._metrics.consecutive_successes < self._config.recovery_threshold
+            ):
+                # Still recovering from degraded
+                new_status = ConnectivityStatus.DEGRADED
+            else:
+                new_status = ConnectivityStatus.CONNECTED
+        else:
+            # Failed heartbeat
+            if self._metrics.consecutive_failures >= self._config.missed_threshold:
+                new_status = ConnectivityStatus.DISCONNECTED
+
+                # Track disconnect time for grace period
+                if self._disconnect_time is None:
+                    self._disconnect_time = time.time()
+            elif self._metrics.consecutive_failures > 0:
+                new_status = ConnectivityStatus.DEGRADED
+            else:
+                new_status = self._status
+
+        if new_status != old_status:
+            self._status = new_status
+            logger.info(f"Connectivity status changed: {old_status.value} -> {new_status.value}")
+
+            # Notify callbacks
+            for callback in self._callbacks:
+                try:
+                    await callback(old_status, new_status)
+                except Exception as e:
+                    logger.error(f"Connectivity callback failed: {e}")
+
+    async def _handle_directive(self, directive: str) -> None:
+        """Handle directive from Orchestrator.
+
+        Args:
+            directive: Directive string ("continue", "demote", "promote", "shutdown")
+        """
+        logger.info(f"Received directive from Orchestrator: {directive}")
+
+        if directive == "demote":
+            # Orchestrator is requesting we demote to STANDBY
+            # This could happen if another instance is being promoted
+            logger.warning("Orchestrator requested demotion")
+            # Note: actual demotion should be handled by caller via callback
+
+        elif directive == "promote":
+            # Orchestrator is requesting we promote to PRIMARY
+            logger.info("Orchestrator requested promotion")
+
+        elif directive == "shutdown":
+            # Orchestrator is requesting graceful shutdown
+            logger.warning("Orchestrator requested shutdown")
+
+        # "continue" means keep doing what you're doing
+
+    def get_status_dict(self) -> dict[str, Any]:
+        """Get current status as dictionary."""
+        return {
+            "processor_id": self._processor_id,
+            "node_id": self._node_id,
+            "role": self._role.value,
+            "connectivity": self._status.value,
+            "is_connected": self.is_connected(),
+            "fencing_epoch": self._fencing_token.epoch if self._fencing_token else None,
+            "metrics": {
+                "total_heartbeats": self._metrics.total_heartbeats,
+                "success_rate": self._metrics.get_success_rate(),
+                "consecutive_failures": self._metrics.consecutive_failures,
+                "avg_latency_ms": self._metrics.avg_latency_ms,
+                "last_latency_ms": self._metrics.last_latency_ms,
+            },
+            "last_state_sync": self._last_state_sync,
+            "uptime_sec": time.time() - self._start_time if self._start_time else 0,
+        }
+
+    def should_demote(self) -> bool:
+        """Check if should demote due to prolonged disconnect.
+
+        Returns:
+            True if disconnect grace period exceeded
+        """
+        if not self._config.auto_demote_on_disconnect:
+            return False
+
+        if self._status != ConnectivityStatus.DISCONNECTED:
+            return False
+
+        if self._disconnect_time is None:
+            return False
+
+        elapsed = time.time() - self._disconnect_time
+        return elapsed >= self._config.demote_grace_period_sec
+
+
+class EdgeHealthReporter:
+    """Reports edge-specific health information to the HealthServer.
+
+    Integrates with the existing HealthServer to provide edge-specific
+    status information in health check responses.
+
+    Usage:
+        reporter = EdgeHealthReporter(
+            heartbeat_manager=heartbeat_mgr,
+            role_manager=role_mgr,
+        )
+
+        # Add to health server
+        health_server.add_health_component("edge", reporter.get_health_status)
+    """
+
+    def __init__(
+        self,
+        heartbeat_manager: HeartbeatManager | None = None,
+        role_manager: Any = None,  # RoleManager, avoiding circular import
+    ):
+        """Initialize health reporter.
+
+        Args:
+            heartbeat_manager: Heartbeat manager instance
+            role_manager: Role manager instance
+        """
+        self._heartbeat = heartbeat_manager
+        self._role_manager = role_manager
+
+    def get_health_status(self) -> dict[str, Any]:
+        """Get edge health status for health endpoint.
+
+        Returns:
+            Dictionary with edge health information
+        """
+        status: dict[str, Any] = {
+            "is_edge": True,
+        }
+
+        if self._heartbeat:
+            status["connectivity"] = self._heartbeat.status.value
+            status["is_connected"] = self._heartbeat.is_connected()
+            metrics = self._heartbeat.metrics
+            status["heartbeat"] = {
+                "success_rate": metrics.get_success_rate(),
+                "avg_latency_ms": metrics.avg_latency_ms,
+                "consecutive_failures": metrics.consecutive_failures,
+            }
+
+        if self._role_manager:
+            status["role"] = self._role_manager.role.value
+            status["can_process"] = self._role_manager.can_process()
+            if self._role_manager.fencing_token:
+                status["fencing_epoch"] = self._role_manager.fencing_token.epoch
+
+        return status
+
+    async def check_health(self) -> tuple[str, str]:
+        """Async health check for HealthServer integration.
+
+        Returns:
+            Tuple of (status, message) where status is "healthy", "degraded", or "unhealthy"
+        """
+        if not self._heartbeat:
+            return "healthy", "Edge health reporter not configured"
+
+        connectivity = self._heartbeat.status
+
+        if connectivity == ConnectivityStatus.CONNECTED:
+            return "healthy", "Connected to Orchestrator"
+        elif connectivity == ConnectivityStatus.DEGRADED:
+            return "degraded", "Degraded connectivity to Orchestrator"
+        elif connectivity == ConnectivityStatus.DISCONNECTED:
+            return "unhealthy", "Disconnected from Orchestrator"
+        else:
+            return "degraded", "Connectivity status unknown"