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

@@ -962,34 +962,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,41 @@ 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
+        return failure
 
     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 +312,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):
         """

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,88 @@ 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
+        replacement_is_pending = False
+        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)
+
+                # Connection.factory() stops tracking the connection before it
+                # returns. Adopt it as pending before doing the blocking
+                # keyspace setup so shutdown can still close it. The shutdown
+                # check and adoption must be atomic to cover the handoff gap.
+                with self._lock:
+                    if self.is_shutdown:
+                        close_replacement = True
+                    else:
+                        close_replacement = False
+                        self._pending_connections.append(replacement)
+                        replacement_is_pending = True
+
+                if close_replacement:
+                    replacement.close()
+                    return
+
+                if self._keyspace:
+                    replacement.set_keyspace_blocking(self._keyspace)
+
+                with self._lock:
+                    if self.is_shutdown:
+                        # shutdown() cleared the pending list and owns closing
+                        # the replacement.
+                        return
+                    else:
+                        self._pending_connections.remove(replacement)
+                        replacement_is_pending = False
+                        self._connections[replacement.features.shard_id] = replacement
+                        replace_finished = True
+        except Exception:
+            shutdown_owns_replacement = False
+            with self._lock:
+                if replacement_is_pending:
+                    try:
+                        self._pending_connections.remove(replacement)
+                    except ValueError:
+                        # shutdown() already drained and will close it.
+                        shutdown_owns_replacement = self.is_shutdown
+                is_shutdown = self.is_shutdown
+
+            if replacement is not None and not shutdown_owns_replacement:
+                replacement.close()
+
+            log.warning("Failed reconnecting %s. Retrying." % (self.host.endpoint,))
+            if 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 +786,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",