Merge pull request #297 from sylwiaszunejko/shutdown_hanging

Close pending connections during shutdown

Author: avelanarius

cassandra/cluster.py

@@ -1691,13 +1691,13 @@ def set_max_connections_per_host(self, host_distance, max_connections):
                 "when using protocol_version 1 or 2.")
         self._max_connections_per_host[host_distance] = max_connections
 
-    def connection_factory(self, endpoint, *args, **kwargs):
+    def connection_factory(self, endpoint, host_conn = None, *args, **kwargs):
         """
         Called to create a new connection with proper configuration.
         Intended for internal use only.
         """
         kwargs = self._make_connection_kwargs(endpoint, kwargs)
-        return self.connection_class.factory(endpoint, self.connect_timeout, *args, **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)

cassandra/connection.py

@@ -865,7 +865,7 @@ def create_timer(cls, timeout, callback):
         raise NotImplementedError()
 
     @classmethod
-    def factory(cls, endpoint, timeout, *args, **kwargs):
+    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
@@ -874,6 +874,10 @@ def factory(cls, endpoint, timeout, *args, **kwargs):
         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:
+                conn.close()
         elapsed = time.time() - start
         conn.connected_event.wait(timeout - elapsed)
         if conn.last_error:

cassandra/pool.py

@@ -404,6 +404,7 @@ def __init__(self, host, host_distance, session):
         self._is_replacing = False
         self._connecting = set()
         self._connections = {}
+        self._pending_connections = []
         # A pool of additional connections which are not used but affect how Scylla
         # assigns shards to them. Scylla tends to assign the shard which has
         # the lowest number of connections. If connections are not distributed
@@ -638,7 +639,9 @@ def shutdown(self):
                 future.cancel()
 
             connections_to_close = self._connections.copy()
+            pending_connections_to_close = self._pending_connections.copy()
             self._connections.clear()
+            self._pending_connections.clear()
 
         # connection.close can call pool.return_connection, which will
         #  obtain self._lock via self._stream_available_condition.
@@ -647,6 +650,10 @@ def shutdown(self):
             log.debug("Closing connection (%s) to %s", id(connection), self.host)
             connection.close()
 
+        for connection in pending_connections_to_close:
+            log.debug("Closing pending connection (%s) to %s", id(connection), self.host)
+            connection.close()
+
         self._close_excess_connections()
 
         trash_conns = None
@@ -714,12 +721,12 @@ def _open_connection_to_missing_shard(self, shard_id):
         log.debug("shard_aware_endpoint=%r", shard_aware_endpoint)
 
         if shard_aware_endpoint:
-            conn = self._session.cluster.connection_factory(shard_aware_endpoint, on_orphaned_stream_released=self.on_orphaned_stream_released,
+            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
         else:
-            conn = self._session.cluster.connection_factory(self.host.endpoint, 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", id(conn), conn.features.shard_id, self.host)
         if self.is_shutdown: