diff --git a/docs/resp3_features.rst b/docs/resp3_features.rst index 3e402cf17f..62b1e96bd8 100644 --- a/docs/resp3_features.rst +++ b/docs/resp3_features.rst @@ -86,10 +86,11 @@ Client-side caching Client-side caching is a technique used to create high performance services. It utilizes the memory on application servers, typically separate from the database nodes, to cache a subset of the data directly on the application side. For more information please check the `Redis client-side caching documentation `_. -Please notice that this feature is available only with RESP3 protocol enabled -in sync clients. redis-py 8.0 and later use RESP3 on the wire by default, and -the examples below pass ``protocol=3`` explicitly to make the requirement clear. -Supported in standalone, Cluster, and Sentinel clients. +Please notice that this feature is available only with RESP3 protocol enabled. +redis-py 8.0 and later use RESP3 on the wire by default, and the examples below +pass ``protocol=3`` explicitly to make the requirement clear. It is supported +by standalone sync and async clients, as well as the sync Cluster and Sentinel +clients. Basic usage: @@ -101,7 +102,15 @@ Enable caching with default configuration: >>> from redis.cache import CacheConfig >>> r = redis.Redis(host='localhost', port=6379, protocol=3, cache_config=CacheConfig()) -The same interface applies to Redis Cluster and Sentinel. +The async client uses the same configuration: + +.. code:: python + + >>> import redis.asyncio as redis + >>> from redis.cache import CacheConfig + >>> r = redis.Redis(host='localhost', port=6379, protocol=3, cache_config=CacheConfig()) + +The same interface applies to the sync Redis Cluster and Sentinel clients. Enable caching with custom cache implementation: diff --git a/redis/asyncio/client.py b/redis/asyncio/client.py index eb9c0bb4d6..18952439c5 100644 --- a/redis/asyncio/client.py +++ b/redis/asyncio/client.py @@ -52,6 +52,7 @@ ) from redis.asyncio.retry import Retry from redis.backoff import ExponentialWithJitterBackoff +from redis.cache import CacheConfig, CacheInterface from redis.client import ( EMPTY_RESPONSE, NEVER_DECODE, @@ -316,6 +317,8 @@ def __init__( credential_provider: CredentialProvider | None = None, protocol: int | None = None, legacy_responses: bool = True, + cache: CacheInterface | None = None, + cache_config: CacheConfig | None = None, event_dispatcher: EventDispatcher | None = None, maint_notifications_config: MaintNotificationsConfig | None = None, ): @@ -461,6 +464,13 @@ def __init__( "ssl_password": ssl_password, } ) + if (cache_config or cache) and check_protocol_version(protocol, 3): + kwargs.update( + { + "cache": cache, + "cache_config": cache_config, + } + ) maint_notifications_enabled = ( maint_notifications_config and maint_notifications_config.enabled ) @@ -492,6 +502,12 @@ def __init__( ) self.connection_pool = connection_pool + + if (cache_config or cache) and not check_protocol_version( + self.connection_pool.get_protocol(), 3 + ): + raise RedisError("Client caching is only supported with RESP version 3") + self.single_connection_client = single_connection_client self.connection: Optional[Connection] = None @@ -544,6 +560,10 @@ def get_encoder(self): """Get the connection pool's encoder""" return self.connection_pool.get_encoder() + def get_cache(self) -> Optional[CacheInterface]: + """Return the client-side cache configured for this connection pool.""" + return self.connection_pool.cache + def get_connection_kwargs(self): """Get the connection's key-word arguments""" return self.connection_pool.connection_kwargs @@ -867,7 +887,7 @@ async def _send_command_parse_response(self, conn, command_name, *args, **option # returns its arity error instead of a client-side IndexError here. key, fieldset_name, values = himport_set return await self._himport_execute_set(conn, key, fieldset_name, values) - await conn.send_command(*args) + await conn.send_command(*args, **options) return await self.parse_response(conn, command_name, **options) async def _himport_reconcile_discards(self, conn): diff --git a/redis/asyncio/connection.py b/redis/asyncio/connection.py index c1b849b3bf..b0df38ed5c 100644 --- a/redis/asyncio/connection.py +++ b/redis/asyncio/connection.py @@ -35,6 +35,8 @@ DB_CLIENT_CONNECTION_STATE, AttributeBuilder, ConnectionState, + CSCReason, + CSCResult, get_pool_name, ) from ..utils import SSL_AVAILABLE, deprecated_function @@ -51,7 +53,7 @@ from ..auth.token import TokenInterface from ..driver_info import DriverInfo, resolve_driver_info from ..event import AsyncAfterConnectionReleasedEvent, EventDispatcher -from ..utils import deprecated_args, format_error_message +from ..utils import deprecated_args, ensure_string, format_error_message # the functionality is available in 3.11.x but has a major issue before # 3.11.3. See https://github.com/redis/redis-py/issues/2633 @@ -74,6 +76,15 @@ ) from redis.asyncio.retry import Retry from redis.backoff import NoBackoff +from redis.cache import ( + CacheEntry, + CacheEntryStatus, + CacheFactory, + CacheFactoryInterface, + CacheInterface, + CacheKey, + CacheProxy, +) from redis.credentials import CredentialProvider, UsernamePasswordCredentialProvider from redis.exceptions import ( AuthenticationError, @@ -94,12 +105,20 @@ _build_moving_connection_kwargs, ) from redis.observability.metrics import CloseReason +from redis.observability.recorder import ( + init_csc_items, + record_csc_eviction, + record_csc_network_saved, + record_csc_request, + register_csc_items_callback, +) from redis.typing import EncodableT from redis.utils import ( DEFAULT_RESP_VERSION, HIREDIS_AVAILABLE, SENTINEL, check_protocol_version, + compare_versions, str_if_bytes, ) @@ -684,6 +703,7 @@ def __init__( self.next_health_check: float = -1 self.encoder = encoder_class(encoding, encoding_errors, decode_responses) self.redis_connect_func = redis_connect_func + self.handshake_metadata = None self._reader: Optional[asyncio.StreamReader] = None self._writer: Optional[asyncio.StreamWriter] = None self._socket_read_size = socket_read_size @@ -970,6 +990,7 @@ async def on_connect_check_health(self, check_health: bool = True) -> None: "HELLO", self.protocol, "AUTH", *auth_args, check_health=False ) response = await self.read_response() + self.handshake_metadata = response if response.get(b"proto") != int(self.protocol) and response.get( "proto" ) != int(self.protocol): @@ -1001,6 +1022,7 @@ async def on_connect_check_health(self, check_health: bool = True) -> None: self._parser.on_connect(self) await self.send_command("HELLO", self.protocol, check_health=check_health) response = await self.read_response() + self.handshake_metadata = response # if response.get(b"proto") != self.protocol and response.get( # "proto" # ) != self.protocol: @@ -1824,6 +1846,321 @@ def parse_url(url: str) -> ConnectKwargs: _CP = TypeVar("_CP", bound="ConnectionPool") +class AsyncCacheProxyConnection: + """Add Redis client-side caching to an asynchronous connection.""" + + DUMMY_CACHE_VALUE = b"foo" + MIN_ALLOWED_VERSION = "7.4.0" + DEFAULT_SERVER_NAME = "redis" + + def __init__( + self, + connection: AbstractConnection, + cache: CacheInterface, + pool_lock: asyncio.Lock, + pool: Optional["ConnectionPool"] = None, + ) -> None: + self._conn = connection + self._cache = cache + self._pool_lock = pool_lock + self._pool = pool + self._cache_lock = asyncio.Lock() + self._current_command_cache_key: CacheKey | None = None + self.register_connect_callback(self._enable_tracking_callback) + + def __getattr__(self, name: str): + return getattr(self._conn, name) + + def __setattr__(self, name: str, value) -> None: + if name.startswith("_") or "_conn" not in self.__dict__: + object.__setattr__(self, name, value) + else: + setattr(self._conn, name, value) + + @property + def retry(self): + return self._conn.retry + + @retry.setter + def retry(self, value) -> None: + self._conn.retry = value + + @property + def is_connected(self) -> bool: + return self._conn.is_connected + + def __repr__(self) -> str: + return repr(self._conn) + + async def connect(self) -> None: + await self._conn.connect() + + server_name = self._conn.handshake_metadata.get(b"server") + if server_name is None: + server_name = self._conn.handshake_metadata.get("server") + server_version = self._conn.handshake_metadata.get(b"version") + if server_version is None: + server_version = self._conn.handshake_metadata.get("version") + if server_version is None or server_name is None: + raise ConnectionError("Cannot retrieve information about server version") + + if ( + ensure_string(server_name) != self.DEFAULT_SERVER_NAME + or compare_versions(ensure_string(server_version), self.MIN_ALLOWED_VERSION) + == 1 + ): + raise ConnectionError( + "To maximize compatibility with all Redis products, client-side " + "caching is supported by Redis 7.4 or later" + ) + + async def disconnect(self, *args, **kwargs) -> None: + async with self._cache_lock: + self._cache.flush() + await self._conn.disconnect(*args, **kwargs) + + async def send_packed_command(self, command, check_health: bool = True) -> None: + async with self._cache_lock: + self._current_command_cache_key = None + await self._conn.send_packed_command(command, check_health=check_health) + + async def send_command(self, *args, **kwargs) -> None: + await self._process_pending_invalidations() + + cache_key = CacheKey(command=args[0], redis_keys=(), redis_args=()) + async with self._cache_lock: + if not self._cache.is_cachable(cache_key): + self._current_command_cache_key = None + is_cacheable = False + else: + is_cacheable = True + + if is_cacheable and kwargs.get("keys") is None: + raise ValueError("Cannot create cache key.") + + if is_cacheable: + self._current_command_cache_key = CacheKey( + command=args[0], + redis_keys=self._normalize_cache_keys(kwargs["keys"]), + redis_args=args, + ) + cached_entry = self._cache.get(self._current_command_cache_key) + + if not is_cacheable: + await self._conn.send_command(*args, **kwargs) + return + + while cached_entry is not None: + if cached_entry.status is CacheEntryStatus.IN_PROGRESS: + if cached_entry.completion_event is None: + async with self._cache_lock: + if ( + self._cache.get(self._current_command_cache_key) + is cached_entry + ): + self._cache.delete_by_cache_keys( + [self._current_command_cache_key] + ) + break + + await cached_entry.completion_event.wait() + async with self._cache_lock: + cached_entry = self._cache.get(self._current_command_cache_key) + if ( + cached_entry is not None + and cached_entry.status is CacheEntryStatus.VALID + ): + return + continue + + if await self._refresh_cached_entry(cached_entry): + async with self._cache_lock: + if self._cache.get(self._current_command_cache_key) is not None: + return + else: + async with self._cache_lock: + if self._cache.get(self._current_command_cache_key) is cached_entry: + self._cache.delete_by_cache_keys( + [self._current_command_cache_key] + ) + break + + cache_key = self._current_command_cache_key + cache_entry = CacheEntry( + cache_key=cache_key, + cache_value=self.DUMMY_CACHE_VALUE, + status=CacheEntryStatus.IN_PROGRESS, + connection_ref=self._conn, + completion_event=asyncio.Event(), + ) + async with self._cache_lock: + self._cache.set(cache_entry) + + try: + await self._conn.send_command(*args, **kwargs) + except BaseException: + async with self._cache_lock: + if cache_key is not None and self._cache.get(cache_key) is cache_entry: + self._cache.delete_by_cache_keys([cache_key]) + self._signal_cache_entry(cache_entry) + if self._current_command_cache_key == cache_key: + self._current_command_cache_key = None + raise + + async def can_read(self, timeout: float = 0) -> bool: + return await self._conn.can_read() + + async def read_response( + self, + disable_decoding: bool = False, + timeout: float | None = None, + *, + disconnect_on_error: bool = True, + push_request: bool | None = False, + ): + async with self._cache_lock: + cache_key = self._current_command_cache_key + if cache_key is not None: + cache_entry = self._cache.get(cache_key) + if ( + cache_entry is not None + and cache_entry.status != CacheEntryStatus.IN_PROGRESS + ): + response = copy.deepcopy(cache_entry.cache_value) + self._current_command_cache_key = None + record_csc_request(result=CSCResult.HIT) + record_csc_network_saved( + bytes_saved=len(response) if hasattr(response, "__len__") else 0 + ) + return response + record_csc_request(result=CSCResult.MISS) + + try: + response = await self._conn.read_response( + disable_decoding=disable_decoding, + timeout=timeout, + disconnect_on_error=disconnect_on_error, + push_request=push_request, + ) + except BaseException: + async with self._cache_lock: + cache_entry = self._cache.get(cache_key) + if ( + cache_entry is not None + and cache_entry.status is CacheEntryStatus.IN_PROGRESS + ): + self._cache.delete_by_cache_keys([cache_key]) + self._signal_cache_entry(cache_entry) + self._current_command_cache_key = None + raise + + async with self._cache_lock: + cache_key = self._current_command_cache_key + if cache_key is None: + return response + if response is None: + cache_entry = self._cache.get(cache_key) + self._cache.delete_by_cache_keys([cache_key]) + if cache_entry is not None: + self._signal_cache_entry(cache_entry) + self._current_command_cache_key = None + return response + + cache_entry = self._cache.get(cache_key) + if cache_entry is not None: + cache_entry.status = CacheEntryStatus.VALID + cache_entry.cache_value = response + self._cache.set(cache_entry) + self._signal_cache_entry(cache_entry) + + self._current_command_cache_key = None + + return response + + async def _enable_tracking_callback(self, connection: AbstractConnection) -> None: + await connection.send_command("CLIENT", "TRACKING", "ON") + await connection.read_response() + connection._parser.set_invalidation_push_handler(self._on_invalidation_callback) + + async def _process_pending_invalidations(self) -> None: + await self._drain_pending_invalidations(self._conn) + + async def _drain_pending_invalidations(self, connection) -> bool: + while await connection.can_read(): + try: + response = await connection.read_response( + push_request=True, + timeout=0, + disconnect_on_error=False, + ) + except TimeoutError: + return False + if response is None: + return False + return True + + async def _refresh_cached_entry(self, cached_entry: CacheEntry) -> bool: + async with self._pool_lock: + is_current_connection = cached_entry.connection_ref is self._conn + is_available_connection = ( + self._pool is None + or is_current_connection + or self._pool._is_connection_available(cached_entry.connection_ref) + ) + if not is_available_connection: + return False + return await self._drain_pending_invalidations(cached_entry.connection_ref) + + @staticmethod + def _normalize_cache_keys(redis_keys) -> tuple: + if isinstance(redis_keys, (str, bytes)): + return (redis_keys,) + return tuple(redis_keys) + + async def _on_invalidation_callback(self, data) -> None: + async with self._cache_lock: + if data[1] is None: + entries = tuple(self._cache.collection.values()) + self._cache.flush() + keys_deleted = 0 + else: + entries = tuple(self._cache.collection.values()) + keys_deleted = len(self._cache.delete_by_redis_keys(data[1])) + + for entry in entries: + if entry.completion_event is not None and ( + data[1] is None or self._matches_redis_keys(entry, data[1]) + ): + entry.completion_event.set() + + if keys_deleted: + record_csc_eviction( + count=keys_deleted, + reason=CSCReason.INVALIDATION, + ) + + @staticmethod + def _signal_cache_entry(entry: CacheEntry) -> None: + if entry.completion_event is not None: + entry.completion_event.set() + + @staticmethod + def _matches_redis_keys(entry: CacheEntry, redis_keys) -> bool: + for redis_key in redis_keys: + candidates = [redis_key] + if isinstance(redis_key, str): + candidates.append(redis_key.encode("utf-8")) + elif isinstance(redis_key, bytes): + try: + candidates.append(redis_key.decode("utf-8")) + except UnicodeDecodeError: + pass + + if any(candidate in entry.cache_key.redis_keys for candidate in candidates): + return True + return False + + class ConnectionPoolInterface(ABC): @abstractmethod def get_protocol(self): @@ -2618,6 +2955,7 @@ def __init__( self, connection_class: Type[AbstractConnection] = Connection, max_connections: Optional[int] = None, + cache_factory: Optional[CacheFactoryInterface] = None, maint_notifications_config: MaintNotificationsConfig | None = None, **connection_kwargs, ): @@ -2628,6 +2966,36 @@ def __init__( self.connection_class = connection_class self._connection_kwargs = connection_kwargs self.max_connections = max_connections + self.cache = None + self._cache_factory = cache_factory + + if connection_kwargs.get("cache_config") or connection_kwargs.get("cache"): + if not check_protocol_version(connection_kwargs.get("protocol"), 3): + raise RedisError("Client caching is only supported with RESP version 3") + + cache = connection_kwargs.get("cache") + if cache is not None: + if not isinstance(cache, CacheInterface): + raise ValueError("Cache must implement CacheInterface") + self.cache = cache + elif self._cache_factory is not None: + cache = self._cache_factory.get_cache() + self.cache = ( + cache if isinstance(cache, CacheProxy) else CacheProxy(cache) + ) + else: + self.cache = CacheFactory( + connection_kwargs.get("cache_config") + ).get_cache() + + init_csc_items() + register_csc_items_callback( + callback=lambda: self.cache.size, + pool_name=get_pool_name(self), + ) + + connection_kwargs.pop("cache", None) + connection_kwargs.pop("cache_config", None) # Resolve the HIMPORT registry. A pre-built ``himport_registry`` (shared, e.g. # from the cluster client) takes precedence; otherwise build a fresh empty one. @@ -2858,7 +3226,19 @@ def make_connection(self): """Create a new connection. Can be overridden by child classes.""" # Note: We don't record IDLE here because async uses a sync make_connection # but async record_connection_count. The recording is handled in get_connection. - return self.connection_class(**self.connection_kwargs) + connection = self.connection_class(**self.connection_kwargs) + if self.cache is not None: + return AsyncCacheProxyConnection( + connection, self.cache, self._lock, pool=self + ) + return connection + + def _is_connection_available(self, connection: AbstractConnection) -> bool: + """Return whether a cached connection is idle in this pool.""" + return any( + candidate is connection or getattr(candidate, "_conn", None) is connection + for candidate in self._available_connections + ) async def ensure_connection(self, connection: AbstractConnection): """Ensure that the connection object is connected and valid""" @@ -3064,7 +3444,8 @@ def set_in_maintenance(self, in_maintenance: bool) -> None: @contextlib.asynccontextmanager async def _maybe_pool_lock(self) -> AsyncIterator[None]: - if self._in_maintenance: + """Serialize pool mutations with cache-owner availability checks.""" + if self.cache is not None or self._in_maintenance: async with self._lock: yield else: diff --git a/redis/cache.py b/redis/cache.py index 1752a4a94b..faf1bacc47 100644 --- a/redis/cache.py +++ b/redis/cache.py @@ -43,11 +43,13 @@ def __init__( cache_value: bytes, status: CacheEntryStatus, connection_ref, + completion_event=None, ): self.cache_key = cache_key self.cache_value = cache_value self.status = status self.connection_ref = connection_ref + self.completion_event = completion_event def __hash__(self): return hash( @@ -278,6 +280,7 @@ def set(self, entry: CacheEntry) -> bool: is_set = self._cache.set(entry) if self.config.is_exceeds_max_size(self.size): + evicted_entry = next(iter(self.collection.values())) # Lazy import to avoid circular dependency from redis.observability.recorder import record_csc_eviction @@ -286,6 +289,8 @@ def set(self, entry: CacheEntry) -> bool: reason=CSCReason.FULL, ) self.eviction_policy.evict_next() + if evicted_entry.completion_event is not None: + evicted_entry.completion_event.set() return is_set diff --git a/redis/connection.py b/redis/connection.py index 7ae5d82e1c..93cabc404c 100644 --- a/redis/connection.py +++ b/redis/connection.py @@ -3021,7 +3021,10 @@ def __init__( self.cache = cache else: if self._cache_factory is not None: - self.cache = CacheProxy(self._cache_factory.get_cache()) + cache = self._cache_factory.get_cache() + self.cache = ( + cache if isinstance(cache, CacheProxy) else CacheProxy(cache) + ) else: self.cache = CacheFactory( self._connection_kwargs.get("cache_config") diff --git a/tests/test_asyncio/test_connection.py b/tests/test_asyncio/test_connection.py index c56b5d9dd4..016c5ab580 100644 --- a/tests/test_asyncio/test_connection.py +++ b/tests/test_asyncio/test_connection.py @@ -15,8 +15,9 @@ _AsyncRESPBase, ) from redis._parsers.hiredis import NOT_ENOUGH_DATA -from redis.asyncio import ConnectionPool, Redis +from redis.asyncio import BlockingConnectionPool, ConnectionPool, Redis from redis.asyncio.connection import ( + AsyncCacheProxyConnection, Connection, SSLConnection, UnixDomainSocketConnection, @@ -24,7 +25,16 @@ ) from redis.asyncio.retry import Retry from redis.backoff import NoBackoff -from redis.exceptions import ConnectionError, InvalidResponse, TimeoutError +from redis.cache import ( + CacheConfig, + CacheEntry, + CacheEntryStatus, + CacheFactory, + CacheKey, + CacheProxy, +) +from redis.observability.attributes import CSCReason +from redis.exceptions import ConnectionError, InvalidResponse, RedisError, TimeoutError from redis.utils import HIREDIS_AVAILABLE from tests.conftest import skip_if_server_version_lt @@ -83,6 +93,395 @@ def test_connection_default_parser_matches_default_protocol(): assert conn.protocol == 3 +async def test_async_redis_accepts_client_side_cache_configuration(): + client = Redis(cache_config=CacheConfig(max_size=10), protocol=3) + + try: + assert client.connection_pool.cache is not None + finally: + await client.aclose() + + +async def test_async_unix_socket_accepts_client_side_cache_configuration(): + client = Redis( + unix_socket_path="unix:///tmp/redis.sock", + cache_config=CacheConfig(max_size=10), + protocol=3, + ) + + try: + assert client.connection_pool.cache is not None + finally: + await client.aclose() + + +async def test_async_redis_rejects_client_side_cache_with_resp2(): + with pytest.raises( + RedisError, match="Client caching is only supported with RESP version 3" + ): + Redis(cache_config=CacheConfig(), protocol=2) + + +def test_async_connection_pool_rejects_client_side_cache_with_resp2(): + with pytest.raises( + RedisError, match="Client caching is only supported with RESP version 3" + ): + ConnectionPool(protocol=2, cache_config=CacheConfig()) + + +async def test_async_connection_pool_creates_cache_proxy_connection(): + pool = ConnectionPool(protocol=3, cache_config=CacheConfig()) + + try: + assert isinstance(pool.make_connection(), AsyncCacheProxyConnection) + finally: + await pool.aclose() + + +async def test_async_cache_proxy_rejects_unsupported_redis_version(): + connection = mock.Mock() + connection.connect = mock.AsyncMock() + connection.handshake_metadata = { + b"server": b"redis", + b"version": b"7.2.4", + } + cache = CacheFactory(CacheConfig()).get_cache() + proxy = AsyncCacheProxyConnection(connection, cache, asyncio.Lock()) + + with pytest.raises(ConnectionError, match="supported by Redis 7.4 or later"): + await proxy.connect() + + +async def test_async_cache_proxy_returns_cached_response_and_handles_invalidation(): + connection = mock.Mock() + connection.can_read = mock.AsyncMock(return_value=False) + connection.send_command = mock.AsyncMock() + connection.read_response = mock.AsyncMock(side_effect=[b"first", b"second"]) + cache = CacheFactory(CacheConfig()).get_cache() + proxy = AsyncCacheProxyConnection(connection, cache, asyncio.Lock()) + + await proxy.send_command("GET", "key", keys=("key",)) + assert await proxy.read_response() == b"first" + + await proxy.send_command("GET", "key", keys=("key",)) + assert await proxy.read_response() == b"first" + assert connection.send_command.await_count == 1 + assert connection.read_response.await_count == 1 + + await proxy._on_invalidation_callback([b"invalidate", [b"key"]]) + await proxy.send_command("GET", "key", keys=("key",)) + assert await proxy.read_response() == b"second" + assert connection.send_command.await_count == 2 + assert connection.read_response.await_count == 2 + + +async def test_async_cache_proxy_waits_for_in_flight_cache_fill(): + first_connection = mock.Mock() + first_connection.can_read = mock.AsyncMock(return_value=False) + first_connection.send_command = mock.AsyncMock() + first_connection.read_response = mock.AsyncMock(return_value=b"first") + second_connection = mock.Mock() + second_connection.can_read = mock.AsyncMock(return_value=False) + second_connection.send_command = mock.AsyncMock() + second_connection.read_response = mock.AsyncMock() + cache = CacheFactory(CacheConfig()).get_cache() + pool_lock = asyncio.Lock() + first_proxy = AsyncCacheProxyConnection(first_connection, cache, pool_lock) + second_proxy = AsyncCacheProxyConnection(second_connection, cache, pool_lock) + + await first_proxy.send_command("GET", "key", keys=("key",)) + waiting_send = asyncio.create_task( + second_proxy.send_command("GET", "key", keys=("key",)) + ) + await asyncio.sleep(0) + + assert waiting_send.done() is False + assert second_connection.send_command.await_count == 0 + + assert await first_proxy.read_response() == b"first" + await waiting_send + assert await second_proxy.read_response() == b"first" + assert second_connection.read_response.await_count == 0 + + +async def test_async_cache_proxy_retries_after_in_flight_fill_fails(): + first_connection = mock.Mock() + first_connection.can_read = mock.AsyncMock(return_value=False) + first_connection.send_command = mock.AsyncMock() + first_connection.read_response = mock.AsyncMock( + side_effect=ConnectionError("read failed") + ) + second_connection = mock.Mock() + second_connection.can_read = mock.AsyncMock(return_value=False) + second_connection.send_command = mock.AsyncMock() + cache = CacheFactory(CacheConfig()).get_cache() + pool_lock = asyncio.Lock() + first_proxy = AsyncCacheProxyConnection(first_connection, cache, pool_lock) + second_proxy = AsyncCacheProxyConnection(second_connection, cache, pool_lock) + + await first_proxy.send_command("GET", "key", keys=("key",)) + waiting_send = asyncio.create_task( + second_proxy.send_command("GET", "key", keys=("key",)) + ) + await asyncio.sleep(0) + + with pytest.raises(ConnectionError, match="read failed"): + await first_proxy.read_response() + + await waiting_send + second_connection.send_command.assert_awaited_once_with("GET", "key", keys=("key",)) + + +async def test_async_cache_proxy_signals_evicted_in_progress_fill(): + cache = CacheFactory(CacheConfig(max_size=1)).get_cache() + first_event = asyncio.Event() + first_entry = CacheEntry( + cache_key=CacheKey(command="GET", redis_keys=("first",), redis_args=()), + cache_value=b"foo", + status=CacheEntryStatus.IN_PROGRESS, + connection_ref=mock.Mock(), + completion_event=first_event, + ) + second_entry = CacheEntry( + cache_key=CacheKey(command="GET", redis_keys=("second",), redis_args=()), + cache_value=b"foo", + status=CacheEntryStatus.IN_PROGRESS, + connection_ref=mock.Mock(), + completion_event=asyncio.Event(), + ) + + cache.set(first_entry) + cache.set(second_entry) + + assert first_event.is_set() + + +async def test_blocking_pool_serializes_cache_owner_checks(): + pool = BlockingConnectionPool( + max_connections=1, protocol=3, cache_config=CacheConfig() + ) + entered = asyncio.Event() + + async def acquire_pool_lock(): + async with pool._maybe_pool_lock(): + entered.set() + + async with pool._lock: + task = asyncio.create_task(acquire_pool_lock()) + await asyncio.sleep(0) + assert entered.is_set() is False + + await task + await pool.aclose() + + +async def test_blocking_pool_skips_lock_without_cache(): + pool = BlockingConnectionPool(max_connections=1) + entered = asyncio.Event() + + async def acquire_pool_lock(): + async with pool._maybe_pool_lock(): + entered.set() + + async with pool._lock: + task = asyncio.create_task(acquire_pool_lock()) + await asyncio.sleep(0) + assert entered.is_set() is True + + await task + await pool.aclose() + + +async def test_async_connection_pool_does_not_double_wrap_custom_cache_factory(): + cache = CacheFactory(CacheConfig()).get_cache() + cache_factory = mock.Mock() + cache_factory.get_cache.return_value = cache + pool = ConnectionPool( + protocol=3, + cache_config=CacheConfig(), + cache_factory=cache_factory, + ) + + try: + assert pool.cache is cache + assert isinstance(pool.cache, CacheProxy) + assert not isinstance(pool.cache._cache, CacheProxy) + finally: + await pool.aclose() + + +async def test_async_cache_proxy_waits_for_replacement_in_progress_fill(): + connection = mock.Mock() + connection.can_read = mock.AsyncMock(return_value=False) + connection.send_command = mock.AsyncMock() + cache = CacheFactory(CacheConfig()).get_cache() + proxy = AsyncCacheProxyConnection(connection, cache, asyncio.Lock()) + replacement_event = asyncio.Event() + replacement_entry = CacheEntry( + cache_key=CacheKey( + command="GET", redis_keys=("key",), redis_args=("GET", "key") + ), + cache_value=b"replacement", + status=CacheEntryStatus.IN_PROGRESS, + connection_ref=connection, + completion_event=replacement_event, + ) + initial_event = mock.Mock() + + async def install_replacement(): + cache.set(replacement_entry) + + initial_event.wait = mock.AsyncMock(side_effect=install_replacement) + initial_entry = CacheEntry( + cache_key=replacement_entry.cache_key, + cache_value=b"foo", + status=CacheEntryStatus.IN_PROGRESS, + connection_ref=connection, + completion_event=initial_event, + ) + cache.set(initial_entry) + + waiting_send = asyncio.create_task(proxy.send_command("GET", "key", keys=("key",))) + await asyncio.sleep(0) + + assert waiting_send.done() is False + replacement_entry.status = CacheEntryStatus.VALID + replacement_event.set() + await waiting_send + assert await proxy.read_response() == b"replacement" + connection.send_command.assert_not_awaited() + + +async def test_async_cache_proxy_records_invalidation_evictions(): + connection = mock.Mock() + connection.can_read = mock.AsyncMock(return_value=False) + connection.send_command = mock.AsyncMock() + connection.read_response = mock.AsyncMock(return_value=b"first") + cache = CacheFactory(CacheConfig()).get_cache() + proxy = AsyncCacheProxyConnection(connection, cache, asyncio.Lock()) + + await proxy.send_command("GET", "key", keys=("key",)) + await proxy.read_response() + + with mock.patch("redis.asyncio.connection.record_csc_eviction") as record: + await proxy._on_invalidation_callback([b"invalidate", [b"key"]]) + + record.assert_called_once_with(count=1, reason=CSCReason.INVALIDATION) + + +async def test_async_cache_proxy_clears_cache_key_before_packed_command(): + connection = mock.Mock() + connection.send_packed_command = mock.AsyncMock() + proxy = AsyncCacheProxyConnection( + connection, CacheFactory(CacheConfig()).get_cache(), asyncio.Lock() + ) + proxy._current_command_cache_key = object() + + await proxy.send_packed_command([b"PING\r\n"]) + + assert proxy._current_command_cache_key is None + + +async def test_async_cache_proxy_does_not_drain_busy_cached_connection(): + first_connection = mock.Mock() + first_connection.can_read = mock.AsyncMock(return_value=False) + first_connection.send_command = mock.AsyncMock() + first_connection.read_response = mock.AsyncMock(return_value=b"first") + second_connection = mock.Mock() + second_connection.can_read = mock.AsyncMock(return_value=False) + second_connection.send_command = mock.AsyncMock() + cache = CacheFactory(CacheConfig()).get_cache() + pool = mock.Mock() + pool._is_connection_available.return_value = False + first_proxy = AsyncCacheProxyConnection( + first_connection, cache, asyncio.Lock(), pool=pool + ) + second_proxy = AsyncCacheProxyConnection( + second_connection, cache, asyncio.Lock(), pool=pool + ) + + await first_proxy.send_command("GET", "key", keys=("key",)) + await first_proxy.read_response() + first_connection.read_response.reset_mock() + + await second_proxy.send_command("GET", "key", keys=("key",)) + + first_connection.read_response.assert_not_awaited() + second_connection.send_command.assert_awaited_once_with("GET", "key", keys=("key",)) + + +async def test_async_cache_proxy_clears_in_progress_entry_when_send_is_cancelled(): + connection = mock.Mock() + connection.can_read = mock.AsyncMock(return_value=False) + connection.send_command = mock.AsyncMock( + side_effect=[asyncio.CancelledError(), None] + ) + cache = CacheFactory(CacheConfig()).get_cache() + proxy = AsyncCacheProxyConnection(connection, cache, asyncio.Lock()) + + with pytest.raises(asyncio.CancelledError): + await proxy.send_command("GET", "key", keys=("key",)) + + assert cache.size == 0 + assert proxy._current_command_cache_key is None + + await proxy.send_command("GET", "key", keys=("key",)) + assert connection.send_command.await_count == 2 + + +async def test_async_cache_proxy_normalizes_scalar_cache_keys(): + connection = mock.Mock() + connection.can_read = mock.AsyncMock(return_value=False) + connection.send_command = mock.AsyncMock() + cache = CacheFactory(CacheConfig()).get_cache() + proxy = AsyncCacheProxyConnection(connection, cache, asyncio.Lock()) + + await proxy.send_command("ZREVRANGE", "myzset", 0, -1, keys="myzset") + + assert next(iter(cache.collection)).redis_keys == ("myzset",) + + +async def test_async_cache_proxy_stops_when_invalidation_read_times_out(): + connection = mock.Mock() + connection.can_read = mock.AsyncMock(return_value=True) + connection.read_response = mock.AsyncMock(return_value=None) + proxy = AsyncCacheProxyConnection( + connection, CacheFactory(CacheConfig()).get_cache(), asyncio.Lock() + ) + + await asyncio.wait_for(proxy._process_pending_invalidations(), timeout=1) + + connection.read_response.assert_awaited_once() + + +async def test_async_client_forwards_command_options_to_connection(): + client = object.__new__(Redis) + client.parse_response = mock.AsyncMock(return_value=b"value") + connection = mock.Mock() + connection.send_command = mock.AsyncMock() + + assert ( + await client._send_command_parse_response( + connection, "GET", "key", keys=("key",) + ) + == b"value" + ) + connection.send_command.assert_awaited_once_with("key", keys=("key",)) + + +@pytest.mark.onlynoncluster +@skip_if_server_version_lt("7.4.0") +async def test_async_client_side_cache_round_trip(create_redis): + client = await create_redis(cache_config=CacheConfig(max_size=10)) + + await client.set("async-cache-key", "first") + assert await client.get("async-cache-key") == b"first" + assert client.get_cache().size == 1 + + await client.set("async-cache-key", "second") + assert await client.get("async-cache-key") == b"second" + + @pytest.mark.parametrize( ("buffer", "eof", "expected"), [ diff --git a/tests/test_connection.py b/tests/test_connection.py index f232615b6c..ca927b3c34 100644 --- a/tests/test_connection.py +++ b/tests/test_connection.py @@ -24,6 +24,7 @@ CacheConfig, CacheEntry, CacheEntryStatus, + CacheFactory, CacheInterface, CacheKey, CacheProxy, @@ -945,6 +946,22 @@ def test_creates_cache_with_custom_cache_factory( assert connection_pool.cache._cache == mock_cache connection_pool.disconnect() + def test_does_not_double_wrap_custom_cache_factory(self): + cache = CacheFactory(CacheConfig()).get_cache() + cache_factory = mock.Mock() + cache_factory.get_cache.return_value = cache + + connection_pool = ConnectionPool( + protocol=3, + cache_config=CacheConfig(), + cache_factory=cache_factory, + ) + + assert connection_pool.cache is cache + assert isinstance(connection_pool.cache, CacheProxy) + assert not isinstance(connection_pool.cache._cache, CacheProxy) + connection_pool.disconnect() + def test_creates_cache_with_given_configuration(self, mock_cache): connection_pool = ConnectionPool( protocol=3, cache_config=CacheConfig(max_size=100)