control-connection: tolerate schema agreement connection loss

Author: dkropachev

cassandra/cluster.py

@@ -3531,6 +3531,7 @@ def __init__(self, cluster, timeout,
         # shutdown) since implementing __del__ disables the cycle detector
         self._cluster = weakref.proxy(cluster)
         self._connection = None
+        self._last_connection_endpoint = None
         self._timeout = timeout
 
         self._schema_event_refresh_window = schema_event_refresh_window
@@ -3564,6 +3565,7 @@ def _set_new_connection(self, conn):
         with self._lock:
             old = self._connection
             self._connection = conn
+            self._last_connection_endpoint = conn.endpoint
 
         if old:
             log.debug("[control connection] Closing old connection %r, replacing with %r", old, conn)
@@ -4094,33 +4096,28 @@ def wait_for_schema_agreement(self, connection=None, preloaded_results=None, wai
 
             if not connection:
                 connection = self._connection
+            local_address = self._schema_agreement_endpoint(connection)
 
             if preloaded_results:
                 log.debug("[control connection] Attempting to use preloaded results for schema agreement")
 
                 peers_result = preloaded_results[0]
                 local_result = preloaded_results[1]
-                schema_mismatches = self._get_schema_mismatches(peers_result, local_result, connection.endpoint)
+                schema_mismatches = self._get_schema_mismatches(peers_result, local_result, local_address)
                 if schema_mismatches is None:
                     return True
 
             log.debug("[control connection] Waiting for schema agreement")
             start = self._time.time()
             elapsed = 0
-            cl = ConsistencyLevel.ONE
             schema_mismatches = None
-            select_peers_query = self._get_peers_query(self.PeersQueryType.PEERS_SCHEMA, connection)
 
             while elapsed < total_timeout:
-                peers_query = QueryMessage(query=maybe_add_timeout_to_query(select_peers_query, self._metadata_request_timeout),
-                                           consistency_level=cl)
-                local_query = QueryMessage(query=maybe_add_timeout_to_query(self._SELECT_SCHEMA_LOCAL, self._metadata_request_timeout),
-                                           consistency_level=cl)
                 try:
                     remaining = total_timeout - elapsed
                     timeout = min(self._timeout, remaining) if self._timeout is not None else remaining
-                    peers_result, local_result = connection.wait_for_responses(
-                        peers_query, local_query, timeout=timeout)
+                    peers_result, local_result, local_address = self._get_schema_agreement_results(
+                        connection, timeout)
                 except OperationTimedOut as timeout:
                     log.debug("[control connection] Timed out waiting for "
                               "response during schema agreement check: %s", timeout)
@@ -4131,9 +4128,13 @@ def wait_for_schema_agreement(self, connection=None, preloaded_results=None, wai
                         log.debug("[control connection] Aborting wait for schema match due to shutdown")
                         return None
                     else:
-                        raise
+                        log.debug("[control connection] Connection lost during schema agreement check")
+                        return False
+                except (ConnectionBusy, ConnectionException, NoConnectionsAvailable) as exc:
+                    log.debug("[control connection] Unable to check schema agreement: %s", exc)
+                    return False
 
-                schema_mismatches = self._get_schema_mismatches(peers_result, local_result, connection.endpoint)
+                schema_mismatches = self._get_schema_mismatches(peers_result, local_result, local_address)
                 if schema_mismatches is None:
                     return True
 
@@ -4142,9 +4143,150 @@ def wait_for_schema_agreement(self, connection=None, preloaded_results=None, wai
                 elapsed = self._time.time() - start
 
             log.warning("Node %s is reporting a schema disagreement: %s",
-                        connection.endpoint, schema_mismatches)
+                        local_address, schema_mismatches)
             return False
 
+    def _get_schema_agreement_results(self, connection, timeout):
+        if self._is_connection_unusable(connection):
+            return self._get_schema_agreement_results_from_session(connection, timeout)
+
+        try:
+            return self._get_schema_agreement_results_from_connection(connection, timeout)
+        except OperationTimedOut:
+            raise
+        except (ConnectionBusy, ConnectionException):
+            if self._is_shutdown:
+                raise
+            self.return_connection(connection)
+            log.debug("[control connection] Falling back to session connection for schema agreement")
+            return self._get_schema_agreement_results_from_session(connection, timeout)
+
+    def _get_schema_agreement_results_from_connection(self, connection, timeout):
+        peers_query, local_query = self._schema_agreement_queries(connection)
+        peers_result, local_result = connection.wait_for_responses(peers_query, local_query, timeout=timeout)
+        return peers_result, local_result, connection.endpoint
+
+    def _get_schema_agreement_results_from_session(self, connection, timeout):
+        endpoint = self._schema_agreement_endpoint(connection)
+        deadline = None if timeout is None else time.time() + timeout
+
+        pool, borrowed_connection, request_id = self._borrow_schema_agreement_connection(endpoint, timeout)
+        peers_query, _ = self._schema_agreement_queries(borrowed_connection)
+        local_address = borrowed_connection.endpoint
+        peers_result = self._wait_for_borrowed_schema_response(
+            pool, borrowed_connection, request_id, peers_query, self._remaining_timeout(deadline))
+
+        pool, borrowed_connection, request_id = self._borrow_schema_agreement_connection(
+            endpoint, self._remaining_timeout(deadline))
+        _, local_query = self._schema_agreement_queries(borrowed_connection)
+        local_result = self._wait_for_borrowed_schema_response(
+            pool, borrowed_connection, request_id, local_query, self._remaining_timeout(deadline))
+
+        return peers_result, local_result, local_address
+
+    def _schema_agreement_queries(self, connection):
+        cl = ConsistencyLevel.ONE
+        select_peers_query = self._get_peers_query(self.PeersQueryType.PEERS_SCHEMA, connection)
+        peers_query = QueryMessage(query=maybe_add_timeout_to_query(select_peers_query, self._metadata_request_timeout),
+                                   consistency_level=cl)
+        local_query = QueryMessage(query=maybe_add_timeout_to_query(self._SELECT_SCHEMA_LOCAL, self._metadata_request_timeout),
+                                   consistency_level=cl)
+        return peers_query, local_query
+
+    def _remaining_timeout(self, deadline):
+        if deadline is None:
+            return None
+        return max(0, deadline - time.time())
+
+    def _borrow_schema_agreement_connection(self, endpoint, timeout):
+        last_error = None
+        for pool in self._schema_agreement_pools(endpoint):
+            try:
+                return (pool,) + pool.borrow_connection(timeout=timeout)
+            except (NoConnectionsAvailable, ConnectionBusy, ConnectionException) as exc:
+                last_error = exc
+                log.debug("[control connection] Unable to borrow connection for schema agreement: %s", exc)
+
+        if last_error:
+            raise last_error
+        raise NoConnectionsAvailable("No session connection available for schema agreement on %s" % (endpoint,))
+
+    def _schema_agreement_pools(self, endpoint):
+        if endpoint is None:
+            return
+
+        host = self._cluster.metadata.get_host(endpoint)
+        for session in tuple(getattr(self._cluster, 'sessions', ())):
+            pools = getattr(session, '_pools', {})
+            pool = pools.get(host) if host else None
+            if pool is None:
+                for pool_host, candidate_pool in tuple(pools.items()):
+                    if getattr(pool_host, 'endpoint', None) == endpoint:
+                        pool = candidate_pool
+                        break
+            if pool and not pool.is_shutdown:
+                yield pool
+
+    def _wait_for_borrowed_schema_response(self, pool, connection, request_id, message, timeout):
+        event = Event()
+        responses = []
+
+        def callback(response):
+            responses.append(response)
+            event.set()
+
+        sent = False
+        orphaned = False
+        try:
+            connection.send_msg(message, request_id, callback)
+            sent = True
+            if not event.wait(timeout):
+                orphaned = self._orphan_borrowed_request(pool, connection, request_id)
+                raise OperationTimedOut(timeout=timeout, in_flight=getattr(connection, 'in_flight', None))
+
+            response = responses[0]
+            if isinstance(response, Exception):
+                if hasattr(response, 'to_exception'):
+                    response = response.to_exception()
+                raise response
+            return response
+        except Exception:
+            if not sent:
+                self._return_request_id_if_unused(connection, request_id)
+            raise
+        finally:
+            if not orphaned:
+                pool.return_connection(connection)
+
+    def _orphan_borrowed_request(self, pool, connection, request_id):
+        try:
+            connection._requests.pop(request_id)
+        except KeyError:
+            return False
+
+        with connection.lock:
+            connection.orphaned_request_ids.add(request_id)
+            if len(connection.orphaned_request_ids) >= connection.orphaned_threshold:
+                connection.orphaned_threshold_reached = True
+
+        pool.return_connection(connection, stream_was_orphaned=True)
+        return True
+
+    def _return_request_id_if_unused(self, connection, request_id):
+        with connection.lock:
+            if request_id in connection._requests or request_id in connection.orphaned_request_ids:
+                return
+            if request_id not in connection.request_ids:
+                connection.request_ids.append(request_id)
+
+    def _schema_agreement_endpoint(self, connection):
+        return getattr(connection, 'endpoint', None) or self._last_connection_endpoint
+
+    def _is_connection_unusable(self, connection):
+        return (connection is None or
+                getattr(connection, 'is_closed', False) or
+                getattr(connection, 'is_defunct', False))
+
     def _get_schema_mismatches(self, peers_result, local_result, local_address):
         peers_result = dict_factory(peers_result.column_names, peers_result.parsed_rows)
 
@@ -4265,7 +4407,8 @@ def get_connections(self):
         return [c] if c else []
 
     def return_connection(self, connection):
-        if connection is self._connection and (connection.is_defunct or connection.is_closed):
+        if (connection is self._connection and
+                (getattr(connection, 'is_defunct', False) or getattr(connection, 'is_closed', False))):
             self.reconnect()
 
 

tests/unit/test_control_connection.py

@@ -14,14 +14,16 @@
 
 import unittest
 
+from collections import deque
 from concurrent.futures import ThreadPoolExecutor
+from threading import Lock
 from unittest.mock import Mock, ANY, call
 
 from cassandra import OperationTimedOut, SchemaTargetType, SchemaChangeType
 from cassandra.protocol import ResultMessage, RESULT_KIND_ROWS
 from cassandra.cluster import ControlConnection, _Scheduler, ProfileManager, EXEC_PROFILE_DEFAULT, ExecutionProfile
 from cassandra.pool import Host
-from cassandra.connection import EndPoint, DefaultEndPoint, DefaultEndPointFactory
+from cassandra.connection import EndPoint, DefaultEndPoint, DefaultEndPointFactory, ConnectionShutdown
 from cassandra.policies import (SimpleConvictionPolicy, RoundRobinPolicy,
                                 ConstantReconnectionPolicy, IdentityTranslator)
 
@@ -106,6 +108,7 @@ class MockCluster(object):
     def __init__(self):
         self.metadata = MockMetadata()
         self.added_hosts = []
+        self.sessions = []
         self.scheduler = Mock(spec=_Scheduler)
         self.executor = Mock(spec=ThreadPoolExecutor)
         self.profile_manager.profiles[EXEC_PROFILE_DEFAULT] = ExecutionProfile(RoundRobinPolicy())
@@ -146,6 +149,7 @@ def _node_meta_results(local_results, peer_results):
 class MockConnection(object):
 
     is_defunct = False
+    is_closed = False
 
     def __init__(self):
         self.endpoint = DefaultEndPoint("192.168.1.0")
@@ -169,6 +173,29 @@ def __init__(self):
         self.wait_for_responses = Mock(return_value=_node_meta_results(self.local_results, self.peer_results))
 
 
+class MockBorrowedConnection(MockConnection):
+
+    is_closed = False
+
+    def __init__(self, responses):
+        super(MockBorrowedConnection, self).__init__()
+        self.responses = deque(responses)
+        self.lock = Lock()
+        self.in_flight = 0
+        self._requests = {}
+        self.request_ids = deque()
+        self.orphaned_request_ids = set()
+        self.orphaned_threshold = 100
+        self.orphaned_threshold_reached = False
+        self.send_msg = Mock(side_effect=self._send_msg)
+
+    def _send_msg(self, message, request_id, callback):
+        self._requests[request_id] = callback
+        callback(self.responses.popleft())
+        self._requests.pop(request_id, None)
+        self.request_ids.append(request_id)
+
+
 class FakeTime(object):
 
     def __init__(self):
@@ -301,6 +328,50 @@ def test_wait_for_schema_agreement_none_timeout(self):
         cc._time = self.time
         assert cc.wait_for_schema_agreement()
 
+    def test_wait_for_schema_agreement_falls_back_to_session_connection(self):
+        """
+        If the control connection is broken, use a session connection from the
+        same host for schema agreement.
+        """
+        self.connection.wait_for_responses.side_effect = ConnectionShutdown("closed")
+        host = self.cluster.metadata.get_host(self.connection.endpoint)
+        borrowed_connection = MockBorrowedConnection(
+            _node_meta_results(self.connection.local_results, self.connection.peer_results))
+
+        pool = Mock()
+        pool.is_shutdown = False
+        pool.borrow_connection.side_effect = [(borrowed_connection, 1), (borrowed_connection, 2)]
+        session = Mock()
+        session._pools = {host: pool}
+        self.cluster.sessions = [session]
+
+        assert self.control_connection.wait_for_schema_agreement()
+        assert pool.borrow_connection.call_count == 2
+        assert pool.return_connection.call_count == 2
+
+    def test_wait_for_schema_agreement_falls_back_when_control_connection_is_gone(self):
+        self.control_connection._last_connection_endpoint = self.connection.endpoint
+        self.control_connection._connection = None
+        host = self.cluster.metadata.get_host(self.connection.endpoint)
+        borrowed_connection = MockBorrowedConnection(
+            _node_meta_results(self.connection.local_results, self.connection.peer_results))
+
+        pool = Mock()
+        pool.is_shutdown = False
+        pool.borrow_connection.side_effect = [(borrowed_connection, 1), (borrowed_connection, 2)]
+        session = Mock()
+        session._pools = {host: pool}
+        self.cluster.sessions = [session]
+
+        assert self.control_connection.wait_for_schema_agreement()
+        assert pool.borrow_connection.call_count == 2
+        assert pool.return_connection.call_count == 2
+
+    def test_wait_for_schema_agreement_tolerates_missing_connection(self):
+        self.control_connection._connection = None
+
+        assert not self.control_connection.wait_for_schema_agreement()
+
     def test_refresh_nodes_and_tokens(self):
         self.control_connection.refresh_node_list_and_token_map()
         meta = self.cluster.metadata