perf: avoid redundant tablet lookup in shard-aware connection selection

When tablets are in use, get_tablet_for_key() was called twice per
request: once in TokenAwarePolicy.make_query_plan() to find the replica,
and again in HostConnection._get_connection_for_routing_key() to
determine the shard.

Stash the tablet found during query planning on the query object
(query._tablet) and pass it through to borrow_connection(), which
skips the second lookup when a tablet is already available. This
eliminates redundant bisect_left calls and associated dict lookups.

Author: mykaul

cassandra/cluster.py

@@ -4613,7 +4613,10 @@ def _query(self, host, message=None, cb=None):
         try:
             # TODO get connectTimeout from cluster settings
             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)
+                connection, request_id = pool.borrow_connection(
+                    timeout=2.0, routing_key=self.query.routing_key,
+                    keyspace=self.query.keyspace, table=self.query.table,
+                    tablet=getattr(self.query, '_tablet', None))
             else:
                 connection, request_id = pool.borrow_connection(timeout=2.0)
             self._connection = connection

cassandra/policies.py

@@ -511,6 +511,10 @@ def make_query_plan(self, working_keyspace=None, query=None):
             child_plan = child.make_query_plan(keyspace, query)
 
             replicas = [host for host in child_plan if host.host_id in replica_dict]
+            # Stash the tablet so that downstream shard-aware
+            # connection selection can reuse it instead of
+            # repeating the bisect lookup.
+            query._tablet = tablet
         else:
             replicas = self._cluster_metadata.get_replicas(keyspace, query.routing_key)
 

cassandra/pool.py

@@ -442,7 +442,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, keyspace=None, table=None):
+    def _get_connection_for_routing_key(self, routing_key=None, keyspace=None, table=None, tablet=None):
         if self.is_shutdown:
             raise ConnectionException(
                 "Pool for %s is shutdown" % (self.host,), self.host)
@@ -456,13 +456,18 @@ def _get_connection_for_routing_key(self, routing_key=None, keyspace=None, table
 
             shard_id = None
             if self.tablets_routing_v1 and table is not None:
-                if keyspace is None:
-                    keyspace = self._keyspace
-
-                tablet = self._session.cluster.metadata._tablets.get_tablet_for_key(keyspace, table, t)
-
+                # Reuse tablet from query planning if available, avoiding
+                # a redundant bisect lookup in the tablet map.
                 if tablet is not None:
                     shard_id = tablet._replica_dict.get(self.host.host_id)
+                else:
+                    if keyspace is None:
+                        keyspace = self._keyspace
+
+                    tablet = self._session.cluster.metadata._tablets.get_tablet_for_key(keyspace, table, t)
+
+                    if tablet is not None:
+                        shard_id = tablet._replica_dict.get(self.host.host_id)
 
             if shard_id is None:
                 shard_id = self.host.sharding_info.shard_id_from_token(t.value)
@@ -505,15 +510,15 @@ def _get_connection_for_routing_key(self, routing_key=None, keyspace=None, table
             return random.choice(active_connections)
         return random.choice(list(self._connections.values()))
 
-    def borrow_connection(self, timeout, routing_key=None, keyspace=None, table=None):
-        conn = self._get_connection_for_routing_key(routing_key, keyspace, table)
+    def borrow_connection(self, timeout, routing_key=None, keyspace=None, table=None, tablet=None):
+        conn = self._get_connection_for_routing_key(routing_key, keyspace, table, tablet)
         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, keyspace, table)
+                conn = self._get_connection_for_routing_key(routing_key, keyspace, table, tablet)
             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