mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-15 04:19:08 +00:00
fix(caching): make Redis Cluster commands wait for a pool slot instead of tripping the circuit breaker (#46488)
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
f95e466bb7
commit
ca2f99edac
6 changed files with 414 additions and 54 deletions
|
|
@ -412,12 +412,19 @@ def _redis_timeout_error_types() -> tuple[type, ...]:
|
|||
``builtins.TimeoutError`` covers ``socket.timeout`` (an alias since py3.10) and, from
|
||||
py3.11, ``asyncio.TimeoutError``; on py3.10 ``asyncio.TimeoutError`` is still its own
|
||||
class, so it is listed explicitly. ``redis.exceptions.TimeoutError`` subclasses neither.
|
||||
``MaxConnectionsError`` is the async cluster client's pool-exhausted signal (its nodes have
|
||||
no blocking pool, so it is raised instantly instead of after a wait): a busy pool, not a
|
||||
dead Redis, like the standalone pool's chained wait timeout.
|
||||
"""
|
||||
try:
|
||||
from redis.exceptions import TimeoutError as RedisTimeoutError
|
||||
from redis import exceptions as redis_exceptions
|
||||
except ImportError:
|
||||
return (TimeoutError, asyncio.TimeoutError)
|
||||
return (RedisTimeoutError, TimeoutError, asyncio.TimeoutError)
|
||||
max_connections_error: Final = cast( # cast-ok: types-redis 4.6 stubs predate MaxConnectionsError
|
||||
"type[Exception]",
|
||||
redis_exceptions.MaxConnectionsError, # pyright: ignore[reportAttributeAccessIssue] # types-redis 4.6 stubs predate it
|
||||
)
|
||||
return (redis_exceptions.TimeoutError, TimeoutError, asyncio.TimeoutError, max_connections_error)
|
||||
|
||||
|
||||
_MAX_EXCEPTION_CAUSE_DEPTH: Final = 20
|
||||
|
|
@ -604,9 +611,14 @@ def _redis_circuit_breaker_guard(method):
|
|||
"""
|
||||
|
||||
@functools.wraps(method)
|
||||
async def wrapper(self, *args, **kwargs):
|
||||
async def wrapper(self: "RedisCache", *args, **kwargs):
|
||||
async def guarded() -> object:
|
||||
return await method(self, *args, **kwargs)
|
||||
|
||||
return await _run_under_circuit_breaker(
|
||||
self._circuit_breaker, method.__name__, lambda: method(self, *args, **kwargs)
|
||||
self._circuit_breaker, # pyright: ignore[reportPrivateUsage] # the guard is RedisCache's own decorator
|
||||
wrapper.__name__,
|
||||
lambda: self._admit_async_call(guarded), # pyright: ignore[reportPrivateUsage] # same decorator
|
||||
)
|
||||
|
||||
return wrapper
|
||||
|
|
@ -690,6 +702,9 @@ class RedisCache(BaseCache):
|
|||
else:
|
||||
super().__init__() # defaults to 60s
|
||||
|
||||
def _admit_async_call(self, call: Callable[[], Awaitable[_RedisCallResult]]) -> Awaitable[_RedisCallResult]:
|
||||
return call()
|
||||
|
||||
def _setup_health_pings(self):
|
||||
"""Setup async and sync health pings for Redis."""
|
||||
# ASYNC HEALTH PING
|
||||
|
|
@ -1070,7 +1085,9 @@ class RedisCache(BaseCache):
|
|||
litellm.in_memory_llm_clients_cache.set_cache(key=script_cache_key, value=executor)
|
||||
return await executor(keys=keys, args=args, client=client)
|
||||
|
||||
return await _run_under_circuit_breaker(self._circuit_breaker, "run_script", execute)
|
||||
return await _run_under_circuit_breaker(
|
||||
self._circuit_breaker, "run_script", lambda: self._admit_async_call(execute)
|
||||
)
|
||||
|
||||
return run_script
|
||||
|
||||
|
|
|
|||
|
|
@ -5,9 +5,17 @@ Key differences:
|
|||
- RedisClient NEEDs to be re-used across requests, adds 3000ms latency if it's re-created
|
||||
"""
|
||||
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
import asyncio
|
||||
import contextvars
|
||||
import os
|
||||
import weakref
|
||||
from collections.abc import Awaitable, Callable
|
||||
from dataclasses import dataclass, field
|
||||
from typing import TYPE_CHECKING, Any, Final, TypeAlias, TypeVar
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.caching.redis_cache import RedisCache
|
||||
from litellm.constants import REDIS_CONNECTION_POOL_TIMEOUT
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from opentelemetry.trace import Span as _Span
|
||||
|
|
@ -23,9 +31,100 @@ else:
|
|||
Span = Any
|
||||
|
||||
|
||||
_T: Final = TypeVar("_T")
|
||||
|
||||
|
||||
def _connection_limit(max_connections: object) -> int | None:
|
||||
try:
|
||||
limit: Final = int(max_connections) # pyright: ignore[reportArgumentType] # config value, str or int
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
return limit if limit > 0 else None
|
||||
|
||||
|
||||
_ClientSemaphores: TypeAlias = weakref.WeakValueDictionary[str, asyncio.Semaphore]
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class _SemaphoreRegistry:
|
||||
"""Weak on both layers: a semaphore binds to its loop once a waiter queues on it, so a strong value would keep a
|
||||
closed loop alive. In-flight ``ClusterConnectionGate.run`` calls hold the semaphore they acquired."""
|
||||
|
||||
_per_loop: "weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, _ClientSemaphores]" = field(
|
||||
default_factory=weakref.WeakKeyDictionary, repr=False
|
||||
)
|
||||
|
||||
def get(self, client_key: str, limit: int) -> asyncio.Semaphore:
|
||||
per_loop: Final = self._per_loop.setdefault(asyncio.get_running_loop(), weakref.WeakValueDictionary())
|
||||
existing: Final = per_loop.get(client_key)
|
||||
if existing is not None:
|
||||
return existing
|
||||
created: Final = asyncio.Semaphore(limit)
|
||||
per_loop[client_key] = created
|
||||
return created
|
||||
|
||||
|
||||
_Admission: TypeAlias = tuple[str, "asyncio.Task[object] | None"]
|
||||
_SEMAPHORES: Final = _SemaphoreRegistry()
|
||||
_ADMITTED: Final[contextvars.ContextVar[frozenset[_Admission]]] = contextvars.ContextVar(
|
||||
"redis_cluster_gate_admitted", default=frozenset()
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ClusterConnectionGate:
|
||||
"""The async ``RedisCluster`` has no blocking pool, so past ``max_connections`` a node raises
|
||||
``MaxConnectionsError`` instead of waiting and the circuit breaker opens on a healthy Redis. This gate makes
|
||||
guarded commands wait for a slot the way the standalone ``BlockingConnectionPool`` does. A command that finds no
|
||||
slot within ``wait_timeout`` runs anyway, so a genuinely saturated pool still reaches the client's own error
|
||||
handling and the breaker."""
|
||||
|
||||
limit: int
|
||||
wait_timeout: float = REDIS_CONNECTION_POOL_TIMEOUT
|
||||
|
||||
@classmethod
|
||||
def from_max_connections(cls, max_connections: object) -> "ClusterConnectionGate | None":
|
||||
limit: Final = _connection_limit(max_connections)
|
||||
return None if limit is None else cls(limit=limit)
|
||||
|
||||
async def run(self, call: Callable[[], Awaitable[_T]], client_key: str) -> _T:
|
||||
admission: Final[_Admission] = (client_key, asyncio.current_task())
|
||||
if admission in _ADMITTED.get():
|
||||
return await call()
|
||||
semaphore: Final = _SEMAPHORES.get(client_key, self.limit)
|
||||
try:
|
||||
await asyncio.wait_for(semaphore.acquire(), timeout=self.wait_timeout)
|
||||
except asyncio.TimeoutError:
|
||||
verbose_logger.debug(
|
||||
"Redis Cluster: no free connection slot within %ss (max_connections=%s), sending anyway",
|
||||
self.wait_timeout,
|
||||
self.limit,
|
||||
)
|
||||
return await call()
|
||||
token: Final = _ADMITTED.set(_ADMITTED.get() | {admission})
|
||||
try:
|
||||
return await call()
|
||||
finally:
|
||||
_ADMITTED.reset(token)
|
||||
semaphore.release()
|
||||
|
||||
|
||||
class RedisClusterCache(RedisCache):
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
def __init__(self, *args, max_connections: object = None, connection_gate: object = None, **kwargs):
|
||||
pool_kwargs: Final = {} if max_connections is None else {"max_connections": max_connections}
|
||||
super().__init__(*args, **pool_kwargs, **kwargs)
|
||||
self._connection_gate: Final = (
|
||||
connection_gate
|
||||
if isinstance(connection_gate, ClusterConnectionGate)
|
||||
else ClusterConnectionGate.from_max_connections(
|
||||
os.environ.get("REDIS_MAX_CONNECTIONS") if max_connections is None else max_connections
|
||||
)
|
||||
)
|
||||
|
||||
def _admit_async_call(self, call: Callable[[], Awaitable[_T]]) -> Awaitable[_T]:
|
||||
if self._connection_gate is None:
|
||||
return call()
|
||||
return self._connection_gate.run(call, self._get_async_client_cache_key())
|
||||
|
||||
def _run_redis_mget_operation(self, keys: list[str]) -> list[Any]:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -489,6 +489,9 @@ async def test_circuit_breaker_disabled_guard_always_calls_method():
|
|||
self._circuit_breaker = RedisCircuitBreaker(failure_threshold=1, recovery_timeout=60, enabled=False)
|
||||
self.call_count = 0
|
||||
|
||||
def _admit_async_call(self, call):
|
||||
return call()
|
||||
|
||||
@_redis_circuit_breaker_guard
|
||||
async def boom(self):
|
||||
self.call_count += 1
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ import time
|
|||
import types
|
||||
from collections.abc import Iterator
|
||||
from datetime import timedelta
|
||||
from typing import Final
|
||||
from typing import Final, cast
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
|
@ -1327,6 +1327,32 @@ async def test_pool_wait_timeout_is_a_timeout_failure_not_hard_connectivity():
|
|||
await client.aclose()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cluster_pool_exhaustion_is_a_timeout_failure_not_hard_connectivity():
|
||||
"""The async cluster client has no blocking pool: a node whose max_connections are all busy raises
|
||||
MaxConnectionsError at once. That is the cluster shape of a busy pool, so it must sit behind the
|
||||
timeout duration gate like the standalone pool's wait timeout instead of opening the breaker."""
|
||||
from redis import exceptions as redis_exceptions
|
||||
|
||||
max_connections_error = cast( # cast-ok: types-redis 4.6 stubs predate MaxConnectionsError
|
||||
"type[Exception]",
|
||||
redis_exceptions.MaxConnectionsError, # pyright: ignore[reportAttributeAccessIssue] # old stubs
|
||||
)
|
||||
|
||||
from litellm.caching.redis_cache import RedisCircuitBreaker, _run_under_circuit_breaker
|
||||
|
||||
breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60, timeout_min_duration=5.0)
|
||||
|
||||
async def pool_exhausted() -> None:
|
||||
raise max_connections_error()
|
||||
|
||||
for _ in range(breaker.failure_threshold * 2):
|
||||
with pytest.raises(max_connections_error):
|
||||
await _run_under_circuit_breaker(breaker, "op", pool_exhausted)
|
||||
|
||||
assert breaker.is_open() is False, "an exhausted client-side pool is a timeout, not a dead Redis"
|
||||
|
||||
|
||||
def test_timeout_classification_follows_the_explicit_cause_chain_only():
|
||||
from redis.exceptions import ConnectionError as RedisConnectionError
|
||||
|
||||
|
|
|
|||
|
|
@ -1,20 +1,22 @@
|
|||
import asyncio
|
||||
from importlib import import_module
|
||||
import gc
|
||||
import json
|
||||
import ssl
|
||||
import weakref
|
||||
from collections.abc import Awaitable, Callable
|
||||
from importlib import import_module
|
||||
from typing import cast
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
from redis.asyncio import Redis, RedisCluster
|
||||
from redis.asyncio.cluster import ClusterNode
|
||||
from redis.asyncio.connection import SSLConnection
|
||||
|
||||
|
||||
from litellm.caching.redis_cache import RedisCache
|
||||
from litellm.caching.redis_cluster_cache import RedisClusterCache
|
||||
from litellm.caching.evicted_client_closer import EvictedClientCloser
|
||||
from litellm.caching.llm_caching_handler import LLMClientCache
|
||||
from litellm.caching.redis_cache import RedisCache
|
||||
from litellm.caching.redis_cluster_cache import ClusterConnectionGate, RedisClusterCache
|
||||
|
||||
|
||||
@patch("litellm._redis.init_redis_cluster")
|
||||
|
|
@ -73,9 +75,7 @@ async def test_redis_cluster_async_batch_get(mock_init_redis_cluster):
|
|||
@patch("litellm._redis.get_redis_connection_pool")
|
||||
@patch("litellm._redis.get_redis_client")
|
||||
@patch.object(import_module("litellm.caching.redis_cache").RedisCache, "_setup_health_pings")
|
||||
def test_cache_init_creates_cluster_cache_from_env_var(
|
||||
mock_health, mock_get_client, mock_get_pool, monkeypatch
|
||||
):
|
||||
def test_cache_init_creates_cluster_cache_from_env_var(mock_health, mock_get_client, mock_get_pool, monkeypatch):
|
||||
"""
|
||||
Test that Cache() creates RedisClusterCache when REDIS_CLUSTER_NODES env var is set.
|
||||
|
||||
|
|
@ -153,9 +153,7 @@ def test_cache_init_creates_redis_cache_without_cluster_config(
|
|||
),
|
||||
],
|
||||
)
|
||||
def test_router_create_redis_cache_cluster_detection(
|
||||
startup_nodes, env_var, expected_cache_type, monkeypatch
|
||||
):
|
||||
def test_router_create_redis_cache_cluster_detection(startup_nodes, env_var, expected_cache_type, monkeypatch):
|
||||
"""
|
||||
Test that Router._create_redis_cache() creates RedisClusterCache when
|
||||
either startup_nodes is in config or REDIS_CLUSTER_NODES env var is set.
|
||||
|
|
@ -416,3 +414,186 @@ def test_evicted_async_client_is_closed_only_for_cluster(cluster: bool) -> None:
|
|||
client.aclose.assert_not_awaited()
|
||||
|
||||
asyncio.run(exercise())
|
||||
|
||||
|
||||
class _InFlightCountingClient:
|
||||
"""Async cluster client double that records how many commands were in flight at once."""
|
||||
|
||||
def __init__(self, command_seconds: float) -> None:
|
||||
self.command_seconds = command_seconds
|
||||
self.in_flight = 0
|
||||
self.peak_in_flight = 0
|
||||
self.completed = 0
|
||||
|
||||
async def set(self, *args: object, **kwargs: object) -> bool:
|
||||
self.in_flight += 1
|
||||
self.peak_in_flight = max(self.peak_in_flight, self.in_flight)
|
||||
await asyncio.sleep(self.command_seconds)
|
||||
self.in_flight -= 1
|
||||
self.completed += 1
|
||||
return True
|
||||
|
||||
|
||||
def _cluster_set_cache_with_client(
|
||||
client: _InFlightCountingClient,
|
||||
max_connections: int | None = None,
|
||||
connection_gate: ClusterConnectionGate | None = None,
|
||||
) -> Callable[[str, str], Awaitable[object]]:
|
||||
with patch("litellm._redis.init_redis_cluster"):
|
||||
cache = RedisClusterCache(
|
||||
startup_nodes=[{"host": "localhost", "port": 7000}],
|
||||
max_connections=max_connections,
|
||||
connection_gate=connection_gate,
|
||||
)
|
||||
setattr(cache, "init_async_client", MagicMock(return_value=client))
|
||||
return cast(
|
||||
"Callable[[str, str], Awaitable[object]]", cache.async_set_cache
|
||||
) # cast-ok: wrapped **kwargs are untyped
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cluster_cache_queues_commands_beyond_max_connections_instead_of_failing() -> None:
|
||||
"""A burst larger than max_connections used to hand every command to the node at once, so redis-py raised
|
||||
MaxConnectionsError on the overflow and the circuit breaker opened. The overflow now waits for a slot."""
|
||||
client = _InFlightCountingClient(command_seconds=0.01)
|
||||
set_cache = _cluster_set_cache_with_client(client, max_connections=3)
|
||||
|
||||
await asyncio.gather(*(set_cache(f"key-{i}", "value") for i in range(20)))
|
||||
|
||||
assert client.completed == 20
|
||||
assert client.peak_in_flight <= 3
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cluster_cache_without_max_connections_does_not_throttle() -> None:
|
||||
client = _InFlightCountingClient(command_seconds=0.01)
|
||||
set_cache = _cluster_set_cache_with_client(client)
|
||||
|
||||
await asyncio.gather(*(set_cache(f"key-{i}", "value") for i in range(20)))
|
||||
|
||||
assert client.completed == 20
|
||||
assert client.peak_in_flight == 20
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cluster_cache_sends_after_the_connection_wait_timeout() -> None:
|
||||
"""A command that cannot get a slot within the pool timeout is still sent, so a saturated pool keeps
|
||||
surfacing through the client's own errors and the breaker rather than stalling forever."""
|
||||
client = _InFlightCountingClient(command_seconds=0.2)
|
||||
set_cache = _cluster_set_cache_with_client(
|
||||
client, connection_gate=ClusterConnectionGate(limit=1, wait_timeout=0.02)
|
||||
)
|
||||
|
||||
await asyncio.wait_for(asyncio.gather(*(set_cache(f"key-{i}", "value") for i in range(3))), timeout=1.0)
|
||||
|
||||
assert client.completed == 3
|
||||
assert client.peak_in_flight == 3
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cluster_caches_sharing_a_client_share_the_connection_cap() -> None:
|
||||
"""init_async_client hands caches with identical settings the same client, so two caches must not each admit
|
||||
a full max_connections worth of commands against that one pool."""
|
||||
client = _InFlightCountingClient(command_seconds=0.01)
|
||||
first = _cluster_set_cache_with_client(client, max_connections=3)
|
||||
second = _cluster_set_cache_with_client(client, max_connections=3)
|
||||
|
||||
await asyncio.gather(
|
||||
*(first(f"first-{i}", "value") for i in range(10)), *(second(f"second-{i}", "value") for i in range(10))
|
||||
)
|
||||
|
||||
assert client.completed == 20
|
||||
assert client.peak_in_flight <= 3
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_guarded_command_nested_in_a_guarded_command_reuses_the_slot() -> None:
|
||||
"""batch_cache_write flushes through async_set_cache_pipeline while it is itself admitted. With one slot the
|
||||
inner command must run on the outer slot instead of waiting for the pool timeout."""
|
||||
gate = ClusterConnectionGate(limit=1, wait_timeout=5.0)
|
||||
|
||||
async def inner() -> str:
|
||||
return "flushed"
|
||||
|
||||
async def outer() -> str:
|
||||
return await gate.run(inner, "client-a")
|
||||
|
||||
assert await asyncio.wait_for(gate.run(outer, "client-a"), timeout=1.0) == "flushed"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_nested_guarded_command_on_another_client_waits_for_that_clients_slot() -> None:
|
||||
gate = ClusterConnectionGate(limit=1, wait_timeout=5.0)
|
||||
client_b_busy = asyncio.Event()
|
||||
release_b = asyncio.Event()
|
||||
peak_b = 0
|
||||
in_flight_b = 0
|
||||
|
||||
async def hold_b() -> None:
|
||||
nonlocal peak_b, in_flight_b
|
||||
in_flight_b += 1
|
||||
peak_b = max(peak_b, in_flight_b)
|
||||
client_b_busy.set()
|
||||
await release_b.wait()
|
||||
in_flight_b -= 1
|
||||
|
||||
async def touch_b() -> None:
|
||||
nonlocal peak_b, in_flight_b
|
||||
in_flight_b += 1
|
||||
peak_b = max(peak_b, in_flight_b)
|
||||
await asyncio.sleep(0)
|
||||
in_flight_b -= 1
|
||||
|
||||
async def outer_on_a() -> None:
|
||||
await client_b_busy.wait()
|
||||
await gate.run(touch_b, "client-b")
|
||||
|
||||
holder = asyncio.create_task(gate.run(hold_b, "client-b"))
|
||||
nested = asyncio.create_task(gate.run(outer_on_a, "client-a"))
|
||||
await client_b_busy.wait()
|
||||
await asyncio.sleep(0.05)
|
||||
release_b.set()
|
||||
await asyncio.gather(holder, nested)
|
||||
|
||||
assert peak_b == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_child_task_spawned_inside_a_guarded_command_does_not_inherit_the_slot() -> None:
|
||||
gate = ClusterConnectionGate(limit=1, wait_timeout=5.0)
|
||||
peak = 0
|
||||
in_flight = 0
|
||||
|
||||
async def command() -> None:
|
||||
nonlocal peak, in_flight
|
||||
in_flight += 1
|
||||
peak = max(peak, in_flight)
|
||||
await asyncio.sleep(0.02)
|
||||
in_flight -= 1
|
||||
|
||||
spawned: list[asyncio.Task[None]] = []
|
||||
|
||||
async def parent() -> None:
|
||||
spawned.append(asyncio.create_task(gate.run(command, "client-a")))
|
||||
await command()
|
||||
|
||||
await gate.run(parent, "client-a")
|
||||
await asyncio.gather(*spawned)
|
||||
|
||||
assert peak == 1
|
||||
|
||||
|
||||
def test_closed_event_loop_is_collected_after_contended_gate_use() -> None:
|
||||
gate = ClusterConnectionGate(limit=1, wait_timeout=5.0)
|
||||
|
||||
async def command() -> None:
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
async def contend() -> "weakref.ref[asyncio.AbstractEventLoop]":
|
||||
await asyncio.gather(gate.run(command, "client-a"), gate.run(command, "client-a"))
|
||||
return weakref.ref(asyncio.get_running_loop())
|
||||
|
||||
loop_ref = asyncio.run(contend())
|
||||
gc.collect()
|
||||
|
||||
assert loop_ref() is None
|
||||
|
|
|
|||
|
|
@ -23,7 +23,7 @@ from __future__ import annotations
|
|||
|
||||
import asyncio
|
||||
from datetime import datetime
|
||||
from typing import Final
|
||||
from typing import Final, cast
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
|
@ -830,31 +830,25 @@ async def test_increment_spend_counters_pipelines_all_scopes_in_one_redis_call(
|
|||
assert counter_cache.in_memory_cache.get_cache(key=key) == 1.5
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_increment_spend_counters_pipeline_failure_invalidates_all_counters(
|
||||
monkeypatch,
|
||||
):
|
||||
"""A failing pipeline must invalidate every pending counter so the next
|
||||
request reseeds from the DB (which already holds this request's cost)
|
||||
instead of trusting a value the write may have partially applied."""
|
||||
from redis.exceptions import MaxConnectionsError
|
||||
_PENDING_SPEND_COUNTER_KEYS: tuple[str, ...] = (
|
||||
"spend:key:hashed",
|
||||
"spend:team:team-1",
|
||||
"spend:team_member:user-1:team-1",
|
||||
"spend:user:user-1",
|
||||
"spend:end_user:eu-1",
|
||||
"spend:tag:tag-a",
|
||||
"spend:tag:tag-b",
|
||||
"spend:org:org-1",
|
||||
)
|
||||
|
||||
|
||||
def _warm_counters_behind_a_failing_pipeline(monkeypatch, pipeline_error: Exception) -> AsyncMock:
|
||||
counter_cache = ps.DualCache()
|
||||
pending_keys = (
|
||||
"spend:key:hashed",
|
||||
"spend:team:team-1",
|
||||
"spend:team_member:user-1:team-1",
|
||||
"spend:user:user-1",
|
||||
"spend:end_user:eu-1",
|
||||
"spend:tag:tag-a",
|
||||
"spend:tag:tag-b",
|
||||
"spend:org:org-1",
|
||||
)
|
||||
for key in pending_keys:
|
||||
for key in _PENDING_SPEND_COUNTER_KEYS:
|
||||
counter_cache.in_memory_cache.set_cache(key=key, value=1.0)
|
||||
fake_redis = AsyncMock()
|
||||
fake_redis.async_get_cache = AsyncMock(return_value=1.0) # counters warm
|
||||
fake_redis.async_increment_pipeline = AsyncMock(side_effect=MaxConnectionsError())
|
||||
fake_redis.async_increment_pipeline = AsyncMock(side_effect=pipeline_error)
|
||||
fake_redis.async_increment = AsyncMock()
|
||||
fake_redis.async_delete_cache = AsyncMock()
|
||||
fake_redis.get_ttl = MagicMock(return_value=None)
|
||||
|
|
@ -862,23 +856,63 @@ async def test_increment_spend_counters_pipeline_failure_invalidates_all_counter
|
|||
monkeypatch.setattr(ps, "spend_counter_cache", counter_cache)
|
||||
monkeypatch.setattr(ps, "user_api_key_cache", ps.DualCache())
|
||||
monkeypatch.setattr(ps, "prisma_client", None)
|
||||
return fake_redis
|
||||
|
||||
with pytest.raises(MaxConnectionsError):
|
||||
await ps.increment_spend_counters(
|
||||
token="hashed",
|
||||
team_id="team-1",
|
||||
user_id="user-1",
|
||||
response_cost=0.5,
|
||||
org_id="org-1",
|
||||
end_user_id="eu-1",
|
||||
tags=["tag-a", "tag-b"],
|
||||
)
|
||||
|
||||
async def _increment_every_spend_counter() -> None:
|
||||
await ps.increment_spend_counters(
|
||||
token="hashed",
|
||||
team_id="team-1",
|
||||
user_id="user-1",
|
||||
response_cost=0.5,
|
||||
org_id="org-1",
|
||||
end_user_id="eu-1",
|
||||
tags=["tag-a", "tag-b"],
|
||||
)
|
||||
|
||||
|
||||
def _assert_every_pending_counter_invalidated(fake_redis: AsyncMock) -> None:
|
||||
assert fake_redis.async_increment.await_count == 0
|
||||
deleted_keys = {call.kwargs["key"] for call in fake_redis.async_delete_cache.await_args_list}
|
||||
assert deleted_keys == set(pending_keys)
|
||||
for key in pending_keys:
|
||||
assert counter_cache.in_memory_cache.get_cache(key=key) is None
|
||||
assert deleted_keys == set(_PENDING_SPEND_COUNTER_KEYS)
|
||||
for key in _PENDING_SPEND_COUNTER_KEYS:
|
||||
assert ps.spend_counter_cache.in_memory_cache.get_cache(key=key) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_increment_spend_counters_pipeline_failure_invalidates_all_counters(
|
||||
monkeypatch,
|
||||
):
|
||||
"""A failing pipeline must invalidate every pending counter so the next
|
||||
request reseeds from the DB (which already holds this request's cost)
|
||||
instead of trusting a value the write may have partially applied."""
|
||||
from redis.exceptions import ConnectionError as RedisConnectionError
|
||||
|
||||
fake_redis = _warm_counters_behind_a_failing_pipeline(monkeypatch, RedisConnectionError("Redis is down"))
|
||||
|
||||
with pytest.raises(RedisConnectionError):
|
||||
await _increment_every_spend_counter()
|
||||
|
||||
_assert_every_pending_counter_invalidated(fake_redis)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_increment_spend_counters_pool_exhaustion_invalidates_all_counters_without_failing(
|
||||
monkeypatch,
|
||||
):
|
||||
"""A cluster node with every connection busy raises MaxConnectionsError at once. That is a busy pool, the
|
||||
same as the standalone pool's wait timeout, so the counters are invalidated and the request still succeeds."""
|
||||
from redis import exceptions as redis_exceptions
|
||||
|
||||
max_connections_error = cast( # cast-ok: types-redis 4.6 stubs predate MaxConnectionsError
|
||||
"type[Exception]",
|
||||
redis_exceptions.MaxConnectionsError, # pyright: ignore[reportAttributeAccessIssue] # old stubs
|
||||
)
|
||||
fake_redis = _warm_counters_behind_a_failing_pipeline(monkeypatch, max_connections_error())
|
||||
|
||||
await _increment_every_spend_counter()
|
||||
|
||||
_assert_every_pending_counter_invalidated(fake_redis)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue