From 56f3bd2b0d579484c0976c9c8edd1aaf29bc8ca7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dawid=20M=C4=99drek?= Date: Wed, 8 Jul 2026 14:02:42 +0200 Subject: [PATCH 1/9] Negotiate the TABLETS_ROUTING_V2 protocol extension Add per-connection negotiation of the TABLETS_ROUTING_V2 extension, the successor to TABLETS_ROUTING_V1. When the server advertises it in the SUPPORTED response, the driver echoes it back during STARTUP to opt in; a driver that negotiates v2 does not negotiate v1. While the feature is experimental the wire name carries the `_EXPERIMENTAL` suffix (TABLETS_ROUTING_V2_EXPERIMENTAL), and the server only advertises it when started with the `strongly-consistent-tables` experimental feature enabled. Also add the trailing tablet_version_block byte to the EXECUTE message body. The server reads exactly one such byte per EXECUTE on a connection that negotiated the extension, so the encoder writes one whenever the connection did -- coalescing an unset value to 0 -- and none otherwise. Later commits fill in the value from the cached tablet version. Deciding this from the connection's negotiated features rather than from the message is what lets one ExecuteMessage be sent, unmodified, on connections that negotiated differently. --- cassandra/protocol.py | 9 +++++- cassandra/protocol_features.py | 25 +++++++++++++--- tests/unit/test_protocol_features.py | 44 +++++++++++++++++++++++++++- 3 files changed, 72 insertions(+), 6 deletions(-) diff --git a/cassandra/protocol.py b/cassandra/protocol.py index 9dfdbf3022..518b0c711d 100644 --- a/cassandra/protocol.py +++ b/cassandra/protocol.py @@ -654,9 +654,11 @@ class ExecuteMessage(_QueryMessage): def __init__(self, query_id, query_params, consistency_level, serial_consistency_level=None, fetch_size=None, paging_state=None, timestamp=None, skip_meta=False, - continuous_paging_options=None, result_metadata_id=None): + continuous_paging_options=None, result_metadata_id=None, + tablet_version_block=None): self.query_id = query_id self.result_metadata_id = result_metadata_id + self.tablet_version_block = tablet_version_block super(ExecuteMessage, self).__init__(query_params, consistency_level, serial_consistency_level, fetch_size, paging_state, timestamp, skip_meta, continuous_paging_options) @@ -689,6 +691,11 @@ def send_body(self, f, protocol_version, protocol_features=None): # responds with full metadata plus the current id. write_string(f, self.result_metadata_id if self.result_metadata_id is not None else b'') self._write_query_params(f, protocol_version, protocol_features) + if protocol_features is not None and protocol_features.tablets_routing_v2: + # A V2 connection makes the server read exactly one trailing byte per + # EXECUTE, so always write one. Coalesce a missing value to 0 to keep + # the frame in sync. + write_byte(f, self.tablet_version_block if self.tablet_version_block is not None else 0) CUSTOM_TYPE = object() diff --git a/cassandra/protocol_features.py b/cassandra/protocol_features.py index 7165117e80..c2bc7ca417 100644 --- a/cassandra/protocol_features.py +++ b/cassandra/protocol_features.py @@ -11,23 +11,31 @@ RATE_LIMIT_ERROR_EXTENSION = "SCYLLA_RATE_LIMIT_ERROR" TABLETS_ROUTING_V1 = "TABLETS_ROUTING_V1" USE_METADATA_ID = "SCYLLA_USE_METADATA_ID" +# The server advertises and expects this exact extension name in SUPPORTED/STARTUP +# (see scylladb transport/cql_protocol_extension.cc). While the feature is gated +# behind the server's `strongly-consistent-tables` experimental flag, the wire +# name carries the `_EXPERIMENTAL` suffix. +TABLETS_ROUTING_V2 = "TABLETS_ROUTING_V2_EXPERIMENTAL" class ProtocolFeatures(object): rate_limit_error = None shard_id = 0 sharding_info = None tablets_routing_v1 = False + tablets_routing_v2 = False lwt_info = None use_metadata_id = False # Keyword-only so that independently developed protocol extensions can add # new fields without conflicting over positional-argument order. - def __init__(self, *, rate_limit_error=None, shard_id=0, sharding_info=None, tablets_routing_v1=False, lwt_info=None, + def __init__(self, *, rate_limit_error=None, shard_id=0, sharding_info=None, + tablets_routing_v1=False, tablets_routing_v2=False, lwt_info=None, use_metadata_id=False): self.rate_limit_error = rate_limit_error self.shard_id = shard_id self.sharding_info = sharding_info self.tablets_routing_v1 = tablets_routing_v1 + self.tablets_routing_v2 = tablets_routing_v2 self.lwt_info = lwt_info self.use_metadata_id = use_metadata_id @@ -36,11 +44,12 @@ def parse_from_supported(supported): rate_limit_error = ProtocolFeatures.maybe_parse_rate_limit_error(supported) shard_id, sharding_info = ProtocolFeatures.parse_sharding_info(supported) tablets_routing_v1 = ProtocolFeatures.parse_tablets_info(supported) + tablets_routing_v2 = ProtocolFeatures.parse_tablets_v2_info(supported) lwt_info = ProtocolFeatures.parse_lwt_info(supported) use_metadata_id = ProtocolFeatures.parse_use_metadata_id(supported) return ProtocolFeatures(rate_limit_error=rate_limit_error, shard_id=shard_id, sharding_info=sharding_info, - tablets_routing_v1=tablets_routing_v1, lwt_info=lwt_info, - use_metadata_id=use_metadata_id) + tablets_routing_v1=tablets_routing_v1, tablets_routing_v2=tablets_routing_v2, + lwt_info=lwt_info, use_metadata_id=use_metadata_id) @staticmethod def maybe_parse_rate_limit_error(supported): @@ -62,7 +71,11 @@ def get_cql_extension_field(vals, key): def add_startup_options(self, options): if self.rate_limit_error is not None: options[RATE_LIMIT_ERROR_EXTENSION] = "" - if self.tablets_routing_v1: + # Only one of TABLETS_ROUTING_V{1,2} should be negotiated + # per connection. Hence the if-else branch. + if self.tablets_routing_v2: + options[TABLETS_ROUTING_V2] = "" + elif self.tablets_routing_v1: options[TABLETS_ROUTING_V1] = "" if self.lwt_info is not None: options[LWT_ADD_METADATA_MARK] = str(self.lwt_info.lwt_meta_bit_mask) @@ -92,6 +105,10 @@ def parse_sharding_info(options): def parse_tablets_info(options): return TABLETS_ROUTING_V1 in options + @staticmethod + def parse_tablets_v2_info(options): + return TABLETS_ROUTING_V2 in options + @staticmethod def parse_use_metadata_id(options): """Return True if the ``SCYLLA_USE_METADATA_ID`` extension is advertised in ``options``.""" diff --git a/tests/unit/test_protocol_features.py b/tests/unit/test_protocol_features.py index 387583680b..915f8b84fd 100644 --- a/tests/unit/test_protocol_features.py +++ b/tests/unit/test_protocol_features.py @@ -2,7 +2,7 @@ import logging -from cassandra.protocol_features import ProtocolFeatures +from cassandra.protocol_features import ProtocolFeatures, TABLETS_ROUTING_V1, TABLETS_ROUTING_V2 LOGGER = logging.getLogger(__name__) @@ -57,3 +57,45 @@ def test_use_metadata_id_not_in_startup_when_not_negotiated(self): startup = {} protocol_features.add_startup_options(startup) assert 'SCYLLA_USE_METADATA_ID' not in startup + + def test_tablets_routing_v2_negotiation(self): + """V2 is detected from SUPPORTED and subsumes V1 in STARTUP options.""" + options = { + TABLETS_ROUTING_V1: [''], + TABLETS_ROUTING_V2: [''], + } + features = ProtocolFeatures.parse_from_supported(options) + assert features.tablets_routing_v1 is True + assert features.tablets_routing_v2 is True + + # V2 subsumes V1: only TABLETS_ROUTING_V2 should appear in startup. + startup = {} + features.add_startup_options(startup) + assert TABLETS_ROUTING_V2 in startup + assert TABLETS_ROUTING_V1 not in startup + + def test_tablets_routing_v1_only(self): + """When server only advertises V1, only V1 is negotiated.""" + options = { + TABLETS_ROUTING_V1: [''], + } + features = ProtocolFeatures.parse_from_supported(options) + assert features.tablets_routing_v1 is True + assert features.tablets_routing_v2 is False + + startup = {} + features.add_startup_options(startup) + assert TABLETS_ROUTING_V1 in startup + assert TABLETS_ROUTING_V2 not in startup + + def test_no_tablets_routing(self): + """When server advertises neither V1 nor V2.""" + options = {} + features = ProtocolFeatures.parse_from_supported(options) + assert features.tablets_routing_v1 is False + assert features.tablets_routing_v2 is False + + startup = {} + features.add_startup_options(startup) + assert TABLETS_ROUTING_V1 not in startup + assert TABLETS_ROUTING_V2 not in startup From 3f8aa317a22ffac2a708a8846a703107c034c58c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dawid=20M=C4=99drek?= Date: Wed, 8 Jul 2026 14:03:11 +0200 Subject: [PATCH 2/9] Track tablet_version and encode the tablet_version_block Store the server-provided 64-bit tablet_version on each cached Tablet and add helpers to encode it into the one-byte tablet_version_block exchanged on the wire. The version stays None until learned: on a cold start, and on a TABLETS_ROUTING_V1 connection, which never reports one. * Tablet.from_row normalizes the version to an unsigned 64-bit value. The server sends an unsigned hash, but the driver deserializes the payload field as a signed long, so the raw value can come back negative; masking to [0, 2**64) keeps the nibble extraction in choose_tablet_version_block consistent with the server's unsigned layout. * choose_tablet_version_block() packs a randomly chosen block index in the high nibble and that block's value in the low nibble, matching the server's locator::compare_tablet_version_block layout. Blocks are indexed from the least significant bits, so block i covers bits [i*4, i*4 + 4) of the version. A random index avoids any shared mutable counter on the hot path while still probing every nibble often enough to detect a server-side version change quickly. * random_tablet_version_block() returns a random byte for cold start, when no version is cached yet. --- cassandra/tablets.py | 46 +++++++++++++++++++++++--- tests/unit/test_tablets.py | 67 +++++++++++++++++++++++++++++++++++++- 2 files changed, 107 insertions(+), 6 deletions(-) diff --git a/cassandra/tablets.py b/cassandra/tablets.py index 96e61a50c2..216d802061 100644 --- a/cassandra/tablets.py +++ b/cassandra/tablets.py @@ -1,5 +1,6 @@ from bisect import bisect_left from operator import attrgetter +from random import getrandbits from threading import Lock from typing import Optional from uuid import UUID @@ -9,6 +10,32 @@ _get_last_token = attrgetter("last_token") +def choose_tablet_version_block(tablet_version: int) -> int: + """ + Encode a tablet_version_block byte from a cached tablet_version. + Picks a block index at random across calls. + Returns an int in [0, 255]. + + The byte layout: the high nibble is the block index, the low nibble is the value + of that block. Blocks are indexed from the least significant bits to the most + significant ones, so block `idx` occupies bits [idx*4, idx*4 + 4). + """ + # Pick the block index in [0, 15]; getrandbits(4) is a fast C call with no + # application-level shared state. + idx = getrandbits(4) + # Extract the 4-bit nibble at block index `idx` (0 = least significant). + shift = idx * 4 + nibble = (tablet_version >> shift) & 0xF + return (idx << 4) | nibble + + +def random_tablet_version_block() -> int: + """ + Generate a random tablet_version_block byte for cold start. + """ + return getrandbits(8) + + class Tablet(object): """ Represents a single ScyllaDB tablet. @@ -18,15 +45,19 @@ class Tablet(object): first_token = 0 last_token = 0 replicas = None + # uint64 hash; None means unknown -- a cold start, or a tablet learned over + # TABLETS_ROUTING_V1, which does not report a version. + tablet_version = None - def __init__(self, first_token=0, last_token=0, replicas=None): + def __init__(self, first_token=0, last_token=0, replicas=None, tablet_version=None): self.first_token = first_token self.last_token = last_token self.replicas = replicas + self.tablet_version = tablet_version def __str__(self): - return "" \ - % (self.first_token, self.last_token, self.replicas) + return "" \ + % (self.first_token, self.last_token, self.replicas, self.tablet_version) __repr__ = __str__ @staticmethod @@ -34,9 +65,14 @@ def _is_valid_tablet(replicas): return replicas is not None and len(replicas) != 0 @staticmethod - def from_row(first_token, last_token, replicas): + def from_row(first_token, last_token, replicas, tablet_version=None): if Tablet._is_valid_tablet(replicas): - tablet = Tablet(first_token, last_token, replicas) + if tablet_version is not None: + # tablet_version is an unsigned 64-bit value, but it is + # deserialized from the wire as a signed LongType; normalize it + # back to unsigned so it matches the server's representation. + tablet_version &= 0xFFFFFFFFFFFFFFFF + tablet = Tablet(first_token, last_token, replicas, tablet_version) return tablet return None diff --git a/tests/unit/test_tablets.py b/tests/unit/test_tablets.py index 7a40e7de4d..87478af46e 100644 --- a/tests/unit/test_tablets.py +++ b/tests/unit/test_tablets.py @@ -1,6 +1,6 @@ import unittest -from cassandra.tablets import Tablets, Tablet +from cassandra.tablets import Tablets, Tablet, choose_tablet_version_block, random_tablet_version_block class TabletsTest(unittest.TestCase): def compare_ranges(self, tablets, ranges): @@ -124,3 +124,68 @@ def __init__(self, v): # Token value 50 is not > first_token (100) of the tablet whose # last_token (200) is >= 50, so no match. self.assertIsNone(tablets.get_tablet_for_key("ks", "tb", Token(50))) + + +class TabletVersionBlockTest(unittest.TestCase): + """Tests for tablet_version_block encoding used by TABLETS_ROUTING_V2.""" + + def _server_block_matches(self, version, block): + """Reimplements the server's locator::compare_tablet_version_block.""" + block_value = block & 0x0F + block_index = (block & 0xF0) >> 4 + hash_block = (version >> (block_index * 4)) & 0x0F + return hash_block == block_value + + def test_choose_tablet_version_block_matches_server(self): + """Every block produced by the driver must match the server's check.""" + version = 0x0123456789ABCDEF + # The index is chosen randomly; sample enough times to exercise many indices. + for _ in range(256): + block = choose_tablet_version_block(version) + self.assertTrue(self._server_block_matches(version, block), + f"Block 0x{block:02X} did not match server check for version 0x{version:016X}") + + def test_choose_tablet_version_block_covers_all_indices(self): + """Over many calls the random index selection should probe every block + index, so that any server-side version change is eventually detected.""" + version = 0xFFFFFFFFFFFFFFFF # All nibbles are 0xF + seen_indices = set() + # 16 indices; 1000 draws makes a missing index astronomically unlikely. + for _ in range(1000): + block = choose_tablet_version_block(version) + seen_indices.add((block >> 4) & 0xF) + self.assertEqual(seen_indices, set(range(16))) + + def test_choose_tablet_version_block_matches_server_for_signed_version(self): + """tablet_version is decoded as a *signed* 64-bit int (LongType), so a + version with the high bit set is stored negative in the driver while the + server treats it as unsigned. The block the driver emits must still match + the server's check computed on the unsigned value (sign-boundary guard).""" + for unsigned in (0x8000000000000000, 0xDEADBEEFCAFEBABE, 0xFFFFFFFFFFFFFFFF): + signed = unsigned - (1 << 64) # how LongType stores a high-bit value + self.assertLess(signed, 0) + # Sample enough times to exercise every one of the 16 block indices. + for _ in range(256): + block = choose_tablet_version_block(signed) + self.assertTrue( + self._server_block_matches(unsigned, block), + f"signed version {signed} (unsigned 0x{unsigned:016X}) produced " + f"block 0x{block:02X} that failed the server check") + + def test_random_tablet_version_block_returns_byte(self): + """Verify random_tablet_version_block returns a value in [0, 255].""" + for _ in range(100): + block = random_tablet_version_block() + self.assertIsInstance(block, int) + self.assertGreaterEqual(block, 0) + self.assertLessEqual(block, 255) + + def test_from_row_stores_tablet_version(self): + """Tablet.from_row stores the tablet_version it is given (the V2 payload field).""" + version = 0xDEADBEEFCAFEBABE + tablet = Tablet.from_row(-100, 100, [("host1", 0), ("host2", 1)], tablet_version=version) + self.assertIsNotNone(tablet) + self.assertEqual(tablet.tablet_version, version) + self.assertEqual(tablet.first_token, -100) + self.assertEqual(tablet.last_token, 100) + self.assertEqual(len(tablet.replicas), 2) From f004e771f48fb71f3dd2e8bdfc740daeb8e0f10a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dawid=20M=C4=99drek?= Date: Wed, 8 Jul 2026 22:44:45 +0200 Subject: [PATCH 3/9] Compute and serialize the tablet_version_block per EXECUTE With TABLETS_ROUTING_V2 the server returns, on a tablet_version mismatch, the tablet's replica set plus the new tablet_version, so the driver can keep its routing cache fresh without the per-response overhead v1 incurs. * Every EXECUTE on a V2 connection carries a tablet_version_block computed from the cached version -- or a random byte on cold start (a token-aware request with no cached version yet), or 0 for a non-token-aware request, which the server never version-checks. * The routing key and its ring token are resolved once per request, in _create_response_future, and handed to both consumers that need them while sending: the tablet_version_block here and shard selection in HostConnection. This keeps cluster-dependent state off the statement, which a caller may share between concurrent requests. The cached tablet is likewise looked up once -- the cache is mutable, so a second lookup could disagree with the first. * 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. 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 the send path looks it up by. * HostConnection.tablets_routing_v1 becomes supports_tablet_routing: shard selection is identical under both versions, since the request goes to this host either way and the pool picks the shard this host owns for the tablet. Refs: SCYLLADB-288 Refs: SCYLLADB-291 --- cassandra/cluster.py | 143 ++++++++++++++++++++++++----- cassandra/pool.py | 27 ++++-- tests/unit/test_response_future.py | 21 +++-- tests/unit/test_tablets.py | 63 +++++++++++++ 4 files changed, 213 insertions(+), 41 deletions(-) diff --git a/cassandra/cluster.py b/cassandra/cluster.py index 88c8d2707a..5d983e4d44 100644 --- a/cassandra/cluster.py +++ b/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,6 +4801,7 @@ 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 @@ -4743,7 +4809,7 @@ class ResponseFuture(object): 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: diff --git a/cassandra/pool.py b/cassandra/pool.py index 176751f60a..39425d8ac7 100644 --- a/cassandra/pool.py +++ b/cassandra/pool.py @@ -389,7 +389,7 @@ class HostConnection(object): # the number below, all excess connections will be closed. max_excess_connections_per_shard_multiplier = 3 - tablets_routing_v1 = False + supports_tablet_routing = False def __init__(self, host, host_distance, session): self.host = host @@ -436,11 +436,13 @@ def __init__(self, host, host_distance, session): if first_connection.features.sharding_info and not self._session.cluster.shard_aware_options.disable: self.host.sharding_info = first_connection.features.sharding_info self._open_connections_for_all_shards(first_connection.features.shard_id) - self.tablets_routing_v1 = first_connection.features.tablets_routing_v1 + + self.supports_tablet_routing = first_connection.features.tablets_routing_v1 \ + or first_connection.features.tablets_routing_v2 log.debug("Finished initializing connection for host %s", self.host) - def _get_connection_for_routing_key(self, routing_key=None, keyspace=None, table=None): + def _get_connection_for_routing_key(self, routing_key=None, keyspace=None, table=None, routing_token=None): if self.is_shutdown: raise ConnectionException( "Pool for %s is shutdown" % (self.host,), self.host) @@ -450,15 +452,20 @@ def _get_connection_for_routing_key(self, routing_key=None, keyspace=None, table shard_id = None if not self._session.cluster.shard_aware_options.disable and self.host.sharding_info and routing_key: - t = self._session.cluster.metadata.token_map.token_class.from_key(routing_key) - - shard_id = None - if self.tablets_routing_v1 and table is not None: + # Reuse the token computed once for this request when available, so + # the routing-key hash runs once per request instead of again here; + # fall back to hashing the routing key directly otherwise. + t = routing_token + if t is None: + t = self._session.cluster.metadata.token_map.token_class.from_key(routing_key) + if self.supports_tablet_routing and table is not None: if keyspace is None: keyspace = self._keyspace tablet = self._session.cluster.metadata._tablets.get_tablet_for_key(keyspace, table, t) + # In both V1 and V2 the request is sent to this host, so we pick + # the shard that this host owns for the tablet. if tablet is not None: for replica in tablet.replicas: if replica[0] == self.host.host_id: @@ -506,15 +513,15 @@ def _get_connection_for_routing_key(self, routing_key=None, keyspace=None, table return random.choice(active_connections) return random.choice(list(self._connections.values())) - def borrow_connection(self, timeout, routing_key=None, keyspace=None, table=None): - conn = self._get_connection_for_routing_key(routing_key, keyspace, table) + def borrow_connection(self, timeout, routing_key=None, keyspace=None, table=None, routing_token=None): + conn = self._get_connection_for_routing_key(routing_key, keyspace, table, routing_token) start = time.time() remaining = timeout last_retry = False while True: if conn.is_closed: # The connection might have been closed in the meantime - if so, try again - conn = self._get_connection_for_routing_key(routing_key, keyspace, table) + conn = self._get_connection_for_routing_key(routing_key, keyspace, table, routing_token) with conn.lock: if (not conn.is_closed or last_retry) and conn.in_flight < conn.max_request_id: # On last retry we ignore connection status, since it is better to return closed connection than diff --git a/tests/unit/test_response_future.py b/tests/unit/test_response_future.py index 232ecf6585..d71943ec04 100644 --- a/tests/unit/test_response_future.py +++ b/tests/unit/test_response_future.py @@ -41,7 +41,6 @@ class ResponseFutureTests(unittest.TestCase): def make_basic_session(self): s = Mock(spec=Session) s.row_factory = lambda col_names, rows: [(col_names, rows)] - s.cluster.control_connection._tablets_routing_v1 = False s.cluster.allow_control_connection_query_fallback = ControlConnectionQueryFallback.Disabled return s @@ -65,6 +64,11 @@ def make_control_connection(self): connection.is_control_connection = True connection.get_request_id.return_value = 7 connection.send_msg.return_value = 128 + # These tests exercise control-connection query fallback, not tablet + # routing; default the tablet features off so _set_result skips + # tablet-payload parsing for the mocked responses. + connection.features.tablets_routing_v2 = False + connection.features.tablets_routing_v1 = False return connection def make_session(self): @@ -94,7 +98,7 @@ def test_result_message(self): rf.send_request() rf.session._pools.get.assert_called_once_with('ip1') - pool.borrow_connection.assert_called_once_with(timeout=ANY, routing_key=ANY, keyspace=ANY, table=ANY) + pool.borrow_connection.assert_called_once_with(timeout=ANY, routing_key=ANY, keyspace=ANY, table=ANY, routing_token=ANY) connection.send_msg.assert_called_once_with(rf.message, 1, cb=ANY, encoder=ProtocolHandler.encode_message, decoder=ProtocolHandler.decode_message, result_metadata=[]) @@ -138,6 +142,9 @@ def test_schema_change_result(self): kind=RESULT_KIND_SCHEMA_CHANGE, schema_change_event=event_results) connection = Mock() + # Skip tablet-payload parsing for this mocked response/connection pair. + connection.features.tablets_routing_v2 = False + connection.features.tablets_routing_v1 = False rf._set_result(None, connection, None, result) session.submit.assert_called_once_with(ANY, ANY, rf, connection, **event_results) @@ -285,7 +292,7 @@ def test_retry_policy_says_retry(self): rf.send_request() rf.session._pools.get.assert_called_once_with('ip1') - pool.borrow_connection.assert_called_once_with(timeout=ANY, routing_key=ANY, keyspace=ANY, table=ANY) + pool.borrow_connection.assert_called_once_with(timeout=ANY, routing_key=ANY, keyspace=ANY, table=ANY, routing_token=ANY) connection.send_msg.assert_called_once_with(rf.message, 1, cb=ANY, encoder=ProtocolHandler.encode_message, decoder=ProtocolHandler.decode_message, result_metadata=[]) result = Mock(spec=UnavailableErrorMessage, info={}) @@ -304,7 +311,7 @@ def test_retry_policy_says_retry(self): # it should try again with the same host since this was # an UnavailableException rf.session._pools.get.assert_called_with(host) - pool.borrow_connection.assert_called_with(timeout=ANY, routing_key=ANY, keyspace=ANY, table=ANY) + pool.borrow_connection.assert_called_with(timeout=ANY, routing_key=ANY, keyspace=ANY, table=ANY, routing_token=ANY) connection.send_msg.assert_called_with(rf.message, 2, cb=ANY, encoder=ProtocolHandler.encode_message, decoder=ProtocolHandler.decode_message, result_metadata=[]) def test_retry_with_different_host(self): @@ -319,7 +326,7 @@ def test_retry_with_different_host(self): rf.send_request() rf.session._pools.get.assert_called_once_with('ip1') - pool.borrow_connection.assert_called_once_with(timeout=ANY, routing_key=ANY, keyspace=ANY, table=ANY) + pool.borrow_connection.assert_called_once_with(timeout=ANY, routing_key=ANY, keyspace=ANY, table=ANY, routing_token=ANY) connection.send_msg.assert_called_once_with(rf.message, 1, cb=ANY, encoder=ProtocolHandler.encode_message, decoder=ProtocolHandler.decode_message, result_metadata=[]) assert ConsistencyLevel.QUORUM == rf.message.consistency_level @@ -338,7 +345,7 @@ def test_retry_with_different_host(self): # it should try with a different host rf.session._pools.get.assert_called_with('ip2') - pool.borrow_connection.assert_called_with(timeout=ANY, routing_key=ANY, keyspace=ANY, table=ANY) + pool.borrow_connection.assert_called_with(timeout=ANY, routing_key=ANY, keyspace=ANY, table=ANY, routing_token=ANY) connection.send_msg.assert_called_with(rf.message, 2, cb=ANY, encoder=ProtocolHandler.encode_message, decoder=ProtocolHandler.decode_message, result_metadata=[]) # the consistency level should be the same @@ -1055,7 +1062,7 @@ def test_single_host_query_plan_exhausted_after_one_retry(self): # Verify initial request was sent rf.session._pools.get.assert_called_once_with(specific_host) - pool.borrow_connection.assert_called_once_with(timeout=ANY, routing_key=ANY, keyspace=ANY, table=ANY) + pool.borrow_connection.assert_called_once_with(timeout=ANY, routing_key=ANY, keyspace=ANY, table=ANY, routing_token=ANY) connection.send_msg.assert_called_once_with(rf.message, 1, cb=ANY, encoder=ProtocolHandler.encode_message, decoder=ProtocolHandler.decode_message, result_metadata=[]) # Simulate a ServerError response (which triggers RETRY_NEXT_HOST by default) diff --git a/tests/unit/test_tablets.py b/tests/unit/test_tablets.py index 87478af46e..f77d163eb8 100644 --- a/tests/unit/test_tablets.py +++ b/tests/unit/test_tablets.py @@ -1,5 +1,9 @@ import unittest +from io import BytesIO +from cassandra import ConsistencyLevel, ProtocolVersion +from cassandra.protocol import ExecuteMessage +from cassandra.protocol_features import ProtocolFeatures from cassandra.tablets import Tablets, Tablet, choose_tablet_version_block, random_tablet_version_block class TabletsTest(unittest.TestCase): @@ -189,3 +193,62 @@ def test_from_row_stores_tablet_version(self): self.assertEqual(tablet.first_token, -100) self.assertEqual(tablet.last_token, 100) self.assertEqual(len(tablet.replicas), 2) + + +class ExecuteMessageSerializationTest(unittest.TestCase): + """ExecuteMessage.send_body decides whether to emit the tablet_version_block + from the serving connection's negotiated protocol features, so the message no + longer has to be copied per send attempt (TABLETS_ROUTING_V2).""" + + # V4 keeps the encoding minimal: no prepared-metadata id, single-byte flags. + PROTOCOL_VERSION = ProtocolVersion.V4 + + def _make_message(self, tablet_version_block): + return ExecuteMessage( + query_id=b"\x01\x02\x03\x04", + query_params=[], + consistency_level=ConsistencyLevel.ONE, + tablet_version_block=tablet_version_block, + ) + + def _encode_body(self, message, protocol_features): + f = BytesIO() + message.send_body(f, self.PROTOCOL_VERSION, protocol_features) + return f.getvalue() + + def test_block_appended_on_v2_connection(self): + """A V2 connection appends exactly one trailing byte carrying the + precomputed tablet_version_block.""" + message = self._make_message(0x7A) + v2 = self._encode_body(message, ProtocolFeatures(tablets_routing_v2=True)) + plain = self._encode_body(message, ProtocolFeatures(tablets_routing_v2=False)) + self.assertEqual(v2, plain + bytes([0x7A])) + + def test_block_absent_without_v2(self): + """No trailing byte when the connection did not negotiate V2, and passing + no features at all is equivalent to V2 being off.""" + message = self._make_message(0x7A) + no_features = self._encode_body(message, None) + v2_off = self._encode_body(message, ProtocolFeatures(tablets_routing_v2=False)) + self.assertEqual(no_features, v2_off) + self.assertNotIn(bytes([0x7A]), no_features[-1:]) + + def test_missing_block_coalesces_to_zero_on_v2(self): + """On a V2 connection the server reads exactly one trailing byte per + EXECUTE, so a message whose block was never computed must still emit a + zero byte to keep the frame in sync.""" + message = self._make_message(None) + v2 = self._encode_body(message, ProtocolFeatures(tablets_routing_v2=True)) + plain = self._encode_body(message, ProtocolFeatures(tablets_routing_v2=False)) + self.assertEqual(v2, plain + bytes([0x00])) + + def test_same_message_encodes_consistently_across_connections(self): + """The same shared message instance yields the V2 or non-V2 framing purely + from the features argument, so it is safe to encode concurrently on + connections with different capabilities without copying it.""" + message = self._make_message(0x3C) + first = self._encode_body(message, ProtocolFeatures(tablets_routing_v2=True)) + second_plain = self._encode_body(message, ProtocolFeatures(tablets_routing_v2=False)) + first_again = self._encode_body(message, ProtocolFeatures(tablets_routing_v2=True)) + self.assertEqual(first, first_again) + self.assertEqual(first, second_plain + bytes([0x3C])) From 25e719f1b59c00ef4f80121d7b842958e88c2e70 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dawid=20M=C4=99drek?= Date: Wed, 8 Jul 2026 16:25:47 +0200 Subject: [PATCH 4/9] Add integration tests for TABLETS_ROUTING_V2 Cover the end-to-end behaviour against a live ScyllaDB started with the `strongly-consistent-tables` experimental feature: v2 negotiation, payload-driven cache population, and the tablet_version_block matching rules (no payload on a matching block, exactly one matching value per index, and v2 taking precedence over v1 on a wrong-shard request). The last of those needs a connection that negotiated both extensions, which the driver never does on its own, so the test patches ProtocolFeatures.add_startup_options. The patch delegates to the real implementation and only adds v1 on top. Enumerating the options itself would silently stop requesting any extension added later while ProtocolFeatures still reported it as negotiated -- that is parsed from SUPPORTED, not from what STARTUP asked for -- and an extension that changes the frame layout, such as SCYLLA_USE_METADATA_ID, would then desynchronize every request on the connection. --- .../standard/test_tablets_routing_v2.py | 407 ++++++++++++++++++ 1 file changed, 407 insertions(+) create mode 100644 tests/integration/standard/test_tablets_routing_v2.py diff --git a/tests/integration/standard/test_tablets_routing_v2.py b/tests/integration/standard/test_tablets_routing_v2.py new file mode 100644 index 0000000000..22a53ed79d --- /dev/null +++ b/tests/integration/standard/test_tablets_routing_v2.py @@ -0,0 +1,407 @@ +""" +End-to-end tests for TABLETS_ROUTING_V2 against a V2-capable Scylla build. + +Unlike the unit tests in tests/unit/test_tablets.py and tests/unit/test_policies.py, +these tests cross the driver<->server boundary: they validate that the driver +negotiates the extension, parses the server's `tablets-routing-v2` payload with +the correct field layout, and that the tablet_version_block it sends actually +matches the server's encoding. + +The whole module is opt-in: the server only advertises the extension when started +with the `strongly-consistent-tables` experimental feature, and it is exchanged on +the wire under the name `TABLETS_ROUTING_V2_EXPERIMENTAL`. When run against a +server that does not advertise it (e.g. a released Scylla), every test is skipped. +""" + +from contextlib import contextmanager + +import pytest + +import cassandra.cqltypes as types +from cassandra import ConsistencyLevel +from cassandra.cluster import Cluster, ExecutionProfile, EXEC_PROFILE_DEFAULT +from cassandra.policies import ConstantReconnectionPolicy, RoundRobinPolicy, TokenAwarePolicy +from cassandra.protocol import ExecuteMessage +from cassandra.protocol_features import ( + ProtocolFeatures, TABLETS_ROUTING_V1, TABLETS_ROUTING_V2, +) + +from tests.integration import PROTOCOL_VERSION, use_cluster + + +def setup_module(module): + try: + # Use a single DC with three racks. + use_cluster('tablets_routing_v2', {"dc1": [1, 1, 1]}, start=True, set_keyspace=False, + configuration_options={ + # `strongly-consistent-tables` is what gates the server's + # advertisement of TABLETS_ROUTING_V2_EXPERIMENTAL. + 'experimental_features': ['udf', 'strongly-consistent-tables'], + }) + except Exception as exc: + pytest.skip("Could not start a Scylla cluster with the " + f"'strongly-consistent-tables' experimental feature: {exc}", + allow_module_level=True) + + +_add_startup_options = ProtocolFeatures.add_startup_options + + +def _startup_with_both_extensions(self, options): + """ + Drop-in replacement for ProtocolFeatures.add_startup_options that negotiates + BOTH tablets_routing_v1 and tablets_routing_v2 on the same connection. + + The real driver makes the two mutually exclusive (V2 wins). Forcing both lets + us prove the server-side precedence rules: scylla checks V2 first and only + falls back to V1 when V2 is not set. + + Every other extension must be negotiated exactly as the driver would, so this + delegates to the real implementation and only adds V1 on top of the V2 it + already requested. Enumerating the options here instead would silently stop + requesting any extension added later, while ProtocolFeatures still reports it + as negotiated (it is parsed from SUPPORTED, not from what STARTUP asked for) -- + and an extension that changes the frame layout, such as + SCYLLA_USE_METADATA_ID, then desynchronizes every request on the connection. + """ + _add_startup_options(self, options) + if self.tablets_routing_v1: + options[TABLETS_ROUTING_V1] = "" + + +class TestTabletsRoutingV2Integration: + @classmethod + def setup_class(cls): + cls.cluster = Cluster(contact_points=["127.0.0.1", "127.0.0.2", "127.0.0.3"], + protocol_version=PROTOCOL_VERSION, + execution_profiles={ + EXEC_PROFILE_DEFAULT: ExecutionProfile( + load_balancing_policy=TokenAwarePolicy(RoundRobinPolicy())) + }, + reconnection_policy=ConstantReconnectionPolicy(1)) + # pytest does not call teardown_class when setup_class raises, so any exit + # from here must shut the Cluster down explicitly or it leaks threads and + # sockets into later test modules. + try: + cls.session = cls.cluster.connect() + # A server without the 'strongly-consistent-tables' experimental + # feature (e.g. a released Scylla) still starts and connects, but it + # neither advertises TABLETS_ROUTING_V2 nor accepts the + # `consistency = 'global'` keyspace that _create_schema needs. Detect + # that here and skip the whole class, instead of letting _create_schema + # fail and erroring every test. + v2_negotiated = cls._v2_negotiated() + if v2_negotiated: + cls._create_schema(cls.session) + except Exception: + cls.cluster.shutdown() + raise + if not v2_negotiated: + cls.cluster.shutdown() + pytest.skip("Server does not support TABLETS_ROUTING_V2_EXPERIMENTAL. " + "It must be started with the 'strongly-consistent-tables' feature " + "and offer support for the protocol extension.") + + @classmethod + def teardown_class(cls): + cls.cluster.shutdown() + + @classmethod + def _create_schema(cls, session): + session.execute("DROP KEYSPACE IF EXISTS test_v2") + session.execute( + """ + CREATE KEYSPACE test_v2 + WITH replication = {'class': 'NetworkTopologyStrategy', 'replication_factor': 2} + AND tablets = {'initial': 8} + """) + session.execute("CREATE TABLE test_v2.t (pk int PRIMARY KEY, v int)") + prepared = session.prepare("INSERT INTO test_v2.t (pk, v) VALUES (?, ?)") + for i in range(50): + session.execute(prepared.bind((i, i))) + + # -- helpers ---------------------------------------------------------------- + + @classmethod + def _v2_negotiated(cls): + connection = cls.session.cluster.control_connection._connection + return bool(connection and connection.features.tablets_routing_v2) + + def _cached_tablet(self, bound): + md = self.session.cluster.metadata + token = md.token_map.token_class.from_key(bound.routing_key) + tablet = md._tablets.get_tablet_for_key(bound.keyspace, bound.table, token) + return tablet, token + + def _ensure_cached(self, bound, attempts=30): + """ + Drive requests until the V2 routing cache is populated for `bound`. + + On a cold start the driver sends a *random* tablet_version_block, which + only matches the server ~1/16 of the time; on a mismatch the server + returns routing info and the cache is filled. We retry until that happens. + """ + for _ in range(attempts): + self.session.execute(bound) + tablet, _token = self._cached_tablet(bound) + if tablet is not None and tablet.tablet_version is not None: + return tablet + raise AssertionError("V2 routing cache was never populated; the server " + "never returned a 'tablets-routing-v2' payload") + + # -- tests ------------------------------------------------------------------ + + def test_v2_is_negotiated(self): + # Every per-host pool must report tablet-routing support, and every live + # connection in it must have negotiated V2 -- V2 gates the per-connection + # EXECUTE framing. + for pool in self.session._pools.values(): + assert pool.supports_tablet_routing is True + for conn in pool._connections.values(): + assert conn.features.tablets_routing_v2 is True + + def test_v2_payload_populates_cache_with_valid_fields(self): + """Test guarding against the payload tuple being decoded out of order.""" + select = self.session.prepare("SELECT v FROM test_v2.t WHERE pk = ?") + bound = select.bind([2]) + + tablet = self._ensure_cached(bound) + _, token = self._cached_tablet(bound) + + # If the tuple were decoded in the wrong order, first_token/last_token + # would actually carry the version / replica list and these invariants + # would not hold. + assert tablet.tablet_version is not None + # tablet_version is an unsigned 64-bit value; a signedness bug in the + # decode path would surface here as a negative or out-of-range number. + assert 0 <= tablet.tablet_version <= 2 ** 64 - 1 + assert tablet.first_token <= tablet.last_token + # get_tablet_for_key matches first_token < token <= last_token. + assert tablet.first_token < token.value and token.value <= tablet.last_token + + # Replicas must be real hosts known to the cluster with sane shard ids. + known_host_ids = {h.host_id for h in self.session.cluster.metadata.all_hosts()} + assert tablet.replicas, "tablet has no replicas" + for host_id, shard in tablet.replicas: + assert host_id in known_host_ids, \ + f"replica host_id {host_id} is not a known host (corrupt payload?)" + assert isinstance(shard, int) and shard >= 0 + + def test_matching_block_yields_no_payload(self): + """Test guarding against a wrong tablet_version_block bit-shift.""" + select = self.session.prepare("SELECT v FROM test_v2.t WHERE pk = ?") + bound = select.bind([7]) + + # Populate the cache so the driver knows the current tablet_version. + self._ensure_cached(bound) + + # The next request carries a block derived from the cached version. If the + # driver's encoding agrees with the server, the versions match and NO + # routing payload is returned. A wrong shift would mismatch and the server + # would keep returning routing info. + result = self.session.execute(bound) + assert result.one() is not None + payload = result.response_future.custom_payload + assert not (payload and 'tablets-routing-v2' in payload), ( + "Server returned routing info despite a cached, up-to-date " + "tablet_version; the driver's tablet_version_block encoding likely " + "disagrees with the server (locator::compare_tablet_version_block)") + + # -- low-level helpers ------------------------------------------------------ + + @staticmethod + def _right_block(version, idx=0): + """ + Build a tablet_version_block that the server will accept as a match for + block `idx` of `version` (high nibble = index, low nibble = that nibble + of the version). + """ + idx &= 0xF + return (idx << 4) | ((version >> (idx * 4)) & 0xF) + + def _send_raw_execute(self, conn, bound, tablet_version_block): + """ + Send an EXECUTE directly on a specific shard connection with a chosen + tablet_version_block (or None to omit the byte entirely, i.e. behave like + the pre-V2 protocol), and return the decoded response message. + + This bypasses ResponseFuture/load balancing so we control exactly which + node+shard the request hits and which byte is on the wire; it also avoids + polluting the driver's tablet cache. + """ + ps = bound.prepared_statement + msg = ExecuteMessage( + ps.query_id, bound.values, ConsistencyLevel.LOCAL_ONE, + serial_consistency_level=None, fetch_size=None, paging_state=None, + timestamp=None, skip_meta=False, + result_metadata_id=ps.result_metadata_id, + tablet_version_block=tablet_version_block) + return conn.wait_for_response(msg, timeout=30) + + def _decode_v2_payload(self, payload): + ctype = types.lookup_casstype( + 'TupleType(LongType, LongType, ListType(TupleType(UUIDType, Int32Type)), LongType)') + info = ctype.from_binary(payload['tablets-routing-v2'], self.cluster.protocol_version) + # LongType decodes as signed, but tablet_version is an unsigned 64-bit + # value; mask it the same way Tablet.from_row does so the decoded value + # matches what the driver cached. + return {'first_token': info[0], 'last_token': info[1], + 'replicas': info[2], + 'tablet_version': info[3] & 0xFFFFFFFFFFFFFFFF} + + def _any_connection(self): + for pool in self.session._pools.values(): + for conn in pool._connections.values(): + return conn + raise AssertionError("no shard connections available") + + @staticmethod + def _all_shard_connections(session): + for host, pool in session._pools.items(): + for shard, conn in pool._connections.items(): + yield host, shard, conn + + @staticmethod + def _wait_for_shard_connections(session, timeout=15): + """Wait until each pool has filled its shard-aware connections (background).""" + import time + deadline = time.time() + timeout + while time.time() < deadline: + if all( + len(pool._connections) >= (min(host.sharding_info.shards_count, 2) + if host.sharding_info else 1) + for host, pool in session._pools.items() + ): + return + time.sleep(0.05) + raise AssertionError(f"Shard-aware connection pools did not fill within {timeout}s") + + @contextmanager + def _cluster_with_v1_and_v2(self): + """ + Yield a (cluster, session) whose connections negotiated BOTH V1 and V2. + Restores the original startup behavior and shuts the cluster down on exit. + """ + original = ProtocolFeatures.add_startup_options + ProtocolFeatures.add_startup_options = _startup_with_both_extensions + + cluster = None + try: + cluster = Cluster(contact_points=["127.0.0.1", "127.0.0.2", "127.0.0.3"], + protocol_version=PROTOCOL_VERSION, + execution_profiles={ + EXEC_PROFILE_DEFAULT: ExecutionProfile( + load_balancing_policy=TokenAwarePolicy(RoundRobinPolicy())) + }, + reconnection_policy=ConstantReconnectionPolicy(1)) + session = cluster.connect('test_v2') + self._wait_for_shard_connections(session) + yield cluster, session + finally: + ProtocolFeatures.add_startup_options = original + if cluster is not None: + cluster.shutdown() + + @staticmethod + def _find_replica_wrong_shard(session, tablet): + """ + Find a connection to a host that *is* a replica of `tablet` but on a shard + that the host does NOT own for it ("right node, wrong shard"). Returns + (host, owner_shard, wrong_shard, conn) or None if no host has >=2 shards. + """ + replica_shard = {host_id: shard for host_id, shard in tablet.replicas} + for host, pool in session._pools.items(): + owner = replica_shard.get(host.host_id) + if owner is None: + continue + for shard, conn in pool._connections.items(): + if shard != owner: + return host, owner, shard, conn + return None + + # -- scenario tests --------------------------------------------------------- + + def test_index0_all_block_values_exactly_one_match(self): + """ + Scenario 1: for block index 0, exactly one of the 16 possible values + matches the server's tablet_version; every other value is reported as a + mismatch carrying that same tablet_version, whose nibble 0 equals the + value that matched. + """ + select = self.session.prepare("SELECT v FROM test_v2.t WHERE pk = ?") + bound = select.bind([11]) + tablet = self._ensure_cached(bound) + version = tablet.tablet_version + + conn = self._any_connection() + + matched_values = [] + reported_versions = [] + for value in range(16): + block = value # index 0 -> high nibble 0, low nibble = value + resp = self._send_raw_execute(conn, bound, block) + payload = resp.custom_payload or {} + if 'tablets-routing-v2' in payload: + reported_versions.append(self._decode_v2_payload(payload)['tablet_version']) + else: + matched_values.append(value) + + # Exactly one value matches: the low nibble of the version. + assert matched_values == [version & 0xF], \ + f"expected exactly one matching block value, got {matched_values}" + # All 15 mismatches report the same tablet_version ... + assert len(reported_versions) == 15 + assert set(reported_versions) == {version} + # ... and that version's block-0 nibble is the value that matched. + assert (version & 0xF) == matched_values[0] + + def test_right_block_to_all_nodes_and_shards_never_returns_payload(self): + """ + Scenario 2: a correct tablet_version_block matches on every node and every + shard (the server's V2 check ignores shard), so no routing payload is ever + returned. + """ + select = self.session.prepare("SELECT v FROM test_v2.t WHERE pk = ?") + bound = select.bind([13]) + tablet = self._ensure_cached(bound) + version = tablet.tablet_version + + sent = 0 + for host, shard, conn in self._all_shard_connections(self.session): + # Vary the block index per shard to also exercise non-zero indices. + block = self._right_block(version, idx=shard) + resp = self._send_raw_execute(conn, bound, block) + payload = resp.custom_payload or {} + assert 'tablets-routing-v2' not in payload, ( + f"host {host} shard {shard} returned a routing payload for a correct " + "tablet_version_block") + sent += 1 + assert sent >= 1, "no shard connections were exercised" + + def test_v2_takes_precedence_over_v1_no_v1_payload_on_wrong_shard(self): + """ + Scenario 3: with BOTH V1 and V2 negotiated, send a correct V2 block to the + wrong shard. The server checks V2 first; since the block matches there is + no payload at all -- crucially no `tablets-routing-v1`, which V1 would have + emitted for a wrong-shard request. + """ + select = self.session.prepare("SELECT v FROM test_v2.t WHERE pk = ?") + bound = select.bind([17]) + tablet = self._ensure_cached(bound) + version = tablet.tablet_version + + with self._cluster_with_v1_and_v2() as (_cluster, session): + target = self._find_replica_wrong_shard(session, tablet) + if target is None: + pytest.skip("need a replica host with >=2 shards to target a wrong shard") + _host, _owner_shard, _wrong_shard, conn = target + assert conn.features.tablets_routing_v1 and conn.features.tablets_routing_v2, \ + "test setup failed: connection did not negotiate both V1 and V2" + + resp = self._send_raw_execute(conn, bound, self._right_block(version)) + payload = resp.custom_payload or {} + assert 'tablets-routing-v1' not in payload, ( + "server emitted V1 routing info despite V2 being negotiated; " + "V2 must take precedence (select_statement.cc)") + # The correct V2 block also means no V2 payload. + assert 'tablets-routing-v2' not in payload From 6d72d07ad3230338f2dee96b2e59ae084077ea0d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dawid=20M=C4=99drek?= Date: Wed, 8 Jul 2026 16:26:07 +0200 Subject: [PATCH 5/9] Document TABLETS_ROUTING_V2 tablet-version tracking Extend the "Tablet Awareness" section of the Scylla-specific guide to cover the V2 protocol extension: the per-connection negotiation and the tablet_version_block byte that lets the server skip re-sending routing information the driver already has. --- docs/scylla-specific.rst | 31 ++++++++++++++++++++++++++++++- 1 file changed, 30 insertions(+), 1 deletion(-) diff --git a/docs/scylla-specific.rst b/docs/scylla-specific.rst index 4f61846b4c..80071d5102 100644 --- a/docs/scylla-specific.rst +++ b/docs/scylla-specific.rst @@ -148,7 +148,7 @@ For more details on paging, see :ref:`query-paging`. Tablet Awareness ---------------- -**scylla-driver** is tablet-aware, which means that it is able to parse `TABLETS_ROUTING_V1` extension to ProtocolFeatures, recieve tablet information sent by Scylla in the `custom_payload` part of the `RESULT` message, and utilize it. +**scylla-driver** is tablet-aware, which means that it is able to parse the `TABLETS_ROUTING_V1` and `TABLETS_ROUTING_V2` extensions to ProtocolFeatures, receive tablet information sent by Scylla in the `custom_payload` part of the `RESULT` message, and utilize it. Thanks to this, queries to tablet-based tables are still shard-aware. Details on the scylla cql protocol extensions @@ -158,6 +158,35 @@ Details on the sending tablet information to the drivers https://github.com/scylladb/scylladb/blob/master/docs/dev/protocol-extensions.md#sending-tablet-info-to-the-drivers +Tablet version tracking +----------------------- + +When the cluster offers it, the driver negotiates ``TABLETS_ROUTING_V2`` in +preference to V1. The negotiation happens per connection, so V2 and V1 +connections can coexist in the same cluster; each connection uses whichever +extension its node offers. V2 adds tablet version tracking on top of V1, +invisible to application code. + +Every tablet now carries a ``tablet_version`` that +changes whenever its replica set is reconfigured. The driver caches the version +it last saw for each tablet and, on every prepared-statement execution over a V2 +connection, appends a single ``tablet_version_block`` byte derived from it. The +server returns updated routing information in the ``custom_payload`` only when +that byte shows the driver's cached view is stale, instead of attaching it to +every response. This keeps the cached routing information fresh while avoiding +the per-response overhead that V1 incurs. + +No configuration is required: as with V1, a ``TokenAwarePolicy`` is all that is +needed. + +.. note:: + + ``TABLETS_ROUTING_V2`` is still experimental: a Scylla node advertises it + (on the wire as ``TABLETS_ROUTING_V2_EXPERIMENTAL``) only when started with + the ``strongly-consistent-tables`` experimental feature enabled. A node + without it offers only ``TABLETS_ROUTING_V1``. + + Prepared Statement Metadata Caching (``SCYLLA_USE_METADATA_ID``) ---------------------------------------------------------------- From 198116794a46182e082efa10fad8bd0d24711206 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dawid=20M=C4=99drek?= Date: Wed, 8 Jul 2026 16:26:53 +0200 Subject: [PATCH 6/9] Record the private consistency mode on keyspace metadata Add KeyspaceMetadata._consistency_mode, derived from the per-keyspace `consistency` option in system_schema.scylla_keyspaces. It is a _ConsistencyMode enum -- EVENTUAL, LOCAL or GLOBAL -- so the mode the server reported is kept verbatim instead of being flattened into a boolean at parse time. The lookup is cached per schema refresh and degrades to EVENTUAL on non-Scylla clusters, on a control connection that did not negotiate TABLETS_ROUTING_V2, and on Scylla versions that lack the table or column. Scylla only implements `global` so far, so a keyspace's tablets have a Raft leader exactly when its mode is GLOBAL; `local` is reserved for a mode that does not exist yet and behaves like `eventual` everywhere. Callers that care compare against _ConsistencyMode.GLOBAL directly, so implementing `local` later only widens those comparisons and leaves the parser and the metadata untouched. A change of mode also invalidates the keyspace's cached tablets, the same way a replication-strategy change does: a tablet cached while the keyspace was eventually consistent carries no leader ordering and must not survive into a strongly-consistent keyspace, where it would be misread as a leader hint. Both names are underscore-prefixed to keep them private: they are not yet stable and we do not want to commit to a public API for them. --- cassandra/metadata.py | 111 +++++++++++++++++++++++++++- tests/unit/test_metadata.py | 142 +++++++++++++++++++++++++++++++++++- 2 files changed, 251 insertions(+), 2 deletions(-) diff --git a/cassandra/metadata.py b/cassandra/metadata.py index 43399b7152..fc9d3a87c0 100644 --- a/cassandra/metadata.py +++ b/cassandra/metadata.py @@ -16,6 +16,7 @@ from bisect import bisect_left from collections import defaultdict from collections.abc import Mapping +from enum import Enum from functools import total_ordering from hashlib import md5 import json @@ -194,7 +195,8 @@ def _update_keyspace(self, keyspace_meta, new_user_types=None): keyspace_meta.functions = old_keyspace_meta.functions keyspace_meta.aggregates = old_keyspace_meta.aggregates keyspace_meta.views = old_keyspace_meta.views - if (keyspace_meta.replication_strategy != old_keyspace_meta.replication_strategy): + if (keyspace_meta.replication_strategy != old_keyspace_meta.replication_strategy or + keyspace_meta._consistency_mode != old_keyspace_meta._consistency_mode): self._keyspace_updated(ks_name) else: self._keyspace_added(ks_name) @@ -733,6 +735,31 @@ def __eq__(self, other): return isinstance(other, LocalStrategy) +class _ConsistencyMode(Enum): + """ + Per-keyspace consistency option reported by ScyllaDB in + ``system_schema.scylla_keyspaces``. The server represents an + eventually-consistent keyspace as ``'eventual'`` or null. + + ScyllaDB only implements ``GLOBAL`` so far, so it is currently the only mode + under which a keyspace's tablets have a Raft leader. ``LOCAL`` is reserved + for a mode that does not exist yet, and until it does it behaves exactly like + ``EVENTUAL`` everywhere in the driver. + """ + EVENTUAL = 'eventual' + LOCAL = 'local' + GLOBAL = 'global' + + @classmethod + def from_string(cls, value: Optional[str]): + # Map the server's string (or a null/unknown value) to a member, + # treating anything unrecognized as eventually consistent. + try: + return cls(value) + except ValueError: + return cls.EVENTUAL + + class KeyspaceMetadata(object): """ A representation of the schema for a single keyspace. @@ -801,6 +828,15 @@ class KeyspaceMetadata(object): A string indicating whether a graph engine is enabled for this keyspace (Core/Classic). """ + _consistency_mode = _ConsistencyMode.EVENTUAL + """ + The consistency mode of the keyspace, derived from the the ``consistency`` + column ScyllaDB stores in ``system_schema.scylla_keyspaces``. + + Private and unstable: it backs leader-aware routing and is not part of the + public API, so the name and semantics may change. + """ + _exc_info = None """ set if metadata parsing failed """ @@ -815,6 +851,7 @@ def __init__(self, name, durable_writes, strategy_class, strategy_options, graph self.aggregates = {} self.views = {} self.graph_engine = graph_engine + self._consistency_mode = _ConsistencyMode.EVENTUAL @property def is_graph_enabled(self): @@ -861,6 +898,19 @@ def as_cql_query(self): ret = "CREATE KEYSPACE %s WITH replication = %s " % ( protect_name(self.name), self.replication_strategy.export_for_schema()) + + # TODO: Replace with `match` once the Python version is bumped to 3.10+. + if self._consistency_mode == _ConsistencyMode.EVENTUAL: + # Eventual consistency is the default. + pass + elif self._consistency_mode == _ConsistencyMode.LOCAL: + # Note: Local consistency mode is not implemented yet. + ret = ret + " AND consistency = 'local'" + elif self._consistency_mode == _ConsistencyMode.GLOBAL: + ret = ret + " AND consistency = 'global'" + else: + raise ValueError("Unknown consistency mode: %s" % self._consistency_mode) + ret = ret + (' AND durable_writes = %s' % ("true" if self.durable_writes else "false")) if self.graph_engine is not None: ret = ret + (" AND graph_engine = '%s'" % self.graph_engine) @@ -2577,6 +2627,11 @@ class SchemaParserV3(SchemaParserV22): _SELECT_AGGREGATES = "SELECT * FROM system_schema.aggregates" _SELECT_VIEWS = "SELECT * FROM system_schema.views" + # ScyllaDB-only: per-keyspace consistency option. The column is null for + # eventually-consistent keyspaces (and the whole table is absent on Cassandra + # and on Scylla versions without strongly-consistent tablets). + _SELECT_SCYLLA_KEYSPACES = "SELECT keyspace_name, consistency FROM system_schema.scylla_keyspaces" + def _is_not_scylla(self): """Check if NOT connected to ScyllaDB by checking for shard awareness.""" return getattr(getattr(self.connection, 'features', None), 'shard_id', None) is None @@ -2610,12 +2665,66 @@ def __init__(self, connection, timeout, fetch_size, metadata_request_timeout): self.indexes_result = [] self.keyspace_table_index_rows = defaultdict(lambda: defaultdict(list)) self.keyspace_view_rows = defaultdict(list) + self._scylla_consistency_cache = None + + def _get_scylla_keyspaces_consistency(self): + """ + Return a ``{keyspace_name: _ConsistencyMode}`` map read from + ``system_schema.scylla_keyspaces``. + + A keyspace absent from the map, or one whose consistency value is not + recognized, is treated as eventually consistent. Returns ``{}`` on + non-Scylla clusters, on a control connection that did not negotiate + ``TABLETS_ROUTING_V2``, or when the table/column is unavailable (older + Scylla). The result is cached per parser instance so it is fetched at + most once per schema refresh. + + A transient failure to read the table (timeout, connection or server + error) propagates like any other schema-refresh query. That aborts the + refresh, so the previously known metadata -- including each keyspace's + consistency mode -- is left in place and retried, rather than every + keyspace being silently reset to eventual (which would disable leader + routing and evict the tablet cache). A missing table/column is not an + error: _query_build_rows absorbs it via expected_failures and yields an + empty map. + """ + if self._scylla_consistency_cache is not None: + return self._scylla_consistency_cache + + if self._is_not_scylla(): + self._scylla_consistency_cache = {} + return self._scylla_consistency_cache + + # The consistency map only feeds V2 leader-aware routing. If the control + # connection did not negotiate TABLETS_ROUTING_V2, there are no + # strongly-consistent tables to route for, so skip the query entirely. + features = getattr(self.connection, 'features', None) + if features is None or not getattr(features, 'tablets_routing_v2', False): + self._scylla_consistency_cache = {} + return self._scylla_consistency_cache + + rows = self._query_build_rows(self._SELECT_SCYLLA_KEYSPACES, lambda row: row) + self._scylla_consistency_cache = {row["keyspace_name"]: _ConsistencyMode.from_string(row.get("consistency")) + for row in rows} + return self._scylla_consistency_cache + + def _set_consistency_mode(self, keyspace_meta: KeyspaceMetadata) -> KeyspaceMetadata: + keyspace_meta._consistency_mode = self._get_scylla_keyspaces_consistency().get( + keyspace_meta.name, _ConsistencyMode.EVENTUAL) + return keyspace_meta + + def get_keyspace(self, keyspaces, keyspace): + keyspace_meta = super(SchemaParserV3, self).get_keyspace(keyspaces, keyspace) + if keyspace_meta is not None: + self._set_consistency_mode(keyspace_meta) + return keyspace_meta def get_all_keyspaces(self): for keyspace_meta in super(SchemaParserV3, self).get_all_keyspaces(): for row in self.keyspace_view_rows[keyspace_meta.name]: view_meta = self._build_view_metadata(row) keyspace_meta._add_view_metadata(view_meta) + self._set_consistency_mode(keyspace_meta) yield keyspace_meta def get_table(self, keyspaces, keyspace, table): diff --git a/tests/unit/test_metadata.py b/tests/unit/test_metadata.py index 15cf283777..7ca3f38bfa 100644 --- a/tests/unit/test_metadata.py +++ b/tests/unit/test_metadata.py @@ -32,10 +32,12 @@ _UnknownStrategy, ColumnMetadata, TableMetadata, IndexMetadata, Function, Aggregate, Metadata, TokenMap, ReplicationFactor, - SchemaParserDSE68) + SchemaParserDSE68, SchemaParserV3, + _ConsistencyMode) from cassandra.policies import SimpleConvictionPolicy from cassandra.pool import Host from cassandra.protocol import QueryMessage +from cassandra.tablets import Tablet from tests.util import assertCountEqual import pytest @@ -522,6 +524,47 @@ def test_comparison_unicode(self): class KeyspaceMetadataTest(unittest.TestCase): + @staticmethod + def _keyspace(consistency_mode=None): + keyspace = KeyspaceMetadata('test', True, 'NetworkTopologyStrategy', dict(dc1=3)) + if consistency_mode is not None: + keyspace._consistency_mode = consistency_mode + return keyspace + + def test_as_cql_query_omits_eventual_consistency(self): + # Eventual consistency is the server's default, so it must not be + # spelled out -- including for a keyspace whose mode was never set, + # which is every keyspace on a non-Scylla cluster. + assert 'consistency' not in self._keyspace().as_cql_query() + assert 'consistency' not in self._keyspace(_ConsistencyMode.EVENTUAL).as_cql_query() + + def test_as_cql_query_includes_consistency_mode(self): + # A recreated keyspace has to keep its consistency mode, or the copy + # silently loses strong consistency. + assert self._keyspace(_ConsistencyMode.GLOBAL).as_cql_query() == ( + "CREATE KEYSPACE test WITH replication = " + "{'class': 'NetworkTopologyStrategy', 'dc1': '3'} " + " AND consistency = 'global' AND durable_writes = true") + assert self._keyspace(_ConsistencyMode.LOCAL).as_cql_query() == ( + "CREATE KEYSPACE test WITH replication = " + "{'class': 'NetworkTopologyStrategy', 'dc1': '3'} " + " AND consistency = 'local' AND durable_writes = true") + + def test_as_cql_query_rejects_unknown_consistency_mode(self): + # The branches above must stay exhaustive: a mode added to + # _ConsistencyMode without a branch here would otherwise be dropped + # silently, and the recreated keyspace would lose the option. Not + # reachable through the parser, which only ever assigns a known member. + keyspace = self._keyspace('not-a-mode') + with pytest.raises(ValueError): + keyspace.as_cql_query() + + def test_export_as_string_includes_consistency_mode(self): + # export_as_string() appends the statement terminator and is what a + # schema dump goes through, so the option has to survive that path too. + exported = self._keyspace(_ConsistencyMode.GLOBAL).export_as_string() + assert "AND consistency = 'global' AND durable_writes = true;" in exported + def test_export_as_string_user_types(self): keyspace_name = 'test' keyspace = KeyspaceMetadata(keyspace_name, True, 'NetworkTopologyStrategy', dict(dc1=3)) @@ -552,6 +595,103 @@ def test_export_as_string_user_types(self): );""" == keyspace.export_as_string() +class KeyspaceConsistencyTabletInvalidationTest(unittest.TestCase): + """ + Metadata._update_keyspace must drop cached tablets when a keyspace's + strong-consistency mode changes, not only when its replication strategy + changes. A tablet cached while the keyspace was eventually consistent has no + leader ordering, so it must not survive an eventual->global flip and then be + misread as a leader hint by TokenAwarePolicy.make_query_plan. + """ + + def _ks_meta(self, strongly_consistent): + meta = KeyspaceMetadata('ks', True, 'NetworkTopologyStrategy', {'replication_factor': '1'}) + meta._consistency_mode = _ConsistencyMode.GLOBAL if strongly_consistent else _ConsistencyMode.EVENTUAL + return meta + + def _add_cached_tablet(self, metadata): + tablet = Tablet(first_token=-100, last_token=100, + replicas=[(uuid.uuid4(), 0)], tablet_version=1) + metadata._tablets.add_tablet('ks', 'tbl', tablet) + + def test_consistency_flip_drops_tablets(self): + metadata = Metadata() + metadata._update_keyspace(self._ks_meta(strongly_consistent=False)) + self._add_cached_tablet(metadata) + assert metadata._tablets.table_has_tablets('ks', 'tbl') + + # Same replication strategy, consistency flips False -> True: the stale + # tablet cache must be dropped. + metadata._update_keyspace(self._ks_meta(strongly_consistent=True)) + assert not metadata._tablets.table_has_tablets('ks', 'tbl') + + def test_no_consistency_change_keeps_tablets(self): + metadata = Metadata() + metadata._update_keyspace(self._ks_meta(strongly_consistent=False)) + self._add_cached_tablet(metadata) + assert metadata._tablets.table_has_tablets('ks', 'tbl') + + # No replication change and no consistency change: cache is preserved. + metadata._update_keyspace(self._ks_meta(strongly_consistent=False)) + assert metadata._tablets.table_has_tablets('ks', 'tbl') + + +class ScyllaKeyspaceConsistencyParsingTest(unittest.TestCase): + """ + SchemaParserV3._set_consistency_mode maps the server's per-keyspace + consistency option to KeyspaceMetadata._consistency_mode. A transient failure + reading the consistency table propagates so the schema refresh aborts and + the previously known metadata is retried, rather than being reset to + eventual. + + Which modes actually get leader-aware routing is TokenAwarePolicy's business + and is covered in tests/unit/test_policies.py. + """ + + def _parser_with_consistency(self, consistency_by_ks): + # Build the parser without a connection: _set_consistency_mode only + # needs the cached consistency result. + parser = SchemaParserV3.__new__(SchemaParserV3) + parser._scylla_consistency_cache = consistency_by_ks + return parser + + def _meta_for(self, parser, ks_name): + meta = KeyspaceMetadata(ks_name, True, 'NetworkTopologyStrategy', {'replication_factor': '1'}) + parser._set_consistency_mode(meta) + return meta + + def test_consistency_mode_is_recorded(self): + # The keyspace stores the mode the server reported verbatim, so 'local' + # stays distinguishable from 'eventual' even though ScyllaDB does not + # implement it yet and the driver treats the two alike. + parser = self._parser_with_consistency( + {'g': _ConsistencyMode.GLOBAL, 'l': _ConsistencyMode.LOCAL, 'e': _ConsistencyMode.EVENTUAL}) + assert self._meta_for(parser, 'g')._consistency_mode is _ConsistencyMode.GLOBAL + assert self._meta_for(parser, 'l')._consistency_mode is _ConsistencyMode.LOCAL + assert self._meta_for(parser, 'e')._consistency_mode is _ConsistencyMode.EVENTUAL + # A keyspace absent from the (successful) read defaults to eventual. + assert self._meta_for(parser, 'absent')._consistency_mode is _ConsistencyMode.EVENTUAL + + def test_read_failure_propagates(self): + # A transient failure reading system_schema.scylla_keyspaces must + # propagate (not be swallowed into "eventual"), so the schema refresh + # aborts and the previously known consistency modes are retried. + parser = SchemaParserV3.__new__(SchemaParserV3) + parser._scylla_consistency_cache = None + parser._is_not_scylla = lambda: False + # The control connection must have negotiated V2 to reach the read; + # otherwise the query is skipped and no failure could propagate. + parser.connection = Mock(features=Mock(tablets_routing_v2=True)) + + def _raise_timeout(*args, **kwargs): + raise cassandra.OperationTimedOut("scylla_keyspaces read timed out") + parser._query_build_rows = _raise_timeout + + meta = KeyspaceMetadata('g', True, 'NetworkTopologyStrategy', {'replication_factor': '1'}) + with pytest.raises(cassandra.OperationTimedOut): + parser._set_consistency_mode(meta) + + class UserTypesTest(unittest.TestCase): def test_as_cql_query(self): From 548197d4f297d0652a04db64da3e6c87054fa666 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dawid=20M=C4=99drek?= Date: Wed, 8 Jul 2026 16:27:23 +0200 Subject: [PATCH 7/9] Route requests to the tablet leader for strongly-consistent tables For a strongly-consistent tablet the TABLETS_ROUTING_V2 server orders the replica set with the Raft leader first (replicas[0]) and keeps it fresh via the tablet_version already tracked in the previous commits. TokenAwarePolicy uses this to send reads and writes for such tables straight to the leader, saving the extra coordinator->leader hop. The leader is yielded first only when the keyspace's consistency mode is GLOBAL -- the only mode Scylla implements, and so the only one whose tablets have a leader -- and when the tablet carries a tablet_version: eventually-consistent tablet tables are assigned a tablet_version too, and a versionless (v1-sourced or stale) tablet must not be mistaken for a leader hint. Requests at consistency level ONE or LOCAL_ONE are left alone. Any single replica satisfies them, so preferring the leader would only concentrate load on it without buying any consistency. The level is read from the statement, so a request that inherits it from an execution profile looks unset here and is routed to the leader anyway; that costs a little leader contention and nothing in correctness, and is tracked separately in scylladb/python-driver#953. The hint stays bounded by the wrapped policy -- the leader is front-run only if the child policy would consider it at all (never a host it reports as IGNORED, nor a cross-datacenter leader under a DCAwareRoundRobinPolicy with no remote hosts); otherwise the usual token-aware (optionally shuffled) ordering applies and the server forwards to the leader as it would without v2. Refs: SCYLLADB-288 Fixes: SCYLLADB-291 --- cassandra/policies.py | 60 +++++- cassandra/pool.py | 3 +- tests/unit/test_policies.py | 398 +++++++++++++++++++++++++++++++++++- 3 files changed, 455 insertions(+), 6 deletions(-) diff --git a/cassandra/policies.py b/cassandra/policies.py index 89702e8c89..30026acda9 100644 --- a/cassandra/policies.py +++ b/cassandra/policies.py @@ -502,15 +502,56 @@ def make_query_plan(self, working_keyspace=None, query=None): yield host return + # Deferred: cassandra.metadata reaches this module through + # cassandra.protocol, so importing it at module scope closes an import + # cycle. Same reason cassandra.connection imports from metadata inside + # its functions. + from cassandra.metadata import _ConsistencyMode + 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 = self._cluster_metadata.token_map.token_class.from_key(query.routing_key) + 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, + # which today means exactly the keyspaces whose consistency mode is + # GLOBAL: it is the only mode ScyllaDB implements so far, so LOCAL + # (reserved, unimplemented) and EVENTUAL both have no leader. This + # comparison has to widen once 'local' consistency exists. + # 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._consistency_mode is _ConsistencyMode.GLOBAL + and tablet.tablet_version is not None and tablet.replicas): + # Even for a leader-eligible tablet, a request at consistency + # level ONE or LOCAL_ONE is satisfied by any single replica, so + # preferring the leader would only concentrate load into a + # hotspot without buying any consistency; spread those instead. + # TODO: Figure out how to obtain the actual effective consisteny + # level here. This can be equal to None, in which case we + # should fall back to the execution profile's consistency + # level. Unfortunately, it's non-trivial to get access to + # it here. + effective_cl = query.consistency_level + prefer_leader = effective_cl not in (ConsistencyLevel.ONE, ConsistencyLevel.LOCAL_ONE) + if prefer_leader: + 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,9 +564,20 @@ 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 + # Note: The leader is always a replica, so we don't need to filter it out here. yield from yield_in_order([host for host in child.make_query_plan(keyspace, query) if host not in replicas]) def on_up(self, *args, **kwargs): diff --git a/cassandra/pool.py b/cassandra/pool.py index 39425d8ac7..c75e9314e5 100644 --- a/cassandra/pool.py +++ b/cassandra/pool.py @@ -465,7 +465,8 @@ def _get_connection_for_routing_key(self, routing_key=None, keyspace=None, table tablet = self._session.cluster.metadata._tablets.get_tablet_for_key(keyspace, table, t) # In both V1 and V2 the request is sent to this host, so we pick - # the shard that this host owns for the tablet. + # the shard that this host owns for the tablet. Leader-aware host + # selection (V2) happens earlier, in the load balancing policy. if tablet is not None: for replica in tablet.replicas: if replica[0] == self.host.host_id: diff --git a/tests/unit/test_policies.py b/tests/unit/test_policies.py index 63a3c3d12d..b0f7630efe 100644 --- a/tests/unit/test_policies.py +++ b/tests/unit/test_policies.py @@ -26,7 +26,7 @@ from cassandra import ConsistencyLevel from cassandra.cluster import Cluster, ControlConnection -from cassandra.metadata import Metadata +from cassandra.metadata import Metadata, _ConsistencyMode from cassandra.policies import (RackAwareRoundRobinPolicy, RoundRobinPolicy, WhiteListRoundRobinPolicy, DCAwareRoundRobinPolicy, TokenAwarePolicy, SimpleConvictionPolicy, HostDistance, ExponentialReconnectionPolicy, @@ -943,6 +943,402 @@ def _assert_shuffle(self, patched_shuffle, cluster, keyspace, routing_key): child_policy.make_query_plan.assert_called_once_with(keyspace, query) assert patched_shuffle.call_count == 1 + def test_leader_aware_routing_with_tablet_version(self): + """ + For a strongly-consistent keyspace, the leader (first replica in the + tablet's replica list) must be yielded first in the query plan, even + when it is not the closest replica. + """ + hosts = [Host(DefaultEndPoint(str(i)), SimpleConvictionPolicy, host_id=uuid.uuid4()) for i in range(4)] + for i, host in enumerate(hosts): + host.set_up() + host.set_location_info("dc1", f"rack{i + 1}") + + # The leader is hosts[2] (first in tablet.replicas). + leader = hosts[2] + other_replica = hosts[3] + tablet = Tablet( + first_token=-100, last_token=100, + replicas=[(leader.host_id, 0), (other_replica.host_id, 1)], + tablet_version=0xDEADBEEF12345678 + ) + + cluster = Mock(spec=Cluster) + cluster.metadata = Mock(spec=Metadata) + cluster.metadata._tablets = Mock(spec=Tablets) + cluster.metadata._tablets.get_tablet_for_key.return_value = tablet + cluster.metadata.get_replicas.return_value = [leader, other_replica] + cluster.metadata.keyspaces = {'ks': Mock(_consistency_mode=_ConsistencyMode.GLOBAL)} + + child_policy = Mock() + # Put the leader last in the child plan and make it the farther replica + # by distance (LOCAL vs LOCAL_RACK). Without leader-first routing, + # other_replica would be yielded before the leader. + child_policy.make_query_plan.return_value = [hosts[0], hosts[1], other_replica, leader] + distances = { + leader: HostDistance.LOCAL, + other_replica: HostDistance.LOCAL_RACK, + hosts[0]: HostDistance.LOCAL, + hosts[1]: HostDistance.LOCAL, + } + child_policy.distance.side_effect = lambda host: distances.get(host, HostDistance.LOCAL) + + # shuffle_replicas=False keeps replica ordering deterministic so the only + # thing that can pull the leader to the front is the leader-first logic. + policy = TokenAwarePolicy(child_policy, shuffle_replicas=False) + policy.populate(cluster, hosts) + + query = Statement(routing_key=b'\x00\x00\x00\x01', keyspace='ks', table='tbl') + # A non-weak consistency level requires the leader; ONE/LOCAL_ONE would + # instead spread the request across replicas (covered separately). + query.consistency_level = ConsistencyLevel.LOCAL_QUORUM + qplan = list(policy.make_query_plan(None, query)) + + # Leader must be first, even though other_replica is closer (LOCAL_RACK + # vs LOCAL) and the leader is last in the child plan. + self.assertEqual(qplan[0], leader) + # The closer replica follows, and the leader appears exactly once. + self.assertEqual(qplan[1], other_replica) + self.assertEqual(qplan.count(leader), 1) + + def test_leader_fallback_when_leader_is_down(self): + """ + When the leader host is down, the driver should fall back to other + replicas without crashing. The leader should NOT appear in the plan. + """ + hosts = [Host(DefaultEndPoint(str(i)), SimpleConvictionPolicy, host_id=uuid.uuid4()) for i in range(4)] + for i, host in enumerate(hosts): + host.set_up() + host.set_location_info("dc1", f"rack{i + 1}") + + leader = hosts[2] + other_replica = hosts[3] + leader.set_down() # Simulate leader being unreachable. + + tablet = Tablet( + first_token=-100, last_token=100, + replicas=[(leader.host_id, 0), (other_replica.host_id, 1)], + tablet_version=0xCAFEBABE00000001 + ) + + cluster = Mock(spec=Cluster) + cluster.metadata = Mock(spec=Metadata) + cluster.metadata._tablets = Mock(spec=Tablets) + cluster.metadata._tablets.get_tablet_for_key.return_value = tablet + cluster.metadata.get_replicas.return_value = [leader, other_replica] + cluster.metadata.keyspaces = {'ks': Mock(_consistency_mode=_ConsistencyMode.GLOBAL)} + + child_policy = Mock() + child_policy.make_query_plan.return_value = hosts + child_policy.distance.return_value = HostDistance.LOCAL + + policy = TokenAwarePolicy(child_policy) + policy.populate(cluster, hosts) + + query = Statement(routing_key=b'\x00\x00\x00\x01', keyspace='ks', table='tbl') + # A non-weak consistency level makes the leader the preferred target. + query.consistency_level = ConsistencyLevel.LOCAL_QUORUM + qplan = list(policy.make_query_plan(None, query)) + + # Leader is down, should not appear in the plan. + self.assertNotIn(leader, qplan) + # Other replica should be first. + self.assertEqual(qplan[0], other_replica) + + def test_no_leader_routing_without_tablet_version(self): + """ + replicas[0] is only leader-ordered for a tablet that came from a + TABLETS_ROUTING_V2 payload. A versionless tablet (tablet_version=None, + e.g. cached from V1 or stale across a consistency flip) has arbitrary + replica order, so leader-first routing must NOT fire even for a + strongly-consistent keyspace. The version gate is the only thing that + should suppress it here. + """ + hosts = [Host(DefaultEndPoint(str(i)), SimpleConvictionPolicy, host_id=uuid.uuid4()) for i in range(4)] + for i, host in enumerate(hosts): + host.set_up() + host.set_location_info("dc1", f"rack{i + 1}") + + first_replica = hosts[2] + second_replica = hosts[3] + # Strongly-consistent keyspace but a versionless (V1-style) tablet. + tablet = Tablet( + first_token=-100, last_token=100, + replicas=[(first_replica.host_id, 0), (second_replica.host_id, 1)], + tablet_version=None + ) + + cluster = Mock(spec=Cluster) + cluster.metadata = Mock(spec=Metadata) + cluster.metadata._tablets = Mock(spec=Tablets) + cluster.metadata._tablets.get_tablet_for_key.return_value = tablet + cluster.metadata.get_replicas.return_value = [first_replica, second_replica] + cluster.metadata.keyspaces = {'ks': Mock(_consistency_mode=_ConsistencyMode.GLOBAL)} + + child_policy = Mock() + # Order the child plan so the second replica comes before replicas[0]; if + # leader-first wrongly triggered, first_replica would be forced to front. + child_policy.make_query_plan.return_value = [second_replica, first_replica, hosts[0], hosts[1]] + child_policy.distance.return_value = HostDistance.LOCAL + + # shuffle_replicas=False keeps replica ordering deterministic so we can + # assert that no leader is forced to the front. + policy = TokenAwarePolicy(child_policy, shuffle_replicas=False) + policy.populate(cluster, hosts) + + query = Statement(routing_key=b'\x00\x00\x00\x01', keyspace='ks', table='tbl') + qplan = list(policy.make_query_plan(None, query)) + + # Leader-first must NOT apply without a tablet_version: ordering follows + # the child plan, so the second replica (not replicas[0]) stays first. + self.assertEqual(qplan[0], second_replica) + self.assertEqual(len(qplan), 4) + + def test_no_leader_routing_for_eventually_consistent_keyspace(self): + """ + A tablet_version is assigned to eventually-consistent tablet tables too + (TABLETS_ROUTING_V2), but the leader concept only exists for + strongly-consistent keyspaces. For an eventually-consistent keyspace the + leader-first optimization must NOT apply even when a tablet_version is + present. + """ + hosts = [Host(DefaultEndPoint(str(i)), SimpleConvictionPolicy, host_id=uuid.uuid4()) for i in range(4)] + for i, host in enumerate(hosts): + host.set_up() + host.set_location_info("dc1", f"rack{i + 1}") + + first_replica = hosts[2] + second_replica = hosts[3] + tablet = Tablet( + first_token=-100, last_token=100, + replicas=[(first_replica.host_id, 0), (second_replica.host_id, 1)], + tablet_version=0xDEADBEEF12345678 + ) + + cluster = Mock(spec=Cluster) + cluster.metadata = Mock(spec=Metadata) + cluster.metadata._tablets = Mock(spec=Tablets) + cluster.metadata._tablets.get_tablet_for_key.return_value = tablet + cluster.metadata.get_replicas.return_value = [first_replica, second_replica] + cluster.metadata.keyspaces = {'ks': Mock(_consistency_mode=_ConsistencyMode.EVENTUAL)} + + child_policy = Mock() + # Order the child plan so the second replica comes before the first; if + # leader-first logic wrongly triggered, first_replica would be forced to + # the front instead. + child_policy.make_query_plan.return_value = [second_replica, first_replica, hosts[0], hosts[1]] + child_policy.distance.return_value = HostDistance.LOCAL + + # shuffle_replicas=False keeps replica ordering deterministic so we can + # assert that no leader is forced to the front. + policy = TokenAwarePolicy(child_policy, shuffle_replicas=False) + policy.populate(cluster, hosts) + + query = Statement(routing_key=b'\x00\x00\x00\x01', keyspace='ks', table='tbl') + qplan = list(policy.make_query_plan(None, query)) + + # Leader-first must NOT apply: ordering follows the child plan, so the + # second replica (not replicas[0]) stays first. + self.assertEqual(qplan[0], second_replica) + self.assertEqual(len(qplan), 4) + + def test_no_leader_routing_when_keyspace_metadata_missing(self): + """ + If keyspace metadata is unavailable (e.g. schema refresh disabled), the + policy must safely fall back to no leader-first routing rather than + crashing or guessing. + """ + hosts = [Host(DefaultEndPoint(str(i)), SimpleConvictionPolicy, host_id=uuid.uuid4()) for i in range(4)] + for i, host in enumerate(hosts): + host.set_up() + host.set_location_info("dc1", f"rack{i + 1}") + + first_replica = hosts[2] + second_replica = hosts[3] + tablet = Tablet( + first_token=-100, last_token=100, + replicas=[(first_replica.host_id, 0), (second_replica.host_id, 1)], + tablet_version=0xDEADBEEF12345678 + ) + + cluster = Mock(spec=Cluster) + cluster.metadata = Mock(spec=Metadata) + cluster.metadata._tablets = Mock(spec=Tablets) + cluster.metadata._tablets.get_tablet_for_key.return_value = tablet + cluster.metadata.get_replicas.return_value = [first_replica, second_replica] + cluster.metadata.keyspaces = {} # no metadata for 'ks' + + child_policy = Mock() + child_policy.make_query_plan.return_value = [second_replica, first_replica, hosts[0], hosts[1]] + child_policy.distance.return_value = HostDistance.LOCAL + + # shuffle_replicas=False keeps replica ordering deterministic so we can + # assert that no leader is forced to the front. + policy = TokenAwarePolicy(child_policy, shuffle_replicas=False) + policy.populate(cluster, hosts) + + query = Statement(routing_key=b'\x00\x00\x00\x01', keyspace='ks', table='tbl') + qplan = list(policy.make_query_plan(None, query)) + + self.assertEqual(qplan[0], second_replica) + self.assertEqual(len(qplan), 4) + + def test_leader_skipped_when_child_policy_ignores_it(self): + """ + The leader is yielded first only if the child policy would actually use + it. If a (custom) child policy reports the leader as IGNORED, leader-first + routing must respect that and not front-run an excluded host. The leader + should not appear in the plan at all. + """ + hosts = [Host(DefaultEndPoint(str(i)), SimpleConvictionPolicy, host_id=uuid.uuid4()) for i in range(4)] + for i, host in enumerate(hosts): + host.set_up() + host.set_location_info("dc1", f"rack{i + 1}") + + leader = hosts[2] + other_replica = hosts[3] + tablet = Tablet( + first_token=-100, last_token=100, + replicas=[(leader.host_id, 0), (other_replica.host_id, 1)], + tablet_version=0xDEADBEEF12345678 + ) + + cluster = Mock(spec=Cluster) + cluster.metadata = Mock(spec=Metadata) + cluster.metadata._tablets = Mock(spec=Tablets) + cluster.metadata._tablets.get_tablet_for_key.return_value = tablet + cluster.metadata.get_replicas.return_value = [leader, other_replica] + cluster.metadata.keyspaces = {'ks': Mock(_consistency_mode=_ConsistencyMode.GLOBAL)} + + child_policy = Mock() + # The child policy yields the leader but reports it as IGNORED, i.e. it + # would never actually route to it. + child_policy.make_query_plan.return_value = [leader, other_replica, hosts[0], hosts[1]] + distances = { + leader: HostDistance.IGNORED, + other_replica: HostDistance.LOCAL, + hosts[0]: HostDistance.LOCAL, + hosts[1]: HostDistance.LOCAL, + } + child_policy.distance.side_effect = lambda host: distances.get(host, HostDistance.LOCAL) + + # shuffle_replicas=False keeps replica ordering deterministic. + policy = TokenAwarePolicy(child_policy, shuffle_replicas=False) + policy.populate(cluster, hosts) + + query = Statement(routing_key=b'\x00\x00\x00\x01', keyspace='ks', table='tbl') + # A non-weak consistency level keeps the leader eligible; the child + # policy reporting it as IGNORED is what must exclude it here. + query.consistency_level = ConsistencyLevel.LOCAL_QUORUM + qplan = list(policy.make_query_plan(None, query)) + + # The IGNORED leader must not be front-run, nor appear at all. + self.assertNotIn(leader, qplan) + self.assertEqual(qplan[0], other_replica) + + def _make_leader_routing_setup(self, *, consistency_mode=_ConsistencyMode.GLOBAL, + tablet_version=0xDEADBEEF12345678, + default_consistency_level=None): + """ + Build a TokenAwarePolicy over a strongly-consistent tablet keyspace in + which the leader (hosts[2]) is the *farther* replica (LOCAL) while + other_replica (hosts[3]) is closer (LOCAL_RACK). With shuffling off, the + only thing that can pull the leader to the front of the plan is the + leader-first hint, so a test can infer whether that hint fired purely + from the resulting order. ``default_consistency_level`` seeds the default + execution profile so tests can exercise the effective-consistency + fallback used when a statement leaves consistency_level unset. + Returns (policy, leader, other_replica). + """ + hosts = [Host(DefaultEndPoint(str(i)), SimpleConvictionPolicy, host_id=uuid.uuid4()) for i in range(4)] + for i, host in enumerate(hosts): + host.set_up() + host.set_location_info("dc1", f"rack{i + 1}") + + leader = hosts[2] + other_replica = hosts[3] + tablet = Tablet( + first_token=-100, last_token=100, + replicas=[(leader.host_id, 0), (other_replica.host_id, 1)], + tablet_version=tablet_version, + ) + + cluster = Mock(spec=Cluster) + cluster.metadata = Mock(spec=Metadata) + cluster.metadata._tablets = Mock(spec=Tablets) + cluster.metadata._tablets.get_tablet_for_key.return_value = tablet + cluster.metadata.get_replicas.return_value = [leader, other_replica] + cluster.metadata.keyspaces = {'ks': Mock(_consistency_mode=consistency_mode)} + cluster.profile_manager.default.consistency_level = default_consistency_level + + child_policy = Mock() + # Leader is last in the child plan and the farther replica (LOCAL vs + # LOCAL_RACK); without leader-first, other_replica is yielded first. + child_policy.make_query_plan.return_value = [hosts[0], hosts[1], other_replica, leader] + distances = { + leader: HostDistance.LOCAL, + other_replica: HostDistance.LOCAL_RACK, + hosts[0]: HostDistance.LOCAL, + hosts[1]: HostDistance.LOCAL, + } + child_policy.distance.side_effect = lambda host: distances.get(host, HostDistance.LOCAL) + + policy = TokenAwarePolicy(child_policy, shuffle_replicas=False) + policy.populate(cluster, hosts) + return policy, leader, other_replica + + def test_no_leader_routing_for_read_with_consistency_level_one_or_local_one(self): + """ + A request at consistency level ONE or LOCAL_ONE is satisfied by any + single replica, so even on a strongly-consistent keyspace the leader + must NOT be forced to the front -- doing so would only create a leader + hotspot. Ordering follows the normal token-aware plan, so the closer + replica stays first and the leader is not front-run. + """ + for cl in (ConsistencyLevel.ONE, ConsistencyLevel.LOCAL_ONE): + policy, leader, other_replica = self._make_leader_routing_setup() + query = Statement(routing_key=b'\x00\x00\x00\x01', keyspace='ks', table='tbl') + query.consistency_level = cl + qplan = list(policy.make_query_plan(None, query)) + + self.assertEqual(qplan[0], other_replica, + "leader must not be front-run for consistency level %s" % cl) + # The leader is still a valid replica -- it just isn't preferred. + self.assertEqual(qplan.count(leader), 1) + + def test_leader_routing_only_for_global_consistency_mode(self): + """ + Only a keyspace whose consistency mode is GLOBAL has a tablet leader. + ScyllaDB does not implement 'local' consistency yet, so LOCAL must be + treated exactly like EVENTUAL and get plain token-aware ordering; this + test has to be revisited once 'local' consistency exists. + """ + for mode in (_ConsistencyMode.LOCAL, _ConsistencyMode.EVENTUAL): + policy, leader, other_replica = self._make_leader_routing_setup(consistency_mode=mode) + query = Statement(routing_key=b'\x00\x00\x00\x01', keyspace='ks', table='tbl') + # A leader-requiring consistency level, so the mode is the only + # thing that can suppress the leader-first hint. + query.consistency_level = ConsistencyLevel.LOCAL_QUORUM + qplan = list(policy.make_query_plan(None, query)) + + self.assertEqual(qplan[0], other_replica, + "leader must not be front-run for consistency mode %s" % mode) + self.assertEqual(qplan.count(leader), 1) + + def test_leader_routing_for_global_consistency_mode(self): + """ + The GLOBAL counterpart of the test above: with the same setup and the + same consistency level, the leader is pulled to the front of the plan + even though it is the farther replica. + """ + policy, leader, other_replica = self._make_leader_routing_setup( + consistency_mode=_ConsistencyMode.GLOBAL) + query = Statement(routing_key=b'\x00\x00\x00\x01', keyspace='ks', table='tbl') + query.consistency_level = ConsistencyLevel.LOCAL_QUORUM + qplan = list(policy.make_query_plan(None, query)) + + self.assertEqual(qplan[0], leader) + self.assertEqual(qplan.count(leader), 1) @patch('cassandra.policies.shuffle') def test_no_shuffle_for_serial_consistency(self, patched_shuffle): From ef82e098710f1f1479e46ed5029342ff0a5849db Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dawid=20M=C4=99drek?= Date: Wed, 8 Jul 2026 16:32:10 +0200 Subject: [PATCH 8/9] Add integration tests for leader-aware routing Extend the TABLETS_ROUTING_V2 integration suite with a strongly-consistent (consistency='global', Raft-backed) keyspace and cover, against a live ScyllaDB: that the driver reads each keyspace's _consistency_mode from system_schema.scylla_keyspaces (statically, and as keyspaces are created and dropped), and that TokenAwarePolicy sends a leader-requiring request for such a table to the Raft leader (replicas[0]). --- .../standard/test_tablets_routing_v2.py | 141 ++++++++++++++++++ 1 file changed, 141 insertions(+) diff --git a/tests/integration/standard/test_tablets_routing_v2.py b/tests/integration/standard/test_tablets_routing_v2.py index 22a53ed79d..9edcdcd64e 100644 --- a/tests/integration/standard/test_tablets_routing_v2.py +++ b/tests/integration/standard/test_tablets_routing_v2.py @@ -20,6 +20,7 @@ import cassandra.cqltypes as types from cassandra import ConsistencyLevel from cassandra.cluster import Cluster, ExecutionProfile, EXEC_PROFILE_DEFAULT +from cassandra.metadata import _ConsistencyMode from cassandra.policies import ConstantReconnectionPolicy, RoundRobinPolicy, TokenAwarePolicy from cassandra.protocol import ExecuteMessage from cassandra.protocol_features import ( @@ -120,6 +121,24 @@ def _create_schema(cls, session): for i in range(50): session.execute(prepared.bind((i, i))) + session.execute("DROP KEYSPACE IF EXISTS test_v2_sc") + session.execute( + """ + CREATE KEYSPACE test_v2_sc + WITH replication = {'class': 'NetworkTopologyStrategy', 'replication_factor': 2} + AND tablets = {'initial': 8} + AND consistency = 'global' + """) + session.execute("CREATE TABLE test_v2_sc.t (pk int PRIMARY KEY, v int)") + prepared_sc = session.prepare("INSERT INTO test_v2_sc.t (pk, v) VALUES (?, ?)") + # Writes to a strongly-consistent (Raft) table are rejected unless they + # use QUORUM/LOCAL_QUORUM. The session default is LOCAL_ONE, so request + # LOCAL_QUORUM for these inserts; it propagates to every statement bound + # from this prepared one. + prepared_sc.consistency_level = ConsistencyLevel.LOCAL_QUORUM + for i in range(50): + session.execute(prepared_sc.bind((i, i))) + # -- helpers ---------------------------------------------------------------- @classmethod @@ -405,3 +424,125 @@ def test_v2_takes_precedence_over_v1_no_v1_payload_on_wrong_shard(self): "V2 must take precedence (select_statement.cc)") # The correct V2 block also means no V2 payload. assert 'tablets-routing-v2' not in payload + + # -- strongly-consistent (leader-aware) routing ----------------------------- + + def test_strongly_consistent_keyspace_metadata(self): + """ + The driver must learn from system_schema.scylla_keyspaces which keyspaces + are strongly consistent: test_v2_sc (consistency='global') reports GLOBAL, + test_v2 (no consistency clause) reports EVENTUAL. A GLOBAL mode is the + precondition for leader-aware routing in + TokenAwarePolicy.make_query_plan. + """ + self.session.cluster.refresh_schema_metadata() + keyspaces = self.session.cluster.metadata.keyspaces + assert keyspaces['test_v2_sc']._consistency_mode is _ConsistencyMode.GLOBAL + assert keyspaces['test_v2']._consistency_mode is _ConsistencyMode.EVENTUAL + + def test_consistency_mode_tracks_dynamic_keyspace_changes(self): + """ + The driver reads system_schema.scylla_keyspaces on every schema refresh + and sets KeyspaceMetadata._consistency_mode for each keyspace. + + The mode is not a one-off computed at connect time -- it has to track the + live schema. Because the driver refreshes its metadata synchronously in + response to a DDL it executes (ResponseFuture handles + RESULT_KIND_SCHEMA_CHANGE by refreshing before returning), a keyspace + created or dropped *after* connecting is immediately reflected in + cluster.metadata.keyspaces, with no sleep or manual refresh required. A + keyspace created by some *other* client is picked up through the very same + refresh path, driven by the control connection's schema-change events + (subject to the schema refresh window); this test drives the changes + through the connected session so the assertions stay deterministic. + """ + metadata = self.session.cluster.metadata + ec_ks = "dyn_ec_ks" # eventually consistent (no consistency clause) + sc_ks = "dyn_sc_ks" # strongly consistent (Raft, consistency='global') + + # Start from a known-clean slate so the test is repeatable. + self.session.execute("DROP KEYSPACE IF EXISTS {0}".format(ec_ks)) + self.session.execute("DROP KEYSPACE IF EXISTS {0}".format(sc_ks)) + try: + assert ec_ks not in metadata.keyspaces + assert sc_ks not in metadata.keyspaces + + # Create an eventually-consistent keyspace after connecting: it shows + # up in the map as EVENTUAL. This exercises the + # "absent from scylla_keyspaces -> eventual" path. + self.session.execute( + "CREATE KEYSPACE {0} WITH replication = " + "{{'class': 'NetworkTopologyStrategy', 'replication_factor': 1}} " + "AND tablets = {{'initial': 1}}".format(ec_ks)) + assert ec_ks in metadata.keyspaces + assert metadata.keyspaces[ec_ks]._consistency_mode is _ConsistencyMode.EVENTUAL + + # Create a strongly-consistent keyspace: same map, mode now GLOBAL. + # (RF is irrelevant here -- the mode only tracks the consistency option.) + self.session.execute( + "CREATE KEYSPACE {0} WITH replication = " + "{{'class': 'NetworkTopologyStrategy', 'replication_factor': 1}} " + "AND tablets = {{'initial': 1}} AND consistency = 'global'".format(sc_ks)) + assert sc_ks in metadata.keyspaces + assert metadata.keyspaces[sc_ks]._consistency_mode is _ConsistencyMode.GLOBAL + # The previously-created keyspace keeps its EVENTUAL mode. + assert metadata.keyspaces[ec_ks]._consistency_mode is _ConsistencyMode.EVENTUAL + + # A full refresh (the bulk get_all_keyspaces path, as opposed to the + # single-keyspace get_keyspace path the DDL above exercised) agrees. + self.session.cluster.refresh_schema_metadata() + assert metadata.keyspaces[ec_ks]._consistency_mode is _ConsistencyMode.EVENTUAL + assert metadata.keyspaces[sc_ks]._consistency_mode is _ConsistencyMode.GLOBAL + + # Dropping a keyspace removes it from the map and leaves the other + # keyspace's mode untouched. + self.session.execute("DROP KEYSPACE {0}".format(sc_ks)) + assert sc_ks not in metadata.keyspaces + assert metadata.keyspaces[ec_ks]._consistency_mode is _ConsistencyMode.EVENTUAL + + self.session.execute("DROP KEYSPACE {0}".format(ec_ks)) + assert ec_ks not in metadata.keyspaces + finally: + self.session.execute("DROP KEYSPACE IF EXISTS {0}".format(ec_ks)) + self.session.execute("DROP KEYSPACE IF EXISTS {0}".format(sc_ks)) + + def test_leader_aware_routing_targets_the_raft_leader(self): + """ + For a strongly-consistent table, the server orders the replica list with + the Raft leader first. Once that payload is cached, a TokenAwarePolicy + must route every leader-requiring request (here a LOCAL_QUORUM read) for + the tablet to replicas[0] (the leader), saving the extra + coordinator->leader hop. This is the strongly-consistent counterpart to + the eventually consistent test_v2 tests above, which never assert *which* + replica is hit. + + A read at ONE/LOCAL_ONE is intentionally *not* pinned to the leader (any + single replica satisfies it); that carve-out is covered by the policy + unit tests. + """ + select = self.session.prepare("SELECT v FROM test_v2_sc.t WHERE pk = ?") + bound = select.bind([2]) + # Leader-first routing only applies to requests that actually need the + # leader, so request LOCAL_QUORUM (a strong read). The session default is + # LOCAL_ONE, which the policy would deliberately spread across replicas. + bound.consistency_level = ConsistencyLevel.LOCAL_QUORUM + + tablet = self._ensure_cached(bound) + assert tablet.replicas, "strongly-consistent tablet has no replicas" + leader_host_id = tablet.replicas[0][0] + + # Leader-first routing only triggers when the keyspace is known to be + # strongly consistent, i.e. when its consistency mode is GLOBAL. + ks_meta = self.session.cluster.metadata.keyspaces['test_v2_sc'] + assert ks_meta._consistency_mode is _ConsistencyMode.GLOBAL + + # With an up-to-date cache the block always matches, so the server returns + # no further payload and replicas[0] stays the leader; every request must + # therefore be coordinated by that leader. + for _ in range(10): + result = self.session.execute(bound) + coordinator = result.response_future.coordinator_host + assert coordinator is not None and coordinator.host_id == leader_host_id, ( + "request coordinated by {} but the Raft leader is replicas[0]={}; " + "leader-aware routing did not target the leader".format( + getattr(coordinator, 'host_id', None), leader_host_id)) From a6039ba1ed7b4177af93f70aa50c715a6470a364 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dawid=20M=C4=99drek?= Date: Wed, 8 Jul 2026 16:32:32 +0200 Subject: [PATCH 9/9] Document leader-aware routing Extend the Scylla-specific guide's TABLETS_ROUTING_V2 section, which the previous docs commit introduced for tablet-version tracking, to cover leader-aware routing: strongly-consistent (Raft-backed) tablet tables have a leader that the driver targets directly to save the coordinator->leader hop, the behaviour is best-effort and bounded by the load-balancing policy, and eventually-consistent tables keep their usual token-aware ordering. Include a table of how ScyllaDB serves each operation on such a table, so the routing distinction has a visible reason: ONE and LOCAL_ONE reads are non-linearizable and take no Raft read barrier, so they keep normal token-aware ordering, while QUORUM and LOCAL_QUORUM reads and writes go through the leader and are routed to it. --- docs/scylla-specific.rst | 55 ++++++++++++++++++++++++++++++++++++---- 1 file changed, 50 insertions(+), 5 deletions(-) diff --git a/docs/scylla-specific.rst b/docs/scylla-specific.rst index 80071d5102..d36ebdd03e 100644 --- a/docs/scylla-specific.rst +++ b/docs/scylla-specific.rst @@ -158,17 +158,17 @@ Details on the sending tablet information to the drivers https://github.com/scylladb/scylladb/blob/master/docs/dev/protocol-extensions.md#sending-tablet-info-to-the-drivers -Tablet version tracking ------------------------ +Tablet version tracking and leader-aware routing +------------------------------------------------ When the cluster offers it, the driver negotiates ``TABLETS_ROUTING_V2`` in preference to V1. The negotiation happens per connection, so V2 and V1 connections can coexist in the same cluster; each connection uses whichever -extension its node offers. V2 adds tablet version tracking on top of V1, +extension its node offers. V2 adds two capabilities on top of V1, both invisible to application code. -Every tablet now carries a ``tablet_version`` that -changes whenever its replica set is reconfigured. The driver caches the version +**Tablet version tracking.** Every tablet now carries a ``tablet_version`` that +changes whenever its replica set or leader changes. The driver caches the version it last saw for each tablet and, on every prepared-statement execution over a V2 connection, appends a single ``tablet_version_block`` byte derived from it. The server returns updated routing information in the ``custom_payload`` only when @@ -176,6 +176,51 @@ that byte shows the driver's cached view is stale, instead of attaching it to every response. This keeps the cached routing information fresh while avoiding the per-response overhead that V1 incurs. +**Leader-aware routing for strongly-consistent tables.** Tables in a +strongly-consistent keyspace -- one created with a ``consistency`` option and +backed by Raft -- have a tablet leader that coordinates operations. For those +tables, the driver sends each request directly to the leader, saving the extra +hop the coordinator would otherwise take to forward it. Reads with consistency +level ``ONE`` or ``LOCAL_ONE`` are an exception to this and retain normal +token-aware replica ordering. Eventually-consistent tables are completely +unaffected and keep their usual token-aware (optionally shuffled) replica +ordering. + +The distinction follows from how ScyllaDB serves each operation on a +strongly-consistent table: + +.. list-table:: + :header-rows: 1 + :widths: 15 25 60 + + * - Operation + - Consistency level + - How it is served + * - Read + - ``ONE``, ``LOCAL_ONE`` + - Non-linearizable: served by any replica, without taking a Raft read + barrier. There is nothing to gain from preferring the leader, so the + driver keeps normal token-aware ordering. + * - Read + - ``QUORUM``, ``LOCAL_QUORUM`` + - Linearizable: goes through the Raft leader, so the driver routes it to + the leader directly. + * - Write + - ``QUORUM``, ``LOCAL_QUORUM`` + - Committed through Raft by the leader, so the driver routes it to the + leader directly. + * - Write + - anything else + - Rejected by the server: strongly-consistent tables accept only + ``QUORUM`` and ``LOCAL_QUORUM`` writes. + +Leader-aware routing is best-effort and bounded by the load-balancing policy: +the leader is only targeted directly if the wrapped policy would consider it in +the first place. For example, a ``DCAwareRoundRobinPolicy`` configured with no +remote hosts will not send cross-datacenter traffic to a leader in another +datacenter; the request goes to a local replica and the server forwards it to +the leader, exactly as it would without V2. + No configuration is required: as with V1, a ``TokenAwarePolicy`` is all that is needed.