fix(csc): avoid draining checked-out cache connections

redis/connection.py

@@ -2058,6 +2058,9 @@ class CacheProxyConnection(MaintNotificationsAbstractConnection, ConnectionInter
     DUMMY_CACHE_VALUE = b"foo"
     MIN_ALLOWED_VERSION = "7.4.0"
     DEFAULT_SERVER_NAME = "redis"
+    # Bound peer reads when a checked-out owner stops processing invalidations.
+    # Only the owner's commands renew this window; peer hits do not extend it.
+    _BUSY_CACHE_GRACE_SECONDS = 0.1
 
     def __init__(
         self,
@@ -2074,6 +2077,11 @@ def __init__(
         self._event_dispatcher = self._conn._event_dispatcher
         self.credential_provider = conn.credential_provider
         self._pool_lock = pool_lock
+        # Checkout and peer invalidation drains share this connection's lock.
+        # Release clears _in_use after the last socket access, before enqueueing.
+        self._in_use_lock = threading.Lock()
+        self._in_use = True
+        self._busy_cache_deadline = 0.0
         self._cache = cache
         self._cache_lock = threading.RLock()
         self._current_command_cache_key = None
@@ -2170,6 +2178,7 @@ def on_connect(self):
         self._conn.on_connect()
 
     def disconnect(self, *args, **kwargs):
+        self._busy_cache_deadline = 0.0
         with self._cache_lock:
             self._cache.flush()
         # Both flags describe an exchange on the socket that is about to die with it. A stale
@@ -2202,6 +2211,7 @@ def send_packed_command(self, command, check_health=True):
         self._pending_caching_reply = False
         self._skip_next_caching = False
         self._pending_cache_hit = _NO_PENDING_HIT
+        self._busy_cache_deadline = 0.0
         self._conn.send_packed_command(command)
 
     def send_command(self, *args, **kwargs):
@@ -2221,6 +2231,7 @@ def send_command(self, *args, **kwargs):
             # slot, and the attempt is tracked by default under optout anyway.
             self._skip_next_caching = False
             self._current_command_cache_key = None
+            self._busy_cache_deadline = 0.0
             self._conn.send_command(*args, **kwargs)
             return
 
@@ -2264,6 +2275,9 @@ def send_command(self, *args, **kwargs):
 
         if not store:
             self._current_command_cache_key = None
+            # PubSub, blocking commands and unknown commands may hold this
+            # connection indefinitely without processing invalidations.
+            self._busy_cache_deadline = 0.0
 
             # One path for keyless and ineligible commands alike, so an eligible read whose
             # command method never plumbed ``keys=`` still gets its optout ``NO``. That is the
@@ -2303,25 +2317,36 @@ def send_command(self, *args, **kwargs):
                         cache_key=self._current_command_cache_key,
                         cache_value=self.DUMMY_CACHE_VALUE,
                         status=CacheEntryStatus.IN_PROGRESS,
-                        connection_ref=self._conn,
+                        connection_ref=self,
                     )
                 )
 
         if candidate is not None:
             # Drain the connection that cached the entry so invalidations do not
             # accumulate on idle connections. This must run outside our cache lock: the
-            # drain takes the pool lock, and invalidations call the owner's callback
-            # under its cache lock. Holding ours would order a cache lock before the
-            # pool lock and another after it, deadlocking connections draining each
-            # other or a pool.disconnect() holding the pool lock while waiting for ours.
-            with self._pool_lock:
-                self._drain_invalidations(candidate.connection_ref)
+            # drain takes the owner's checkout lock, and invalidations call the
+            # owner's callback under its cache lock. Holding ours could deadlock
+            # two connections draining each other.
+            owner = candidate.connection_ref
+            can_serve = False
+            if isinstance(owner, CacheProxyConnection):
+                with owner._in_use_lock:
+                    if owner is self or not owner._in_use:
+                        self._drain_invalidations(owner._conn)
+                        can_serve = True
+                    elif time.monotonic() < owner._busy_cache_deadline:
+                        # Never touch a checked-out owner's socket. Recent cache
+                        # reads can share its value until its grace period expires.
+                        # Hits do not move ownership or extend that deadline.
+                        can_serve = True
+            # Custom entries without pool ownership information cannot be drained
+            # safely and must use a fresh read on our connection.
 
             with self._cache_lock:
                 # Re-check: the drained entry may have been invalidated or replaced.
                 entry = self._cache.get(self._current_command_cache_key)
 
-                # A hit requires this very entry object to have been drained. A
+                # A hit requires this very entry object to have been checked. A
                 # replacement filled during the drain may belong to another connection
                 # with pending invalidations and is not trusted, even if it is VALID by
                 # the time we check here.
@@ -2336,7 +2361,7 @@ def send_command(self, *args, **kwargs):
                 #   and never read the reply we did send;
                 # - an invalidation removes it, so a re-derived check would find nothing
                 #   and read a reply for a command we never sent.
-                if entry is candidate:
+                if can_serve and entry is candidate:
                     cached_value = entry.cache_value
                     self._current_command_cache_key = None
                     is_hit = True
@@ -2347,14 +2372,18 @@ def send_command(self, *args, **kwargs):
                             cache_key=self._current_command_cache_key,
                             cache_value=self.DUMMY_CACHE_VALUE,
                             status=CacheEntryStatus.IN_PROGRESS,
-                            connection_ref=self._conn,
+                            connection_ref=self,
                         )
                     )
 
         if is_hit:
             self._pending_cache_hit = _clone_cache_value(cached_value)
             return
 
+        # A wire reply may be followed by invalidations we have not read yet.
+        # Keep peers on the fresh-read path until this exchange completes.
+        self._busy_cache_deadline = 0.0
+
         # Send command over socket only if it's allowed
         # read-only command that not yet cached. Under optin the server remembers nothing
         # unless ``CLIENT CACHING YES`` comes right before the read, so this - the one place
@@ -2466,14 +2495,20 @@ def read_response(
 
             cache_entry = self._cache.get(self._current_command_cache_key)
 
+            # Peers use independent cache locks. Start the sharing window before
+            # making the reply visible as VALID to avoid unnecessary misses.
+            self._busy_cache_deadline = (
+                time.monotonic() + self._BUSY_CACHE_GRACE_SECONDS
+            )
+
             # Promote only the placeholder this connection staked. It may have been
             # invalidated in the meantime, or replaced by another connection's placeholder
             # for the same key - promoting that one would bind our reply to its
             # ``connection_ref``, and a later hit would drain that connection's own reply.
             if (
                 cache_entry is not None
                 and cache_entry.status == CacheEntryStatus.IN_PROGRESS
-                and cache_entry.connection_ref is self._conn
+                and cache_entry.connection_ref is self
             ):
                 cache_entry.cache_value = cloned_response
                 cache_entry.status = CacheEntryStatus.VALID
@@ -2509,7 +2544,7 @@ def _drop_own_placeholder(self) -> None:
             if (
                 entry is not None
                 and entry.status == CacheEntryStatus.IN_PROGRESS
-                and entry.connection_ref is self._conn
+                and entry.connection_ref is self
             ):
                 self._cache.delete_by_cache_keys([key])
 
@@ -2782,7 +2817,7 @@ def _enable_tracking_callback(self, conn: ConnectionInterface) -> None:
                     if (
                         entry is not None
                         and entry.status == CacheEntryStatus.IN_PROGRESS
-                        and entry.connection_ref is self._conn
+                        and entry.connection_ref is self
                     ):
                         own_placeholder = entry
                 self._cache.flush()
@@ -2829,6 +2864,7 @@ def _drain_invalidations(self, conn: ConnectionInterface):
 
     def _process_pending_invalidations(self):
         self._drain_invalidations(self._conn)
+        self._busy_cache_deadline = time.monotonic() + self._BUSY_CACHE_GRACE_SECONDS
 
     def _on_invalidation_callback(self, data: List[Union[str, Optional[List[bytes]]]]):
         with self._cache_lock:
@@ -4185,6 +4221,11 @@ def get_connection(self, command_name=None, *keys, **options) -> "Connection":
                 is_created = True
             self._in_use_connections.add(connection)
 
+        if isinstance(connection, CacheProxyConnection):
+            # Wait only for this connection's drain, without holding the pool lock.
+            with connection._in_use_lock:
+                connection._in_use = True
+
         # Record state transition for observability
         pool_name = get_pool_name(self)
         if is_created:
@@ -4299,6 +4340,8 @@ def release(self, connection: "Connection") -> None:
                             f"{connection}, {connection.extract_connection_details()}"
                         )
                     connection.disconnect()
+                if isinstance(connection, CacheProxyConnection):
+                    connection._in_use = False
                 self._available_connections.append(connection)
                 self._event_dispatcher.dispatch(
                     AfterConnectionReleasedEvent(connection)
@@ -4687,6 +4730,11 @@ def get_connection(self, command_name=None, *keys, **options):
                     pass
                 self._locked = False
 
+        if isinstance(connection, CacheProxyConnection):
+            # Join this connection's drain before connect/can_read uses the socket.
+            with connection._in_use_lock:
+                connection._in_use = True
+
         # Record state transition for observability
         pool_name = get_pool_name(self)
         if is_created:
@@ -4771,6 +4819,8 @@ def release(self, connection):
                     )
                 connection.disconnect()
             # Put the connection back into the pool.
+            if isinstance(connection, CacheProxyConnection):
+                connection._in_use = False
             pool_name = get_pool_name(self)
             try:
                 self.pool.put_nowait(connection)