Use tablets in token and shard awareness

Add mechanism to parse system.tablets periodically. In TokenAwarePolicy
check if keyspace uses tablets if so try to use them to find replicas.

Make shard awareness work when using tablets.

Everything is wrapped in experimental setting, because tablets are still
experimental in ScyllaDB and changes in the tablets format are possible.

Author: sylwiaszunejko

cassandra/cluster.py

@@ -41,7 +41,7 @@
 import weakref
 from weakref import WeakValueDictionary
 
-from cassandra import (ConsistencyLevel, AuthenticationFailed,
+from cassandra import (ConsistencyLevel, AuthenticationFailed, InvalidRequest,
                        OperationTimedOut, UnsupportedOperation,
                        SchemaTargetType, DriverException, ProtocolVersion,
                        UnresolvableContactPoints)
@@ -51,6 +51,7 @@
                                   EndPoint, DefaultEndPoint, DefaultEndPointFactory,
                                   ContinuousPagingState, SniEndPointFactory, ConnectionBusy)
 from cassandra.cqltypes import UserType
+import cassandra.cqltypes as types
 from cassandra.encoder import Encoder
 from cassandra.protocol import (QueryMessage, ResultMessage,
                                 ErrorMessage, ReadTimeoutErrorMessage,
@@ -79,6 +80,7 @@
                              named_tuple_factory, dict_factory, tuple_factory, FETCH_SIZE_UNSET,
                              HostTargetingStatement)
 from cassandra.marshal import int64_pack
+from cassandra.tablets import Tablet, Tablets
 from cassandra.timestamps import MonotonicTimestampGenerator
 from cassandra.compat import Mapping
 from cassandra.util import _resolve_contact_points_to_string_map, Version
@@ -1397,6 +1399,8 @@ def __init__(self,
             schema_metadata_enabled, token_metadata_enabled,
             schema_meta_page_size=schema_metadata_page_size)
 
+        self._init_tablets()
+
         if client_id is None:
             self.client_id = uuid.uuid4()
         if application_name is not None:
@@ -2388,7 +2392,10 @@ def _prepare_all_queries(self, host):
     def add_prepared(self, query_id, prepared_statement):
         with self._prepared_statement_lock:
             self._prepared_statements[query_id] = prepared_statement
-
+    
+    # Experimental, this interface and use may change
+    def _init_tablets(self):
+        self.metadata._tablets = Tablets({})
 
 class Session(object):
     """
@@ -3541,6 +3548,7 @@ class PeersQueryType(object):
     _schema_meta_page_size = 1000
 
     _uses_peers_v2 = True
+    _tablets_routing_v1 = False
 
     # for testing purposes
     _time = time
@@ -3674,6 +3682,8 @@ def _try_connect(self, host):
         # If sharding information is available, it's a ScyllaDB cluster, so do not use peers_v2 table.
         if connection.features.sharding_info is not None:
             self._uses_peers_v2 = False
+        
+        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
@@ -4600,7 +4610,10 @@ def _query(self, host, message=None, cb=None):
         connection = None
         try:
             # TODO get connectTimeout from cluster settings
-            connection, request_id = pool.borrow_connection(timeout=2.0, routing_key=self.query.routing_key if self.query else None)
+            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)
+            else:
+                connection, request_id = pool.borrow_connection(timeout=2.0)
             self._connection = connection
             result_meta = self.prepared_statement.result_metadata if self.prepared_statement else []
 
@@ -4719,6 +4732,19 @@ 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 = types.lookup_casstype('TupleType(LongType, LongType, ListType(TupleType(UUIDType, Int32Type)))')
+                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 isinstance(response, ResultMessage):
                 if response.kind == RESULT_KIND_SET_KEYSPACE:
                     session = getattr(self, 'session', None)

cassandra/metadata.py

@@ -126,6 +126,7 @@ def __init__(self):
         self._hosts = {}
         self._host_id_by_endpoint = {}
         self._hosts_lock = RLock()
+        self._tablets = None
 
     def export_schema_as_string(self):
         """

cassandra/policies.py

@@ -335,6 +335,7 @@ class TokenAwarePolicy(LoadBalancingPolicy):
 
     _child_policy = None
     _cluster_metadata = None
+    _tablets_routing_v1 = False
     shuffle_replicas = False
     """
     Yield local replicas in a random order.
@@ -346,6 +347,7 @@ def __init__(self, child_policy, shuffle_replicas=False):
 
     def populate(self, cluster, hosts):
         self._cluster_metadata = cluster.metadata
+        self._tablets_routing_v1 = cluster.control_connection._tablets_routing_v1
         self._child_policy.populate(cluster, hosts)
 
     def check_supported(self):
@@ -376,7 +378,19 @@ def make_query_plan(self, working_keyspace=None, query=None):
                 for host in child.make_query_plan(keyspace, query):
                     yield host
             else:
-                replicas = self._cluster_metadata.get_replicas(keyspace, routing_key)
+                replicas = []
+                if self._tablets_routing_v1:
+                    tablet = self._cluster_metadata._tablets.get_tablet_for_key(keyspace, query.table, self._cluster_metadata.token_map.token_class.from_key(routing_key))
+
+                    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]
+
+                if replicas == []:
+                    replicas = self._cluster_metadata.get_replicas(keyspace, routing_key)
+
                 if self.shuffle_replicas:
                     shuffle(replicas)
                 for replica in replicas:

cassandra/pool.py

@@ -392,6 +392,8 @@ class HostConnection(object):
     # the number below, all excess connections will be closed.
     max_excess_connections_per_shard_multiplier = 3
 
+    tablets_routing_v1 = False
+
     def __init__(self, host, host_distance, session):
         self.host = host
         self.host_distance = host_distance
@@ -436,10 +438,11 @@ 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
 
         log.debug("Finished initializing connection for host %s", self.host)
 
-    def _get_connection_for_routing_key(self, routing_key=None):
+    def _get_connection_for_routing_key(self, routing_key=None, keyspace=None, table=None):
         if self.is_shutdown:
             raise ConnectionException(
                 "Pool for %s is shutdown" % (self.host,), self.host)
@@ -450,7 +453,22 @@ def _get_connection_for_routing_key(self, routing_key=None):
         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 = self.host.sharding_info.shard_id_from_token(t.value)
+            
+            shard_id = None
+            if self.tablets_routing_v1 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)
+
+                if tablet is not None:
+                    for replica in tablet.replicas:
+                        if replica[0] == self.host.host_id:
+                            shard_id = replica[1]
+                            break
+
+            if shard_id is None:
+                shard_id = self.host.sharding_info.shard_id_from_token(t.value)
 
         conn = self._connections.get(shard_id)
 
@@ -496,15 +514,15 @@ def _get_connection_for_routing_key(self, routing_key=None):
             return random.choice(active_connections)
         return random.choice(list(self._connections.values()))
 
-    def borrow_connection(self, timeout, routing_key=None):
-        conn = self._get_connection_for_routing_key(routing_key)
+    def borrow_connection(self, timeout, routing_key=None, keyspace=None, table=None):
+        conn = self._get_connection_for_routing_key(routing_key, keyspace, table)
         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)
+                conn = self._get_connection_for_routing_key(routing_key, keyspace, table)
             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

cassandra/query.py

@@ -253,6 +253,13 @@ class Statement(object):
     .. versionadded:: 2.1.3
     """
 
+    table = None
+    """
+    The string name of the table this query acts on. This is used when the tablet
+    experimental feature is enabled and in the same time :class`~.TokenAwarePolicy`
+    is configured in the profile load balancing policy.
+    """
+
     custom_payload = None
     """
     :ref:`custom_payload` to be passed to the server.
@@ -272,7 +279,7 @@ class Statement(object):
 
     def __init__(self, retry_policy=None, consistency_level=None, routing_key=None,
                  serial_consistency_level=None, fetch_size=FETCH_SIZE_UNSET, keyspace=None, custom_payload=None,
-                 is_idempotent=False):
+                 is_idempotent=False, table=None):
         if retry_policy and not hasattr(retry_policy, 'on_read_timeout'):  # just checking one method to detect positional parameter errors
             raise ValueError('retry_policy should implement cassandra.policies.RetryPolicy')
         if retry_policy is not None:
@@ -286,6 +293,8 @@ def __init__(self, retry_policy=None, consistency_level=None, routing_key=None,
             self.fetch_size = fetch_size
         if keyspace is not None:
             self.keyspace = keyspace
+        if table is not None:
+            self.table = table
         if custom_payload is not None:
             self.custom_payload = custom_payload
         self.is_idempotent = is_idempotent
@@ -548,6 +557,7 @@ def __init__(self, prepared_statement, retry_policy=None, consistency_level=None
         meta = prepared_statement.column_metadata
         if meta:
             self.keyspace = meta[0].keyspace_name
+            self.table = meta[0].table_name
 
         Statement.__init__(self, retry_policy, consistency_level, routing_key,
                            serial_consistency_level, fetch_size, keyspace, custom_payload,

cassandra/tablets.py

@@ -0,0 +1,107 @@
+# Experimental, this interface and use may change
+from threading import Lock
+
+class Tablet(object):
+    """
+    Represents a single ScyllaDB tablet.
+    It stores information about each replica, its host and shard, 
+    and the token interval in the format (first_token, last_token].
+    """
+    first_token = 0
+    last_token = 0
+    replicas = None
+
+    def __init__(self, first_token = 0, last_token = 0, replicas = None):
+        self.first_token = first_token
+        self.last_token = last_token
+        self.replicas = replicas
+
+    def __str__(self):
+        return "<Tablet: first_token=%s last_token=%s replicas=%s>" \
+               % (self.first_token, self.last_token, self.replicas)
+    __repr__ = __str__
+
+    @staticmethod
+    def _is_valid_tablet(replicas):
+        return replicas is not None and len(replicas) != 0
+
+    @staticmethod
+    def from_row(first_token, last_token, replicas):
+        if Tablet._is_valid_tablet(replicas):
+            tablet = Tablet(first_token, last_token,replicas)
+            return tablet
+        return None
+
+# Experimental, this interface and use may change
+class Tablets(object):
+    _lock = None
+    _tablets = {}
+
+    def __init__(self, tablets):
+        self._tablets = tablets
+        self._lock = Lock()
+    
+    def get_tablet_for_key(self, keyspace, table, t):
+        tablet = self._tablets.get((keyspace, table), [])
+        if tablet == []:
+            return None
+        
+        id = bisect_left(tablet, t.value, key = lambda tablet: tablet.last_token)
+        if id < len(tablet) and t.value > tablet[id].first_token:
+            return tablet[id]
+        return None
+
+    def add_tablet(self, keyspace, table, tablet):
+        with self._lock:
+            tablets_for_table = self._tablets.setdefault((keyspace, table), [])
+
+            # find first overlaping range 
+            start = bisect_left(tablets_for_table, tablet.first_token, key = lambda t: t.first_token)
+            if start > 0 and tablets_for_table[start - 1].last_token > tablet.first_token:
+                start = start - 1
+
+            # find last overlaping range 
+            end = bisect_left(tablets_for_table, tablet.last_token, key = lambda t: t.last_token)
+            if end < len(tablets_for_table) and tablets_for_table[end].first_token >= tablet.last_token:
+                end = end - 1
+
+            if start <= end:
+                del tablets_for_table[start:end + 1]
+
+            tablets_for_table.insert(start, tablet)
+
+# bisect.bisect_left implementation from Python 3.11, needed untill support for
+# Python < 3.10 is dropped, it is needed to use `key` to extract last_token from
+# Tablet list - better solution performance-wise than materialize list of last_tokens
+def bisect_left(a, x, lo=0, hi=None, *, key=None):
+    """Return the index where to insert item x in list a, assuming a is sorted.
+
+    The return value i is such that all e in a[:i] have e < x, and all e in
+    a[i:] have e >= x.  So if x already appears in the list, a.insert(i, x) will
+    insert just before the leftmost x already there.
+
+    Optional args lo (default 0) and hi (default len(a)) bound the
+    slice of a to be searched.
+    """
+
+    if lo < 0:
+        raise ValueError('lo must be non-negative')
+    if hi is None:
+        hi = len(a)
+    # Note, the comparison uses "<" to match the
+    # __lt__() logic in list.sort() and in heapq.
+    if key is None:
+        while lo < hi:
+            mid = (lo + hi) // 2
+            if a[mid] < x:
+                lo = mid + 1
+            else:
+                hi = mid
+    else:
+        while lo < hi:
+            mid = (lo + hi) // 2
+            if key(a[mid]) < x:
+                lo = mid + 1
+            else:
+                hi = mid
+    return lo

tests/unit/test_policies.py

@@ -526,6 +526,7 @@ class TokenAwarePolicyTest(unittest.TestCase):
     def test_wrap_round_robin(self):
         cluster = Mock(spec=Cluster)
         cluster.metadata = Mock(spec=Metadata)
+        cluster.control_connection._tablets_routing_v1 = False
         hosts = [Host(DefaultEndPoint(str(i)), SimpleConvictionPolicy) for i in range(4)]
         for host in hosts:
             host.set_up()
@@ -557,6 +558,7 @@ def get_replicas(keyspace, packed_key):
     def test_wrap_dc_aware(self):
         cluster = Mock(spec=Cluster)
         cluster.metadata = Mock(spec=Metadata)
+        cluster.control_connection._tablets_routing_v1 = False
         hosts = [Host(DefaultEndPoint(str(i)), SimpleConvictionPolicy) for i in range(4)]
         for host in hosts:
             host.set_up()
@@ -685,6 +687,7 @@ def test_statement_keyspace(self):
 
         cluster = Mock(spec=Cluster)
         cluster.metadata = Mock(spec=Metadata)
+        cluster.control_connection._tablets_routing_v1 = False
         replicas = hosts[2:]
         cluster.metadata.get_replicas.return_value = replicas
 
@@ -775,6 +778,7 @@ def _assert_shuffle(self, patched_shuffle, keyspace, routing_key):
 
         cluster = Mock(spec=Cluster)
         cluster.metadata = Mock(spec=Metadata)
+        cluster.control_connection._tablets_routing_v1 = False
         replicas = hosts[2:]
         cluster.metadata.get_replicas.return_value = replicas
 
@@ -1448,6 +1452,7 @@ def test_query_plan_deferred_to_child(self):
 
     def test_wrap_token_aware(self):
         cluster = Mock(spec=Cluster)
+        cluster.control_connection._tablets_routing_v1 = False
         hosts = [Host(DefaultEndPoint("127.0.0.{}".format(i)), SimpleConvictionPolicy) for i in range(1, 6)]
         for host in hosts:
             host.set_up()