perf(observability): claim the pool sample atomically so a burst spawns one task

Checking whether a sample was due and consuming that interval were separate
steps, so every caller in a concurrent burst of database calls saw the same
due-ness and dispatched its own task. Only one did real work, but the rest
still allocated a task on the auth hot path.

try_claim does both in one synchronous step, and the dispatched task now takes
the sample it was already granted instead of claiming again, which would have
failed and left no sample taken at all.

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:

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__:

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)}"