Implement TABLETS_ROUTING_V2

cassandra/cluster.py

@@ -70,7 +70,7 @@
                                 RESULT_KIND_SET_KEYSPACE, RESULT_KIND_ROWS,
                                 RESULT_KIND_SCHEMA_CHANGE, ProtocolHandler,
                                 RESULT_KIND_VOID, ProtocolException)
-from cassandra.metadata import Metadata, protect_name, murmur3, _NodeInfo
+from cassandra.metadata import Metadata, Token, protect_name, murmur3, _NodeInfo
 from cassandra.policies import (TokenAwarePolicy, DCAwareRoundRobinPolicy, SimpleConvictionPolicy,
                                 ExponentialReconnectionPolicy, HostDistance,
                                 RetryPolicy, IdentityTranslator, NoSpeculativeExecutionPlan,
@@ -84,7 +84,7 @@
                              named_tuple_factory, dict_factory, tuple_factory, FETCH_SIZE_UNSET,
                              HostTargetingStatement)
 from cassandra.marshal import int64_pack
-from cassandra.tablets import Tablet
+from cassandra.tablets import Tablet, choose_tablet_version_block, random_tablet_version_block
 from cassandra.timestamps import MonotonicTimestampGenerator
 from cassandra.util import _resolve_contact_points_to_string_map, Version, maybe_add_timeout_to_query
 
@@ -3051,6 +3051,21 @@ def _create_response_future(self, query, parameters, trace, custom_payload,
         # bound statements carry cached result metadata (set in the BoundStatement branch).
         bound_result_metadata = _NOT_SET
 
+        # Compute the ring token once, here on the request path, and pass it
+        # explicitly to the two consumers that run while sending: the
+        # tablet_version_block below and shard selection in the pool (via the
+        # ResponseFuture). The token is a pure function of the routing key and
+        # the cluster's partitioner, so computing it here keeps cluster-dependent
+        # state off the statement and avoids races when a statement is shared
+        # across concurrent requests. The load balancing policy computes its own
+        # token from the same routing key when ordering replicas.
+        routing_token = None
+        routing_key = query.routing_key
+        if routing_key is not None:
+            token_map = self.cluster.metadata.token_map
+            if token_map is not None:
+                routing_token = token_map.token_class.from_key(routing_key)
+
         if isinstance(query, SimpleStatement):
             query_string = query.query_string
             statement_keyspace = query.keyspace if ProtocolVersion.uses_keyspace_flag(self._protocol_version) else None
@@ -3076,13 +3091,19 @@ def _create_response_future(self, query, parameters, trace, custom_payload,
             # decode page 2+ against.
             result_metadata, result_metadata_id = prepared_statement.result_metadata_and_id
             bound_result_metadata = result_metadata
+
+            # The tablet_version_block value is connection-independent, so compute
+            # it once here instead of copying the message per send attempt. The
+            # serializer emits it only when the serving connection negotiated
+            # TABLETS_ROUTING_V2 (see ExecuteMessage.send_body).
             message = ExecuteMessage(
                 prepared_statement.query_id, query.values, cl,
                 serial_cl, fetch_size, paging_state, timestamp,
                 skip_meta=bool(result_metadata) and result_metadata_id is not None
                           and continuous_paging_options is None,
                 continuous_paging_options=continuous_paging_options,
-                result_metadata_id=result_metadata_id)
+                result_metadata_id=result_metadata_id,
+                tablet_version_block=self._compute_tablet_version_block(query, routing_key, routing_token))
         elif isinstance(query, BatchStatement):
             if self._protocol_version < 2:
                 raise UnsupportedOperation(
@@ -3109,7 +3130,54 @@ def _create_response_future(self, query, parameters, trace, custom_payload,
             self, message, query, timeout, metrics=self._metrics,
             prepared_statement=prepared_statement, retry_policy=retry_policy, row_factory=row_factory,
             load_balancer=load_balancing_policy, start_time=start_time, speculative_execution_plan=spec_exec_plan,
-            continuous_paging_state=None, host=host, bound_result_metadata=bound_result_metadata)
+            continuous_paging_state=None, host=host, bound_result_metadata=bound_result_metadata,
+            routing_token=routing_token)
+
+    def _compute_tablet_version_block(self, query, routing_key: Optional[bytes],
+                                      routing_token: Optional[Token]) -> int:
+        """
+        Compute the tablet_version_block byte for a BoundStatement.
+
+        Always returns an int in [0, 255]. A non-token-aware query (no routing
+        key) can never resolve to a tablet, so the server never version-checks
+        it; we send 0 and skip the work. Otherwise, when no cached tablet is
+        known for the routing key (unknown keyspace/table, vnode table, cold
+        cache, or a missing token map) a random block is returned; the server
+        treats that as a version miss and replies with fresh routing info.
+
+        ``routing_key`` and ``routing_token`` are the statement's routing key and
+        the ring token derived from it, both resolved once per request by the
+        caller (see :meth:`_create_response_future`) and passed in so the send
+        path has a single source of truth for them. ``routing_token`` is ``None``
+        both when there is no routing key and when no token map was available, so
+        telling those two cases apart needs the routing key as well -- taking it
+        as an argument rather than re-reading ``query.routing_key`` keeps the two
+        values here guaranteed to describe the same statement.
+
+        This is computed once per request at message construction; the value is
+        connection-independent, and the serializer emits it only on connections
+        that negotiated TABLETS_ROUTING_V2 (see ExecuteMessage.send_body).
+        """
+        if routing_key is None:
+            # Non-token-aware query: the server won't version-check it, so skip
+            # generating random bits and just send 0.
+            return 0
+
+        keyspace = query.keyspace or self.keyspace
+        table = query.table
+        if not keyspace or not table or routing_token is None:
+            return random_tablet_version_block()
+
+        # A single lookup: get_tablet_for_key already reports a table with no
+        # cached tablets (a vnode table, or a tablet table on cold start) as
+        # None, and going through the mutable cache twice would leave a window
+        # for the tablet to disappear between the checks.
+        tablet = self.cluster.metadata._tablets.get_tablet_for_key(keyspace, table, routing_token)
+        if tablet is None or tablet.tablet_version is None:
+            # A version miss on the server, which replies with fresh routing info.
+            return random_tablet_version_block()
+
+        return choose_tablet_version_block(tablet.tablet_version)
 
     def get_execution_profile(self, name):
         """
@@ -3787,7 +3855,6 @@ class PeersQueryType(object):
     _schema_meta_page_size = 1000
 
     _uses_peers_v2 = True
-    _tablets_routing_v1 = False
 
     # for testing purposes
     _time = time
@@ -3921,8 +3988,6 @@ def _try_connect(self, endpoint):
         self._metadata_request_timeout = None if connection.features.sharding_info is None or not self._cluster.metadata_request_timeout \
             else datetime.timedelta(seconds=self._cluster.metadata_request_timeout)
 
-        self._tablets_routing_v1 = connection.features.tablets_routing_v1
-
         # use weak references in both directions
         # _clear_watcher will be called when this ControlConnection is about to be finalized
         # _watch_callback will get the actual callback from the Connection and relay it to
@@ -4736,14 +4801,15 @@ class ResponseFuture(object):
     _host = None
     _control_connection_query_attempted = False
     _TABLET_ROUTING_CTYPE = None
+    _TABLET_ROUTING_V2_CTYPE = None
     _bound_result_metadata = None
 
     _warned_timeout = False
 
     def __init__(self, session, message, query, timeout, metrics=None, prepared_statement=None,
                  retry_policy=RetryPolicy(), row_factory=None, load_balancer=None, start_time=None,
                  speculative_execution_plan=None, continuous_paging_state=None, host=None,
-                 bound_result_metadata=_NOT_SET):
+                 bound_result_metadata=_NOT_SET, routing_token=None):
         self.session = session
         # TODO: normalize handling of retry policy and row factory
         self.row_factory = row_factory or session.row_factory
@@ -4763,6 +4829,7 @@ def __init__(self, session, message, query, timeout, metrics=None, prepared_stat
         self._callback_lock = Lock()
         self._start_time = start_time or time.time()
         self._host = host
+        self._routing_token = routing_token
         self._control_connection_query_attempted = False
         self._spec_execution_plan = speculative_execution_plan or self._spec_execution_plan
         self._make_query_plan()
@@ -5033,7 +5100,12 @@ def _query(self, host, message=None, cb=None):
         try:
             # TODO get connectTimeout from cluster settings
             if self.query:
-                connection, request_id = pool.borrow_connection(timeout=2.0, routing_key=self.query.routing_key, keyspace=self.query.keyspace, table=self.query.table)
+                # Pass the ring token computed once for this request so the pool
+                # can select the shard without re-hashing the routing key.
+                connection, request_id = pool.borrow_connection(
+                    timeout=2.0, routing_key=self.query.routing_key,
+                    keyspace=self.query.keyspace, table=self.query.table,
+                    routing_token=self._routing_token)
             else:
                 connection, request_id = pool.borrow_connection(timeout=2.0)
             self._connection = connection
@@ -5144,6 +5216,27 @@ def _reprepare(self, prepare_message, host, connection, pool):
             # try to submit the original prepared statement on some other host
             self.send_request()
 
+    def _cache_tablet_from_payload(self, payload_key, ctype):
+        """
+        Parse a tablets-routing ``custom_payload`` entry and cache the Tablet.
+
+        ``ctype`` is the tuple type for the negotiated extension. The V1 and V2
+        layouts differ only by a trailing ``tablet_version`` field, and
+        ``Tablet.from_row`` accepts that as an optional final argument, so
+        unpacking the decoded tuple positionally serves both. The tablet is
+        cached under the effective keyspace (the statement's, else the
+        session's) so a prepared statement executed in a session keyspace lands
+        under the same key ``_compute_tablet_version_block`` looks it up by;
+        otherwise that lookup always misses.
+        """
+        info = self._custom_payload.get(payload_key)
+        protocol = self.session.cluster.protocol_version
+        tablet = Tablet.from_row(*ctype.from_binary(info, protocol))
+        keyspace = self.query.keyspace or self.session.keyspace
+        table = self.query.table
+        if tablet and keyspace and table:
+            self.session.cluster.metadata._tablets.add_tablet(keyspace, table, tablet)
+
     def _set_result(self, host, connection, pool, response):
         try:
             self.coordinator_host = host
@@ -5159,21 +5252,23 @@ def _set_result(self, host, connection, pool, response):
             self._warnings = getattr(response, 'warnings', None)
             self._custom_payload = getattr(response, 'custom_payload', None)
 
-            if self._custom_payload and self.session.cluster.control_connection._tablets_routing_v1 and 'tablets-routing-v1' in self._custom_payload:
-                protocol = self.session.cluster.protocol_version
-                info = self._custom_payload.get('tablets-routing-v1')
-                ctype = ResponseFuture._TABLET_ROUTING_CTYPE
-                if ctype is None:
-                    ctype = types.lookup_casstype('TupleType(LongType, LongType, ListType(TupleType(UUIDType, Int32Type)))')
-                    ResponseFuture._TABLET_ROUTING_CTYPE = ctype
-                tablet_routing_info = ctype.from_binary(info, protocol)
-                first_token = tablet_routing_info[0]
-                last_token = tablet_routing_info[1]
-                tablet_replicas = tablet_routing_info[2]
-                tablet = Tablet.from_row(first_token, last_token, tablet_replicas)
-                keyspace = self.query.keyspace
-                table = self.query.table
-                self.session.cluster.metadata._tablets.add_tablet(keyspace, table, tablet)
+            if self._custom_payload and connection is not None:
+                # Parse the routing payload according to what the connection that
+                # *served this request* negotiated, not the control connection:
+                # different nodes may negotiate different extensions, and each
+                # payload key matches the extension its own connection negotiated.
+                if connection.features.tablets_routing_v2 and 'tablets-routing-v2' in self._custom_payload:
+                    ctype = ResponseFuture._TABLET_ROUTING_V2_CTYPE
+                    if ctype is None:
+                        ctype = types.lookup_casstype('TupleType(LongType, LongType, ListType(TupleType(UUIDType, Int32Type)), LongType)')
+                        ResponseFuture._TABLET_ROUTING_V2_CTYPE = ctype
+                    self._cache_tablet_from_payload('tablets-routing-v2', ctype)
+                elif connection.features.tablets_routing_v1 and 'tablets-routing-v1' in self._custom_payload:
+                    ctype = ResponseFuture._TABLET_ROUTING_CTYPE
+                    if ctype is None:
+                        ctype = types.lookup_casstype('TupleType(LongType, LongType, ListType(TupleType(UUIDType, Int32Type)))')
+                        ResponseFuture._TABLET_ROUTING_CTYPE = ctype
+                    self._cache_tablet_from_payload('tablets-routing-v1', ctype)
 
             if isinstance(response, ResultMessage):
                 if response.kind == RESULT_KIND_SET_KEYSPACE: