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..bcc4fd062d 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,25 @@ 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() + # 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() + def connect(self): if self._is_shutdown: return @@ -4002,6 +4057,127 @@ 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() + + 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() + 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_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: + 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.""" + 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 _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 = {} @@ -4196,11 +4372,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: @@ -5021,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 @@ -5063,8 +5247,19 @@ def _on_timeout(self, _attempts=0): conn_in_flight = None if self._connection is not None: + control_connection_request = \ + self._connection.is_control_connection and \ + self._control_connection_query_attempted + if control_connection_request: + self._orphan_control_connection_request( + self._connection, self._req_id) try: - self._connection._requests.pop(self._req_id) + # 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. # We should still raise OperationTimedOut to reject the future so that the main event thread will not @@ -5080,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 @@ -5093,7 +5292,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 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: @@ -5123,6 +5323,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 +5346,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 +5376,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 +5407,163 @@ 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() - - 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) + 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 _handle_control_connection_response(self, connection, cb, response): + 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: + 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, request_id, 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: + 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. + + 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: + orphan_key = (connection, request_id) + expected_callback = self._control_connection_requests.get( + orphan_key) + if expected_callback is None: + return False + try: + callback, decoder, result_metadata = \ + connection._requests[request_id] + except KeyError: + return False + + if callback is not expected_callback: + return False + 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 + + 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 + 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) + 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, + provisional_session=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, + 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, + provisional_session=None): if message is None: message = self.message @@ -5229,23 +5578,43 @@ 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) - self._connection = connection + if request_id is None: + request_id = self._borrow_control_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) - 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. + 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) + control_connection._finish_application_session_claim( + provisional_session, True) + if record_size: + self.request_encoded_size = encoded_size + if record_attempt: + self.attempted_hosts.append(host) return request_id except NoConnectionsAvailable as exc: log.debug("Control connection is at capacity") @@ -5260,10 +5629,81 @@ 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, provisional_session) 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 + 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 + + try: + request_id = self._borrow_control_connection(connection) + except NoConnectionsAvailable as exc: + log.debug("Control connection is at capacity") + self._errors[host] = exc + 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 + control_connection._finish_application_session_claim( + provisional_session, False) + if self._metrics is not None: + self._metrics.on_connection_error() + return 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, + provisional_session=provisional_session) + + return self._send_control_connection_message( + message=message, cb=cb, connection=connection, host=host, + request_id=request_id, + provisional_session=provisional_session) + def _query(self, host, message=None, cb=None): if message is None: message = self.message @@ -5281,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: @@ -5298,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: @@ -5317,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/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..1515143bd5 100644 --- a/tests/unit/test_response_future.py +++ b/tests/unit/test_response_future.py @@ -12,15 +12,18 @@ # 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 Event, 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.datastax.graph import SimpleGraphStatement from cassandra.protocol import (ReadTimeoutErrorMessage, WriteTimeoutErrorMessage, UnavailableErrorMessage, ResultMessage, QueryMessage, ExecuteMessage, @@ -40,8 +43,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 +72,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 +441,587 @@ 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_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() + 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_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 = [] + 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 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: + 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',)])) + + # 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() - assert connection.keyspace == 'newks' - assert session.keyspace == 'newks' - assert rf.result().current_rows == [] + 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_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() + 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() @@ -560,6 +1125,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) @@ -612,6 +1178,118 @@ 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_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 = \ + 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 @@ -649,9 +1327,108 @@ 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_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 + 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