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

@@ -27,6 +27,8 @@
 from itertools import groupby, count, chain
 import json
 import logging
+import os
+import threading
 from warnings import warn
 from random import random
 import re
@@ -73,7 +75,7 @@
                                 NeverRetryPolicy)
 from cassandra.pool import (Host, _ReconnectionHandler, _HostReconnectionHandler,
                             HostConnectionPool, HostConnection,
-                            NoConnectionsAvailable)
+                            NoConnectionsAvailable, Tablet)
 from cassandra.query import (SimpleStatement, PreparedStatement, BoundStatement,
                              BatchStatement, bind_params, QueryTrace, TraceUnavailable,
                              named_tuple_factory, dict_factory, tuple_factory, FETCH_SIZE_UNSET,
@@ -139,6 +141,9 @@ def _is_gevent_monkey_patched():
     except ImportError:
         from cassandra.io.asyncorereactor import AsyncoreConnection as DefaultConnection  # NOQA
 
+# If true tablet experimental feature is enabled, its interface and use may change
+EXPERIMENTAL_TABLETS = os.environ.get('EXPERIMENTAL_TABLETS', '').lower() == 'true'
+
 # Forces load of utf8 encoding module to avoid deadlock that occurs
 # if code that is being imported tries to import the module in a seperate
 # thread.
@@ -938,6 +943,12 @@ def default_retry_policy(self, policy):
     establish connection pools. This can cause a rush of connections and queries if not mitigated with this factor.
     """
 
+    tablet_refresh_time = 60
+    """
+    The time at which the system.tablets table is polled periodically.
+    This field is experimental and may be changed / removed in the future.
+    """
+
     prepare_on_all_hosts = True
     """
     Specifies whether statements should be prepared on all hosts, or just one.
@@ -1124,6 +1135,7 @@ def __init__(self,
                  schema_metadata_page_size=1000,
                  address_translator=None,
                  status_event_refresh_window=2,
+                 tablet_refresh_time=60, #  This field is experimental and may be changed / removed in the future
                  prepare_on_all_hosts=True,
                  reprepare_on_up=True,
                  execution_profiles=None,
@@ -1364,6 +1376,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.tablet_refresh_time = tablet_refresh_time
         self.connect_timeout = connect_timeout
         self.prepare_on_all_hosts = prepare_on_all_hosts
         self.reprepare_on_up = reprepare_on_up
@@ -1416,7 +1429,7 @@ 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,
+            self.status_event_refresh_window, self.tablet_refresh_time,
             schema_metadata_enabled, token_metadata_enabled,
             schema_meta_page_size=schema_metadata_page_size)
 
@@ -2381,6 +2394,15 @@ 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_support(self, connection, timeout):
+        try:
+            sel_tablets = "SELECT * FROM system.tablets"
+            tablets_query = QueryMessage(query=sel_tablets, consistency_level=ConsistencyLevel.ONE)
+            connection.wait_for_response(tablets_query, timeout=timeout)
+        except:
+            return False
+        return True
 
 class Session(object):
     """
@@ -3513,6 +3535,8 @@ class ControlConnection(object):
     _SELECT_PEERS_NO_TOKENS_V2 = "SELECT host_id, peer, peer_port, data_center, rack, native_address, native_port, release_version, schema_version FROM system.peers_v2"
     _SELECT_SCHEMA_PEERS_V2 = "SELECT host_id, peer, peer_port, native_address, native_port, schema_version FROM system.peers_v2"
 
+    _SELECT_TABLETS = "SELECT * FROM system.tablets"
+
     _MINIMUM_NATIVE_ADDRESS_DSE_VERSION = Version("6.0.0")
 
     class PeersQueryType(object):
@@ -3527,6 +3551,7 @@ class PeersQueryType(object):
     _schema_event_refresh_window = None
     _topology_event_refresh_window = None
     _status_event_refresh_window = None
+    _tablet_refresh_time = 60
 
     _schema_meta_enabled = True
     _token_meta_enabled = True
@@ -3541,6 +3566,7 @@ def __init__(self, cluster, timeout,
                  schema_event_refresh_window,
                  topology_event_refresh_window,
                  status_event_refresh_window,
+                 tablet_refresh_time=60,
                  schema_meta_enabled=True,
                  token_meta_enabled=True,
                  schema_meta_page_size=1000):
@@ -3553,6 +3579,7 @@ def __init__(self, cluster, timeout,
         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._tablet_refresh_time = tablet_refresh_time
         self._schema_meta_enabled = schema_meta_enabled
         self._token_meta_enabled = token_meta_enabled
         self._schema_meta_page_size = schema_meta_page_size
@@ -3616,7 +3643,7 @@ def _reconnect_internal(self):
                 raise DriverException("[control connection] Reconnection in progress during shutdown")
 
         raise NoHostAvailable("Unable to connect to any servers", errors)
-
+    
     def _try_connect(self, host):
         """
         Creates a new Connection, registers for pushed events, and refreshes
@@ -3656,8 +3683,12 @@ def _try_connect(self, host):
                 "TOPOLOGY_CHANGE": partial(_watch_callback, self_weakref, '_handle_topology_change'),
                 "STATUS_CHANGE": partial(_watch_callback, self_weakref, '_handle_status_change'),
                 "SCHEMA_CHANGE": partial(_watch_callback, self_weakref, '_handle_schema_change')
+                # "TABLET_CHANGE": partial(_watch_callback, self_weakref, '_handle_tablet_change')
             }, register_timeout=self._timeout)
 
+            if EXPERIMENTAL_TABLETS and self._cluster.check_tablets_support(connection, self._timeout):
+                self._cluster.scheduler.schedule_unique(self._tablet_refresh_time, self._refresh_tablets, connection)
+
             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)
@@ -3679,6 +3710,11 @@ def _try_connect(self, host):
             shared_results = (peers_result, local_result)
             self._refresh_node_list_and_token_map(connection, preloaded_results=shared_results)
             self._refresh_schema(connection, preloaded_results=shared_results, schema_agreement_wait=-1)
+
+            if EXPERIMENTAL_TABLETS and self._cluster.check_tablets_support(connection, self._timeout):
+                tablets_query = QueryMessage(query=self._SELECT_TABLETS, consistency_level=ConsistencyLevel.ONE)
+                tablets_result = connection.wait_for_response(tablets_query, timeout=self._timeout)
+                self._refresh_tablets(connection, preloaded_result=tablets_result)
         except Exception:
             connection.close()
             raise
@@ -3964,12 +4000,53 @@ 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 _refresh_tablets(self, connection, preloaded_result=None,
+                                         force_token_rebuild=False):
+        if preloaded_result:
+            log.debug("[control connection] Refreshing tablet list using preloaded result")
+            tablets_result = preloaded_result
+        else:
+            sel_tablets = self._SELECT_TABLETS
+            tablets_query = QueryMessage(query=sel_tablets, consistency_level=ConsistencyLevel.ONE)
+            tablets_result = connection.wait_for_response(
+                tablets_query, timeout=self._timeout)
+        tablets_result = dict_factory(tablets_result.column_names, tablets_result.parsed_rows)
+
+        tablets = []
+        for row in tablets_result:
+            if not self._is_valid_tablet(row):
+                log.warning(
+                    "Found an invalid row for tablet (%s). Ignoring tablet." %
+                    _NodeInfo.get_broadcast_rpc_address(row))
+                continue
+
+            tablet = Tablet()
+            tablet.keyspace_name = row.get("keyspace_name")
+            tablet.table_name = row.get("table_name")
+            tablet.table_id = row.get("table_id")
+            tablet.tablet_count = row.get("tablet_count")
+            tablet.last_token = row.get("last_token")
+            tablet.new_replicas = row.get("new_replicas")
+            tablet.replicas = row.get("replicas")
+            tablet.stage = row.get("stage")
+
+            tablets.append(tablet)
+        
+        self._cluster.metadata.set_tablets(tablets)
+        return ""
+    
     @staticmethod
     def _is_valid_peer(row):
         return bool(_NodeInfo.get_broadcast_rpc_address(row) and row.get("host_id") and
                     row.get("data_center") and row.get("rack") and
                     ('tokens' not in row or row.get('tokens')))
 
+    # Experimental, this interface and use may change
+    @staticmethod
+    def _is_valid_tablet(row):
+        return bool(row.get("replicas") is not None and len(row.get("replicas")) != 0 and row.get("table_name") is not None)
+
     def _update_location_info(self, host, datacenter, rack):
         if host.datacenter == datacenter and host.rack == rack:
             return False
@@ -4571,7 +4648,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 []
 

cassandra/metadata.py

@@ -126,6 +126,8 @@ def __init__(self):
         self._hosts = {}
         self._host_id_by_endpoint = {}
         self._hosts_lock = RLock()
+        self._tablets = []
+        self._tablets_lock = RLock()
 
     def export_schema_as_string(self):
         """
@@ -391,6 +393,22 @@ def all_hosts_items(self):
         with self._hosts_lock:
             return list(self._hosts.items())
 
+    # Experimental, this interface and use may change
+    def all_tablets(self):
+        """
+        Returns a list of all known :class`.Tablet` instances in the cluster.
+        """
+        with self._tablets_lock:
+            return self._tablets
+        
+    # Experimental, this interface and use may change
+    def set_tablets(self, tablets):
+        with self._tablets_lock:
+            self._tablets = tablets
+    
+    # Experimental, this interface and use may change
+    def use_tablets(self):
+        return self.all_tablets() != []
 
 REPLICATION_STRATEGY_CLASS_PREFIX = "org.apache.cassandra.locator."
 

cassandra/policies.py

@@ -14,6 +14,7 @@
 
 from itertools import islice, cycle, groupby, repeat
 import logging
+import os
 from random import randint, shuffle
 from threading import Lock
 import socket
@@ -376,7 +377,36 @@ 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)
+                if self._cluster_metadata.use_tablets():
+                    tablets = self._cluster_metadata.all_tablets()
+
+                    l = -1
+                    for i, tablet in enumerate(tablets): 
+                        if tablet.keyspace_name == keyspace and tablet.table_name == query.table:
+                            l = i
+                            break
+                    
+                    if l != -1:
+                        r = l + tablets[l].tablet_count
+                        t = self._cluster_metadata.token_map.token_class.from_key(routing_key)
+                        while l < r:
+                            m = (l + r) // 2
+                            if tablets[m].last_token < t.value:
+                                l = m + 1
+                            else:
+                                r = m
+
+                        replicas = []
+                        replicas_mapped = set(map(lambda r: r[0], tablets[l].replicas))
+                        child_plan = child.make_query_plan(keyspace, query)
+                        for host in child_plan:
+                            if host.host_id in replicas_mapped:
+                                replicas.append(host)
+                    else:
+                        replicas = self._cluster_metadata.get_replicas(keyspace, routing_key)
+                else:
+                    replicas = self._cluster_metadata.get_replicas(keyspace, routing_key)
+
                 if self.shuffle_replicas:
                     shuffle(replicas)
                 for replica in replicas:

cassandra/pool.py

@@ -18,6 +18,7 @@
 from concurrent.futures import Future
 from functools import total_ordering
 import logging
+import os
 import socket
 import time
 import random
@@ -252,6 +253,32 @@ def __repr__(self):
         dc = (" %s" % (self._datacenter,)) if self._datacenter else ""
         return "<%s: %s%s>" % (self.__class__.__name__, self.endpoint, dc)
 
+# Experimental, this interface and use may change
+class Tablet(object):
+    """
+    Represents a single ScyllaDB tablet.
+    """
+    keyspace_name = ""
+    table_id = None
+    last_token = 0
+    table_name = ""
+    tablet_count = 0
+    new_replicas = None
+    replicas = None
+    stage = ""
+
+    lock = None
+
+    def __init__(self, keyspace_name = "", table_id = None, last_token = 0, table_name = "", tablet_count = 0, new_replicas = None, replicas = None, stage = ""):
+        self.keyspace_name = keyspace_name
+        self.table_id = table_id
+        self.table_name = table_name
+        self.tablet_count = tablet_count
+        self.last_token = last_token
+        self.new_replicas = new_replicas
+        self.replicas = replicas
+        self.stage = stage
+        self.lock = RLock()
 
 class _ReconnectionHandler(object):
     """
@@ -439,7 +466,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 +477,32 @@ 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)
+            
+            if self._session.cluster._load_balancing_policy._cluster_metadata.use_tablets() and table is not None:
+                tablets = self._session.cluster._load_balancing_policy._cluster_metadata.all_tablets()
+
+                if keyspace is None:
+                    keyspace = self._keyspace
+
+                l = -1
+                for i, tablet in enumerate(tablets): 
+                    if tablet.keyspace_name == keyspace and tablet.table_name == table: 
+                        l = i
+                        break
+                
+                if l != -1:
+                    r = l + tablets[l].tablet_count
+                    while l < r:
+                        m = (l + r) // 2
+                        if tablets[m].last_token < t.value:
+                            l = m + 1
+                        else:
+                            r = m
+                    shard_id = tablets[l].replicas[0][1]
+                else:
+                    shard_id = self.host.sharding_info.shard_id_from_token(t.value)
+            else:
+                shard_id = self.host.sharding_info.shard_id_from_token(t.value)
 
         conn = self._connections.get(shard_id)
 
@@ -496,15 +548,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