feat: split ConnectionImpl and SessionPool cleanly (#1022)

* feat: split `ConnectionImpl` and `SessionPool` cleanly

Remove the circular dependency between `ConnectionImpl` and
`SessionPool`. The latter can now make RPCs directly instead of
having to call back through the `SessionManager` interface.

The `SessionHolder` deleter also now calls directly into `SessionPool`
so we don't need to forward those calls from `ConnectionImpl`.

I also removed the tests where `BatchCreateSessions` returned
`RESOURCE_EXHAUSTED` since we learned from spanner team that will never
happen (since there are no longer server-side session limits), and the tests
conflicted with our retry policy.

Part of #307

Author: mr-salty

google/cloud/spanner/internal/connection_impl.cc

@@ -21,9 +21,7 @@
 #include "google/cloud/spanner/read_partition.h"
 #include "google/cloud/grpc_utils/grpc_error_delegate.h"
 #include "google/cloud/internal/make_unique.h"
-#include <google/spanner/v1/spanner.pb.h>
 #include <limits>
-#include <memory>
 
 namespace google {
 namespace cloud {
@@ -111,7 +109,8 @@ ConnectionImpl::ConnectionImpl(Database db, std::shared_ptr<SpannerStub> stub,
       stub_(std::move(stub)),
       retry_policy_(std::move(retry_policy)),
       backoff_policy_(std::move(backoff_policy)),
-      session_pool_(this) {}
+      session_pool_(std::make_shared<SessionPool>(
+          db_, stub_, retry_policy_->clone(), backoff_policy_->clone())) {}
 
 RowStream ConnectionImpl::Read(ReadParams params) {
   return internal::Visit(
@@ -286,7 +285,7 @@ RowStream ConnectionImpl::ReadImpl(SessionHolder& session,
                                    spanner_proto::TransactionSelector& s,
                                    ReadParams params) {
   if (!session) {
-    auto session_or = AllocateSession();
+    auto session_or = session_pool_->Allocate();
     if (!session_or) {
       return RowStream(
           google::cloud::internal::make_unique<StatusOnlyResultSetSource>(
@@ -348,7 +347,7 @@ StatusOr<std::vector<ReadPartition>> ConnectionImpl::PartitionReadImpl(
   if (!session) {
     // Since the session may be sent to other machines, it should not be
     // returned to the pool when the Transaction is destroyed.
-    auto session_or = AllocateSession(/*dissociate_from_pool=*/true);
+    auto session_or = session_pool_->Allocate(/*dissociate_from_pool=*/true);
     if (!session_or) {
       return std::move(session_or).status();
     }
@@ -401,7 +400,7 @@ StatusOr<ResultType> ConnectionImpl::ExecuteSqlImpl(
         google::spanner::v1 ::ExecuteSqlRequest& request)> const&
         retry_resume_fn) {
   if (!session) {
-    auto session_or = AllocateSession();
+    auto session_or = session_pool_->Allocate();
     if (!session_or) {
       return std::move(session_or).status();
     }
@@ -560,7 +559,7 @@ StatusOr<std::vector<QueryPartition>> ConnectionImpl::PartitionQueryImpl(
   if (!session) {
     // Since the session may be sent to other machines, it should not be
     // returned to the pool when the Transaction is destroyed.
-    auto session_or = AllocateSession(/*dissociate_from_pool=*/true);
+    auto session_or = session_pool_->Allocate(/*dissociate_from_pool=*/true);
     if (!session_or) {
       return std::move(session_or).status();
     }
@@ -606,7 +605,7 @@ StatusOr<BatchDmlResult> ConnectionImpl::ExecuteBatchDmlImpl(
     SessionHolder& session, spanner_proto::TransactionSelector& s,
     std::int64_t seqno, ExecuteBatchDmlParams params) {
   if (!session) {
-    auto session_or = AllocateSession();
+    auto session_or = session_pool_->Allocate();
     if (!session_or) {
       return std::move(session_or).status();
     }
@@ -649,7 +648,7 @@ StatusOr<PartitionedDmlResult> ConnectionImpl::ExecutePartitionedDmlImpl(
     SessionHolder& session, spanner_proto::TransactionSelector& s,
     std::int64_t seqno, ExecutePartitionedDmlParams params) {
   if (!session) {
-    auto session_or = AllocateSession();
+    auto session_or = session_pool_->Allocate();
     if (!session_or) {
       return std::move(session_or).status();
     }
@@ -703,7 +702,7 @@ StatusOr<CommitResult> ConnectionImpl::CommitImpl(
     SessionHolder& session, spanner_proto::TransactionSelector& s,
     CommitParams params) {
   if (!session) {
-    auto session_or = AllocateSession();
+    auto session_or = session_pool_->Allocate();
     if (!session_or) {
       return std::move(session_or).status();
     }
@@ -743,7 +742,7 @@ StatusOr<CommitResult> ConnectionImpl::CommitImpl(
 Status ConnectionImpl::RollbackImpl(SessionHolder& session,
                                     spanner_proto::TransactionSelector& s) {
   if (!session) {
-    auto session_or = AllocateSession();
+    auto session_or = session_pool_->Allocate();
     if (!session_or) {
       return std::move(session_or).status();
     }
@@ -771,55 +770,6 @@ Status ConnectionImpl::RollbackImpl(SessionHolder& session,
       request, __func__);
 }
 
-StatusOr<SessionHolder> ConnectionImpl::AllocateSession(
-    bool dissociate_from_pool) {
-  auto session = session_pool_.Allocate(dissociate_from_pool);
-  if (!session.ok()) {
-    return std::move(session).status();
-  }
-
-  if (dissociate_from_pool) {
-    // Uses the default deleter; the Session is not returned to the pool.
-    return {*std::move(session)};
-  }
-
-  std::weak_ptr<ConnectionImpl> connection = shared_from_this();
-  return SessionHolder(session->release(), [connection](Session* session) {
-    auto shared_connection = connection.lock();
-    // If `connection` is still alive, release the `Session` to its pool;
-    // otherwise just delete the `Session`.
-    if (shared_connection) {
-      shared_connection->ReleaseSession(session);
-    } else {
-      delete session;
-    }
-  });
-}
-
-StatusOr<std::vector<std::unique_ptr<Session>>> ConnectionImpl::CreateSessions(
-    int num_sessions) {
-  spanner_proto::BatchCreateSessionsRequest request;
-  request.set_database(db_.FullName());
-  request.set_session_count(std::int32_t{num_sessions});
-  auto response = RetryLoop(
-      retry_policy_->clone(), backoff_policy_->clone(), true,
-      [this](grpc::ClientContext& context,
-             spanner_proto::BatchCreateSessionsRequest const& request) {
-        return stub_->BatchCreateSessions(context, request);
-      },
-      request, __func__);
-  if (!response) {
-    return response.status();
-  }
-  std::vector<std::unique_ptr<Session>> sessions;
-  sessions.reserve(response->session_size());
-  for (auto& session : *response->mutable_session()) {
-    sessions.push_back(google::cloud::internal::make_unique<Session>(
-        std::move(*session.mutable_name())));
-  }
-  return {std::move(sessions)};
-}
-
 }  // namespace internal
 }  // namespace SPANNER_CLIENT_NS
 }  // namespace spanner

google/cloud/spanner/internal/connection_impl.h

@@ -25,6 +25,7 @@
 #include "google/cloud/spanner/version.h"
 #include "google/cloud/status.h"
 #include "google/cloud/status_or.h"
+#include <google/spanner/v1/spanner.pb.h>
 #include <cstdint>
 #include <memory>
 #include <mutex>
@@ -46,9 +47,6 @@ std::unique_ptr<BackoffPolicy> DefaultConnectionBackoffPolicy();
 /**
  * Factory method to construct a `ConnectionImpl`.
  *
- * @note `ConnectionImpl` relies on `std::enable_shared_from_this`; the
- * factory method ensures it can only be constructed using a `std::shared_ptr`
- *
  * @note In tests we can use a mock stub and custom (or mock) policies.
  */
 class ConnectionImpl;
@@ -63,9 +61,7 @@ std::shared_ptr<ConnectionImpl> MakeConnection(
  * Spanner instance. See `MakeConnection()` for a factory function that creates
  * and returns instances of this class.
  */
-class ConnectionImpl : public Connection,
-                       public SessionManager,
-                       public std::enable_shared_from_this<ConnectionImpl> {
+class ConnectionImpl : public Connection {
  public:
   RowStream Read(ReadParams) override;
   StatusOr<std::vector<ReadPartition>> PartitionRead(
@@ -139,29 +135,6 @@ class ConnectionImpl : public Connection,
   Status RollbackImpl(SessionHolder& session,
                       google::spanner::v1::TransactionSelector& s);
 
-  /**
-   * Get a session from the pool, or create one if the pool is empty.
-   * @returns an error if session creation fails; always returns a valid
-   * `SessionHolder` (never `nullptr`) on success.
-   *
-   * The `SessionHolder` usually returns the session to the pool when it is
-   * destroyed. However, if `dissociate_from_pool` is true the session will not
-   * be returned to the session pool. This is used in partitioned operations,
-   * since we don't know when all parties are done using the session.
-   */
-  StatusOr<SessionHolder> AllocateSession(bool dissociate_from_pool = false);
-
-  // Forwards calls for the `SessionPool`; used in the `SessionHolder` deleter
-  // so it can hold a `weak_ptr` to `ConnectionImpl` (it's already reference
-  // counted, and manages the lifetime of `SessionPool`).
-  void ReleaseSession(Session* session) {
-    session_pool_.Release(std::unique_ptr<Session>(session));
-  }
-
-  // `SessionManager` methods; used by the `SessionPool`
-  StatusOr<std::vector<std::unique_ptr<Session>>> CreateSessions(
-      int num_sessions) override;
-
   template <typename ResultType>
   StatusOr<ResultType> ExecuteSqlImpl(
       SessionHolder& session, google::spanner::v1::TransactionSelector& s,
@@ -176,7 +149,6 @@ class ConnectionImpl : public Connection,
       SessionHolder& session, google::spanner::v1::TransactionSelector& s,
       std::int64_t seqno, SqlParams params,
       google::spanner::v1::ExecuteSqlRequest::QueryMode query_mode);
-
   template <typename ResultType>
   StatusOr<ResultType> CommonDmlImpl(
       SessionHolder& session, google::spanner::v1::TransactionSelector& s,
@@ -187,7 +159,7 @@ class ConnectionImpl : public Connection,
   std::shared_ptr<SpannerStub> stub_;
   std::shared_ptr<RetryPolicy> retry_policy_;
   std::shared_ptr<BackoffPolicy> backoff_policy_;
-  SessionPool session_pool_;
+  std::shared_ptr<SessionPool> session_pool_;
 };
 
 }  // namespace internal

google/cloud/spanner/internal/session_pool.cc

@@ -14,19 +14,29 @@
 
 #include "google/cloud/spanner/internal/session_pool.h"
 #include "google/cloud/spanner/internal/connection_impl.h"
+#include "google/cloud/spanner/internal/retry_loop.h"
 #include "google/cloud/spanner/internal/session.h"
+#include "google/cloud/internal/make_unique.h"
 #include "google/cloud/status.h"
 #include <algorithm>
-#include <memory>
 
 namespace google {
 namespace cloud {
 namespace spanner {
 inline namespace SPANNER_CLIENT_NS {
 namespace internal {
 
-SessionPool::SessionPool(SessionManager* manager, SessionPoolOptions options)
-    : manager_(manager), options_(options) {
+namespace spanner_proto = ::google::spanner::v1;
+
+SessionPool::SessionPool(Database db, std::shared_ptr<SpannerStub> stub,
+                         std::unique_ptr<RetryPolicy> retry_policy,
+                         std::unique_ptr<BackoffPolicy> backoff_policy,
+                         SessionPoolOptions options)
+    : db_(std::move(db)),
+      stub_(std::move(stub)),
+      retry_policy_(std::move(retry_policy)),
+      backoff_policy_(std::move(backoff_policy)),
+      options_(options) {
   // Ensure the options have sensible values.
   options_.min_sessions = (std::max)(options_.min_sessions, 0);
   options_.max_sessions = (std::max)(options_.max_sessions, 1);
@@ -43,7 +53,7 @@ SessionPool::SessionPool(SessionManager* manager, SessionPoolOptions options)
   if (options_.min_sessions == 0) {
     return;
   }
-  auto sessions = manager_->CreateSessions(options_.min_sessions);
+  auto sessions = CreateSessions(options_.min_sessions);
   if (sessions.ok()) {
     std::unique_lock<std::mutex> lk(mu_);
     total_sessions_ += static_cast<int>(sessions->size());
@@ -53,8 +63,7 @@ SessionPool::SessionPool(SessionManager* manager, SessionPoolOptions options)
   }
 }
 
-StatusOr<std::unique_ptr<Session>> SessionPool::Allocate(
-    bool dissociate_from_pool) {
+StatusOr<SessionHolder> SessionPool::Allocate(bool dissociate_from_pool) {
   std::unique_lock<std::mutex> lk(mu_);
   for (;;) {
     if (!sessions_.empty()) {
@@ -64,7 +73,7 @@ StatusOr<std::unique_ptr<Session>> SessionPool::Allocate(
       if (dissociate_from_pool) {
         --total_sessions_;
       }
-      return {std::move(session)};
+      return {MakeSessionHolder(std::move(session), dissociate_from_pool)};
     }
 
     // If the pool is at its max size, fail or wait until someone returns a
@@ -98,7 +107,7 @@ StatusOr<std::unique_ptr<Session>> SessionPool::Allocate(
     create_in_progress_ = true;
     lk.unlock();
     // TODO(#307) do we need to limit the call rate here?
-    auto sessions = manager_->CreateSessions(sessions_to_create);
+    auto sessions = CreateSessions(sessions_to_create);
     lk.lock();
     create_in_progress_ = false;
     if (!sessions.ok() || sessions->empty()) {
@@ -128,21 +137,64 @@ StatusOr<std::unique_ptr<Session>> SessionPool::Allocate(
       // Wake up everyone that was waiting for a session.
       cond_.notify_all();
     }
-    return {std::move(session)};
+    return {MakeSessionHolder(std::move(session), dissociate_from_pool)};
   }
 }
 
-void SessionPool::Release(std::unique_ptr<Session> session) {
+void SessionPool::Release(Session* session) {
   std::unique_lock<std::mutex> lk(mu_);
   bool notify = sessions_.empty();
-  sessions_.push_back(std::move(session));
+  sessions_.emplace_back(session);
   // If sessions_ was empty, wake up someone who was waiting for a session.
   if (notify) {
     lk.unlock();
     cond_.notify_one();
   }
 }
 
+StatusOr<std::vector<std::unique_ptr<Session>>> SessionPool::CreateSessions(
+    int num_sessions) {
+  spanner_proto::BatchCreateSessionsRequest request;
+  request.set_database(db_.FullName());
+  request.set_session_count(std::int32_t{num_sessions});
+  auto response = RetryLoop(
+      retry_policy_->clone(), backoff_policy_->clone(), true,
+      [this](grpc::ClientContext& context,
+             spanner_proto::BatchCreateSessionsRequest const& request) {
+        return stub_->BatchCreateSessions(context, request);
+      },
+      request, __func__);
+  if (!response) {
+    return response.status();
+  }
+  std::vector<std::unique_ptr<Session>> sessions;
+  sessions.reserve(response->session_size());
+  for (auto& session : *response->mutable_session()) {
+    sessions.push_back(google::cloud::internal::make_unique<Session>(
+        std::move(*session.mutable_name())));
+  }
+  return {std::move(sessions)};
+}
+
+SessionHolder SessionPool::MakeSessionHolder(std::unique_ptr<Session> session,
+                                             bool dissociate_from_pool) {
+  if (dissociate_from_pool) {
+    // Uses the default deleter; the `Session` is not returned to the pool.
+    return {std::move(session)};
+  }
+  std::weak_ptr<SessionPool> pool = shared_from_this();
+  return SessionHolder(session.release(), [pool](Session* session) {
+    auto shared_pool = pool.lock();
+    // If `pool` is still alive, release the `Session` to it; otherwise just
+    // delete the `Session`.
+    if (shared_pool) {
+      shared_pool->Release(session);
+    } else {
+      delete session;
+    }
+  });
+}
+
 }  // namespace internal
 }  // namespace SPANNER_CLIENT_NS
 }  // namespace spanner