pool: fence failures from retired connections

cassandra/_lifecycle.py

@@ -20,7 +20,7 @@
 
 from collections import deque
 from concurrent.futures import Future
-from contextlib import contextmanager, nullcontext
+from contextlib import contextmanager
 from functools import partial, wraps
 import logging
 from threading import local, RLock
@@ -56,9 +56,73 @@ def owns_locked(self):
                     self.host, self.generation, self.expected_pool, self.endpoint))
 
     def is_current(self):
+        with self.locked():
+            return self.owns_locked()
+
+    @contextmanager
+    def locked(self):
         with self.session._lock:
             with self.host.lock:
-                return self.owns_locked()
+                yield
+
+
+class ConnectionFailureReservation:
+    """Retain the pool and shard slot that own a connection failure."""
+
+    def __init__(self, pool, connection, replacing=False):
+        self.pool = pool
+        self.connection = connection
+        self.session = pool._session_ref()
+        self.host = pool.host
+        self.replacing = replacing
+        self.retired = False
+
+    @contextmanager
+    def locked(self):
+        if self.session is None:
+            yield
+            return
+        while True:
+            with self.session._lock:
+                with self.host.lock:
+                    # Unpublished candidates can hold their work lock while
+                    # publication waits for Session under the membership lock.
+                    # A completed retirement no longer needs shard ownership
+                    # inspected; its receipt retains only Session/Host guards.
+                    if (self.retired or
+                            self.session._pools.get(self.host) is not self.pool):
+                        yield
+                        return
+                    if self.pool._lock.acquire(False):
+                        try:
+                            yield
+                        finally:
+                            self.pool._lock.release()
+                        return
+            # Never retain work while reacquiring Session/Host. Active socket
+            # failures may await I/O under work before they can retire a pool.
+            with self.pool._lock:
+                pass
+
+    def owns_locked(self):
+        if (self.session is None or
+                self.session._pools.get(self.host) is not self.pool):
+            return False
+        owns_slot = self.retired
+        if not self.retired and not self.pool.is_shutdown:
+            current = self.pool._connections.get(self.connection.features.shard_id)
+            owns_slot = (current is self.connection or
+                         (self.replacing and current is None and
+                          self.pool._replacement_connections.get(
+                              self.connection.features.shard_id) is self.connection))
+        return (not self.session.is_shutdown and
+                not self.session.cluster.is_shutdown and
+                not self.host._is_removed and
+                self.host.endpoint == self.pool.endpoint and owns_slot)
+
+    def is_current(self):
+        with self.locked():
+            return self.owns_locked()
 
 
 class LifecycleCoordinator:
@@ -196,27 +260,26 @@ def observe_health(self, host):
     def begin_down(self, host, health_generation, connected=False,
                    expect_down=False, pool_failure=None):
         """Return (generation, retry_only), or None if DOWN is discounted."""
-        failure_lock = (pool_failure.session._lock
-                        if pool_failure is not None else nullcontext())
+        failure_lock = (pool_failure.locked()
+                        if pool_failure is not None else host.lock)
         with failure_lock:
-            with host.lock:
-                if (self.closed or host._is_removed or connected or
-                        health_generation != host._health_generation or
-                        (pool_failure is not None and not pool_failure.owns_locked())):
-                    return None
-                host._health_generation += 1
-                self.cancel_preparation(host)
-                host._readiness_generation = (
-                    getattr(host, '_readiness_generation', 0) + 1)
-                host._currently_handling_node_up = False
-                was_up = host.is_up
-                host.set_down()
-                if (not was_up and (host.is_currently_reconnecting() or
-                                    host._currently_handling_node_down)):
-                    return None
-                host._currently_handling_node_down = True
-                host._down_event_generation += 1
-                return host._down_event_generation, not was_up and not expect_down
+            if (self.closed or host._is_removed or connected or
+                    health_generation != host._health_generation or
+                    (pool_failure is not None and not pool_failure.owns_locked())):
+                return None
+            host._health_generation += 1
+            self.cancel_preparation(host)
+            host._readiness_generation = (
+                getattr(host, '_readiness_generation', 0) + 1)
+            host._currently_handling_node_up = False
+            was_up = host.is_up
+            host.set_down()
+            if (not was_up and (host.is_currently_reconnecting() or
+                                host._currently_handling_node_down)):
+                return None
+            host._currently_handling_node_down = True
+            host._down_event_generation += 1
+            return host._down_event_generation, not was_up and not expect_down
 
     def dispatch_down(self, host, health_generation, is_host_addition,
                       connected, expect_host_to_be_down, on_add_reconnection,
@@ -242,8 +305,8 @@ def dispatch_down(self, host, health_generation, is_host_addition,
                         on_add_reconnection=on_add_reconnection)
                 if future is None:
                     self.finish_down(host, down_event_generation)
-                # Restarting the host reconnector does not dispatch the DOWN
-                # callbacks that reconnect the control connection.
+                # Restarting the host reconnector does not dispatch
+                # DOWN callbacks that reconnect the control connection.
                 return False
 
             log.warning("Host %s has been marked down", host)

cassandra/cluster.py

@@ -2655,6 +2655,12 @@ def on_down(self, host, is_host_addition, expect_host_to_be_down=False,
         if (self._discount_down_events and
                 self.profile_manager.distance(host) != HostDistance.IGNORED):
             for session in tuple(self.sessions):
+                # This failure will shut down its pool, including any healthy
+                # sibling shards. Only independent pools can discount DOWN.
+                if (isinstance(_pool_failure,
+                               _lifecycle.ConnectionFailureReservation) and
+                        session is _pool_failure.session):
+                    continue
                 pool_state = session.get_pool_state().get(host)
                 if pool_state:
                     connected |= pool_state['open_count'] > 0

cassandra/pool.py

@@ -30,7 +30,7 @@
 except ImportError:
     from cassandra.util import WeakSet  # NOQA
 
-from cassandra import AuthenticationFailed
+from cassandra import AuthenticationFailed, _lifecycle
 from cassandra.connection import (ConnectionException, ConnectionSendError, EndPoint,
                                   DefaultEndPoint, UnixSocketEndPoint)
 from cassandra.policies import HostDistance
@@ -468,11 +468,12 @@ def __init__(self, host, host_distance, session):
         self.endpoint = host.endpoint
         self.host_distance = host_distance
         self._session = weakref.proxy(session)
+        self._session_ref = weakref.ref(session)
         self._lock = Lock()
         self._membership_lock = Lock()
         # this is used in conjunction with the connection streams. Not using the connection lock because the connection can be replaced in the lifetime of the pool.
         self._stream_available_condition = Condition(Lock())
-        self._is_replacing = False
+        self._replacement_connections = {}
         self._connecting = set()
         self._connections = {}
         self._pending_connections = []
@@ -654,41 +655,62 @@ def return_connection(self, connection, stream_was_orphaned=False,
                 self._stream_available_condition.notify()
 
         if connection.is_defunct or connection.is_closed:
-            if connection.signaled_error and not replace_only and \
-                    (not self.shutdown_on_error or
-                     (isinstance(connection.last_error, ConnectionSendError) and
-                      not heartbeat_failed)):
+            # Claim the transport before pool work: a redundant reactor
+            # callback must return while that work waits for its response.
+            with connection.lock:
+                if connection.signaled_error and not replace_only and \
+                        (not self.shutdown_on_error or
+                         (isinstance(connection.last_error, ConnectionSendError) and
+                          not heartbeat_failed)):
+                    return
+                report_failure = not replace_only and not connection.signaled_error
+                connection.signaled_error = True
+            reservation = _lifecycle.ConnectionFailureReservation(self, connection)
+            with reservation.locked():
+                current = reservation.owns_locked()
+            if not current:
+                with self._lock:
+                    self._trash.discard(connection)
+                # A stale completion still releases its stream above, but its
+                # failed transport cannot affect the current pool or shard.
+                connection.close()
                 return
 
             is_down = False
-            if replace_only:
-                # A local send-queue failure says nothing about host health.
-                # Mark the connection handled so later callbacks do not report
-                # this failure to the host conviction policy.
-                connection.signaled_error = True
-            elif not connection.signaled_error:
+            if report_failure:
                 log.debug("Defunct or closed connection (%s) returned to pool, potentially "
                           "marking host %s as down", id(connection), self.host)
                 is_down = self.host.signal_connection_failure(connection.last_error)
-                connection.signaled_error = True
 
             if self.shutdown_on_error and not is_down and not replace_only:
                 is_down = True
 
             if is_down:
-                self.shutdown()
-                self._session.cluster.on_down(self.host, is_host_addition=False)
+                # Retire before inspecting other sessions: concurrent failures
+                # must not discount each other's doomed sibling connections.
+                self.shutdown(_failure=reservation)
+                if reservation.is_current():
+                    self._session.cluster.on_down(
+                        self.host, is_host_addition=False,
+                        _pool_failure=reservation)
+                elif not reservation.retired:
+                    with self._lock:
+                        self._trash.discard(connection)
+                    connection.close()
             else:
+                schedule_replacement = False
+                with reservation.locked():
+                    if reservation.owns_locked():
+                        with self._membership_lock:
+                            self._connections.pop(connection.features.shard_id)
+                        self._replacement_connections[
+                            connection.features.shard_id] = connection
+                        schedule_replacement = True
+                if not schedule_replacement:
+                    with self._lock:
+                        self._trash.discard(connection)
                 connection.close()
-                with self._lock:
-                    if self.is_shutdown:
-                        return
-                    with self._membership_lock:
-                        self._connections.pop(
-                            connection.features.shard_id, None)
-                    if self._is_replacing:
-                        return
-                    self._is_replacing = True
+                if schedule_replacement:
                     self._session.submit(self._replace, connection)
         elif connection in self._trash:
             with connection.lock:
@@ -712,11 +734,25 @@ def on_orphaned_stream_released(self):
             self._stream_available_condition.notify()
 
     def _replace(self, connection):
+        reservation = _lifecycle.ConnectionFailureReservation(
+            self, connection, replacing=True)
+        if not reservation.is_current():
+            with self._lock:
+                shard_id = connection.features.shard_id
+                if self._replacement_connections.get(shard_id) is connection:
+                    self._replacement_connections.pop(shard_id)
+            return
         failed_connection = None
         with self._lock:
-            if self.is_shutdown:
+            shard_id = connection.features.shard_id
+            current = self._connections.get(shard_id)
+            if (self.is_shutdown or
+                    (current is not None and current is not connection) or
+                    (current is None and
+                     self._replacement_connections.get(shard_id) is not connection)):
                 return
 
+            self._replacement_connections[shard_id] = connection
             log.debug("Replacing connection (%s) to %s", id(connection), self.host)
             try:
                 with self._membership_lock:
@@ -725,41 +761,52 @@ def _replace(self, connection):
                     self._connecting.add(connection.features.shard_id)
                     self._session.submit(self._open_connection_to_missing_shard, connection.features.shard_id)
                 else:
-                    connection = self._session.cluster.connection_factory(self.endpoint,
-                                                                          on_orphaned_stream_released=self.on_orphaned_stream_released)
+                    new_connection = self._session.cluster.connection_factory(
+                        self.endpoint,
+                        on_orphaned_stream_released=self.on_orphaned_stream_released)
                     try:
-                        self._sync_keyspace(connection)
+                        self._sync_keyspace(new_connection)
                     except Exception:
-                        failed_connection = connection
+                        failed_connection = new_connection
                         raise
             except Exception:
                 log.warning("Failed reconnecting %s. Retrying." % (self.endpoint,))
+                self._replacement_connections[shard_id] = connection
                 self._session.submit(self._replace, connection)
             else:
-                self._is_replacing = False
+                self._replacement_connections.pop(shard_id, None)
                 with self._stream_available_condition:
                     self._stream_available_condition.notify()
         if failed_connection is not None:
             failed_connection.close()
 
-    def shutdown(self):
+    def shutdown(self, _failure=None):
         log.debug("Shutting down connections to %s", self.host)
-        with self._lock:
+        shutdown_lock = _failure.locked() if _failure is not None else self._lock
+        with shutdown_lock:
+            if (_failure is not None and not _failure.owns_locked()):
+                return
             with self._membership_lock:
                 if self.is_shutdown:
                     return
+                if _failure is not None:
+                    # Only this reservation may retain failure authority after
+                    # its own guarded retirement clears the shard slots.
+                    _failure.retired = True
                 self.is_shutdown = True
                 connections_to_close = self._connections.copy()
                 self._connections.clear()
             with self._stream_available_condition:
                 self._stream_available_condition.notify_all()
 
-            for future in self._shard_connections_futures:
-                future.cancel()
-
+            futures_to_cancel = tuple(self._shard_connections_futures)
+            self._replacement_connections.clear()
             pending_connections_to_close = self._pending_connections.copy()
             self._pending_connections.clear()
 
+        for future in futures_to_cancel:
+            future.cancel()
+
         # connection.close can call pool.return_connection, which will
         #  obtain self._lock via self._stream_available_condition.
         # So, it never should be called within self._lock context