Trixy7/trixy_core/network/keepalive.py

# -*- coding: utf-8 -*-
"""
KeepAlive-Manager für Netzwerkverbindungen.

Bietet automatisches Ping/Pong-Scheduling und
Verbindungsüberwachung für langlebige Connections.
"""

from __future__ import annotations

import asyncio
import logging
from dataclasses import dataclass, field
from datetime import datetime, timedelta
from enum import Enum, auto
from typing import Any, Awaitable, Callable, Protocol
from weakref import WeakValueDictionary


class ConnectionState(Enum):
    """Verbindungszustand."""
    CONNECTED = auto()
    DISCONNECTED = auto()
    STALE = auto()           # Keine Antwort auf Ping
    RECONNECTING = auto()


class Pingable(Protocol):
    """Protokoll für pingbare Verbindungen."""

    async def send_ping(self) -> None:
        """Sendet einen Ping."""
        ...

    async def handle_pong(self, data: Any = None) -> None:
        """Verarbeitet empfangenen Pong."""
        ...


@dataclass
class KeepAliveConfig:
    """Konfiguration für KeepAlive-Verhalten."""
    ping_interval: float = 30.0        # Sekunden zwischen Pings
    pong_timeout: float = 10.0         # Max. Wartezeit auf Pong
    max_missed_pongs: int = 3          # Verlorene Pongs vor Disconnect
    jitter: float = 0.1                # Variation im Intervall (0.0-1.0)
    initial_delay: float = 5.0         # Verzögerung vor erstem Ping
    adaptive: bool = True              # Intervall anpassen basierend auf RTT
    min_interval: float = 10.0         # Minimales Intervall bei adaptiv
    max_interval: float = 120.0        # Maximales Intervall bei adaptiv


@dataclass
class ConnectionStats:
    """Statistiken für eine Verbindung."""
    pings_sent: int = 0
    pongs_received: int = 0
    missed_pongs: int = 0
    last_ping_at: datetime | None = None
    last_pong_at: datetime | None = None
    last_rtt_ms: float = 0.0
    avg_rtt_ms: float = 0.0
    min_rtt_ms: float = float("inf")
    max_rtt_ms: float = 0.0
    state: ConnectionState = ConnectionState.CONNECTED
    state_changed_at: datetime = field(default_factory=datetime.now)
    reconnect_count: int = 0


class KeepAliveSession:
    """
    KeepAlive-Sitzung für eine einzelne Verbindung.

    Verwaltet Ping/Pong-Zyklen und überwacht den
    Verbindungszustand.
    """

    def __init__(
        self,
        connection_id: str,
        ping_func: Callable[[], Awaitable[None]],
        config: KeepAliveConfig | None = None,
        on_timeout: Callable[[str], Awaitable[None]] | None = None,
        on_disconnect: Callable[[str], Awaitable[None]] | None = None,
        logger: logging.Logger | None = None,
    ) -> None:
        """
        Initialisiert die KeepAlive-Sitzung.

        Args:
            connection_id: Eindeutige Verbindungs-ID
            ping_func: Async-Funktion zum Senden von Pings
            config: KeepAlive-Konfiguration
            on_timeout: Callback bei Timeout
            on_disconnect: Callback bei Disconnect
            logger: Logger-Instanz
        """
        self.connection_id = connection_id
        self._ping_func = ping_func
        self.config = config or KeepAliveConfig()
        self._on_timeout = on_timeout
        self._on_disconnect = on_disconnect
        self.logger = logger or logging.getLogger(__name__)

        self.stats = ConnectionStats()
        self._task: asyncio.Task | None = None
        self._waiting_for_pong = False
        self._ping_sent_time: datetime | None = None
        self._rtt_samples: list[float] = []
        self._current_interval = self.config.ping_interval

    @property
    def is_running(self) -> bool:
        """Prüft ob KeepAlive aktiv ist."""
        return self._task is not None and not self._task.done()

    @property
    def state(self) -> ConnectionState:
        """Aktueller Verbindungszustand."""
        return self.stats.state

    async def start(self) -> None:
        """Startet KeepAlive-Schleife."""
        if self.is_running:
            return

        self.stats.state = ConnectionState.CONNECTED
        self.stats.state_changed_at = datetime.now()
        self._task = asyncio.create_task(self._keepalive_loop())
        self.logger.debug(f"KeepAlive gestartet für {self.connection_id}")

    async def stop(self) -> None:
        """Stoppt KeepAlive-Schleife."""
        if self._task:
            self._task.cancel()
            try:
                await self._task
            except asyncio.CancelledError:
                pass
            self._task = None
        self.logger.debug(f"KeepAlive gestoppt für {self.connection_id}")

    def handle_pong(self, timestamp: datetime | None = None) -> None:
        """
        Verarbeitet empfangenen Pong.

        Args:
            timestamp: Zeitpunkt des Pong-Empfangs
        """
        if not self._waiting_for_pong:
            return

        self._waiting_for_pong = False
        self.stats.pongs_received += 1
        self.stats.last_pong_at = timestamp or datetime.now()

        # RTT berechnen
        if self._ping_sent_time:
            rtt_ms = (self.stats.last_pong_at - self._ping_sent_time).total_seconds() * 1000
            self.stats.last_rtt_ms = rtt_ms
            self.stats.min_rtt_ms = min(self.stats.min_rtt_ms, rtt_ms)
            self.stats.max_rtt_ms = max(self.stats.max_rtt_ms, rtt_ms)

            # Durchschnitt aktualisieren
            self._rtt_samples.append(rtt_ms)
            if len(self._rtt_samples) > 20:
                self._rtt_samples.pop(0)
            self.stats.avg_rtt_ms = sum(self._rtt_samples) / len(self._rtt_samples)

            # Adaptives Intervall
            if self.config.adaptive:
                self._update_interval()

        # Zustand aktualisieren
        if self.stats.state == ConnectionState.STALE:
            self.stats.state = ConnectionState.CONNECTED
            self.stats.state_changed_at = datetime.now()
            self.stats.missed_pongs = 0

    def _update_interval(self) -> None:
        """Passt Ping-Intervall basierend auf RTT an."""
        avg_rtt = self.stats.avg_rtt_ms

        if avg_rtt < 50:
            # Sehr schnelle Verbindung - längeres Intervall
            self._current_interval = min(
                self._current_interval * 1.2,
                self.config.max_interval
            )
        elif avg_rtt > 500:
            # Langsame Verbindung - kürzeres Intervall
            self._current_interval = max(
                self._current_interval * 0.8,
                self.config.min_interval
            )

    async def _keepalive_loop(self) -> None:
        """Haupt-KeepAlive-Schleife."""
        # Initiale Verzögerung
        await asyncio.sleep(self.config.initial_delay)

        while True:
            try:
                # Jitter berechnen
                jitter_range = self._current_interval * self.config.jitter
                import random
                actual_interval = self._current_interval + random.uniform(
                    -jitter_range, jitter_range
                )

                await asyncio.sleep(actual_interval)

                # Ping senden
                self._waiting_for_pong = True
                self._ping_sent_time = datetime.now()
                self.stats.pings_sent += 1
                self.stats.last_ping_at = self._ping_sent_time

                try:
                    await self._ping_func()
                except Exception as e:
                    self.logger.warning(f"Ping fehlgeschlagen für {self.connection_id}: {e}")
                    continue

                # Auf Pong warten
                try:
                    await asyncio.wait_for(
                        self._wait_for_pong(),
                        timeout=self.config.pong_timeout
                    )
                except asyncio.TimeoutError:
                    await self._handle_timeout()

            except asyncio.CancelledError:
                break
            except Exception as e:
                self.logger.error(f"KeepAlive-Fehler für {self.connection_id}: {e}")
                await asyncio.sleep(5)  # Kurze Pause bei Fehler

    async def _wait_for_pong(self) -> None:
        """Wartet auf Pong-Empfang."""
        while self._waiting_for_pong:
            await asyncio.sleep(0.1)

    async def _handle_timeout(self) -> None:
        """Behandelt Pong-Timeout."""
        self._waiting_for_pong = False
        self.stats.missed_pongs += 1

        self.logger.warning(
            f"Pong-Timeout für {self.connection_id} "
            f"({self.stats.missed_pongs}/{self.config.max_missed_pongs})"
        )

        # Timeout-Callback
        if self._on_timeout:
            try:
                await self._on_timeout(self.connection_id)
            except Exception as e:
                self.logger.error(f"Timeout-Callback Fehler: {e}")

        # Zustand aktualisieren
        if self.stats.state != ConnectionState.STALE:
            self.stats.state = ConnectionState.STALE
            self.stats.state_changed_at = datetime.now()

        # Zu viele verlorene Pongs?
        if self.stats.missed_pongs >= self.config.max_missed_pongs:
            self.stats.state = ConnectionState.DISCONNECTED
            self.stats.state_changed_at = datetime.now()

            if self._on_disconnect:
                try:
                    await self._on_disconnect(self.connection_id)
                except Exception as e:
                    self.logger.error(f"Disconnect-Callback Fehler: {e}")


class KeepAliveManager:
    """
    Zentraler Manager für KeepAlive-Sitzungen.

    Verwaltet mehrere Verbindungen und bietet eine
    einheitliche Schnittstelle für KeepAlive-Operationen.
    """

    def __init__(
        self,
        default_config: KeepAliveConfig | None = None,
        logger: logging.Logger | None = None,
    ) -> None:
        """
        Initialisiert den KeepAlive-Manager.

        Args:
            default_config: Standard-Konfiguration für neue Sitzungen
            logger: Logger-Instanz
        """
        self.default_config = default_config or KeepAliveConfig()
        self.logger = logger or logging.getLogger(__name__)
        self._sessions: dict[str, KeepAliveSession] = {}
        self._on_timeout_handlers: list[Callable[[str], Awaitable[None]]] = []
        self._on_disconnect_handlers: list[Callable[[str], Awaitable[None]]] = []

    @property
    def session_count(self) -> int:
        """Anzahl aktiver Sitzungen."""
        return len(self._sessions)

    @property
    def active_count(self) -> int:
        """Anzahl laufender Sitzungen."""
        return sum(1 for s in self._sessions.values() if s.is_running)

    def on_timeout(self, handler: Callable[[str], Awaitable[None]]) -> None:
        """Registriert globalen Timeout-Handler."""
        self._on_timeout_handlers.append(handler)

    def on_disconnect(self, handler: Callable[[str], Awaitable[None]]) -> None:
        """Registriert globalen Disconnect-Handler."""
        self._on_disconnect_handlers.append(handler)

    async def _global_timeout_handler(self, connection_id: str) -> None:
        """Ruft alle Timeout-Handler auf."""
        for handler in self._on_timeout_handlers:
            try:
                await handler(connection_id)
            except Exception as e:
                self.logger.error(f"Timeout-Handler Fehler: {e}")

    async def _global_disconnect_handler(self, connection_id: str) -> None:
        """Ruft alle Disconnect-Handler auf."""
        for handler in self._on_disconnect_handlers:
            try:
                await handler(connection_id)
            except Exception as e:
                self.logger.error(f"Disconnect-Handler Fehler: {e}")

    async def register(
        self,
        connection_id: str,
        ping_func: Callable[[], Awaitable[None]],
        config: KeepAliveConfig | None = None,
        auto_start: bool = True,
    ) -> KeepAliveSession:
        """
        Registriert eine Verbindung für KeepAlive.

        Args:
            connection_id: Eindeutige Verbindungs-ID
            ping_func: Funktion zum Senden von Pings
            config: Optional: Benutzerdefinierte Konfiguration
            auto_start: Automatisch starten

        Returns:
            KeepAlive-Sitzung
        """
        if connection_id in self._sessions:
            await self._sessions[connection_id].stop()

        session = KeepAliveSession(
            connection_id=connection_id,
            ping_func=ping_func,
            config=config or self.default_config,
            on_timeout=self._global_timeout_handler,
            on_disconnect=self._global_disconnect_handler,
            logger=self.logger,
        )

        self._sessions[connection_id] = session

        if auto_start:
            await session.start()

        return session

    async def unregister(self, connection_id: str) -> bool:
        """
        Entfernt Verbindung aus KeepAlive.

        Args:
            connection_id: Verbindungs-ID

        Returns:
            True wenn entfernt
        """
        if connection_id not in self._sessions:
            return False

        session = self._sessions.pop(connection_id)
        await session.stop()
        return True

    def get_session(self, connection_id: str) -> KeepAliveSession | None:
        """Holt Sitzung nach ID."""
        return self._sessions.get(connection_id)

    def handle_pong(
        self,
        connection_id: str,
        timestamp: datetime | None = None,
    ) -> None:
        """
        Verarbeitet Pong für Verbindung.

        Args:
            connection_id: Verbindungs-ID
            timestamp: Pong-Zeitstempel
        """
        session = self._sessions.get(connection_id)
        if session:
            session.handle_pong(timestamp)

    def get_stats(self, connection_id: str) -> ConnectionStats | None:
        """Holt Statistiken für Verbindung."""
        session = self._sessions.get(connection_id)
        return session.stats if session else None

    def get_all_stats(self) -> dict[str, ConnectionStats]:
        """Holt Statistiken für alle Verbindungen."""
        return {cid: s.stats for cid, s in self._sessions.items()}

    def get_stale_connections(self) -> list[str]:
        """Findet Verbindungen im STALE-Zustand."""
        return [
            cid for cid, session in self._sessions.items()
            if session.state == ConnectionState.STALE
        ]

    def get_disconnected(self) -> list[str]:
        """Findet getrennte Verbindungen."""
        return [
            cid for cid, session in self._sessions.items()
            if session.state == ConnectionState.DISCONNECTED
        ]

    async def start_all(self) -> None:
        """Startet alle registrierten Sitzungen."""
        for session in self._sessions.values():
            if not session.is_running:
                await session.start()

    async def stop_all(self) -> None:
        """Stoppt alle Sitzungen."""
        for session in self._sessions.values():
            await session.stop()

    async def cleanup(self) -> int:
        """
        Entfernt getrennte Verbindungen.

        Returns:
            Anzahl entfernter Verbindungen
        """
        to_remove = self.get_disconnected()
        for cid in to_remove:
            await self.unregister(cid)
        return len(to_remove)

    def get_summary(self) -> dict[str, Any]:
        """Liefert Übersicht über alle Verbindungen."""
        states = {}
        for session in self._sessions.values():
            state_name = session.state.name
            states[state_name] = states.get(state_name, 0) + 1

        return {
            "total_sessions": self.session_count,
            "active_sessions": self.active_count,
            "states": states,
            "total_pings_sent": sum(
                s.stats.pings_sent for s in self._sessions.values()
            ),
            "total_pongs_received": sum(
                s.stats.pongs_received for s in self._sessions.values()
            ),
        }