From fb4b1e728a18a56f8ec93d72ade3c237c6bbddb6 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Thu, 3 Sep 2026 12:16:08 -0700 Subject: [PATCH] test(team-race): wait on pg_locks instead of a fixed sleep The three race tests claimed to pin the interleaving deterministically, but lock_acquired.set() ran as the first statement of the task coroutine, before the awaited endpoint call, so it only signalled that the task had started. The real synchronisation was `await asyncio.sleep(0.2)` followed by `assert not task.done()`, which is a timing assumption on a box running four xdist workers against one Postgres. Take the blocking connection's advisory lock key straight out of pg_locks, then poll from the unblocked watcher connection until a non-granted lock on that same key appears. That is the condition the sleep was standing in for, and it holds however slow the machine is. If the endpoint returns without ever queueing, the helper now fails with the endpoint's own exception chained on instead of a bare assert. The two tests reported failing in CI used sleep(0.2); the third used sleep(0.3) and was not reported, which is consistent with the margin being the cause. --- .../test_team_delete_member_add_race.py | 74 +++++++++++++------ 1 file changed, 53 insertions(+), 21 deletions(-) diff --git a/tests/proxy_admin_ui_tests/test_team_delete_member_add_race.py b/tests/proxy_admin_ui_tests/test_team_delete_member_add_race.py index 7577570be48..9af742b0dfe 100644 --- a/tests/proxy_admin_ui_tests/test_team_delete_member_add_race.py +++ b/tests/proxy_admin_ui_tests/test_team_delete_member_add_race.py @@ -9,14 +9,16 @@ be forced by a sequential script: it needs one request to be genuinely mid-fligh other commits. A mocked prisma cannot arbitrate that either, since the property under test is whether Postgres's own advisory lock actually serializes the two requests. -These tests pin the interleaving the same way test_access_group_team_sync.py does: a second -real connection holds the team's advisory lock in its own transaction, so the function under -test is provably blocked on it rather than hoping a sleep lands in the right gap. +These tests pin the interleaving without a timing assumption: a second real connection holds +the team's advisory lock in its own transaction, and the test then waits for Postgres itself +to report the endpoint queued behind that exact lock. A sleep can only guess whether the +endpoint has reached the lock yet; pg_locks answers it. """ import asyncio import json import os +import time import uuid from contextlib import asynccontextmanager from datetime import timedelta @@ -39,6 +41,48 @@ _DELETE_SEEDED = 'DELETE FROM "LiteLLM_TeamMembership" WHERE team_id = $1' _DELETE_USER = 'DELETE FROM "LiteLLM_UserTable" WHERE user_id = $1' _DELETE_TEAM = 'DELETE FROM "LiteLLM_TeamTable" WHERE team_id = $1' _LOCK_SQL = "SELECT pg_advisory_xact_lock(hashtext($1)) IS NULL AS locked" +_HELD_LOCK_KEY_SQL = ( + "SELECT classid::bigint AS classid, objid::bigint AS objid FROM pg_locks " + "WHERE locktype = 'advisory' AND granted AND pid = pg_backend_pid()" +) +_LOCK_WAITER_SQL = ( + "SELECT count(*)::int AS waiters FROM pg_locks " + "WHERE locktype = 'advisory' AND NOT granted " + "AND classid::bigint = $1 AND objid::bigint = $2" +) +_LOCK_WAIT_TIMEOUT_SECONDS = 20.0 +_LOCK_POLL_SECONDS = 0.01 + + +async def _hold_team_lock(held, team_id: str) -> tuple[int, int]: + """Take the team's advisory lock and return its pg_locks key. + + Reading the key back off our own backend avoids re-deriving hashtext()'s signed + 32-bit split here, and pins the watcher to this lock rather than to any advisory + lock another xdist worker happens to hold on the same database.""" + await held.query_raw(_LOCK_SQL, team_id) + rows = await held.query_raw(_HELD_LOCK_KEY_SQL) + assert len(rows) == 1, f"expected exactly one advisory lock on the blocking connection, got {rows}" + return rows[0]["classid"], rows[0]["objid"] + + +async def _await_lock_contention(watcher, lock_key: tuple[int, int], task, what: str) -> None: + """Block until Postgres reports `task` queued behind the held lock. + + This is the assertion that the endpoint serializes on the team's advisory lock, and it + is what a fixed sleep was standing in for: the endpoint is only provably waiting once a + non-granted advisory lock on the same key exists. `watcher` must be a connection that is + not itself blocked, so it can observe the queue.""" + classid, objid = lock_key + deadline = time.monotonic() + _LOCK_WAIT_TIMEOUT_SECONDS + while time.monotonic() < deadline: + if task.done(): + raise AssertionError(f"{what} returned without waiting on the team's advisory lock") from task.exception() + rows = await watcher.query_raw(_LOCK_WAITER_SQL, classid, objid) + if rows[0]["waiters"]: + return + await asyncio.sleep(_LOCK_POLL_SECONDS) + raise AssertionError(f"{what} never queued on the team's advisory lock within {_LOCK_WAIT_TIMEOUT_SECONDS}s") def _race_ids() -> tuple[str, str]: @@ -110,10 +154,8 @@ async def test_member_add_blocked_by_delete_writes_no_dangling_reference(): blocker = Prisma() await blocker.connect() - lock_acquired = asyncio.Event() async def add_member(): - lock_acquired.set() await _add_team_members_to_team( data=TeamMemberAddRequest( team_id=team_id, @@ -128,11 +170,9 @@ async def test_member_add_blocked_by_delete_writes_no_dangling_reference(): try: async with blocker.tx(timeout=timedelta(seconds=30)) as held: - await held.query_raw(_LOCK_SQL, team_id) + lock_key = await _hold_team_lock(held, team_id) task = asyncio.create_task(add_member()) - await lock_acquired.wait() - await asyncio.sleep(0.2) - assert not task.done(), "member_add did not wait on the team's advisory lock" + await _await_lock_contention(db, lock_key, task, "member_add") # the delete wins the race: strip the team row while the lock is held await held.execute_raw(_DELETE_TEAM, team_id) @@ -185,10 +225,8 @@ async def test_member_delete_blocked_by_member_add_removes_from_the_fresh_roster blocker = Prisma() await blocker.connect() - lock_acquired = asyncio.Event() async def run_delete(): - lock_acquired.set() return await team_member_delete( data=TeamMemberDeleteRequest(team_id=team_id, user_id=user_id), user_api_key_dict=_admin_auth(), @@ -196,11 +234,9 @@ async def test_member_delete_blocked_by_member_add_removes_from_the_fresh_roster try: async with blocker.tx(timeout=timedelta(seconds=30)) as held: - await held.query_raw(_LOCK_SQL, team_id) + lock_key = await _hold_team_lock(held, team_id) task = asyncio.create_task(run_delete()) - await lock_acquired.wait() - await asyncio.sleep(0.2) - assert not task.done(), "member_delete did not wait on the team's advisory lock" + await _await_lock_contention(db, lock_key, task, "member_delete") # member_add wins the race: it adds `other_user` while holding the lock await held.litellm_teamtable.update( @@ -265,10 +301,8 @@ async def test_delete_blocked_by_member_add_sweeps_the_fresh_reference(): blocker = Prisma() await blocker.connect() - lock_acquired = asyncio.Event() async def run_delete(): - lock_acquired.set() return await delete_team( data=DeleteTeamRequest(team_ids=[team_id]), http_request=MagicMock(), @@ -278,11 +312,9 @@ async def test_delete_blocked_by_member_add_sweeps_the_fresh_reference(): try: async with blocker.tx(timeout=timedelta(seconds=30)) as held: - await held.query_raw(_LOCK_SQL, team_id) + lock_key = await _hold_team_lock(held, team_id) task = asyncio.create_task(run_delete()) - await lock_acquired.wait() - await asyncio.sleep(0.3) - assert not task.done(), "delete_team did not wait on the team's advisory lock" + await _await_lock_contention(db, lock_key, task, "delete_team") # member_add wins the race: write the reference while holding the lock await held.litellm_usertable.upsert(