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

Author: scotthart

CHANGELOG.md

@@ -4,7 +4,11 @@
 breaking changes in the upcoming 3.x release. This release is scheduled for
 2024-12 or 2025-01.
 
-## v2.22.1 - TBD
+## v2.22.1 - 2025-03
+
+### [Spanner](/google/cloud/spanner/README.md)
+
+- fix(spanner): update session bookkeeping for session NotFound (#15009)
 
 ## v2.22.0 - 2024-03
 

CMakeLists.txt

@@ -23,7 +23,7 @@ project(
     google-cloud-cpp
     VERSION 2.22.1
     LANGUAGES CXX)
-set(PROJECT_VERSION_PRE_RELEASE "rc")
+set(PROJECT_VERSION_PRE_RELEASE "")
 
 if (NOT "${PROJECT_VERSION_PRE_RELEASE}" STREQUAL "")
     set(PROJECT_VERSION "${PROJECT_VERSION}-${PROJECT_VERSION_PRE_RELEASE}")

google/cloud/internal/version_info.h

@@ -21,6 +21,6 @@
 #define GOOGLE_CLOUD_CPP_VERSION_MINOR 22
 // NOLINTNEXTLINE(modernize-macro-to-enum)
 #define GOOGLE_CLOUD_CPP_VERSION_PATCH 1
-#define GOOGLE_CLOUD_CPP_VERSION_PRE_RELEASE "rc"
+#define GOOGLE_CLOUD_CPP_VERSION_PRE_RELEASE ""
 
 #endif  // GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_INTERNAL_VERSION_INFO_H

google/cloud/spanner/internal/session_pool.cc

@@ -184,6 +184,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;
     }
   }
@@ -287,6 +288,21 @@ Status SessionPool::CreateSessions(
   return return_status;
 }
 
+int SessionPool::total_sessions() const {
+  std::lock_guard<std::mutex> lk(mu_);
+  return total_sessions_;
+}
+
+void SessionPool::DecrementSessionCount(
+    std::unique_lock<std::mutex> const&,
+    google::cloud::spanner_internal::Session const& session) {
+  --total_sessions_;
+  auto const& channel = session.channel();
+  if (channel) {
+    --channel->session_count;
+  }
+}
+
 StatusOr<SessionHolder> SessionPool::Allocate(bool dissociate_from_pool) {
   // We choose to ignore the internal::CurrentOptions() here as it is
   // non-deterministic when RPCs to create sessions are actually made.
@@ -299,11 +315,7 @@ StatusOr<SessionHolder> SessionPool::Allocate(bool dissociate_from_pool) {
       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)};
     }
@@ -367,11 +379,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

@@ -99,6 +99,15 @@ 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.
+   *
+   * @note This function should only be used for testing as other threads
+   *   could be modifying the underlying value immediately after it returns.
+   */
+  int total_sessions() const;
+
  private:
   friend std::shared_ptr<SessionPool> MakeSessionPool(
       spanner::Database, std::vector<std::shared_ptr<SpannerStub>>,
@@ -179,6 +188,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 const& session);
+
   spanner::Database const db_;
   google::cloud::CompletionQueue cq_;
   Options const opts_;
@@ -188,7 +201,7 @@ class SessionPool : public std::enable_shared_from_this<SessionPool> {
   int const max_pool_size_;
   std::mt19937 random_generator_;
 
-  std::mutex mu_;
+  mutable std::mutex mu_;
   std::condition_variable cond_;
   std::vector<std::unique_ptr<Session>> sessions_;  // GUARDED_BY(mu_)
   int total_sessions_ = 0;                          // GUARDED_BY(mu_)

google/cloud/spanner/internal/session_pool_test.cc

@@ -601,21 +601,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