session: avoid stale endpoint pool reuse

Author: dkropachev

cassandra/cluster.py

@@ -4303,6 +4303,7 @@ def update_created_pools(self):
             with self._lock:
                 with host.lock:
                     host_endpoint = host.endpoint
+                    host_is_up = host.is_up
                 pool = self._get_pool_by_host_identity(
                     host, expected_endpoint=host_endpoint)
                 any_pool = pool or self._get_pool_by_host_identity(host)
@@ -4320,8 +4321,10 @@ def update_created_pools(self):
                         # on_up() keeps host.is_up False until this future succeeds.
                         future = self._reuse_or_invalidate_pool_creation(
                             host, pool_creation_future)
-                    elif host.is_up in (True, None):
+                    elif host_is_up in (True, None):
                         future = self.add_or_renew_pool(host, False)
+                    elif any_pool is not None:
+                        future = self.remove_pool(host)
                 elif pool_creation_future is not None:
                     future = self._reuse_or_invalidate_pool_creation(
                         host, pool_creation_future)
@@ -5662,7 +5665,13 @@ def _query(self, host, message=None, cb=None):
         if message is None:
             message = self.message
 
-        pool = self.session._get_pool_by_host_identity(host)
+        if isinstance(host, Host):
+            with host.lock:
+                expected_endpoint = host.endpoint
+            pool = self.session._get_pool_by_host_identity(
+                host, expected_endpoint=expected_endpoint)
+        else:
+            pool = self.session._get_pool_by_host_identity(host)
         if not pool:
             self._errors[host] = ConnectionException("Host has been marked down or removed")
             return None

tests/unit/test_cluster.py

@@ -323,6 +323,17 @@ def test_schedule_unique_keeps_client_route_events_for_distinct_ports(self, *_):
 
 class SessionPoolRaceTest(unittest.TestCase):
 
+    class _DuplicatePoolEntries(object):
+
+        def __init__(self, entries):
+            self._entries = entries
+
+        def __len__(self):
+            return len(self._entries)
+
+        def items(self):
+            return list(self._entries)
+
     @staticmethod
     def _make_host(address):
         return Host(address, SimpleConvictionPolicy, host_id=uuid.uuid4())
@@ -511,6 +522,22 @@ def make_pool(host, distance, pool_session, endpoint=None):
         assert session._pools == {}
         created_pools[0].shutdown.assert_called_once_with()
 
+    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()
+        cluster, session, executor = self._make_cluster_and_session([host])
+        stale_pool = self._make_pool(host, HostDistance.LOCAL, session)
+        session._pools[host] = stale_pool
+
+        host.endpoint = DefaultEndPoint("127.0.0.2")
+
+        futures = session.update_created_pools()
+
+        assert len(futures) == 1
+        executor.run_next()
+        assert session._pools == {}
+        stale_pool.shutdown.assert_called_once_with()
+
     def test_update_created_pools_replaces_pool_after_endpoint_change(self):
         host = self._make_host("127.0.0.1")
         old_endpoint = host.endpoint
@@ -628,11 +655,13 @@ def test_remove_pool_expected_endpoint_preserves_replacement_pool(self):
         old_endpoint = host.endpoint
         cluster, session, executor = self._make_cluster_and_session([host])
         stale_pool = self._make_pool(host, HostDistance.LOCAL, session)
-        session._pools[host] = stale_pool
 
         host.endpoint = DefaultEndPoint("127.0.0.2")
         replacement_pool = self._make_pool(host, HostDistance.LOCAL, session)
-        session._pools[host] = replacement_pool
+        session._pools = self._DuplicatePoolEntries([
+            (host, stale_pool),
+            (host, replacement_pool),
+        ])
 
         assert len(session._pools) == 2
 

tests/unit/test_response_future.py

@@ -20,7 +20,7 @@
 from unittest.mock import Mock, MagicMock, ANY
 
 from cassandra import ConsistencyLevel, Unavailable, SchemaTargetType, SchemaChangeType, OperationTimedOut
-from cassandra.cluster import Session, ResponseFuture, NoHostAvailable, ProtocolVersion
+from cassandra.cluster import Cluster, Session, ResponseFuture, NoHostAvailable, ProtocolVersion
 from cassandra.connection import Connection, ConnectionException, DefaultEndPoint
 from cassandra.protocol import (ReadTimeoutErrorMessage, WriteTimeoutErrorMessage,
                                 UnavailableErrorMessage, ResultMessage, QueryMessage,
@@ -700,6 +700,32 @@ def test_timeout_returns_orphan_to_original_pool_after_endpoint_swap(self):
         old_pool.return_connection.assert_called_once_with(
             connection, stream_was_orphaned=True)
 
+    def test_query_does_not_borrow_stale_pool_after_endpoint_swap(self):
+        session = self.make_basic_session()
+        host = Host(DefaultEndPoint('127.0.0.1'), SimpleConvictionPolicy,
+                    host_id=uuid.uuid4())
+        stale_pool = self.make_pool()
+        stale_pool.host = host
+        stale_pool.endpoint = host.endpoint
+        stale_pool.is_shutdown = False
+        connection = Mock(spec=Connection)
+        stale_pool.borrow_connection.return_value = (connection, 1)
+
+        session._lock = RLock()
+        session._pools = {host: stale_pool}
+        session._endpoints_match = Session._endpoints_match.__get__(session, Session)
+        session._pool_matches_expected = Session._pool_matches_expected.__get__(session, Session)
+        session._get_pool_by_host_identity = Session._get_pool_by_host_identity.__get__(session, Session)
+        session.cluster._endpoints_match.side_effect = Cluster._endpoints_match
+        session.cluster._default_load_balancing_policy.make_query_plan.return_value = [host]
+        host.endpoint = DefaultEndPoint('127.0.0.2')
+
+        rf = self.make_response_future(session)
+
+        assert not rf.send_request()
+        stale_pool.borrow_connection.assert_not_called()
+        assert isinstance(rf._errors[host], ConnectionException)
+
     def test_single_host_query_plan_exhausted_after_one_retry(self):
         """
         Test that when a specific host is provided, the query plan is properly