Merge pull request #2 from jul-stas/shard-awareness

Shard awareness

Author: haaawk

CMakeLists.txt

@@ -1,4 +1,4 @@
-cmake_minimum_required(VERSION 2.8.12)
+cmake_minimum_required(VERSION 3.1)
 project(cassandra C CXX)
 
 set(CASS_ROOT_DIR ${CMAKE_CURRENT_SOURCE_DIR})
@@ -43,10 +43,12 @@ option(CASS_USE_KERBEROS "Use Kerberos" OFF)
 option(CASS_USE_LIBSSH2 "Use libssh2 for integration tests" OFF)
 option(CASS_USE_OPENSSL "Use OpenSSL" ON)
 option(CASS_USE_STATIC_LIBS "Link static libraries when building executables" OFF)
-option(CASS_USE_STD_ATOMIC "Use C++11 atomics library" OFF)
+option(CASS_USE_STD_ATOMIC "Use std::atomic library" ON)
 option(CASS_USE_ZLIB "Use zlib" ON)
 option(CASS_USE_TIMERFD "Use timerfd (Linux only)" ON)
 
+set(CASS_CPP_STANDARD "11" CACHE STRING "C++ standard (11, 14, 17, etc.)")
+
 # Handle testing dependencies
 if(CASS_BUILD_TESTS)
   # Enable integration and unit tests
@@ -161,12 +163,10 @@ endif()
 # Top-level compiler flags
 #------------------------
 
+set (CMAKE_CXX_STANDARD ${CASS_CPP_STANDARD})
+
 if("${CMAKE_CXX_COMPILER_ID}" STREQUAL "Clang" OR
    "${CMAKE_CXX_COMPILER_ID}" STREQUAL "GNU")
-  # Enable C++11 support to use std::atomic
-  if(CASS_USE_STD_ATOMIC)
-    set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11")
-  endif()
 
   # OpenSSL is deprecated on later versions of Mac OS X. The long-term solution
   # is to provide a CommonCryto implementation.

driver_config.hpp.in

@@ -4,6 +4,7 @@
 #cmakedefine HAVE_KERBEROS
 #cmakedefine HAVE_OPENSSL
 #cmakedefine HAVE_STD_ATOMIC
+#cmakedefine CASS_CPP_STANDARD @CASS_CPP_STANDARD@
 #cmakedefine HAVE_BOOST_ATOMIC
 #cmakedefine HAVE_NOSIGPIPE
 #cmakedefine HAVE_SIGTIMEDWAIT

LICENSE.txt renamed to licenses/apache-2.0.txt

@@ -0,0 +1,23 @@
+Boost Software License - Version 1.0 - August 17th, 2003
+
+Permission is hereby granted, free of charge, to any person or organization
+obtaining a copy of the software and accompanying documentation covered by
+this license (the "Software") to use, reproduce, display, distribute,
+execute, and transmit the Software, and to prepare derivative works of the
+Software, and to permit third-parties to whom the Software is furnished to
+do so, all subject to the following:
+
+The copyright notices in the Software and this entire statement, including
+the above license grant, this restriction and the following disclaimer,
+must be included in all copies of the Software, in whole or in part, and
+all derivative works of the Software, unless such copies or derivative
+works are solely in the form of machine-executable object code generated by
+a source language processor.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE, TITLE AND NON-INFRINGEMENT. IN NO EVENT
+SHALL THE COPYRIGHT HOLDERS OR ANYONE DISTRIBUTING THE SOFTWARE BE LIABLE
+FOR ANY DAMAGES OR OTHER LIABILITY, WHETHER IN CONTRACT, TORT OR OTHERWISE,
+ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+DEALINGS IN THE SOFTWARE.

licenses/boost-1.0.txt

@@ -36,6 +36,15 @@ else()
   endif()
 endif()
 
+# Determine `optional` library to include
+if(CMAKE_CXX_STANDARD LESS 17)
+  message(STATUS "Using akrzemi's `optional` implementation")
+  list(APPEND SOURCES optional/optional_akrzemi.hpp)
+else()
+  message(STATUS "Using std::optional library")
+  list(APPEND SOURCES optional/optional_std.hpp)
+endif()
+
 add_subdirectory(third_party/curl)
 add_subdirectory(third_party/hdr_histogram)
 add_subdirectory(third_party/http-parser)

src/CMakeLists.txt

@@ -25,6 +25,8 @@
 #include "speculative_execution.hpp"
 #include "utils.hpp"
 
+#include <iostream>
+
 using namespace datastax;
 using namespace datastax::internal::core;
 
@@ -235,6 +237,10 @@ Cluster::Cluster(const ControlConnection::Ptr& connection, ClusterListener* list
     , local_dc_(local_dc)
     , supported_options_(supported_options)
     , is_recording_events_(settings.disable_events_on_startup) {
+  static const auto optimized_msg = "===== Using optimized driver!!! =====\n";
+  std::cout << optimized_msg;
+  LOG_INFO(optimized_msg);
+
   inc_ref();
   connection_->set_listener(this);
 

src/cluster.cpp

@@ -14,13 +14,14 @@
   limitations under the License.
 */
 
+#ifndef DATASTAX_INTERNAL_CONNECTION_HPP
+#define DATASTAX_INTERNAL_CONNECTION_HPP
+
 #include "event_response.hpp"
 #include "request_callback.hpp"
 #include "socket.hpp"
 #include "stream_manager.hpp"
-
-#ifndef DATASTAX_INTERNAL_CONNECTION_HPP
-#define DATASTAX_INTERNAL_CONNECTION_HPP
+#include "optional.hpp"
 
 namespace datastax { namespace internal { namespace core {
 
@@ -195,6 +196,11 @@ class Connection : public RefCounted<Connection> {
    */
   void set_listener(ConnectionListener* listener = NULL);
 
+  int32_t shard_id() const { return shard_id_; }
+  void set_shard_id(int32_t shard_id) {
+    shard_id_ = shard_id;
+  }
+
   /**
    * Start heartbeats to keep the connection alive and to detect a network or
    * server-side failure.
@@ -241,6 +247,8 @@ class Connection : public RefCounted<Connection> {
   ProtocolVersion protocol_version_;
   String keyspace_;
 
+  int32_t shard_id_ = 0;
+
   unsigned int idle_timeout_secs_;
   unsigned int heartbeat_interval_secs_;
   bool heartbeat_outstanding_;

src/connection.hpp

@@ -22,6 +22,7 @@
 #include "utils.hpp"
 
 #include <algorithm>
+#include <numeric>
 
 using namespace datastax;
 using namespace datastax::internal::core;
@@ -77,34 +78,70 @@ ConnectionPool::ConnectionPool(const Connection::Vec& connections, ConnectionPoo
   set_pointer_keys(reconnection_schedules_);
   set_pointer_keys(to_flush_);
 
+  if (host->sharding_info()) {
+    const auto hosts_shard_cnt = host->sharding_info()->get_shards_count();
+    connections_by_shard_.resize(hosts_shard_cnt);
+    num_connections_per_shard_ = settings_.num_connections_per_host / hosts_shard_cnt
+        + (settings_.num_connections_per_host % hosts_shard_cnt ? 1u : 0u);
+  } else {
+    connections_by_shard_.resize(1);
+    num_connections_per_shard_ = settings_.num_connections_per_host;
+  }
+
   for (Connection::Vec::const_iterator it = connections.begin(), end = connections.end(); it != end;
        ++it) {
     const Connection::Ptr& connection(*it);
     if (!connection->is_closing()) {
-      add_connection(PooledConnection::Ptr(new PooledConnection(this, connection)));
+      if (connections_by_shard_[connection->shard_id()].size() < num_connections_per_shard_) {
+        add_connection(PooledConnection::Ptr(new PooledConnection(this, connection)));
+      } else {
+        connection->close();
+      }
     }
   }
 
   notify_up_or_down();
 
   // We had non-critical errors or some connections closed
-  assert(connections.size() <= settings_.num_connections_per_host);
-  size_t needed = settings_.num_connections_per_host - connections_.size();
+  size_t needed = num_connections_per_shard_ * connections_by_shard_.size()
+      - std::accumulate(connections_by_shard_.begin(), connections_by_shard_.end(), 0u,
+          [] (size_t acc, const PooledConnection::Vec& v) {
+            return acc + v.size();
+          });
   for (size_t i = 0; i < needed; ++i) {
     schedule_reconnect();
   }
 }
 
-PooledConnection::Ptr ConnectionPool::find_least_busy() const {
+PooledConnection::Ptr ConnectionPool::find_least_busy(int64_t token) const {
+  if (token == CASS_INT64_MIN || !host_->sharding_info()) {
+    // We got a placeholder token, or a sensible token that is useless without the sharding info.
+    // In both cases we return the least busy connection of the *entire pool* (or NULL).
+    PooledConnection::Ptr least_busy; // NULL by default
+    for (const auto& shard_pool : connections_by_shard_) {
+      for (const auto& conn : shard_pool) {
+        if (!conn->is_closing()) {
+          least_busy = least_busy ? std::min(least_busy, conn, least_busy_comp) : conn;
+        }
+      }
+    }
+    return least_busy;
+  }
+
+  // Otherwise, find the least busy connection pointing to the right shard (if possible)
+  const auto& shard_pool = connections_by_shard_[host_->sharding_info()->shard_id(token)];
   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();
+      std::min_element(shard_pool.begin(), shard_pool.end(), least_busy_comp);
+  if (it == shard_pool.end() || (*it)->is_closing()) {
+    return find_least_busy(CASS_INT64_MIN);
   }
   return *it;
 }
 
-bool ConnectionPool::has_connections() const { return !connections_.empty(); }
+bool ConnectionPool::has_connections() const {
+  return std::any_of(connections_by_shard_.begin(), connections_by_shard_.end(),
+      [] (const PooledConnection::Vec& v) { return !v.empty(); });
+}
 
 void ConnectionPool::flush() {
   for (DenseHashSet<PooledConnection*>::const_iterator it = to_flush_.begin(),
@@ -142,8 +179,9 @@ void ConnectionPool::close_connection(PooledConnection* connection, Protected) {
   if (metrics_) {
     metrics_->total_connections.dec();
   }
-  connections_.erase(std::remove(connections_.begin(), connections_.end(), connection),
-                     connections_.end());
+  auto& shard_pool = connections_by_shard_[connection->shard_id()];
+  shard_pool.erase(std::remove(shard_pool.begin(), shard_pool.end(), connection),
+                     shard_pool.end());
   to_flush_.erase(connection);
 
   if (close_state_ != CLOSE_STATE_OPEN) {
@@ -161,16 +199,16 @@ void ConnectionPool::add_connection(const PooledConnection::Ptr& connection) {
   if (metrics_) {
     metrics_->total_connections.inc();
   }
-  connections_.push_back(connection);
+  const size_t new_connections_shard = connection->shard_id();
+  LOG_INFO("add_connection: to host %s to shard %ld", host_->address_string().c_str(), new_connections_shard);
+  connections_by_shard_[new_connections_shard].push_back(connection);
 }
 
 void ConnectionPool::notify_up_or_down() {
-  if ((notify_state_ == NOTIFY_STATE_NEW || notify_state_ == NOTIFY_STATE_UP) &&
-      connections_.empty()) {
+  if ((notify_state_ == NOTIFY_STATE_NEW || notify_state_ == NOTIFY_STATE_UP) && !has_connections()) {
     notify_state_ = NOTIFY_STATE_DOWN;
     listener_->on_pool_down(host_->address());
-  } else if ((notify_state_ == NOTIFY_STATE_NEW || notify_state_ == NOTIFY_STATE_DOWN) &&
-             !connections_.empty()) {
+  } else if ((notify_state_ == NOTIFY_STATE_NEW || notify_state_ == NOTIFY_STATE_DOWN) && has_connections()) {
     notify_state_ = NOTIFY_STATE_UP;
     listener_->on_pool_up(host_->address());
   }
@@ -211,11 +249,12 @@ void ConnectionPool::internal_close() {
     // Make copies of connection/connector data structures to prevent iterator
     // invalidation.
 
-    PooledConnection::Vec connections(connections_);
-    for (PooledConnection::Vec::iterator it = connections.begin(), end = connections.end();
-         it != end; ++it) {
-      (*it)->close();
-    }
+    auto connections_per_shards = connections_by_shard_;
+    std::for_each(connections_per_shards.begin(), connections_per_shards.end(), [] (PooledConnection::Vec& v) {
+      for (auto& c : v) {
+        c->close();
+      }
+    });
 
     DelayedConnector::Vec pending_connections(pending_connections_);
     for (DelayedConnector::Vec::iterator it = pending_connections.begin(),
@@ -232,8 +271,7 @@ void ConnectionPool::internal_close() {
 void ConnectionPool::maybe_closed() {
   // Remove the pool once all current connections and pending connections
   // are terminated.
-  if (close_state_ == CLOSE_STATE_WAITING_FOR_CONNECTIONS && connections_.empty() &&
-      pending_connections_.empty()) {
+  if (close_state_ == CLOSE_STATE_WAITING_FOR_CONNECTIONS && !has_connections() && pending_connections_.empty()) {
     close_state_ = CLOSE_STATE_CLOSED;
     // Only mark DOWN if it's UP otherwise we might get multiple DOWN events
     // when connecting the pool.
@@ -263,9 +301,18 @@ void ConnectionPool::on_reconnect(DelayedConnector* connector) {
   }
 
   if (connector->is_ok()) {
-    add_connection(
-        PooledConnection::Ptr(new PooledConnection(this, connector->release_connection())));
-    notify_up_or_down();
+    PooledConnection::Ptr pooled_conn {new PooledConnection(this, connector->release_connection())};
+    const size_t new_connections_shard = pooled_conn->shard_id();
+    if (connections_by_shard_.size() > new_connections_shard
+        && connections_by_shard_[new_connections_shard].size() < num_connections_per_shard_) {
+      add_connection(pooled_conn);
+      notify_up_or_down();
+    } else {
+      LOG_INFO("Reconnection to host %s connected us to shard %ld, reconnecting again",
+               address().to_string().c_str(), new_connections_shard);
+      pooled_conn->close();
+      schedule_reconnect(schedule.release());
+    }
   } else if (!connector->is_canceled()) {
     if (connector->is_critical_error()) {
       LOG_ERROR("Closing established connection pool to host %s because of the following error: %s",

src/connection_pool.cpp

@@ -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.
@@ -213,6 +213,9 @@ class ConnectionPool : public RefCounted<ConnectionPool> {
 private:
   void notify_up_or_down();
   void notify_critical_error(Connector::ConnectionError code, const String& message);
+
+  /** Adds connection to the pool. It's the caller's responsibility
+   *  to keep track of the connections count. */
   void add_connection(const PooledConnection::Ptr& connection);
   void schedule_reconnect(ReconnectionSchedule* schedule = NULL);
   void internal_close();
@@ -232,7 +235,8 @@ class ConnectionPool : public RefCounted<ConnectionPool> {
 
   CloseState close_state_;
   NotifyState notify_state_;
-  PooledConnection::Vec connections_;
+  std::vector<PooledConnection::Vec> connections_by_shard_; /// Index is the shard ID
+  size_t num_connections_per_shard_;
   DelayedConnector::Vec pending_connections_;
   DenseHashSet<PooledConnection*> to_flush_;
 };

src/connection_pool.hpp

@@ -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 {