From f6d743fb18ba341168388003a48bbe98fd62ae08 Mon Sep 17 00:00:00 2001 From: harish-berri Date: Tue, 19 May 2026 00:32:08 +0000 Subject: [PATCH] feat(proxy): enhance PrismaClient health watchdog with reconnect cooldown and tests - Introduced a reconnect cooldown for the health watchdog to prevent rapid reconnection attempts during database errors. - Updated health watchdog parameters to allow for configurable intervals and timeouts via environment variables. - Enhanced logging to provide clearer information on the watchdog's operation. - Added tests to verify the reconnect behavior during cooldown periods and ensure proper functionality of the health watchdog. --- litellm/proxy/utils.py | 26 ++++++++++++---- .../proxy/db/test_prisma_self_heal.py | 30 +++++++++++++++++++ 2 files changed, 51 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index be8bf4e0920..b9190644427 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -2851,21 +2851,26 @@ class PrismaClient: self._db_reconnect_lock = asyncio.Lock() self._db_health_watchdog_task: Optional[asyncio.Task] = None self._db_last_reconnect_attempt_ts: float = 0.0 + self._db_last_health_watchdog_reconnect_attempt_ts: float = 0.0 self._db_reconnect_cooldown_seconds: int = max( 1, int(os.getenv("PRISMA_RECONNECT_COOLDOWN_SECONDS", "15")) ) self._db_health_watchdog_interval_seconds: int = max( - 5, int(os.getenv("PRISMA_HEALTH_WATCHDOG_INTERVAL_SECONDS", "30")) + 5, int(os.getenv("PRISMA_HEALTH_WATCHDOG_INTERVAL_SECONDS", "120")) ) self._db_health_watchdog_enabled: bool = ( str_to_bool(os.getenv("PRISMA_HEALTH_WATCHDOG_ENABLED", "true")) is True ) self._db_health_watchdog_probe_timeout_seconds: float = max( 0.5, - float(os.getenv("PRISMA_HEALTH_WATCHDOG_PROBE_TIMEOUT_SECONDS", "5.0")), + float(os.getenv("PRISMA_HEALTH_WATCHDOG_PROBE_TIMEOUT_SECONDS", "15")), + ) + self._db_health_watchdog_reconnect_cooldown_seconds: int = max( + 1, + int(os.getenv("PRISMA_HEALTH_WATCHDOG_RECONNECT_COOLDOWN_SECONDS", "120")), ) self._db_watchdog_reconnect_timeout_seconds: float = max( - 1.0, float(os.getenv("PRISMA_WATCHDOG_RECONNECT_TIMEOUT_SECONDS", "30.0")) + 1.0, float(os.getenv("PRISMA_WATCHDOG_RECONNECT_TIMEOUT_SECONDS", "60.0")) ) self._db_auth_reconnect_timeout_seconds: float = max( 0.5, float(os.getenv("PRISMA_AUTH_RECONNECT_TIMEOUT_SECONDS", "2.0")) @@ -4772,9 +4777,9 @@ class PrismaClient: self._db_health_watchdog_loop() ) verbose_proxy_logger.info( - "Started Prisma DB health watchdog (interval=%ss, reconnect_cooldown=%ss, probe_timeout=%ss, reconnect_timeout=%ss)", + "Started Prisma DB health watchdog (interval=%ss, watchdog_reconnect_cooldown=%ss, probe_timeout=%ss, reconnect_timeout=%ss)", self._db_health_watchdog_interval_seconds, - self._db_reconnect_cooldown_seconds, + self._db_health_watchdog_reconnect_cooldown_seconds, self._db_health_watchdog_probe_timeout_seconds, self._db_watchdog_reconnect_timeout_seconds, ) @@ -4807,8 +4812,19 @@ class PrismaClient: if isinstance( e, asyncio.TimeoutError ) or PrismaDBExceptionHandler.is_database_connection_error(e): + now = time.time() + if ( + now - self._db_last_health_watchdog_reconnect_attempt_ts + < self._db_health_watchdog_reconnect_cooldown_seconds + ): + verbose_proxy_logger.debug( + "Skipping DB health watchdog reconnect due to watchdog cooldown." + ) + continue + self._db_last_health_watchdog_reconnect_attempt_ts = now await self.attempt_db_reconnect( reason="db_health_watchdog_connection_error", + force=True, timeout_seconds=self._db_watchdog_reconnect_timeout_seconds, ) else: diff --git a/tests/test_litellm/proxy/db/test_prisma_self_heal.py b/tests/test_litellm/proxy/db/test_prisma_self_heal.py index 3f9ba6af3af..7b171e33ff9 100644 --- a/tests/test_litellm/proxy/db/test_prisma_self_heal.py +++ b/tests/test_litellm/proxy/db/test_prisma_self_heal.py @@ -271,6 +271,7 @@ async def test_db_health_watchdog_should_trigger_reconnect_on_db_error( client.attempt_db_reconnect.assert_awaited_once_with( reason="db_health_watchdog_connection_error", + force=True, timeout_seconds=7.0, ) @@ -302,10 +303,39 @@ async def test_db_health_watchdog_should_trigger_reconnect_on_probe_timeout( client.attempt_db_reconnect.assert_awaited_once_with( reason="db_health_watchdog_connection_error", + force=True, timeout_seconds=9.0, ) +@pytest.mark.asyncio +async def test_db_health_watchdog_should_skip_reconnect_during_watchdog_cooldown( + mock_proxy_logging, +): + client = PrismaClient( + database_url="mock://test", proxy_logging_obj=mock_proxy_logging + ) + client.db.query_raw = AsyncMock(side_effect=Exception("db connection dropped")) + client.attempt_db_reconnect = AsyncMock(return_value=True) + client._db_health_watchdog_interval_seconds = 1 + client._db_health_watchdog_reconnect_cooldown_seconds = 3600 + client._db_last_health_watchdog_reconnect_attempt_ts = time.time() + + with ( + patch( + "litellm.proxy.utils.asyncio.sleep", + AsyncMock(side_effect=[None, asyncio.CancelledError()]), + ), + patch( + "litellm.proxy.db.exception_handler.PrismaDBExceptionHandler.is_database_connection_error", + return_value=True, + ), + ): + await client._db_health_watchdog_loop() + + client.attempt_db_reconnect.assert_not_called() + + @pytest.mark.asyncio async def test_db_health_watchdog_start_stop_lifecycle(mock_proxy_logging): client = PrismaClient(