Bound network probe timeouts

macfleet/pool/network.py

@@ -7,6 +7,7 @@
 from __future__ import annotations
 
 import asyncio
+import math
 import socket
 import subprocess
 import time
@@ -15,6 +16,17 @@
 from typing import Optional
 
 
+def _validate_positive_finite(name: str, value: float) -> float:
+    if (
+        isinstance(value, bool)
+        or not isinstance(value, (int, float))
+        or not math.isfinite(float(value))
+        or value <= 0
+    ):
+        raise ValueError(f"{name} must be a positive finite number")
+    return float(value)
+
+
 class LinkType(Enum):
     """Network link types, ordered by typical bandwidth."""
     WIFI = "wifi"
@@ -248,19 +260,26 @@ def _parse_ifconfig() -> list[NetworkLink]:
 
 
 async def measure_latency(host: str, port: int, timeout: float = 5.0) -> float:
-    """Measure TCP round-trip latency to a peer in milliseconds."""
+    """Measure TCP connection latency to a peer in milliseconds.
+
+    ``timeout`` bounds the complete probe, including connection cleanup.
+    Returns -1 when the peer cannot complete the probe in time.
+    """
+    timeout = _validate_positive_finite("timeout", timeout)
+    writer = None
     try:
-        start = time.monotonic()
-        _, writer = await asyncio.wait_for(
-            asyncio.open_connection(host, port),
-            timeout=timeout,
-        )
-        elapsed = (time.monotonic() - start) * 1000.0
-        writer.close()
-        await writer.wait_closed()
-        return elapsed
+        async with asyncio.timeout(timeout):
+            start = time.monotonic()
+            _, writer = await asyncio.open_connection(host, port)
+            elapsed = (time.monotonic() - start) * 1000.0
+            writer.close()
+            await writer.wait_closed()
+            return elapsed
     except (OSError, asyncio.TimeoutError):
         return -1.0
+    finally:
+        if writer is not None:
+            writer.close()
 
 
 async def measure_bandwidth(
@@ -271,52 +290,51 @@ async def measure_bandwidth(
 ) -> float:
     """Measure TCP bandwidth to a peer in Mbps.
 
-    Sends a payload, waits for the receiver to close the connection
-    (signaling it has drained the bytes), and divides total elapsed
-    time by payload size. The wait_closed() bound is the closest the
-    stdlib affords to a real ACK — without it, writer.drain() only
-    confirms the kernel send buffer received the bytes, not the wire.
-    Requires a listening peer that accepts and discards data.
+    Sends a payload, then waits for the receiver to close the connection
+    after draining the bytes. ``timeout`` bounds the complete probe: connect,
+    transfer, receiver acknowledgement, and connection cleanup. Returns 0
+    when any phase fails or the receiver does not acknowledge completion.
+    Requires a listening peer that reads until EOF and then closes.
     """
+    payload_mb = _validate_positive_finite("payload_mb", payload_mb)
+    timeout = _validate_positive_finite("timeout", timeout)
     payload_bytes = int(payload_mb * 1024 * 1024)
+    if payload_bytes == 0:
+        raise ValueError("payload_mb must represent at least one byte")
     chunk_size = 65536
+    writer = None
 
     try:
-        reader, writer = await asyncio.wait_for(
-            asyncio.open_connection(host, port),
-            timeout=timeout,
-        )
-
-        data = b"\x00" * chunk_size
-        start = time.monotonic()
-        sent = 0
-        while sent < payload_bytes:
-            to_send = min(chunk_size, payload_bytes - sent)
-            writer.write(data[:to_send])
-            await writer.drain()
-            sent += to_send
-
-        # Half-close the writer side so the peer sees EOF, then wait
-        # for the connection to fully close — this is our acknowledgement
-        # signal that everything we sent has actually been read on the
-        # other end (or the receiver hung up). Either way, the elapsed
-        # time more closely reflects wire bandwidth than the local
-        # send-buffer drain time alone.
-        if writer.can_write_eof():
+        async with asyncio.timeout(timeout):
+            reader, writer = await asyncio.open_connection(host, port)
+
+            data = b"\x00" * chunk_size
+            start = time.monotonic()
+            sent = 0
+            while sent < payload_bytes:
+                to_send = min(chunk_size, payload_bytes - sent)
+                writer.write(data[:to_send])
+                await writer.drain()
+                sent += to_send
+
+            # Half-close so the peer sees EOF. A full close from the peer is
+            # the acknowledgement that its handler finished draining data.
+            if not writer.can_write_eof():
+                return 0.0
             writer.write_eof()
-        try:
-            await asyncio.wait_for(reader.read(), timeout=timeout)
-        except asyncio.TimeoutError:
-            pass
-        elapsed = time.monotonic() - start
-        writer.close()
-        await writer.wait_closed()
-
-        if elapsed > 0:
-            return (payload_bytes * 8) / (elapsed * 1_000_000)  # Mbps
-        return 0.0
+            await reader.read()
+            elapsed = time.monotonic() - start
+            writer.close()
+            await writer.wait_closed()
+
+            if elapsed > 0:
+                return (payload_bytes * 8) / (elapsed * 1_000_000)  # Mbps
+            return 0.0
     except (OSError, asyncio.TimeoutError):
         return 0.0
+    finally:
+        if writer is not None:
+            writer.close()
 
 
 def get_network_topology() -> NetworkTopology:

tests/test_pool/test_network.py

@@ -1,5 +1,10 @@
 """Tests for network detection and link scoring."""
 
+import asyncio
+import time
+
+import pytest
+
 from macfleet.pool import network
 from macfleet.pool.network import (
     LinkType,
@@ -96,3 +101,169 @@ def test_get_topology(self):
         topo = get_network_topology()
         assert isinstance(topo, NetworkTopology)
         assert topo.hostname != ""
+
+
+class TestNetworkProbes:
+    @pytest.mark.parametrize(
+        "timeout",
+        [0, -1, True, float("nan"), float("inf"), "soon"],
+    )
+    async def test_latency_rejects_invalid_timeout(self, timeout):
+        with pytest.raises(ValueError, match="timeout"):
+            await network.measure_latency("127.0.0.1", 9, timeout=timeout)
+
+    @pytest.mark.parametrize(
+        "timeout",
+        [0, -1, True, float("nan"), float("inf"), "soon"],
+    )
+    async def test_bandwidth_rejects_invalid_timeout(self, timeout):
+        with pytest.raises(ValueError, match="timeout"):
+            await network.measure_bandwidth("127.0.0.1", 9, timeout=timeout)
+
+    @pytest.mark.parametrize(
+        "payload_mb",
+        [0, -1, True, float("nan"), float("inf"), "large"],
+    )
+    async def test_bandwidth_rejects_invalid_payload_size(self, payload_mb):
+        with pytest.raises(ValueError, match="payload_mb"):
+            await network.measure_bandwidth("127.0.0.1", 9, payload_mb=payload_mb)
+
+    async def test_latency_timeout_bounds_connection_close(self, monkeypatch):
+        writer = _ProbeWriter(wait_closed_delay=None)
+
+        async def open_connection(host, port):
+            return _ProbeReader(), writer
+
+        monkeypatch.setattr(network.asyncio, "open_connection", open_connection)
+
+        started = time.monotonic()
+        result = await asyncio.wait_for(
+            network.measure_latency("peer", 5000, timeout=0.01),
+            timeout=0.2,
+        )
+
+        assert result == -1.0
+        assert time.monotonic() - started < 0.15
+        assert writer.closed
+
+    async def test_bandwidth_requires_receiver_acknowledgement(self, monkeypatch):
+        writer = _ProbeWriter()
+
+        async def open_connection(host, port):
+            return _ProbeReader(read_delay=None), writer
+
+        monkeypatch.setattr(network.asyncio, "open_connection", open_connection)
+
+        result = await network.measure_bandwidth(
+            "peer",
+            5000,
+            payload_mb=0.001,
+            timeout=0.01,
+        )
+
+        assert result == 0.0
+        assert writer.closed
+
+    async def test_bandwidth_timeout_is_end_to_end(self, monkeypatch):
+        writer = _ProbeWriter()
+
+        async def open_connection(host, port):
+            await asyncio.sleep(0.03)
+            return _ProbeReader(read_delay=0.03), writer
+
+        monkeypatch.setattr(network.asyncio, "open_connection", open_connection)
+
+        started = time.monotonic()
+        result = await network.measure_bandwidth(
+            "peer",
+            5000,
+            payload_mb=0.001,
+            timeout=0.05,
+        )
+
+        assert result == 0.0
+        assert time.monotonic() - started < 0.09
+        assert writer.closed
+
+    async def test_bandwidth_timeout_bounds_stalled_drain(self, monkeypatch):
+        writer = _ProbeWriter(drain_delay=None)
+
+        async def open_connection(host, port):
+            return _ProbeReader(), writer
+
+        monkeypatch.setattr(network.asyncio, "open_connection", open_connection)
+
+        result = await asyncio.wait_for(
+            network.measure_bandwidth(
+                "peer",
+                5000,
+                payload_mb=0.001,
+                timeout=0.01,
+            ),
+            timeout=0.2,
+        )
+
+        assert result == 0.0
+        assert writer.closed
+
+    async def test_bandwidth_reports_measurement_after_acknowledgement(self, monkeypatch):
+        writer = _ProbeWriter()
+
+        async def open_connection(host, port):
+            return _ProbeReader(), writer
+
+        monkeypatch.setattr(network.asyncio, "open_connection", open_connection)
+
+        result = await network.measure_bandwidth(
+            "peer",
+            5000,
+            payload_mb=0.001,
+            timeout=0.1,
+        )
+
+        assert result > 0.0
+        assert writer.bytes_written == int(0.001 * 1024 * 1024)
+        assert writer.eof_written
+        assert writer.closed
+
+
+class _ProbeReader:
+    def __init__(self, read_delay=0.0):
+        self.read_delay = read_delay
+
+    async def read(self):
+        if self.read_delay is None:
+            await asyncio.Future()
+        await asyncio.sleep(self.read_delay)
+        return b""
+
+
+class _ProbeWriter:
+    def __init__(self, drain_delay=0.0, wait_closed_delay=0.0):
+        self.drain_delay = drain_delay
+        self.wait_closed_delay = wait_closed_delay
+        self.bytes_written = 0
+        self.eof_written = False
+        self.closed = False
+
+    def write(self, data):
+        self.bytes_written += len(data)
+
+    async def drain(self):
+        if self.drain_delay is None:
+            await asyncio.Future()
+        await asyncio.sleep(self.drain_delay)
+
+    def can_write_eof(self):
+        return True
+
+    def write_eof(self):
+        self.eof_written = True
+
+    def close(self):
+        self.closed = True
+
+    async def wait_closed(self):
+        if self.wait_closed_delay is None:
+            await asyncio.Future()
+        await asyncio.sleep(self.wait_closed_delay)