fix(proxy-extras): hold only Prisma's lock key during migrate deploy so it never waits on the request log index build (#46660)

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm-proxy-extras/litellm_proxy_extras/migration_lock.py

@@ -18,6 +18,25 @@ MIGRATION_LOCK_KEY: Final = int.from_bytes(b"llm_mig2", "big")
 PRISMA_ADVISORY_LOCK_KEY: Final = 72707369
 _SERVER_ANSWERED_MARKERS: Final = ("FATAL:", "SSL", "certificate")
 
+
+@dataclass(frozen=True, slots=True)
+class AdvisoryLock:
+    key: int
+    name: str
+    holder: str
+
+
+COORDINATOR_LOCK: Final = AdvisoryLock(
+    MIGRATION_LOCK_KEY,
+    "v2 migration coordinator lock",
+    "another v2 migration resolver",
+)
+PRISMA_LOCK: Final = AdvisoryLock(
+    PRISMA_ADVISORY_LOCK_KEY,
+    "Prisma migration lock",
+    "another Prisma migration",
+)
+
 if TYPE_CHECKING:
     import psycopg
 
@@ -100,7 +119,10 @@ class MigrationCoordinator:
 
 
 @contextmanager
-def migration_lock(database_url: str) -> Generator[MigrationCoordinator, None, None]:
+def migration_lock(
+    database_url: str,
+    lock: AdvisoryLock = COORDINATOR_LOCK,
+) -> Generator[MigrationCoordinator, None, None]:
     import psycopg
 
     wait_seconds: Final = migration_lock_timeout()
@@ -112,11 +134,11 @@ def migration_lock(database_url: str) -> Generator[MigrationCoordinator, None, N
     try:
         with connection:
             coordinator: Final = MigrationCoordinator(connection)
-            logger.info("Waiting for the v2 migration coordinator lock (up to %ss)", wait_seconds)
+            logger.info("Waiting for the %s (up to %ss)", lock.name, wait_seconds)
             while time.monotonic() < deadline:
                 with connection.transaction():
-                    if _try_lock(connection):
-                        logger.info("Acquired the v2 migration coordinator lock")
+                    if _try_lock(connection, lock.key):
+                        logger.info("Acquired the %s", lock.name)
 
                         yield coordinator
                         coordinator.check_connection()
@@ -125,18 +147,17 @@ def migration_lock(database_url: str) -> Generator[MigrationCoordinator, None, N
     except psycopg.Error as exc:
         raise RuntimeError(_coordination_failure(database_url, exc, connecting=False)) from None
     raise RuntimeError(
-        f"Timed out waiting for another v2 migration resolver after {wait_seconds}s. "
+        f"Timed out waiting for {lock.holder} after {wait_seconds}s. "
         f"Check the running migration or increase {MIGRATION_LOCK_TIMEOUT_ENV_VAR}."
     )
 
 
 @contextmanager
 def prisma_migration_lock(database_url: str) -> Generator[None, None, None]:
-    """Hold the coordinator lock and Prisma's own lock key as transaction-level locks on one
-    pinned connection while a Prisma command runs with its session advisory lock disabled, so a
-    transaction pooler can never strand that lock on an idle backend after the command exits."""
-    with migration_lock(database_url) as coordinator:
-        coordinator.acquire_prisma_lock()
+    """Hold only Prisma's key as an xact lock on a pinned connection with Prisma's session lock disabled.
+
+    This keeps deploy independent of the coordinator key held by request-log index builds."""
+    with migration_lock(database_url, PRISMA_LOCK) as coordinator:
         coordinator.release_statement_snapshot()
         yield
 

tests/e2e/migrations/test_startup.py

@@ -97,11 +97,9 @@ class TestMigrationStartup:
             with containers.start(
                 database, (COMPLETE,), environment={"LITELLM_MIGRATION_LOCK_TIMEOUT": "12"}
             ) as replica:
-                until(
-                    "a replica holding the coordinator lock while Prisma's lock is held",
-                    lambda: bool(database.holders(COORDINATOR_LOCK)),
-                )
-                failed((replica,), "Timed out waiting for Prisma's lock")
+                until("Prisma lock contention", lambda: "Waiting for the Prisma migration lock" in replica.logs())
+                failed((replica,), "Timed out waiting for another Prisma migration")
+                assert database.holders(COORDINATOR_LOCK) == ()
                 assert database.history() == history
                 assert not database.exists("migration_effect")
         with containers.start(database, (COMPLETE,)) as restarted:

tests/integration/database/test_prisma_migration_lock.py

@@ -51,8 +51,31 @@ def test_prisma_migration_lock_holds_prismas_key_on_a_pinned_connection_and_leav
         with prisma_migration_lock(os.environ["DATABASE_URL"]):
             held: Final = _granted_advisory_keys(peer)
             assert PRISMA_ADVISORY_LOCK_KEY in held
-            assert MIGRATION_LOCK_KEY in held
+            assert MIGRATION_LOCK_KEY not in held
             assert not _can_take_prismas_session_lock(peer)
 
         assert _granted_advisory_keys(peer).isdisjoint({PRISMA_ADVISORY_LOCK_KEY, MIGRATION_LOCK_KEY})
         assert _can_take_prismas_session_lock(peer)
+
+
+def test_prisma_migration_lock_does_not_wait_for_the_coordinator_lock(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setenv("LITELLM_MIGRATION_LOCK_TIMEOUT", "3")
+    with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as peer:
+        peer.execute("SELECT pg_advisory_lock(%s)", (MIGRATION_LOCK_KEY,))
+
+        with prisma_migration_lock(os.environ["DATABASE_URL"]):
+            assert PRISMA_ADVISORY_LOCK_KEY in _granted_advisory_keys(peer)
+
+
+def test_prisma_migration_lock_times_out_when_another_prisma_migration_holds_its_key(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setenv("LITELLM_MIGRATION_LOCK_TIMEOUT", "3")
+    with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as peer:
+        peer.execute("SELECT pg_advisory_lock(%s)", (PRISMA_ADVISORY_LOCK_KEY,))
+
+        with pytest.raises(RuntimeError, match="Timed out waiting for another Prisma migration"):
+            with prisma_migration_lock(os.environ["DATABASE_URL"]):
+                pytest.fail("Prisma migration lock was acquired while another migration held it")