fix: sessions were not always removed from checkedOutSessions (#1438)

Sessions were added to a set of checked out sessions when one was checked out from
the session pool. When the session was released back to the pool, the session would
normally be removed from the set of checked out sessions. The latter would not always
happen if the application that checked out the session did not use the session for
any reads or writes.

Co-authored-by: Thiago Nunes <thiagotnunes@google.com>

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();
+  }
 }