Trixy7/trixy_core/network/pool/health.py

# -*- coding: utf-8 -*-
"""
Health-Check für Verbindungen.

Bietet Ping-basierte und protokollspezifische Health-Checks.
"""

import asyncio
import time
from dataclasses import dataclass
from enum import IntEnum, auto
from typing import Callable, Awaitable, Any

from trixy_core.network.pool.connection import Connection


class HealthStatus(IntEnum):
    """Gesundheitszustand einer Verbindung."""
    HEALTHY = auto()     # Verbindung ist gesund
    DEGRADED = auto()    # Verbindung funktioniert, aber langsam
    UNHEALTHY = auto()   # Verbindung ist nicht gesund
    UNKNOWN = auto()     # Status unbekannt


@dataclass
class HealthCheckConfig:
    """
    Konfiguration für Health-Checks.

    Attributes:
        interval: Prüfungsintervall in Sekunden
        timeout: Timeout für einzelne Prüfung
        healthy_threshold: Anzahl erfolgreicher Prüfungen für HEALTHY
        unhealthy_threshold: Anzahl fehlgeschlagener Prüfungen für UNHEALTHY
        degraded_latency: Latenz ab der DEGRADED gemeldet wird
    """
    interval: float = 30.0
    timeout: float = 5.0
    healthy_threshold: int = 2
    unhealthy_threshold: int = 3
    degraded_latency: float = 1.0  # 1 Sekunde


@dataclass
class HealthCheckResult:
    """Ergebnis eines Health-Checks."""
    status: HealthStatus
    latency: float
    message: str = ""
    timestamp: float = 0.0

    def __post_init__(self) -> None:
        if self.timestamp == 0.0:
            self.timestamp = time.time()


# Type alias für Health-Check-Funktionen
HealthCheckFunc = Callable[[Connection], Awaitable[HealthCheckResult]]


class HealthChecker:
    """
    Health-Checker für Verbindungen.

    Führt periodische Gesundheitsprüfungen durch und verfolgt den Status.

    Example:
        checker = HealthChecker(HealthCheckConfig(
            interval=30.0,
            timeout=5.0
        ))

        # Custom-Check registrieren
        async def ping_check(conn: Connection) -> HealthCheckResult:
            start = time.monotonic()
            await conn.write(b"PING\\n")
            response = await conn.readline()
            latency = time.monotonic() - start

            if response == b"PONG\\n":
                return HealthCheckResult(HealthStatus.HEALTHY, latency)
            return HealthCheckResult(HealthStatus.UNHEALTHY, latency, "Invalid response")

        checker.set_check_function(ping_check)

        # Einzelprüfung
        result = await checker.check(connection)
        print(f"Status: {result.status.name}, Latenz: {result.latency*1000:.2f}ms")
    """

    def __init__(self, config: HealthCheckConfig | None = None) -> None:
        """
        Initialisiert den Health-Checker.

        Args:
            config: Health-Check-Konfiguration
        """
        self._config = config or HealthCheckConfig()
        self._check_func: HealthCheckFunc | None = None
        self._history: dict[str, list[HealthCheckResult]] = {}
        self._status: dict[str, HealthStatus] = {}

    def set_check_function(self, func: HealthCheckFunc) -> None:
        """
        Setzt die Health-Check-Funktion.

        Args:
            func: Async-Funktion die einen HealthCheckResult zurückgibt
        """
        self._check_func = func

    async def check(self, conn: Connection) -> HealthCheckResult:
        """
        Führt einen Health-Check durch.

        Args:
            conn: Die zu prüfende Verbindung

        Returns:
            HealthCheckResult
        """
        if self._check_func:
            try:
                result = await asyncio.wait_for(
                    self._check_func(conn),
                    timeout=self._config.timeout
                )
            except asyncio.TimeoutError:
                result = HealthCheckResult(
                    status=HealthStatus.UNHEALTHY,
                    latency=self._config.timeout,
                    message="Health-Check Timeout"
                )
            except Exception as e:
                result = HealthCheckResult(
                    status=HealthStatus.UNHEALTHY,
                    latency=0.0,
                    message=str(e)
                )
        else:
            # Standard-Check: Socket-Prüfung
            result = await self._default_check(conn)

        # Historie aktualisieren
        self._update_history(conn.id, result)

        return result

    async def _default_check(self, conn: Connection) -> HealthCheckResult:
        """Standard Health-Check basierend auf Socket-Status."""
        start = time.monotonic()

        try:
            # Prüfe ob Verbindung noch offen
            if conn.at_eof():
                return HealthCheckResult(
                    status=HealthStatus.UNHEALTHY,
                    latency=0.0,
                    message="Connection EOF"
                )

            # Prüfe Socket-Status
            if not conn.is_usable:
                return HealthCheckResult(
                    status=HealthStatus.UNHEALTHY,
                    latency=0.0,
                    message=f"Connection state: {conn.state.name}"
                )

            latency = time.monotonic() - start

            # Latenz-basierte Bewertung
            if latency > self._config.degraded_latency:
                return HealthCheckResult(
                    status=HealthStatus.DEGRADED,
                    latency=latency,
                    message="High latency"
                )

            return HealthCheckResult(
                status=HealthStatus.HEALTHY,
                latency=latency
            )

        except Exception as e:
            return HealthCheckResult(
                status=HealthStatus.UNHEALTHY,
                latency=time.monotonic() - start,
                message=str(e)
            )

    def _update_history(self, conn_id: str, result: HealthCheckResult) -> None:
        """Aktualisiert die Check-Historie."""
        if conn_id not in self._history:
            self._history[conn_id] = []

        history = self._history[conn_id]
        history.append(result)

        # Nur letzte N Ergebnisse behalten
        max_history = max(
            self._config.healthy_threshold,
            self._config.unhealthy_threshold
        ) * 2
        if len(history) > max_history:
            history[:] = history[-max_history:]

        # Status aktualisieren
        self._update_status(conn_id)

    def _update_status(self, conn_id: str) -> None:
        """Aktualisiert den aggregierten Status."""
        history = self._history.get(conn_id, [])

        if not history:
            self._status[conn_id] = HealthStatus.UNKNOWN
            return

        # Zähle letzte Ergebnisse
        recent = history[-max(
            self._config.healthy_threshold,
            self._config.unhealthy_threshold
        ):]

        healthy_count = sum(
            1 for r in recent if r.status == HealthStatus.HEALTHY
        )
        unhealthy_count = sum(
            1 for r in recent if r.status == HealthStatus.UNHEALTHY
        )

        if unhealthy_count >= self._config.unhealthy_threshold:
            self._status[conn_id] = HealthStatus.UNHEALTHY
        elif healthy_count >= self._config.healthy_threshold:
            self._status[conn_id] = HealthStatus.HEALTHY
        else:
            # Prüfe auf DEGRADED
            degraded_count = sum(
                1 for r in recent if r.status == HealthStatus.DEGRADED
            )
            if degraded_count > 0:
                self._status[conn_id] = HealthStatus.DEGRADED
            else:
                self._status[conn_id] = HealthStatus.UNKNOWN

    def get_status(self, conn_id: str) -> HealthStatus:
        """
        Gibt den aggregierten Status einer Verbindung zurück.

        Args:
            conn_id: Verbindungs-ID

        Returns:
            Aggregierter HealthStatus
        """
        return self._status.get(conn_id, HealthStatus.UNKNOWN)

    def get_history(self, conn_id: str) -> list[HealthCheckResult]:
        """
        Gibt die Check-Historie einer Verbindung zurück.

        Args:
            conn_id: Verbindungs-ID

        Returns:
            Liste der letzten Ergebnisse
        """
        return self._history.get(conn_id, []).copy()

    def get_average_latency(self, conn_id: str, samples: int = 10) -> float:
        """
        Berechnet die durchschnittliche Latenz.

        Args:
            conn_id: Verbindungs-ID
            samples: Anzahl zu berücksichtigender Samples

        Returns:
            Durchschnittliche Latenz in Sekunden
        """
        history = self._history.get(conn_id, [])
        if not history:
            return 0.0

        recent = history[-samples:]
        latencies = [r.latency for r in recent if r.latency > 0]

        if not latencies:
            return 0.0

        return sum(latencies) / len(latencies)

    def clear(self, conn_id: str | None = None) -> None:
        """
        Löscht die Historie.

        Args:
            conn_id: Optional spezifische Verbindung, sonst alle
        """
        if conn_id:
            self._history.pop(conn_id, None)
            self._status.pop(conn_id, None)
        else:
            self._history.clear()
            self._status.clear()

    def get_statistics(self) -> dict[str, Any]:
        """Gibt Statistiken zurück."""
        total = len(self._status)
        healthy = sum(1 for s in self._status.values() if s == HealthStatus.HEALTHY)
        degraded = sum(1 for s in self._status.values() if s == HealthStatus.DEGRADED)
        unhealthy = sum(1 for s in self._status.values() if s == HealthStatus.UNHEALTHY)

        return {
            "total_connections": total,
            "healthy": healthy,
            "degraded": degraded,
            "unhealthy": unhealthy,
            "unknown": total - healthy - degraded - unhealthy
        }


class PeriodicHealthChecker:
    """
    Periodischer Health-Checker für mehrere Verbindungen.

    Führt automatisch Health-Checks in regelmäßigen Intervallen durch.
    """

    def __init__(
        self,
        checker: HealthChecker,
        on_unhealthy: Callable[[Connection], Awaitable[None]] | None = None
    ) -> None:
        """
        Initialisiert den periodischen Checker.

        Args:
            checker: Der zu verwendende HealthChecker
            on_unhealthy: Callback bei unhealthy Verbindung
        """
        self._checker = checker
        self._on_unhealthy = on_unhealthy
        self._connections: list[Connection] = []
        self._task: asyncio.Task | None = None
        self._running = False

    def add_connection(self, conn: Connection) -> None:
        """Fügt eine Verbindung zur Überwachung hinzu."""
        if conn not in self._connections:
            self._connections.append(conn)

    def remove_connection(self, conn: Connection) -> None:
        """Entfernt eine Verbindung von der Überwachung."""
        if conn in self._connections:
            self._connections.remove(conn)
            self._checker.clear(conn.id)

    async def start(self) -> None:
        """Startet die periodischen Checks."""
        if self._running:
            return

        self._running = True
        self._task = asyncio.create_task(self._check_loop())

    async def stop(self) -> None:
        """Stoppt die periodischen Checks."""
        self._running = False

        if self._task:
            self._task.cancel()
            try:
                await self._task
            except asyncio.CancelledError:
                pass
            self._task = None

    async def _check_loop(self) -> None:
        """Hauptschleife für periodische Checks."""
        while self._running:
            try:
                await asyncio.sleep(self._checker._config.interval)

                for conn in list(self._connections):
                    try:
                        result = await self._checker.check(conn)

                        if result.status == HealthStatus.UNHEALTHY:
                            if self._on_unhealthy:
                                await self._on_unhealthy(conn)
                    except Exception:
                        pass

            except asyncio.CancelledError:
                break
            except Exception:
                pass