fix cache proxy connection unbalanced command send and read

Author: tzongw

redis/connection.py

@@ -1365,7 +1365,7 @@ def __init__(
         self._pool_lock = pool_lock
         self._cache = cache
         self._cache_lock = threading.RLock()
-        self._current_command_cache_key = None
+        self._current_command_cache_entry = None
         self._current_options = None
         self.register_connect_callback(self._enable_tracking_callback)
 
@@ -1453,42 +1453,49 @@ def send_command(self, *args, **kwargs):
             if not self._cache.is_cachable(
                 CacheKey(command=args[0], redis_keys=(), redis_args=())
             ):
-                self._current_command_cache_key = None
+                self._current_command_cache_entry = None
                 self._conn.send_command(*args, **kwargs)
                 return
 
         if kwargs.get("keys") is None:
             raise ValueError("Cannot create cache key.")
 
-        # Creates cache key.
-        self._current_command_cache_key = CacheKey(
+        cache_key = CacheKey(
             command=args[0], redis_keys=tuple(kwargs.get("keys")), redis_args=args
         )
+        self._current_command_cache_entry = None
 
         with self._cache_lock:
-            # We have to trigger invalidation processing in case if
-            # it was cached by another connection to avoid
-            # queueing invalidations in stale connections.
-            if self._cache.get(self._current_command_cache_key):
-                entry = self._cache.get(self._current_command_cache_key)
-
-                if entry.connection_ref != self._conn:
+            cache_entry = self._cache.get(cache_key)
+            if cache_entry is not None and cache_entry.status == CacheEntryStatus.VALID:
+                # We have to trigger invalidation processing in case if
+                # it was cached by another connection to avoid
+                # queueing invalidations in stale connections.
+                if cache_entry.connection_ref != self._conn:
                     with self._pool_lock:
-                        while entry.connection_ref.can_read():
-                            entry.connection_ref.read_response(push_request=True)
-
-                return
+                        while cache_entry.connection_ref.can_read():
+                            cache_entry.connection_ref.read_response(push_request=True)
+                    # Check if entry still exists.
+                    if self._cache.get(cache_key) is not None:
+                        self._current_command_cache_entry = cache_entry
+                        return
+                    cache_entry = None
+                else:
+                    self._current_command_cache_entry = cache_entry
+                    return
 
-            # Set temporary entry value to prevent
-            # race condition from another connection.
-            self._cache.set(
-                CacheEntry(
-                    cache_key=self._current_command_cache_key,
+            if cache_entry is None:
+                # Creates cache entry.
+                cache_entry = CacheEntry(
+                    cache_key=cache_key,
                     cache_value=self.DUMMY_CACHE_VALUE,
                     status=CacheEntryStatus.IN_PROGRESS,
                     connection_ref=self._conn,
                 )
-            )
+                # Set temporary entry value to prevent
+                # race condition from another connection.
+                self._cache.set(cache_entry)
+                self._current_command_cache_entry = cache_entry
 
         # Send command over socket only if it's allowed
         # read-only command that not yet cached.
@@ -1501,17 +1508,15 @@ def read_response(
         self, disable_decoding=False, *, disconnect_on_error=True, push_request=False
     ):
         with self._cache_lock:
-            # Check if command response exists in a cache and it's not in progress.
+            # Check if command response cache entry exists and it's valid.
             if (
-                self._current_command_cache_key is not None
-                and self._cache.get(self._current_command_cache_key) is not None
-                and self._cache.get(self._current_command_cache_key).status
-                != CacheEntryStatus.IN_PROGRESS
+                self._current_command_cache_entry is not None
+                and self._current_command_cache_entry.status == CacheEntryStatus.VALID
             ):
                 res = copy.deepcopy(
-                    self._cache.get(self._current_command_cache_key).cache_value
+                    self._current_command_cache_entry.cache_value
                 )
-                self._current_command_cache_key = None
+                self._current_command_cache_entry = None
                 return res
 
         response = self._conn.read_response(
@@ -1522,23 +1527,24 @@ def read_response(
 
         with self._cache_lock:
             # Prevent not-allowed command from caching.
-            if self._current_command_cache_key is None:
+            if self._current_command_cache_entry is None:
                 return response
             # If response is None prevent from caching.
+            cache_key = self._current_command_cache_entry.cache_key
             if response is None:
-                self._cache.delete_by_cache_keys([self._current_command_cache_key])
+                self._cache.delete_by_cache_keys([cache_key])
                 return response
 
-            cache_entry = self._cache.get(self._current_command_cache_key)
+            cache_entry = self._cache.get(cache_key)
 
             # Cache only responses that still valid
             # and wasn't invalidated by another connection in meantime.
-            if cache_entry is not None:
+            if cache_entry is self._current_command_cache_entry:
                 cache_entry.status = CacheEntryStatus.VALID
                 cache_entry.cache_value = response
                 self._cache.set(cache_entry)
 
-            self._current_command_cache_key = None
+            self._current_command_cache_entry = None
 
         return response
 

tests/test_connection.py

@@ -526,7 +526,7 @@ def test_read_response_returns_cached_reply(self, mock_cache, mock_connection):
         )
         proxy_connection.send_command(*["GET", "foo"], **{"keys": ["foo"]})
         assert proxy_connection.read_response() == b"bar"
-        assert proxy_connection._current_command_cache_key is None
+        assert proxy_connection._current_command_cache_entry is None
         assert proxy_connection.read_response() == b"bar"
 
         mock_cache.set.assert_has_calls(