From 8bc237128757b0fce30ae9a01c6546c6ec2d052f Mon Sep 17 00:00:00 2001 From: s09x <185440786+s09x@users.noreply.github.com> Date: Fri, 29 May 2026 11:58:25 +0100 Subject: [PATCH] fix: wait for daemon health before reclaiming port --- .../hindsight_embed/daemon_embed_manager.py | 93 ++++++++++++++++--- hindsight-embed/tests/test_daemon_client.py | 92 +++++++++++++++++- 2 files changed, 169 insertions(+), 16 deletions(-) diff --git a/hindsight-embed/hindsight_embed/daemon_embed_manager.py b/hindsight-embed/hindsight_embed/daemon_embed_manager.py index 376bd1cb78..8057b9a606 100644 --- a/hindsight-embed/hindsight_embed/daemon_embed_manager.py +++ b/hindsight-embed/hindsight_embed/daemon_embed_manager.py @@ -6,6 +6,7 @@ """ import logging +import math import os import platform import re @@ -30,12 +31,42 @@ # Suppress noisy httpx logs logging.getLogger("httpx").setLevel(logging.WARNING) + +def _parse_float_env(name: str, default: float) -> float: + """Parse a float environment variable, falling back on invalid values.""" + try: + return float(os.getenv(name, str(default))) + except ValueError: + return default + + +def _safe_non_negative_float(value: float, fallback: float) -> float: + """Return a finite non-negative float, or fallback for invalid values.""" + return value if math.isfinite(value) and value >= 0 else fallback + + +def _safe_positive_float(value: float, fallback: float) -> float: + """Return a finite positive float, or fallback for invalid values.""" + return value if math.isfinite(value) and value > 0 else fallback + + # Constants # Allow CI/Windows to extend the startup budget — pg0-embedded's Windows wheel # unpacks and runs initdb on first boot, which takes noticeably longer on cold # runners than POSIX. DAEMON_STARTUP_TIMEOUT = int(os.getenv("HINDSIGHT_EMBED_DAEMON_STARTUP_TIMEOUT", "180")) DEFAULT_DAEMON_IDLE_TIMEOUT = 0 # 0 = disabled (no auto-exit) +# When another process is concurrently starting the daemon, the TCP port can be +# bound before /health returns 200. Give that warming daemon a short grace window +# before treating the listener as stale/foreign and attempting to reclaim it. +PORT_HEALTH_GRACE_TIMEOUT = _safe_non_negative_float( + _parse_float_env("HINDSIGHT_EMBED_PORT_HEALTH_GRACE_TIMEOUT", 30.0), + 30.0, +) +PORT_HEALTH_CHECK_INTERVAL = _safe_positive_float( + _parse_float_env("HINDSIGHT_EMBED_PORT_HEALTH_CHECK_INTERVAL", 0.5), + 0.5, +) def _detach_popen_kwargs(log_handle: IO[bytes]) -> dict: @@ -219,6 +250,40 @@ def _kill_process(pid: int) -> bool: return True # Already gone return False + @staticmethod + def _port_health_ok(port: int) -> bool: + """Return True when the listener on port responds like initialized Hindsight.""" + try: + with httpx.Client(timeout=2) as client: + response = client.get(f"http://127.0.0.1:{port}/health") + if response.status_code != 200: + return False + try: + health = response.json() + except Exception: + return False + return health.get("status") == "healthy" and health.get("database") == "connected" + except Exception: + return False + + def _wait_for_port_health(self, port: int, timeout: float | None = None) -> bool: + """Wait briefly for a just-bound daemon port to become healthy.""" + timeout = _safe_non_negative_float( + PORT_HEALTH_GRACE_TIMEOUT if timeout is None else timeout, + 0.0, + ) + interval = _safe_positive_float(PORT_HEALTH_CHECK_INTERVAL, 0.5) + deadline = time.monotonic() + timeout + while True: + if not self._is_port_in_use(port): + return False + if self._port_health_ok(port): + return True + remaining = deadline - time.monotonic() + if remaining <= 0: + return False + time.sleep(min(interval, remaining)) + def _clear_port(self, port: int) -> bool: """ Ensure the port is free before starting a daemon. @@ -229,29 +294,27 @@ def _clear_port(self, port: int) -> bool: The caller's "start" is effectively a no-op: the daemon is already up. Killing it would race concurrent starts (one process kills the other's freshly-started daemon, both rush to rebind the port). - * Port occupied but /health is unreachable / non-200 → treat as a stale - hindsight daemon (or foreign process) and attempt to reclaim by killing - the PID listening on the port. This preserves the original intent of - clearing stale daemons from version upgrades. + * Port occupied but /health is unreachable, non-200, or does not return + Hindsight's initialized health payload → treat as a stale daemon (or + foreign process) and attempt to reclaim by killing the PID listening + on the port. This preserves the original intent of clearing stale + daemons from version upgrades. * Kill failed, or non-hindsight process occupying the port → False. """ if not self._is_port_in_use(port): return True - # Port is occupied — check if it's a healthy hindsight daemon. - health_ok = False - try: - with httpx.Client(timeout=2) as client: - response = client.get(f"http://127.0.0.1:{port}/health") - health_ok = response.status_code == 200 - except Exception: - health_ok = False - - if health_ok: + # Port is occupied — check if it's a healthy Hindsight daemon. In + # concurrent startup races the socket can bind before /health returns + # 200, so wait briefly before deciding the listener is stale/foreign. + if self._wait_for_port_health(port): logger.debug(f"Port {port} already serving a healthy hindsight daemon; reusing it") return True - # Unhealthy — attempt to reclaim by killing the listener. + # Unhealthy after grace window — attempt to reclaim by killing the listener. + if not self._is_port_in_use(port): + return True + pid = self._find_pid_on_port(port) if pid is None: logger.warning(f"Port {port} is in use by another process") diff --git a/hindsight-embed/tests/test_daemon_client.py b/hindsight-embed/tests/test_daemon_client.py index b904559788..74df3e9c66 100644 --- a/hindsight-embed/tests/test_daemon_client.py +++ b/hindsight-embed/tests/test_daemon_client.py @@ -11,6 +11,7 @@ from hindsight_embed import daemon_client from hindsight_embed.daemon_embed_manager import DaemonEmbedManager + @pytest.fixture def config(): """Default config for tests.""" @@ -207,7 +208,10 @@ def test_port_occupied_by_healthy_hindsight_is_reused(self): mock_client = MagicMock() mock_client.__enter__ = Mock(return_value=mock_client) mock_client.__exit__ = Mock(return_value=False) - mock_client.get.return_value = Mock(status_code=200) + mock_client.get.return_value = Mock( + status_code=200, + json=Mock(return_value={"status": "healthy", "database": "connected"}), + ) mock_httpx_cls.return_value = mock_client assert manager._clear_port(9555) is True @@ -220,6 +224,7 @@ def test_port_occupied_by_non_hindsight_returns_false(self): with ( patch.object(DaemonEmbedManager, "_is_port_in_use", return_value=True), patch("httpx.Client") as mock_httpx_cls, + patch("hindsight_embed.daemon_embed_manager.PORT_HEALTH_GRACE_TIMEOUT", 0.0), ): mock_client = MagicMock() mock_client.__enter__ = Mock(return_value=mock_client) @@ -235,6 +240,7 @@ def test_port_occupied_health_non_200_returns_false(self): with ( patch.object(DaemonEmbedManager, "_is_port_in_use", return_value=True), patch("httpx.Client") as mock_httpx_cls, + patch("hindsight_embed.daemon_embed_manager.PORT_HEALTH_GRACE_TIMEOUT", 0.0), ): mock_client = MagicMock() mock_client.__enter__ = Mock(return_value=mock_client) @@ -251,6 +257,7 @@ def test_unhealthy_daemon_pid_not_found_returns_false(self): patch.object(DaemonEmbedManager, "_is_port_in_use", return_value=True), patch("httpx.Client") as mock_httpx_cls, patch.object(DaemonEmbedManager, "_find_pid_on_port", return_value=None), + patch("hindsight_embed.daemon_embed_manager.PORT_HEALTH_GRACE_TIMEOUT", 0.0), ): mock_client = MagicMock() mock_client.__enter__ = Mock(return_value=mock_client) @@ -268,6 +275,7 @@ def test_unhealthy_daemon_kill_fails_returns_false(self): patch("httpx.Client") as mock_httpx_cls, patch.object(DaemonEmbedManager, "_find_pid_on_port", return_value=12345), patch.object(DaemonEmbedManager, "_kill_process", return_value=False), + patch("hindsight_embed.daemon_embed_manager.PORT_HEALTH_GRACE_TIMEOUT", 0.0), ): mock_client = MagicMock() mock_client.__enter__ = Mock(return_value=mock_client) @@ -285,6 +293,75 @@ def test_unhealthy_daemon_kill_succeeds_returns_true(self): patch("httpx.Client") as mock_httpx_cls, patch.object(DaemonEmbedManager, "_find_pid_on_port", return_value=12345), patch.object(DaemonEmbedManager, "_kill_process", return_value=True), + patch("hindsight_embed.daemon_embed_manager.PORT_HEALTH_GRACE_TIMEOUT", 0.0), + ): + mock_client = MagicMock() + mock_client.__enter__ = Mock(return_value=mock_client) + mock_client.__exit__ = Mock(return_value=False) + mock_client.get.return_value = Mock(status_code=503) + mock_httpx_cls.return_value = mock_client + + assert manager._clear_port(9555) is True + + def test_port_occupied_by_warming_hindsight_is_reused(self): + """Port bound before /health is ready — wait briefly and reuse when healthy.""" + manager = DaemonEmbedManager() + with ( + patch.object(DaemonEmbedManager, "_is_port_in_use", return_value=True), + patch("httpx.Client") as mock_httpx_cls, + patch.object(DaemonEmbedManager, "_find_pid_on_port") as mock_find_pid, + patch.object(DaemonEmbedManager, "_kill_process") as mock_kill, + patch("hindsight_embed.daemon_embed_manager.PORT_HEALTH_GRACE_TIMEOUT", 1.0), + patch("hindsight_embed.daemon_embed_manager.PORT_HEALTH_CHECK_INTERVAL", 0.01), + ): + mock_client = MagicMock() + mock_client.__enter__ = Mock(return_value=mock_client) + mock_client.__exit__ = Mock(return_value=False) + mock_client.get.side_effect = [ + Mock(status_code=503), + Mock(status_code=503), + Mock( + status_code=200, + json=Mock(return_value={"status": "healthy", "database": "connected"}), + ), + ] + mock_httpx_cls.return_value = mock_client + + assert manager._clear_port(9555) is True + assert mock_client.get.call_count == 3 + mock_find_pid.assert_not_called() + mock_kill.assert_not_called() + + def test_port_occupied_by_foreign_health_200_returns_false(self): + """HTTP 200 alone is not enough to identify the listener as Hindsight.""" + manager = DaemonEmbedManager() + with ( + patch.object(DaemonEmbedManager, "_is_port_in_use", return_value=True), + patch("httpx.Client") as mock_httpx_cls, + patch.object(DaemonEmbedManager, "_find_pid_on_port", return_value=None) as mock_find_pid, + patch("hindsight_embed.daemon_embed_manager.PORT_HEALTH_GRACE_TIMEOUT", 0.0), + ): + mock_client = MagicMock() + mock_client.__enter__ = Mock(return_value=mock_client) + mock_client.__exit__ = Mock(return_value=False) + mock_client.get.return_value = Mock( + status_code=200, + json=Mock(return_value={"status": "ok"}), + ) + mock_httpx_cls.return_value = mock_client + + assert manager._clear_port(9555) is False + mock_find_pid.assert_called_once_with(9555) + + def test_port_cleared_during_grace_returns_true(self): + """If a stale listener exits during the grace wait, the port is already clear.""" + manager = DaemonEmbedManager() + with ( + patch.object(DaemonEmbedManager, "_is_port_in_use", side_effect=[True, False, False]), + patch("httpx.Client") as mock_httpx_cls, + patch.object(DaemonEmbedManager, "_find_pid_on_port") as mock_find_pid, + patch("hindsight_embed.daemon_embed_manager.PORT_HEALTH_GRACE_TIMEOUT", 1.0), + patch("hindsight_embed.daemon_embed_manager.PORT_HEALTH_CHECK_INTERVAL", 0.01), ): mock_client = MagicMock() mock_client.__enter__ = Mock(return_value=mock_client) @@ -293,6 +370,19 @@ def test_unhealthy_daemon_kill_succeeds_returns_true(self): mock_httpx_cls.return_value = mock_client assert manager._clear_port(9555) is True + mock_find_pid.assert_not_called() + + def test_invalid_port_health_timeout_is_bounded(self): + """Invalid grace timeout values must not create an unbounded wait.""" + manager = DaemonEmbedManager() + with ( + patch.object(DaemonEmbedManager, "_is_port_in_use", return_value=True), + patch.object(DaemonEmbedManager, "_port_health_ok", return_value=False) as mock_health, + patch("hindsight_embed.daemon_embed_manager.time.sleep") as mock_sleep, + ): + assert manager._wait_for_port_health(9555, timeout=float("nan")) is False + mock_health.assert_called_once_with(9555) + mock_sleep.assert_not_called() class TestStartDaemonSerialization: