session: refresh stale pool creation by endpoint

Author: dkropachev

cassandra/cluster.py

@@ -4073,46 +4073,46 @@ def callback(pool, errors):
                     else:
                         with host.lock:
                             endpoint_changed = not self._endpoints_match(host.endpoint, creation_endpoint)
-                        if endpoint_changed:
-                            log.debug(
-                                "Discarding stale connection pool for host %s; endpoint changed from %s",
-                                host, creation_endpoint)
-                            self._invalidate_pool_creation(host, expected_endpoint=creation_endpoint)
-                            discard_pool = True
-                        else:
-                            # Rebuild by identity so endpoint hash changes do not
-                            # leave stale pool entries behind.
-                            retained_pools = {}
-                            for pool_host, host_pool in self._pools.items():
-                                if pool_host is host:
-                                    previous_pools.append(host_pool)
-                                else:
-                                    retained_pools[pool_host] = host_pool
-
-                            # Keep the current metadata host keyed by identity.
-                            metadata_host = host
-                            if isinstance(self.cluster.metadata, Metadata):
-                                metadata_host = self.cluster.metadata.get_host_by_host_id(host.host_id)
-
-                            target_host = metadata_host if metadata_host is not None else host
-                            target_host_matches = False
-                            for pool_host in tuple(retained_pools):
-                                if pool_host is target_host:
-                                    target_host_matches = True
-                                elif pool_host == target_host:
-                                    previous_pools.append(retained_pools.pop(pool_host))
-
-                            if target_host_matches:
-                                reuse_existing_pool = True
+                            if endpoint_changed:
+                                log.debug(
+                                    "Discarding stale connection pool for host %s; endpoint changed from %s",
+                                    host, creation_endpoint)
+                                self._invalidate_pool_creation(host, expected_endpoint=creation_endpoint)
+                                discard_pool = True
                             else:
-                                source_host = new_pool.host
-                                if (source_host is not target_host and
-                                        target_host.sharding_info is None):
-                                    target_host.sharding_info = source_host.sharding_info
-                                new_pool.host = target_host
-                                retained_pools[target_host] = new_pool
-                            self._pools = retained_pools
-                            self._clear_pool_creation(host, creation_epoch)
+                                # Rebuild by identity so endpoint hash changes do not
+                                # leave stale pool entries behind.
+                                retained_pools = {}
+                                for pool_host, host_pool in self._pools.items():
+                                    if pool_host is host:
+                                        previous_pools.append(host_pool)
+                                    else:
+                                        retained_pools[pool_host] = host_pool
+
+                                # Keep the current metadata host keyed by identity.
+                                metadata_host = host
+                                if isinstance(self.cluster.metadata, Metadata):
+                                    metadata_host = self.cluster.metadata.get_host_by_host_id(host.host_id)
+
+                                target_host = metadata_host if metadata_host is not None else host
+                                target_host_matches = False
+                                for pool_host in tuple(retained_pools):
+                                    if pool_host is target_host:
+                                        target_host_matches = True
+                                    elif pool_host == target_host:
+                                        previous_pools.append(retained_pools.pop(pool_host))
+
+                                if target_host_matches:
+                                    reuse_existing_pool = True
+                                else:
+                                    source_host = new_pool.host
+                                    if (source_host is not target_host and
+                                            target_host.sharding_info is None):
+                                        target_host.sharding_info = source_host.sharding_info
+                                    new_pool.host = target_host
+                                    retained_pools[target_host] = new_pool
+                                self._pools = retained_pools
+                                self._clear_pool_creation(host, creation_epoch)
 
             if reuse_existing_pool:
                 log.debug("Reusing existing connection pool for host %s", host)
@@ -4141,7 +4141,13 @@ def callback(pool, errors):
         with self._lock:
             state = self._get_pool_creation_state(host)
             if state.creation_epoch is not None:
-                return state.future
+                with host.lock:
+                    endpoint_changed = not self._endpoints_match(
+                        host.endpoint, state.endpoint)
+                if not endpoint_changed:
+                    return state.future
+                self._invalidate_pool_creation(
+                    host, expected_endpoint=state.endpoint)
 
             creation_epoch = state.advance()
             state.creation_epoch = creation_epoch
@@ -4156,7 +4162,8 @@ def callback(pool, errors):
                 state.future = future
             return future
 
-    def remove_pool(self, host, expected_host=None, expected_endpoint=None):
+    def remove_pool(self, host, expected_host=None, expected_endpoint=None,
+                    expected_pool=None):
         removed_pools = []
         cleanup_context = _POOL_CLEANUP_EPOCH.get()
         with self._lock:
@@ -4187,11 +4194,12 @@ def remove_pool(self, host, expected_host=None, expected_endpoint=None):
                 if expected_endpoint is not None:
                     remove_all = self._endpoints_match(host.endpoint, expected_endpoint)
 
-                if remove_all:
-                    self._invalidate_pool_creation(host)
-                else:
-                    self._invalidate_pool_creation(
-                        host, expected_endpoint=expected_endpoint)
+                if expected_pool is None:
+                    if remove_all:
+                        self._invalidate_pool_creation(host)
+                    else:
+                        self._invalidate_pool_creation(
+                            host, expected_endpoint=expected_endpoint)
 
                 retained_pools = {}
                 for pool_host, host_pool in self._pools.items():
@@ -4201,7 +4209,8 @@ def remove_pool(self, host, expected_host=None, expected_endpoint=None):
 
                     matches = self._pool_matches_expected(
                         host_pool, expected_host=expected_host,
-                        expected_endpoint=None if remove_all else expected_endpoint)
+                        expected_endpoint=None if remove_all else expected_endpoint,
+                        expected_pool=expected_pool)
                     if matches:
                         removed_pools.append(host_pool)
                     else:
@@ -4230,7 +4239,10 @@ def _shutdown_removed_pools(pools):
         for pool in pools:
             pool.shutdown()
 
-    def _pool_matches_expected(self, pool, expected_host=None, expected_endpoint=None):
+    def _pool_matches_expected(self, pool, expected_host=None,
+                               expected_endpoint=None, expected_pool=None):
+        if expected_pool is not None and pool is not expected_pool:
+            return False
         if expected_host is not None and pool.host is not expected_host:
             return False
         if expected_endpoint is not None:

cassandra/pool.py

@@ -597,7 +597,8 @@ def return_connection(self, connection, stream_was_orphaned=False):
                 # Drop only this stale pool; endpoint reuse may already belong to
                 # a replacement host instance.
                 future = self._session.remove_pool(
-                    self.host, expected_host=self.host, expected_endpoint=self.endpoint)
+                    self.host, expected_host=self.host,
+                    expected_endpoint=self.endpoint, expected_pool=self)
                 if future:
                     future.add_done_callback(lambda f: self._session.update_created_pools())
                 with self._lock:
@@ -645,7 +646,8 @@ def on_orphaned_stream_released(self):
 
     def _remove_stale_pool(self, expected_endpoint):
         future = self._session.remove_pool(
-            self.host, expected_host=self.host, expected_endpoint=expected_endpoint)
+            self.host, expected_host=self.host,
+            expected_endpoint=expected_endpoint, expected_pool=self)
         if future:
             future.add_done_callback(lambda f: self._session.update_created_pools())
 

tests/unit/test_cluster.py

@@ -26,7 +26,7 @@
 from cassandra.cluster import _Scheduler, Session, Cluster, default_lbp_factory, \
     ExecutionProfile, _ConfigMode, EXEC_PROFILE_DEFAULT
 from cassandra.connection import ClientRoutesEndPoint, ConnectionException, DefaultEndPoint, SniEndPoint
-from cassandra.pool import Host, _HostReconnectionHandler
+from cassandra.pool import Host, HostConnection, _HostReconnectionHandler
 from cassandra.policies import HostDistance, RetryPolicy, RoundRobinPolicy, DowngradingConsistencyRetryPolicy, SimpleConvictionPolicy
 from cassandra.query import SimpleStatement, named_tuple_factory, tuple_factory
 from tests.unit.utils import mock_session_pools
@@ -323,6 +323,28 @@ def test_schedule_unique_keeps_client_route_events_for_distinct_ports(self, *_):
 
 class SessionPoolRaceTest(unittest.TestCase):
 
+    class _EndpointSwapIfPoolUnpublishedOnFirstExitLock(object):
+
+        def __init__(self, host, new_endpoint, pool_is_published):
+            self._lock = RLock()
+            self._host = host
+            self._new_endpoint = new_endpoint
+            self._pool_is_published = pool_is_published
+            self._exits = 0
+            self.pool_was_published_on_first_exit = None
+
+        def __enter__(self):
+            self._lock.acquire()
+            return self
+
+        def __exit__(self, exc_type, exc_value, traceback):
+            self._lock.release()
+            self._exits += 1
+            if self._exits == 1:
+                self.pool_was_published_on_first_exit = self._pool_is_published()
+                if not self.pool_was_published_on_first_exit:
+                    self._host.endpoint = self._new_endpoint
+
     class _DuplicatePoolEntries(object):
 
         def __init__(self, entries):
@@ -522,6 +544,64 @@ def make_pool(host, distance, pool_session, endpoint=None):
         assert session._pools == {}
         created_pools[0].shutdown.assert_called_once_with()
 
+    def test_add_or_renew_pool_invalidates_creation_after_endpoint_change(self):
+        host = self._make_host("127.0.0.1")
+        old_endpoint = host.endpoint
+        new_endpoint = DefaultEndPoint("127.0.0.2")
+        _, session, executor = self._make_cluster_and_session([host])
+        created_pools = []
+
+        def make_pool(host, distance, pool_session, endpoint=None):
+            pool = self._make_pool(host, distance, pool_session, endpoint)
+            created_pools.append(pool)
+            return pool
+
+        with patch("cassandra.cluster.HostConnection", side_effect=make_pool):
+            old_future = session.add_or_renew_pool(
+                host, is_host_addition=False)
+            host.endpoint = new_endpoint
+
+            new_future = session.add_or_renew_pool(
+                host, is_host_addition=False)
+
+            assert new_future is not old_future
+            assert len(executor.submissions) == 2
+
+            executor.run_next()
+            executor.run_next()
+
+        assert old_future.result() is False
+        assert new_future.result() is True
+        assert created_pools[0].endpoint == old_endpoint
+        assert created_pools[1].endpoint == new_endpoint
+        created_pools[0].shutdown.assert_called_once_with()
+        created_pools[1].shutdown.assert_not_called()
+        assert session._pools[host] is created_pools[1]
+
+    def test_pool_creation_publishes_before_endpoint_lock_is_released(self):
+        host = self._make_host("127.0.0.1")
+        new_endpoint = DefaultEndPoint("127.0.0.2")
+        cluster, session, executor = self._make_cluster_and_session([host])
+        created_pools = []
+
+        def make_pool(host, distance, pool_session, endpoint=None):
+            pool = self._make_pool(host, distance, pool_session, endpoint)
+            created_pools.append(pool)
+            return pool
+
+        with patch("cassandra.cluster.HostConnection", side_effect=make_pool):
+            future = session.add_or_renew_pool(
+                host, is_host_addition=False)
+            host.lock = self._EndpointSwapIfPoolUnpublishedOnFirstExitLock(
+                host, new_endpoint, lambda: bool(session._pools))
+
+            executor.run_next()
+
+        assert host.lock.pool_was_published_on_first_exit is True
+        assert future.result() is True
+        assert session._pools[host] is created_pools[0]
+        created_pools[0].shutdown.assert_not_called()
+
     def test_update_created_pools_removes_stale_pool_for_down_host_after_endpoint_change(self):
         host = self._make_host("127.0.0.1")
         host.set_down()
@@ -701,6 +781,28 @@ def make_pool(host, distance, pool_session, endpoint=None):
         assert created_pools[0].endpoint == host.endpoint
         created_pools[0].shutdown.assert_not_called()
 
+    def test_stale_host_connection_cleanup_after_endpoint_flip_back_preserves_current_pool(self):
+        host = self._make_host("127.0.0.1")
+        old_endpoint = host.endpoint
+        replacement_endpoint = DefaultEndPoint("127.0.0.2")
+        cluster, session, executor = self._make_cluster_and_session([host])
+
+        host.endpoint = replacement_endpoint
+        host.endpoint = old_endpoint
+        current_pool = self._make_pool(
+            host, HostDistance.LOCAL, session, endpoint=old_endpoint)
+        session._pools[host] = current_pool
+
+        stale_pool = HostConnection.__new__(HostConnection)
+        stale_pool.host = host
+        stale_pool.endpoint = old_endpoint
+        stale_pool._session = session
+
+        stale_pool._remove_stale_pool(old_endpoint)
+
+        assert session._get_pool_by_host_identity(host) is current_pool
+        current_pool.shutdown.assert_not_called()
+
     def test_add_or_renew_pool_tags_pool_with_creation_endpoint(self):
         host = self._make_host("127.0.0.1")
         old_endpoint = host.endpoint