fix: sessions were not always removed from checkedOutSessions

google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java

@@ -1250,25 +1250,27 @@ public void prepareReadWriteTransaction() {
 
     @Override
     public void close() {
-      synchronized (lock) {
-        leakedException = null;
-        checkedOutSessions.remove(this);
-      }
-      PooledSession delegate = getOrNull();
-      if (delegate != null) {
-        delegate.close();
+      try {
+        asyncClose().get();
+      } catch (InterruptedException e) {
+        throw SpannerExceptionFactory.propagateInterrupt(e);
+      } catch (ExecutionException e) {
+        throw SpannerExceptionFactory.asSpannerException(e.getCause());
       }
     }
 
     @Override
     public ApiFuture<Empty> asyncClose() {
-      synchronized (lock) {
-        leakedException = null;
-        checkedOutSessions.remove(this);
-      }
-      PooledSession delegate = getOrNull();
-      if (delegate != null) {
-        return delegate.asyncClose();
+      try {
+        PooledSession delegate = getOrNull();
+        if (delegate != null) {
+          return delegate.asyncClose();
+        }
+      } finally {
+        synchronized (lock) {
+          leakedException = null;
+          checkedOutSessions.remove(this);
+        }
       }
       return ApiFutures.immediateFuture(Empty.getDefaultInstance());
     }
@@ -1777,7 +1779,8 @@ private enum Position {
   private final Set<PooledSession> allSessions = new HashSet<>();
 
   @GuardedBy("lock")
-  private final Set<PooledSessionFuture> checkedOutSessions = new HashSet<>();
+  @VisibleForTesting
+  final Set<PooledSessionFuture> checkedOutSessions = new HashSet<>();
 
   private final SessionConsumer sessionConsumer = new SessionConsumerImpl();
 

google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java

@@ -69,6 +69,7 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
+import java.util.Set;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
@@ -534,13 +535,18 @@ public void testAsyncTransactionManagerCommitWithTag() {
 
   @Test
   public void singleUse() {
-    DatabaseClient client =
-        spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE));
+    DatabaseClientImpl client =
+        (DatabaseClientImpl)
+            spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE));
+    Set<PooledSessionFuture> checkedOut = client.pool.checkedOutSessions;
+    assertThat(checkedOut).isEmpty();
     try (ResultSet rs = client.singleUse().executeQuery(SELECT1)) {
       assertThat(rs.next()).isTrue();
+      assertThat(checkedOut).hasSize(1);
       assertThat(rs.getLong(0)).isEqualTo(1L);
       assertThat(rs.next()).isFalse();
     }
+    assertThat(checkedOut).isEmpty();
   }
 
   @Test
@@ -2097,4 +2103,71 @@ public void testAsyncTransactionManagerCommitWithPriority() {
     assertNotNull(request.getRequestOptions());
     assertEquals(Priority.PRIORITY_HIGH, request.getRequestOptions().getPriority());
   }
+
+  @Test
+  public void singleUseNoAction_ClearsCheckedOutSession() {
+    DatabaseClientImpl client =
+        (DatabaseClientImpl)
+            spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE));
+    Set<PooledSessionFuture> checkedOut = client.pool.checkedOutSessions;
+    assertThat(checkedOut).isEmpty();
+
+    // Getting a single use read-only transaction and not using it should not cause any sessions
+    // to be stuck in the map of checked out sessions.
+    client.singleUse().close();
+
+    assertThat(checkedOut).isEmpty();
+  }
+
+  @Test
+  public void singleUseReadOnlyTransactionNoAction_ClearsCheckedOutSession() {
+    DatabaseClientImpl client =
+        (DatabaseClientImpl)
+            spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE));
+    Set<PooledSessionFuture> checkedOut = client.pool.checkedOutSessions;
+    assertThat(checkedOut).isEmpty();
+
+    client.singleUseReadOnlyTransaction().close();
+
+    assertThat(checkedOut).isEmpty();
+  }
+
+  @Test
+  public void readWriteTransactionNoAction_ClearsCheckedOutSession() {
+    DatabaseClientImpl client =
+        (DatabaseClientImpl)
+            spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE));
+    Set<PooledSessionFuture> checkedOut = client.pool.checkedOutSessions;
+    assertThat(checkedOut).isEmpty();
+
+    client.readWriteTransaction();
+
+    assertThat(checkedOut).isEmpty();
+  }
+
+  @Test
+  public void readOnlyTransactionNoAction_ClearsCheckedOutSession() {
+    DatabaseClientImpl client =
+        (DatabaseClientImpl)
+            spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE));
+    Set<PooledSessionFuture> checkedOut = client.pool.checkedOutSessions;
+    assertThat(checkedOut).isEmpty();
+
+    client.readOnlyTransaction().close();
+
+    assertThat(checkedOut).isEmpty();
+  }
+
+  @Test
+  public void transactionManagerNoAction_ClearsCheckedOutSession() {
+    DatabaseClientImpl client =
+        (DatabaseClientImpl)
+            spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE));
+    Set<PooledSessionFuture> checkedOut = client.pool.checkedOutSessions;
+    assertThat(checkedOut).isEmpty();
+
+    client.transactionManager().close();
+
+    assertThat(checkedOut).isEmpty();
+  }
 }