connection: tolerate clean startup close

cassandra/cluster.py

@@ -53,7 +53,8 @@
 from cassandra.connection import (ClientRoutesEndPointFactory, ConnectionException, ConnectionShutdown,
                                   ConnectionHeartbeat, ProtocolVersionUnsupported,
                                   EndPoint, DefaultEndPoint, DefaultEndPointFactory,
-                                  SniEndPointFactory, ConnectionBusy, locally_supported_compressions)
+                                  SniEndPointFactory, ConnectionBusy, locally_supported_compressions,
+                                  _startup_close_error)
 from cassandra.cqltypes import UserType
 import cassandra.cqltypes as types
 from cassandra.encoder import Encoder
@@ -1746,14 +1747,26 @@ def add_execution_profile(self, name, profile, pool_wait_timeout=5):
     def connection_factory(self, endpoint, host_conn = None, *args, **kwargs):
         """
         Called to create a new connection with proper configuration.
+
+        This preserves ``Connection.factory``'s return contract, including
+        a closed connection returned after a clean server close during startup.
+        Cluster-owned callers validate the result before adopting it.
+
         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)
 
     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)
+        # Keep the raw factory result observable. The host reconnector owns the
+        # decision to park a clean startup close.
+        return partial(
+            self.connection_class.factory,
+            host.endpoint,
+            self.connect_timeout,
+            *args,
+            **kwargs)
 
     def _make_connection_kwargs(self, endpoint, kwargs_dict):
         if self._auth_provider_callable:
@@ -2074,9 +2087,14 @@ def on_down_potentially_blocking(self, host, is_host_addition):
 
         self._start_reconnector(host, is_host_addition)
 
-    def on_down(self, host, is_host_addition, expect_host_to_be_down=False):
+    def on_down(self, host, is_host_addition, expect_host_to_be_down=False,
+                force=False):
         """
         Intended for internal use only.
+
+        ``force`` bypasses the open-pool discount when the caller has
+        authoritative evidence that the host cannot accept new CQL
+        connections.
         """
         if self.is_shutdown or self.allow_control_connection_query_fallback == ControlConnectionQueryFallback.SkipPoolCreation:
             return
@@ -2086,7 +2104,8 @@ def on_down(self, host, is_host_addition, expect_host_to_be_down=False):
 
             # ignore down signals if we have open pools to the host
             # this is to avoid closing pools when a control connection host became isolated
-            if self._discount_down_events and self.profile_manager.distance(host) != HostDistance.IGNORED:
+            if not force and self._discount_down_events and \
+                    self.profile_manager.distance(host) != HostDistance.IGNORED:
                 connected = False
                 for session in tuple(self.sessions):
                     pool_states = session.get_pool_state()
@@ -2190,10 +2209,17 @@ def on_remove(self, host):
         if reconnection_handler:
             reconnection_handler.cancel()
 
-    def signal_connection_failure(self, host, connection_exc, is_host_addition, expect_host_to_be_down=False):
+    def signal_connection_failure(self, host, connection_exc, is_host_addition,
+                                  expect_host_to_be_down=False, force=False):
+        # ``force`` is reserved for authoritative failures that cannot leave a
+        # usable pool, and bypasses both conviction and open-pool discounting.
         is_down = host.signal_connection_failure(connection_exc)
-        if is_down:
-            self.on_down(host, is_host_addition, expect_host_to_be_down)
+        if is_down or force:
+            self.on_down(
+                host,
+                is_host_addition,
+                expect_host_to_be_down,
+                force=force)
         return is_down
 
     def add_host(self, endpoint, datacenter=None, rack=None, signal=True, refresh_nodes=True, host_id=None):
@@ -2412,6 +2438,8 @@ def _prepare_all_queries(self, host):
         connection = None
         try:
             connection = self.connection_factory(host.endpoint)
+            if connection.is_closed:
+                raise _startup_close_error(connection, host.endpoint)
             statements = list(self._prepared_statements.values())
             if ProtocolVersion.uses_keyspace_flag(self.protocol_version):
                 # V5 protocol and higher, no need to set the keyspace
@@ -3822,6 +3850,15 @@ def _set_new_connection(self, conn):
             log.debug("[control connection] Closing old connection %r, replacing with %r", old, conn)
             old.close()
 
+        # A successful control reconnection is authoritative evidence that its
+        # host can serve CQL again. This also resumes hosts whose regular
+        # reconnector was parked after a clean maintenance-mode startup close.
+        if old is not None:
+            host = self._cluster.metadata.get_host(conn.endpoint)
+            if host is not None and not host.is_up:
+                self._cluster.scheduler.schedule_unique(
+                    0, self._cluster.on_up, host)
+
     def _try_connect_to_hosts(self):
         errors = {}
 
@@ -3876,6 +3913,8 @@ def _try_connect(self, endpoint):
                 if self._is_shutdown:
                     connection.close()
                     raise DriverException("Reconnecting during shutdown")
+                if connection.is_closed:
+                    raise _startup_close_error(connection, endpoint)
                 break
             except ProtocolVersionUnsupported as e:
                 self._cluster.protocol_downgrade(endpoint, e.startup_version)

cassandra/connection.py

@@ -533,6 +533,24 @@ class ConnectionShutdown(ConnectionException):
     pass
 
 
+class _ConnectionClosedDuringStartup(ConnectionShutdown):
+    """
+    Internal owner-side signal for a clean close returned by
+    ``Connection.factory``.
+    """
+    pass
+
+
+def _startup_close_error(connection, endpoint=None):
+    """Build an owner-side error without changing the factory return contract."""
+    connection_endpoint = getattr(connection, 'endpoint', None)
+    if connection_endpoint is not None:
+        endpoint = connection_endpoint
+    return _ConnectionClosedDuringStartup(
+        "Connection to %s was closed during the startup handshake" % (endpoint,),
+        endpoint)
+
+
 class ProtocolVersionUnsupported(ConnectionException):
     """
     Server rejected startup message due to unsupported protocol version
@@ -962,34 +980,48 @@ def handle_fork(cls):
     def create_timer(cls, timeout, callback):
         raise NotImplementedError()
 
+    def _is_clean_close_error(self, exc):
+        return not self.is_defunct and isinstance(exc, ConnectionShutdown)
+
     @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 conn._is_clean_close_error(conn.last_error):
+                    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/io/twistedreactor.py

@@ -24,6 +24,7 @@
 
 from twisted.internet import reactor, protocol
 from twisted.internet.endpoints import connectProtocol, TCP4ClientEndpoint, SSL4ClientEndpoint
+from twisted.internet.error import ConnectionDone
 from twisted.internet.interfaces import IOpenSSLClientConnectionCreator
 from twisted.python.failure import Failure
 from zope.interface import implementer
@@ -198,6 +199,9 @@ def create_timer(cls, timeout, callback):
         cls._loop.add_timer(timer)
         return timer
 
+    def _is_clean_close_error(self, exc):
+        return isinstance(exc, ConnectionDone) or Connection._is_clean_close_error(self, exc)
+
     def __init__(self, *args, **kwargs):
         """
         Initialization method.
@@ -209,7 +213,10 @@ def __init__(self, *args, **kwargs):
         """
         Connection.__init__(self, *args, **kwargs)
 
-        self.is_closed = True
+        # A scheduled or in-progress endpoint connection is still closeable.
+        # Keeping it logically open lets pool shutdown cancel it before
+        # connectionMade.
+        self.is_closed = False
         self.connector = None
         self.transport = None
 
@@ -226,9 +233,12 @@ def _check_pyopenssl(self):
 
     def add_connection(self):
         """
-        Convenience function to connect and store the resulting
-        connector.
+        Convenience function to connect and store the resulting Deferred.
         """
+        with self.lock:
+            if self.is_closed:
+                return
+
         host, port = self.endpoint.resolve()
         if self.ssl_context or self.ssl_options:
             # Can't use optionsForClientTLS here because it *forces* hostname verification.
@@ -257,16 +267,42 @@ def add_connection(self):
                 port,
                 timeout=self.connect_timeout
             )
-        connectProtocol(endpoint, TwistedConnectionProtocol(self))
+
+        # connectProtocol returns a cancellable Deferred. Keep the final
+        # shutdown check and assignment under the lock so close() either stops
+        # this attempt before it starts or observes the Deferred and cancels it.
+        with self.lock:
+            if self.is_closed:
+                return
+            self.connector = connectProtocol(
+                endpoint, TwistedConnectionProtocol(self))
+            self.connector.addErrback(self._handle_connect_failure)
+
+    def _handle_connect_failure(self, failure):
+        # Pool shutdown intentionally cancels the endpoint Deferred. Consume
+        # that failure instead of leaving an unhandled cancellation in Twisted.
+        with self.lock:
+            if self.is_closed:
+                return None
+        self.defunct(failure.value)
+        return None
 
     def client_connection_made(self, transport):
         """
         Called by twisted protocol when a connection attempt has
         succeeded.
         """
         with self.lock:
-            self.is_closed = False
-        self.transport = transport
+            if self.is_closed:
+                close_transport = True
+            else:
+                close_transport = False
+                self.transport = transport
+
+        if close_transport:
+            reactor.callFromThread(transport.connector.disconnect)
+            return
+
         self._send_options_message()
 
     def close(self):
@@ -277,18 +313,29 @@ def close(self):
             if self.is_closed:
                 return
             self.is_closed = True
+            connector = self.connector
+            transport = self.transport
+
+            shutdown_error = None
+            if not self.is_defunct:
+                msg = "Connection to %s was closed" % self.endpoint
+                if self.last_error:
+                    msg += ": %s" % (self.last_error,)
+                shutdown_error = ConnectionShutdown(msg)
+                if not self.connected_event.is_set():
+                    self.last_error = shutdown_error
+                # Wake Connection.factory before waiting for reactor cleanup.
+                self.connected_event.set()
 
         log.debug("Closing connection (%s) to %s", id(self), self.endpoint)
-        reactor.callFromThread(self.transport.connector.disconnect)
+        if transport is not None:
+            reactor.callFromThread(transport.connector.disconnect)
+        elif connector is not None:
+            reactor.callFromThread(connector.cancel)
         log.debug("Closed socket to %s", self.endpoint)
 
-        if not self.is_defunct:
-            msg = "Connection to %s was closed" % self.endpoint
-            if self.last_error:
-                msg += ": %s" % (self.last_error,)
-            self.error_all_requests(ConnectionShutdown(msg))
-            # don't leave in-progress operations hanging
-            self.connected_event.set()
+        if shutdown_error is not None:
+            self.error_all_requests(shutdown_error)
 
     def handle_read(self):
         """