From 499c334fce491fa67fc1db92e5e1c160cd00b33d Mon Sep 17 00:00:00 2001 From: ryan Date: Wed, 16 Sep 2026 03:35:34 +0000 Subject: [PATCH] fix(proxy): lock each team in sorted order before the member spend upsert Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/db/db_spend_update_writer.py | 12 +++++------ .../proxy/db/test_db_spend_update_writer.py | 21 +++++++++++-------- 2 files changed, 17 insertions(+), 16 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index f617302d8ad..09ce6b6f451 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -164,12 +164,9 @@ def _spend_update_tx(prisma_client: PrismaClient) -> _SpendTransactionManager: # The per-team advisory lock the team endpoints hold while changing a roster (TEAM_ADVISORY_LOCK_SQL), -# taken in sorted order so the roster check below cannot interleave with their writes. A row lock would -# deadlock with the access-group endpoints, which lock a team row after an access-group lock. -_TEAM_ADVISORY_LOCKS_SQL: Final = """ -SELECT pg_advisory_xact_lock(hashtext(teams.team_id)) -FROM (SELECT DISTINCT team_id FROM unnest($1::text[]) AS team_id ORDER BY team_id) AS teams -""" +# so the roster check below cannot interleave with their writes. A row lock would deadlock with the +# access-group endpoints, which lock a team row after an access-group lock. +_TEAM_ADVISORY_LOCK_SQL: Final = "SELECT pg_advisory_xact_lock(hashtext($1)) IS NULL AS locked" # One statement adds every member's cost to their membership row. A missing row is created only # while the user is still on the team's roster, so a spend flush landing after a removal never @@ -195,7 +192,8 @@ async def _write_team_member_spend(transaction: _SpendTransaction, spend_by_memb # keeping lock order consistent across pods to prevent deadlocks keys: Final = tuple(sorted(spend_by_member_key)) team_ids: Final = [key.split("::")[1] for key in keys] - _ = await transaction.execute_raw(_TEAM_ADVISORY_LOCKS_SQL, team_ids) + for team_id in dict.fromkeys(team_ids): + _ = await transaction.execute_raw(_TEAM_ADVISORY_LOCK_SQL, team_id) _ = await transaction.execute_raw( _TEAM_MEMBER_SPEND_SQL, [key.split("::")[3] for key in keys], diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 9f3bc5acd92..2ac9fc8eccd 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -17,7 +17,7 @@ from redis.exceptions import DataError import litellm from litellm.proxy._types import Litellm_EntityType from litellm.proxy.db.db_spend_update_writer import ( - _TEAM_ADVISORY_LOCKS_SQL, + _TEAM_ADVISORY_LOCK_SQL, _TEAM_MEMBER_SPEND_SQL, DBSpendUpdateWriter, ) @@ -973,11 +973,10 @@ async def test_commit_spend_updates_to_db_writes_team_member_spend_in_one_roster ) lock_call, spend_call = mock_transaction.execute_raw.await_args_list - lock_statement, locked_team_ids = lock_call.args - assert lock_statement is _TEAM_ADVISORY_LOCKS_SQL - assert locked_team_ids == [team_id] - assert "pg_advisory_xact_lock(hashtext(teams.team_id))" in lock_statement - assert "ORDER BY team_id" in lock_statement + lock_statement, locked_team_id = lock_call.args + assert lock_statement is _TEAM_ADVISORY_LOCK_SQL + assert locked_team_id == team_id + assert "pg_advisory_xact_lock(hashtext($1))" in lock_statement statement, user_ids, team_ids, costs = spend_call.args assert statement is _TEAM_MEMBER_SPEND_SQL assert (user_ids, team_ids, costs) == ([user_id], [team_id], [response_cost]) @@ -993,7 +992,7 @@ async def test_commit_spend_updates_to_db_orders_team_member_rows_by_team_then_u """ The member spend statement touches rows in the order of its input arrays, so the batch is handed over sorted by (team_id, user_id), with each cost kept next to its member, and - the advisory lock statement receives the same team ids, so concurrent pods lock in the + each distinct team is locked once, in that same order, so concurrent pods lock in the same order and cannot deadlock. """ db_writer = DBSpendUpdateWriter() @@ -1016,9 +1015,13 @@ async def test_commit_spend_updates_to_db_orders_team_member_rows_by_team_then_u ), ) - lock_call, spend_call = mock_transaction.execute_raw.await_args_list + *lock_calls, spend_call = mock_transaction.execute_raw.await_args_list _statement, user_ids, team_ids, costs = spend_call.args - assert lock_call.args == (_TEAM_ADVISORY_LOCKS_SQL, ["team_a", "team_a", "team_b", "team_c"]) + assert [lock_call.args for lock_call in lock_calls] == [ + (_TEAM_ADVISORY_LOCK_SQL, "team_a"), + (_TEAM_ADVISORY_LOCK_SQL, "team_b"), + (_TEAM_ADVISORY_LOCK_SQL, "team_c"), + ] assert list(zip(team_ids, user_ids, costs)) == [ ("team_a", "user_x", 0.3), ("team_a", "user_y", 0.2),