connection: tolerate clean startup close

Author: dkropachev

cassandra/cluster.py

@@ -345,6 +345,14 @@ def _execution_profile_to_string(name):
     return '"%s"' % (name,)
 
 
+def _raise_if_connection_closed_during_startup(conn, endpoint):
+    if conn.is_closed:
+        endpoint = getattr(conn, 'endpoint', endpoint)
+        raise ConnectionShutdown(
+            "Connection to %s was closed during the startup handshake" % (endpoint,),
+            endpoint)
+
+
 class ExecutionProfile(object):
     load_balancing_policy = None
     """
@@ -1749,11 +1757,22 @@ def connection_factory(self, endpoint, host_conn = None, *args, **kwargs):
         Intended for internal use only.
         """
         kwargs = self._make_connection_kwargs(endpoint, kwargs)
-        return self.connection_class.factory(endpoint, self.connect_timeout, host_conn, *args, **kwargs)
+        conn = self.connection_class.factory(endpoint, self.connect_timeout, host_conn, *args, **kwargs)
+        _raise_if_connection_closed_during_startup(conn, endpoint)
+        return conn
 
     def _make_connection_factory(self, host, *args, **kwargs):
-        kwargs = self._make_connection_kwargs(host.endpoint, kwargs)
-        return partial(self.connection_class.factory, host.endpoint, self.connect_timeout, *args, **kwargs)
+        endpoint = host.endpoint
+        connection_class = self.connection_class
+        connect_timeout = self.connect_timeout
+        kwargs = self._make_connection_kwargs(endpoint, kwargs)
+
+        def connection_factory():
+            conn = connection_class.factory(endpoint, connect_timeout, *args, **kwargs)
+            _raise_if_connection_closed_during_startup(conn, endpoint)
+            return conn
+
+        return connection_factory
 
     def _make_connection_kwargs(self, endpoint, kwargs_dict):
         if self._auth_provider_callable:

cassandra/connection.py

@@ -965,31 +965,42 @@ 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, returns a closed connection if the server closes during
+        startup, or raises an exception otherwise.
+
+        Pool callers pass ``host_conn`` so the connection can be tracked while
+        startup is in progress and closed during pool shutdown.
         """
         start = time.time()
         kwargs['connect_timeout'] = timeout
         conn = cls(endpoint, *args, **kwargs)
-        if host_conn is not None:
-            host_conn._pending_connections.append(conn)
-            if host_conn.is_shutdown:
+        pending_registered = False
+        try:
+            if host_conn is not None:
+                host_conn._pending_connections.append(conn)
+                pending_registered = True
+                if host_conn.is_shutdown:
+                    conn.close()
+            elapsed = time.time() - start
+            conn.connected_event.wait(timeout - elapsed)
+            if conn.last_error:
+                if conn.is_unsupported_proto_version:
+                    raise ProtocolVersionUnsupported(endpoint, conn.protocol_version)
+                if conn.is_closed and not conn.is_defunct and isinstance(conn.last_error, ConnectionShutdown):
+                    return conn
+                raise conn.last_error
+            elif not conn.connected_event.is_set():
                 conn.close()
-        elapsed = time.time() - start
-        conn.connected_event.wait(timeout - elapsed)
-        if conn.last_error:
-            if conn.is_unsupported_proto_version:
-                raise ProtocolVersionUnsupported(endpoint, conn.protocol_version)
-            raise conn.last_error
-        elif not conn.connected_event.is_set():
-            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)
-        else:
+                raise OperationTimedOut("Timed out creating connection (%s seconds)" % timeout,
+                                        timeout=timeout)
             return conn
+        finally:
+            if pending_registered:
+                try:
+                    host_conn._pending_connections.remove(conn)
+                except ValueError:
+                    pass
 
     def _build_ssl_context_from_options(self):
 

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,
+            host_conn=self,
+            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
@@ -593,30 +596,60 @@ def on_orphaned_stream_released(self):
             self._stream_available_condition.notify()
 
     def _replace(self, connection):
-        with self._lock:
-            if self.is_shutdown:
-                return
+        replacement = None
+        replace_finished = False
 
-            log.debug("Replacing connection (%s) to %s", id(connection), self.host)
-            try:
+        try:
+            with self._lock:
+                if self.is_shutdown:
+                    return
+
+                log.debug("Replacing connection (%s) to %s", id(connection), self.host)
                 if connection.features.shard_id in self._connections:
                     del self._connections[connection.features.shard_id]
-                if self.host.sharding_info and not self._session.cluster.shard_aware_options.disable:
+
+                needs_replacement = not (
+                    self.host.sharding_info and not self._session.cluster.shard_aware_options.disable)
+
+                if not needs_replacement:
                     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)
-                    if self._keyspace:
-                        connection.set_keyspace_blocking(self._keyspace)
-                    self._connections[connection.features.shard_id] = connection
-            except Exception:
-                log.warning("Failed reconnecting %s. Retrying." % (self.host.endpoint,))
-                self._session.submit(self._replace, connection)
-            else:
-                self._is_replacing = False
-                with self._stream_available_condition:
-                    self._stream_available_condition.notify()
+                    replace_finished = True
+
+            if needs_replacement:
+                replacement = self._session.cluster.connection_factory(
+                    self.host.endpoint,
+                    host_conn=self,
+                    on_orphaned_stream_released=self.on_orphaned_stream_released)
+                if self._keyspace:
+                    replacement.set_keyspace_blocking(self._keyspace)
+
+                with self._lock:
+                    if self.is_shutdown:
+                        close_replacement = True
+                    else:
+                        close_replacement = False
+                        self._connections[replacement.features.shard_id] = replacement
+                        replacement = None
+                        replace_finished = True
+
+                if close_replacement:
+                    replacement.close()
+                    return
+        except Exception:
+            log.warning("Failed reconnecting %s. Retrying." % (self.host.endpoint,))
+            with self._lock:
+                if self.is_shutdown:
+                    return
+            self._session.submit(self._replace, connection)
+        else:
+            if replace_finished:
+                with self._lock:
+                    if self.is_shutdown:
+                        return
+                    self._is_replacing = False
+                    with self._stream_available_condition:
+                        self._stream_available_condition.notify()
 
     def shutdown(self):
         log.debug("Shutting down connections to %s", self.host)
@@ -725,15 +758,21 @@ def _open_connection_to_missing_shard(self, shard_id):
         log.debug("shard_aware_endpoint=%r", shard_aware_endpoint)
         if shard_aware_endpoint:
             try:
-                conn = self._session.cluster.connection_factory(shard_aware_endpoint, host_conn=self, on_orphaned_stream_released=self.on_orphaned_stream_released,
-                                                                shard_id=shard_id,
-                                                                total_shards=self.host.sharding_info.shards_count)
+                conn = self._session.cluster.connection_factory(
+                    shard_aware_endpoint,
+                    host_conn=self,
+                    on_orphaned_stream_released=self.on_orphaned_stream_released,
+                    shard_id=shard_id,
+                    total_shards=self.host.sharding_info.shards_count)
                 conn.original_endpoint = self.host.endpoint
             except Exception as exc:
                 log.error("Failed to open connection to %s, on shard_id=%i: %s", self.host, shard_id, exc)
                 raise
         else:
-            conn = self._session.cluster.connection_factory(self.host.endpoint, host_conn=self, on_orphaned_stream_released=self.on_orphaned_stream_released)
+            conn = self._session.cluster.connection_factory(
+                self.host.endpoint,
+                host_conn=self,
+                on_orphaned_stream_released=self.on_orphaned_stream_released)
 
         log.debug(
             "Received a connection %s for shard_id=%i on host %s",

tests/integration/standard/test_maintenance_mode_connection.py

@@ -0,0 +1,152 @@
+# Copyright DataStax, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import socket
+import unittest
+from threading import Event, Thread
+
+import pytest
+
+from cassandra.cluster import NoHostAvailable
+from cassandra.connection import ConnectionShutdown, DefaultEndPoint
+from tests.integration import TestCluster
+
+
+class MaintenanceModeCqlServer(object):
+    """
+    Minimal CQL listener that accepts a startup attempt and then closes the
+    socket without replying, matching Scylla's maintenance-mode failure shape.
+    """
+
+    def __init__(self, max_connections=1):
+        self._closed = False
+        self._max_connections = max_connections
+        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(max_connections)
+        self._sock.settimeout(0.2)
+        self.port = self._sock.getsockname()[1]
+        self.frames = []
+        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:
+            while len(self.frames) < self._max_connections and not self._closed:
+                try:
+                    client, _ = self._sock.accept()
+                except socket.timeout:
+                    continue
+
+                with client:
+                    client.settimeout(2)
+                    frame = b''
+                    while len(frame) < 9:
+                        chunk = client.recv(9 - len(frame))
+                        if not chunk:
+                            break
+                        frame += chunk
+                    self.frames.append(frame)
+                    self.received_frame.set()
+        except Exception as exc:
+            if not self._closed:
+                self.error = exc
+        finally:
+            self.received_frame.set()
+
+    def close(self):
+        self._closed = True
+        try:
+            self._sock.close()
+        except socket.error:
+            # Best-effort test cleanup; close can race with the listener thread.
+            pass
+        self.thread.join(2)
+
+
+class MaintenanceModeConnectionTest(unittest.TestCase):
+
+    def setUp(self):
+        self.cluster = TestCluster(contact_points=[], connect_timeout=2)
+        self.cluster.connection_class.initialize_reactor()
+
+    def tearDown(self):
+        self.cluster.shutdown()
+
+    def test_startup_close_is_observable_low_level_but_rejected_by_cluster_factory(self):
+        """
+        Exercise the real reactor/factory path against a socket that behaves
+        like a node rejecting regular CQL traffic while in maintenance mode.
+        """
+        endpoint, server = self._new_endpoint_and_server()
+        try:
+            conn = self.cluster.connection_class.factory(
+                endpoint,
+                self.cluster.connect_timeout,
+                **self.cluster._make_connection_kwargs(endpoint, {}))
+
+            assert conn.is_closed
+            self._assert_server_saw_options_frame(server)
+        finally:
+            server.close()
+
+        endpoint, server = self._new_endpoint_and_server()
+        try:
+            with pytest.raises(ConnectionShutdown) as exc_info:
+                self.cluster.connection_factory(endpoint)
+
+            assert "closed during the startup handshake" in str(exc_info.value)
+            self._assert_server_saw_options_frame(server)
+        finally:
+            server.close()
+
+    def test_cluster_connect_reports_startup_close_as_unavailable_host(self):
+        endpoint, server = self._new_endpoint_and_server(max_connections=4)
+        cluster = TestCluster(
+            contact_points=[endpoint],
+            connect_timeout=2,
+            control_connection_timeout=2)
+        cluster.connection_class.initialize_reactor()
+
+        try:
+            with pytest.raises(NoHostAvailable) as exc_info:
+                cluster.connect()
+
+            errors = exc_info.value.errors
+            assert errors
+            assert any(isinstance(exc, ConnectionShutdown) for exc in errors.values())
+            assert any("closed during the startup handshake" in str(exc)
+                       for exc in errors.values())
+            self._assert_server_saw_options_frame(server)
+        finally:
+            cluster.shutdown()
+            server.close()
+
+    def _new_endpoint_and_server(self, max_connections=1):
+        server = MaintenanceModeCqlServer(max_connections=max_connections)
+        assert server.ready.wait(2)
+        return DefaultEndPoint('127.0.0.1', server.port), server
+
+    def _assert_server_saw_options_frame(self, server):
+        assert server.received_frame.wait(2)
+        assert server.error is None
+        assert server.frames
+        assert len(server.frames[0]) >= 5
+        assert server.frames[0][4] == 0x05  # OPTIONS