diff --git a/litellm/proxy/db/db_pool_metrics.py b/litellm/proxy/db/db_pool_metrics.py index 20161614582..256a3dec2a0 100644 --- a/litellm/proxy/db/db_pool_metrics.py +++ b/litellm/proxy/db/db_pool_metrics.py @@ -117,22 +117,36 @@ class DBPoolMetricsSampler: self._pending_baseline: float = 0.0 def is_due(self) -> bool: - """Whether enough time has passed to justify reading the engine again.""" + """Whether a sample is due. Read-only; does not claim.""" if self._last_sampled_at is None: return True return (self._monotonic() - self._last_sampled_at) >= self._min_interval_seconds - async def maybe_sample(self, resolve_client: Callable[[], SupportsPoolSample | None]) -> DBPoolMetricsUpdate | None: - """Read the engine's counters, or ``None`` if not yet due. + def try_claim(self) -> bool: + """Claim the next sample, or return False if one is not due yet. - The interval is consumed before the client is resolved, so a deployment - with nothing to sample throttles the same as one that does. Never raises - and never blocks for long; a pool sample is diagnostic. + Check and set happen with no await between them, so a burst of callers + on the event loop produces exactly one claim and therefore one sample + rather than one per caller. """ if not self.is_due(): - return None + return False self._last_sampled_at = self._monotonic() + return True + async def maybe_sample(self, resolve_client: Callable[[], SupportsPoolSample | None]) -> DBPoolMetricsUpdate | None: + """Claim a sample and take it, or return ``None`` if one is not due.""" + if not self.try_claim(): + return None + return await self.sample(resolve_client) + + async def sample(self, resolve_client: Callable[[], SupportsPoolSample | None]) -> DBPoolMetricsUpdate | None: + """Take a sample the caller has already claimed. + + Never raises and never blocks for long; a pool sample is diagnostic, and + an engine that cannot answer one is already the subject of a louder + alarm. + """ try: client: Final = resolve_client() if client is None: diff --git a/litellm/proxy/db/log_db_metrics.py b/litellm/proxy/db/log_db_metrics.py index 3074f435d43..96bd1d3fa18 100644 --- a/litellm/proxy/db/log_db_metrics.py +++ b/litellm/proxy/db/log_db_metrics.py @@ -57,7 +57,7 @@ async def _sample_db_pool_metrics() -> None: succeeded, so nothing here may turn a working query into a failed one. """ try: - update: Final = await _pool_metrics_sampler.maybe_sample(_resolve_pool_client) + update: Final = await _pool_metrics_sampler.sample(_resolve_pool_client) if update is None: return logger: Final = _prometheus_logger() @@ -130,8 +130,9 @@ def log_db_metrics(func): # Dispatched, never awaited. Awaiting here would add a suspension # point after the query already succeeded, so a client disconnect in # that window would discard a completed result and skip the success - # hook below. `is_due` keeps this to one task per sample interval. - if _pool_metrics_sampler.is_due(): + # hook below. The claim is atomic, so a burst spawns one task rather + # than one per caller. + if _pool_metrics_sampler.try_claim(): asyncio.create_task(_sample_db_pool_metrics()) if "PROXY" not in func.__name__: diff --git a/tests/test_litellm/proxy/db/test_log_db_metrics.py b/tests/test_litellm/proxy/db/test_log_db_metrics.py index 2fcf6815e9c..7002d44cf8a 100644 --- a/tests/test_litellm/proxy/db/test_log_db_metrics.py +++ b/tests/test_litellm/proxy/db/test_log_db_metrics.py @@ -185,4 +185,79 @@ async def test_a_proxy_without_prometheus_never_reads_the_query_engine(): await _sample_db_pool_metrics() assert reads == [] - assert _pool_metrics_sampler.is_due() is False, "the interval must still be consumed" + + +@pytest.mark.asyncio +async def test_the_decorated_path_actually_samples_end_to_end(): + """The decorator claims the interval and the dispatched task takes the + sample. If both claimed, the task's claim would fail and no sample would + ever be taken, while every throttle test kept passing.""" + import asyncio + + from litellm.proxy.db.db_pool_metrics import DBPoolSample + from litellm.proxy.db.log_db_metrics import _pool_metrics_sampler, log_db_metrics + + sample = DBPoolSample( + busy_connections=2.0, + idle_connections=1.0, + open_connections=3.0, + pending_acquirers=0.0, + acquire_wait_seconds_total=0.0, + acquire_count_total=0, + query_duration_seconds_total=0.0, + query_count_total=0, + ) + reads = [] + + async def get_pool_sample(): + reads.append(1) + return sample + + prisma_client = MagicMock() + prisma_client.db.get_pool_sample = get_pool_sample + logger = MagicMock() + proxy_logging_obj = MagicMock() + proxy_logging_obj.service_logging_obj.async_service_success_hook = AsyncMock() + + @log_db_metrics + async def fake_db_call(**kwargs): + return "rows" + + _pool_metrics_sampler._last_sampled_at = None + with ( + patch("litellm.proxy.db.log_db_metrics._prometheus_logger", return_value=logger), + patch("litellm.proxy.proxy_server.prisma_client", prisma_client), + patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_obj), + ): + assert await fake_db_call(table_name="x") == "rows" + for _ in range(50): + await asyncio.sleep(0.01) + if reads: + break + + assert reads == [1], f"the dispatched task must take exactly one sample, got {len(reads)}" + logger.record_db_pool_sample.assert_called_once() + + +@pytest.mark.asyncio +async def test_a_burst_of_database_calls_claims_only_one_sample(): + """The decorator gates task creation on the claim, so a burst of concurrent + database calls must produce one sample rather than one per caller. + + Coroutines only interleave at await points, so this holds as long as the + claim stays synchronous. Introducing an await between the due check and the + timestamp write is what would break it. + """ + import asyncio + + from litellm.proxy.db.db_pool_metrics import DBPoolMetricsSampler + + sampler = DBPoolMetricsSampler(min_interval_seconds=10.0) + + async def caller() -> bool: + await asyncio.sleep(0) # force a real scheduling point before claiming + return sampler.try_claim() + + claims = await asyncio.gather(*[caller() for _ in range(50)]) + + assert sum(claims) == 1, f"exactly one caller may claim the interval, got {sum(claims)}"