fix(spanner): update session bookkeeping for session NotFound (#15009)

Author: scotthart

google/cloud/spanner/internal/session_pool.cc

@@ -199,6 +199,7 @@ void SessionPool::Erase(std::string const& session_name) {
       target = std::move(session);  // deferred deletion
       session = std::move(sessions_.back());
       sessions_.pop_back();
+      DecrementSessionCount(lk, *target);
       break;
     }
   }
@@ -369,6 +370,16 @@ std::shared_ptr<SpannerStub> SessionPool::GetStub(Session const& session) {
   return GetStub(std::unique_lock<std::mutex>(mu_));
 }
 
+void SessionPool::DecrementSessionCount(
+    std::unique_lock<std::mutex> const&,
+    google::cloud::spanner_internal::Session& session) {
+  --total_sessions_;
+  auto const& channel = session.channel();
+  if (channel) {
+    --channel->session_count;
+  }
+}
+
 StatusOr<SessionHolder> SessionPool::Allocate(std::unique_lock<std::mutex> lk,
                                               bool dissociate_from_pool) {
   // We choose to ignore the internal::CurrentOptions() here as it is
@@ -381,11 +392,7 @@ StatusOr<SessionHolder> SessionPool::Allocate(std::unique_lock<std::mutex> lk,
       auto session = std::move(sessions_.back());
       sessions_.pop_back();
       if (dissociate_from_pool) {
-        --total_sessions_;
-        auto const& channel = session->channel();
-        if (channel) {
-          --channel->session_count;
-        }
+        DecrementSessionCount(lk, *session);
       }
       return {MakeSessionHolder(std::move(session), dissociate_from_pool)};
     }
@@ -442,11 +449,7 @@ void SessionPool::Release(std::unique_ptr<Session> session) {
   if (session->is_bad()) {
     // Once we have support for background processing, we may want to signal
     // that to replenish this bad session.
-    --total_sessions_;
-    auto const& channel = session->channel();
-    if (channel) {
-      --channel->session_count;
-    }
+    DecrementSessionCount(lk, *session);
     return;
   }
   session->update_last_use_time();

google/cloud/spanner/internal/session_pool.h

@@ -115,6 +115,12 @@ class SessionPool : public std::enable_shared_from_this<SessionPool> {
    */
   std::shared_ptr<SpannerStub> GetStub(Session const& session);
 
+  /**
+   * Returns the number of sessions in the session pool plus the number of
+   * sessions allocated to running transactions.
+   */
+  int total_sessions() const { return total_sessions_; }
+
  private:
   friend std::shared_ptr<SessionPool> MakeSessionPool(
       spanner::Database, std::vector<std::shared_ptr<SpannerStub>>,
@@ -207,6 +213,10 @@ class SessionPool : public std::enable_shared_from_this<SessionPool> {
   // Remove the named session from the pool (if it is present).
   void Erase(std::string const& session_name);
 
+  // Performs the necessary bookkeeping when a session is removed from use.
+  void DecrementSessionCount(std::unique_lock<std::mutex> const& lk,
+                             Session& session);
+
   spanner::Database const db_;
   google::cloud::CompletionQueue cq_;
   Options const opts_;
@@ -220,9 +230,11 @@ class SessionPool : public std::enable_shared_from_this<SessionPool> {
   std::condition_variable cond_;
   SessionHolder multiplexed_session_;               // GUARDED_BY(mu_)
   std::vector<std::unique_ptr<Session>> sessions_;  // GUARDED_BY(mu_)
-  int total_sessions_ = 0;                          // GUARDED_BY(mu_)
-  int create_calls_in_progress_ = 0;                // GUARDED_BY(mu_)
-  int num_waiting_for_session_ = 0;                 // 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.
+  int total_sessions_ = 0;            // GUARDED_BY(mu_)
+  int create_calls_in_progress_ = 0;  // GUARDED_BY(mu_)
+  int num_waiting_for_session_ = 0;   // GUARDED_BY(mu_)
 
   // Lower bound on all `sessions_[i]->last_use_time()` values.
   Session::Clock::time_point last_use_time_lower_bound_ =

google/cloud/spanner/internal/session_pool_test.cc

@@ -856,21 +856,26 @@ TEST_F(SessionPoolTest, SessionRefreshNotFound) {
     // Wait for "s2" to need refreshing before releasing "s1".
     clock->AdvanceTime(keep_alive_interval * 2);
   }
+  EXPECT_EQ(pool->total_sessions(), 2);
 
   // Simulate completion of pending operations, which will result in
   // a call to RefreshExpiringSessions(). This should refresh "s2" and
   // satisfy the AsyncExecuteSql() expectation, which fails the call.
   impl->SimulateCompletion(true);
+  EXPECT_EQ(pool->total_sessions(), 1);
 
   // We should still be able to allocate session "s1".
   auto s1 = pool->Allocate();
   ASSERT_STATUS_OK(s1);
   EXPECT_EQ("s1", (*s1)->session_name());
+  EXPECT_EQ(pool->total_sessions(), 1);
+
   // However "s2" will be gone now, so a new allocation will produce
   // "s3", satisfying the final BatchCreateSessions() expectation.
   auto s3 = pool->Allocate();
   ASSERT_STATUS_OK(s3);
   EXPECT_EQ("s3", (*s3)->session_name());
+  EXPECT_EQ(pool->total_sessions(), 2);
 
   // Cancel all pending operations, satisfying any remaining futures. When
   // compiling with exceptions disabled the destructors eventually invoke