feat: add SessionPoolOptions with labels to MakeConnection() (#1109)

Moves `SessionPoolOptions` out of the `internal` namespace,
and adds a `std::map<std::string, std::string> labels` field.
Plumbs those labels through to the `BatchCreateSessions()`
call.

Fixes #421.

Author: devbww

google/cloud/spanner/CMakeLists.txt

@@ -193,6 +193,7 @@ add_library(
     retry_policy.h
     row.cc
     row.h
+    session_pool_options.h
     sql_statement.cc
     sql_statement.h
     timestamp.h

google/cloud/spanner/client.cc

@@ -219,24 +219,29 @@ StatusOr<PartitionedDmlResult> Client::ExecutePartitionedDml(
   return conn_->ExecutePartitionedDml({std::move(statement)});
 }
 
-std::shared_ptr<Connection> MakeConnection(Database const& db,
-                                           ConnectionOptions const& options) {
-  return MakeConnection(db, options, internal::DefaultConnectionRetryPolicy(),
+std::shared_ptr<Connection> MakeConnection(
+    Database const& db, ConnectionOptions const& connection_options,
+    SessionPoolOptions session_pool_options) {
+  return MakeConnection(db, connection_options, std::move(session_pool_options),
+                        internal::DefaultConnectionRetryPolicy(),
                         internal::DefaultConnectionBackoffPolicy());
 }
 
 std::shared_ptr<Connection> MakeConnection(
-    Database const& db, ConnectionOptions const& options,
+    Database const& db, ConnectionOptions const& connection_options,
+    SessionPoolOptions session_pool_options,
     std::unique_ptr<RetryPolicy> retry_policy,
     std::unique_ptr<BackoffPolicy> backoff_policy) {
   std::vector<std::shared_ptr<internal::SpannerStub>> stubs;
-  int num_channels = std::min(options.num_channels(), 1);
+  int num_channels = std::min(connection_options.num_channels(), 1);
   stubs.reserve(num_channels);
   for (int channel_id = 0; channel_id < num_channels; ++channel_id) {
-    stubs.push_back(internal::CreateDefaultSpannerStub(options, channel_id));
+    stubs.push_back(
+        internal::CreateDefaultSpannerStub(connection_options, channel_id));
   }
-  return internal::MakeConnection(db, std::move(stubs), std::move(retry_policy),
-                                  std::move(backoff_policy));
+  return internal::MakeConnection(
+      db, std::move(stubs), std::move(session_pool_options),
+      std::move(retry_policy), std::move(backoff_policy));
 }
 
 }  // namespace SPANNER_CLIENT_NS

google/cloud/spanner/client.h

@@ -29,6 +29,7 @@
 #include "google/cloud/spanner/read_partition.h"
 #include "google/cloud/spanner/results.h"
 #include "google/cloud/spanner/retry_policy.h"
+#include "google/cloud/spanner/session_pool_options.h"
 #include "google/cloud/spanner/sql_statement.h"
 #include "google/cloud/spanner/transaction.h"
 #include "google/cloud/optional.h"
@@ -534,23 +535,28 @@ class Client {
  * @see `Connection`
  *
  * @param db See `Database`.
- * @param options (optional) configure the `Connection` created by this
- *     function.
+ * @param connection_options (optional) configure the `Connection` created by
+ *     this function.
+ * @param session_pool_options (optional) configure the `SessionPool` created
+ *     by the `Connection`.
  */
 std::shared_ptr<Connection> MakeConnection(
-    Database const& db, ConnectionOptions const& options = ConnectionOptions());
+    Database const& db,
+    ConnectionOptions const& connection_options = ConnectionOptions(),
+    SessionPoolOptions session_pool_options = SessionPoolOptions());
 
 /**
- * @copydoc MakeConnection(Database const&, ConnectionOptions const&)
+ * @copydoc MakeConnection(Database const&, ConnectionOptions const&, SessionPoolOptions)
  *
- * @param retry_policy override the default `RetryPolicy`, controls for how long
- *     does the returned `Connection` object retry requests on transient
+ * @param retry_policy override the default `RetryPolicy`, controls how long
+ *     the returned `Connection` object retries requests on transient
  *     failures.
- * @param backoff_policy override the default `BackoffPolicy`, controls for how
- *     long does the `Connection` object waits before retrying a failed request.
+ * @param backoff_policy override the default `BackoffPolicy`, controls how
+ *     long the `Connection` object waits before retrying a failed request.
  */
 std::shared_ptr<Connection> MakeConnection(
-    Database const& db, ConnectionOptions const& options,
+    Database const& db, ConnectionOptions const& connection_options,
+    SessionPoolOptions session_pool_options,
     std::unique_ptr<RetryPolicy> retry_policy,
     std::unique_ptr<BackoffPolicy> backoff_policy);
 

google/cloud/spanner/client_test.cc

@@ -389,6 +389,9 @@ TEST(ClientTest, MakeConnectionOptionalArguments) {
 
   conn = MakeConnection(db, ConnectionOptions());
   EXPECT_NE(conn, nullptr);
+
+  conn = MakeConnection(db, ConnectionOptions(), SessionPoolOptions());
+  EXPECT_NE(conn, nullptr);
 }
 
 TEST(ClientTest, CommitMutatorSuccess) {

google/cloud/spanner/internal/connection_impl.cc

@@ -78,22 +78,25 @@ std::unique_ptr<BackoffPolicy> DefaultConnectionBackoffPolicy() {
 
 std::shared_ptr<ConnectionImpl> MakeConnection(
     Database db, std::vector<std::shared_ptr<SpannerStub>> stubs,
+    SessionPoolOptions session_pool_options,
     std::unique_ptr<RetryPolicy> retry_policy,
     std::unique_ptr<BackoffPolicy> backoff_policy) {
-  return std::shared_ptr<ConnectionImpl>(
-      new ConnectionImpl(std::move(db), std::move(stubs),
-                         std::move(retry_policy), std::move(backoff_policy)));
+  return std::shared_ptr<ConnectionImpl>(new ConnectionImpl(
+      std::move(db), std::move(stubs), std::move(session_pool_options),
+      std::move(retry_policy), std::move(backoff_policy)));
 }
 
 ConnectionImpl::ConnectionImpl(Database db,
                                std::vector<std::shared_ptr<SpannerStub>> stubs,
+                               SessionPoolOptions session_pool_options,
                                std::unique_ptr<RetryPolicy> retry_policy,
                                std::unique_ptr<BackoffPolicy> backoff_policy)
     : db_(std::move(db)),
       retry_policy_prototype_(std::move(retry_policy)),
       backoff_policy_prototype_(std::move(backoff_policy)),
       session_pool_(std::make_shared<SessionPool>(
-          db_, std::move(stubs), retry_policy_prototype_->clone(),
+          db_, std::move(stubs), std::move(session_pool_options),
+          retry_policy_prototype_->clone(),
           backoff_policy_prototype_->clone())) {}
 
 RowStream ConnectionImpl::Read(ReadParams params) {

google/cloud/spanner/internal/connection_impl.h

@@ -52,6 +52,7 @@ std::unique_ptr<BackoffPolicy> DefaultConnectionBackoffPolicy();
 class ConnectionImpl;
 std::shared_ptr<ConnectionImpl> MakeConnection(
     Database db, std::vector<std::shared_ptr<SpannerStub>> stubs,
+    SessionPoolOptions session_pool_options = SessionPoolOptions{},
     std::unique_ptr<RetryPolicy> retry_policy = DefaultConnectionRetryPolicy(),
     std::unique_ptr<BackoffPolicy> backoff_policy =
         DefaultConnectionBackoffPolicy());
@@ -82,9 +83,10 @@ class ConnectionImpl : public Connection {
  private:
   // Only the factory method can construct instances of this class.
   friend std::shared_ptr<ConnectionImpl> MakeConnection(
-      Database, std::vector<std::shared_ptr<SpannerStub>>,
+      Database, std::vector<std::shared_ptr<SpannerStub>>, SessionPoolOptions,
       std::unique_ptr<RetryPolicy>, std::unique_ptr<BackoffPolicy>);
   ConnectionImpl(Database db, std::vector<std::shared_ptr<SpannerStub>> stubs,
+                 SessionPoolOptions session_pool_options,
                  std::unique_ptr<RetryPolicy> retry_policy,
                  std::unique_ptr<BackoffPolicy> backoff_policy);
 

google/cloud/spanner/internal/connection_impl_test.cc

@@ -92,7 +92,7 @@ std::shared_ptr<Connection> MakeLimitedRetryConnection(
     Database const& db,
     std::shared_ptr<spanner_testing::MockSpannerStub> mock) {
   return MakeConnection(
-      db, {std::move(mock)},
+      db, {std::move(mock)}, SessionPoolOptions{},
       LimitedErrorCountRetryPolicy(/*maximum_failures=*/2).clone(),
       ExponentialBackoffPolicy(/*initial_delay=*/std::chrono::microseconds(1),
                                /*maximum_delay=*/std::chrono::microseconds(1),

google/cloud/spanner/internal/session_pool.cc

@@ -47,13 +47,14 @@ SessionPoolOptions SanitizeOptions(SessionPoolOptions options,
 
 SessionPool::SessionPool(Database db,
                          std::vector<std::shared_ptr<SpannerStub>> stubs,
+                         SessionPoolOptions options,
                          std::unique_ptr<RetryPolicy> retry_policy,
-                         std::unique_ptr<BackoffPolicy> backoff_policy,
-                         SessionPoolOptions options)
+                         std::unique_ptr<BackoffPolicy> backoff_policy)
     : db_(std::move(db)),
+      options_(
+          SanitizeOptions(std::move(options), static_cast<int>(stubs.size()))),
       retry_policy_prototype_(std::move(retry_policy)),
       backoff_policy_prototype_(std::move(backoff_policy)),
-      options_(SanitizeOptions(options, static_cast<int>(stubs.size()))),
       max_pool_size_(options_.max_sessions_per_channel *
                      static_cast<int>(stubs.size())) {
   if (stubs.empty()) {
@@ -88,7 +89,7 @@ SessionPool::SessionPool(Database db,
     }
     // Just ignore failures; we'll try again when the caller requests a
     // session, and we'll be in a position to return an error at that time.
-    (void)CreateSessions(lk, channel, num_sessions);
+    (void)CreateSessions(lk, channel, options_.labels, num_sessions);
   }
 }
 
@@ -135,7 +136,8 @@ StatusOr<SessionHolder> SessionPool::Allocate(bool dissociate_from_pool) {
     int sessions_to_create =
         (std::min)(options_.min_sessions + 1,
                    options_.max_sessions_per_channel - channel.session_count);
-    auto create_status = CreateSessions(lk, channel, sessions_to_create);
+    auto create_status =
+        CreateSessions(lk, channel, options_.labels, sessions_to_create);
     if (!create_status.ok()) {
       return create_status;
     }
@@ -173,12 +175,15 @@ void SessionPool::Release(Session* session) {
 //
 // Requires `lk` has locked `mu_` prior to this call. `lk` will be dropped
 // while the RPC is in progress and then reacquired.
-Status SessionPool::CreateSessions(std::unique_lock<std::mutex>& lk,
-                                   ChannelInfo& channel, int num_sessions) {
+Status SessionPool::CreateSessions(
+    std::unique_lock<std::mutex>& lk, ChannelInfo& channel,
+    std::map<std::string, std::string> const& labels, int num_sessions) {
   create_in_progress_ = true;
   lk.unlock();
   spanner_proto::BatchCreateSessionsRequest request;
   request.set_database(db_.FullName());
+  request.mutable_session_template()->mutable_labels()->insert(labels.begin(),
+                                                               labels.end());
   request.set_session_count(std::int32_t{num_sessions});
   auto const& stub = channel.stub;
   auto response = RetryLoop(

google/cloud/spanner/internal/session_pool.h

@@ -20,10 +20,10 @@
 #include "google/cloud/spanner/internal/session.h"
 #include "google/cloud/spanner/internal/spanner_stub.h"
 #include "google/cloud/spanner/retry_policy.h"
+#include "google/cloud/spanner/session_pool_options.h"
 #include "google/cloud/spanner/version.h"
 #include "google/cloud/status_or.h"
 #include <google/spanner/v1/spanner.pb.h>
-#include <chrono>
 #include <condition_variable>
 #include <cstddef>
 #include <map>
@@ -39,34 +39,6 @@ namespace spanner {
 inline namespace SPANNER_CLIENT_NS {
 namespace internal {
 
-// What action to take if the session pool is exhausted.
-enum class ActionOnExhaustion { BLOCK, FAIL };
-
-struct SessionPoolOptions {
-  // The minimum number of sessions to keep in the pool.
-  // Values <= 0 are treated as 0.
-  // This value will be reduced if it exceeds the overall limit on the number
-  // of sessions (`max_sessions_per_channel` * number of channels).
-  int min_sessions = 0;
-
-  // The maximum number of sessions to create on each channel.
-  // Values <= 1 are treated as 1.
-  int max_sessions_per_channel = 100;
-
-  // The maximum number of sessions that can be in the pool in an idle state.
-  // Values <= 0 are treated as 0.
-  int max_idle_sessions = 0;
-
-  // Decide whether to block or fail on pool exhaustion.
-  ActionOnExhaustion action_on_exhaustion = ActionOnExhaustion::BLOCK;
-
-  // This is the interval at which we refresh sessions so they don't get
-  // collected by the backend GC. The GC collects objects older than 60
-  // minutes, so any duration below that (less some slack to allow the calls
-  // to be made to refresh the sessions) should suffice.
-  std::chrono::minutes keep_alive_interval = std::chrono::minutes(55);
-};
-
 /**
  * Maintains a pool of `Session` objects.
  *
@@ -90,9 +62,9 @@ class SessionPool : public std::enable_shared_from_this<SessionPool> {
    * create them. `stubs` must not be empty.
    */
   SessionPool(Database db, std::vector<std::shared_ptr<SpannerStub>> stubs,
+              SessionPoolOptions options,
               std::unique_ptr<RetryPolicy> retry_policy,
-              std::unique_ptr<BackoffPolicy> backoff_policy,
-              SessionPoolOptions options = SessionPoolOptions());
+              std::unique_ptr<BackoffPolicy> backoff_policy);
 
   /**
    * Allocate a `Session` from the pool, creating a new one if necessary.
@@ -130,6 +102,7 @@ class SessionPool : public std::enable_shared_from_this<SessionPool> {
   void Release(Session* session);
 
   Status CreateSessions(std::unique_lock<std::mutex>& lk, ChannelInfo& channel,
+                        std::map<std::string, std::string> const& labels,
                         int num_sessions);  // EXCLUSIVE_LOCKS_REQUIRED(mu_)
 
   SessionHolder MakeSessionHolder(std::unique_ptr<Session> session,
@@ -138,9 +111,9 @@ class SessionPool : public std::enable_shared_from_this<SessionPool> {
   void UpdateNextChannelForCreateSessions();  // EXCLUSIVE_LOCKS_REQUIRED(mu_)
 
   Database const db_;
+  SessionPoolOptions const options_;
   std::unique_ptr<RetryPolicy const> retry_policy_prototype_;
   std::unique_ptr<BackoffPolicy const> backoff_policy_prototype_;
-  SessionPoolOptions const options_;
   int const max_pool_size_;
 
   std::mutex mu_;