From c7163a80ddfd34da743b06d770d60aba05b292eb Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 9 Sep 2026 10:48:06 -0700 Subject: [PATCH] perf(proxy): collapse per-worker SGR upserts into one statement per flush (#40362) Each proxy worker flushed one Prisma upsert per active (date, category, route) bucket every interval, so the Postgres primary saw workers x routes statements per interval across the deployment. A flush now builds a single multi-row INSERT ... ON CONFLICT DO UPDATE, and with use_redis_transaction_buffer on the workers push snapshots to a Redis list that one lease-holding pod folds and commits, so the whole deployment costs one statement per interval. The leader keeps popping until the list is empty so a deployment wider than the dequeue cap cannot build a backlog, and rows that fail both the commit and the Redis re-queue fall back to the leader's own accumulator instead of being lost. Resolves LIT-7371 Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 1 + litellm/proxy/db/gateway_request_tracking.py | 220 ++++++++-- litellm/proxy/proxy_server.py | 14 +- .../proxy/db/test_gateway_request_tracking.py | 391 +++++++++++++++--- 4 files changed, 537 insertions(+), 89 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 108a914e9c1..f9389d22dea 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -335,6 +335,7 @@ DEFAULT_SSL_CIPHERS: Final = os.getenv( ########### v2 Architecture constants for managing writing updates to the database ########### REDIS_UPDATE_BUFFER_KEY: Final = "litellm_spend_update_buffer" +REDIS_GATEWAY_REQUESTS_BUFFER_KEY: Final = "litellm_gateway_requests_buffer" REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_spend_update_buffer" REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_team_spend_update_buffer" REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_org_spend_update_buffer" diff --git a/litellm/proxy/db/gateway_request_tracking.py b/litellm/proxy/db/gateway_request_tracking.py index bebd74e877c..c9ace68db33 100644 --- a/litellm/proxy/db/gateway_request_tracking.py +++ b/litellm/proxy/db/gateway_request_tracking.py @@ -10,13 +10,28 @@ strings rather than passing the raw path through. Nothing a caller sends can add a key, so the fold and the table it commits to are bounded by (days x routes) however much traffic arrives, and the response path carries no unbounded queue that would block once full. + +A flush commits its whole snapshot as one multi-row ``INSERT ... ON CONFLICT DO +UPDATE`` rather than one upsert per key, so a worker costs the primary one +statement per interval however many routes it served. With +``use_redis_transaction_buffer`` on, workers instead push their snapshot to a +Redis list and one lock-holding pod folds every entry and writes the table, so +the deployment as a whole costs the primary one statement per interval. """ -from dataclasses import asdict +import json +from collections.abc import AsyncIterator, Iterable from datetime import datetime, timezone -from typing import TYPE_CHECKING, Final +from itertools import chain +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, TypeAlias + +from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger +from litellm.caching import RedisCache +from litellm.constants import MAX_REDIS_BUFFER_DEQUEUE_COUNT, REDIS_GATEWAY_REQUESTS_BUFFER_KEY +from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager from litellm.proxy.middleware.billable_request_metrics_middleware import BillableCategory from litellm.types.proxy.gateway_requests import ( GatewayRequestCounts, @@ -28,6 +43,15 @@ if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient _EMPTY: Final = GatewayRequestCounts(successful_requests=0, failed_requests=0) +_TABLE: Final = '"LiteLLM_DailyGatewayRequests"' +_COLUMNS_PER_ROW: Final = 5 +_UTC_NOW: Final = "(NOW() AT TIME ZONE 'UTC')" +GATEWAY_REQUESTS_JOB_NAME: Final = "update_gateway_requests_job" + +_BufferedRows: TypeAlias = tuple[tuple[str, str, str, int, int], ...] +_BUFFERED_ROWS: Final = TypeAdapter(_BufferedRows) +_BUFFERED_ENTRIES: Final = TypeAdapter(tuple[str | bytes, ...]) +_NO_COUNTS: Final[GatewayRequestSnapshot] = MappingProxyType({}) def _utc_date() -> str: @@ -59,20 +83,54 @@ class GatewayRequestAccumulator: route) however long the database is unreachable. This buys at-least-once, not exactly-once, and the cost is worth stating. - The batch commits inside its context manager's ``__aexit__``, so a failure - raised after the transaction committed (a connection dropped while reading - the acknowledgement) restores counts that are already persisted, and the - next flush increments them a second time. Exactly-once would need a dedup - key the upserts could ignore on replay. For a traffic-volume metric a rare + The statement commits on the server before its acknowledgement is read, so + a failure raised after the commit (a connection dropped while reading the + acknowledgement) restores counts that are already persisted, and the next + flush increments them a second time. Exactly-once would need a dedup key + the upsert could ignore on replay. For a traffic-volume metric a rare overcount on a dropped acknowledgement beats losing a whole interval to every database blip, so the trade is deliberate. """ - for key, counts in snapshot.items(): - existing = self._counts.get(key, _EMPTY) - self._counts[key] = GatewayRequestCounts( - successful_requests=existing.successful_requests + counts.successful_requests, - failed_requests=existing.failed_requests + counts.failed_requests, - ) + self._counts = dict(fold_counts(chain(self._counts.items(), snapshot.items()))) # mutable-ok: fold replaced + + +def fold_counts(items: Iterable[tuple[GatewayRequestKey, GatewayRequestCounts]]) -> GatewayRequestSnapshot: + """Sum counts key-wise; the result stays bounded by (date x category x route).""" + folded: Final[dict[GatewayRequestKey, GatewayRequestCounts]] = {} # mutable-ok: local fold returned once + for key, counts in items: + existing = folded.get(key, _EMPTY) + folded[key] = GatewayRequestCounts( + successful_requests=existing.successful_requests + counts.successful_requests, + failed_requests=existing.failed_requests + counts.failed_requests, + ) + return folded + + +def build_gateway_requests_upsert(snapshot: GatewayRequestSnapshot) -> tuple[str, tuple[str | int, ...]]: + """ + One ``INSERT ... ON CONFLICT DO UPDATE`` that increments every (date, category, + route) in the snapshot. Rows are ordered by the conflict key so concurrent + writers lock rows in the same order and cannot deadlock. + """ + ordered: Final = sorted(snapshot.items(), key=lambda item: (item[0].date, item[0].category, item[0].route)) + rows: Final = ", ".join( + f"(${base + 1}::text, ${base + 2}::text, ${base + 3}::text, ${base + 4}::bigint, ${base + 5}::bigint, {_UTC_NOW})" + for base in range(0, len(ordered) * _COLUMNS_PER_ROW, _COLUMNS_PER_ROW) + ) + sql: Final = ( + f'INSERT INTO {_TABLE} ("date", "category", "route", "successful_requests", "failed_requests", "updated_at")\n' + f"VALUES {rows}\n" + 'ON CONFLICT ("date", "category", "route") DO UPDATE SET\n' + f' "successful_requests" = {_TABLE}."successful_requests" + EXCLUDED."successful_requests",\n' + f' "failed_requests" = {_TABLE}."failed_requests" + EXCLUDED."failed_requests",\n' + f' "updated_at" = {_UTC_NOW}' + ) + params: Final[tuple[str | int, ...]] = tuple( + value + for key, counts in ordered + for value in (key.date, key.category, key.route, counts.successful_requests, counts.failed_requests) + ) + return sql, params async def commit_gateway_requests_to_db( @@ -80,50 +138,130 @@ async def commit_gateway_requests_to_db( prisma_client: "PrismaClient", snapshot: GatewayRequestSnapshot, ) -> None: - """Upsert one incrementing row per (date, category, route).""" + """Increment every (date, category, route) in the snapshot with a single statement.""" if not snapshot: return - ordered: Final = sorted(snapshot.items(), key=lambda item: (item[0].date, item[0].category, item[0].route)) + sql, params = build_gateway_requests_upsert(snapshot) + await prisma_client.db.execute_raw(sql, *params) # pyright: ignore[reportAny] # untyped prisma client - # pyright: ignore[reportAny] on both lines -- prisma's generated client is untyped, - # so .db and every table action off it resolve to Any at this boundary. The dict - # literals below are the shape prisma's generated inputs require. - async with prisma_client.db.batch_() as batcher: # pyright: ignore[reportAny] # untyped prisma client - for key, counts in ordered: - columns = asdict(key) - batcher.litellm_dailygatewayrequests.upsert( # pyright: ignore[reportAny] # untyped prisma client - where={"date_category_route": columns}, # mutable-ok: prisma input is dict-shaped - data={ # mutable-ok: prisma input is dict-shaped - "create": { # mutable-ok: prisma input is dict-shaped - **columns, - "successful_requests": counts.successful_requests, - "failed_requests": counts.failed_requests, - }, - "update": { # mutable-ok: prisma input is dict-shaped - "successful_requests": {"increment": counts.successful_requests}, # mutable-ok: as above - "failed_requests": {"increment": counts.failed_requests}, # mutable-ok: as above - }, - }, + verbose_proxy_logger.debug( + "Gateway request tracking - committed %d aggregated rows in one statement", len(snapshot) + ) + + +class GatewayRequestRedisBuffer: + """ + Folds every worker's snapshot through one Redis list so a single pod per + interval writes the table, mirroring the spend writer's transaction buffer. + + Each entry is one worker's snapshot as JSON rows; the lock holder pops them, + sums them, and commits one statement. A commit failure pushes the summed + rows back so the next holder retries, keeping the at-least-once guarantee. + If that push fails too, the rows go back to the holder's own accumulator so + they ride along with its next flush instead of vanishing with the pop. + """ + + def __init__(self, *, redis_cache: RedisCache, pod_lock_manager: PodLockManager) -> None: + self._redis_cache: Final = redis_cache + self._pod_lock_manager: Final = pod_lock_manager + + async def push(self, snapshot: GatewayRequestSnapshot) -> None: + if not snapshot: + return + rows: Final[_BufferedRows] = tuple( + (key.date, key.category, key.route, counts.successful_requests, counts.failed_requests) + for key, counts in snapshot.items() + ) + await self._redis_cache.async_rpush(key=REDIS_GATEWAY_REQUESTS_BUFFER_KEY, values=(json.dumps(rows),)) + + async def _pop_batch(self) -> tuple[str | bytes, ...]: + popped: Final[object] = await self._redis_cache.async_lpop( # pyright: ignore[reportAny] # redis returns Any + key=REDIS_GATEWAY_REQUESTS_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT + ) + if not popped: + return () + return _BUFFERED_ENTRIES.validate_python(popped if isinstance(popped, list) else (popped,)) + + async def _pop_all(self) -> AsyncIterator[str | bytes]: + while True: + batch = await self._pop_batch() + for entry in batch: + yield entry + if len(batch) < MAX_REDIS_BUFFER_DEQUEUE_COUNT: + return + + async def pop(self) -> GatewayRequestSnapshot: + entries: Final = tuple([entry async for entry in self._pop_all()]) + return fold_counts( + ( + GatewayRequestKey(date=date, category=category, route=route), + GatewayRequestCounts(successful_requests=succeeded, failed_requests=failed), ) + for entry in entries + for date, category, route, succeeded, failed in _BUFFERED_ROWS.validate_json(entry) + ) - verbose_proxy_logger.debug("Gateway request tracking - committed %d aggregated rows", len(ordered)) + async def commit_if_leader(self, prisma_client: "PrismaClient") -> GatewayRequestSnapshot: + """ + Drain the list and write it as one statement, but only on the pod holding the job lock. + + The lock is a lease, never released: the holder re-enters it on every flush and + keeps committing alone until the TTL lapses, so the primary sees one statement + per flush interval deployment-wide instead of one per worker. + + Returns the popped rows that could be neither committed nor re-queued, for the + caller to keep in memory. Empty on success. + """ + if not await self._pod_lock_manager.acquire_lock(cronjob_id=GATEWAY_REQUESTS_JOB_NAME): + return _NO_COUNTS + buffered: Final = await self.pop() + try: + await commit_gateway_requests_to_db(prisma_client=prisma_client, snapshot=buffered) + except Exception: # noqa: BLE001 -- a failed commit must not stop the scheduler + verbose_proxy_logger.warning( + "Gateway request tracking - failed to commit %d buffered rows, re-queuing to Redis for the next flush", + len(buffered), + exc_info=True, + ) + return await self._requeue(buffered) + return _NO_COUNTS + + async def _requeue(self, snapshot: GatewayRequestSnapshot) -> GatewayRequestSnapshot: + try: + await self.push(snapshot) + except Exception: # noqa: BLE001 -- the rows go back to the caller's accumulator instead + verbose_proxy_logger.warning( + "Gateway request tracking - Redis re-queue failed, keeping %d rows in memory for the next flush", + len(snapshot), + exc_info=True, + ) + return snapshot + return _NO_COUNTS async def flush_gateway_requests( prisma_client: "PrismaClient", accumulator: GatewayRequestAccumulator, + redis_buffer: GatewayRequestRedisBuffer | None = None, ) -> None: """ Scheduler entrypoint. Never raises: a metering failure must not kill the job. + With ``redis_buffer`` the snapshot goes to Redis and only the lease holder + writes to Postgres. Shutdown passes no buffer so a departing worker writes its + own counts directly instead of parking them behind a lease it may not hold. + ``CancelledError`` is deliberately not caught, so a flush cancelled during shutdown drops its snapshot rather than restoring counts onto an accumulator the process is about to discard. """ snapshot: Final = accumulator.drain() try: - await commit_gateway_requests_to_db(prisma_client=prisma_client, snapshot=snapshot) + if redis_buffer is None: + await commit_gateway_requests_to_db(prisma_client=prisma_client, snapshot=snapshot) + else: + await redis_buffer.push(snapshot) except Exception: # noqa: BLE001 -- a failed flush must not stop the scheduler accumulator.restore(snapshot) verbose_proxy_logger.warning( @@ -131,3 +269,13 @@ async def flush_gateway_requests( len(snapshot), exc_info=True, ) + return + if redis_buffer is None: + return + try: + accumulator.restore(await redis_buffer.commit_if_leader(prisma_client)) + except Exception: # noqa: BLE001 -- entries still in Redis are drained by the next flush + verbose_proxy_logger.warning( + "Gateway request tracking - leader drain failed, buffered rows stay in Redis for the next flush", + exc_info=True, + ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c241e66049b..09e43eb74e1 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -424,6 +424,7 @@ from litellm.proxy.db.exception_handler import ( ) from litellm.proxy.db.gateway_request_tracking import ( GatewayRequestAccumulator, + GatewayRequestRedisBuffer, flush_gateway_requests, ) from litellm.proxy.db.proxy_worker_heartbeat import ( @@ -2355,6 +2356,17 @@ open_telemetry_logger: OpenTelemetry | None = None gateway_request_accumulator: Final = GatewayRequestAccumulator() ### INITIALIZE GLOBAL LOGGING OBJECT ### proxy_logging_obj: ProxyLogging = ProxyLogging(user_api_key_cache=user_api_key_cache, premium_user=premium_user) + + +def _gateway_request_redis_buffer() -> GatewayRequestRedisBuffer | None: + """Shares the spend writer's transaction-buffer Redis and pod lock when use_redis_transaction_buffer is on.""" + writer: Final = proxy_logging_obj.db_spend_update_writer + redis_cache: Final = writer.redis_update_buffer.redis_cache + if redis_cache is None or not writer.redis_update_buffer._should_commit_spend_updates_to_redis(): + return None + return GatewayRequestRedisBuffer(redis_cache=redis_cache, pod_lock_manager=writer.pod_lock_manager) + + ### REDIS QUEUE ### async_result: Final = None celery_app_conn: Final = None @@ -9633,7 +9645,7 @@ class ProxyStartupEvent: flush_gateway_requests, "interval", seconds=batch_writing_interval, - args=(prisma_client, gateway_request_accumulator), + args=(prisma_client, gateway_request_accumulator, _gateway_request_redis_buffer()), id="update_gateway_requests_job", replace_existing=True, misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, diff --git a/tests/test_litellm/proxy/db/test_gateway_request_tracking.py b/tests/test_litellm/proxy/db/test_gateway_request_tracking.py index 93a11a914cb..045261e2d53 100644 --- a/tests/test_litellm/proxy/db/test_gateway_request_tracking.py +++ b/tests/test_litellm/proxy/db/test_gateway_request_tracking.py @@ -8,8 +8,11 @@ from datetime import datetime, timezone import pytest +from litellm.constants import MAX_REDIS_BUFFER_DEQUEUE_COUNT, REDIS_GATEWAY_REQUESTS_BUFFER_KEY from litellm.proxy.db.gateway_request_tracking import ( + GATEWAY_REQUESTS_JOB_NAME, GatewayRequestAccumulator, + GatewayRequestRedisBuffer, commit_gateway_requests_to_db, flush_gateway_requests, ) @@ -83,40 +86,47 @@ def test_drain_snapshot_is_not_mutated_by_later_records(): # ── commit ──────────────────────────────────────────────────────────────────── -class FakeTable: - def __init__(self) -> None: - self.upserts: list[dict] = [] - - def upsert(self, *, where: dict, data: dict) -> None: - self.upserts.append({"where": where, "data": data}) - - -class FakeBatcher: - def __init__(self, table: FakeTable) -> None: - self.litellm_dailygatewayrequests = table - - async def __aenter__(self) -> "FakeBatcher": - return self - - async def __aexit__(self, *args: object) -> bool: - return False - - class FakeDB: - def __init__(self, table: FakeTable) -> None: - self._table = table + def __init__(self) -> None: + self.statements: list[tuple[str, tuple[object, ...]]] = [] - def batch_(self) -> FakeBatcher: - return FakeBatcher(self._table) + async def execute_raw(self, query: str, *args: object) -> int: + self.statements.append((query, args)) + return len(args) // 5 class FakePrismaClient: def __init__(self) -> None: - self.table = FakeTable() - self.db = FakeDB(self.table) + self.db = FakeDB() -def test_commit_upserts_one_incrementing_row_per_key(): +def _rows_written(client: FakePrismaClient) -> list[tuple[object, ...]]: + """Every (date, category, route, successful, failed) tuple the database received, in statement order.""" + return [params[i : i + 5] for _, params in client.db.statements for i in range(0, len(params), 5)] + + +def test_commit_increments_with_a_single_statement_for_the_whole_snapshot(): + """One statement per flush is the whole point: the previous per-key upsert cost + the primary (workers x routes) statements per interval.""" + client = FakePrismaClient() + snapshot = { + GatewayRequestKey(date="2026-08-01", category="llm", route=route): ( + GatewayRequestCounts(successful_requests=7, failed_requests=2) + ) + for route in ("/chat/completions", "/embeddings", "/responses", "/v1/messages", "/mcp") + } + + asyncio.run(commit_gateway_requests_to_db(prisma_client=client, snapshot=snapshot)) + + assert len(client.db.statements) == 1 + sql, params = client.db.statements[0] + assert sql.count("ON CONFLICT") == 1 + assert sql.count("(NOW() AT TIME ZONE 'UTC'))") == 5 + assert len(params) == 25 + + +def test_commit_sql_adds_to_the_existing_row_instead_of_replacing_it(): + """A worker only knows its own share; the SQL must add EXCLUDED onto the stored count.""" client = FakePrismaClient() snapshot = { GatewayRequestKey(date="2026-08-01", category="llm", route="/chat/completions"): ( @@ -126,20 +136,36 @@ def test_commit_upserts_one_incrementing_row_per_key(): asyncio.run(commit_gateway_requests_to_db(prisma_client=client, snapshot=snapshot)) - assert len(client.table.upserts) == 1 - written = client.table.upserts[0] - assert written["where"] == { - "date_category_route": { - "date": "2026-08-01", - "category": "llm", - "route": "/chat/completions", - } + sql, params = client.db.statements[0] + assert 'INSERT INTO "LiteLLM_DailyGatewayRequests"' in sql + assert 'ON CONFLICT ("date", "category", "route") DO UPDATE SET' in sql + assert ( + '"successful_requests" = "LiteLLM_DailyGatewayRequests"."successful_requests" + EXCLUDED."successful_requests"' + in sql + ) + assert '"failed_requests" = "LiteLLM_DailyGatewayRequests"."failed_requests" + EXCLUDED."failed_requests"' in sql + assert params == ("2026-08-01", "llm", "/chat/completions", 7, 2) + + +def test_commit_placeholders_line_up_with_params(): + """$n positions are generated per row; a drift here silently swaps a route for a count.""" + client = FakePrismaClient() + snapshot = { + GatewayRequestKey(date="2026-08-01", category="llm", route="/chat/completions"): ( + GatewayRequestCounts(successful_requests=1, failed_requests=0) + ), + GatewayRequestKey(date="2026-08-01", category="mcp", route="/mcp"): ( + GatewayRequestCounts(successful_requests=0, failed_requests=3) + ), } - assert written["data"]["update"] == { - "successful_requests": {"increment": 7}, - "failed_requests": {"increment": 2}, - } - assert written["data"]["create"]["successful_requests"] == 7 + + asyncio.run(commit_gateway_requests_to_db(prisma_client=client, snapshot=snapshot)) + + sql, params = client.db.statements[0] + assert "($1::text, $2::text, $3::text, $4::bigint, $5::bigint," in sql + assert "($6::text, $7::text, $8::text, $9::bigint, $10::bigint," in sql + assert "$11" not in sql + assert params == ("2026-08-01", "llm", "/chat/completions", 1, 0, "2026-08-01", "mcp", "/mcp", 0, 3) def test_commit_is_deterministically_ordered(): @@ -154,17 +180,14 @@ def test_commit_is_deterministically_ordered(): asyncio.run(commit_gateway_requests_to_db(prisma_client=client, snapshot=snapshot)) - written_order = [ - (row["where"]["date_category_route"]["date"], row["where"]["date_category_route"]["category"]) - for row in client.table.upserts - ] + written_order = [(row[0], row[1]) for row in _rows_written(client)] assert written_order == [("2026-08-01", "llm"), ("2026-08-01", "mcp"), ("2026-08-02", "llm")] def test_commit_skips_the_database_entirely_when_nothing_accumulated(): client = FakePrismaClient() asyncio.run(commit_gateway_requests_to_db(prisma_client=client, snapshot={})) - assert client.table.upserts == [] + assert client.db.statements == [] # ── flush ───────────────────────────────────────────────────────────────────── @@ -177,12 +200,12 @@ def test_flush_drains_and_commits(): asyncio.run(flush_gateway_requests(client, acc)) - assert len(client.table.upserts) == 1 + assert len(client.db.statements) == 1 assert acc.drain() == {} class ExplodingDB: - def batch_(self): + async def execute_raw(self, query: str, *args: object) -> int: raise RuntimeError("db gone") @@ -208,10 +231,7 @@ def test_failed_flush_keeps_counts_for_the_next_attempt(): client = FakePrismaClient() asyncio.run(flush_gateway_requests(client, acc)) - assert client.table.upserts[0]["data"]["update"] == { - "successful_requests": {"increment": 1}, - "failed_requests": {"increment": 1}, - } + assert _rows_written(client) == [(_today(), "llm", "/chat/completions", 1, 1)] def test_restored_counts_merge_with_requests_recorded_meanwhile(): @@ -223,5 +243,272 @@ def test_restored_counts_merge_with_requests_recorded_meanwhile(): client = FakePrismaClient() asyncio.run(flush_gateway_requests(client, acc)) - assert len(client.table.upserts) == 1 - assert client.table.upserts[0]["data"]["update"]["successful_requests"] == {"increment": 2} + assert _rows_written(client) == [(_today(), "llm", "/chat/completions", 2, 0)] + + +class ExplodingDBWithInFlightRequest: + """Fails the write after a request has been recorded while it was in flight.""" + + def __init__(self, accumulator: GatewayRequestAccumulator) -> None: + self.accumulator = accumulator + + async def execute_raw(self, query: str, *args: object) -> int: + _record(self.accumulator, 500) + raise RuntimeError("db gone") + + +class ExplodingClientWithInFlightRequest: + def __init__(self, accumulator: GatewayRequestAccumulator) -> None: + self.db = ExplodingDBWithInFlightRequest(accumulator) + + +def test_restore_keeps_requests_recorded_while_the_failed_write_was_in_flight(): + acc = GatewayRequestAccumulator() + _record(acc, 200) + asyncio.run(flush_gateway_requests(ExplodingClientWithInFlightRequest(acc), acc)) + + client = FakePrismaClient() + asyncio.run(flush_gateway_requests(client, acc)) + + assert _rows_written(client) == [(_today(), "llm", "/chat/completions", 1, 1)] + + +# ── redis buffer ────────────────────────────────────────────────────────────── + + +class FakeRedis: + def __init__(self) -> None: + self.lists: dict[str, list[str]] = {} + + async def async_rpush(self, key: str, values: list[str]) -> int: + self.lists.setdefault(key, []).extend(values) + return len(self.lists[key]) + + async def async_lpop(self, key: str, count: int) -> list[str] | None: + queue = self.lists.get(key, []) + if not queue: + return None + popped, self.lists[key] = queue[:count], queue[count:] + return popped + + +class FakePodLock: + def __init__(self, *, leader: bool) -> None: + self.leader = leader + self.held: list[str] = [] + self.released: list[str] = [] + + async def acquire_lock(self, cronjob_id: str) -> bool: + self.held.append(cronjob_id) + return self.leader + + async def release_lock(self, cronjob_id: str) -> None: + self.released.append(cronjob_id) + + +class FakeLease: + """Redis-side view of the job lock: SET NX by pod id, re-entrant for the holder, freed only by release or TTL.""" + + def __init__(self) -> None: + self.holder: str | None = None + + +class FakeLeasePodLock: + def __init__(self, lease: FakeLease, pod_id: str) -> None: + self.lease = lease + self.pod_id = pod_id + + async def acquire_lock(self, cronjob_id: str) -> bool: + if self.lease.holder is None: + self.lease.holder = self.pod_id + return self.lease.holder == self.pod_id + + async def release_lock(self, cronjob_id: str) -> None: + if self.lease.holder == self.pod_id: + self.lease.holder = None + + +def _buffer(redis: FakeRedis, *, leader: bool) -> tuple[GatewayRequestRedisBuffer, FakePodLock]: + lock = FakePodLock(leader=leader) + return GatewayRequestRedisBuffer(redis_cache=redis, pod_lock_manager=lock), lock # pyright: ignore[reportArgumentType] # duck-typed fakes + + +def test_non_leader_workers_push_to_redis_and_never_touch_the_database(): + redis = FakeRedis() + client = FakePrismaClient() + for _ in range(3): + acc = GatewayRequestAccumulator() + _record(acc, 200) + buffer, _ = _buffer(redis, leader=False) + asyncio.run(flush_gateway_requests(client, acc, buffer)) + + assert client.db.statements == [] + assert len(redis.lists[REDIS_GATEWAY_REQUESTS_BUFFER_KEY]) == 3 + + +def test_leader_folds_every_workers_snapshot_into_one_statement(): + """Fifty workers each flushing the same routes must cost the primary one statement, not fifty.""" + redis = FakeRedis() + client = FakePrismaClient() + for _ in range(50): + acc = GatewayRequestAccumulator() + _record(acc, 200) + _record(acc, 500, route="/responses") + buffer, _ = _buffer(redis, leader=False) + asyncio.run(flush_gateway_requests(client, acc, buffer)) + + leader_acc = GatewayRequestAccumulator() + _record(leader_acc, 200) + leader, lock = _buffer(redis, leader=True) + asyncio.run(flush_gateway_requests(client, leader_acc, leader)) + + assert len(client.db.statements) == 1 + assert _rows_written(client) == [ + (_today(), "llm", "/chat/completions", 51, 0), + (_today(), "llm", "/responses", 0, 50), + ] + assert redis.lists[REDIS_GATEWAY_REQUESTS_BUFFER_KEY] == [] + assert lock.held == [GATEWAY_REQUESTS_JOB_NAME] + assert lock.released == [] + + +def test_leader_keeps_the_lease_so_staggered_pods_cost_one_statement_per_interval(): + """Pods flush on their own clocks; without the lease each one would win the lock in turn and commit alone.""" + redis = FakeRedis() + client = FakePrismaClient() + lease = FakeLease() + pods = tuple( + GatewayRequestRedisBuffer(redis_cache=redis, pod_lock_manager=FakeLeasePodLock(lease, f"pod-{i}")) # pyright: ignore[reportArgumentType] # duck-typed fakes + for i in range(4) + ) + + for _interval in range(3): + for pod in pods: + acc = GatewayRequestAccumulator() + _record(acc, 200) + asyncio.run(flush_gateway_requests(client, acc, pod)) + + assert lease.holder == "pod-0" + assert len(client.db.statements) == 3 + assert [row[3] for row in _rows_written(client)] == [1, 4, 4] + assert len(redis.lists[REDIS_GATEWAY_REQUESTS_BUFFER_KEY]) == 3 + + +def test_leader_drains_a_backlog_deeper_than_one_capped_pop(): + """More workers than MAX_REDIS_BUFFER_DEQUEUE_COUNT must not leave a growing tail queued behind the cap.""" + redis = FakeRedis() + client = FakePrismaClient() + workers = MAX_REDIS_BUFFER_DEQUEUE_COUNT * 2 + 1 + for _ in range(workers): + acc = GatewayRequestAccumulator() + _record(acc, 200) + buffer, _ = _buffer(redis, leader=False) + asyncio.run(flush_gateway_requests(client, acc, buffer)) + + leader, _ = _buffer(redis, leader=True) + asyncio.run(flush_gateway_requests(client, GatewayRequestAccumulator(), leader)) + + assert len(client.db.statements) == 1 + assert _rows_written(client) == [(_today(), "llm", "/chat/completions", workers, 0)] + assert redis.lists[REDIS_GATEWAY_REQUESTS_BUFFER_KEY] == [] + + +def test_leader_with_nothing_buffered_writes_nothing(): + redis = FakeRedis() + client = FakePrismaClient() + leader, lock = _buffer(redis, leader=True) + + asyncio.run(flush_gateway_requests(client, GatewayRequestAccumulator(), leader)) + + assert client.db.statements == [] + assert lock.released == [] + + +def test_leader_requeues_to_redis_when_the_database_commit_fails(): + """Counts popped from Redis are gone from every worker; a failed commit must put them back.""" + redis = FakeRedis() + acc = GatewayRequestAccumulator() + _record(acc, 200) + _record(acc, 200) + leader, lock = _buffer(redis, leader=True) + + asyncio.run(flush_gateway_requests(ExplodingClient(), acc, leader)) + + assert len(redis.lists[REDIS_GATEWAY_REQUESTS_BUFFER_KEY]) == 1 + assert lock.released == [] + assert acc.drain() == {} + + client = FakePrismaClient() + retry, _ = _buffer(redis, leader=True) + asyncio.run(flush_gateway_requests(client, GatewayRequestAccumulator(), retry)) + assert _rows_written(client) == [(_today(), "llm", "/chat/completions", 2, 0)] + + +class ExplodingRedis(FakeRedis): + async def async_rpush(self, key: str, values: list[str]) -> int: + raise RuntimeError("redis gone") + + +class UnreadableRedis(FakeRedis): + async def async_lpop(self, key: str, count: int) -> list[str] | None: + raise RuntimeError("redis gone mid-flush") + + +class UnwritableRedis(FakeRedis): + """Pops succeed, pushes fail: a Redis that went read-only between the leader's pop and its re-queue.""" + + async def async_rpush(self, key: str, values: list[str]) -> int: + raise RuntimeError("redis read-only") + + +def test_leader_keeps_popped_counts_in_memory_when_both_the_database_and_the_requeue_fail(): + """The pop removed the only copy; if Redis will not take it back the leader itself must carry it.""" + redis = FakeRedis() + worker_acc = GatewayRequestAccumulator() + _record(worker_acc, 200) + _record(worker_acc, 200) + worker, _ = _buffer(redis, leader=False) + asyncio.run(flush_gateway_requests(FakePrismaClient(), worker_acc, worker)) + + degraded = UnwritableRedis() + degraded.lists = redis.lists + leader_acc = GatewayRequestAccumulator() + leader, _ = _buffer(degraded, leader=True) + asyncio.run(flush_gateway_requests(ExplodingClient(), leader_acc, leader)) + assert degraded.lists[REDIS_GATEWAY_REQUESTS_BUFFER_KEY] == [] + + client = FakePrismaClient() + retry, _ = _buffer(redis, leader=True) + asyncio.run(flush_gateway_requests(client, leader_acc, retry)) + assert _rows_written(client) == [(_today(), "llm", "/chat/completions", 2, 0)] + + +def test_leader_whose_redis_read_fails_leaves_the_pushed_rows_for_the_next_flush(): + """The scheduler job must not raise, and nothing is popped so nothing needs restoring anywhere.""" + redis = UnreadableRedis() + acc = GatewayRequestAccumulator() + _record(acc, 200) + client = FakePrismaClient() + leader, _ = _buffer(redis, leader=True) + + asyncio.run(flush_gateway_requests(client, acc, leader)) + + assert client.db.statements == [] + assert acc.drain() == {} + assert len(redis.lists[REDIS_GATEWAY_REQUESTS_BUFFER_KEY]) == 1 + + +def test_failed_redis_push_keeps_counts_locally_for_the_next_flush(): + acc = GatewayRequestAccumulator() + _record(acc, 200) + _record(acc, 500) + buffer, lock = _buffer(ExplodingRedis(), leader=True) + + asyncio.run(flush_gateway_requests(FakePrismaClient(), acc, buffer)) + + assert lock.held == [] + assert acc.drain() == { + GatewayRequestKey(date=_today(), category="llm", route="/chat/completions"): ( + GatewayRequestCounts(successful_requests=1, failed_requests=1) + ) + }