diff --git a/hbd/client/main.py b/hbd/client/main.py index e7b5392..e883cb4 100644 --- a/hbd/client/main.py +++ b/hbd/client/main.py @@ -25,6 +25,7 @@ from .. import __version__ # Import plugin system from .plugin import PluginRegistry, PluginLoader, InfoPlugin, MonitorPlugin +from .connectivity import Connectivity, set_shared # Constants PORT = 50003 @@ -281,7 +282,16 @@ async def handle_update(conn: AsyncConnection, _msg: dict): # pyright: ignore[r stop() -async def heartbeat_sender(conn: AsyncConnection, interval: int): +def build_heartbeat_msg(conn, interval: int, connectivity: Optional[Connectivity] = None) -> dict: + """HTB payload; carries p4/p6 presence flags when the probe knows them.""" + msg = {"acks": conn.ackcount, "rtt": conn.rtts[-1], "interval": interval} + if connectivity is not None: + msg.update(connectivity.flags()) + return msg + + +async def heartbeat_sender(conn: AsyncConnection, interval: int, + connectivity: Optional[Connectivity] = None): """Send periodic heartbeats, retrying the connection if it is not open. IPv6 connections that fail to open before their first successful send are @@ -291,6 +301,7 @@ async def heartbeat_sender(conn: AsyncConnection, interval: int): Args: conn: Connection to send on interval: Heartbeat interval in seconds + connectivity: Optional probe whose presence flags are attached to every heartbeat """ logger = logging.getLogger("hbc.heartbeat") IPV6_EARLY_FAIL_LIMIT = 3 @@ -328,11 +339,7 @@ async def heartbeat_sender(conn: AsyncConnection, interval: int): continue try: - msg = { - "acks": conn.ackcount, - "rtt": conn.rtts[-1], - "interval": interval - } + msg = build_heartbeat_msg(conn, interval, connectivity) await conn.sendto(msg, "HTB") except asyncio.CancelledError: @@ -620,8 +627,17 @@ async def async_main(args, config): # Start async tasks # Heartbeat senders (one per connection) + # Per-protocol internet presence probe (shared with the network_monitor plugin) + net_cfg = dict((config.get("plugins") or {}).get("network_monitor") + or config.get("network_monitor") or {}) + connectivity = None + if net_cfg.get("probe_enabled", True): + connectivity = Connectivity.from_config(net_cfg) + set_shared(connectivity) + active_tasks.append(asyncio.create_task(connectivity.run(shutdown_event))) + for conn in connections: - task = asyncio.create_task(heartbeat_sender(conn, interval)) + task = asyncio.create_task(heartbeat_sender(conn, interval, connectivity)) active_tasks.append(task) # Plugin collector (uses all connections, but we'll use first one) diff --git a/hbd/client/plugins/network_monitor.py b/hbd/client/plugins/network_monitor.py index dcd243b..9ac654f 100644 --- a/hbd/client/plugins/network_monitor.py +++ b/hbd/client/plugins/network_monitor.py @@ -13,6 +13,7 @@ except ImportError: psutil = None from hbd.client.plugin import MonitorPlugin +from hbd.client.connectivity import get_shared logger = logging.getLogger(__name__) @@ -32,6 +33,10 @@ class NetworkMonitorPlugin(MonitorPlugin): interfaces: List of interfaces to monitor (default: all) include_connections: Include connection statistics (default: True) include_addresses: Include interface addresses (default: False) + probe_enabled: Run the per-protocol internet presence probe (default: True) + probe_v4: IPv4 address to ping (default: 1.1.1.1) + probe_v6: IPv6 address to ping (default: 2606:4700:4700::1111) + probe_interval: Seconds between ICMP probes while healthy (default: 60) """ name = "network_monitor" @@ -103,6 +108,7 @@ class NetworkMonitorPlugin(MonitorPlugin): - etc. - addresses: Interface address information (if include_addresses) - Dict keyed by interface name with address details + - connectivity: per-protocol presence {ipv4, ipv6} -> {present, addr, gateway, rtt} """ if psutil is None: logger.error("psutil not available") @@ -228,6 +234,11 @@ class NetworkMonitorPlugin(MonitorPlugin): except Exception as e: logger.warning(f"Could not collect interface stats: {e}") + + # Per-protocol internet presence, address and gateway (from the shared probe) + shared = get_shared() + if shared is not None: + metrics['connectivity'] = shared.as_metric() return metrics diff --git a/tests/test_client_heartbeat.py b/tests/test_client_heartbeat.py new file mode 100644 index 0000000..1083317 --- /dev/null +++ b/tests/test_client_heartbeat.py @@ -0,0 +1,28 @@ +from hbd.client.main import build_heartbeat_msg + + +class _Conn: + ackcount = 3 + rtts = [0, 12.5] + + +class _Conn2: + def __init__(self, flags): + self._flags = flags + + def flags(self): + return self._flags + + +def test_heartbeat_without_connectivity_is_unchanged(): + assert build_heartbeat_msg(_Conn(), 60) == {"acks": 3, "rtt": 12.5, "interval": 60} + + +def test_heartbeat_carries_presence_flags(): + msg = build_heartbeat_msg(_Conn(), 60, _Conn2({"p4": True, "p6": False})) + assert msg == {"acks": 3, "rtt": 12.5, "interval": 60, "p4": True, "p6": False} + + +def test_heartbeat_omits_unknown_flags(): + msg = build_heartbeat_msg(_Conn(), 60, _Conn2({})) + assert "p4" not in msg and "p6" not in msg diff --git a/tests/test_network_monitor_connectivity.py b/tests/test_network_monitor_connectivity.py new file mode 100644 index 0000000..3b0bcdd --- /dev/null +++ b/tests/test_network_monitor_connectivity.py @@ -0,0 +1,26 @@ +import asyncio + +from hbd.client import connectivity +from hbd.client.plugins.network_monitor import NetworkMonitorPlugin + + +class _FakeShared: + def as_metric(self): + return {"ipv4": {"present": True, "addr": "192.0.2.5", "gateway": "192.0.2.1", "rtt": 9.0}, + "ipv6": {"present": False, "addr": None, "gateway": None, "rtt": None}} + + +def test_collect_includes_connectivity_when_probe_running(): + connectivity.set_shared(_FakeShared()) + try: + data = asyncio.run(NetworkMonitorPlugin()._collect_metrics()) + finally: + connectivity.set_shared(None) + assert data["connectivity"]["ipv4"]["present"] is True + assert data["connectivity"]["ipv6"]["present"] is False + + +def test_collect_omits_connectivity_without_probe(): + connectivity.set_shared(None) + data = asyncio.run(NetworkMonitorPlugin()._collect_metrics()) + assert "connectivity" not in data