diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index fab613ab8a7..d21b3714fed 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -16,6 +16,8 @@ import inspect import itertools import json import logging +import math +import re import threading import time import weakref @@ -24,6 +26,7 @@ from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence from contextvars import ContextVar from dataclasses import dataclass from datetime import timedelta +from decimal import Decimal from types import FrameType, MappingProxyType from typing import TYPE_CHECKING, Any, Final, Protocol, TypeVar, cast @@ -140,10 +143,34 @@ _INCREMENT_IF_PRESENT_LUA: Final = ( ) _LUA_COUNT: Final = TypeAdapter(int) +_PLAIN_DECIMAL: Final = re.compile(r"-?[0-9]+(?:\.[0-9]+)?(?:[eE][-+]?[0-9]+)?") +_PIPELINE_COMMAND_NUMBER: Final = re.compile(r"Command # ([0-9]+) \(") _LUA_OPTIONAL_TOTAL: Final[TypeAdapter[float | None]] = TypeAdapter(float | None) _OPTIONAL_COUNTS: Final = TypeAdapter(tuple[int | None, ...]) +def _as_incrbyfloat_writes(value: float) -> str: + """``value`` as INCRBYFLOAT writes it into a missing key: 17 decimal places with the trailing zeros cut.""" + text: Final = format(Decimal(repr(value)), ".17f").rstrip("0").rstrip(".") + return "0" if text == "-0" else text + + +def _failed_after_the_first_command(error: Exception) -> bool: + """Whether redis-py names a pipeline command after the first as the one that failed, so the first one applied.""" + message: Final = error.args[0] if error.args else None + number: Final = _PIPELINE_COMMAND_NUMBER.match(message) if isinstance(message, str) else None + return number is not None and int(number.group(1)) > 1 + + +def _plain_float(reply: object) -> float | None: + """The finite float a stored value spells in plain decimal or exponent form, or None for anything else.""" + text: Final = reply.decode(errors="replace") if isinstance(reply, bytes) else reply + if not isinstance(text, str) or _PLAIN_DECIMAL.fullmatch(text) is None: + return None + number: Final = float(text) + return number if math.isfinite(number) else None + + def _decoded_counts(values: Sequence[bytes | str | None]) -> tuple[int | None, ...]: return _OPTIONAL_COUNTS.validate_python( tuple(value.decode("utf-8") if isinstance(value, bytes) else value for value in values) @@ -723,6 +750,11 @@ class RedisCache(BaseCache): def _admit_async_call(self, call: Callable[[], Awaitable[_RedisCallResult]]) -> Awaitable[_RedisCallResult]: return call() + @property + def circuit_breaker_closed(self) -> bool: + """Whether calls reach Redis without a recovery probe. Reading it, unlike making a call, never starts one.""" + return not self._circuit_breaker.enabled or self._circuit_breaker.state == RedisCircuitBreaker.CLOSED + def _setup_health_pings(self): """Setup async and sync health pings for Redis.""" # ASYNC HEALTH PING @@ -1482,15 +1514,69 @@ class RedisCache(BaseCache): await _redis_client.expire(key, ttl) return float(result) + async def _set_if_missing_and_get(self, key: str, value: float, ttl: int | None) -> tuple[bool, float] | None: + """SET NX GET and its TTL command in one round trip: whether the key was created and the value it holds, or + None when the server refused SET NX GET rather than failed to answer, or the key holds a value this does not + read as INCRBYFLOAT would. SET NX GET is one command, so a refused one wrote nothing. A refused TTL command + raises, as it does in main's add.""" + from redis.exceptions import ClusterError, RedisError + from redis.exceptions import ConnectionError as RedisConnectionError + from redis.exceptions import TimeoutError as RedisTimeoutError + + if not math.isfinite(value): + return None + stored: Final = _as_incrbyfloat_writes(value) + try: + async with self._async_commands().pipeline(transaction=False) as pipe: + pipe.set(name=key, value=stored, nx=True, get=True) + if ttl is not None: + pipe.expire(key, ttl) + replies: Final[Sequence[object]] = await pipe.execute() # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType] # redis generics + except RedisError as e: + failed_to_answer: Final = isinstance(e, (RedisConnectionError, RedisTimeoutError, ClusterError)) + if failed_to_answer or _failed_after_the_first_command(e): + raise + return None + if replies[0] is None: + return True, float(stored) + held: Final = _plain_float(replies[0]) # pyright: ignore[reportUnknownArgumentType] # redis generics + return None if held is None else (False, held) + + async def _set_if_missing_with_ttl( + self, + _redis_client: "Redis[bytes]", + key: str, + value: float, + ttl: int | None, + refresh_ttl: bool, + on_added: Callable[[], None] | None, + ) -> float: + """The value the key holds after setting it if it was missing. A server that refuses the commands this sends, + such as a user denied SET or a Redis before 7.0, and a key holding a value INCRBYFLOAT might read differently, + get a plain add instead. ``on_added`` runs when this call added the value.""" + written: Final = await self._set_if_missing_and_get(key=key, value=value, ttl=ttl) + current: Final = ( + await self._incrbyfloat_with_ttl(_redis_client, key=key, value=value, ttl=ttl, refresh_ttl=refresh_ttl) + if written is None + else written[1] + ) + if on_added is not None and (written is None or written[0]): + on_added() + return current + @_redis_circuit_breaker_guard async def async_increment( self, - key, + key: str, value: float, ttl: int | None = None, parent_otel_span: Span | None = None, refresh_ttl: bool = False, + only_if_missing: bool = False, + on_added: Callable[[], None] | None = None, ) -> float: + """Add ``value`` to the key and return its new value. With ``only_if_missing``, a key that exists keeps its + value, which is returned, and only the TTL is refreshed, and ``on_added`` runs if this call added the value.""" from redis.asyncio import Redis _redis_client: Final[Redis] = self.init_async_client() @@ -1498,8 +1584,19 @@ class RedisCache(BaseCache): _used_ttl: Final = self.get_ttl(ttl=ttl) key = self.check_and_fix_namespace(key=key) try: - result: Final = await self._incrbyfloat_with_ttl( - _redis_client, key=key, value=value, ttl=_used_ttl, refresh_ttl=refresh_ttl + result: Final = await ( + self._set_if_missing_with_ttl( + _redis_client, # pyright: ignore[reportUnknownArgumentType] # redis generics + key=key, + value=value, + ttl=_used_ttl, + refresh_ttl=refresh_ttl, + on_added=on_added, + ) + if only_if_missing + else self._incrbyfloat_with_ttl( + _redis_client, key=key, value=value, ttl=_used_ttl, refresh_ttl=refresh_ttl + ) ) ## LOGGING ## diff --git a/litellm/constants.py b/litellm/constants.py index 8d05115e76e..737dc28de69 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1942,6 +1942,8 @@ SPEND_LOG_QUEUE_POLL_INTERVAL: Final = get_env_float("SPEND_LOG_QUEUE_POLL_INTER RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS: Final = max(1, get_env_int("RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS", 3)) RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL: Final = get_env_float("RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL", 0.2) SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE: Final = get_env_int("SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE", 10000) +SPEND_COUNTER_COLD_SEED_JOIN_SECONDS: Final = 0.05 +SPEND_COUNTER_INCRBYFLOAT_ROUNDING: Final = 1e-16 PROXY_DB_LOOKUP_MAX_CONCURRENCY: Final = max(1, get_env_int("PROXY_DB_LOOKUP_MAX_CONCURRENCY", 25)) PROXY_DB_LOOKUP_DEADLINE_SECONDS: Final = max(0.1, get_env_float("PROXY_DB_LOOKUP_DEADLINE_SECONDS", 10.0)) PROXY_DB_LOOKUP_STALL_WINDOW_SECONDS: Final = max(0.0, get_env_float("PROXY_DB_LOOKUP_STALL_WINDOW_SECONDS", 30.0)) diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index 721999b1868..f27c67497f9 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -13,22 +13,31 @@ memory in long-lived deployments. """ import asyncio +import functools from collections import OrderedDict -from collections.abc import Mapping +from collections.abc import Awaitable, Callable, Coroutine, Mapping +from contextvars import ContextVar, copy_context +from dataclasses import dataclass from datetime import datetime, timezone from types import MappingProxyType from typing import TYPE_CHECKING, ClassVar, Final, Optional from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger -from litellm.constants import SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE +from litellm.constants import ( + SPEND_COUNTER_COLD_SEED_JOIN_SECONDS, + SPEND_COUNTER_INCRBYFLOAT_ROUNDING, + SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE, +) from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.proxy._types import Litellm_EntityType from litellm.proxy.db.db_lookup_gate import bounded_db_lookup, db_lookup_gate from litellm.proxy.spend_tracking.spend_counter_batch import ( SPEND_COUNTERS_TARGET, + SpendCounterRead, read_batched_spend_counter, record_spend_counter_value, + spend_counter_tick, ) from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.project_repository import ProjectRepository @@ -68,6 +77,93 @@ _WINDOW_SPEND_LOG_FIELDS: Final[Mapping[str, str]] = MappingProxyType( ) +@dataclass(slots=True) +class _CachedSpendSeed: + """This process adding its cached spend to one cold counter: the add in flight, the tick it started at, the tick + the last write of its spend landed at, and that spend.""" + + started_at: int = 0 + landed_at: int = 0 + landed_spend: float = 0.0 + in_flight: asyncio.Future[float | None] | None = None + + def on_this_loop(self) -> bool: + return self.in_flight is None or self.in_flight.get_loop() is asyncio.get_running_loop() + + def run(self, add: Callable[[], Awaitable[float | None]]) -> Coroutine[object, object, float | None]: + """Make ``add`` the add in flight from this call on, and return the coroutine that runs it and publishes its + result to the requests waiting on ``in_flight``, which ends cancelled when ``add`` fails or is cancelled.""" + in_flight: Final[asyncio.Future[float | None]] = asyncio.get_running_loop().create_future() + self.started_at = spend_counter_tick() + self.in_flight = in_flight + + async def add_and_publish() -> float | None: + try: + added: Final = await add() + in_flight.set_result(added) + return added + finally: + if self.in_flight is in_flight: + self.in_flight = None + in_flight.cancel() + + return add_and_publish() + + +@functools.lru_cache(maxsize=SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE) +def _cached_spend_seed(spend_counter_cache: "DualCache", counter_key: str) -> _CachedSpendSeed: + return _CachedSpendSeed() + + +_cold_seed_waited: Final[ContextVar[float]] = ContextVar("cold_spend_seed_waited", default=0.0) + + +async def _wait_for_cold_seed(in_flight: asyncio.Future[float | None]) -> None: + """Wait on ``in_flight`` for what this context has left of ``SPEND_COUNTER_COLD_SEED_JOIN_SECONDS``.""" + loop: Final = asyncio.get_running_loop() + waited: Final = _cold_seed_waited.get() + started: Final = loop.time() + await asyncio.wait((in_flight,), timeout=max(0.0, SPEND_COUNTER_COLD_SEED_JOIN_SECONDS - waited)) + _cold_seed_waited.set(waited + loop.time() - started) + + +_late_writes: Final[set[asyncio.Task[float | None]]] = set() # mutable-ok: asyncio only weak-refs pending tasks +_late_write_abandoned: Final[ContextVar[asyncio.Event | None]] = ContextVar("late_spend_write_abandoned", default=None) + + +def late_write_abandoned() -> bool: + """Whether this runs in a late write whose request was cancelled, which main's add would not have outlived.""" + abandoned: Final = _late_write_abandoned.get() + return abandoned is not None and abandoned.is_set() + + +def _late_write_finished(task: asyncio.Task[float | None]) -> None: + _late_writes.discard(task) + if not task.cancelled(): + task.exception() + + +async def _landing_even_if_cancelled(write: Coroutine[object, object, float | None]) -> float | None: + """Await ``write`` in a task of its own, so cancelling the caller leaves it to land as a command already sent does, + and marks it abandoned from then on, so ``late_write_abandoned`` holds inside it.""" + abandoned: Final = asyncio.Event() + context: Final = copy_context() + context.run(_late_write_abandoned.set, abandoned) + task: Final = context.run(asyncio.ensure_future, write) + _late_writes.add(task) + task.add_done_callback(_late_write_finished) + try: + return await asyncio.shield(task) + except asyncio.CancelledError: + abandoned.set() + raise + + +def _redis_breaker_closed(spend_counter_cache: "DualCache") -> bool: + redis_cache: Final = spend_counter_cache.redis_cache + return redis_cache is None or redis_cache.circuit_breaker_closed + + def _as_utc(value: datetime) -> datetime: return value if value.tzinfo is not None else value.replace(tzinfo=timezone.utc) @@ -112,15 +208,97 @@ class SpendCounterReseed: SpendCounterReseed._locks.popitem(last=False) return lock + @staticmethod + async def add_cached_spend_once( + spend_counter_cache: "DualCache", + counter_key: str, + read: SpendCounterRead, + spend: float, + add: Callable[[], Awaitable[float | None]], + set_if_missing: Callable[[Callable[[], None]], Awaitable[float | None]], + ) -> None: + """Run ``add``, this process adding ``spend`` to a cold counter, once for all of this process's requests that + read the counter cold before that add landed. A request joins an add that started after its read came back and + gets its result, and starts or joins another add when that one fails, is cancelled or wrote nothing after its + read went out. A request whose read went out before an add landed, and one still waiting on an add after + ``SPEND_COUNTER_COLD_SEED_JOIN_SECONDS`` across its counters, runs ``set_if_missing``, which calls its argument + if it added the spend, in place of ``add``, then runs ``add`` too if it did not and the counter holds less than + the spend it should already hold, beyond what INCRBYFLOAT's 17 decimal places can round away. + Those two steps finish even if the request is cancelled, and stay the add in flight until they do, but once it is + cancelled a failure in them leaves the counter as main's cancelled add does, without the DEL. A request on + another event loop than the add in flight, one that finds the Redis circuit breaker not closed, and each other + process, still run ``add``.""" + seed: Final = _cached_spend_seed(spend_counter_cache, counter_key) + + def spend_landed() -> None: + seed.landed_at = spend_counter_tick() + seed.landed_spend = spend + + async def add_spend() -> float | None: + added: Final = await add() + spend_landed() + return added + + async def write_if_short() -> float | None: + least: Final = min(spend, seed.landed_spend) if seed.landed_at >= read.sent_at else spend + set_added: Final = asyncio.Event() + + def on_added() -> None: + set_added.set() + spend_landed() + + current: Final = await set_if_missing(on_added) + if current is None or set_added.is_set(): + return current + return await add_spend() if least - current > SPEND_COUNTER_INCRBYFLOAT_ROUNDING else current + + async def add_if_short() -> float | None: + return await _landing_even_if_cancelled(seed.run(write_if_short)) + + async def add_unless_covered() -> float | None: + in_flight: Final = seed.in_flight + if not seed.on_this_loop() or not _redis_breaker_closed(spend_counter_cache): + return await add() + if in_flight is None: + return await (seed.run(add_spend) if seed.landed_at < read.sent_at else add_if_short()) + covers: Final = seed.started_at > read.returned_at + await _wait_for_cold_seed(in_flight) + if not in_flight.done(): + return await add_if_short() + if covers and not in_flight.cancelled() and seed.landed_at >= read.sent_at: + return in_flight.result() + return await add_unless_covered() + + added: Final = await add_unless_covered() + if added is not None: + record_spend_counter_value(counter_key, added) + @staticmethod @with_service_target(SPEND_COUNTERS_TARGET) - async def increment_in_memory(spend_counter_cache: "DualCache", counter_key: str, increment: float) -> float | None: - """Apply local deltas after an in-flight reseed establishes the spend balance.""" + async def increment_in_memory( + spend_counter_cache: "DualCache", + counter_key: str, + increment: float, + only_if_missing: bool = False, + on_added: Callable[[], None] | None = None, + ) -> float | None: + """Apply local deltas after an in-flight reseed establishes the spend balance. With ``only_if_missing``, a + counter that exists keeps its value, which is returned, and ``on_added`` runs if this call added to it.""" lock: Final = await SpendCounterReseed._get_lock(counter_key) async with lock: - return await spend_counter_cache.async_increment_cache( + current: Final[float | None] = ( # pyright: ignore[reportUnknownVariableType] # untyped in-memory cache + spend_counter_cache.in_memory_cache.get_cache(key=counter_key) # pyright: ignore[reportUnknownMemberType] # untyped in-memory cache + if only_if_missing + else None + ) + if current is not None: + return float(current) + added: Final = await spend_counter_cache.async_increment_cache( key=counter_key, value=increment, local_only=True, refresh_ttl=True ) + if on_added is not None: + on_added() + return added @staticmethod async def from_db(prisma_client: Optional["PrismaClient"], counter_key: str) -> float | None: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f837e75039e..17c5552e2ef 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -682,7 +682,7 @@ from litellm.proxy.db.request_error_tracking import ( flush_request_errors, request_error_accumulator, ) -from litellm.proxy.db.spend_counter_reseed import END_USER_COUNTER_PREFIX, SpendCounterReseed +from litellm.proxy.db.spend_counter_reseed import END_USER_COUNTER_PREFIX, SpendCounterReseed, late_write_abandoned from litellm.proxy.discovery_endpoints import ( agent_skills_discovery_router, ui_discovery_endpoints_router, @@ -981,6 +981,7 @@ from litellm.proxy.spend_tracking.spend_counter_batch import ( read_batched_spend_counter, record_spend_counter_value, spend_counter_batch_scope, + spend_counter_read, ) from litellm.proxy.spend_tracking.spend_management_endpoints import ( router as spend_management_router, @@ -4299,6 +4300,7 @@ async def ensure_spend_counter_initialized( ) -> None: is_warm: Final = await _is_spend_counter_cache_warm(counter_key=counter_key) if is_warm is False: + read: Final = spend_counter_read(counter_key) # Shares the per-counter lock with get_current_spend. db_spend: Final = await SpendCounterReseed.coalesced( prisma_client=prisma_client, @@ -4310,7 +4312,16 @@ async def ensure_spend_counter_initialized( # DB unavailable - fall back to in-process cache (may be stale). base_spend: Final = await _get_source_cache_base_spend(source_cache_key=source_cache_key) if base_spend > 0: - await increment_spend_counter_cache(counter_key=counter_key, increment=base_spend) + await SpendCounterReseed.add_cached_spend_once( + spend_counter_cache=spend_counter_cache, + counter_key=counter_key, + read=read, + spend=base_spend, + add=lambda: increment_spend_counter_cache(counter_key=counter_key, increment=base_spend), + set_if_missing=lambda on_added: increment_spend_counter_cache( + counter_key=counter_key, increment=base_spend, only_if_missing=True, on_added=on_added + ), + ) _ensure_spend_counter_initialized: Final = ensure_spend_counter_initialized @@ -4413,16 +4424,32 @@ async def refresh_spend_counter_ttl(counter_key: str) -> bool: @with_service_target(SPEND_COUNTERS_TARGET) -async def increment_spend_counter_cache(counter_key: str, increment: float) -> float | None: +async def increment_spend_counter_cache( + counter_key: str, + increment: float, + only_if_missing: bool = False, + on_added: Callable[[], None] | None = None, +) -> float | None: if spend_counter_cache.redis_cache is not None: try: - current_value: Final = await spend_counter_cache.redis_cache.async_increment( - key=counter_key, - value=increment, - refresh_ttl=True, + current_value: Final = await ( + spend_counter_cache.redis_cache.async_increment( + key=counter_key, + value=increment, + refresh_ttl=True, + only_if_missing=True, + on_added=on_added, + ) + if only_if_missing + else spend_counter_cache.redis_cache.async_increment( + key=counter_key, + value=increment, + refresh_ttl=True, + ) ) except Exception: - await invalidate_spend_counter(counter_key=counter_key) + if not late_write_abandoned(): + await invalidate_spend_counter(counter_key=counter_key) raise spend_counter_cache.in_memory_cache.set_cache( key=counter_key, @@ -4432,7 +4459,11 @@ async def increment_spend_counter_cache(counter_key: str, increment: float) -> f return current_value return await SpendCounterReseed.increment_in_memory( - spend_counter_cache=spend_counter_cache, counter_key=counter_key, increment=increment + spend_counter_cache=spend_counter_cache, + counter_key=counter_key, + increment=increment, + only_if_missing=only_if_missing, + on_added=on_added, ) diff --git a/litellm/proxy/spend_tracking/spend_counter_batch.py b/litellm/proxy/spend_tracking/spend_counter_batch.py index 7e0821e091e..8ff709086bb 100644 --- a/litellm/proxy/spend_tracking/spend_counter_batch.py +++ b/litellm/proxy/spend_tracking/spend_counter_batch.py @@ -1,11 +1,12 @@ """One Redis MGET per phase (admission, reservation, post-call) for the spend counters it reads, not one GET each.""" import asyncio +import itertools from collections.abc import Iterator, Mapping, Sequence from contextvars import ContextVar, Token from dataclasses import dataclass from types import MappingProxyType, TracebackType -from typing import Final +from typing import Final, TypeAlias from pydantic import TypeAdapter @@ -22,6 +23,24 @@ from litellm.proxy.common_utils.user_api_key_cache import ( _CounterValues: Final = TypeAdapter(dict[str, float | None]) _NO_VALUES: Final[Mapping[str, float | None]] = MappingProxyType({}) SPEND_COUNTERS_TARGET: Final = "spend_counters" +_ticks: Final = itertools.count(1) + + +def spend_counter_tick() -> int: + """The next number in a per-process sequence that orders spend counter reads and cold seeds.""" + return next(_ticks) + + +@dataclass(frozen=True, slots=True) +class SpendCounterRead: + """When a read of a spend counter went out and when its reply came back, as spend counter ticks.""" + + sent_at: int + returned_at: int + + +_NO_READS: Final[Mapping[str, SpendCounterRead]] = MappingProxyType({}) +_PendingMGet: TypeAlias = tuple[BatchResult[Mapping[str, object]], int] @dataclass(frozen=True, slots=True) @@ -39,7 +58,17 @@ class SpendCounterBatch: hook declares whatever is bound but unread, so whoever flushes first (the auth object prefetch, usually) carries the spend counters in the same round trip.""" - __slots__ = ("_fetched", "_inflight", "_keys", "_loaded", "_lock", "_open", "_redis_cache", "_request_batch") + __slots__ = ( + "_fetched", + "_inflight", + "_keys", + "_loaded", + "_lock", + "_open", + "_reads", + "_redis_cache", + "_request_batch", + ) def __init__(self, redis_cache: RedisCache) -> None: self._redis_cache: Final = redis_cache @@ -48,7 +77,8 @@ class SpendCounterBatch: self._keys: frozenset[str] = frozenset() self._fetched: frozenset[str] = frozenset() self._loaded: Mapping[str, float | None] = _NO_VALUES - self._inflight: Final[list[BatchResult[Mapping[str, object]]]] = [] # mutable-ok: drained by _load + self._reads: Mapping[str, SpendCounterRead] = _NO_READS + self._inflight: Final[list[_PendingMGet]] = [] # mutable-ok: drained by _load self._request_batch: Final[RedisBatch | None] = active_request_redis_batch(redis_cache) if self._request_batch is not None: self._request_batch.add_flush_hook(self._declare_pending) @@ -78,6 +108,10 @@ class SpendCounterBatch: return None return loaded[counter_key], True + def counter_read(self, counter_key: str) -> SpendCounterRead | None: + """When the MGET behind this batch's value for the counter went out and came back, None without that value.""" + return self._reads.get(counter_key) if self._open and counter_key in self._loaded else None + def record(self, counter_key: str, value: float) -> None: """A write returned the counter's new value; later reads in this scope see it instead of the MGET value.""" if not self._open: @@ -93,6 +127,7 @@ class SpendCounterBatch: self._keys = self._keys - key self._fetched = self._fetched - key self._loaded = MappingProxyType({k: v for k, v in self._loaded.items() if k != counter_key}) + self._reads = MappingProxyType({k: v for k, v in self._reads.items() if k != counter_key}) async def _load(self) -> Mapping[str, float | None]: async with self._lock: @@ -103,10 +138,15 @@ class SpendCounterBatch: pending: Final = self._keys - self._fetched if pending: self._fetched = self._fetched | pending + sent_at: Final = spend_counter_tick() fetched: Final = await self._fetch(pending) - self._loaded = MappingProxyType({**fetched, **self._loaded}) + self._merge(fetched, SpendCounterRead(sent_at=sent_at, returned_at=spend_counter_tick())) return self._loaded + def _merge(self, fetched: Mapping[str, float | None], read: SpendCounterRead) -> None: + self._loaded = MappingProxyType({**fetched, **self._loaded}) + self._reads = MappingProxyType({**dict.fromkeys(fetched, read), **self._reads}) + def _declare_pending(self) -> None: """Flush hook: put every bound-but-unread counter on the request pipeline that is about to go out.""" if self._request_batch is None or not self._open: @@ -114,19 +154,21 @@ class SpendCounterBatch: pending: Final = self._keys - self._fetched if pending: self._fetched = self._fetched | pending + sent_at: Final = spend_counter_tick() with service_target(SPEND_COUNTERS_TARGET): - self._inflight.append(self._request_batch.mget(sorted(pending))) + read: Final = self._request_batch.mget(sorted(pending)) + self._inflight.append((read, sent_at)) async def _collect_inflight(self) -> None: - results: Final = tuple(self._inflight) + reads: Final = tuple(self._inflight) self._inflight.clear() - for result in results: + for read, sent_at in reads: try: - fetched: Mapping[str, float | None] = _CounterValues.validate_python(await result) + fetched: Mapping[str, float | None] = _CounterValues.validate_python(await read) except Exception as e: # noqa: BLE001 # per-key reads take over and apply their own Redis fallback verbose_proxy_logger.debug("spend counter batch read failed, falling back to per-key reads: %s", e) continue - self._loaded = MappingProxyType({**fetched, **self._loaded}) + self._merge(fetched, SpendCounterRead(sent_at=sent_at, returned_at=spend_counter_tick())) async def _fetch(self, keys: frozenset[str]) -> Mapping[str, float | None]: try: @@ -145,6 +187,17 @@ def active_spend_counter_batch() -> SpendCounterBatch | None: return _active_batch.get() +def spend_counter_read(counter_key: str) -> SpendCounterRead: + """When this context's read of the counter went out and came back: the open batch's MGET for it, else now, since + ``SpendCounterReseed.coalesced`` reads the counter again after a read outside the batch.""" + batch: Final = _active_batch.get() + batched: Final = batch.counter_read(counter_key) if batch is not None else None + if batched is not None: + return batched + now: Final = spend_counter_tick() + return SpendCounterRead(sent_at=now, returned_at=now) + + class spend_counter_batch_scope: """Reads inside the scope share one MGET for the keys bound here or by ``bind_*`` calls inside it. Opened inside a scope whose batch is still open, it binds into that batch so both phases share the MGET.""" diff --git a/tests/integration/spend/test_end_user_budget_enforcement.py b/tests/integration/spend/test_end_user_budget_enforcement.py index 92259d463b9..56a8774f77a 100644 --- a/tests/integration/spend/test_end_user_budget_enforcement.py +++ b/tests/integration/spend/test_end_user_budget_enforcement.py @@ -1,17 +1,20 @@ import json +import os import re import uuid from collections.abc import Callable, Mapping +from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass from typing import Final import httpx import pytest from integration._support.client import Gateway, Scenario, eventually, object_value, string_value -from integration._support.database import read_rows +from integration._support.database import read_rows, write_rows from integration._support.upstream import delete_scenario, register_scenario from integration.cost_calculation.cost_tracking_case import JsonResponse from pydantic import BaseModel, JsonValue +from redis import Redis PROVIDER_AUTHORIZATION: Final = "Bearer integration-provider-key" CANNED_CONTENT: Final = "Hello! This is a mock response from the fake OpenAI endpoint." @@ -649,3 +652,141 @@ def test_streamed_call_is_refused_as_json_once_over_budget(gateway: Gateway) -> assert "data:" not in denied.text, denied.text _assert_budget_refusal(denied, end_user, 0.05) assert _observed_requests(upstream) == [] + + +_COLD_BURST: Final = 40 +_COLD_SEED_SPEND: Final = 6.0 +_COLD_BUDGET: Final = 10.0 +_CHAT_COST: Final = 20 * 0.001 + 20 * 0.002 + + +def _serves_model(gateway: Gateway, model: str) -> bool: + response: Final = gateway.request("GET", "/model/info") + return response.status_code == 200 and any( + object_value(entry)["model_name"] == model for entry in object_value(response.json())["data"] + ) + + +def _counter_cache() -> Redis: + return Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) + + +def _counter(cache: Redis, end_user: str) -> float | None: + raw: Final = cache.get(f"spend:end_user:{end_user}") + return None if raw is None else float(raw) + + +def _lose_counter(cache: Redis, end_user: str) -> None: + cache.delete(f"spend:end_user:{end_user}") + + +def _set_end_user_row_spend(end_user: str, spend: float) -> None: + write_rows('UPDATE "LiteLLM_EndUserTable" SET spend = %s WHERE user_id = %s', (str(spend), end_user)) + + +def _await_row_spend(end_user: str, spend: float) -> None: + eventually( + lambda: _end_user_row(end_user), + lambda rows: len(rows) == 1 and float(str(rows[0]["spend"])) == pytest.approx(spend), + seconds=70, + ) + + +def _refused_spend(denied: httpx.Response, end_user: str) -> float: + error: Final = _denial_error(denied) + assert error["type"] == "budget_exceeded", denied.text + matched: Final = _REFUSAL_MESSAGE.match(string_value(error["message"])) + assert matched is not None, denied.text + assert matched.group("user") == end_user + return float(matched.group("spend")) + + +@pytest.mark.parametrize("spelling", ("user", "responses_user", "messages_metadata_user_id")) +def test_overlapping_cold_starts_count_the_end_user_spend_once(gateway: Gateway, spelling: str) -> None: + chosen: Final = SPELLINGS[spelling] + with ( + gateway.scenario() as scenario, + httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, + _counter_cache() as cache, + ): + deployment: Final = _deployment(scenario, chosen) + key: Final = scenario.key(models=[deployment.model]) + end_user: Final = f"integration-end-user-{uuid.uuid4().hex}" + _new_customer(gateway, end_user, max_budget=_COLD_BUDGET, spend=_COLD_SEED_SPEND) + scenario.cleanups.callback(gateway.request, "POST", "/customer/delete", {"user_ids": [end_user]}) + assert _customer_info(gateway, end_user).spend == _COLD_SEED_SPEND + assert _counter(cache, end_user) is None + upstream.delete("/__observations").raise_for_status() + + with ThreadPoolExecutor(max_workers=_COLD_BURST) as pool: + burst: Final = tuple( + pool.map( + lambda index: _send(gateway, chosen, deployment, key, end_user, f"cold {index}"), + range(_COLD_BURST), + ) + ) + refused: Final = [response.text for response in burst if response.status_code != 200] + assert refused == [], refused + follow_up: Final = _send(gateway, chosen, deployment, key, end_user, "cold follow-up") + _assert_served(chosen, follow_up) + assert len(_observed_requests(upstream)) == _COLD_BURST + 1 + + expected_spend: Final = _COLD_SEED_SPEND + (_COLD_BURST + 1) * _CHAT_COST + _await_row_spend(end_user, expected_spend) + assert _customer_info(gateway, end_user).spend == pytest.approx(expected_spend) + counted: Final = eventually( + lambda: _counter(cache, end_user), + lambda value: value is not None and value == pytest.approx(expected_spend), + seconds=30, + return_last_on_timeout=True, + ) + assert counted == pytest.approx(expected_spend) + + +def test_a_lost_counter_keeps_spend_this_proxy_has_not_written_yet(gateway: Gateway) -> None: + with gateway.scenario() as scenario, _counter_cache() as cache: + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.25) + key: Final = scenario.key(models=[model]) + end_user: Final = f"integration-end-user-{uuid.uuid4().hex}" + call_cost: Final = 20 * 0.001 + 20 * 0.25 + _new_customer(gateway, end_user, max_budget=_COLD_BUDGET, spend=_COLD_SEED_SPEND) + scenario.cleanups.callback(gateway.request, "POST", "/customer/delete", {"user_ids": [end_user]}) + served: Final = _chat(gateway, model, key, end_user) + assert served.status_code == 200, served.text + _await_row_spend(end_user, _COLD_SEED_SPEND + call_cost) + _set_end_user_row_spend(end_user, _COLD_SEED_SPEND) + _lose_counter(cache, end_user) + + denied: Final = _chat(gateway, model, key, end_user) + + assert _refused_spend(denied, end_user) == pytest.approx(_COLD_SEED_SPEND + call_cost) + + +def test_a_lost_counter_keeps_spend_another_proxy_has_not_written_yet(gateway: Gateway, peer: Gateway) -> None: + with gateway.scenario() as scenario, _counter_cache() as cache: + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.149) + eventually(lambda: _serves_model(peer, model), bool, seconds=30) + key: Final = scenario.key(models=[model]) + end_user: Final = f"integration-end-user-{uuid.uuid4().hex}" + call_cost: Final = 20 * 0.001 + 20 * 0.149 + _new_customer(gateway, end_user, max_budget=_COLD_BUDGET, spend=_COLD_SEED_SPEND) + scenario.cleanups.callback(gateway.request, "POST", "/customer/delete", {"user_ids": [end_user]}) + first: Final = _chat(peer, model, key, end_user) + assert first.status_code == 200, first.text + _await_row_spend(end_user, _COLD_SEED_SPEND + call_cost) + _set_end_user_row_spend(end_user, _COLD_SEED_SPEND) + _lose_counter(cache, end_user) + second: Final = _chat(peer, model, key, end_user) + assert second.status_code == 200, second.text + expected_spend: Final = _COLD_SEED_SPEND + 2 * call_cost + counted: Final = eventually( + lambda: _counter(cache, end_user), + lambda value: value is not None and value == pytest.approx(expected_spend), + seconds=30, + return_last_on_timeout=True, + ) + + denied: Final = _chat(gateway, model, key, end_user) + + assert counted == pytest.approx(expected_spend) + assert _denial_error(denied)["type"] == "budget_exceeded", denied.text diff --git a/tests/unit/caching/test_redis_cache.py b/tests/unit/caching/test_redis_cache.py index 0fbeb05e154..504c7d18196 100644 --- a/tests/unit/caching/test_redis_cache.py +++ b/tests/unit/caching/test_redis_cache.py @@ -8,6 +8,10 @@ from typing import Final, cast from unittest.mock import AsyncMock, MagicMock, patch import pytest +from fakeredis import FakeAsyncRedis, FakeServer +from redis.exceptions import ClusterDownError, NoPermissionError, RedisError, ResponseError +from redis.exceptions import ConnectionError as RedisConnectionError +from redis.exceptions import TimeoutError as RedisTimeoutError import litellm from litellm._service_logger import ServiceLogging @@ -1033,8 +1037,9 @@ def test_prometheus_services_logger_starts_from_the_breakers_already_built(monke monkeypatch.setattr(litellm, "service_callback", ["prometheus_system"]) gc.collect() ServiceLogging() - assert REGISTRY.get_sample_value("litellm_redis_circuit_breaker_state", {"state": "closed"}) == ( - redis_breaker_state_counts()["closed"] + assert ( + REGISTRY.get_sample_value("litellm_redis_circuit_breaker_state", {"state": "closed"}) + == (redis_breaker_state_counts()["closed"]) ) @@ -1474,6 +1479,411 @@ async def test_async_increment_pipelines_the_ttl_command( assert client.round_trips == expected_round_trips +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("existing", "expected"), + [ + pytest.param(None, 1.5, id="missing key: set to the value"), + pytest.param(0.25, 0.25, id="key below the value: kept"), + pytest.param(4.0, 4.0, id="key above the value: kept"), + ], +) +async def test_async_increment_only_if_missing_sets_a_missing_key_and_refreshes_the_ttl( + monkeypatch: pytest.MonkeyPatch, redis_no_ping: object, existing: float | None, expected: float +): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache: Final = RedisCache(namespace="ns") + client: Final = FakeAsyncRedis(server=FakeServer()) + if existing is not None: + await client.set("ns:spend:key:k", existing, ex=5) + + with patch.object(redis_cache, "init_async_client", return_value=client): + current: Final = await redis_cache.async_increment( + key="spend:key:k", value=1.5, ttl=60, refresh_ttl=True, only_if_missing=True + ) + after_add: Final = await redis_cache.async_increment(key="spend:key:k", value=2.0, ttl=60, refresh_ttl=True) + + assert current == expected + assert after_add == expected + 2.0 + assert await client.ttl("ns:spend:key:k") == 60 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("value", "stored"), + [(6.0, b"6"), (1e-05, b"0.00001"), (6.000000002, b"6.000000002"), (0.1, b"0.1"), (1e-20, b"0"), (-0.0, b"0")], +) +async def test_async_increment_only_if_missing_stores_a_missing_key_the_way_an_add_stores_it( + monkeypatch: pytest.MonkeyPatch, redis_no_ping: object, value: float, stored: bytes +): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache: Final = RedisCache(namespace="ns") + client: Final = FakeAsyncRedis(server=FakeServer()) + + with patch.object(redis_cache, "init_async_client", return_value=client): + current: Final = await redis_cache.async_increment(key="spend:key:k", value=value, ttl=60, only_if_missing=True) + + assert await client.get("ns:spend:key:k") == stored + assert current == float(stored) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("existing", "reported"), [pytest.param(None, ["created"], id="created"), pytest.param(7.0, [], id="kept")] +) +async def test_async_increment_only_if_missing_reports_only_a_key_it_created( + monkeypatch: pytest.MonkeyPatch, redis_no_ping: object, existing: float | None, reported: list[str] +): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache: Final = RedisCache(namespace="ns") + client: Final = FakeAsyncRedis(server=FakeServer()) + if existing is not None: + await client.set("ns:spend:key:k", existing) + created: Final[list[str]] = [] + + with patch.object(redis_cache, "init_async_client", return_value=client): + await redis_cache.async_increment( + key="spend:key:k", value=1.5, ttl=60, only_if_missing=True, on_added=lambda: created.append("created") + ) + + assert created == reported + + +class _RaisingPipeline: + """A pipeline whose execute raises the first error in its replies, as redis-py's does by default.""" + + def __init__(self, replies: list[object]) -> None: + self.replies = replies + self.commands: list[tuple[object, ...]] = [] + + async def __aenter__(self) -> "_RaisingPipeline": + return self + + async def __aexit__(self, *exc_info: object) -> None: + return None + + def set(self, name: str, value: str, nx: bool = False, get: bool = False) -> None: + self.commands.append(("SET", name, value, nx, get)) + + def incrbyfloat(self, name: str, amount: float) -> None: + self.commands.append(("INCRBYFLOAT", name, amount)) + + def expire(self, name: str, ttl: int) -> None: + self.commands.append(("EXPIRE", name, ttl)) + + async def execute(self) -> list[object]: + first_error: Final = next((reply for reply in self.replies if isinstance(reply, Exception)), None) + if first_error is not None: + raise first_error + return self.replies + + +class _ClusterPipelineAfterADroppedConnection(_RaisingPipeline): + """The first attempt reads a dropped connection. As in redis-py 5.3.1's ClusterPipeline.execute, raising is what + drives the retry, and ``raise_on_error=False`` hands the error back as a reply instead.""" + + async def execute(self, raise_on_error: bool = True) -> list[object]: + if not raise_on_error: + return [RedisConnectionError("Connection closed by server."), self.replies[1]] + return self.replies + + +@pytest.mark.asyncio +async def test_async_increment_only_if_missing_keeps_the_cluster_retry_of_a_dropped_connection( + monkeypatch: pytest.MonkeyPatch, redis_no_ping: object +): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache: Final = RedisCache(namespace="ns") + client: Final = MagicMock() + client.pipeline.return_value = _ClusterPipelineAfterADroppedConnection([None, True]) + + with patch.object(redis_cache, "init_async_client", return_value=client): + current: Final = await redis_cache.async_increment( + key="spend:key:k", value=1.5, ttl=60, refresh_ttl=True, only_if_missing=True + ) + + assert current == 1.5 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "refusal", + [ + pytest.param(NoPermissionError("no permissions to run the 'set' command"), id="denied"), + pytest.param(ResponseError("syntax error"), id="redis-before-7-has-no-set-nx-get"), + pytest.param(ResponseError("unknown command 'SET'"), id="renamed"), + pytest.param(RedisError("SET command doesn't exist in Redis commands"), id="renamed-on-a-cluster-client"), + ], +) +async def test_async_increment_only_if_missing_adds_the_value_when_the_server_refuses_the_set( + monkeypatch: pytest.MonkeyPatch, redis_no_ping: object, refusal: Exception +): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache: Final = RedisCache(namespace="ns") + add: Final = _RaisingPipeline([7.5, True]) + client: Final = MagicMock() + client.pipeline.side_effect = [_RaisingPipeline([refusal, True]), add] + failure_hook: Final = AsyncMock() + added: Final = MagicMock() + + with ( + patch.object(redis_cache, "init_async_client", return_value=client), + patch.object(redis_cache.service_logger_obj, "async_service_failure_hook", failure_hook), + ): + current: Final = await redis_cache.async_increment( + key="spend:key:k", value=1.5, ttl=60, refresh_ttl=True, only_if_missing=True, on_added=added + ) + await asyncio.sleep(0) + + assert current == 7.5 + added.assert_called_once_with() + assert add.commands == [("INCRBYFLOAT", "ns:spend:key:k", 1.5), ("EXPIRE", "ns:spend:key:k", 60)] + failure_hook.assert_not_awaited() + assert redis_cache._circuit_breaker._hard_failure_count == 0 + + +@pytest.mark.asyncio +async def test_async_increment_only_if_missing_reports_a_write_redis_refuses_with_the_add_and_its_error_text( + monkeypatch: pytest.MonkeyPatch, redis_no_ping: object +): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache: Final = RedisCache(namespace="ns") + add_error: Final = ResponseError("Command # 1 (INCRBYFLOAT ns:spend:key:k 1.5) of pipeline caused error") + client: Final = MagicMock() + client.pipeline.side_effect = [ + _RaisingPipeline([ResponseError("Command # 1 (SET) of pipeline caused error"), True]), + _RaisingPipeline([add_error, True]), + ] + failure_hook: Final = AsyncMock() + added: Final = MagicMock() + + with ( + patch.object(redis_cache, "init_async_client", return_value=client), + patch.object(redis_cache.service_logger_obj, "async_service_failure_hook", failure_hook), + pytest.raises(ResponseError) as raised, + ): + await redis_cache.async_increment( + key="spend:key:k", value=1.5, ttl=60, refresh_ttl=True, only_if_missing=True, on_added=added + ) + await asyncio.sleep(0) + + assert raised.value is add_error + added.assert_not_called() + failure_hook.assert_awaited_once() + assert failure_hook.await_args.kwargs["error"] is add_error + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "failure", + [ + pytest.param(RedisConnectionError("Connection closed by server."), id="dropped-connection"), + pytest.param(RedisTimeoutError("Timeout reading from socket"), id="stalled-node"), + pytest.param(ClusterDownError("CLUSTERDOWN The cluster is down"), id="cluster-down"), + ], +) +async def test_async_increment_only_if_missing_sends_no_add_after_a_failure_the_client_already_retried( + monkeypatch: pytest.MonkeyPatch, redis_no_ping: object, failure: Exception +): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache: Final = RedisCache(namespace="ns") + client: Final = MagicMock() + client.pipeline.return_value = _RaisingPipeline([failure, True]) + + with patch.object(redis_cache, "init_async_client", return_value=client), pytest.raises(type(failure)): + await redis_cache.async_increment(key="spend:key:k", value=1.5, ttl=60, refresh_ttl=True, only_if_missing=True) + + assert client.pipeline.call_count == 1 + + +class _Keyspace: + """String keys as Redis holds them, for pipelines that run each command in order and raise the first error after + running the rest, as a non-transaction pipeline does. Commands named in ``refused`` are refused, a pipeline that + has run ``writes_before_full`` write commands has the rest refused as Redis does past ``maxmemory``, and pipelines + after the first ``reached`` never reach Redis, as when the connection drops.""" + + def __init__( + self, refused: frozenset[str] = frozenset(), writes_before_full: int | None = None, reached: int | None = None + ) -> None: + self.values: Final[dict[str, str]] = {} # mutable-ok: stands in for the Redis keyspace + self.refused = refused + self.writes_before_full = writes_before_full + self.reached = reached + self.pipelines = 0 + + def pipeline(self, **_: object) -> "_KeyspacePipeline": + self.pipelines += 1 + return _KeyspacePipeline(self, lost=self.reached is not None and self.pipelines > self.reached) + + +class _KeyspacePipeline(_RaisingPipeline): + def __init__(self, keyspace: _Keyspace, lost: bool) -> None: + super().__init__([]) + self.keyspace = keyspace + self.lost = lost + self.writes_left = keyspace.writes_before_full + + def setnx(self, name: str, value: float) -> None: + self.commands.append(("SETNX", name, value)) + + def _write(self, command: str) -> None: + if self.writes_left is not None: + if self.writes_left == 0: + raise ResponseError("OOM command not allowed when used memory > 'maxmemory'.") + self.writes_left -= 1 + + def _run(self, command: tuple[object, ...]) -> object: + name: Final = str(command[1]) + held: Final = self.keyspace.values.get(name) + if command[0] in self.keyspace.refused: + raise NoPermissionError(f"this user has no permissions to run the '{str(command[0]).lower()}' command") + if command[0] == "EXPIRE": + return held is not None + self._write(str(command[0])) + if command[0] == "SET": + if held is None: + self.keyspace.values[name] = str(command[2]) + return None if held is None else held.encode() + if command[0] == "SETNX": + if held is None: + self.keyspace.values[name] = repr(command[2]) + return held is None + total: Final = float(held or 0) + float(cast(float, command[2])) + self.keyspace.values[name] = repr(total) + return total + + async def execute(self) -> list[object]: + if self.lost: + raise RedisConnectionError("Connection closed by server.") + replies: Final[list[object]] = [] + for number, command in enumerate(self.commands, start=1): + try: + replies.append(self._run(command)) + except RedisError as e: + words: str = " ".join(str(part) for part in command) + e.args = (f"Command # {number} ({words}) of pipeline caused error: {e.args}",) + replies.append(e) + first_error: Final = next((reply for reply in replies if isinstance(reply, Exception)), None) + if first_error is not None: + raise first_error + return replies + + +@pytest.mark.asyncio +async def test_async_increment_only_if_missing_leaves_a_missing_key_missing_when_the_set_is_refused_and_the_add_lost( + monkeypatch: pytest.MonkeyPatch, redis_no_ping: object +): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache: Final = RedisCache(namespace="ns") + keyspace: Final = _Keyspace(refused=frozenset(("SET", "SETNX")), reached=1) + + with patch.object(redis_cache, "init_async_client", return_value=keyspace), pytest.raises(RedisConnectionError): + await redis_cache.async_increment(key="spend:key:k", value=6.0, ttl=60, refresh_ttl=True, only_if_missing=True) + + assert keyspace.values == {} + + +@pytest.mark.asyncio +async def test_async_increment_only_if_missing_sets_a_missing_key_once_when_redis_runs_out_of_memory_after_the_set( + monkeypatch: pytest.MonkeyPatch, redis_no_ping: object +): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache: Final = RedisCache(namespace="ns") + keyspace: Final = _Keyspace(writes_before_full=1) + added: Final = MagicMock() + + with patch.object(redis_cache, "init_async_client", return_value=keyspace): + current: Final = await redis_cache.async_increment( + key="spend:key:k", value=6.0, ttl=60, refresh_ttl=True, only_if_missing=True, on_added=added + ) + + assert current == 6.0 + assert float(keyspace.values["ns:spend:key:k"]) == 6.0 + assert keyspace.pipelines == 1 + added.assert_called_once_with() + + +@pytest.mark.asyncio +async def test_async_increment_only_if_missing_raises_a_refused_ttl_command_as_mains_add_does_without_adding_again( + monkeypatch: pytest.MonkeyPatch, redis_no_ping: object +): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache: Final = RedisCache(namespace="ns") + keyspace: Final = _Keyspace(refused=frozenset(("EXPIRE",))) + added: Final = MagicMock() + + with patch.object(redis_cache, "init_async_client", return_value=keyspace), pytest.raises(NoPermissionError): + await redis_cache.async_increment( + key="spend:key:k", value=6.0, ttl=60, refresh_ttl=True, only_if_missing=True, on_added=added + ) + + assert keyspace.values == {"ns:spend:key:k": "6"} + assert keyspace.pipelines == 1 + added.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("held", ["not-a-number", "0x10", "6_0", " 6", "inf", "nan", "1e999"]) +async def test_async_increment_only_if_missing_leaves_a_value_it_cannot_read_like_incrbyfloat_to_mains_add( + monkeypatch: pytest.MonkeyPatch, redis_no_ping: object, held: str +): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache: Final = RedisCache(namespace="ns") + add: Final = _RaisingPipeline([7.5, True]) + client: Final = MagicMock() + client.pipeline.side_effect = [_RaisingPipeline([held.encode(), True]), add] + + with patch.object(redis_cache, "init_async_client", return_value=client): + current: Final = await redis_cache.async_increment( + key="spend:key:k", value=1.5, ttl=60, refresh_ttl=True, only_if_missing=True + ) + + assert current == 7.5 + assert add.commands == [("INCRBYFLOAT", "ns:spend:key:k", 1.5), ("EXPIRE", "ns:spend:key:k", 60)] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("value", [float("inf"), float("nan")]) +async def test_async_increment_only_if_missing_sends_a_value_incrbyfloat_refuses_as_mains_add( + monkeypatch: pytest.MonkeyPatch, redis_no_ping: object, value: float +): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache: Final = RedisCache(namespace="ns") + add: Final = _RaisingPipeline([ResponseError("increment would produce NaN or Infinity"), True]) + client: Final = MagicMock() + client.pipeline.return_value = add + + with patch.object(redis_cache, "init_async_client", return_value=client), pytest.raises(ResponseError): + await redis_cache.async_increment( + key="spend:key:k", value=value, ttl=60, refresh_ttl=True, only_if_missing=True + ) + + assert client.pipeline.call_count == 1 + assert add.commands[0][0] == "INCRBYFLOAT" + + +@pytest.mark.asyncio +async def test_async_increment_only_if_missing_reports_a_redis_failure_like_an_increment( + monkeypatch: pytest.MonkeyPatch, redis_no_ping: object +): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache: Final = RedisCache(namespace="ns") + client: Final = MagicMock() + client.pipeline.side_effect = ConnectionError("connection reset") + failure_hook: Final = AsyncMock() + + with ( + patch.object(redis_cache, "init_async_client", return_value=client), + patch.object(redis_cache.service_logger_obj, "async_service_failure_hook", failure_hook), + pytest.raises(ConnectionError), + ): + await redis_cache.async_increment(key="spend:key:k", value=1.5, ttl=60, refresh_ttl=True, only_if_missing=True) + await asyncio.sleep(0) + + assert failure_hook.await_args.kwargs["call_type"] == "async_increment" + assert redis_cache._circuit_breaker._hard_failure_count == 1 + + class _SetRecordingPipeline: def __init__(self) -> None: self.sets: list[tuple[str, str, timedelta | None]] = [] diff --git a/tests/unit/proxy/proxy_server/test_spend_counters.py b/tests/unit/proxy/proxy_server/test_spend_counters.py index d273ef8f42f..d63950ebbdf 100644 --- a/tests/unit/proxy/proxy_server/test_spend_counters.py +++ b/tests/unit/proxy/proxy_server/test_spend_counters.py @@ -22,7 +22,10 @@ Pins covered: from __future__ import annotations import asyncio -from collections.abc import Awaitable, Callable +import contextlib +import gc +import threading +from collections.abc import Awaitable, Callable, Collection from datetime import datetime from typing import Final, Literal, cast from unittest.mock import AsyncMock, MagicMock, call @@ -36,11 +39,13 @@ from opentelemetry.trace import StatusCode import litellm import litellm.proxy.proxy_server as ps +from litellm._internal_context import current_service_target, service_target from litellm._service_logger import ServiceLogging -from litellm._internal_context import current_service_target +from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache -from litellm.caching.redis_cache import RedisCache, RedisCircuitBreaker -from litellm.constants import GLOBAL_PROXY_SPEND_CACHE_KEY +from litellm.caching.redis_batch import request_redis_batch_scope +from litellm.caching.redis_cache import RedisCache, RedisCircuitBreaker, RedisCircuitBreakerOpenError +from litellm.constants import AUTH_OBJECTS_TARGET, GLOBAL_PROXY_SPEND_CACHE_KEY from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig from litellm.models.end_user import LiteLLM_EndUserTable from litellm.models.tag import LiteLLM_TagTable @@ -49,7 +54,14 @@ from litellm.models.user import LiteLLM_UserTable from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_local from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, end_user_cache_key, tag_cache_key +from litellm.proxy.db import spend_counter_reseed from litellm.proxy.management_endpoints.team_endpoints import refresh_cached_team +from litellm.proxy.spend_tracking.spend_counter_batch import ( + SPEND_COUNTERS_TARGET, + active_spend_counter_batch, + read_batched_spend_counter, + spend_counter_batch_scope, +) from .conftest import normalize @@ -1805,10 +1817,10 @@ class _FakeRedisRunningIncrementIfPresent(FakeAsyncRedis): class _RedisCacheOnClient(RedisCache): - def __init__(self, client: FakeAsyncRedis) -> None: + def __init__(self, client: FakeAsyncRedis, breaker: RedisCircuitBreaker | None = None) -> None: self._injected_client: Final = client self.namespace = None - self._circuit_breaker = RedisCircuitBreaker(failure_threshold=5, recovery_timeout=30) + self._circuit_breaker = breaker or RedisCircuitBreaker(failure_threshold=5, recovery_timeout=30) self.service_logger_obj = ServiceLogging() self.default_ttl = None @@ -1817,7 +1829,9 @@ class _RedisCacheOnClient(RedisCache): class _FakeRedisHoldingOneReply(FakeAsyncRedis): - def __init__(self, server: FakeServer, held_increment: float, applied: asyncio.Event, released: asyncio.Event) -> None: + def __init__( + self, server: FakeServer, held_increment: float, applied: asyncio.Event, released: asyncio.Event + ) -> None: super().__init__(server=server) self._held_increment: Final = held_increment self._applied: Final = applied @@ -1970,9 +1984,7 @@ async def test_update_cache_traces_the_global_spend_redis_call_under_the_request for span in exporter.get_finished_spans() if span.name == "redis" and span.attributes is not None and span.parent is not None ) - assert redis_spans == ( - ("async_increment_if_present", expected_status, request_span.get_span_context().span_id), - ) + assert redis_spans == (("async_increment_if_present", expected_status, request_span.get_span_context().span_id),) class _UnwritableKeyUserApiKeyCache(UserApiKeyCache): @@ -2013,3 +2025,1600 @@ async def test_update_cache_still_adds_spend_to_the_others_when_one_cached_objec 1.25, 10.25, ) + + +_COLD_SEED_SOURCE: Final = "end_user_id:customer-1" + + +class _SharedRedis: + """The counter commands a cold start sends, yielding like a network call so overlapping requests interleave.""" + + circuit_breaker_closed: Final = True + + def __init__(self) -> None: + self.values: Final[dict[str, float]] = {} # mutable-ok: stands in for the Redis keyspace + self.failing = False + + async def async_batch_get_cache(self, key_list: list[str], **_: object) -> dict[str, float | None]: + await asyncio.sleep(0) + return {key: self.values.get(key) for key in key_list} + + async def async_get_cache(self, key: str, **_: object) -> float | None: + await asyncio.sleep(0) + return self.values.get(key) + + def _apply( + self, + key: str, + value: float, + only_if_missing: bool = False, + on_added: Callable[[], None] | None = None, + **_: object, + ) -> float: + if not only_if_missing or key not in self.values: + self.values[key] = self.values.get(key, 0.0) + value + if on_added is not None: + on_added() + return self.values[key] + + async def async_increment(self, key: str, value: float, **kwargs: object) -> float: + await asyncio.sleep(0) + if self.failing: + raise ConnectionError("redis down") + return self._apply(key, value, **kwargs) + + async def async_delete_cache(self, key: str) -> None: + await asyncio.sleep(0) + if not self.failing: + self.values.pop(key, None) + + +def _process( + monkeypatch: pytest.MonkeyPatch, redis: _SharedRedis | RedisCache | None, cached_spend: float +) -> DualCache: + counters: Final = DualCache(in_memory_cache=InMemoryCache(), redis_cache=cast(RedisCache, redis)) + objects: Final = DualCache(in_memory_cache=InMemoryCache()) + objects.in_memory_cache.set_cache(key=_COLD_SEED_SOURCE, value={"spend": cached_spend}) + monkeypatch.setattr(ps, "spend_counter_cache", counters) + monkeypatch.setattr(ps, "user_api_key_cache", objects) + monkeypatch.setattr(ps, "prisma_client", object()) + return counters + + +async def _cold_start(redis: _SharedRedis | RedisCache | None, counter: str) -> None: + with spend_counter_batch_scope(cast(RedisCache, redis) if redis is not None else None, frozenset((counter,))): + await ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=_COLD_SEED_SOURCE) + + +async def _read_cold(counter: str) -> None: + batch: Final = active_spend_counter_batch() + assert batch is not None + assert await batch.read(counter) == (None, True) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("with_redis", (True, False)) +async def test_overlapping_cold_starts_in_one_process_add_its_cached_spend_once( + monkeypatch: pytest.MonkeyPatch, with_redis: bool +) -> None: + counter: Final = f"spend:end_user:once-{with_redis}" + redis: Final = _SharedRedis() if with_redis else None + counters: Final = _process(monkeypatch, redis, cached_spend=6.0) + + await asyncio.gather(*(_cold_start(redis, counter) for _ in range(5))) + + seeded: Final[float | None] = ( + redis.values.get(counter) if redis is not None else counters.in_memory_cache.get_cache(key=counter) + ) + assert seeded == 6.0 + + +@pytest.mark.asyncio +async def test_a_request_that_read_the_counter_before_this_process_seeded_it_does_not_add_again( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:stale-reader" + redis: Final = _SharedRedis() + _process(monkeypatch, redis, cached_spend=6.0) + read: Final = asyncio.Event() + + async def _stale_request() -> None: + with spend_counter_batch_scope(cast(RedisCache, redis), frozenset((counter,))): + await _read_cold(counter) + read.set() + await seeded.wait() + await ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=_COLD_SEED_SOURCE) + assert await read_batched_spend_counter(counter) == (6.0, True) + + seeded: Final = asyncio.Event() + + async def _first_request() -> None: + await read.wait() + await _cold_start(redis, counter) + seeded.set() + + await asyncio.gather(_stale_request(), _first_request()) + + assert redis.values[counter] == 6.0 + + +@pytest.mark.asyncio +async def test_a_counter_lost_again_after_this_process_seeded_it_is_seeded_again( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:lost-again" + redis: Final = _SharedRedis() + _process(monkeypatch, redis, cached_spend=6.0) + read: Final = asyncio.Event() + seeded: Final = asyncio.Event() + + async def _stale_request() -> None: + with spend_counter_batch_scope(cast(RedisCache, redis), frozenset((counter,))): + await _read_cold(counter) + read.set() + await seeded.wait() + del redis.values[counter] + await ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=_COLD_SEED_SOURCE) + + async def _first_request() -> None: + await read.wait() + await _cold_start(redis, counter) + seeded.set() + + await asyncio.gather(_stale_request(), _first_request()) + + assert redis.values[counter] == 6.0 + + +@pytest.mark.asyncio +async def test_another_process_seeding_after_this_request_read_the_counter_still_sums( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:two-processes" + redis: Final = _SharedRedis() + read: Final = asyncio.Event() + a_seeded: Final = asyncio.Event() + + async def _process_b() -> None: + with spend_counter_batch_scope(cast(RedisCache, redis), frozenset((counter,))): + await _read_cold(counter) + read.set() + await a_seeded.wait() + _process(monkeypatch, redis, cached_spend=5.0) + await ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=_COLD_SEED_SOURCE) + + async def _process_a() -> None: + await read.wait() + _process(monkeypatch, redis, cached_spend=3.0) + await _cold_start(redis, counter) + a_seeded.set() + + await asyncio.gather(_process_b(), _process_a()) + + assert redis.values[counter] == 8.0 + + +@pytest.mark.asyncio +async def test_seeds_that_keep_failing_fail_every_cold_start_and_the_next_cold_start_seeds( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:failed-seed" + redis: Final = _SharedRedis() + _process(monkeypatch, redis, cached_spend=6.0) + redis.failing = True + + outcomes: Final = await asyncio.gather(*(_cold_start(redis, counter) for _ in range(3)), return_exceptions=True) + redis.failing = False + await _cold_start(redis, counter) + + assert [type(outcome) for outcome in outcomes] == [ConnectionError] * 3 + assert redis.values[counter] == 6.0 + + +class _HeldReplyRedis(_SharedRedis): + """Applies an increment at once but holds its reply until released, like a slow network round trip.""" + + def __init__(self) -> None: + super().__init__() + self.applied: Final = asyncio.Event() + self.release: Final = asyncio.Event() + + async def async_increment(self, key: str, value: float, **kwargs: object) -> float: + total: Final = self._apply(key, value, **kwargs) + self.applied.set() + await self.release.wait() + return total + + +@pytest.mark.asyncio +async def test_a_counter_deleted_before_an_earlier_add_returns_is_seeded_by_a_later_cold_reader( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:deleted-mid-add" + redis: Final = _HeldReplyRedis() + _process(monkeypatch, redis, cached_spend=6.0) + + first: Final = asyncio.ensure_future(_cold_start(redis, counter)) + await redis.applied.wait() + del redis.values[counter] + later: Final = asyncio.ensure_future(_cold_start(redis, counter)) + for _ in range(20): + await asyncio.sleep(0) + redis.release.set() + await asyncio.gather(first, later) + + assert redis.values[counter] == 6.0 + + +def test_a_cold_start_on_another_event_loop_adds_on_its_own(monkeypatch: pytest.MonkeyPatch) -> None: + counter: Final = "spend:end_user:two-loops" + monkeypatch.setattr(spend_counter_reseed, "SPEND_COUNTER_COLD_SEED_JOIN_SECONDS", 60.0) + held: Final = threading.Event() + release: Final = threading.Event() + + class _ThreadHeldRedis(_SharedRedis): + async def async_increment(self, key: str, value: float, **kwargs: object) -> float: + if threading.current_thread() is not threading.main_thread(): + held.set() + while not release.is_set(): + await asyncio.sleep(0.001) + return await super().async_increment(key, value, **kwargs) + + redis: Final = _ThreadHeldRedis() + _process(monkeypatch, redis, cached_spend=6.0) + other_loop: Final = threading.Thread(target=lambda: asyncio.run(_cold_start(redis, counter)), daemon=True) + other_loop.start() + assert held.wait(5) + + try: + asyncio.run(_cold_start(redis, counter)) + finally: + release.set() + other_loop.join(5) + + assert redis.values[counter] == 12.0 + + +@pytest.mark.asyncio +async def test_a_request_that_read_the_counter_cold_after_this_process_seeded_it_still_adds_its_cached_spend( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:read-after-seed" + redis: Final = _SharedRedis() + _process(monkeypatch, redis, cached_spend=6.0) + opened: Final = asyncio.Event() + seeded: Final = asyncio.Event() + + async def _late_reader() -> None: + with spend_counter_batch_scope(cast(RedisCache, redis), frozenset((counter,))): + opened.set() + await seeded.wait() + del redis.values[counter] + await _read_cold(counter) + redis.values[counter] = 0.5 + await ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=_COLD_SEED_SOURCE) + + async def _first_request() -> None: + await opened.wait() + await _cold_start(redis, counter) + seeded.set() + + await asyncio.gather(_late_reader(), _first_request()) + + assert redis.values[counter] == 6.5 + + +@pytest.mark.asyncio +async def test_a_request_whose_pipelined_read_returned_before_this_process_seeded_the_counter_does_not_add_again( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:pipelined-reader" + client: Final = FakeAsyncRedis(server=FakeServer()) + redis: Final = _RedisCacheOnClient(client) + _process(monkeypatch, redis, cached_spend=6.0) + read: Final = asyncio.Event() + seeded: Final = asyncio.Event() + + async def _pipelined_request() -> None: + with request_redis_batch_scope() as batches, spend_counter_batch_scope(redis, frozenset((counter,))): + await batches.batch(redis).flush() + read.set() + await seeded.wait() + await ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=_COLD_SEED_SOURCE) + + async def _first_request() -> None: + await read.wait() + await _cold_start(redis, counter) + seeded.set() + + await asyncio.gather(_pipelined_request(), _first_request()) + + assert float(await client.get(counter) or 0) == 6.0 + + +class _DroppedFirstIncrementRedis(_SharedRedis): + """Resets the connection that carries the first increment and answers every later one.""" + + def __init__(self) -> None: + super().__init__() + self.dropped = False + + async def async_increment(self, key: str, value: float, **kwargs: object) -> float: + if not self.dropped: + self.dropped = True + await asyncio.sleep(0) + raise ConnectionError("connection reset by peer") + return await super().async_increment(key, value, **kwargs) + + +@pytest.mark.asyncio +async def test_requests_that_joined_a_seed_whose_connection_reset_seed_again_and_only_its_starter_fails( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:reset-seed" + redis: Final = _DroppedFirstIncrementRedis() + _process(monkeypatch, redis, cached_spend=6.0) + + outcomes: Final = await asyncio.gather(*(_cold_start(redis, counter) for _ in range(4)), return_exceptions=True) + + failures: Final = [outcome for outcome in outcomes if outcome is not None] + assert [type(failure) for failure in failures] == [ConnectionError], outcomes + assert redis.values[counter] == 6.0 + + +class _StalledFirstIncrementRedis(_SharedRedis): + """Never answers the first increment until released, like a connection that stopped responding.""" + + def __init__(self) -> None: + super().__init__() + self.stalled: Final = asyncio.Event() + self.release: Final = asyncio.Event() + + async def async_increment(self, key: str, value: float, **kwargs: object) -> float: + if not self.stalled.is_set(): + self.stalled.set() + await self.release.wait() + raise TimeoutError("redis did not answer") + return await super().async_increment(key, value, **kwargs) + + +@pytest.mark.asyncio +async def test_cold_starts_stop_waiting_on_a_stalled_seed_and_set_the_missing_counter_once( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:stalled-seed" + redis: Final = _StalledFirstIncrementRedis() + _process(monkeypatch, redis, cached_spend=6.0) + stalled: Final = asyncio.ensure_future(_cold_start(redis, counter)) + await redis.stalled.wait() + + try: + await asyncio.wait_for(asyncio.gather(*(_cold_start(redis, counter) for _ in range(3))), timeout=10) + seeded_while_stalled: Final = redis.values.get(counter) + finally: + redis.release.set() + + with pytest.raises(TimeoutError): + await stalled + assert seeded_while_stalled == 6.0 + + +class _FakeRedisHoldingIncrements(FakeAsyncRedis): + def __init__(self, server: FakeServer) -> None: + super().__init__(server=server) + self.holding: Final = asyncio.Event() + self.release: Final = asyncio.Event() + + async def incrbyfloat(self, name: str, amount: float = 1.0) -> float: + self.holding.set() + await self.release.wait() + return await super().incrbyfloat(name, amount) + + +@pytest.mark.asyncio +async def test_a_cold_start_after_the_redis_breaker_opened_fails_fast_instead_of_using_an_add_in_flight( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:breaker-open" + breaker: Final = RedisCircuitBreaker(failure_threshold=1, recovery_timeout=30) + client: Final = _FakeRedisHoldingIncrements(FakeServer()) + redis: Final = _RedisCacheOnClient(client, breaker) + _process(monkeypatch, redis, cached_spend=6.0) + read: Final = asyncio.Event() + opened: Final = asyncio.Event() + + async def _read_before_the_seed_started() -> None: + with spend_counter_batch_scope(redis, frozenset((counter,))): + await _read_cold(counter) + read.set() + await opened.wait() + await ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=_COLD_SEED_SOURCE) + + later: Final = asyncio.ensure_future(_read_before_the_seed_started()) + await read.wait() + first: Final = asyncio.ensure_future(_cold_start(redis, counter)) + await client.holding.wait() + breaker.record_failure() + opened.set() + for _ in range(100): + await asyncio.sleep(0) + client.release.set() + + outcomes: Final = await asyncio.gather(first, later, return_exceptions=True) + + assert outcomes[0] is None and isinstance(outcomes[1], RedisCircuitBreakerOpenError), outcomes + + +class _HeldFirstMGetRedis(_SharedRedis): + """Reads the first MGET's keys at once but holds its reply until released, like a slow network round trip.""" + + def __init__(self) -> None: + super().__init__() + self.mget_read: Final = asyncio.Event() + self.release_mget: Final = asyncio.Event() + + async def async_batch_get_cache(self, key_list: list[str], **_: object) -> dict[str, float | None]: + values: Final = {key: self.values.get(key) for key in key_list} + if not self.mget_read.is_set(): + self.mget_read.set() + await self.release_mget.wait() + return values + + +@pytest.mark.asyncio +async def test_a_seed_that_lands_while_a_cold_read_is_on_its_way_back_is_not_added_again( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:seed-during-read" + redis: Final = _HeldFirstMGetRedis() + _process(monkeypatch, redis, cached_spend=6.0) + + async def _first_request() -> None: + await redis.mget_read.wait() + await _cold_start(redis, counter) + redis.release_mget.set() + + await asyncio.gather(_cold_start(redis, counter), _first_request()) + + assert redis.values[counter] == 6.0 + + +class _HeldGetRedis(_SharedRedis): + """Answers the first GET with the value its key had when it arrived, holding the reply until released.""" + + def __init__(self) -> None: + super().__init__() + self.get_arrived: Final = asyncio.Event() + self.release_get: Final = asyncio.Event() + + async def async_get_cache(self, key: str, **_: object) -> float | None: + value: Final = self.values.get(key) + if not self.get_arrived.is_set(): + self.get_arrived.set() + await self.release_get.wait() + return value + + +@pytest.mark.asyncio +async def test_a_seed_that_lands_while_an_unbatched_cold_read_is_on_its_way_back_is_not_added_again( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:seed-during-unbatched-read" + redis: Final = _HeldGetRedis() + _process(monkeypatch, redis, cached_spend=6.0) + + async def _first_request() -> None: + await redis.get_arrived.wait() + await _cold_start(redis, counter) + redis.release_get.set() + + await asyncio.gather( + ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=_COLD_SEED_SOURCE), + _first_request(), + ) + + assert redis.values[counter] == 6.0 + + +class _QueuedReadRedis(_HeldReplyRedis): + """Holds the first read before Redis runs it, while increments apply at once and hold their replies.""" + + def __init__(self) -> None: + super().__init__() + self.read_sent: Final = asyncio.Event() + self.run_read: Final = asyncio.Event() + + async def _queue_the_first_read(self) -> None: + if not self.read_sent.is_set(): + self.read_sent.set() + await self.run_read.wait() + + async def async_batch_get_cache(self, key_list: list[str], **_: object) -> dict[str, float | None]: + await self._queue_the_first_read() + return {key: self.values.get(key) for key in key_list} + + async def async_get_cache(self, key: str, **_: object) -> float | None: + await self._queue_the_first_read() + return self.values.get(key) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("batched", (True, False)) +async def test_a_cold_read_sent_before_an_add_started_does_not_join_it_when_the_counter_was_deleted_after_it_landed( + monkeypatch: pytest.MonkeyPatch, batched: bool +) -> None: + counter: Final = f"spend:end_user:deleted-while-read-queued-{batched}" + redis: Final = _QueuedReadRedis() + _process(monkeypatch, redis, cached_spend=6.0) + + reader: Final = asyncio.ensure_future( + _cold_start(redis, counter) + if batched + else ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=_COLD_SEED_SOURCE) + ) + await redis.read_sent.wait() + first: Final = asyncio.ensure_future(_cold_start(redis, counter)) + await redis.applied.wait() + del redis.values[counter] + redis.run_read.set() + for _ in range(20): + await asyncio.sleep(0) + redis.release.set() + await asyncio.gather(reader, first) + + assert redis.values.get(counter) == 6.0 + + +@pytest.mark.asyncio +async def test_a_counter_lost_and_re_created_below_the_seed_gets_the_cached_spend_on_top( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:re-created-below" + redis: Final = _SharedRedis() + _process(monkeypatch, redis, cached_spend=6.0) + read: Final = asyncio.Event() + seeded: Final = asyncio.Event() + + async def _stale_request() -> None: + with spend_counter_batch_scope(cast(RedisCache, redis), frozenset((counter,))): + await _read_cold(counter) + read.set() + await seeded.wait() + redis.values[counter] = 0.5 + await ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=_COLD_SEED_SOURCE) + assert await read_batched_spend_counter(counter) == (6.5, True) + + async def _first_request() -> None: + await read.wait() + await _cold_start(redis, counter) + seeded.set() + + await asyncio.gather(_stale_request(), _first_request()) + + assert redis.values[counter] == 6.5 + + +@pytest.mark.asyncio +async def test_a_counter_holding_the_seed_is_left_alone_by_a_request_whose_source_cache_moved_past_it( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:source-moved" + redis: Final = _SharedRedis() + counters: Final = _process(monkeypatch, redis, cached_spend=6.0) + read: Final = asyncio.Event() + seeded: Final = asyncio.Event() + + async def _stale_request() -> None: + with spend_counter_batch_scope(cast(RedisCache, redis), frozenset((counter,))): + await _read_cold(counter) + read.set() + await seeded.wait() + ps.user_api_key_cache.in_memory_cache.set_cache(key=_COLD_SEED_SOURCE, value={"spend": 6.25}) + await ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=_COLD_SEED_SOURCE) + + async def _first_request() -> None: + await read.wait() + await _cold_start(redis, counter) + seeded.set() + + await asyncio.gather(_stale_request(), _first_request()) + + assert redis.values[counter] == 6.0 + assert counters.in_memory_cache.get_cache(key=counter) == 6.0 + + +@pytest.mark.asyncio +async def test_a_request_that_left_a_peers_counter_alone_does_not_cost_a_later_read_its_add( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:left-alone-then-fresh" + redis: Final = _SharedRedis() + _process(monkeypatch, redis, cached_spend=6.0) + stale_read: Final = asyncio.Event() + seeded: Final = asyncio.Event() + fresh_read: Final = asyncio.Event() + left_alone: Final = asyncio.Event() + + async def _stale_request() -> None: + with spend_counter_batch_scope(cast(RedisCache, redis), frozenset((counter,))): + await _read_cold(counter) + stale_read.set() + await fresh_read.wait() + await ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=_COLD_SEED_SOURCE) + left_alone.set() + + async def _first_request() -> None: + await stale_read.wait() + await _cold_start(redis, counter) + seeded.set() + + async def _fresh_request() -> None: + await seeded.wait() + with spend_counter_batch_scope(cast(RedisCache, redis), frozenset((counter,))): + del redis.values[counter] + await _read_cold(counter) + redis.values[counter] = 7.0 + fresh_read.set() + await left_alone.wait() + await ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=_COLD_SEED_SOURCE) + + await asyncio.gather(_stale_request(), _first_request(), _fresh_request()) + + assert redis.values[counter] == 13.0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("read_back", "expected"), + [ + pytest.param(0.00123, 0.00123, id="rounded-by-incrbyfloat-is-left-alone"), + pytest.param(0.0012299, 0.0024599000000000002, id="short-by-more-than-rounding-gets-the-add"), + ], +) +async def test_a_counter_reading_back_just_under_the_spend_is_only_topped_up_when_really_short( + monkeypatch: pytest.MonkeyPatch, read_back: float, expected: float +) -> None: + counter: Final = "spend:end_user:read-back-rounding" + redis: Final = _SharedRedis() + _process(monkeypatch, redis, cached_spend=0.0012300000000000002) + read: Final = asyncio.Event() + seeded: Final = asyncio.Event() + + async def _stale_request() -> None: + with spend_counter_batch_scope(cast(RedisCache, redis), frozenset((counter,))): + await _read_cold(counter) + read.set() + await seeded.wait() + redis.values[counter] = read_back + await ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=_COLD_SEED_SOURCE) + + async def _first_request() -> None: + await read.wait() + await _cold_start(redis, counter) + seeded.set() + + await asyncio.gather(_stale_request(), _first_request()) + + assert redis.values[counter] == pytest.approx(expected, abs=1e-15) + + +@pytest.mark.asyncio +async def test_a_counter_set_while_the_stalled_seed_is_cancelled_is_not_added_to_by_a_request_that_read_it_cold( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:set-then-stalled-cancelled" + monkeypatch.setattr(spend_counter_reseed, "SPEND_COUNTER_COLD_SEED_JOIN_SECONDS", 0.01) + redis: Final = _HeldIncrementRedis(held_increments=1) + _process(monkeypatch, redis, cached_spend=6.0) + read: Final = asyncio.Event() + set_by_a_waiter: Final = asyncio.Event() + + async def _read_before_the_set() -> None: + with spend_counter_batch_scope(cast(RedisCache, redis), frozenset((counter,))): + await _read_cold(counter) + read.set() + await set_by_a_waiter.wait() + await ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=_COLD_SEED_SOURCE) + + reader: Final = asyncio.ensure_future(_read_before_the_set()) + await read.wait() + stalled: Final = asyncio.ensure_future(_cold_start(redis, counter)) + await redis.arrived[0].wait() + await _cold_start(redis, counter) + stalled.cancel() + with contextlib.suppress(asyncio.CancelledError): + await stalled + set_by_a_waiter.set() + await reader + + assert redis.values[counter] == 6.0 + + +@pytest.mark.asyncio +async def test_a_request_done_waiting_on_a_stalled_seed_adds_when_the_counter_holds_less_than_its_cached_spend( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:below-own-spend" + monkeypatch.setattr(spend_counter_reseed, "SPEND_COUNTER_COLD_SEED_JOIN_SECONDS", 0.01) + redis: Final = _HeldIncrementRedis(held_increments=1) + _process(monkeypatch, redis, cached_spend=6.0) + stalled: Final = asyncio.ensure_future(_cold_start(redis, counter)) + await redis.arrived[0].wait() + + with spend_counter_batch_scope(cast(RedisCache, redis), frozenset((counter,))): + await _read_cold(counter) + redis.values[counter] = 0.5 + await ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=_COLD_SEED_SOURCE) + while_stalled: Final = redis.values[counter] + redis.releases[0].set() + await stalled + + assert while_stalled == 6.5 + + +@pytest.mark.asyncio +async def test_a_counter_without_redis_set_only_if_missing_keeps_the_value_it_holds( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:no-redis-set-if-missing" + monkeypatch.setattr(ps, "spend_counter_cache", DualCache(in_memory_cache=InMemoryCache())) + + created: Final[list[str]] = [] # mutable-ok: the created calls the test asserts on + first: Final = await ps.increment_spend_counter_cache( + counter, 6.0, only_if_missing=True, on_added=lambda: created.append("first") + ) + second: Final = await ps.increment_spend_counter_cache( + counter, 4.0, only_if_missing=True, on_added=lambda: created.append("second") + ) + added: Final = await ps.increment_spend_counter_cache(counter, 1.0) + + assert (first, second, added) == (6.0, 6.0, 7.0) + assert created == ["first"] + + +@pytest.mark.asyncio +async def test_a_counter_set_by_a_request_done_waiting_counts_as_the_seed_for_a_later_request( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:set-counts-as-seed" + monkeypatch.setattr(spend_counter_reseed, "SPEND_COUNTER_COLD_SEED_JOIN_SECONDS", 0.01) + redis: Final = _HeldIncrementRedis(held_increments=1) + _process(monkeypatch, redis, cached_spend=6.0) + read: Final = asyncio.Event() + set_by_a_waiter: Final = asyncio.Event() + + async def _stale_request() -> None: + with spend_counter_batch_scope(cast(RedisCache, redis), frozenset((counter,))): + await _read_cold(counter) + read.set() + await set_by_a_waiter.wait() + redis.values[counter] = 0.5 + await ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=_COLD_SEED_SOURCE) + + stale: Final = asyncio.ensure_future(_stale_request()) + await read.wait() + stalled: Final = asyncio.ensure_future(_cold_start(redis, counter)) + await redis.arrived[0].wait() + await _cold_start(redis, counter) + set_by_a_waiter.set() + await stale + after_the_stale_request: Final = redis.values[counter] + redis.releases[0].set() + await stalled + + assert after_the_stale_request == 6.5 + + +class _CommandCountingRedis(_SharedRedis): + def __init__(self) -> None: + super().__init__() + self.commands: Final[list[str]] = [] # mutable-ok: the commands the test asserts on + + async def async_batch_get_cache(self, key_list: list[str], **_: object) -> dict[str, float | None]: + self.commands.append("mget") + return await super().async_batch_get_cache(key_list) + + async def async_get_cache(self, key: str, **_: object) -> float | None: + self.commands.append("get") + return await super().async_get_cache(key) + + async def async_increment(self, key: str, value: float, **kwargs: object) -> float: + self.commands.append("set_if_missing" if kwargs.get("only_if_missing") else "add") + return await super().async_increment(key, value, **kwargs) + + +@pytest.mark.asyncio +async def test_a_request_that_read_before_a_seed_landed_sends_one_write_and_keeps_its_value_for_later_reads( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:one-write" + redis: Final = _CommandCountingRedis() + _process(monkeypatch, redis, cached_spend=6.0) + read: Final = asyncio.Event() + seeded: Final = asyncio.Event() + + async def _stale_request() -> None: + with spend_counter_batch_scope(cast(RedisCache, redis), frozenset((counter,))): + await _read_cold(counter) + read.set() + await seeded.wait() + redis.commands.clear() + await ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=_COLD_SEED_SOURCE) + assert await read_batched_spend_counter(counter) == (6.0, True) + assert await ps.get_current_spend(counter_key=counter, fallback_spend=6.0) == 6.0 + + async def _first_request() -> None: + await read.wait() + await _cold_start(redis, counter) + seeded.set() + + await asyncio.gather(_stale_request(), _first_request()) + + assert redis.commands == ["set_if_missing"] + + +@pytest.mark.asyncio +async def test_a_request_waits_on_other_requests_adds_for_the_limit_in_total_across_its_counters( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(spend_counter_reseed, "SPEND_COUNTER_COLD_SEED_JOIN_SECONDS", 0.3) + keys: Final = tuple(f"spend:end_user:shared-wait-{index}" for index in range(4)) + redis: Final = _HeldIncrementRedis(held_increments=len(keys)) + _process(monkeypatch, redis, cached_spend=6.0) + stalled: Final = [asyncio.ensure_future(_cold_start(redis, key)) for key in keys] + for arrived in redis.arrived: + await arrived.wait() + + clock: Final = [0.0] # mutable-ok: the time the event loop reports to the wait + limits: Final[list[float]] = [] # mutable-ok: the limit of each wait the test asserts on + + real_wait: Final = asyncio.wait + + async def _waits_on_an_add_take_a_tenth_of_a_second( + waited_on: Collection[asyncio.Future[object]], timeout: float + ) -> tuple[set[asyncio.Future[object]], set[asyncio.Future[object]]]: + if any(isinstance(future, asyncio.Task) for future in waited_on): + return await real_wait(waited_on, timeout=timeout) + limits.append(timeout) + clock[0] += 0.1 + return set(), set(waited_on) + + monkeypatch.setattr(asyncio.get_running_loop(), "time", lambda: clock[0]) + monkeypatch.setattr(asyncio, "wait", _waits_on_an_add_take_a_tenth_of_a_second) + with spend_counter_batch_scope(cast(RedisCache, redis), frozenset(keys)): + for key in keys: + await ps.ensure_spend_counter_initialized(counter_key=key, source_cache_key=_COLD_SEED_SOURCE) + monkeypatch.undo() + for release in redis.releases: + release.set() + await asyncio.gather(*stalled) + + assert limits == pytest.approx([0.3, 0.2, 0.1, 0.0]) + + +class _LostAndRecreatedAfterFailedMGetRedis(_SharedRedis): + """Holds the first MGET until released, then loses the counter and resets that MGET. Once the request has read the + counter cold with its two GETs, another process's post-call increment lands before the request's next command.""" + + def __init__(self, counter: str) -> None: + super().__init__() + self.counter: Final = counter + self.mget_sent: Final = asyncio.Event() + self.release_mget: Final = asyncio.Event() + self.cold_gets = 0 + + def _other_process_increments_after_the_cold_gets(self) -> None: + if self.cold_gets == 2: + self.cold_gets += 1 + self.values[self.counter] = self.values.get(self.counter, 0.0) + 0.3 + + async def async_batch_get_cache(self, key_list: list[str], **_: object) -> dict[str, float | None]: + if not self.mget_sent.is_set(): + self.mget_sent.set() + await self.release_mget.wait() + del self.values[self.counter] + raise ConnectionError("connection reset") + return await super().async_batch_get_cache(key_list) + + async def async_get_cache(self, key: str, **_: object) -> float | None: + self._other_process_increments_after_the_cold_gets() + if self.release_mget.is_set() and self.cold_gets < 2: + self.cold_gets += 1 + return await super().async_get_cache(key) + + async def async_increment(self, key: str, value: float, **kwargs: object) -> float: + self._other_process_increments_after_the_cold_gets() + return await super().async_increment(key, value, **kwargs) + + +@pytest.mark.asyncio +async def test_a_request_whose_mget_failed_and_whose_gets_went_out_after_this_process_seeded_adds_its_cached_spend( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:failed-mget" + redis: Final = _LostAndRecreatedAfterFailedMGetRedis(counter) + _process(monkeypatch, redis, cached_spend=6.0) + + request: Final = asyncio.ensure_future(_cold_start(redis, counter)) + await redis.mget_sent.wait() + await _cold_start(redis, counter) + redis.release_mget.set() + await request + + assert redis.values[counter] == pytest.approx(6.3) + + +class _AppliedHeldThenResetRedis(_HeldReplyRedis): + """Applies an increment at once, holds its reply until released, then resets the connection instead of replying.""" + + def __init__(self) -> None: + super().__init__() + self.deleted = False + + async def async_increment(self, key: str, value: float, **kwargs: object) -> float: + await super().async_increment(key, value, **kwargs) + raise ConnectionError("connection reset") + + async def async_delete_cache(self, key: str) -> None: + self.deleted = True + await super().async_delete_cache(key) + + +@pytest.mark.asyncio +async def test_a_cold_seed_whose_request_was_cancelled_does_not_delete_the_counter_later( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:cancelled-seed" + redis: Final = _AppliedHeldThenResetRedis() + _process(monkeypatch, redis, cached_spend=6.0) + + request: Final = asyncio.ensure_future(_cold_start(redis, counter)) + await redis.applied.wait() + request.cancel() + await asyncio.wait((request,)) + redis.release.set() + for _ in range(20): + await asyncio.sleep(0) + + assert request.cancelled() + assert redis.deleted is False + assert redis.values[counter] == 6.0 + + +@pytest.mark.asyncio +async def test_requests_waiting_on_a_cold_seed_whose_request_was_cancelled_add_their_cached_spend_once( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:cancelled-seed-with-waiters" + monkeypatch.setattr(spend_counter_reseed, "SPEND_COUNTER_COLD_SEED_JOIN_SECONDS", 60.0) + redis: Final = _StalledFirstIncrementRedis() + _process(monkeypatch, redis, cached_spend=6.0) + read: Final = asyncio.Event() + started: Final = asyncio.Event() + + async def _read_before_the_seed_started() -> None: + with spend_counter_batch_scope(cast(RedisCache, redis), frozenset((counter,))): + await _read_cold(counter) + read.set() + await started.wait() + await ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=_COLD_SEED_SOURCE) + + joined_before_start: Final = asyncio.ensure_future(_read_before_the_seed_started()) + await read.wait() + seed: Final = asyncio.ensure_future(_cold_start(redis, counter)) + await redis.stalled.wait() + started.set() + joined_after_start: Final = asyncio.ensure_future(_cold_start(redis, counter)) + for _ in range(20): + await asyncio.sleep(0) + seed.cancel() + await asyncio.wait_for(asyncio.gather(joined_before_start, joined_after_start), timeout=5) + + assert seed.cancelled() + assert redis.values[counter] == 6.0 + + +class _HeldIncrementRedis(_SharedRedis): + """Holds each increment before Redis runs it until its release event in ``releases`` is set, counting reads.""" + + def __init__(self, held_increments: int) -> None: + super().__init__() + self.reads = 0 + self.arrived: Final = tuple(asyncio.Event() for _ in range(held_increments)) + self.releases: Final = tuple(asyncio.Event() for _ in range(held_increments)) + self.increments = 0 + + async def async_batch_get_cache(self, key_list: list[str], **_: object) -> dict[str, float | None]: + self.reads += 1 + return await super().async_batch_get_cache(key_list) + + async def async_get_cache(self, key: str, **_: object) -> float | None: + self.reads += 1 + return await super().async_get_cache(key) + + async def async_increment(self, key: str, value: float, **kwargs: object) -> float: + index: Final = self.increments + self.increments += 1 + if index < len(self.arrived): + self.arrived[index].set() + await self.releases[index].wait() + return await super().async_increment(key, value, **kwargs) + + +@pytest.mark.asyncio +async def test_a_request_that_read_the_counter_before_a_seed_started_takes_its_value_without_reading_again( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:joined-seed-value" + redis: Final = _HeldIncrementRedis(held_increments=1) + _process(monkeypatch, redis, cached_spend=6.0) + read: Final = asyncio.Event() + started: Final = asyncio.Event() + + async def _joined_request() -> None: + with spend_counter_batch_scope(cast(RedisCache, redis), frozenset((counter,))): + await _read_cold(counter) + read.set() + await started.wait() + reads_before_the_seed_landed: Final = redis.reads + ensuring: Final = asyncio.ensure_future( + ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=_COLD_SEED_SOURCE) + ) + for _ in range(10): + await asyncio.sleep(0) + redis.releases[0].set() + await ensuring + assert await read_batched_spend_counter(counter) == (6.0, True) + assert redis.reads == reads_before_the_seed_landed + + joined: Final = asyncio.ensure_future(_joined_request()) + await read.wait() + first: Final = asyncio.ensure_future(_cold_start(redis, counter)) + await redis.arrived[0].wait() + started.set() + await asyncio.gather(joined, first) + + assert redis.values[counter] == 6.0 + + +@pytest.mark.asyncio +async def test_a_request_that_read_the_counter_while_a_seed_was_in_flight_does_not_add_again_after_it_lands( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:read-during-seed" + redis: Final = _HeldIncrementRedis(held_increments=1) + _process(monkeypatch, redis, cached_spend=6.0) + + first: Final = asyncio.ensure_future(_cold_start(redis, counter)) + await redis.arrived[0].wait() + during: Final = asyncio.ensure_future(_cold_start(redis, counter)) + for _ in range(10): + await asyncio.sleep(0) + redis.releases[0].set() + await asyncio.gather(first, during) + + assert redis.values[counter] == 6.0 + + +@pytest.mark.asyncio +async def test_a_seed_that_lands_late_in_the_wait_is_not_added_again(monkeypatch: pytest.MonkeyPatch) -> None: + counter: Final = "spend:end_user:late-landing-seed" + monkeypatch.setattr(spend_counter_reseed, "SPEND_COUNTER_COLD_SEED_JOIN_SECONDS", 0.5) + redis: Final = _HeldIncrementRedis(held_increments=1) + _process(monkeypatch, redis, cached_spend=6.0) + + first: Final = asyncio.ensure_future(_cold_start(redis, counter)) + await redis.arrived[0].wait() + during: Final = asyncio.ensure_future(_cold_start(redis, counter)) + await asyncio.sleep(0.4) + redis.releases[0].set() + await asyncio.gather(first, during) + + assert redis.values[counter] == 6.0 + + +@pytest.mark.asyncio +async def test_a_request_joining_an_add_that_started_after_its_read_came_back_sends_no_write( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:joined-no-write" + monkeypatch.setattr(spend_counter_reseed, "SPEND_COUNTER_COLD_SEED_JOIN_SECONDS", 60.0) + redis: Final = _HeldIncrementRedis(held_increments=1) + _process(monkeypatch, redis, cached_spend=6.0) + read: Final = asyncio.Event() + + async def _joining_request() -> None: + with spend_counter_batch_scope(cast(RedisCache, redis), frozenset((counter,))): + await _read_cold(counter) + read.set() + await redis.arrived[0].wait() + await ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=_COLD_SEED_SOURCE) + + joining: Final = asyncio.ensure_future(_joining_request()) + await read.wait() + first: Final = asyncio.ensure_future(_cold_start(redis, counter)) + await redis.arrived[0].wait() + for _ in range(100): + await asyncio.sleep(0) + redis.releases[0].set() + await asyncio.gather(joining, first) + + assert (redis.values[counter], redis.increments) == (6.0, 1) + + +class _HeldSecondGetRedis(_HeldIncrementRedis): + """Also answers the second GET with the value its key had when it arrived, holding the reply until released.""" + + def __init__(self) -> None: + super().__init__(held_increments=1) + self.gets = 0 + self.second_get_arrived: Final = asyncio.Event() + self.release_second_get: Final = asyncio.Event() + + async def async_get_cache(self, key: str, **_: object) -> float | None: + self.gets += 1 + value: Final = await super().async_get_cache(key) + if self.gets == 2: + self.second_get_arrived.set() + await self.release_second_get.wait() + return value + + +@pytest.mark.asyncio +async def test_a_seed_that_lands_while_the_reseed_read_of_an_unbatched_request_is_out_is_not_added_again( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:seed-during-reseed-read" + redis: Final = _HeldSecondGetRedis() + _process(monkeypatch, redis, cached_spend=6.0) + + first: Final = asyncio.ensure_future(_cold_start(redis, counter)) + await redis.arrived[0].wait() + unbatched: Final = asyncio.ensure_future( + ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=_COLD_SEED_SOURCE) + ) + await redis.second_get_arrived.wait() + redis.releases[0].set() + await first + redis.release_second_get.set() + await unbatched + + assert redis.values[counter] == 6.0 + + +@pytest.mark.asyncio +async def test_a_stalled_seed_finishing_does_not_hide_the_newer_seed_still_in_flight( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:older-seed-finishes" + redis: Final = _HeldIncrementRedis(held_increments=2) + _process(monkeypatch, redis, cached_spend=6.0) + + older: Final = asyncio.ensure_future(_cold_start(redis, counter)) + await redis.arrived[0].wait() + newer: Final = asyncio.ensure_future(_cold_start(redis, counter)) + await redis.arrived[1].wait() + redis.releases[0].set() + await older + del redis.values[counter] + late: Final = asyncio.ensure_future(_cold_start(redis, counter)) + for _ in range(10): + await asyncio.sleep(0) + commands_while_the_newer_seed_was_in_flight: Final = redis.increments + redis.releases[1].set() + await asyncio.gather(newer, late) + + assert commands_while_the_newer_seed_was_in_flight == 2 + assert redis.values[counter] == 6.0 + + +class _TargetRecordingRedis(_SharedRedis): + def __init__(self) -> None: + super().__init__() + self.targets: Final[list[str | None]] = [] # mutable-ok: the targets the test asserts on + + async def async_get_cache(self, key: str, **_: object) -> float | None: + self.targets.append(current_service_target()) + return 1.0 + + +@pytest.mark.asyncio +async def test_the_cold_check_read_reports_under_spend_counters_inside_another_phases_target( + monkeypatch: pytest.MonkeyPatch, +) -> None: + redis: Final = _TargetRecordingRedis() + _process(monkeypatch, redis, cached_spend=6.0) + + with service_target(AUTH_OBJECTS_TARGET): + await ps.ensure_spend_counter_initialized(counter_key="spend:key:target", source_cache_key=_COLD_SEED_SOURCE) + + assert redis.targets == [SPEND_COUNTERS_TARGET] + + +class _HeldSetRedis(_SharedRedis): + """Holds each set on the wire, before Redis runs it, until ``release_set`` is set.""" + + def __init__(self) -> None: + super().__init__() + self.set_sent: Final = asyncio.Event() + self.release_set: Final = asyncio.Event() + + async def async_increment(self, key: str, value: float, **kwargs: object) -> float: + if kwargs.get("only_if_missing"): + self.set_sent.set() + await self.release_set.wait() + return await super().async_increment(key, value, **kwargs) + + +async def _a_request_read_cold_before_this_process_seeded( + redis: _SharedRedis, counter: str, seeded: asyncio.Event, go: asyncio.Event +) -> None: + with spend_counter_batch_scope(cast(RedisCache, redis), frozenset((counter,))): + await _read_cold(counter) + seeded.set() + await go.wait() + await ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=_COLD_SEED_SOURCE) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("peer_value", "expected"), + [ + pytest.param(7.0, 13.0, id="the-set-left-a-peers-counter-alone-so-the-fresh-read-adds"), + pytest.param(None, 6.0, id="the-set-created-the-counter-so-the-fresh-read-joins-it"), + ], +) +async def test_a_fresh_cold_read_joining_a_set_only_skips_its_add_when_that_set_wrote( + monkeypatch: pytest.MonkeyPatch, peer_value: float | None, expected: float +) -> None: + counter: Final = "spend:end_user:fresh-read-joins-a-set" + monkeypatch.setattr(spend_counter_reseed, "SPEND_COUNTER_COLD_SEED_JOIN_SECONDS", 60.0) + redis: Final = _HeldSetRedis() + _process(monkeypatch, redis, cached_spend=6.0) + read_before_seed: Final = asyncio.Event() + go: Final = asyncio.Event() + stale: Final = asyncio.ensure_future( + _a_request_read_cold_before_this_process_seeded(redis, counter, read_before_seed, go) + ) + await read_before_seed.wait() + await _cold_start(redis, counter) + del redis.values[counter] + + async def _fresh_request() -> None: + with spend_counter_batch_scope(cast(RedisCache, redis), frozenset((counter,))): + await _read_cold(counter) + if peer_value is not None: + redis.values[counter] = peer_value + go.set() + await redis.set_sent.wait() + await ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=_COLD_SEED_SOURCE) + + fresh: Final = asyncio.ensure_future(_fresh_request()) + await redis.set_sent.wait() + for _ in range(20): + await asyncio.sleep(0) + redis.release_set.set() + await asyncio.wait_for(asyncio.gather(stale, fresh), timeout=5) + + assert redis.values[counter] == expected + + +@pytest.mark.asyncio +async def test_a_request_cancelled_with_its_set_on_the_wire_still_adds_to_a_short_counter( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:cancelled-with-the-set-out" + redis: Final = _HeldSetRedis() + _process(monkeypatch, redis, cached_spend=6.0) + read_before_seed: Final = asyncio.Event() + go: Final = asyncio.Event() + stale: Final = asyncio.ensure_future( + _a_request_read_cold_before_this_process_seeded(redis, counter, read_before_seed, go) + ) + await read_before_seed.wait() + await _cold_start(redis, counter) + redis.values[counter] = 0.25 + go.set() + await redis.set_sent.wait() + + stale.cancel() + await asyncio.wait((stale,)) + redis.release_set.set() + for _ in range(20): + await asyncio.sleep(0) + + assert stale.cancelled() + assert redis.values[counter] == 6.25 + + +class _AddLandsAfterSetRedis(_HeldSetRedis): + """Also holds a plain add on the wire while ``set_landed`` is clear, so a set still in flight lands first.""" + + def __init__(self) -> None: + super().__init__() + self.set_landed: Final = asyncio.Event() + self.set_landed.set() + + async def async_increment(self, key: str, value: float, **kwargs: object) -> float: + if not kwargs.get("only_if_missing"): + await self.set_landed.wait() + added: Final = await super().async_increment(key, value, **kwargs) + if kwargs.get("only_if_missing"): + self.set_landed.set() + return added + + +@pytest.mark.asyncio +async def test_a_set_still_on_the_wire_after_its_request_was_cancelled_is_joined_by_a_newer_cold_read( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:cancelled-set-joined" + monkeypatch.setattr(spend_counter_reseed, "SPEND_COUNTER_COLD_SEED_JOIN_SECONDS", 60.0) + redis: Final = _AddLandsAfterSetRedis() + _process(monkeypatch, redis, cached_spend=6.0) + read_before_seed: Final = asyncio.Event() + go: Final = asyncio.Event() + stale: Final = asyncio.ensure_future( + _a_request_read_cold_before_this_process_seeded(redis, counter, read_before_seed, go) + ) + await read_before_seed.wait() + await _cold_start(redis, counter) + del redis.values[counter] + redis.set_landed.clear() + go.set() + await redis.set_sent.wait() + stale.cancel() + await asyncio.wait((stale,)) + + newer: Final = asyncio.ensure_future(_cold_start(redis, counter)) + for _ in range(20): + await asyncio.sleep(0) + redis.release_set.set() + await asyncio.wait_for(newer, timeout=5) + + assert stale.cancelled() + assert redis.values[counter] == 6.0 + + +@pytest.mark.asyncio +async def test_a_fresh_read_released_with_a_stale_read_joins_the_set_the_stale_read_starts( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:released-together" + monkeypatch.setattr(spend_counter_reseed, "SPEND_COUNTER_COLD_SEED_JOIN_SECONDS", 60.0) + redis: Final = _AddLandsAfterSetRedis() + _process(monkeypatch, redis, cached_spend=6.0) + stale_read: Final = asyncio.Event() + fresh_read: Final = asyncio.Event() + go: Final = asyncio.Event() + stale: Final = asyncio.ensure_future( + _a_request_read_cold_before_this_process_seeded(redis, counter, stale_read, go) + ) + await stale_read.wait() + await _cold_start(redis, counter) + del redis.values[counter] + redis.set_landed.clear() + fresh: Final = asyncio.ensure_future( + _a_request_read_cold_before_this_process_seeded(redis, counter, fresh_read, go) + ) + await fresh_read.wait() + owner: Final = asyncio.ensure_future(_cold_start(redis, counter)) + for _ in range(20): + await asyncio.sleep(0) + go.set() + for _ in range(20): + await asyncio.sleep(0) + + owner.cancel() + await redis.set_sent.wait() + for _ in range(20): + await asyncio.sleep(0) + redis.release_set.set() + await asyncio.wait_for(asyncio.gather(stale, fresh), timeout=5) + + assert redis.values[counter] == 6.0 + + +class _SetFailsOnReleaseRedis(_HeldSetRedis): + """Holds each set on the wire and fails it once released, while a DEL still reaches Redis.""" + + async def async_increment(self, key: str, value: float, **kwargs: object) -> float: + if kwargs.get("only_if_missing"): + self.set_sent.set() + await self.release_set.wait() + raise ConnectionError("Timeout reading from socket") + return await super().async_increment(key, value, **kwargs) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("cancelled", "counter_left", "in_memory_left"), + [ + pytest.param(True, 9.0, 6.0, id="cancelled-keeps-the-counter-as-mains-cancelled-add-does"), + pytest.param(False, None, None, id="still-waiting-deletes-it-as-mains-failed-add-does"), + ], +) +async def test_a_set_failing_after_its_request_was_cancelled_keeps_the_counter_and_logs_no_stray_task_error( + monkeypatch: pytest.MonkeyPatch, cancelled: bool, counter_left: float | None, in_memory_left: float | None +) -> None: + counter: Final = "spend:end_user:cancelled-set-fails" + redis: Final = _SetFailsOnReleaseRedis() + _process(monkeypatch, redis, cached_spend=6.0) + read_before_seed: Final = asyncio.Event() + go: Final = asyncio.Event() + stale: Final = asyncio.ensure_future( + _a_request_read_cold_before_this_process_seeded(redis, counter, read_before_seed, go) + ) + await read_before_seed.wait() + await _cold_start(redis, counter) + go.set() + await redis.set_sent.wait() + if cancelled: + stale.cancel() + await asyncio.wait((stale,)) + stray: Final[list[object]] = [] # mutable-ok: collects what the loop's exception handler is given + asyncio.get_running_loop().set_exception_handler(lambda _, context: stray.append(context)) + redis.values[counter] = 9.0 + redis.release_set.set() + await asyncio.wait((stale,)) + for _ in range(20): + await asyncio.sleep(0) + gc.collect() + await asyncio.sleep(0) + + assert redis.values.get(counter) == counter_left + assert ps.spend_counter_cache.in_memory_cache.get_cache(key=counter) == in_memory_left + assert stale.cancelled() or isinstance(stale.exception(), ConnectionError) + assert stray == [] + + +@pytest.mark.asyncio +async def test_only_the_late_write_sees_that_its_request_was_cancelled(monkeypatch: pytest.MonkeyPatch) -> None: + counter: Final = "spend:end_user:abandoned-stays-in-the-write" + redis: Final = _HeldSetRedis() + _process(monkeypatch, redis, cached_spend=6.0) + seen: Final[list[tuple[str, bool]]] = [] # mutable-ok: records who saw the request as gone + + original_increment: Final = redis.async_increment + + async def increment_noting_abandoned(key: str, value: float, **kwargs: object) -> float: + current: Final = await original_increment(key, value, **kwargs) + if kwargs.get("only_if_missing"): + seen.append(("write", spend_counter_reseed.late_write_abandoned())) + return current + + monkeypatch.setattr(redis, "async_increment", increment_noting_abandoned) + read_before_seed: Final = asyncio.Event() + go: Final = asyncio.Event() + + async def request() -> None: + try: + await _a_request_read_cold_before_this_process_seeded(redis, counter, read_before_seed, go) + except asyncio.CancelledError: + seen.append(("request", spend_counter_reseed.late_write_abandoned())) + raise + + stale: Final = asyncio.ensure_future(request()) + await read_before_seed.wait() + await _cold_start(redis, counter) + go.set() + await redis.set_sent.wait() + stale.cancel() + await asyncio.wait((stale,)) + redis.release_set.set() + for _ in range(20): + await asyncio.sleep(0) + + assert seen == [("request", False), ("write", True)] + + +@pytest.mark.asyncio +async def test_a_request_joining_an_add_that_is_cancelled_writes_the_counter_itself( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:joined-add-cancelled" + redis: Final = _HeldIncrementRedis(held_increments=2) + redis.releases[0].set() + _process(monkeypatch, redis, cached_spend=6.0) + monkeypatch.setattr(spend_counter_reseed, "SPEND_COUNTER_COLD_SEED_JOIN_SECONDS", 60.0) + read_before_seed: Final = asyncio.Event() + go: Final = asyncio.Event() + older: Final = asyncio.ensure_future( + _a_request_read_cold_before_this_process_seeded(redis, counter, read_before_seed, go) + ) + await read_before_seed.wait() + await _cold_start(redis, counter) + del redis.values[counter] + fresh: Final = asyncio.ensure_future(_cold_start(redis, counter)) + await redis.arrived[1].wait() + go.set() + for _ in range(20): + await asyncio.sleep(0) + + fresh.cancel() + await asyncio.wait((fresh,)) + await older + + assert fresh.cancelled() + assert redis.values[counter] == 6.0 + + +@pytest.mark.asyncio +async def test_a_set_finding_a_short_counter_adds_even_if_another_requests_add_landed_during_that_set( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:add-landed-during-set" + redis: Final = _HeldIncrementRedis(held_increments=2) + _process(monkeypatch, redis, cached_spend=6.0) + monkeypatch.setattr(spend_counter_reseed, "SPEND_COUNTER_COLD_SEED_JOIN_SECONDS", 0.01) + first: Final = asyncio.ensure_future(_cold_start(redis, counter)) + await redis.arrived[0].wait() + second: Final = asyncio.ensure_future(_cold_start(redis, counter)) + await redis.arrived[1].wait() + redis.releases[0].set() + await first + redis.values[counter] = 0.3 + + redis.releases[1].set() + await second + + assert redis.values[counter] == pytest.approx(6.3) + + +class _RefusesTheSetRedis(_SharedRedis): + """A server that refuses SETNX, where ``async_increment`` falls back to a plain add.""" + + async def async_increment(self, key: str, value: float, **kwargs: object) -> float: + return await super().async_increment(key, value, **{**kwargs, "only_if_missing": False}) + + +@pytest.mark.asyncio +async def test_a_late_request_whose_set_became_a_plain_add_does_not_add_again_to_a_counter_below_zero( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:refused-set-negative-counter" + redis: Final = _RefusesTheSetRedis() + _process(monkeypatch, redis, cached_spend=6.0) + read_before_seed: Final = asyncio.Event() + go: Final = asyncio.Event() + stale: Final = asyncio.ensure_future( + _a_request_read_cold_before_this_process_seeded(redis, counter, read_before_seed, go) + ) + await read_before_seed.wait() + await _cold_start(redis, counter) + redis.values[counter] = -1.5 + go.set() + await asyncio.wait_for(stale, timeout=5) + + assert redis.values[counter] == 4.5 + + +@pytest.mark.asyncio +async def test_a_stale_read_on_a_real_client_seeds_a_missing_counter_with_its_ttl_and_keeps_an_existing_one( + monkeypatch: pytest.MonkeyPatch, +) -> None: + counter: Final = "spend:end_user:real-client" + client: Final = FakeAsyncRedis(server=FakeServer()) + redis: Final = _RedisCacheOnClient(client) + redis.default_ttl = 120 + _process(monkeypatch, redis, cached_spend=6.0) + read: Final = asyncio.Event() + seeded: Final = asyncio.Event() + + async def _stale_request() -> None: + with spend_counter_batch_scope(redis, frozenset((counter,))): + await _read_cold(counter) + read.set() + await seeded.wait() + await client.delete(counter) + await ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=_COLD_SEED_SOURCE) + + async def _first_request() -> None: + await read.wait() + await _cold_start(redis, counter) + seeded.set() + + await asyncio.gather(_stale_request(), _first_request()) + + assert float(await client.get(counter) or 0) == 6.0 + assert 0 < await client.ttl(counter) <= 120