diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 8b39eb8c49..0591d5b04d 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -3,6 +3,7 @@ import copy import hashlib import json import os +import signal import smtplib import threading import time @@ -2293,6 +2294,11 @@ class PrismaClient: 0.0, float(os.getenv("PRISMA_AUTH_RECONNECT_LOCK_TIMEOUT_SECONDS", "0.1")), ) + self._engine_pidfd: int = -1 + self._engine_pid: int = 0 + self._watching_engine: bool = False + self._engine_confirmed_dead: bool = False + self._sigchld_installed: bool = False verbose_proxy_logger.debug("Success - Created Prisma Client") def get_request_status( @@ -3560,31 +3566,383 @@ class PrismaClient: ) raise e + def _get_engine_pid(self) -> int: + """ + Get the PID of the Prisma query engine subprocess. + + Primary: access Prisma internals directly. + Fallback: scan /proc for prisma-query-engine (Linux only). + Returns 0 if not found or on non-Linux platforms. + """ + try: + engine = self.db._original_prisma._engine # type: ignore[attr-defined] + if engine is not None and engine.process is not None: + return engine.process.pid + except (AttributeError, TypeError): + pass + + try: + for entry in os.scandir("/proc"): + if not entry.name.isdigit(): + continue + try: + with open(f"/proc/{entry.name}/cmdline", "rb") as f: + cmdline = f.read().decode("utf-8", errors="replace") + if "prisma-query-engine" not in cmdline: + continue + with open(f"/proc/{entry.name}/stat", "r") as f: + stat = f.read() + last_paren = stat.rfind(")") + state = stat[last_paren + 2] if last_paren >= 0 else "?" + if state not in ("Z", "X", "x"): + return int(entry.name) + except ( + FileNotFoundError, + PermissionError, + ProcessLookupError, + IndexError, + ValueError, + ): + continue + except FileNotFoundError: + pass + return 0 + + def _is_engine_alive(self) -> bool: + """ + Check whether the tracked engine PID is still a live (non-zombie) process. + + Returns True if /proc is unavailable (non-Linux: assume alive). + Returns False if the process is gone or in zombie/dead state. + """ + if self._engine_pid <= 0: + return True # unknown — assume alive + try: + with open(f"/proc/{self._engine_pid}/stat", "r") as f: + stat = f.read() + last_paren = stat.rfind(")") + if last_paren < 0 or last_paren + 2 >= len(stat): + return True # parse failed — assume alive + state = stat[last_paren + 2] + return state not in ("Z", "X", "x") + except FileNotFoundError: + return False # process gone + except (PermissionError, ProcessLookupError, OSError): + return True # cannot read — assume alive + + @staticmethod + def _reap_all_zombies() -> set: + """Reap ALL zombie child processes via waitpid(-1, WNOHANG). + + Returns a set of reaped PIDs. As PID 1 in Docker (or any + process that spawns children), we must reap ALL terminated + children to prevent zombie accumulation. + """ + reaped: set = set() + while True: + try: + pid, _ = os.waitpid(-1, os.WNOHANG) + if pid == 0: + break + reaped.add(pid) + except ChildProcessError: + break + return reaped + + def _install_sigchld_handler(self) -> bool: + """Install SIGCHLD handler on the asyncio event loop. + + SIGCHLD is delivered by the kernel the instant any child process + exits, is killed, or enters zombie state. This gives + sub-millisecond detection with zero CPU overhead (no polling, + no file descriptors to manage). + + Returns True if installed, False on failure (non-Unix, no + running event loop, restricted environment, etc.). + """ + if self._sigchld_installed: + return True + try: + loop = asyncio.get_running_loop() + loop.add_signal_handler(signal.SIGCHLD, self._on_sigchld) + self._sigchld_installed = True + return True + except (RuntimeError, OSError, ValueError) as e: + verbose_proxy_logger.debug("Could not install SIGCHLD handler: %s", e) + return False + + def _remove_sigchld_handler(self) -> None: + """Remove SIGCHLD handler from the event loop.""" + if not self._sigchld_installed: + return + try: + loop = asyncio.get_running_loop() + loop.remove_signal_handler(signal.SIGCHLD) + except (RuntimeError, OSError, ValueError): + pass + self._sigchld_installed = False + + def _on_sigchld(self) -> None: + """SIGCHLD received -- reap all zombies and check if engine died. + + This fires the instant any child process exits or becomes a + zombie. We reap ALL children (fulfilling PID-1 responsibility) + then check if the tracked Prisma engine was among the dead. + """ + reaped = self._reap_all_zombies() + if not reaped: + return + if ( + self._engine_pid > 0 + and self._engine_pid in reaped + and not self._engine_confirmed_dead + ): + verbose_proxy_logger.error( + "prisma-query-engine PID %s reaped via SIGCHLD; triggering reconnect.", + self._engine_pid, + ) + self._engine_confirmed_dead = True + self._cleanup_engine_watcher() + asyncio.create_task( + self.attempt_db_reconnect( + reason="engine_process_death", + force=True, + ) + ) + elif reaped: + verbose_proxy_logger.debug("Reaped non-engine zombie PIDs: %s", reaped) + + def _try_pidfd_watch(self, pid: int) -> bool: + """ + Watch engine PID via pidfd_open + asyncio event loop reader. + + Returns True if pidfd watch was set up, False if unavailable or failed. + Broad OSError catch handles both ENOSYS and SECCOMP-blocked syscalls. + """ + if not hasattr(os, "pidfd_open"): + return False + fd = -1 + try: + fd = os.pidfd_open(pid, 0) # type: ignore[attr-defined] + asyncio.get_running_loop().add_reader(fd, self._on_pidfd_readable) + self._engine_pidfd = fd + return True + except OSError: + if fd >= 0: + os.close(fd) + return False + + def _on_pidfd_readable(self) -> None: + """pidfd became readable: engine process exited or became zombie. + + Sets _engine_confirmed_dead BEFORE cleanup so _run_reconnect_cycle + takes the heavy path (recreate Prisma client + re-arm watcher). + """ + if self._engine_confirmed_dead: + # Already handled by SIGCHLD -- just clean up pidfd resources. + if self._engine_pidfd >= 0: + try: + asyncio.get_running_loop().remove_reader(self._engine_pidfd) + except Exception: + pass + try: + os.close(self._engine_pidfd) + except OSError: + pass + self._engine_pidfd = -1 + return + dead_pid = self._engine_pid + verbose_proxy_logger.error( + "prisma-query-engine PID %s exited (pidfd event); triggering reconnect.", + dead_pid, + ) + self._engine_confirmed_dead = True + self._reap_all_zombies() + self._cleanup_engine_watcher() + asyncio.create_task( + self.attempt_db_reconnect( + reason="engine_process_death", + force=True, + ) + ) + + async def _poll_engine_proc(self) -> None: + """Last-resort fallback: poll /proc//stat every 1s. + + Only used when BOTH SIGCHLD handler and pidfd_open are unavailable + (e.g., non-Linux platforms or heavily restricted containers). + Prefer SIGCHLD (instant, zero-overhead) or pidfd (event-driven). + """ + while self._watching_engine and self._engine_pid > 0: + try: + with open(f"/proc/{self._engine_pid}/stat", "r") as f: + stat = f.read() + last_paren = stat.rfind(")") + if last_paren < 0 or last_paren + 2 >= len(stat): + state = "?" + else: + state = stat[last_paren + 2] + if state in ("Z", "X", "x"): + verbose_proxy_logger.error( + "prisma-query-engine PID %s in state '%s'; triggering reconnect.", + self._engine_pid, + state, + ) + self._engine_confirmed_dead = True + self._reap_all_zombies() + self._cleanup_engine_watcher() + await self.attempt_db_reconnect( + reason="engine_process_death", + force=True, + ) + return + except FileNotFoundError: + verbose_proxy_logger.error( + "prisma-query-engine PID %s disappeared; triggering reconnect.", + self._engine_pid, + ) + self._engine_confirmed_dead = True + self._reap_all_zombies() + self._cleanup_engine_watcher() + await self.attempt_db_reconnect( + reason="engine_process_death", + force=True, + ) + return + except (PermissionError, ProcessLookupError): + verbose_proxy_logger.debug( + "Cannot read /proc/%s/stat; stopping engine poll.", + self._engine_pid, + ) + self._cleanup_engine_watcher() + return + await asyncio.sleep(1) + + def _cleanup_engine_watcher(self) -> None: + """Clean up pidfd reader or stop /proc polling and reset engine tracking state.""" + self._watching_engine = False + if self._engine_pidfd >= 0: + try: + asyncio.get_running_loop().remove_reader(self._engine_pidfd) + except Exception: + pass + try: + os.close(self._engine_pidfd) + except OSError: + pass + self._engine_pidfd = -1 + self._engine_pid = 0 + + async def _start_engine_watcher(self) -> None: + """ + Start watching the Prisma query engine process for death. + + Detection priority (all instant, kernel-level when available): + 1. SIGCHLD signal handler -- instant notification from the kernel when + ANY child changes state. Zero CPU overhead. Works on all Unix/Linux. + Also reaps orphan zombies (PID 1 responsibility in Docker). + 2. pidfd_open (Linux 5.3+) -- event-driven fd for targeted engine + monitoring. Supplementary to SIGCHLD. + 3. /proc//stat polling (1s) -- last-resort fallback only when both + SIGCHLD and pidfd are unavailable (non-Linux or restricted envs). + """ + if self._watching_engine or self._engine_pidfd >= 0: + return + pid = self._get_engine_pid() + if pid == 0: + verbose_proxy_logger.debug("Could not find prisma-query-engine PID; engine death detection unavailable.") + return + self._engine_pid = pid + self._engine_confirmed_dead = False + verbose_proxy_logger.info("Found prisma-query-engine at PID %s.", pid) + # Primary: SIGCHLD -- instant kernel notification for ALL child state changes. + sigchld_ok = self._install_sigchld_handler() + # Supplementary: pidfd -- targeted event-driven watch on this specific PID. + pidfd_ok = self._try_pidfd_watch(pid) + if sigchld_ok and pidfd_ok: + verbose_proxy_logger.info( + "Watching engine PID %s via SIGCHLD + pidfd (dual event-driven).", pid, + ) + elif sigchld_ok: + verbose_proxy_logger.info( + "Watching engine PID %s via SIGCHLD (event-driven).", pid, + ) + elif pidfd_ok: + verbose_proxy_logger.info( + "Watching engine PID %s via pidfd (event-driven).", pid, + ) + else: + verbose_proxy_logger.info( + "Watching engine PID %s via /proc polling (1s fallback).", pid, + ) + self._watching_engine = True + asyncio.create_task(self._poll_engine_proc()) + + def _stop_engine_watcher(self) -> None: + """Stop watching the engine process and clean up all resources.""" + self._remove_sigchld_handler() + self._cleanup_engine_watcher() + self._engine_confirmed_dead = False + verbose_proxy_logger.debug("Stopped engine process watcher.") + async def _run_reconnect_cycle( self, timeout_seconds: Optional[float] = None ) -> None: """ - Run a reconnect cycle with direct db operations and a single overall timeout - budget to avoid long retries on hot paths (e.g. auth). + Run a reconnect cycle with a single overall timeout budget. + + Uses the _engine_confirmed_dead flag (set by SIGCHLD / pidfd / poll + handlers) to choose between heavy reconnect (engine dead -- recreate + Prisma client, re-arm watcher) and lightweight reconnect (network + blip -- disconnect, connect, SELECT 1). + + The flag-based approach fixes the race condition where + _cleanup_engine_watcher resets _engine_pid to 0 before this method + could check it, causing the heavy path to never execute. """ - async def _do_direct_reconnect() -> None: - try: - await self.db.disconnect() - except Exception as disconnect_err: - verbose_proxy_logger.debug( - "Prisma DB disconnect before reconnect failed (ignored): %s", - disconnect_err, - ) - - await self.db.connect() - await self.db.query_raw("SELECT 1") - effective_timeout = ( - timeout_seconds - if timeout_seconds is not None - else self._db_watchdog_reconnect_timeout_seconds + timeout_seconds if timeout_seconds is not None else self._db_watchdog_reconnect_timeout_seconds ) - await asyncio.wait_for(_do_direct_reconnect(), timeout=effective_timeout) + + engine_is_dead = self._engine_confirmed_dead or ( + self._engine_pid > 0 and not self._is_engine_alive() + ) + + if engine_is_dead: + dead_pid = self._engine_pid + verbose_proxy_logger.warning( + "prisma-query-engine PID %s is dead; performing heavy reconnect.", + dead_pid, + ) + self._reap_all_zombies() + self._cleanup_engine_watcher() + self._engine_confirmed_dead = False + + async def _do_heavy_reconnect() -> None: + db_url = os.getenv("DATABASE_URL", "") + if not db_url: + verbose_proxy_logger.error("DATABASE_URL not set; cannot recreate Prisma client.") + raise RuntimeError("DATABASE_URL not set") + await self.db.recreate_prisma_client(db_url) + await self._start_engine_watcher() + + await asyncio.wait_for(_do_heavy_reconnect(), timeout=effective_timeout) + else: + verbose_proxy_logger.debug("Performing lightweight Prisma DB reconnect (engine alive or unknown).") + + async def _do_direct_reconnect() -> None: + try: + await self.db.disconnect() + except Exception as disconnect_err: + verbose_proxy_logger.debug( + "Prisma DB disconnect before reconnect failed (ignored): %s", + disconnect_err, + ) + + await self.db.connect() + await self.db.query_raw("SELECT 1") + + await asyncio.wait_for(_do_direct_reconnect(), timeout=effective_timeout) async def attempt_db_reconnect( self, @@ -3703,9 +4061,10 @@ class PrismaClient: self._db_reconnect_lock.release() async def start_db_health_watchdog_task(self) -> None: - """ - Start a background task that probes DB health and attempts reconnect on failure. - """ + """Start background tasks that monitor DB health: + - A periodic SELECT 1 probe that triggers reconnect on network/connection failure. + - A process-level watcher that detects engine death via SIGCHLD (instant, + kernel-level), pidfd (event-driven), or /proc polling (last-resort).""" if self._db_health_watchdog_enabled is not True: verbose_proxy_logger.debug( "Prisma DB health watchdog disabled via PRISMA_HEALTH_WATCHDOG_ENABLED" @@ -3723,11 +4082,11 @@ class PrismaClient: self._db_health_watchdog_probe_timeout_seconds, self._db_watchdog_reconnect_timeout_seconds, ) + await self._start_engine_watcher() async def stop_db_health_watchdog_task(self) -> None: - """ - Stop DB health watchdog task gracefully. - """ + """Stop DB health watchdog task and engine watcher gracefully.""" + self._stop_engine_watcher() if self._db_health_watchdog_task is None: return self._db_health_watchdog_task.cancel() diff --git a/tests/litellm/proxy/test_prisma_engine_watchdog.py b/tests/litellm/proxy/test_prisma_engine_watchdog.py new file mode 100644 index 0000000000..3c4afb47cb --- /dev/null +++ b/tests/litellm/proxy/test_prisma_engine_watchdog.py @@ -0,0 +1,431 @@ +""" +Tests for PrismaClient engine watchdog: death detection and automatic reconnect. + +Covers: +- Engine PID discovery and liveness check +- Process disappears from /proc → reconnect triggered via attempt_db_reconnect +- Process becomes zombie in /proc → reconnect triggered via attempt_db_reconnect +- pidfd handler → schedules attempt_db_reconnect even when lock is held +- SIGCHLD handler → reaps all zombies, triggers reconnect if engine reaped +- _run_reconnect_cycle branches: heavy path (engine dead) vs lightweight path (engine alive) +- _engine_confirmed_dead flag ensures heavy reconnect even after _engine_pid reset +- Successful heavy reconnect → watcher re-armed for new process +- Missing DATABASE_URL → graceful RuntimeError in reconnect cycle +- Shutdown → polling loop exits cleanly, SIGCHLD handler removed +""" + +import asyncio +import os +import signal +import time +from unittest.mock import AsyncMock, MagicMock, mock_open, patch + +import pytest + +from litellm.proxy.utils import PrismaClient, ProxyLogging + + +@pytest.fixture(autouse=True) +def mock_prisma_binary(): + """Mock prisma.Prisma to avoid requiring generated Prisma binaries for unit tests.""" + import sys + + mock_module = MagicMock() + with patch.dict(sys.modules, {"prisma": mock_module}): + yield + + +@pytest.fixture +def mock_proxy_logging(): + proxy_logging = AsyncMock(spec=ProxyLogging) + proxy_logging.failure_handler = AsyncMock() + return proxy_logging + + +@pytest.fixture +def engine_client(mock_proxy_logging) -> PrismaClient: + """ + Minimal PrismaClient fixture for engine watchdog tests. + Uses the real constructor pattern from PR #21706 (database_url). + """ + client = PrismaClient(database_url="mock://test", proxy_logging_obj=mock_proxy_logging) + client.db = MagicMock() + client.db.recreate_prisma_client = AsyncMock() + client.db.disconnect = AsyncMock(return_value=None) + client.db.connect = AsyncMock(return_value=None) + client.db.query_raw = AsyncMock(return_value=[{"result": 1}]) + return client + + +# --------------------------------------------------------------------------- +# _is_engine_alive +# --------------------------------------------------------------------------- + + +def test_is_engine_alive_returns_true_when_pid_unknown(engine_client): + """_is_engine_alive returns True when no engine PID is tracked.""" + engine_client._engine_pid = 0 + assert engine_client._is_engine_alive() is True + + +def test_is_engine_alive_returns_false_when_process_gone(engine_client): + """_is_engine_alive returns False when /proc//stat is missing.""" + engine_client._engine_pid = 9999 + with patch("builtins.open", side_effect=FileNotFoundError): + assert engine_client._is_engine_alive() is False + + +def test_is_engine_alive_returns_false_for_zombie(engine_client): + """_is_engine_alive returns False when process is in zombie state.""" + engine_client._engine_pid = 1234 + zombie_stat = "1234 (prisma-query-engine) Z 1\n" + with patch("builtins.open", mock_open(read_data=zombie_stat)): + assert engine_client._is_engine_alive() is False + + +def test_is_engine_alive_returns_true_for_running_process(engine_client): + """_is_engine_alive returns True when process is in sleeping state.""" + engine_client._engine_pid = 1234 + alive_stat = "1234 (prisma-query-engine) S 1\n" + with patch("builtins.open", mock_open(read_data=alive_stat)): + assert engine_client._is_engine_alive() is True + + +# --------------------------------------------------------------------------- +# _poll_engine_proc — calls attempt_db_reconnect on death +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_poll_missing_process_triggers_reconnect(engine_client) -> None: + """Polling loop triggers attempt_db_reconnect when the engine process disappears.""" + engine_client._engine_pid = 1234 + engine_client._watching_engine = True + engine_client.attempt_db_reconnect = AsyncMock(return_value=True) + + with patch("builtins.open", side_effect=FileNotFoundError): + await engine_client._poll_engine_proc() + + engine_client.attempt_db_reconnect.assert_awaited_once_with( + reason="engine_process_death", + force=True, + ) + + +@pytest.mark.asyncio +async def test_poll_zombie_process_triggers_reconnect(engine_client) -> None: + """Polling loop triggers attempt_db_reconnect when the engine enters zombie state.""" + engine_client._engine_pid = 1234 + engine_client._watching_engine = True + engine_client.attempt_db_reconnect = AsyncMock(return_value=True) + + zombie_stat = "1234 (prisma-query-engine) Z 1\n" + with patch("builtins.open", mock_open(read_data=zombie_stat)): + await engine_client._poll_engine_proc() + + engine_client.attempt_db_reconnect.assert_awaited_once_with( + reason="engine_process_death", + force=True, + ) + + +@pytest.mark.asyncio +async def test_stop_loop_halts_polling(engine_client) -> None: + """Polling loop exits cleanly when _stop_engine_watcher is called.""" + engine_client._engine_pid = 1234 + engine_client._watching_engine = True + + alive_stat = "1234 (prisma-query-engine) S 1\n" + + async def stop_during_sleep(_duration: float) -> None: + engine_client._stop_engine_watcher() + + with ( + patch("builtins.open", mock_open(read_data=alive_stat)), + patch("asyncio.sleep", side_effect=stop_during_sleep), + ): + await engine_client._poll_engine_proc() + + assert engine_client._watching_engine is False + assert engine_client._engine_pid == 0 + + +# --------------------------------------------------------------------------- +# _on_pidfd_readable — calls attempt_db_reconnect +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_pidfd_readable_schedules_reconnect(engine_client) -> None: + """pidfd handler schedules attempt_db_reconnect via asyncio.create_task.""" + engine_client._engine_pid = 1234 + engine_client.attempt_db_reconnect = AsyncMock(return_value=True) + + created_coros = [] + + def capture_task(coro): + created_coros.append(coro) + return MagicMock() + + with patch("asyncio.create_task", side_effect=capture_task): + engine_client._on_pidfd_readable() + + # Run the captured coroutine to completion + assert len(created_coros) == 1 + await created_coros[0] + + engine_client.attempt_db_reconnect.assert_awaited_once_with( + reason="engine_process_death", + force=True, + ) + + +@pytest.mark.asyncio +async def test_pidfd_schedules_reconnect_task_when_lock_held(engine_client) -> None: + """pidfd handler schedules reconnect task even when _db_reconnect_lock is held.""" + engine_client._engine_pid = 1234 + + created_coros = [] + + def capture_task(coro): + created_coros.append(coro) + return MagicMock() + + async with engine_client._db_reconnect_lock: + with patch("asyncio.create_task", side_effect=capture_task): + engine_client._on_pidfd_readable() + + for coro in created_coros: + coro.close() + + assert len(created_coros) == 1 + + +# --------------------------------------------------------------------------- +# _run_reconnect_cycle — engine liveness branching +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_run_reconnect_cycle_uses_heavy_path_when_engine_dead( + engine_client, +) -> None: + """_run_reconnect_cycle calls recreate_prisma_client when engine is dead.""" + engine_client._engine_pid = 1234 + engine_client._start_engine_watcher = AsyncMock() + + with ( + patch.object(engine_client, "_is_engine_alive", return_value=False), + patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}), + patch("os.waitpid", side_effect=ChildProcessError), + ): + await engine_client._run_reconnect_cycle(timeout_seconds=5.0) + + engine_client.db.recreate_prisma_client.assert_awaited_once_with("postgresql://test") + engine_client._start_engine_watcher.assert_awaited_once() + engine_client.db.connect.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_run_reconnect_cycle_uses_heavy_path_when_confirmed_dead( + engine_client, +) -> None: + """_run_reconnect_cycle takes heavy path when _engine_confirmed_dead is set. + + This is the critical race-condition fix: SIGCHLD/pidfd handlers set + _engine_confirmed_dead BEFORE _cleanup_engine_watcher resets _engine_pid + to 0, so the heavy path executes even after cleanup. + """ + engine_client._engine_pid = 0 # Already reset by cleanup! + engine_client._engine_confirmed_dead = True # But flag survives cleanup + engine_client._start_engine_watcher = AsyncMock() + + with ( + patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}), + patch("os.waitpid", side_effect=ChildProcessError), + ): + await engine_client._run_reconnect_cycle(timeout_seconds=5.0) + + engine_client.db.recreate_prisma_client.assert_awaited_once_with("postgresql://test") + engine_client._start_engine_watcher.assert_awaited_once() + engine_client.db.connect.assert_not_awaited() + assert engine_client._engine_confirmed_dead is False # Reset after use + + +@pytest.mark.asyncio +async def test_run_reconnect_cycle_uses_lightweight_path_when_engine_alive( + engine_client, +) -> None: + """_run_reconnect_cycle uses disconnect/connect when engine is alive.""" + engine_client._engine_pid = 1234 + + with patch.object(engine_client, "_is_engine_alive", return_value=True): + await engine_client._run_reconnect_cycle(timeout_seconds=5.0) + + engine_client.db.connect.assert_awaited_once() + engine_client.db.query_raw.assert_awaited_once_with("SELECT 1") + engine_client.db.recreate_prisma_client.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_run_reconnect_cycle_uses_lightweight_path_when_pid_unknown( + engine_client, +) -> None: + """_run_reconnect_cycle uses lightweight path when engine PID is not tracked.""" + engine_client._engine_pid = 0 + + await engine_client._run_reconnect_cycle(timeout_seconds=5.0) + + engine_client.db.connect.assert_awaited_once() + engine_client.db.query_raw.assert_awaited_once_with("SELECT 1") + engine_client.db.recreate_prisma_client.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_run_reconnect_cycle_heavy_path_raises_without_database_url( + engine_client, +) -> None: + """Heavy reconnect raises RuntimeError when DATABASE_URL is not set.""" + engine_client._engine_pid = 1234 + + with ( + patch.object(engine_client, "_is_engine_alive", return_value=False), + patch.dict(os.environ, {}, clear=True), + patch("os.waitpid", side_effect=ChildProcessError), + ): + with pytest.raises(RuntimeError, match="DATABASE_URL not set"): + await engine_client._run_reconnect_cycle(timeout_seconds=5.0) + + engine_client.db.recreate_prisma_client.assert_not_awaited() + + +# --------------------------------------------------------------------------- +# start/stop lifecycle integration +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_start_watchdog_task_also_starts_engine_watcher( + engine_client, +) -> None: + """start_db_health_watchdog_task() also starts engine watcher.""" + engine_client._start_engine_watcher = AsyncMock() + + loop = asyncio.get_running_loop() + dummy_task = loop.create_task(asyncio.sleep(3600)) + + def fake_create_task(coro): + coro.close() + return dummy_task + + with patch("asyncio.create_task", side_effect=fake_create_task): + await engine_client.start_db_health_watchdog_task() + + engine_client._start_engine_watcher.assert_awaited_once() + dummy_task.cancel() + try: + await dummy_task + except asyncio.CancelledError: + pass + + +@pytest.mark.asyncio +async def test_stop_watchdog_task_also_stops_engine_watcher( + engine_client, +) -> None: + """stop_db_health_watchdog_task() also stops engine watcher and SIGCHLD handler.""" + engine_client._stop_engine_watcher = MagicMock() + + loop = asyncio.get_running_loop() + dummy_task = loop.create_task(asyncio.sleep(3600)) + engine_client._db_health_watchdog_task = dummy_task + + await engine_client.stop_db_health_watchdog_task() + + engine_client._stop_engine_watcher.assert_called_once() + assert engine_client._db_health_watchdog_task is None + + +# --------------------------------------------------------------------------- +# SIGCHLD handler +# --------------------------------------------------------------------------- + + +def test_sigchld_handler_reaps_engine_and_triggers_reconnect(engine_client): + """SIGCHLD handler detects engine death, reaps zombies, triggers reconnect.""" + engine_client._engine_pid = 1234 + created_coros = [] + + def capture_task(coro): + created_coros.append(coro) + return MagicMock() + + # Simulate waitpid returning the engine PID then raising ChildProcessError + with ( + patch("os.waitpid", side_effect=[(1234, 0), ChildProcessError]), + patch("asyncio.create_task", side_effect=capture_task), + ): + engine_client._on_sigchld() + + assert engine_client._engine_confirmed_dead is True + assert engine_client._engine_pid == 0 # cleanup ran + assert len(created_coros) == 1 + # Clean up the coroutine + created_coros[0].close() + + +def test_sigchld_handler_ignores_non_engine_zombies(engine_client): + """SIGCHLD handler reaps non-engine zombies without triggering reconnect.""" + engine_client._engine_pid = 1234 + + # Simulate reaping PID 5555 (not the engine) + with patch("os.waitpid", side_effect=[(5555, 0), ChildProcessError]): + engine_client._on_sigchld() + + assert engine_client._engine_confirmed_dead is False + assert engine_client._engine_pid == 1234 # unchanged + + +def test_sigchld_handler_no_double_trigger(engine_client): + """SIGCHLD handler does not trigger reconnect if already confirmed dead.""" + engine_client._engine_pid = 1234 + engine_client._engine_confirmed_dead = True # Already handled + + with ( + patch("os.waitpid", side_effect=[(1234, 0), ChildProcessError]), + patch("asyncio.create_task") as mock_create_task, + ): + engine_client._on_sigchld() + + mock_create_task.assert_not_called() + + +def test_install_sigchld_handler_success(engine_client): + """SIGCHLD handler installs on a running event loop.""" + mock_loop = MagicMock() + with patch("asyncio.get_running_loop", return_value=mock_loop): + assert engine_client._install_sigchld_handler() is True + + assert engine_client._sigchld_installed is True + mock_loop.add_signal_handler.assert_called_once_with( + signal.SIGCHLD, engine_client._on_sigchld + ) + + +def test_install_sigchld_handler_no_loop(engine_client): + """SIGCHLD handler returns False when no event loop is running.""" + with patch("asyncio.get_running_loop", side_effect=RuntimeError): + assert engine_client._install_sigchld_handler() is False + + assert engine_client._sigchld_installed is False + + +def test_remove_sigchld_handler(engine_client): + """SIGCHLD handler is properly removed.""" + engine_client._sigchld_installed = True + mock_loop = MagicMock() + with patch("asyncio.get_running_loop", return_value=mock_loop): + engine_client._remove_sigchld_handler() + + assert engine_client._sigchld_installed is False + mock_loop.remove_signal_handler.assert_called_once_with(signal.SIGCHLD)