Support client side caching with ConnectionPool (#3099)

* sync * async * fixs connection mocks * fix async connection mock * fix test_asyncio/test_connection.py::test_single_connection * add test for cache blacklist and flushdb at the end of each test * fix review comments

Author: dvora-h

redis/cache.py renamed to redis/_cache.py

@@ -178,7 +178,11 @@ class _LocalCache:
  """
 
  def __init__(
- self, max_size: int, ttl: int, eviction_policy: EvictionPolicy, **kwargs
+ self,
+ max_size: int = 100,
+ ttl: int = 0,
+ eviction_policy: EvictionPolicy = DEFAULT_EVICTION_POLICY,
+ **kwargs,
  ):
  self.max_size = max_size
  self.ttl = ttl

redis/asyncio/client.py

@@ -25,6 +25,12 @@
  cast,
 )
 
+from redis._cache import (
+ DEFAULT_BLACKLIST,
+ DEFAULT_EVICTION_POLICY,
+ DEFAULT_WHITELIST,
+ _LocalCache,
+)
 from redis._parsers.helpers import (
  _RedisCallbacks,
  _RedisCallbacksRESP2,
@@ -39,12 +45,6 @@
 )
 from redis.asyncio.lock import Lock
 from redis.asyncio.retry import Retry
-from redis.cache import (
- DEFAULT_BLACKLIST,
- DEFAULT_EVICTION_POLICY,
- DEFAULT_WHITELIST,
- _LocalCache,
-)
 from redis.client import (
  EMPTY_RESPONSE,
  NEVER_DECODE,
@@ -67,7 +67,7 @@
  TimeoutError,
  WatchError,
 )
-from redis.typing import ChannelT, EncodableT, KeysT, KeyT, ResponseT
+from redis.typing import ChannelT, EncodableT, KeyT
 from redis.utils import (
  HIREDIS_AVAILABLE,
  _set_info_logger,
@@ -294,6 +294,13 @@ def __init__(
  "lib_version": lib_version,
  "redis_connect_func": redis_connect_func,
  "protocol": protocol,
+ "cache_enable": cache_enable,
+ "client_cache": client_cache,
+ "cache_max_size": cache_max_size,
+ "cache_ttl": cache_ttl,
+ "cache_eviction_policy": cache_eviction_policy,
+ "cache_blacklist": cache_blacklist,
+ "cache_whitelist": cache_whitelist,
  }
  # based on input, setup appropriate connection args
  if unix_socket_path is not None:
@@ -350,16 +357,6 @@ def __init__(
  # on a set of redis commands
  self._single_conn_lock = asyncio.Lock()
 
- self.client_cache = client_cache
- if cache_enable:
- self.client_cache = _LocalCache(
- cache_max_size, cache_ttl, cache_eviction_policy
- )
- if self.client_cache is not None:
- self.cache_blacklist = cache_blacklist
- self.cache_whitelist = cache_whitelist
- self.client_cache_initialized = False
-
  def __repr__(self):
  return (
  f"<{self.__class__.__module__}.{self.__class__.__name__}"
@@ -374,10 +371,6 @@ async def initialize(self: _RedisT) -> _RedisT:
  async with self._single_conn_lock:
  if self.connection is None:
  self.connection = await self.connection_pool.get_connection("_")
- if self.client_cache is not None:
- self.connection._parser.set_invalidation_push_handler(
- self._cache_invalidation_process
- )
  return self
 
  def set_response_callback(self, command: str, callback: ResponseCallbackT):
@@ -596,8 +589,6 @@ async def aclose(self, close_connection_pool: Optional[bool] = None) -> None:
  close_connection_pool is None and self.auto_close_connection_pool
  ):
  await self.connection_pool.disconnect()
- if self.client_cache:
- self.client_cache.flush()
 
  @deprecated_function(version="5.0.1", reason="Use aclose() instead", name="close")
  async def close(self, close_connection_pool: Optional[bool] = None) -> None:
@@ -626,89 +617,28 @@ async def _disconnect_raise(self, conn: Connection, error: Exception):
  ):
  raise error
 
- def _cache_invalidation_process(
- self, data: List[Union[str, Optional[List[str]]]]
- ) -> None:
- """
- Invalidate (delete) all redis commands associated with a specific key.
- `data` is a list of strings, where the first string is the invalidation message
- and the second string is the list of keys to invalidate.
- (if the list of keys is None, then all keys are invalidated)
- """
- if data[1] is not None:
- for key in data[1]:
- self.client_cache.invalidate(str_if_bytes(key))
- else:
- self.client_cache.flush()
-
- async def _get_from_local_cache(self, command: str):
- """
- If the command is in the local cache, return the response
- """
- if (
- self.client_cache is None
- or command[0] in self.cache_blacklist
- or command[0] not in self.cache_whitelist
- ):
- return None
- while not self.connection._is_socket_empty():
- await self.connection.read_response(push_request=True)
- return self.client_cache.get(command)
-
- def _add_to_local_cache(
- self, command: Tuple[str], response: ResponseT, keys: List[KeysT]
- ):
- """
- Add the command and response to the local cache if the command
- is allowed to be cached
- """
- if (
- self.client_cache is not None
- and (self.cache_blacklist == [] or command[0] not in self.cache_blacklist)
- and (self.cache_whitelist == [] or command[0] in self.cache_whitelist)
- ):
- self.client_cache.set(command, response, keys)
-
- def delete_from_local_cache(self, command: str):
- """
- Delete the command from the local cache
- """
- try:
- self.client_cache.delete(command)
- except AttributeError:
- pass
-
  # COMMAND EXECUTION AND PROTOCOL PARSING
  async def execute_command(self, *args, **options):
  """Execute a command and return a parsed response"""
  await self.initialize()
  command_name = args[0]
  keys = options.pop("keys", None) # keys are used only for client side caching
- response_from_cache = await self._get_from_local_cache(args)
+ pool = self.connection_pool
+ conn = self.connection or await pool.get_connection(command_name, **options)
+ response_from_cache = await conn._get_from_local_cache(args)
  if response_from_cache is not None:
  return response_from_cache
  else:
- pool = self.connection_pool
- conn = self.connection or await pool.get_connection(command_name, **options)
-
  if self.single_connection_client:
  await self._single_conn_lock.acquire()
  try:
- if self.client_cache is not None and not self.client_cache_initialized:
- await conn.retry.call_with_retry(
- lambda: self._send_command_parse_response(
- conn, "CLIENT", *("CLIENT", "TRACKING", "ON")
- ),
- lambda error: self._disconnect_raise(conn, error),
- )
- self.client_cache_initialized = True
  response = await conn.retry.call_with_retry(
  lambda: self._send_command_parse_response(
  conn, command_name, *args, **options
  ),
  lambda error: self._disconnect_raise(conn, error),
  )
- self._add_to_local_cache(args, response, keys)
+ conn._add_to_local_cache(args, response, keys)
  return response
  finally:
  if self.single_connection_client:

redis/asyncio/connection.py

@@ -47,9 +47,15 @@
  ResponseError,
  TimeoutError,
 )
-from redis.typing import EncodableT
+from redis.typing import EncodableT, KeysT, ResponseT
 from redis.utils import HIREDIS_AVAILABLE, get_lib_version, str_if_bytes
 
+from .._cache import (
+ DEFAULT_BLACKLIST,
+ DEFAULT_EVICTION_POLICY,
+ DEFAULT_WHITELIST,
+ _LocalCache,
+)
 from .._parsers import (
  BaseParser,
  Encoder,
@@ -114,6 +120,9 @@ class AbstractConnection:
  "encoder",
  "ssl_context",
  "protocol",
+ "client_cache",
+ "cache_blacklist",
+ "cache_whitelist",
  "_reader",
  "_writer",
  "_parser",
@@ -148,6 +157,13 @@ def __init__(
  encoder_class: Type[Encoder] = Encoder,
  credential_provider: Optional[CredentialProvider] = None,
  protocol: Optional[int] = 2,
+ cache_enable: bool = False,
+ client_cache: Optional[_LocalCache] = None,
+ cache_max_size: int = 100,
+ cache_ttl: int = 0,
+ cache_eviction_policy: str = DEFAULT_EVICTION_POLICY,
+ cache_blacklist: List[str] = DEFAULT_BLACKLIST,
+ cache_whitelist: List[str] = DEFAULT_WHITELIST,
  ):
  if (username or password) and credential_provider is not None:
  raise DataError(
@@ -205,6 +221,14 @@ def __init__(
  if p < 2 or p > 3:
  raise ConnectionError("protocol must be either 2 or 3")
  self.protocol = protocol
+ if cache_enable:
+ _cache = _LocalCache(cache_max_size, cache_ttl, cache_eviction_policy)
+ else:
+ _cache = None
+ self.client_cache = client_cache if client_cache is not None else _cache
+ if self.client_cache is not None:
+ self.cache_blacklist = cache_blacklist
+ self.cache_whitelist = cache_whitelist
 
  def __del__(self, _warnings: Any = warnings):
  # For some reason, the individual streams don't get properly garbage
@@ -395,6 +419,11 @@ async def on_connect(self) -> None:
  # if a database is specified, switch to it. Also pipeline this
  if self.db:
  await self.send_command("SELECT", self.db)
+ # if client caching is enabled, start tracking
+ if self.client_cache:
+ await self.send_command("CLIENT", "TRACKING", "ON")
+ await self.read_response()
+ self._parser.set_invalidation_push_handler(self._cache_invalidation_process)
 
  # read responses from pipeline
  for _ in (sent for sent in (self.lib_name, self.lib_version) if sent):
@@ -429,6 +458,9 @@ async def disconnect(self, nowait: bool = False) -> None:
  raise TimeoutError(
  f"Timed out closing connection after {self.socket_connect_timeout}"
  ) from None
+ finally:
+ if self.client_cache:
+ self.client_cache.flush()
 
  async def _send_ping(self):
  """Send PING, expect PONG in return"""
@@ -646,10 +678,62 @@ def pack_commands(self, commands: Iterable[Iterable[EncodableT]]) -> List[bytes]
  output.append(SYM_EMPTY.join(pieces))
  return output
 
- def _is_socket_empty(self):
+ def _socket_is_empty(self):
  """Check if the socket is empty"""
  return not self._reader.at_eof()
 
+ def _cache_invalidation_process(
+ self, data: List[Union[str, Optional[List[str]]]]
+ ) -> None:
+ """
+ Invalidate (delete) all redis commands associated with a specific key.
+ `data` is a list of strings, where the first string is the invalidation message
+ and the second string is the list of keys to invalidate.
+ (if the list of keys is None, then all keys are invalidated)
+ """
+ if data[1] is not None:
+ self.client_cache.flush()
+ else:
+ for key in data[1]:
+ self.client_cache.invalidate(str_if_bytes(key))
+
+ async def _get_from_local_cache(self, command: str):
+ """
+ If the command is in the local cache, return the response
+ """
+ if (
+ self.client_cache is None
+ or command[0] in self.cache_blacklist
+ or command[0] not in self.cache_whitelist
+ ):
+ return None
+ while not self._socket_is_empty():
+ await self.read_response(push_request=True)
+ return self.client_cache.get(command)
+
+ def _add_to_local_cache(
+ self, command: Tuple[str], response: ResponseT, keys: List[KeysT]
+ ):
+ """
+ Add the command and response to the local cache if the command
+ is allowed to be cached
+ """
+ if (
+ self.client_cache is not None
+ and (self.cache_blacklist == [] or command[0] not in self.cache_blacklist)
+ and (self.cache_whitelist == [] or command[0] in self.cache_whitelist)
+ ):
+ self.client_cache.set(command, response, keys)
+
+ def delete_from_local_cache(self, command: str):
+ """
+ Delete the command from the local cache
+ """
+ try:
+ self.client_cache.delete(command)
+ except AttributeError:
+ pass
+
 
 class Connection(AbstractConnection):
  "Manages TCP communication to and from a Redis server"