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
@@ -938,6 +940,11 @@ def default_retry_policy(self, policy):
     establish connection pools. This can cause a rush of connections and queries if not mitigated with this factor.
     """
 
+    experimental_tablet_feature_enabled = False
+    """
+    If true tablet experimental feature is enabled, its interface and use may change.
+    """
+
     prepare_on_all_hosts = True
     """
     Specifies whether statements should be prepared on all hosts, or just one.
@@ -1124,6 +1131,7 @@ def __init__(self,
                  schema_metadata_page_size=1000,
                  address_translator=None,
                  status_event_refresh_window=2,
+                 experimental_tablet_feature_enabled=False,  #  This field is experimental and may be changed / removed in the future
                  prepare_on_all_hosts=True,
                  reprepare_on_up=True,
                  execution_profiles=None,
@@ -1341,6 +1349,7 @@ def __init__(self,
         self.schema_event_refresh_window = schema_event_refresh_window
         self.topology_event_refresh_window = topology_event_refresh_window
         self.status_event_refresh_window = status_event_refresh_window
+        self.experimental_tablet_feature_enabled = experimental_tablet_feature_enabled
         self.connect_timeout = connect_timeout
         self.prepare_on_all_hosts = prepare_on_all_hosts
         self.reprepare_on_up = reprepare_on_up
@@ -1393,9 +1402,8 @@ def __init__(self,
         self.control_connection = ControlConnection(
             self, self.control_connection_timeout,
             self.schema_event_refresh_window, self.topology_event_refresh_window,
-            self.status_event_refresh_window,
-            schema_metadata_enabled, token_metadata_enabled,
-            schema_meta_page_size=schema_metadata_page_size)
+            self.status_event_refresh_window, schema_metadata_enabled,
+            token_metadata_enabled, schema_meta_page_size=schema_metadata_page_size)
 
         if client_id is None:
             self.client_id = uuid.uuid4()
@@ -2389,6 +2397,9 @@ 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 check_tablets_enabled(self):
+        return self.experimental_tablet_feature_enabled
 
 class Session(object):
     """
@@ -3687,6 +3698,9 @@ def _try_connect(self, host):
                 "SCHEMA_CHANGE": partial(_watch_callback, self_weakref, '_handle_schema_change')
             }, register_timeout=self._timeout)
 
+            if self._cluster.check_tablets_enabled():
+                self._init_tablets()
+
             sel_peers = self._get_peers_query(self.PeersQueryType.PEERS, connection)
             sel_local = self._SELECT_LOCAL if self._token_meta_enabled else self._SELECT_LOCAL_NO_TOKENS
             peers_query = QueryMessage(query=sel_peers, consistency_level=ConsistencyLevel.ONE)
@@ -3993,6 +4007,34 @@ def _refresh_node_list_and_token_map(self, connection, preloaded_results=None,
             log.debug("[control connection] Rebuilding token map due to topology changes")
             self._cluster.metadata.rebuild_token_map(partitioner, token_map)
 
+    # Experimental, this interface and use may change
+    def _init_tablets(self):
+        self._cluster.metadata._tablets = Tablets({})
+
+    # Experimental, this interface and use may change
+    def add_tablet(self, keyspace, table, tablet):
+        tablets = self._cluster.metadata._tablets._tablets
+        tablets_for_table = tablets.setdefault((keyspace, table), [])
+        start = None
+        end = len(tablets_for_table) - 1
+        for id, t in enumerate(tablets_for_table):
+            if t.last_token > tablet.first_token and t.first_token < tablet.last_token:
+                if start == None:
+                    start = id
+                end = id
+        if start != None:
+            del tablets_for_table[start:end + 1]
+
+        insert_idx = len(tablets_for_table)
+        for id, t in enumerate(tablets_for_table):
+            if t.first_token > tablet.first_token:
+                insert_idx = id
+                break
+        tablets_for_table.insert(insert_idx, tablet)
+
+        tablets[(keyspace, table)] = tablets_for_table
+        self._cluster.metadata._tablets = Tablets(tablets)
+
     @staticmethod
     def _is_valid_peer(row):
         return bool(_NodeInfo.get_broadcast_rpc_address(row) and row.get("host_id") and
@@ -4600,7 +4642,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 +4764,17 @@ 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.check_tablets_enabled():
+                ctype = types.lookup_casstype('ListType(TupleType(UUIDType, Int32Type))')
+                tablet_replicas = ctype.from_binary(self._custom_payload.get('tablet_replicas'), 3)
+                ctype_token = types.lookup_casstype('LongType')
+                first_token = ctype_token.from_binary(self._custom_payload.get('token_range')[:8], 3)
+                last_token = ctype_token.from_binary(self._custom_payload.get('token_range')[8:], 3)
+                tablet = Tablet.from_row(first_token, last_token, tablet_replicas)
+                keyspace = self.query.keyspace
+                table = self.query.table
+                self.session.cluster.control_connection.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
+    _experimental_tablets_enabled = 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._experimental_tablets_enabled = cluster.check_tablets_enabled()
         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._experimental_tablets_enabled:
+                    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

@@ -439,7 +439,7 @@ def __init__(self, host, host_distance, session):
 
         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 +450,23 @@ 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._session.cluster.check_tablets_enabled() and table is not None:
+                if keyspace is None:
+                    keyspace = self._keyspace
+
+                tablet = self._session.cluster._load_balancing_policy._cluster_metadata._tablets.get_tablet_for_key(keyspace, table, t)
+
+                if tablet is not None:
+                    for replica in tablet.replicas:
+                        log.info(replica)
+                        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 +512,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,77 @@
+# Experimental, this interface and use may change
+class Tablet(object):
+    """
+    Represents a single ScyllaDB tablet.
+    """
+    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
+
+    @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):
+    _tablets = {}
+
+    def __init__(self, tablets):
+        self._tablets = tablets
+    
+    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
+
+# 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