From 18001e8c38257e09269baa7d0f43a9d768bea671 Mon Sep 17 00:00:00 2001 From: Dmitry Kropachev Date: Fri, 23 Oct 2026 19:36:13 -0400 Subject: [PATCH 1/2] lifecycle: fence DOWN inspection against qualified recovery A DOWN event can observe an empty pool, let recovery finish during a policy or pool-state callback, and then tear down the healthy host. Capture the host health generation before inspection and validate it under the host lock at DOWN admission. Retain pool-failure reservations and admit later genuine failures with a fresh observation. Qualify recovery health through explicit readiness without pool work or initial publication of a usable connection under current pool ownership. Keep membership, Session, Host and the selected Connection protected through publication. Ignored ADD, empty or failed candidates and later UP/ADD aggregate bookkeeping must not manufacture health evidence. Use a short membership guard for publication and member/shutdown changes. Synchronize new members before insertion and pair keyspace targets with member snapshots, keeping connection creation and USE waits outside the guard so background work cannot block initial publication. Cover healthy and failed publication, actual UP/ADD aggregation, ignored ADD execution, current failure reservations, later genuine failure and publication progress during contention and unanswered corrective USE. Keep later asynchronous shard health and socket ownership fencing under #858 and #859. Refs: #925 Refs: #921 --- cassandra/_lifecycle.py | 24 +- cassandra/cluster.py | 74 ++-- cassandra/pool.py | 86 ++-- tests/unit/test_cluster.py | 18 +- tests/unit/test_lifecycle_coordination.py | 516 ++++++++++++++++++++-- 5 files changed, 606 insertions(+), 112 deletions(-) diff --git a/cassandra/_lifecycle.py b/cassandra/_lifecycle.py index 316d59ad27..ccb7910423 100644 --- a/cassandra/_lifecycle.py +++ b/cassandra/_lifecycle.py @@ -186,15 +186,25 @@ def finish_up(host, generation=None): getattr(host, '_readiness_generation', 0)): host._currently_handling_node_up = False - def begin_down(self, host, connected=False, expect_down=False, pool_failure=None): + def observe_health(self, host): + """Capture ownership before inspecting policies or session pools.""" + with host.lock: + if self.closed or host._is_removed: + return None + return host._health_generation + + 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()) 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) @@ -208,13 +218,15 @@ def begin_down(self, host, connected=False, expect_down=False, pool_failure=None host._down_event_generation += 1 return host._down_event_generation, not was_up and not expect_down - def dispatch_down(self, host, is_host_addition, connected, expect_host_to_be_down, - on_add_reconnection, pool_failure=None): + def dispatch_down(self, host, health_generation, is_host_addition, + connected, expect_host_to_be_down, on_add_reconnection, + pool_failure=None): # Publish DOWN state and FIFO admission together. Do not drain on # this caller; the executor (or another drainer) owns delivery. with self.transaction(drain=False): admission = self.begin_down( - host, connected=connected, expect_down=expect_host_to_be_down, + host, health_generation, connected=connected, + expect_down=expect_host_to_be_down, pool_failure=pool_failure) if admission is None: return False @@ -434,7 +446,7 @@ def claim_addition_notification(self, host, generation): return True def mark_ready(self, host, set_up=True, readiness_generation=None, - addition_token=None): + addition_token=None, advance_health=True): if (not self.readiness_is_current(host, readiness_generation) or (addition_token is not None and not self.addition_is_current(host, addition_token))): @@ -448,6 +460,8 @@ def mark_ready(self, host, set_up=True, readiness_generation=None, return False if set_up: host.is_up = True + if advance_health: + host._health_generation += 1 host._addition_token = None host._endpoint_add_retry_pending = False host._lifecycle_pending = False diff --git a/cassandra/cluster.py b/cassandra/cluster.py index 85c0d8e73a..1962f95a65 100644 --- a/cassandra/cluster.py +++ b/cassandra/cluster.py @@ -2384,7 +2384,9 @@ def _on_up_future_completed(self, host, futures, results, lock, finished_future, return log.info("Connection pools established for node %s", host) - if not self._lifecycle.mark_ready(host, readiness_generation=generation): + if not self._lifecycle.mark_ready( + host, readiness_generation=generation, + advance_health=False): return self._notify_pending_add(host, generation) for listener in self.listeners: @@ -2642,8 +2644,13 @@ def on_down(self, host, is_host_addition, expect_host_to_be_down=False, if self.is_shutdown or self.allow_control_connection_query_fallback == ControlConnectionQueryFallback.SkipPoolCreation: return False + health_generation = self._lifecycle.observe_health(host) + if health_generation is None: + return False + # Custom distance policies may refresh topology. Do not invoke them - # under the Host state lock needed by terminal publication. + # under the Host state lock needed by terminal publication. Recovery + # or pool publication during inspection invalidates this observation. connected = False if (self._discount_down_events and self.profile_manager.distance(host) != HostDistance.IGNORED): @@ -2653,7 +2660,8 @@ def on_down(self, host, is_host_addition, expect_host_to_be_down=False, connected |= pool_state['open_count'] > 0 return self._lifecycle.dispatch_down( - host, is_host_addition, connected, expect_host_to_be_down, + host, health_generation, is_host_addition, connected, + expect_host_to_be_down, on_add_reconnection, _pool_failure) def on_add(self, host, refresh_nodes=True, @@ -2801,7 +2809,8 @@ def future_completed(future): log.warning("Connection pool could not be created, not marking node %s up", host) else: self._finalize_add(host, generation=generation, - addition_token=addition_token) + addition_token=addition_token, + advance_health=False) return if not reconcile_pools_on_failure: @@ -2909,10 +2918,10 @@ def future_completed(future): @_lifecycle.event def _finalize_add(self, host, set_up=True, generation=None, - addition_token=None): + addition_token=None, advance_health=True): if not self._lifecycle.mark_ready( host, set_up, readiness_generation=generation, - addition_token=addition_token): + addition_token=addition_token, advance_health=advance_health): return self._notify_pending_add(host, generation) @@ -4824,28 +4833,37 @@ def dispatch_authentication_failure(reservation): publish_pool = False stale_candidate = False stale_result = False - with self._lock: - with host.lock: - if not self._owns_pool_generation_locked( - host, generation, expected_pool, endpoint): - stale_candidate = True - stale_result = \ - self._stale_pool_creation_result_locked( - host, expected_pool, intent, distance, - ignored_intent_is_success=True) - else: - keyspace = self.keyspace - if new_pool._keyspace == keyspace: - current_intent = \ - self._pool_intents_locked().get(host) - if (current_intent is not intent and - expected_pool is not None and - not expected_pool.is_shutdown): - new_pool.host_distance = \ - expected_pool.host_distance - self._pools[host] = new_pool - self._pool_generations[host] = object() - publish_pool = True + # This short guard excludes membership changes, not + # background connection creation or keyspace waits under + # the pool work lock. + with new_pool._membership_lock: + with self._lock: + with host.lock: + if not self._owns_pool_generation_locked( + host, generation, expected_pool, endpoint): + stale_candidate = True + stale_result = \ + self._stale_pool_creation_result_locked( + host, expected_pool, intent, distance, + ignored_intent_is_success=True) + else: + keyspace = self.keyspace + if new_pool._keyspace == keyspace: + current_intent = \ + self._pool_intents_locked().get(host) + if (current_intent is not intent and + expected_pool is not None and + not expected_pool.is_shutdown): + new_pool.host_distance = \ + expected_pool.host_distance + with ( + new_pool._lock_usable_connection() + ) as connected: + self._pools[host] = new_pool + self._pool_generations[host] = object() + if connected: + host._health_generation += 1 + publish_pool = True if stale_candidate: log.debug( diff --git a/cassandra/pool.py b/cassandra/pool.py index e23423beb2..8e49258114 100644 --- a/cassandra/pool.py +++ b/cassandra/pool.py @@ -16,6 +16,7 @@ Connection pooling and host management. """ from concurrent.futures import Future +from contextlib import contextmanager from functools import total_ordering import logging import time @@ -160,6 +161,7 @@ class Host(object): _currently_handling_node_up = False _currently_handling_node_down = False _down_event_generation = 0 + _health_generation = 0 _lifecycle_pending = False _pending_host_addition = False _pending_host_addition_notification = False @@ -467,6 +469,7 @@ def __init__(self, host, host_distance, session): self.host_distance = host_distance self._session = weakref.proxy(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 @@ -499,7 +502,8 @@ def __init__(self, host, host_distance, session): log.debug("Initializing connection for host %s", self.host) first_connection = session.cluster.connection_factory(self.endpoint, on_orphaned_stream_released=self.on_orphaned_stream_released) log.debug("First connection created to %s for shard_id=%s", self.host, first_connection.features.shard_id) - self._connections[first_connection.features.shard_id] = first_connection + with self._membership_lock: + self._connections[first_connection.features.shard_id] = first_connection self._keyspace = session.keyspace if self._keyspace: @@ -679,7 +683,9 @@ def return_connection(self, connection, stream_was_orphaned=False, with self._lock: if self.is_shutdown: return - self._connections.pop(connection.features.shard_id, None) + with self._membership_lock: + self._connections.pop( + connection.features.shard_id, None) if self._is_replacing: return self._is_replacing = True @@ -713,22 +719,17 @@ def _replace(self, connection): log.debug("Replacing connection (%s) to %s", id(connection), self.host) try: - if connection.features.shard_id in self._connections: - del self._connections[connection.features.shard_id] + with self._membership_lock: + self._connections.pop(connection.features.shard_id, None) if self.host.sharding_info and not self._session.cluster.shard_aware_options.disable: 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) - if self._keyspace: - connection.set_keyspace_blocking(self._keyspace) - self._connections[connection.features.shard_id] = connection try: - self._sync_keyspace(connection) + self._sync_keyspace(connection, publish=True) except Exception: - if self._connections.get(connection.features.shard_id) is connection: - del self._connections[connection.features.shard_id] failed_connection = connection raise except Exception: @@ -744,19 +745,19 @@ def _replace(self, connection): def shutdown(self): log.debug("Shutting down connections to %s", self.host) with self._lock: - if self.is_shutdown: - return - else: + with self._membership_lock: + if self.is_shutdown: + return 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() - connections_to_close = self._connections.copy() pending_connections_to_close = self._pending_connections.copy() - self._connections.clear() self._pending_connections.clear() # connection.close can call pool.return_connection, which will @@ -907,17 +908,9 @@ def _open_connection_to_missing_shard(self, shard_id): conn.features.shard_id, self.host ) - if self._keyspace: - conn.set_keyspace_blocking(self._keyspace) - self._connections[conn.features.shard_id] = conn try: - self._sync_keyspace(conn) + self._sync_keyspace(conn, publish=True) except Exception as exc: - if self._connections.get(conn.features.shard_id) is conn: - if old_conn is None: - del self._connections[conn.features.shard_id] - else: - self._connections[conn.features.shard_id] = old_conn self._connecting.discard(shard_id) sync_error = exc @@ -1028,12 +1021,21 @@ def _open_connections_for_all_shards(self, skip_shard_id=None): for conn in self._trash: conn.close() - def _sync_keyspace(self, conn): - """Bring a newly published connection up to the pool's current keyspace.""" - while self._keyspace and conn.keyspace != self._keyspace: - target = self._keyspace - conn.set_keyspace_blocking(target) - if self._keyspace == target: + def _sync_keyspace(self, conn, publish=False): + """ + Synchronize before publishing membership, without a network-wait + guard. + """ + while True: + with self._membership_lock: + target = self._keyspace + if target and conn.keyspace != target: + conn.set_keyspace_blocking(target) + with self._membership_lock: + if self._keyspace != target: + continue + if publish: + self._connections[conn.features.shard_id] = conn break with self._stream_available_condition: self._stream_available_condition.notify_all() @@ -1044,11 +1046,12 @@ def _set_keyspace_for_all_conns(self, keyspace, callback): connections have been set, `callback` will be called with two arguments: this pool, and a list of any errors that occurred. """ - # A connection published after this snapshot rechecks the new value in - # _sync_keyspace. Do not take self._lock here: the publishing thread may - # hold it while waiting for a keyspace reply on this event-loop thread. - self._keyspace = keyspace - remaining_callbacks = set(self._connections.values()) + # Pair the target and member snapshot with final synchronized + # insertion. Never take the broad work lock here: it can await + # this reply thread. + with self._membership_lock: + self._keyspace = keyspace + remaining_callbacks = set(self._connections.values()) remaining_callbacks_lock = Lock() errors = [] @@ -1101,6 +1104,19 @@ def get_connections(self): connections = self._connections return list(connections.values()) if connections else [] + @contextmanager + def _lock_usable_connection(self): + # Publication owns _membership_lock, then Session and Host locks. + # Host must precede Connection, matching ResponseFuture's send path. + if not self.is_shutdown: + for connection in self._connections.values(): + with connection.lock: + if not (connection.is_closed or connection.is_defunct): + # Keep this member usable through health publication. + yield True + return + yield False + def get_state(self): in_flights = [c.in_flight for c in list(self._connections.values())] orphan_requests = [c.orphaned_request_ids for c in list(self._connections.values())] diff --git a/tests/unit/test_cluster.py b/tests/unit/test_cluster.py index 62ee34cf0b..5178b14509 100644 --- a/tests/unit/test_cluster.py +++ b/tests/unit/test_cluster.py @@ -14,6 +14,7 @@ import unittest from concurrent.futures import Future, ThreadPoolExecutor +from contextlib import nullcontext import gc import logging import multiprocessing @@ -1024,7 +1025,8 @@ def create_pool(pool_host, distance, session): pool = Mock( host=pool_host, host_distance=distance, - is_shutdown=False, _keyspace=None) + is_shutdown=False, _keyspace=None, _membership_lock=RLock()) + pool._lock_usable_connection.return_value = nullcontext(True) pool.get_state.return_value = {"open_count": 1} created_pools.append(pool) return pool @@ -1255,7 +1257,9 @@ def submit(fn, *args, **kwargs): new_pool = Mock( host=host, host_distance=HostDistance.LOCAL, - is_shutdown=False, _keyspace="old_keyspace") + is_shutdown=False, _keyspace="old_keyspace", + _membership_lock=RLock()) + new_pool._lock_usable_connection.return_value = nullcontext(True) keyspace_error = ConnectionException("keyspace update failed") def fail_keyspace_update(_keyspace, callback): @@ -1363,7 +1367,8 @@ def submit(fn, *args, **kwargs): new_pool = Mock( host=host, host_distance=HostDistance.LOCAL, - is_shutdown=False, _keyspace=None) + is_shutdown=False, _keyspace=None, _membership_lock=RLock()) + new_pool._lock_usable_connection.return_value = nullcontext(True) def create_pool(*args, **kwargs): cluster.remove_host(host, trigger_reconciliation=False) @@ -2874,7 +2879,8 @@ def _host(self, address="127.0.0.1"): def _pool(self, host, keyspace=None, distance=HostDistance.LOCAL): pool = Mock( host=host, host_distance=distance, is_shutdown=False, - _keyspace=keyspace) + _keyspace=keyspace, _membership_lock=RLock()) + pool._lock_usable_connection.return_value = nullcontext(True) pool.shutdown.side_effect = lambda: setattr( pool, 'is_shutdown', True) return pool @@ -5216,7 +5222,9 @@ def test_pool_renewal_uses_pool_host_not_retained_dict_key(self): "127.0.0.2", SimpleConvictionPolicy, host_id=host_id) old_pool = Mock(host=original_host) pool_state = {"open_count": 1} - new_pool = Mock(host=current_host, _keyspace=None) + new_pool = Mock( + host=current_host, _keyspace=None, _membership_lock=RLock()) + new_pool._lock_usable_connection.return_value = nullcontext(True) new_pool.get_state.return_value = pool_state session = new_session_with_pool_state({original_host: old_pool}) session.keyspace = None diff --git a/tests/unit/test_lifecycle_coordination.py b/tests/unit/test_lifecycle_coordination.py index 5b78c34d0d..0fee02480a 100644 --- a/tests/unit/test_lifecycle_coordination.py +++ b/tests/unit/test_lifecycle_coordination.py @@ -14,16 +14,20 @@ from concurrent.futures import Future from threading import Event, Lock, Thread -from unittest.mock import Mock +from unittest.mock import Mock, patch import pytest from cassandra.cluster import Cluster, ControlConnection -from cassandra._lifecycle import LifecycleCoordinator -from cassandra.connection import DefaultEndPoint +from cassandra._lifecycle import LifecycleCoordinator, PoolFailureReservation +from cassandra.connection import ( + Connection, ConnectionException, DefaultEndPoint) from cassandra.policies import HostDistance +from cassandra.pool import HostConnection +from cassandra.protocol_features import ProtocolFeatures from tests.unit.test_control_connection import ( HOST_ID_1, HOST_ID_2, HOST_ID_3, MockConnection, _node_meta_results) +from tests.unit.utils import new_session_with_pool_state @pytest.fixture @@ -725,63 +729,497 @@ def test_duplicate_up_does_not_schedule_another_pool_creation(topology): assert not host._currently_handling_node_up -class _ExpectedDownAdmissionFailure(Exception): - pass +@pytest.mark.parametrize('initial_state', [None, False, True]) +def test_ignored_add_during_down_inspection_does_not_establish_health( + topology, initial_state): + cluster, _ = topology + host = cluster.metadata.get_host_by_host_id(HOST_ID_3) + host.is_up = initial_state + cluster.control_connection.on_add = Mock() + cluster.on_down_potentially_blocking = Mock(return_value=Future()) + cluster._restart_reconnector = Mock(return_value=Future()) + listener = Mock() + cluster.register_listener(listener) + + def distance(current): + assert current is host + cluster.profile_manager.distance.side_effect = None + cluster.profile_manager.distance.return_value = HostDistance.IGNORED + cluster.on_add(host, refresh_nodes=False) + listener.on_add.assert_called_once_with(host) + assert host.is_up is initial_state + return HostDistance.IGNORED + + cluster.profile_manager.distance.side_effect = distance + + dispatched = cluster.on_down(host, is_host_addition=False) + + cluster.profile_manager.distance.assert_called() + listener.on_add.assert_called_once_with(host) + assert host.is_up is False + assert dispatched is (initial_state is True) + if initial_state: + cluster.on_down_potentially_blocking.assert_called_once() + cluster._restart_reconnector.assert_not_called() + else: + cluster.on_down_potentially_blocking.assert_not_called() + cluster._restart_reconnector.assert_called_once() @pytest.mark.parametrize('inspection', ['distance', 'pool_state']) -@pytest.mark.parametrize('postcondition', ['host_up', 'not_dispatched', 'no_down_handler']) -@pytest.mark.xfail( - strict=True, raises=_ExpectedDownAdmissionFailure, - reason='#925: DOWN admission does not fence health inspection against recovery') -def test_down_health_inspection_cannot_undo_completed_recovery( - topology, inspection, postcondition): +@pytest.mark.parametrize('event', ['up', 'add']) +@pytest.mark.parametrize('healthy', [False, True]) +def test_down_health_inspection_respects_aggregate_pool_health( + topology, inspection, event, healthy): cluster, _ = topology host = cluster.metadata.get_host_by_host_id(HOST_ID_3) host.set_down() + session = new_session_with_pool_state() + session.cluster = cluster + session._profile_manager = cluster.profile_manager + session.keyspace = None + session.is_shutdown = False + session.shutdown = Mock() + session.update_created_pools = Mock() + cluster.sessions = (session,) cluster.control_connection.on_up = Mock() + cluster.control_connection.on_add = Mock() + cluster._prepare_host = lambda host, success, failure: success() cluster.on_down_potentially_blocking = Mock(return_value=Future()) - session = Mock() - ready = Future() - session.add_or_renew_pool.return_value = ready - session.get_pool_state.return_value = {host: {'open_count': 0}} - cluster.sessions.add(session) - cluster.on_up(host) - assert not host.is_up - recovered = [] + connection = _PublicationConnection(host.endpoint, compression=False) + constructing = Event() + release = Event() + completed = Event() + listener = Mock() + getattr(listener, 'on_' + event).side_effect = lambda _: completed.set() + cluster.register_listener(listener) - def finish_recovery(): - ready.set_result(True) - recovered.append(host.is_up) + def construct(*args, **kwargs): + constructing.set() + assert release.wait(5) + return connection + + cluster.connection_factory = Mock(side_effect=construct) + owns = session._owns_pool_generation_locked + + def check_ownership(*args): + if cluster.connection_factory.called and not healthy: + connection.defunct(ConnectionException('candidate failed')) + return owns(*args) + + session._owns_pool_generation_locked = check_ownership + if event == 'up': + cluster.on_up(host) + else: + cluster.on_add(host, refresh_nodes=False) + assert constructing.wait(5) + + def finish(): + release.set() + assert completed.wait(5) + assert host.is_up if inspection == 'distance': def distance(_host): - finish_recovery() - # No pool-state inspection follows an IGNORED decision, so this - # also exercises admission without a pool-failure reservation. + finish() return HostDistance.IGNORED cluster.profile_manager.distance.side_effect = distance else: def pool_state(): - # Finish recovery after capturing an empty pool observation but - # before the older DOWN operation admits its destructive work. observed = {host: {'open_count': 0}} - finish_recovery() + finish() return observed - session.get_pool_state.side_effect = pool_state + session.get_pool_state = Mock(side_effect=pool_state) + + try: + assert cluster.on_down(host, False) is (not healthy) + assert host.is_up is healthy + finally: + release.set() + if healthy: + cluster.on_down_potentially_blocking.assert_not_called() + cluster.profile_manager.distance.side_effect = None + session.get_pool_state = Mock(return_value={host: {'open_count': 0}}) + assert cluster.on_down(host, False) + cluster.on_down_potentially_blocking.assert_called_once() + + +@pytest.mark.parametrize('event', ['up', 'add']) +@pytest.mark.parametrize('ignored', [False, True]) +def test_readiness_without_pool_work_retains_health_authority( + topology, event, ignored): + cluster, _ = topology + host = cluster.metadata.get_host_by_host_id(HOST_ID_3) + host.set_down() + cluster.sessions = () + if ignored: + cluster.profile_manager.distance.return_value = HostDistance.IGNORED + cluster._prepare_host = lambda host, success, failure: success() + cluster.control_connection.on_up = Mock() + cluster.control_connection.on_add = Mock() + before = host._health_generation + if event == 'up': + cluster.on_up(host) + else: + cluster.on_add(host, refresh_nodes=False, set_up_if_ignored=True) + assert host.is_up + assert host._health_generation > before + + +@pytest.mark.parametrize('inspection', ['distance', 'pool_state']) +@pytest.mark.parametrize('pool_failure', [False, True]) +@pytest.mark.parametrize('pool_distance, connect_remote, connection_healthy', [ + pytest.param(HostDistance.REMOTE, False, True, id='disabled_remote'), + pytest.param(HostDistance.REMOTE, True, True, id='healthy_remote'), + pytest.param(HostDistance.LOCAL, False, True, id='healthy_local'), + pytest.param(HostDistance.REMOTE, True, False, id='closed_remote'), + pytest.param(HostDistance.LOCAL, False, False, id='closed_local'), +]) +def test_down_health_inspection_cannot_undo_pool_publication( + topology, inspection, pool_failure, pool_distance, connect_remote, + connection_healthy): + cluster, _ = topology + host = cluster.metadata.get_host_by_host_id(HOST_ID_3) + cluster.on_down_potentially_blocking = Mock(return_value=Future()) + session = new_session_with_pool_state() + session.cluster = cluster + session._profile_manager = cluster.profile_manager + session.keyspace = None + session.is_shutdown = False + session.shutdown = Mock() + cluster.sessions = (session,) + cluster.connect_to_remote_hosts = connect_remote + cluster.profile_manager.distance.return_value = pool_distance + connection = Mock( + features=ProtocolFeatures(shard_id=0), + is_closed=not connection_healthy, + is_defunct=False, lock=Lock()) + cluster.connection_factory = Mock(return_value=connection) + has_connections = pool_distance != HostDistance.REMOTE or connect_remote + establishes_health = has_connections and connection_healthy + + def publish_pool(): + assert session.add_or_renew_pool( + host, is_host_addition=False).result(timeout=5) is True + + if inspection == 'distance': + def distance(_host): + cluster.profile_manager.distance.side_effect = None + publish_pool() + return HostDistance.IGNORED + cluster.profile_manager.distance.side_effect = distance + else: + def pool_state(): + publish_pool() + return {host: {'open_count': 0}} + session.get_pool_state = Mock(side_effect=pool_state) + + # Keep a failing session's reservation current while another session + # publishes a pool. The host observation must fence both admission paths. + reservation = None + if pool_failure: + failed = new_session_with_pool_state() + failed.is_shutdown = False + intent = object() + generation = object() + failed._pool_intents[host] = intent + failed._pool_generations[host] = generation + reservation = PoolFailureReservation( + failed, host, generation, None, intent, host.endpoint) + + dispatched = cluster.on_down( + host, is_host_addition=False, _pool_failure=reservation) + + candidate = session._pools[host] + assert candidate.open_count == int(establishes_health) + assert cluster.connection_factory.call_count == int(has_connections) + assert not candidate.is_shutdown + assert host.is_up is establishes_health + assert dispatched is (not establishes_health) + if establishes_health: + cluster.on_down_potentially_blocking.assert_not_called() + else: + cluster.on_down_potentially_blocking.assert_called_once() + if reservation is not None: + assert reservation.is_current() + + if not establishes_health: + return + + # A later genuine failure must still be admitted with a fresh observation. + cluster.profile_manager.distance.side_effect = None + cluster.profile_manager.distance.return_value = HostDistance.LOCAL + session.get_pool_state = Mock(return_value={host: {'open_count': 0}}) + assert cluster.on_down( + host, is_host_addition=False, _pool_failure=reservation) + assert host.is_up is False + cluster.on_down_potentially_blocking.assert_called_once() - dispatched = cluster.on_down(host, is_host_addition=False) - assert recovered == [True] +class _PublicationConnection(Connection): + """Real connection state transitions without a network transport.""" + + def close(self): + with self.lock: + self.is_closed = True + + +@pytest.mark.parametrize('failure', ['close', 'defunct']) +def test_failed_candidate_cannot_fence_down_before_publication( + topology, failure): + cluster, _ = topology + host = cluster.metadata.get_host_by_host_id(HOST_ID_3) + session = new_session_with_pool_state() + session.cluster = cluster + session._profile_manager = cluster.profile_manager + session.keyspace = None + session.is_shutdown = False + session.shutdown = Mock() + cluster.sessions = (session,) + cluster.on_down_potentially_blocking = Mock(return_value=Future()) + connection = _PublicationConnection(host.endpoint, compression=False) + cluster.connection_factory = Mock(return_value=connection) + publishing = Event() + failed = Event() + errors = [] + owns = session._owns_pool_generation_locked + + def fail(): + try: + assert publishing.wait(5) + if failure == 'close': + connection.close() + else: + connection.defunct(ConnectionException('candidate failed')) + except BaseException as exc: + errors.append(exc) + finally: + failed.set() + + def check_ownership(*args): + if cluster.connection_factory.called and not publishing.is_set(): + # The candidate exists, and publication owns Session/Host state. + # A positive observation made before these locks is now obsolete. + publishing.set() + assert failed.wait(5) + return owns(*args) + + session._owns_pool_generation_locked = check_ownership + + def inspect(): + observed = {host: {'open_count': 0}} + assert session.add_or_renew_pool(host, False).result(timeout=5) is True + return observed + + session.get_pool_state = Mock(side_effect=inspect) + worker = Thread(target=fail) + worker.start() try: - if postcondition == 'host_up': - assert host.is_up, 'an obsolete DOWN observation undid successful recovery' - elif postcondition == 'not_dispatched': - assert not dispatched - elif postcondition == 'no_down_handler': - cluster.on_down_potentially_blocking.assert_not_called() - except AssertionError as error: - raise _ExpectedDownAdmissionFailure from error + assert cluster.on_down(host, False) + assert host.is_up is False + assert session._pools[host].open_count == 0 + cluster.on_down_potentially_blocking.assert_called_once() + finally: + publishing.set() + worker.join(5) + assert not worker.is_alive() + assert not errors + + +def test_publication_guard_releases_connection_and_membership(topology): + cluster, _ = topology + host = cluster.metadata.get_host_by_host_id(HOST_ID_3) + session = new_session_with_pool_state() + session.cluster = cluster + session.keyspace = None + connection = _PublicationConnection(host.endpoint, compression=False) + cluster.connection_factory = Mock(return_value=connection) + pool = HostConnection(host, HostDistance.LOCAL, session) + attempting = Event() + failed = Event() + + def fail(): + attempting.set() + connection.defunct(ConnectionException('failure after publication')) + pool.shutdown() + failed.set() + + worker = Thread(target=fail) + try: + with pool._membership_lock: + with session._lock: + with host.lock: + with pool._lock_usable_connection() as connected: + assert connected + worker.start() + assert attempting.wait(5) + assert not failed.wait(0.05) + assert not connection.is_defunct + assert pool._connections.get( + connection.features.shard_id) is connection + assert failed.wait(5) + finally: + if worker.ident is not None: + worker.join(5) + pool.shutdown() + assert not worker.is_alive() + assert connection.is_defunct + assert pool.is_shutdown + + +def test_publication_contention_keeps_session_available(topology): + cluster, _ = topology + host = cluster.metadata.get_host_by_host_id(HOST_ID_3) + session = new_session_with_pool_state() + session.cluster = cluster + session._profile_manager = cluster.profile_manager + session.keyspace = None + session.is_shutdown = False + session.shutdown = Mock() + connection = _PublicationConnection(host.endpoint, compression=False) + cluster.connection_factory = Mock(return_value=connection) + pool = HostConnection(host, HostDistance.LOCAL, session) + membership = pool._membership_lock + attempting = Event() + replied = Event() + + class ObservedLock: + def __enter__(self): + attempting.set() + membership.acquire() + + def __exit__(self, *args): + membership.release() + + pool._membership_lock = ObservedLock() + reply_thread = None + errors = [] + + def reply(): + try: + session._set_keyspace_for_all_pools( + None, lambda errors: replied.set()) + except BaseException as exc: + errors.append(exc) + + try: + membership.acquire() + try: + with patch('cassandra.cluster.HostConnection', return_value=pool): + publication = session.add_or_renew_pool(host, False) + assert attempting.wait(5) + reply_thread = Thread(target=reply) + reply_thread.start() + assert replied.wait(2) + finally: + membership.release() + if reply_thread is not None: + reply_thread.join(5) + assert not reply_thread.is_alive() + assert not errors + assert publication.result(timeout=5) is True + finally: + pool.shutdown() + + +@pytest.mark.parametrize('healthy', [False, True]) +def test_publication_progress_with_unsynchronized_background_shard( + topology, healthy): + from cassandra.protocol import ResultMessage, RESULT_KIND_SET_KEYSPACE + from cassandra.shard_info import _ShardingInfo + + cluster, _ = topology + host = cluster.metadata.get_host_by_host_id(HOST_ID_3) + session = new_session_with_pool_state() + session.cluster = cluster + session._profile_manager = cluster.profile_manager + session.keyspace = 'ks' + session.is_shutdown = False + session.shutdown = Mock() + sharding = _ShardingInfo(0, 2, '', '', 0, None, None) + first = _PublicationConnection(host.endpoint, compression=False) + first.features = ProtocolFeatures(shard_id=0, sharding_info=sharding) + second = _PublicationConnection(host.endpoint, compression=False) + second.features = ProtocolFeatures(shard_id=1, sharding_info=sharding) + requests = [] + pending = [] + response_lock = Lock() + cleanup_mode = Event() + initial_use = Event() + current_use = Event() + work = [] + submit = session.submit + + def schedule(fn, *args, **kwargs): + future = submit(fn, *args, **kwargs) + if getattr(fn, '__name__', '') == '_open_connection_to_missing_shard': + work.append(future) + return future + + session.submit = schedule + + def respond(message, request_id, cb, **kwargs): + cb(ResultMessage(RESULT_KIND_SET_KEYSPACE)) + return 0 + + def unanswered(message, request_id, cb, **kwargs): + with response_lock: + requests.append(cb) + immediate = cleanup_mode.is_set() + if not immediate: + pending.append(cb) + (initial_use if len(requests) == 1 else current_use).set() + if immediate: + cb(ResultMessage(RESULT_KIND_SET_KEYSPACE)) + return 0 + + def respond_one(): + with response_lock: + callback = pending.pop(0) if pending else None + assert callback is not None or cleanup_mode.is_set() + if callback is not None: + callback(ResultMessage(RESULT_KIND_SET_KEYSPACE)) + + first.send_msg = respond + second.send_msg = unanswered + cluster.connection_factory = Mock(side_effect=[first, second]) + cluster.shard_aware_options.disable_shardaware_port = True + candidates = [] + + def construct(*args, **kwargs): + pool = HostConnection(*args, **kwargs) + candidates.append(pool) + assert initial_use.wait(5) + # Force the background socket to await a corrective USE after its + # original target succeeds. It must not qualify as usable membership. + session.keyspace = 'newks' + pool._set_keyspace_for_all_conns('newks', lambda pool, errors: None) + respond_one() + assert current_use.wait(5) + if not healthy: + first.close() + return pool + + before = host._health_generation + try: + with patch('cassandra.cluster.HostConnection', side_effect=construct): + publication = session.add_or_renew_pool(host, False) + assert current_use.wait(5) + assert publication.result(timeout=2) is True + assert host._health_generation == before + int(healthy) + finally: + with response_lock: + cleanup_mode.set() + callbacks = list(pending) + pending.clear() + for callback in callbacks: + callback(ResultMessage(RESULT_KIND_SET_KEYSPACE)) + for future in work: + future.result(timeout=5) + for pool in candidates: + pool.shutdown() + assert second.keyspace == 'newks' + assert host._health_generation == before + int(healthy) def test_refresh_publishes_nodes_and_tokens_before_queued_listener(topology): From 1973a31ca5f6a092df5d4615da98806d8b5fb4db Mon Sep 17 00:00:00 2001 From: Dmitry Kropachev Date: Sat, 24 Oct 2026 23:54:04 -0400 Subject: [PATCH 2/2] pool: always publish synchronized shard connections Remove the unused private publication flag so direct keyspace-sync tests exercise the same membership update as both production callers. --- cassandra/pool.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/cassandra/pool.py b/cassandra/pool.py index 8e49258114..a9c35cc375 100644 --- a/cassandra/pool.py +++ b/cassandra/pool.py @@ -728,7 +728,7 @@ def _replace(self, connection): connection = self._session.cluster.connection_factory(self.endpoint, on_orphaned_stream_released=self.on_orphaned_stream_released) try: - self._sync_keyspace(connection, publish=True) + self._sync_keyspace(connection) except Exception: failed_connection = connection raise @@ -909,7 +909,7 @@ def _open_connection_to_missing_shard(self, shard_id): self.host ) try: - self._sync_keyspace(conn, publish=True) + self._sync_keyspace(conn) except Exception as exc: self._connecting.discard(shard_id) sync_error = exc @@ -1021,7 +1021,7 @@ def _open_connections_for_all_shards(self, skip_shard_id=None): for conn in self._trash: conn.close() - def _sync_keyspace(self, conn, publish=False): + def _sync_keyspace(self, conn): """ Synchronize before publishing membership, without a network-wait guard. @@ -1034,8 +1034,7 @@ def _sync_keyspace(self, conn, publish=False): with self._membership_lock: if self._keyspace != target: continue - if publish: - self._connections[conn.features.shard_id] = conn + self._connections[conn.features.shard_id] = conn break with self._stream_available_condition: self._stream_available_condition.notify_all()