From df0a475b820455e685575b4a9b0920de1a10055d Mon Sep 17 00:00:00 2001 From: Dmitry Kropachev Date: Mon, 21 Sep 2026 22:37:14 -0400 Subject: [PATCH 1/4] control-connection: preserve fallback session keyspace Bind application use of the shared control connection to one keyspace so fallback queries from different sessions cannot leak USE state into one another. Reject explicit USE and select the bound keyspace before sending application requests. Reclaim the binding only after its owners are gone and its fallback requests have drained. Track that traffic independently from control traffic, unwind a new claim when no request is sent, and keep protocol encoding outside the binding lock. Fixes #1013. --- CHANGELOG.rst | 27 + cassandra/cluster.py | 353 ++++++++++-- .../test_control_connection_query_fallback.py | 90 +++ tests/unit/test_cluster.py | 9 +- tests/unit/test_response_future.py | 516 +++++++++++++++++- 5 files changed, 944 insertions(+), 51 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index bd2f2b28be..5ef0aff5b7 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -29,6 +29,33 @@ Features statements skip re-sending result metadata on EXECUTE, and the driver automatically refreshes cached metadata when the server detects a schema change (DRIVER-153) +Bug Fixes +--------- +* Prevent sessions from leaking keyspace state when application queries fall back to + the shared control connection. The first fallback session binds that connection to + its keyspace (including no keyspace); other fallback sessions with a different + keyspace are rejected while that binding is held. The binding is released once + every session holding it has been shut down or garbage collected and its fallback + requests have drained, so a later session can take it over. Taking it over from a + session without a keyspace is rejected when the previous holder left the shared + connection in a keyspace, since CQL offers no way back to "no keyspace". An explicit + ``USE`` statement -- including the one ``Session.set_keyspace()`` executes -- is now + rejected with ``InvalidRequest`` on the fallback path, since it would change the + keyspace of the shared connection under every other session using it; the keyspace + has to be chosen when the session is created (#1013). +* Fix the client-side timeout never firing for a request on the control-connection + fallback path while the driver's own ``USE`` is in flight. The ``USE`` is sent + without recording a host attempt, and ``_on_speculative_execute`` checked its + "no attempt recorded yet" guard before the deadline, so with a speculative + execution policy configured a repeatedly failing ``USE`` rescheduled the callback + every 10ms instead of ever timing the request out. The deadline is now checked + first (#1013). +* Fix ``ResponseFuture._req_id`` being left pointing at the driver's own ``USE`` + instead of the request actually in flight, when the ``SET_KEYSPACE`` reply is + processed before ``send_msg()`` returns. On a later timeout the driver then + orphaned the wrong stream id and left the real request in the connection's + request map (#1013). + Others ------ * ``DCAwareRoundRobinPolicy.local_dc`` is now read-only. It is set by the constructor, diff --git a/cassandra/cluster.py b/cassandra/cluster.py index 9169936072..181fdf47b8 100644 --- a/cassandra/cluster.py +++ b/cassandra/cluster.py @@ -45,7 +45,7 @@ from cassandra import (ConsistencyLevel, AuthenticationFailed, OperationTimedOut, UnsupportedOperation, SchemaTargetType, DriverException, ProtocolVersion, - UnresolvableContactPoints, DependencyException) + UnresolvableContactPoints, DependencyException, InvalidRequest) from cassandra.auth import _proxy_execute_key, PlainTextAuthProvider from cassandra.client_routes import ClientRoutesChangeType, ClientRoutesConfig, _ClientRoutesHandler from cassandra.connection import (ClientRoutesEndPointFactory, ConnectionException, ConnectionShutdown, @@ -595,6 +595,22 @@ class ControlConnectionQueryFallback(enum.Enum): ``SkipPoolCreation`` disables node-pool creation for the session and uses the control-connection fallback path for application queries. + The first session using fallback binds the shared control connection to its + keyspace, including :const:`None`. Other sessions may use fallback only + while their keyspace matches that binding; a different keyspace is + rejected. The binding is released once every session holding it has been + shut down or garbage collected and its fallback requests have drained, + after which a later session may take it over. A session without a keyspace + cannot take over a binding that left the shared connection in a keyspace, + because CQL offers no way back to "no keyspace"; that case is rejected with + :class:`.InvalidRequest`. + + An explicit ``USE`` statement -- including the one + :meth:`.Session.set_keyspace` executes -- is rejected with + :class:`.InvalidRequest` on the fallback path, because it would change the + keyspace of the shared connection under every other session using it. The + keyspace has to be chosen when the session is created. + The fallback path is not used for requests targeted to an explicit host. """ @@ -926,6 +942,8 @@ def default_retry_policy(self, policy): ``Disabled`` keeps the old behavior. ``Fallback`` enables control-connection fallback when no usable node pools exist. ``SkipPoolCreation`` skips node-pool creation and uses the control connection fallback path. + The first fallback session binds the shared control connection to its keyspace; + later fallback sessions must use the same keyspace. This fallback is still not used for requests targeted to an explicit host. """ @@ -2797,6 +2815,18 @@ def __init__(self, cluster, hosts, keyspace=None): msg += " using keyspace '%s'" % self.keyspace raise NoHostAvailable(msg, [h.address for h in hosts]) + # Only SkipPoolCreation is a fallback Session for its whole life, so + # only it takes the binding up front. Under Fallback the pools may well + # come up moments later, and _query_control_connection() takes the + # binding if and when a query actually needs the shared connection; + # claiming it here would fail connect() over a transient blip and then + # hold the binding for the Session's lifetime. + if fallback_mode is ControlConnectionQueryFallback.SkipPoolCreation: + control_connection = self.cluster.control_connection + conflict = control_connection._attach_application_session(self.keyspace, self) + if conflict is not None: + raise InvalidRequest(conflict) + self.session_id = uuid.uuid4() if self.cluster.column_encryption_policy is not None: @@ -3591,6 +3621,12 @@ def set_keyspace(self, keyspace): """ Set the default keyspace for all queries made through this Session. This operation blocks until complete. + + Raises :class:`.InvalidRequest` when the session uses the + control-connection fallback path (see + :class:`.ControlConnectionQueryFallback`), where the keyspace of the + shared connection cannot be changed; create a Session with the wanted + keyspace instead. """ self.execute('USE %s' % (protect_name(keyspace),)) @@ -3981,6 +4017,20 @@ def __init__(self, cluster, timeout, self._event_schedule_times = {} + # The first fallback Session binds application use of the shared + # control connection to one keyspace (including None). Keeping that + # binding stable avoids connection-level USE state leaking between + # Sessions without adding a dispatcher to this exceptional path. + # The binding is released once every Session holding it is gone and its + # requests have drained, so a later Session can take it over. Response + # callbacks keep this lock while handing one fallback request off to + # the next; it must be re-entrant because some connections deliver a + # response synchronously from send_msg(). + self._application_query_lock = RLock() + self._application_keyspace = _NOT_SET + self._application_sessions = WeakSet() + self._application_requests_in_flight = 0 + def connect(self): if self._is_shutdown: return @@ -4002,6 +4052,73 @@ def _set_new_connection(self, conn): log.debug("[control connection] Closing old connection %r, replacing with %r", old, conn) old.close() + def _attach_application_session(self, keyspace, session): + """Bind application use of the control connection to ``keyspace``. + + The first fallback ``Session`` takes the binding; other Sessions may + share it while they use the same keyspace. Once every Session holding + the binding has been shut down or collected and its fallback requests + have drained, the binding is reclaimed and a later Session can take it + over. + + Returns ``None`` when the binding was taken or shared, otherwise a + message explaining the conflict. + """ + with self._application_query_lock: + self._prune_application_sessions() + + # A Session already holding the binding is not in conflict with + # itself: it may rebind to the keyspace it switched to. + if any(other is not session for other in self._application_sessions): + if self._application_keyspace != keyspace: + return ("Control-connection fallback is already attached to " + "keyspace %r; cannot use it from a Session using " + "keyspace %r" % (self._application_keyspace, keyspace)) + elif keyspace is None: + # Reclaiming from a gone Session cannot undo the USE it left on + # the shared connection: CQL has no way back to "no keyspace". + leftover = self._leftover_application_keyspace() + if leftover is not None: + return ("Control-connection fallback was attached to keyspace " + "%r by a Session that is gone, and the shared " + "connection cannot be reset to no keyspace; create a " + "Session using keyspace %r instead" % (leftover, leftover)) + + self._application_keyspace = keyspace + self._application_sessions.add(session) + return None + + def _discard_application_session(self, session): + """Undo an attachment when its first fallback request was not sent.""" + with self._application_query_lock: + self._application_sessions.discard(session) + if not self._application_sessions and \ + not self._application_requests_in_flight: + self._application_keyspace = _NOT_SET + + def _prune_application_sessions(self): + """Drop shut-down owners after shared-connection requests have drained.""" + if self._application_requests_in_flight: + return + for session in tuple(self._application_sessions): + if session.is_shutdown: + self._application_sessions.discard(session) + + def _leftover_application_keyspace(self): + """The keyspace a released binding left the shared connection in, if any.""" + if self._application_keyspace is _NOT_SET or self._application_keyspace is None: + return None + connection = self._connection + # A reconnect replaces the connection, and a fresh one starts out with + # no keyspace, so the leftover USE state went away with the old one. + if connection is None or connection.keyspace is None: + return None + return connection.keyspace + + def _get_application_keyspace(self): + with self._application_query_lock: + return self._application_keyspace + def _try_connect_to_hosts(self): errors = {} @@ -4196,11 +4313,12 @@ def shutdown(self): return else: self._is_shutdown = True - log.debug("Shutting down control connection") - if self._connection: - self._connection.close() - self._connection = None + connection = self._connection + self._connection = None + + if connection: + connection.close() def refresh_schema(self, force=False, **kwargs): try: @@ -5123,6 +5241,18 @@ def _on_speculative_execute(self): self._timer = None if not self._event.is_set(): + # Check the deadline before the PYTHON-836 guard below. That guard + # only exists to keep speculative queries from running ahead of the + # main thread's first query; it must not swallow an expired client + # timeout. The driver's own USE on the control-connection fallback + # path is sent with record_attempt=False, so attempted_hosts stays + # empty for its whole round trip - and if that USE keeps failing and + # retrying, the 0.01s reschedule below would spin forever instead of + # ever timing the request out. + if self._time_remaining is not None and self._time_remaining <= 0: + self._on_timeout() + return + # PYTHON-836, the speculative queries must be after # the query is sent from the main thread, otherwise the # query from the main thread may raise NoHostAvailable @@ -5134,10 +5264,6 @@ def _on_speculative_execute(self): self._timer = self.session.cluster.connection_class.create_timer(0.01, self._on_speculative_execute) return - if self._time_remaining is not None: - if self._time_remaining <= 0: - self._on_timeout() - return self.send_request(error_no_hosts=False) self._start_timer() @@ -5168,8 +5294,12 @@ def send_request(self, error_no_hosts=True): if error_no_hosts: if self._fallback_to_control_connection(): req_id = self._query_control_connection() + if req_id is _NOT_SET: + return True if req_id is not None: - self._req_id = req_id + # _send_control_connection_message() already recorded the + # id of the message actually in flight. Re-assigning here + # would clobber it with the USE id on the keyspace path. return True self._set_final_exception(NoHostAvailable( @@ -5195,26 +5325,103 @@ def _fallback_to_control_connection(self): return not self._has_usable_node_pool() def _borrow_control_connection(self, connection): - with connection.lock: - if connection.in_flight >= connection.max_request_id: - raise NoConnectionsAvailable("All request IDs are currently in use") - connection.in_flight += 1 - return connection.get_request_id() + control_connection = self.session.cluster.control_connection + with control_connection._application_query_lock: + with connection.lock: + if connection.in_flight >= connection.max_request_id: + raise NoConnectionsAvailable("All request IDs are currently in use") + connection.in_flight += 1 + request_id = connection.get_request_id() + control_connection._application_requests_in_flight += 1 + return request_id - def _release_control_connection_request(self, connection, request_id): - with connection.lock: - connection.in_flight -= 1 - connection.request_ids.append(request_id) - connection._requests.pop(request_id, None) + def _release_control_connection_request(self, connection, request_id, + application_session=None): + control_connection = self.session.cluster.control_connection + with control_connection._application_query_lock: + with connection.lock: + connection.in_flight -= 1 + connection.request_ids.append(request_id) + connection._requests.pop(request_id, None) + control_connection._application_requests_in_flight -= 1 + if application_session is not None: + control_connection._discard_application_session(application_session) def _handle_control_connection_response(self, connection, cb, response): + control_connection = self.session.cluster.control_connection with connection.lock: connection.in_flight -= 1 - cb(response) - - def _query_control_connection(self, message=None, cb=None, connection=None, host=None): - self._control_connection_query_attempted = True + # Keep this request counted while its callback may synchronously hand + # off to another physical fallback request. This closes the zero-in- + # flight window between the driver's USE response and the application + # query it sends next without holding the binding lock during encoding. + try: + cb(response) + finally: + with control_connection._application_query_lock: + control_connection._application_requests_in_flight -= 1 + def _is_keyspace_change_query(self, message=None): + message = self.message if message is None else message + if not isinstance(message, QueryMessage): + return False + query = getattr(message.query, 'query_string', message.query) + if isinstance(query, (bytes, bytearray)): + try: + query = query.decode('utf8') + except UnicodeDecodeError: + return False + return isinstance(query, str) and \ + re.match(r'^(?:\s|(?:--|//)[^\r\n]*(?:\r\n?|\n|$)|/\*(?:[^*]|\*(?!/))*\*/)*USE\b', + query, re.IGNORECASE) is not None + + def _control_connection_failed(self): + self._set_final_exception(NoHostAvailable( + "Unable to complete the operation against any hosts", self._errors)) + + def _set_control_connection_keyspace(self, connection, host, keyspace, + message=None, cb=None, request_id=_NOT_SET, + release_session_on_failure=None): + use_message = QueryMessage( + query='USE %s' % protect_name(keyspace), + consistency_level=ConsistencyLevel.ONE) + + def keyspace_set(response): + if isinstance(response, ResultMessage) and response.kind == RESULT_KIND_SET_KEYSPACE: + connection.keyspace = response.new_keyspace + if self._send_control_connection_message( + message=message, cb=cb, connection=connection, host=host) is None: + self._control_connection_failed() + elif isinstance(response, ErrorMessage): + self._set_final_exception(response.to_exception()) + elif isinstance(response, ConnectionException): + self._errors[host] = response + # Known limitation: this retry goes around the RetryPolicy and + # _query_retries, because both are keyed to the user's statement + # and this is the driver's own USE. It also has no backoff, so a + # flapping control connection re-sends USE until the client-side + # timeout fires instead of failing fast. Routing it through the + # retry machinery shared with the pooled path would be the fix. + self.session.submit(self._retry_task, False, host) + elif isinstance(response, Exception): + self._set_final_exception(response) + else: + self._set_final_exception(ConnectionException( + "Unexpected response while setting the control-connection keyspace: %r" % + (response,), connection.endpoint)) + + # Returns None on failure; send_request() turns that into the final + # NoHostAvailable. Setting it here too would fire every errback twice. + return self._send_control_connection_message( + message=use_message, cb=keyspace_set, connection=connection, + host=host, record_attempt=False, record_size=False, + request_id=request_id, + release_session_on_failure=release_session_on_failure) + + def _send_control_connection_message(self, message=None, cb=None, connection=None, + host=None, record_attempt=True, record_size=True, + request_id=_NOT_SET, + release_session_on_failure=None): if message is None: message = self.message @@ -5229,10 +5436,13 @@ def _query_control_connection(self, message=None, cb=None, connection=None, host host = self.session.cluster.get_control_connection_host() or connection.endpoint self._current_host = host - request_id = None + if request_id is _NOT_SET: + request_id = None request_sent = False + previous_req_id = self._req_id try: - request_id = self._borrow_control_connection(connection) + if request_id is None: + request_id = self._borrow_control_connection(connection) self._connection = connection result_meta = self._bound_result_metadata if cb is None: @@ -5240,12 +5450,22 @@ def _query_control_connection(self, message=None, cb=None, connection=None, host cb = partial(self._handle_control_connection_response, connection, cb) log.debug("No usable node pools; falling back to control connection for host %s", host) - self.request_encoded_size = connection.send_msg(message, request_id, cb=cb, - encoder=self._protocol_handler.encode_message, - decoder=self._protocol_handler.decode_message, - result_metadata=result_meta) + # Record the stream id before sending, not after. The reply can be + # handled re-entrantly - a SET_KEYSPACE reply sends the real message + # from inside this very call - and that nested send has to be the one + # whose id survives in _req_id. Assigning after send_msg() would put + # the already-completed USE id there instead, so a later timeout would + # orphan the wrong stream and leave the real request in _requests. + self._req_id = request_id + encoded_size = connection.send_msg(message, request_id, cb=cb, + encoder=self._protocol_handler.encode_message, + decoder=self._protocol_handler.decode_message, + result_metadata=result_meta) + if record_size: + self.request_encoded_size = encoded_size request_sent = True - self.attempted_hosts.append(host) + if record_attempt: + self.attempted_hosts.append(host) return request_id except NoConnectionsAvailable as exc: log.debug("Control connection is at capacity") @@ -5260,10 +5480,79 @@ def _query_control_connection(self, message=None, cb=None, connection=None, host self._metrics.on_connection_error() finally: if request_id is not None and not request_sent: - self._release_control_connection_request(connection, request_id) + # Only roll back if nothing else claimed _req_id in the meantime, + # so a nested send that did go out keeps its id. + if self._req_id == request_id: + self._req_id = previous_req_id + self._release_control_connection_request( + connection, request_id, release_session_on_failure) return None + def _query_control_connection(self, message=None, cb=None, connection=None, host=None): + self._control_connection_query_attempted = True + control_connection = self.session.cluster.control_connection + if control_connection is None: + self._errors['control connection'] = ConnectionException( + "Control connection is not connected") + return None + + if self._is_keyspace_change_query(message): + self._set_final_exception(InvalidRequest( + "Cannot change keyspace while using control-connection fallback; " + "create a Session with the attached keyspace instead")) + return _NOT_SET + + # Hold this from binding through borrowing a stream. The send itself is + # outside the lock so protocol encoding cannot stall response handling. + with control_connection._application_query_lock: + if connection is None: + connection = control_connection._connection + if connection is None: + self._errors['control connection'] = ConnectionException( + "Control connection is not connected") + return None + + keyspace = self.session.keyspace + new_application_session = \ + self.session not in control_connection._application_sessions + conflict = control_connection._attach_application_session(keyspace, self.session) + if conflict is not None: + self._set_final_exception(InvalidRequest(conflict)) + return _NOT_SET + + if host is None: + host = self.session.cluster.get_control_connection_host() or connection.endpoint + + try: + request_id = self._borrow_control_connection(connection) + except NoConnectionsAvailable as exc: + log.debug("Control connection is at capacity") + self._errors[host] = exc + if new_application_session: + control_connection._discard_application_session(self.session) + return None + except Exception as exc: + log.debug("Error borrowing control connection", exc_info=True) + self._errors[host] = exc + if new_application_session: + control_connection._discard_application_session(self.session) + if self._metrics is not None: + self._metrics.on_connection_error() + return None + + release_session_on_failure = self.session if new_application_session else None + if keyspace is not None and connection.keyspace != keyspace: + return self._set_control_connection_keyspace( + connection, host, keyspace, message=message, cb=cb, + request_id=request_id, + release_session_on_failure=release_session_on_failure) + + return self._send_control_connection_message( + message=message, cb=cb, connection=connection, host=host, + request_id=request_id, + release_session_on_failure=release_session_on_failure) + def _query(self, host, message=None, cb=None): if message is None: message = self.message diff --git a/tests/integration/standard/test_control_connection_query_fallback.py b/tests/integration/standard/test_control_connection_query_fallback.py index b5481f15bc..fed04461d3 100644 --- a/tests/integration/standard/test_control_connection_query_fallback.py +++ b/tests/integration/standard/test_control_connection_query_fallback.py @@ -16,6 +16,7 @@ import pytest +from cassandra import InvalidRequest from cassandra.cluster import ControlConnectionQueryFallback, NoHostAvailable from tests.integration import TestCluster, local, remove_cluster, use_cluster @@ -104,3 +105,92 @@ def test_no_node_pool_fallback_executes_queries_without_creating_pools(self): "SELECT release_version, rpc_address FROM system.local WHERE key='local'").one() assert str(row.rpc_address) == _UNREACHABLE_BROADCAST_RPC_ADDRESS assert row.release_version + + def _bootstrap_keyspaces(self, *keyspaces, tables=()): + bootstrap_cluster = TestCluster( + allow_control_connection_query_fallback=ControlConnectionQueryFallback.SkipPoolCreation, + connect_timeout=1, + ) + try: + setup_session = bootstrap_cluster.connect() + for keyspace in keyspaces: + setup_session.execute("DROP KEYSPACE IF EXISTS {}".format(keyspace)) + setup_session.execute( + "CREATE KEYSPACE {} WITH replication = " + "{{'class': 'NetworkTopologyStrategy', 'replication_factor': 1}}".format(keyspace)) + for table in tables: + setup_session.execute(table) + finally: + bootstrap_cluster.shutdown() + + def test_shared_control_connection_accepts_only_one_session_keyspace(self): + self._bootstrap_keyspaces( + 'fallback_ks_one', 'fallback_ks_two', + tables=("CREATE TABLE fallback_ks_one.items (id int PRIMARY KEY, value text)",)) + + self.cluster = TestCluster( + allow_control_connection_query_fallback=ControlConnectionQueryFallback.SkipPoolCreation, + connect_timeout=1, + ) + session_one = self.cluster.connect('fallback_ks_one') + session_one_peer = self.cluster.connect('fallback_ks_one') + control_connection = self.cluster.control_connection._connection + + assert list(session_one.get_pools()) == [] + assert list(session_one_peer.get_pools()) == [] + with pytest.raises(InvalidRequest, match='already attached'): + self.cluster.connect('fallback_ks_two') + with pytest.raises(InvalidRequest, match='already attached'): + self.cluster.connect() + + insert_one = session_one.execute_async( + "INSERT INTO items (id, value) VALUES (1, 'one')") + insert_two = session_one_peer.execute_async( + "INSERT INTO items (id, value) VALUES (2, 'two')") + insert_one.result() + insert_two.result() + + prepared_one = session_one.prepare( + "INSERT INTO items (id, value) VALUES (?, ?)") + prepared_insert_one = session_one.execute_async(prepared_one, (2, 'prepared-one')) + prepared_insert_one.result() + + select_one = session_one_peer.execute_async( + "SELECT value FROM items WHERE id IN (1, 2)") + + assert {row.value for row in select_one.result()} == {'one', 'prepared-one'} + assert self.cluster.control_connection._connection is control_connection + + def test_shared_control_connection_keyspace_is_reclaimed_after_shutdown(self): + self._bootstrap_keyspaces('fallback_ks_one', 'fallback_ks_two') + + self.cluster = TestCluster( + allow_control_connection_query_fallback=ControlConnectionQueryFallback.SkipPoolCreation, + connect_timeout=1, + ) + session_one = self.cluster.connect('fallback_ks_one') + control_connection = self.cluster.control_connection._connection + + assert list(session_one.get_pools()) == [] + session_one.execute("SELECT key FROM system.local WHERE key='local'") + assert control_connection.keyspace == 'fallback_ks_one' + + # while the binding is held, another keyspace cannot use the fallback + with pytest.raises(InvalidRequest, match='already attached'): + self.cluster.connect('fallback_ks_two') + + session_one.shutdown() + + # the binding is released with its only holder, so it can be taken over + session_two = self.cluster.connect('fallback_ks_two') + session_two.execute("SELECT key FROM system.local WHERE key='local'") + assert control_connection.keyspace == 'fallback_ks_two' + + session_two.shutdown() + + # ...but not by a session without a keyspace: the shared connection is + # still in 'fallback_ks_two' and CQL cannot unset it + with pytest.raises(InvalidRequest, match='cannot be reset to no keyspace'): + self.cluster.connect() + + assert self.cluster.control_connection._connection is control_connection diff --git a/tests/unit/test_cluster.py b/tests/unit/test_cluster.py index 755a3f4888..0aad0cfa69 100644 --- a/tests/unit/test_cluster.py +++ b/tests/unit/test_cluster.py @@ -29,7 +29,7 @@ from cassandra import ConsistencyLevel, DriverException, Timeout, Unavailable, RequestExecutionException, ReadTimeout, WriteTimeout, CoordinationFailure, ReadFailure, WriteFailure, FunctionFailure, AlreadyExists,\ InvalidRequest, Unauthorized, AuthenticationFailed, OperationTimedOut, UnsupportedOperation, RequestValidationException, ConfigurationException, ProtocolVersion from cassandra.cluster import _Scheduler, Session, Cluster, ResultSet, SchemaAgreementScope, ControlConnectionQueryFallback, default_lbp_factory, \ - ExecutionProfile, _ConfigMode, EXEC_PROFILE_DEFAULT + ExecutionProfile, _ConfigMode, EXEC_PROFILE_DEFAULT, _NOT_SET from cassandra.connection import (Connection, ConnectionBusy, ConnectionException, DefaultEndPoint) from cassandra.ssl_session_cache import SSLSessionCache @@ -231,6 +231,11 @@ def test_control_connection_query_fallback_no_node_pool_mode_skips_pool_creation assert session._pools == {} assert session.update_created_pools() == set() + same_keyspace_session = Session(cluster, [host]) + assert same_keyspace_session.keyspace is None + with pytest.raises(InvalidRequest, match='already attached'): + Session(cluster, [host], keyspace='different') + def test_control_connection_query_fallback_fallback_tolerates_empty_initial_pools(self): cluster = Cluster( allow_control_connection_query_fallback=ControlConnectionQueryFallback.Fallback, @@ -246,6 +251,8 @@ def test_control_connection_query_fallback_fallback_tolerates_empty_initial_pool mocked_add_or_renew_pool.assert_called_once_with(host, is_host_addition=False) assert session._initial_connect_futures == {future} assert session._pools == {} + assert set(cluster.control_connection._application_sessions) == set() + assert cluster.control_connection._get_application_keyspace() is _NOT_SET def test_compression_autodisabled_without_libraries(self): with patch.dict('cassandra.cluster.locally_supported_compressions', {}, clear=True): diff --git a/tests/unit/test_response_future.py b/tests/unit/test_response_future.py index d71943ec04..f2572bd29a 100644 --- a/tests/unit/test_response_future.py +++ b/tests/unit/test_response_future.py @@ -12,15 +12,17 @@ # See the License for the specific language governing permissions and # limitations under the License. +import time import unittest from collections import deque -from threading import RLock +from threading import RLock, Thread from unittest.mock import Mock, MagicMock, ANY, patch -from cassandra import ConsistencyLevel, Unavailable, SchemaTargetType, SchemaChangeType, OperationTimedOut -from cassandra.cluster import Session, ResponseFuture, NoHostAvailable, ProtocolVersion, ControlConnectionQueryFallback -from cassandra.connection import Connection, ConnectionException +from cassandra import ConsistencyLevel, InvalidRequest, Unavailable, SchemaTargetType, SchemaChangeType, OperationTimedOut +from cassandra.cluster import (Session, ResponseFuture, NoHostAvailable, ProtocolVersion, + ControlConnection, ControlConnectionQueryFallback, _NOT_SET) +from cassandra.connection import Connection, ConnectionBusy, ConnectionException from cassandra.protocol import (ReadTimeoutErrorMessage, WriteTimeoutErrorMessage, UnavailableErrorMessage, ResultMessage, QueryMessage, ExecuteMessage, @@ -40,8 +42,15 @@ class ResponseFutureTests(unittest.TestCase): def make_basic_session(self): s = Mock(spec=Session) + s.keyspace = None + s.is_shutdown = False s.row_factory = lambda col_names, rows: [(col_names, rows)] s.cluster.allow_control_connection_query_fallback = ControlConnectionQueryFallback.Disabled + s.cluster.control_connection = ControlConnection( + s.cluster, timeout=1, + schema_event_refresh_window=0, + topology_event_refresh_window=0, + status_event_refresh_window=0) return s def make_pool(self): @@ -62,6 +71,7 @@ def make_control_connection(self): connection.orphaned_threshold = 75 connection.orphaned_threshold_reached = False connection.is_control_connection = True + connection.keyspace = None connection.get_request_id.return_value = 7 connection.send_msg.return_value = 128 # These tests exercise control-connection query fallback, not tablet @@ -430,33 +440,402 @@ def test_control_connection_fallback_disabled_by_default(self): with pytest.raises(NoHostAvailable): rf.result() - def test_control_connection_fallback_updates_connection_keyspace(self): + def test_control_connection_fallback_does_not_bind_without_connection(self): session = self.make_basic_session() - session.cluster.allow_control_connection_query_fallback = ControlConnectionQueryFallback.Fallback - session.cluster._default_load_balancing_policy.make_query_plan.return_value = ['ip1'] + session.cluster.allow_control_connection_query_fallback = \ + ControlConnectionQueryFallback.Fallback + session.cluster._default_load_balancing_policy.make_query_plan.return_value = [] session._pools = {} + control_connection = session.cluster.control_connection + control_connection._connection = None - def set_keyspace_for_all_pools(keyspace, callback): - session.keyspace = keyspace - callback({}) + rf = self.make_response_future(session) + assert not rf.send_request() - session._set_keyspace_for_all_pools.side_effect = set_keyspace_for_all_pools + assert not control_connection._application_sessions + assert control_connection._get_application_keyspace() is _NOT_SET + with pytest.raises(NoHostAvailable): + rf.result() + + def test_control_connection_fallback_rejects_use(self): + for query_string in ( + "USE newks", + "-- select another keyspace\nUSE newks", + "-- select another keyspace\rUSE newks", + "// select another keyspace\nUSE newks", + "/* select another keyspace */ USE newks", + b"USE newks", + b"-- select another keyspace\rUSE newks"): + with self.subTest(query_string=query_string): + session = self.make_basic_session() + session.cluster.allow_control_connection_query_fallback = \ + ControlConnectionQueryFallback.Fallback + session.cluster._default_load_balancing_policy.make_query_plan.return_value = ['ip1'] + session._pools = {} + + connection = self.make_control_connection() + session.cluster.control_connection._connection = connection + + query = SimpleStatement(query_string) + rf = ResponseFuture( + session, + QueryMessage(query=query.query_string, consistency_level=ConsistencyLevel.ONE), + query, 1) + assert rf.send_request() + + connection.send_msg.assert_not_called() + assert rf._req_id is None + with pytest.raises(InvalidRequest, match='Cannot change keyspace'): + rf.result() + + # the rejected USE must not have claimed the binding + control_connection = session.cluster.control_connection + assert control_connection._get_application_keyspace() is _NOT_SET + assert not control_connection._application_sessions + + def test_control_connection_fallback_accepts_stream_id_zero(self): + session = self.make_basic_session() + session.cluster.allow_control_connection_query_fallback = \ + ControlConnectionQueryFallback.SkipPoolCreation + session.cluster._default_load_balancing_policy.make_query_plan.return_value = [] + session._pools = {} + session.keyspace = 'ks' connection = self.make_control_connection() - connection.keyspace = 'oldks' + connection.get_request_id.return_value = 0 session.cluster.control_connection._connection = connection + + rf = self.make_response_future(session) + assert rf.send_request() + assert rf._req_id == 0 + assert rf._final_exception is None + assert connection.send_msg.call_args[0][0].query == 'USE ks' + + def test_control_connection_fallback_binds_first_session_keyspace(self): + session1 = self.make_basic_session() + session1.cluster.allow_control_connection_query_fallback = ControlConnectionQueryFallback.SkipPoolCreation + session1.cluster._default_load_balancing_policy.make_query_plan.return_value = [] + session1._pools = {} + session1.keyspace = 'ks1' + + session2 = self.make_basic_session() + session2.cluster = session1.cluster + session2._pools = {} + session2.keyspace = 'ks1' + + connection = self.make_control_connection() + connection.get_request_id.side_effect = [7, 8, 9] + session1.cluster.control_connection._connection = connection control_host = Mock(endpoint=connection.endpoint) - session.cluster.get_control_connection_host.return_value = control_host + session1.cluster.get_control_connection_host.return_value = control_host + + rf1 = self.make_response_future(session1) + rf2 = self.make_response_future(session2) + assert rf1.send_request() + + assert connection.send_msg.call_count == 1 + assert connection.send_msg.call_args_list[0][0][0].query == 'USE ks1' + + connection.send_msg.call_args_list[0][1]['cb']( + Mock(spec=ResultMessage, kind=RESULT_KIND_SET_KEYSPACE, new_keyspace='ks1')) + assert connection.send_msg.call_count == 2 + assert connection.send_msg.call_args_list[1][0][0] is rf1.message + + connection.send_msg.call_args_list[1][1]['cb']( + self.make_mock_response(['value'], [('one',)])) + assert rf2.send_request() + assert connection.send_msg.call_count == 3 + assert connection.send_msg.call_args_list[2][0][0] is rf2.message + connection.send_msg.call_args_list[2][1]['cb']( + self.make_mock_response(['value'], [('two',)])) + + assert rf1.result().one() == (['value'], [('one',)]) + assert rf2.result().one() == (['value'], [('two',)]) + assert connection.keyspace == 'ks1' + assert session1.cluster.control_connection._get_application_keyspace() == 'ks1' + + def test_control_connection_fallback_rejects_different_session_keyspace(self): + for first_keyspace, second_keyspace in ( + ('ks1', 'ks2'), ('ks1', None), (None, 'ks1')): + with self.subTest(first_keyspace=first_keyspace, + second_keyspace=second_keyspace): + session1 = self.make_basic_session() + session1.cluster.allow_control_connection_query_fallback = \ + ControlConnectionQueryFallback.SkipPoolCreation + session1.cluster._default_load_balancing_policy.make_query_plan.return_value = [] + session1._pools = {} + session1.keyspace = first_keyspace + + session2 = self.make_basic_session() + session2.cluster = session1.cluster + session2._pools = {} + session2.keyspace = second_keyspace + + connection = self.make_control_connection() + session1.cluster.control_connection._connection = connection + assert session1.cluster.control_connection._attach_application_session( + first_keyspace, session1) is None + + rf2 = self.make_response_future(session2) + assert rf2.send_request() + + connection.send_msg.assert_not_called() + assert rf2._req_id is None + with pytest.raises(InvalidRequest, match='already attached'): + rf2.result() + + def test_control_connection_fallback_rebinds_same_session_new_keyspace(self): + session1 = self.make_basic_session() + control_connection = session1.cluster.control_connection + connection = self.make_control_connection() + control_connection._connection = connection + + assert control_connection._attach_application_session('ks1', session1) is None + connection.keyspace = 'ks1' + + # the session holding the binding is not in conflict with itself: it + # may rebind to the keyspace it switched to + assert control_connection._attach_application_session('ks2', session1) is None + assert control_connection._get_application_keyspace() == 'ks2' + + session2 = self.make_basic_session() + session2.cluster = session1.cluster + conflict = control_connection._attach_application_session('ks1', session2) + assert conflict is not None and 'already attached' in conflict + + def test_control_connection_fallback_req_id_tracks_real_message(self): + session = self._make_fallback_session(keyspace='ks1') + connection = self.make_control_connection() + connection.get_request_id.side_effect = [7, 8] + session.cluster.control_connection._connection = connection + session.cluster.get_control_connection_host.return_value = Mock(endpoint=connection.endpoint) + + def send_msg(message, request_id, cb=None, **kwargs): + lock = session.cluster.control_connection._application_query_lock + lock_available = [] + + def probe_lock(): + acquired = lock.acquire(False) + lock_available.append(acquired) + if acquired: + lock.release() + + probe = Thread(target=probe_lock) + probe.start() + probe.join() + assert lock_available == [True] + # the SET_KEYSPACE reply lands before send_msg() returns, so the real + # query is sent re-entrantly from inside this very call + if getattr(message, 'query', None) == 'USE ks1': + connection.keyspace = 'ks1' + cb(Mock(spec=ResultMessage, kind=RESULT_KIND_SET_KEYSPACE, + new_keyspace='ks1')) + return 128 + + connection.send_msg.side_effect = send_msg rf = self.make_response_future(session) assert rf.send_request() - result = Mock(spec=ResultMessage, kind=RESULT_KIND_SET_KEYSPACE, new_keyspace='newks') - connection.send_msg.call_args[1]['cb'](result) + # 7 is the USE's stream id, 8 the real query's. _req_id has to be 8: on a + # later timeout it is the id popped from _requests and orphaned, so the + # USE id here would leave the real request in flight and unorphaned. + assert [c[0][1] for c in connection.send_msg.call_args_list] == [7, 8] + assert rf._req_id == 8 + + def test_control_connection_fallback_req_id_restored_when_send_fails(self): + session = self._make_fallback_session(keyspace=None) + connection = self.make_control_connection() + session.cluster.control_connection._connection = connection + session.cluster.get_control_connection_host.return_value = Mock(endpoint=connection.endpoint) + connection.send_msg.side_effect = ConnectionBusy('no streams') + + rf = self.make_response_future(session) + rf._req_id = 42 + # nothing could be sent, so the future ends as NoHostAvailable + assert not rf.send_request() + + # a request that never went out must not leave its id behind for + # _on_timeout() to orphan + assert rf._req_id == 42 + assert not session.cluster.control_connection._application_sessions + assert session.cluster.control_connection._get_application_keyspace() is _NOT_SET + assert session.cluster.control_connection._application_requests_in_flight == 0 + + def test_speculative_execute_honours_expired_deadline_without_attempts(self): + session = self.make_basic_session() + session.cluster._default_load_balancing_policy.make_query_plan.return_value = [] + rf = self.make_response_future(session) + rf.timeout = 1 + rf._start_time = time.time() - 5 + rf.attempted_hosts = [] + rf._on_timeout = Mock() + session.cluster.connection_class.create_timer.reset_mock() + + rf._on_speculative_execute() + + # the PYTHON-836 "no attempt recorded yet" guard must not swallow an + # already-expired client timeout: the driver's own USE is sent with + # record_attempt=False, so a retrying USE would otherwise reschedule + # this callback every 10ms forever and the request would never time out + rf._on_timeout.assert_called_once_with() + session.cluster.connection_class.create_timer.assert_not_called() + + def _make_fallback_session(self, cluster=None, keyspace=None): + session = self.make_basic_session() + if cluster is not None: + session.cluster = cluster + session.cluster.allow_control_connection_query_fallback = \ + ControlConnectionQueryFallback.SkipPoolCreation + session.cluster._default_load_balancing_policy.make_query_plan.return_value = [] + session._pools = {} + session.keyspace = keyspace + return session + + def test_control_connection_fallback_rebinds_after_owner_shutdown(self): + session1 = self._make_fallback_session(keyspace='ks1') + connection = self.make_control_connection() + connection.get_request_id.side_effect = [7, 8, 9] + session1.cluster.control_connection._connection = connection + session1.cluster.get_control_connection_host.return_value = Mock(endpoint=connection.endpoint) + + assert self.make_response_future(session1).send_request() + assert connection.send_msg.call_args_list[0][0][0].query == 'USE ks1' + connection.send_msg.call_args_list[0][1]['cb']( + Mock(spec=ResultMessage, kind=RESULT_KIND_SET_KEYSPACE, new_keyspace='ks1')) + connection.send_msg.call_args_list[1][1]['cb']( + self.make_mock_response(['value'], [('one',)])) - assert connection.keyspace == 'newks' - assert session.keyspace == 'newks' - assert rf.result().current_rows == [] + # the binding is released once its only owner is gone and its request + # has drained + session1.is_shutdown = True + + session2 = self._make_fallback_session(cluster=session1.cluster, keyspace='ks2') + rf2 = self.make_response_future(session2) + assert rf2.send_request() + + assert rf2._final_exception is None + assert connection.send_msg.call_args_list[-1][0][0].query == 'USE ks2' + assert session1.cluster.control_connection._get_application_keyspace() == 'ks2' + + def test_control_connection_fallback_reclaim_ignores_other_control_requests(self): + session1 = self._make_fallback_session(keyspace='ks1') + control_connection = session1.cluster.control_connection + connection = self.make_control_connection() + control_connection._connection = connection + + assert control_connection._attach_application_session('ks1', session1) is None + session1.is_shutdown = True + # Heartbeats and metadata refreshes also contribute to this connection- + # wide counter; they must not keep a fallback binding alive. + connection.in_flight = 1 + + session2 = self._make_fallback_session(cluster=session1.cluster, keyspace='ks2') + assert control_connection._attach_application_session('ks2', session2) is None + assert control_connection._get_application_keyspace() == 'ks2' + + def test_control_connection_fallback_keeps_shutdown_owner_until_requests_drain(self): + session1 = self._make_fallback_session(keyspace='ks1') + connection = self.make_control_connection() + connection.get_request_id.side_effect = [7, 8, 9] + session1.cluster.control_connection._connection = connection + session1.cluster.get_control_connection_host.return_value = Mock(endpoint=connection.endpoint) + + rf1 = self.make_response_future(session1) + assert rf1.send_request() + assert connection.send_msg.call_args_list[0][0][0].query == 'USE ks1' + session1.is_shutdown = True + + # The driver's USE is still in flight, so shutdown cannot release the + # binding to a Session using another keyspace. + session2 = self._make_fallback_session(cluster=session1.cluster, keyspace='ks2') + rf2 = self.make_response_future(session2) + assert rf2.send_request() + with pytest.raises(InvalidRequest, match='already attached'): + rf2.result() + + connection.send_msg.call_args_list[0][1]['cb']( + Mock(spec=ResultMessage, kind=RESULT_KIND_SET_KEYSPACE, new_keyspace='ks1')) + assert connection.send_msg.call_args_list[1][0][0] is rf1.message + + # USE drained, but application query remains in flight and keeps the + # shutdown owner bound. + session3 = self._make_fallback_session(cluster=session1.cluster, keyspace='ks2') + rf3 = self.make_response_future(session3) + assert rf3.send_request() + with pytest.raises(InvalidRequest, match='already attached'): + rf3.result() + + connection.send_msg.call_args_list[1][1]['cb']( + self.make_mock_response(['value'], [('one',)])) + + # Once both physical requests drain, another keyspace can take over. + session4 = self._make_fallback_session(cluster=session1.cluster, keyspace='ks2') + rf4 = self.make_response_future(session4) + assert rf4.send_request() + assert rf4._final_exception is None + assert connection.send_msg.call_args_list[2][0][0].query == 'USE ks2' + assert session1.cluster.control_connection._get_application_keyspace() == 'ks2' + + def test_control_connection_fallback_reclaim_without_keyspace_rejected(self): + session1 = self._make_fallback_session(keyspace='ks1') + connection = self.make_control_connection() + session1.cluster.control_connection._connection = connection + assert session1.cluster.control_connection._attach_application_session( + 'ks1', session1) is None + connection.keyspace = 'ks1' + session1.is_shutdown = True + + # the shared connection is still in 'ks1' and CQL cannot unset it + session2 = self._make_fallback_session(cluster=session1.cluster, keyspace=None) + rf2 = self.make_response_future(session2) + assert rf2.send_request() + + connection.send_msg.assert_not_called() + with pytest.raises(InvalidRequest, match='cannot be reset to no keyspace'): + rf2.result() + + def test_control_connection_fallback_reports_connection_leftover_keyspace(self): + session1 = self._make_fallback_session(keyspace='ks2') + control_connection = session1.cluster.control_connection + connection = self.make_control_connection() + control_connection._connection = connection + assert control_connection._attach_application_session('ks2', session1) is None + # The desired binding can move ahead of the USE that changes the + # physical connection. Report the connection's actual state. + connection.keyspace = 'ks1' + session1.is_shutdown = True + + session2 = self._make_fallback_session(cluster=session1.cluster, keyspace=None) + rf2 = self.make_response_future(session2) + assert rf2.send_request() + + with pytest.raises(InvalidRequest, match="keyspace 'ks1'"): + rf2.result() + + def test_control_connection_fallback_reclaim_without_keyspace_after_reconnect(self): + session1 = self._make_fallback_session(keyspace='ks1') + connection = self.make_control_connection() + session1.cluster.control_connection._connection = connection + assert session1.cluster.control_connection._attach_application_session( + 'ks1', session1) is None + connection.keyspace = 'ks1' + session1.is_shutdown = True + + # a reconnect leaves a fresh connection with no keyspace, so the + # leftover USE state is gone and the binding is freely reclaimable + reconnected = self.make_control_connection() + session1.cluster.control_connection._connection = reconnected + session1.cluster.get_control_connection_host.return_value = \ + Mock(endpoint=reconnected.endpoint) + + session2 = self._make_fallback_session(cluster=session1.cluster, keyspace=None) + rf2 = self.make_response_future(session2) + assert rf2.send_request() + + assert rf2._final_exception is None + assert reconnected.send_msg.call_args_list[0][0][0] is rf2.message + assert session1.cluster.control_connection._get_application_keyspace() is None def test_control_connection_fallback_when_no_usable_pools(self): session = self.make_basic_session() @@ -612,6 +991,107 @@ def test_control_connection_fallback_reprepares_prepared_statement(self): assert connection.in_flight == 0 assert rf.result()[0] == expected_result + def test_control_connection_fallback_reprepare_send_failure_retries(self): + session = self.make_basic_session() + session.cluster.allow_control_connection_query_fallback = \ + ControlConnectionQueryFallback.Fallback + session.cluster._default_load_balancing_policy.make_query_plan.return_value = [] + session._pools = {} + connection = self.make_control_connection() + connection.send_msg.side_effect = ConnectionBusy() + session.cluster.control_connection._connection = connection + host = Mock(endpoint=connection.endpoint) + + rf = self.make_response_future(session) + rf.send_request = Mock() + rf._reprepare(PrepareMessage("SELECT * FROM foo"), host, connection, None) + + rf.send_request.assert_called_once_with() + + def test_control_connection_fallback_rejects_indented_use(self): + session = self.make_basic_session() + session.cluster.allow_control_connection_query_fallback = \ + ControlConnectionQueryFallback.Fallback + session.cluster._default_load_balancing_policy.make_query_plan.return_value = [] + session._pools = {} + + indent = '\n' + ' ' * 26 + rf = self.make_response_future(session) + + # A deeply indented statement must be classified without the leading + # whitespace sending the matcher into exponential backtracking. + indented_use = QueryMessage(query=indent + 'USE newks', + consistency_level=ConsistencyLevel.ONE) + indented_select = QueryMessage(query=indent + 'SELECT * FROM foo', + consistency_level=ConsistencyLevel.ONE) + + start = time.time() + assert rf._is_keyspace_change_query(indented_use) + assert not rf._is_keyspace_change_query(indented_select) + assert time.time() - start < 2 + + def test_control_connection_fallback_use_failure_reports_error_once(self): + session = self.make_basic_session() + session.cluster.allow_control_connection_query_fallback = \ + ControlConnectionQueryFallback.Fallback + session.cluster._default_load_balancing_policy.make_query_plan.return_value = [] + session._pools = {} + session.keyspace = 'ks' + + connection = self.make_control_connection() + connection.keyspace = None + connection.send_msg.side_effect = ConnectionBusy() + session.cluster.control_connection._connection = connection + + rf = self.make_response_future(session) + errback = Mock() + rf.add_errback(errback) + rf.send_request() + + # the USE message could not be sent; the failure is reported exactly once + assert errback.call_count == 1 + assert connection.in_flight == 0 + assert not session.cluster.control_connection._application_sessions + assert session.cluster.control_connection._get_application_keyspace() is _NOT_SET + assert session.cluster.control_connection._application_requests_in_flight == 0 + with pytest.raises(NoHostAvailable): + rf.result() + + def test_control_connection_fallback_reprepare_sets_keyspace_first(self): + session = self.make_basic_session() + session.cluster.allow_control_connection_query_fallback = \ + ControlConnectionQueryFallback.Fallback + session.cluster._default_load_balancing_policy.make_query_plan.return_value = [] + session._pools = {} + session.keyspace = 'ks' + + connection = self.make_control_connection() + connection.keyspace = None + connection.get_request_id.side_effect = [7, 8] + session.cluster.control_connection._connection = connection + host = Mock(endpoint=connection.endpoint) + + rf = self.make_response_future(session) + prepare_message = PrepareMessage("SELECT * FROM foo") + rf._reprepare(prepare_message, host, connection, None) + + # the control connection is switched to the session keyspace first + assert connection.send_msg.call_count == 1 + assert connection.send_msg.call_args_list[0][0][0].query == 'USE ks' + + connection.send_msg.call_args_list[0][1]['cb']( + Mock(spec=ResultMessage, kind=RESULT_KIND_SET_KEYSPACE, new_keyspace='ks')) + + # ...and only then is the PrepareMessage sent, with the reprepare callback + assert connection.send_msg.call_count == 2 + assert connection.send_msg.call_args_list[1][0][0] is prepare_message + assert connection.send_msg.call_args_list[1][0][1] == 8 + + prepared_response = Mock(spec=ResultMessage, kind=RESULT_KIND_PREPARED) + connection.send_msg.call_args_list[1][1]['cb'](prepared_response) + session.submit.assert_called_once_with( + rf._execute_after_prepare, host, connection, None, prepared_response) + def test_control_connection_fallback_not_used_when_pool_can_serve(self): session = self.make_basic_session() session.cluster.allow_control_connection_query_fallback = ControlConnectionQueryFallback.Fallback From 515d72ec5607dedc7a3ac5c6c2cbbcd26b2f69fd Mon Sep 17 00:00:00 2001 From: Dmitry Kropachev Date: Mon, 21 Sep 2026 23:18:21 -0400 Subject: [PATCH 2/4] control-connection: harden fallback binding lifecycle Release logical fallback accounting when a request times out while retaining an orphan barrier until its late response can no longer change the physical keyspace. Prevent active self-rebinding and validate reclaimed bindings against the connection's actual keyspace. Limit CQL USE detection to CQL statements so valid graph queries are not rejected. --- cassandra/cluster.py | 97 +++++++++++++++++++--- tests/unit/test_response_future.py | 127 +++++++++++++++++++++++++++++ 2 files changed, 211 insertions(+), 13 deletions(-) diff --git a/cassandra/cluster.py b/cassandra/cluster.py index 181fdf47b8..02b4cd4d58 100644 --- a/cassandra/cluster.py +++ b/cassandra/cluster.py @@ -4030,6 +4030,7 @@ def __init__(self, cluster, timeout, self._application_keyspace = _NOT_SET self._application_sessions = WeakSet() self._application_requests_in_flight = 0 + self._application_orphaned_requests = set() def connect(self): if self._is_shutdown: @@ -4067,14 +4068,26 @@ def _attach_application_session(self, keyspace, session): with self._application_query_lock: self._prune_application_sessions() - # A Session already holding the binding is not in conflict with - # itself: it may rebind to the keyspace it switched to. - if any(other is not session for other in self._application_sessions): - if self._application_keyspace != keyspace: - return ("Control-connection fallback is already attached to " - "keyspace %r; cannot use it from a Session using " - "keyspace %r" % (self._application_keyspace, keyspace)) - elif keyspace is None: + other_sessions = any( + other is not session for other in self._application_sessions) + current_connection_has_orphans = \ + self._current_connection_has_application_orphans() + binding_is_busy = self._application_requests_in_flight or \ + current_connection_has_orphans + + # Even the Session that owns the binding cannot change it while a + # request is active. A recovered node pool may have changed the + # Session keyspace while an earlier control-connection USE or query + # is still outstanding, and allowing the next fallback request to + # rebind would put both keyspaces on the shared connection at once. + if self._application_keyspace is not _NOT_SET and \ + self._application_keyspace != keyspace and \ + (other_sessions or binding_is_busy): + return ("Control-connection fallback is already attached to " + "keyspace %r; cannot use it from a Session using " + "keyspace %r" % (self._application_keyspace, keyspace)) + + if keyspace is None: # Reclaiming from a gone Session cannot undo the USE it left on # the shared connection: CQL has no way back to "no keyspace". leftover = self._leftover_application_keyspace() @@ -4093,7 +4106,8 @@ def _discard_application_session(self, session): with self._application_query_lock: self._application_sessions.discard(session) if not self._application_sessions and \ - not self._application_requests_in_flight: + not self._application_requests_in_flight and \ + not self._current_connection_has_application_orphans(): self._application_keyspace = _NOT_SET def _prune_application_sessions(self): @@ -4106,8 +4120,6 @@ def _prune_application_sessions(self): def _leftover_application_keyspace(self): """The keyspace a released binding left the shared connection in, if any.""" - if self._application_keyspace is _NOT_SET or self._application_keyspace is None: - return None connection = self._connection # A reconnect replaces the connection, and a fresh one starts out with # no keyspace, so the leftover USE state went away with the old one. @@ -4115,10 +4127,27 @@ def _leftover_application_keyspace(self): return None return connection.keyspace + def _current_connection_has_application_orphans(self): + return any( + connection is self._connection + for connection, _ in self._application_orphaned_requests) + def _get_application_keyspace(self): with self._application_query_lock: return self._application_keyspace + def _handle_orphaned_application_response(self, connection, request_id, + response): + """Retire a timed-out fallback stream without resuming its request.""" + try: + if isinstance(response, ResultMessage) and \ + response.kind == RESULT_KIND_SET_KEYSPACE: + connection.keyspace = response.new_keyspace + finally: + with self._application_query_lock: + self._application_orphaned_requests.discard( + (connection, request_id)) + def _try_connect_to_hosts(self): errors = {} @@ -5181,8 +5210,14 @@ def _on_timeout(self, _attempts=0): conn_in_flight = None if self._connection is not None: + orphaned_control_request = \ + self._connection.is_control_connection and \ + self._control_connection_query_attempted and \ + self._orphan_control_connection_request( + self._connection, self._req_id) try: - self._connection._requests.pop(self._req_id) + if not orphaned_control_request: + self._connection._requests.pop(self._req_id) # PYTHON-1044 # This request might have been removed from the connection after the latter was defunct by heartbeat. # We should still raise OperationTimedOut to reject the future so that the main event thread will not @@ -5211,7 +5246,8 @@ def _on_timeout(self, _attempts=0): self._connection.orphaned_threshold_reached = True pool.return_connection(self._connection, stream_was_orphaned=True) - elif self._connection.is_control_connection: + elif self._connection.is_control_connection and \ + not orphaned_control_request: with self._connection.lock: self._connection.orphaned_request_ids.add(self._req_id) if len(self._connection.orphaned_request_ids) >= self._connection.orphaned_threshold: @@ -5361,8 +5397,43 @@ def _handle_control_connection_response(self, connection, cb, response): with control_connection._application_query_lock: control_connection._application_requests_in_flight -= 1 + def _orphan_control_connection_request(self, connection, request_id): + """Detach a timed-out fallback request while retaining a safe barrier. + + The normal response callback owns the application in-flight count. A + timeout removes that callback, so replace it with a cleanup-only + callback and release the count here. The orphan remains a binding + barrier until its late response arrives because a timed-out ``USE`` can + still change the physical connection keyspace. + """ + control_connection = self.session.cluster.control_connection + with control_connection._application_query_lock: + with connection.lock: + try: + _, decoder, result_metadata = connection._requests[request_id] + except KeyError: + return False + + orphan_key = (connection, request_id) + if orphan_key in control_connection._application_orphaned_requests: + return True + callback = partial( + control_connection._handle_orphaned_application_response, + connection, request_id) + connection._requests[request_id] = \ + (callback, decoder, result_metadata) + connection.orphaned_request_ids.add(request_id) + if len(connection.orphaned_request_ids) >= connection.orphaned_threshold: + connection.orphaned_threshold_reached = True + + control_connection._application_requests_in_flight -= 1 + control_connection._application_orphaned_requests.add(orphan_key) + return True + def _is_keyspace_change_query(self, message=None): message = self.message if message is None else message + if isinstance(self.query, GraphStatement): + return False if not isinstance(message, QueryMessage): return False query = getattr(message.query, 'query_string', message.query) diff --git a/tests/unit/test_response_future.py b/tests/unit/test_response_future.py index f2572bd29a..e1929f6cba 100644 --- a/tests/unit/test_response_future.py +++ b/tests/unit/test_response_future.py @@ -23,6 +23,7 @@ from cassandra.cluster import (Session, ResponseFuture, NoHostAvailable, ProtocolVersion, ControlConnection, ControlConnectionQueryFallback, _NOT_SET) from cassandra.connection import Connection, ConnectionBusy, ConnectionException +from cassandra.datastax.graph import SimpleGraphStatement from cassandra.protocol import (ReadTimeoutErrorMessage, WriteTimeoutErrorMessage, UnavailableErrorMessage, ResultMessage, QueryMessage, ExecuteMessage, @@ -603,6 +604,37 @@ def test_control_connection_fallback_rebinds_same_session_new_keyspace(self): conflict = control_connection._attach_application_session('ks1', session2) assert conflict is not None and 'already attached' in conflict + def test_control_connection_fallback_blocks_self_rebind_while_request_active(self): + session = self._make_fallback_session(keyspace='ks1') + control_connection = session.cluster.control_connection + connection = self.make_control_connection() + connection.keyspace = 'ks1' + connection.get_request_id.side_effect = [7, 8] + control_connection._connection = connection + session.cluster.get_control_connection_host.return_value = \ + Mock(endpoint=connection.endpoint) + + rf1 = self.make_response_future(session) + assert rf1.send_request() + assert control_connection._application_requests_in_flight == 1 + + # Simulate set_keyspace() succeeding through a recovered node pool + # while the old-keyspace fallback query is still outstanding. + session.keyspace = 'ks2' + rf2 = self.make_response_future(session) + assert rf2.send_request() + with pytest.raises(InvalidRequest, match='already attached'): + rf2.result() + assert connection.send_msg.call_count == 1 + + connection.send_msg.call_args_list[0][1]['cb']( + self.make_mock_response(['value'], [('one',)])) + + # Once the old request drains, the sole owner may safely rebind. + rf3 = self.make_response_future(session) + assert rf3.send_request() + assert connection.send_msg.call_args_list[1][0][0].query == 'USE ks2' + def test_control_connection_fallback_req_id_tracks_real_message(self): session = self._make_fallback_session(keyspace='ks1') connection = self.make_control_connection() @@ -813,6 +845,35 @@ def test_control_connection_fallback_reports_connection_leftover_keyspace(self): with pytest.raises(InvalidRequest, match="keyspace 'ks1'"): rf2.result() + def test_control_connection_fallback_send_failure_preserves_physical_keyspace(self): + session1 = self._make_fallback_session(keyspace='ks1') + control_connection = session1.cluster.control_connection + connection = self.make_control_connection() + connection.keyspace = 'ks1' + connection.send_msg.side_effect = ConnectionBusy('no streams') + control_connection._connection = connection + assert control_connection._attach_application_session( + 'ks1', session1) is None + session1.is_shutdown = True + + # The new owner is provisionally attached, but its USE ks2 cannot be + # sent and the logical attachment is discarded. + session2 = self._make_fallback_session( + cluster=session1.cluster, keyspace='ks2') + rf2 = self.make_response_future(session2) + assert not rf2.send_request() + assert control_connection._get_application_keyspace() is _NOT_SET + assert connection.keyspace == 'ks1' + + # The discarded logical binding must not hide the physical ks1 state. + session3 = self._make_fallback_session( + cluster=session1.cluster, keyspace=None) + rf3 = self.make_response_future(session3) + assert rf3.send_request() + with pytest.raises(InvalidRequest, match="keyspace 'ks1'"): + rf3.result() + assert connection.send_msg.call_count == 1 + def test_control_connection_fallback_reclaim_without_keyspace_after_reconnect(self): session1 = self._make_fallback_session(keyspace='ks1') connection = self.make_control_connection() @@ -939,6 +1000,7 @@ def test_control_connection_fallback_reprepares_prepared_statement(self): session = self.make_basic_session() session.cluster.allow_control_connection_query_fallback = ControlConnectionQueryFallback.Fallback session.cluster.protocol_version = ProtocolVersion.V4 + session.keyspace = "FooKeyspace" session.cluster._default_load_balancing_policy.make_query_plan.return_value = ['ip1'] session._pools = {} session.submit.side_effect = lambda fn, *args, **kwargs: fn(*args, **kwargs) @@ -1030,6 +1092,17 @@ def test_control_connection_fallback_rejects_indented_use(self): assert not rf._is_keyspace_change_query(indented_select) assert time.time() - start < 2 + def test_control_connection_fallback_does_not_treat_graph_use_as_cql(self): + session = self.make_basic_session() + session.cluster._default_load_balancing_policy.make_query_plan.return_value = [] + graph_query = SimpleGraphStatement( + 'use(TimeCategory) { g.V().count() }') + message = QueryMessage(query=graph_query.query, + consistency_level=ConsistencyLevel.ONE) + rf = ResponseFuture(session, message, graph_query, 1) + + assert not rf._is_keyspace_change_query() + def test_control_connection_fallback_use_failure_reports_error_once(self): session = self.make_basic_session() session.cluster.allow_control_connection_query_fallback = \ @@ -1129,9 +1202,63 @@ def send_msg(message, request_id, cb, **kwargs): assert 7 in connection.orphaned_request_ids assert connection.in_flight == 1 + assert session.cluster.control_connection._application_requests_in_flight == 0 + assert (connection, 7) in \ + session.cluster.control_connection._application_orphaned_requests with pytest.raises(OperationTimedOut): rf.result() + def test_control_connection_fallback_timeout_barrier_ends_on_late_use(self): + session1 = self._make_fallback_session(keyspace='ks1') + control_connection = session1.cluster.control_connection + connection = self.make_control_connection() + connection.get_request_id.side_effect = [7, 8] + control_connection._connection = connection + session1.cluster.get_control_connection_host.return_value = \ + Mock(endpoint=connection.endpoint) + + def send_msg(message, request_id, cb, **kwargs): + connection._requests[request_id] = \ + (cb, kwargs.get('decoder'), kwargs.get('result_metadata')) + return 128 + + connection.send_msg.side_effect = send_msg + + rf1 = self.make_response_future(session1) + assert rf1.send_request() + assert connection.send_msg.call_args_list[0][0][0].query == 'USE ks1' + rf1._on_timeout() + session1.is_shutdown = True + + # Active accounting is released at timeout, but the orphaned USE still + # prevents a different keyspace from sharing the physical connection. + assert control_connection._application_requests_in_flight == 0 + session2 = self._make_fallback_session( + cluster=session1.cluster, keyspace='ks2') + rf2 = self.make_response_future(session2) + assert rf2.send_request() + with pytest.raises(InvalidRequest, match='already attached'): + rf2.result() + + # Process the late SET_KEYSPACE response through the cleanup-only + # callback installed by the timeout. It must not send rf1's query. + late_cb, _, _ = connection._requests.pop(7) + with connection.lock: + connection.in_flight -= 1 + connection.orphaned_request_ids.remove(7) + late_cb(Mock(spec=ResultMessage, kind=RESULT_KIND_SET_KEYSPACE, + new_keyspace='ks1')) + assert connection.send_msg.call_count == 1 + assert not control_connection._application_orphaned_requests + assert connection.keyspace == 'ks1' + + # With the late USE retired, a new owner can safely switch keyspaces. + session3 = self._make_fallback_session( + cluster=session1.cluster, keyspace='ks2') + rf3 = self.make_response_future(session3) + assert rf3.send_request() + assert connection.send_msg.call_args_list[1][0][0].query == 'USE ks2' + def test_control_connection_fallback_timeout_without_metadata_host_uses_connection_endpoint(self): session = self.make_basic_session() session.cluster.allow_control_connection_query_fallback = ControlConnectionQueryFallback.Fallback From 33fb777e66635642cafcf18c7c94bc5c66323dc7 Mon Sep 17 00:00:00 2001 From: Dmitry Kropachev Date: Tue, 22 Sep 2026 08:09:22 -0400 Subject: [PATCH 3/4] control-connection: preserve concurrent fallback ownership Keep provisional session claims until all concurrent initial sends resolve, so a failed first send cannot release ownership claimed by another request. Cover the overlap where one send fails while another is pending, then verify different keyspaces remain rejected and the original session can continue. --- cassandra/cluster.py | 65 ++++++++++++++++++++------ tests/unit/test_response_future.py | 75 +++++++++++++++++++++++++++++- 2 files changed, 124 insertions(+), 16 deletions(-) diff --git a/cassandra/cluster.py b/cassandra/cluster.py index 02b4cd4d58..a52d2af8f4 100644 --- a/cassandra/cluster.py +++ b/cassandra/cluster.py @@ -4029,6 +4029,10 @@ def __init__(self, cluster, timeout, self._application_query_lock = RLock() self._application_keyspace = _NOT_SET self._application_sessions = WeakSet() + # A new Session is provisional until one of its fallback sends + # succeeds. Count concurrent attempts so one failure cannot discard an + # attachment another attempt is about to confirm. + self._application_session_claims = {} self._application_requests_in_flight = 0 self._application_orphaned_requests = set() @@ -4104,12 +4108,38 @@ def _attach_application_session(self, keyspace, session): def _discard_application_session(self, session): """Undo an attachment when its first fallback request was not sent.""" with self._application_query_lock: + self._application_session_claims.pop(session, None) self._application_sessions.discard(session) if not self._application_sessions and \ not self._application_requests_in_flight and \ not self._current_connection_has_application_orphans(): self._application_keyspace = _NOT_SET + def _begin_application_session_claim(self, session, new_session): + """Track a send that can confirm a provisional Session attachment.""" + with self._application_query_lock: + if new_session: + self._application_session_claims[session] = 0 + if session not in self._application_session_claims: + return None + self._application_session_claims[session] += 1 + return session + + def _finish_application_session_claim(self, session, request_sent): + """Confirm an attachment, or discard it after all initial sends fail.""" + if session is None: + return + with self._application_query_lock: + claims = self._application_session_claims.get(session) + if claims is None: + return + if request_sent: + del self._application_session_claims[session] + elif claims > 1: + self._application_session_claims[session] = claims - 1 + else: + self._discard_application_session(session) + def _prune_application_sessions(self): """Drop shut-down owners after shared-connection requests have drained.""" if self._application_requests_in_flight: @@ -5372,7 +5402,7 @@ def _borrow_control_connection(self, connection): return request_id def _release_control_connection_request(self, connection, request_id, - application_session=None): + provisional_session=None): control_connection = self.session.cluster.control_connection with control_connection._application_query_lock: with connection.lock: @@ -5380,8 +5410,8 @@ def _release_control_connection_request(self, connection, request_id, connection.request_ids.append(request_id) connection._requests.pop(request_id, None) control_connection._application_requests_in_flight -= 1 - if application_session is not None: - control_connection._discard_application_session(application_session) + control_connection._finish_application_session_claim( + provisional_session, False) def _handle_control_connection_response(self, connection, cb, response): control_connection = self.session.cluster.control_connection @@ -5452,7 +5482,7 @@ def _control_connection_failed(self): def _set_control_connection_keyspace(self, connection, host, keyspace, message=None, cb=None, request_id=_NOT_SET, - release_session_on_failure=None): + provisional_session=None): use_message = QueryMessage( query='USE %s' % protect_name(keyspace), consistency_level=ConsistencyLevel.ONE) @@ -5487,12 +5517,12 @@ def keyspace_set(response): message=use_message, cb=keyspace_set, connection=connection, host=host, record_attempt=False, record_size=False, request_id=request_id, - release_session_on_failure=release_session_on_failure) + provisional_session=provisional_session) def _send_control_connection_message(self, message=None, cb=None, connection=None, host=None, record_attempt=True, record_size=True, request_id=_NOT_SET, - release_session_on_failure=None): + provisional_session=None): if message is None: message = self.message @@ -5532,9 +5562,12 @@ def _send_control_connection_message(self, message=None, cb=None, connection=Non encoder=self._protocol_handler.encode_message, decoder=self._protocol_handler.decode_message, result_metadata=result_meta) + request_sent = True + control_connection = self.session.cluster.control_connection + control_connection._finish_application_session_claim( + provisional_session, True) if record_size: self.request_encoded_size = encoded_size - request_sent = True if record_attempt: self.attempted_hosts.append(host) return request_id @@ -5556,7 +5589,7 @@ def _send_control_connection_message(self, message=None, cb=None, connection=Non if self._req_id == request_id: self._req_id = previous_req_id self._release_control_connection_request( - connection, request_id, release_session_on_failure) + connection, request_id, provisional_session) return None @@ -5591,6 +5624,9 @@ def _query_control_connection(self, message=None, cb=None, connection=None, host if conflict is not None: self._set_final_exception(InvalidRequest(conflict)) return _NOT_SET + provisional_session = \ + control_connection._begin_application_session_claim( + self.session, new_application_session) if host is None: host = self.session.cluster.get_control_connection_host() or connection.endpoint @@ -5600,29 +5636,28 @@ def _query_control_connection(self, message=None, cb=None, connection=None, host except NoConnectionsAvailable as exc: log.debug("Control connection is at capacity") self._errors[host] = exc - if new_application_session: - control_connection._discard_application_session(self.session) + control_connection._finish_application_session_claim( + provisional_session, False) return None except Exception as exc: log.debug("Error borrowing control connection", exc_info=True) self._errors[host] = exc - if new_application_session: - control_connection._discard_application_session(self.session) + control_connection._finish_application_session_claim( + provisional_session, False) if self._metrics is not None: self._metrics.on_connection_error() return None - release_session_on_failure = self.session if new_application_session else None if keyspace is not None and connection.keyspace != keyspace: return self._set_control_connection_keyspace( connection, host, keyspace, message=message, cb=cb, request_id=request_id, - release_session_on_failure=release_session_on_failure) + provisional_session=provisional_session) return self._send_control_connection_message( message=message, cb=cb, connection=connection, host=host, request_id=request_id, - release_session_on_failure=release_session_on_failure) + provisional_session=provisional_session) def _query(self, host, message=None, cb=None): if message is None: diff --git a/tests/unit/test_response_future.py b/tests/unit/test_response_future.py index e1929f6cba..87bee8f8ca 100644 --- a/tests/unit/test_response_future.py +++ b/tests/unit/test_response_future.py @@ -16,7 +16,7 @@ import unittest from collections import deque -from threading import RLock, Thread +from threading import Event, RLock, Thread from unittest.mock import Mock, MagicMock, ANY, patch from cassandra import ConsistencyLevel, InvalidRequest, Unavailable, SchemaTargetType, SchemaChangeType, OperationTimedOut @@ -694,6 +694,79 @@ def test_control_connection_fallback_req_id_restored_when_send_fails(self): assert session.cluster.control_connection._get_application_keyspace() is _NOT_SET assert session.cluster.control_connection._application_requests_in_flight == 0 + def test_control_connection_fallback_concurrent_send_preserves_session(self): + session = self._make_fallback_session(keyspace='ks1') + session.cluster.allow_control_connection_query_fallback = \ + ControlConnectionQueryFallback.Fallback + control_connection = session.cluster.control_connection + connection = self.make_control_connection() + connection.keyspace = 'ks1' + connection.get_request_id.side_effect = [7, 8, 9] + control_connection._connection = connection + session.cluster.get_control_connection_host.return_value = \ + Mock(endpoint=connection.endpoint) + + first_send_started = Event() + second_send_started = Event() + release_first_send = Event() + release_second_send = Event() + + def send_msg(message, request_id, **kwargs): + if request_id == 7: + first_send_started.set() + assert release_first_send.wait(5) + raise ConnectionBusy('no streams') + if request_id == 8: + second_send_started.set() + assert release_second_send.wait(5) + return 128 + + connection.send_msg.side_effect = send_msg + rf1 = self.make_response_future(session) + first_result = [] + first_thread = Thread(target=lambda: first_result.append(rf1.send_request())) + first_thread.start() + assert first_send_started.wait(5) + + rf2 = self.make_response_future(session) + second_result = [] + second_thread = Thread( + target=lambda: second_result.append(rf2.send_request())) + second_thread.start() + assert second_send_started.wait(5) + + # Even before the second send returns successfully, its pending claim + # prevents the first failure from discarding their shared Session. + release_first_send.set() + first_thread.join(5) + assert not first_thread.is_alive() + assert first_result == [False] + assert session in control_connection._application_sessions + + release_second_send.set() + second_thread.join(5) + assert not second_thread.is_alive() + assert second_result == [True] + + # The successful concurrent request claimed the provisional binding. + # Finishing the failed first send must not release its live Session. + connection.send_msg.call_args_list[1][1]['cb']( + self.make_mock_response(['value'], [('two',)])) + assert control_connection._application_requests_in_flight == 0 + assert session in control_connection._application_sessions + assert control_connection._get_application_keyspace() == 'ks1' + + other_session = self._make_fallback_session( + cluster=session.cluster, keyspace='ks2') + other_rf = self.make_response_future(other_session) + assert other_rf.send_request() + with pytest.raises(InvalidRequest, match='already attached'): + other_rf.result() + + rf3 = self.make_response_future(session) + assert rf3.send_request() + assert connection.send_msg.call_args_list[2][0][0] is rf3.message + def test_speculative_execute_honours_expired_deadline_without_attempts(self): session = self.make_basic_session() session.cluster._default_load_balancing_policy.make_query_plan.return_value = [] From ae648309d2f7fd4991f0dc73746000f7af3da07c Mon Sep 17 00:00:00 2001 From: Dmitry Kropachev Date: Tue, 22 Sep 2026 11:40:03 -0400 Subject: [PATCH 4/4] cluster: protect timeout stream ownership Record pooled request IDs before sending so a speculative timeout can detach the callback for the stream actually in flight. Track active control-connection fallback callbacks and verify their identity before orphaning a stream. This keeps recycled control request IDs and fallback accounting untouched after the original response has completed. --- cassandra/cluster.py | 97 ++++++++++++++++++++++++------ tests/unit/test_response_future.py | 97 ++++++++++++++++++++++++++++++ 2 files changed, 174 insertions(+), 20 deletions(-) diff --git a/cassandra/cluster.py b/cassandra/cluster.py index a52d2af8f4..bcc4fd062d 100644 --- a/cassandra/cluster.py +++ b/cassandra/cluster.py @@ -5198,6 +5198,13 @@ def __init__(self, session, message, query, timeout, metrics=None, prepared_stat self._errors = {} self._callbacks = [] self._errbacks = [] + # Maps fallback (connection, request id) pairs to the callback this + # future installed. Request ids may be reused as soon as a response + # callback returns, while the future can remain pending for a retry, + # reprepare, or schema agreement. Keep the callback identity as well as + # the id so a later timeout cannot detach an unrelated request that + # reused the same stream. + self._control_connection_requests = {} self.attempted_hosts = [] self._start_timer() self._continuous_paging_state = continuous_paging_state @@ -5240,13 +5247,18 @@ def _on_timeout(self, _attempts=0): conn_in_flight = None if self._connection is not None: - orphaned_control_request = \ + control_connection_request = \ self._connection.is_control_connection and \ - self._control_connection_query_attempted and \ + self._control_connection_query_attempted + if control_connection_request: self._orphan_control_connection_request( self._connection, self._req_id) try: - if not orphaned_control_request: + # A completed fallback stream may already have been reused by + # control traffic while this future waits for follow-up work. + # _orphan_control_connection_request() verifies ownership; if + # it does not own the stream, leave the current request alone. + if not control_connection_request: self._connection._requests.pop(self._req_id) # PYTHON-1044 # This request might have been removed from the connection after the latter was defunct by heartbeat. @@ -5263,7 +5275,11 @@ def _on_timeout(self, _attempts=0): # Capture connection stats before pool.return_connection() can alter state conn_in_flight = self._connection.in_flight - pool = self.session._pools.get(self._current_host) + # Fallback requests never belong to a Session pool. A pool may + # recover while the future waits for retry/schema work, but the + # completed fallback stream must not be returned through that pool. + pool = None if control_connection_request else \ + self.session._pools.get(self._current_host) if pool and not pool.is_shutdown: # Do not return the stream ID to the pool yet. We cannot reuse it # because the node might still be processing the query and will @@ -5277,7 +5293,7 @@ def _on_timeout(self, _attempts=0): pool.return_connection(self._connection, stream_was_orphaned=True) elif self._connection.is_control_connection and \ - not orphaned_control_request: + not control_connection_request: with self._connection.lock: self._connection.orphaned_request_ids.add(self._req_id) if len(self._connection.orphaned_request_ids) >= self._connection.orphaned_threshold: @@ -5405,15 +5421,31 @@ def _release_control_connection_request(self, connection, request_id, provisional_session=None): control_connection = self.session.cluster.control_connection with control_connection._application_query_lock: - with connection.lock: - connection.in_flight -= 1 - connection.request_ids.append(request_id) - connection._requests.pop(request_id, None) - control_connection._application_requests_in_flight -= 1 + request_key = (connection, request_id) + active_request = self._control_connection_requests.pop( + request_key, None) + orphaned_request = request_key in \ + control_connection._application_orphaned_requests + if orphaned_request: + control_connection._application_orphaned_requests.discard( + request_key) + + # A timeout may have taken ownership while send_msg() was still + # encoding and then send_msg() may fail before pushing any bytes. + # Retire either form of ownership exactly once in that case. + if active_request is not None or orphaned_request: + with connection.lock: + connection.in_flight -= 1 + connection.request_ids.append(request_id) + connection._requests.pop(request_id, None) + connection.orphaned_request_ids.discard(request_id) + if active_request is not None: + control_connection._application_requests_in_flight -= 1 control_connection._finish_application_session_claim( provisional_session, False) - def _handle_control_connection_response(self, connection, cb, response): + def _handle_control_connection_response(self, connection, request_id, cb, + response): control_connection = self.session.cluster.control_connection with connection.lock: connection.in_flight -= 1 @@ -5425,7 +5457,10 @@ def _handle_control_connection_response(self, connection, cb, response): cb(response) finally: with control_connection._application_query_lock: - control_connection._application_requests_in_flight -= 1 + request_key = (connection, request_id) + if self._control_connection_requests.pop( + request_key, None) is not None: + control_connection._application_requests_in_flight -= 1 def _orphan_control_connection_request(self, connection, request_id): """Detach a timed-out fallback request while retaining a safe barrier. @@ -5439,14 +5474,19 @@ def _orphan_control_connection_request(self, connection, request_id): control_connection = self.session.cluster.control_connection with control_connection._application_query_lock: with connection.lock: + orphan_key = (connection, request_id) + expected_callback = self._control_connection_requests.get( + orphan_key) + if expected_callback is None: + return False try: - _, decoder, result_metadata = connection._requests[request_id] + callback, decoder, result_metadata = \ + connection._requests[request_id] except KeyError: return False - orphan_key = (connection, request_id) - if orphan_key in control_connection._application_orphaned_requests: - return True + if callback is not expected_callback: + return False callback = partial( control_connection._handle_orphaned_application_response, connection, request_id) @@ -5456,6 +5496,7 @@ def _orphan_control_connection_request(self, connection, request_id): if len(connection.orphaned_request_ids) >= connection.orphaned_threshold: connection.orphaned_threshold_reached = True + self._control_connection_requests.pop(orphan_key, None) control_connection._application_requests_in_flight -= 1 control_connection._application_orphaned_requests.add(orphan_key) return True @@ -5544,11 +5585,17 @@ def _send_control_connection_message(self, message=None, cb=None, connection=Non try: if request_id is None: request_id = self._borrow_control_connection(connection) - self._connection = connection result_meta = self._bound_result_metadata if cb is None: cb = partial(self._set_result, host, connection, None) - cb = partial(self._handle_control_connection_response, connection, cb) + cb = partial(self._handle_control_connection_response, connection, + request_id, cb) + + control_connection = self.session.cluster.control_connection + with control_connection._application_query_lock: + self._control_connection_requests[(connection, request_id)] = cb + self._connection = connection + self._req_id = request_id log.debug("No usable node pools; falling back to control connection for host %s", host) # Record the stream id before sending, not after. The reply can be @@ -5557,13 +5604,11 @@ def _send_control_connection_message(self, message=None, cb=None, connection=Non # whose id survives in _req_id. Assigning after send_msg() would put # the already-completed USE id there instead, so a later timeout would # orphan the wrong stream and leave the real request in _requests. - self._req_id = request_id encoded_size = connection.send_msg(message, request_id, cb=cb, encoder=self._protocol_handler.encode_message, decoder=self._protocol_handler.decode_message, result_metadata=result_meta) request_sent = True - control_connection = self.session.cluster.control_connection control_connection._finish_application_session_claim( provisional_session, True) if record_size: @@ -5676,6 +5721,9 @@ def _query(self, host, message=None, cb=None): self._current_host = host connection = None + previous_req_id = self._req_id + request_id = None + request_sent = False try: # TODO get connectTimeout from cluster settings if self.query: @@ -5693,10 +5741,15 @@ def _query(self, host, message=None, cb=None): if cb is None: cb = partial(self._set_result, host, connection, pool) + # Record the stream before send_msg() starts. Encoding or pushing a + # message can overlap the deadline timer; the timeout must be able + # to detach the callback send_msg() has installed for this stream. + self._req_id = request_id self.request_encoded_size = connection.send_msg(message, request_id, cb=cb, encoder=self._protocol_handler.encode_message, decoder=self._protocol_handler.decode_message, result_metadata=result_meta) + request_sent = True self.attempted_hosts.append(host) return request_id except NoConnectionsAvailable as exc: @@ -5712,6 +5765,10 @@ def _query(self, host, message=None, cb=None): self._metrics.on_connection_error() if connection: pool.return_connection(connection) + finally: + if request_id is not None and not request_sent and \ + self._req_id == request_id: + self._req_id = previous_req_id return None diff --git a/tests/unit/test_response_future.py b/tests/unit/test_response_future.py index 87bee8f8ca..1515143bd5 100644 --- a/tests/unit/test_response_future.py +++ b/tests/unit/test_response_future.py @@ -786,6 +786,58 @@ def test_speculative_execute_honours_expired_deadline_without_attempts(self): rf._on_timeout.assert_called_once_with() session.cluster.connection_class.create_timer.assert_not_called() + def test_speculative_timeout_during_first_send_orphans_real_stream(self): + session = self.make_basic_session() + host = Mock(endpoint='ip1') + session.cluster._default_load_balancing_policy.make_query_plan.return_value = [host] + + pool = Mock(is_shutdown=False) + connection = Mock(spec=Connection) + connection.lock = RLock() + connection._requests = {} + connection.in_flight = 1 + connection.is_control_connection = False + connection.orphaned_request_ids = set() + connection.orphaned_threshold = 75 + connection.orphaned_threshold_reached = False + pool.borrow_connection.return_value = (connection, 11) + session._pools = {host: pool} + + send_started = Event() + release_send = Event() + + def send_msg(message, request_id, cb, **kwargs): + connection._requests[request_id] = \ + (cb, kwargs.get('decoder'), kwargs.get('result_metadata')) + send_started.set() + assert release_send.wait(5) + return 128 + + connection.send_msg.side_effect = send_msg + rf = self.make_response_future(session) + send_result = [] + send_thread = Thread(target=lambda: send_result.append(rf.send_request())) + send_thread.start() + assert send_started.wait(5) + + # send_request() has not returned to copy its local request id yet, but + # the timeout must still detach the callback send_msg() installed. + assert not rf.attempted_hosts + assert rf._req_id == 11 + rf._start_time = time.time() - 5 + rf._on_speculative_execute() + + release_send.set() + send_thread.join(5) + assert not send_thread.is_alive() + assert send_result == [True] + assert 11 not in connection._requests + assert 11 in connection.orphaned_request_ids + pool.return_connection.assert_called_once_with( + connection, stream_was_orphaned=True) + with pytest.raises(OperationTimedOut): + rf.result() + def _make_fallback_session(self, cluster=None, keyspace=None): session = self.make_basic_session() if cluster is not None: @@ -1281,6 +1333,51 @@ def send_msg(message, request_id, cb, **kwargs): with pytest.raises(OperationTimedOut): rf.result() + def test_control_connection_timeout_does_not_orphan_reused_stream(self): + session = self._make_fallback_session() + connection = self.make_control_connection() + session.cluster.control_connection._connection = connection + + def send_msg(message, request_id, cb, **kwargs): + connection._requests[request_id] = \ + (cb, kwargs.get('decoder'), kwargs.get('result_metadata')) + return 128 + + connection.send_msg.side_effect = send_msg + rf = self.make_response_future(session) + assert rf.send_request() + + # Complete the physical fallback request with a schema change. The + # future remains pending while schema agreement runs, but stream 7 is + # no longer owned by it and may be reused by control traffic. + response_cb, _, _ = connection._requests.pop(7) + response_cb(Mock(spec=ResultMessage, + kind=RESULT_KIND_SCHEMA_CHANGE, + schema_change_event={})) + assert not rf._event.is_set() + assert not rf._control_connection_requests + assert session.cluster.control_connection._application_requests_in_flight == 0 + + unrelated_cb = Mock() + unrelated_decoder = Mock() + unrelated_metadata = Mock() + unrelated_request = \ + (unrelated_cb, unrelated_decoder, unrelated_metadata) + connection._requests[7] = unrelated_request + connection.in_flight = 1 + recovered_pool = Mock(is_shutdown=False) + session._pools = {rf._current_host: recovered_pool} + + rf._on_timeout() + + assert connection._requests[7] is unrelated_request + assert not connection.orphaned_request_ids + recovered_pool.return_connection.assert_not_called() + assert session.cluster.control_connection._application_requests_in_flight == 0 + assert not session.cluster.control_connection._application_orphaned_requests + with pytest.raises(OperationTimedOut): + rf.result() + def test_control_connection_fallback_timeout_barrier_ends_on_late_use(self): session1 = self._make_fallback_session(keyspace='ks1') control_connection = session1.cluster.control_connection