Improvement: add Cluster.eager_prepare_scope for DC/rack-scoped eager prepare - prepare should be done only on local rack / local DC / etc.

CHANGELOG.rst

@@ -6,6 +6,12 @@ Features
 * Negotiate and implement the ``SCYLLA_USE_METADATA_ID`` protocol extension: prepared
   statements skip re-sending result metadata on EXECUTE, and the driver automatically
   refreshes cached metadata when the server detects a schema change (DRIVER-153)
+* Add ``Cluster.eager_prepare_scope``, an ``EagerPrepareScope`` value controlling which
+  hosts are eligible for the eager preparation performed by ``prepare_on_all_hosts`` and
+  ``reprepare_on_up``. Defaults to ``EagerPrepareScope.ALL`` (today's behavior, unchanged);
+  applications on large, multi-DC clusters can narrow it to ``LOCAL_DC``, ``LOCAL_RACK``, or
+  ``NONE`` to avoid eagerly preparing statements on rarely-queried remote hosts
+  (scylla-drivers#127).
 
 Others
 ------

cassandra/cluster.py

@@ -650,6 +650,41 @@
     SkipPoolCreation = "SkipPoolCreation"
 
 
+class EagerPrepareScope(enum.Enum):
+    """
+    Controls which hosts are eligible for the eager statement preparation performed by
+    :attr:`Cluster.prepare_on_all_hosts` and :attr:`Cluster.reprepare_on_up`, via
+    :attr:`Cluster.eager_prepare_scope`.
+
+    Levels are ordered from narrowest to widest: ``NONE`` makes no host eligible
+    (statements are only ever prepared lazily, on first use against a given host, the
+    same as setting :attr:`Cluster.prepare_on_all_hosts` and
+    :attr:`Cluster.reprepare_on_up` to :const:`False`); ``LOCAL_RACK`` includes only
+    hosts at :attr:`~.HostDistance.LOCAL_RACK`; ``LOCAL_DC`` additionally includes
+    hosts at :attr:`~.HostDistance.LOCAL`; ``ALL`` additionally includes hosts at
+    :attr:`~.HostDistance.REMOTE` (i.e. every host with an open connection pool, which
+    was the only behavior available before this setting was introduced).
+    """
+
+    NONE = 0
+    LOCAL_RACK = 1
+    LOCAL_DC = 2
+    ALL = 3
+
+    def includes_distance(self, distance):
+        """
+        Whether a host at the given :class:`~.HostDistance` is eligible for eager
+        preparation under this scope.
+        """
+        if self is EagerPrepareScope.NONE:
+            return False
+        if self is EagerPrepareScope.LOCAL_RACK:
+            return distance == HostDistance.LOCAL_RACK
+        if self is EagerPrepareScope.LOCAL_DC:
+            return distance in (HostDistance.LOCAL_RACK, HostDistance.LOCAL)
+        return distance != HostDistance.IGNORED
+
+
 class Cluster(object):
     """
     The main class to use when interacting with a Cassandra cluster.
@@ -1046,6 +1081,8 @@
     This can reasonably be disabled on long-running applications with numerous clients preparing statements on startup,
     where a randomized initial condition of the load balancing policy can be expected to distribute prepares from
     different clients across the cluster.
+
+    Which hosts count as "all hosts" is controlled by :attr:`.eager_prepare_scope`.
     """
 
     reprepare_on_up = True
@@ -1055,6 +1092,22 @@
     May be used to avoid overwhelming a node on return, or if it is supposed that the node was only marked down due to
     network. If statements are not reprepared, they are prepared on the first execution, causing
     an extra roundtrip for one or more client requests.
+
+    Whether a given node coming up is eligible for this is controlled by :attr:`.eager_prepare_scope`.
+    """
+
+    eager_prepare_scope = EagerPrepareScope.ALL
+    """
+    An :class:`.EagerPrepareScope` value controlling which hosts are eligible for the eager
+    preparation performed by :attr:`.prepare_on_all_hosts` and :attr:`.reprepare_on_up`.
+
+    Defaults to :attr:`EagerPrepareScope.ALL`, preserving the historical behavior of eagerly
+    preparing on every host with an open connection pool, including hosts at
+    :attr:`~.HostDistance.REMOTE` (kept connected only as a cross-DC fallback via
+    :attr:`.connect_to_remote_hosts` and ``used_hosts_per_remote_dc``). Applications on large,
+    multi-DC clusters that don't want to pay the eager-prepare cost on rarely-queried remote hosts
+    can narrow this to :attr:`EagerPrepareScope.LOCAL_DC` or :attr:`EagerPrepareScope.LOCAL_RACK`;
+    those hosts still fall back to lazy, on-first-use preparation, so correctness is unaffected.
     """
 
     connect_timeout = 5
@@ -1267,7 +1320,8 @@
                  column_encryption_policy=None,
                  application_info:Optional[ApplicationInfoBase]=None,
                  client_routes_config:Optional[ClientRoutesConfig]=None,
-                 allow_control_connection_query_fallback:Optional[ControlConnectionQueryFallback]=ControlConnectionQueryFallback.Disabled
+                 allow_control_connection_query_fallback:Optional[ControlConnectionQueryFallback]=ControlConnectionQueryFallback.Disabled,
+                 eager_prepare_scope=EagerPrepareScope.ALL,
                  ):
         """
         ``executor_threads`` defines the number of threads in a pool for handling asynchronous tasks such as
@@ -1289,6 +1343,9 @@
             raise TypeError(
                 "allow_control_connection_query_fallback must be a ControlConnectionQueryFallback value")
 
+        if not isinstance(eager_prepare_scope, EagerPrepareScope):
+            raise TypeError("eager_prepare_scope must be an EagerPrepareScope value")
+
         if connection_class is not None:
             self.connection_class = connection_class
 
@@ -1530,6 +1587,7 @@
         self.connect_timeout = connect_timeout
         self.prepare_on_all_hosts = prepare_on_all_hosts
         self.reprepare_on_up = reprepare_on_up
+        self.eager_prepare_scope = eager_prepare_scope
         self.monitor_reporting_enabled = monitor_reporting_enabled
         self.monitor_reporting_interval = monitor_reporting_interval
         self.shard_aware_options = ShardAwareOptions(opts=shard_aware_options)
@@ -1995,8 +2053,9 @@
                 log.debug("Now that host %s is up, cancelling the reconnection handler", host)
                 reconnector.cancel()
 
-            if self.profile_manager.distance(host) != HostDistance.IGNORED:
-                self._prepare_all_queries(host)
+            distance = self.profile_manager.distance(host)
+            if distance != HostDistance.IGNORED:
+                self._prepare_all_queries(host, distance)
                 log.debug("Done preparing all queries for host %s, ", host)
 
             for session in tuple(self.sessions):
@@ -2113,7 +2172,7 @@
 
         distance = self.profile_manager.distance(host)
         if distance != HostDistance.IGNORED:
-            self._prepare_all_queries(host)
+            self._prepare_all_queries(host, distance)
             log.debug("Done preparing queries for new host %r", host)
 
         if distance == HostDistance.IGNORED:
@@ -2403,10 +2462,18 @@
                     log.debug("Got unexpected response when preparing "
                               "statement on host %s: %r", host, response)
 
-    def _prepare_all_queries(self, host):
+    def _prepare_all_queries(self, host, distance=None):
         if not self._prepared_statements or not self.reprepare_on_up:
             return
 
+        if distance is None:
+            distance = self.profile_manager.distance(host)
+
+        if not self.eager_prepare_scope.includes_distance(distance):
+            log.debug("Not preparing known prepared statements against host %s: "
+                      "outside eager_prepare_scope %s", host, self.eager_prepare_scope)
+            return
+
         log.debug("Preparing all known prepared statements against host %s", host)
         connection = None
         try:
@@ -3242,31 +3309,40 @@
 
     def prepare_on_all_hosts(self, query, excluded_host, keyspace=None):
         """
-        Prepare the given query on all hosts, excluding ``excluded_host``.
+        Prepare the given query on all hosts within :attr:`Cluster.eager_prepare_scope`,
+        excluding ``excluded_host``.
         Intended for internal use only.
         """
+        scope = self.cluster.eager_prepare_scope
         futures = []
         for host in tuple(self._pools.keys()):
-            if host != excluded_host and host.is_up:
-                future = ResponseFuture(self, PrepareMessage(query=query, keyspace=keyspace),
-                                            None, self.default_timeout)
+            if host == excluded_host or not host.is_up:
+                continue
 
-                # we don't care about errors preparing against specific hosts,
-                # since we can always prepare them as needed when the prepared
-                # statement is used.  Just log errors and continue on.
-                try:
-                    request_id = future._query(host)
-                except Exception:
-                    log.exception("Error preparing query for host %s:", host)
-                    continue
+            if not scope.includes_distance(self._profile_manager.distance(host)):
+                log.debug("Not preparing query for host %s: outside eager_prepare_scope %s",
+                          host, scope)
+                continue
 
-                if request_id is None:
-                    # the error has already been logged by ResponsFuture
-                    log.debug("Failed to prepare query for host %s: %r",
-                              host, future._errors.get(host))
-                    continue
+            future = ResponseFuture(self, PrepareMessage(query=query, keyspace=keyspace),
+                                        None, self.default_timeout)
+
+            # we don't care about errors preparing against specific hosts,
+            # since we can always prepare them as needed when the prepared
+            # statement is used.  Just log errors and continue on.
+            try:
+                request_id = future._query(host)
+            except Exception:
+                log.exception("Error preparing query for host %s:", host)
+                continue
+
+            if request_id is None:
+                # the error has already been logged by ResponsFuture
+                log.debug("Failed to prepare query for host %s: %r",
+                          host, future._errors.get(host))
+                continue
 
-                futures.append((host, future))
+            futures.append((host, future))
 
         for host, future in futures:
             try:
@@ -4634,7 +4710,7 @@
                         self._scheduled_tasks.discard(task)
                         fn, args, kwargs = task
                         kwargs = dict(kwargs)
                         future = self._executor.submit(fn, *args, **kwargs)
                         future.add_done_callback(self._log_if_failed)
                     else:
                         self._queue.put_nowait((run_at, i, task))

docs/api/cassandra/cluster.rst

@@ -64,6 +64,9 @@ Clusters and Sessions
 
    .. autoattribute:: reprepare_on_up
 
+   .. autoattribute:: eager_prepare_scope
+      :annotation: = EagerPrepareScope.ALL
+
    .. autoattribute:: connect_timeout
 
    .. autoattribute:: schema_metadata_enabled
@@ -111,6 +114,9 @@ Clusters and Sessions
 .. autoclass:: ControlConnectionQueryFallback
    :members:
 
+.. autoclass:: EagerPrepareScope
+   :members:
+
 .. autoclass:: ExecutionProfile (load_balancing_policy=<object object>, retry_policy=None, consistency_level=ConsistencyLevel.LOCAL_ONE, serial_consistency_level=None, request_timeout=10.0, row_factory=<function named_tuple_factory>, speculative_execution_policy=None)
    :members:
    :exclude-members: consistency_level

tests/unit/test_cluster.py

@@ -20,11 +20,12 @@
 
 from unittest.mock import patch, Mock
 import uuid
+from weakref import WeakValueDictionary
 
 from cassandra import ConsistencyLevel, DriverException, Timeout, Unavailable, RequestExecutionException, ReadTimeout, WriteTimeout, CoordinationFailure, ReadFailure, WriteFailure, FunctionFailure, AlreadyExists,\
     InvalidRequest, Unauthorized, AuthenticationFailed, OperationTimedOut, UnsupportedOperation, RequestValidationException, ConfigurationException, ProtocolVersion
 from cassandra.cluster import _Scheduler, Session, Cluster, ResultSet, SchemaAgreementScope, ControlConnectionQueryFallback, default_lbp_factory, \
-    ExecutionProfile, _ConfigMode, EXEC_PROFILE_DEFAULT
+    ExecutionProfile, _ConfigMode, EXEC_PROFILE_DEFAULT, EagerPrepareScope
 from cassandra.connection import ConnectionBusy, ConnectionException
 from cassandra.pool import Host
 from cassandra.policies import HostDistance, RetryPolicy, RoundRobinPolicy, DowngradingConsistencyRetryPolicy, SimpleConvictionPolicy
@@ -187,6 +188,106 @@ def test_port_range(self):
             with pytest.raises(ValueError):
                 cluster = Cluster(contact_points=['127.0.0.1'], port=invalid_port)
 
+    def test_eager_prepare_scope_default(self):
+        assert Cluster().eager_prepare_scope is EagerPrepareScope.ALL
+
+    def test_eager_prepare_scope_rejects_non_enum_values(self):
+        for invalid_value in (True, False, 'ALL', 1, None):
+            with pytest.raises(TypeError):
+                Cluster(eager_prepare_scope=invalid_value)
+
+    def test_eager_prepare_scope_includes_distance(self):
+        assert EagerPrepareScope.NONE.includes_distance(HostDistance.LOCAL_RACK) is False
+        assert EagerPrepareScope.NONE.includes_distance(HostDistance.LOCAL) is False
+        assert EagerPrepareScope.NONE.includes_distance(HostDistance.REMOTE) is False
+
+        assert EagerPrepareScope.LOCAL_RACK.includes_distance(HostDistance.LOCAL_RACK) is True
+        assert EagerPrepareScope.LOCAL_RACK.includes_distance(HostDistance.LOCAL) is False
+        assert EagerPrepareScope.LOCAL_RACK.includes_distance(HostDistance.REMOTE) is False
+
+        assert EagerPrepareScope.LOCAL_DC.includes_distance(HostDistance.LOCAL_RACK) is True
+        assert EagerPrepareScope.LOCAL_DC.includes_distance(HostDistance.LOCAL) is True
+        assert EagerPrepareScope.LOCAL_DC.includes_distance(HostDistance.REMOTE) is False
+
+        assert EagerPrepareScope.ALL.includes_distance(HostDistance.LOCAL_RACK) is True
+        assert EagerPrepareScope.ALL.includes_distance(HostDistance.LOCAL) is True
+        assert EagerPrepareScope.ALL.includes_distance(HostDistance.REMOTE) is True
+
+        for scope in EagerPrepareScope:
+            assert scope.includes_distance(HostDistance.IGNORED) is False
+
+    def test_prepare_all_queries_skips_hosts_outside_eager_prepare_scope(self):
+        cluster = Cluster(
+            allow_control_connection_query_fallback=ControlConnectionQueryFallback.SkipPoolCreation,
+            monitor_reporting_enabled=False,
+            eager_prepare_scope=EagerPrepareScope.LOCAL_DC,
+        )
+        prepared_statement = Mock(keyspace=None)
+        cluster._prepared_statements = WeakValueDictionary({'query_id': prepared_statement})
+        cluster.profile_manager.distance = Mock(return_value=HostDistance.REMOTE)
+        host = Host("127.0.0.1", SimpleConvictionPolicy, host_id=uuid.uuid4())
+
+        with patch.object(Cluster, "connection_factory") as mocked_connection_factory:
+            cluster._prepare_all_queries(host)
+
+        mocked_connection_factory.assert_not_called()
+
+        cluster.profile_manager.distance = Mock(return_value=HostDistance.LOCAL)
+        with patch.object(Cluster, "connection_factory") as mocked_connection_factory:
+            cluster._prepare_all_queries(host)
+
+        mocked_connection_factory.assert_called_once_with(host.endpoint)
+
+    def test_prepare_all_queries_reuses_caller_provided_distance(self):
+        cluster = Cluster(
+            allow_control_connection_query_fallback=ControlConnectionQueryFallback.SkipPoolCreation,
+            monitor_reporting_enabled=False,
+            eager_prepare_scope=EagerPrepareScope.LOCAL_DC,
+        )
+        prepared_statement = Mock(keyspace=None)
+        cluster._prepared_statements = WeakValueDictionary({'query_id': prepared_statement})
+        cluster.profile_manager.distance = Mock(
+            side_effect=AssertionError("distance() should not be recomputed when already provided by the caller"))
+        host = Host("127.0.0.1", SimpleConvictionPolicy, host_id=uuid.uuid4())
+
+        with patch.object(Cluster, "connection_factory") as mocked_connection_factory:
+            cluster._prepare_all_queries(host, HostDistance.LOCAL)
+
+        mocked_connection_factory.assert_called_once_with(host.endpoint)
+
+    def test_prepare_on_all_hosts_skips_hosts_outside_eager_prepare_scope(self):
+        cluster = Cluster(
+            allow_control_connection_query_fallback=ControlConnectionQueryFallback.SkipPoolCreation,
+            monitor_reporting_enabled=False,
+            eager_prepare_scope=EagerPrepareScope.LOCAL_DC,
+        )
+        session = Session(cluster, [])
+
+        excluded_host = Host("127.0.0.1", SimpleConvictionPolicy, host_id=uuid.uuid4())
+        excluded_host.set_up()
+        remote_host = Host("127.0.0.2", SimpleConvictionPolicy, host_id=uuid.uuid4())
+        remote_host.set_up()
+        local_host = Host("127.0.0.3", SimpleConvictionPolicy, host_id=uuid.uuid4())
+        local_host.set_up()
+        down_host = Host("127.0.0.4", SimpleConvictionPolicy, host_id=uuid.uuid4())
+
+        session._pools = {
+            excluded_host: Mock(),
+            remote_host: Mock(),
+            local_host: Mock(),
+            down_host: Mock(),
+        }
+        session._profile_manager.distance = Mock(
+            side_effect=lambda h: HostDistance.REMOTE if h is remote_host else HostDistance.LOCAL)
+
+        mock_future = Mock()
+        mock_future._query.return_value = 1
+        with patch('cassandra.cluster.ResponseFuture', Mock(return_value=mock_future)):
+            session.prepare_on_all_hosts('SELECT 1', excluded_host)
+
+        queried_hosts = [call.args[0] for call in mock_future._query.call_args_list]
+        assert queried_hosts == [local_host]
+
     def test_control_connection_query_fallback_modes(self):
         assert Cluster().allow_control_connection_query_fallback is ControlConnectionQueryFallback.Disabled
         with pytest.raises(TypeError):