shard aware: shard aware unique port (advenced shard aware)

shard aware port in now advertised OPTIONS messge,
and we need to replace the connection with the new
host/port

* fixing tests to match the advenced shard awareness

now that we could have two host listed (one with 9042 port, and one
with 19042), we need to make the test a bit less prune to failure cause
of that change

Author: fruch

cassandra/c_shard_info.pyx

@@ -22,15 +22,19 @@ cdef class ShardingInfo():
     cdef readonly str partitioner
     cdef readonly str sharding_algorithm
     cdef readonly int sharding_ignore_msb
+    cdef readonly int shard_aware_port
+    cdef readonly int shard_aware_port_ssl
 
     cdef object __weakref__
 
-    def __init__(self, shard_id, shards_count, partitioner, sharding_algorithm, sharding_ignore_msb):
+    def __init__(self, shard_id, shards_count, partitioner, sharding_algorithm, sharding_ignore_msb, shard_aware_port,
+                 shard_aware_port_ssl):
         self.shards_count = int(shards_count)
         self.partitioner = partitioner
         self.sharding_algorithm = sharding_algorithm
         self.sharding_ignore_msb = int(sharding_ignore_msb)
-
+        self.shard_aware_port = int(shard_aware_port) if shard_aware_port else 0
+        self.shard_aware_port_ssl = int(shard_aware_port_ssl) if shard_aware_port_ssl else 0
 
     @staticmethod
     def parse_sharding_info(message):
@@ -39,12 +43,15 @@ cdef class ShardingInfo():
         partitioner = message.options.get('SCYLLA_PARTITIONER', [''])[0] or None
         sharding_algorithm = message.options.get('SCYLLA_SHARDING_ALGORITHM', [''])[0] or None
         sharding_ignore_msb = message.options.get('SCYLLA_SHARDING_IGNORE_MSB', [''])[0] or None
+        shard_aware_port = message.options.get('SCYLLA_SHARD_AWARE_PORT', [''])[0] or None
+        shard_aware_port_ssl = message.options.get('SCYLLA_SHARD_AWARE_PORT_SSL', [''])[0] or None
 
         if not (shard_id or shards_count or partitioner == "org.apache.cassandra.dht.Murmur3Partitioner" or
             sharding_algorithm == "biased-token-round-robin" or sharding_ignore_msb):
             return 0, None
 
-        return int(shard_id), ShardingInfo(shard_id, shards_count, partitioner, sharding_algorithm, sharding_ignore_msb)
+        return int(shard_id), ShardingInfo(shard_id, shards_count, partitioner, sharding_algorithm, sharding_ignore_msb,
+                                           shard_aware_port, shard_aware_port_ssl)
 
 
     def shard_id_from_token(self, int64_t token_input):

cassandra/cluster.py

@@ -1734,14 +1734,20 @@ def get_connection_holders(self):
         holders.append(self.control_connection)
         return holders
 
+    def get_all_pools(self):
+        pools = []
+        for s in tuple(self.sessions):
+            pools.extend(s.get_pools())
+        return pools
+
     def is_shard_aware(self):
-        return bool(self.get_connection_holders()[:-1][0].host.sharding_info)
+        return bool(self.get_all_pools()[0].host.sharding_info)
 
     def shard_aware_stats(self):
         if self.is_shard_aware():
             return {str(pool.host.endpoint): {'shards_count': pool.host.sharding_info.shards_count,
                                               'connected': len(pool._connections.keys())}
-                    for pool in self.get_connection_holders()[:-1]}
+                    for pool in self.get_all_pools()}
 
     def shutdown(self):
         """
@@ -3756,7 +3762,7 @@ def _refresh_node_list_and_token_map(self, connection, preloaded_results=None,
             partitioner = local_row.get("partitioner")
             tokens = local_row.get("tokens")
 
-            host = self._cluster.metadata.get_host(connection.endpoint)
+            host = self._cluster.metadata.get_host(connection.original_endpoint)
             if host:
                 datacenter = local_row.get("data_center")
                 rack = local_row.get("rack")
@@ -4049,9 +4055,8 @@ def _get_peers_query(self, peers_query_type, connection=None):
                 query_template = (self._SELECT_SCHEMA_PEERS_TEMPLATE
                                   if peers_query_type == self.PeersQueryType.PEERS_SCHEMA
                                   else self._SELECT_PEERS_NO_TOKENS_TEMPLATE)
-
-                host_release_version = self._cluster.metadata.get_host(connection.endpoint).release_version
-                host_dse_version = self._cluster.metadata.get_host(connection.endpoint).dse_version
+                host_release_version = self._cluster.metadata.get_host(connection.original_endpoint).release_version
+                host_dse_version = self._cluster.metadata.get_host(connection.original_endpoint).dse_version
                 uses_native_address_query = (
                     host_dse_version and Version(host_dse_version) >= self._MINIMUM_NATIVE_ADDRESS_DSE_VERSION)
 

cassandra/connection.py

@@ -28,7 +28,8 @@
 import time
 import ssl
 import weakref
-
+import random
+import itertools
 
 if 'gevent.monkey' in sys.modules:
     from gevent.queue import Queue, Empty
@@ -116,6 +117,10 @@ def decompress(byts):
 HEADER_DIRECTION_TO_CLIENT = 0x80
 HEADER_DIRECTION_MASK = 0x80
 
+# shard aware default for opening per shard connection
+DEFAULT_LOCAL_PORT_LOW = 49152
+DEFAULT_LOCAL_PORT_HIGH = 65535
+
 frame_header_v1_v2 = struct.Struct('>BbBi')
 frame_header_v3 = struct.Struct('>BhBi')
 
@@ -666,6 +671,17 @@ def reset_cql_frame_buffer(self):
             self.reset_io_buffer()
 
 
+class ShardawarePortGenerator:
+    @classmethod
+    def generate(cls, shard_id, total_shards):
+        start = random.randrange(DEFAULT_LOCAL_PORT_LOW, DEFAULT_LOCAL_PORT_HIGH)
+        available_ports = itertools.chain(range(start, DEFAULT_LOCAL_PORT_HIGH), range(DEFAULT_LOCAL_PORT_LOW, start))
+
+        for port in available_ports:
+            if port % total_shards == shard_id:
+                yield port
+
+
 class Connection(object):
 
     CALLBACK_ERR_THREAD_THRESHOLD = 100
@@ -762,7 +778,7 @@ def __init__(self, host='127.0.0.1', port=9042, authenticator=None,
                  ssl_options=None, sockopts=None, compression=True,
                  cql_version=None, protocol_version=ProtocolVersion.MAX_SUPPORTED, is_control_connection=False,
                  user_type_map=None, connect_timeout=None, allow_beta_protocol_version=False, no_compact=False,
-                 ssl_context=None, owning_pool=None):
+                 ssl_context=None, owning_pool=None, shard_id=None, total_shards=None):
 
         # TODO next major rename host to endpoint and remove port kwarg.
         self.endpoint = host if isinstance(host, EndPoint) else DefaultEndPoint(host, port)
@@ -812,6 +828,9 @@ def __init__(self, host='127.0.0.1', port=9042, authenticator=None,
 
         self.lock = RLock()
         self.connected_event = Event()
+        self.shard_id = shard_id
+        self.total_shards = total_shards
+        self.original_endpoint = self.endpoint
 
     @property
     def host(self):
@@ -874,6 +893,15 @@ def _wrap_socket_from_context(self):
         self._socket = self.ssl_context.wrap_socket(self._socket, **ssl_options)
 
     def _initiate_connection(self, sockaddr):
+        if self.shard_id is not None:
+            for port in ShardawarePortGenerator.generate(self.shard_id, self.total_shards):
+                try:
+                    self._socket.bind(('', port))
+                    break
+                except Exception as ex:
+                    log.debug("port=%d couldn't bind cause: %s", port, str(ex))
+            log.debug(f'connection (%r) port=%d should be shard_id=%d', id(self), port, port % self.total_shards)
+
         self._socket.connect(sockaddr)
 
     def _match_hostname(self):
@@ -894,6 +922,7 @@ def _get_socket_addresses(self):
     def _connect_socket(self):
         sockerr = None
         addresses = self._get_socket_addresses()
+        port = None
         for (af, socktype, proto, _, sockaddr) in addresses:
             try:
                 self._socket = self._socket_impl.socket(af, socktype, proto)

cassandra/metadata.py

@@ -134,8 +134,8 @@ def export_schema_as_string(self):
 
     def refresh(self, connection, timeout, target_type=None, change_type=None, **kwargs):
 
-        server_version = self.get_host(connection.endpoint).release_version
-        dse_version = self.get_host(connection.endpoint).dse_version
+        server_version = self.get_host(connection.original_endpoint).release_version
+        dse_version = self.get_host(connection.original_endpoint).dse_version
         parser = get_schema_parser(connection, server_version, dse_version, timeout)
 
         if not target_type:

cassandra/pool.py

@@ -21,6 +21,7 @@
 import socket
 import time
 import random
+import copy
 from threading import Lock, RLock, Condition
 import weakref
 try:
@@ -412,6 +413,7 @@ def __init__(self, host, host_distance, session):
         # so that we can dispose of them.
         self._trash = set()
         self._shard_connections_futures = []
+        self.advanced_shardaware_block_until = 0
 
         if host_distance == HostDistance.IGNORED:
             log.debug("Not opening connection to ignored host %s", self.host)
@@ -431,7 +433,7 @@ def __init__(self, host, host_distance, session):
 
         if first_connection.sharding_info:
             self.host.sharding_info = first_connection.sharding_info
-            self._open_connections_for_all_shards()
+            self._open_connections_for_all_shards(first_connection.shard_id)
 
         log.debug("Finished initializing connection for host %s", self.host)
 
@@ -645,6 +647,24 @@ def _close_excess_connections(self):
             log.debug("Closing excess connection (%s) to %s", id(c), self.host)
             c.close()
 
+    def disable_advanced_shard_aware(self, secs):
+        log.warning("disabling advanced_shard_aware for %i seconds, could be that this client is behind NAT?", secs)
+        self.advanced_shardaware_block_until = max(time.time() + secs, self.advanced_shardaware_block_until)
+
+    def _get_shard_aware_endpoint(self):
+        if self.advanced_shardaware_block_until and self.advanced_shardaware_block_until < time.time():
+            return None
+
+        endpoint = None
+        if self._session.cluster.ssl_options and self.host.sharding_info.shard_aware_port_ssl:
+            endpoint = copy.copy(self.host.endpoint)
+            endpoint._port = self.host.sharding_info.shard_aware_port_ssl
+        elif self.host.sharding_info.shard_aware_port:
+            endpoint = copy.copy(self.host.endpoint)
+            endpoint._port = self.host.sharding_info.shard_aware_port
+
+        return endpoint
+
     def _open_connection_to_missing_shard(self, shard_id):
         """
         Creates a new connection, checks its shard_id and populates our shard
@@ -666,13 +686,28 @@ def _open_connection_to_missing_shard(self, shard_id):
         with self._lock:
             if self.is_shutdown:
                 return
+        shard_aware_endpoint = self._get_shard_aware_endpoint()
+        log.debug("shard_aware_endpoint=%r", shard_aware_endpoint)
+
+        if shard_aware_endpoint:
+            conn = self._session.cluster.connection_factory(shard_aware_endpoint, owning_pool=self,
+                                                            shard_id=shard_id,
+                                                            total_shards=self.host.sharding_info.shards_count)
+            conn.original_endpoint = self.host.endpoint
+        else:
+            conn = self._session.cluster.connection_factory(self.host.endpoint, owning_pool=self)
 
-        conn = self._session.cluster.connection_factory(self.host.endpoint)
         log.debug("Received a connection %s for shard_id=%i on host %s", id(conn), conn.shard_id, self.host)
         if self.is_shutdown:
             log.debug("Pool for host %s is in shutdown, closing the new connection (%s)", self.host, id(conn))
             conn.close()
             return
+
+        if shard_aware_endpoint and shard_id != conn.shard_id:
+            # connection didn't land on expected shared
+            # assuming behind a NAT, disabling advanced shard aware for a while
+            self.disable_advanced_shard_aware(10 * 60)
+
         old_conn = self._connections.get(conn.shard_id)
         if old_conn is None or old_conn.orphaned_threshold_reached:
             log.debug(
@@ -776,7 +811,7 @@ def _open_connection_to_missing_shard(self, shard_id):
                 conn.close()
         self._connecting.discard(shard_id)
 
-    def _open_connections_for_all_shards(self):
+    def _open_connections_for_all_shards(self, skip_shard_id=None):
         """
         Loop over all the shards and try to open a connection to each one.
         """
@@ -785,6 +820,8 @@ def _open_connections_for_all_shards(self):
                 return
 
             for shard_id in range(self.host.sharding_info.shards_count):
+                if skip_shard_id and skip_shard_id == shard_id:
+                    continue
                 future = self._session.submit(self._open_connection_to_missing_shard, shard_id)
                 if isinstance(future, Future):
                     self._connecting.add(shard_id)

cassandra/shard_info.py

@@ -20,11 +20,13 @@
 
 class _ShardingInfo(object):
 
-    def __init__(self, shard_id, shards_count, partitioner, sharding_algorithm, sharding_ignore_msb):
+    def __init__(self, shard_id, shards_count, partitioner, sharding_algorithm, sharding_ignore_msb, shard_aware_port, shard_aware_port_ssl):
         self.shards_count = int(shards_count)
         self.partitioner = partitioner
         self.sharding_algorithm = sharding_algorithm
         self.sharding_ignore_msb = int(sharding_ignore_msb)
+        self.shard_aware_port = int(shard_aware_port) if shard_aware_port else None
+        self.shard_aware_port_ssl = int(shard_aware_port_ssl) if shard_aware_port_ssl else None
 
     @staticmethod
     def parse_sharding_info(message):
@@ -33,13 +35,16 @@ def parse_sharding_info(message):
         partitioner = message.options.get('SCYLLA_PARTITIONER', [''])[0] or None
         sharding_algorithm = message.options.get('SCYLLA_SHARDING_ALGORITHM', [''])[0] or None
         sharding_ignore_msb = message.options.get('SCYLLA_SHARDING_IGNORE_MSB', [''])[0] or None
+        shard_aware_port = message.options.get('SCYLLA_SHARD_AWARE_PORT', [''])[0] or None
+        shard_aware_port_ssl = message.options.get('SCYLLA_SHARD_AWARE_PORT_SSL', [''])[0] or None
         log.debug("Parsing sharding info from message options %s", message.options)
 
         if not (shard_id or shards_count or partitioner == "org.apache.cassandra.dht.Murmur3Partitioner" or
             sharding_algorithm == "biased-token-round-robin" or sharding_ignore_msb):
             return 0, None
 
-        return int(shard_id), _ShardingInfo(shard_id, shards_count, partitioner, sharding_algorithm, sharding_ignore_msb)
+        return int(shard_id), _ShardingInfo(shard_id, shards_count, partitioner, sharding_algorithm, sharding_ignore_msb,
+                                            shard_aware_port, shard_aware_port_ssl)
 
     def shard_id_from_token(self, token):
         """

docs/scylla_specific.rst

@@ -8,10 +8,13 @@ Shard Awareness
 As a result, latency is significantly reduced because there is no need to pass data between the shards.
 
 Details on the scylla cql protocol extensions
-https://github.com/scylladb/scylla/blob/master/docs/design-notes/protocol-extensions.md
+https://github.com/scylladb/scylla/blob/master/docs/design-notes/protocol-extensions.md#intranode-sharding
 
 For using it you only need to enable ``TokenAwarePolicy`` on the ``Cluster``
 
+See the configuration of ``native_shard_aware_transport_port`` and ``native_shard_aware_transport_port_ssl`` on scylla.yaml:
+https://github.com/scylladb/scylla/blob/master/docs/design-notes/protocols.md#cql-client-protocol
+
 .. code:: python
 
     from cassandra.cluster import Cluster

tests/integration/standard/test_cluster.py

@@ -1514,7 +1514,7 @@ def test_prepare_on_ignored_hosts(self):
         # the length of mock_calls will vary, but all should use the unignored
         # address
         for c in cluster.connection_factory.mock_calls:
-            self.assertEqual(call(DefaultEndPoint(unignored_address)), c)
+            self.assertEqual(unignored_address, c.args[0].address)
         cluster.shutdown()
 
 

tests/integration/standard/test_shard_aware.py

@@ -15,6 +15,7 @@
 import time
 import random
 from subprocess import run
+import logging
 
 try:
     from concurrent.futures import ThreadPoolExecutor, as_completed
@@ -28,10 +29,12 @@
 
 from cassandra.cluster import Cluster
 from cassandra.policies import TokenAwarePolicy, RoundRobinPolicy, ConstantReconnectionPolicy
-from cassandra import OperationTimedOut
+from cassandra import OperationTimedOut, ConsistencyLevel
 
 from tests.integration import use_cluster, get_node, PROTOCOL_VERSION
 
+LOGGER = logging.getLogger(__name__)
+
 
 def setup_module():
     os.environ['SCYLLA_EXT_OPTS'] = "--smp 4 --memory 2048M"
@@ -41,12 +44,12 @@ def setup_module():
 class TestShardAwareIntegration(unittest.TestCase):
     @classmethod
     def setup_class(cls):
-        cls.cluster = Cluster(protocol_version=PROTOCOL_VERSION, load_balancing_policy=TokenAwarePolicy(RoundRobinPolicy()),
+        cls.cluster = Cluster(contact_points=["127.0.0.1"], protocol_version=PROTOCOL_VERSION,
+                              load_balancing_policy=TokenAwarePolicy(RoundRobinPolicy()),
                               reconnection_policy=ConstantReconnectionPolicy(1))
         cls.session = cls.cluster.connect()
-
-        print(cls.cluster.is_shard_aware())
-        print(cls.cluster.shard_aware_stats())
+        LOGGER.info(cls.cluster.is_shard_aware())
+        LOGGER.info(cls.cluster.shard_aware_stats())
 
     @classmethod
     def teardown_class(cls):
@@ -56,7 +59,7 @@ def verify_same_shard_in_tracing(self, results, shard_name):
         traces = results.get_query_trace()
         events = traces.events
         for event in events:
-            print(event.thread_name, event.description)
+            LOGGER.info("%s %s", event.thread_name, event.description)
         for event in events:
             self.assertEqual(event.thread_name, shard_name)
         self.assertIn('querying locally', "\n".join([event.description for event in events]))
@@ -65,7 +68,7 @@ def verify_same_shard_in_tracing(self, results, shard_name):
         traces = self.session.execute("SELECT * FROM system_traces.events WHERE session_id = %s", (trace_id,))
         events = [event for event in traces]
         for event in events:
-            print(event.thread, event.activity)
+            LOGGER.info("%s %s", event.thread, event.activity)
         for event in events:
             self.assertEqual(event.thread, shard_name)
         self.assertIn('querying locally', "\n".join([event.activity for event in events]))