Added routing to shards on top of TokenAware policy

At this point metrics indicate that cross-shard ops are reduced,
but the implementation is still raw. The reasons:
When connection-to-the-right-shard is being searched among per-host
connections, it is done by linear search. Connecting logic does not
attempt to reconnect until all shards are hit. Topology change
events are not accounted for.

Author: jul-stas

src/connection_pool.cpp

@@ -95,13 +95,25 @@ ConnectionPool::ConnectionPool(const Connection::Vec& connections, ConnectionPoo
   }
 }
 
-PooledConnection::Ptr ConnectionPool::find_least_busy() const {
-  PooledConnection::Vec::const_iterator it =
-      std::min_element(connections_.begin(), connections_.end(), least_busy_comp);
-  if (it == connections_.end() || (*it)->is_closing()) {
-    return PooledConnection::Ptr();
+PooledConnection::Ptr ConnectionPool::find_least_busy(int64_t token) const {
+  if (token == CASS_INT64_MIN) {
+    PooledConnection::Vec::const_iterator it =
+        std::min_element(connections_.begin(), connections_.end(), least_busy_comp);
+    if (it == connections_.end() || (*it)->is_closing()) {
+      return PooledConnection::Ptr();
+    }
+    return *it;
+  }
+
+  const auto desired_shard_num = host_->sharding_info()->shard_id(token);
+  const auto conn_it = std::find_if(connections_.begin(), connections_.end(), [desired_shard_num] (const PooledConnection::Ptr& c) {
+    return c->shard_id() == desired_shard_num;
+  });
+  if (conn_it != connections_.end()) {
+    return *conn_it;
   }
-  return *it;
+
+  return find_least_busy(CASS_INT64_MIN);
 }
 
 bool ConnectionPool::has_connections() const { return !connections_.empty(); }

src/connection_pool.hpp

@@ -135,7 +135,7 @@ class ConnectionPool : public RefCounted<ConnectionPool> {
    *
    * @return The least busy connection or null if no connection is available.
    */
-  PooledConnection::Ptr find_least_busy() const;
+  PooledConnection::Ptr find_least_busy(int64_t token) const;
 
   /**
    * Determine if the pool has any valid connections.

src/connection_pool_manager.cpp

@@ -62,12 +62,12 @@ ConnectionPoolManager::ConnectionPoolManager(const ConnectionPool::Map& pools, u
   }
 }
 
-PooledConnection::Ptr ConnectionPoolManager::find_least_busy(const Address& address) const {
+PooledConnection::Ptr ConnectionPoolManager::find_least_busy(const Address& address, int64_t token) const {
   ConnectionPool::Map::const_iterator it = pools_.find(address);
   if (it == pools_.end()) {
     return PooledConnection::Ptr();
   }
-  return it->second->find_least_busy();
+  return it->second->find_least_busy(token);
 }
 
 bool ConnectionPoolManager::has_connections(const Address& address) const {

src/connection_pool_manager.hpp

@@ -84,7 +84,14 @@ class ConnectionPoolManager
    * @return The least busy connection for a host or null if no connections are
    * available.
    */
-  PooledConnection::Ptr find_least_busy(const Address& address) const;
+  PooledConnection::Ptr find_least_busy(const Address& address, int64_t token) const;
+
+  /**
+   * Non-token-aware version of `find_least_busy()`, kept only for testing.
+   */
+  PooledConnection::Ptr find_least_busy(const Address& address) const {
+    return find_least_busy(address, CASS_INT64_MIN);
+  }
 
   /**
    * Determine if a pool has any valid connections.

src/pooled_connection.hpp

@@ -79,6 +79,8 @@ class PooledConnection
    */
   bool is_closing() const;
 
+  int32_t shard_id() const { return connection_->shard_id(); }
+
 public:
   const String& keyspace() const { return connection_->keyspace(); } // Test only
 

src/request_handler.cpp

@@ -29,6 +29,7 @@
 #include "result_response.hpp"
 #include "row.hpp"
 #include "session.hpp"
+#include "token_map_impl.hpp"
 
 #include <uv.h>
 
@@ -357,8 +358,16 @@ void RequestHandler::internal_retry(RequestExecution* request_execution) {
 
   bool is_done = false;
   while (!is_done && request_execution->current_host()) {
+    int64_t token = CASS_INT64_MIN;
+    const RoutableRequest* routable_req = dynamic_cast<const RoutableRequest*>(request());
+    if (routable_req) {
+      String routing_key;
+      routable_req->get_routing_key(&routing_key);
+      token = Murmur3Partitioner::hash(routing_key);
+    }
+
     PooledConnection::Ptr connection =
-        manager_->find_least_busy(request_execution->current_host()->address());
+        manager_->find_least_busy(request_execution->current_host()->address(), token);
     if (connection) {
       int32_t result = connection->write(request_execution);
 

src/request_processor.cpp

@@ -378,7 +378,7 @@ bool RequestProcessor::on_prepare_all(const RequestHandler::Ptr& request_handler
     PrepareAllCallback::Ptr prepare_all_callback(
         new PrepareAllCallback(address, prepare_all_handler));
 
-    PooledConnection::Ptr connection(connection_pool_manager_->find_least_busy(address));
+    PooledConnection::Ptr connection(connection_pool_manager_->find_least_busy(address, CASS_INT64_MIN));
     if (connection) {
       connection->write(prepare_all_callback.get());
     }
@@ -580,7 +580,7 @@ bool RequestProcessor::write_wait_callback(const RequestHandler::Ptr& request_ha
                                            const Host::Ptr& current_host,
                                            const RequestCallback::Ptr& callback) {
   PooledConnection::Ptr connection(
-      connection_pool_manager_->find_least_busy(current_host->address()));
+      connection_pool_manager_->find_least_busy(current_host->address(), CASS_INT64_MIN));
   if (connection && connection->write(callback.get()) > 0) {
     // Stop the original request timer now that we have a response and
     // are waiting for the maximum wait time of the handler.