From 8340ddea64323f4181684bd2d49bc42d857b0c02 Mon Sep 17 00:00:00 2001 From: yassin Date: Sat, 29 Aug 2026 09:50:14 +0000 Subject: [PATCH] fix(db): release reconnect lock when a waiting reconnect is cancelled Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/utils.py | 20 ++++-- .../test_prisma_client_reconnect.py | 63 ++++++++++--------- 2 files changed, 48 insertions(+), 35 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 9fbe1b4bd06..76580a2a7b0 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -5580,12 +5580,8 @@ class PrismaClient: return True acquire_task: Final = asyncio.create_task(_acquire_reconnect_lock()) - done, _pending = await asyncio.wait( - {acquire_task}, - timeout=lock_timeout_seconds, - return_when=asyncio.FIRST_COMPLETED, - ) - if acquire_task not in done: + + async def _abandon_acquire_task() -> None: acquire_task.cancel() try: await acquire_task @@ -5600,6 +5596,18 @@ class PrismaClient: self._db_reconnect_lock.release() except RuntimeError: pass + + try: + done, _pending = await asyncio.wait( + {acquire_task}, + timeout=lock_timeout_seconds, + return_when=asyncio.FIRST_COMPLETED, + ) + except asyncio.CancelledError: + await asyncio.shield(_abandon_acquire_task()) + raise + if acquire_task not in done: + await _abandon_acquire_task() verbose_proxy_logger.debug( "Skipping DB reconnect attempt due to lock acquisition timeout. reason=%s timeout=%ss", reason, diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_reconnect.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_reconnect.py index 719d7cc73f5..41b0eb3cf95 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_reconnect.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_reconnect.py @@ -89,9 +89,7 @@ async def test_run_reconnect_cycle_direct_path_recreates_when_probe_fails( prisma_client._cleanup_engine_watcher = MagicMock() writer = MagicMock() - writer.query_raw = AsyncMock( - side_effect=[ConnectionError("probe failed"), [{"?column?": 1}]] - ) + writer.query_raw = AsyncMock(side_effect=[ConnectionError("probe failed"), [{"?column?": 1}]]) monkeypatch.setattr( PrismaClient, "writer_db", @@ -171,9 +169,7 @@ async def test_run_reconnect_cycle_passes_writer_generation_to_recreate( writer = MagicMock() writer._engine_generation = 7 - writer.query_raw = AsyncMock( - side_effect=[ConnectionError("probe failed"), [{"?column?": 1}]] - ) + writer.query_raw = AsyncMock(side_effect=[ConnectionError("probe failed"), [{"?column?": 1}]]) monkeypatch.setattr( PrismaClient, "writer_db", @@ -229,9 +225,7 @@ async def test_attempt_reconnect_inside_lock_runs_cycle_and_resets_counter( prisma_client._consecutive_reconnect_failures = 2 prisma_client._run_reconnect_cycle = AsyncMock() - ok = await prisma_client._attempt_reconnect_inside_lock( - force=True, reason="test", timeout_seconds=1 - ) + ok = await prisma_client._attempt_reconnect_inside_lock(force=True, reason="test", timeout_seconds=1) pinned = { "returned": ok, "cycle_called": prisma_client._run_reconnect_cycle.await_count, @@ -254,9 +248,7 @@ async def test_attempt_reconnect_inside_lock_skips_when_in_cooldown( prisma_client._db_last_reconnect_attempt_ts = time.time() prisma_client._run_reconnect_cycle = AsyncMock() - ok = await prisma_client._attempt_reconnect_inside_lock( - force=False, reason="test", timeout_seconds=1 - ) + ok = await prisma_client._attempt_reconnect_inside_lock(force=False, reason="test", timeout_seconds=1) assert ok is False assert prisma_client._run_reconnect_cycle.await_count == 0 @@ -269,9 +261,7 @@ async def test_attempt_reconnect_inside_lock_increments_failure_counter_on_error prisma_client._consecutive_reconnect_failures = 0 prisma_client._run_reconnect_cycle = AsyncMock(side_effect=RuntimeError("boom")) - ok = await prisma_client._attempt_reconnect_inside_lock( - force=True, reason="failing_test", timeout_seconds=1 - ) + ok = await prisma_client._attempt_reconnect_inside_lock(force=True, reason="failing_test", timeout_seconds=1) assert ok is False assert prisma_client._consecutive_reconnect_failures == 1 @@ -316,9 +306,7 @@ async def test_attempt_db_reconnect_lock_timeout_returns_false( by replacing ``asyncio.wait`` with a callable that returns the loser task as still-pending after it's already been completed elsewhere. """ - completed_task: asyncio.Task[bool] = asyncio.get_running_loop().create_task( - _no_op_returning_true() - ) + completed_task: asyncio.Task[bool] = asyncio.get_running_loop().create_task(_no_op_returning_true()) # Ensure the inner task has finished before attempt_db_reconnect sees it. await completed_task @@ -329,7 +317,7 @@ async def test_attempt_db_reconnect_lock_timeout_returns_false( monkeypatch.setattr( asyncio, "create_task", - lambda coro, *a, **kw: (coro.close() or completed_task), + lambda coro, *a, **kw: coro.close() or completed_task, ) prisma_client._db_last_reconnect_attempt_ts = 0.0 @@ -465,9 +453,7 @@ async def test_db_health_watchdog_loop_triggers_reconnect_on_timeout( await prisma_client._db_health_watchdog_loop() pinned = { "reconnect_called": prisma_client.attempt_db_reconnect.await_count, - "reconnect_reason": prisma_client.attempt_db_reconnect.await_args.kwargs[ - "reason" - ], + "reconnect_reason": prisma_client.attempt_db_reconnect.await_args.kwargs["reason"], "wait_for_calls": call_count["n"], "loop_exited_clean": True, } @@ -522,10 +508,7 @@ async def test_iam_refresh_racing_reconnect_recreates_engine_only_once( from litellm.proxy.db.prisma_client import PrismaWrapper def token_db_url(created: datetime) -> str: - token = ( - f"host/?X-Amz-Date={created.strftime('%Y%m%dT%H%M%SZ')}" - f"&X-Amz-Expires=900&X-Amz-Signature=abc" - ) + token = f"host/?X-Amz-Date={created.strftime('%Y%m%dT%H%M%SZ')}&X-Amz-Expires=900&X-Amz-Signature=abc" return f"postgresql://user:{urllib.parse.quote(token, safe='')}@host:5432/db" # Old engine (PID 111) carries an expired token; in-flight queries on it @@ -577,9 +560,7 @@ async def test_iam_refresh_racing_reconnect_recreates_engine_only_once( # In-flight transport-error path fires while the refresh holds the # wrapper's reconnection lock mid-recreate. reconnect_task = asyncio.create_task( - prisma_client.attempt_db_reconnect( - reason="in_flight_transport_error", force=True - ) + prisma_client.attempt_db_reconnect(reason="in_flight_transport_error", force=True) ) await asyncio.sleep(0.05) release_connect.set() @@ -1096,3 +1077,27 @@ async def test_unrelated_reconnect_failure_does_not_erase_the_burst_record( "cycles_after": prisma_client._run_reconnect_cycle.await_count, } assert pinned == {"cycles_before": 2, "cycles_after": 2} + + +@pytest.mark.asyncio +async def test_attempt_db_reconnect_cancelled_while_waiting_does_not_strand_lock( + prisma_client: PrismaClient, +) -> None: + """A reconnect cancelled while waiting on the lock (e.g. the readiness + probe deadline firing) must abandon its lock-acquisition task instead of + leaving it to grab the lock later with no owner to release it.""" + prisma_client._db_last_reconnect_attempt_ts = 0.0 + prisma_client._attempt_reconnect_inside_lock = AsyncMock(return_value=True) + + await prisma_client._db_reconnect_lock.acquire() + waiting_reconnect: Final = asyncio.create_task( + prisma_client.attempt_db_reconnect(reason="probe_deadline", lock_timeout_seconds=30.0) + ) + await asyncio.sleep(0.05) + waiting_reconnect.cancel() + with pytest.raises(asyncio.CancelledError): + await waiting_reconnect + + prisma_client._db_reconnect_lock.release() + await asyncio.sleep(0.05) + assert prisma_client._db_reconnect_lock.locked() is False