diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 32078a3635..bbc549782d 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -3,7 +3,6 @@ import copy import hashlib import json import os -import signal import smtplib import threading import time @@ -2296,11 +2295,6 @@ 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( @@ -3568,383 +3562,31 @@ 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 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. + Run a reconnect cycle with direct db operations and a single overall timeout + budget to avoid long retries on hot paths (e.g. auth). """ + 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 ) - - 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) + await asyncio.wait_for(_do_direct_reconnect(), timeout=effective_timeout) async def attempt_db_reconnect( self, @@ -4076,10 +3718,9 @@ class PrismaClient: self._db_reconnect_lock.release() async def start_db_health_watchdog_task(self) -> None: - """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).""" + """ + Start a background task that probes DB health and attempts reconnect on failure. + """ if self._db_health_watchdog_enabled is not True: verbose_proxy_logger.debug( "Prisma DB health watchdog disabled via PRISMA_HEALTH_WATCHDOG_ENABLED" @@ -4097,11 +3738,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 and engine watcher gracefully.""" - self._stop_engine_watcher() + """ + Stop DB health watchdog task gracefully. + """ 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 deleted file mode 100644 index 3c4afb47cb..0000000000 --- a/tests/litellm/proxy/test_prisma_engine_watchdog.py +++ /dev/null @@ -1,431 +0,0 @@ -""" -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)