heartbeat/hbd/client/connectivity.py

"""Per-protocol internet presence probe.

Tier 1 (every tier1_interval seconds, sends no packets): a UDP connect() to the probe
address reveals whether a route exists and which source address would be used.
Tier 2 (every probe_interval seconds, sooner after a failure or a change of address or
gateway): one ICMP echo to the probe address. Absence needs fail_limit consecutive
failures (or no route at all); presence needs one success.

The result is read by the heartbeat sender (p4/p6 flags on every HTB) and by the
network_monitor plugin (connectivity metric) via the shared instance.
"""

import asyncio
import logging
import socket
import sys
import time
from dataclasses import asdict, dataclass, replace
from typing import Any, Awaitable, Callable, Dict, Optional, Set, Tuple

from hbd.client import gateway as gateway_mod
from hbd.client.pingutil import PingUnavailable, run_ping

logger = logging.getLogger("hbc.connectivity")

FAMILIES = ("ipv4", "ipv6")
FLAG_KEYS = {"ipv4": "p4", "ipv6": "p6"}
DEFAULT_PROBES = {"ipv4": "1.1.1.1", "ipv6": "2606:4700:4700::1111"}


class ProbeUnavailable(Exception):
    """The address family cannot be probed on this host (e.g. no IPv6 support)."""


@dataclass
class FamilyState:
    present: Optional[bool] = None      # None = unknown (probe cannot run yet/at all)
    addr: Optional[str] = None
    gateway: Optional[str] = None
    rtt: Optional[float] = None


def local_address(family: str, target: str) -> Optional[str]:
    """Source address the kernel would use to reach target, or None if there is no route.

    A UDP connect() sends nothing. Raises ProbeUnavailable if the family is unsupported.
    """
    af = socket.AF_INET if family == "ipv4" else socket.AF_INET6
    try:
        sock = socket.socket(af, socket.SOCK_DGRAM)
    except OSError as e:
        raise ProbeUnavailable(str(e)) from e
    try:
        sock.connect((target, 9))
        return str(sock.getsockname()[0])
    except OSError:
        return None
    finally:
        sock.close()


async def ping_rtt(target: str, family: str) -> Optional[float]:
    """One ICMP echo; returns RTT in ms or None if no reply."""
    result = await run_ping(target, 1, 2, 4 if family == "ipv4" else 6)
    rtt = result["rtt_avg"]
    return None if rtt == float("inf") else rtt


class Connectivity:
    def __init__(
        self,
        probes: Optional[Dict[str, str]] = None,
        probe_interval: float = 60.0,
        retry_interval: float = 10.0,
        tier1_interval: float = 5.0,
        fail_limit: int = 2,
        local_address_fn: Callable[[str, str], Optional[str]] = local_address,
        ping_fn: Callable[[str, str], Awaitable[Optional[float]]] = ping_rtt,
        gateway_fn: Callable[[str], Awaitable[Optional[str]]] = gateway_mod.default_gateway,
        clock: Callable[[], float] = time.monotonic,
    ):
        self.probes = {**DEFAULT_PROBES, **(probes or {})}
        self.probe_interval = probe_interval
        self.retry_interval = retry_interval
        self.tier1_interval = tier1_interval
        self.fail_limit = fail_limit
        self._local_address = local_address_fn
        self._ping = ping_fn
        self._gateway_fn = gateway_fn
        self._clock = clock
        # Reading the gateway is a cheap /proc read on Linux but a subprocess elsewhere,
        # so elsewhere it is refreshed only when a tier-2 probe is due.
        self._cheap_gateway = sys.platform.startswith("linux")
        self.states: Dict[str, FamilyState] = {f: FamilyState() for f in FAMILIES}
        self._fails: Dict[str, int] = {f: 0 for f in FAMILIES}
        self._next_ping: Dict[str, float] = {f: 0.0 for f in FAMILIES}
        self._last_key: Dict[str, Optional[Tuple[Optional[str], Optional[str]]]] = {
            f: None for f in FAMILIES}
        self._ping_disabled: Set[str] = set()

    @classmethod
    def from_config(cls, cfg: Optional[Dict[str, Any]]) -> "Connectivity":
        cfg = cfg or {}
        probes = {}
        if cfg.get("probe_v4"):
            probes["ipv4"] = str(cfg["probe_v4"])
        if cfg.get("probe_v6"):
            probes["ipv6"] = str(cfg["probe_v6"])
        return cls(probes=probes, probe_interval=float(cfg.get("probe_interval", 60)))

    # ---- readers (cheap; safe to call on every heartbeat) --------------------------
    def snapshot(self) -> Dict[str, FamilyState]:
        return {f: replace(s) for f, s in self.states.items()}

    def flags(self) -> Dict[str, bool]:
        """HTB flags for families whose presence is known."""
        return {FLAG_KEYS[f]: s.present for f, s in self.states.items() if s.present is not None}

    def as_metric(self) -> Dict[str, Dict[str, Any]]:
        return {f: asdict(s) for f, s in self.states.items()}

    # ---- probing -------------------------------------------------------------------
    async def refresh(self) -> None:
        for fam in FAMILIES:
            try:
                await self._refresh_family(fam)
            except Exception:
                logger.exception("connectivity probe failed for %s", fam)

    async def run(self, stop: asyncio.Event) -> None:
        while not stop.is_set():
            await self.refresh()
            try:
                await asyncio.wait_for(stop.wait(), timeout=self.tier1_interval)
            except asyncio.TimeoutError:
                pass

    async def _gateway(self, fam: str) -> Optional[str]:
        try:
            return await self._gateway_fn(fam)
        except Exception:
            return None

    async def _refresh_family(self, fam: str) -> None:
        st = self.states[fam]
        try:
            addr = self._local_address(fam, self.probes[fam])
        except ProbeUnavailable:
            st.present = None
            return

        if addr is None:                       # tier 1: no route => absent immediately
            self._fails[fam] = 0
            self._next_ping[fam] = 0.0
            self._last_key[fam] = None
            st.present, st.addr, st.gateway, st.rtt = False, None, None, None
            return

        now = self._clock()
        gw = await self._gateway(fam) if self._cheap_gateway else st.gateway
        key = (addr, gw)
        due = key != self._last_key[fam] or now >= self._next_ping[fam]
        if due and not self._cheap_gateway:
            gw = await self._gateway(fam)
            key = (addr, gw)
        st.addr, st.gateway = addr, gw
        self._last_key[fam] = key
        if fam in self._ping_disabled:
            st.present = None          # cannot probe: unknown, never absent, while a route exists
        elif due:
            await self._tier2(fam, st, now)

    async def _tier2(self, fam: str, st: FamilyState, now: float) -> None:
        try:
            rtt = await self._ping(self.probes[fam], fam)
        except PingUnavailable:
            logger.warning("ping unavailable; %s presence stays unknown", fam)
            self._ping_disabled.add(fam)
            st.present = None
            return
        except (asyncio.TimeoutError, OSError):
            rtt = None
        if rtt is None:
            self._fails[fam] += 1
            self._next_ping[fam] = now + self.retry_interval
            if self._fails[fam] >= self.fail_limit:
                st.present, st.rtt = False, None
        else:
            self._fails[fam] = 0
            self._next_ping[fam] = now + self.probe_interval
            st.present, st.rtt = True, rtt


# ---- shared instance (heartbeat sender and network_monitor plugin) -----------------
_shared: Optional[Connectivity] = None


def set_shared(conn: Optional[Connectivity]) -> None:
    global _shared
    _shared = conn


def get_shared() -> Optional[Connectivity]:
    return _shared