scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-12 03:48:56 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 57

test(migrations): assert lock-timeout retry and exhaustion instead of an unbounded wait (#45955)

Browse source 
* test(migrations): assert lock-timeout retry and exhaustion instead of an unbounded wait

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

* test(migrations): accumulate follower lock wait observations

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

---------

Co-authored-by: mateo <mateo@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot] 2026-10-11 17:36:47 +00:00 • committed by GitHub
parent c2fc628d36
commit 464a4dba9d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 121 additions and 12 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
.circleci/scripts/run_migration_tests.py
View file

@ -10,7 +10,7 @@ from typing import Final
from xml.etree import ElementTree
SUITES: Final = {
"startup": (("test_startup.py",), 12),
"startup": (("test_startup.py",), 13),
"recovery": (("test_recovery.py",), 15),
"legacy": (("test_legacy.py", "test_pooling.py"), 11),
"upgrade": (("test_upgrade.py", "test_rolling_upgrade.py"), 5),

19
tests/e2e/migrations/checks.py
View file

@ -39,6 +39,25 @@ def assert_completed(database: Database, migration: Migration = COMPLETE) -> Non
assert database.query("SELECT id FROM migration_effect") == ((1,),)
@step("Check that migration {migration.name} rolled back lock timeouts and completed exactly once")
def assert_completed_after_lock_timeouts(database: Database, migration: Migration) -> None:
rows: Final = database.query(
"SELECT finished_at IS NOT NULL, rolled_back_at IS NOT NULL, applied_steps_count FROM "
"_prisma_migrations WHERE migration_name = %s ORDER BY started_at",
(migration.name,),
)
finished: Final = tuple(row for row in rows if row[0])
rolled_back: Final = tuple(row for row in rows if row[1])
assert finished == ((True, False, 1),), "Expected exactly one completed migration attempt"
assert rolled_back and all(row == (False, True, 0) for row in rolled_back), (
"Expected every failed migration attempt to be rolled back"
)
assert all(is_finished != is_rolled_back for is_finished, is_rolled_back, _ in rows), (
"Expected every migration attempt to be either completed or rolled back"
)
assert database.query("SELECT id FROM migration_effect") == ((1,),)
@step("Apply the test migration by hand and record it in _prisma_migrations")
def confirmed_history(database: Database) -> str:
database.execute(COMPLETE_SQL)

10
tests/e2e/migrations/containers.py
View file

@ -41,6 +41,7 @@ def until(description: str, condition: Callable[[], bool], seconds: float = 150)
@dataclass(frozen=True, slots=True)
class Replica:
name: str
identity: str
transport: HttpTransport
output: Path
@ -171,7 +172,7 @@ class Containers:
"model_list: []\ngeneral_settings:\n master_key: os.environ/LITELLM_MASTER_KEY\n"
)
env: Final = {
"DATABASE_URL": prisma_url(database.container_url, database.schema),
"DATABASE_URL": prisma_url(database.container_url, database.schema, application_name=name),
"LITELLM_MASTER_KEY": MASTER_KEY,
"LITELLM_SALT_KEY": MASTER_KEY,
"LITELLM_LOCAL_MODEL_COST_MAP": "True",
@ -209,7 +210,12 @@ class Containers:
"4000",
)
port: Final = int(docker("port", name, "4000/tcp").rsplit(":", 1)[1])
replica: Final = Replica(name, HttpTransport(f"http://127.0.0.1:{port}", MASTER_KEY, 15), directory)
replica: Final = Replica(
name=name,
identity=name,
transport=HttpTransport(f"http://127.0.0.1:{port}", MASTER_KEY, 15),
output=directory,
)
yield replica
finally:
try:

37
tests/e2e/migrations/database.py
View file

@ -3,6 +3,7 @@ from __future__ import annotations
from collections.abc import Generator
from contextlib import contextmanager
from dataclasses import dataclass
from itertools import chain
from typing import Final, LiteralString
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
from uuid import uuid4
@ -24,10 +25,19 @@ def connect_url(url: str, name: str) -> str:
return urlunsplit(urlsplit(url)._replace(path=f"/{name}", query=""))
def prisma_url(url: str, schema: str) -> str:
def prisma_url(url: str, schema: str, application_name: str | None = None) -> str:
parsed: Final = urlsplit(url)
query: Final = tuple((key, value) for key, value in parse_qsl(parsed.query) if key != "schema")
return urlunsplit(parsed._replace(query=urlencode((*query, ("schema", schema)))))
query: Final = tuple(
(key, value)
for key, value in parse_qsl(parsed.query)
if key != "schema" and (application_name is None or key != "application_name")
)
parameters: Final = query + (
(("schema", schema),)
if application_name is None
else (("schema", schema), ("application_name", application_name))
)
return urlunsplit(parsed._replace(query=urlencode(parameters)))
@dataclass(frozen=True, slots=True)
@ -79,6 +89,27 @@ class Database:
(key >> 32, key & 0xFFFFFFFF),
)
@step("List replica identities waiting on any of the advisory locks")
def waiting_identities(self, keys: tuple[int, ...]) -> frozenset[str]:
if not keys:
return frozenset()
key_pairs: Final = tuple((key >> 32, key & 0xFFFFFFFF) for key in keys)
key_predicates: Final = sql.SQL(" OR ").join(
sql.SQL("(waiting_lock.classid = %s AND waiting_lock.objid = %s)") for _ in key_pairs
)
parameters: Final = tuple(chain.from_iterable(key_pairs))
statement: Final = sql.SQL(
"SELECT DISTINCT activity.application_name "
"FROM pg_stat_activity AS activity "
"JOIN pg_locks AS waiting_lock ON waiting_lock.pid = activity.pid "
"WHERE activity.datname = current_database() "
"AND waiting_lock.locktype = 'advisory' AND NOT waiting_lock.granted "
"AND waiting_lock.database = (SELECT oid FROM pg_database WHERE datname = current_database()) "
"AND ({})"
).format(key_predicates)
identities: Final = self.query(statement, parameters)
return frozenset(row[0] for row in identities if isinstance(row[0], str))
@step("Hold an advisory lock on the test database")
@contextmanager
def lock(self, key: int = GATE_KEY) -> Generator[None]:

65
tests/e2e/migrations/test_startup.py
View file

@ -1,17 +1,31 @@
from contextlib import ExitStack
from itertools import accumulate, repeat
from typing import Final
import pytest
from .checks import COMPLETE, FATAL, GATED, assert_completed, start_replicas
from .checks import (
COMPLETE,
FATAL,
GATED,
assert_completed,
assert_completed_after_lock_timeouts,
start_replicas,
)
from .containers import Containers, failed, ready, until, waiting
from .database import PRISMA_LOCK, Database, Databases, restricted_user
from .database import GATE_KEY, PRISMA_LOCK, Database, Databases, restricted_user
from .startup_models import Migration
from e2e_metadata import Domain, Subject, meta
pytestmark: Final = [pytest.mark.e2e, pytest.mark.migration_startup]
def union_waiting_identities(
observed: frozenset[str], snapshot: frozenset[str]
) -> frozenset[str]:
return observed | snapshot
class TestMigrationStartup:
@pytest.mark.parametrize("replicas,v2", ((1, True), (3, True), (1, False)))
@meta(
@ -42,16 +56,34 @@ class TestMigrationStartup:
domain=Domain.DB,
)
)
def test_waiters_survive_prolonged_contention(self, containers: Containers, database: Database) -> None:
def test_waiters_retry_lock_timeouts_until_contention_clears(
self, containers: Containers, database: Database
) -> None:
with ExitStack() as stack:
with database.lock():
owner: Final = stack.enter_context(containers.start(database, (GATED,)))
until("owner blocked in migration SQL", lambda: bool(database.blocked()))
followers: Final = start_replicas(stack, containers, database, (GATED,), count=2)
until("both followers attempted Prisma locking", lambda: len(database.blocked(PRISMA_LOCK)) == 2)
waiting((owner, *followers), 120)
waiting_identities: Final = accumulate(
(database.waiting_identities((PRISMA_LOCK, GATE_KEY)) for _ in repeat(None)),
union_waiting_identities,
initial=frozenset(),
)
follower_identities: Final = frozenset(follower.identity for follower in followers)
until(
"both followers entered migration locking",
lambda: follower_identities <= next(waiting_identities),
)
until(
"a replica logged a migration lock-timeout retry",
lambda: any(
"timed out waiting for a table lock" in replica.logs()
for replica in (owner, *followers)
),
)
waiting((owner, *followers), 20)
ready((owner, *followers), database)
assert_completed(database, GATED)
assert_completed_after_lock_timeouts(database, GATED)
@meta(
Subject(
@ -72,6 +104,27 @@ class TestMigrationStartup:
ready((restarted,), database)
assert_completed(database)
@meta(
Subject(
domain=Domain.DB,
)
)
def test_lock_timeout_exhaustion_then_restart(self, containers: Containers, database: Database) -> None:
with database.lock():
with containers.start(
database, (GATED,), environment={"LITELLM_MIGRATION_DDL_LOCK_TIMEOUT": "2"}
) as replica:
until("owner blocked in migration SQL", lambda: bool(database.blocked()))
failed((replica,), "attempts that made no progress")
assert not database.exists("migration_effect")
assert database.query(
"SELECT count(*) FROM _prisma_migrations WHERE migration_name = %s AND finished_at IS NOT NULL",
(GATED.name,),
) == ((0,),)
with containers.start(database, (GATED,)) as restarted:
ready((restarted,), database)
assert_completed_after_lock_timeouts(database, GATED)
@meta(
Subject(
domain=Domain.DB,