connection: tolerate clean startup close

Author: dkropachev

cassandra/cluster.py

@@ -1753,6 +1753,7 @@ def connection_factory(self, endpoint, host_conn = None, *args, **kwargs):
 
     def _make_connection_factory(self, host, *args, **kwargs):
         kwargs = self._make_connection_kwargs(host.endpoint, kwargs)
+        kwargs['_raise_on_startup_close'] = True
         return partial(self.connection_class.factory, host.endpoint, self.connect_timeout, *args, **kwargs)
 
     def _make_connection_kwargs(self, endpoint, kwargs_dict):

cassandra/connection.py

@@ -965,10 +965,16 @@ def create_timer(cls, timeout, callback):
     @classmethod
     def factory(cls, endpoint, timeout, host_conn = None, *args, **kwargs):
         """
-        A factory function which returns connections which have
-        succeeded in connecting and are ready for service (or
-        raises an exception otherwise).
+        A factory function which returns a connection once startup has
+        completed, or raises an exception otherwise.
+
+        Direct callers may receive a closed connection when the server accepts
+        the socket and then cleanly closes it during the startup handshake.
+        Callers that own pool startup or reconnection probes pass ``host_conn``
+        or ``_raise_on_startup_close`` and receive :exc:`ConnectionShutdown`
+        for that clean startup close instead.
         """
+        raise_on_startup_close = kwargs.pop('_raise_on_startup_close', False)
         start = time.time()
         kwargs['connect_timeout'] = timeout
         conn = cls(endpoint, *args, **kwargs)
@@ -986,8 +992,10 @@ def factory(cls, endpoint, timeout, host_conn = None, *args, **kwargs):
             conn.close()
             raise OperationTimedOut("Timed out creating connection (%s seconds)" % timeout,
                                     timeout=timeout)
-        elif conn.is_closed:
-            raise ConnectionShutdown("Connection to %s was closed by server" % conn.endpoint)
+        elif conn.is_closed and (host_conn is not None or raise_on_startup_close):
+            raise ConnectionShutdown(
+                "Connection to %s was closed during the startup handshake" % (endpoint,),
+                endpoint)
         else:
             return conn
 

cassandra/pool.py

@@ -426,7 +426,10 @@ def __init__(self, host, host_distance, session):
             return
 
         log.debug("Initializing connection for host %s", self.host)
-        first_connection = session.cluster.connection_factory(self.host.endpoint, on_orphaned_stream_released=self.on_orphaned_stream_released)
+        first_connection = session.cluster.connection_factory(
+            self.host.endpoint,
+            _raise_on_startup_close=True,
+            on_orphaned_stream_released=self.on_orphaned_stream_released)
         log.debug("First connection created to %s for shard_id=%i", self.host, first_connection.features.shard_id)
         self._connections[first_connection.features.shard_id] = first_connection
         self._keyspace = session.keyspace
@@ -605,8 +608,10 @@ def _replace(self, connection):
                     self._connecting.add(connection.features.shard_id)
                     self._session.submit(self._open_connection_to_missing_shard, connection.features.shard_id)
                 else:
-                    connection = self._session.cluster.connection_factory(self.host.endpoint,
-                                                                          on_orphaned_stream_released=self.on_orphaned_stream_released)
+                    connection = self._session.cluster.connection_factory(
+                        self.host.endpoint,
+                        _raise_on_startup_close=True,
+                        on_orphaned_stream_released=self.on_orphaned_stream_released)
                     if self._keyspace:
                         connection.set_keyspace_blocking(self._keyspace)
                     self._connections[connection.features.shard_id] = connection

tests/unit/test_cluster.py

@@ -276,8 +276,22 @@ def test_connection_factory_passes_compression_kwarg(self):
                 assert conn == 'connection'
                 assert factory.call_count == 1
                 assert factory.call_args.kwargs['compression'] == expected
+                assert '_raise_on_startup_close' not in factory.call_args.kwargs
                 assert cluster.compression == expected
 
+    def test_reconnection_factory_rejects_startup_close(self):
+        endpoint = Mock(address='127.0.0.1')
+        host = Mock(endpoint=endpoint)
+
+        with patch.object(Cluster.connection_class, 'factory', autospec=True, return_value='connection') as factory:
+            cluster = Cluster()
+            conn_factory = cluster._make_connection_factory(host)
+            conn = conn_factory()
+
+        assert conn == 'connection'
+        assert factory.call_count == 1
+        assert factory.call_args.kwargs['_raise_on_startup_close'] is True
+
 
 class SchedulerTest(unittest.TestCase):
     # TODO: this suite could be expanded; for now just adding a test covering a ticket

tests/unit/test_connection.py

@@ -12,10 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import itertools
+import socket
 import unittest
 from io import BytesIO
 import time
-from threading import Lock
+from threading import Event, Lock, Thread
 from unittest.mock import Mock, ANY, call, patch
 
 from cassandra import OperationTimedOut
@@ -383,6 +384,142 @@ def test_wait_for_responses_shutdown_includes_last_error(self):
         assert "already closed" in error_message
         assert "Bad file descriptor" in error_message
 
+    def test_factory_returns_maintenance_mode_clean_startup_close(self):
+        """
+        Maintenance mode accepts regular CQL sockets and closes them during
+        startup. Direct callers keep that clean close observable, but
+        pool-owned startup connections and reconnection probes must be rejected
+        so unusable connections cannot be stored or reported as host recovery.
+        """
+
+        class MaintenanceModeCqlServer(object):
+            def __init__(self):
+                self._sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+                self._sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+                self._sock.bind(('127.0.0.1', 0))
+                self._sock.listen(1)
+                self._sock.settimeout(2)
+                self.port = self._sock.getsockname()[1]
+                self.first_frame = b''
+                self.ready = Event()
+                self.received_frame = Event()
+                self.error = None
+                self.thread = Thread(target=self._run)
+                self.thread.daemon = True
+                self.thread.start()
+
+            def _run(self):
+                self.ready.set()
+                try:
+                    client, _ = self._sock.accept()
+                    with client:
+                        client.settimeout(2)
+                        while len(self.first_frame) < 9:
+                            chunk = client.recv(9 - len(self.first_frame))
+                            if not chunk:
+                                break
+                            self.first_frame += chunk
+                except Exception as exc:
+                    self.error = exc
+                finally:
+                    self.received_frame.set()
+
+            def close(self):
+                self._sock.close()
+                self.thread.join(2)
+
+        class MaintenanceModeConnection(Connection):
+            def __init__(self, *args, **kwargs):
+                super(MaintenanceModeConnection, self).__init__(*args, **kwargs)
+                self._reader = None
+                self._connect_socket()
+                self._send_options_message()
+                self._reader = Thread(target=self._read_until_server_closes)
+                self._reader.daemon = True
+                self._reader.start()
+
+            def push(self, data):
+                self._socket.sendall(data)
+
+            def close(self):
+                with self.lock:
+                    if self.is_closed:
+                        return
+                    self.is_closed = True
+
+                if self._socket:
+                    self._socket.close()
+
+                if not self.is_defunct:
+                    self.error_all_requests(
+                        ConnectionShutdown("Connection to %s was closed" % self.endpoint))
+                    self.connected_event.set()
+
+            def _read_until_server_closes(self):
+                try:
+                    while True:
+                        data = self._socket.recv(self.in_buffer_size)
+                        if not data:
+                            self.close()
+                            return
+                        self._iobuf.write(data)
+                        self.process_io_buffer()
+                except socket.error as exc:
+                    if not self.is_closed:
+                        self.defunct(exc)
+
+        server = MaintenanceModeCqlServer()
+        try:
+            assert server.ready.wait(2)
+            conn = MaintenanceModeConnection.factory(
+                DefaultEndPoint('127.0.0.1', server.port), timeout=2)
+
+            assert conn.is_closed
+            assert server.received_frame.wait(2)
+            assert server.error is None
+            assert server.first_frame[4] == 0x05  # OPTIONS
+        finally:
+            server.close()
+
+        server = MaintenanceModeCqlServer()
+        try:
+            assert server.ready.wait(2)
+            host_conn = Mock()
+            host_conn.is_shutdown = False
+            host_conn._pending_connections = []
+
+            with pytest.raises(ConnectionShutdown) as exc_info:
+                MaintenanceModeConnection.factory(
+                    DefaultEndPoint('127.0.0.1', server.port),
+                    timeout=2,
+                    host_conn=host_conn)
+
+            assert "closed during the startup handshake" in str(exc_info.value)
+            assert server.received_frame.wait(2)
+            assert server.error is None
+            assert server.first_frame[4] == 0x05  # OPTIONS
+            assert len(host_conn._pending_connections) == 1
+            assert host_conn._pending_connections[0].is_closed
+        finally:
+            server.close()
+
+        server = MaintenanceModeCqlServer()
+        try:
+            assert server.ready.wait(2)
+
+            with pytest.raises(ConnectionShutdown) as exc_info:
+                MaintenanceModeConnection.factory(
+                    DefaultEndPoint('127.0.0.1', server.port),
+                    timeout=2,
+                    _raise_on_startup_close=True)
+
+            assert "closed during the startup handshake" in str(exc_info.value)
+            assert server.received_frame.wait(2)
+            assert server.error is None
+            assert server.first_frame[4] == 0x05  # OPTIONS
+        finally:
+            server.close()
+
 
 @patch('cassandra.connection.ConnectionHeartbeat._raise_if_stopped')
 class ConnectionHeartbeatTest(unittest.TestCase):

tests/unit/test_host_connection_pool.py

@@ -50,7 +50,10 @@ def test_borrow_and_return(self):
         session.cluster.connection_factory.return_value = conn
 
         pool = self.PoolImpl(host, HostDistance.LOCAL, session)
-        session.cluster.connection_factory.assert_called_once_with(host.endpoint, on_orphaned_stream_released=pool.on_orphaned_stream_released)
+        session.cluster.connection_factory.assert_called_once_with(
+            host.endpoint,
+            _raise_on_startup_close=True,
+            on_orphaned_stream_released=pool.on_orphaned_stream_released)
 
         c, request_id = pool.borrow_connection(timeout=0.01)
         assert c is conn
@@ -69,7 +72,10 @@ def test_failed_wait_for_connection(self):
         session.cluster.connection_factory.return_value = conn
 
         pool = self.PoolImpl(host, HostDistance.LOCAL, session)
-        session.cluster.connection_factory.assert_called_once_with(host.endpoint, on_orphaned_stream_released=pool.on_orphaned_stream_released)
+        session.cluster.connection_factory.assert_called_once_with(
+            host.endpoint,
+            _raise_on_startup_close=True,
+            on_orphaned_stream_released=pool.on_orphaned_stream_released)
 
         pool.borrow_connection(timeout=0.01)
         assert 1 == conn.in_flight
@@ -89,7 +95,10 @@ def test_successful_wait_for_connection(self):
         session.cluster.connection_factory.return_value = conn
 
         pool = self.PoolImpl(host, HostDistance.LOCAL, session)
-        session.cluster.connection_factory.assert_called_once_with(host.endpoint, on_orphaned_stream_released=pool.on_orphaned_stream_released)
+        session.cluster.connection_factory.assert_called_once_with(
+            host.endpoint,
+            _raise_on_startup_close=True,
+            on_orphaned_stream_released=pool.on_orphaned_stream_released)
 
         pool.borrow_connection(timeout=0.01)
         assert 1 == conn.in_flight
@@ -114,7 +123,10 @@ def test_spawn_when_at_max(self):
         session.cluster.connection_factory.return_value = conn
 
         pool = self.PoolImpl(host, HostDistance.LOCAL, session)
-        session.cluster.connection_factory.assert_called_once_with(host.endpoint, on_orphaned_stream_released=pool.on_orphaned_stream_released)
+        session.cluster.connection_factory.assert_called_once_with(
+            host.endpoint,
+            _raise_on_startup_close=True,
+            on_orphaned_stream_released=pool.on_orphaned_stream_released)
 
         pool.borrow_connection(timeout=0.01)
         assert 1 == conn.in_flight
@@ -138,7 +150,10 @@ def test_return_defunct_connection(self):
         session.cluster.connection_factory.return_value = conn
 
         pool = self.PoolImpl(host, HostDistance.LOCAL, session)
-        session.cluster.connection_factory.assert_called_once_with(host.endpoint, on_orphaned_stream_released=pool.on_orphaned_stream_released)
+        session.cluster.connection_factory.assert_called_once_with(
+            host.endpoint,
+            _raise_on_startup_close=True,
+            on_orphaned_stream_released=pool.on_orphaned_stream_released)
 
         pool.borrow_connection(timeout=0.01)
         conn.is_defunct = True
@@ -160,7 +175,10 @@ def test_return_defunct_connection_on_down_host(self):
         session.cluster.shard_aware_options = ShardAwareOptions()
 
         pool = self.PoolImpl(host, HostDistance.LOCAL, session)
-        session.cluster.connection_factory.assert_called_once_with(host.endpoint, on_orphaned_stream_released=pool.on_orphaned_stream_released)
+        session.cluster.connection_factory.assert_called_once_with(
+            host.endpoint,
+            _raise_on_startup_close=True,
+            on_orphaned_stream_released=pool.on_orphaned_stream_released)
 
         pool.borrow_connection(timeout=0.01)
         conn.is_defunct = True
@@ -187,7 +205,10 @@ def test_return_closed_connection(self):
         session.cluster.connection_factory.return_value = conn
 
         pool = self.PoolImpl(host, HostDistance.LOCAL, session)
-        session.cluster.connection_factory.assert_called_once_with(host.endpoint, on_orphaned_stream_released=pool.on_orphaned_stream_released)
+        session.cluster.connection_factory.assert_called_once_with(
+            host.endpoint,
+            _raise_on_startup_close=True,
+            on_orphaned_stream_released=pool.on_orphaned_stream_released)
 
         pool.borrow_connection(timeout=0.01)
         conn.is_closed = True
@@ -230,6 +251,32 @@ class HostConnectionTests(_PoolTests):
     PoolImpl = HostConnection
     uses_single_connection = True
 
+    def test_replace_uses_startup_close_flag_without_pending_registration(self):
+        host = Mock(spec=Host, address='ip1')
+        host.sharding_info = None
+        session = self.make_session()
+        first_conn = HashableMock(spec=Connection, in_flight=0, is_defunct=False,
+                                  is_closed=False, max_request_id=100)
+        first_conn.features = ProtocolFeatures(shard_id=0)
+        replacement_conn = HashableMock(spec=Connection, in_flight=0, is_defunct=False,
+                                        is_closed=False, max_request_id=100)
+        replacement_conn.features = ProtocolFeatures(shard_id=0)
+        session.cluster.connection_factory.side_effect = [first_conn, replacement_conn]
+
+        pool = self.PoolImpl(host, HostDistance.LOCAL, session)
+        session.cluster.connection_factory.reset_mock()
+        pool._is_replacing = True
+
+        pool._replace(first_conn)
+
+        session.cluster.connection_factory.assert_called_once_with(
+            host.endpoint,
+            _raise_on_startup_close=True,
+            on_orphaned_stream_released=pool.on_orphaned_stream_released)
+        assert pool._pending_connections == []
+        assert pool._connections[0] is replacement_conn
+        assert not pool._is_replacing
+
     def test_fast_shutdown(self):
         class MockSession(MagicMock):
             is_shutdown = False