#1837: Fixed TSAN issues in C++ driver in PartitionAwarePolicy.

Summary:
Reworked updating of partitions in the metadata object for PartitionAwarePolicy:
 - the partitions updating timer (PeriodicTask) was moved from PartitionAwarePolicy into Session object. (Using Session's event loop.)
 - Instead of internal `ControlMultipleRequestCallback::execute_query` now using `Session::execute ` async API.

Test Plan: ybd tsan --cxx-test integration-tests_cassandra_cpp_driver-test --gtest_filter CppCassandraDriverTest.BatchWriteDuringSoftMemoryLimit

Reviewers: mikhail

Reviewed By: mikhail

Subscribers: yql

Differential Revision: https://phabricator.dev.yugabyte.com/D9241

Author: OlegLoginov

src/config.hpp

@@ -298,6 +298,10 @@ class Config {
 
   bool partition_aware_routing() const { return partition_aware_routing_; }
 
+  unsigned partition_refresh_frequency_secs() const {
+    return partition_refresh_frequency_secs_;
+  }
+
   void set_partition_aware_routing(bool is_partition_aware,
                                    unsigned refresh_frequency_secs = CASS_DEFAULT_METADATA_REFRESH_FREQUENCY_SECS) {
     partition_aware_routing_ = is_partition_aware;

src/control_connection.cpp

@@ -509,7 +509,8 @@ void ControlConnection::on_query_hosts(ControlConnection* control_connection,
   session->purge_hosts(is_initial_connection);
 
   if (control_connection->use_schema_ ||
-      control_connection->token_aware_routing_) {
+      control_connection->token_aware_routing_ ||
+      control_connection->partition_aware_routing_) {
     control_connection->query_meta_schema();
   } else if (is_initial_connection) {
     control_connection->state_ = CONTROL_STATE_READY;
@@ -520,11 +521,8 @@ void ControlConnection::on_query_hosts(ControlConnection* control_connection,
   }
 }
 
-void ControlConnection::refresh_partitions() {
-  SharedRefPtr<ControlMultipleRequestCallback<UnusedData> > callback(
-        new ControlMultipleRequestCallback<UnusedData>(this, ControlConnection::on_query_meta_schema, UnusedData()));
-
-  callback->execute_query("partitions", YB_SELECT_PARTITIONS);
+std::string ControlConnection::get_yb_select_partitions_statement() {
+  return std::string(YB_SELECT_PARTITIONS);
 }
 
 //TODO: query and callbacks should be in Metadata

src/control_connection.hpp

@@ -46,6 +46,8 @@ class ControlConnection : public Connection::Listener {
                                               const Value* rpc_value,
                                               Address* output);
 
+  static std::string get_yb_select_partitions_statement();
+
   enum State {
     CONTROL_STATE_NEW,
     CONTROL_STATE_READY,
@@ -73,8 +75,6 @@ class ControlConnection : public Connection::Listener {
   void on_up(const Address& address);
   void on_down(const Address& address);
 
-  void refresh_partitions();
-
 private:
   template<class T>
   class ControlMultipleRequestCallback : public MultipleRequestCallback {

src/metadata.cpp

@@ -890,14 +890,9 @@ void Metadata::update_aggregates(int protocol_version, const VersionNumber& cass
 }
 
 void Metadata::update_partitions(int protocol_version, const VersionNumber& cassandra_version, ResultResponse* result) {
+  ScopedMutex l(&mutex_);
   schema_snapshot_version_++;
-
-  if (is_front_buffer()) {
-    ScopedMutex l(&mutex_);
-    updating_->update_partitions(protocol_version, cassandra_version, cache_, result);
-  } else {
-    updating_->update_partitions(protocol_version, cassandra_version, cache_, result);
-  }
+  updating_->update_partitions(protocol_version, cassandra_version, cache_, result);
 }
 
 void Metadata::drop_keyspace(const std::string& keyspace_name) {

src/partition_aware_policy.cpp

@@ -179,25 +179,6 @@ void PartitionAwarePolicy::init(const Host::Ptr& connected_host,
   ChainedLoadBalancingPolicy::init(connected_host, hosts, random);
 }
 
-void PartitionAwarePolicy::register_handles(uv_loop_t* loop) {
-  ChainedLoadBalancingPolicy::register_handles(loop);
-
-  refresh_metadata_task_ = PeriodicTask::start(loop,
-                                               refresh_frequency_secs_*1000,
-                                               this,
-                                               PartitionAwarePolicy::on_work,
-                                               PartitionAwarePolicy::on_after_work);
-}
-
-void PartitionAwarePolicy::close_handles() {
-  if (refresh_metadata_task_) {
-    PeriodicTask::stop(refresh_metadata_task_);
-    refresh_metadata_task_.reset();
-  }
-
-  ChainedLoadBalancingPolicy::close_handles();
-}
-
 QueryPlan* PartitionAwarePolicy::new_query_plan(const std::string& keyspace,
                                                 RequestHandler* request_handler) {
   QueryPlan* const child_plan = child_policy_->new_query_plan(keyspace, request_handler);
@@ -315,15 +296,4 @@ Host::Ptr PartitionAwarePolicy::PartitionAwareQueryPlan::compute_next() {
   return Host::Ptr();
 }
 
-void PartitionAwarePolicy::on_work(PeriodicTask* task) {
-  PartitionAwarePolicy* const policy = static_cast<PartitionAwarePolicy*>(task->data());
-  if (policy->control_connection() != NULL) {
-    policy->control_connection()->refresh_partitions();
-  }
-}
-
-void PartitionAwarePolicy::on_after_work(PeriodicTask*) {
-  // No-op.
-}
-
 } // namespace cass

src/partition_aware_policy.hpp

@@ -29,9 +29,6 @@ class CASS_EXPORT PartitionAwarePolicy: public ChainedLoadBalancingPolicy {
 
   virtual void init(const Host::Ptr& connected_host, const HostMap& hosts, Random* random);
 
-  virtual void register_handles(uv_loop_t* loop);
-  virtual void close_handles();
-
   virtual QueryPlan* new_query_plan(const std::string& keyspace,
                                     RequestHandler* request_handler);
 
@@ -74,10 +71,6 @@ class CASS_EXPORT PartitionAwarePolicy: public ChainedLoadBalancingPolicy {
     size_t remaining_;
   };
 
-  static void on_work(PeriodicTask* task);
-  static void on_after_work(PeriodicTask*);
-
-  PeriodicTask::Ptr refresh_metadata_task_;
   CopyOnWriteHostVec hosts_;
   int index_;
   unsigned refresh_frequency_secs_;

src/session.cpp

@@ -20,9 +20,11 @@
 #include "cluster.hpp"
 #include "config.hpp"
 #include "constants.hpp"
+#include "error_response.hpp"
 #include "execute_request.hpp"
 #include "logger.hpp"
 #include "prepare_request.hpp"
+#include "query_request.hpp"
 #include "scoped_lock.hpp"
 #include "statement.hpp"
 #include "timer.hpp"
@@ -158,13 +160,15 @@ Session::Session()
   uv_mutex_init(&state_mutex_);
   uv_mutex_init(&hosts_mutex_);
   uv_mutex_init(&keyspace_mutex_);
+  uv_mutex_init(&refresh_metadata_future_mutex_);
 }
 
 Session::~Session() {
   join();
   uv_mutex_destroy(&state_mutex_);
   uv_mutex_destroy(&hosts_mutex_);
   uv_mutex_destroy(&keyspace_mutex_);
+  uv_mutex_destroy(&refresh_metadata_future_mutex_);
 }
 
 void Session::clear(const Config& config) {
@@ -173,6 +177,10 @@ void Session::clear(const Config& config) {
   metrics_.reset(new Metrics(config_.thread_count_io() + 1));
   connect_future_.reset();
   close_future_.reset();
+  {
+    ScopedMutex lock_future(&refresh_metadata_future_mutex_);
+    refresh_metadata_future_.reset();
+  }
   { // Lock hosts
     ScopedMutex l(&hosts_mutex_);
     hosts_.clear();
@@ -416,6 +424,76 @@ void Session::notify_connected() {
     connect_future_->set();
     connect_future_.reset();
   }
+  {
+    ScopedMutex lock_future(&refresh_metadata_future_mutex_);
+    if (refresh_metadata_future_) {
+      refresh_metadata_future_->set();
+      refresh_metadata_future_.reset();
+    }
+  }
+
+  if (config().partition_aware_routing()) {
+    refresh_metadata_task_ = PeriodicTask::start(loop(),
+                                                 config().partition_refresh_frequency_secs()*1000,
+                                                 this,
+                                                 Session::on_refresh_metadata,
+                                                 Session::on_after_refresh_metadata);
+  }
+}
+
+void Session::on_refresh_metadata(PeriodicTask* task) {
+  Session* const session = static_cast<Session*>(task->data());
+  {
+    ScopedMutex l(&session->state_mutex_);
+    if (session->state_.load(MEMORY_ORDER_RELAXED) != SESSION_STATE_CONNECTED) {
+      return; // The session is finished.
+    }
+  }
+
+  ScopedMutex lock_future(&session->refresh_metadata_future_mutex_);
+  if (!session->refresh_metadata_future_) {
+    session->refresh_metadata_future_.reset(new ResponseFuture());
+    session->refresh_metadata_future_->set_callback(&Session::refresh_metadata_callback, session);
+
+    cass::QueryRequest* const query_request =
+        new cass::QueryRequest(ControlConnection::get_yb_select_partitions_statement(), 0);
+    RequestHandler::Ptr request_handler(new RequestHandler(
+        QueryRequest::ConstPtr(query_request), session->refresh_metadata_future_, session));
+    session->execute(request_handler);
+  }
+}
+
+void Session::on_after_refresh_metadata(PeriodicTask* task) {
+  // No-op.
+}
+
+void Session::refresh_metadata_callback(CassFuture* future, void* data) {
+  ResponseFuture::Ptr rf(static_cast<cass::ResponseFuture*>(future->from()));
+  Session* const session = static_cast<Session*>(data);
+
+  ScopedMutex lock_future(&session->refresh_metadata_future_mutex_);
+  session->refresh_metadata_future_.reset();
+
+  {
+    ScopedMutex l(&session->state_mutex_);
+    if (session->state_.load(MEMORY_ORDER_RELAXED) != SESSION_STATE_CONNECTED) {
+      return; // The session is finished.
+    }
+  }
+
+  if (!rf->ready()) {
+    return; // ResponseFuture is not ready.
+  }
+
+  cass::Response::Ptr response(rf->response());
+  if (!response || check_error_or_invalid_response("Session", CQL_OPCODE_RESULT, response.get())) {
+    return; // Error or null response.
+  }
+
+  const int protocol_version = session->control_connection_.protocol_version();
+  const VersionNumber& cassandra_version = session->control_connection_.cassandra_version();
+  ResultResponse* const partitions_result = static_cast<ResultResponse*>(response.get());
+  session->metadata().update_partitions(protocol_version, cassandra_version, partitions_result);
 }
 
 void Session::notify_connect_error(CassError code, const std::string& message) {
@@ -445,12 +523,23 @@ void Session::notify_closed() {
     close_future_->set();
     close_future_.reset();
   }
+
+  ScopedMutex lock_future(&refresh_metadata_future_mutex_);
+  if (refresh_metadata_future_) {
+    refresh_metadata_future_->set();
+    refresh_metadata_future_.reset();
+  }
 }
 
 void Session::close_handles() {
   EventThread<SessionEvent>::close_handles();
   request_queue_->close_handles();
   config_.load_balancing_policy()->close_handles();
+
+  if (refresh_metadata_task_) {
+    PeriodicTask::stop(refresh_metadata_task_);
+    refresh_metadata_task_.reset();
+  }
 }
 
 void Session::on_run() {

src/session.hpp

@@ -216,6 +216,11 @@ class Session : public EventThread<SessionEvent>, public RequestListener {
                                           const std::string& result_metadata_id,
                                           const ResultResponse::ConstPtr& result_response);
 
+  static void on_refresh_metadata(PeriodicTask* task);
+  static void on_after_refresh_metadata(PeriodicTask* task);
+
+  static void refresh_metadata_callback(CassFuture* future, void* data);
+
 private:
   typedef std::vector<IOWorker::Ptr > IOWorkerVec;
 
@@ -229,6 +234,10 @@ class Session : public EventThread<SessionEvent>, public RequestListener {
   Future::Ptr connect_future_;
   Future::Ptr close_future_;
 
+  PeriodicTask::Ptr refresh_metadata_task_;
+  uv_mutex_t refresh_metadata_future_mutex_;
+  ResponseFuture::Ptr refresh_metadata_future_;
+
   HostMap hosts_;
   uv_mutex_t hosts_mutex_;