Move sharding info to ProtocolFeatures

Sharding is a protocol extention, now sharing-related info is a part of
ProtocolFeatures class, also _ShardingInfo.parse_sharding_info is moved to
ProtocolFeatures to have all features strings in one place.

Author: sylwiaszunejko

cassandra/c_shard_info.pyx

@@ -36,24 +36,6 @@ cdef class ShardingInfo():
         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):
-        shard_id = message.options.get('SCYLLA_SHARD', [''])[0] or None
-        shards_count = message.options.get('SCYLLA_NR_SHARDS', [''])[0] or None
-        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,
-                                           shard_aware_port, shard_aware_port_ssl)
-
-    
     def shard_id_from_token(self, int64_t token_input):
         cdef uint64_t biased_token = token_input + (<uint64_t>1 << 63);
         biased_token <<= self.sharding_ignore_msb;

cassandra/connection.py

@@ -767,9 +767,6 @@ class Connection(object):
 
     _owning_pool = None
 
-    shard_id = 0
-    sharding_info = None
-
     _is_checksumming_enabled = False
 
     _on_orphaned_stream_released = None
@@ -835,7 +832,7 @@ 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.features = ProtocolFeatures(shard_id=shard_id)
         self.total_shards = total_shards
         self.original_endpoint = self.endpoint
 
@@ -900,8 +897,8 @@ 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):
+        if self.features.shard_id is not None:
+            for port in ShardawarePortGenerator.generate(self.features.shard_id, self.total_shards):
                 try:
                     self._socket.bind(('', port))
                     break
@@ -1322,7 +1319,7 @@ def _send_options_message(self):
 
     @defunct_on_error
     def _handle_options_response(self, options_response):
-        self.shard_id, self.sharding_info = ShardingInfo.parse_sharding_info(options_response)
+        self.features = ProtocolFeatures.parse_from_supported(options_response.options)
         if self.is_defunct:
             return
 
@@ -1342,10 +1339,8 @@ def _handle_options_response(self, options_response):
         remote_supported_compressions = options_response.options['COMPRESSION']
         self._product_type = options_response.options.get('PRODUCT_TYPE', [None])[0]
 
-        protocol_features = ProtocolFeatures.parse_from_supported(options_response.options)
         options = {}
-        protocol_features.add_startup_options(options)
-        self.features = protocol_features
+        self.features.add_startup_options(options)
 
         if self.cql_version:
             if self.cql_version not in supported_cql_versions:

cassandra/pool.py

@@ -427,15 +427,15 @@ def __init__(self, host, host_distance, session):
 
         log.debug("Initializing connection for host %s", self.host)
         first_connection = session.cluster.connection_factory(self.host.endpoint, on_orphaned_stream_released=self.on_orphaned_stream_released)
-        log.debug("First connection created to %s for shard_id=%i", self.host, first_connection.shard_id)
-        self._connections[first_connection.shard_id] = first_connection
+        log.debug("First connection created to %s for shard_id=%i", self.host, first_connection.features.shard_id)
+        self._connections[first_connection.features.shard_id] = first_connection
         self._keyspace = session.keyspace
 
         if self._keyspace:
             first_connection.set_keyspace_blocking(self._keyspace)
-        if first_connection.sharding_info and not self._session.cluster.shard_aware_options.disable:
-            self.host.sharding_info = first_connection.sharding_info
-            self._open_connections_for_all_shards(first_connection.shard_id)
+        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)
 
         log.debug("Finished initializing connection for host %s", self.host)
 
@@ -556,7 +556,7 @@ def return_connection(self, connection, stream_was_orphaned=False):
                 with self._lock:
                     if self.is_shutdown:
                         return
-                    self._connections.pop(connection.shard_id, None)
+                    self._connections.pop(connection.features.shard_id, None)
                     if self._is_replacing:
                         return
                     self._is_replacing = True
@@ -587,17 +587,17 @@ def _replace(self, connection):
 
             log.debug("Replacing connection (%s) to %s", id(connection), self.host)
             try:
-                if connection.shard_id in self._connections.keys():
-                    del self._connections[connection.shard_id]
+                if connection.features.shard_id in self._connections.keys():
+                    del self._connections[connection.features.shard_id]
                 if self.host.sharding_info and not self._session.cluster.shard_aware_options.disable:
-                    self._connecting.add(connection.shard_id)
-                    self._session.submit(self._open_connection_to_missing_shard, connection.shard_id)
+                    self._connecting.add(connection.features.shard_id)
+                    self._session.submit(self._open_connection_to_missing_shard, connection.features.shard_id)
                 else:
                     connection = self._session.cluster.connection_factory(self.host.endpoint,
                                                                           on_orphaned_stream_released=self.on_orphaned_stream_released)
                     if self._keyspace:
                         connection.set_keyspace_blocking(self._keyspace)
-                    self._connections[connection.shard_id] = connection
+                    self._connections[connection.features.shard_id] = connection
             except Exception:
                 log.warning("Failed reconnecting %s. Retrying." % (self.host.endpoint,))
                 self._session.submit(self._replace, connection)
@@ -703,59 +703,59 @@ def _open_connection_to_missing_shard(self, shard_id):
         else:
             conn = self._session.cluster.connection_factory(self.host.endpoint, on_orphaned_stream_released=self.on_orphaned_stream_released)
 
-        log.debug("Received a connection %s for shard_id=%i on host %s", id(conn), conn.shard_id, self.host)
+        log.debug("Received a connection %s for shard_id=%i on host %s", id(conn), conn.features.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:
+        if shard_aware_endpoint and shard_id != conn.features.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)
+        old_conn = self._connections.get(conn.features.shard_id)
         if old_conn is None or old_conn.orphaned_threshold_reached:
             log.debug(
                 "New connection (%s) created to shard_id=%i on host %s",
                 id(conn),
-                conn.shard_id,
+                conn.features.shard_id,
                 self.host
             )
             old_conn = None
             with self._lock:
                 if self.is_shutdown:
                     conn.close()
                     return
-                if conn.shard_id in self._connections.keys():
+                if conn.features.shard_id in self._connections.keys():
                     # Move the current connection to the trash and use the new one from now on
-                    old_conn = self._connections[conn.shard_id]
+                    old_conn = self._connections[conn.features.shard_id]
                     log.debug(
                         "Replacing overloaded connection (%s) with (%s) for shard %i for host %s",
                         id(old_conn),
                         id(conn),
-                        conn.shard_id,
+                        conn.features.shard_id,
                         self.host
                     )
                 if self._keyspace:
                     conn.set_keyspace_blocking(self._keyspace)
 
-                self._connections[conn.shard_id] = conn
+                self._connections[conn.features.shard_id] = conn
             if old_conn is not None:
                 remaining = old_conn.in_flight - len(old_conn.orphaned_request_ids)
                 if remaining == 0:
                     log.debug(
                         "Immediately closing the old connection (%s) for shard %i on host %s",
                         id(old_conn),
-                        old_conn.shard_id,
+                        old_conn.features.shard_id,
                         self.host
                     )
                     old_conn.close()
                 else:
                     log.debug(
                         "Moving the connection (%s) for shard %i to trash on host %s, %i requests remaining",
                         id(old_conn),
-                        old_conn.shard_id,
+                        old_conn.features.shard_id,
                         self.host,
                         remaining,
                     )
@@ -800,7 +800,7 @@ def _open_connection_to_missing_shard(self, shard_id):
             log.debug(
                 "Putting a connection %s to shard %i to the excess pool of host %s",
                 id(conn),
-                conn.shard_id,
+                conn.features.shard_id,
                 self.host
             )
             close_connection = False

cassandra/protocol_features.py

@@ -1,19 +1,27 @@
 import logging
 
+from cassandra.shard_info import _ShardingInfo
+
 log = logging.getLogger(__name__)
 
 
 RATE_LIMIT_ERROR_EXTENSION = "SCYLLA_RATE_LIMIT_ERROR"
 
 class ProtocolFeatures(object):
     rate_limit_error = None
+    shard_id = 0
+    sharding_info = None
 
-    def __init__(self, rate_limit_error=None):
+    def __init__(self, rate_limit_error=None, shard_id=0, sharding_info=None):
         self.rate_limit_error = rate_limit_error
+        self.shard_id = shard_id
+        self.sharding_info = sharding_info
 
     @staticmethod
     def parse_from_supported(supported):
-        return ProtocolFeatures(rate_limit_error = ProtocolFeatures.maybe_parse_rate_limit_error(supported))
+        rate_limit_error = ProtocolFeatures.maybe_parse_rate_limit_error(supported)
+        shard_id, sharding_info = ProtocolFeatures.parse_sharding_info(supported)
+        return ProtocolFeatures(rate_limit_error, shard_id, sharding_info)
 
     @staticmethod
     def maybe_parse_rate_limit_error(supported):
@@ -36,3 +44,22 @@ def add_startup_options(self, options):
         if self.rate_limit_error is not None:
             options[RATE_LIMIT_ERROR_EXTENSION] = ""
 
+    @staticmethod
+    def parse_sharding_info(options):
+        shard_id = options.get('SCYLLA_SHARD', [''])[0] or None
+        shards_count = options.get('SCYLLA_NR_SHARDS', [''])[0] or None
+        partitioner = options.get('SCYLLA_PARTITIONER', [''])[0] or None
+        sharding_algorithm = options.get('SCYLLA_SHARDING_ALGORITHM', [''])[0] or None
+        sharding_ignore_msb = options.get('SCYLLA_SHARDING_IGNORE_MSB', [''])[0] or None
+        shard_aware_port = options.get('SCYLLA_SHARD_AWARE_PORT', [''])[0] or None
+        shard_aware_port_ssl = options.get('SCYLLA_SHARD_AWARE_PORT_SSL', [''])[0] or None
+        log.debug("Parsing sharding info from message options %s", 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,
+                                            shard_aware_port, shard_aware_port_ssl)
+
+

cassandra/shard_info.py

@@ -28,24 +28,6 @@ def __init__(self, shard_id, shards_count, partitioner, sharding_algorithm, shar
         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):
-        shard_id = message.options.get('SCYLLA_SHARD', [''])[0] or None
-        shards_count = message.options.get('SCYLLA_NR_SHARDS', [''])[0] or None
-        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,
-                                            shard_aware_port, shard_aware_port_ssl)
-
     def shard_id_from_token(self, token):
         """
         Convert a Murmur3 token to shard_id based on the number of shards on the host

tests/unit/test_host_connection_pool.py

@@ -14,6 +14,7 @@
 from concurrent.futures import ThreadPoolExecutor
 import logging
 import time
+from cassandra.protocol_features import ProtocolFeatures
 
 from cassandra.shard_info import _ShardingInfo
 
@@ -300,11 +301,11 @@ def mock_connection_factory(self, *args, **kwargs):
                 connection.is_shutdown = False
                 connection.is_defunct = False
                 connection.is_closed = False
-                connection.shard_id = self.connection_counter
+                connection.features = ProtocolFeatures(shard_id=self.connection_counter, 
+                                                       sharding_info=_ShardingInfo(shard_id=1, shards_count=14,
+                                                                    partitioner="", sharding_algorithm="", sharding_ignore_msb=0,
+                                                                    shard_aware_port="", shard_aware_port_ssl=""))
                 self.connection_counter += 1
-                connection.sharding_info = _ShardingInfo(shard_id=1, shards_count=14,
-                                                         partitioner="", sharding_algorithm="", sharding_ignore_msb=0,
-                                                         shard_aware_port="", shard_aware_port_ssl="")
 
                 return connection
 

tests/unit/test_shard_aware.py

@@ -25,6 +25,7 @@
 from cassandra.pool import HostConnection, HostDistance
 from cassandra.connection import ShardingInfo, DefaultEndPoint
 from cassandra.metadata import Murmur3Token
+from cassandra.protocol_features import ProtocolFeatures
 
 LOGGER = logging.getLogger(__name__)
 
@@ -43,7 +44,7 @@ class OptionsHolder(object):
                 'SCYLLA_SHARDING_ALGORITHM': ['biased-token-round-robin'],
                 'SCYLLA_SHARDING_IGNORE_MSB': ['12']
             }
-        shard_id, shard_info = ShardingInfo.parse_sharding_info(OptionsHolder())
+        shard_id, shard_info = ProtocolFeatures.parse_sharding_info(OptionsHolder().options)
 
         self.assertEqual(shard_id, 1)
         self.assertEqual(shard_info.shard_id_from_token(Murmur3Token.from_key(b"a").value), 4)
@@ -88,12 +89,10 @@ def mock_connection_factory(self, *args, **kwargs):
                 connection.is_defunct = False
                 connection.is_closed = False
                 connection.orphaned_threshold_reached = False
-                connection.endpoint = args[0]
-                connection.shard_id = kwargs.get('shard_id', self.connection_counter)
+                connection.endpoint = args[0] 
+                sharding_info = ShardingInfo(shard_id=1, shards_count=4, partitioner="", sharding_algorithm="", sharding_ignore_msb=0, shard_aware_port=19042, shard_aware_port_ssl=19045)
+                connection.features = ProtocolFeatures(shard_id=kwargs.get('shard_id', self.connection_counter), sharding_info=sharding_info)
                 self.connection_counter += 1
-                connection.sharding_info = ShardingInfo(shard_id=1, shards_count=4,
-                                                         partitioner="", sharding_algorithm="", sharding_ignore_msb=0,
-                                                         shard_aware_port=19042, shard_aware_port_ssl=19045)
 
                 return connection
 
@@ -107,7 +106,7 @@ def mock_connection_factory(self, *args, **kwargs):
                 f.result()
             assert len(pool._connections) == 4
             for shard_id, connection in pool._connections.items():
-                assert connection.shard_id == shard_id
+                assert connection.features.shard_id == shard_id
                 if shard_id == 0:
                     assert connection.endpoint == DefaultEndPoint("1.2.3.4")
                 else: