fix: only close and return sessions once

olavloite · olavloite · commit a2fc47b92f9f · 2025-04-30T18:39:29.000+02:00
Closing a pooled session multiple times would cause it to be added to the pool
multiple times. This fix prevents this by keeping track of the state of a session
that has been checked out.
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java
@@ -1314,6 +1314,7 @@ default void addListener(Runnable listener, Executor exec) {}
   class PooledSessionFuture extends SimpleForwardingListenableFuture<PooledSession>
       implements SessionFuture {
 
+    private boolean closed;
     private volatile LeakedSessionException leakedException;
     private final AtomicBoolean inUse = new AtomicBoolean();
     private final CountDownLatch initialized = new CountDownLatch(1);
@@ -1331,6 +1332,7 @@ void clearLeakedException() {
     }
 
     private void markCheckedOut() {
+
       if (options.isTrackStackTraceOfSessionCheckout()) {
         this.leakedException = new LeakedSessionException();
         synchronized (SessionPool.this.lock) {
@@ -1520,6 +1522,13 @@ public void close() {
 
     @Override
     public ApiFuture<Empty> asyncClose() {
+      synchronized (this) {
+        // Don't add the session twice to the pool if a resource is being closed multiple times.
+        if (closed) {
+          return ApiFutures.immediateFuture(Empty.getDefaultInstance());
+        }
+        closed = true;
+      }
       try {
         PooledSession delegate = getOrNull();
         if (delegate != null) {
@@ -2709,9 +2718,7 @@ private Tuple<PooledSession, Integer> findSessionToKeepAlive(
     return null;
   }
 
-  /**
-   * @return true if this {@link SessionPool} is still valid.
-   */
+  /** @return true if this {@link SessionPool} is still valid. */
   boolean isValid() {
     synchronized (lock) {
       return closureFuture == null && resourceNotFoundException == null;
@@ -3142,6 +3149,13 @@ int totalSessions() {
     }
   }
 
+  @VisibleForTesting
+  int numSessionsInPool() {
+    synchronized (lock) {
+      return sessions.size();
+    }
+  }
+
   private ApiFuture<Empty> closeSessionAsync(final PooledSession sess) {
     ApiFuture<Empty> res = sess.delegate.asyncClose();
     res.addListener(
diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncTransactionManagerTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncTransactionManagerTest.java
@@ -1237,6 +1237,43 @@ public void testAbandonedAsyncTransactionManager_rollbackFails() throws Exceptio
     assertTrue(gotException);
   }
 
+  @Test
+  public void testRollbackAndCloseEmptyTransaction() throws Exception {
+    assumeFalse(spannerWithEmptySessionPool
+        .getOptions()
+        .getSessionPoolOptions()
+        .getUseMultiplexedSessionForRW());
+
+    DatabaseClientImpl client = (DatabaseClientImpl) clientWithEmptySessionPool();
+
+    // Create a transaction manager and start a transaction. This should create a session and
+    // check it out of the pool.
+    AsyncTransactionManager manager = client.transactionManagerAsync();
+    manager.beginAsync().get();
+    assertEquals(0, client.pool.numSessionsInPool());
+    assertEquals(1, client.pool.totalSessions());
+
+    // Rolling back an empty transaction will return the session to the pool.
+    manager.rollbackAsync().get();
+    assertEquals(1, client.pool.numSessionsInPool());
+    // Closing the transaction manager should not cause the session to be added to the pool again.
+    manager.close();
+    // The total number of sessions does not change.
+    assertEquals(1, client.pool.numSessionsInPool());
+
+    // Check out 2 sessions. Make sure that the pool really created a new session, and did not
+    // return the same session twice.
+    AsyncTransactionManager manager1 = client.transactionManagerAsync();
+    AsyncTransactionManager manager2 = client.transactionManagerAsync();
+    manager1.beginAsync().get();
+    manager2.beginAsync().get();
+    assertEquals(2, client.pool.totalSessions());
+    assertEquals(0, client.pool.numSessionsInPool());
+    manager1.close();
+    manager2.close();
+    assertEquals(2, client.pool.numSessionsInPool());
+  }
+
   private boolean isMultiplexedSessionsEnabled() {
     if (spanner.getOptions() == null || spanner.getOptions().getSessionPoolOptions() == null) {
       return false;
