impl(spanner): add support for Multiplexed Session maintenance (#15226)

Author: scotthart

google/cloud/spanner/internal/connection_impl_test.cc

@@ -21,7 +21,15 @@ GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
 
 SessionHolder MakeDissociatedSessionHolder(std::string session_name) {
   return SessionHolder(
-      new Session(std::move(session_name), /*channel=*/nullptr),
+      new Session(std::move(session_name), Session::Mode::kDisassociated),
+      std::default_delete<Session>());
+}
+
+SessionHolder MakeMultiplexedSessionHolder(
+    std::string session_name, std::shared_ptr<Session::Clock> clock) {
+  return SessionHolder(
+      new Session(std::move(session_name), Session::Mode::kMultiplexed,
+                  std::move(clock)),
       std::default_delete<Session>());
 }
 

google/cloud/spanner/internal/session.cc

@@ -39,14 +39,23 @@ GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
  */
 class Session {
  public:
+  enum class Mode { kPooled, kDisassociated, kMultiplexed };
+
   using Clock = ::google::cloud::internal::SteadyClock;
   Session(std::string session_name, std::shared_ptr<Channel> channel,
           std::shared_ptr<Clock> clock = std::make_shared<Clock>())
       : session_name_(std::move(session_name)),
         channel_(std::move(channel)),
         is_bad_(false),
         clock_(std::move(clock)),
-        last_use_time_(clock_->Now()) {}
+        last_use_time_(clock_->Now()),
+        creation_time_(clock_->Now()) {}
+
+  explicit Session(std::string session_name, Mode mode,
+                   std::shared_ptr<Clock> clock = std::make_shared<Clock>())
+      : Session(std::move(session_name), nullptr, std::move(clock)) {
+    mode_ = mode;
+  }
 
   // Not copyable or moveable.
   Session(Session const&) = delete;
@@ -60,6 +69,9 @@ class Session {
   void set_bad() { is_bad_.store(true, std::memory_order_relaxed); }
   bool is_bad() const { return is_bad_.load(std::memory_order_relaxed); }
 
+  bool is_disassociated() const { return mode_ == Mode::kDisassociated; }
+  bool is_multiplexed() const { return mode_ == Mode::kMultiplexed; }
+
  private:
   // Give `SessionPool` access to the private methods below.
   friend class SessionPool;
@@ -71,11 +83,15 @@ class Session {
   Clock::time_point last_use_time() const { return last_use_time_; }
   void update_last_use_time() { last_use_time_ = clock_->Now(); }
 
+  Clock::time_point creation_time() const { return creation_time_; }
+
   std::string const session_name_;
   std::shared_ptr<Channel> const channel_;
   std::atomic<bool> is_bad_;
   std::shared_ptr<Clock> clock_;
   Clock::time_point last_use_time_;
+  Clock::time_point creation_time_;
+  Mode mode_ = Mode::kPooled;
 };
 
 /**
@@ -94,6 +110,14 @@ using SessionHolder = std::shared_ptr<Session>;
  */
 SessionHolder MakeDissociatedSessionHolder(std::string session_name);
 
+/**
+ * Returns a `SessionHolder` for a new `Session` that is not associated with
+ * any pool; it just deletes the `Session`. This is for use with Multiplexed
+ * Sessions that do not require a pool.
+ */
+SessionHolder MakeMultiplexedSessionHolder(
+    std::string session_name, std::shared_ptr<Session::Clock> clock);
+
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
 }  // namespace spanner_internal
 }  // namespace cloud

google/cloud/spanner/internal/session.h

@@ -63,12 +63,28 @@ SessionPool::SessionPool(spanner::Database db,
           opts_.get<spanner::SessionPoolMaxSessionsPerChannelOption>() *
           static_cast<int>(stubs.size())),
       random_generator_(std::random_device()()),
+      multiplexed_session_replacement_interval_(
+          opts_.has<
+              spanner_internal::MultiplexedSessionReplacementIntervalOption>()
+              ? opts_.get<spanner_internal::
+                              MultiplexedSessionReplacementIntervalOption>()
+              : std::chrono::hours(24 * 7)),
+      multiplexed_session_background_interval_(
+          opts_.has<spanner_internal::
+                        MultiplexedSessionBackgroundWorkIntervalOption>()
+              ? opts_.get<spanner_internal::
+                              MultiplexedSessionBackgroundWorkIntervalOption>()
+              : std::chrono::minutes(10)),
       channels_(stubs.size()) {
   if (stubs.empty()) {
     google::cloud::internal::ThrowInvalidArgument(
         "SessionPool requires a non-empty set of stubs");
   }
 
+  multiplexed_session_ = internal::NotFoundError(
+      "no multiplexed session is currently available; try again later",
+      GCP_ERROR_INFO());
+
   for (auto i = 0U; i < stubs.size(); ++i) {
     channels_[i] = std::make_shared<Channel>(std::move(stubs[i]));
   }
@@ -78,13 +94,19 @@ SessionPool::SessionPool(spanner::Database db,
 
 void SessionPool::Initialize() {
   internal::OptionsSpan span(opts_);
-  CreateMultiplexedSession();
-  auto const min_sessions = opts_.get<spanner::SessionPoolMinSessionsOption>();
-  if (min_sessions > 0) {
+  if (opts_.has<spanner_experimental::EnableMultiplexedSessionOption>()) {
     std::unique_lock<std::mutex> lk(mu_);
-    Grow(lk, min_sessions, WaitForSessionAllocation::kWait);
+    CreateMultiplexedSession(lk);
+    ScheduleMultiplexedBackgroundWork(multiplexed_session_background_interval_);
+  } else {
+    auto const min_sessions =
+        opts_.get<spanner::SessionPoolMinSessionsOption>();
+    if (min_sessions > 0) {
+      std::unique_lock<std::mutex> lk(mu_);
+      Grow(lk, min_sessions, WaitForSessionAllocation::kWait);
+    }
+    ScheduleBackgroundWork(std::chrono::seconds(5));
   }
-  ScheduleBackgroundWork(std::chrono::seconds(5));
 }
 
 SessionPool::~SessionPool() {
@@ -98,21 +120,60 @@ SessionPool::~SessionPool() {
   // they must not have successfully finished a call to `lock()` on the
   // `weak_ptr` to `this` they hold. Any in-progress or subsequent `lock()`
   // will now return `nullptr`, in which case no work is done.
-  current_timer_.cancel();
+  if (opts_.has<spanner_experimental::EnableMultiplexedSessionOption>()) {
+    current_multiplexed_timer_.cancel();
+  } else {
+    current_timer_.cancel();
+  }
 
   // Send fire-and-forget `AsyncDeleteSession()` calls for all sessions.
-  if (HasValidMultiplexedSession(std::unique_lock<std::mutex>(mu_))) {
-    AsyncDeleteSession(cq_, GetStub(*multiplexed_session_),
-                       multiplexed_session_->session_name())
-        .then([](auto result) { auto status = result.get(); });
-  }
+  // Multiplexed Sessions do not require an explicit Delete call.
   for (auto const& session : sessions_) {
     if (session->is_bad()) continue;
     AsyncDeleteSession(cq_, GetStub(*session), session->session_name())
         .then([](auto result) { auto status = result.get(); });
   }
 }
 
+void SessionPool::ScheduleMultiplexedBackgroundWork(
+    std::chrono::seconds relative_time) {
+  std::weak_ptr<SessionPool> pool = shared_from_this();
+  current_multiplexed_timer_ =
+      cq_.MakeRelativeTimer(relative_time)
+          .then([pool](future<StatusOr<std::chrono::system_clock::time_point>>
+                           result) {
+            if (result.get().ok()) {
+              if (auto shared_pool = pool.lock()) {
+                shared_pool->DoMultiplexedBackgroundWork();
+              }
+            }
+          });
+}
+
+void SessionPool::ReplaceMultiplexedSession() {
+  auto now = clock_->Now();
+  auto refresh_limit = now - multiplexed_session_replacement_interval_;
+  std::unique_lock<std::mutex> lk(mu_);
+  if (create_calls_in_progress_ == 0 &&
+      (*multiplexed_session_)->creation_time() <= refresh_limit) {
+    ++create_calls_in_progress_;
+    auto stub = GetStub(std::move(lk));
+    std::weak_ptr<SessionPool> pool = shared_from_this();
+    CreateMultiplexedSessionAsync(std::move(stub))
+        .then([pool](future<StatusOr<google::spanner::v1::Session>> response) {
+          if (auto shared_pool = pool.lock()) {
+            shared_pool->HandleMultiplexedCreateSessionDone(
+                std::move(response).get());
+          }
+        });
+  }
+}
+
+void SessionPool::DoMultiplexedBackgroundWork() {
+  ReplaceMultiplexedSession();
+  ScheduleMultiplexedBackgroundWork(multiplexed_session_background_interval_);
+}
+
 void SessionPool::ScheduleBackgroundWork(std::chrono::seconds relative_time) {
   std::weak_ptr<SessionPool> pool = shared_from_this();
   current_timer_ =
@@ -136,7 +197,6 @@ void SessionPool::DoBackgroundWork() {
 // Ensure the pool size conforms to what was specified in the `SessionOptions`,
 // creating or deleting sessions as necessary.
 void SessionPool::MaintainPoolSize() {
-  CreateMultiplexedSession();
   auto const min_sessions = opts_.get<spanner::SessionPoolMinSessionsOption>();
   std::unique_lock<std::mutex> lk(mu_);
   if (create_calls_in_progress_ == 0 && total_sessions_ < min_sessions) {
@@ -205,22 +265,30 @@ void SessionPool::Erase(std::string const& session_name) {
   }
 }
 
-Status SessionPool::CreateMultiplexedSession() {
+Status SessionPool::HandleMultiplexedCreateSessionDone(
+    StatusOr<google::spanner::v1::Session> response) {
   std::unique_lock<std::mutex> lk(mu_);
-  if (!HasValidMultiplexedSession(lk)) {
+  --create_calls_in_progress_;
+  if (!response.ok()) {
+    multiplexed_session_ = response.status();
+    return response.status();
+  }
+
+  multiplexed_session_ = MakeMultiplexedSessionHolder(response->name(), clock_);
+  return {};
+}
+
+Status SessionPool::CreateMultiplexedSession(std::unique_lock<std::mutex>& lk) {
+  if (create_calls_in_progress_ == 0) {
+    create_calls_in_progress_++;
     auto stub = GetStub(std::move(lk));
-    auto name = CreateMultiplexedSession(std::move(stub));
-    if (!name) return name.status();
-    auto session = std::make_shared<Session>(*std::move(name),
-                                             /*channel=*/nullptr, clock_);
-    std::unique_lock<std::mutex> lk(mu_);
-    multiplexed_session_ = std::move(session);
+    return CreateMultiplexedSessionSync(std::move(stub));
   }
   return Status{};
 }
 
-StatusOr<std::string> SessionPool::CreateMultiplexedSession(
-    std::shared_ptr<SpannerStub> stub) const {
+Status SessionPool::CreateMultiplexedSessionSync(
+    std::shared_ptr<SpannerStub> stub) {
   google::spanner::v1::CreateSessionRequest request;
   request.set_database(db_.FullName());
   auto* session = request.mutable_session();
@@ -241,13 +309,33 @@ StatusOr<std::string> SessionPool::CreateMultiplexedSession(
         return stub->CreateSession(context, options, request);
       },
       opts_, request, __func__);
-  if (!response) return std::move(response).status();
-  return response->name();
+  return HandleMultiplexedCreateSessionDone(std::move(response));
 }
 
-bool SessionPool::HasValidMultiplexedSession(
-    std::unique_lock<std::mutex> const&) const {
-  return multiplexed_session_ && !multiplexed_session_->is_bad();
+future<StatusOr<google::spanner::v1::Session>>
+SessionPool::CreateMultiplexedSessionAsync(std::shared_ptr<SpannerStub> stub) {
+  google::spanner::v1::CreateSessionRequest request;
+  request.set_database(db_.FullName());
+  auto* session = request.mutable_session();
+  auto const& labels = opts_.get<spanner::SessionPoolLabelsOption>();
+  if (!labels.empty()) {
+    session->mutable_labels()->insert(labels.begin(), labels.end());
+  }
+  auto const& role = opts_.get<spanner::SessionCreatorRoleOption>();
+  if (!role.empty()) session->set_creator_role(role);
+  session->set_multiplexed(true);
+
+  return google::cloud::internal::AsyncRetryLoop(
+      retry_policy_prototype_->clone(), backoff_policy_prototype_->clone(),
+      google::cloud::Idempotency::kIdempotent, cq_,
+      [&stub](CompletionQueue cq, std::shared_ptr<grpc::ClientContext> context,
+              internal::ImmutableOptions options,
+              google::spanner::v1::CreateSessionRequest const& request) {
+        RouteToLeader(*context);  // always for CreateSession()
+        return stub->AsyncCreateSession(cq, std::move(context),
+                                        std::move(options), request);
+      },
+      internal::SaveCurrentOptions(), std::move(request), __func__);
 }
 
 /*
@@ -354,10 +442,12 @@ StatusOr<SessionHolder> SessionPool::Allocate(bool dissociate_from_pool) {
 }
 
 StatusOr<SessionHolder> SessionPool::Multiplexed() {
-  std::unique_lock<std::mutex> lk(mu_);
-  // If we don't have a multiplexed session (yet), use a regular one.
-  if (!HasValidMultiplexedSession(lk)) return Allocate(std::move(lk), false);
-  return multiplexed_session_;
+  if (opts_.has<spanner_experimental::EnableMultiplexedSessionOption>()) {
+    std::unique_lock<std::mutex> lk(mu_);
+    return multiplexed_session_;
+  }
+  return internal::FailedPreconditionError(
+      "multiplexed sessions are not enabled", GCP_ERROR_INFO());
 }
 
 std::shared_ptr<SpannerStub> SessionPool::GetStub(Session const& session) {

google/cloud/spanner/internal/session_pool.cc

@@ -51,6 +51,17 @@ struct SessionPoolClockOption {
   using Type = std::shared_ptr<Session::Clock>;
 };
 
+// Frequency of when the existing multiplexed sessions is replaced with a new
+// multiplexed sessions.
+struct MultiplexedSessionReplacementIntervalOption {
+  using Type = std::chrono::hours;
+};
+
+// Frequency of when background work is performed.
+struct MultiplexedSessionBackgroundWorkIntervalOption {
+  using Type = std::chrono::minutes;
+};
+
 class SessionPool;
 
 /**
@@ -167,10 +178,13 @@ class SessionPool : public std::enable_shared_from_this<SessionPool> {
     --num_waiting_for_session_;
   }
 
-  Status CreateMultiplexedSession();  // LOCKS_EXCLUDED(mu_)
-  StatusOr<std::string> CreateMultiplexedSession(
-      std::shared_ptr<SpannerStub>) const;
-  bool HasValidMultiplexedSession(std::unique_lock<std::mutex> const&) const;
+  Status CreateMultiplexedSession(
+      std::unique_lock<std::mutex>& lk);  // EXCLUSIVE_LOCKS_REQUIRED(mu_)
+  Status CreateMultiplexedSessionSync(std::shared_ptr<SpannerStub>);
+  future<StatusOr<google::spanner::v1::Session>> CreateMultiplexedSessionAsync(
+      std::shared_ptr<SpannerStub>);
+  Status HandleMultiplexedCreateSessionDone(
+      StatusOr<google::spanner::v1::Session> response);
 
   Status Grow(std::unique_lock<std::mutex>& lk, int sessions_to_create,
               WaitForSessionAllocation wait);  // EXCLUSIVE_LOCKS_REQUIRED(mu_)
@@ -208,6 +222,10 @@ class SessionPool : public std::enable_shared_from_this<SessionPool> {
       std::shared_ptr<Channel> const& channel,
       StatusOr<google::spanner::v1::BatchCreateSessionsResponse> response);
 
+  void ScheduleMultiplexedBackgroundWork(std::chrono::seconds relative_time);
+  void DoMultiplexedBackgroundWork();
+  void ReplaceMultiplexedSession();
+
   void ScheduleBackgroundWork(std::chrono::seconds relative_time);
   void DoBackgroundWork();
   void MaintainPoolSize();
@@ -228,10 +246,12 @@ class SessionPool : public std::enable_shared_from_this<SessionPool> {
   std::shared_ptr<Session::Clock> clock_;
   int const max_pool_size_;
   std::mt19937 random_generator_;
+  std::chrono::hours multiplexed_session_replacement_interval_;
+  std::chrono::minutes multiplexed_session_background_interval_;
 
   mutable std::mutex mu_;
   std::condition_variable cond_;
-  SessionHolder multiplexed_session_;               // GUARDED_BY(mu_)
+  StatusOr<SessionHolder> multiplexed_session_;     // GUARDED_BY(mu_)
   std::vector<std::unique_ptr<Session>> sessions_;  // GUARDED_BY(mu_)
   // total_sessions_ tracks the number of sessions in the pool, a.k.a.
   // sessions_.size(), plus the sessions that have been allocated.
@@ -243,6 +263,7 @@ class SessionPool : public std::enable_shared_from_this<SessionPool> {
   Session::Clock::time_point last_use_time_lower_bound_ =
       clock_->Now();  // GUARDED_BY(mu_)
 
+  future<void> current_multiplexed_timer_;
   future<void> current_timer_;
 
   // `channels_` is guaranteed to be non-empty and will not be resized after