Merge pull request scylladb#251 from sylwiaszunejko/partition_rate_error

Add support for ScyllaDB's per partition rate limiting's new error

Author: Lorak-mmk

cassandra/__init__.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from enum import Enum
 import logging
 
 
@@ -728,3 +729,21 @@ class UnresolvableContactPoints(DriverException):
     contact points, only when lookup fails for all hosts
     """
     pass
+
+
+class OperationType(Enum):
+    Read = 0
+    Write = 1
+
+class RateLimitReached(ConfigurationException):
+    '''
+    Rate limit was exceeded for a partition affected by the request.
+    '''
+    op_type = None
+    rejected_by_coordinator = False
+
+    def __init__(self, op_type=None, rejected_by_coordinator=False):
+        self.op_type = op_type
+        self.rejected_by_coordinator = rejected_by_coordinator
+        message = f"[request_error_rate_limit_reached OpType={op_type.name} RejectedByCoordinator={rejected_by_coordinator}]"
+        Exception.__init__(self, message)

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

@@ -31,6 +31,8 @@
 import random
 import itertools
 
+from cassandra.protocol_features import ProtocolFeatures
+
 if 'gevent.monkey' in sys.modules:
     from gevent.queue import Queue, Empty
 else:
@@ -765,13 +767,12 @@ class Connection(object):
 
     _owning_pool = None
 
-    shard_id = 0
-    sharding_info = None
-
     _is_checksumming_enabled = False
 
     _on_orphaned_stream_released = None
 
+    features = None
+
     @property
     def _iobuf(self):
         # backward compatibility, to avoid any change in the reactors
@@ -831,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
 
@@ -896,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
@@ -1263,7 +1264,7 @@ def process_msg(self, header, body):
                     return
 
         try:
-            response = decoder(header.version, self.user_type_map, stream_id,
+            response = decoder(header.version, self.features, self.user_type_map, stream_id,
                                header.flags, header.opcode, body, self.decompressor, result_metadata)
         except Exception as exc:
             log.exception("Error decoding response from Cassandra. "
@@ -1318,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
 
@@ -1338,6 +1339,9 @@ 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]
 
+        options = {}
+        self.features.add_startup_options(options)
+
         if self.cql_version:
             if self.cql_version not in supported_cql_versions:
                 raise ProtocolError(
@@ -1388,13 +1392,14 @@ def _handle_options_response(self, options_response):
                     self._compressor, self.decompressor = \
                         locally_supported_compressions[compression_type]
 
-        self._send_startup_message(compression_type, no_compact=self.no_compact)
+        self._send_startup_message(compression_type, no_compact=self.no_compact, extra_options=options)
 
     @defunct_on_error
-    def _send_startup_message(self, compression=None, no_compact=False):
+    def _send_startup_message(self, compression=None, no_compact=False, extra_options=None):
         log.debug("Sending StartupMessage on %s", self)
         opts = {'DRIVER_NAME': DRIVER_NAME,
-                'DRIVER_VERSION': DRIVER_VERSION}
+                'DRIVER_VERSION': DRIVER_VERSION,
+                **extra_options}
         if compression:
             opts['COMPRESSION'] = compression
         if no_compact:

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.py

@@ -22,9 +22,9 @@
 from six.moves import range
 import io
 
-from cassandra import ProtocolVersion
+from cassandra import OperationType, ProtocolVersion
 from cassandra import type_codes, DriverException
-from cassandra import (Unavailable, WriteTimeout, ReadTimeout,
+from cassandra import (Unavailable, WriteTimeout, RateLimitReached, ReadTimeout,
                        WriteFailure, ReadFailure, FunctionFailure,
                        AlreadyExists, InvalidRequest, Unauthorized,
                        UnsupportedOperation, UserFunctionDescriptor,
@@ -126,10 +126,13 @@ def __init__(self, code, message, info):
         self.info = info
 
     @classmethod
-    def recv_body(cls, f, protocol_version, *args):
+    def recv_body(cls, f, protocol_version, protocol_features, *args):
         code = read_int(f)
         msg = read_string(f)
-        subcls = error_classes.get(code, cls)
+        if code == protocol_features.rate_limit_error:
+            subcls = RateLimitReachedException
+        else:
+            subcls = error_classes.get(code, cls)
         extra_info = subcls.recv_error_info(f, protocol_version)
         return subcls(code=code, message=msg, info=extra_info)
 
@@ -390,6 +393,19 @@ def recv_error_info(f, protocol_version):
     def to_exception(self):
         return AlreadyExists(**self.info)
 
+class RateLimitReachedException(ConfigurationException):
+    summary= 'Rate limit was exceeded for a partition affected by the request'
+    error_code = 0x4321
+
+    @staticmethod
+    def recv_error_info(f, protocol_version):
+        return {
+            'op_type': OperationType(read_byte(f)),
+            'rejected_by_coordinator': read_byte(f) != 0
+        }
+
+    def to_exception(self):
+        return RateLimitReached(**self.info)
 
 class ClientWriteError(RequestExecutionException):
     summary = 'Client write failure.'
@@ -738,7 +754,7 @@ def recv(self, f, protocol_version, user_type_map, result_metadata):
             raise DriverException("Unknown RESULT kind: %d" % self.kind)
 
     @classmethod
-    def recv_body(cls, f, protocol_version, user_type_map, result_metadata):
+    def recv_body(cls, f, protocol_version, protocol_features, user_type_map, result_metadata):
         kind = read_int(f)
         msg = cls(kind)
         msg.recv(f, protocol_version, user_type_map, result_metadata)
@@ -1147,7 +1163,7 @@ def _write_header(f, version, flags, stream_id, opcode, length):
         write_int(f, length)
 
     @classmethod
-    def decode_message(cls, protocol_version, user_type_map, stream_id, flags, opcode, body,
+    def decode_message(cls, protocol_version, protocol_features, user_type_map, stream_id, flags, opcode, body,
                        decompressor, result_metadata):
         """
         Decodes a native protocol message body
@@ -1193,7 +1209,7 @@ def decode_message(cls, protocol_version, user_type_map, stream_id, flags, opcod
             log.warning("Unknown protocol flags set: %02x. May cause problems.", flags)
 
         msg_class = cls.message_types_by_opcode[opcode]
-        msg = msg_class.recv_body(body, protocol_version, user_type_map, result_metadata)
+        msg = msg_class.recv_body(body, protocol_version, protocol_features, user_type_map, result_metadata)
         msg.stream_id = stream_id
         msg.trace_id = trace_id
         msg.custom_payload = custom_payload