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 index fbfc472bf5..95cc167fe1 100644 --- 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 @@ -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 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 allSessions = new HashSet<>(); @GuardedBy("lock") - private final Set checkedOutSessions = new HashSet<>(); + @VisibleForTesting + final Set checkedOutSessions = new HashSet<>(); private final SessionConsumer sessionConsumer = new SessionConsumerImpl(); diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java index 0a1efc2a48..c2d6dc4cd2 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java +++ b/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 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 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 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 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 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 checkedOut = client.pool.checkedOutSessions; + assertThat(checkedOut).isEmpty(); + + client.transactionManager().close(); + + assertThat(checkedOut).isEmpty(); + } }