Initial shard aware driver

fruch · fruch · commit ac58615fd51a · 2020-03-11T18:28:45.000+02:00
diff --git a/cassandra/cluster.py b/cassandra/cluster.py
@@ -4384,7 +4384,7 @@ 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)
+            connection, request_id = pool.borrow_connection(timeout=2.0, routing_key=self.query.routing_key if self.query else None)
             self._connection = connection
             result_meta = self.prepared_statement.result_metadata if self.prepared_statement else []
 
diff --git a/cassandra/connection.py b/cassandra/connection.py
@@ -44,6 +44,7 @@
                                 RegisterMessage, ReviseRequestMessage)
 from cassandra.util import OrderedDict
 
+MIN_LONG = -(2 ** 63)
 
 log = logging.getLogger(__name__)
 
@@ -599,6 +600,39 @@ def int_from_buf_item(i):
 else:
     int_from_buf_item = ord
 
+class ShardingInfo(object):
+
+    def __init__(self, shard_id, shards_count, partitioner, sharding_algorithm, sharding_ignore_msb):
+        self.shards_count = int(shards_count)
+        self.partitioner = partitioner
+        self.sharding_algorithm = sharding_algorithm
+        self.sharding_ignore_msb = int(sharding_ignore_msb)
+
+    @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
+
+        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)
+
+    def shard_id(self, t):
+        token = t.value
+        token += MIN_LONG
+        token <<= self.sharding_ignore_msb
+        tokLo = token & 0xffffffff
+        tokHi = (token >> 32) & 0xffffffff
+        mul1 = tokLo * self.shards_count
+        mul2 = tokHi * self.shards_count
+        _sum = (mul1 >> 32) + mul2
+        output = _sum >> 32
+        return output
 
 class Connection(object):
 
@@ -666,6 +700,9 @@ class Connection(object):
     _check_hostname = False
     _product_type = None
 
+    shard_id = 0
+    sharding_info = None
+
     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,
@@ -1126,6 +1163,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)
         if self.is_defunct:
             return
 
diff --git a/cassandra/pool.py b/cassandra/pool.py
@@ -20,6 +20,7 @@
 import logging
 import socket
 import time
+import random
 from threading import Lock, RLock, Condition
 import weakref
 try:
@@ -123,6 +124,8 @@ class Host(object):
 
     _currently_handling_node_up = False
 
+    sharding_info = None
+
     def __init__(self, endpoint, conviction_policy_factory, datacenter=None, rack=None, host_id=None):
         if endpoint is None:
             raise ValueError("endpoint may not be None")
@@ -339,7 +342,6 @@ class HostConnection(object):
     shutdown_on_error = False
 
     _session = None
-    _connection = None
     _lock = None
     _keyspace = None
 
@@ -351,6 +353,7 @@ def __init__(self, host, host_distance, session):
         # this is used in conjunction with the connection streams. Not using the connection lock because the connection can be replaced in the lifetime of the pool.
         self._stream_available_condition = Condition(self._lock)
         self._is_replacing = False
+        self._connections = dict()
 
         if host_distance == HostDistance.IGNORED:
             log.debug("Not opening connection to ignored host %s", self.host)
@@ -360,18 +363,45 @@ def __init__(self, host, host_distance, session):
             return
 
         log.debug("Initializing connection for host %s", self.host)
-        self._connection = session.cluster.connection_factory(host.endpoint)
+        first_connection = session.cluster.connection_factory(host.endpoint)
+        log.debug("first connection created for shard_id=%i", first_connection.shard_id)
+        self._connections[first_connection.shard_id] = first_connection
         self._keyspace = session.keyspace
+
         if self._keyspace:
-            self._connection.set_keyspace_blocking(self._keyspace)
+            first_connection.set_keyspace_blocking(self._keyspace)
+
+        if first_connection.sharding_info:
+            self.host.sharding_info = weakref.proxy(first_connection.sharding_info)
+            for _ in range(first_connection.sharding_info.shards_count * 2):
+                conn = self._session.cluster.connection_factory(self.host.endpoint)
+                if conn.shard_id not in self._connections.keys():
+                    log.debug("new connection created for shard_id=%i", conn.shard_id)
+                    self._connections[conn.shard_id] = conn
+                    if self._keyspace:
+                        self._connections[conn.shard_id].set_keyspace_blocking(self._keyspace)
+
+                if len(self._connections.keys()) == first_connection.sharding_info.shards_count:
+                    break
+            if not len(self._connections.keys()) == first_connection.sharding_info.shards_count:
+                raise NoConnectionsAvailable("not enough shard connection opened")
+
         log.debug("Finished initializing connection for host %s", self.host)
 
-    def borrow_connection(self, timeout):
+    def borrow_connection(self, timeout, routing_key=None):
         if self.is_shutdown:
             raise ConnectionException(
                 "Pool for %s is shutdown" % (self.host,), self.host)
 
-        conn = self._connection
+        shard_id = 0
+        if self.host.sharding_info:
+            if routing_key:
+                t = self._session.cluster.metadata.token_map.token_class.from_key(routing_key)
+                shard_id =self.host.sharding_info.shard_id(t)
+            else:
+                shard_id = random.randint(0, self.host.sharding_info.shards_count - 1)
+
+        conn = self._connections.get(shard_id)
         if not conn:
             raise NoConnectionsAvailable()
 
@@ -416,7 +446,7 @@ def return_connection(self, connection):
             if is_down:
                 self.shutdown()
             else:
-                self._connection = None
+                del self._connections[connection.shard_id]
                 with self._lock:
                     if self._is_replacing:
                         return
@@ -433,7 +463,7 @@ def _replace(self, connection):
             conn = self._session.cluster.connection_factory(self.host.endpoint)
             if self._keyspace:
                 conn.set_keyspace_blocking(self._keyspace)
-            self._connection = conn
+            self._connections[connection.shard_id] = conn
         except Exception:
             log.warning("Failed reconnecting %s. Retrying." % (self.host.endpoint,))
             self._session.submit(self._replace, connection)
@@ -450,36 +480,48 @@ def shutdown(self):
                 self.is_shutdown = True
             self._stream_available_condition.notify_all()
 
-        if self._connection:
-            self._connection.close()
-            self._connection = None
+        if self._connections:
+            for c in self._connections.values():
+                c.close()
+            self._connections = dict()
 
     def _set_keyspace_for_all_conns(self, keyspace, callback):
-        if self.is_shutdown or not self._connection:
+        """
+        Asynchronously sets the keyspace for all connections.  When all
+        connections have been set, `callback` will be called with two
+        arguments: this pool, and a list of any errors that occurred.
+        """
+        remaining_callbacks = set(self._connections.values())
+        errors = []
+
+        if not remaining_callbacks:
+            callback(self, errors)
             return
 
         def connection_finished_setting_keyspace(conn, error):
             self.return_connection(conn)
-            errors = [] if not error else [error]
-            callback(self, errors)
+            remaining_callbacks.remove(conn)
+            if error:
+                errors.append(error)
+
+            if not remaining_callbacks:
+                callback(self, errors)
 
         self._keyspace = keyspace
-        self._connection.set_keyspace_async(keyspace, connection_finished_setting_keyspace)
+        for conn in self._connections.values():
+            conn.set_keyspace_async(keyspace, connection_finished_setting_keyspace)
 
     def get_connections(self):
-        c = self._connection
-        return [c] if c else []
+        c = self._connections
+        return list(self._connections.values()) if c else []
 
     def get_state(self):
-        connection = self._connection
-        open_count = 1 if connection and not (connection.is_closed or connection.is_defunct) else 0
-        in_flights = [connection.in_flight] if connection else []
-        return {'shutdown': self.is_shutdown, 'open_count': open_count, 'in_flights': in_flights}
+        in_flights = [c.in_flight for c in self._connections.values()]
+        return {'shutdown': self.is_shutdown, 'open_count': self.open_count, 'in_flights': in_flights}
 
     @property
     def open_count(self):
-        connection = self._connection
-        return 1 if connection and not (connection.is_closed or connection.is_defunct) else 0
+        return sum([1 if c and not (c.is_closed or c.is_defunct) else 0 for c in self._connections.values()])
 
 _MAX_SIMULTANEOUS_CREATION = 1
 _MIN_TRASH_INTERVAL = 10
@@ -522,7 +564,7 @@ def __init__(self, host, host_distance, session):
         self.open_count = core_conns
         log.debug("Finished initializing new connection pool for host %s", self.host)
 
-    def borrow_connection(self, timeout):
+    def borrow_connection(self, timeout, routing_key=None):
         if self.is_shutdown:
             raise ConnectionException(
                 "Pool for %s is shutdown" % (self.host,), self.host)
diff --git a/tests/integration/standard/test_connection.py b/tests/integration/standard/test_connection.py
@@ -172,8 +172,8 @@ def fetch_connections(self, host, cluster):
                     if conn._connections is not None and len(conn._connections) > 0:
                         connections.append(conn._connections)
                 else:
-                    if conn._connection is not None:
-                        connections.append(conn._connection)
+                    if conn._connections and len(conn._connections.values()) > 0:
+                        connections.append(conn._connections.values())
         return connections
 
     def wait_for_connections(self, host, cluster):
diff --git a/tests/unit/test_response_future.py b/tests/unit/test_response_future.py
@@ -74,7 +74,7 @@ def test_result_message(self):
         rf.send_request()
 
         rf.session._pools.get.assert_called_once_with('ip1')
-        pool.borrow_connection.assert_called_once_with(timeout=ANY)
+        pool.borrow_connection.assert_called_once_with(timeout=ANY, routing_key=ANY)
 
         connection.send_msg.assert_called_once_with(rf.message, 1, cb=ANY, encoder=ProtocolHandler.encode_message, decoder=ProtocolHandler.decode_message, result_metadata=[])
 
@@ -256,7 +256,7 @@ def test_retry_policy_says_retry(self):
         rf.send_request()
 
         rf.session._pools.get.assert_called_once_with('ip1')
-        pool.borrow_connection.assert_called_once_with(timeout=ANY)
+        pool.borrow_connection.assert_called_once_with(timeout=ANY, routing_key=ANY)
         connection.send_msg.assert_called_once_with(rf.message, 1, cb=ANY, encoder=ProtocolHandler.encode_message, decoder=ProtocolHandler.decode_message, result_metadata=[])
 
         result = Mock(spec=UnavailableErrorMessage, info={})
@@ -275,7 +275,7 @@ def test_retry_policy_says_retry(self):
         # it should try again with the same host since this was
         # an UnavailableException
         rf.session._pools.get.assert_called_with(host)
-        pool.borrow_connection.assert_called_with(timeout=ANY)
+        pool.borrow_connection.assert_called_with(timeout=ANY, routing_key=ANY)
         connection.send_msg.assert_called_with(rf.message, 2, cb=ANY, encoder=ProtocolHandler.encode_message, decoder=ProtocolHandler.decode_message, result_metadata=[])
 
     def test_retry_with_different_host(self):
@@ -290,7 +290,7 @@ def test_retry_with_different_host(self):
         rf.send_request()
 
         rf.session._pools.get.assert_called_once_with('ip1')
-        pool.borrow_connection.assert_called_once_with(timeout=ANY)
+        pool.borrow_connection.assert_called_once_with(timeout=ANY, routing_key=ANY)
         connection.send_msg.assert_called_once_with(rf.message, 1, cb=ANY, encoder=ProtocolHandler.encode_message, decoder=ProtocolHandler.decode_message, result_metadata=[])
         self.assertEqual(ConsistencyLevel.QUORUM, rf.message.consistency_level)
 
@@ -309,7 +309,7 @@ def test_retry_with_different_host(self):
 
         # it should try with a different host
         rf.session._pools.get.assert_called_with('ip2')
-        pool.borrow_connection.assert_called_with(timeout=ANY)
+        pool.borrow_connection.assert_called_with(timeout=ANY, routing_key=ANY)
         connection.send_msg.assert_called_with(rf.message, 2, cb=ANY, encoder=ProtocolHandler.encode_message, decoder=ProtocolHandler.decode_message, result_metadata=[])
 
         # the consistency level should be the same
