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

redis/connection.py

@@ -2056,6 +2056,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.
+    # Active readers renew this window without extra network requests.
+    _BUSY_CACHE_GRACE_SECONDS = 0.1
 
     def __init__(
         self,
@@ -2072,6 +2075,10 @@ def __init__(
         self._event_dispatcher = self._conn._event_dispatcher
         self.credential_provider = conn.credential_provider
         self._pool_lock = pool_lock
+        # Only a pool release makes the connection available for another
+        # connection's invalidation drain. Pools update this under _pool_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
@@ -2168,6 +2175,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
@@ -2200,6 +2208,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):
@@ -2219,6 +2228,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
 
@@ -2262,6 +2272,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
@@ -2296,12 +2309,23 @@ def send_command(self, *args, **kwargs):
             # (which returns non-push replies too) would consume either one.
             entry = self._cache.get(self._current_command_cache_key)
             if entry is not None and entry.status != CacheEntryStatus.IN_PROGRESS:
-                with self._pool_lock:
-                    self._drain_invalidations(entry.connection_ref)
-
-                # Re-check: the entry may have been invalidated during the drain, or filled
-                # in by the connection that was fetching it.
-                entry = self._cache.get(self._current_command_cache_key)
+                owner = entry.connection_ref
+                if isinstance(owner, CacheProxyConnection):
+                    with owner._pool_lock:
+                        if owner is self or not owner._in_use:
+                            self._drain_invalidations(owner._conn)
+                            # The drain may have invalidated the cached value.
+                            entry = self._cache.get(self._current_command_cache_key)
+                        elif time.monotonic() >= owner._busy_cache_deadline:
+                            # Never touch a checked-out owner's socket. Recent
+                            # cache reads can share its value, but a pending wire
+                            # reply or an owner that stopped reading needs a fresh
+                            # read on our connection. Hits do not move ownership.
+                            entry = None
+                else:
+                    # A custom entry without pool ownership information cannot
+                    # be drained safely.
+                    entry = None
 
             # Whether this read is served locally is decided here, once, and the value is
             # carried to ``read_response`` - which must not re-derive it from cache state.
@@ -2330,14 +2354,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
@@ -2456,13 +2484,16 @@ def read_response(
             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
                 self._cache.set(cache_entry)
 
             self._current_command_cache_key = None
+            self._busy_cache_deadline = (
+                time.monotonic() + self._BUSY_CACHE_GRACE_SECONDS
+            )
 
         return response
 
@@ -2492,7 +2523,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])
 
@@ -2765,7 +2796,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()
@@ -2812,6 +2843,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:
@@ -4133,6 +4165,8 @@ def get_connection(self, command_name=None, *keys, **options) -> "Connection":
                 connection = self.make_connection()
                 is_created = True
             self._in_use_connections.add(connection)
+            if isinstance(connection, CacheProxyConnection):
+                connection._in_use = True
 
         # Record state transition for observability
         pool_name = get_pool_name(self)
@@ -4248,6 +4282,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)
@@ -4612,6 +4648,12 @@ def get_connection(self, command_name=None, *keys, **options):
                     pass
                 self._locked = False
 
+        if isinstance(connection, CacheProxyConnection):
+            # Queue checkout is independent of _lock. Join the invalidation
+            # drain's critical section before connect/can_read can use the socket.
+            with self._lock:
+                connection._in_use = True
+
         # Record state transition for observability
         pool_name = get_pool_name(self)
         if is_created:
@@ -4696,6 +4738,9 @@ def release(self, connection):
                     )
                 connection.disconnect()
             # Put the connection back into the pool.
+            if isinstance(connection, CacheProxyConnection):
+                with self._lock:
+                    connection._in_use = False
             pool_name = get_pool_name(self)
             try:
                 self.pool.put_nowait(connection)