Route requests to the tablet leader for strongly-consistent tables

With TABLETS_ROUTING_V2 the server returns, on a tablet_version
mismatch, the tablet's replica set with the Raft leader first plus the
new tablet_version. Use this to route strongly-consistent reads and
writes straight to the leader:

* Every EXECUTE on a V2 connection carries a tablet_version_block
  computed from the cached version (or a random byte on cold start and
  for non-single-partition requests). The block is attached per
  connection, keyed off the connection that actually serves the request,
  so a rolling upgrade that mixes v1/v2 connections never desyncs the
  frame.
* On the response, the routing payload is parsed according to what the
  serving connection negotiated; the v2 tuple additionally carries the
  tablet_version, which is stored back on the tablet.
* TokenAwarePolicy yields the leader (the first replica) first, but only
  for strongly-consistent keyspaces -- a tablet_version is assigned to
  eventually-consistent tablet tables too and must not be mistaken for a
  leader hint.

To keep this hot path cheap, the ring token is memoized once per
statement (Statement.routing_token) and reused by the load balancing
policy, the shard selection in the pool, and the tablet_version_block
computation, instead of re-hashing the routing key three times.

Author: dawmd

cassandra/cluster.py

@@ -85,7 +85,7 @@
                              named_tuple_factory, dict_factory, tuple_factory, FETCH_SIZE_UNSET,
                              HostTargetingStatement)
 from cassandra.marshal import int64_pack
-from cassandra.tablets import Tablet, Tablets
+from cassandra.tablets import Tablet, Tablets, 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
 
@@ -3059,6 +3059,8 @@ def _create_response_future(self, query, parameters, trace, custom_payload,
                 continuous_paging_options, statement_keyspace)
         elif isinstance(query, BoundStatement):
             prepared_statement = query.prepared_statement
+            # The tablet_version_block is filled in per-target-host at send time
+            # (see ResponseFuture._query), because V2 is negotiated per connection.
             message = ExecuteMessage(
                 prepared_statement.query_id, query.values, cl,
                 serial_cl, fetch_size, paging_state, timestamp,
@@ -3093,6 +3095,42 @@ def _create_response_future(self, query, parameters, trace, custom_payload,
             load_balancer=load_balancing_policy, start_time=start_time, speculative_execution_plan=spec_exec_plan,
             continuous_paging_state=None, host=host)
 
+    def _compute_tablet_version_block(self, query):
+        """
+        Compute the tablet_version_block byte for a BoundStatement.
+
+        Always returns an int in [0, 255]. 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. Callers invoke this
+        only for connections that negotiated TABLETS_ROUTING_V2.
+        """
+        routing_key = query.routing_key
+        if routing_key is None:
+            return random_tablet_version_block()
+
+        keyspace = query.keyspace or self.keyspace
+        table = query.table
+        if not keyspace or not table:
+            return random_tablet_version_block()
+
+        # Skip the Murmur3 token hash + tablet lookup when we have no cached
+        # tablets for this table (vnode tables, or tablet tables on cold start);
+        # both correctly fall back to a random block below.
+        if not self.cluster.metadata._tablets.table_has_tablets(keyspace, table):
+            return random_tablet_version_block()
+
+        token_map = self.cluster.metadata.token_map
+        if token_map is None:
+            return random_tablet_version_block()
+
+        t = query.routing_token(token_map.token_class)
+        tablet = self.cluster.metadata._tablets.get_tablet_for_key(keyspace, table, t)
+        if tablet is None or tablet.tablet_version is None:
+            return random_tablet_version_block()
+
+        return choose_tablet_version_block(tablet.tablet_version)
+
     def get_execution_profile(self, name):
         """
         Returns the execution profile associated with the provided ``name``.
@@ -3764,7 +3802,6 @@ class PeersQueryType(object):
     _schema_meta_page_size = 1000
 
     _uses_peers_v2 = True
-    _tablets_routing_v1 = False
 
     # for testing purposes
     _time = time
@@ -3898,8 +3935,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
@@ -4713,6 +4748,7 @@ class ResponseFuture(object):
     _host = None
     _control_connection_query_attempted = False
     _TABLET_ROUTING_CTYPE = None
+    _TABLET_ROUTING_V2_CTYPE = None
 
     _warned_timeout = False
 
@@ -4930,6 +4966,36 @@ def _handle_control_connection_response(self, connection, cb, response):
             connection.in_flight -= 1
         cb(response)
 
+    def _prepare_message_for_connection(self, message, connection):
+        """
+        Return the message to send on ``connection``, attaching the
+        tablet_version_block to ExecuteMessages according to the capability that
+        *this specific connection* negotiated.
+
+        Keying off the borrowed connection (already in hand at every call site)
+        is both necessary and sufficient: a connection that negotiated
+        TABLETS_ROUTING_V2 always gets the block -- even if the pool was created,
+        and any cached flag latched, before the cluster feature was enabled
+        (e.g. mid rolling-upgrade) -- while a non-V2 connection never gets one,
+        even if a sibling shard connection in the same pool already negotiated
+        V2. The server reads the trailing byte only on V2 connections, so
+        attaching it to a non-V2 connection would leave an unread trailing byte
+        and desync the frame; a pool-level flag cannot get this right.
+
+        ExecuteMessage is copied per send because ``self.message`` is shared
+        across speculative executions and retries that may run concurrently on
+        different threads; mutating tablet_version_block on the shared instance
+        would race with another in-flight send encoding the same object.
+        """
+        if not isinstance(message, ExecuteMessage):
+            return message
+        message = copy(message)
+        if connection.features.tablets_routing_v2:
+            message.tablet_version_block = self.session._compute_tablet_version_block(self.query)
+        else:
+            message.tablet_version_block = None
+        return message
+
     def _query_control_connection(self, message=None, cb=None, connection=None, host=None):
         self._control_connection_query_attempted = True
 
@@ -4957,6 +5023,9 @@ def _query_control_connection(self, message=None, cb=None, connection=None, host
                 cb = partial(self._set_result, host, connection, None)
             cb = partial(self._handle_control_connection_response, connection, cb)
 
+            # The control connection may also be a V2 connection, so it needs the
+            # trailing tablet_version_block byte just like a pooled send.
+            message = self._prepare_message_for_connection(message, connection)
             log.debug("No usable node pools; falling back to control connection for host %s", host)
             self.request_encoded_size = connection.send_msg(message, request_id, cb=cb,
                                                             encoder=self._protocol_handler.encode_message,
@@ -5002,7 +5071,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 statement so the pool can reuse the ring token it
+                # memoized for this request instead of 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,
+                    query=self.query)
             else:
                 connection, request_id = pool.borrow_connection(timeout=2.0)
             self._connection = connection
@@ -5011,6 +5085,8 @@ def _query(self, host, message=None, cb=None):
             if cb is None:
                 cb = partial(self._set_result, host, connection, pool)
 
+            message = self._prepare_message_for_connection(message, connection)
+
             self.request_encoded_size = connection.send_msg(message, request_id, cb=cb,
                                                             encoder=self._protocol_handler.encode_message,
                                                             decoder=self._protocol_handler.decode_message,
@@ -5113,6 +5189,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
@@ -5128,21 +5225,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:
+                # during a rolling upgrade connections may differ, 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:
@@ -5546,7 +5645,7 @@ def add_callback(self, fn, *args, **kwargs):
 
         Note: in the case that the result is not available when the callback is added,
         the callback is executed by IO event thread. This means that the callback
-        should not block or attempt further synchronous requests, because no further
+        should not block or attempt further synchronous requests because no further
         IO will be processed until the callback returns.
 
         **Important**: if the callback you attach results in an exception being

cassandra/policies.py

@@ -503,14 +503,33 @@ def make_query_plan(self, working_keyspace=None, query=None):
             return
 
         replicas = []
-        tablet = self._cluster_metadata._tablets.get_tablet_for_key(
-            keyspace, query.table, self._cluster_metadata.token_map.token_class.from_key(query.routing_key))
+        leader_host = None
+        token = query.routing_token(self._cluster_metadata.token_map.token_class)
+        tablet = self._cluster_metadata._tablets.get_tablet_for_key(keyspace, query.table, token)
 
         if tablet is not None:
             replicas_mapped = set(map(lambda r: r[0], tablet.replicas))
             child_plan = child.make_query_plan(keyspace, query)
 
             replicas = [host for host in child_plan if host.host_id in replicas_mapped]
+
+            # The leader concept only exists for strongly-consistent keyspaces.
+            # TABLETS_ROUTING_V2 assigns a tablet_version to *every* tablet table
+            # (eventually- and strongly-consistent alike), so the version alone
+            # must not be used to infer a leader. Conversely, replicas[0] is only
+            # leader-ordered for a tablet that came from a V2 payload, so a
+            # versionless tablet (V1-sourced, or stale across a consistency flip)
+            # must not be treated as a leader hint either. Require both a
+            # strongly-consistent keyspace and a versioned tablet; otherwise keep
+            # normal token-aware/shuffled ordering.
+            ks_meta = self._cluster_metadata.keyspaces.get(keyspace)
+            if (ks_meta is not None and ks_meta.strongly_consistent
+                    and tablet.tablet_version is not None and tablet.replicas):
+                leader_host_id = tablet.replicas[0][0]
+                for host in replicas:
+                    if host.host_id == leader_host_id:
+                        leader_host = host
+                        break
         else:
             replicas = self._cluster_metadata.get_replicas(keyspace, query.routing_key)
 
@@ -523,10 +542,21 @@ def yield_in_order(hosts):
                     if replica.is_up and child.distance(replica) == distance:
                         yield replica
 
-        # yield replicas: local_rack, local, remote
-        yield from yield_in_order(replicas)
+        # If we have a leader hint, yield it first -- but respect the child
+        # policy's own filter: never front-run a host the child policy would
+        # exclude (e.g. one a custom policy reports as IGNORED).
+        if (leader_host is not None and leader_host.is_up
+                and child.distance(leader_host) != HostDistance.IGNORED):
+            yield leader_host
+
+        # yield replicas: local_rack, local, remote (skipping leader already yielded)
+        for host in yield_in_order(replicas):
+            if host is not leader_host:
+                yield host
         # yield rest of the cluster: local_rack, local, remote
-        yield from yield_in_order([host for host in child.make_query_plan(keyspace, query) if host not in replicas])
+        for host in yield_in_order([host for host in child.make_query_plan(keyspace, query) if host not in replicas]):
+            if host is not leader_host:
+                yield host
 
     def on_up(self, *args, **kwargs):
         return self._child_policy.on_up(*args, **kwargs)