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
@@ -79,6 +81,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 +941,17 @@ 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_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.
+    """
+
+    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 +1138,8 @@ def __init__(self,
                  schema_metadata_page_size=1000,
                  address_translator=None,
                  status_event_refresh_window=2,
+                 experimental_tablet_refresh_time=60, #  This field is experimental and may be changed / removed in the future
+                 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,
@@ -1364,6 +1380,8 @@ 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_refresh_time = experimental_tablet_refresh_time
+        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
@@ -1416,7 +1434,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.experimental_tablet_refresh_time,
             schema_metadata_enabled, token_metadata_enabled,
             schema_meta_page_size=schema_metadata_page_size)
 
@@ -2381,6 +2399,19 @@ 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 = ControlConnection._SELECT_TABLETS
+            tablets_query = QueryMessage(query=sel_tablets, consistency_level=ConsistencyLevel.ONE)
+            connection.wait_for_response(tablets_query, timeout=timeout)
+        except:
+            return False
+        return True
+
+    # Experimental, this interface and use may change
+    def check_tablets_enabled(self):
+        return self.experimental_tablet_feature_enabled
 
 class Session(object):
     """
@@ -3513,6 +3544,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 +3560,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 +3575,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 +3588,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
@@ -3658,6 +3694,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() 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 +3718,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 self._cluster.check_tablets_enabled() 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 +4008,59 @@ 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
+
+            keyspace_name = row.get("keyspace_name")
+            table_name = row.get("table_name")
+
+            tablet = Tablet()
+            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")
+
+            prev = tablets.get((keyspace_name, table_name), [])
+            prev.append(tablet)
+            tablets[(keyspace_name, table_name)] = prev
+        
+        if self._cluster.metadata._tablets is None:
+            self._cluster.metadata._tablets = Tablets()
+        self._cluster.metadata._tablets.set_tablets(tablets)
+        self._cluster.scheduler.schedule_unique(self._tablet_refresh_time, self._refresh_tablets, connection)
+        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 +4662,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

@@ -23,7 +23,7 @@
 import six
 from six.moves import zip
 import sys
-from threading import RLock
+from threading import RLock, Lock
 import struct
 import random
 import itertools
@@ -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

@@ -12,8 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import bisect
 from itertools import islice, cycle, groupby, repeat
 import logging
+import os
 from random import randint, shuffle
 from threading import Lock
 import socket
@@ -335,6 +337,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 +349,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 +380,22 @@ 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._experimental_tablets_enabled and self._cluster_metadata._tablets is not None:
+                    replicas = self._cluster_metadata._tablets.get_replicas_for_key(keyspace, query.table, self._cluster_metadata.token_map.token_class.from_key(routing_key))
+
+                    if replicas != []:
+                        replicas_mapped = set(map(lambda r: r[0], replicas))
+                        child_plan = child.make_query_plan(keyspace, query)
+
+                        replicas = []
+                        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

@@ -15,9 +15,11 @@
 """
 Connection pooling and host management.
 """
+import bisect
 from concurrent.futures import Future
 from functools import total_ordering
 import logging
+import os
 import socket
 import time
 import random
@@ -439,7 +441,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 +452,19 @@ 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.check_tablets_enabled() and self._session.cluster._load_balancing_policy._cluster_metadata._tablets is not None and table is not None:
+                if keyspace is None:
+                    keyspace = self._keyspace
+
+                replicas = self._session.cluster._load_balancing_policy._cluster_metadata._tablets.get_replicas_for_key(keyspace, table, t)
+
+                if replicas != []:
+                    shard_id = 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 +510,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,