Replace unique_ptr to db::ConnectionContextBase object with a shared pointer in ConnectOperation and Connection classes.

Summary: `ConnectionOperation::connection_context_` field accumulates some logging information while the connection operation is in flighjt. This diff preserves `connection_context_` pointers in operation and connection objects so the connection context can be accessed through any of them at the end of the operation.

Reviewed By: jkedgar

Differential Revision: D33663744

fbshipit-source-id: b35ba2dfb40c602b6a6815b988b8141278558c88

Author: maxgeorg

squangle/mysql_client/AsyncConnectionPool.cpp

@@ -296,16 +296,12 @@ void AsyncConnectionPool::openNewConnection(
   // Sanity check
   DCHECK(pool_op != nullptr);
   conn_storage_.queueOperation(poolKey, pool_op);
-  // Copy the ConnectionContext from the incoming operation. These contexts
+  // Propagate the ConnectionContext from the incoming operation. These contexts
   // contain application specific logging that will be lost if not passed to
   // the new ConnectOperation that is spawned to fulfill the pool miss.
-  // Creating a copy ensures that both operations are logged with the
-  // expected additional logging
-  std::unique_ptr<db::ConnectionContextBase> context;
-  if (pool_op->connection_context_) {
-    context = pool_op->connection_context_->copy();
-  }
-  tryRequestNewConnection(poolKey, std::move(context));
+  // Propagating the context pointer ensures that both operations are logged
+  // with the expected additional logging
+  tryRequestNewConnection(poolKey, pool_op->connection_context_);
 }
 
 void AsyncConnectionPool::reuseConnWithChangeUser(
@@ -562,7 +558,7 @@ void AsyncConnectionPool::connectionSpotFreed(const PoolKey& pool_key) {
 
 void AsyncConnectionPool::tryRequestNewConnection(
     const PoolKey& pool_key,
-    std::unique_ptr<db::ConnectionContextBase> context) {
+    std::shared_ptr<db::ConnectionContextBase> context) {
   // Only called internally, this doesn't need to check if it's shutting
   // down
   DCHECK_EQ(std::this_thread::get_id(), mysql_client_->threadId());
@@ -581,7 +577,7 @@ void AsyncConnectionPool::tryRequestNewConnection(
     auto connOp = mysql_client_->beginConnection(pool_key.connKey);
     connOp->setConnectionOptions(pool_key.connOptions);
     if (!context) {
-      context = std::make_unique<db::ConnectionContextBase>();
+      context = std::make_shared<db::ConnectionContextBase>();
     }
     connOp->setConnectionContext(std::move(context));
     auto pool_ptr = getSelfWeakPointer();

squangle/mysql_client/AsyncConnectionPool.h

@@ -569,7 +569,7 @@ class AsyncConnectionPool
   // these connections.
   void tryRequestNewConnection(
       const PoolKey& pool_key,
-      std::unique_ptr<db::ConnectionContextBase> context = nullptr);
+      std::shared_ptr<db::ConnectionContextBase> context = nullptr);
 
   // Used for when we fail to open a requested connection. In case of mysql
   // failure (e.g. bad password) we propagate the error to all queued

squangle/mysql_client/Connection.h

@@ -531,7 +531,7 @@ class Connection {
     }
   }
 
-  void setConnectionContext(std::unique_ptr<db::ConnectionContextBase>&& e) {
+  void setConnectionContext(std::shared_ptr<db::ConnectionContextBase> e) {
     connection_context_ = std::move(e);
   }
 
@@ -543,7 +543,7 @@ class Connection {
   bool killOnQueryTimeout_ = false;
 
   // Context information for logging purposes.
-  std::unique_ptr<db::ConnectionContextBase> connection_context_;
+  std::shared_ptr<db::ConnectionContextBase> connection_context_;
 
   // Unowned pointer to the client we're from.
   MysqlClientBase* mysql_client_;

squangle/mysql_client/MysqlConnectionHolder.cpp

@@ -31,16 +31,13 @@ MysqlConnectionHolder::MysqlConnectionHolder(
     const ConnectionKey* connKey)
     : client_(from_holder->client_),
       conn_key_((connKey) ? *connKey : from_holder->conn_key_),
-      conn_context_(nullptr),
+      conn_context_(from_holder->conn_context_),
       creation_time_(from_holder->creation_time_),
       last_activity_time_(from_holder->last_activity_time_),
       connection_opened_(from_holder->connection_opened_),
       can_reuse_(from_holder->can_reuse_) {
   mysql_ = from_holder->stealMysql();
   client_->activeConnectionAdded(&conn_key_);
-  if (from_holder->conn_context_) {
-    conn_context_ = from_holder->conn_context_->copy();
-  }
 }
 
 bool MysqlConnectionHolder::inTransaction() {

squangle/mysql_client/MysqlConnectionHolder.h

@@ -94,7 +94,7 @@ class MysqlConnectionHolder {
   }
 
   void setConnectionContext(
-      std::unique_ptr<db::ConnectionContextBase> conn_context) {
+      std::shared_ptr<db::ConnectionContextBase> conn_context) {
     conn_context_ = std::move(conn_context);
   }
 
@@ -129,7 +129,7 @@ class MysqlConnectionHolder {
   // notification of completed operations.
   MYSQL* mysql_;
   const ConnectionKey conn_key_;
-  std::unique_ptr<db::ConnectionContextBase> conn_context_;
+  std::shared_ptr<db::ConnectionContextBase> conn_context_;
   Timepoint creation_time_;
   Timepoint last_activity_time_;
   bool connection_opened_ = false;

squangle/mysql_client/Operation.cpp

@@ -733,10 +733,7 @@ void ConnectOperation::socketActionable() {
         }
       }
       conn()->socketHandler()->changeHandlerFD(socket);
-      if (connection_context_) {
-        conn()->mysqlConnection()->setConnectionContext(
-            connection_context_->copy());
-      }
+      conn()->mysqlConnection()->setConnectionContext(connection_context_);
       conn()->mysqlConnection()->connectionOpened();
       attemptSucceeded(OperationResult::Succeeded);
     } else {
@@ -883,7 +880,7 @@ void ConnectOperation::specializedCompleteOperation() {
 
   conn()->setConnectionOptions(conn_options_);
   conn()->setKillOnQueryTimeout(getKillOnQueryTimeout());
-  conn()->setConnectionContext(std::move(connection_context_));
+  conn()->setConnectionContext(connection_context_);
 
   conn()->notify();
 

squangle/mysql_client/Operation.h

@@ -793,7 +793,7 @@ class Operation : public std::enable_shared_from_this<Operation> {
 
   folly::Optional<folly::dynamic> user_data_;
   ObserverCallback observer_callback_;
-  std::unique_ptr<db::ConnectionContextBase> connection_context_;
+  std::shared_ptr<db::ConnectionContextBase> connection_context_;
 
   MysqlClientBase* mysql_client_;
 
@@ -862,7 +862,7 @@ class ConnectOperation : public Operation {
   // operation will create.
   ConnectOperation* setDefaultQueryTimeout(Duration t);
   ConnectOperation* setConnectionContext(
-      std::unique_ptr<db::ConnectionContextBase>&& e) {
+      std::shared_ptr<db::ConnectionContextBase> e) {
     CHECK_THROW(
         state_ == OperationState::Unstarted, db::OperationStateException);
     connection_context_ = std::move(e);
@@ -1022,7 +1022,7 @@ class ConnectOperation : public Operation {
   ConnectionOptions conn_options_;
 
   // Context information for logging purposes.
-  std::unique_ptr<db::ConnectionContextBase> connection_context_;
+  std::shared_ptr<db::ConnectionContextBase> connection_context_;
 
  private:
   void specializedRunImpl();