mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
fix(proxy): give user-key objects their own in-memory cache partition (#40713)
Key objects share the 200-entry UserApiKeyCache in-memory store with teams, end users, tags and memberships, so churn in those objects evicts hot keys and forces a LiteLLM_VerificationToken lookup on the next request. Route bare hashed-token keys to a dedicated InMemoryCache inside UserApiKeyCache while keeping Redis, TTL, serialization and invalidation shared Co-authored-by: yassin <yassin@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
729ea6b832
commit
3df127b439
8 changed files with 296 additions and 32 deletions
|
|
@ -376,7 +376,9 @@ class DualCache(BaseCache):
|
|||
)
|
||||
|
||||
# async_batch_set_cache
|
||||
async def async_set_cache_pipeline(self, cache_list: list, local_only: bool = False, **kwargs):
|
||||
async def async_set_cache_pipeline(
|
||||
self, cache_list: Sequence[tuple[str, object]], local_only: bool = False, **kwargs
|
||||
):
|
||||
"""
|
||||
Batch write values to the cache
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -196,9 +196,7 @@ class AuthCacheInvalidationSubscriber:
|
|||
for additional_cache in self._additional_in_memory_caches:
|
||||
additional_cache.set_cache(parsed.cache_key, parsed.new_value, ttl=parsed.ttl)
|
||||
return
|
||||
in_memory_cache: Final = self._user_api_key_cache.in_memory_cache
|
||||
if in_memory_cache is not None:
|
||||
in_memory_cache.delete_cache(parsed.cache_key)
|
||||
self._user_api_key_cache.in_memory_cache_for(parsed.cache_key).delete_cache(parsed.cache_key)
|
||||
for additional_cache in self._additional_in_memory_caches:
|
||||
additional_cache.delete_cache(parsed.cache_key)
|
||||
|
||||
|
|
|
|||
|
|
@ -147,8 +147,11 @@ async def memory_usage_in_mem_cache(
|
|||
llm_router.cache.in_memory_cache.ttl_dict
|
||||
)
|
||||
|
||||
num_items_in_user_api_key_cache: Final = len(user_api_key_cache.in_memory_cache.cache_dict) + len(
|
||||
user_api_key_cache.in_memory_cache.ttl_dict
|
||||
num_items_in_user_api_key_cache: Final = (
|
||||
len(user_api_key_cache.in_memory_cache.cache_dict)
|
||||
+ len(user_api_key_cache.in_memory_cache.ttl_dict)
|
||||
+ len(user_api_key_cache.key_object_cache.in_memory_cache.cache_dict)
|
||||
+ len(user_api_key_cache.key_object_cache.in_memory_cache.ttl_dict)
|
||||
)
|
||||
|
||||
num_items_in_proxy_logging_obj_cache: Final = len(
|
||||
|
|
@ -189,6 +192,8 @@ async def memory_usage_in_mem_cache_items(
|
|||
return {
|
||||
"user_api_key_cache": user_api_key_cache.in_memory_cache.cache_dict,
|
||||
"user_api_key_ttl": user_api_key_cache.in_memory_cache.ttl_dict,
|
||||
"user_key_object_cache": user_api_key_cache.key_object_cache.in_memory_cache.cache_dict,
|
||||
"user_key_object_ttl": user_api_key_cache.key_object_cache.in_memory_cache.ttl_dict,
|
||||
"llm_router_cache": llm_router_in_memory_cache_dict,
|
||||
"llm_router_ttl": llm_router_in_memory_ttl_dict,
|
||||
"proxy_logging_obj_cache": proxy_logging_obj.internal_usage_cache.dual_cache.in_memory_cache.cache_dict,
|
||||
|
|
@ -294,7 +299,9 @@ async def get_memory_summary(
|
|||
|
||||
try:
|
||||
# User API key cache
|
||||
user_cache_items: Final = len(user_api_key_cache.in_memory_cache.cache_dict)
|
||||
user_cache_items: Final = len(user_api_key_cache.in_memory_cache.cache_dict) + len(
|
||||
user_api_key_cache.key_object_cache.in_memory_cache.cache_dict
|
||||
)
|
||||
total_cache_items += user_cache_items
|
||||
caches["user_api_keys"] = {
|
||||
"count": user_cache_items,
|
||||
|
|
@ -429,10 +436,16 @@ def _get_cache_memory_stats(
|
|||
cache_stats: Final[dict[str, object]] = {}
|
||||
try:
|
||||
# User API key cache
|
||||
user_cache_size: Final = sys.getsizeof(user_api_key_cache.in_memory_cache.cache_dict)
|
||||
user_ttl_size: Final = sys.getsizeof(user_api_key_cache.in_memory_cache.ttl_dict)
|
||||
key_object_in_memory_cache: Final = user_api_key_cache.key_object_cache.in_memory_cache
|
||||
user_cache_size: Final = sys.getsizeof(user_api_key_cache.in_memory_cache.cache_dict) + sys.getsizeof(
|
||||
key_object_in_memory_cache.cache_dict
|
||||
)
|
||||
user_ttl_size: Final = sys.getsizeof(user_api_key_cache.in_memory_cache.ttl_dict) + sys.getsizeof(
|
||||
key_object_in_memory_cache.ttl_dict
|
||||
)
|
||||
cache_stats["user_api_key_cache"] = {
|
||||
"num_items": len(user_api_key_cache.in_memory_cache.cache_dict),
|
||||
"num_items": len(user_api_key_cache.in_memory_cache.cache_dict)
|
||||
+ len(key_object_in_memory_cache.cache_dict),
|
||||
"cache_dict_size_bytes": user_cache_size,
|
||||
"ttl_dict_size_bytes": user_ttl_size,
|
||||
"total_size_mb": round((user_cache_size + user_ttl_size) / (1024 * 1024), 2),
|
||||
|
|
|
|||
|
|
@ -1,11 +1,15 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from collections.abc import Sequence
|
||||
from typing import TYPE_CHECKING, Any, Final, TypeVar, cast, overload
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.caching.redis_cache import RedisCache
|
||||
from litellm.constants import DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL
|
||||
from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec
|
||||
|
||||
|
|
@ -14,6 +18,13 @@ if TYPE_CHECKING:
|
|||
|
||||
T = TypeVar("T", bound=BaseModel)
|
||||
|
||||
_HASHED_TOKEN_CACHE_KEY: Final = re.compile(r"[0-9a-f]{64}")
|
||||
|
||||
|
||||
def is_user_key_cache_key(key: str) -> bool:
|
||||
"""Only user-key objects are cached under a bare ``hash_token`` digest; every other object uses a prefixed key."""
|
||||
return _HASHED_TOKEN_CACHE_KEY.fullmatch(key) is not None
|
||||
|
||||
|
||||
class UserApiKeyCache(DualCache):
|
||||
"""
|
||||
|
|
@ -36,10 +47,50 @@ class UserApiKeyCache(DualCache):
|
|||
``async_set_cache_pipeline`` applies the same untyped Codec pass as omitting
|
||||
``model_type`` on ``async_set_cache`` (so ``BaseModel`` rows are dumped before Redis).
|
||||
|
||||
User-key objects (see ``is_user_key_cache_key``) live in their own in-memory partition,
|
||||
``key_object_cache``, so churn in the other management objects cannot evict them. Both
|
||||
partitions share the same Redis backend and TTL settings.
|
||||
|
||||
``get_cache`` / ``async_get_cache`` overloads and implementations must be contiguous
|
||||
(no other methods in between) so mypy resolves ``@overload`` + implementation correctly.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_memory_cache: InMemoryCache | None = None,
|
||||
redis_cache: RedisCache | None = None,
|
||||
default_in_memory_ttl: float | None = None,
|
||||
default_redis_ttl: float | None = None,
|
||||
key_object_in_memory_cache: InMemoryCache | None = None,
|
||||
) -> None:
|
||||
super().__init__(
|
||||
in_memory_cache=in_memory_cache,
|
||||
redis_cache=redis_cache,
|
||||
default_in_memory_ttl=default_in_memory_ttl,
|
||||
default_redis_ttl=default_redis_ttl,
|
||||
)
|
||||
self.key_object_cache: Final = DualCache(
|
||||
in_memory_cache=key_object_in_memory_cache or InMemoryCache(),
|
||||
redis_cache=redis_cache,
|
||||
default_in_memory_ttl=default_in_memory_ttl,
|
||||
default_redis_ttl=default_redis_ttl,
|
||||
)
|
||||
|
||||
def in_memory_cache_for(self, key: str) -> InMemoryCache:
|
||||
return self.key_object_cache.in_memory_cache if is_user_key_cache_key(key) else self.in_memory_cache
|
||||
|
||||
def update_cache_ttl(self, default_in_memory_ttl: float | None, default_redis_ttl: float | None) -> None:
|
||||
super().update_cache_ttl(default_in_memory_ttl=default_in_memory_ttl, default_redis_ttl=default_redis_ttl)
|
||||
self.key_object_cache.update_cache_ttl(
|
||||
default_in_memory_ttl=default_in_memory_ttl, default_redis_ttl=default_redis_ttl
|
||||
)
|
||||
|
||||
def attach_redis_cache(
|
||||
self, redis_cache: RedisCache | None = None, *, default_redis_ttl: float | None = None
|
||||
) -> None:
|
||||
super().attach_redis_cache(redis_cache, default_redis_ttl=default_redis_ttl)
|
||||
self.key_object_cache.attach_redis_cache(redis_cache, default_redis_ttl=default_redis_ttl)
|
||||
|
||||
@overload
|
||||
def get_cache(
|
||||
self,
|
||||
|
|
@ -71,7 +122,11 @@ class UserApiKeyCache(DualCache):
|
|||
) -> object:
|
||||
if model_type is None and "model_type" in kwargs:
|
||||
model_type = cast(type[BaseModel] | None, kwargs.pop("model_type", None))
|
||||
cached: Final = super().get_cache(key=key, parent_otel_span=parent_otel_span, local_only=local_only, **kwargs)
|
||||
cached: Final = (
|
||||
self.key_object_cache.get_cache(key=key, parent_otel_span=parent_otel_span, local_only=local_only, **kwargs)
|
||||
if is_user_key_cache_key(key)
|
||||
else super().get_cache(key=key, parent_otel_span=parent_otel_span, local_only=local_only, **kwargs)
|
||||
)
|
||||
if model_type is None:
|
||||
return cached
|
||||
if cached is None:
|
||||
|
|
@ -117,8 +172,14 @@ class UserApiKeyCache(DualCache):
|
|||
) -> object:
|
||||
if model_type is None and "model_type" in kwargs:
|
||||
model_type = cast(type[BaseModel] | None, kwargs.pop("model_type", None))
|
||||
cached: Final = await super().async_get_cache(
|
||||
key=key, parent_otel_span=parent_otel_span, local_only=local_only, **kwargs
|
||||
cached: Final = (
|
||||
await self.key_object_cache.async_get_cache(
|
||||
key=key, parent_otel_span=parent_otel_span, local_only=local_only, **kwargs
|
||||
)
|
||||
if is_user_key_cache_key(key)
|
||||
else await super().async_get_cache(
|
||||
key=key, parent_otel_span=parent_otel_span, local_only=local_only, **kwargs
|
||||
)
|
||||
)
|
||||
if model_type is None:
|
||||
return cached
|
||||
|
|
@ -137,20 +198,49 @@ class UserApiKeyCache(DualCache):
|
|||
def set_cache(self, key: str | None, value: object, local_only: bool = False, **kwargs: object):
|
||||
model_type: Final = cast(type[BaseModel] | None, kwargs.pop("model_type", None))
|
||||
payload: Final[object] = CacheCodec.serialize(value, model_type=model_type)
|
||||
if key is not None and is_user_key_cache_key(key):
|
||||
return self.key_object_cache.set_cache(key=key, value=payload, local_only=local_only, **kwargs)
|
||||
return super().set_cache(key=key, value=payload, local_only=local_only, **kwargs)
|
||||
|
||||
async def async_set_cache(self, key: str | None, value: object, local_only: bool = False, **kwargs: object):
|
||||
model_type: Final = cast(type[BaseModel] | None, kwargs.pop("model_type", None))
|
||||
payload: Final[object] = CacheCodec.serialize(value, model_type=model_type)
|
||||
if key is not None and is_user_key_cache_key(key):
|
||||
return await self.key_object_cache.async_set_cache(key=key, value=payload, local_only=local_only, **kwargs)
|
||||
return await super().async_set_cache(key=key, value=payload, local_only=local_only, **kwargs)
|
||||
|
||||
async def async_set_cache_pipeline(self, cache_list: list, local_only: bool = False, **kwargs: object) -> None:
|
||||
def delete_cache(self, key: str) -> None:
|
||||
if is_user_key_cache_key(key):
|
||||
self.key_object_cache.delete_cache(key)
|
||||
return
|
||||
super().delete_cache(key)
|
||||
|
||||
async def async_delete_cache(self, key: str) -> None:
|
||||
if is_user_key_cache_key(key):
|
||||
await self.key_object_cache.async_delete_cache(key)
|
||||
return
|
||||
await super().async_delete_cache(key)
|
||||
|
||||
def flush_cache(self) -> None:
|
||||
super().flush_cache()
|
||||
self.key_object_cache.in_memory_cache.flush_cache()
|
||||
|
||||
async def async_set_cache_pipeline(
|
||||
self, cache_list: Sequence[tuple[str, object]], local_only: bool = False, **kwargs: object
|
||||
) -> None:
|
||||
"""
|
||||
Batch writes with the same Codec boundary as ``async_set_cache`` without
|
||||
``model_type``: ``BaseModel`` values become JSON-safe dicts; dicts/scalars unchanged.
|
||||
"""
|
||||
normalized: Final = [(key, CacheCodec.serialize(value, model_type=None)) for key, value in cache_list]
|
||||
return await super().async_set_cache_pipeline(cache_list=normalized, local_only=local_only, **kwargs)
|
||||
normalized: Final = tuple((key, CacheCodec.serialize(value, model_type=None)) for key, value in cache_list)
|
||||
key_object_entries: Final = tuple(entry for entry in normalized if is_user_key_cache_key(entry[0]))
|
||||
other_entries: Final = tuple(entry for entry in normalized if not is_user_key_cache_key(entry[0]))
|
||||
if key_object_entries:
|
||||
await self.key_object_cache.async_set_cache_pipeline(
|
||||
cache_list=key_object_entries, local_only=local_only, **kwargs
|
||||
)
|
||||
if other_entries:
|
||||
await super().async_set_cache_pipeline(cache_list=other_entries, local_only=local_only, **kwargs)
|
||||
|
||||
|
||||
#: Value cached under ``user_object_permission_id_cache_key`` when the user links no permission row,
|
||||
|
|
|
|||
|
|
@ -3729,7 +3729,7 @@ async def delete_key_fn(
|
|||
)
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
"/keys/delete - cache after delete: %s", user_api_key_cache.in_memory_cache.cache_dict
|
||||
"/keys/delete - cache after delete: %s", user_api_key_cache.key_object_cache.in_memory_cache.cache_dict
|
||||
)
|
||||
|
||||
asyncio.create_task(
|
||||
|
|
|
|||
|
|
@ -7150,7 +7150,7 @@ async def test_extract_user_id_rehydrates_cross_replica_dict_cache(proxy_globals
|
|||
|
||||
key = "sk-alice-key"
|
||||
cache = UserApiKeyCache()
|
||||
cache.in_memory_cache.set_cache(hash_token(key), {"token": hash_token(key), "user_id": "alice"})
|
||||
cache.set_cache(hash_token(key), {"token": hash_token(key), "user_id": "alice"})
|
||||
proxy_globals.user_api_key_cache = cache
|
||||
proxy_globals.prisma_client = object()
|
||||
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
import asyncio
|
||||
import hashlib
|
||||
import json
|
||||
from typing import Iterable, List, Optional, Tuple
|
||||
from unittest.mock import patch
|
||||
|
|
@ -7,6 +8,7 @@ import pytest
|
|||
from redis.asyncio import Redis
|
||||
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import (
|
||||
AUTH_CACHE_INVALIDATION_CHANNEL,
|
||||
AuthCacheInvalidationSubscriber,
|
||||
|
|
@ -145,6 +147,35 @@ async def test_subscriber_deletes_local_cache_entry_on_message() -> None:
|
|||
assert pubsub.subscribed_channels == [AUTH_CACHE_INVALIDATION_CHANNEL]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_subscriber_deletes_key_object_partition_entry_on_message() -> None:
|
||||
"""
|
||||
LIT-7563 moved user-key objects into their own in-memory partition; a key
|
||||
invalidation broadcast must still evict the hashed-token entry there, or a
|
||||
deleted key keeps authenticating on other workers until its TTL expires.
|
||||
"""
|
||||
hashed_token = hashlib.sha256(b"sk-lit7563-hot-key").hexdigest()
|
||||
cache = UserApiKeyCache()
|
||||
cache.set_cache(hashed_token, UserAPIKeyAuth(token=hashed_token), model_type=UserAPIKeyAuth)
|
||||
assert cache.get_cache(hashed_token, model_type=UserAPIKeyAuth) is not None
|
||||
|
||||
pubsub = _QueuePubSub(initial_messages=[_invalidation_message(hashed_token)])
|
||||
subscriber = AuthCacheInvalidationSubscriber(
|
||||
redis_cache=_FakeRedisCache(client=_ScriptedPubSubRedisClient(pubsubs=[pubsub])),
|
||||
user_api_key_cache=cache,
|
||||
)
|
||||
subscriber.start()
|
||||
try:
|
||||
for _ in range(200):
|
||||
if cache.get_cache(hashed_token, model_type=UserAPIKeyAuth) is None:
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
finally:
|
||||
await subscriber.stop()
|
||||
|
||||
assert cache.get_cache(hashed_token, model_type=UserAPIKeyAuth) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_subscriber_deletes_additional_in_memory_cache_entry_on_message() -> None:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import hashlib
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
|
|
@ -10,10 +11,14 @@ from litellm.constants import DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL
|
|||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.common_utils.user_api_key_cache import (
|
||||
UserApiKeyCache,
|
||||
end_user_cache_key,
|
||||
get_management_object_ttl,
|
||||
is_user_key_cache_key,
|
||||
)
|
||||
from litellm.proxy.proxy_server import UserAPIKeyCacheTTLEnum
|
||||
|
||||
HASHED_TOKEN = hashlib.sha256(b"sk-lit7563-hot-key").hexdigest()
|
||||
|
||||
|
||||
class CapturingInMemoryCache(InMemoryCache):
|
||||
"""Records ``ttl`` passed into ``set_cache`` (what DualCache injects)."""
|
||||
|
|
@ -204,9 +209,7 @@ class TestUserApiKeyCache:
|
|||
|
||||
# Bypass UserApiKeyCache.serialize: CacheCodec rejects non-dict cached values
|
||||
# for dict-based models (deserialize returns None).
|
||||
await cache.in_memory_cache.async_set_cache(
|
||||
key="k", value="invalid-payload-not-a-dict"
|
||||
)
|
||||
await cache.in_memory_cache.async_set_cache(key="k", value="invalid-payload-not-a-dict")
|
||||
|
||||
value = await cache.async_get_cache("k", model_type=UserAPIKeyAuth)
|
||||
assert value is None
|
||||
|
|
@ -224,6 +227,141 @@ class TestUserApiKeyCache:
|
|||
fake.set_cache("k2", {"ok": NotSerializable()})
|
||||
|
||||
|
||||
class TestUserKeyObjectPartition:
|
||||
"""
|
||||
Regression for LIT-7563: user-key objects share one 200-entry ``InMemoryCache`` with
|
||||
every other management object, so end-user / team / tag churn evicts hot keys and
|
||||
forces a ``LiteLLM_VerificationToken`` lookup on the next request.
|
||||
"""
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("key", "expected"),
|
||||
[
|
||||
(HASHED_TOKEN, True),
|
||||
(HASHED_TOKEN.upper(), False),
|
||||
(f"team_id:{HASHED_TOKEN}", False),
|
||||
(end_user_cache_key("u1"), False),
|
||||
("sk-lit7563-hot-key", False),
|
||||
],
|
||||
)
|
||||
def test_is_user_key_cache_key(self, key: str, expected: bool):
|
||||
assert is_user_key_cache_key(key) is expected
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_management_object_churn_does_not_evict_key_object(self):
|
||||
cache = UserApiKeyCache(in_memory_cache=InMemoryCache(max_size_in_memory=2))
|
||||
await cache.async_set_cache(HASHED_TOKEN, _make_key_obj(HASHED_TOKEN), model_type=UserAPIKeyAuth, ttl=100)
|
||||
for i in range(2):
|
||||
await cache.async_set_cache(end_user_cache_key(f"u{i}"), {"user_id": f"u{i}"}, ttl=200)
|
||||
|
||||
key_obj = await cache.async_get_cache(HASHED_TOKEN, model_type=UserAPIKeyAuth)
|
||||
assert key_obj is not None
|
||||
assert key_obj.token == HASHED_TOKEN
|
||||
assert cache.get_cache(end_user_cache_key("u1")) == {"user_id": "u1"}
|
||||
assert HASHED_TOKEN not in cache.in_memory_cache.cache_dict
|
||||
|
||||
def test_sync_write_and_read_route_to_key_object_partition(self):
|
||||
cache = UserApiKeyCache(in_memory_cache=InMemoryCache(max_size_in_memory=2))
|
||||
cache.set_cache(HASHED_TOKEN, _make_key_obj(HASHED_TOKEN), model_type=UserAPIKeyAuth, ttl=100)
|
||||
for i in range(2):
|
||||
cache.set_cache(end_user_cache_key(f"u{i}"), {"user_id": f"u{i}"}, ttl=200)
|
||||
|
||||
key_obj = cache.get_cache(HASHED_TOKEN, model_type=UserAPIKeyAuth)
|
||||
assert key_obj is not None
|
||||
assert key_obj.token == HASHED_TOKEN
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_redis_hit_backfills_key_object_partition_with_configured_ttl(self):
|
||||
redis = FakeRedisCache()
|
||||
writer = UserApiKeyCache(redis_cache=redis, default_in_memory_ttl=30)
|
||||
await writer.async_set_cache(HASHED_TOKEN, _make_key_obj(HASHED_TOKEN), model_type=UserAPIKeyAuth)
|
||||
|
||||
key_partition = CapturingInMemoryCache()
|
||||
reader = UserApiKeyCache(redis_cache=redis, default_in_memory_ttl=30, key_object_in_memory_cache=key_partition)
|
||||
key_obj = await reader.async_get_cache(HASHED_TOKEN, model_type=UserAPIKeyAuth)
|
||||
|
||||
assert key_obj is not None
|
||||
assert key_obj.token == HASHED_TOKEN
|
||||
assert key_partition.last_ttl == 30
|
||||
assert HASHED_TOKEN not in reader.in_memory_cache.cache_dict
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_cache_ttl_applies_to_key_object_partition(self):
|
||||
key_partition = CapturingInMemoryCache()
|
||||
cache = UserApiKeyCache(default_in_memory_ttl=60, key_object_in_memory_cache=key_partition)
|
||||
cache.update_cache_ttl(default_in_memory_ttl=7, default_redis_ttl=7)
|
||||
|
||||
await cache.async_set_cache(HASHED_TOKEN, _make_key_obj(HASHED_TOKEN), model_type=UserAPIKeyAuth)
|
||||
|
||||
assert key_partition.last_ttl == 7
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_attach_redis_cache_applies_to_key_object_partition(self):
|
||||
redis = FakeRedisCache()
|
||||
cache = UserApiKeyCache()
|
||||
cache.attach_redis_cache(redis)
|
||||
|
||||
await cache.async_set_cache(HASHED_TOKEN, _make_key_obj(HASHED_TOKEN), model_type=UserAPIKeyAuth)
|
||||
|
||||
other_worker = UserApiKeyCache(redis_cache=redis)
|
||||
key_obj = await other_worker.async_get_cache(HASHED_TOKEN, model_type=UserAPIKeyAuth)
|
||||
assert key_obj is not None
|
||||
assert key_obj.token == HASHED_TOKEN
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_removes_key_object_from_partition_and_redis(self):
|
||||
redis = FakeRedisCache()
|
||||
cache = UserApiKeyCache(redis_cache=redis)
|
||||
await cache.async_set_cache(HASHED_TOKEN, _make_key_obj(HASHED_TOKEN), model_type=UserAPIKeyAuth)
|
||||
assert await cache.async_get_cache(HASHED_TOKEN, model_type=UserAPIKeyAuth) is not None
|
||||
|
||||
cache.delete_cache(HASHED_TOKEN)
|
||||
|
||||
assert await cache.async_get_cache(HASHED_TOKEN, model_type=UserAPIKeyAuth) is None
|
||||
assert await redis.async_get_cache(HASHED_TOKEN) is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_delete_removes_key_object_from_partition_and_redis(self):
|
||||
redis = FakeRedisCache()
|
||||
cache = UserApiKeyCache(redis_cache=redis)
|
||||
await cache.async_set_cache(HASHED_TOKEN, _make_key_obj(HASHED_TOKEN), model_type=UserAPIKeyAuth)
|
||||
|
||||
await cache.async_delete_cache(HASHED_TOKEN)
|
||||
|
||||
assert await cache.async_get_cache(HASHED_TOKEN, model_type=UserAPIKeyAuth) is None
|
||||
assert await redis.async_get_cache(HASHED_TOKEN) is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pipeline_write_routes_each_entry_to_its_partition(self):
|
||||
cache = UserApiKeyCache(in_memory_cache=InMemoryCache(max_size_in_memory=2))
|
||||
await cache.async_set_cache_pipeline(
|
||||
[(HASHED_TOKEN, _make_key_obj(HASHED_TOKEN))]
|
||||
+ [(end_user_cache_key(f"u{i}"), {"user_id": f"u{i}"}) for i in range(2)],
|
||||
ttl=100,
|
||||
)
|
||||
|
||||
key_obj = await cache.async_get_cache(HASHED_TOKEN, model_type=UserAPIKeyAuth)
|
||||
assert key_obj is not None
|
||||
assert key_obj.token == HASHED_TOKEN
|
||||
assert HASHED_TOKEN not in cache.in_memory_cache.cache_dict
|
||||
assert cache.get_cache(end_user_cache_key("u1")) == {"user_id": "u1"}
|
||||
|
||||
def test_flush_clears_key_object_partition(self):
|
||||
cache = UserApiKeyCache()
|
||||
cache.set_cache(HASHED_TOKEN, _make_key_obj(HASHED_TOKEN), model_type=UserAPIKeyAuth)
|
||||
cache.set_cache(end_user_cache_key("u1"), {"user_id": "u1"})
|
||||
|
||||
cache.flush_cache()
|
||||
|
||||
assert cache.get_cache(HASHED_TOKEN, model_type=UserAPIKeyAuth) is None
|
||||
assert cache.get_cache(end_user_cache_key("u1")) is None
|
||||
|
||||
def test_in_memory_cache_for_routes_by_key(self):
|
||||
cache = UserApiKeyCache()
|
||||
assert cache.in_memory_cache_for(HASHED_TOKEN) is cache.key_object_cache.in_memory_cache
|
||||
assert cache.in_memory_cache_for(end_user_cache_key("u1")) is cache.in_memory_cache
|
||||
|
||||
|
||||
class TestManagementObjectTTL:
|
||||
"""
|
||||
Regression for LIT-3338: ``general_settings.user_api_key_cache_ttl`` (which the
|
||||
|
|
@ -238,19 +376,13 @@ class TestManagementObjectTTL:
|
|||
def test_falls_back_to_constant_when_no_default_configured(self):
|
||||
cache = UserApiKeyCache()
|
||||
assert cache.default_in_memory_ttl is None
|
||||
assert (
|
||||
get_management_object_ttl(cache)
|
||||
== DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL
|
||||
)
|
||||
assert get_management_object_ttl(cache) == DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL
|
||||
|
||||
def test_resolves_on_a_plain_dual_cache(self):
|
||||
# Many call sites are typed UserApiKeyCache but exercised in tests with a
|
||||
# bare DualCache; the resolver must work on the base type, not just the subclass.
|
||||
assert get_management_object_ttl(DualCache(default_in_memory_ttl=300)) == 300
|
||||
assert (
|
||||
get_management_object_ttl(DualCache())
|
||||
== DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL
|
||||
)
|
||||
assert get_management_object_ttl(DualCache()) == DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_management_write_uses_configured_ttl_over_constant(self):
|
||||
|
|
@ -260,9 +392,7 @@ class TestManagementObjectTTL:
|
|||
redis_cache=FakeRedisCache(),
|
||||
default_in_memory_ttl=300,
|
||||
)
|
||||
assert get_management_object_ttl(cache) != (
|
||||
DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL
|
||||
)
|
||||
assert get_management_object_ttl(cache) != (DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL)
|
||||
|
||||
await cache.async_set_cache(
|
||||
"team_id:abc",
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue