From 89a99b2fb569ffd0e04e726ece9b538ca4bdd5f3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= Date: Sat, 27 Sep 2025 16:12:36 +0200 Subject: [PATCH] chore: remove session pool implementation Removes the SessionPool and related classes and deprecates the SessionPoolOptions class. This simplifies large parts of the code, as the only possible code path is the use of multiplexed sessions. --- .github/sync-repo-settings.yaml | 5 - ...tegration-multiplexed-sessions-enabled.cfg | 48 - .../integration-regular-sessions-enabled.cfg | 48 - ...tractMultiplexedSessionDatabaseClient.java | 5 - .../cloud/spanner/DatabaseClientImpl.java | 191 +- .../DelayedMultiplexedSessionTransaction.java | 17 +- .../MultiplexedSessionDatabaseClient.java | 304 +- .../google/cloud/spanner/SessionClient.java | 11 +- .../com/google/cloud/spanner/SessionImpl.java | 5 + .../com/google/cloud/spanner/SessionPool.java | 3516 ----------------- .../SessionPoolAsyncTransactionManager.java | 287 -- .../cloud/spanner/SessionPoolOptions.java | 28 +- .../cloud/spanner/SessionReference.java | 10 +- .../com/google/cloud/spanner/SpannerImpl.java | 40 +- .../google/cloud/spanner/AsyncRunnerTest.java | 7 - .../spanner/AsyncTransactionManagerTest.java | 240 +- .../cloud/spanner/BackendExhaustedTest.java | 210 - .../cloud/spanner/BaseSessionPoolTest.java | 187 - .../spanner/BatchCreateSessionsTest.java | 237 -- .../cloud/spanner/ChannelUsageTest.java | 34 +- .../CloseSpannerWithOpenResultSetTest.java | 195 - .../cloud/spanner/DatabaseClientImplTest.java | 1602 +------- .../cloud/spanner/DefaultBenchmark.java | 10 - .../spanner/ITSessionPoolIntegrationTest.java | 176 - .../cloud/spanner/InlineBeginBenchmark.java | 14 +- .../IntegrationTestWithClosedSessionsEnv.java | 138 - .../spanner/LongRunningSessionsBenchmark.java | 328 -- .../cloud/spanner/MockSpannerServiceImpl.java | 24 +- ...edSessionDatabaseClientMockServerTest.java | 681 ---- .../spanner/MultiplexedSessionsBenchmark.java | 4 - .../google/cloud/spanner/ReadAsyncTest.java | 22 - .../RetryOnInvalidatedSessionTest.java | 1797 --------- .../spanner/RetryableInternalErrorTest.java | 8 +- .../cloud/spanner/SelectRandomBenchmark.java | 7 +- .../cloud/spanner/SessionPoolBenchmark.java | 265 -- .../cloud/spanner/SessionPoolLeakTest.java | 252 -- .../SessionPoolMaintainerBenchmark.java | 244 -- .../SessionPoolMaintainerMockServerTest.java | 181 - .../spanner/SessionPoolMaintainerTest.java | 398 -- .../cloud/spanner/SessionPoolStressTest.java | 296 -- .../google/cloud/spanner/SessionPoolTest.java | 2336 ----------- .../spanner/SessionPoolUnbalancedTest.java | 241 -- .../google/cloud/spanner/SpannerImplTest.java | 59 - .../spanner/TransactionRunnerImplTest.java | 2 +- .../spanner/connection/ConnectionTest.java | 84 - .../cloud/spanner/it/ITClosedSessionTest.java | 286 -- 46 files changed, 183 insertions(+), 14897 deletions(-) delete mode 100644 .kokoro/presubmit/integration-multiplexed-sessions-enabled.cfg delete mode 100644 .kokoro/presubmit/integration-regular-sessions-enabled.cfg delete mode 100644 google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java delete mode 100644 google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPoolAsyncTransactionManager.java delete mode 100644 google-cloud-spanner/src/test/java/com/google/cloud/spanner/BackendExhaustedTest.java delete mode 100644 google-cloud-spanner/src/test/java/com/google/cloud/spanner/BaseSessionPoolTest.java delete mode 100644 google-cloud-spanner/src/test/java/com/google/cloud/spanner/BatchCreateSessionsTest.java delete mode 100644 google-cloud-spanner/src/test/java/com/google/cloud/spanner/CloseSpannerWithOpenResultSetTest.java delete mode 100644 google-cloud-spanner/src/test/java/com/google/cloud/spanner/ITSessionPoolIntegrationTest.java delete mode 100644 google-cloud-spanner/src/test/java/com/google/cloud/spanner/IntegrationTestWithClosedSessionsEnv.java delete mode 100644 google-cloud-spanner/src/test/java/com/google/cloud/spanner/LongRunningSessionsBenchmark.java delete mode 100644 google-cloud-spanner/src/test/java/com/google/cloud/spanner/RetryOnInvalidatedSessionTest.java delete mode 100644 google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolBenchmark.java delete mode 100644 google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolLeakTest.java delete mode 100644 google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolMaintainerBenchmark.java delete mode 100644 google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolMaintainerMockServerTest.java delete mode 100644 google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolMaintainerTest.java delete mode 100644 google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolStressTest.java delete mode 100644 google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java delete mode 100644 google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolUnbalancedTest.java delete mode 100644 google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITClosedSessionTest.java diff --git a/.github/sync-repo-settings.yaml b/.github/sync-repo-settings.yaml index 39ea467f10b..6b48c7d1174 100644 --- a/.github/sync-repo-settings.yaml +++ b/.github/sync-repo-settings.yaml @@ -14,7 +14,6 @@ branchProtectionRules: - units (8) - units (11) - 'Kokoro - Test: Integration' - - 'Kokoro - Test: Integration with Multiplexed Sessions' - cla/google - checkstyle - compile (8) @@ -154,7 +153,6 @@ branchProtectionRules: - units (8) - units (11) - 'Kokoro - Test: Integration' - - 'Kokoro - Test: Integration with Multiplexed Sessions' - cla/google - checkstyle - compile (8) @@ -173,7 +171,6 @@ branchProtectionRules: - units (8) - units (11) - 'Kokoro - Test: Integration' - - 'Kokoro - Test: Integration with Multiplexed Sessions' - cla/google - checkstyle - compile (8) @@ -194,7 +191,6 @@ branchProtectionRules: - units (8) - units (11) - 'Kokoro - Test: Integration' - - 'Kokoro - Test: Integration with Multiplexed Sessions' - cla/google - checkstyle - compile (8) @@ -215,7 +211,6 @@ branchProtectionRules: - units (8) - units (11) - 'Kokoro - Test: Integration' - - 'Kokoro - Test: Integration with Multiplexed Sessions' - cla/google - checkstyle - compile (8) diff --git a/.kokoro/presubmit/integration-multiplexed-sessions-enabled.cfg b/.kokoro/presubmit/integration-multiplexed-sessions-enabled.cfg deleted file mode 100644 index 800e2a21558..00000000000 --- a/.kokoro/presubmit/integration-multiplexed-sessions-enabled.cfg +++ /dev/null @@ -1,48 +0,0 @@ -# Format: //devtools/kokoro/config/proto/build.proto - -# Configure the docker image for kokoro-trampoline. -env_vars: { - key: "TRAMPOLINE_IMAGE" - value: "gcr.io/cloud-devrel-kokoro-resources/java8" -} - -env_vars: { - key: "JOB_TYPE" - value: "integration-multiplexed-sessions-enabled" -} - -# TODO: remove this after we've migrated all tests and scripts -env_vars: { - key: "GCLOUD_PROJECT" - value: "gcloud-devel" -} - -env_vars: { - key: "GOOGLE_CLOUD_PROJECT" - value: "gcloud-devel" -} - -env_vars: { - key: "GOOGLE_APPLICATION_CREDENTIALS" - value: "secret_manager/java-it-service-account" -} - -env_vars: { - key: "SECRET_MANAGER_KEYS" - value: "java-it-service-account" -} - -env_vars: { - key: "GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS" - value: "true" -} - -env_vars: { - key: "GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_PARTITIONED_OPS" - value: "true" -} - -env_vars: { - key: "GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_FOR_RW" - value: "true" -} diff --git a/.kokoro/presubmit/integration-regular-sessions-enabled.cfg b/.kokoro/presubmit/integration-regular-sessions-enabled.cfg deleted file mode 100644 index c2ba18efac0..00000000000 --- a/.kokoro/presubmit/integration-regular-sessions-enabled.cfg +++ /dev/null @@ -1,48 +0,0 @@ -# Format: //devtools/kokoro/config/proto/build.proto - -# Configure the docker image for kokoro-trampoline. -env_vars: { - key: "TRAMPOLINE_IMAGE" - value: "gcr.io/cloud-devrel-kokoro-resources/java8" -} - -env_vars: { - key: "JOB_TYPE" - value: "integration" -} - -# TODO: remove this after we've migrated all tests and scripts -env_vars: { - key: "GCLOUD_PROJECT" - value: "gcloud-devel" -} - -env_vars: { - key: "GOOGLE_CLOUD_PROJECT" - value: "gcloud-devel" -} - -env_vars: { - key: "GOOGLE_APPLICATION_CREDENTIALS" - value: "secret_manager/java-it-service-account" -} - -env_vars: { - key: "SECRET_MANAGER_KEYS" - value: "java-it-service-account" -} - -env_vars: { - key: "GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS" - value: "false" -} - -env_vars: { - key: "GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_PARTITIONED_OPS" - value: "false" -} - -env_vars: { - key: "GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_FOR_RW" - value: "false" -} diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractMultiplexedSessionDatabaseClient.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractMultiplexedSessionDatabaseClient.java index f9b04136ecb..7d083db211a 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractMultiplexedSessionDatabaseClient.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractMultiplexedSessionDatabaseClient.java @@ -26,11 +26,6 @@ */ abstract class AbstractMultiplexedSessionDatabaseClient implements DatabaseClient { - @Override - public String getDatabaseRole() { - throw new UnsupportedOperationException(); - } - @Override public Timestamp writeAtLeastOnce(Iterable mutations) throws SpannerException { return writeAtLeastOnceWithOptions(mutations).getCommitTimestamp(); diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClientImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClientImpl.java index 40dbd710bd5..a8a8e8edf7c 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClientImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClientImpl.java @@ -20,10 +20,10 @@ import com.google.cloud.Timestamp; import com.google.cloud.spanner.Options.TransactionOption; import com.google.cloud.spanner.Options.UpdateOption; -import com.google.cloud.spanner.SessionPool.PooledSessionFuture; import com.google.cloud.spanner.SpannerImpl.ClosedException; import com.google.cloud.spanner.Statement.StatementFactory; import com.google.common.annotations.VisibleForTesting; +import com.google.common.util.concurrent.Futures; import com.google.common.util.concurrent.ListenableFuture; import com.google.spanner.v1.BatchWriteResponse; import io.opentelemetry.api.common.Attributes; @@ -34,7 +34,6 @@ import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; import java.util.concurrent.atomic.AtomicInteger; -import java.util.function.BiFunction; import javax.annotation.Nullable; class DatabaseClientImpl implements DatabaseClient { @@ -42,60 +41,21 @@ class DatabaseClientImpl implements DatabaseClient { private static final String READ_ONLY_TRANSACTION = "CloudSpanner.ReadOnlyTransaction"; private static final String PARTITION_DML_TRANSACTION = "CloudSpanner.PartitionDMLTransaction"; private final TraceWrapper tracer; - private Attributes commonAttributes; + private final Attributes commonAttributes; @VisibleForTesting final String clientId; - @VisibleForTesting final SessionPool pool; @VisibleForTesting final MultiplexedSessionDatabaseClient multiplexedSessionDatabaseClient; - @VisibleForTesting final boolean useMultiplexedSessionPartitionedOps; - @VisibleForTesting final boolean useMultiplexedSessionForRW; @VisibleForTesting final int dbId; private final AtomicInteger nthRequest; private final Map clientIdToOrdinalMap; - final boolean useMultiplexedSessionBlindWrite; - - @VisibleForTesting - DatabaseClientImpl(SessionPool pool, TraceWrapper tracer) { - this( - "", - pool, - /* useMultiplexedSessionBlindWrite= */ false, - /* multiplexedSessionDatabaseClient= */ null, - /* useMultiplexedSessionPartitionedOps= */ false, - tracer, - /* useMultiplexedSessionForRW= */ false, - Attributes.empty()); - } - - @VisibleForTesting - DatabaseClientImpl(String clientId, SessionPool pool, TraceWrapper tracer) { - this( - clientId, - pool, - /* useMultiplexedSessionBlindWrite= */ false, - /* multiplexedSessionDatabaseClient= */ null, - /* useMultiplexedSessionPartitionedOps= */ false, - tracer, - /* useMultiplexedSessionForRW= */ false, - Attributes.empty()); - } - DatabaseClientImpl( String clientId, - SessionPool pool, - boolean useMultiplexedSessionBlindWrite, - @Nullable MultiplexedSessionDatabaseClient multiplexedSessionDatabaseClient, - boolean useMultiplexedSessionPartitionedOps, + MultiplexedSessionDatabaseClient multiplexedSessionDatabaseClient, TraceWrapper tracer, - boolean useMultiplexedSessionForRW, Attributes commonAttributes) { this.clientId = clientId; - this.pool = pool; - this.useMultiplexedSessionBlindWrite = useMultiplexedSessionBlindWrite; this.multiplexedSessionDatabaseClient = multiplexedSessionDatabaseClient; - this.useMultiplexedSessionPartitionedOps = useMultiplexedSessionPartitionedOps; this.tracer = tracer; - this.useMultiplexedSessionForRW = useMultiplexedSessionForRW; this.commonAttributes = commonAttributes; this.clientIdToOrdinalMap = new HashMap(); @@ -113,55 +73,14 @@ synchronized int dbIdFromClientId(String clientId) { return id; } - @VisibleForTesting - PooledSessionFuture getSession() { - return pool.getSession(); - } - @VisibleForTesting DatabaseClient getMultiplexedSession() { - if (canUseMultiplexedSessions()) { - return this.multiplexedSessionDatabaseClient; - } - return pool.getMultiplexedSessionWithFallback(); - } - - @VisibleForTesting - DatabaseClient getMultiplexedSessionForRW() { - if (canUseMultiplexedSessionsForRW()) { - return getMultiplexedSession(); - } - return getSession(); - } - - private MultiplexedSessionDatabaseClient getMultiplexedSessionDatabaseClient() { - return canUseMultiplexedSessions() ? this.multiplexedSessionDatabaseClient : null; - } - - private boolean canUseMultiplexedSessions() { - return this.multiplexedSessionDatabaseClient != null - && this.multiplexedSessionDatabaseClient.isMultiplexedSessionsSupported(); - } - - private boolean canUseMultiplexedSessionsForRW() { - return this.useMultiplexedSessionForRW - && this.multiplexedSessionDatabaseClient != null - && this.multiplexedSessionDatabaseClient.isMultiplexedSessionsForRWSupported(); - } - - private boolean canUseMultiplexedSessionsForPartitionedOps() { - return this.useMultiplexedSessionPartitionedOps - && this.multiplexedSessionDatabaseClient != null - && this.multiplexedSessionDatabaseClient.isMultiplexedSessionsForPartitionedOpsSupported(); + return this.multiplexedSessionDatabaseClient; } @Override public Dialect getDialect() { - MultiplexedSessionDatabaseClient client = getMultiplexedSessionDatabaseClient(); - if (client != null) { - return client.getDialect(); - } - return pool.getDialect(); + return this.multiplexedSessionDatabaseClient.getDialect(); } private final AbstractLazyInitializer statementFactorySupplier = @@ -191,7 +110,7 @@ public StatementFactory getStatementFactory() { @Override @Nullable public String getDatabaseRole() { - return pool.getDatabaseRole(); + return multiplexedSessionDatabaseClient.getDatabaseRole(); } @Override @@ -205,14 +124,7 @@ public CommitResponse writeWithOptions( throws SpannerException { ISpan span = tracer.spanBuilder(READ_WRITE_TRANSACTION, commonAttributes, options); try (IScope s = tracer.withSpan(span)) { - if (canUseMultiplexedSessionsForRW() && getMultiplexedSessionDatabaseClient() != null) { - return getMultiplexedSessionDatabaseClient().writeWithOptions(mutations, options); - } - - return runWithSessionRetry( - (session, reqId) -> { - return session.writeWithOptions(mutations, withReqId(reqId, options)); - }); + return multiplexedSessionDatabaseClient.writeWithOptions(mutations, options); } catch (RuntimeException e) { span.setStatus(e); throw e; @@ -232,13 +144,7 @@ public CommitResponse writeAtLeastOnceWithOptions( throws SpannerException { ISpan span = tracer.spanBuilder(READ_WRITE_TRANSACTION, commonAttributes, options); try (IScope s = tracer.withSpan(span)) { - if (useMultiplexedSessionBlindWrite && getMultiplexedSessionDatabaseClient() != null) { - return getMultiplexedSessionDatabaseClient() - .writeAtLeastOnceWithOptions(mutations, options); - } - return runWithSessionRetry( - (session, reqId) -> - session.writeAtLeastOnceWithOptions(mutations, withReqId(reqId, options))); + return multiplexedSessionDatabaseClient.writeAtLeastOnceWithOptions(mutations, options); } catch (RuntimeException e) { span.setStatus(e); throw e; @@ -262,12 +168,7 @@ public ServerStream batchWriteAtLeastOnce( throws SpannerException { ISpan span = tracer.spanBuilder(READ_WRITE_TRANSACTION, commonAttributes, options); try (IScope s = tracer.withSpan(span)) { - if (canUseMultiplexedSessionsForRW() && getMultiplexedSessionDatabaseClient() != null) { - return getMultiplexedSessionDatabaseClient().batchWriteAtLeastOnce(mutationGroups, options); - } - return runWithSessionRetry( - (session, reqId) -> - session.batchWriteAtLeastOnce(mutationGroups, withReqId(reqId, options))); + return multiplexedSessionDatabaseClient.batchWriteAtLeastOnce(mutationGroups, options); } catch (RuntimeException e) { span.setStatus(e); throw e; @@ -352,7 +253,7 @@ public ReadOnlyTransaction readOnlyTransaction(TimestampBound bound) { public TransactionRunner readWriteTransaction(TransactionOption... options) { ISpan span = tracer.spanBuilder(READ_WRITE_TRANSACTION, commonAttributes, options); try (IScope s = tracer.withSpan(span)) { - return getMultiplexedSessionForRW().readWriteTransaction(options); + return multiplexedSessionDatabaseClient.readWriteTransaction(options); } catch (RuntimeException e) { span.setStatus(e); span.end(); @@ -364,7 +265,7 @@ public TransactionRunner readWriteTransaction(TransactionOption... options) { public TransactionManager transactionManager(TransactionOption... options) { ISpan span = tracer.spanBuilder(READ_WRITE_TRANSACTION, commonAttributes, options); try (IScope s = tracer.withSpan(span)) { - return getMultiplexedSessionForRW().transactionManager(options); + return multiplexedSessionDatabaseClient.transactionManager(options); } catch (RuntimeException e) { span.setStatus(e); span.end(); @@ -376,7 +277,7 @@ public TransactionManager transactionManager(TransactionOption... options) { public AsyncRunner runAsync(TransactionOption... options) { ISpan span = tracer.spanBuilder(READ_WRITE_TRANSACTION, commonAttributes, options); try (IScope s = tracer.withSpan(span)) { - return getMultiplexedSessionForRW().runAsync(options); + return multiplexedSessionDatabaseClient.runAsync(options); } catch (RuntimeException e) { span.setStatus(e); span.end(); @@ -388,7 +289,7 @@ public AsyncRunner runAsync(TransactionOption... options) { public AsyncTransactionManager transactionManagerAsync(TransactionOption... options) { ISpan span = tracer.spanBuilder(READ_WRITE_TRANSACTION, commonAttributes, options); try (IScope s = tracer.withSpan(span)) { - return getMultiplexedSessionForRW().transactionManagerAsync(options); + return multiplexedSessionDatabaseClient.transactionManagerAsync(options); } catch (RuntimeException e) { span.setStatus(e); span.end(); @@ -398,25 +299,11 @@ public AsyncTransactionManager transactionManagerAsync(TransactionOption... opti @Override public long executePartitionedUpdate(final Statement stmt, final UpdateOption... options) { - - if (canUseMultiplexedSessionsForPartitionedOps()) { - try { - return getMultiplexedSession().executePartitionedUpdate(stmt, options); - } catch (SpannerException e) { - if (!multiplexedSessionDatabaseClient.maybeMarkUnimplementedForPartitionedOps(e)) { - throw e; - } - } - } - return executePartitionedUpdateWithPooledSession(stmt, options); + return multiplexedSessionDatabaseClient.executePartitionedUpdate(stmt, options); } private Future getDialectAsync() { - MultiplexedSessionDatabaseClient client = getMultiplexedSessionDatabaseClient(); - if (client != null) { - return client.getDialectAsync(); - } - return pool.getDialectAsync(); + return multiplexedSessionDatabaseClient.getDialectAsync(); } private UpdateOption[] withReqId( @@ -447,53 +334,13 @@ private TransactionOption[] withReqId( return allOptions; } - private long executePartitionedUpdateWithPooledSession( - final Statement stmt, final UpdateOption... options) { - ISpan span = tracer.spanBuilder(PARTITION_DML_TRANSACTION, commonAttributes); - try (IScope s = tracer.withSpan(span)) { - return runWithSessionRetry( - (session, reqId) -> { - return session.executePartitionedUpdate(stmt, withReqId(reqId, options)); - }); - } catch (RuntimeException e) { - span.setStatus(e); - span.end(); - throw e; - } - } - - @VisibleForTesting - T runWithSessionRetry(BiFunction callable) { - PooledSessionFuture session = getSession(); - XGoogSpannerRequestId reqId = - XGoogSpannerRequestId.of( - this.dbId, Long.valueOf(session.getChannel()), this.nextNthRequest(), 1); - while (true) { - try { - return callable.apply(session, reqId); - } catch (SessionNotFoundException e) { - session = - (PooledSessionFuture) - pool.getPooledSessionReplacementHandler().replaceSession(e, session); - reqId = - XGoogSpannerRequestId.of( - this.dbId, Long.valueOf(session.getChannel()), this.nextNthRequest(), 1); - } - } - } - boolean isValid() { - return pool.isValid() - && (multiplexedSessionDatabaseClient == null - || multiplexedSessionDatabaseClient.isValid() - || !multiplexedSessionDatabaseClient.isMultiplexedSessionsSupported()); + return multiplexedSessionDatabaseClient.isValid(); } ListenableFuture closeAsync(ClosedException closedException) { - if (this.multiplexedSessionDatabaseClient != null) { - // This method is non-blocking. - this.multiplexedSessionDatabaseClient.close(); - } - return pool.closeAsync(closedException); + // This method is non-blocking. + this.multiplexedSessionDatabaseClient.close(); + return Futures.immediateFuture(null); } } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DelayedMultiplexedSessionTransaction.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DelayedMultiplexedSessionTransaction.java index debb07d7af4..81e29cfda48 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DelayedMultiplexedSessionTransaction.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DelayedMultiplexedSessionTransaction.java @@ -44,17 +44,19 @@ class DelayedMultiplexedSessionTransaction extends AbstractMultiplexedSessionDat private final ISpan span; private final ApiFuture sessionFuture; - private final SessionPool sessionPool; DelayedMultiplexedSessionTransaction( MultiplexedSessionDatabaseClient client, ISpan span, - ApiFuture sessionFuture, - SessionPool sessionPool) { + ApiFuture sessionFuture) { this.client = client; this.span = span; this.sessionFuture = sessionFuture; - this.sessionPool = sessionPool; + } + + @Override + public String getDatabaseRole() { + return this.client.getDatabaseRole(); } @Override @@ -192,12 +194,7 @@ public TransactionRunner readWriteTransaction(TransactionOption... options) { this.sessionFuture, sessionReference -> new MultiplexedSessionTransaction( - client, - span, - sessionReference, - NO_CHANNEL_HINT, - /* singleUse= */ false, - this.sessionPool) + client, span, sessionReference, NO_CHANNEL_HINT, /* singleUse= */ false) .readWriteTransaction(options), MoreExecutors.directExecutor())); } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java index ef1f4a88945..06723c7a9cb 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java @@ -17,7 +17,6 @@ package com.google.cloud.spanner; import static com.google.cloud.spanner.SessionImpl.NO_CHANNEL_HINT; -import static com.google.cloud.spanner.SpannerExceptionFactory.newSpannerException; import com.google.api.core.ApiFuture; import com.google.api.core.ApiFutures; @@ -27,15 +26,10 @@ import com.google.cloud.spanner.Options.TransactionOption; import com.google.cloud.spanner.Options.UpdateOption; import com.google.cloud.spanner.SessionClient.SessionConsumer; -import com.google.cloud.spanner.SessionPool.SessionPoolTransactionRunner; import com.google.cloud.spanner.SpannerException.ResourceNotFoundException; import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Preconditions; -import com.google.common.util.concurrent.MoreExecutors; import com.google.spanner.v1.BatchWriteResponse; -import com.google.spanner.v1.BeginTransactionRequest; -import com.google.spanner.v1.RequestOptions; -import com.google.spanner.v1.Transaction; import java.time.Clock; import java.time.Duration; import java.time.Instant; @@ -49,99 +43,22 @@ import java.util.concurrent.ScheduledFuture; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; -import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicLong; import java.util.concurrent.atomic.AtomicReference; -/** - * {@link TransactionRunner} that automatically handles "UNIMPLEMENTED" errors with the message - * "Transaction type read_write not supported with multiplexed sessions" by switching from a - * multiplexed session to a regular session and then restarts the transaction. - */ -class MultiplexedSessionTransactionRunner implements TransactionRunner { - private final SessionPool sessionPool; - private final TransactionRunnerImpl transactionRunnerForMultiplexedSession; - private SessionPoolTransactionRunner transactionRunnerForRegularSession; - private final TransactionOption[] options; - private boolean isUsingMultiplexedSession = true; - - public MultiplexedSessionTransactionRunner( - SessionImpl multiplexedSession, SessionPool sessionPool, TransactionOption... options) { - this.sessionPool = sessionPool; - this.transactionRunnerForMultiplexedSession = - new TransactionRunnerImpl( - multiplexedSession, options); // Uses multiplexed session initially - multiplexedSession.setActive(this.transactionRunnerForMultiplexedSession); - this.options = options; - } - - private TransactionRunner getRunner() { - if (this.isUsingMultiplexedSession) { - return this.transactionRunnerForMultiplexedSession; - } else { - if (this.transactionRunnerForRegularSession == null) { - this.transactionRunnerForRegularSession = - new SessionPoolTransactionRunner<>( - sessionPool.getSession(), - sessionPool.getPooledSessionReplacementHandler(), - options); - } - return this.transactionRunnerForRegularSession; - } - } - - @Override - public T run(TransactionCallable callable) { - while (true) { - try { - return getRunner().run(callable); - } catch (SpannerException e) { - if (e.getErrorCode() == ErrorCode.UNIMPLEMENTED - && verifyUnimplementedErrorMessageForRWMux(e)) { - this.isUsingMultiplexedSession = false; // Fallback to regular session - } else { - throw e; // Other errors propagate - } - } - } - } - - @Override - public Timestamp getCommitTimestamp() { - return getRunner().getCommitTimestamp(); - } - - @Override - public CommitResponse getCommitResponse() { - return getRunner().getCommitResponse(); - } - - @Override - public TransactionRunner allowNestedTransaction() { - getRunner().allowNestedTransaction(); - return this; - } - - private boolean verifyUnimplementedErrorMessageForRWMux(SpannerException spannerException) { - if (spannerException.getCause() == null) { - return false; - } - if (spannerException.getCause().getMessage() == null) { - return false; - } - return spannerException - .getCause() - .getMessage() - .contains("Transaction type read_write not supported with multiplexed sessions"); - } -} - /** * {@link DatabaseClient} implementation that uses a single multiplexed session to execute * transactions. */ final class MultiplexedSessionDatabaseClient extends AbstractMultiplexedSessionDatabaseClient { + @VisibleForTesting + static final Statement DETERMINE_DIALECT_STATEMENT = + Statement.newBuilder( + "select option_value " + + "from information_schema.database_options " + + "where option_name='database_dialect'") + .build(); /** * Represents a single transaction on a multiplexed session. This can be both a single-use or @@ -160,7 +77,6 @@ static class MultiplexedSessionTransaction extends SessionImpl { private final int singleUseChannelHint; private boolean done; - private final SessionPool pool; MultiplexedSessionTransaction( MultiplexedSessionDatabaseClient client, @@ -168,22 +84,11 @@ static class MultiplexedSessionTransaction extends SessionImpl { SessionReference sessionReference, int singleUseChannelHint, boolean singleUse) { - this(client, span, sessionReference, singleUseChannelHint, singleUse, null); - } - - MultiplexedSessionTransaction( - MultiplexedSessionDatabaseClient client, - ISpan span, - SessionReference sessionReference, - int singleUseChannelHint, - boolean singleUse, - SessionPool pool) { super(client.sessionClient.getSpanner(), sessionReference, singleUseChannelHint); this.client = client; this.singleUse = singleUse; this.singleUseChannelHint = singleUseChannelHint; this.client.numSessionsAcquired.incrementAndGet(); - this.pool = pool; setCurrentSpan(span); } @@ -197,15 +102,6 @@ void onError(SpannerException spannerException) { // synchronizing, as it does not really matter exactly which error is set. this.client.resourceNotFoundException.set((ResourceNotFoundException) spannerException); } - // Mark multiplexed sessions for RW as unimplemented and fall back to regular sessions if - // UNIMPLEMENTED with error message "Transaction type read_write not supported with - // multiplexed sessions" is returned. - this.client.maybeMarkUnimplementedForRW(spannerException); - // Mark multiplexed sessions for Partitioned Ops as unimplemented and fall back to regular - // sessions if - // UNIMPLEMENTED with error message "Partitioned operations are not supported with multiplexed - // sessions". - this.client.maybeMarkUnimplementedForPartitionedOps(spannerException); } @Override @@ -231,11 +127,6 @@ public CommitResponse writeAtLeastOnceWithOptions( return response; } - @Override - public TransactionRunner readWriteTransaction(TransactionOption... options) { - return new MultiplexedSessionTransactionRunner(this, pool, options); - } - @Override void onTransactionDone() { boolean markedDone = false; @@ -283,12 +174,6 @@ public void close() { /** The current multiplexed session that is used by this client. */ private final AtomicReference> multiplexedSessionReference; - /** - * The Transaction response returned by the BeginTransaction request with read-write when a - * multiplexed session is created during client initialization. - */ - private final SettableApiFuture readWriteBeginTransactionReferenceFuture; - /** The expiration date/time of the current multiplexed session. */ private final AtomicReference expirationDate; @@ -309,26 +194,6 @@ public void close() { private final AtomicLong numSessionsReleased = new AtomicLong(); - /** - * This flag is set to true if the server return UNIMPLEMENTED when we try to create a multiplexed - * session. TODO: Remove once this is guaranteed to be available. - */ - private final AtomicBoolean unimplemented = new AtomicBoolean(false); - - /** - * This flag is set to true if the server return UNIMPLEMENTED when a read-write transaction is - * executed on a multiplexed session. TODO: Remove once this is guaranteed to be available. - */ - @VisibleForTesting final AtomicBoolean unimplementedForRW = new AtomicBoolean(false); - - /** - * This flag is set to true if the server return UNIMPLEMENTED when partitioned transaction is - * executed on a multiplexed session. TODO: Remove once this is guaranteed to be available. - */ - @VisibleForTesting final AtomicBoolean unimplementedForPartitionedOps = new AtomicBoolean(false); - - private SessionPool pool; - MultiplexedSessionDatabaseClient(SessionClient sessionClient) { this(sessionClient, Clock.systemUTC()); } @@ -356,7 +221,6 @@ public void close() { this.tracer = sessionClient.getSpanner().getTracer(); final SettableApiFuture initialSessionReferenceFuture = SettableApiFuture.create(); - this.readWriteBeginTransactionReferenceFuture = SettableApiFuture.create(); this.multiplexedSessionReference = new AtomicReference<>(initialSessionReferenceFuture); this.sessionClient.asyncCreateMultiplexedSession( new SessionConsumer() { @@ -366,21 +230,6 @@ public void onSessionReady(SessionImpl session) { // only start the maintainer if we actually managed to create a session in the first // place. maintainer.start(); - - // initiate a begin transaction request to verify if read-write transactions are - // supported using multiplexed sessions. - if (sessionClient - .getSpanner() - .getOptions() - .getSessionPoolOptions() - .getUseMultiplexedSessionForRW() - && !sessionClient - .getSpanner() - .getOptions() - .getSessionPoolOptions() - .getSkipVerifyBeginTransactionForMuxRW()) { - verifyBeginTransactionWithRWOnMultiplexedSessionAsync(session.getName()); - } if (sessionClient .getSpanner() .getOptions() @@ -392,9 +241,17 @@ public void onSessionReady(SessionImpl session) { @Override public void onSessionCreateFailure(Throwable t, int createFailureForSessionCount) { - // Mark multiplexes sessions as unimplemented and fall back to regular sessions if - // UNIMPLEMENTED is returned. - maybeMarkUnimplemented(t); + SpannerException spannerException = SpannerExceptionFactory.asSpannerException(t); + if (MultiplexedSessionDatabaseClient.this.resourceNotFoundException.get() == null + && (spannerException instanceof DatabaseNotFoundException + || spannerException instanceof InstanceNotFoundException + || spannerException instanceof SessionNotFoundException)) { + // This could in theory set this field more than once, but we don't want to bother + // with + // synchronizing, as it does not really matter exactly which error is set. + MultiplexedSessionDatabaseClient.this.resourceNotFoundException.set( + (ResourceNotFoundException) spannerException); + } initialSessionReferenceFuture.setException(t); } }); @@ -403,10 +260,6 @@ public void onSessionCreateFailure(Throwable t, int createFailureForSessionCount initialSessionReferenceFuture); } - void setPool(SessionPool pool) { - this.pool = pool; - } - private static void maybeWaitForSessionCreation( SessionPoolOptions sessionPoolOptions, ApiFuture future) { Duration waitDuration = sessionPoolOptions.getWaitForMinSessions(); @@ -426,88 +279,6 @@ private static void maybeWaitForSessionCreation( } } - private void maybeMarkUnimplemented(Throwable t) { - SpannerException spannerException = SpannerExceptionFactory.asSpannerException(t); - if (spannerException.getErrorCode() == ErrorCode.UNIMPLEMENTED) { - unimplemented.set(true); - } - } - - private void maybeMarkUnimplementedForRW(SpannerException spannerException) { - if (spannerException.getErrorCode() == ErrorCode.UNIMPLEMENTED - && verifyErrorMessage( - spannerException, - "Transaction type read_write not supported with multiplexed sessions")) { - unimplementedForRW.set(true); - } - } - - boolean maybeMarkUnimplementedForPartitionedOps(SpannerException spannerException) { - if (spannerException.getErrorCode() == ErrorCode.UNIMPLEMENTED - && verifyErrorMessage( - spannerException, - "Transaction type partitioned_dml not supported with multiplexed sessions")) { - unimplementedForPartitionedOps.set(true); - return true; - } - return false; - } - - static boolean verifyErrorMessage(SpannerException spannerException, String message) { - if (spannerException.getCause() == null) { - return false; - } - if (spannerException.getCause().getMessage() == null) { - return false; - } - return spannerException.getCause().getMessage().contains(message); - } - - private void verifyBeginTransactionWithRWOnMultiplexedSessionAsync(String sessionName) { - // TODO: Remove once this is guaranteed to be available. - // annotate the explict BeginTransactionRequest with a transaction tag - // "multiplexed-rw-background-begin-txn" to avoid storing this request on mock spanner. - // this is to safeguard other mock spanner tests whose BeginTransaction request count will - // otherwise increase by 1. Modifying the unit tests do not seem valid since this code is - // temporary and will be removed once the read-write on multiplexed session looks stable at - // backend. - BeginTransactionRequest.Builder requestBuilder = - BeginTransactionRequest.newBuilder() - .setSession(sessionName) - .setOptions( - SessionImpl.createReadWriteTransactionOptions( - Options.fromTransactionOptions(), /* previousTransactionId= */ null)) - .setRequestOptions( - RequestOptions.newBuilder() - .setTransactionTag("multiplexed-rw-background-begin-txn") - .build()); - final BeginTransactionRequest request = requestBuilder.build(); - final ApiFuture requestFuture; - requestFuture = - sessionClient - .getSpanner() - .getRpc() - .beginTransactionAsync(request, /* options= */ null, /* routeToLeader= */ true); - requestFuture.addListener( - () -> { - try { - Transaction txn = requestFuture.get(); - if (txn.getId().isEmpty()) { - throw newSpannerException( - ErrorCode.INTERNAL, "Missing id in transaction\n" + sessionName); - } - readWriteBeginTransactionReferenceFuture.set(txn); - } catch (Exception e) { - SpannerException spannerException = SpannerExceptionFactory.newSpannerException(e); - // Mark multiplexed sessions for RW as unimplemented and fall back to regular sessions - // if UNIMPLEMENTED is returned. - maybeMarkUnimplementedForRW(spannerException); - readWriteBeginTransactionReferenceFuture.setException(e); - } - }, - MoreExecutors.directExecutor()); - } - boolean isValid() { return resourceNotFoundException.get() == null; } @@ -520,18 +291,6 @@ AtomicLong getNumSessionsReleased() { return this.numSessionsReleased; } - boolean isMultiplexedSessionsSupported() { - return !this.unimplemented.get(); - } - - boolean isMultiplexedSessionsForRWSupported() { - return !this.unimplementedForRW.get(); - } - - boolean isMultiplexedSessionsForPartitionedOpsSupported() { - return !this.unimplementedForPartitionedOps.get(); - } - void close() { synchronized (this) { if (!this.isClosed) { @@ -557,17 +316,6 @@ SessionReference getCurrentSessionReference() { } } - @VisibleForTesting - Transaction getReadWriteBeginTransactionReference() { - try { - return this.readWriteBeginTransactionReferenceFuture.get(); - } catch (ExecutionException executionException) { - throw SpannerExceptionFactory.asSpannerException(executionException.getCause()); - } catch (InterruptedException interruptedException) { - throw SpannerExceptionFactory.propagateInterrupt(interruptedException); - } - } - /** * Returns true if the multiplexed session has been created. This client can be used before the * session has been created, and will in that case use a delayed transaction that contains a @@ -597,8 +345,7 @@ private MultiplexedSessionTransaction createDirectMultiplexedSessionTransaction( // any special handling of such errors. multiplexedSessionReference.get().get(), singleUse ? getSingleUseChannelHint() : NO_CHANNEL_HINT, - singleUse, - this.pool); + singleUse); } catch (ExecutionException executionException) { throw SpannerExceptionFactory.asSpannerException(executionException.getCause()); } catch (InterruptedException interruptedException) { @@ -608,7 +355,7 @@ private MultiplexedSessionTransaction createDirectMultiplexedSessionTransaction( private DelayedMultiplexedSessionTransaction createDelayedMultiplexSessionTransaction() { return new DelayedMultiplexedSessionTransaction( - this, tracer.getCurrentSpan(), multiplexedSessionReference.get(), this.pool); + this, tracer.getCurrentSpan(), multiplexedSessionReference.get()); } private int getSingleUseChannelHint() { @@ -633,8 +380,7 @@ private int getSingleUseChannelHint() { new AbstractLazyInitializer() { @Override protected Dialect initialize() { - try (ResultSet dialectResultSet = - singleUse().executeQuery(SessionPool.DETERMINE_DIALECT_STATEMENT)) { + try (ResultSet dialectResultSet = singleUse().executeQuery(DETERMINE_DIALECT_STATEMENT)) { if (dialectResultSet.next()) { return Dialect.fromName(dialectResultSet.getString(0)); } @@ -661,6 +407,11 @@ Future getDialectAsync() { } } + @Override + public String getDatabaseRole() { + return this.sessionClient.getSpanner().getOptions().getDatabaseRole(); + } + @Override public Timestamp write(Iterable mutations) throws SpannerException { return createMultiplexedSessionTransaction(/* singleUse= */ false).write(mutations); @@ -811,9 +562,6 @@ public void onSessionCreateFailure(Throwable t, int createFailureForSessionCount // ignore any errors during re-creation of the multiplexed session. This means that // we continue to use the session that has passed its expiration date for now, and // that a new attempt at creating a new session will be done in 10 minutes from now. - // The only exception to this rule is if the server returns UNIMPLEMENTED. In that - // case we invalidate the client and fall back to regular sessions. - maybeMarkUnimplemented(t); } }); } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionClient.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionClient.java index 20c86bdf25b..5823f87e9f9 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionClient.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionClient.java @@ -249,6 +249,7 @@ SessionImpl createSession() { SessionReference sessionReference = new SessionReference( session.getName(), + spanner.getOptions().getDatabaseRole(), session.getCreateTime(), session.getMultiplexed(), optionMap(SessionOption.channelHint(channelId))); @@ -307,7 +308,11 @@ SessionImpl createMultiplexedSession() { new SessionImpl( spanner, new SessionReference( - session.getName(), session.getCreateTime(), session.getMultiplexed(), null)); + session.getName(), + spanner.getOptions().getDatabaseRole(), + session.getCreateTime(), + session.getMultiplexed(), + null)); sessionImpl.setRequestIdCreator(this); span.addAnnotation( String.format("Request for %d multiplexed session returned %d session", 1, 1)); @@ -444,6 +449,7 @@ private List internalBatchCreateSessions( spanner, new SessionReference( session.getName(), + spanner.getOptions().getDatabaseRole(), session.getCreateTime(), session.getMultiplexed(), optionMap(SessionOption.channelHint(channelHint)))); @@ -464,7 +470,8 @@ SessionImpl sessionWithId(String name) { synchronized (this) { options = optionMap(SessionOption.channelHint(sessionChannelCounter++)); } - SessionImpl sessionImpl = new SessionImpl(spanner, new SessionReference(name, options)); + SessionImpl sessionImpl = + new SessionImpl(spanner, new SessionReference(name, /* databaseRole= */ null, options)); sessionImpl.setRequestIdCreator(this); return sessionImpl; } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java index 68c37561c9d..e9a68bfcee6 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java @@ -173,6 +173,11 @@ public String getName() { return sessionReference.getName(); } + @Override + public String getDatabaseRole() { + return sessionReference.getDatabaseRole(); + } + /** * Updates the session reference with the fallback session. This should only be used for updating * session reference with regular session in case of unimplemented error in multiplexed session. 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 deleted file mode 100644 index 42a67a66296..00000000000 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java +++ /dev/null @@ -1,3516 +0,0 @@ -/* - * Copyright 2017 Google LLC - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package com.google.cloud.spanner; - -import static com.google.cloud.spanner.MetricRegistryConstants.COUNT; -import static com.google.cloud.spanner.MetricRegistryConstants.GET_SESSION_TIMEOUTS; -import static com.google.cloud.spanner.MetricRegistryConstants.IS_MULTIPLEXED; -import static com.google.cloud.spanner.MetricRegistryConstants.MAX_ALLOWED_SESSIONS; -import static com.google.cloud.spanner.MetricRegistryConstants.MAX_ALLOWED_SESSIONS_DESCRIPTION; -import static com.google.cloud.spanner.MetricRegistryConstants.MAX_IN_USE_SESSIONS; -import static com.google.cloud.spanner.MetricRegistryConstants.MAX_IN_USE_SESSIONS_DESCRIPTION; -import static com.google.cloud.spanner.MetricRegistryConstants.METRIC_PREFIX; -import static com.google.cloud.spanner.MetricRegistryConstants.NUM_ACQUIRED_SESSIONS; -import static com.google.cloud.spanner.MetricRegistryConstants.NUM_ACQUIRED_SESSIONS_DESCRIPTION; -import static com.google.cloud.spanner.MetricRegistryConstants.NUM_IN_USE_SESSIONS; -import static com.google.cloud.spanner.MetricRegistryConstants.NUM_READ_SESSIONS; -import static com.google.cloud.spanner.MetricRegistryConstants.NUM_RELEASED_SESSIONS; -import static com.google.cloud.spanner.MetricRegistryConstants.NUM_RELEASED_SESSIONS_DESCRIPTION; -import static com.google.cloud.spanner.MetricRegistryConstants.NUM_SESSIONS_AVAILABLE; -import static com.google.cloud.spanner.MetricRegistryConstants.NUM_SESSIONS_BEING_PREPARED; -import static com.google.cloud.spanner.MetricRegistryConstants.NUM_SESSIONS_IN_POOL; -import static com.google.cloud.spanner.MetricRegistryConstants.NUM_SESSIONS_IN_POOL_DESCRIPTION; -import static com.google.cloud.spanner.MetricRegistryConstants.NUM_SESSIONS_IN_USE; -import static com.google.cloud.spanner.MetricRegistryConstants.NUM_WRITE_SESSIONS; -import static com.google.cloud.spanner.MetricRegistryConstants.SESSIONS_TIMEOUTS_DESCRIPTION; -import static com.google.cloud.spanner.MetricRegistryConstants.SESSIONS_TYPE; -import static com.google.cloud.spanner.MetricRegistryConstants.SPANNER_DEFAULT_LABEL_VALUES; -import static com.google.cloud.spanner.MetricRegistryConstants.SPANNER_LABEL_KEYS; -import static com.google.cloud.spanner.MetricRegistryConstants.SPANNER_LABEL_KEYS_WITH_MULTIPLEXED_SESSIONS; -import static com.google.cloud.spanner.MetricRegistryConstants.SPANNER_LABEL_KEYS_WITH_TYPE; -import static com.google.cloud.spanner.SpannerExceptionFactory.asSpannerException; -import static com.google.cloud.spanner.SpannerExceptionFactory.newSpannerException; -import static com.google.common.base.Preconditions.checkState; - -import com.google.api.core.ApiFuture; -import com.google.api.core.ApiFutures; -import com.google.api.core.SettableApiFuture; -import com.google.api.gax.core.ExecutorProvider; -import com.google.api.gax.rpc.ServerStream; -import com.google.cloud.Timestamp; -import com.google.cloud.Tuple; -import com.google.cloud.grpc.GrpcTransportOptions; -import com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory; -import com.google.cloud.spanner.Options.QueryOption; -import com.google.cloud.spanner.Options.ReadOption; -import com.google.cloud.spanner.Options.TransactionOption; -import com.google.cloud.spanner.Options.UpdateOption; -import com.google.cloud.spanner.SessionClient.SessionConsumer; -import com.google.cloud.spanner.SessionPoolOptions.InactiveTransactionRemovalOptions; -import com.google.cloud.spanner.SpannerException.ResourceNotFoundException; -import com.google.cloud.spanner.SpannerImpl.ClosedException; -import com.google.cloud.spanner.spi.v1.SpannerRpc; -import com.google.common.annotations.VisibleForTesting; -import com.google.common.base.MoreObjects; -import com.google.common.base.Preconditions; -import com.google.common.base.Ticker; -import com.google.common.cache.Cache; -import com.google.common.cache.CacheBuilder; -import com.google.common.collect.ImmutableList; -import com.google.common.util.concurrent.ForwardingListenableFuture; -import com.google.common.util.concurrent.ForwardingListenableFuture.SimpleForwardingListenableFuture; -import com.google.common.util.concurrent.ListenableFuture; -import com.google.common.util.concurrent.MoreExecutors; -import com.google.common.util.concurrent.SettableFuture; -import com.google.protobuf.Empty; -import com.google.spanner.v1.BatchWriteResponse; -import com.google.spanner.v1.ResultSetStats; -import io.opencensus.metrics.DerivedLongCumulative; -import io.opencensus.metrics.DerivedLongGauge; -import io.opencensus.metrics.LabelValue; -import io.opencensus.metrics.MetricOptions; -import io.opencensus.metrics.MetricRegistry; -import io.opencensus.metrics.Metrics; -import io.opentelemetry.api.OpenTelemetry; -import io.opentelemetry.api.common.Attributes; -import io.opentelemetry.api.common.AttributesBuilder; -import io.opentelemetry.api.metrics.Meter; -import java.io.PrintWriter; -import java.io.StringWriter; -import java.time.Duration; -import java.time.Instant; -import java.util.ArrayList; -import java.util.HashSet; -import java.util.Iterator; -import java.util.LinkedList; -import java.util.List; -import java.util.Queue; -import java.util.Random; -import java.util.Set; -import java.util.concurrent.CountDownLatch; -import java.util.concurrent.ExecutionException; -import java.util.concurrent.Executor; -import java.util.concurrent.Future; -import java.util.concurrent.ScheduledExecutorService; -import java.util.concurrent.ScheduledFuture; -import java.util.concurrent.TimeUnit; -import java.util.concurrent.TimeoutException; -import java.util.concurrent.atomic.AtomicBoolean; -import java.util.concurrent.atomic.AtomicInteger; -import java.util.concurrent.atomic.AtomicLong; -import java.util.function.Function; -import java.util.logging.Level; -import java.util.logging.Logger; -import javax.annotation.Nonnull; -import javax.annotation.Nullable; -import javax.annotation.concurrent.GuardedBy; - -/** - * Maintains a pool of sessions. This class itself is thread safe and is meant to be used - * concurrently across multiple threads. - */ -class SessionPool { - - private static final Logger logger = Logger.getLogger(SessionPool.class.getName()); - private final TraceWrapper tracer; - static final String WAIT_FOR_SESSION = "SessionPool.WaitForSession"; - - /** - * If the {@link SessionPoolOptions#getWaitForMinSessions()} duration is greater than zero, waits - * for the creation of at least {@link SessionPoolOptions#getMinSessions()} in the pool using the - * given duration. If the waiting times out, a {@link SpannerException} with the {@link - * ErrorCode#DEADLINE_EXCEEDED} is thrown. - */ - void maybeWaitOnMinSessions() { - final long timeoutNanos = options.getWaitForMinSessions().toNanos(); - if (timeoutNanos <= 0) { - return; - } - - try { - if (!waitOnMinSessionsLatch.await(timeoutNanos, TimeUnit.NANOSECONDS)) { - final long timeoutMillis = options.getWaitForMinSessions().toMillis(); - throw SpannerExceptionFactory.newSpannerException( - ErrorCode.DEADLINE_EXCEEDED, - "Timed out after waiting " + timeoutMillis + "ms for session pool creation"); - } - } catch (InterruptedException e) { - throw SpannerExceptionFactory.propagateInterrupt(e); - } - } - - private abstract static class CachedResultSetSupplier - implements com.google.common.base.Supplier { - - private ResultSet cached; - - abstract ResultSet load(); - - ResultSet reload() { - return cached = load(); - } - - @Override - public ResultSet get() { - if (cached == null) { - cached = load(); - } - return cached; - } - } - - /** - * Wrapper around {@code ReadContext} that releases the session to the pool once the call is - * finished, if it is a single use context. - */ - private static class AutoClosingReadContext - implements ReadContext { - /** - * {@link AsyncResultSet} implementation that keeps track of the async operations that are still - * running for this {@link ReadContext} and that should finish before the {@link ReadContext} - * releases its session back into the pool. - */ - private class AutoClosingReadContextAsyncResultSetImpl extends AsyncResultSetImpl { - private AutoClosingReadContextAsyncResultSetImpl( - ExecutorProvider executorProvider, ResultSet delegate, int bufferRows) { - super(executorProvider, delegate, bufferRows); - } - - @Override - public ApiFuture setCallback(Executor exec, ReadyCallback cb) { - Runnable listener = - () -> { - synchronized (lock) { - if (asyncOperationsCount.decrementAndGet() == 0 && closed) { - // All async operations for this read context have finished. - AutoClosingReadContext.this.close(); - } - } - }; - try { - asyncOperationsCount.incrementAndGet(); - addListener(listener); - return super.setCallback(exec, cb); - } catch (Throwable t) { - removeListener(listener); - asyncOperationsCount.decrementAndGet(); - throw t; - } - } - } - - private final Function readContextDelegateSupplier; - private T readContextDelegate; - private final SessionPool sessionPool; - private final SessionReplacementHandler sessionReplacementHandler; - private final boolean isSingleUse; - private final AtomicInteger asyncOperationsCount = new AtomicInteger(); - - private final Object lock = new Object(); - - @GuardedBy("lock") - private boolean sessionUsedForQuery = false; - - @GuardedBy("lock") - private I session; - - @GuardedBy("lock") - private boolean closed; - - @GuardedBy("lock") - private boolean delegateClosed; - - private AutoClosingReadContext( - Function delegateSupplier, - SessionPool sessionPool, - SessionReplacementHandler sessionReplacementHandler, - I session, - boolean isSingleUse) { - this.readContextDelegateSupplier = delegateSupplier; - this.sessionPool = sessionPool; - this.sessionReplacementHandler = sessionReplacementHandler; - this.session = session; - this.isSingleUse = isSingleUse; - } - - T getReadContextDelegate() { - synchronized (lock) { - if (readContextDelegate == null) { - while (true) { - try { - this.readContextDelegate = readContextDelegateSupplier.apply(this.session); - break; - } catch (SessionNotFoundException e) { - replaceSessionIfPossible(e); - } - } - } - } - return readContextDelegate; - } - - private ResultSet wrap(final CachedResultSetSupplier resultSetSupplier) { - return new ForwardingResultSet(resultSetSupplier) { - private boolean beforeFirst = true; - - @Override - public boolean next() throws SpannerException { - while (true) { - try { - return internalNext(); - } catch (SessionNotFoundException e) { - while (true) { - // Keep the replace-if-possible outside the try-block to let the exception bubble up - // if it's too late to replace the session. - replaceSessionIfPossible(e); - try { - replaceDelegate(resultSetSupplier.reload()); - break; - } catch (SessionNotFoundException snfe) { - e = snfe; - // retry on yet another session. - } - } - } - } - } - - private boolean internalNext() { - try { - boolean ret = super.next(); - if (beforeFirst) { - synchronized (lock) { - session.get().markUsed(); - beforeFirst = false; - sessionUsedForQuery = true; - } - } - if (!ret && isSingleUse) { - close(); - } - return ret; - } catch (SessionNotFoundException e) { - throw e; - } catch (SpannerException e) { - synchronized (lock) { - if (!closed && isSingleUse) { - session.get().setLastException(e); - AutoClosingReadContext.this.close(); - } - } - throw e; - } - } - - @Override - public void close() { - try { - super.close(); - } finally { - if (isSingleUse) { - AutoClosingReadContext.this.close(); - } - } - } - }; - } - - private void replaceSessionIfPossible(SessionNotFoundException notFound) { - synchronized (lock) { - if (isSingleUse || !sessionUsedForQuery) { - // This class is only used by read-only transactions, so we know that we only need a - // read-only session. - session = sessionReplacementHandler.replaceSession(notFound, session); - readContextDelegate = readContextDelegateSupplier.apply(session); - } else { - throw notFound; - } - } - } - - @Override - public ResultSet read( - final String table, - final KeySet keys, - final Iterable columns, - final ReadOption... options) { - return wrap( - new CachedResultSetSupplier() { - @Override - ResultSet load() { - return getReadContextDelegate().read(table, keys, columns, options); - } - }); - } - - @Override - public AsyncResultSet readAsync( - final String table, - final KeySet keys, - final Iterable columns, - final ReadOption... options) { - Options readOptions = Options.fromReadOptions(options); - final int bufferRows = - readOptions.hasBufferRows() - ? readOptions.bufferRows() - : AsyncResultSetImpl.DEFAULT_BUFFER_SIZE; - return new AutoClosingReadContextAsyncResultSetImpl( - sessionPool.sessionClient.getSpanner().getAsyncExecutorProvider(), - wrap( - new CachedResultSetSupplier() { - @Override - ResultSet load() { - return getReadContextDelegate().read(table, keys, columns, options); - } - }), - bufferRows); - } - - @Override - public ResultSet readUsingIndex( - final String table, - final String index, - final KeySet keys, - final Iterable columns, - final ReadOption... options) { - return wrap( - new CachedResultSetSupplier() { - @Override - ResultSet load() { - return getReadContextDelegate().readUsingIndex(table, index, keys, columns, options); - } - }); - } - - @Override - public AsyncResultSet readUsingIndexAsync( - final String table, - final String index, - final KeySet keys, - final Iterable columns, - final ReadOption... options) { - Options readOptions = Options.fromReadOptions(options); - final int bufferRows = - readOptions.hasBufferRows() - ? readOptions.bufferRows() - : AsyncResultSetImpl.DEFAULT_BUFFER_SIZE; - return new AutoClosingReadContextAsyncResultSetImpl( - sessionPool.sessionClient.getSpanner().getAsyncExecutorProvider(), - wrap( - new CachedResultSetSupplier() { - @Override - ResultSet load() { - return getReadContextDelegate() - .readUsingIndex(table, index, keys, columns, options); - } - }), - bufferRows); - } - - @Override - @Nullable - public Struct readRow(String table, Key key, Iterable columns) { - try { - while (true) { - try { - synchronized (lock) { - session.get().markUsed(); - } - return getReadContextDelegate().readRow(table, key, columns); - } catch (SessionNotFoundException e) { - replaceSessionIfPossible(e); - } - } - } finally { - synchronized (lock) { - sessionUsedForQuery = true; - } - if (isSingleUse) { - close(); - } - } - } - - @Override - public ApiFuture readRowAsync(String table, Key key, Iterable columns) { - try (AsyncResultSet rs = readAsync(table, KeySet.singleKey(key), columns)) { - return AbstractReadContext.consumeSingleRowAsync(rs); - } - } - - @Override - @Nullable - public Struct readRowUsingIndex(String table, String index, Key key, Iterable columns) { - try { - while (true) { - try { - synchronized (lock) { - session.get().markUsed(); - } - return getReadContextDelegate().readRowUsingIndex(table, index, key, columns); - } catch (SessionNotFoundException e) { - replaceSessionIfPossible(e); - } - } - } finally { - synchronized (lock) { - sessionUsedForQuery = true; - } - if (isSingleUse) { - close(); - } - } - } - - @Override - public ApiFuture readRowUsingIndexAsync( - String table, String index, Key key, Iterable columns) { - try (AsyncResultSet rs = readUsingIndexAsync(table, index, KeySet.singleKey(key), columns)) { - return AbstractReadContext.consumeSingleRowAsync(rs); - } - } - - @Override - public ResultSet executeQuery(final Statement statement, final QueryOption... options) { - return wrap( - new CachedResultSetSupplier() { - @Override - ResultSet load() { - return getReadContextDelegate().executeQuery(statement, options); - } - }); - } - - @Override - public AsyncResultSet executeQueryAsync( - final Statement statement, final QueryOption... options) { - Options queryOptions = Options.fromQueryOptions(options); - final int bufferRows = - queryOptions.hasBufferRows() - ? queryOptions.bufferRows() - : AsyncResultSetImpl.DEFAULT_BUFFER_SIZE; - return new AutoClosingReadContextAsyncResultSetImpl( - sessionPool.sessionClient.getSpanner().getAsyncExecutorProvider(), - wrap( - new CachedResultSetSupplier() { - @Override - ResultSet load() { - return getReadContextDelegate().executeQuery(statement, options); - } - }), - bufferRows); - } - - @Override - public ResultSet analyzeQuery(final Statement statement, final QueryAnalyzeMode queryMode) { - return wrap( - new CachedResultSetSupplier() { - @Override - ResultSet load() { - return getReadContextDelegate().analyzeQuery(statement, queryMode); - } - }); - } - - @Override - public void close() { - synchronized (lock) { - if (closed && delegateClosed) { - return; - } - closed = true; - if (asyncOperationsCount.get() == 0) { - if (readContextDelegate != null) { - readContextDelegate.close(); - } - session.close(); - delegateClosed = true; - } - } - } - } - - private static class AutoClosingReadTransaction - extends AutoClosingReadContext implements ReadOnlyTransaction { - - AutoClosingReadTransaction( - Function txnSupplier, - SessionPool sessionPool, - SessionReplacementHandler sessionReplacementHandler, - I session, - boolean isSingleUse) { - super(txnSupplier, sessionPool, sessionReplacementHandler, session, isSingleUse); - } - - @Override - public Timestamp getReadTimestamp() { - return getReadContextDelegate().getReadTimestamp(); - } - } - - interface SessionReplacementHandler { - T replaceSession(SessionNotFoundException notFound, T sessionFuture); - - T denyListSession(RetryOnDifferentGrpcChannelException retryException, T sessionFuture); - } - - class PooledSessionReplacementHandler implements SessionReplacementHandler { - @Override - public PooledSessionFuture replaceSession( - SessionNotFoundException e, PooledSessionFuture session) { - if (!options.isFailIfSessionNotFound() && session.get().isAllowReplacing()) { - synchronized (lock) { - numSessionsInUse--; - numSessionsReleased++; - checkedOutSessions.remove(session); - markedCheckedOutSessions.remove(session); - } - session.leakedException = null; - invalidateSession(session.get()); - return getSession(); - } else { - throw e; - } - } - - @Override - public PooledSessionFuture denyListSession( - RetryOnDifferentGrpcChannelException retryException, PooledSessionFuture session) { - // The feature was not enabled when the session pool was created. - if (denyListedChannels == null) { - throw SpannerExceptionFactory.asSpannerException(retryException.getCause()); - } - - int channel = session.get().getChannel(); - synchronized (lock) { - // Calculate the size manually by iterating over the possible keys. We do this because the - // size of a cache can be stale, and manually checking for each possible key will make sure - // we get the correct value, and it will update the cache. - int currentSize = 0; - for (int i = 0; i < numChannels; i++) { - if (denyListedChannels.getIfPresent(i) != null) { - currentSize++; - } - } - if (currentSize < numChannels - 1) { - denyListedChannels.put(channel, DENY_LISTED); - } else { - // We have now deny-listed all channels. Give up and just throw the original error. - throw SpannerExceptionFactory.asSpannerException(retryException.getCause()); - } - } - session.get().releaseToPosition = Position.LAST; - session.close(); - return getSession(); - } - } - - interface SessionNotFoundHandler { - /** - * Handles the given {@link SessionNotFoundException} by possibly converting it to a different - * exception that should be thrown. - */ - SpannerException handleSessionNotFound(SessionNotFoundException notFound); - } - - static class SessionPoolResultSet extends ForwardingResultSet { - private final SessionNotFoundHandler handler; - - private SessionPoolResultSet(SessionNotFoundHandler handler, ResultSet delegate) { - super(delegate); - this.handler = Preconditions.checkNotNull(handler); - } - - @Override - public boolean next() { - try { - return super.next(); - } catch (SessionNotFoundException e) { - throw handler.handleSessionNotFound(e); - } - } - } - - static class AsyncSessionPoolResultSet extends ForwardingAsyncResultSet { - private final SessionNotFoundHandler handler; - - private AsyncSessionPoolResultSet(SessionNotFoundHandler handler, AsyncResultSet delegate) { - super(delegate); - this.handler = Preconditions.checkNotNull(handler); - } - - @Override - public ApiFuture setCallback(Executor executor, final ReadyCallback callback) { - return super.setCallback( - executor, - resultSet -> { - try { - return callback.cursorReady(resultSet); - } catch (SessionNotFoundException e) { - throw handler.handleSessionNotFound(e); - } - }); - } - - @Override - public boolean next() { - try { - return super.next(); - } catch (SessionNotFoundException e) { - throw handler.handleSessionNotFound(e); - } - } - - @Override - public CursorState tryNext() { - try { - return super.tryNext(); - } catch (SessionNotFoundException e) { - throw handler.handleSessionNotFound(e); - } - } - } - - /** - * {@link TransactionContext} that is used in combination with an {@link - * AutoClosingTransactionManager}. This {@link TransactionContext} handles {@link - * SessionNotFoundException}s by replacing the underlying session with a fresh one, and then - * throws an {@link AbortedException} to trigger the retry-loop that has been created by the - * caller. - */ - static class SessionPoolTransactionContext implements TransactionContext { - private final SessionNotFoundHandler handler; - final TransactionContext delegate; - - SessionPoolTransactionContext(SessionNotFoundHandler handler, TransactionContext delegate) { - this.handler = Preconditions.checkNotNull(handler); - this.delegate = delegate; - } - - @Override - public ResultSet read( - String table, KeySet keys, Iterable columns, ReadOption... options) { - return new SessionPoolResultSet(handler, delegate.read(table, keys, columns, options)); - } - - @Override - public AsyncResultSet readAsync( - String table, KeySet keys, Iterable columns, ReadOption... options) { - return new AsyncSessionPoolResultSet( - handler, delegate.readAsync(table, keys, columns, options)); - } - - @Override - public ResultSet readUsingIndex( - String table, String index, KeySet keys, Iterable columns, ReadOption... options) { - return new SessionPoolResultSet( - handler, delegate.readUsingIndex(table, index, keys, columns, options)); - } - - @Override - public AsyncResultSet readUsingIndexAsync( - String table, String index, KeySet keys, Iterable columns, ReadOption... options) { - return new AsyncSessionPoolResultSet( - handler, delegate.readUsingIndexAsync(table, index, keys, columns, options)); - } - - @Override - public Struct readRow(String table, Key key, Iterable columns) { - try { - return delegate.readRow(table, key, columns); - } catch (SessionNotFoundException e) { - throw handler.handleSessionNotFound(e); - } - } - - @Override - public ApiFuture readRowAsync(String table, Key key, Iterable columns) { - try (AsyncResultSet rs = readAsync(table, KeySet.singleKey(key), columns)) { - return ApiFutures.catching( - AbstractReadContext.consumeSingleRowAsync(rs), - SessionNotFoundException.class, - input -> { - throw handler.handleSessionNotFound(input); - }, - MoreExecutors.directExecutor()); - } - } - - @Override - public void buffer(Mutation mutation) { - delegate.buffer(mutation); - } - - @Override - public ApiFuture bufferAsync(Mutation mutation) { - return delegate.bufferAsync(mutation); - } - - @Override - public Struct readRowUsingIndex(String table, String index, Key key, Iterable columns) { - try { - return delegate.readRowUsingIndex(table, index, key, columns); - } catch (SessionNotFoundException e) { - throw handler.handleSessionNotFound(e); - } - } - - @Override - public ApiFuture readRowUsingIndexAsync( - String table, String index, Key key, Iterable columns) { - try (AsyncResultSet rs = readUsingIndexAsync(table, index, KeySet.singleKey(key), columns)) { - return ApiFutures.catching( - AbstractReadContext.consumeSingleRowAsync(rs), - SessionNotFoundException.class, - input -> { - throw handler.handleSessionNotFound(input); - }, - MoreExecutors.directExecutor()); - } - } - - @Override - public void buffer(Iterable mutations) { - delegate.buffer(mutations); - } - - @Override - public ApiFuture bufferAsync(Iterable mutations) { - return delegate.bufferAsync(mutations); - } - - @SuppressWarnings("deprecation") - @Override - public ResultSetStats analyzeUpdate( - Statement statement, QueryAnalyzeMode analyzeMode, UpdateOption... options) { - try (ResultSet resultSet = analyzeUpdateStatement(statement, analyzeMode, options)) { - return resultSet.getStats(); - } - } - - @Override - public ResultSet analyzeUpdateStatement( - Statement statement, QueryAnalyzeMode analyzeMode, UpdateOption... options) { - try { - return delegate.analyzeUpdateStatement(statement, analyzeMode, options); - } catch (SessionNotFoundException e) { - throw handler.handleSessionNotFound(e); - } - } - - @Override - public long executeUpdate(Statement statement, UpdateOption... options) { - try { - return delegate.executeUpdate(statement, options); - } catch (SessionNotFoundException e) { - throw handler.handleSessionNotFound(e); - } - } - - @Override - public ApiFuture executeUpdateAsync(Statement statement, UpdateOption... options) { - return ApiFutures.catching( - delegate.executeUpdateAsync(statement, options), - SessionNotFoundException.class, - input -> { - throw handler.handleSessionNotFound(input); - }, - MoreExecutors.directExecutor()); - } - - @Override - public long[] batchUpdate(Iterable statements, UpdateOption... options) { - try { - return delegate.batchUpdate(statements, options); - } catch (SessionNotFoundException e) { - throw handler.handleSessionNotFound(e); - } - } - - @Override - public ApiFuture batchUpdateAsync( - Iterable statements, UpdateOption... options) { - return ApiFutures.catching( - delegate.batchUpdateAsync(statements, options), - SessionNotFoundException.class, - input -> { - throw handler.handleSessionNotFound(input); - }, - MoreExecutors.directExecutor()); - } - - @Override - public ResultSet executeQuery(Statement statement, QueryOption... options) { - return new SessionPoolResultSet(handler, delegate.executeQuery(statement, options)); - } - - @Override - public AsyncResultSet executeQueryAsync(Statement statement, QueryOption... options) { - return new AsyncSessionPoolResultSet(handler, delegate.executeQueryAsync(statement, options)); - } - - @Override - public ResultSet analyzeQuery(Statement statement, QueryAnalyzeMode queryMode) { - return new SessionPoolResultSet(handler, delegate.analyzeQuery(statement, queryMode)); - } - - @Override - public void close() { - delegate.close(); - } - } - - private static class AutoClosingTransactionManager - implements TransactionManager, SessionNotFoundHandler { - private TransactionManager delegate; - private T session; - private final SessionReplacementHandler sessionReplacementHandler; - private final TransactionOption[] options; - private boolean closed; - private boolean restartedAfterSessionNotFound; - - AutoClosingTransactionManager( - T session, - SessionReplacementHandler sessionReplacementHandler, - TransactionOption... options) { - this.session = session; - this.options = options; - this.sessionReplacementHandler = sessionReplacementHandler; - } - - @Override - public TransactionContext begin() { - this.delegate = session.get().transactionManager(options); - // This cannot throw a SessionNotFoundException, as it does not call the BeginTransaction RPC. - // Instead, the BeginTransaction will be included with the first statement of the transaction. - return internalBegin(); - } - - @Override - public TransactionContext begin(AbortedException exception) { - // For regular sessions, the input exception is ignored and the behavior is equivalent to - // calling {@link #begin()}. - return begin(); - } - - private TransactionContext internalBegin() { - TransactionContext res = new SessionPoolTransactionContext(this, delegate.begin()); - session.get().markUsed(); - return res; - } - - @Override - public SpannerException handleSessionNotFound(SessionNotFoundException notFoundException) { - session = sessionReplacementHandler.replaceSession(notFoundException, session); - CachedSession cachedSession = session.get(); - delegate = cachedSession.getDelegate().transactionManager(options); - restartedAfterSessionNotFound = true; - return createAbortedExceptionWithMinimalRetryDelay(notFoundException); - } - - private static SpannerException createAbortedExceptionWithMinimalRetryDelay( - SessionNotFoundException notFoundException) { - return SpannerExceptionFactory.newSpannerException( - ErrorCode.ABORTED, - notFoundException.getMessage(), - SpannerExceptionFactory.createAbortedExceptionWithRetryDelay( - notFoundException.getMessage(), notFoundException, 0, 1)); - } - - @Override - public void commit() { - try { - delegate.commit(); - } catch (SessionNotFoundException e) { - throw handleSessionNotFound(e); - } finally { - if (getState() != TransactionState.ABORTED) { - close(); - } - } - } - - @Override - public void rollback() { - try { - delegate.rollback(); - } finally { - close(); - } - } - - @Override - public TransactionContext resetForRetry() { - while (true) { - try { - if (restartedAfterSessionNotFound) { - TransactionContext res = new SessionPoolTransactionContext(this, delegate.begin()); - restartedAfterSessionNotFound = false; - return res; - } else { - return new SessionPoolTransactionContext(this, delegate.resetForRetry()); - } - } catch (SessionNotFoundException e) { - session = sessionReplacementHandler.replaceSession(e, session); - CachedSession cachedSession = session.get(); - delegate = cachedSession.getDelegate().transactionManager(options); - restartedAfterSessionNotFound = true; - } - } - } - - @Override - public Timestamp getCommitTimestamp() { - return delegate.getCommitTimestamp(); - } - - @Override - public CommitResponse getCommitResponse() { - return delegate.getCommitResponse(); - } - - @Override - public void close() { - if (closed) { - return; - } - closed = true; - try { - if (delegate != null) { - delegate.close(); - } - } finally { - session.close(); - } - } - - @Override - public TransactionState getState() { - if (restartedAfterSessionNotFound) { - return TransactionState.ABORTED; - } else { - return delegate == null ? null : delegate.getState(); - } - } - } - - /** - * {@link TransactionRunner} that automatically handles {@link SessionNotFoundException}s by - * replacing the underlying session and then restarts the transaction. - */ - static final class SessionPoolTransactionRunner - implements TransactionRunner { - - private I session; - private final SessionReplacementHandler sessionReplacementHandler; - private final TransactionOption[] options; - private TransactionRunner runner; - - SessionPoolTransactionRunner( - I session, - SessionReplacementHandler sessionReplacementHandler, - TransactionOption... options) { - this.session = session; - this.options = options; - this.sessionReplacementHandler = sessionReplacementHandler; - } - - private TransactionRunner getRunner() { - if (this.runner == null) { - this.runner = session.get().readWriteTransaction(options); - } - return runner; - } - - @Override - @Nullable - public T run(TransactionCallable callable) { - try { - T result; - while (true) { - try { - result = getRunner().run(callable); - break; - } catch (SessionNotFoundException e) { - session = sessionReplacementHandler.replaceSession(e, session); - CachedSession cachedSession = session.get(); - runner = cachedSession.getDelegate().readWriteTransaction(); - } catch (RetryOnDifferentGrpcChannelException retryException) { - // This error is thrown by the RetryOnDifferentGrpcChannelErrorHandler in the specific - // case that a transaction failed with a DEADLINE_EXCEEDED error. This is an - // experimental feature that is disabled by default, and that can be removed in a - // future version. - session = sessionReplacementHandler.denyListSession(retryException, session); - CachedSession cachedSession = session.get(); - runner = cachedSession.getDelegate().readWriteTransaction(); - } - } - session.get().markUsed(); - return result; - } catch (SpannerException e) { - //noinspection ThrowableNotThrown - session.get().setLastException(e); - throw e; - } finally { - session.close(); - } - } - - @Override - public Timestamp getCommitTimestamp() { - return getRunner().getCommitTimestamp(); - } - - @Override - public CommitResponse getCommitResponse() { - return getRunner().getCommitResponse(); - } - - @Override - public TransactionRunner allowNestedTransaction() { - getRunner().allowNestedTransaction(); - return this; - } - } - - private static class SessionPoolAsyncRunner implements AsyncRunner { - private volatile I session; - private final SessionReplacementHandler sessionReplacementHandler; - private final TransactionOption[] options; - private SettableApiFuture commitResponse; - - private SessionPoolAsyncRunner( - I session, - SessionReplacementHandler sessionReplacementHandler, - TransactionOption... options) { - this.session = session; - this.options = options; - this.sessionReplacementHandler = sessionReplacementHandler; - } - - @Override - public ApiFuture runAsync(final AsyncWork work, Executor executor) { - commitResponse = SettableApiFuture.create(); - final SettableApiFuture res = SettableApiFuture.create(); - executor.execute( - () -> { - SpannerException exception = null; - R r = null; - AsyncRunner runner = null; - while (true) { - SpannerException se = null; - try { - runner = session.get().runAsync(options); - r = runner.runAsync(work, MoreExecutors.directExecutor()).get(); - break; - } catch (ExecutionException e) { - se = asSpannerException(e.getCause()); - } catch (InterruptedException e) { - se = SpannerExceptionFactory.propagateInterrupt(e); - } catch (Throwable t) { - se = SpannerExceptionFactory.newSpannerException(t); - } finally { - if (se instanceof SessionNotFoundException) { - try { - // The replaceSession method will re-throw the SessionNotFoundException if the - // session cannot be replaced with a new one. - session = - sessionReplacementHandler.replaceSession( - (SessionNotFoundException) se, session); - } catch (SessionNotFoundException e) { - exception = e; - break; - } - } else { - exception = se; - break; - } - } - } - session.get().markUsed(); - session.close(); - setCommitResponse(runner); - if (exception != null) { - res.setException(exception); - } else { - res.set(r); - } - }); - return res; - } - - private void setCommitResponse(AsyncRunner delegate) { - try { - commitResponse.set(delegate.getCommitResponse().get()); - } catch (Throwable t) { - commitResponse.setException(t); - } - } - - @Override - public ApiFuture getCommitTimestamp() { - checkState(commitResponse != null, "runAsync() has not yet been called"); - return ApiFutures.transform( - commitResponse, CommitResponse::getCommitTimestamp, MoreExecutors.directExecutor()); - } - - @Override - public ApiFuture getCommitResponse() { - checkState(commitResponse != null, "runAsync() has not yet been called"); - return commitResponse; - } - } - - // Exception class used just to track the stack trace at the point when a session was handed out - // from the pool. - final class LeakedSessionException extends RuntimeException { - private static final long serialVersionUID = 1451131180314064914L; - - private LeakedSessionException() { - super("Session was checked out from the pool at " + clock.instant()); - } - - private LeakedSessionException(String message) { - super(message); - } - } - - private enum SessionState { - AVAILABLE, - BUSY, - CLOSING, - } - - private PooledSessionFuture createPooledSessionFuture( - ListenableFuture future, ISpan span) { - return new PooledSessionFuture(future, span); - } - - /** Wrapper class for the {@link SessionFuture} implementations. */ - interface SessionFutureWrapper extends DatabaseClient { - - /** Method to resolve {@link SessionFuture} implementation for different use-cases. */ - T get(); - - default Dialect getDialect() { - return get().getDialect(); - } - - default String getDatabaseRole() { - return get().getDatabaseRole(); - } - - default Timestamp write(Iterable mutations) throws SpannerException { - return get().write(mutations); - } - - default CommitResponse writeWithOptions( - Iterable mutations, TransactionOption... options) throws SpannerException { - return get().writeWithOptions(mutations, options); - } - - default Timestamp writeAtLeastOnce(Iterable mutations) throws SpannerException { - return get().writeAtLeastOnce(mutations); - } - - default CommitResponse writeAtLeastOnceWithOptions( - Iterable mutations, TransactionOption... options) throws SpannerException { - return get().writeAtLeastOnceWithOptions(mutations, options); - } - - default ServerStream batchWriteAtLeastOnce( - Iterable mutationGroups, TransactionOption... options) - throws SpannerException { - return get().batchWriteAtLeastOnce(mutationGroups, options); - } - - default ReadContext singleUse() { - return get().singleUse(); - } - - default ReadContext singleUse(TimestampBound bound) { - return get().singleUse(bound); - } - - default ReadOnlyTransaction singleUseReadOnlyTransaction() { - return get().singleUseReadOnlyTransaction(); - } - - default ReadOnlyTransaction singleUseReadOnlyTransaction(TimestampBound bound) { - return get().singleUseReadOnlyTransaction(bound); - } - - default ReadOnlyTransaction readOnlyTransaction() { - return get().readOnlyTransaction(); - } - - default ReadOnlyTransaction readOnlyTransaction(TimestampBound bound) { - return get().readOnlyTransaction(bound); - } - - default TransactionRunner readWriteTransaction(TransactionOption... options) { - return get().readWriteTransaction(options); - } - - default TransactionManager transactionManager(TransactionOption... options) { - return get().transactionManager(options); - } - - default AsyncRunner runAsync(TransactionOption... options) { - return get().runAsync(options); - } - - default AsyncTransactionManager transactionManagerAsync(TransactionOption... options) { - return get().transactionManagerAsync(options); - } - - default long executePartitionedUpdate(Statement stmt, UpdateOption... options) { - return get().executePartitionedUpdate(stmt, options); - } - } - - class PooledSessionFutureWrapper implements SessionFutureWrapper { - PooledSessionFuture pooledSessionFuture; - - public PooledSessionFutureWrapper(PooledSessionFuture pooledSessionFuture) { - this.pooledSessionFuture = pooledSessionFuture; - } - - @Override - public PooledSessionFuture get() { - return this.pooledSessionFuture; - } - } - - interface SessionFuture extends Session { - - /** - * We need to do this because every implementation of {@link SessionFuture} today extends {@link - * SimpleForwardingListenableFuture}. The get() method in parent {@link - * java.util.concurrent.Future} classes specifies checked exceptions in method signature. - * - *

This method is a workaround we don't have to handle checked exceptions specified by other - * interfaces. - */ - CachedSession get(); - - default void addListener(Runnable listener, Executor exec) {} - } - - class PooledSessionFuture extends SimpleForwardingListenableFuture - implements SessionFuture { - - private boolean closed; - private volatile LeakedSessionException leakedException; - private final AtomicBoolean inUse = new AtomicBoolean(); - private final CountDownLatch initialized = new CountDownLatch(1); - private final ISpan span; - - @VisibleForTesting - PooledSessionFuture(ListenableFuture delegate, ISpan span) { - super(delegate); - this.span = span; - } - - @VisibleForTesting - void clearLeakedException() { - this.leakedException = null; - } - - private void markCheckedOut() { - - if (options.isTrackStackTraceOfSessionCheckout()) { - this.leakedException = new LeakedSessionException(); - synchronized (SessionPool.this.lock) { - SessionPool.this.markedCheckedOutSessions.add(this); - } - } - } - - @Override - public Timestamp write(Iterable mutations) throws SpannerException { - return writeWithOptions(mutations).getCommitTimestamp(); - } - - @Override - public CommitResponse writeWithOptions( - Iterable mutations, TransactionOption... options) throws SpannerException { - try { - return get().writeWithOptions(mutations, options); - } finally { - close(); - } - } - - @Override - public Timestamp writeAtLeastOnce(Iterable mutations) throws SpannerException { - return writeAtLeastOnceWithOptions(mutations).getCommitTimestamp(); - } - - @Override - public CommitResponse writeAtLeastOnceWithOptions( - Iterable mutations, TransactionOption... options) throws SpannerException { - try { - return get().writeAtLeastOnceWithOptions(mutations, options); - } finally { - close(); - } - } - - @Override - public ServerStream batchWriteAtLeastOnce( - Iterable mutationGroups, TransactionOption... options) - throws SpannerException { - try { - return get().batchWriteAtLeastOnce(mutationGroups, options); - } finally { - close(); - } - } - - @Override - public ReadContext singleUse() { - try { - return new AutoClosingReadContext<>( - session -> { - PooledSession ps = session.get(); - return ps.delegate.singleUse(); - }, - SessionPool.this, - pooledSessionReplacementHandler, - this, - true); - } catch (Exception e) { - close(); - throw e; - } - } - - @Override - public ReadContext singleUse(final TimestampBound bound) { - try { - return new AutoClosingReadContext<>( - session -> { - PooledSession ps = session.get(); - return ps.delegate.singleUse(bound); - }, - SessionPool.this, - pooledSessionReplacementHandler, - this, - true); - } catch (Exception e) { - close(); - throw e; - } - } - - @Override - public ReadOnlyTransaction singleUseReadOnlyTransaction() { - return internalReadOnlyTransaction( - session -> { - PooledSession ps = session.get(); - return ps.delegate.singleUseReadOnlyTransaction(); - }, - true); - } - - @Override - public ReadOnlyTransaction singleUseReadOnlyTransaction(final TimestampBound bound) { - return internalReadOnlyTransaction( - session -> { - PooledSession ps = session.get(); - return ps.delegate.singleUseReadOnlyTransaction(bound); - }, - true); - } - - @Override - public ReadOnlyTransaction readOnlyTransaction() { - return internalReadOnlyTransaction( - session -> { - PooledSession ps = session.get(); - return ps.delegate.readOnlyTransaction(); - }, - false); - } - - @Override - public ReadOnlyTransaction readOnlyTransaction(final TimestampBound bound) { - return internalReadOnlyTransaction( - session -> { - PooledSession ps = session.get(); - return ps.delegate.readOnlyTransaction(bound); - }, - false); - } - - private ReadOnlyTransaction internalReadOnlyTransaction( - Function transactionSupplier, - boolean isSingleUse) { - try { - return new AutoClosingReadTransaction<>( - transactionSupplier, - SessionPool.this, - pooledSessionReplacementHandler, - this, - isSingleUse); - } catch (Exception e) { - close(); - throw e; - } - } - - @Override - public TransactionRunner readWriteTransaction(TransactionOption... options) { - return new SessionPoolTransactionRunner<>(this, pooledSessionReplacementHandler, options); - } - - @Override - public TransactionManager transactionManager(TransactionOption... options) { - return new AutoClosingTransactionManager<>(this, pooledSessionReplacementHandler, options); - } - - @Override - public AsyncRunner runAsync(TransactionOption... options) { - return new SessionPoolAsyncRunner<>(this, pooledSessionReplacementHandler, options); - } - - @Override - public AsyncTransactionManager transactionManagerAsync(TransactionOption... options) { - return new SessionPoolAsyncTransactionManager<>( - pooledSessionReplacementHandler, this, options); - } - - @Override - public long executePartitionedUpdate(Statement stmt, UpdateOption... options) { - try { - return get(true).executePartitionedUpdate(stmt, options); - } finally { - close(); - } - } - - @Override - public String getName() { - return get().getName(); - } - - @Override - public void close() { - try { - asyncClose().get(); - } catch (InterruptedException e) { - throw SpannerExceptionFactory.propagateInterrupt(e); - } catch (ExecutionException e) { - throw asSpannerException(e.getCause()); - } - } - - @Override - public ApiFuture 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) { - return delegate.asyncClose(); - } - } finally { - synchronized (lock) { - leakedException = null; - checkedOutSessions.remove(this); - markedCheckedOutSessions.remove(this); - } - } - return ApiFutures.immediateFuture(Empty.getDefaultInstance()); - } - - private PooledSession getOrNull() { - try { - return get(); - } catch (Throwable t) { - return null; - } - } - - @Override - public PooledSession get() { - return get(false); - } - - PooledSession get(final boolean eligibleForLongRunning) { - if (inUse.compareAndSet(false, true)) { - PooledSession res = null; - try { - res = super.get(); - } catch (Throwable e) { - // ignore the exception as it will be handled by the call to super.get() below. - } - if (res != null) { - res.markBusy(span); - span.addAnnotation("Using Session", "sessionId", res.getName()); - synchronized (lock) { - incrementNumSessionsInUse(); - checkedOutSessions.add(this); - } - res.eligibleForLongRunning = eligibleForLongRunning; - } - initialized.countDown(); - } - try { - initialized.await(); - return super.get(); - } catch (ExecutionException e) { - throw SpannerExceptionFactory.newSpannerException(e.getCause()); - } catch (InterruptedException e) { - throw SpannerExceptionFactory.propagateInterrupt(e); - } - } - - public int getChannel() { - return get().getChannel(); - } - } - - interface CachedSession extends Session { - - SessionImpl getDelegate(); - - void markBusy(ISpan span); - - void markUsed(); - - SpannerException setLastException(SpannerException exception); - - AsyncTransactionManagerImpl transactionManagerAsync(TransactionOption... options); - - void setAllowReplacing(boolean b); - } - - class PooledSession implements CachedSession { - - @VisibleForTesting final SessionImpl delegate; - private volatile SpannerException lastException; - private volatile boolean allowReplacing = true; - - /** - * This ensures that the session is added at a random position in the pool the first time it is - * actually added to the pool. - */ - @GuardedBy("lock") - private Position releaseToPosition = initialReleasePosition; - - /** - * Property to mark if the session is eligible to be long-running. This can only be true if the - * session is executing certain types of transactions (for ex - Partitioned DML) which can be - * long-running. By default, most transaction types are not expected to be long-running and - * hence this value is false. - */ - private volatile boolean eligibleForLongRunning = false; - - /** - * Property to mark if the session is no longer part of the session pool. For ex - A session - * which is long-running gets cleaned up and removed from the pool. - */ - private volatile boolean isRemovedFromPool = false; - - /** - * Property to mark if a leaked session exception is already logged. Given a session maintainer - * thread runs repeatedly at a defined interval, this property allows us to ensure that an - * exception is logged only once per leaked session. This is to avoid noisy repeated logs around - * session leaks for long-running sessions. - */ - private volatile boolean isLeakedExceptionLogged = false; - - @GuardedBy("lock") - private SessionState state; - - private PooledSession(SessionImpl delegate) { - this.delegate = Preconditions.checkNotNull(delegate); - this.state = SessionState.AVAILABLE; - - // initialise the lastUseTime field for each session. - this.markUsed(); - } - - int getChannel() { - Long channelHint = (Long) delegate.getOptions().get(SpannerRpc.Option.CHANNEL_HINT); - return channelHint == null - ? 0 - : (int) (channelHint % sessionClient.getSpanner().getOptions().getNumChannels()); - } - - @Override - public String toString() { - return getName(); - } - - @VisibleForTesting - @Override - public void setAllowReplacing(boolean allowReplacing) { - this.allowReplacing = allowReplacing; - } - - @VisibleForTesting - void setEligibleForLongRunning(boolean eligibleForLongRunning) { - this.eligibleForLongRunning = eligibleForLongRunning; - } - - @Override - public Timestamp write(Iterable mutations) throws SpannerException { - return writeWithOptions(mutations).getCommitTimestamp(); - } - - @Override - public CommitResponse writeWithOptions( - Iterable mutations, TransactionOption... options) throws SpannerException { - try { - markUsed(); - return delegate.writeWithOptions(mutations, options); - } catch (SpannerException e) { - throw lastException = e; - } - } - - @Override - public Timestamp writeAtLeastOnce(Iterable mutations) throws SpannerException { - return writeAtLeastOnceWithOptions(mutations).getCommitTimestamp(); - } - - @Override - public CommitResponse writeAtLeastOnceWithOptions( - Iterable mutations, TransactionOption... options) throws SpannerException { - try { - markUsed(); - return delegate.writeAtLeastOnceWithOptions(mutations, options); - } catch (SpannerException e) { - throw lastException = e; - } - } - - @Override - public ServerStream batchWriteAtLeastOnce( - Iterable mutationGroups, TransactionOption... options) - throws SpannerException { - try { - markUsed(); - return delegate.batchWriteAtLeastOnce(mutationGroups, options); - } catch (SpannerException e) { - throw lastException = e; - } - } - - @Override - public long executePartitionedUpdate(Statement stmt, UpdateOption... options) - throws SpannerException { - try { - markUsed(); - return delegate.executePartitionedUpdate(stmt, options); - } catch (SpannerException e) { - throw lastException = e; - } - } - - @Override - public ReadContext singleUse() { - return delegate.singleUse(); - } - - @Override - public ReadContext singleUse(TimestampBound bound) { - return delegate.singleUse(bound); - } - - @Override - public ReadOnlyTransaction singleUseReadOnlyTransaction() { - return delegate.singleUseReadOnlyTransaction(); - } - - @Override - public ReadOnlyTransaction singleUseReadOnlyTransaction(TimestampBound bound) { - return delegate.singleUseReadOnlyTransaction(bound); - } - - @Override - public ReadOnlyTransaction readOnlyTransaction() { - return delegate.readOnlyTransaction(); - } - - @Override - public ReadOnlyTransaction readOnlyTransaction(TimestampBound bound) { - return delegate.readOnlyTransaction(bound); - } - - @Override - public TransactionRunner readWriteTransaction(TransactionOption... options) { - return delegate.readWriteTransaction(options); - } - - @Override - public AsyncRunner runAsync(TransactionOption... options) { - return delegate.runAsync(options); - } - - @Override - public AsyncTransactionManagerImpl transactionManagerAsync(TransactionOption... options) { - return delegate.transactionManagerAsync(options); - } - - @Override - public ApiFuture asyncClose() { - close(); - return ApiFutures.immediateFuture(Empty.getDefaultInstance()); - } - - @Override - public void close() { - synchronized (lock) { - numSessionsInUse--; - numSessionsReleased++; - } - if ((lastException != null && isSessionNotFound(lastException)) || isRemovedFromPool) { - invalidateSession(this); - } else { - if (isDatabaseOrInstanceNotFound(lastException)) { - // Mark this session pool as no longer valid and then release the session into the pool as - // there is nothing we can do with it anyways. - synchronized (lock) { - SessionPool.this.resourceNotFoundException = - MoreObjects.firstNonNull( - SessionPool.this.resourceNotFoundException, - (ResourceNotFoundException) lastException); - } - } - lastException = null; - isRemovedFromPool = false; - if (state != SessionState.CLOSING) { - state = SessionState.AVAILABLE; - } - releaseSession(this, false); - } - } - - @Override - public String getName() { - return delegate.getName(); - } - - private void keepAlive() { - markUsed(); - final ISpan previousSpan = delegate.getCurrentSpan(); - delegate.setCurrentSpan(tracer.getBlankSpan()); - try (ResultSet resultSet = - delegate - .singleUse(TimestampBound.ofMaxStaleness(60, TimeUnit.SECONDS)) - .executeQuery(Statement.newBuilder("SELECT 1").build())) { - resultSet.next(); - } finally { - delegate.setCurrentSpan(previousSpan); - } - } - - private void determineDialectAsync(final SettableFuture dialect) { - Preconditions.checkNotNull(dialect); - executor.submit( - () -> { - try { - dialect.set(determineDialect()); - } catch (Throwable t) { - // Catch-all as we want to propagate all exceptions to anyone who might be interested - // in the database dialect, and there's nothing sensible that we can do with it here. - dialect.setException(t); - } finally { - releaseSession(this, false); - } - }); - } - - private Dialect determineDialect() { - try (ResultSet dialectResultSet = - delegate.singleUse().executeQuery(DETERMINE_DIALECT_STATEMENT)) { - if (dialectResultSet.next()) { - return Dialect.fromName(dialectResultSet.getString(0)); - } else { - throw SpannerExceptionFactory.newSpannerException( - ErrorCode.NOT_FOUND, "No dialect found for database"); - } - } - } - - @Override - public SessionImpl getDelegate() { - return this.delegate; - } - - @Override - public void markBusy(ISpan span) { - this.delegate.setCurrentSpan(span); - this.state = SessionState.BUSY; - } - - private void markClosing() { - this.state = SessionState.CLOSING; - } - - @Override - public void markUsed() { - delegate.markUsed(clock.instant()); - } - - @Override - public SpannerException setLastException(SpannerException exception) { - this.lastException = exception; - return exception; - } - - boolean isAllowReplacing() { - return this.allowReplacing; - } - - @Override - public TransactionManager transactionManager(TransactionOption... options) { - return delegate.transactionManager(options); - } - } - - private final class WaiterFuture extends ForwardingListenableFuture { - private static final long MAX_SESSION_WAIT_TIMEOUT = 240_000L; - private final SettableFuture waiter = SettableFuture.create(); - - @Override - @Nonnull - protected ListenableFuture delegate() { - return waiter; - } - - private void put(PooledSession session) { - waiter.set(session); - } - - private void put(SpannerException e) { - waiter.setException(e); - } - - @Override - public PooledSession get() { - long currentTimeout = options.getInitialWaitForSessionTimeoutMillis(); - while (true) { - ISpan span = tracer.spanBuilder(WAIT_FOR_SESSION); - try (IScope ignore = tracer.withSpan(span)) { - PooledSession s = - pollUninterruptiblyWithTimeout(currentTimeout, options.getAcquireSessionTimeout()); - if (s == null) { - // Set the status to DEADLINE_EXCEEDED and retry. - numWaiterTimeouts.incrementAndGet(); - tracer.getCurrentSpan().setStatus(ErrorCode.DEADLINE_EXCEEDED); - currentTimeout = Math.min(currentTimeout * 2, MAX_SESSION_WAIT_TIMEOUT); - } else { - return s; - } - } catch (Exception e) { - if (e instanceof SpannerException - && ErrorCode.RESOURCE_EXHAUSTED.equals(((SpannerException) e).getErrorCode())) { - numWaiterTimeouts.incrementAndGet(); - tracer.getCurrentSpan().setStatus(ErrorCode.RESOURCE_EXHAUSTED); - } - span.setStatus(e); - throw e; - } finally { - span.end(); - } - } - } - - private PooledSession pollUninterruptiblyWithTimeout( - long timeoutMillis, Duration acquireSessionTimeout) { - boolean interrupted = false; - try { - while (true) { - try { - return acquireSessionTimeout == null - ? waiter.get(timeoutMillis, TimeUnit.MILLISECONDS) - : waiter.get(acquireSessionTimeout.toMillis(), TimeUnit.MILLISECONDS); - } catch (InterruptedException e) { - interrupted = true; - } catch (TimeoutException e) { - if (acquireSessionTimeout != null) { - SpannerException exception = - SpannerExceptionFactory.newSpannerException( - ErrorCode.RESOURCE_EXHAUSTED, - "Timed out after waiting " - + acquireSessionTimeout.toMillis() - + "ms for acquiring session. To mitigate error" - + " SessionPoolOptions#setAcquireSessionTimeout(Duration) to set a higher" - + " timeout or increase the number of sessions in the session pool.\n" - + createCheckedOutSessionsStackTraces()); - if (waiter.setException(exception)) { - // Only throw the exception if setting it on the waiter was successful. The - // waiter.setException(..) method returns false if some other thread in the meantime - // called waiter.set(..), which means that a session became available between the - // time that the TimeoutException was thrown and now. - throw exception; - } - } - return null; - } catch (ExecutionException e) { - throw SpannerExceptionFactory.newSpannerException(e.getCause()); - } - } - } finally { - if (interrupted) { - Thread.currentThread().interrupt(); - } - } - } - } - - /** - * Background task to maintain the pool. Tasks: - * - *

    - *
  • Removes idle sessions from the pool. Sessions that go above MinSessions that have not - * been used for the last 55 minutes will be removed from the pool. These will automatically - * be garbage collected by the backend. - *
  • Keeps alive sessions that have not been used for a user configured time in order to keep - * MinSessions sessions alive in the pool at any time. The keep-alive traffic is smeared out - * over a window of 10 minutes to avoid bursty traffic. - *
  • Removes unexpected long running transactions from the pool. Only certain transaction - * types (for ex - Partitioned DML / Batch Reads) can be long running. This tasks checks the - * sessions which have been inactive for a longer than usual duration (for ex - 60 minutes) - * and removes such sessions from the pool. - *
- */ - final class PoolMaintainer { - - // Length of the window in millis over which we keep track of maximum number of concurrent - // sessions in use. - private final Duration windowLength = Duration.ofMillis(TimeUnit.MINUTES.toMillis(10)); - // Frequency of the timer loop. - @VisibleForTesting final long loopFrequency = options.getLoopFrequency(); - // Number of loop iterations in which we need to close all the sessions waiting for closure. - @VisibleForTesting final long numClosureCycles = windowLength.toMillis() / loopFrequency; - private final Duration keepAliveMillis = - Duration.ofMillis(TimeUnit.MINUTES.toMillis(options.getKeepAliveIntervalMinutes())); - // Number of loop iterations in which we need to keep alive all the sessions - @VisibleForTesting final long numKeepAliveCycles = keepAliveMillis.toMillis() / loopFrequency; - - /** - * Variable maintaining the last execution time of the long-running transaction cleanup task. - * - *

The long-running transaction cleanup needs to be performed every X minutes. The X minutes - * recurs multiple times within the invocation of the pool maintainer thread. For ex - If the - * main thread runs every 10s and the long-running transaction clean-up needs to be performed - * every 2 minutes, then we need to keep a track of when was the last time that this task - * executed and makes sure we only execute it every 2 minutes and not every 10 seconds. - */ - @VisibleForTesting Instant lastExecutionTime; - - /** - * The previous numSessionsAcquired seen by the maintainer. This is used to calculate the - * transactions per second, which again is used to determine whether to randomize the order of - * the session pool. - */ - private long prevNumSessionsAcquired; - - boolean closed = false; - - @GuardedBy("lock") - ScheduledFuture scheduledFuture; - - @GuardedBy("lock") - boolean running; - - void init() { - lastExecutionTime = clock.instant(); - - // Scheduled pool maintenance worker. - synchronized (lock) { - scheduledFuture = - executor.scheduleAtFixedRate( - this::maintainPool, loopFrequency, loopFrequency, TimeUnit.MILLISECONDS); - } - } - - void close() { - synchronized (lock) { - if (!closed) { - closed = true; - scheduledFuture.cancel(false); - if (!running) { - decrementPendingClosures(1); - } - } - } - } - - boolean isClosed() { - synchronized (lock) { - return closed; - } - } - - // Does various pool maintenance activities. - void maintainPool() { - Instant currTime; - synchronized (lock) { - if (SessionPool.this.isClosed()) { - return; - } - running = true; - if (loopFrequency >= 1000L) { - SessionPool.this.transactionsPerSecond = - (SessionPool.this.numSessionsAcquired - prevNumSessionsAcquired) - / (loopFrequency / 1000L); - } - this.prevNumSessionsAcquired = SessionPool.this.numSessionsAcquired; - - currTime = clock.instant(); - // Reset the start time for recording the maximum number of sessions in the pool - if (currTime.isAfter(SessionPool.this.lastResetTime.plus(Duration.ofMinutes(10)))) { - SessionPool.this.maxSessionsInUse = SessionPool.this.numSessionsInUse; - SessionPool.this.lastResetTime = currTime; - } - } - - removeIdleSessions(currTime); - // Now go over all the remaining sessions and see if they need to be kept alive explicitly. - keepAliveSessions(currTime); - replenishPool(); - synchronized (lock) { - running = false; - if (SessionPool.this.isClosed()) { - decrementPendingClosures(1); - } - } - removeLongRunningSessions(currTime); - } - - private void removeIdleSessions(Instant currTime) { - synchronized (lock) { - // Determine the minimum last use time for a session to be deemed to still be alive. Remove - // all sessions that have a lastUseTime before that time, unless it would cause us to go - // below MinSessions. - Instant minLastUseTime = currTime.minus(options.getRemoveInactiveSessionAfterDuration()); - Iterator iterator = sessions.descendingIterator(); - while (iterator.hasNext()) { - PooledSession session = iterator.next(); - if (session.delegate.getLastUseTime() != null - && session.delegate.getLastUseTime().isBefore(minLastUseTime)) { - if (session.state != SessionState.CLOSING) { - boolean isRemoved = removeFromPool(session); - if (isRemoved) { - numIdleSessionsRemoved++; - if (idleSessionRemovedListener != null) { - idleSessionRemovedListener.apply(session); - } - } - iterator.remove(); - } - } - } - } - } - - private void keepAliveSessions(Instant currTime) { - long numSessionsToKeepAlive = 0; - synchronized (lock) { - if (numSessionsInUse >= (options.getMinSessions() + options.getMaxIdleSessions())) { - // At least MinSessions are in use, so we don't have to ping any sessions. - return; - } - // In each cycle only keep alive a subset of sessions to prevent burst of traffic. - numSessionsToKeepAlive = - (long) - Math.ceil( - (double) - ((options.getMinSessions() + options.getMaxIdleSessions()) - - numSessionsInUse) - / numKeepAliveCycles); - } - // Now go over all the remaining sessions and see if they need to be kept alive explicitly. - Instant keepAliveThreshold = currTime.minus(keepAliveMillis); - - // Keep chugging till there is no session that needs to be kept alive. - while (numSessionsToKeepAlive > 0) { - Tuple sessionToKeepAlive; - synchronized (lock) { - sessionToKeepAlive = findSessionToKeepAlive(sessions, keepAliveThreshold, 0); - } - if (sessionToKeepAlive == null) { - break; - } - try { - logger.log(Level.FINE, "Keeping alive session " + sessionToKeepAlive.x().getName()); - numSessionsToKeepAlive--; - sessionToKeepAlive.x().keepAlive(); - releaseSession(sessionToKeepAlive); - } catch (SpannerException e) { - handleException(e, sessionToKeepAlive); - } - } - } - - private void replenishPool() { - synchronized (lock) { - // If we have gone below min pool size, create that many sessions. - int sessionCount = options.getMinSessions() - (totalSessions() + numSessionsBeingCreated); - if (sessionCount > 0) { - createSessions(getAllowedCreateSessions(sessionCount), false); - } - } - } - - // cleans up sessions which are unexpectedly long-running. - void removeLongRunningSessions(Instant currentTime) { - try { - if (SessionPool.this.isClosed()) { - return; - } - final InactiveTransactionRemovalOptions inactiveTransactionRemovalOptions = - options.getInactiveTransactionRemovalOptions(); - final Instant minExecutionTime = - lastExecutionTime.plus(inactiveTransactionRemovalOptions.getExecutionFrequency()); - if (currentTime.isBefore(minExecutionTime)) { - return; - } - lastExecutionTime = currentTime; // update this only after we have decided to execute task - if (options.closeInactiveTransactions() - || options.warnInactiveTransactions() - || options.warnAndCloseInactiveTransactions()) { - removeLongRunningSessions(currentTime, inactiveTransactionRemovalOptions); - } - } catch (final Throwable t) { - logger.log(Level.WARNING, "Failed removing long running transactions", t); - } - } - - private void removeLongRunningSessions( - final Instant currentTime, - final InactiveTransactionRemovalOptions inactiveTransactionRemovalOptions) { - synchronized (lock) { - final double usedSessionsRatio = getRatioOfSessionsInUse(); - if (usedSessionsRatio > inactiveTransactionRemovalOptions.getUsedSessionsRatioThreshold()) { - Iterator iterator = checkedOutSessions.iterator(); - while (iterator.hasNext()) { - final PooledSessionFuture sessionFuture = iterator.next(); - // the below get() call on future object is non-blocking since checkedOutSessions - // collection is populated only when the get() method in {@code PooledSessionFuture} is - // called. - final PooledSession session = (PooledSession) sessionFuture.get(); - final Duration durationFromLastUse = - Duration.between(session.getDelegate().getLastUseTime(), currentTime); - if (!session.eligibleForLongRunning - && durationFromLastUse.compareTo( - inactiveTransactionRemovalOptions.getIdleTimeThreshold()) - > 0) { - if ((options.warnInactiveTransactions() || options.warnAndCloseInactiveTransactions()) - && !session.isLeakedExceptionLogged) { - if (options.warnAndCloseInactiveTransactions()) { - logger.log( - Level.WARNING, - String.format("Removing long-running session => %s", session.getName()), - sessionFuture.leakedException); - session.isLeakedExceptionLogged = true; - } else if (options.warnInactiveTransactions()) { - logger.log( - Level.WARNING, - String.format( - "Detected long-running session => %s. To automatically remove" - + " long-running sessions, set SessionOption" - + " ActionOnInactiveTransaction to WARN_AND_CLOSE by invoking" - + " setWarnAndCloseIfInactiveTransactions() method.", - session.getName()), - sessionFuture.leakedException); - session.isLeakedExceptionLogged = true; - } - } - if ((options.closeInactiveTransactions() - || options.warnAndCloseInactiveTransactions()) - && session.state != SessionState.CLOSING) { - final boolean isRemoved = removeFromPool(session); - if (isRemoved) { - session.isRemovedFromPool = true; - numLeakedSessionsRemoved++; - if (longRunningSessionRemovedListener != null) { - longRunningSessionRemovedListener.apply(session); - } - } - iterator.remove(); - } - } - } - } - } - } - } - - enum Position { - FIRST, - LAST, - RANDOM - } - - /** - * This statement is (currently) used to determine the dialect of the database that is used by the - * session pool. This statement is subject to change when the INFORMATION_SCHEMA contains a table - * where the dialect of the database can be read directly, and any tests that want to detect the - * specific 'determine dialect statement' should rely on this constant instead of the actual - * value. - */ - @VisibleForTesting - static final Statement DETERMINE_DIALECT_STATEMENT = - Statement.newBuilder( - "select option_value " - + "from information_schema.database_options " - + "where option_name='database_dialect'") - .build(); - - private final SessionPoolOptions options; - private final SettableFuture dialect = SettableFuture.create(); - private final String databaseRole; - private final SessionClient sessionClient; - private final int numChannels; - private final ScheduledExecutorService executor; - private final ExecutorFactory executorFactory; - - final PoolMaintainer poolMaintainer; - private final Clock clock; - - /** - * initialReleasePosition determines where in the pool sessions are added when they are released - * into the pool the first time. This is always RANDOM in production, but some tests use FIRST to - * be able to verify the order of sessions in the pool. Using RANDOM ensures that we do not get an - * unbalanced session pool where all sessions belonging to one gRPC channel are added to the same - * region in the pool. - */ - private final Position initialReleasePosition; - - private final Object lock = new Object(); - private final Random random = new Random(); - - @GuardedBy("lock") - private boolean detectDialectStarted; - - @GuardedBy("lock") - private int pendingClosure; - - @GuardedBy("lock") - private SettableFuture closureFuture; - - @GuardedBy("lock") - private ClosedException closedException; - - @GuardedBy("lock") - private ResourceNotFoundException resourceNotFoundException; - - @GuardedBy("lock") - private final LinkedList sessions = new LinkedList<>(); - - @GuardedBy("lock") - private final Queue waiters = new LinkedList<>(); - - @GuardedBy("lock") - private int numSessionsBeingCreated = 0; - - @GuardedBy("lock") - private int numSessionsInUse = 0; - - @GuardedBy("lock") - private int maxSessionsInUse = 0; - - @GuardedBy("lock") - private Instant lastResetTime = Clock.INSTANCE.instant(); - - @GuardedBy("lock") - private long numSessionsAcquired = 0; - - @GuardedBy("lock") - private long numSessionsReleased = 0; - - @GuardedBy("lock") - private long numIdleSessionsRemoved = 0; - - @GuardedBy("lock") - private long transactionsPerSecond = 0L; - - @GuardedBy("lock") - private long numLeakedSessionsRemoved = 0; - - private final AtomicLong numWaiterTimeouts = new AtomicLong(); - - @GuardedBy("lock") - private final Set allSessions = new HashSet<>(); - - @GuardedBy("lock") - @VisibleForTesting - final Set checkedOutSessions = new HashSet<>(); - - @GuardedBy("lock") - private final Set markedCheckedOutSessions = new HashSet<>(); - - private final SessionConsumer sessionConsumer = new SessionConsumerImpl(); - - @VisibleForTesting Function idleSessionRemovedListener; - - @VisibleForTesting Function longRunningSessionRemovedListener; - private final CountDownLatch waitOnMinSessionsLatch; - private final PooledSessionReplacementHandler pooledSessionReplacementHandler = - new PooledSessionReplacementHandler(); - - private static final Object DENY_LISTED = new Object(); - private final Cache denyListedChannels; - - /** - * Create a session pool with the given options and for the given database. It will also start - * eagerly creating sessions if {@link SessionPoolOptions#getMinSessions()} is greater than 0. - * Return pool is immediately ready for use, though getting a session might block for sessions to - * be created. - */ - static SessionPool createPool( - SpannerOptions spannerOptions, - SessionClient sessionClient, - TraceWrapper tracer, - List labelValues, - Attributes attributes, - AtomicLong numMultiplexedSessionsAcquired, - AtomicLong numMultiplexedSessionsReleased) { - final SessionPoolOptions sessionPoolOptions = spannerOptions.getSessionPoolOptions(); - - // A clock instance is passed in {@code SessionPoolOptions} in order to allow mocking via tests. - final Clock poolMaintainerClock = sessionPoolOptions.getPoolMaintainerClock(); - return createPool( - sessionPoolOptions, - spannerOptions.getDatabaseRole(), - ((GrpcTransportOptions) spannerOptions.getTransportOptions()).getExecutorFactory(), - sessionClient, - poolMaintainerClock == null ? new Clock() : poolMaintainerClock, - Position.RANDOM, - Metrics.getMetricRegistry(), - tracer, - labelValues, - spannerOptions.getOpenTelemetry(), - attributes, - numMultiplexedSessionsAcquired, - numMultiplexedSessionsReleased); - } - - static SessionPool createPool( - SessionPoolOptions poolOptions, - ExecutorFactory executorFactory, - SessionClient sessionClient, - TraceWrapper tracer, - OpenTelemetry openTelemetry) { - return createPool( - poolOptions, - executorFactory, - sessionClient, - new Clock(), - Position.RANDOM, - tracer, - openTelemetry); - } - - static SessionPool createPool( - SessionPoolOptions poolOptions, - ExecutorFactory executorFactory, - SessionClient sessionClient, - Clock clock, - Position initialReleasePosition, - TraceWrapper tracer, - OpenTelemetry openTelemetry) { - return createPool( - poolOptions, - null, - executorFactory, - sessionClient, - clock, - initialReleasePosition, - Metrics.getMetricRegistry(), - tracer, - SPANNER_DEFAULT_LABEL_VALUES, - openTelemetry, - null, - new AtomicLong(), - new AtomicLong()); - } - - static SessionPool createPool( - SessionPoolOptions poolOptions, - String databaseRole, - ExecutorFactory executorFactory, - SessionClient sessionClient, - Clock clock, - Position initialReleasePosition, - MetricRegistry metricRegistry, - TraceWrapper tracer, - List labelValues, - OpenTelemetry openTelemetry, - Attributes attributes, - AtomicLong numMultiplexedSessionsAcquired, - AtomicLong numMultiplexedSessionsReleased) { - SessionPool pool = - new SessionPool( - poolOptions, - databaseRole, - executorFactory, - executorFactory.get(), - sessionClient, - clock, - initialReleasePosition, - metricRegistry, - tracer, - labelValues, - openTelemetry, - attributes, - numMultiplexedSessionsAcquired, - numMultiplexedSessionsReleased); - pool.initPool(); - return pool; - } - - private SessionPool( - SessionPoolOptions options, - String databaseRole, - ExecutorFactory executorFactory, - ScheduledExecutorService executor, - SessionClient sessionClient, - Clock clock, - Position initialReleasePosition, - MetricRegistry metricRegistry, - TraceWrapper tracer, - List labelValues, - OpenTelemetry openTelemetry, - Attributes attributes, - AtomicLong numMultiplexedSessionsAcquired, - AtomicLong numMultiplexedSessionsReleased) { - this.options = options; - this.databaseRole = databaseRole; - this.executorFactory = executorFactory; - this.executor = executor; - this.sessionClient = sessionClient; - this.numChannels = sessionClient.getSpanner().getOptions().getNumChannels(); - this.clock = clock; - this.initialReleasePosition = initialReleasePosition; - this.poolMaintainer = new PoolMaintainer(); - this.tracer = tracer; - this.initOpenCensusMetricsCollection( - metricRegistry, - labelValues, - numMultiplexedSessionsAcquired, - numMultiplexedSessionsReleased); - this.initOpenTelemetryMetricsCollection( - openTelemetry, attributes, numMultiplexedSessionsAcquired, numMultiplexedSessionsReleased); - this.waitOnMinSessionsLatch = - options.getMinSessions() > 0 ? new CountDownLatch(1) : new CountDownLatch(0); - this.denyListedChannels = - RetryOnDifferentGrpcChannelErrorHandler.isEnabled() - ? CacheBuilder.newBuilder() - .expireAfterWrite(java.time.Duration.ofMinutes(1)) - .maximumSize(this.numChannels) - .concurrencyLevel(1) - .ticker( - new Ticker() { - @Override - public long read() { - return TimeUnit.NANOSECONDS.convert( - clock.instant().toEpochMilli(), TimeUnit.MILLISECONDS); - } - }) - .build() - : null; - } - - /** - * @return the {@link Dialect} of the underlying database. This method will block until the - * dialect is available. It will potentially execute one or two RPCs to get the dialect if - * necessary: One to create a session if there are no sessions in the pool (yet), and one to - * query the database for the dialect that is used. It is recommended that clients that always - * need to know the dialect set {@link - * SessionPoolOptions.Builder#setAutoDetectDialect(boolean)} to true. This will ensure that - * the dialect is fetched automatically in a background task when a session pool is created. - */ - Dialect getDialect() { - boolean mustDetectDialect = false; - synchronized (lock) { - if (!detectDialectStarted) { - mustDetectDialect = true; - detectDialectStarted = true; - } - } - if (mustDetectDialect) { - try (PooledSessionFuture session = getSession()) { - dialect.set(((PooledSession) session.get()).determineDialect()); - } - } - try { - return dialect.get(60L, TimeUnit.SECONDS); - } catch (ExecutionException executionException) { - throw asSpannerException(executionException); - } catch (InterruptedException interruptedException) { - throw SpannerExceptionFactory.propagateInterrupt(interruptedException); - } catch (TimeoutException timeoutException) { - throw SpannerExceptionFactory.propagateTimeout(timeoutException); - } - } - - Future getDialectAsync() { - return executor.submit(this::getDialect); - } - - PooledSessionReplacementHandler getPooledSessionReplacementHandler() { - return pooledSessionReplacementHandler; - } - - @Nullable - public String getDatabaseRole() { - return databaseRole; - } - - @VisibleForTesting - int getNumberOfSessionsInUse() { - synchronized (lock) { - return numSessionsInUse; - } - } - - @VisibleForTesting - int getMaxSessionsInUse() { - synchronized (lock) { - return maxSessionsInUse; - } - } - - @VisibleForTesting - double getRatioOfSessionsInUse() { - synchronized (lock) { - final int maxSessions = options.getMaxSessions(); - if (maxSessions == 0) { - return 0; - } - return (double) numSessionsInUse / maxSessions; - } - } - - boolean removeFromPool(PooledSession session) { - synchronized (lock) { - if (isClosed()) { - decrementPendingClosures(1); - return false; - } - session.markClosing(); - allSessions.remove(session); - return true; - } - } - - long numIdleSessionsRemoved() { - synchronized (lock) { - return numIdleSessionsRemoved; - } - } - - @VisibleForTesting - long numLeakedSessionsRemoved() { - synchronized (lock) { - return numLeakedSessionsRemoved; - } - } - - @VisibleForTesting - int getNumberOfSessionsInPool() { - synchronized (lock) { - return sessions.size(); - } - } - - @VisibleForTesting - int getNumberOfSessionsBeingCreated() { - synchronized (lock) { - return numSessionsBeingCreated; - } - } - - @VisibleForTesting - int getTotalSessionsPlusNumSessionsBeingCreated() { - synchronized (lock) { - return numSessionsBeingCreated + allSessions.size(); - } - } - - @VisibleForTesting - long getNumWaiterTimeouts() { - return numWaiterTimeouts.get(); - } - - private void initPool() { - synchronized (lock) { - poolMaintainer.init(); - if (options.getMinSessions() > 0) { - createSessions(options.getMinSessions(), true); - } - } - } - - private boolean isClosed() { - synchronized (lock) { - return closureFuture != null; - } - } - - private void handleException(SpannerException e, Tuple session) { - if (isSessionNotFound(e)) { - invalidateSession(session.x()); - } else { - releaseSession(session); - } - } - - private boolean isSessionNotFound(SpannerException e) { - return e.getErrorCode() == ErrorCode.NOT_FOUND && e.getMessage().contains("Session not found"); - } - - private boolean isDatabaseOrInstanceNotFound(SpannerException e) { - return e instanceof DatabaseNotFoundException || e instanceof InstanceNotFoundException; - } - - private void invalidateSession(PooledSession session) { - synchronized (lock) { - if (isClosed()) { - decrementPendingClosures(1); - return; - } - allSessions.remove(session); - // replenish the pool. - createSessions(getAllowedCreateSessions(1), false); - } - } - - private Tuple findSessionToKeepAlive( - Queue queue, Instant keepAliveThreshold, int numAlreadyChecked) { - int numChecked = 0; - Iterator iterator = queue.iterator(); - while (iterator.hasNext() - && (numChecked + numAlreadyChecked) - < (options.getMinSessions() + options.getMaxIdleSessions() - numSessionsInUse)) { - PooledSession session = iterator.next(); - if (session.delegate.getLastUseTime() != null - && session.delegate.getLastUseTime().isBefore(keepAliveThreshold)) { - iterator.remove(); - return Tuple.of(session, numChecked); - } - numChecked++; - } - return null; - } - - /** - * @return true if this {@link SessionPool} is still valid. - */ - boolean isValid() { - synchronized (lock) { - return closureFuture == null && resourceNotFoundException == null; - } - } - - /** - * Returns a multiplexed session. The method fallbacks to a regular session if {@link - * SessionPoolOptions#getUseMultiplexedSession} is not set. - */ - PooledSessionFutureWrapper getMultiplexedSessionWithFallback() throws SpannerException { - return new PooledSessionFutureWrapper(getSession()); - } - - /** - * Returns a session to be used for requests to spanner. This method is always non-blocking and - * returns a {@link PooledSessionFuture}. In case the pool is exhausted and {@link - * SessionPoolOptions#isFailIfPoolExhausted()} has been set, it will throw an exception. Returned - * session must be closed by calling {@link Session#close()}. - * - *

Implementation strategy: - * - *

    - *
  1. If a read session is available, return that. - *
  2. Otherwise if a session can be created, fire a creation request. - *
  3. Wait for a session to become available. Note that this can be unblocked either by a - * session being returned to the pool or a new session being created. - *
- */ - PooledSessionFuture getSession() throws SpannerException { - ISpan span = tracer.getCurrentSpan(); - span.addAnnotation("Acquiring session"); - WaiterFuture waiter = null; - PooledSession sess = null; - synchronized (lock) { - if (closureFuture != null) { - span.addAnnotation("Pool has been closed"); - throw new IllegalStateException("Pool has been closed", closedException); - } - if (resourceNotFoundException != null) { - span.addAnnotation("Database has been deleted"); - throw SpannerExceptionFactory.newSpannerException( - ErrorCode.NOT_FOUND, - String.format( - "The session pool has been invalidated because a previous RPC returned 'Database" - + " not found': %s", - resourceNotFoundException.getMessage()), - resourceNotFoundException); - } - if (denyListedChannels != null - && denyListedChannels.size() > 0 - && denyListedChannels.size() < numChannels) { - // There are deny-listed channels. Get a session that is not affiliated with a deny-listed - // channel. - for (PooledSession session : sessions) { - if (denyListedChannels.getIfPresent(session.getChannel()) == null) { - sessions.remove(session); - sess = session; - break; - } - // Size is cached and can change after calling getIfPresent. - if (denyListedChannels.size() == 0) { - break; - } - } - } - if (sess == null) { - sess = sessions.poll(); - } - if (sess == null) { - span.addAnnotation("No session available"); - maybeCreateSession(); - waiter = new WaiterFuture(); - waiters.add(waiter); - } else { - span.addAnnotation("Acquired session"); - } - return checkoutSession(span, sess, waiter); - } - } - - private PooledSessionFuture checkoutSession( - final ISpan span, final PooledSession readySession, WaiterFuture waiter) { - ListenableFuture sessionFuture; - if (waiter != null) { - logger.log( - Level.FINE, - "No session available in the pool. Blocking for one to become available/created"); - span.addAnnotation("Waiting for a session to come available"); - sessionFuture = waiter; - } else { - SettableFuture fut = SettableFuture.create(); - fut.set(readySession); - sessionFuture = fut; - } - PooledSessionFuture res = createPooledSessionFuture(sessionFuture, span); - res.markCheckedOut(); - return res; - } - - private void incrementNumSessionsInUse() { - synchronized (lock) { - if (maxSessionsInUse < ++numSessionsInUse) { - maxSessionsInUse = numSessionsInUse; - } - numSessionsAcquired++; - } - } - - private void maybeCreateSession() { - ISpan span = tracer.getCurrentSpan(); - boolean throwResourceExhaustedException = false; - synchronized (lock) { - if (numWaiters() >= numSessionsBeingCreated) { - if (canCreateSession()) { - span.addAnnotation("Creating sessions"); - createSessions(getAllowedCreateSessions(options.getIncStep()), false); - } else if (options.isFailIfPoolExhausted()) { - throwResourceExhaustedException = true; - } - } - } - if (!throwResourceExhaustedException) { - return; - } - span.addAnnotation("Pool exhausted. Failing"); - - String message = - "No session available in the pool. Maximum number of sessions in the pool can be overridden" - + " by invoking SessionPoolOptions#Builder#setMaxSessions. Client can be made to block" - + " rather than fail by setting SessionPoolOptions#Builder#setBlockIfPoolExhausted.\n" - + createCheckedOutSessionsStackTraces(); - throw newSpannerException(ErrorCode.RESOURCE_EXHAUSTED, message); - } - - private StringBuilder createCheckedOutSessionsStackTraces() { - List currentlyCheckedOutSessions; - synchronized (lock) { - currentlyCheckedOutSessions = new ArrayList<>(this.markedCheckedOutSessions); - } - - // Create the error message without holding the lock, as we are potentially looping through a - // large set, and analyzing a large number of stack traces. - StringBuilder stackTraces = - new StringBuilder("MinSessions: ") - .append(options.getMinSessions()) - .append("\nMaxSessions: ") - .append(options.getMaxSessions()) - .append("\nThere are currently ") - .append(currentlyCheckedOutSessions.size()) - .append(" sessions checked out:\n\n"); - if (options.isTrackStackTraceOfSessionCheckout()) { - for (PooledSessionFuture session : currentlyCheckedOutSessions) { - if (session.leakedException != null) { - StringWriter writer = new StringWriter(); - PrintWriter printWriter = new PrintWriter(writer); - session.leakedException.printStackTrace(printWriter); - stackTraces.append(writer).append("\n\n"); - } - } - } - return stackTraces; - } - - private void releaseSession(Tuple sessionWithPosition) { - releaseSession(sessionWithPosition.x(), false, sessionWithPosition.y()); - } - - private void releaseSession(PooledSession session, boolean isNewSession) { - releaseSession(session, isNewSession, null); - } - - /** Releases a session back to the pool. This might cause one of the waiters to be unblocked. */ - private void releaseSession( - PooledSession session, boolean isNewSession, @Nullable Integer position) { - Preconditions.checkNotNull(session); - synchronized (lock) { - if (closureFuture != null) { - return; - } - if (waiters.isEmpty()) { - // There are no pending waiters. - // Add to a random position if the transactions per second is high or the head of the - // session pool already contains many sessions with the same channel as this one. - if (session.releaseToPosition != Position.RANDOM && shouldRandomize()) { - session.releaseToPosition = Position.RANDOM; - } else if (session.releaseToPosition == Position.FIRST && isUnbalanced(session)) { - session.releaseToPosition = Position.RANDOM; - } else if (session.releaseToPosition == Position.RANDOM - && !isNewSession - && checkedOutSessions.size() <= 2) { - // Do not randomize if there are few other sessions checked out and this session has been - // used. This ensures that this session will be re-used for the next transaction, which is - // more efficient. - session.releaseToPosition = options.getReleaseToPosition(); - } - if (position != null) { - // Make sure we use a valid position, as the number of sessions could have changed in the - // meantime. - int actualPosition = Math.min(position, sessions.size()); - sessions.add(actualPosition, session); - } else if (session.releaseToPosition == Position.RANDOM && !sessions.isEmpty()) { - // A session should only be added at a random position the first time it is added to - // the pool or if the pool was deemed unbalanced. All following releases into the pool - // should normally happen at the default release position (unless the pool is again deemed - // to be unbalanced and the insertion would happen at the front of the pool). - session.releaseToPosition = options.getReleaseToPosition(); - int pos = random.nextInt(sessions.size() + 1); - sessions.add(pos, session); - } else if (session.releaseToPosition == Position.LAST) { - sessions.addLast(session); - } else { - sessions.addFirst(session); - } - session.releaseToPosition = options.getReleaseToPosition(); - } else { - waiters.poll().put(session); - } - } - } - - /** - * Returns true if the position where we return the session should be random if: - * - *
    - *
  1. The current TPS is higher than the configured threshold. - *
  2. AND the number of sessions checked out is larger than the number of channels. - *
- * - * The second check prevents the session pool from being randomized when the application is - * running many small, quick queries using a small number of parallel threads. This can cause a - * high TPS, without actually having a high degree of parallelism. - */ - @VisibleForTesting - boolean shouldRandomize() { - return this.options.getRandomizePositionQPSThreshold() > 0 - && this.transactionsPerSecond >= this.options.getRandomizePositionQPSThreshold() - && this.numSessionsInUse >= this.numChannels; - } - - private boolean isUnbalanced(PooledSession session) { - int channel = session.getChannel(); - int numChannels = sessionClient.getSpanner().getOptions().getNumChannels(); - return isUnbalanced(channel, this.sessions, this.checkedOutSessions, numChannels); - } - - /** - * Returns true if the given list of sessions is considered unbalanced when compared to the - * sessionChannel that is about to be added to the pool. - * - *

The method returns true if all the following is true: - * - *

    - *
  1. The list of sessions is not empty. - *
  2. The number of checked out sessions is > 2. - *
  3. The number of channels being used by the pool is > 1. - *
  4. And at least one of the following is true: - *
      - *
    1. The first numChannels sessions in the list of sessions contains more than 2 - * sessions that use the same channel as the one being added. - *
    2. The list of currently checked out sessions contains more than 2 times the the - * number of sessions with the same channel as the one being added than it should in - * order for it to be perfectly balanced. Perfectly balanced in this case means that - * the list should preferably contain size/numChannels sessions of each channel. - *
    - *
- * - * @param channelOfSessionBeingAdded the channel number being used by the session that is about to - * be released into the pool - * @param sessions the list of all sessions in the pool - * @param checkedOutSessions the currently checked out sessions of the pool - * @param numChannels the number of channels in use - * @return true if the pool is considered unbalanced, and false otherwise - */ - @VisibleForTesting - static boolean isUnbalanced( - int channelOfSessionBeingAdded, - List sessions, - Set checkedOutSessions, - int numChannels) { - // Do not re-balance the pool if the number of checked out sessions is low, as it is - // better to re-use sessions as much as possible in a low-QPS scenario. - if (sessions.isEmpty() || checkedOutSessions.size() <= 2) { - return false; - } - if (numChannels == 1) { - return false; - } - - // Ideally, the first numChannels sessions in the pool should contain exactly one session for - // each channel. - // Check if the first numChannels sessions at the head of the pool already contain more than 2 - // sessions that use the same channel as this one. If so, we re-balance. - // We also re-balance the pool in the specific case that the pool uses 2 channels and the first - // two sessions use those two channels. - int maxSessionsAtHeadOfPool = Math.min(numChannels, 3); - int count = 0; - for (int i = 0; i < Math.min(numChannels, sessions.size()); i++) { - PooledSession otherSession = sessions.get(i); - if (channelOfSessionBeingAdded == otherSession.getChannel()) { - count++; - if (count >= maxSessionsAtHeadOfPool) { - return true; - } - } - } - // Ideally, the use of a channel in the checked out sessions is exactly - // numCheckedOut / numChannels - // We check whether we are more than a factor two away from that perfect distribution. - // If we are, then we re-balance. - count = 0; - int checkedOutThreshold = Math.max(2, 2 * checkedOutSessions.size() / numChannels); - for (PooledSessionFuture otherSession : checkedOutSessions) { - if (otherSession.isDone() && channelOfSessionBeingAdded == otherSession.get().getChannel()) { - count++; - if (count > checkedOutThreshold) { - return true; - } - } - } - return false; - } - - private void handleCreateSessionsFailure(SpannerException e, int count) { - synchronized (lock) { - for (int i = 0; i < count; i++) { - if (!waiters.isEmpty()) { - waiters.poll().put(e); - } else { - break; - } - } - if (!dialect.isDone()) { - dialect.setException(e); - } - if (isDatabaseOrInstanceNotFound(e)) { - setResourceNotFoundException((ResourceNotFoundException) e); - poolMaintainer.close(); - } - } - } - - void setResourceNotFoundException(ResourceNotFoundException e) { - this.resourceNotFoundException = MoreObjects.firstNonNull(this.resourceNotFoundException, e); - } - - private void decrementPendingClosures(int count) { - pendingClosure -= count; - if (pendingClosure == 0) { - closureFuture.set(null); - } - } - - /** - * Close all the sessions. Once this method is invoked {@link #getSession()} will start throwing - * {@code IllegalStateException}. The returned future blocks till all the sessions created in this - * pool have been closed. - */ - ListenableFuture closeAsync(ClosedException closedException) { - ListenableFuture retFuture = null; - synchronized (lock) { - if (closureFuture != null) { - throw new IllegalStateException("Close has already been invoked", this.closedException); - } - this.closedException = closedException; - // Fail all pending waiters. - WaiterFuture waiter = waiters.poll(); - while (waiter != null) { - waiter.put(newSpannerException(ErrorCode.INTERNAL, "Client has been closed")); - waiter = waiters.poll(); - } - closureFuture = SettableFuture.create(); - retFuture = closureFuture; - - pendingClosure = totalSessions() + numSessionsBeingCreated; - - if (!poolMaintainer.isClosed()) { - pendingClosure += 1; // For pool maintenance thread - poolMaintainer.close(); - } - - sessions.clear(); - for (PooledSessionFuture session : checkedOutSessions) { - if (session.leakedException != null) { - if (options.isFailOnSessionLeak()) { - throw session.leakedException; - } else { - logger.log(Level.WARNING, "Leaked session", session.leakedException); - } - } else { - String message = - "Leaked session. Call" - + " SessionOptions.Builder#setTrackStackTraceOfSessionCheckout(true) to start" - + " tracking the call stack trace of the thread that checked out the session."; - if (options.isFailOnSessionLeak()) { - throw new LeakedSessionException(message); - } else { - logger.log(Level.WARNING, message); - } - } - } - for (final PooledSession session : ImmutableList.copyOf(allSessions)) { - if (session.state != SessionState.CLOSING) { - closeSessionAsync(session); - } - } - - // Nothing to be closed, mark as complete - if (pendingClosure == 0) { - closureFuture.set(null); - } - } - - retFuture.addListener(() -> executorFactory.release(executor), MoreExecutors.directExecutor()); - return retFuture; - } - - private int numWaiters() { - synchronized (lock) { - return waiters.size(); - } - } - - @VisibleForTesting - int totalSessions() { - synchronized (lock) { - return allSessions.size(); - } - } - - @VisibleForTesting - int numSessionsInPool() { - synchronized (lock) { - return sessions.size(); - } - } - - private ApiFuture closeSessionAsync(final PooledSession sess) { - ApiFuture res = sess.delegate.asyncClose(); - res.addListener( - () -> { - synchronized (lock) { - allSessions.remove(sess); - if (isClosed()) { - decrementPendingClosures(1); - return; - } - // Create a new session if needed to unblock some waiter. - if (numWaiters() > numSessionsBeingCreated) { - createSessions( - getAllowedCreateSessions(numWaiters() - numSessionsBeingCreated), false); - } - } - }, - MoreExecutors.directExecutor()); - return res; - } - - /** - * Returns the minimum of the wanted number of sessions that the caller wants to create and the - * actual max number that may be created at this moment. - */ - private int getAllowedCreateSessions(int wantedSessions) { - synchronized (lock) { - return Math.min( - wantedSessions, options.getMaxSessions() - (totalSessions() + numSessionsBeingCreated)); - } - } - - private boolean canCreateSession() { - synchronized (lock) { - return totalSessions() + numSessionsBeingCreated < options.getMaxSessions(); - } - } - - private void createSessions(final int sessionCount, boolean distributeOverChannels) { - logger.log(Level.FINE, String.format("Creating %d sessions", sessionCount)); - synchronized (lock) { - numSessionsBeingCreated += sessionCount; - try { - // Create a batch of sessions. The actual session creation can be split into multiple gRPC - // calls and the session consumer consumes the returned sessions as they become available. - // The batchCreateSessions method automatically spreads the sessions evenly over all - // available channels. - sessionClient.asyncBatchCreateSessions( - sessionCount, distributeOverChannels, sessionConsumer); - } catch (Throwable t) { - // Expose this to customer via a metric. - numSessionsBeingCreated -= sessionCount; - if (isClosed()) { - decrementPendingClosures(sessionCount); - } - handleCreateSessionsFailure(newSpannerException(t), sessionCount); - } - } - } - - /** - * {@link SessionConsumer} that receives the created sessions from a {@link SessionClient} and - * releases these into the pool. The session pool only needs one instance of this, as all sessions - * should be returned to the same pool regardless of what triggered the creation of the sessions. - */ - class SessionConsumerImpl implements SessionConsumer { - /** Release a new session to the pool. */ - @Override - public void onSessionReady(SessionImpl session) { - PooledSession pooledSession = null; - boolean closeSession = false; - synchronized (lock) { - int minSessions = options.getMinSessions(); - pooledSession = new PooledSession(session); - numSessionsBeingCreated--; - if (closureFuture != null) { - closeSession = true; - } else { - Preconditions.checkState(totalSessions() <= options.getMaxSessions() - 1); - allSessions.add(pooledSession); - if (allSessions.size() >= minSessions) { - waitOnMinSessionsLatch.countDown(); - } - if (options.isAutoDetectDialect() - && !detectDialectStarted - && !options.getUseMultiplexedSession()) { - // Get the dialect of the underlying database if that has not yet been done. Note that - // this method will release the session into the pool once it is done. - detectDialectStarted = true; - pooledSession.determineDialectAsync(SessionPool.this.dialect); - } else { - // Release the session to a random position in the pool to prevent the case that a batch - // of sessions that are affiliated with the same channel are all placed sequentially in - // the pool. - releaseSession(pooledSession, true); - } - } - } - if (closeSession) { - closeSessionAsync(pooledSession); - } - } - - /** - * Informs waiters for a session that session creation failed. The exception will propagate to - * the waiters as a {@link SpannerException}. - */ - @Override - public void onSessionCreateFailure(Throwable t, int createFailureForSessionCount) { - synchronized (lock) { - numSessionsBeingCreated -= createFailureForSessionCount; - if (numSessionsBeingCreated == 0) { - // Don't continue to block if no more sessions are being created. - waitOnMinSessionsLatch.countDown(); - } - if (isClosed()) { - decrementPendingClosures(createFailureForSessionCount); - } - handleCreateSessionsFailure(newSpannerException(t), createFailureForSessionCount); - } - } - } - - /** - * Initializes and creates Spanner session relevant metrics using OpenCensus. When coupled with an - * exporter, it allows users to monitor client behavior. - */ - private void initOpenCensusMetricsCollection( - MetricRegistry metricRegistry, - List labelValues, - AtomicLong numMultiplexedSessionsAcquired, - AtomicLong numMultiplexedSessionsReleased) { - if (!SpannerOptions.isEnabledOpenCensusMetrics()) { - return; - } - DerivedLongGauge maxInUseSessionsMetric = - metricRegistry.addDerivedLongGauge( - METRIC_PREFIX + MAX_IN_USE_SESSIONS, - MetricOptions.builder() - .setDescription(MAX_IN_USE_SESSIONS_DESCRIPTION) - .setUnit(COUNT) - .setLabelKeys(SPANNER_LABEL_KEYS) - .build()); - - DerivedLongGauge maxAllowedSessionsMetric = - metricRegistry.addDerivedLongGauge( - METRIC_PREFIX + MAX_ALLOWED_SESSIONS, - MetricOptions.builder() - .setDescription(MAX_ALLOWED_SESSIONS_DESCRIPTION) - .setUnit(COUNT) - .setLabelKeys(SPANNER_LABEL_KEYS) - .build()); - - DerivedLongCumulative sessionsTimeouts = - metricRegistry.addDerivedLongCumulative( - METRIC_PREFIX + GET_SESSION_TIMEOUTS, - MetricOptions.builder() - .setDescription(SESSIONS_TIMEOUTS_DESCRIPTION) - .setUnit(COUNT) - .setLabelKeys(SPANNER_LABEL_KEYS) - .build()); - - DerivedLongCumulative numAcquiredSessionsMetric = - metricRegistry.addDerivedLongCumulative( - METRIC_PREFIX + NUM_ACQUIRED_SESSIONS, - MetricOptions.builder() - .setDescription(NUM_ACQUIRED_SESSIONS_DESCRIPTION) - .setUnit(COUNT) - .setLabelKeys(SPANNER_LABEL_KEYS_WITH_MULTIPLEXED_SESSIONS) - .build()); - - DerivedLongCumulative numReleasedSessionsMetric = - metricRegistry.addDerivedLongCumulative( - METRIC_PREFIX + NUM_RELEASED_SESSIONS, - MetricOptions.builder() - .setDescription(NUM_RELEASED_SESSIONS_DESCRIPTION) - .setUnit(COUNT) - .setLabelKeys(SPANNER_LABEL_KEYS_WITH_MULTIPLEXED_SESSIONS) - .build()); - - DerivedLongGauge numSessionsInPoolMetric = - metricRegistry.addDerivedLongGauge( - METRIC_PREFIX + NUM_SESSIONS_IN_POOL, - MetricOptions.builder() - .setDescription(NUM_SESSIONS_IN_POOL_DESCRIPTION) - .setUnit(COUNT) - .setLabelKeys(SPANNER_LABEL_KEYS_WITH_TYPE) - .build()); - - // The value of a maxSessionsInUse is observed from a callback function. This function is - // invoked whenever metrics are collected. - maxInUseSessionsMetric.removeTimeSeries(labelValues); - maxInUseSessionsMetric.createTimeSeries( - labelValues, this, sessionPool -> sessionPool.maxSessionsInUse); - - // The value of a maxSessions is observed from a callback function. This function is invoked - // whenever metrics are collected. - maxAllowedSessionsMetric.removeTimeSeries(labelValues); - maxAllowedSessionsMetric.createTimeSeries( - labelValues, options, SessionPoolOptions::getMaxSessions); - - // The value of a numWaiterTimeouts is observed from a callback function. This function is - // invoked whenever metrics are collected. - sessionsTimeouts.removeTimeSeries(labelValues); - sessionsTimeouts.createTimeSeries(labelValues, this, SessionPool::getNumWaiterTimeouts); - - List labelValuesWithRegularSessions = new ArrayList<>(labelValues); - List labelValuesWithMultiplexedSessions = new ArrayList<>(labelValues); - labelValuesWithMultiplexedSessions.add(LabelValue.create("true")); - labelValuesWithRegularSessions.add(LabelValue.create("false")); - - numAcquiredSessionsMetric.removeTimeSeries(labelValuesWithRegularSessions); - numAcquiredSessionsMetric.createTimeSeries( - labelValuesWithRegularSessions, this, sessionPool -> sessionPool.numSessionsAcquired); - numAcquiredSessionsMetric.removeTimeSeries(labelValuesWithMultiplexedSessions); - numAcquiredSessionsMetric.createTimeSeries( - labelValuesWithMultiplexedSessions, this, unused -> numMultiplexedSessionsAcquired.get()); - - numReleasedSessionsMetric.removeTimeSeries(labelValuesWithRegularSessions); - numReleasedSessionsMetric.createTimeSeries( - labelValuesWithRegularSessions, this, sessionPool -> sessionPool.numSessionsReleased); - numReleasedSessionsMetric.removeTimeSeries(labelValuesWithMultiplexedSessions); - numReleasedSessionsMetric.createTimeSeries( - labelValuesWithMultiplexedSessions, this, unused -> numMultiplexedSessionsReleased.get()); - - List labelValuesWithBeingPreparedType = new ArrayList<>(labelValues); - labelValuesWithBeingPreparedType.add(NUM_SESSIONS_BEING_PREPARED); - numSessionsInPoolMetric.removeTimeSeries(labelValuesWithBeingPreparedType); - numSessionsInPoolMetric.createTimeSeries( - labelValuesWithBeingPreparedType, - this, - // TODO: Remove metric. - ignored -> 0L); - - List labelValuesWithInUseType = new ArrayList<>(labelValues); - labelValuesWithInUseType.add(NUM_IN_USE_SESSIONS); - numSessionsInPoolMetric.removeTimeSeries(labelValuesWithInUseType); - numSessionsInPoolMetric.createTimeSeries( - labelValuesWithInUseType, this, sessionPool -> sessionPool.numSessionsInUse); - - List labelValuesWithReadType = new ArrayList<>(labelValues); - labelValuesWithReadType.add(NUM_READ_SESSIONS); - numSessionsInPoolMetric.removeTimeSeries(labelValuesWithReadType); - numSessionsInPoolMetric.createTimeSeries( - labelValuesWithReadType, this, sessionPool -> sessionPool.sessions.size()); - - List labelValuesWithWriteType = new ArrayList<>(labelValues); - labelValuesWithWriteType.add(NUM_WRITE_SESSIONS); - numSessionsInPoolMetric.removeTimeSeries(labelValuesWithWriteType); - numSessionsInPoolMetric.createTimeSeries( - labelValuesWithWriteType, - this, - // TODO: Remove metric. - ignored -> 0L); - } - - /** - * Initializes and creates Spanner session relevant metrics using OpenTelemetry. When coupled with - * an exporter, it allows users to monitor client behavior. - */ - private void initOpenTelemetryMetricsCollection( - OpenTelemetry openTelemetry, - Attributes attributes, - AtomicLong numMultiplexedSessionsAcquired, - AtomicLong numMultiplexedSessionsReleased) { - if (openTelemetry == null || !SpannerOptions.isEnabledOpenTelemetryMetrics()) { - return; - } - - Meter meter = openTelemetry.getMeter(MetricRegistryConstants.INSTRUMENTATION_SCOPE); - meter - .gaugeBuilder(MAX_ALLOWED_SESSIONS) - .setDescription(MAX_ALLOWED_SESSIONS_DESCRIPTION) - .setUnit(COUNT) - .buildWithCallback( - measurement -> { - // Although Max sessions is a constant value, OpenTelemetry requires to define this as - // a callback. - measurement.record(options.getMaxSessions(), attributes); - }); - - meter - .gaugeBuilder(MAX_IN_USE_SESSIONS) - .setDescription(MAX_IN_USE_SESSIONS_DESCRIPTION) - .setUnit(COUNT) - .buildWithCallback( - measurement -> { - measurement.record(this.maxSessionsInUse, attributes); - }); - - AttributesBuilder attributesBuilder; - if (attributes != null) { - attributesBuilder = attributes.toBuilder(); - } else { - attributesBuilder = Attributes.builder(); - } - Attributes attributesInUseSessions = - attributesBuilder.put(SESSIONS_TYPE, NUM_SESSIONS_IN_USE).build(); - Attributes attributesAvailableSessions = - attributesBuilder.put(SESSIONS_TYPE, NUM_SESSIONS_AVAILABLE).build(); - meter - .upDownCounterBuilder(NUM_SESSIONS_IN_POOL) - .setDescription(NUM_SESSIONS_IN_POOL_DESCRIPTION) - .setUnit(COUNT) - .buildWithCallback( - measurement -> { - measurement.record(this.numSessionsInUse, attributesInUseSessions); - measurement.record(this.sessions.size(), attributesAvailableSessions); - }); - - AttributesBuilder attributesBuilderIsMultiplexed; - if (attributes != null) { - attributesBuilderIsMultiplexed = attributes.toBuilder(); - } else { - attributesBuilderIsMultiplexed = Attributes.builder(); - } - Attributes attributesRegularSession = - attributesBuilderIsMultiplexed.put(IS_MULTIPLEXED, false).build(); - Attributes attributesMultiplexedSession = - attributesBuilderIsMultiplexed.put(IS_MULTIPLEXED, true).build(); - meter - .counterBuilder(GET_SESSION_TIMEOUTS) - .setDescription(SESSIONS_TIMEOUTS_DESCRIPTION) - .setUnit(COUNT) - .buildWithCallback( - measurement -> { - measurement.record(this.getNumWaiterTimeouts(), attributes); - }); - - meter - .counterBuilder(NUM_ACQUIRED_SESSIONS) - .setDescription(NUM_ACQUIRED_SESSIONS_DESCRIPTION) - .setUnit(COUNT) - .buildWithCallback( - measurement -> { - measurement.record(this.numSessionsAcquired, attributesRegularSession); - measurement.record( - numMultiplexedSessionsAcquired.get(), attributesMultiplexedSession); - }); - - meter - .counterBuilder(NUM_RELEASED_SESSIONS) - .setDescription(NUM_RELEASED_SESSIONS_DESCRIPTION) - .setUnit(COUNT) - .buildWithCallback( - measurement -> { - measurement.record(this.numSessionsReleased, attributesRegularSession); - measurement.record( - numMultiplexedSessionsReleased.get(), attributesMultiplexedSession); - }); - } -} diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPoolAsyncTransactionManager.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPoolAsyncTransactionManager.java deleted file mode 100644 index 5e48d1b78bc..00000000000 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPoolAsyncTransactionManager.java +++ /dev/null @@ -1,287 +0,0 @@ -/* - * Copyright 2020 Google LLC - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package com.google.cloud.spanner; - -import com.google.api.core.ApiFuture; -import com.google.api.core.ApiFutureCallback; -import com.google.api.core.ApiFutures; -import com.google.api.core.SettableApiFuture; -import com.google.cloud.Timestamp; -import com.google.cloud.spanner.Options.TransactionOption; -import com.google.cloud.spanner.SessionPool.SessionFuture; -import com.google.cloud.spanner.SessionPool.SessionNotFoundHandler; -import com.google.cloud.spanner.SessionPool.SessionReplacementHandler; -import com.google.cloud.spanner.TransactionContextFutureImpl.CommittableAsyncTransactionManager; -import com.google.cloud.spanner.TransactionManager.TransactionState; -import com.google.common.base.Preconditions; -import com.google.common.util.concurrent.MoreExecutors; -import javax.annotation.concurrent.GuardedBy; - -class SessionPoolAsyncTransactionManager - implements CommittableAsyncTransactionManager, SessionNotFoundHandler { - private final Object lock = new Object(); - - @GuardedBy("lock") - private TransactionState txnState; - - @GuardedBy("lock") - private AbortedException abortedException; - - private final SessionReplacementHandler sessionReplacementHandler; - private final TransactionOption[] options; - private volatile I session; - private volatile SettableApiFuture delegate; - private boolean restartedAfterSessionNotFound; - - SessionPoolAsyncTransactionManager( - SessionReplacementHandler sessionReplacementHandler, - I session, - TransactionOption... options) { - this.options = options; - this.sessionReplacementHandler = sessionReplacementHandler; - createTransaction(session); - } - - private void createTransaction(I session) { - this.session = session; - this.delegate = SettableApiFuture.create(); - this.session.addListener( - () -> { - try { - delegate.set( - SessionPoolAsyncTransactionManager.this - .session - .get() - .transactionManagerAsync(options)); - } catch (Throwable t) { - delegate.setException(t); - } - }, - MoreExecutors.directExecutor()); - } - - @Override - public SpannerException handleSessionNotFound(SessionNotFoundException notFound) { - // Restart the entire transaction with a new session and throw an AbortedException to force the - // client application to retry. - createTransaction(sessionReplacementHandler.replaceSession(notFound, session)); - restartedAfterSessionNotFound = true; - return SpannerExceptionFactory.newSpannerException( - ErrorCode.ABORTED, notFound.getMessage(), notFound); - } - - @Override - public void close() { - SpannerApiFutures.get(closeAsync()); - } - - @Override - public ApiFuture closeAsync() { - final SettableApiFuture res = SettableApiFuture.create(); - ApiFutures.addCallback( - delegate, - new ApiFutureCallback() { - @Override - public void onFailure(Throwable t) { - session.close(); - } - - @Override - public void onSuccess(AsyncTransactionManagerImpl result) { - ApiFutures.addCallback( - result.closeAsync(), - new ApiFutureCallback() { - @Override - public void onFailure(Throwable t) { - session.close(); - res.setException(t); - } - - @Override - public void onSuccess(Void result) { - session.close(); - res.set(result); - } - }, - MoreExecutors.directExecutor()); - } - }, - MoreExecutors.directExecutor()); - return res; - } - - @Override - public TransactionContextFuture beginAsync() { - synchronized (lock) { - Preconditions.checkState(txnState == null, "begin can only be called once"); - txnState = TransactionState.STARTED; - } - final SettableApiFuture delegateTxnFuture = SettableApiFuture.create(); - ApiFutures.addCallback( - delegate, - new ApiFutureCallback() { - @Override - public void onFailure(Throwable t) { - delegateTxnFuture.setException(t); - } - - @Override - public void onSuccess(AsyncTransactionManagerImpl result) { - ApiFutures.addCallback( - result.beginAsync(), - new ApiFutureCallback() { - @Override - public void onFailure(Throwable t) { - delegateTxnFuture.setException(t); - } - - @Override - public void onSuccess(TransactionContext result) { - delegateTxnFuture.set( - new SessionPool.SessionPoolTransactionContext( - SessionPoolAsyncTransactionManager.this, result)); - } - }, - MoreExecutors.directExecutor()); - } - }, - MoreExecutors.directExecutor()); - return new TransactionContextFutureImpl(this, delegateTxnFuture); - } - - @Override - public TransactionContextFuture beginAsync(AbortedException exception) { - // For regular sessions, the input exception is ignored and the behavior is equivalent to - // calling {@link #beginAsync()}. - return beginAsync(); - } - - @Override - public void onError(Throwable t) { - if (t instanceof AbortedException) { - synchronized (lock) { - txnState = TransactionState.ABORTED; - abortedException = (AbortedException) t; - } - } - } - - @Override - public ApiFuture commitAsync() { - synchronized (lock) { - Preconditions.checkState( - txnState == TransactionState.STARTED || txnState == TransactionState.ABORTED, - "commit can only be invoked if the transaction is in progress. Current state: " - + txnState); - if (txnState == TransactionState.ABORTED) { - return ApiFutures.immediateFailedFuture(abortedException); - } - txnState = TransactionState.COMMITTED; - } - return ApiFutures.transformAsync( - delegate, - input -> { - final SettableApiFuture res = SettableApiFuture.create(); - ApiFutures.addCallback( - input.commitAsync(), - new ApiFutureCallback() { - @Override - public void onFailure(Throwable t) { - synchronized (lock) { - if (t instanceof AbortedException) { - txnState = TransactionState.ABORTED; - abortedException = (AbortedException) t; - } else { - txnState = TransactionState.COMMIT_FAILED; - } - } - res.setException(t); - } - - @Override - public void onSuccess(Timestamp result) { - res.set(result); - } - }, - MoreExecutors.directExecutor()); - return res; - }, - MoreExecutors.directExecutor()); - } - - @Override - public ApiFuture rollbackAsync() { - synchronized (lock) { - Preconditions.checkState( - txnState == TransactionState.STARTED, - "rollback can only be called if the transaction is in progress"); - txnState = TransactionState.ROLLED_BACK; - } - return ApiFutures.transformAsync( - delegate, - input -> { - ApiFuture res = input.rollbackAsync(); - res.addListener(() -> session.close(), MoreExecutors.directExecutor()); - return res; - }, - MoreExecutors.directExecutor()); - } - - @Override - public TransactionContextFuture resetForRetryAsync() { - synchronized (lock) { - Preconditions.checkState( - txnState == TransactionState.ABORTED || restartedAfterSessionNotFound, - "resetForRetry can only be called after the transaction aborted."); - txnState = TransactionState.STARTED; - } - return new TransactionContextFutureImpl( - this, - ApiFutures.transform( - ApiFutures.transformAsync( - delegate, - input -> { - if (restartedAfterSessionNotFound) { - restartedAfterSessionNotFound = false; - return input.beginAsync(); - } - return input.resetForRetryAsync(); - }, - MoreExecutors.directExecutor()), - input -> - new SessionPool.SessionPoolTransactionContext( - SessionPoolAsyncTransactionManager.this, input), - MoreExecutors.directExecutor())); - } - - @Override - public TransactionState getState() { - synchronized (lock) { - return txnState; - } - } - - public ApiFuture getCommitResponse() { - synchronized (lock) { - Preconditions.checkState( - txnState == TransactionState.COMMITTED, - "commit can only be invoked if the transaction was successfully committed"); - } - return ApiFutures.transformAsync( - delegate, AsyncTransactionManagerImpl::getCommitResponse, MoreExecutors.directExecutor()); - } -} diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPoolOptions.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPoolOptions.java index 605f96da74b..51d0ca3d476 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPoolOptions.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPoolOptions.java @@ -21,7 +21,6 @@ import com.google.api.core.InternalApi; import com.google.api.core.ObsoleteApi; -import com.google.cloud.spanner.SessionPool.Position; import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Preconditions; import java.time.Duration; @@ -29,7 +28,14 @@ import java.util.Objects; /** Options for the session pool used by {@code DatabaseClient}. */ +@Deprecated public class SessionPoolOptions { + enum Position { + FIRST, + LAST, + RANDOM + } + // Default number of channels * 100. private static final int DEFAULT_MAX_SESSIONS = 400; private static final int DEFAULT_MIN_SESSIONS = 100; @@ -583,10 +589,10 @@ public static class Builder { /** * Capture the call stack of the thread that checked out a session of the pool. This will - * pre-create a {@link com.google.cloud.spanner.SessionPool.LeakedSessionException} already when - * a session is checked out. This can be disabled by users, for example if their monitoring - * systems log the pre-created exception. If disabled, the {@link - * com.google.cloud.spanner.SessionPool.LeakedSessionException} will only be created when an + * pre-create a com.google.cloud.spanner.SessionPool.LeakedSessionException already when a + * session is checked out. This can be disabled by users, for example if their monitoring + * systems log the pre-created exception. If disabled, the + * com.google.cloud.spanner.SessionPool.LeakedSessionException will only be created when an * actual session leak is detected. The stack trace of the exception will in that case not * contain the call stack of when the session was checked out. */ @@ -945,8 +951,8 @@ Builder setInitialWaitForSessionTimeoutMillis(long timeout) { } /** - * If a session has been invalidated by the server, the {@link SessionPool} will by default - * retry the session. Set this option to throw an exception instead of retrying. + * If a session has been invalidated by the server, the SessionPool will by default retry the + * session. Set this option to throw an exception instead of retrying. */ @VisibleForTesting Builder setFailIfSessionNotFound() { @@ -962,8 +968,8 @@ Builder setFailOnSessionLeak() { /** * Sets whether the session pool should capture the call stack trace when a session is checked - * out of the pool. This will internally prepare a {@link - * com.google.cloud.spanner.SessionPool.LeakedSessionException} that will only be thrown if the + * out of the pool. This will internally prepare a + * com.google.cloud.spanner.SessionPool.LeakedSessionException that will only be thrown if the * session is actually leaked. This makes it easier to debug session leaks, as the stack trace * of the thread that checked out the session will be available in the exception. * @@ -1015,8 +1021,8 @@ public Builder setAcquireSessionTimeout(org.threeten.bp.Duration acquireSessionT } /** - * If greater than zero, we wait for said duration when no sessions are available in the {@link - * SessionPool}. The default is a 60s timeout. Set the value to null to disable the timeout. + * If greater than zero, we wait for said duration when no sessions are available in the + * SessionPool. The default is a 60s timeout. Set the value to null to disable the timeout. */ public Builder setAcquireSessionTimeoutDuration(Duration acquireSessionTimeout) { try { diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionReference.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionReference.java index e96be9effaa..1fd6c303ede 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionReference.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionReference.java @@ -33,15 +33,17 @@ class SessionReference { private final String name; private final DatabaseId databaseId; + @Nullable private final String databaseRole; private final Map options; private volatile Instant lastUseTime; @Nullable private final Instant createTime; private final boolean isMultiplexed; - SessionReference(String name, Map options) { + SessionReference(String name, @Nullable String databaseRole, Map options) { this.options = options; this.name = checkNotNull(name); this.databaseId = SessionId.of(name).getDatabaseId(); + this.databaseRole = databaseRole; this.lastUseTime = Instant.now(); this.createTime = null; this.isMultiplexed = false; @@ -49,12 +51,14 @@ class SessionReference { SessionReference( String name, + @Nullable String databaseRole, com.google.protobuf.Timestamp createTime, boolean isMultiplexed, Map options) { this.options = options; this.name = checkNotNull(name); this.databaseId = SessionId.of(name).getDatabaseId(); + this.databaseRole = databaseRole; this.lastUseTime = Instant.now(); this.createTime = convert(createTime); this.isMultiplexed = isMultiplexed; @@ -64,6 +68,10 @@ public String getName() { return name; } + public String getDatabaseRole() { + return databaseRole; + } + public DatabaseId getDatabaseId() { return databaseId; } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java index 8f5baca64f6..4dd5be0edd8 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java @@ -298,25 +298,9 @@ public DatabaseClient getDatabaseClient(DatabaseId db) { useMultiplexedSession ? multiplexedSessionDatabaseClient.getNumSessionsReleased() : new AtomicLong(); - SessionPool pool = - SessionPool.createPool( - getOptions(), - SpannerImpl.this.getSessionClient(db), - this.tracer, - labelValues, - attributesBuilder.build(), - numMultiplexedSessionsAcquired, - numMultiplexedSessionsReleased); - pool.maybeWaitOnMinSessions(); DatabaseClientImpl dbClient = createDatabaseClient( - clientId, - pool, - getOptions().getSessionPoolOptions().getUseMultiplexedSessionBlindWrite(), - multiplexedSessionDatabaseClient, - getOptions().getSessionPoolOptions().getUseMultiplexedSessionPartitionedOps(), - useMultiplexedSessionForRW, - this.tracer.createCommonAttributes(db)); + clientId, multiplexedSessionDatabaseClient, this.tracer.createCommonAttributes(db)); dbClients.put(db, dbClient); return dbClient; } @@ -326,27 +310,9 @@ public DatabaseClient getDatabaseClient(DatabaseId db) { @VisibleForTesting DatabaseClientImpl createDatabaseClient( String clientId, - SessionPool pool, - boolean useMultiplexedSessionBlindWrite, - @Nullable MultiplexedSessionDatabaseClient multiplexedSessionClient, - boolean useMultiplexedSessionPartitionedOps, - boolean useMultiplexedSessionForRW, + MultiplexedSessionDatabaseClient multiplexedSessionClient, Attributes commonAttributes) { - if (multiplexedSessionClient != null) { - // Set the session pool in the multiplexed session client. - // This is required to handle fallback to regular sessions for in-progress transactions that - // use multiplexed sessions but fail with UNIMPLEMENTED errors. - multiplexedSessionClient.setPool(pool); - } - return new DatabaseClientImpl( - clientId, - pool, - useMultiplexedSessionBlindWrite, - multiplexedSessionClient, - useMultiplexedSessionPartitionedOps, - tracer, - useMultiplexedSessionForRW, - commonAttributes); + return new DatabaseClientImpl(clientId, multiplexedSessionClient, tracer, commonAttributes); } @Override diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncRunnerTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncRunnerTest.java index 56d33c54878..5e34fcc7e71 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncRunnerTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncRunnerTest.java @@ -545,9 +545,6 @@ public void closeTransactionBeforeEndOfAsyncQuery() throws Exception { final SettableApiFuture finished = SettableApiFuture.create(); DatabaseClientImpl clientImpl = (DatabaseClientImpl) client(); - // There should currently not be any sessions checked out of the pool. - assertThat(clientImpl.pool.getNumberOfSessionsInUse()).isEqualTo(0); - AsyncRunner runner = clientImpl.runAsync(); final CountDownLatch dataReceived = new CountDownLatch(1); final CountDownLatch dataChecked = new CountDownLatch(1); @@ -592,9 +589,6 @@ public void closeTransactionBeforeEndOfAsyncQuery() throws Exception { // Wait until at least one row has been fetched. At that moment there should be one session // checked out. dataReceived.await(); - if (!isMultiplexedSessionsEnabledForRW()) { - assertThat(clientImpl.pool.getNumberOfSessionsInUse()).isEqualTo(1); - } assertThat(res.isDone()).isFalse(); dataChecked.countDown(); // Get the data from the transaction. @@ -605,7 +599,6 @@ public void closeTransactionBeforeEndOfAsyncQuery() throws Exception { assertThat(finished.get()).isTrue(); assertThat(resultList).containsExactly("k1", "k2", "k3"); assertThat(res.get()).isNull(); - assertThat(clientImpl.pool.getNumberOfSessionsInUse()).isEqualTo(0); } @Test 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 index 59657026e57..0e06f9554c2 100644 --- 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 @@ -41,7 +41,6 @@ import com.google.cloud.spanner.MockSpannerServiceImpl.SimulatedExecutionTime; import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult; import com.google.cloud.spanner.Options.ReadOption; -import com.google.cloud.spanner.SessionPool.SessionPoolTransactionContext; import com.google.cloud.spanner.TransactionRunnerImpl.TransactionContextImpl; import com.google.common.collect.ImmutableList; import com.google.common.collect.Iterables; @@ -180,9 +179,6 @@ public void asyncTransactionManager_shouldRollbackOnCloseAsync() throws Exceptio AsyncTransactionManager manager = client().transactionManagerAsync(); TransactionContext txn = manager.beginAsync().get(); txn.executeUpdateAsync(UPDATE_STATEMENT).get(); - if (txn instanceof SessionPoolTransactionContext) { - txn = ((SessionPoolTransactionContext) txn).delegate; - } TransactionContextImpl impl = (TransactionContextImpl) txn; final TransactionSelector selector = impl.getTransactionSelector(); @@ -363,18 +359,6 @@ public void asyncTransactionManagerFireAndForgetInvalidUpdate() throws Exception } } } - ImmutableList> expectedRequests = - ImmutableList.of( - BatchCreateSessionsRequest.class, - // The first update that fails. This will cause a transaction retry. - ExecuteSqlRequest.class, - // The retry will use an explicit BeginTransaction call. - BeginTransactionRequest.class, - // The first update will again fail, but now there is a transaction id, so the - // transaction can continue. - ExecuteSqlRequest.class, - ExecuteSqlRequest.class, - CommitRequest.class); ImmutableList> expectedRequestsWithMultiplexedSessionForRW = ImmutableList.of( CreateSessionRequest.class, @@ -387,14 +371,8 @@ public void asyncTransactionManagerFireAndForgetInvalidUpdate() throws Exception ExecuteSqlRequest.class, ExecuteSqlRequest.class, CommitRequest.class); - if (isMultiplexedSessionsEnabledForRW()) { - assertThat(mockSpanner.getRequestTypes()) - .containsExactlyElementsIn(expectedRequestsWithMultiplexedSessionForRW); - } else if (isMultiplexedSessionsEnabled()) { - assertThat(mockSpanner.getRequestTypes()).containsAtLeastElementsIn(expectedRequests); - } else { - assertThat(mockSpanner.getRequestTypes()).containsExactlyElementsIn(expectedRequests); - } + assertThat(mockSpanner.getRequestTypes()) + .containsExactlyElementsIn(expectedRequestsWithMultiplexedSessionForRW); } @Test @@ -533,25 +511,14 @@ public void asyncTransactionManagerUpdateAbortedWithoutGettingResult() throws Ex // The server may receive 1 or 2 commit requests depending on whether the call to // commitAsync() already knows that the transaction has aborted. If it does, it will not // attempt to call the Commit RPC and instead directly propagate the Aborted error. - if (isMultiplexedSessionsEnabledForRW()) { - assertThat(mockSpanner.getRequestTypes()) - .containsAtLeast( - CreateSessionRequest.class, - ExecuteSqlRequest.class, - // The retry will use a BeginTransaction RPC. - BeginTransactionRequest.class, - ExecuteSqlRequest.class, - CommitRequest.class); - } else { - assertThat(mockSpanner.getRequestTypes()) - .containsAtLeast( - BatchCreateSessionsRequest.class, - ExecuteSqlRequest.class, - // The retry will use a BeginTransaction RPC. - BeginTransactionRequest.class, - ExecuteSqlRequest.class, - CommitRequest.class); - } + assertThat(mockSpanner.getRequestTypes()) + .containsAtLeast( + CreateSessionRequest.class, + ExecuteSqlRequest.class, + // The retry will use a BeginTransaction RPC. + BeginTransactionRequest.class, + ExecuteSqlRequest.class, + CommitRequest.class); break; } catch (AbortedException e) { transactionContextFuture = manager.resetForRetryAsync(); @@ -599,22 +566,9 @@ public void asyncTransactionManagerWaitsUntilAsyncUpdateHasFinished() throws Exc executor) .commitAsync() .get(); - if (isMultiplexedSessionsEnabledForRW()) { - assertThat(mockSpanner.getRequestTypes()) - .containsExactly( - CreateSessionRequest.class, ExecuteSqlRequest.class, CommitRequest.class); - } else if (isMultiplexedSessionsEnabled()) { - assertThat(mockSpanner.getRequestTypes()) - .containsExactly( - CreateSessionRequest.class, - BatchCreateSessionsRequest.class, - ExecuteSqlRequest.class, - CommitRequest.class); - } else { - assertThat(mockSpanner.getRequestTypes()) - .containsExactly( - BatchCreateSessionsRequest.class, ExecuteSqlRequest.class, CommitRequest.class); - } + assertThat(mockSpanner.getRequestTypes()) + .containsExactly( + CreateSessionRequest.class, ExecuteSqlRequest.class, CommitRequest.class); break; } catch (AbortedException e) { txn = mgr.resetForRetryAsync(); @@ -735,14 +689,8 @@ public void asyncTransactionManagerFireAndForgetInvalidBatchUpdate() throws Exce ExecuteBatchDmlRequest.class, ExecuteBatchDmlRequest.class, CommitRequest.class); - if (isMultiplexedSessionsEnabledForRW()) { - assertThat(mockSpanner.getRequestTypes()) - .containsExactlyElementsIn(expectedRequestsWithMultiplexedSessionsRW); - } else if (isMultiplexedSessionsEnabled()) { - assertThat(mockSpanner.getRequestTypes()).containsAtLeastElementsIn(expectedRequests); - } else { - assertThat(mockSpanner.getRequestTypes()).containsExactlyElementsIn(expectedRequests); - } + assertThat(mockSpanner.getRequestTypes()) + .containsExactlyElementsIn(expectedRequestsWithMultiplexedSessionsRW); } @Test @@ -789,14 +737,8 @@ public void asyncTransactionManagerBatchUpdateAborted() throws Exception { BeginTransactionRequest.class, ExecuteBatchDmlRequest.class, CommitRequest.class); - if (isMultiplexedSessionsEnabledForRW()) { - assertThat(mockSpanner.getRequestTypes()) - .containsExactlyElementsIn(expectedRequestsWithMultiplexedSessionsRW); - } else if (isMultiplexedSessionsEnabled()) { - assertThat(mockSpanner.getRequestTypes()).containsAtLeastElementsIn(expectedRequests); - } else { - assertThat(mockSpanner.getRequestTypes()).containsExactlyElementsIn(expectedRequests); - } + assertThat(mockSpanner.getRequestTypes()) + .containsExactlyElementsIn(expectedRequestsWithMultiplexedSessionsRW); } @Test @@ -844,14 +786,8 @@ public void asyncTransactionManagerBatchUpdateAbortedBeforeFirstStatement() thro BeginTransactionRequest.class, ExecuteBatchDmlRequest.class, CommitRequest.class); - if (isMultiplexedSessionsEnabledForRW()) { - assertThat(mockSpanner.getRequestTypes()) - .containsExactlyElementsIn(expectedRequestsWithMultiplexedSessionsRW); - } else if (isMultiplexedSessionsEnabled()) { - assertThat(mockSpanner.getRequestTypes()).containsAtLeastElementsIn(expectedRequests); - } else { - assertThat(mockSpanner.getRequestTypes()).containsExactlyElementsIn(expectedRequests); - } + assertThat(mockSpanner.getRequestTypes()) + .containsExactlyElementsIn(expectedRequestsWithMultiplexedSessionsRW); } @Test @@ -916,14 +852,8 @@ public void asyncTransactionManagerWithBatchUpdateCommitAborted() throws Excepti BeginTransactionRequest.class, ExecuteBatchDmlRequest.class, CommitRequest.class); - if (isMultiplexedSessionsEnabledForRW()) { - assertThat(mockSpanner.getRequestTypes()) - .containsExactlyElementsIn(expectedRequestsWithMultiplexedSessionsRW); - } else if (isMultiplexedSessionsEnabled()) { - assertThat(mockSpanner.getRequestTypes()).containsAtLeastElementsIn(expectedRequests); - } else { - assertThat(mockSpanner.getRequestTypes()).containsExactlyElementsIn(expectedRequests); - } + assertThat(mockSpanner.getRequestTypes()) + .containsExactlyElementsIn(expectedRequestsWithMultiplexedSessionsRW); } @Test @@ -960,51 +890,25 @@ public void asyncTransactionManagerBatchUpdateAbortedWithoutGettingResult() thro } assertThat(attempt.get()).isEqualTo(2); List> requests = mockSpanner.getRequestTypes(); - // Remove the CreateSession requests for multiplexed sessions, as those are not relevant for - // this test if multiplexed session for read-write is not enabled. - if (!isMultiplexedSessionsEnabledForRW()) { - requests.removeIf(request -> request == CreateSessionRequest.class); - } int size = Iterables.size(requests); assertThat(size).isIn(Range.closed(5, 6)); if (size == 5) { - if (isMultiplexedSessionsEnabledForRW()) { - assertThat(requests) - .containsExactly( - CreateSessionRequest.class, - ExecuteBatchDmlRequest.class, - BeginTransactionRequest.class, - ExecuteBatchDmlRequest.class, - CommitRequest.class); - } else { - assertThat(requests) - .containsExactly( - BatchCreateSessionsRequest.class, - ExecuteBatchDmlRequest.class, - BeginTransactionRequest.class, - ExecuteBatchDmlRequest.class, - CommitRequest.class); - } + assertThat(requests) + .containsExactly( + CreateSessionRequest.class, + ExecuteBatchDmlRequest.class, + BeginTransactionRequest.class, + ExecuteBatchDmlRequest.class, + CommitRequest.class); } else { - if (isMultiplexedSessionsEnabledForRW()) { - assertThat(requests) - .containsExactly( - CreateSessionRequest.class, - ExecuteBatchDmlRequest.class, - CommitRequest.class, - BeginTransactionRequest.class, - ExecuteBatchDmlRequest.class, - CommitRequest.class); - } else { - assertThat(requests) - .containsExactly( - BatchCreateSessionsRequest.class, - ExecuteBatchDmlRequest.class, - CommitRequest.class, - BeginTransactionRequest.class, - ExecuteBatchDmlRequest.class, - CommitRequest.class); - } + assertThat(requests) + .containsExactly( + CreateSessionRequest.class, + ExecuteBatchDmlRequest.class, + CommitRequest.class, + BeginTransactionRequest.class, + ExecuteBatchDmlRequest.class, + CommitRequest.class); } } @@ -1038,14 +942,8 @@ public void asyncTransactionManagerWithBatchUpdateCommitFails() { ImmutableList> expectedRequestsWithMultiplexedSessionsRW = ImmutableList.of( CreateSessionRequest.class, ExecuteBatchDmlRequest.class, CommitRequest.class); - if (isMultiplexedSessionsEnabledForRW()) { - assertThat(mockSpanner.getRequestTypes()) - .containsExactlyElementsIn(expectedRequestsWithMultiplexedSessionsRW); - } else if (isMultiplexedSessionsEnabled()) { - assertThat(mockSpanner.getRequestTypes()).containsAtLeastElementsIn(expectedRequests); - } else { - assertThat(mockSpanner.getRequestTypes()).containsExactlyElementsIn(expectedRequests); - } + assertThat(mockSpanner.getRequestTypes()) + .containsExactlyElementsIn(expectedRequestsWithMultiplexedSessionsRW); } @Test @@ -1075,14 +973,8 @@ public void asyncTransactionManagerWaitsUntilAsyncBatchUpdateHasFinished() throw ImmutableList> expectedRequestsWithMultiplexedSessionsRW = ImmutableList.of( CreateSessionRequest.class, ExecuteBatchDmlRequest.class, CommitRequest.class); - if (isMultiplexedSessionsEnabledForRW()) { - assertThat(mockSpanner.getRequestTypes()) - .containsExactlyElementsIn(expectedRequestsWithMultiplexedSessionsRW); - } else if (isMultiplexedSessionsEnabled()) { - assertThat(mockSpanner.getRequestTypes()).containsAtLeastElementsIn(expectedRequests); - } else { - assertThat(mockSpanner.getRequestTypes()).containsExactlyElementsIn(expectedRequests); - } + assertThat(mockSpanner.getRequestTypes()) + .containsExactlyElementsIn(expectedRequestsWithMultiplexedSessionsRW); } @Test @@ -1244,56 +1136,4 @@ 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; - } - return spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSession(); - } - - private boolean isMultiplexedSessionsEnabledForRW() { - if (spanner.getOptions() == null || spanner.getOptions().getSessionPoolOptions() == null) { - return false; - } - return spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW(); - } } diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/BackendExhaustedTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/BackendExhaustedTest.java deleted file mode 100644 index 76cc20cb4de..00000000000 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/BackendExhaustedTest.java +++ /dev/null @@ -1,210 +0,0 @@ -/* - * Copyright 2020 Google LLC - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package com.google.cloud.spanner; - -import static com.google.cloud.spanner.DisableDefaultMtlsProvider.disableDefaultMtlsProvider; -import static com.google.common.truth.Truth.assertThat; - -import com.google.api.gax.grpc.testing.LocalChannelProvider; -import com.google.cloud.NoCredentials; -import com.google.cloud.grpc.GrpcTransportOptions; -import com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory; -import com.google.cloud.spanner.MockSpannerServiceImpl.SimulatedExecutionTime; -import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult; -import com.google.protobuf.ListValue; -import com.google.spanner.v1.ResultSetMetadata; -import com.google.spanner.v1.StructType; -import com.google.spanner.v1.StructType.Field; -import com.google.spanner.v1.TypeCode; -import io.grpc.Server; -import io.grpc.Status; -import io.grpc.inprocess.InProcessServerBuilder; -import java.util.concurrent.Executors; -import java.util.concurrent.ScheduledExecutorService; -import java.util.concurrent.ScheduledThreadPoolExecutor; -import java.util.concurrent.TimeUnit; -import org.junit.After; -import org.junit.AfterClass; -import org.junit.Before; -import org.junit.BeforeClass; -import org.junit.Test; -import org.junit.runner.RunWith; -import org.junit.runners.JUnit4; - -/** - * Tests that a degraded backend that can no longer create any new sessions will not cause an - * application that already has a healthy session pool to stop functioning. - */ -@RunWith(JUnit4.class) -public class BackendExhaustedTest { - private static final String TEST_PROJECT = "my-project"; - private static final String TEST_INSTANCE = "my-instance"; - private static final String TEST_DATABASE = "my-database"; - private static MockSpannerServiceImpl mockSpanner; - private static Server server; - private static LocalChannelProvider channelProvider; - private static final Statement UPDATE_STATEMENT = - Statement.of("UPDATE FOO SET BAR=1 WHERE BAZ=2"); - private static final Statement INVALID_UPDATE_STATEMENT = - Statement.of("UPDATE NON_EXISTENT_TABLE SET BAR=1 WHERE BAZ=2"); - private static final long UPDATE_COUNT = 1L; - private static final Statement SELECT1 = Statement.of("SELECT 1 AS COL1"); - private static final ResultSetMetadata SELECT1_METADATA = - ResultSetMetadata.newBuilder() - .setRowType( - StructType.newBuilder() - .addFields( - Field.newBuilder() - .setName("COL1") - .setType( - com.google.spanner.v1.Type.newBuilder() - .setCode(TypeCode.INT64) - .build()) - .build()) - .build()) - .build(); - private static final com.google.spanner.v1.ResultSet SELECT1_RESULTSET = - com.google.spanner.v1.ResultSet.newBuilder() - .addRows( - ListValue.newBuilder() - .addValues(com.google.protobuf.Value.newBuilder().setStringValue("1").build()) - .build()) - .setMetadata(SELECT1_METADATA) - .build(); - private Spanner spanner; - private DatabaseClientImpl client; - - @BeforeClass - public static void startStaticServer() throws Exception { - disableDefaultMtlsProvider(); - mockSpanner = new MockSpannerServiceImpl(); - mockSpanner.setAbortProbability(0.0D); // We don't want any unpredictable aborted transactions. - mockSpanner.putStatementResult(StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT)); - mockSpanner.putStatementResult(StatementResult.query(SELECT1, SELECT1_RESULTSET)); - mockSpanner.putStatementResult( - StatementResult.exception( - INVALID_UPDATE_STATEMENT, - Status.INVALID_ARGUMENT.withDescription("invalid statement").asRuntimeException())); - - String uniqueName = InProcessServerBuilder.generateName(); - server = - InProcessServerBuilder.forName(uniqueName) - // We need to use a real executor for timeouts to occur. - .scheduledExecutorService(new ScheduledThreadPoolExecutor(1)) - .addService(mockSpanner) - .build() - .start(); - channelProvider = LocalChannelProvider.create(uniqueName); - } - - @AfterClass - public static void stopServer() throws InterruptedException { - // Force a shutdown as there are still requests stuck in the server. - server.shutdownNow(); - server.awaitTermination(); - } - - @Before - public void setUp() throws Exception { - SpannerOptions options = - SpannerOptions.newBuilder() - .setProjectId(TEST_PROJECT) - .setChannelProvider(channelProvider) - .setCredentials(NoCredentials.getInstance()) - .build(); - ExecutorFactory executorFactory = - ((GrpcTransportOptions) options.getTransportOptions()).getExecutorFactory(); - ScheduledThreadPoolExecutor executor = (ScheduledThreadPoolExecutor) executorFactory.get(); - options = - options.toBuilder() - .setSessionPoolOption( - SessionPoolOptions.newBuilder() - .setMinSessions(executor.getCorePoolSize()) - .setMaxSessions(executor.getCorePoolSize() * 3) - .build()) - .build(); - executorFactory.release(executor); - - spanner = options.getService(); - client = - (DatabaseClientImpl) - spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); - // Wait until the session pool has initialized. - while (client.pool.getNumberOfSessionsInPool() - < spanner.getOptions().getSessionPoolOptions().getMinSessions()) { - Thread.sleep(1L); - } - } - - @After - public void tearDown() { - mockSpanner.reset(); - mockSpanner.removeAllExecutionTimes(); - // This test case force-closes the Spanner instance as it would otherwise wait - // forever on the BatchCreateSessions requests that are 'stuck'. - try { - ((SpannerImpl) spanner).close(10L, TimeUnit.MILLISECONDS); - } catch (SpannerException e) { - // ignore any errors during close as they are expected. - } - } - - @Test - public void test() throws Exception { - // Simulate very heavy load on the server by effectively stopping session creation. - mockSpanner.setBatchCreateSessionsExecutionTime( - SimulatedExecutionTime.ofMinimumAndRandomTime(Integer.MAX_VALUE, 0)); - // Create an executor that can handle twice as many requests as the minimum number of sessions - // in the pool and then start that many read requests. That will initiate the creation of - // additional sessions. - ScheduledExecutorService executor = - Executors.newScheduledThreadPool( - spanner.getOptions().getSessionPoolOptions().getMinSessions() * 2); - // Also temporarily freeze the server to ensure that the requests that can be served will - // continue to be in-flight and keep the sessions in the pool checked out. - mockSpanner.freeze(); - for (int i = 0; i < spanner.getOptions().getSessionPoolOptions().getMinSessions() * 2; i++) { - executor.submit(new ReadRunnable()); - } - // Now schedule as many write requests as there can be sessions in the pool. - for (int i = 0; i < spanner.getOptions().getSessionPoolOptions().getMaxSessions(); i++) { - executor.submit(new WriteRunnable()); - } - // Now unfreeze the server and verify that all requests can be served using the sessions that - // were already present in the pool. - mockSpanner.unfreeze(); - executor.shutdown(); - assertThat(executor.awaitTermination(10, TimeUnit.SECONDS)).isTrue(); - } - - private final class ReadRunnable implements Runnable { - @Override - public void run() { - try (ResultSet rs = client.singleUse().executeQuery(SELECT1)) { - while (rs.next()) {} - } - } - } - - private final class WriteRunnable implements Runnable { - @Override - public void run() { - TransactionRunner runner = client.readWriteTransaction(); - runner.run(transaction -> transaction.executeUpdate(UPDATE_STATEMENT)); - } - } -} diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/BaseSessionPoolTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/BaseSessionPoolTest.java deleted file mode 100644 index 939114a7f60..00000000000 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/BaseSessionPoolTest.java +++ /dev/null @@ -1,187 +0,0 @@ -/* - * Copyright 2017 Google LLC - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package com.google.cloud.spanner; - -import static org.mockito.Mockito.any; -import static org.mockito.Mockito.anyLong; -import static org.mockito.Mockito.doReturn; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.spy; -import static org.mockito.Mockito.when; - -import com.google.api.core.ApiFuture; -import com.google.api.core.ApiFutures; -import com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory; -import com.google.cloud.spanner.Options.TransactionOption; -import com.google.cloud.spanner.spi.v1.SpannerRpc.Option; -import com.google.protobuf.Empty; -import com.google.protobuf.Timestamp; -import java.util.HashMap; -import java.util.Map; -import java.util.concurrent.ScheduledExecutorService; -import java.util.concurrent.ScheduledFuture; -import java.util.concurrent.ScheduledThreadPoolExecutor; -import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicLong; - -abstract class BaseSessionPoolTest { - ScheduledExecutorService mockExecutor; - int sessionIndex; - AtomicLong channelHint = new AtomicLong(0L); - - final class TestExecutorFactory implements ExecutorFactory { - - @Override - public ScheduledExecutorService get() { - ScheduledExecutorService realExecutor = new ScheduledThreadPoolExecutor(2); - mockExecutor = spy(realExecutor); - @SuppressWarnings("rawtypes") - ScheduledFuture mockFuture = mock(ScheduledFuture.class); - // To prevent maintenance loop from running. - doReturn(mockFuture) - .when(mockExecutor) - .scheduleAtFixedRate(any(Runnable.class), anyLong(), anyLong(), any(TimeUnit.class)); - return mockExecutor; - } - - @Override - public void release(ScheduledExecutorService executor) { - try { - executor.shutdown(); - } catch (Throwable ignore) { - } - } - } - - @SuppressWarnings("unchecked") - SessionImpl mockSession() { - final SessionImpl session = mock(SessionImpl.class); - Map options = new HashMap<>(); - options.put(Option.CHANNEL_HINT, channelHint.getAndIncrement()); - when(session.getOptions()).thenReturn(options); - when(session.getName()) - .thenReturn( - "projects/dummy/instances/dummy/database/dummy/sessions/session" + sessionIndex); - when(session.asyncClose()).thenReturn(ApiFutures.immediateFuture(Empty.getDefaultInstance())); - when(session.writeWithOptions(any(Iterable.class))) - .thenReturn(new CommitResponse(com.google.spanner.v1.CommitResponse.getDefaultInstance())); - when(session.writeAtLeastOnceWithOptions(any(Iterable.class))) - .thenReturn(new CommitResponse(com.google.spanner.v1.CommitResponse.getDefaultInstance())); - sessionIndex++; - return session; - } - - SessionImpl mockMultiplexedSession() { - final SessionImpl session = mock(SessionImpl.class); - Map options = new HashMap<>(); - when(session.getIsMultiplexed()).thenReturn(true); - when(session.getOptions()).thenReturn(options); - when(session.getName()) - .thenReturn( - "projects/dummy/instances/dummy/database/dummy/sessions/session" + sessionIndex); - when(session.asyncClose()).thenReturn(ApiFutures.immediateFuture(Empty.getDefaultInstance())); - when(session.writeWithOptions(any(Iterable.class))) - .thenReturn(new CommitResponse(com.google.spanner.v1.CommitResponse.getDefaultInstance())); - when(session.writeAtLeastOnceWithOptions(any(Iterable.class))) - .thenReturn(new CommitResponse(com.google.spanner.v1.CommitResponse.getDefaultInstance())); - sessionIndex++; - return session; - } - - SessionImpl buildMockSession(SpannerImpl spanner, ReadContext context) { - Map options = new HashMap<>(); - options.put(Option.CHANNEL_HINT, channelHint.getAndIncrement()); - final SessionImpl session = - new SessionImpl( - spanner, - new SessionReference( - "projects/dummy/instances/dummy/databases/dummy/sessions/session" + sessionIndex, - options)) { - @Override - public ReadContext singleUse(TimestampBound bound) { - // The below stubs are added so that we can mock keep-alive. - return context; - } - - @Override - public ApiFuture asyncClose() { - return ApiFutures.immediateFuture(Empty.getDefaultInstance()); - } - - @Override - public CommitResponse writeAtLeastOnceWithOptions( - Iterable mutations, TransactionOption... transactionOptions) - throws SpannerException { - return new CommitResponse(com.google.spanner.v1.CommitResponse.getDefaultInstance()); - } - - @Override - public CommitResponse writeWithOptions( - Iterable mutations, TransactionOption... options) throws SpannerException { - return new CommitResponse(com.google.spanner.v1.CommitResponse.getDefaultInstance()); - } - }; - sessionIndex++; - return session; - } - - SessionImpl buildMockMultiplexedSession( - SpannerImpl spanner, ReadContext context, Timestamp creationTime) { - Map options = new HashMap<>(); - final SessionImpl session = - new SessionImpl( - spanner, - new SessionReference( - "projects/dummy/instances/dummy/databases/dummy/sessions/session" + sessionIndex, - creationTime, - true, - options)) { - @Override - public ReadContext singleUse(TimestampBound bound) { - // The below stubs are added so that we can mock keep-alive. - return context; - } - - @Override - public ApiFuture asyncClose() { - return ApiFutures.immediateFuture(Empty.getDefaultInstance()); - } - - @Override - public CommitResponse writeAtLeastOnceWithOptions( - Iterable mutations, TransactionOption... transactionOptions) - throws SpannerException { - return new CommitResponse(com.google.spanner.v1.CommitResponse.getDefaultInstance()); - } - - @Override - public CommitResponse writeWithOptions( - Iterable mutations, TransactionOption... options) throws SpannerException { - return new CommitResponse(com.google.spanner.v1.CommitResponse.getDefaultInstance()); - } - }; - sessionIndex++; - return session; - } - - void runMaintenanceLoop(FakeClock clock, SessionPool pool, long numCycles) { - for (int i = 0; i < numCycles; i++) { - pool.poolMaintainer.maintainPool(); - clock.currentTimeMillis.addAndGet(pool.poolMaintainer.loopFrequency); - } - } -} diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/BatchCreateSessionsTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/BatchCreateSessionsTest.java deleted file mode 100644 index 4ff06755494..00000000000 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/BatchCreateSessionsTest.java +++ /dev/null @@ -1,237 +0,0 @@ -/* - * Copyright 2019 Google LLC - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package com.google.cloud.spanner; - -import static com.google.cloud.spanner.DisableDefaultMtlsProvider.disableDefaultMtlsProvider; -import static org.hamcrest.CoreMatchers.equalTo; -import static org.hamcrest.CoreMatchers.is; -import static org.hamcrest.MatcherAssert.assertThat; - -import com.google.api.gax.grpc.testing.LocalChannelProvider; -import com.google.cloud.NoCredentials; -import com.google.cloud.spanner.MockSpannerServiceImpl.SimulatedExecutionTime; -import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult; -import com.google.common.base.Stopwatch; -import com.google.protobuf.ListValue; -import com.google.spanner.v1.ResultSetMetadata; -import com.google.spanner.v1.StructType; -import com.google.spanner.v1.StructType.Field; -import com.google.spanner.v1.TypeCode; -import io.grpc.Server; -import io.grpc.Status; -import io.grpc.inprocess.InProcessServerBuilder; -import java.util.concurrent.TimeUnit; -import org.junit.AfterClass; -import org.junit.Before; -import org.junit.BeforeClass; -import org.junit.Test; -import org.junit.runner.RunWith; -import org.junit.runners.JUnit4; - -@RunWith(JUnit4.class) -public class BatchCreateSessionsTest { - private static final Statement SELECT1AND2 = - Statement.of("SELECT 1 AS COL1 UNION ALL SELECT 2 AS COL1"); - private static final ResultSetMetadata SELECT1AND2_METADATA = - ResultSetMetadata.newBuilder() - .setRowType( - StructType.newBuilder() - .addFields( - Field.newBuilder() - .setName("COL1") - .setType( - com.google.spanner.v1.Type.newBuilder() - .setCode(TypeCode.INT64) - .build()) - .build()) - .build()) - .build(); - private static final com.google.spanner.v1.ResultSet SELECT1_RESULTSET = - com.google.spanner.v1.ResultSet.newBuilder() - .addRows( - ListValue.newBuilder() - .addValues(com.google.protobuf.Value.newBuilder().setStringValue("1").build()) - .build()) - .addRows( - ListValue.newBuilder() - .addValues(com.google.protobuf.Value.newBuilder().setStringValue("2").build()) - .build()) - .setMetadata(SELECT1AND2_METADATA) - .build(); - - private static MockSpannerServiceImpl mockSpanner; - private static Server server; - private static LocalChannelProvider channelProvider; - - @BeforeClass - public static void startStaticServer() throws Exception { - disableDefaultMtlsProvider(); - mockSpanner = new MockSpannerServiceImpl(); - mockSpanner.setAbortProbability(0.0D); // We don't want any unpredictable aborted transactions. - mockSpanner.putStatementResult(StatementResult.query(SELECT1AND2, SELECT1_RESULTSET)); - - String uniqueName = InProcessServerBuilder.generateName(); - server = - InProcessServerBuilder.forName(uniqueName) - .directExecutor() - .addService(mockSpanner) - .build() - .start(); - channelProvider = LocalChannelProvider.create(uniqueName); - } - - @AfterClass - public static void stopServer() throws InterruptedException { - server.shutdown(); - server.awaitTermination(); - } - - @Before - public void setUp() { - mockSpanner.reset(); - mockSpanner.removeAllExecutionTimes(); - } - - private Spanner createSpanner(int minSessions, int maxSessions) { - SessionPoolOptions sessionPoolOptions = - SessionPoolOptions.newBuilder() - .setMinSessions(minSessions) - .setMaxSessions(maxSessions) - .build(); - return SpannerOptions.newBuilder() - .setProjectId("[PROJECT]") - .setChannelProvider(channelProvider) - .setSessionPoolOption(sessionPoolOptions) - .setCredentials(NoCredentials.getInstance()) - .build() - .getService(); - } - - @Test - public void testCreatedMinSessions() throws InterruptedException { - int minSessions = 1000; - int maxSessions = 4000; - try (Spanner spanner = createSpanner(minSessions, maxSessions)) { - DatabaseClientImpl client = - (DatabaseClientImpl) - spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]")); - Stopwatch watch = Stopwatch.createStarted(); - while (client.pool.totalSessions() < minSessions && watch.elapsed(TimeUnit.SECONDS) < 10) { - Thread.sleep(10L); - } - assertThat(client.pool.totalSessions(), is(equalTo(minSessions))); - } - } - - @Test - public void testClosePoolWhileInitializing() throws InterruptedException { - int minSessions = 10_000; - int maxSessions = 10_000; - DatabaseClientImpl client; - // Freeze the server to prevent it from creating sessions before we want to. - mockSpanner.freeze(); - try (Spanner spanner = createSpanner(minSessions, maxSessions)) { - // Create a database client which will create a session pool. - // No sessions will be created at the moment as the server is frozen. - client = - (DatabaseClientImpl) - spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]")); - // Make sure session creation takes a little time to avoid all sessions being created at once. - mockSpanner.setBatchCreateSessionsExecutionTime( - SimulatedExecutionTime.ofMinimumAndRandomTime(10, 0)); - // Unfreeze the server to allow session creation to start. - mockSpanner.unfreeze(); - // Wait until at least one batch of sessions has been created. - Stopwatch watch = Stopwatch.createStarted(); - while (client.pool.totalSessions() == 0 && watch.elapsed(TimeUnit.SECONDS) < 10) { - Thread.sleep(1L); - } - // Close the Spanner instance which will start to delete sessions while the session pool is - // still being initialized. - } - // Verify that all sessions have been deleted. - assertThat(client.pool.totalSessions(), is(equalTo(0))); - } - - @Test - public void testSpannerReturnsAllAvailableSessionsAndThenNoSessions() - throws InterruptedException { - int minSessions = 1000; - int maxSessions = 1000; - // Set a maximum number of sessions that will be created by the server. - // After this the server will return an error when batchCreateSessions is called. - // This error is not propagated to the client. - int maxServerSessions = 550; - DatabaseClientImpl client; - mockSpanner.setMaxTotalSessions(maxServerSessions); - try (Spanner spanner = createSpanner(minSessions, maxSessions)) { - // Create a database client which will create a session pool. - client = - (DatabaseClientImpl) - spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]")); - Stopwatch watch = Stopwatch.createStarted(); - while (client.pool.totalSessions() < maxServerSessions - && watch.elapsed(TimeUnit.SECONDS) < 10) { - Thread.sleep(10L); - } - assertThat(client.pool.totalSessions(), is(equalTo(maxServerSessions))); - // Wait until the pool has given up creating sessions. - watch = watch.reset(); - watch.start(); - while (client.pool.getNumberOfSessionsBeingCreated() > 0 - && watch.elapsed(TimeUnit.SECONDS) < 10) { - Thread.sleep(10L); - } - // Remove the max server sessions limit. - mockSpanner.setMaxTotalSessions(Integer.MAX_VALUE); - // Wait a little. No more sessions should be created, as the previous requests have given up, - // and no new sessions have been requested from the pool. - Thread.sleep(20L); - assertThat(client.pool.totalSessions(), is(equalTo(maxServerSessions))); - } - // Verify that all sessions have been deleted. - assertThat(client.pool.totalSessions(), is(equalTo(0))); - } - - @Test - public void testSpannerReturnsFailedPrecondition() throws InterruptedException { - int minSessions = 100; - int maxSessions = 1000; - int expectedSessions; - DatabaseClientImpl client; - // Make the first BatchCreateSessions return an error. - mockSpanner.addException(Status.FAILED_PRECONDITION.asRuntimeException()); - try (Spanner spanner = createSpanner(minSessions, maxSessions)) { - // Create a database client which will create a session pool. - client = - (DatabaseClientImpl) - spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]")); - // Wait for the pool to be initialized. - // The first session creation request will fail. - expectedSessions = minSessions - minSessions / spanner.getOptions().getNumChannels(); - Stopwatch watch = Stopwatch.createStarted(); - while (client.pool.totalSessions() < expectedSessions - && watch.elapsed(TimeUnit.SECONDS) < 10) { - Thread.sleep(10L); - } - // Wait a little to allow any additional session creation to finish. - Thread.sleep(20L); - } - // Verify that all sessions have been deleted. - assertThat(client.pool.totalSessions(), is(equalTo(0))); - } -} diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ChannelUsageTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ChannelUsageTest.java index a06eeb91662..8177ee25279 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ChannelUsageTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ChannelUsageTest.java @@ -106,8 +106,6 @@ public static Collection data() { private static MockSpannerServiceImpl mockSpanner; private static Server server; private static InetSocketAddress address; - private static final Set batchCreateSessionLocalIps = - ConcurrentHashMap.newKeySet(); private static final Set executeSqlLocalIps = ConcurrentHashMap.newKeySet(); private static Level originalLogLevel; @@ -147,10 +145,6 @@ public ServerCall.Listener interceptCall( .findFirst() .orElse(null); if (key != null) { - if (call.getMethodDescriptor() - .equals(SpannerGrpc.getBatchCreateSessionsMethod())) { - batchCreateSessionLocalIps.add(attributes.get(key)); - } if (call.getMethodDescriptor() .equals(SpannerGrpc.getExecuteStreamingSqlMethod())) { executeSqlLocalIps.add(attributes.get(key)); @@ -185,7 +179,6 @@ public static void resetLogging() { @After public void reset() { mockSpanner.reset(); - batchCreateSessionLocalIps.clear(); executeSqlLocalIps.clear(); } @@ -215,27 +208,11 @@ private SpannerOptions createSpannerOptions() { return builder.build(); } - @Test - public void testCreatesNumChannels() { - try (Spanner spanner = createSpannerOptions().getService()) { - assumeFalse( - "GRPC-GCP is currently not supported with multiplexed sessions", - isMultiplexedSessionsEnabled(spanner) && enableGcpPool); - DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of("p", "i", "d")); - try (ResultSet resultSet = client.singleUse().executeQuery(SELECT1)) { - while (resultSet.next()) {} - } - } - assertEquals(numChannels, batchCreateSessionLocalIps.size()); - } - @Test public void testUsesAllChannels() throws InterruptedException { - final int multiplier = 2; + final int multiplier = 10; try (Spanner spanner = createSpannerOptions().getService()) { - assumeFalse( - "GRPC-GCP is currently not supported with multiplexed sessions", - isMultiplexedSessionsEnabled(spanner)); + assumeFalse("GRPC-GCP is currently not supported with multiplexed sessions", enableGcpPool); DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of("p", "i", "d")); ListeningExecutorService executor = MoreExecutors.listeningDecorator(Executors.newFixedThreadPool(numChannels * multiplier)); @@ -265,11 +242,4 @@ public void testUsesAllChannels() throws InterruptedException { } assertEquals(numChannels, executeSqlLocalIps.size()); } - - private boolean isMultiplexedSessionsEnabled(Spanner spanner) { - if (spanner.getOptions() == null || spanner.getOptions().getSessionPoolOptions() == null) { - return false; - } - return spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSession(); - } } diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/CloseSpannerWithOpenResultSetTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/CloseSpannerWithOpenResultSetTest.java deleted file mode 100644 index 7988e53f84f..00000000000 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/CloseSpannerWithOpenResultSetTest.java +++ /dev/null @@ -1,195 +0,0 @@ -/* - * Copyright 2024 Google LLC - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package com.google.cloud.spanner; - -import static org.junit.Assert.assertEquals; -import static org.junit.Assert.assertNotNull; -import static org.junit.Assert.assertThrows; -import static org.junit.Assert.assertTrue; -import static org.junit.Assume.assumeFalse; - -import com.google.cloud.NoCredentials; -import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult; -import com.google.cloud.spanner.connection.AbstractMockServerTest; -import com.google.cloud.spanner.spi.v1.GapicSpannerRpc; -import com.google.spanner.v1.DeleteSessionRequest; -import com.google.spanner.v1.ExecuteSqlRequest; -import io.grpc.ManagedChannelBuilder; -import io.grpc.Status; -import java.time.Duration; -import java.util.ArrayList; -import java.util.List; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.Executors; -import java.util.concurrent.Future; -import java.util.concurrent.ThreadLocalRandom; -import java.util.concurrent.TimeUnit; -import java.util.stream.Collectors; -import org.junit.After; -import org.junit.AfterClass; -import org.junit.BeforeClass; -import org.junit.Test; -import org.junit.runner.RunWith; -import org.junit.runners.JUnit4; - -@RunWith(JUnit4.class) -public class CloseSpannerWithOpenResultSetTest extends AbstractMockServerTest { - - Spanner createSpanner() { - return SpannerOptions.newBuilder() - .setProjectId("p") - .setHost(String.format("http://localhost:%d", getPort())) - .setChannelConfigurator(ManagedChannelBuilder::usePlaintext) - .setCredentials(NoCredentials.getInstance()) - .setSessionPoolOption( - SessionPoolOptions.newBuilder() - .setWaitForMinSessionsDuration(Duration.ofSeconds(5L)) - .build()) - .build() - .getService(); - } - - @BeforeClass - public static void setWatchdogTimeout() { - System.setProperty("com.google.cloud.spanner.watchdogTimeoutSeconds", "1"); - } - - @AfterClass - public static void clearWatchdogTimeout() { - System.clearProperty("com.google.cloud.spanner.watchdogTimeoutSeconds"); - } - - @After - public void cleanup() { - mockSpanner.unfreeze(); - mockSpanner.clearRequests(); - } - - @Test - public void testBatchClient_closedSpannerWithOpenResultSet_streamsAreCancelled() { - Spanner spanner = createSpanner(); - assumeFalse(spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSession()); - - BatchClient client = spanner.getBatchClient(DatabaseId.of("p", "i", "d")); - try (BatchReadOnlyTransaction transaction = - client.batchReadOnlyTransaction(TimestampBound.strong()); - ResultSet resultSet = transaction.executeQuery(SELECT_RANDOM_STATEMENT)) { - mockSpanner.freezeAfterReturningNumRows(1); - // This can sometimes fail, as the mock server may not always actually return the first row. - try { - assertTrue(resultSet.next()); - } catch (SpannerException exception) { - assertEquals(ErrorCode.DEADLINE_EXCEEDED, exception.getErrorCode()); - return; - } - ((SpannerImpl) spanner).close(1, TimeUnit.MILLISECONDS); - // This should return an error as the stream is cancelled. - SpannerException exception = - assertThrows( - SpannerException.class, - () -> { //noinspection StatementWithEmptyBody - while (resultSet.next()) {} - }); - assertEquals(ErrorCode.CANCELLED, exception.getErrorCode()); - } - } - - @Test - public void testNormalDatabaseClient_closedSpannerWithOpenResultSet_sessionsAreDeleted() - throws Exception { - Spanner spanner = createSpanner(); - assumeFalse(spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSession()); - - DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of("p", "i", "d")); - try (ReadOnlyTransaction transaction = client.readOnlyTransaction(TimestampBound.strong()); - ResultSet resultSet = transaction.executeQuery(SELECT_RANDOM_STATEMENT)) { - mockSpanner.freezeAfterReturningNumRows(1); - // This can sometimes fail, as the mock server may not always actually return the first row. - try { - assertTrue(resultSet.next()); - } catch (SpannerException exception) { - assertEquals(ErrorCode.DEADLINE_EXCEEDED, exception.getErrorCode()); - return; - } - List executeSqlRequests = - mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).stream() - .filter(request -> request.getSql().equals(SELECT_RANDOM_STATEMENT.getSql())) - .collect(Collectors.toList()); - assertEquals(1, executeSqlRequests.size()); - ExecutorService service = Executors.newSingleThreadExecutor(); - service.submit(spanner::close); - // Verify that the session that is used by this transaction is deleted. - // That will automatically cancel the query. - mockSpanner.waitForRequestsToContain( - request -> - request instanceof DeleteSessionRequest - && ((DeleteSessionRequest) request) - .getName() - .equals(executeSqlRequests.get(0).getSession()), - /* timeoutMillis= */ 1000L); - service.shutdownNow(); - } - } - - @Test - public void testStreamsAreCleanedUp() throws Exception { - String invalidSql = "select * from foo"; - Statement invalidStatement = Statement.of(invalidSql); - mockSpanner.putStatementResult( - StatementResult.exception( - invalidStatement, - Status.NOT_FOUND.withDescription("Table not found: foo").asRuntimeException())); - int numThreads = 16; - int numQueries = 32; - try (Spanner spanner = createSpanner()) { - BatchClient client = spanner.getBatchClient(DatabaseId.of("p", "i", "d")); - ExecutorService service = Executors.newFixedThreadPool(numThreads); - List> futures = new ArrayList<>(numQueries); - for (int n = 0; n < numQueries; n++) { - futures.add( - service.submit( - () -> { - try (BatchReadOnlyTransaction transaction = - client.batchReadOnlyTransaction(TimestampBound.strong())) { - if (ThreadLocalRandom.current().nextInt(10) < 2) { - try (ResultSet resultSet = transaction.executeQuery(invalidStatement)) { - SpannerException exception = - assertThrows(SpannerException.class, resultSet::next); - assertEquals(ErrorCode.NOT_FOUND, exception.getErrorCode()); - } - } else { - try (ResultSet resultSet = - transaction.executeQuery(SELECT_RANDOM_STATEMENT)) { - while (resultSet.next()) { - assertNotNull(resultSet.getCurrentRowAsStruct()); - } - } - } - } - })); - } - service.shutdown(); - for (Future fut : futures) { - fut.get(); - } - assertTrue(service.awaitTermination(1L, TimeUnit.MINUTES)); - // Verify that all response observers have been unregistered. - assertEquals( - 0, ((GapicSpannerRpc) ((SpannerImpl) spanner).getRpc()).getNumActiveResponseObservers()); - } - } -} 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 9ca5c17330a..827476d57c1 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 @@ -28,17 +28,10 @@ import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; -import static org.junit.Assert.assertNotEquals; import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertNull; import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; -import static org.junit.Assume.assumeFalse; -import static org.junit.Assume.assumeTrue; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; import com.google.api.core.ApiFuture; import com.google.api.core.ApiFutures; @@ -56,11 +49,7 @@ import com.google.cloud.spanner.Options.RpcLockHint; import com.google.cloud.spanner.Options.RpcOrderBy; import com.google.cloud.spanner.Options.RpcPriority; -import com.google.cloud.spanner.Options.TransactionOption; import com.google.cloud.spanner.ReadContext.QueryAnalyzeMode; -import com.google.cloud.spanner.SessionPool.PooledSessionFuture; -import com.google.cloud.spanner.SessionPoolOptions.ActionOnInactiveTransaction; -import com.google.cloud.spanner.SessionPoolOptions.InactiveTransactionRemovalOptions; import com.google.cloud.spanner.SingerProto.Genre; import com.google.cloud.spanner.SingerProto.SingerInfo; import com.google.cloud.spanner.SpannerException.ResourceNotFoundException; @@ -83,7 +72,6 @@ import com.google.spanner.v1.BeginTransactionRequest; import com.google.spanner.v1.CommitRequest; import com.google.spanner.v1.CreateSessionRequest; -import com.google.spanner.v1.DeleteSessionRequest; import com.google.spanner.v1.DirectedReadOptions; import com.google.spanner.v1.DirectedReadOptions.IncludeReplicas; import com.google.spanner.v1.DirectedReadOptions.ReplicaSelection; @@ -112,11 +100,8 @@ import io.grpc.StatusRuntimeException; import io.grpc.inprocess.InProcessServerBuilder; import io.grpc.protobuf.lite.ProtoLiteUtils; -import io.opencensus.trace.Tracing; -import io.opentelemetry.api.OpenTelemetry; import java.nio.charset.StandardCharsets; import java.time.Duration; -import java.time.Instant; import java.util.ArrayList; import java.util.Arrays; import java.util.Base64; @@ -131,7 +116,6 @@ import java.util.concurrent.ScheduledThreadPoolExecutor; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; -import java.util.concurrent.atomic.AtomicLong; import java.util.function.Function; import java.util.logging.Level; import java.util.logging.Logger; @@ -294,1078 +278,6 @@ public void tearDown() { mockSpanner.removeAllExecutionTimes(); } - @Test - public void - testPoolMaintainer_whenInactiveTransactionAndSessionIsNotFoundOnBackend_removeSessionsFromPool() { - assumeFalse( - "Session pool maintainer test skipped for multiplexed sessions", - isMultiplexedSessionsEnabledForRW()); - FakeClock poolMaintainerClock = new FakeClock(); - InactiveTransactionRemovalOptions inactiveTransactionRemovalOptions = - InactiveTransactionRemovalOptions.newBuilder() - .setIdleTimeThreshold( - Duration.ofSeconds( - 2L)) // any session not used for more than 2s will be long-running - .setActionOnInactiveTransaction(ActionOnInactiveTransaction.CLOSE) - .setExecutionFrequency(Duration.ofSeconds(1)) // execute thread every 1s - .build(); - SessionPoolOptions sessionPoolOptions = - SessionPoolOptions.newBuilder() - .setMinSessions(1) - .setMaxSessions(1) // to ensure there is 1 session and pool is 100% utilized - .setInactiveTransactionRemovalOptions(inactiveTransactionRemovalOptions) - .setLoopFrequency(1000L) // main thread runs every 1s - .setPoolMaintainerClock(poolMaintainerClock) - .build(); - try (Spanner spanner = - SpannerOptions.newBuilder() - .setProjectId(TEST_PROJECT) - .setDatabaseRole(TEST_DATABASE_ROLE) - .setChannelProvider(channelProvider) - .setCredentials(NoCredentials.getInstance()) - .setSessionPoolOption(sessionPoolOptions) - .build() - .getService()) { - DatabaseClientImpl client = - (DatabaseClientImpl) - spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); - Instant initialExecutionTime = client.pool.poolMaintainer.lastExecutionTime; - - poolMaintainerClock.currentTimeMillis.addAndGet(Duration.ofMinutes(3).toMillis()); - - try (TransactionManager manager = client.transactionManager()) { - TransactionContext transaction = manager.begin(); - mockSpanner.setCommitExecutionTime( - SimulatedExecutionTime.ofException( - mockSpanner.createSessionNotFoundException("TEST_SESSION_NAME"))); - while (true) { - try { - transaction.executeUpdate(UPDATE_STATEMENT); - - // Simulate a delay of 3 minutes to ensure that the below transaction is a long-running - // one. - // As per this test, anything which takes more than 2s is long-running - poolMaintainerClock.currentTimeMillis.addAndGet(Duration.ofMinutes(3).toMillis()); - // force trigger pool maintainer to check for long-running sessions - client.pool.poolMaintainer.maintainPool(); - - manager.commit(); - assertNotNull(manager.getCommitTimestamp()); - break; - } catch (AbortedException e) { - transaction = manager.resetForRetry(); - } - mockSpanner.setCommitExecutionTime(SimulatedExecutionTime.ofMinimumAndRandomTime(0, 0)); - } - } - Instant endExecutionTime = client.pool.poolMaintainer.lastExecutionTime; - - // first session executed update, session found to be long-running and cleaned up. - // During commit, SessionNotFound exception from backend caused replacement of session and - // transaction needs to be retried. - // On retry, session again found to be long-running and cleaned up. - // During commit, there was no exception from backend. - - assertNotEquals( - endExecutionTime, - initialExecutionTime); // if session clean up task runs then these timings won't match - assertEquals(2, client.pool.numLeakedSessionsRemoved()); - assertTrue(client.pool.getNumberOfSessionsInPool() <= client.pool.totalSessions()); - } - } - - @Test - public void - testPoolMaintainer_whenInactiveTransactionAndSessionExistsOnBackend_removeSessionsFromPool() { - assumeFalse( - "Session leaks tests are skipped for multiplexed sessions", - isMultiplexedSessionsEnabledForRW()); - FakeClock poolMaintainerClock = new FakeClock(); - InactiveTransactionRemovalOptions inactiveTransactionRemovalOptions = - InactiveTransactionRemovalOptions.newBuilder() - .setIdleTimeThreshold( - Duration.ofSeconds( - 2L)) // any session not used for more than 2s will be long-running - .setActionOnInactiveTransaction(ActionOnInactiveTransaction.CLOSE) - .setExecutionFrequency(Duration.ofSeconds(1)) // execute thread every 1s - .build(); - SessionPoolOptions sessionPoolOptions = - SessionPoolOptions.newBuilder() - .setMinSessions(1) - .setMaxSessions(1) // to ensure there is 1 session and pool is 100% utilized - .setInactiveTransactionRemovalOptions(inactiveTransactionRemovalOptions) - .setLoopFrequency(1000L) // main thread runs every 1s - .setPoolMaintainerClock(poolMaintainerClock) - .build(); - try (Spanner spanner = - SpannerOptions.newBuilder() - .setProjectId(TEST_PROJECT) - .setDatabaseRole(TEST_DATABASE_ROLE) - .setChannelProvider(channelProvider) - .setCredentials(NoCredentials.getInstance()) - .setSessionPoolOption(sessionPoolOptions) - .build() - .getService()) { - DatabaseClientImpl client = - (DatabaseClientImpl) - spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); - Instant initialExecutionTime = client.pool.poolMaintainer.lastExecutionTime; - - poolMaintainerClock.currentTimeMillis.addAndGet(Duration.ofMinutes(3).toMillis()); - - try (TransactionManager manager = client.transactionManager()) { - TransactionContext transaction = manager.begin(); - while (true) { - try { - transaction.executeUpdate(UPDATE_STATEMENT); - - // Simulate a delay of 3 minutes to ensure that the below transaction is a long-running - // one. - // As per this test, anything which takes more than 2s is long-running - poolMaintainerClock.currentTimeMillis.addAndGet(Duration.ofMinutes(3).toMillis()); - // force trigger pool maintainer to check for long-running sessions - client.pool.poolMaintainer.maintainPool(); - - manager.commit(); - assertNotNull(manager.getCommitTimestamp()); - break; - } catch (AbortedException e) { - transaction = manager.resetForRetry(); - } - } - } - Instant endExecutionTime = client.pool.poolMaintainer.lastExecutionTime; - - // first session executed update, session found to be long-running and cleaned up. - // During commit, SessionNotFound exception from backend caused replacement of session and - // transaction needs to be retried. - // On retry, session again found to be long-running and cleaned up. - // During commit, there was no exception from backend. - assertNotEquals( - endExecutionTime, - initialExecutionTime); // if session clean up task runs then these timings won't match - assertEquals(1, client.pool.numLeakedSessionsRemoved()); - assertTrue(client.pool.getNumberOfSessionsInPool() <= client.pool.totalSessions()); - } - } - - @Test - public void testPoolMaintainer_whenLongRunningPartitionedUpdateRequest_takeNoAction() { - FakeClock poolMaintainerClock = new FakeClock(); - InactiveTransactionRemovalOptions inactiveTransactionRemovalOptions = - InactiveTransactionRemovalOptions.newBuilder() - .setIdleTimeThreshold( - Duration.ofSeconds( - 2L)) // any session not used for more than 2s will be long-running - .setActionOnInactiveTransaction(ActionOnInactiveTransaction.CLOSE) - .setExecutionFrequency(Duration.ofSeconds(1)) // execute thread every 1s - .build(); - SessionPoolOptions sessionPoolOptions = - SessionPoolOptions.newBuilder() - .setMinSessions(1) - .setMaxSessions(1) // to ensure there is 1 session and pool is 100% utilized - .setInactiveTransactionRemovalOptions(inactiveTransactionRemovalOptions) - .setLoopFrequency(1000L) // main thread runs every 1s - .setPoolMaintainerClock(poolMaintainerClock) - .build(); - - try (Spanner spanner = - SpannerOptions.newBuilder() - .setProjectId(TEST_PROJECT) - .setDatabaseRole(TEST_DATABASE_ROLE) - .setChannelProvider(channelProvider) - .setCredentials(NoCredentials.getInstance()) - .setSessionPoolOption(sessionPoolOptions) - .build() - .getService()) { - DatabaseClientImpl client = - (DatabaseClientImpl) - spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); - Instant initialExecutionTime = client.pool.poolMaintainer.lastExecutionTime; - poolMaintainerClock.currentTimeMillis.addAndGet(Duration.ofMinutes(3).toMillis()); - - client.executePartitionedUpdate(UPDATE_STATEMENT); - - // Simulate a delay of 3 minutes to ensure that the below transaction is a long-running one. - // As per this test, anything which takes more than 2s is long-running - poolMaintainerClock.currentTimeMillis.addAndGet(Duration.ofMinutes(3).toMillis()); - - // force trigger pool maintainer to check for long-running sessions - client.pool.poolMaintainer.maintainPool(); - - Instant endExecutionTime = client.pool.poolMaintainer.lastExecutionTime; - - assertNotEquals( - endExecutionTime, - initialExecutionTime); // if session clean up task runs then these timings won't match - assertEquals(0, client.pool.numLeakedSessionsRemoved()); - assertTrue(client.pool.getNumberOfSessionsInPool() <= client.pool.totalSessions()); - } - } - - /** - * PDML transaction is expected to be long-running. This is indicated through session flag - * eligibleForLongRunning = true . For all other transactions which are not expected to be - * long-running eligibleForLongRunning = false. - * - *

Below tests uses a session for PDML transaction. Post that, the same session is used for - * executeUpdate(). Both transactions are long-running. The test verifies that - * eligibleForLongRunning = false for the second transaction, and it's identified as a - * long-running transaction. - */ - @Test - public void testPoolMaintainer_whenPDMLFollowedByInactiveTransaction_removeSessionsFromPool() { - assumeFalse( - "Session leaks tests are skipped for multiplexed sessions", - isMultiplexedSessionsEnabledForRW()); - FakeClock poolMaintainerClock = new FakeClock(); - InactiveTransactionRemovalOptions inactiveTransactionRemovalOptions = - InactiveTransactionRemovalOptions.newBuilder() - .setIdleTimeThreshold( - Duration.ofSeconds( - 2L)) // any session not used for more than 2s will be long-running - .setActionOnInactiveTransaction(ActionOnInactiveTransaction.CLOSE) - .setExecutionFrequency(Duration.ofSeconds(1)) // execute thread every 1s - .build(); - SessionPoolOptions sessionPoolOptions = - SessionPoolOptions.newBuilder() - .setMinSessions(1) - .setMaxSessions(1) // to ensure there is 1 session and pool is 100% utilized - .setInactiveTransactionRemovalOptions(inactiveTransactionRemovalOptions) - .setLoopFrequency(1000L) // main thread runs every 1s - .setPoolMaintainerClock(poolMaintainerClock) - .build(); - try (Spanner spanner = - SpannerOptions.newBuilder() - .setProjectId(TEST_PROJECT) - .setDatabaseRole(TEST_DATABASE_ROLE) - .setChannelProvider(channelProvider) - .setCredentials(NoCredentials.getInstance()) - .setSessionPoolOption(sessionPoolOptions) - .build() - .getService()) { - DatabaseClientImpl client = - (DatabaseClientImpl) - spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); - Instant initialExecutionTime = client.pool.poolMaintainer.lastExecutionTime; - - poolMaintainerClock.currentTimeMillis.addAndGet(Duration.ofMinutes(3).toMillis()); - - client.executePartitionedUpdate(UPDATE_STATEMENT); - - // Simulate a delay of 3 minutes to ensure that the below transaction is a long-running one. - // As per this test, anything which takes more than 2s is long-running - poolMaintainerClock.currentTimeMillis.addAndGet(Duration.ofMinutes(3).toMillis()); - - // force trigger pool maintainer to check for long-running sessions - client.pool.poolMaintainer.maintainPool(); - - Instant endExecutionTime = client.pool.poolMaintainer.lastExecutionTime; - - assertNotEquals( - endExecutionTime, - initialExecutionTime); // if session clean up task runs then these timings won't match - assertEquals(0, client.pool.numLeakedSessionsRemoved()); - assertTrue(client.pool.getNumberOfSessionsInPool() <= client.pool.totalSessions()); - - poolMaintainerClock.currentTimeMillis.addAndGet(Duration.ofMinutes(3).toMillis()); - - try (TransactionManager manager = client.transactionManager()) { - TransactionContext transaction = manager.begin(); - while (true) { - try { - transaction.executeUpdate(UPDATE_STATEMENT); - - // Simulate a delay of 3 minutes to ensure that the below transaction is a long-running - // one. - // As per this test, anything which takes more than 2s is long-running - poolMaintainerClock.currentTimeMillis.addAndGet(Duration.ofMinutes(3).toMillis()); - // force trigger pool maintainer to check for long-running sessions - client.pool.poolMaintainer.maintainPool(); - - manager.commit(); - assertNotNull(manager.getCommitTimestamp()); - break; - } catch (AbortedException e) { - transaction = manager.resetForRetry(); - } - } - } - endExecutionTime = client.pool.poolMaintainer.lastExecutionTime; - - // first session executed update, session found to be long-running and cleaned up. - // During commit, SessionNotFound exception from backend caused replacement of session and - // transaction needs to be retried. - // On retry, session again found to be long-running and cleaned up. - // During commit, there was no exception from backend. - assertNotEquals( - endExecutionTime, - initialExecutionTime); // if session clean up task runs then these timings won't match - assertEquals(1, client.pool.numLeakedSessionsRemoved()); - assertTrue(client.pool.getNumberOfSessionsInPool() <= client.pool.totalSessions()); - } - } - - @Test - public void - testPoolMaintainer_whenLongRunningReadsUsingTransactionRunner_retainSessionForTransaction() - throws Exception { - FakeClock poolMaintainerClock = new FakeClock(); - InactiveTransactionRemovalOptions inactiveTransactionRemovalOptions = - InactiveTransactionRemovalOptions.newBuilder() - .setIdleTimeThreshold( - Duration.ofSeconds( - 3L)) // any session not used for more than 3s will be long-running - .setActionOnInactiveTransaction(ActionOnInactiveTransaction.CLOSE) - .setExecutionFrequency(Duration.ofSeconds(1)) // execute thread every 1s - .build(); - SessionPoolOptions sessionPoolOptions = - SessionPoolOptions.newBuilder() - .setMinSessions(1) - .setMaxSessions(1) // to ensure there is 1 session and pool is 100% utilized - .setInactiveTransactionRemovalOptions(inactiveTransactionRemovalOptions) - .setLoopFrequency(1000L) // main thread runs every 1s - .setPoolMaintainerClock(poolMaintainerClock) - .build(); - try (Spanner spanner = - SpannerOptions.newBuilder() - .setProjectId(TEST_PROJECT) - .setDatabaseRole(TEST_DATABASE_ROLE) - .setChannelProvider(channelProvider) - .setCredentials(NoCredentials.getInstance()) - .setSessionPoolOption(sessionPoolOptions) - .build() - .getService()) { - DatabaseClientImpl client = - (DatabaseClientImpl) - spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); - Instant initialExecutionTime = client.pool.poolMaintainer.lastExecutionTime; - - TransactionRunner runner = client.readWriteTransaction(); - runner.run( - transaction -> { - try (ResultSet resultSet = - transaction.read( - READ_TABLE_NAME, - KeySet.singleKey(Key.of(1L)), - READ_COLUMN_NAMES, - Options.priority(RpcPriority.HIGH))) { - consumeResults(resultSet); - } - poolMaintainerClock.currentTimeMillis.addAndGet(Duration.ofMillis(1050).toMillis()); - - try (ResultSet resultSet = - transaction.read( - READ_TABLE_NAME, - KeySet.singleKey(Key.of(1L)), - READ_COLUMN_NAMES, - Options.priority(RpcPriority.HIGH))) { - consumeResults(resultSet); - } - poolMaintainerClock.currentTimeMillis.addAndGet(Duration.ofMillis(2050).toMillis()); - - // force trigger pool maintainer to check for long-running sessions - client.pool.poolMaintainer.maintainPool(); - - return null; - }); - - Instant endExecutionTime = client.pool.poolMaintainer.lastExecutionTime; - - assertNotEquals( - endExecutionTime, - initialExecutionTime); // if session clean up task runs then these timings won't match - assertEquals(0, client.pool.numLeakedSessionsRemoved()); - } - } - - @Test - public void - testPoolMaintainer_whenLongRunningQueriesUsingTransactionRunner_retainSessionForTransaction() { - FakeClock poolMaintainerClock = new FakeClock(); - InactiveTransactionRemovalOptions inactiveTransactionRemovalOptions = - InactiveTransactionRemovalOptions.newBuilder() - .setIdleTimeThreshold( - Duration.ofSeconds( - 3L)) // any session not used for more than 3s will be long-running - .setActionOnInactiveTransaction(ActionOnInactiveTransaction.CLOSE) - .setExecutionFrequency(Duration.ofSeconds(1)) // execute thread every 1s - .build(); - SessionPoolOptions sessionPoolOptions = - SessionPoolOptions.newBuilder() - .setMinSessions(1) - .setMaxSessions(1) // to ensure there is 1 session and pool is 100% utilized - .setInactiveTransactionRemovalOptions(inactiveTransactionRemovalOptions) - .setLoopFrequency(1000L) // main thread runs every 1s - .setPoolMaintainerClock(poolMaintainerClock) - .build(); - try (Spanner spanner = - SpannerOptions.newBuilder() - .setProjectId(TEST_PROJECT) - .setDatabaseRole(TEST_DATABASE_ROLE) - .setChannelProvider(channelProvider) - .setCredentials(NoCredentials.getInstance()) - .setSessionPoolOption(sessionPoolOptions) - .build() - .getService()) { - DatabaseClientImpl client = - (DatabaseClientImpl) - spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); - Instant initialExecutionTime = client.pool.poolMaintainer.lastExecutionTime; - - TransactionRunner runner = client.readWriteTransaction(); - runner.run( - transaction -> { - try (ResultSet resultSet = transaction.executeQuery(SELECT1)) { - consumeResults(resultSet); - } - poolMaintainerClock.currentTimeMillis.addAndGet(Duration.ofMillis(1050).toMillis()); - - try (ResultSet resultSet = transaction.executeQuery(SELECT1)) { - consumeResults(resultSet); - } - poolMaintainerClock.currentTimeMillis.addAndGet(Duration.ofMillis(2050).toMillis()); - - // force trigger pool maintainer to check for long-running sessions - client.pool.poolMaintainer.maintainPool(); - - return null; - }); - - Instant endExecutionTime = client.pool.poolMaintainer.lastExecutionTime; - - assertNotEquals( - endExecutionTime, - initialExecutionTime); // if session clean up task runs then these timings won't match - assertEquals(0, client.pool.numLeakedSessionsRemoved()); - } - } - - @Test - public void - testPoolMaintainer_whenLongRunningUpdatesUsingTransactionManager_retainSessionForTransaction() { - FakeClock poolMaintainerClock = new FakeClock(); - InactiveTransactionRemovalOptions inactiveTransactionRemovalOptions = - InactiveTransactionRemovalOptions.newBuilder() - .setIdleTimeThreshold( - Duration.ofSeconds( - 3L)) // any session not used for more than 3s will be long-running - .setActionOnInactiveTransaction(ActionOnInactiveTransaction.CLOSE) - .setExecutionFrequency(Duration.ofSeconds(1)) // execute thread every 1s - .build(); - SessionPoolOptions sessionPoolOptions = - SessionPoolOptions.newBuilder() - .setMinSessions(1) - .setMaxSessions(1) // to ensure there is 1 session and pool is 100% utilized - .setInactiveTransactionRemovalOptions(inactiveTransactionRemovalOptions) - .setLoopFrequency(1000L) // main thread runs every 1s - .setPoolMaintainerClock(poolMaintainerClock) - .build(); - try (Spanner spanner = - SpannerOptions.newBuilder() - .setProjectId(TEST_PROJECT) - .setDatabaseRole(TEST_DATABASE_ROLE) - .setChannelProvider(channelProvider) - .setCredentials(NoCredentials.getInstance()) - .setSessionPoolOption(sessionPoolOptions) - .build() - .getService()) { - DatabaseClientImpl client = - (DatabaseClientImpl) - spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); - Instant initialExecutionTime = client.pool.poolMaintainer.lastExecutionTime; - - try (TransactionManager manager = client.transactionManager()) { - TransactionContext transaction = manager.begin(); - while (true) { - try { - transaction.executeUpdate(UPDATE_STATEMENT); - poolMaintainerClock.currentTimeMillis.addAndGet(Duration.ofMillis(1050).toMillis()); - - transaction.executeUpdate(UPDATE_STATEMENT); - poolMaintainerClock.currentTimeMillis.addAndGet(Duration.ofMillis(2050).toMillis()); - - // force trigger pool maintainer to check for long-running sessions - client.pool.poolMaintainer.maintainPool(); - - manager.commit(); - assertNotNull(manager.getCommitTimestamp()); - break; - } catch (AbortedException e) { - transaction = manager.resetForRetry(); - } - } - } - Instant endExecutionTime = client.pool.poolMaintainer.lastExecutionTime; - - assertNotEquals( - endExecutionTime, - initialExecutionTime); // if session clean up task runs then these timings won't match - assertEquals(0, client.pool.numLeakedSessionsRemoved()); - assertTrue(client.pool.getNumberOfSessionsInPool() <= client.pool.totalSessions()); - } - } - - @Test - public void - testPoolMaintainer_whenLongRunningReadsUsingTransactionManager_retainSessionForTransaction() { - FakeClock poolMaintainerClock = new FakeClock(); - InactiveTransactionRemovalOptions inactiveTransactionRemovalOptions = - InactiveTransactionRemovalOptions.newBuilder() - .setIdleTimeThreshold( - Duration.ofSeconds( - 3L)) // any session not used for more than 3s will be long-running - .setActionOnInactiveTransaction(ActionOnInactiveTransaction.CLOSE) - .setExecutionFrequency(Duration.ofSeconds(1)) // execute thread every 1s - .build(); - SessionPoolOptions sessionPoolOptions = - SessionPoolOptions.newBuilder() - .setMinSessions(1) - .setMaxSessions(1) // to ensure there is 1 session and pool is 100% utilized - .setInactiveTransactionRemovalOptions(inactiveTransactionRemovalOptions) - .setLoopFrequency(1000L) // main thread runs every 1s - .setPoolMaintainerClock(poolMaintainerClock) - .build(); - try (Spanner spanner = - SpannerOptions.newBuilder() - .setProjectId(TEST_PROJECT) - .setDatabaseRole(TEST_DATABASE_ROLE) - .setChannelProvider(channelProvider) - .setCredentials(NoCredentials.getInstance()) - .setSessionPoolOption(sessionPoolOptions) - .build() - .getService()) { - DatabaseClientImpl client = - (DatabaseClientImpl) - spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); - Instant initialExecutionTime = client.pool.poolMaintainer.lastExecutionTime; - - try (TransactionManager manager = client.transactionManager()) { - TransactionContext transaction = manager.begin(); - while (true) { - try { - try (ResultSet resultSet = - transaction.read( - READ_TABLE_NAME, - KeySet.singleKey(Key.of(1L)), - READ_COLUMN_NAMES, - Options.priority(RpcPriority.HIGH))) { - consumeResults(resultSet); - } - poolMaintainerClock.currentTimeMillis.addAndGet(Duration.ofMillis(1050).toMillis()); - - try (ResultSet resultSet = - transaction.read( - READ_TABLE_NAME, - KeySet.singleKey(Key.of(1L)), - READ_COLUMN_NAMES, - Options.priority(RpcPriority.HIGH))) { - consumeResults(resultSet); - } - poolMaintainerClock.currentTimeMillis.addAndGet(Duration.ofMillis(2050).toMillis()); - - // force trigger pool maintainer to check for long-running sessions - client.pool.poolMaintainer.maintainPool(); - - manager.commit(); - assertNotNull(manager.getCommitTimestamp()); - break; - } catch (AbortedException e) { - transaction = manager.resetForRetry(); - } - } - } - Instant endExecutionTime = client.pool.poolMaintainer.lastExecutionTime; - - assertNotEquals( - endExecutionTime, - initialExecutionTime); // if session clean up task runs then these timings won't match - assertEquals(0, client.pool.numLeakedSessionsRemoved()); - assertTrue(client.pool.getNumberOfSessionsInPool() <= client.pool.totalSessions()); - } - } - - @Test - public void - testPoolMaintainer_whenLongRunningReadRowUsingTransactionManager_retainSessionForTransaction() { - FakeClock poolMaintainerClock = new FakeClock(); - InactiveTransactionRemovalOptions inactiveTransactionRemovalOptions = - InactiveTransactionRemovalOptions.newBuilder() - .setIdleTimeThreshold( - Duration.ofSeconds( - 3L)) // any session not used for more than 3s will be long-running - .setActionOnInactiveTransaction(ActionOnInactiveTransaction.CLOSE) - .setExecutionFrequency(Duration.ofSeconds(1)) // execute thread every 1s - .build(); - SessionPoolOptions sessionPoolOptions = - SessionPoolOptions.newBuilder() - .setMinSessions(1) - .setMaxSessions(1) // to ensure there is 1 session and pool is 100% utilized - .setInactiveTransactionRemovalOptions(inactiveTransactionRemovalOptions) - .setLoopFrequency(1000L) // main thread runs every 1s - .setPoolMaintainerClock(poolMaintainerClock) - .build(); - try (Spanner spanner = - SpannerOptions.newBuilder() - .setProjectId(TEST_PROJECT) - .setDatabaseRole(TEST_DATABASE_ROLE) - .setChannelProvider(channelProvider) - .setCredentials(NoCredentials.getInstance()) - .setSessionPoolOption(sessionPoolOptions) - .build() - .getService()) { - DatabaseClientImpl client = - (DatabaseClientImpl) - spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); - Instant initialExecutionTime = client.pool.poolMaintainer.lastExecutionTime; - - try (TransactionManager manager = client.transactionManager()) { - TransactionContext transaction = manager.begin(); - while (true) { - try { - transaction.readRow(READ_TABLE_NAME, Key.of(1L), READ_COLUMN_NAMES); - - poolMaintainerClock.currentTimeMillis.addAndGet(Duration.ofMillis(1050).toMillis()); - - transaction.readRow(READ_TABLE_NAME, Key.of(1L), READ_COLUMN_NAMES); - - poolMaintainerClock.currentTimeMillis.addAndGet(Duration.ofMillis(2050).toMillis()); - - // force trigger pool maintainer to check for long-running sessions - client.pool.poolMaintainer.maintainPool(); - - manager.commit(); - assertNotNull(manager.getCommitTimestamp()); - break; - } catch (AbortedException e) { - transaction = manager.resetForRetry(); - } - } - } - Instant endExecutionTime = client.pool.poolMaintainer.lastExecutionTime; - - assertNotEquals( - endExecutionTime, - initialExecutionTime); // if session clean up task runs then these timings won't match - assertEquals(0, client.pool.numLeakedSessionsRemoved()); - assertTrue(client.pool.getNumberOfSessionsInPool() <= client.pool.totalSessions()); - } - } - - @Test - public void - testPoolMaintainer_whenLongRunningAnalyzeUpdateStatementUsingTransactionManager_retainSessionForTransaction() { - FakeClock poolMaintainerClock = new FakeClock(); - InactiveTransactionRemovalOptions inactiveTransactionRemovalOptions = - InactiveTransactionRemovalOptions.newBuilder() - .setIdleTimeThreshold( - Duration.ofSeconds( - 3L)) // any session not used for more than 3s will be long-running - .setActionOnInactiveTransaction(ActionOnInactiveTransaction.CLOSE) - .setExecutionFrequency(Duration.ofSeconds(1)) // execute thread every 1s - .build(); - SessionPoolOptions sessionPoolOptions = - SessionPoolOptions.newBuilder() - .setMinSessions(1) - .setMaxSessions(1) // to ensure there is 1 session and pool is 100% utilized - .setInactiveTransactionRemovalOptions(inactiveTransactionRemovalOptions) - .setLoopFrequency(1000L) // main thread runs every 1s - .setPoolMaintainerClock(poolMaintainerClock) - .build(); - try (Spanner spanner = - SpannerOptions.newBuilder() - .setProjectId(TEST_PROJECT) - .setDatabaseRole(TEST_DATABASE_ROLE) - .setChannelProvider(channelProvider) - .setCredentials(NoCredentials.getInstance()) - .setSessionPoolOption(sessionPoolOptions) - .build() - .getService()) { - DatabaseClientImpl client = - (DatabaseClientImpl) - spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); - Instant initialExecutionTime = client.pool.poolMaintainer.lastExecutionTime; - - try (TransactionManager manager = client.transactionManager()) { - TransactionContext transaction = manager.begin(); - while (true) { - try { - try (ResultSet resultSet = - transaction.analyzeUpdateStatement(UPDATE_STATEMENT, QueryAnalyzeMode.PROFILE)) { - consumeResults(resultSet); - } - poolMaintainerClock.currentTimeMillis.addAndGet(Duration.ofMillis(1050).toMillis()); - - try (ResultSet resultSet = - transaction.analyzeUpdateStatement(UPDATE_STATEMENT, QueryAnalyzeMode.PROFILE)) { - consumeResults(resultSet); - } - poolMaintainerClock.currentTimeMillis.addAndGet(Duration.ofMillis(2050).toMillis()); - - // force trigger pool maintainer to check for long-running sessions - client.pool.poolMaintainer.maintainPool(); - - manager.commit(); - assertNotNull(manager.getCommitTimestamp()); - break; - } catch (AbortedException e) { - transaction = manager.resetForRetry(); - } - } - } - Instant endExecutionTime = client.pool.poolMaintainer.lastExecutionTime; - - assertNotEquals( - endExecutionTime, - initialExecutionTime); // if session clean up task runs then these timings won't match - assertEquals(0, client.pool.numLeakedSessionsRemoved()); - assertTrue(client.pool.getNumberOfSessionsInPool() <= client.pool.totalSessions()); - } - } - - @Test - public void - testPoolMaintainer_whenLongRunningBatchUpdatesUsingTransactionManager_retainSessionForTransaction() { - FakeClock poolMaintainerClock = new FakeClock(); - InactiveTransactionRemovalOptions inactiveTransactionRemovalOptions = - InactiveTransactionRemovalOptions.newBuilder() - .setIdleTimeThreshold( - Duration.ofSeconds( - 3L)) // any session not used for more than 3s will be long-running - .setActionOnInactiveTransaction(ActionOnInactiveTransaction.CLOSE) - .setExecutionFrequency(Duration.ofSeconds(1)) // execute thread every 1s - .build(); - SessionPoolOptions sessionPoolOptions = - SessionPoolOptions.newBuilder() - .setMinSessions(1) - .setMaxSessions(1) // to ensure there is 1 session and pool is 100% utilized - .setInactiveTransactionRemovalOptions(inactiveTransactionRemovalOptions) - .setLoopFrequency(1000L) // main thread runs every 1s - .setPoolMaintainerClock(poolMaintainerClock) - .build(); - try (Spanner spanner = - SpannerOptions.newBuilder() - .setProjectId(TEST_PROJECT) - .setDatabaseRole(TEST_DATABASE_ROLE) - .setChannelProvider(channelProvider) - .setCredentials(NoCredentials.getInstance()) - .setSessionPoolOption(sessionPoolOptions) - .build() - .getService()) { - DatabaseClientImpl client = - (DatabaseClientImpl) - spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); - Instant initialExecutionTime = client.pool.poolMaintainer.lastExecutionTime; - - try (TransactionManager manager = client.transactionManager()) { - TransactionContext transaction = manager.begin(); - while (true) { - try { - transaction.batchUpdate(Lists.newArrayList(UPDATE_STATEMENT)); - - poolMaintainerClock.currentTimeMillis.addAndGet(Duration.ofMillis(1050).toMillis()); - - transaction.batchUpdate(Lists.newArrayList(UPDATE_STATEMENT)); - - poolMaintainerClock.currentTimeMillis.addAndGet(Duration.ofMillis(2050).toMillis()); - - // force trigger pool maintainer to check for long-running sessions - client.pool.poolMaintainer.maintainPool(); - - manager.commit(); - assertNotNull(manager.getCommitTimestamp()); - break; - } catch (AbortedException e) { - transaction = manager.resetForRetry(); - } - } - } - Instant endExecutionTime = client.pool.poolMaintainer.lastExecutionTime; - - assertNotEquals( - endExecutionTime, - initialExecutionTime); // if session clean up task runs then these timings won't match - assertEquals(0, client.pool.numLeakedSessionsRemoved()); - assertTrue(client.pool.getNumberOfSessionsInPool() <= client.pool.totalSessions()); - } - } - - @Test - public void - testPoolMaintainer_whenLongRunningBatchUpdatesAsyncUsingTransactionManager_retainSessionForTransaction() { - FakeClock poolMaintainerClock = new FakeClock(); - InactiveTransactionRemovalOptions inactiveTransactionRemovalOptions = - InactiveTransactionRemovalOptions.newBuilder() - .setIdleTimeThreshold( - Duration.ofSeconds( - 3L)) // any session not used for more than 3s will be long-running - .setActionOnInactiveTransaction(ActionOnInactiveTransaction.CLOSE) - .setExecutionFrequency(Duration.ofSeconds(1)) // execute thread every 1s - .build(); - SessionPoolOptions sessionPoolOptions = - SessionPoolOptions.newBuilder() - .setMinSessions(1) - .setMaxSessions(1) // to ensure there is 1 session and pool is 100% utilized - .setInactiveTransactionRemovalOptions(inactiveTransactionRemovalOptions) - .setLoopFrequency(1000L) // main thread runs every 1s - .setPoolMaintainerClock(poolMaintainerClock) - .build(); - try (Spanner spanner = - SpannerOptions.newBuilder() - .setProjectId(TEST_PROJECT) - .setDatabaseRole(TEST_DATABASE_ROLE) - .setChannelProvider(channelProvider) - .setCredentials(NoCredentials.getInstance()) - .setSessionPoolOption(sessionPoolOptions) - .build() - .getService()) { - DatabaseClientImpl client = - (DatabaseClientImpl) - spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); - Instant initialExecutionTime = client.pool.poolMaintainer.lastExecutionTime; - - try (TransactionManager manager = client.transactionManager()) { - TransactionContext transaction = manager.begin(); - while (true) { - try { - transaction.batchUpdateAsync(Lists.newArrayList(UPDATE_STATEMENT)); - - poolMaintainerClock.currentTimeMillis.addAndGet(Duration.ofMillis(1050).toMillis()); - - transaction.batchUpdateAsync(Lists.newArrayList(UPDATE_STATEMENT)); - - poolMaintainerClock.currentTimeMillis.addAndGet(Duration.ofMillis(2050).toMillis()); - - // force trigger pool maintainer to check for long-running sessions - client.pool.poolMaintainer.maintainPool(); - - manager.commit(); - assertNotNull(manager.getCommitTimestamp()); - break; - } catch (AbortedException e) { - transaction = manager.resetForRetry(); - } - } - } - Instant endExecutionTime = client.pool.poolMaintainer.lastExecutionTime; - - assertNotEquals( - endExecutionTime, - initialExecutionTime); // if session clean up task runs then these timings won't match - assertEquals(0, client.pool.numLeakedSessionsRemoved()); - assertTrue(client.pool.getNumberOfSessionsInPool() <= client.pool.totalSessions()); - } - } - - @Test - public void - testPoolMaintainer_whenLongRunningExecuteQueryUsingTransactionManager_retainSessionForTransaction() { - FakeClock poolMaintainerClock = new FakeClock(); - InactiveTransactionRemovalOptions inactiveTransactionRemovalOptions = - InactiveTransactionRemovalOptions.newBuilder() - .setIdleTimeThreshold( - Duration.ofSeconds( - 3L)) // any session not used for more than 3s will be long-running - .setActionOnInactiveTransaction(ActionOnInactiveTransaction.CLOSE) - .setExecutionFrequency(Duration.ofSeconds(1)) // execute thread every 1s - .build(); - SessionPoolOptions sessionPoolOptions = - SessionPoolOptions.newBuilder() - .setMinSessions(1) - .setMaxSessions(1) // to ensure there is 1 session and pool is 100% utilized - .setInactiveTransactionRemovalOptions(inactiveTransactionRemovalOptions) - .setLoopFrequency(1000L) // main thread runs every 1s - .setPoolMaintainerClock(poolMaintainerClock) - .build(); - try (Spanner spanner = - SpannerOptions.newBuilder() - .setProjectId(TEST_PROJECT) - .setDatabaseRole(TEST_DATABASE_ROLE) - .setChannelProvider(channelProvider) - .setCredentials(NoCredentials.getInstance()) - .setSessionPoolOption(sessionPoolOptions) - .build() - .getService()) { - DatabaseClientImpl client = - (DatabaseClientImpl) - spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); - Instant initialExecutionTime = client.pool.poolMaintainer.lastExecutionTime; - - try (TransactionManager manager = client.transactionManager()) { - TransactionContext transaction = manager.begin(); - while (true) { - try { - try (ResultSet resultSet = transaction.executeQuery(SELECT1)) { - consumeResults(resultSet); - } - poolMaintainerClock.currentTimeMillis.addAndGet(Duration.ofMillis(1050).toMillis()); - - try (ResultSet resultSet = transaction.executeQuery(SELECT1)) { - consumeResults(resultSet); - } - poolMaintainerClock.currentTimeMillis.addAndGet(Duration.ofMillis(2050).toMillis()); - - // force trigger pool maintainer to check for long-running sessions - client.pool.poolMaintainer.maintainPool(); - - manager.commit(); - assertNotNull(manager.getCommitTimestamp()); - break; - } catch (AbortedException e) { - transaction = manager.resetForRetry(); - } - } - } - Instant endExecutionTime = client.pool.poolMaintainer.lastExecutionTime; - - assertNotEquals( - endExecutionTime, - initialExecutionTime); // if session clean up task runs then these timings won't match - assertEquals(0, client.pool.numLeakedSessionsRemoved()); - assertTrue(client.pool.getNumberOfSessionsInPool() <= client.pool.totalSessions()); - } - } - - @Test - public void - testPoolMaintainer_whenLongRunningExecuteQueryAsyncUsingTransactionManager_retainSessionForTransaction() { - FakeClock poolMaintainerClock = new FakeClock(); - InactiveTransactionRemovalOptions inactiveTransactionRemovalOptions = - InactiveTransactionRemovalOptions.newBuilder() - .setIdleTimeThreshold( - Duration.ofSeconds( - 3L)) // any session not used for more than 3s will be long-running - .setActionOnInactiveTransaction(ActionOnInactiveTransaction.CLOSE) - .setExecutionFrequency(Duration.ofSeconds(1)) // execute thread every 1s - .build(); - SessionPoolOptions sessionPoolOptions = - SessionPoolOptions.newBuilder() - .setMinSessions(1) - .setMaxSessions(1) // to ensure there is 1 session and pool is 100% utilized - .setInactiveTransactionRemovalOptions(inactiveTransactionRemovalOptions) - .setLoopFrequency(1000L) // main thread runs every 1s - .setPoolMaintainerClock(poolMaintainerClock) - .build(); - try (Spanner spanner = - SpannerOptions.newBuilder() - .setProjectId(TEST_PROJECT) - .setDatabaseRole(TEST_DATABASE_ROLE) - .setChannelProvider(channelProvider) - .setCredentials(NoCredentials.getInstance()) - .setSessionPoolOption(sessionPoolOptions) - .build() - .getService()) { - DatabaseClientImpl client = - (DatabaseClientImpl) - spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); - Instant initialExecutionTime = client.pool.poolMaintainer.lastExecutionTime; - - try (TransactionManager manager = client.transactionManager()) { - TransactionContext transaction = manager.begin(); - while (true) { - try { - try (ResultSet resultSet = transaction.executeQueryAsync(SELECT1)) { - consumeResults(resultSet); - } - poolMaintainerClock.currentTimeMillis.addAndGet(Duration.ofMillis(1050).toMillis()); - - try (ResultSet resultSet = transaction.executeQueryAsync(SELECT1)) { - consumeResults(resultSet); - } - poolMaintainerClock.currentTimeMillis.addAndGet(Duration.ofMillis(2050).toMillis()); - - // force trigger pool maintainer to check for long-running sessions - client.pool.poolMaintainer.maintainPool(); - - manager.commit(); - assertNotNull(manager.getCommitTimestamp()); - break; - } catch (AbortedException e) { - transaction = manager.resetForRetry(); - } - } - } - Instant endExecutionTime = client.pool.poolMaintainer.lastExecutionTime; - - assertNotEquals( - endExecutionTime, - initialExecutionTime); // if session clean up task runs then these timings won't match - assertEquals(0, client.pool.numLeakedSessionsRemoved()); - assertTrue(client.pool.getNumberOfSessionsInPool() <= client.pool.totalSessions()); - } - } - - @Test - public void - testPoolMaintainer_whenLongRunningAnalyzeQueryUsingTransactionManager_retainSessionForTransaction() { - FakeClock poolMaintainerClock = new FakeClock(); - InactiveTransactionRemovalOptions inactiveTransactionRemovalOptions = - InactiveTransactionRemovalOptions.newBuilder() - .setIdleTimeThreshold( - Duration.ofSeconds( - 3L)) // any session not used for more than 3s will be long-running - .setActionOnInactiveTransaction(ActionOnInactiveTransaction.CLOSE) - .setExecutionFrequency(Duration.ofSeconds(1)) // execute thread every 1s - .build(); - SessionPoolOptions sessionPoolOptions = - SessionPoolOptions.newBuilder() - .setMinSessions(1) - .setMaxSessions(1) // to ensure there is 1 session and pool is 100% utilized - .setInactiveTransactionRemovalOptions(inactiveTransactionRemovalOptions) - .setLoopFrequency(1000L) // main thread runs every 1s - .setPoolMaintainerClock(poolMaintainerClock) - .build(); - try (Spanner spanner = - SpannerOptions.newBuilder() - .setProjectId(TEST_PROJECT) - .setDatabaseRole(TEST_DATABASE_ROLE) - .setChannelProvider(channelProvider) - .setCredentials(NoCredentials.getInstance()) - .setSessionPoolOption(sessionPoolOptions) - .build() - .getService()) { - DatabaseClientImpl client = - (DatabaseClientImpl) - spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); - Instant initialExecutionTime = client.pool.poolMaintainer.lastExecutionTime; - - try (TransactionManager manager = client.transactionManager()) { - TransactionContext transaction = manager.begin(); - while (true) { - try { - try (ResultSet resultSet = - transaction.analyzeQuery(SELECT1, QueryAnalyzeMode.PROFILE)) { - consumeResults(resultSet); - } - poolMaintainerClock.currentTimeMillis.addAndGet(Duration.ofMillis(1050).toMillis()); - - try (ResultSet resultSet = - transaction.analyzeQuery(SELECT1, QueryAnalyzeMode.PROFILE)) { - consumeResults(resultSet); - } - poolMaintainerClock.currentTimeMillis.addAndGet(Duration.ofMillis(2050).toMillis()); - - // force trigger pool maintainer to check for long-running sessions - client.pool.poolMaintainer.maintainPool(); - - manager.commit(); - assertNotNull(manager.getCommitTimestamp()); - break; - } catch (AbortedException e) { - transaction = manager.resetForRetry(); - } - } - } - Instant endExecutionTime = client.pool.poolMaintainer.lastExecutionTime; - - assertNotEquals( - endExecutionTime, - initialExecutionTime); // if session clean up task runs then these timings won't match - assertEquals(0, client.pool.numLeakedSessionsRemoved()); - assertTrue(client.pool.getNumberOfSessionsInPool() <= client.pool.totalSessions()); - } - } - @Test public void testWrite() { DatabaseClient client = @@ -2360,17 +1272,11 @@ public void singleUse() { 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(); - if (!isMultiplexedSessionsEnabled()) { - assertThat(checkedOut).hasSize(1); - } assertThat(rs.getLong(0)).isEqualTo(1L); assertThat(rs.next()).isFalse(); } - assertThat(checkedOut).isEmpty(); } @Test @@ -2909,17 +1815,11 @@ public void testPartitionedDmlDoesNotTimeout() { })); assertEquals(ErrorCode.DEADLINE_EXCEEDED, e.getErrorCode()); - DatabaseClientImpl dbImpl = ((DatabaseClientImpl) client); - int channelId = 0; - try (Session session = dbImpl.getSession()) { - channelId = ((PooledSessionFuture) session).getChannel(); - } - int dbId = dbImpl.dbId; long NON_DETERMINISTIC = XGoogSpannerRequestIdTest.NON_DETERMINISTIC; XGoogSpannerRequestIdTest.MethodAndRequestId[] wantStreamingValues = { XGoogSpannerRequestIdTest.ofMethodAndRequestId( "google.spanner.v1.Spanner/ExecuteStreamingSql", - new XGoogSpannerRequestId(NON_DETERMINISTIC, channelId, 6, 1)), + new XGoogSpannerRequestId(NON_DETERMINISTIC, NON_DETERMINISTIC, 6, 1)), }; if (false) { // TODO(@odeke-em): enable in next PRs. xGoogReqIdInterceptor.checkExpectedStreamingXGoogRequestIds(wantStreamingValues); @@ -2928,13 +1828,13 @@ public void testPartitionedDmlDoesNotTimeout() { XGoogSpannerRequestIdTest.MethodAndRequestId[] wantUnaryValues = { XGoogSpannerRequestIdTest.ofMethodAndRequestId( "google.spanner.v1.Spanner/BeginTransaction", - new XGoogSpannerRequestId(NON_DETERMINISTIC, channelId, 7, 1)), + new XGoogSpannerRequestId(NON_DETERMINISTIC, NON_DETERMINISTIC, 7, 1)), XGoogSpannerRequestIdTest.ofMethodAndRequestId( "google.spanner.v1.Spanner/CreateSession", new XGoogSpannerRequestId(NON_DETERMINISTIC, 0, 1, 1)), XGoogSpannerRequestIdTest.ofMethodAndRequestId( "google.spanner.v1.Spanner/ExecuteSql", - new XGoogSpannerRequestId(NON_DETERMINISTIC, channelId, 8, 1)), + new XGoogSpannerRequestId(NON_DETERMINISTIC, NON_DETERMINISTIC, 8, 1)), }; if (false) { // TODO(@odeke-em): enable in next PRs. xGoogReqIdInterceptor.checkExpectedUnaryXGoogRequestIdsAsSuffixes(wantUnaryValues); @@ -3022,17 +1922,11 @@ public void testPartitionedDmlWithHigherTimeout() { assertThat(e.getErrorCode()).isEqualTo(ErrorCode.DEADLINE_EXCEEDED); assertThat(updateCount).isEqualTo(UPDATE_COUNT); - DatabaseClientImpl dbImpl = ((DatabaseClientImpl) client); - int channelId = 0; - try (Session session = dbImpl.getSession()) { - channelId = ((PooledSessionFuture) session).getChannel(); - } - int dbId = dbImpl.dbId; long NON_DETERMINISTIC = XGoogSpannerRequestIdTest.NON_DETERMINISTIC; XGoogSpannerRequestIdTest.MethodAndRequestId[] wantStreamingValues = { XGoogSpannerRequestIdTest.ofMethodAndRequestId( "google.spanner.v1.Spanner/ExecuteStreamingSql", - new XGoogSpannerRequestId(NON_DETERMINISTIC, channelId, 6, 1)), + new XGoogSpannerRequestId(NON_DETERMINISTIC, NON_DETERMINISTIC, 6, 1)), }; if (false) { // TODO(@odeke-em): enable in next PRs. @@ -3042,13 +1936,13 @@ public void testPartitionedDmlWithHigherTimeout() { XGoogSpannerRequestIdTest.MethodAndRequestId[] wantUnaryValues = { XGoogSpannerRequestIdTest.ofMethodAndRequestId( "google.spanner.v1.Spanner/BeginTransaction", - new XGoogSpannerRequestId(NON_DETERMINISTIC, channelId, 7, 1)), + new XGoogSpannerRequestId(NON_DETERMINISTIC, NON_DETERMINISTIC, 7, 1)), XGoogSpannerRequestIdTest.ofMethodAndRequestId( "google.spanner.v1.Spanner/CreateSession", new XGoogSpannerRequestId(NON_DETERMINISTIC, 0, 1, 1)), XGoogSpannerRequestIdTest.ofMethodAndRequestId( "google.spanner.v1.Spanner/ExecuteSql", - new XGoogSpannerRequestId(NON_DETERMINISTIC, channelId, 8, 1)), + new XGoogSpannerRequestId(NON_DETERMINISTIC, NON_DETERMINISTIC, 8, 1)), }; if (false) { // TODO(@odeke-em): enable in next PRs. xGoogReqIdInterceptor.checkExpectedUnaryXGoogRequestIdsAsSuffixes(wantUnaryValues); @@ -3090,7 +1984,7 @@ public void testDatabaseOrInstanceDoesNotExistOnInitialization() throws Exceptio .setCredentials(NoCredentials.getInstance()) .build() .getService()) { - mockSpanner.setBatchCreateSessionsExecutionTime( + mockSpanner.setCreateSessionExecutionTime( SimulatedExecutionTime.ofStickyException(exception)); DatabaseClientImpl dbClient = (DatabaseClientImpl) @@ -3099,13 +1993,12 @@ public void testDatabaseOrInstanceDoesNotExistOnInitialization() throws Exceptio // Wait until session creation has finished. Stopwatch watch = Stopwatch.createStarted(); while (watch.elapsed(TimeUnit.SECONDS) < 5 - && dbClient.pool.getNumberOfSessionsBeingCreated() > 0) { + && dbClient.multiplexedSessionDatabaseClient.isValid()) { //noinspection BusyWait Thread.sleep(1L); } // All session creation should fail and stop trying. - assertThat(dbClient.pool.getNumberOfSessionsInPool()).isEqualTo(0); - assertThat(dbClient.pool.getNumberOfSessionsBeingCreated()).isEqualTo(0); + assertFalse(dbClient.isValid()); mockSpanner.reset(); mockSpanner.removeAllExecutionTimes(); } @@ -3189,57 +2082,6 @@ public void testDatabaseOrInstanceDoesNotExistOnCreate() { } } - @Test - public void testDatabaseOrInstanceDoesNotExistOnReplenish() throws Exception { - StatusRuntimeException[] exceptions = - new StatusRuntimeException[] { - SpannerExceptionFactoryTest.newStatusResourceNotFoundException( - "Database", SpannerExceptionFactory.DATABASE_RESOURCE_TYPE, DATABASE_NAME), - SpannerExceptionFactoryTest.newStatusResourceNotFoundException( - "Instance", SpannerExceptionFactory.INSTANCE_RESOURCE_TYPE, INSTANCE_NAME) - }; - for (StatusRuntimeException exception : exceptions) { - try (Spanner spanner = - SpannerOptions.newBuilder() - .setProjectId(TEST_PROJECT) - .setChannelProvider(channelProvider) - .setCredentials(NoCredentials.getInstance()) - .build() - .getService()) { - mockSpanner.setBatchCreateSessionsExecutionTime( - SimulatedExecutionTime.ofStickyException(exception)); - DatabaseClientImpl dbClient = - (DatabaseClientImpl) - spanner.getDatabaseClient( - DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); - // Wait until session creation has finished. - Stopwatch watch = Stopwatch.createStarted(); - while (watch.elapsed(TimeUnit.SECONDS) < 5 - && dbClient.pool.getNumberOfSessionsBeingCreated() > 0) { - //noinspection BusyWait - Thread.sleep(1L); - } - // All session creation should fail and stop trying. - assertThat(dbClient.pool.getNumberOfSessionsInPool()).isEqualTo(0); - assertThat(dbClient.pool.getNumberOfSessionsBeingCreated()).isEqualTo(0); - // Force a maintainer run. This should schedule new session creation. - dbClient.pool.poolMaintainer.maintainPool(); - // Wait until the replenish has finished. - watch.reset().start(); - while (watch.elapsed(TimeUnit.SECONDS) < 5 - && dbClient.pool.getNumberOfSessionsBeingCreated() > 0) { - //noinspection BusyWait - Thread.sleep(1L); - } - // All session creation from replenishPool should fail and stop trying. - assertThat(dbClient.pool.getNumberOfSessionsInPool()).isEqualTo(0); - assertThat(dbClient.pool.getNumberOfSessionsBeingCreated()).isEqualTo(0); - } - mockSpanner.reset(); - mockSpanner.removeAllExecutionTimes(); - } - } - /** * Test showing that when a database is deleted while it is in use by a database client and then * re-created with the same name, will continue to return {@link DatabaseNotFoundException}s until @@ -3269,7 +2111,7 @@ public void testDatabaseOrInstanceIsDeletedAndThenRecreated() throws Exception { // Wait until all sessions have been created and prepared. Stopwatch watch = Stopwatch.createStarted(); while (watch.elapsed(TimeUnit.SECONDS) < 5 - && (dbClient.pool.getNumberOfSessionsBeingCreated() > 0)) { + && (dbClient.multiplexedSessionDatabaseClient.getCurrentSessionReference() == null)) { //noinspection BusyWait Thread.sleep(1L); } @@ -3342,8 +2184,6 @@ public void testGetInvalidatedClientMultipleTimes() { for (StatusRuntimeException exception : exceptions) { mockSpanner.setCreateSessionExecutionTime( SimulatedExecutionTime.ofStickyException(exception)); - mockSpanner.setBatchCreateSessionsExecutionTime( - SimulatedExecutionTime.ofStickyException(exception)); try (Spanner spanner = SpannerOptions.newBuilder() .setProjectId(TEST_PROJECT) @@ -3358,22 +2198,15 @@ public void testGetInvalidatedClientMultipleTimes() { spanner.getDatabaseClient( DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); for (int useClient = 0; useClient < 2; useClient++) { - // Using the same client multiple times should continue to return the same - // ResourceNotFoundException, even though the session pool has been invalidated. + // The multiplexed session client tries to create a new session at every attempt. assertThrows( ResourceNotFoundException.class, () -> dbClient.singleUse().executeQuery(SELECT1).next()); - if (spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSession()) { - // We should only receive 1 CreateSession request. The query should never be executed, - // as the session creation fails before it gets to executing a query. - assertEquals(1, mockSpanner.countRequestsOfType(CreateSessionRequest.class)); - assertEquals(0, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)); - } else { - // The server should only receive one BatchCreateSessions request for each run as we - // have set MinSessions=0. - assertThat(mockSpanner.getRequests()).hasSize(run + 1); - assertThat(dbClient.pool.isValid()).isFalse(); - } + // We should only receive 1 CreateSession request per attempt. + // The query should never be executed, as the session creation fails before it gets to + // executing a query. + assertEquals(run + 1, mockSpanner.countRequestsOfType(CreateSessionRequest.class)); + assertEquals(0, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)); } } } @@ -3391,27 +2224,19 @@ public void testAllowNestedTransactions() throws InterruptedException { final int minSessions = spanner.getOptions().getSessionPoolOptions().getMinSessions(); Stopwatch watch = Stopwatch.createStarted(); while (watch.elapsed(TimeUnit.SECONDS) < 5 - && client.pool.getNumberOfSessionsInPool() < minSessions) { + && client.multiplexedSessionDatabaseClient.getCurrentSessionReference() == null) { //noinspection BusyWait Thread.sleep(1L); } - assertThat(client.pool.getNumberOfSessionsInPool()).isEqualTo(minSessions); - int expectedMinSessions = - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW() - ? minSessions - : minSessions - 1; Long res = client .readWriteTransaction() .allowNestedTransaction() .run( transaction -> { - assertThat(client.pool.getNumberOfSessionsInPool()) - .isEqualTo(expectedMinSessions); return transaction.executeUpdate(UPDATE_STATEMENT); }); assertThat(res).isEqualTo(UPDATE_COUNT); - assertThat(client.pool.getNumberOfSessionsInPool()).isEqualTo(minSessions); } @Test @@ -3427,37 +2252,22 @@ public void testNestedTransactionsUsingTwoDatabases() throws InterruptedExceptio final int minSessions = spanner.getOptions().getSessionPoolOptions().getMinSessions(); Stopwatch watch = Stopwatch.createStarted(); while (watch.elapsed(TimeUnit.SECONDS) < 5 - && (client1.pool.getNumberOfSessionsInPool() < minSessions - || client2.pool.getNumberOfSessionsInPool() < minSessions)) { + && (client1.multiplexedSessionDatabaseClient.getCurrentSessionReference() == null + || client2.multiplexedSessionDatabaseClient.getCurrentSessionReference() == null)) { //noinspection BusyWait Thread.sleep(1L); } - assertThat(client1.pool.getNumberOfSessionsInPool()).isEqualTo(minSessions); - assertThat(client2.pool.getNumberOfSessionsInPool()).isEqualTo(minSessions); - // When read-write transaction uses multiplexed sessions, then sessions are not checked out from - // the session pool. - int expectedMinSessions = isMultiplexedSessionsEnabledForRW() ? minSessions : minSessions - 1; Long res = client1 .readWriteTransaction() .allowNestedTransaction() .run( transaction -> { - // Client1 should have 1 session checked out. - // Client2 should have 0 sessions checked out. - assertThat(client1.pool.getNumberOfSessionsInPool()) - .isEqualTo(expectedMinSessions); - assertThat(client2.pool.getNumberOfSessionsInPool()).isEqualTo(minSessions); Long add = client2 .readWriteTransaction() .run( transaction1 -> { - // Both clients should now have 1 session checked out. - assertThat(client1.pool.getNumberOfSessionsInPool()) - .isEqualTo(expectedMinSessions); - assertThat(client2.pool.getNumberOfSessionsInPool()) - .isEqualTo(expectedMinSessions); try (ResultSet rs = transaction1.executeQuery(SELECT1)) { if (rs.next()) { return rs.getLong(0); @@ -3474,9 +2284,6 @@ public void testNestedTransactionsUsingTwoDatabases() throws InterruptedExceptio } }); assertThat(res).isEqualTo(2L); - // All sessions should now be checked back in to the pools. - assertThat(client1.pool.getNumberOfSessionsInPool()).isEqualTo(minSessions); - assertThat(client2.pool.getNumberOfSessionsInPool()).isEqualTo(minSessions); } @Test @@ -3602,16 +2409,8 @@ public void testBackendPartitionQueryOptions() { // statistics package and directed read options. List requests = mockSpanner.getRequests(); assert requests.size() >= 2 : "required to have at least 2 requests"; - if (isMultiplexedSessionsEnabled()) { - assertThat(requests.get(requests.size() - 1)).isInstanceOf(ExecuteSqlRequest.class); - } else { - assertThat(requests.get(requests.size() - 1)).isInstanceOf(DeleteSessionRequest.class); - assertThat(requests.get(requests.size() - 2)).isInstanceOf(ExecuteSqlRequest.class); - } - ExecuteSqlRequest executeSqlRequest = - (ExecuteSqlRequest) - requests.get( - isMultiplexedSessionsEnabled() ? requests.size() - 1 : requests.size() - 2); + assertThat(requests.get(requests.size() - 1)).isInstanceOf(ExecuteSqlRequest.class); + ExecuteSqlRequest executeSqlRequest = (ExecuteSqlRequest) requests.get(requests.size() - 1); assertThat(executeSqlRequest.getQueryOptions()).isNotNull(); assertThat(executeSqlRequest.getQueryOptions().getOptimizerVersion()).isEqualTo("1"); assertThat(executeSqlRequest.getQueryOptions().getOptimizerStatisticsPackage()) @@ -3659,16 +2458,8 @@ public void testBackendPartitionQueryOptions() { // statistics package and directed read options. List requests = mockSpanner.getRequests(); assert requests.size() >= 2 : "required to have at least 2 requests"; - if (isMultiplexedSessionsEnabled()) { - assertThat(requests.get(requests.size() - 1)).isInstanceOf(ExecuteSqlRequest.class); - } else { - assertThat(requests.get(requests.size() - 1)).isInstanceOf(DeleteSessionRequest.class); - assertThat(requests.get(requests.size() - 2)).isInstanceOf(ExecuteSqlRequest.class); - } - ExecuteSqlRequest executeSqlRequest = - (ExecuteSqlRequest) - requests.get( - isMultiplexedSessionsEnabled() ? requests.size() - 1 : requests.size() - 2); + assertThat(requests.get(requests.size() - 1)).isInstanceOf(ExecuteSqlRequest.class); + ExecuteSqlRequest executeSqlRequest = (ExecuteSqlRequest) requests.get(requests.size() - 1); assertThat(executeSqlRequest.getQueryOptions()).isNotNull(); assertThat(executeSqlRequest.getQueryOptions().getOptimizerVersion()).isEqualTo("1"); assertThat(executeSqlRequest.getQueryOptions().getOptimizerStatisticsPackage()) @@ -3712,16 +2503,8 @@ public void testBackendPartitionReadOptions() { // statistics package and directed read options. List requests = mockSpanner.getRequests(); assert requests.size() >= 2 : "required to have at least 2 requests"; - if (isMultiplexedSessionsEnabled()) { - assertThat(requests.get(requests.size() - 1)).isInstanceOf(ReadRequest.class); - } else { - assertThat(requests.get(requests.size() - 1)).isInstanceOf(DeleteSessionRequest.class); - assertThat(requests.get(requests.size() - 2)).isInstanceOf(ReadRequest.class); - } - ReadRequest readRequest = - (ReadRequest) - requests.get( - isMultiplexedSessionsEnabled() ? requests.size() - 1 : requests.size() - 2); + assertThat(requests.get(requests.size() - 1)).isInstanceOf(ReadRequest.class); + ReadRequest readRequest = (ReadRequest) requests.get(requests.size() - 1); assertThat(readRequest.getDirectedReadOptions()).isEqualTo(DIRECTED_READ_OPTIONS1); } } @@ -3762,16 +2545,8 @@ public void testBackendPartitionReadOptions() { // statistics package and directed read options. List requests = mockSpanner.getRequests(); assert requests.size() >= 2 : "required to have at least 2 requests"; - if (isMultiplexedSessionsEnabled()) { - assertThat(requests.get(requests.size() - 1)).isInstanceOf(ReadRequest.class); - } else { - assertThat(requests.get(requests.size() - 1)).isInstanceOf(DeleteSessionRequest.class); - assertThat(requests.get(requests.size() - 2)).isInstanceOf(ReadRequest.class); - } - ReadRequest readRequest = - (ReadRequest) - requests.get( - isMultiplexedSessionsEnabled() ? requests.size() - 1 : requests.size() - 2); + assertThat(requests.get(requests.size() - 1)).isInstanceOf(ReadRequest.class); + ReadRequest readRequest = (ReadRequest) requests.get(requests.size() - 1); assertThat(readRequest.getDirectedReadOptions()).isEqualTo(DIRECTED_READ_OPTIONS2); } } @@ -3988,34 +2763,8 @@ public void testSpecificTimeout() { }); } - @Test - public void testBatchCreateSessionsFailure_shouldNotPropagateToCloseMethod() { - assumeFalse( - "BatchCreateSessions RPC is not invoked for multiplexed sessions", - isMultiplexedSessionsEnabled()); - try { - // Simulate session creation failures on the backend. - mockSpanner.setBatchCreateSessionsExecutionTime( - SimulatedExecutionTime.ofStickyException( - Status.FAILED_PRECONDITION.asRuntimeException())); - DatabaseClient client = - spannerWithEmptySessionPool.getDatabaseClient( - DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); - // This will not cause any failure as getting a session from the pool is guaranteed to be - // non-blocking, and any exceptions will be delayed until actual query execution. - try (ResultSet rs = client.singleUse().executeQuery(SELECT1)) { - SpannerException e = assertThrows(SpannerException.class, rs::next); - assertThat(e.getErrorCode()).isEqualTo(ErrorCode.FAILED_PRECONDITION); - } - } finally { - mockSpanner.setBatchCreateSessionsExecutionTime(SimulatedExecutionTime.none()); - } - } - @Test public void testCreateSessionsFailure_shouldNotPropagateToCloseMethod() { - assumeTrue( - "CreateSessions is not invoked for regular sessions", isMultiplexedSessionsEnabled()); try { // Simulate session creation failures on the backend. mockSpanner.setCreateSessionExecutionTime( @@ -4036,61 +2785,6 @@ public void testCreateSessionsFailure_shouldNotPropagateToCloseMethod() { } } - @Test - public void testReadWriteTransaction_usesOptions() { - SessionPool pool = mock(SessionPool.class); - PooledSessionFuture session = mock(PooledSessionFuture.class); - when(pool.getSession()).thenReturn(session); - TransactionOption option = mock(TransactionOption.class); - - TraceWrapper traceWrapper = - new TraceWrapper(Tracing.getTracer(), OpenTelemetry.noop().getTracer(""), false); - - DatabaseClientImpl client = new DatabaseClientImpl(pool, traceWrapper); - client.readWriteTransaction(option); - - verify(session).readWriteTransaction(option); - } - - @Test - public void testTransactionManager_usesOptions() { - SessionPool pool = mock(SessionPool.class); - PooledSessionFuture session = mock(PooledSessionFuture.class); - when(pool.getSession()).thenReturn(session); - TransactionOption option = mock(TransactionOption.class); - - DatabaseClientImpl client = new DatabaseClientImpl(pool, mock(TraceWrapper.class)); - try (TransactionManager ignore = client.transactionManager(option)) { - verify(session).transactionManager(option); - } - } - - @Test - public void testRunAsync_usesOptions() { - SessionPool pool = mock(SessionPool.class); - PooledSessionFuture session = mock(PooledSessionFuture.class); - when(pool.getSession()).thenReturn(session); - TransactionOption option = mock(TransactionOption.class); - - DatabaseClientImpl client = new DatabaseClientImpl(pool, mock(TraceWrapper.class)); - client.runAsync(option); - - verify(session).runAsync(option); - } - - @Test - public void testTransactionManagerAsync_usesOptions() { - SessionPool pool = mock(SessionPool.class); - PooledSessionFuture session = mock(PooledSessionFuture.class); - when(pool.getSession()).thenReturn(session); - TransactionOption option = mock(TransactionOption.class); - - DatabaseClientImpl client = new DatabaseClientImpl(pool, mock(TraceWrapper.class)); - try (AsyncTransactionManager ignore = client.transactionManagerAsync(option)) { - verify(session).transactionManagerAsync(option); - } - } - @Test public void testExecuteQueryWithPriority() { DatabaseClient client = @@ -4355,73 +3049,6 @@ public void testAsyncTransactionManagerCommitWithMaxCommitDelay() { request.getMaxCommitDelay()); } - @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(); - } - @Test public void transactionContextFailsIfUsedMultipleTimes() { DatabaseClient client = @@ -5338,97 +3965,6 @@ public void testSelectHasXGoogRequestIdHeader() { } } - @Test - public void testSessionPoolExhaustedError_containsStackTraces() { - assumeFalse( - "Session pool tests are skipped for multiplexed sessions", - isMultiplexedSessionsEnabledForRW()); - try (Spanner spanner = - SpannerOptions.newBuilder() - .setProjectId(TEST_PROJECT) - .setChannelProvider(channelProvider) - .setCredentials(NoCredentials.getInstance()) - .setSessionPoolOption( - SessionPoolOptions.newBuilder() - .setFailIfPoolExhausted() - .setMinSessions(2) - .setMaxSessions(4) - .setWaitForMinSessionsDuration(Duration.ofSeconds(10L)) - .build()) - .build() - .getService()) { - DatabaseClient client = - spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); - List transactions = new ArrayList<>(); - // Deliberately leak 4 sessions. - for (int i = 0; i < 4; i++) { - // Get a transaction manager without doing anything with it. This will reserve a session - // from the pool, but not increase the number of sessions marked as in use. - transactions.add(client.transactionManager()); - } - // Trying to get yet another transaction will fail. - // NOTE: This fails directly, because we have set the setFailIfPoolExhausted() option. - SpannerException spannerException = - assertThrows(SpannerException.class, client::transactionManager); - assertEquals(ErrorCode.RESOURCE_EXHAUSTED, spannerException.getErrorCode()); - assertTrue( - spannerException.getMessage(), - spannerException.getMessage().contains("There are currently 4 sessions checked out:")); - assertTrue( - spannerException.getMessage(), - spannerException.getMessage().contains("Session was checked out from the pool at")); - - SessionPool pool = ((DatabaseClientImpl) client).pool; - // Verify that there are no sessions in the pool. - assertEquals(0, pool.getNumberOfSessionsInPool()); - // Verify that the sessions have not (yet) been marked as in use. - assertEquals(0, pool.getNumberOfSessionsInUse()); - assertEquals(0, pool.getMaxSessionsInUse()); - // Verify that we have 4 sessions in the pool. - assertEquals(4, pool.getTotalSessionsPlusNumSessionsBeingCreated()); - - // Release the sessions back into the pool. - for (TransactionManager transaction : transactions) { - transaction.close(); - } - // Wait up to 100 milliseconds for the sessions to actually all be in the pool, as there are - // two possible ways that the session pool handles the above: - // 1. The pool starts to create 4 sessions. - // 2. It then hands out whatever session has been created to one of the waiters. - // 3. The waiting process then executes its transaction, and when finished, the session is - // given to any other process waiting at that moment. - // The above means that although there will always be 4 sessions created, it could in theory - // be that not all of them are used, as it could be that a transaction finishes before the - // creation of session 2, 3, or 4 finished, and then the existing session is re-used. - Stopwatch watch = Stopwatch.createStarted(); - while (pool.getNumberOfSessionsInPool() < 4 && watch.elapsed(TimeUnit.MILLISECONDS) < 100) { - Thread.yield(); - } - // Closing the transactions should return the sessions to the pool. - assertEquals(4, pool.getNumberOfSessionsInPool()); - - DatabaseClientImpl dbClient = (DatabaseClientImpl) client; - int channelId = 0; - try (Session session = dbClient.getSession()) { - channelId = ((PooledSessionFuture) session).getChannel(); - } - int dbId = dbClient.dbId; - XGoogSpannerRequestIdTest.MethodAndRequestId[] wantStreamingValues = {}; - - xGoogReqIdInterceptor.checkExpectedStreamingXGoogRequestIds(wantStreamingValues); - long NON_DETERMINISTIC = XGoogSpannerRequestIdTest.NON_DETERMINISTIC; - - XGoogSpannerRequestIdTest.MethodAndRequestId[] wantUnaryValues = { - XGoogSpannerRequestIdTest.ofMethodAndRequestId( - "google.spanner.v1.Spanner/CreateSession", - new XGoogSpannerRequestId(NON_DETERMINISTIC, 0, 1, 1)), - }; - if (false) { // TODO(@odeke-em): enable in next PRs. - xGoogReqIdInterceptor.checkExpectedUnaryXGoogRequestIdsAsSuffixes(wantUnaryValues); - } - } - } - static void assertAsString(String expected, ResultSet resultSet, int col) { assertEquals(expected, resultSet.getValue(col).getAsString()); assertEquals(ImmutableList.of(expected), resultSet.getValue(col).getAsStringList()); @@ -5737,88 +4273,4 @@ private ListValue getRows(Dialect dialect) { return valuesBuilder.build(); } - - private boolean isMultiplexedSessionsEnabled() { - if (spanner.getOptions() == null || spanner.getOptions().getSessionPoolOptions() == null) { - return false; - } - return spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSession(); - } - - private boolean isMultiplexedSessionsEnabledForRW() { - if (spanner.getOptions() == null || spanner.getOptions().getSessionPoolOptions() == null) { - return false; - } - return spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW(); - } - - @Test - public void testdbIdFromClientId() { - SessionPool pool = mock(SessionPool.class); - PooledSessionFuture session = mock(PooledSessionFuture.class); - when(pool.getSession()).thenReturn(session); - TransactionOption option = mock(TransactionOption.class); - DatabaseClientImpl client = new DatabaseClientImpl(pool, mock(TraceWrapper.class)); - - for (int i = 0; i < 10; i++) { - String dbId = String.format("%d", i); - int id = client.dbIdFromClientId(dbId); - assertEquals(id, i + 2); // There was already 1 dbId after new DatabaseClientImpl. - } - } - - @Test - public void testrunWithSessionRetry_withRequestId() { - // Tests that DatabaseClientImpl.runWithSessionRetry correctly returns a XGoogSpannerRequestId - // and correctly increases its nthRequest ordinal number and that attempts stay at 1, given - // a fresh session returned on SessionNotFoundException. - SessionPool pool = mock(SessionPool.class); - PooledSessionFuture sessionFut = mock(PooledSessionFuture.class); - when(pool.getSession()).thenReturn(sessionFut); - SessionPool.PooledSession pooledSession = mock(SessionPool.PooledSession.class); - when(sessionFut.get()).thenReturn(pooledSession); - SessionPool.PooledSessionReplacementHandler sessionReplacementHandler = - mock(SessionPool.PooledSessionReplacementHandler.class); - when(pool.getPooledSessionReplacementHandler()).thenReturn(sessionReplacementHandler); - when(sessionReplacementHandler.replaceSession(any(), any())).thenReturn(sessionFut); - DatabaseClientImpl client = new DatabaseClientImpl(pool, mock(TraceWrapper.class)); - - // 1. Run with no fail runs a single attempt. - final AtomicInteger nCalls = new AtomicInteger(0); - client.runWithSessionRetry( - (session, reqId) -> { - assertEquals(reqId.getAttempt(), 1); - nCalls.incrementAndGet(); - return 1; - }); - assertEquals(nCalls.get(), 1); - - // Reset the call counter. - nCalls.set(0); - - // 2. Run with SessionNotFoundException and ensure that a fresh requestId is returned each time. - SessionNotFoundException excSessionNotFound = - SpannerExceptionFactoryTest.newSessionNotFoundException( - "projects/p/instances/i/databases/d/sessions/s"); - - final AtomicLong priorNthRequest = new AtomicLong(client.getNthRequest()); - client.runWithSessionRetry( - (session, reqId) -> { - // Monotonically increasing priorNthRequest. - assertEquals(reqId.getNthRequest() - priorNthRequest.get(), 1); - priorNthRequest.set(reqId.getNthRequest()); - - // Attempts stay at 1 since with a SessionNotFound exception, - // a fresh requestId is generated. - assertEquals(reqId.getAttempt(), 1); - - if (nCalls.addAndGet(1) < 4) { - throw excSessionNotFound; - } - - return 1; - }); - - assertEquals(nCalls.get(), 4); - } } diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DefaultBenchmark.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DefaultBenchmark.java index 35712cd5b4e..28580f53365 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DefaultBenchmark.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DefaultBenchmark.java @@ -17,7 +17,6 @@ package com.google.cloud.spanner; import static com.google.cloud.spanner.BenchmarkingUtilityScripts.collectResults; -import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertNotNull; @@ -114,9 +113,6 @@ public void teardown() throws Exception { @Benchmark public void burstQueries(final BenchmarkState server) throws Exception { final DatabaseClientImpl client = server.client; - SessionPool pool = client.pool; - assertThat(pool.totalSessions()) - .isEqualTo(server.spanner.getOptions().getSessionPoolOptions().getMinSessions()); ListeningScheduledExecutorService service = MoreExecutors.listeningDecorator(Executors.newScheduledThreadPool(PARALLEL_THREADS)); @@ -132,9 +128,6 @@ public void burstQueries(final BenchmarkState server) throws Exception { @Benchmark public void burstQueriesAndWrites(final BenchmarkState server) throws Exception { final DatabaseClientImpl client = server.client; - SessionPool pool = client.pool; - assertThat(pool.totalSessions()) - .isEqualTo(server.spanner.getOptions().getSessionPoolOptions().getMinSessions()); ListeningScheduledExecutorService service = MoreExecutors.listeningDecorator(Executors.newScheduledThreadPool(PARALLEL_THREADS)); @@ -154,9 +147,6 @@ public void burstQueriesAndWrites(final BenchmarkState server) throws Exception @Benchmark public void burstUpdates(final BenchmarkState server) throws Exception { final DatabaseClientImpl client = server.client; - SessionPool pool = client.pool; - assertThat(pool.totalSessions()) - .isEqualTo(server.spanner.getOptions().getSessionPoolOptions().getMinSessions()); ListeningScheduledExecutorService service = MoreExecutors.listeningDecorator(Executors.newScheduledThreadPool(PARALLEL_THREADS)); diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ITSessionPoolIntegrationTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ITSessionPoolIntegrationTest.java deleted file mode 100644 index ab6dfea4d61..00000000000 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ITSessionPoolIntegrationTest.java +++ /dev/null @@ -1,176 +0,0 @@ -/* - * Copyright 2017 Google LLC - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package com.google.cloud.spanner; - -import static com.google.cloud.spanner.testing.ExperimentalHostHelper.isExperimentalHost; -import static com.google.common.truth.Truth.assertThat; -import static org.junit.Assume.assumeFalse; - -import com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory; -import com.google.cloud.spanner.SessionPool.PooledSessionFuture; -import io.opencensus.trace.Tracing; -import io.opentelemetry.api.OpenTelemetry; -import java.util.ArrayList; -import java.util.List; -import java.util.concurrent.CountDownLatch; -import java.util.concurrent.ScheduledExecutorService; -import java.util.concurrent.ScheduledThreadPoolExecutor; -import java.util.concurrent.TimeUnit; -import org.junit.Before; -import org.junit.BeforeClass; -import org.junit.ClassRule; -import org.junit.Test; -import org.junit.experimental.categories.Category; -import org.junit.runner.RunWith; -import org.junit.runners.JUnit4; - -/** - * Integration tests for read and query. - * - *

See also {@code it/WriteIntegrationTest}, which provides coverage of writing and reading back - * all Cloud Spanner types. - */ -@Category(SerialIntegrationTest.class) -@RunWith(JUnit4.class) -public class ITSessionPoolIntegrationTest { - @ClassRule public static IntegrationTestEnv env = new IntegrationTestEnv(); - private static final String TABLE_NAME = "TestTable"; - - private static Database db; - private SessionPool pool; - - @BeforeClass - public static void setUpDatabase() { - assumeFalse("Only Multiplexed Sessions are supported on this host", isExperimentalHost()); - db = - env.getTestHelper() - .createTestDatabase( - "CREATE TABLE TestTable (" - + " Key STRING(MAX) NOT NULL," - + " StringValue STRING(MAX)," - + ") PRIMARY KEY (Key)", - "CREATE INDEX TestTableByValue ON TestTable(StringValue)"); - - // Includes k0..k14. Note that strings k{10,14} sort between k1 and k2. - List mutations = new ArrayList<>(); - for (int i = 0; i < 15; ++i) { - mutations.add( - Mutation.newInsertOrUpdateBuilder(TABLE_NAME) - .set("Key") - .to("k" + i) - .set("StringValue") - .to("v" + i) - .build()); - } - env.getTestHelper().getDatabaseClient(db).write(mutations); - } - - @Before - public void setUp() { - SessionPoolOptions options = - SessionPoolOptions.newBuilder().setMinSessions(1).setMaxSessions(2).build(); - pool = - SessionPool.createPool( - options, - new ExecutorFactory() { - - @Override - public void release(ScheduledExecutorService executor) { - executor.shutdown(); - } - - @Override - public ScheduledExecutorService get() { - return new ScheduledThreadPoolExecutor(2); - } - }, - ((SpannerImpl) env.getTestHelper().getClient()).getSessionClient(db.getId()), - new TraceWrapper(Tracing.getTracer(), OpenTelemetry.noop().getTracer(""), false), - OpenTelemetry.noop()); - } - - @Test - public void sessionCreation() { - try (PooledSessionFuture session = pool.getSession()) { - assertThat(session.get()).isNotNull(); - } - - try (PooledSessionFuture session = pool.getSession(); - PooledSessionFuture session2 = pool.getSession()) { - assertThat(session.get()).isNotNull(); - assertThat(session2.get()).isNotNull(); - } - } - - @Test - public void poolExhaustion() throws Exception { - Session session1 = pool.getSession().get(); - Session session2 = pool.getSession().get(); - final CountDownLatch latch = new CountDownLatch(1); - new Thread( - () -> { - try (Session session3 = pool.getSession().get()) { - latch.countDown(); - } - }) - .start(); - assertThat(latch.await(5, TimeUnit.SECONDS)).isFalse(); - session1.close(); - session2.close(); - latch.await(); - } - - @Test - public void multipleWaiters() throws Exception { - Session session1 = pool.getSession().get(); - Session session2 = pool.getSession().get(); - int numSessions = 5; - final CountDownLatch latch = new CountDownLatch(numSessions); - for (int i = 0; i < numSessions; i++) { - new Thread( - () -> { - try (Session session = pool.getSession().get()) { - latch.countDown(); - } - }) - .start(); - } - session1.close(); - session2.close(); - // Everyone should get session pretty quickly. - assertThat(latch.await(1, TimeUnit.SECONDS)).isTrue(); - } - - @Test - public void closeQuicklyDoesNotBlockIndefinitely() throws Exception { - pool.closeAsync(new SpannerImpl.ClosedException()).get(); - } - - @Test - public void closeAfterInitialCreateDoesNotBlockIndefinitely() throws Exception { - pool.getSession().close(); - pool.closeAsync(new SpannerImpl.ClosedException()).get(); - } - - @Test - public void closeWhenSessionsActiveFinishes() throws Exception { - pool.getSession().get(); - // This will log a warning that a session has been leaked, as the session that we retrieved in - // the previous statement was never returned to the pool. - pool.closeAsync(new SpannerImpl.ClosedException()).get(); - } -} diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/InlineBeginBenchmark.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/InlineBeginBenchmark.java index 1448ebbc96a..c3063f4d6c5 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/InlineBeginBenchmark.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/InlineBeginBenchmark.java @@ -16,8 +16,6 @@ package com.google.cloud.spanner; -import static com.google.common.truth.Truth.assertThat; - import com.google.api.gax.rpc.TransportChannelProvider; import com.google.cloud.NoCredentials; import com.google.common.base.Stopwatch; @@ -103,8 +101,7 @@ public void setup() throws Exception { spanner.getDatabaseClient(DatabaseId.of(options.getProjectId(), instance, database)); Stopwatch watch = Stopwatch.createStarted(); // Wait until the session pool has initialized. - while (client.pool.getNumberOfSessionsInPool() - < spanner.getOptions().getSessionPoolOptions().getMinSessions()) { + while (client.multiplexedSessionDatabaseClient.getCurrentSessionReference() == null) { Thread.sleep(1L); if (watch.elapsed(TimeUnit.SECONDS) > 10L) { break; @@ -143,9 +140,6 @@ public void teardown() throws Exception { public void burstRead(final BenchmarkState server) throws Exception { int totalQueries = server.spanner.getOptions().getSessionPoolOptions().getMaxSessions() * 8; int parallelThreads = server.spanner.getOptions().getSessionPoolOptions().getMaxSessions() * 2; - SessionPool pool = server.client.pool; - assertThat(pool.totalSessions()) - .isEqualTo(server.spanner.getOptions().getSessionPoolOptions().getMinSessions()); ListeningScheduledExecutorService service = MoreExecutors.listeningDecorator(Executors.newScheduledThreadPool(parallelThreads)); @@ -173,9 +167,6 @@ public void burstRead(final BenchmarkState server) throws Exception { public void burstWrite(final BenchmarkState server) throws Exception { int totalWrites = server.spanner.getOptions().getSessionPoolOptions().getMaxSessions() * 8; int parallelThreads = server.spanner.getOptions().getSessionPoolOptions().getMaxSessions() * 2; - SessionPool pool = server.client.pool; - assertThat(pool.totalSessions()) - .isEqualTo(server.spanner.getOptions().getSessionPoolOptions().getMinSessions()); ListeningScheduledExecutorService service = MoreExecutors.listeningDecorator(Executors.newScheduledThreadPool(parallelThreads)); @@ -201,9 +192,6 @@ public void burstReadAndWrite(final BenchmarkState server) throws Exception { int totalWrites = server.spanner.getOptions().getSessionPoolOptions().getMaxSessions() * 4; int totalReads = server.spanner.getOptions().getSessionPoolOptions().getMaxSessions() * 4; int parallelThreads = server.spanner.getOptions().getSessionPoolOptions().getMaxSessions() * 2; - SessionPool pool = server.client.pool; - assertThat(pool.totalSessions()) - .isEqualTo(server.spanner.getOptions().getSessionPoolOptions().getMinSessions()); ListeningScheduledExecutorService service = MoreExecutors.listeningDecorator(Executors.newScheduledThreadPool(parallelThreads)); diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/IntegrationTestWithClosedSessionsEnv.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/IntegrationTestWithClosedSessionsEnv.java deleted file mode 100644 index f852fc2903f..00000000000 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/IntegrationTestWithClosedSessionsEnv.java +++ /dev/null @@ -1,138 +0,0 @@ -/* - * Copyright 2019 Google LLC - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package com.google.cloud.spanner; - -import com.google.cloud.spanner.SessionPool.PooledSession; -import com.google.cloud.spanner.SessionPool.PooledSessionFuture; -import com.google.cloud.spanner.SessionPool.SessionFutureWrapper; -import com.google.cloud.spanner.testing.RemoteSpannerHelper; -import io.opentelemetry.api.common.Attributes; - -/** - * Subclass of {@link IntegrationTestEnv} that allows the user to specify when the underlying - * session of a {@link PooledSession} should be closed. This can be used to ensure that the - * recreation of sessions that have been invalidated by the server works. - */ -public class IntegrationTestWithClosedSessionsEnv extends IntegrationTestEnv { - private static class RemoteSpannerHelperWithClosedSessions extends RemoteSpannerHelper { - private RemoteSpannerHelperWithClosedSessions( - SpannerOptions options, InstanceId instanceId, Spanner client) { - super(options, instanceId, client); - } - } - - @Override - RemoteSpannerHelper createTestHelper(SpannerOptions options, InstanceId instanceId) { - SpannerWithClosedSessionsImpl spanner = new SpannerWithClosedSessionsImpl(options); - return new RemoteSpannerHelperWithClosedSessions(options, instanceId, spanner); - } - - private static class SpannerWithClosedSessionsImpl extends SpannerImpl { - SpannerWithClosedSessionsImpl(SpannerOptions options) { - super(options); - } - - @Override - DatabaseClientImpl createDatabaseClient( - String clientId, - SessionPool pool, - boolean useMultiplexedSessionBlindWriteIgnore, - MultiplexedSessionDatabaseClient ignore, - boolean useMultiplexedSessionPartitionedOpsIgnore, - boolean useMultiplexedSessionForRWIgnore, - Attributes attributes) { - return new DatabaseClientWithClosedSessionImpl(clientId, pool, tracer); - } - } - - /** - * {@link DatabaseClient} that allows the user to specify when an underlying session of a {@link - * PooledSession} should be closed. - */ - public static class DatabaseClientWithClosedSessionImpl extends DatabaseClientImpl { - private boolean invalidateNextSession = false; - private boolean allowReplacing = true; - - DatabaseClientWithClosedSessionImpl(String clientId, SessionPool pool, TraceWrapper tracer) { - super(clientId, pool, tracer); - } - - /** Invalidate the next session that is checked out from the pool. */ - public void invalidateNextSession() { - invalidateNextSession = true; - } - - /** Sets whether invalidated sessions should be replaced or not. */ - public void setAllowSessionReplacing(boolean allow) { - this.allowReplacing = allow; - } - - @Override - PooledSessionFuture getSession() { - PooledSessionFuture session = super.getSession(); - if (invalidateNextSession) { - session.get().delegate.close(); - session.get().setAllowReplacing(false); - awaitDeleted(session.get().delegate); - session.get().setAllowReplacing(allowReplacing); - invalidateNextSession = false; - } - session.get().setAllowReplacing(allowReplacing); - return session; - } - - @Override - SessionFutureWrapper getMultiplexedSession() { - SessionFutureWrapper session = (SessionFutureWrapper) super.getMultiplexedSession(); - if (invalidateNextSession) { - session.get().get().getDelegate().close(); - session.get().get().setAllowReplacing(false); - awaitDeleted(session.get().get().getDelegate()); - session.get().get().setAllowReplacing(allowReplacing); - invalidateNextSession = false; - } - session.get().get().setAllowReplacing(allowReplacing); - return session; - } - - /** - * Deleting a session server side takes some time. This method checks and waits until the - * session really has been deleted. - */ - private void awaitDeleted(Session session) { - // Wait until the session has actually been deleted. - while (true) { - try (ResultSet rs = session.singleUse().executeQuery(Statement.of("SELECT 1"))) { - while (rs.next()) { - // Do nothing. - } - Thread.sleep(500L); - } catch (SpannerException e) { - if (e.getErrorCode() == ErrorCode.NOT_FOUND - && (e.getMessage().contains("Session not found") - || e.getMessage().contains("Session was concurrently deleted"))) { - break; - } else { - throw e; - } - } catch (InterruptedException e) { - break; - } - } - } - } -} diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/LongRunningSessionsBenchmark.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/LongRunningSessionsBenchmark.java deleted file mode 100644 index 58eb423a5db..00000000000 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/LongRunningSessionsBenchmark.java +++ /dev/null @@ -1,328 +0,0 @@ -/* - * Copyright 2023 Google LLC - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package com.google.cloud.spanner; - -import static com.google.common.truth.Truth.assertThat; - -import com.google.api.gax.rpc.TransportChannelProvider; -import com.google.cloud.NoCredentials; -import com.google.cloud.spanner.MockSpannerServiceImpl.SimulatedExecutionTime; -import com.google.cloud.spanner.SessionPoolOptions.ActionOnInactiveTransaction; -import com.google.cloud.spanner.SessionPoolOptions.InactiveTransactionRemovalOptions; -import com.google.common.util.concurrent.Futures; -import com.google.common.util.concurrent.ListenableFuture; -import com.google.common.util.concurrent.ListeningScheduledExecutorService; -import com.google.common.util.concurrent.MoreExecutors; -import com.google.spanner.v1.BatchCreateSessionsRequest; -import java.time.Duration; -import java.util.ArrayList; -import java.util.List; -import java.util.Random; -import java.util.concurrent.Executors; -import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicInteger; -import org.openjdk.jmh.annotations.AuxCounters; -import org.openjdk.jmh.annotations.Benchmark; -import org.openjdk.jmh.annotations.BenchmarkMode; -import org.openjdk.jmh.annotations.Fork; -import org.openjdk.jmh.annotations.Level; -import org.openjdk.jmh.annotations.Measurement; -import org.openjdk.jmh.annotations.Mode; -import org.openjdk.jmh.annotations.OutputTimeUnit; -import org.openjdk.jmh.annotations.Param; -import org.openjdk.jmh.annotations.Scope; -import org.openjdk.jmh.annotations.Setup; -import org.openjdk.jmh.annotations.State; -import org.openjdk.jmh.annotations.TearDown; -import org.openjdk.jmh.annotations.Warmup; - -/** - * Benchmarks for long-running sessions scenarios. The simulated execution times are based on - * reasonable estimates and are primarily intended to keep the benchmarks comparable with each other - * before and after changes have been made to the pool. The benchmarks are bound to the Maven - * profile `benchmark` and can be executed like this: - * mvn clean test -DskipTests -Pbenchmark -Dbenchmark.name=LongRunningSessionsBenchmark - * - */ -@BenchmarkMode(Mode.AverageTime) -@Fork(value = 1, warmups = 0) -@Measurement(batchSize = 1, iterations = 1, timeUnit = TimeUnit.MILLISECONDS) -@Warmup(batchSize = 0, iterations = 0) -@OutputTimeUnit(TimeUnit.SECONDS) -public class LongRunningSessionsBenchmark { - private static final String TEST_PROJECT = "my-project"; - private static final String TEST_INSTANCE = "my-instance"; - private static final String TEST_DATABASE = "my-database"; - private static final int HOLD_SESSION_TIME = 100; - private static final int LONG_HOLD_SESSION_TIME = 10000; // 10 seconds - private static final int RND_WAIT_TIME_BETWEEN_REQUESTS = 100; - private static final Random RND = new Random(); - - @State(Scope.Thread) - @AuxCounters(org.openjdk.jmh.annotations.AuxCounters.Type.EVENTS) - public static class BenchmarkState { - private StandardBenchmarkMockServer mockServer; - private Spanner spanner; - private DatabaseClientImpl client; - private AtomicInteger longRunningSessions; - - @Param({"100"}) - int minSessions; - - @Param({"400"}) - int maxSessions; - - @Param({"4"}) - int numChannels; - - /** AuxCounter for number of RPCs. */ - public int numBatchCreateSessionsRpcs() { - return mockServer.countRequests(BatchCreateSessionsRequest.class); - } - - /** AuxCounter for number of sessions created. */ - public int sessionsCreated() { - return mockServer.getMockSpanner().numSessionsCreated(); - } - - @Setup(Level.Invocation) - public void setup() throws Exception { - mockServer = new StandardBenchmarkMockServer(); - longRunningSessions = new AtomicInteger(); - TransportChannelProvider channelProvider = mockServer.start(); - - /** - * This ensures that the background thread responsible for cleaning long-running sessions - * executes every 10s. Any transaction for which session has not been used for more than 2s - * will be treated as long-running. - */ - InactiveTransactionRemovalOptions inactiveTransactionRemovalOptions = - InactiveTransactionRemovalOptions.newBuilder() - .setActionOnInactiveTransaction(ActionOnInactiveTransaction.WARN_AND_CLOSE) - .setExecutionFrequency(Duration.ofSeconds(10)) - .setIdleTimeThreshold(Duration.ofSeconds(2)) - .build(); - SpannerOptions options = - SpannerOptions.newBuilder() - .setProjectId(TEST_PROJECT) - .setChannelProvider(channelProvider) - .setNumChannels(numChannels) - .setCredentials(NoCredentials.getInstance()) - .setSessionPoolOption( - SessionPoolOptions.newBuilder() - .setMinSessions(minSessions) - .setMaxSessions(maxSessions) - .setWaitForMinSessionsDuration(Duration.ofSeconds(5)) - .setInactiveTransactionRemovalOptions(inactiveTransactionRemovalOptions) - .build()) - .build(); - - spanner = options.getService(); - client = - (DatabaseClientImpl) - spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); - } - - @TearDown(Level.Invocation) - public void teardown() throws Exception { - spanner.close(); - mockServer.shutdown(); - } - } - - /** - * Measures the time needed to execute a burst of read requests. - * - *

Some read requests will be long-running and will cause session leaks. Such sessions will be - * removed by the session maintenance background task if SessionPool Option - * ActionOnInactiveTransaction is set as WARN_AND_CLOSE. - * - * @param server - * @throws Exception - */ - @Benchmark - public void burstRead(final BenchmarkState server) throws Exception { - int totalQueries = server.maxSessions * 8; - int parallelThreads = server.maxSessions * 2; - final DatabaseClient client = - server.spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); - SessionPool pool = ((DatabaseClientImpl) client).pool; - assertThat(pool.totalSessions()).isEqualTo(server.minSessions); - - ListeningScheduledExecutorService service = - MoreExecutors.listeningDecorator(Executors.newScheduledThreadPool(parallelThreads)); - List> futures = new ArrayList<>(totalQueries); - for (int i = 0; i < totalQueries; i++) { - futures.add( - service.submit( - () -> { - Thread.sleep(RND.nextInt(RND_WAIT_TIME_BETWEEN_REQUESTS)); - try (ResultSet rs = - client.singleUse().executeQuery(StandardBenchmarkMockServer.SELECT1)) { - while (rs.next()) { - // introduce random sleep times to have long-running sessions - randomWait(server); - } - return null; - } - })); - } - // explicitly run the maintenance cycle to clean up any dangling long-running sessions. - pool.poolMaintainer.maintainPool(); - - Futures.allAsList(futures).get(); - service.shutdown(); - assertNumLeakedSessionsRemoved(server, pool); - } - - /** - * Measures the time needed to execute a burst of write requests (PDML). - * - *

Some write requests will be long-running. The test asserts that no sessions are removed by - * the session maintenance background task with SessionPool Option ActionOnInactiveTransaction set - * as WARN_AND_CLOSE. This is because PDML writes are expected to be long-running. - * - * @param server - * @throws Exception - */ - @Benchmark - public void burstWrite(final BenchmarkState server) throws Exception { - int totalWrites = server.maxSessions * 8; - int parallelThreads = server.maxSessions * 2; - final DatabaseClient client = - server.spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); - SessionPool pool = ((DatabaseClientImpl) client).pool; - assertThat(pool.totalSessions()).isEqualTo(server.minSessions); - - ListeningScheduledExecutorService service = - MoreExecutors.listeningDecorator(Executors.newScheduledThreadPool(parallelThreads)); - List> futures = new ArrayList<>(totalWrites); - for (int i = 0; i < totalWrites; i++) { - futures.add( - service.submit( - () -> { - // introduce random sleep times so that some sessions are long-running sessions - randomWaitForMockServer(server); - client.executePartitionedUpdate(StandardBenchmarkMockServer.UPDATE_STATEMENT); - })); - } - // explicitly run the maintenance cycle to clean up any dangling long-running sessions. - pool.poolMaintainer.maintainPool(); - - Futures.allAsList(futures).get(); - service.shutdown(); - assertThat(pool.numLeakedSessionsRemoved()) - .isEqualTo(0); // no sessions should be cleaned up in case of partitioned updates. - } - - /** - * Measures the time needed to execute a burst of read and write requests. - * - *

Some read requests will be long-running and will cause session leaks. Such sessions will be - * removed by the session maintenance background task if SessionPool Option - * ActionOnInactiveTransaction is set as WARN_AND_CLOSE. - * - *

Some write requests will be long-running. The test asserts that no sessions are removed by - * the session maintenance background task with SessionPool Option ActionOnInactiveTransaction set - * as WARN_AND_CLOSE. This is because PDML writes are expected to be long-running. - * - * @param server - * @throws Exception - */ - @Benchmark - public void burstReadAndWrite(final BenchmarkState server) throws Exception { - int totalWrites = server.maxSessions * 4; - int totalReads = server.maxSessions * 4; - int parallelThreads = server.maxSessions * 2; - final DatabaseClient client = - server.spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); - SessionPool pool = ((DatabaseClientImpl) client).pool; - assertThat(pool.totalSessions()).isEqualTo(server.minSessions); - - ListeningScheduledExecutorService service = - MoreExecutors.listeningDecorator(Executors.newScheduledThreadPool(parallelThreads)); - List> futures = new ArrayList<>(totalReads + totalWrites); - for (int i = 0; i < totalWrites; i++) { - futures.add( - service.submit( - () -> { - // introduce random sleep times so that some sessions are long-running sessions - randomWaitForMockServer(server); - client.executePartitionedUpdate(StandardBenchmarkMockServer.UPDATE_STATEMENT); - })); - } - for (int i = 0; i < totalReads; i++) { - futures.add( - service.submit( - () -> { - Thread.sleep(RND.nextInt(RND_WAIT_TIME_BETWEEN_REQUESTS)); - try (ResultSet rs = - client.singleUse().executeQuery(StandardBenchmarkMockServer.SELECT1)) { - while (rs.next()) { - // introduce random sleep times to have long-running sessions - randomWait(server); - } - return null; - } - })); - } - // explicitly run the maintenance cycle to clean up any dangling long-running sessions. - pool.poolMaintainer.maintainPool(); - - Futures.allAsList(futures).get(); - service.shutdown(); - assertNumLeakedSessionsRemoved(server, pool); - } - - private void randomWait(final BenchmarkState server) throws InterruptedException { - if (RND.nextBoolean()) { - server.longRunningSessions.incrementAndGet(); - Thread.sleep(LONG_HOLD_SESSION_TIME); - } else { - Thread.sleep(HOLD_SESSION_TIME); - } - } - - private void randomWaitForMockServer(final BenchmarkState server) { - if (RND.nextBoolean()) { - server.longRunningSessions.incrementAndGet(); - server - .mockServer - .getMockSpanner() - .setExecuteStreamingSqlExecutionTime( - SimulatedExecutionTime.ofMinimumAndRandomTime(LONG_HOLD_SESSION_TIME, 0)); - } else { - server - .mockServer - .getMockSpanner() - .setExecuteStreamingSqlExecutionTime( - SimulatedExecutionTime.ofMinimumAndRandomTime(HOLD_SESSION_TIME, 0)); - } - } - - private void assertNumLeakedSessionsRemoved(final BenchmarkState server, final SessionPool pool) { - final SessionPoolOptions sessionPoolOptions = - server.spanner.getOptions().getSessionPoolOptions(); - assertThat(server.longRunningSessions.get()).isNotEqualTo(0); - if (sessionPoolOptions.warnAndCloseInactiveTransactions() - || sessionPoolOptions.closeInactiveTransactions()) { - assertThat(pool.numLeakedSessionsRemoved()).isGreaterThan(0); - } else if (sessionPoolOptions.warnInactiveTransactions()) { - assertThat(pool.numLeakedSessionsRemoved()).isEqualTo(0); - } - } -} diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MockSpannerServiceImpl.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MockSpannerServiceImpl.java index a47aecdccc4..e6928dbcb37 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MockSpannerServiceImpl.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MockSpannerServiceImpl.java @@ -20,7 +20,6 @@ import com.google.cloud.ByteArray; import com.google.cloud.Date; import com.google.cloud.spanner.AbstractResultSet.LazyByteArray; -import com.google.cloud.spanner.SessionPool.SessionPoolTransactionContext; import com.google.cloud.spanner.TransactionRunnerImpl.TransactionContextImpl; import com.google.common.base.Optional; import com.google.common.base.Preconditions; @@ -301,7 +300,7 @@ public static StatementResult exception(Statement statement, StatusRuntimeExcept /** Creates a result for the query that detects the dialect that is used for the database. */ public static StatementResult detectDialectResult(Dialect resultDialect) { return StatementResult.query( - SessionPool.DETERMINE_DIALECT_STATEMENT, + MultiplexedSessionDatabaseClient.DETERMINE_DIALECT_STATEMENT, ResultSet.newBuilder() .setMetadata( ResultSetMetadata.newBuilder() @@ -581,9 +580,10 @@ private static void checkStreamException( private double abortProbability = 0.0010D; /** - * Flip this switch to true if you want the {@link SessionPool#DETERMINE_DIALECT_STATEMENT} - * statement to be included in the recorded requests on the mock server. It is ignored by default - * to prevent tests that do not expect this request to suddenly start failing. + * Flip this switch to true if you want the {@link + * MultiplexedSessionDatabaseClient#DETERMINE_DIALECT_STATEMENT} statement to be included in the + * recorded requests on the mock server. It is ignored by default to prevent tests that do not + * expect this request to suddenly start failing. */ private boolean includeDetermineDialectStatementInRequests = false; @@ -746,9 +746,10 @@ public void setAbortProbability(double probability) { } /** - * Set this to true if you want the {@link SessionPool#DETERMINE_DIALECT_STATEMENT} statement to - * be included in the recorded requests on the mock server. It is ignored by default to prevent - * tests that do not expect this request to suddenly start failing. + * Set this to true if you want the {@link + * MultiplexedSessionDatabaseClient#DETERMINE_DIALECT_STATEMENT} statement to be included in the + * recorded requests on the mock server. It is ignored by default to prevent tests that do not + * expect this request to suddenly start failing. */ public void setIncludeDetermineDialectStatementInRequests(boolean include) { this.includeDetermineDialectStatementInRequests = include; @@ -760,9 +761,6 @@ public void setIncludeDetermineDialectStatementInRequests(boolean include) { */ public void abortTransaction(TransactionContext transactionContext) { Preconditions.checkNotNull(transactionContext); - if (transactionContext instanceof SessionPoolTransactionContext) { - transactionContext = ((SessionPoolTransactionContext) transactionContext).delegate; - } if (transactionContext instanceof TransactionContextImpl) { TransactionContextImpl impl = (TransactionContextImpl) transactionContext; ByteString id = @@ -1223,7 +1221,9 @@ public void executeBatchDml( public void executeStreamingSql( ExecuteSqlRequest request, StreamObserver responseObserver) { if (includeDetermineDialectStatementInRequests - || !request.getSql().equals(SessionPool.DETERMINE_DIALECT_STATEMENT.getSql())) { + || !request + .getSql() + .equals(MultiplexedSessionDatabaseClient.DETERMINE_DIALECT_STATEMENT.getSql())) { requests.add(request); } Preconditions.checkNotNull(request.getSession()); diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientMockServerTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientMockServerTest.java index 0448656475a..0ad4c6b82bb 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientMockServerTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientMockServerTest.java @@ -42,7 +42,6 @@ import com.google.cloud.spanner.Options.RpcPriority; import com.google.cloud.spanner.TransactionRunnerImpl.TransactionContextImpl; import com.google.cloud.spanner.connection.RandomResultSetGenerator; -import com.google.common.base.Stopwatch; import com.google.common.collect.ImmutableList; import com.google.common.collect.Lists; import com.google.common.util.concurrent.MoreExecutors; @@ -210,134 +209,6 @@ public void testMaintainerMaintainsMultipleClients() { } } - @Test - public void testUnimplementedErrorOnCreation_fallsBackToRegularSessions() { - mockSpanner.setCreateSessionExecutionTime( - SimulatedExecutionTime.ofException( - Status.UNIMPLEMENTED - .withDescription("Multiplexed sessions are not implemented") - .asRuntimeException())); - DatabaseClientImpl client = - (DatabaseClientImpl) spanner.getDatabaseClient(DatabaseId.of("p", "i", "d")); - // Get the current session reference. This will block until the CreateSession RPC has failed. - assertNotNull(client.multiplexedSessionDatabaseClient); - SpannerException spannerException = - assertThrows( - SpannerException.class, - client.multiplexedSessionDatabaseClient::getCurrentSessionReference); - assertEquals(ErrorCode.UNIMPLEMENTED, spannerException.getErrorCode()); - try (ResultSet resultSet = client.singleUse().executeQuery(STATEMENT)) { - //noinspection StatementWithEmptyBody - while (resultSet.next()) { - // ignore - } - } - // Verify that we received one ExecuteSqlRequest, and that it used a regular session. - assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)); - List requests = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class); - - Session session = mockSpanner.getSession(requests.get(0).getSession()); - assertNotNull(session); - assertFalse(session.getMultiplexed()); - - assertNotNull(client.multiplexedSessionDatabaseClient); - assertEquals(0L, client.multiplexedSessionDatabaseClient.getNumSessionsAcquired().get()); - assertEquals(0L, client.multiplexedSessionDatabaseClient.getNumSessionsReleased().get()); - } - - @Test - public void - testUnimplementedErrorOnCreation_firstReceivesError_secondFallsBackToRegularSessions() { - mockSpanner.setCreateSessionExecutionTime( - SimulatedExecutionTime.ofException( - Status.UNIMPLEMENTED - .withDescription("Multiplexed sessions are not implemented") - .asRuntimeException())); - // Freeze the mock server to ensure that the CreateSession RPC does not return an error or any - // other result just yet. - mockSpanner.freeze(); - // Get a database client using multiplexed sessions. The CreateSession RPC will be blocked as - // long as the mock server is frozen. - DatabaseClientImpl client = - (DatabaseClientImpl) spanner.getDatabaseClient(DatabaseId.of("p", "i", "d")); - // Try to execute a query. This is all non-blocking until the call to ResultSet#next(). - try (ResultSet resultSet = client.singleUse().executeQuery(STATEMENT)) { - // Unfreeze the mock server to get the error from the backend. This query will then fail. - mockSpanner.unfreeze(); - SpannerException spannerException = assertThrows(SpannerException.class, resultSet::next); - assertEquals(ErrorCode.UNIMPLEMENTED, spannerException.getErrorCode()); - } - // The next query will fall back to regular sessions and succeed. - try (ResultSet resultSet = client.singleUse().executeQuery(STATEMENT)) { - //noinspection StatementWithEmptyBody - while (resultSet.next()) { - // ignore - } - } - // Verify that we received one ExecuteSqlRequest, and that it used a regular session. - assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)); - List requests = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class); - - Session session = mockSpanner.getSession(requests.get(0).getSession()); - assertNotNull(session); - assertFalse(session.getMultiplexed()); - - assertNotNull(client.multiplexedSessionDatabaseClient); - assertEquals(0L, client.multiplexedSessionDatabaseClient.getNumSessionsAcquired().get()); - assertEquals(0L, client.multiplexedSessionDatabaseClient.getNumSessionsReleased().get()); - } - - @Test - public void testMaintainerInvalidatesMultiplexedSessionClientIfUnimplemented() { - DatabaseClientImpl client = - (DatabaseClientImpl) spanner.getDatabaseClient(DatabaseId.of("p", "i", "d")); - // The first query should succeed. - try (ResultSet resultSet = client.singleUse().executeQuery(STATEMENT)) { - //noinspection StatementWithEmptyBody - while (resultSet.next()) { - // ignore - } - } - // Now ensure that CreateSession returns UNIMPLEMENTED. This error should be recognized by the - // maintainer and invalidate the MultiplexedSessionDatabaseClient. New queries will fall back to - // regular sessions. - mockSpanner.setCreateSessionExecutionTime( - SimulatedExecutionTime.ofException( - Status.UNIMPLEMENTED - .withDescription("Multiplexed sessions are not implemented") - .asRuntimeException())); - // Wait until the client sees that MultiplexedSessions are not supported. - assertNotNull(client.multiplexedSessionDatabaseClient); - Stopwatch stopwatch = Stopwatch.createStarted(); - while (client.multiplexedSessionDatabaseClient.isMultiplexedSessionsSupported() - && stopwatch.elapsed().compareTo(Duration.ofSeconds(5)) < 0) { - Thread.yield(); - } - // Queries should fall back to regular sessions. - try (ResultSet resultSet = client.singleUse().executeQuery(STATEMENT)) { - //noinspection StatementWithEmptyBody - while (resultSet.next()) { - // ignore - } - } - // Verify that we received two ExecuteSqlRequests, and that the first one used a multiplexed - // session, and that the second used a regular session. - assertEquals(2, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)); - List requests = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class); - - Session session1 = mockSpanner.getSession(requests.get(0).getSession()); - assertNotNull(session1); - assertTrue(session1.getMultiplexed()); - - Session session2 = mockSpanner.getSession(requests.get(1).getSession()); - assertNotNull(session2); - assertFalse(session2.getMultiplexed()); - - assertNotNull(client.multiplexedSessionDatabaseClient); - assertEquals(1L, client.multiplexedSessionDatabaseClient.getNumSessionsAcquired().get()); - assertEquals(1L, client.multiplexedSessionDatabaseClient.getNumSessionsReleased().get()); - } - @Test public void testWriteAtLeastOnceAborted() { DatabaseClientImpl client = @@ -1474,381 +1345,6 @@ public void testMutationOnlyUsingTransactionManagerAsyncAbortedDuringBeginTransa spanner.close(); } - // Tests the behavior of the server-side kill switch for read-write multiplexed sessions.. - @Test - public void testInitialBeginTransactionWithRW_receivesUnimplemented_fallsBackToRegularSession() { - mockSpanner.setBeginTransactionExecutionTime( - SimulatedExecutionTime.ofException( - Status.UNIMPLEMENTED - .withDescription( - "Transaction type read_write not supported with multiplexed sessions") - .asRuntimeException())); - DatabaseClientImpl client = - (DatabaseClientImpl) spanner.getDatabaseClient(DatabaseId.of("p", "i", "d")); - - assertNotNull(client.multiplexedSessionDatabaseClient); - - // Wait until the client sees that MultiplexedSessions are not supported for read-write. - // Get the begin transaction reference. This will block until the BeginTransaction RPC with - // read-write has failed. - SpannerException spannerException = - assertThrows( - SpannerException.class, - client.multiplexedSessionDatabaseClient::getReadWriteBeginTransactionReference); - assertEquals(ErrorCode.UNIMPLEMENTED, spannerException.getErrorCode()); - assertTrue(client.multiplexedSessionDatabaseClient.unimplementedForRW.get()); - - // read-write transaction should fallback to regular sessions - client - .readWriteTransaction() - .run( - transaction -> { - try (ResultSet resultSet = transaction.executeQuery(STATEMENT)) { - //noinspection StatementWithEmptyBody - while (resultSet.next()) { - // ignore - } - } - return null; - }); - - // Verify that we received one ExecuteSqlRequest, and it uses a regular session due to fallback. - List executeSqlRequests = - mockSpanner.getRequestsOfType(ExecuteSqlRequest.class); - assertEquals(1, executeSqlRequests.size()); - // Verify the requests are not executed using multiplexed sessions - Session session2 = mockSpanner.getSession(executeSqlRequests.get(0).getSession()); - assertNotNull(session2); - assertFalse(session2.getMultiplexed()); - } - - // Tests the behavior of the server-side kill switch for read-write multiplexed sessions. - @Test - public void - testInitialBeginTransactionWithPDML_receivesUnimplemented_fallsBackToRegularSession() { - mockSpanner.setBeginTransactionExecutionTime( - SimulatedExecutionTime.ofExceptions( - Arrays.asList( - Status.UNIMPLEMENTED - .withDescription( - "Transaction type partitioned_dml not supported with multiplexed sessions") - .asRuntimeException(), - Status.UNIMPLEMENTED - .withDescription( - "Transaction type partitioned_dml not supported with multiplexed sessions") - .asRuntimeException()))); - DatabaseClientImpl client = - (DatabaseClientImpl) spanner.getDatabaseClient(DatabaseId.of("p", "i", "d")); - - assertNotNull(client.multiplexedSessionDatabaseClient); - - // Partitioned Ops transaction should fallback to regular sessions - assertEquals(UPDATE_COUNT, client.executePartitionedUpdate(UPDATE_STATEMENT)); - - // Verify that we received one ExecuteSqlRequest, and it uses a regular session due to fallback. - List executeSqlRequests = - mockSpanner.getRequestsOfType(ExecuteSqlRequest.class); - assertEquals(1, executeSqlRequests.size()); - // Verify the requests are not executed using multiplexed sessions - Session session2 = mockSpanner.getSession(executeSqlRequests.get(0).getSession()); - assertNotNull(session2); - assertFalse(session2.getMultiplexed()); - assertTrue(client.multiplexedSessionDatabaseClient.unimplementedForPartitionedOps.get()); - } - - /** - * Tests the behavior of the server-side kill switch for partitioned query multiplexed sessions. 2 - * PartitionQueryRequest should be received. First with Multiplexed session and second with - * regular session. - */ - @Test - public void testPartitionedQuery_receivesUnimplemented_fallsBackToRegularSession() { - try { - mockSpanner.setPartitionQueryExecutionTime( - SimulatedExecutionTime.ofException( - Status.INVALID_ARGUMENT - .withDescription( - "Partitioned operations are not supported with multiplexed sessions") - .asRuntimeException())); - BatchClientImpl client = - (BatchClientImpl) spanner.getBatchClient(DatabaseId.of("p", "i", "d")); - - try (BatchReadOnlyTransaction transaction = - client.batchReadOnlyTransaction(TimestampBound.strong())) { - transaction.partitionQuery(PartitionOptions.getDefaultInstance(), STATEMENT); - - // Verify that we received one PartitionQueryRequest. - List partitionQueryRequests = - mockSpanner.getRequestsOfType(PartitionQueryRequest.class); - assertEquals(2, partitionQueryRequests.size()); - // Verify the requests were executed using multiplexed sessions - Session session = mockSpanner.getSession(partitionQueryRequests.get(0).getSession()); - assertNotNull(session); - assertTrue(session.getMultiplexed()); - assertTrue(BatchClientImpl.unimplementedForPartitionedOps.get()); - - session = mockSpanner.getSession(partitionQueryRequests.get(1).getSession()); - assertNotNull(session); - assertFalse(session.getMultiplexed()); - } - } finally { - BatchClientImpl.unimplementedForPartitionedOps.set(false); - } - } - - /** - * Tests the behavior of the server-side kill switch for partitioned query multiplexed sessions. - * The BatchReadOnlyTransaction is initiated using BatchTransactionId. 2 PartitionQueryRequest - * should be received. First with Multiplexed session and second with regular session. - */ - @Test - public void - testPartitionedQueryWithTransactionId_receivesUnimplemented_fallsBackToRegularSession() { - try { - mockSpanner.setPartitionQueryExecutionTime( - SimulatedExecutionTime.ofException( - Status.INVALID_ARGUMENT - .withDescription( - "Partitioned operations are not supported with multiplexed sessions") - .asRuntimeException())); - BatchClientImpl client = - (BatchClientImpl) spanner.getBatchClient(DatabaseId.of("p", "i", "d")); - - try (BatchReadOnlyTransaction transaction = - client.batchReadOnlyTransaction(TimestampBound.strong())) { - - try (BatchReadOnlyTransaction transaction1 = - client.batchReadOnlyTransaction(transaction.getBatchTransactionId())) { - transaction1.partitionQuery(PartitionOptions.getDefaultInstance(), STATEMENT); - - // Verify that we received one PartitionQueryRequest. - List partitionQueryRequests = - mockSpanner.getRequestsOfType(PartitionQueryRequest.class); - assertEquals(2, partitionQueryRequests.size()); - // Verify the requests were executed using multiplexed sessions - Session session = mockSpanner.getSession(partitionQueryRequests.get(0).getSession()); - assertNotNull(session); - assertTrue(session.getMultiplexed()); - assertTrue(BatchClientImpl.unimplementedForPartitionedOps.get()); - - session = mockSpanner.getSession(partitionQueryRequests.get(1).getSession()); - assertNotNull(session); - assertFalse(session.getMultiplexed()); - - List beginTransactionRequests = - mockSpanner.getRequestsOfType(BeginTransactionRequest.class); - assertEquals(2, beginTransactionRequests.size()); - - session = mockSpanner.getSession(beginTransactionRequests.get(0).getSession()); - assertNotNull(session); - assertTrue(session.getMultiplexed()); - - session = mockSpanner.getSession(beginTransactionRequests.get(1).getSession()); - assertNotNull(session); - assertFalse(session.getMultiplexed()); - assertEquals( - transaction.getBatchTransactionId().getTimestamp(), - transaction1.getBatchTransactionId().getTimestamp()); - } - } - } finally { - BatchClientImpl.unimplementedForPartitionedOps.set(false); - } - } - - @Test - public void - testReadWriteUnimplementedErrorDuringInitialBeginTransactionRPC_firstRetriedWithRegularSession_secondFallsBackToRegularSessions() { - // This test simulates the following scenario, - // 1. The server-side flag for RW multiplexed sessions is disabled. - // 2. Application starts. The initial BeginTransaction RPC during client initialization will - // fail with UNIMPLEMENTED error. - // 3. Read-write transaction initialized before the BeginTransaction RPC response will fail with - // UNIMPLEMENTED error. - // 4. Read-write transaction initialized after the BeginTransaction RPC response will fallback - // to regular sessions. - mockSpanner.setBeginTransactionExecutionTime( - SimulatedExecutionTime.ofException( - Status.UNIMPLEMENTED - .withDescription( - "Transaction type read_write not supported with multiplexed sessions") - .asRuntimeException())); - mockSpanner.setExecuteStreamingSqlExecutionTime( - SimulatedExecutionTime.ofException( - Status.UNIMPLEMENTED - .withDescription( - "Transaction type read_write not supported with multiplexed sessions") - .asRuntimeException())); - // Freeze the mock server to ensure that the BeginTransaction with read-write on multiplexed - // session RPC does not return an error or any - // other result just yet. - mockSpanner.freeze(); - // Get a database client using multiplexed sessions. The BeginTransaction RPC to validation - // read-write on multiplexed session will be blocked as - // long as the mock server is frozen. - DatabaseClientImpl client = - (DatabaseClientImpl) spanner.getDatabaseClient(DatabaseId.of("p", "i", "d")); - - // Get the runner so that the read-write transaction is executed via multiplexed session. - TransactionRunner runner = client.readWriteTransaction(); - - // Unfreeze the mock server to get the error from the backend. The above read-write transaction - // will then fail. - mockSpanner.unfreeze(); - - // The ExecuteStreamingSql call fails with UNIMPLEMENTED error, but the retry should happen - // internally with regular session. - runner.run( - transaction -> { - ResultSet resultSet = transaction.executeQuery(STATEMENT); - //noinspection StatementWithEmptyBody - while (resultSet.next()) { - // ignore - } - return null; - }); - assertNotNull(runner.getCommitTimestamp()); - assertNotNull(runner.getCommitResponse()); - - // Wait until the client sees that MultiplexedSessions are not supported for read-write. - assertNotNull(client.multiplexedSessionDatabaseClient); - SpannerException spannerException = - assertThrows( - SpannerException.class, - client.multiplexedSessionDatabaseClient::getReadWriteBeginTransactionReference); - assertEquals(ErrorCode.UNIMPLEMENTED, spannerException.getErrorCode()); - assertTrue(client.multiplexedSessionDatabaseClient.unimplementedForRW.get()); - - // The next read-write transaction will fall back to regular sessions and succeed. - client - .readWriteTransaction() - .run( - transaction -> { - try (ResultSet resultSet = transaction.executeQuery(STATEMENT)) { - //noinspection StatementWithEmptyBody - while (resultSet.next()) { - // ignore - } - } - return null; - }); - - // Verify that two ExecuteSqlRequests were received: the first using a multiplexed session and - // the second using a regular session. - assertEquals(3, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)); - List requests = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class); - - // The ExecuteSqlRequest of the first read-write transaction should use multiplexed session. - Session session1 = mockSpanner.getSession(requests.get(0).getSession()); - assertNotNull(session1); - assertTrue(session1.getMultiplexed()); - - // Retry of the ExecuteSqlRequest of the first read-write transaction should use regular - // session. - Session session2 = mockSpanner.getSession(requests.get(1).getSession()); - assertNotNull(session2); - assertFalse(session2.getMultiplexed()); - - // The ExecuteSqlRequest of the second read-write transaction should use regular session. - Session session3 = mockSpanner.getSession(requests.get(2).getSession()); - assertNotNull(session3); - assertFalse(session3.getMultiplexed()); - - assertNotNull(client.multiplexedSessionDatabaseClient); - assertEquals(1L, client.multiplexedSessionDatabaseClient.getNumSessionsAcquired().get()); - assertEquals(1L, client.multiplexedSessionDatabaseClient.getNumSessionsReleased().get()); - } - - @Test - public void - testReadWriteUnimplemented_firstRetriedWithRegularSession_secondFallsBackToRegularSessions() { - // This test simulates the following scenario, - // 1. The server side flag for read-write multiplexed session is not disabled. When an - // application starts, the initial BeginTransaction RPC with read-write will succeed. - // 2. After time t, the server side flag for read-write multiplexed session is disabled. After - // this a read-write transaction executed with multiplexed sessions should fail with - // UNIMPLEMENTED error. - // 3. All read-write transactions in the application after the initial failure should fallback - // to using regular sessions. - mockSpanner.setExecuteStreamingSqlExecutionTime( - SimulatedExecutionTime.ofException( - Status.UNIMPLEMENTED - .withDescription( - "Transaction type read_write not supported with multiplexed sessions") - .asRuntimeException())); - - DatabaseClientImpl client = - (DatabaseClientImpl) spanner.getDatabaseClient(DatabaseId.of("p", "i", "d")); - - // Wait until the initial BeginTransaction RPC with read-write is complete. - assertNotNull(client.multiplexedSessionDatabaseClient); - Transaction txn = - client.multiplexedSessionDatabaseClient.getReadWriteBeginTransactionReference(); - assertNotNull(txn); - assertNotNull(txn.getId()); - assertFalse(client.multiplexedSessionDatabaseClient.unimplementedForRW.get()); - - // Initially, the first attempt executes an ExecuteSqlRequest using multiplexed sessions, but it - // fails with UNIMPLEMENTED. - // On retry, the request should automatically switch to regular sessions, ensuring the - // transaction completes successfully. - client - .readWriteTransaction() - .run( - transaction -> { - ResultSet resultSet = transaction.executeQuery(STATEMENT); - //noinspection StatementWithEmptyBody - while (resultSet.next()) { - // ignore - } - return null; - }); - - // Verify that the previous failed transaction during first attempt has marked multiplexed - // session client to be - // unimplemented for read-write. - assertTrue(client.multiplexedSessionDatabaseClient.unimplementedForRW.get()); - - // The next read-write transaction will automatically fall back to regular sessions and succeed. - client - .readWriteTransaction() - .run( - transaction -> { - try (ResultSet resultSet = transaction.executeQuery(STATEMENT)) { - //noinspection StatementWithEmptyBody - while (resultSet.next()) { - // ignore - } - } - return null; - }); - - // Verify that two ExecuteSqlRequests were received: the first using a multiplexed session and - // the second using a regular session. - assertEquals(3, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)); - List requests = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class); - - // The ExecuteSqlRequest of the first read-write transaction should use multiplexed session. - Session session1 = mockSpanner.getSession(requests.get(0).getSession()); - assertNotNull(session1); - assertTrue(session1.getMultiplexed()); - - // Retry of the ExecuteSqlRequest of the first read-write transaction should use regular - // session. - Session session2 = mockSpanner.getSession(requests.get(1).getSession()); - assertNotNull(session2); - assertFalse(session2.getMultiplexed()); - - // The ExecuteSqlRequest of the second read-write transaction should use regular session. - Session session3 = mockSpanner.getSession(requests.get(1).getSession()); - assertNotNull(session3); - assertFalse(session3.getMultiplexed()); - - assertNotNull(client.multiplexedSessionDatabaseClient); - assertEquals(1L, client.multiplexedSessionDatabaseClient.getNumSessionsAcquired().get()); - assertEquals(1L, client.multiplexedSessionDatabaseClient.getNumSessionsReleased().get()); - } - @Test public void testOtherUnimplementedError_ReadWriteTransactionStillUsesMultiplexedSession() { mockSpanner.setExecuteStreamingSqlExecutionTime( @@ -1860,22 +1356,11 @@ public void testOtherUnimplementedError_ReadWriteTransactionStillUsesMultiplexed DatabaseClientImpl client = (DatabaseClientImpl) spanner.getDatabaseClient(DatabaseId.of("p", "i", "d")); - // Wait until the initial BeginTransaction RPC with read-write is complete. - assertNotNull(client.multiplexedSessionDatabaseClient); - Transaction txn = - client.multiplexedSessionDatabaseClient.getReadWriteBeginTransactionReference(); - assertNotNull(txn); - assertNotNull(txn.getId()); - assertFalse(client.multiplexedSessionDatabaseClient.unimplementedForRW.get()); - // Try to execute a query using single use transaction. try (ResultSet resultSet = client.singleUse().executeQuery(STATEMENT)) { SpannerException spannerException = assertThrows(SpannerException.class, resultSet::next); assertEquals(ErrorCode.UNIMPLEMENTED, spannerException.getErrorCode()); } - // Verify other UNIMPLEMENTED errors does not turn off read-write transactions to use - // multiplexed sessions. - assertFalse(client.multiplexedSessionDatabaseClient.unimplementedForRW.get()); // The read-write transaction should use multiplexed sessions and succeed. client @@ -2006,172 +1491,6 @@ public void testBatchWriteAtLeastOnce() { assertEquals(1L, client.multiplexedSessionDatabaseClient.getNumSessionsReleased().get()); } - @Test - public void - testReadWriteUnimplementedError_DuringExplicitBegin_RetriedWithRegularSessionForInFlightTransaction() { - // Test scenario: - // 1. The first attempt does an inline begin using a multiplexed session with an invalid - // statement, resulting in failure due to invalid syntax. - // 2. A retry occurs with an explicit begin using a multiplexed session, but we assume the - // backend flag is turned OFF, leading to UNIMPLEMENTED errors. - // 3. Upon encountering the UNIMPLEMENTED error, the entire transaction callable is retried - // using regular sessions, but the inline begin fails again. - // 4. A final retry executes the explicit BeginTransaction on a regular session. - Spanner spanner = setupSpannerBySkippingBeginTransactionVerificationForMux(); - mockSpanner.setBeginTransactionExecutionTime( - SimulatedExecutionTime.ofException( - Status.UNIMPLEMENTED - .withDescription( - "Transaction type read_write not supported with multiplexed sessions") - .asRuntimeException())); - - DatabaseClientImpl client = - (DatabaseClientImpl) spanner.getDatabaseClient(DatabaseId.of("p", "i", "d")); - TransactionRunner runner = client.readWriteTransaction(); - Long updateCount = - runner.run( - transaction -> { - // This update statement carries the BeginTransaction, but fails. This will - // cause the entire transaction to be retried with an explicit - // BeginTransaction RPC to ensure all statements in the transaction are - // actually executed against the same transaction. - SpannerException e = - assertThrows( - SpannerException.class, - () -> transaction.executeUpdate(INVALID_UPDATE_STATEMENT)); - assertEquals(ErrorCode.INVALID_ARGUMENT, e.getErrorCode()); - return transaction.executeUpdate(UPDATE_STATEMENT); - }); - - assertThat(updateCount).isEqualTo(1L); - List beginTransactionRequests = - mockSpanner.getRequestsOfType(BeginTransactionRequest.class); - assertEquals(2, beginTransactionRequests.size()); - - // Verify the first BeginTransaction request is executed using multiplexed sessions. - assertTrue( - mockSpanner.getSession(beginTransactionRequests.get(0).getSession()).getMultiplexed()); - - // Verify the second BeginTransaction request is executed using regular sessions. - assertFalse( - mockSpanner.getSession(beginTransactionRequests.get(1).getSession()).getMultiplexed()); - } - - @Test - public void - testReadWriteUnimplementedError_RetriedWithRegularSessionForInFlightTransaction_RetriedWithSessionNotFound() { - // Test scenario: - // 1. The initial attempt performs an inline begin using a multiplexed session, but with the - // backend flag assumed to be OFF, resulting in an UNIMPLEMENTED error. - // 2. Upon encountering the UNIMPLEMENTED error, the entire transaction callable is retried - // using regular sessions. However, the Commit request fails due to a SessionNotFound error. - // 3. A final retry is triggered to handle the SessionNotFound error by selecting a new session - // from the pool, leading to a successful transaction. - Spanner spanner = setupSpannerBySkippingBeginTransactionVerificationForMux(); - - // The first ExecuteSql request that does an inline begin with multiplexed sessions fail with - // UNIMPLEMENTED error. - mockSpanner.setExecuteSqlExecutionTime( - SimulatedExecutionTime.ofException( - Status.UNIMPLEMENTED - .withDescription( - "Transaction type read_write not supported with multiplexed sessions") - .asRuntimeException())); - - // The first Commit request fails with SessionNotFound exception. The first time this commit is - // called with be using regular sessions. - // This is done to verify if SessionNotFound errors on regular sessions are handled. - mockSpanner.setCommitExecutionTime( - SimulatedExecutionTime.ofException( - mockSpanner.createSessionNotFoundException("TEST_SESSION_NAME"))); - - DatabaseClientImpl client = - (DatabaseClientImpl) spanner.getDatabaseClient(DatabaseId.of("p", "i", "d")); - TransactionRunner runner = client.readWriteTransaction(); - Long updateCount = - runner.run( - transaction -> { - return transaction.executeUpdate(UPDATE_STATEMENT); - }); - - assertThat(updateCount).isEqualTo(1L); - List executeSqlRequests = - mockSpanner.getRequestsOfType(ExecuteSqlRequest.class); - assertEquals(3, executeSqlRequests.size()); - - // Verify the first BeginTransaction request is executed using multiplexed sessions. - assertTrue(mockSpanner.getSession(executeSqlRequests.get(0).getSession()).getMultiplexed()); - - // Verify the second BeginTransaction request is executed using regular sessions. - assertFalse(mockSpanner.getSession(executeSqlRequests.get(1).getSession()).getMultiplexed()); - - // Verify the second BeginTransaction request is executed using regular sessions. - assertFalse(mockSpanner.getSession(executeSqlRequests.get(2).getSession()).getMultiplexed()); - - // Verify that after the first regular session failed with SessionNotFoundException, a new - // regular session is picked up to re-run the transaction. - assertNotEquals(executeSqlRequests.get(1).getSession(), executeSqlRequests.get(2).getSession()); - } - - @Test - public void - testReadWriteUnimplementedError_FirstSucceedsWithMux_SecondRetriedWithRegularSessionDueToUnimplementedError() { - // Test scenario: - // 1. The first read-write transaction successfully performs an inline begin using a multiplexed - // session. - // 2. The second read-write transaction attempts to execute with a multiplexed session, but - // since the backend flag is assumed to be OFF, it encounters an UNIMPLEMENTED error. - // 3. Upon encountering the UNIMPLEMENTED error, the entire transaction callable for the second - // read-write transaction is retried using a regular session. - - Spanner spanner = setupSpannerBySkippingBeginTransactionVerificationForMux(); - - DatabaseClientImpl client = - (DatabaseClientImpl) spanner.getDatabaseClient(DatabaseId.of("p", "i", "d")); - - // First read-write transaction attempt succeeds. - TransactionRunner runner = client.readWriteTransaction(); - Long updateCount = - runner.run( - transaction -> { - return transaction.executeUpdate(UPDATE_STATEMENT); - }); - - assertThat(updateCount).isEqualTo(1L); - - // The ExecuteSql request is forced to fail with UNIMPLEMENTED error. - mockSpanner.setExecuteSqlExecutionTime( - SimulatedExecutionTime.ofException( - Status.UNIMPLEMENTED - .withDescription( - "Transaction type read_write not supported with multiplexed sessions") - .asRuntimeException())); - - // Second read-write transaction on mux fails with UNIMPLEMENTED error, and then retried using - // regular session. - TransactionRunner runner1 = client.readWriteTransaction(); - Long updateCount1 = - runner1.run( - transaction -> { - return transaction.executeUpdate(UPDATE_STATEMENT); - }); - - assertThat(updateCount1).isEqualTo(1L); - - List executeSqlRequests = - mockSpanner.getRequestsOfType(ExecuteSqlRequest.class); - assertEquals(3, executeSqlRequests.size()); - - // Verify the first BeginTransaction request is executed using multiplexed sessions. - assertTrue(mockSpanner.getSession(executeSqlRequests.get(0).getSession()).getMultiplexed()); - - // Verify the second BeginTransaction request is executed using multiplexed sessions. - assertTrue(mockSpanner.getSession(executeSqlRequests.get(1).getSession()).getMultiplexed()); - - // Verify the second BeginTransaction request is executed using regular sessions. - assertFalse(mockSpanner.getSession(executeSqlRequests.get(2).getSession()).getMultiplexed()); - } - @Test public void testRWTransactionWithTransactionManager_CommitAborted_SetsTransactionId_AndUsedInNewInstance() { diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionsBenchmark.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionsBenchmark.java index c6f7e22f280..f71fdfe37a3 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionsBenchmark.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionsBenchmark.java @@ -17,7 +17,6 @@ package com.google.cloud.spanner; import static com.google.cloud.spanner.BenchmarkingUtilityScripts.collectResults; -import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertNotNull; @@ -118,9 +117,6 @@ public void teardown() throws Exception { @Benchmark public void burstQueries(final BenchmarkState server) throws Exception { final DatabaseClientImpl client = server.client; - SessionPool pool = client.pool; - assertThat(pool.totalSessions()) - .isEqualTo(server.spanner.getOptions().getSessionPoolOptions().getMinSessions()); ListeningScheduledExecutorService service = MoreExecutors.listeningDecorator(Executors.newScheduledThreadPool(PARALLEL_THREADS)); diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadAsyncTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadAsyncTest.java index 3704b118906..ede202ced3a 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadAsyncTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadAsyncTest.java @@ -231,9 +231,6 @@ public void closeTransactionBeforeEndOfAsyncQuery() throws Exception { ApiFuture closed; DatabaseClientImpl clientImpl = (DatabaseClientImpl) client; - // There should currently not be any sessions checked out of the pool. - assertThat(clientImpl.pool.getNumberOfSessionsInUse()).isEqualTo(0); - final CountDownLatch dataReceived = new CountDownLatch(1); try (ReadOnlyTransaction tx = client.readOnlyTransaction()) { try (AsyncResultSet rs = @@ -264,22 +261,6 @@ public void closeTransactionBeforeEndOfAsyncQuery() throws Exception { // Wait until at least one row has been fetched. At that moment there should be one session // checked out. dataReceived.await(); - - if (isMultiplexedSessionsEnabled()) { - assertThat(clientImpl.pool.getNumberOfSessionsInUse()).isEqualTo(0); - } else { - assertThat(clientImpl.pool.getNumberOfSessionsInUse()).isEqualTo(1); - } - } - // The read-only transaction is now closed, but the ready callback will continue to receive - // data. As it tries to put the data into a synchronous queue and the underlying buffer can also - // only hold 1 row, the async result set has not yet finished. The read-only transaction will - // release the session back into the pool when all async statements have finished. The number of - // sessions in use is therefore still 1. - if (isMultiplexedSessionsEnabled()) { - assertThat(clientImpl.pool.getNumberOfSessionsInUse()).isEqualTo(0); - } else { - assertThat(clientImpl.pool.getNumberOfSessionsInUse()).isEqualTo(1); } List resultList = new ArrayList<>(); do { @@ -287,10 +268,7 @@ public void closeTransactionBeforeEndOfAsyncQuery() throws Exception { } while (!finished.isDone() || results.size() > 0); assertThat(finished.get()).isTrue(); assertThat(resultList).containsExactly("k1", "k2", "k3"); - // The session will be released back into the pool by the asynchronous result set when it has - // returned all rows. As this is done in the background, it could take a couple of milliseconds. closed.get(); - assertThat(clientImpl.pool.getNumberOfSessionsInUse()).isEqualTo(0); } @Test diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/RetryOnInvalidatedSessionTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/RetryOnInvalidatedSessionTest.java deleted file mode 100644 index 6a0f16c02b1..00000000000 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/RetryOnInvalidatedSessionTest.java +++ /dev/null @@ -1,1797 +0,0 @@ -/* - * Copyright 2019 Google LLC - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package com.google.cloud.spanner; - -import static com.google.cloud.spanner.DisableDefaultMtlsProvider.disableDefaultMtlsProvider; -import static com.google.cloud.spanner.SpannerApiFutures.get; -import static com.google.common.truth.Truth.assertThat; -import static org.junit.Assert.assertThrows; -import static org.junit.Assert.fail; -import static org.junit.Assume.assumeFalse; - -import com.google.api.core.ApiFuture; -import com.google.api.core.ApiFutures; -import com.google.api.gax.core.NoCredentialsProvider; -import com.google.api.gax.grpc.testing.LocalChannelProvider; -import com.google.cloud.NoCredentials; -import com.google.cloud.spanner.AsyncResultSet.CallbackResponse; -import com.google.cloud.spanner.AsyncTransactionManager.AsyncTransactionStep; -import com.google.cloud.spanner.AsyncTransactionManager.CommitTimestampFuture; -import com.google.cloud.spanner.AsyncTransactionManager.TransactionContextFuture; -import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult; -import com.google.cloud.spanner.v1.SpannerClient; -import com.google.cloud.spanner.v1.SpannerClient.ListSessionsPagedResponse; -import com.google.cloud.spanner.v1.SpannerSettings; -import com.google.common.base.Function; -import com.google.common.base.Stopwatch; -import com.google.common.util.concurrent.MoreExecutors; -import com.google.protobuf.ListValue; -import com.google.spanner.v1.ResultSetMetadata; -import com.google.spanner.v1.StructType; -import com.google.spanner.v1.StructType.Field; -import com.google.spanner.v1.TypeCode; -import io.grpc.Server; -import io.grpc.inprocess.InProcessServerBuilder; -import java.util.ArrayList; -import java.util.Arrays; -import java.util.Collection; -import java.util.Collections; -import java.util.List; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.Executors; -import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicInteger; -import java.util.concurrent.atomic.AtomicLong; -import java.util.function.Supplier; -import org.junit.AfterClass; -import org.junit.Before; -import org.junit.BeforeClass; -import org.junit.Test; -import org.junit.runner.RunWith; -import org.junit.runners.Parameterized; -import org.junit.runners.Parameterized.Parameter; -import org.junit.runners.Parameterized.Parameters; - -@RunWith(Parameterized.class) -public class RetryOnInvalidatedSessionTest { - private static final class ToLongTransformer implements Function { - @Override - public Long apply(StructReader input) { - return input.getLong(0); - } - } - - private static final ToLongTransformer TO_LONG = new ToLongTransformer(); - - @Parameter(0) - public boolean failOnInvalidatedSession; - - @Parameters(name = "fail on invalidated session = {0}") - public static Collection data() { - List params = new ArrayList<>(); - params.add(new Object[] {false}); - params.add(new Object[] {true}); - return params; - } - - private static final ResultSetMetadata READ_METADATA = - ResultSetMetadata.newBuilder() - .setRowType( - StructType.newBuilder() - .addFields( - Field.newBuilder() - .setName("BAR") - .setType( - com.google.spanner.v1.Type.newBuilder() - .setCode(TypeCode.INT64) - .build()) - .build()) - .build()) - .build(); - private static final com.google.spanner.v1.ResultSet READ_RESULTSET = - com.google.spanner.v1.ResultSet.newBuilder() - .addRows( - ListValue.newBuilder() - .addValues(com.google.protobuf.Value.newBuilder().setStringValue("1").build()) - .build()) - .addRows( - ListValue.newBuilder() - .addValues(com.google.protobuf.Value.newBuilder().setStringValue("2").build()) - .build()) - .setMetadata(READ_METADATA) - .build(); - private static final com.google.spanner.v1.ResultSet READ_ROW_RESULTSET = - com.google.spanner.v1.ResultSet.newBuilder() - .addRows( - ListValue.newBuilder() - .addValues(com.google.protobuf.Value.newBuilder().setStringValue("1").build()) - .build()) - .setMetadata(READ_METADATA) - .build(); - private static final Statement SELECT1AND2 = - Statement.of("SELECT 1 AS COL1 UNION ALL SELECT 2 AS COL1"); - private static final ResultSetMetadata SELECT1AND2_METADATA = - ResultSetMetadata.newBuilder() - .setRowType( - StructType.newBuilder() - .addFields( - Field.newBuilder() - .setName("COL1") - .setType( - com.google.spanner.v1.Type.newBuilder() - .setCode(TypeCode.INT64) - .build()) - .build()) - .build()) - .build(); - private static final com.google.spanner.v1.ResultSet SELECT1_RESULTSET = - com.google.spanner.v1.ResultSet.newBuilder() - .addRows( - ListValue.newBuilder() - .addValues(com.google.protobuf.Value.newBuilder().setStringValue("1").build()) - .build()) - .addRows( - ListValue.newBuilder() - .addValues(com.google.protobuf.Value.newBuilder().setStringValue("2").build()) - .build()) - .setMetadata(SELECT1AND2_METADATA) - .build(); - private static final Statement UPDATE_STATEMENT = - Statement.of("UPDATE FOO SET BAR=1 WHERE BAZ=2"); - private static final long UPDATE_COUNT = 1L; - private static MockSpannerServiceImpl mockSpanner; - private static Server server; - private static LocalChannelProvider channelProvider; - private static SpannerClient spannerClient; - private static Spanner spanner; - private static DatabaseClient client; - private static ExecutorService executor; - - @BeforeClass - public static void startStaticServer() throws Exception { - disableDefaultMtlsProvider(); - mockSpanner = new MockSpannerServiceImpl(); - mockSpanner.setAbortProbability(0.0D); // We don't want any unpredictable aborted transactions. - mockSpanner.putStatementResult( - StatementResult.read( - "FOO", KeySet.all(), Collections.singletonList("BAR"), READ_RESULTSET)); - mockSpanner.putStatementResult( - StatementResult.read( - "FOO", - KeySet.singleKey(Key.of()), - Collections.singletonList("BAR"), - READ_ROW_RESULTSET)); - mockSpanner.putStatementResult(StatementResult.query(SELECT1AND2, SELECT1_RESULTSET)); - mockSpanner.putStatementResult(StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT)); - - String uniqueName = InProcessServerBuilder.generateName(); - server = - InProcessServerBuilder.forName(uniqueName) - .directExecutor() - .addService(mockSpanner) - .build() - .start(); - channelProvider = LocalChannelProvider.create(uniqueName); - - SpannerSettings settings = - SpannerSettings.newBuilder() - .setTransportChannelProvider(channelProvider) - .setCredentialsProvider(NoCredentialsProvider.create()) - .build(); - spannerClient = SpannerClient.create(settings); - executor = Executors.newSingleThreadExecutor(); - } - - @AfterClass - public static void stopServer() throws InterruptedException { - spannerClient.close(); - server.shutdown(); - server.awaitTermination(); - executor.shutdown(); - } - - @Before - public void setUp() throws InterruptedException { - mockSpanner.reset(); - if (spanner == null - || spanner.getOptions().getSessionPoolOptions().isFailIfSessionNotFound() - != failOnInvalidatedSession) { - if (spanner != null) { - spanner.close(); - } - SessionPoolOptions.Builder builder = SessionPoolOptions.newBuilder().setFailOnSessionLeak(); - if (failOnInvalidatedSession) { - builder.setFailIfSessionNotFound(); - } - // This prevents repeated retries for a large number of sessions in the pool. - builder.setMinSessions(1); - SessionPoolOptions sessionPoolOptions = builder.build(); - spanner = - SpannerOptions.newBuilder() - .setProjectId("[PROJECT]") - .setChannelProvider(channelProvider) - .setSessionPoolOption(sessionPoolOptions) - .setCredentials(NoCredentials.getInstance()) - .build() - .getService(); - client = spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]")); - } - invalidateSessionPool(client, spanner.getOptions().getSessionPoolOptions().getMinSessions()); - } - - private static void invalidateSessionPool(DatabaseClient client, int minSessions) - throws InterruptedException { - // Wait for all sessions to have been created, and then delete them. - Stopwatch watch = Stopwatch.createStarted(); - while (((DatabaseClientImpl) client).pool.totalSessions() < minSessions) { - if (watch.elapsed(TimeUnit.SECONDS) > 5L) { - fail(String.format("Failed to create MinSessions=%d", minSessions)); - } - Thread.sleep(1L); - } - - ListSessionsPagedResponse response = - spannerClient.listSessions("projects/[PROJECT]/instances/[INSTANCE]/databases/[DATABASE]"); - for (com.google.spanner.v1.Session session : response.iterateAll()) { - spannerClient.deleteSession(session.getName()); - } - } - - private T assertThrowsSessionNotFoundIfShouldFail(Supplier supplier) { - if (failOnInvalidatedSession) { - assertThrows(SessionNotFoundException.class, () -> supplier.get()); - return null; - } else { - return supplier.get(); - } - } - - @Test - public void singleUseSelect() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSession()); - // This call will receive an invalidated session that will be replaced on the first call to - // rs.next(). - try (ReadContext context = client.singleUse()) { - try (ResultSet rs = context.executeQuery(SELECT1AND2)) { - assertThrowsSessionNotFoundIfShouldFail(() -> rs.next()); - } - } - } - - @Test - public void singleUseSelectAsync() throws Exception { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSession()); - ApiFuture> list; - try (AsyncResultSet rs = client.singleUse().executeQueryAsync(SELECT1AND2)) { - list = rs.toListAsync(TO_LONG, executor); - assertThrowsSessionNotFoundIfShouldFail(() -> get(list)); - } - } - - @Test - public void singleUseRead() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSession()); - try (ReadContext context = client.singleUse()) { - try (ResultSet rs = context.read("FOO", KeySet.all(), Collections.singletonList("BAR"))) { - assertThrowsSessionNotFoundIfShouldFail(() -> rs.next()); - } - } - } - - @Test - public void singleUseReadUsingIndex() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSession()); - try (ReadContext context = client.singleUse()) { - try (ResultSet rs = - context.readUsingIndex("FOO", "IDX", KeySet.all(), Collections.singletonList("BAR"))) { - assertThrowsSessionNotFoundIfShouldFail(() -> rs.next()); - } - } - } - - @Test - public void singleUseReadRow() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSession()); - try (ReadContext context = client.singleUse()) { - assertThrowsSessionNotFoundIfShouldFail( - () -> context.readRow("FOO", Key.of(), Collections.singletonList("BAR"))); - } - } - - @Test - public void singleUseReadRowUsingIndex() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSession()); - try (ReadContext context = client.singleUse()) { - assertThrowsSessionNotFoundIfShouldFail( - () -> - context.readRowUsingIndex("FOO", "IDX", Key.of(), Collections.singletonList("BAR"))); - } - } - - @Test - public void singleUseReadOnlyTransactionSelect() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSession()); - try (ReadContext context = client.singleUseReadOnlyTransaction()) { - try (ResultSet rs = context.executeQuery(SELECT1AND2)) { - assertThrowsSessionNotFoundIfShouldFail(() -> rs.next()); - } - } - } - - @Test - public void singleUseReadOnlyTransactionRead() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSession()); - try (ReadContext context = client.singleUseReadOnlyTransaction()) { - try (ResultSet rs = context.read("FOO", KeySet.all(), Collections.singletonList("BAR"))) { - assertThrowsSessionNotFoundIfShouldFail(() -> rs.next()); - } - } - } - - @Test - public void singlUseReadOnlyTransactionReadUsingIndex() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSession()); - try (ReadContext context = client.singleUseReadOnlyTransaction()) { - try (ResultSet rs = - context.readUsingIndex("FOO", "IDX", KeySet.all(), Collections.singletonList("BAR"))) { - assertThrowsSessionNotFoundIfShouldFail(() -> rs.next()); - } - } - } - - @Test - public void singleUseReadOnlyTransactionReadRow() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSession()); - try (ReadContext context = client.singleUseReadOnlyTransaction()) { - assertThrowsSessionNotFoundIfShouldFail( - () -> context.readRow("FOO", Key.of(), Collections.singletonList("BAR"))); - } - } - - @Test - public void singleUseReadOnlyTransactionReadRowUsingIndex() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSession()); - try (ReadContext context = client.singleUseReadOnlyTransaction()) { - assertThrowsSessionNotFoundIfShouldFail( - () -> - context.readRowUsingIndex("FOO", "IDX", Key.of(), Collections.singletonList("BAR"))); - } - } - - @Test - public void readOnlyTransactionSelect() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSession()); - try (ReadContext context = client.readOnlyTransaction()) { - try (ResultSet rs = context.executeQuery(SELECT1AND2)) { - assertThrowsSessionNotFoundIfShouldFail(() -> rs.next()); - } - } - } - - @Test - public void readOnlyTransactionRead() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSession()); - try (ReadContext context = client.readOnlyTransaction()) { - try (ResultSet rs = context.read("FOO", KeySet.all(), Collections.singletonList("BAR"))) { - assertThrowsSessionNotFoundIfShouldFail(() -> rs.next()); - } - } - } - - @Test - public void readOnlyTransactionReadUsingIndex() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSession()); - try (ReadContext context = client.readOnlyTransaction()) { - try (ResultSet rs = - context.readUsingIndex("FOO", "IDX", KeySet.all(), Collections.singletonList("BAR"))) { - assertThrowsSessionNotFoundIfShouldFail(() -> rs.next()); - } - } - } - - @Test - public void readOnlyTransactionReadRow() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSession()); - try (ReadContext context = client.readOnlyTransaction()) { - assertThrowsSessionNotFoundIfShouldFail( - () -> context.readRow("FOO", Key.of(), Collections.singletonList("BAR"))); - } - } - - @Test - public void readOnlyTransactionReadRowUsingIndex() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSession()); - try (ReadContext context = client.readOnlyTransaction()) { - assertThrowsSessionNotFoundIfShouldFail( - () -> - context.readRowUsingIndex("FOO", "IDX", Key.of(), Collections.singletonList("BAR"))); - } - } - - @Test - public void readOnlyTransactionSelectNonRecoverable() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSession()); - try (ReadContext context = client.readOnlyTransaction()) { - try (ResultSet rs = context.executeQuery(SELECT1AND2)) { - assertThrowsSessionNotFoundIfShouldFail(() -> rs.next()); - } - // Invalidate the session pool while in a transaction. This is not recoverable. - invalidateSessionPool(client, spanner.getOptions().getSessionPoolOptions().getMinSessions()); - try (ResultSet rs = context.executeQuery(SELECT1AND2)) { - assertThrows(SessionNotFoundException.class, () -> rs.next()); - } - } - } - - @Test - public void readOnlyTransactionReadNonRecoverable() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSession()); - try (ReadContext context = client.readOnlyTransaction()) { - try (ResultSet rs = context.read("FOO", KeySet.all(), Collections.singletonList("BAR"))) { - assertThrowsSessionNotFoundIfShouldFail(() -> rs.next()); - } - invalidateSessionPool(client, spanner.getOptions().getSessionPoolOptions().getMinSessions()); - try (ResultSet rs = context.read("FOO", KeySet.all(), Collections.singletonList("BAR"))) { - assertThrows(SessionNotFoundException.class, () -> rs.next()); - } - } - } - - @Test - public void readOnlyTransactionReadUsingIndexNonRecoverable() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSession()); - try (ReadContext context = client.readOnlyTransaction()) { - try (ResultSet rs = - context.readUsingIndex("FOO", "IDX", KeySet.all(), Collections.singletonList("BAR"))) { - assertThrowsSessionNotFoundIfShouldFail(() -> rs.next()); - } - invalidateSessionPool(client, spanner.getOptions().getSessionPoolOptions().getMinSessions()); - try (ResultSet rs = - context.readUsingIndex("FOO", "IDX", KeySet.all(), Collections.singletonList("BAR"))) { - assertThrows(SessionNotFoundException.class, () -> rs.next()); - } - } - } - - @Test - public void readOnlyTransactionReadRowNonRecoverable() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSession()); - try (ReadContext context = client.readOnlyTransaction()) { - assertThrowsSessionNotFoundIfShouldFail( - () -> context.readRow("FOO", Key.of(), Collections.singletonList("BAR"))); - invalidateSessionPool(client, spanner.getOptions().getSessionPoolOptions().getMinSessions()); - assertThrows( - SessionNotFoundException.class, - () -> context.readRow("FOO", Key.of(), Collections.singletonList("BAR"))); - } - } - - @Test - public void readOnlyTransactionReadRowUsingIndexNonRecoverable() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSession()); - try (ReadContext context = client.readOnlyTransaction()) { - assertThrowsSessionNotFoundIfShouldFail( - () -> - context.readRowUsingIndex("FOO", "IDX", Key.of(), Collections.singletonList("BAR"))); - invalidateSessionPool(client, spanner.getOptions().getSessionPoolOptions().getMinSessions()); - assertThrows( - SessionNotFoundException.class, - () -> - context.readRowUsingIndex("FOO", "IDX", Key.of(), Collections.singletonList("BAR"))); - } - } - - @Test - public void readWriteTransactionReadOnlySessionInPool() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - SessionPoolOptions.Builder builder = SessionPoolOptions.newBuilder(); - if (failOnInvalidatedSession) { - builder.setFailIfSessionNotFound(); - } - try (Spanner spanner = - SpannerOptions.newBuilder() - .setProjectId("[PROJECT]") - .setChannelProvider(channelProvider) - .setSessionPoolOption(builder.build()) - .setCredentials(NoCredentials.getInstance()) - .build() - .getService()) { - DatabaseClient client = - spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]")); - invalidateSessionPool(client, spanner.getOptions().getSessionPoolOptions().getMinSessions()); - TransactionRunner runner = client.readWriteTransaction(); - assertThrowsSessionNotFoundIfShouldFail( - () -> - runner.run( - transaction -> { - try (ResultSet rs = transaction.executeQuery(SELECT1AND2)) { - while (rs.next()) {} - } - return null; - })); - } - } - - @Test - public void readWriteTransactionSelect() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - TransactionRunner runner = client.readWriteTransaction(); - assertThrowsSessionNotFoundIfShouldFail( - () -> - runner.run( - transaction -> { - try (ResultSet rs = transaction.executeQuery(SELECT1AND2)) { - while (rs.next()) {} - } - return null; - })); - } - - @Test - public void readWriteTransactionRead() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - TransactionRunner runner = client.readWriteTransaction(); - assertThrowsSessionNotFoundIfShouldFail( - () -> - runner.run( - transaction -> { - try (ResultSet rs = - transaction.read("FOO", KeySet.all(), Collections.singletonList("BAR"))) { - while (rs.next()) {} - } - return null; - })); - } - - @Test - public void readWriteTransactionReadWithOptimisticLock() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - TransactionRunner runner = client.readWriteTransaction(Options.optimisticLock()); - assertThrowsSessionNotFoundIfShouldFail( - () -> - runner.run( - transaction -> { - try (ResultSet rs = - transaction.read("FOO", KeySet.all(), Collections.singletonList("BAR"))) { - while (rs.next()) {} - } - return null; - })); - } - - @Test - public void readWriteTransactionReadUsingIndex() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - TransactionRunner runner = client.readWriteTransaction(); - assertThrowsSessionNotFoundIfShouldFail( - () -> - runner.run( - transaction -> { - try (ResultSet rs = - transaction.readUsingIndex( - "FOO", "IDX", KeySet.all(), Collections.singletonList("BAR"))) { - while (rs.next()) {} - } - return null; - })); - } - - @Test - public void readWriteTransactionReadRow() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - TransactionRunner runner = client.readWriteTransaction(); - assertThrowsSessionNotFoundIfShouldFail( - () -> - runner.run( - transaction -> - transaction.readRow("FOO", Key.of(), Collections.singletonList("BAR")))); - } - - @Test - public void readWriteTransactionReadRowUsingIndex() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - TransactionRunner runner = client.readWriteTransaction(); - assertThrowsSessionNotFoundIfShouldFail( - () -> - runner.run( - transaction -> - transaction.readRowUsingIndex( - "FOO", "IDX", Key.of(), Collections.singletonList("BAR")))); - } - - @Test - public void readWriteTransactionUpdate() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - TransactionRunner runner = client.readWriteTransaction(); - assertThrowsSessionNotFoundIfShouldFail( - () -> runner.run(transaction -> transaction.executeUpdate(UPDATE_STATEMENT))); - } - - @Test - public void readWriteTransactionBatchUpdate() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - TransactionRunner runner = client.readWriteTransaction(); - assertThrowsSessionNotFoundIfShouldFail( - () -> - runner.run( - transaction -> - transaction.batchUpdate(Collections.singletonList(UPDATE_STATEMENT)))); - } - - @Test - public void readWriteTransactionBuffer() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - TransactionRunner runner = client.readWriteTransaction(); - assertThrowsSessionNotFoundIfShouldFail( - () -> - runner.run( - transaction -> { - transaction.buffer(Mutation.newInsertBuilder("FOO").set("BAR").to(1L).build()); - return null; - })); - } - - @Test - public void readWriteTransactionSelectInvalidatedDuringTransaction() { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - TransactionRunner runner = client.readWriteTransaction(); - final AtomicInteger attempt = new AtomicInteger(); - assertThrowsSessionNotFoundIfShouldFail( - () -> - runner.run( - transaction -> { - attempt.incrementAndGet(); - try (ResultSet rs = transaction.executeQuery(SELECT1AND2)) { - while (rs.next()) {} - } - if (attempt.get() == 1) { - invalidateSessionPool( - client, spanner.getOptions().getSessionPoolOptions().getMinSessions()); - } - try (ResultSet rs = transaction.executeQuery(SELECT1AND2)) { - while (rs.next()) {} - } - assertThat(attempt.get()).isGreaterThan(1); - return null; - })); - } - - @Test - public void readWriteTransactionReadInvalidatedDuringTransaction() { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - TransactionRunner runner = client.readWriteTransaction(); - final AtomicInteger attempt = new AtomicInteger(); - assertThrowsSessionNotFoundIfShouldFail( - () -> - runner.run( - transaction -> { - attempt.incrementAndGet(); - try (ResultSet rs = - transaction.read("FOO", KeySet.all(), Collections.singletonList("BAR"))) { - while (rs.next()) {} - } - if (attempt.get() == 1) { - invalidateSessionPool( - client, spanner.getOptions().getSessionPoolOptions().getMinSessions()); - } - try (ResultSet rs = - transaction.read("FOO", KeySet.all(), Collections.singletonList("BAR"))) { - while (rs.next()) {} - } - assertThat(attempt.get()).isGreaterThan(1); - return null; - })); - } - - @Test - public void readWriteTransactionReadUsingIndexInvalidatedDuringTransaction() { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - TransactionRunner runner = client.readWriteTransaction(); - final AtomicInteger attempt = new AtomicInteger(); - assertThrowsSessionNotFoundIfShouldFail( - () -> - runner.run( - transaction -> { - attempt.incrementAndGet(); - try (ResultSet rs = - transaction.readUsingIndex( - "FOO", "IDX", KeySet.all(), Collections.singletonList("BAR"))) { - while (rs.next()) {} - } - if (attempt.get() == 1) { - invalidateSessionPool( - client, spanner.getOptions().getSessionPoolOptions().getMinSessions()); - } - try (ResultSet rs = - transaction.readUsingIndex( - "FOO", "IDX", KeySet.all(), Collections.singletonList("BAR"))) { - while (rs.next()) {} - } - assertThat(attempt.get()).isGreaterThan(1); - return null; - })); - } - - @Test - public void readWriteTransactionReadRowInvalidatedDuringTransaction() { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - TransactionRunner runner = client.readWriteTransaction(); - final AtomicInteger attempt = new AtomicInteger(); - assertThrowsSessionNotFoundIfShouldFail( - () -> - runner.run( - transaction -> { - attempt.incrementAndGet(); - Struct row = - transaction.readRow("FOO", Key.of(), Collections.singletonList("BAR")); - assertThat(row.getLong(0)).isEqualTo(1L); - if (attempt.get() == 1) { - invalidateSessionPool( - client, spanner.getOptions().getSessionPoolOptions().getMinSessions()); - } - transaction.readRow("FOO", Key.of(), Collections.singletonList("BAR")); - assertThat(attempt.get()).isGreaterThan(1); - return null; - })); - } - - @Test - public void readWriteTransactionReadRowUsingIndexInvalidatedDuringTransaction() { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - TransactionRunner runner = client.readWriteTransaction(); - final AtomicInteger attempt = new AtomicInteger(); - assertThrowsSessionNotFoundIfShouldFail( - () -> - runner.run( - transaction -> { - attempt.incrementAndGet(); - Struct row = - transaction.readRowUsingIndex( - "FOO", "IDX", Key.of(), Collections.singletonList("BAR")); - assertThat(row.getLong(0)).isEqualTo(1L); - if (attempt.get() == 1) { - invalidateSessionPool( - client, spanner.getOptions().getSessionPoolOptions().getMinSessions()); - } - transaction.readRowUsingIndex( - "FOO", "IDX", Key.of(), Collections.singletonList("BAR")); - assertThat(attempt.get()).isGreaterThan(1); - return null; - })); - } - - @SuppressWarnings("resource") - @Test - public void transactionManagerReadOnlySessionInPool() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - try (TransactionManager manager = client.transactionManager()) { - TransactionContext transaction = manager.begin(); - while (true) { - try { - try (ResultSet rs = transaction.executeQuery(SELECT1AND2)) { - assertThrowsSessionNotFoundIfShouldFail(() -> rs.next()); - } - manager.commit(); - break; - } catch (AbortedException e) { - transaction = assertThrowsSessionNotFoundIfShouldFail(() -> manager.resetForRetry()); - if (transaction == null) { - break; - } - } - } - } - } - - @SuppressWarnings("resource") - @Test - public void transactionManagerSelect() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - try (TransactionManager manager = client.transactionManager()) { - TransactionContext transaction = manager.begin(); - while (true) { - try { - try (ResultSet rs = transaction.executeQuery(SELECT1AND2)) { - assertThrowsSessionNotFoundIfShouldFail(() -> rs.next()); - } - manager.commit(); - break; - } catch (AbortedException e) { - transaction = assertThrowsSessionNotFoundIfShouldFail(() -> manager.resetForRetry()); - if (transaction == null) { - break; - } - } - } - } - } - - @SuppressWarnings("resource") - @Test - public void transactionManagerRead() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - try (TransactionManager manager = client.transactionManager()) { - TransactionContext transaction = manager.begin(); - while (true) { - try { - try (ResultSet rs = - transaction.read("FOO", KeySet.all(), Collections.singletonList("BAR"))) { - assertThrowsSessionNotFoundIfShouldFail(() -> rs.next()); - } - manager.commit(); - break; - } catch (AbortedException e) { - transaction = assertThrowsSessionNotFoundIfShouldFail(() -> manager.resetForRetry()); - if (transaction == null) { - break; - } - } - } - } - } - - @SuppressWarnings("resource") - @Test - public void transactionManagerReadUsingIndex() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - try (TransactionManager manager = client.transactionManager()) { - TransactionContext transaction = manager.begin(); - while (true) { - try { - try (ResultSet rs = - transaction.readUsingIndex( - "FOO", "IDX", KeySet.all(), Collections.singletonList("BAR"))) { - assertThrowsSessionNotFoundIfShouldFail(() -> rs.next()); - } - manager.commit(); - break; - } catch (AbortedException e) { - transaction = assertThrowsSessionNotFoundIfShouldFail(() -> manager.resetForRetry()); - if (transaction == null) { - break; - } - } - } - } - } - - @Test - public void transactionManagerReadRow() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - try (TransactionManager manager = client.transactionManager()) { - TransactionContext transaction = manager.begin(); - while (true) { - try { - TransactionContext context = transaction; - assertThrowsSessionNotFoundIfShouldFail( - () -> context.readRow("FOO", Key.of(), Collections.singletonList("BAR"))); - manager.commit(); - break; - } catch (AbortedException e) { - transaction = assertThrowsSessionNotFoundIfShouldFail(() -> manager.resetForRetry()); - if (transaction == null) { - break; - } - } - } - } - } - - @Test - public void transactionManagerReadRowUsingIndex() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - try (TransactionManager manager = client.transactionManager()) { - TransactionContext transaction = manager.begin(); - while (true) { - try { - TransactionContext context = transaction; - assertThrowsSessionNotFoundIfShouldFail( - () -> - context.readRowUsingIndex( - "FOO", "IDX", Key.of(), Collections.singletonList("BAR"))); - manager.commit(); - break; - } catch (AbortedException e) { - transaction = assertThrowsSessionNotFoundIfShouldFail(() -> manager.resetForRetry()); - if (transaction == null) { - break; - } - } - } - } - } - - @Test - public void transactionManagerUpdate() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - try (TransactionManager manager = client.transactionManager(Options.commitStats())) { - TransactionContext transaction = manager.begin(); - while (true) { - try { - TransactionContext context = transaction; - assertThrowsSessionNotFoundIfShouldFail(() -> context.executeUpdate(UPDATE_STATEMENT)); - manager.commit(); - break; - } catch (AbortedException e) { - transaction = assertThrowsSessionNotFoundIfShouldFail(() -> manager.resetForRetry()); - if (transaction == null) { - break; - } - } - } - } - } - - @Test - public void transactionManagerAborted_thenSessionNotFoundOnBeginTransaction() - throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - int attempt = 0; - try (TransactionManager manager = client.transactionManager()) { - TransactionContext transaction = manager.begin(); - while (true) { - try { - attempt++; - if (attempt == 1) { - mockSpanner.abortNextStatement(); - } - if (attempt == 2) { - invalidateSessionPool( - client, spanner.getOptions().getSessionPoolOptions().getMinSessions()); - } - TransactionContext context = transaction; - assertThrowsSessionNotFoundIfShouldFail(() -> context.executeUpdate(UPDATE_STATEMENT)); - manager.commit(); - // The actual number of attempts depends on when the transaction manager will actually get - // a valid session, as we invalidate the entire session pool. - assertThat(attempt).isAtLeast(3); - break; - } catch (AbortedException e) { - transaction = assertThrowsSessionNotFoundIfShouldFail(() -> manager.resetForRetry()); - if (transaction == null) { - break; - } - } - } - } - } - - @Test - public void transactionManagerBatchUpdate() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - try (TransactionManager manager = client.transactionManager()) { - TransactionContext transaction = manager.begin(); - while (true) { - try { - TransactionContext context = transaction; - assertThrowsSessionNotFoundIfShouldFail( - () -> context.batchUpdate(Collections.singletonList(UPDATE_STATEMENT))); - manager.commit(); - break; - } catch (AbortedException e) { - transaction = assertThrowsSessionNotFoundIfShouldFail(() -> manager.resetForRetry()); - if (transaction == null) { - manager.close(); - break; - } - } - } - } - } - - @SuppressWarnings("resource") - @Test - public void transactionManagerBuffer() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - try (TransactionManager manager = client.transactionManager()) { - TransactionContext transaction = manager.begin(); - while (true) { - transaction.buffer(Mutation.newInsertBuilder("FOO").set("BAR").to(1L).build()); - try { - manager.commit(); - break; - } catch (AbortedException e) { - transaction = assertThrowsSessionNotFoundIfShouldFail(() -> manager.resetForRetry()); - if (transaction == null) { - break; - } - } - } - assertThat(manager.getCommitTimestamp()).isNotNull(); - assertThat(failOnInvalidatedSession).isFalse(); - } catch (SessionNotFoundException e) { - assertThat(failOnInvalidatedSession).isTrue(); - } - } - - @SuppressWarnings("resource") - @Test - public void transactionManagerSelectInvalidatedDuringTransaction() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - SessionPoolOptions.Builder builder = SessionPoolOptions.newBuilder(); - if (failOnInvalidatedSession) { - builder.setFailIfSessionNotFound(); - } - try (Spanner spanner = - SpannerOptions.newBuilder() - .setProjectId("[PROJECT]") - .setChannelProvider(channelProvider) - .setSessionPoolOption(builder.build()) - .setCredentials(NoCredentials.getInstance()) - .build() - .getService()) { - DatabaseClient client = - spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]")); - try (TransactionManager manager = client.transactionManager()) { - int attempts = 0; - TransactionContext transaction = manager.begin(); - while (true) { - attempts++; - try { - try (ResultSet rs = transaction.executeQuery(SELECT1AND2)) { - while (rs.next()) {} - } - if (attempts == 1) { - invalidateSessionPool( - client, spanner.getOptions().getSessionPoolOptions().getMinSessions()); - } - try (ResultSet rs = transaction.executeQuery(SELECT1AND2)) { - if (assertThrowsSessionNotFoundIfShouldFail(() -> rs.next()) == null) { - break; - } - } - manager.commit(); - assertThat(attempts).isGreaterThan(1); - break; - } catch (AbortedException e) { - transaction = assertThrowsSessionNotFoundIfShouldFail(() -> manager.resetForRetry()); - } - } - } - } - } - - @SuppressWarnings("resource") - @Test - public void transactionManagerReadInvalidatedDuringTransaction() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - SessionPoolOptions.Builder builder = SessionPoolOptions.newBuilder(); - if (failOnInvalidatedSession) { - builder.setFailIfSessionNotFound(); - } - try (Spanner spanner = - SpannerOptions.newBuilder() - .setProjectId("[PROJECT]") - .setChannelProvider(channelProvider) - .setSessionPoolOption(builder.build()) - .setCredentials(NoCredentials.getInstance()) - .build() - .getService()) { - DatabaseClient client = - spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]")); - try (TransactionManager manager = client.transactionManager()) { - int attempts = 0; - TransactionContext transaction = manager.begin(); - while (true) { - attempts++; - try { - try (ResultSet rs = - transaction.read("FOO", KeySet.all(), Collections.singletonList("BAR"))) { - while (rs.next()) {} - } - if (attempts == 1) { - invalidateSessionPool( - client, spanner.getOptions().getSessionPoolOptions().getMinSessions()); - } - try (ResultSet rs = - transaction.read("FOO", KeySet.all(), Collections.singletonList("BAR"))) { - if (assertThrowsSessionNotFoundIfShouldFail(() -> rs.next()) == null) { - break; - } - } - manager.commit(); - break; - } catch (AbortedException e) { - transaction = manager.resetForRetry(); - } - } - } - } - } - - @SuppressWarnings("resource") - @Test - public void transactionManagerReadUsingIndexInvalidatedDuringTransaction() - throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - SessionPoolOptions.Builder builder = SessionPoolOptions.newBuilder(); - if (failOnInvalidatedSession) { - builder.setFailIfSessionNotFound(); - } - try (Spanner spanner = - SpannerOptions.newBuilder() - .setProjectId("[PROJECT]") - .setChannelProvider(channelProvider) - .setSessionPoolOption(builder.build()) - .setCredentials(NoCredentials.getInstance()) - .build() - .getService()) { - DatabaseClient client = - spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]")); - try (TransactionManager manager = client.transactionManager()) { - int attempts = 0; - TransactionContext transaction = manager.begin(); - while (true) { - attempts++; - try { - try (ResultSet rs = - transaction.readUsingIndex( - "FOO", "IDX", KeySet.all(), Collections.singletonList("BAR"))) { - while (rs.next()) {} - } - if (attempts == 1) { - invalidateSessionPool( - client, spanner.getOptions().getSessionPoolOptions().getMinSessions()); - } - try (ResultSet rs = - transaction.readUsingIndex( - "FOO", "IDX", KeySet.all(), Collections.singletonList("BAR"))) { - if (assertThrowsSessionNotFoundIfShouldFail(() -> rs.next()) == null) { - break; - } - } - manager.commit(); - break; - } catch (AbortedException e) { - transaction = manager.resetForRetry(); - } - } - } - } - } - - @SuppressWarnings("resource") - @Test - public void transactionManagerReadRowInvalidatedDuringTransaction() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - SessionPoolOptions.Builder builder = SessionPoolOptions.newBuilder(); - if (failOnInvalidatedSession) { - builder.setFailIfSessionNotFound(); - } - try (Spanner spanner = - SpannerOptions.newBuilder() - .setProjectId("[PROJECT]") - .setChannelProvider(channelProvider) - .setSessionPoolOption(builder.build()) - .setCredentials(NoCredentials.getInstance()) - .build() - .getService()) { - DatabaseClient client = - spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]")); - try (TransactionManager manager = client.transactionManager()) { - int attempts = 0; - TransactionContext transaction = manager.begin(); - while (true) { - attempts++; - try { - Struct row = transaction.readRow("FOO", Key.of(), Collections.singletonList("BAR")); - assertThat(row.getLong(0)).isEqualTo(1L); - if (attempts == 1) { - invalidateSessionPool( - client, spanner.getOptions().getSessionPoolOptions().getMinSessions()); - } - TransactionContext context = transaction; - if (assertThrowsSessionNotFoundIfShouldFail( - () -> context.readRow("FOO", Key.of(), Collections.singletonList("BAR"))) - == null) { - break; - } - manager.commit(); - break; - } catch (AbortedException e) { - transaction = manager.resetForRetry(); - } - } - } - } - } - - @SuppressWarnings("resource") - @Test - public void transactionManagerReadRowUsingIndexInvalidatedDuringTransaction() - throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - SessionPoolOptions.Builder builder = SessionPoolOptions.newBuilder(); - if (failOnInvalidatedSession) { - builder.setFailIfSessionNotFound(); - } - try (Spanner spanner = - SpannerOptions.newBuilder() - .setProjectId("[PROJECT]") - .setChannelProvider(channelProvider) - .setSessionPoolOption(builder.build()) - .setCredentials(NoCredentials.getInstance()) - .build() - .getService()) { - DatabaseClient client = - spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]")); - try (TransactionManager manager = client.transactionManager()) { - int attempts = 0; - TransactionContext transaction = manager.begin(); - while (true) { - attempts++; - try { - Struct row = - transaction.readRowUsingIndex( - "FOO", "IDX", Key.of(), Collections.singletonList("BAR")); - assertThat(row.getLong(0)).isEqualTo(1L); - if (attempts == 1) { - invalidateSessionPool( - client, spanner.getOptions().getSessionPoolOptions().getMinSessions()); - } - TransactionContext context = transaction; - if (assertThrowsSessionNotFoundIfShouldFail( - () -> - context.readRowUsingIndex( - "FOO", "IDX", Key.of(), Collections.singletonList("BAR"))) - == null) { - break; - } - manager.commit(); - break; - } catch (AbortedException e) { - transaction = manager.resetForRetry(); - } - } - } - } - } - - @Test - public void partitionedDml() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionPartitionedOps()); - assertThrowsSessionNotFoundIfShouldFail( - () -> client.executePartitionedUpdate(UPDATE_STATEMENT)); - } - - @Test - public void write() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - assertThrowsSessionNotFoundIfShouldFail( - () -> client.write(Collections.singletonList(Mutation.delete("FOO", KeySet.all())))); - } - - @Test - public void writeAtLeastOnce() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSession()); - assertThrowsSessionNotFoundIfShouldFail( - () -> - client.writeAtLeastOnce( - Collections.singletonList(Mutation.delete("FOO", KeySet.all())))); - } - - @Test - public void asyncRunnerSelect() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - asyncRunner_withReadFunction(input -> input.executeQueryAsync(SELECT1AND2)); - } - - @Test - public void asyncRunnerRead() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - asyncRunner_withReadFunction( - input -> input.readAsync("FOO", KeySet.all(), Collections.singletonList("BAR"))); - } - - @Test - public void asyncRunnerReadUsingIndex() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - asyncRunner_withReadFunction( - input -> - input.readUsingIndexAsync( - "FOO", "IDX", KeySet.all(), Collections.singletonList("BAR"))); - } - - private void asyncRunner_withReadFunction( - final Function readFunction) throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - final ExecutorService queryExecutor = Executors.newSingleThreadExecutor(); - try { - AsyncRunner runner = client.runAsync(); - final AtomicLong counter = new AtomicLong(); - assertThrowsSessionNotFoundIfShouldFail( - () -> - get( - runner.runAsync( - txn -> { - AsyncResultSet rs = readFunction.apply(txn); - ApiFuture fut = - rs.setCallback( - queryExecutor, - resultSet -> { - while (true) { - switch (resultSet.tryNext()) { - case OK: - counter.incrementAndGet(); - break; - case DONE: - return CallbackResponse.DONE; - case NOT_READY: - return CallbackResponse.CONTINUE; - } - } - }); - return ApiFutures.transform( - fut, input -> counter.get(), MoreExecutors.directExecutor()); - }, - executor))); - } finally { - queryExecutor.shutdown(); - } - } - - @Test - public void asyncRunnerReadRow() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - AsyncRunner runner = client.runAsync(); - assertThrowsSessionNotFoundIfShouldFail( - () -> - get( - runner.runAsync( - txn -> txn.readRowAsync("FOO", Key.of(), Collections.singletonList("BAR")), - executor))); - } - - @Test - public void asyncRunnerReadRowUsingIndex() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - AsyncRunner runner = client.runAsync(); - assertThrowsSessionNotFoundIfShouldFail( - () -> - get( - runner.runAsync( - txn -> - txn.readRowUsingIndexAsync( - "FOO", "IDX", Key.of(), Collections.singletonList("BAR")), - executor))); - } - - @Test - public void asyncRunnerUpdate() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - AsyncRunner runner = client.runAsync(); - assertThrowsSessionNotFoundIfShouldFail( - () -> get(runner.runAsync(txn -> txn.executeUpdateAsync(UPDATE_STATEMENT), executor))); - } - - @Test - public void asyncRunnerBatchUpdate() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - AsyncRunner runner = client.runAsync(); - assertThrowsSessionNotFoundIfShouldFail( - () -> - get( - runner.runAsync( - txn -> txn.batchUpdateAsync(Arrays.asList(UPDATE_STATEMENT, UPDATE_STATEMENT)), - executor))); - } - - @Test - public void asyncRunnerBuffer() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - AsyncRunner runner = client.runAsync(); - assertThrowsSessionNotFoundIfShouldFail( - () -> - get( - runner.runAsync( - txn -> { - txn.buffer(Mutation.newInsertBuilder("FOO").set("BAR").to(1L).build()); - return ApiFutures.immediateFuture(null); - }, - executor))); - } - - @Test - public void asyncTransactionManagerAsyncSelect() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - asyncTransactionManager_readAsync(input -> input.executeQueryAsync(SELECT1AND2)); - } - - @Test - public void asyncTransactionManagerAsyncRead() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - asyncTransactionManager_readAsync( - input -> input.readAsync("FOO", KeySet.all(), Collections.singletonList("BAR"))); - } - - @Test - public void asyncTransactionManagerAsyncReadUsingIndex() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - asyncTransactionManager_readAsync( - input -> - input.readUsingIndexAsync( - "FOO", "idx", KeySet.all(), Collections.singletonList("BAR"))); - } - - private void asyncTransactionManager_readAsync( - final Function fn) throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - final ExecutorService queryExecutor = Executors.newSingleThreadExecutor(); - try (AsyncTransactionManager manager = client.transactionManagerAsync()) { - TransactionContextFuture context = manager.beginAsync(); - while (true) { - try { - final AtomicLong counter = new AtomicLong(); - AsyncTransactionStep count = - context.then( - (transaction, ignored) -> { - AsyncResultSet rs = fn.apply(transaction); - ApiFuture fut = - rs.setCallback( - queryExecutor, - resultSet -> { - while (true) { - switch (resultSet.tryNext()) { - case OK: - counter.incrementAndGet(); - break; - case DONE: - return CallbackResponse.DONE; - case NOT_READY: - return CallbackResponse.CONTINUE; - } - } - }); - return ApiFutures.transform( - fut, input -> counter.get(), MoreExecutors.directExecutor()); - }, - executor); - CommitTimestampFuture ts = count.commitAsync(); - assertThrowsSessionNotFoundIfShouldFail(() -> get(ts)); - break; - } catch (AbortedException e) { - context = manager.resetForRetryAsync(); - } - } - } finally { - queryExecutor.shutdown(); - } - } - - @Test - public void asyncTransactionManagerSelect() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - asyncTransactionManager_readSync(input -> input.executeQuery(SELECT1AND2)); - } - - @Test - public void asyncTransactionManagerRead() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - asyncTransactionManager_readSync( - input -> input.read("FOO", KeySet.all(), Collections.singletonList("BAR"))); - } - - @Test - public void asyncTransactionManagerReadUsingIndex() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - asyncTransactionManager_readSync( - input -> - input.readUsingIndex("FOO", "idx", KeySet.all(), Collections.singletonList("BAR"))); - } - - private void asyncTransactionManager_readSync(final Function fn) - throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - final ExecutorService queryExecutor = Executors.newSingleThreadExecutor(); - try (AsyncTransactionManager manager = client.transactionManagerAsync()) { - TransactionContextFuture context = manager.beginAsync(); - while (true) { - try { - AsyncTransactionStep count = - context.then( - (transaction, ignored) -> { - long counter = 0L; - try (ResultSet rs = fn.apply(transaction)) { - while (rs.next()) { - counter++; - } - } - return ApiFutures.immediateFuture(counter); - }, - executor); - CommitTimestampFuture ts = count.commitAsync(); - assertThrowsSessionNotFoundIfShouldFail(() -> get(ts)); - break; - } catch (AbortedException e) { - context = manager.resetForRetryAsync(); - } - } - } finally { - queryExecutor.shutdown(); - } - } - - @Test - public void asyncTransactionManagerReadRow() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - asyncTransactionManager_readRowFunction( - input -> - ApiFutures.immediateFuture( - input.readRow("FOO", Key.of("foo"), Collections.singletonList("BAR")))); - } - - @Test - public void asyncTransactionManagerReadRowUsingIndex() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - asyncTransactionManager_readRowFunction( - input -> - ApiFutures.immediateFuture( - input.readRowUsingIndex( - "FOO", "idx", Key.of("foo"), Collections.singletonList("BAR")))); - } - - @Test - public void asyncTransactionManagerReadRowAsync() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - asyncTransactionManager_readRowFunction( - input -> input.readRowAsync("FOO", Key.of("foo"), Collections.singletonList("BAR"))); - } - - @Test - public void asyncTransactionManagerReadRowUsingIndexAsync() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - asyncTransactionManager_readRowFunction( - input -> - input.readRowUsingIndexAsync( - "FOO", "idx", Key.of("foo"), Collections.singletonList("BAR"))); - } - - private void asyncTransactionManager_readRowFunction( - final Function> fn) throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - final ExecutorService queryExecutor = Executors.newSingleThreadExecutor(); - try (AsyncTransactionManager manager = client.transactionManagerAsync()) { - TransactionContextFuture context = manager.beginAsync(); - while (true) { - try { - AsyncTransactionStep row = - context.then((transaction, ignored) -> fn.apply(transaction), executor); - CommitTimestampFuture ts = row.commitAsync(); - assertThrowsSessionNotFoundIfShouldFail(() -> get(ts)); - break; - } catch (AbortedException e) { - context = manager.resetForRetryAsync(); - } - } - } finally { - queryExecutor.shutdown(); - } - } - - @Test - public void asyncTransactionManagerUpdateAsync() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - asyncTransactionManager_updateFunction( - input -> input.executeUpdateAsync(UPDATE_STATEMENT), UPDATE_COUNT); - } - - @Test - public void asyncTransactionManagerUpdate() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - asyncTransactionManager_updateFunction( - input -> ApiFutures.immediateFuture(input.executeUpdate(UPDATE_STATEMENT)), UPDATE_COUNT); - } - - @Test - public void asyncTransactionManagerBatchUpdateAsync() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - asyncTransactionManager_updateFunction( - input -> input.batchUpdateAsync(Arrays.asList(UPDATE_STATEMENT, UPDATE_STATEMENT)), - new long[] {UPDATE_COUNT, UPDATE_COUNT}); - } - - @Test - public void asyncTransactionManagerBatchUpdate() throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - asyncTransactionManager_updateFunction( - input -> - ApiFutures.immediateFuture( - input.batchUpdate(Arrays.asList(UPDATE_STATEMENT, UPDATE_STATEMENT))), - new long[] {UPDATE_COUNT, UPDATE_COUNT}); - } - - private void asyncTransactionManager_updateFunction( - final Function> fn, T expected) throws InterruptedException { - assumeFalse( - "Multiplexed session do not throw a SessionNotFound errors. ", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW()); - try (AsyncTransactionManager manager = client.transactionManagerAsync()) { - TransactionContextFuture transaction = manager.beginAsync(); - while (true) { - try { - AsyncTransactionStep res = - transaction.then((txn, input) -> fn.apply(txn), executor); - CommitTimestampFuture ts = res.commitAsync(); - assertThrowsSessionNotFoundIfShouldFail(() -> get(ts)); - break; - } catch (AbortedException e) { - transaction = manager.resetForRetryAsync(); - } - } - } - } -} diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/RetryableInternalErrorTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/RetryableInternalErrorTest.java index 3106bd16526..2e9d4185cb9 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/RetryableInternalErrorTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/RetryableInternalErrorTest.java @@ -23,7 +23,7 @@ import com.google.cloud.NoCredentials; import com.google.cloud.spanner.MockSpannerServiceImpl.SimulatedExecutionTime; import com.google.cloud.spanner.connection.AbstractMockServerTest; -import com.google.spanner.v1.BatchCreateSessionsRequest; +import com.google.spanner.v1.CreateSessionRequest; import com.google.spanner.v1.ExecuteSqlRequest; import io.grpc.ManagedChannelBuilder; import io.grpc.Status; @@ -36,7 +36,7 @@ public class RetryableInternalErrorTest extends AbstractMockServerTest { @Test public void testTranslateInternalException() { - mockSpanner.setBatchCreateSessionsExecutionTime( + mockSpanner.setCreateSessionExecutionTime( SimulatedExecutionTime.ofException( Status.INTERNAL .withDescription("Authentication backend internal server error. Please retry.") @@ -69,9 +69,9 @@ public void testTranslateInternalException() { assertTrue(resultSet.next()); assertFalse(resultSet.next()); } - // Verify that both the BatchCreateSessions call and the ExecuteStreamingSql call were + // Verify that both the CreateSession call and the ExecuteStreamingSql call were // retried. - assertEquals(2, mockSpanner.countRequestsOfType(BatchCreateSessionsRequest.class)); + assertEquals(2, mockSpanner.countRequestsOfType(CreateSessionRequest.class)); assertEquals(2, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)); // Clear the requests before the next test. mockSpanner.clearRequests(); diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SelectRandomBenchmark.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SelectRandomBenchmark.java index e18cddd3bf5..e1f93d334dc 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SelectRandomBenchmark.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SelectRandomBenchmark.java @@ -16,8 +16,6 @@ package com.google.cloud.spanner; -import static com.google.common.truth.Truth.assertThat; - import com.google.api.gax.rpc.TransportChannelProvider; import com.google.cloud.NoCredentials; import com.google.common.util.concurrent.Futures; @@ -99,8 +97,7 @@ public void setup() throws Exception { (DatabaseClientImpl) spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); // Wait until the session pool has initialized. - while (client.pool.getNumberOfSessionsInPool() - < spanner.getOptions().getSessionPoolOptions().getMinSessions()) { + while (client.multiplexedSessionDatabaseClient.getCurrentSessionReference() == null) { Thread.sleep(1L); } } @@ -119,8 +116,6 @@ public void burstRead(final BenchmarkState server) throws Exception { int parallelThreads = server.maxSessions * 2; final DatabaseClient client = server.spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); - SessionPool pool = ((DatabaseClientImpl) client).pool; - assertThat(pool.totalSessions()).isEqualTo(server.minSessions); ListeningScheduledExecutorService service = MoreExecutors.listeningDecorator(Executors.newScheduledThreadPool(parallelThreads)); diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolBenchmark.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolBenchmark.java deleted file mode 100644 index 4415ba7d707..00000000000 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolBenchmark.java +++ /dev/null @@ -1,265 +0,0 @@ -/* - * Copyright 2020 Google LLC - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package com.google.cloud.spanner; - -import static com.google.common.truth.Truth.assertThat; - -import com.google.api.gax.rpc.TransportChannelProvider; -import com.google.cloud.NoCredentials; -import com.google.common.util.concurrent.Futures; -import com.google.common.util.concurrent.ListenableFuture; -import com.google.common.util.concurrent.ListeningScheduledExecutorService; -import com.google.common.util.concurrent.MoreExecutors; -import com.google.spanner.v1.BatchCreateSessionsRequest; -import java.util.ArrayList; -import java.util.List; -import java.util.Random; -import java.util.concurrent.Executors; -import java.util.concurrent.TimeUnit; -import org.openjdk.jmh.annotations.AuxCounters; -import org.openjdk.jmh.annotations.Benchmark; -import org.openjdk.jmh.annotations.BenchmarkMode; -import org.openjdk.jmh.annotations.Fork; -import org.openjdk.jmh.annotations.Level; -import org.openjdk.jmh.annotations.Measurement; -import org.openjdk.jmh.annotations.Mode; -import org.openjdk.jmh.annotations.OutputTimeUnit; -import org.openjdk.jmh.annotations.Param; -import org.openjdk.jmh.annotations.Scope; -import org.openjdk.jmh.annotations.Setup; -import org.openjdk.jmh.annotations.State; -import org.openjdk.jmh.annotations.TearDown; -import org.openjdk.jmh.annotations.Warmup; - -/** - * Benchmarks for common session pool scenarios. The simulated execution times are based on - * reasonable estimates and are primarily intended to keep the benchmarks comparable with each other - * before and after changes have been made to the pool. The benchmarks are bound to the Maven - * profile `benchmark` and can be executed like this: - * mvn clean test -DskipTests -Pbenchmark -Dbenchmark.name=SessionPoolBenchmark - * - */ -@BenchmarkMode(Mode.AverageTime) -@Fork(value = 1, warmups = 0) -@Measurement(batchSize = 1, iterations = 1, timeUnit = TimeUnit.MILLISECONDS) -@Warmup(batchSize = 0, iterations = 0) -@OutputTimeUnit(TimeUnit.MILLISECONDS) -public class SessionPoolBenchmark { - private static final String TEST_PROJECT = "my-project"; - private static final String TEST_INSTANCE = "my-instance"; - private static final String TEST_DATABASE = "my-database"; - private static final int HOLD_SESSION_TIME = 100; - private static final int RND_WAIT_TIME_BETWEEN_REQUESTS = 10; - private static final Random RND = new Random(); - - @State(Scope.Thread) - @AuxCounters(org.openjdk.jmh.annotations.AuxCounters.Type.EVENTS) - public static class BenchmarkState { - private StandardBenchmarkMockServer mockServer; - private Spanner spanner; - private DatabaseClientImpl client; - - @Param({"100"}) - int minSessions; - - @Param({"400"}) - int maxSessions; - - @Param({"1", "10", "20", "25", "30", "40", "50", "100"}) - int incStep; - - @Param({"4"}) - int numChannels; - - @Param({"0.2"}) - float writeFraction; - - /** AuxCounter for number of RPCs. */ - public int numBatchCreateSessionsRpcs() { - return mockServer.countRequests(BatchCreateSessionsRequest.class); - } - - /** AuxCounter for number of sessions created. */ - public int sessionsCreated() { - return mockServer.getMockSpanner().numSessionsCreated(); - } - - @Setup(Level.Invocation) - public void setup() throws Exception { - mockServer = new StandardBenchmarkMockServer(); - TransportChannelProvider channelProvider = mockServer.start(); - - SpannerOptions options = - SpannerOptions.newBuilder() - .setProjectId(TEST_PROJECT) - .setChannelProvider(channelProvider) - .setNumChannels(numChannels) - .setCredentials(NoCredentials.getInstance()) - .setSessionPoolOption( - SessionPoolOptions.newBuilder() - .setMinSessions(minSessions) - .setMaxSessions(maxSessions) - .setIncStep(incStep) - .setWriteSessionsFraction(writeFraction) - .build()) - .build(); - - spanner = options.getService(); - client = - (DatabaseClientImpl) - spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); - // Wait until the session pool has initialized. - while (client.pool.getNumberOfSessionsInPool() - < spanner.getOptions().getSessionPoolOptions().getMinSessions()) { - Thread.sleep(1L); - } - } - - @TearDown(Level.Invocation) - public void teardown() throws Exception { - spanner.close(); - mockServer.shutdown(); - } - - int expectedStepsToMax() { - int remainder = (maxSessions - minSessions) % incStep == 0 ? 0 : 1; - return numChannels + ((maxSessions - minSessions) / incStep) + remainder; - } - } - - /** Measures the time needed to execute a burst of read requests. */ - @Benchmark - public void burstRead(final BenchmarkState server) throws Exception { - int totalQueries = server.maxSessions * 8; - int parallelThreads = server.maxSessions * 2; - final DatabaseClient client = - server.spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); - SessionPool pool = ((DatabaseClientImpl) client).pool; - assertThat(pool.totalSessions()).isEqualTo(server.minSessions); - - ListeningScheduledExecutorService service = - MoreExecutors.listeningDecorator(Executors.newScheduledThreadPool(parallelThreads)); - List> futures = new ArrayList<>(totalQueries); - for (int i = 0; i < totalQueries; i++) { - futures.add( - service.submit( - () -> { - Thread.sleep(RND.nextInt(RND_WAIT_TIME_BETWEEN_REQUESTS)); - try (ResultSet rs = - client.singleUse().executeQuery(StandardBenchmarkMockServer.SELECT1)) { - while (rs.next()) { - Thread.sleep(RND.nextInt(HOLD_SESSION_TIME)); - } - return null; - } - })); - } - Futures.allAsList(futures).get(); - service.shutdown(); - } - - /** Measures the time needed to execute a burst of write requests. */ - @Benchmark - public void burstWrite(final BenchmarkState server) throws Exception { - int totalWrites = server.maxSessions * 8; - int parallelThreads = server.maxSessions * 2; - final DatabaseClient client = - server.spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); - SessionPool pool = ((DatabaseClientImpl) client).pool; - assertThat(pool.totalSessions()).isEqualTo(server.minSessions); - - ListeningScheduledExecutorService service = - MoreExecutors.listeningDecorator(Executors.newScheduledThreadPool(parallelThreads)); - List> futures = new ArrayList<>(totalWrites); - for (int i = 0; i < totalWrites; i++) { - futures.add( - service.submit( - () -> { - Thread.sleep(RND.nextInt(RND_WAIT_TIME_BETWEEN_REQUESTS)); - TransactionRunner runner = client.readWriteTransaction(); - return runner.run( - transaction -> - transaction.executeUpdate(StandardBenchmarkMockServer.UPDATE_STATEMENT)); - })); - } - Futures.allAsList(futures).get(); - service.shutdown(); - } - - /** Measures the time needed to execute a burst of read and write requests. */ - @Benchmark - public void burstReadAndWrite(final BenchmarkState server) throws Exception { - int totalWrites = server.maxSessions * 4; - int totalReads = server.maxSessions * 4; - int parallelThreads = server.maxSessions * 2; - final DatabaseClient client = - server.spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); - SessionPool pool = ((DatabaseClientImpl) client).pool; - assertThat(pool.totalSessions()).isEqualTo(server.minSessions); - - ListeningScheduledExecutorService service = - MoreExecutors.listeningDecorator(Executors.newScheduledThreadPool(parallelThreads)); - List> futures = new ArrayList<>(totalReads + totalWrites); - for (int i = 0; i < totalWrites; i++) { - futures.add( - service.submit( - () -> { - Thread.sleep(RND.nextInt(RND_WAIT_TIME_BETWEEN_REQUESTS)); - TransactionRunner runner = client.readWriteTransaction(); - return runner.run( - transaction -> - transaction.executeUpdate(StandardBenchmarkMockServer.UPDATE_STATEMENT)); - })); - } - for (int i = 0; i < totalReads; i++) { - futures.add( - service.submit( - () -> { - Thread.sleep(RND.nextInt(RND_WAIT_TIME_BETWEEN_REQUESTS)); - try (ResultSet rs = - client.singleUse().executeQuery(StandardBenchmarkMockServer.SELECT1)) { - while (rs.next()) { - Thread.sleep(RND.nextInt(HOLD_SESSION_TIME)); - } - return null; - } - })); - } - Futures.allAsList(futures).get(); - service.shutdown(); - } - - /** Measures the time needed to acquire MaxSessions session sequentially. */ - @Benchmark - public void steadyIncrease(BenchmarkState server) { - final DatabaseClient client = - server.spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); - SessionPool pool = ((DatabaseClientImpl) client).pool; - assertThat(pool.totalSessions()).isEqualTo(server.minSessions); - - // Checkout maxSessions sessions by starting maxSessions read-only transactions sequentially. - List transactions = new ArrayList<>(server.maxSessions); - for (int i = 0; i < server.maxSessions; i++) { - ReadOnlyTransaction tx = client.readOnlyTransaction(); - tx.executeQuery(StandardBenchmarkMockServer.SELECT1); - transactions.add(tx); - } - for (ReadOnlyTransaction tx : transactions) { - tx.close(); - } - } -} diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolLeakTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolLeakTest.java deleted file mode 100644 index 080091e6615..00000000000 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolLeakTest.java +++ /dev/null @@ -1,252 +0,0 @@ -/* - * Copyright 2019 Google LLC - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * https://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package com.google.cloud.spanner; - -import static com.google.cloud.spanner.DisableDefaultMtlsProvider.disableDefaultMtlsProvider; -import static org.hamcrest.CoreMatchers.equalTo; -import static org.hamcrest.CoreMatchers.is; -import static org.hamcrest.MatcherAssert.assertThat; -import static org.junit.Assert.assertEquals; -import static org.junit.Assert.assertThrows; -import static org.junit.Assume.assumeFalse; - -import com.google.api.gax.grpc.testing.LocalChannelProvider; -import com.google.cloud.NoCredentials; -import com.google.cloud.spanner.MockSpannerServiceImpl.SimulatedExecutionTime; -import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult; -import com.google.cloud.spanner.SessionPool.LeakedSessionException; -import com.google.protobuf.ListValue; -import com.google.protobuf.Value; -import com.google.spanner.v1.ResultSetMetadata; -import com.google.spanner.v1.StructType; -import com.google.spanner.v1.StructType.Field; -import com.google.spanner.v1.Type; -import com.google.spanner.v1.TypeCode; -import io.grpc.Server; -import io.grpc.StatusRuntimeException; -import io.grpc.inprocess.InProcessServerBuilder; -import java.util.concurrent.ScheduledThreadPoolExecutor; -import org.junit.After; -import org.junit.AfterClass; -import org.junit.Before; -import org.junit.BeforeClass; -import org.junit.Test; -import org.junit.runner.RunWith; -import org.junit.runners.JUnit4; - -@RunWith(JUnit4.class) -public class SessionPoolLeakTest { - private static final StatusRuntimeException FAILED_PRECONDITION = - io.grpc.Status.FAILED_PRECONDITION - .withDescription("Non-retryable test exception.") - .asRuntimeException(); - private static MockSpannerServiceImpl mockSpanner; - private static Server server; - private static LocalChannelProvider channelProvider; - private Spanner spanner; - private DatabaseClient client; - private SessionPool pool; - - @BeforeClass - public static void startStaticServer() throws Exception { - disableDefaultMtlsProvider(); - mockSpanner = new MockSpannerServiceImpl(); - mockSpanner.setAbortProbability(0.0D); // We don't want any unpredictable aborted transactions. - String uniqueName = InProcessServerBuilder.generateName(); - server = - InProcessServerBuilder.forName(uniqueName) - .scheduledExecutorService(new ScheduledThreadPoolExecutor(1)) - .addService(mockSpanner) - .build() - .start(); - channelProvider = LocalChannelProvider.create(uniqueName); - } - - @AfterClass - public static void stopServer() throws InterruptedException { - server.shutdown(); - server.awaitTermination(); - } - - @Before - public void setUp() { - mockSpanner.reset(); - mockSpanner.removeAllExecutionTimes(); - SpannerOptions.Builder builder = - SpannerOptions.newBuilder() - .setProjectId("[PROJECT]") - .setChannelProvider(channelProvider) - .setCredentials(NoCredentials.getInstance()); - // Make sure the session pool is empty by default, does not contain any sessions, - // contains at most 2 sessions, and creates sessions in steps of 1. - builder.setSessionPoolOption( - SessionPoolOptions.newBuilder().setMinSessions(0).setMaxSessions(2).setIncStep(1).build()); - spanner = builder.build().getService(); - client = spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]")); - pool = ((DatabaseClientImpl) client).pool; - } - - @After - public void tearDown() { - spanner.close(); - } - - @Test - public void testIgnoreLeakedSession() { - for (boolean trackStackTraceofSessionCheckout : new boolean[] {true, false}) { - SessionPoolOptions sessionPoolOptions = - SessionPoolOptions.newBuilder() - .setMinSessions(0) - .setMaxSessions(2) - .setIncStep(1) - .setFailOnSessionLeak() - .setTrackStackTraceOfSessionCheckout(trackStackTraceofSessionCheckout) - .build(); - assumeFalse( - "Session Leaks do not occur with Multiplexed Sessions", - sessionPoolOptions.getUseMultiplexedSession()); - SpannerOptions.Builder builder = - SpannerOptions.newBuilder() - .setProjectId("[PROJECT]") - .setChannelProvider(channelProvider) - .setCredentials(NoCredentials.getInstance()); - builder.setSessionPoolOption(sessionPoolOptions); - Spanner spanner = builder.build().getService(); - DatabaseClient client = - spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]")); - mockSpanner.putStatementResult( - StatementResult.query( - Statement.of("SELECT 1"), - com.google.spanner.v1.ResultSet.newBuilder() - .setMetadata( - ResultSetMetadata.newBuilder() - .setRowType( - StructType.newBuilder() - .addFields( - Field.newBuilder() - .setName("c") - .setType( - Type.newBuilder().setCode(TypeCode.INT64).build()) - .build()) - .build()) - .build()) - .addRows( - ListValue.newBuilder() - .addValues(Value.newBuilder().setStringValue("1").build()) - .build()) - .build())); - - // Start a read-only transaction without closing it before closing the Spanner instance. - // This will cause a session leak. - ReadOnlyTransaction transaction = client.readOnlyTransaction(); - try (ResultSet resultSet = transaction.executeQuery(Statement.of("SELECT 1"))) { - //noinspection StatementWithEmptyBody - while (resultSet.next()) { - // ignore - } - } - LeakedSessionException exception = assertThrows(LeakedSessionException.class, spanner::close); - // The top of the stack trace will be "markCheckedOut" if we keep track of the point where the - // session was checked out, while it will be "closeAsync" if we don't. In the latter case, we - // get the stack trace of the method that tries to close the Spanner instance, while in the - // former the stack trace will contain the method that checked out the session. - assertEquals( - trackStackTraceofSessionCheckout ? "markCheckedOut" : "closeAsync", - exception.getStackTrace()[0].getMethodName()); - } - } - - @Test - public void testReadWriteTransactionExceptionOnCreateSession() { - assumeFalse( - "Session Leaks do not occur with Multiplexed Sessions", - isMultiplexedSessionsEnabledForRW()); - readWriteTransactionTest( - () -> - mockSpanner.setBatchCreateSessionsExecutionTime( - SimulatedExecutionTime.ofException(FAILED_PRECONDITION)), - 0); - } - - @Test - public void testReadWriteTransactionExceptionOnBegin() { - assumeFalse( - "Session Leaks do not occur with Multiplexed Sessions", - isMultiplexedSessionsEnabledForRW()); - readWriteTransactionTest( - () -> - mockSpanner.setBeginTransactionExecutionTime( - SimulatedExecutionTime.ofException(FAILED_PRECONDITION)), - 1); - } - - private void readWriteTransactionTest( - Runnable setup, int expectedNumberOfSessionsAfterExecution) { - assertEquals(0, pool.getNumberOfSessionsInPool()); - setup.run(); - SpannerException e = - assertThrows( - SpannerException.class, () -> client.readWriteTransaction().run(transaction -> null)); - assertEquals(ErrorCode.FAILED_PRECONDITION, e.getErrorCode()); - assertEquals(expectedNumberOfSessionsAfterExecution, pool.getNumberOfSessionsInPool()); - } - - @Test - public void testTransactionManagerExceptionOnCreateSession() { - assumeFalse( - "Session Leaks do not occur with Multiplexed Sessions", - isMultiplexedSessionsEnabledForRW()); - transactionManagerTest( - () -> - mockSpanner.setBatchCreateSessionsExecutionTime( - SimulatedExecutionTime.ofException(FAILED_PRECONDITION)), - 0); - } - - @Test - public void testTransactionManagerExceptionOnBegin() { - assumeFalse( - "Session Leaks do not occur with Multiplexed Sessions", - isMultiplexedSessionsEnabledForRW()); - assertThat(pool.getNumberOfSessionsInPool(), is(equalTo(0))); - mockSpanner.setBeginTransactionExecutionTime( - SimulatedExecutionTime.ofException(FAILED_PRECONDITION)); - try (TransactionManager txManager = client.transactionManager()) { - // This should not cause an error, as the actual BeginTransaction will be included with the - // first statement of the transaction. - txManager.begin(); - } - assertThat(pool.getNumberOfSessionsInPool(), is(equalTo(1))); - } - - private void transactionManagerTest(Runnable setup, int expectedNumberOfSessionsAfterExecution) { - assertEquals(0, pool.getNumberOfSessionsInPool()); - setup.run(); - try (TransactionManager txManager = client.transactionManager()) { - SpannerException e = assertThrows(SpannerException.class, txManager::begin); - assertEquals(ErrorCode.FAILED_PRECONDITION, e.getErrorCode()); - } - assertEquals(expectedNumberOfSessionsAfterExecution, pool.getNumberOfSessionsInPool()); - } - - private boolean isMultiplexedSessionsEnabledForRW() { - if (spanner.getOptions() == null || spanner.getOptions().getSessionPoolOptions() == null) { - return false; - } - return spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW(); - } -} diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolMaintainerBenchmark.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolMaintainerBenchmark.java deleted file mode 100644 index 0370f5420e2..00000000000 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolMaintainerBenchmark.java +++ /dev/null @@ -1,244 +0,0 @@ -/* - * Copyright 2020 Google LLC - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package com.google.cloud.spanner; - -import static com.google.common.truth.Truth.assertThat; - -import com.google.api.gax.rpc.TransportChannelProvider; -import com.google.cloud.NoCredentials; -import com.google.common.util.concurrent.Futures; -import com.google.common.util.concurrent.ListenableFuture; -import com.google.common.util.concurrent.ListeningScheduledExecutorService; -import com.google.common.util.concurrent.MoreExecutors; -import com.google.spanner.v1.BatchCreateSessionsRequest; -import com.google.spanner.v1.BeginTransactionRequest; -import com.google.spanner.v1.DeleteSessionRequest; -import java.time.Duration; -import java.util.ArrayList; -import java.util.List; -import java.util.Random; -import java.util.concurrent.Executors; -import java.util.concurrent.TimeUnit; -import org.openjdk.jmh.annotations.AuxCounters; -import org.openjdk.jmh.annotations.Benchmark; -import org.openjdk.jmh.annotations.BenchmarkMode; -import org.openjdk.jmh.annotations.Fork; -import org.openjdk.jmh.annotations.Level; -import org.openjdk.jmh.annotations.Measurement; -import org.openjdk.jmh.annotations.Mode; -import org.openjdk.jmh.annotations.OutputTimeUnit; -import org.openjdk.jmh.annotations.Param; -import org.openjdk.jmh.annotations.Scope; -import org.openjdk.jmh.annotations.Setup; -import org.openjdk.jmh.annotations.State; -import org.openjdk.jmh.annotations.TearDown; -import org.openjdk.jmh.annotations.Warmup; - -/** - * Benchmarks for the SessionPoolMaintainer. Run these benchmarks from the command line like this: - * - * mvn clean test -DskipTests -Pbenchmark -Dbenchmark.name=SessionPoolMaintainerBenchmark - * - */ -@BenchmarkMode(Mode.AverageTime) -@Fork(value = 1, warmups = 0) -@Measurement(batchSize = 1, iterations = 1, timeUnit = TimeUnit.MILLISECONDS) -@Warmup(batchSize = 0, iterations = 0) -@OutputTimeUnit(TimeUnit.MILLISECONDS) -public class SessionPoolMaintainerBenchmark { - private static final String TEST_PROJECT = "my-project"; - private static final String TEST_INSTANCE = "my-instance"; - private static final String TEST_DATABASE = "my-database"; - private static final int HOLD_SESSION_TIME = 10; - private static final int RND_WAIT_TIME_BETWEEN_REQUESTS = 100; - private static final Random RND = new Random(); - - @State(Scope.Thread) - @AuxCounters(org.openjdk.jmh.annotations.AuxCounters.Type.EVENTS) - public static class MockServer { - private StandardBenchmarkMockServer mockServer; - private Spanner spanner; - private DatabaseClientImpl client; - - /** - * The tests set the session idle timeout to an extremely low value to force timeouts and - * sessions to be evicted from the pool. This is not intended to replicate a realistic scenario, - * only to detect whether certain changes to the client library might cause the number of RPCs - * or the execution time to change drastically. - */ - @Param({"100"}) - long idleTimeout; - - /** AuxCounter for number of create RPCs. */ - public int numBatchCreateSessionsRpcs() { - return mockServer.countRequests(BatchCreateSessionsRequest.class); - } - - /** AuxCounter for number of delete RPCs. */ - public int numDeleteSessionRpcs() { - return mockServer.countRequests(DeleteSessionRequest.class); - } - - /** AuxCounter for number of begin tx RPCs. */ - public int numBeginTransactionRpcs() { - return mockServer.countRequests(BeginTransactionRequest.class); - } - - @Setup(Level.Invocation) - public void setup() throws Exception { - mockServer = new StandardBenchmarkMockServer(); - TransportChannelProvider channelProvider = mockServer.start(); - - SpannerOptions options = - SpannerOptions.newBuilder() - .setProjectId(TEST_PROJECT) - .setChannelProvider(channelProvider) - .setCredentials(NoCredentials.getInstance()) - .setSessionPoolOption( - SessionPoolOptions.newBuilder() - // Set idle timeout and loop frequency to very low values. - .setRemoveInactiveSessionAfterDuration(Duration.ofMillis(idleTimeout)) - .setLoopFrequency(idleTimeout / 10) - .build()) - .build(); - - spanner = options.getService(); - client = - (DatabaseClientImpl) - spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); - // Wait until the session pool has initialized. - while (client.pool.getNumberOfSessionsInPool() - < spanner.getOptions().getSessionPoolOptions().getMinSessions()) { - Thread.sleep(1L); - } - } - - @TearDown(Level.Invocation) - public void teardown() throws Exception { - spanner.close(); - mockServer.shutdown(); - } - } - - /** Measures the time and RPCs needed to execute read requests. */ - @Benchmark - public void read(final MockServer server) throws Exception { - int min = server.spanner.getOptions().getSessionPoolOptions().getMinSessions(); - int max = server.spanner.getOptions().getSessionPoolOptions().getMaxSessions(); - int totalQueries = max * 4; - int parallelThreads = min; - final DatabaseClient client = - server.spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); - SessionPool pool = ((DatabaseClientImpl) client).pool; - assertThat(pool.totalSessions()).isEqualTo(min); - - ListeningScheduledExecutorService service = - MoreExecutors.listeningDecorator(Executors.newScheduledThreadPool(parallelThreads)); - List> futures = new ArrayList<>(totalQueries); - for (int i = 0; i < totalQueries; i++) { - futures.add( - service.submit( - () -> { - Thread.sleep(RND.nextInt(RND_WAIT_TIME_BETWEEN_REQUESTS)); - try (ResultSet rs = - client.singleUse().executeQuery(StandardBenchmarkMockServer.SELECT1)) { - while (rs.next()) { - Thread.sleep(RND.nextInt(HOLD_SESSION_TIME)); - } - return null; - } - })); - } - Futures.allAsList(futures).get(); - service.shutdown(); - } - - /** Measures the time and RPCs needed to execute write requests. */ - @Benchmark - public void write(final MockServer server) throws Exception { - int min = server.spanner.getOptions().getSessionPoolOptions().getMinSessions(); - int max = server.spanner.getOptions().getSessionPoolOptions().getMaxSessions(); - int totalWrites = max * 4; - int parallelThreads = max; - final DatabaseClient client = - server.spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); - SessionPool pool = ((DatabaseClientImpl) client).pool; - assertThat(pool.totalSessions()).isEqualTo(min); - - ListeningScheduledExecutorService service = - MoreExecutors.listeningDecorator(Executors.newScheduledThreadPool(parallelThreads)); - List> futures = new ArrayList<>(totalWrites); - for (int i = 0; i < totalWrites; i++) { - futures.add( - service.submit( - () -> { - Thread.sleep(RND.nextInt(RND_WAIT_TIME_BETWEEN_REQUESTS)); - TransactionRunner runner = client.readWriteTransaction(); - return runner.run( - transaction -> - transaction.executeUpdate(StandardBenchmarkMockServer.UPDATE_STATEMENT)); - })); - } - Futures.allAsList(futures).get(); - service.shutdown(); - } - - /** Measures the time and RPCs needed to execute read and write requests. */ - @Benchmark - public void readAndWrite(final MockServer server) throws Exception { - int min = server.spanner.getOptions().getSessionPoolOptions().getMinSessions(); - int max = server.spanner.getOptions().getSessionPoolOptions().getMaxSessions(); - int totalWrites = max * 2; - int totalReads = max * 2; - int parallelThreads = max; - final DatabaseClient client = - server.spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); - SessionPool pool = ((DatabaseClientImpl) client).pool; - assertThat(pool.totalSessions()).isEqualTo(min); - - ListeningScheduledExecutorService service = - MoreExecutors.listeningDecorator(Executors.newScheduledThreadPool(parallelThreads)); - List> futures = new ArrayList<>(totalReads + totalWrites); - for (int i = 0; i < totalWrites; i++) { - futures.add( - service.submit( - () -> { - Thread.sleep(RND.nextInt(RND_WAIT_TIME_BETWEEN_REQUESTS)); - TransactionRunner runner = client.readWriteTransaction(); - return runner.run( - transaction -> - transaction.executeUpdate(StandardBenchmarkMockServer.UPDATE_STATEMENT)); - })); - } - for (int i = 0; i < totalReads; i++) { - futures.add( - service.submit( - () -> { - Thread.sleep(RND.nextInt(RND_WAIT_TIME_BETWEEN_REQUESTS)); - try (ResultSet rs = - client.singleUse().executeQuery(StandardBenchmarkMockServer.SELECT1)) { - while (rs.next()) { - Thread.sleep(RND.nextInt(HOLD_SESSION_TIME)); - } - return null; - } - })); - } - Futures.allAsList(futures).get(); - service.shutdown(); - } -} diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolMaintainerMockServerTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolMaintainerMockServerTest.java deleted file mode 100644 index 99a773eeb0f..00000000000 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolMaintainerMockServerTest.java +++ /dev/null @@ -1,181 +0,0 @@ -/* - * Copyright 2024 Google LLC - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package com.google.cloud.spanner; - -import static org.junit.Assert.assertEquals; -import static org.junit.Assert.assertFalse; -import static org.junit.Assert.assertTrue; -import static org.junit.Assume.assumeFalse; - -import com.google.cloud.NoCredentials; -import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult; -import com.google.common.base.Stopwatch; -import com.google.protobuf.ListValue; -import com.google.protobuf.Value; -import com.google.spanner.v1.BatchCreateSessionsRequest; -import com.google.spanner.v1.ExecuteSqlRequest; -import com.google.spanner.v1.ResultSetMetadata; -import com.google.spanner.v1.StructType; -import com.google.spanner.v1.StructType.Field; -import com.google.spanner.v1.Type; -import com.google.spanner.v1.TypeCode; -import java.time.Duration; -import java.util.concurrent.TimeUnit; -import org.junit.Before; -import org.junit.BeforeClass; -import org.junit.Test; -import org.junit.runner.RunWith; -import org.junit.runners.JUnit4; - -@RunWith(JUnit4.class) -public class SessionPoolMaintainerMockServerTest extends AbstractMockServerTest { - private final FakeClock clock = new FakeClock(); - - @BeforeClass - public static void setupResults() { - mockSpanner.putStatementResult( - StatementResult.query( - Statement.of("SELECT 1"), - com.google.spanner.v1.ResultSet.newBuilder() - .setMetadata( - ResultSetMetadata.newBuilder() - .setRowType( - StructType.newBuilder() - .addFields( - Field.newBuilder() - .setName("C") - .setType(Type.newBuilder().setCode(TypeCode.INT64).build()) - .build()) - .build()) - .build()) - .addRows( - ListValue.newBuilder() - .addValues(Value.newBuilder().setStringValue("1").build()) - .build()) - .build())); - } - - @Before - public void createSpannerInstance() { - clock.currentTimeMillis.set(System.currentTimeMillis()); - spanner = - SpannerOptions.newBuilder() - .setProjectId("p") - .setChannelProvider(channelProvider) - .setCredentials(NoCredentials.getInstance()) - .setSessionPoolOption( - SessionPoolOptions.newBuilder() - .setPoolMaintainerClock(clock) - .setWaitForMinSessionsDuration(Duration.ofSeconds(10L)) - .setFailOnSessionLeak() - .build()) - .build() - .getService(); - } - - @Test - public void testMaintain() { - int minSessions = spanner.getOptions().getSessionPoolOptions().getMinSessions(); - DatabaseClientImpl client = - (DatabaseClientImpl) spanner.getDatabaseClient(DatabaseId.of("p", "i", "d")); - assertEquals(minSessions, mockSpanner.getSessions().size()); - assertEquals(0, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)); - clock.currentTimeMillis.addAndGet(Duration.ofMinutes(35).toMillis()); - client.pool.poolMaintainer.maintainPool(); - assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)); - client.pool.poolMaintainer.maintainPool(); - assertEquals(2, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)); - clock.currentTimeMillis.addAndGet(Duration.ofMinutes(21).toMillis()); - - // Most sessions are considered idle and are removed. Freeze the mock Spanner server to prevent - // the replenish action to fill the pool again before we check the number of sessions in the - // pool. - mockSpanner.freeze(); - client.pool.poolMaintainer.maintainPool(); - assertEquals(2, client.pool.totalSessions()); - mockSpanner.unfreeze(); - - // The pool should be replenished. - client.pool.poolMaintainer.maintainPool(); - assertEquals(minSessions, client.pool.getTotalSessionsPlusNumSessionsBeingCreated()); - Stopwatch watch = Stopwatch.createStarted(); - //noinspection StatementWithEmptyBody - while (client.pool.totalSessions() < minSessions - && watch.elapsed(TimeUnit.MILLISECONDS) - < spanner.getOptions().getSessionPoolOptions().getWaitForMinSessions().toMillis()) { - // wait for the pool to be replenished. - } - assertEquals(minSessions, client.pool.totalSessions()); - } - - @Test - public void testSessionNotFoundIsRetried() { - assumeFalse( - "Session not found errors are not relevant for multiplexed sessions", - spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSession()); - - int minSessions = spanner.getOptions().getSessionPoolOptions().getMinSessions(); - DatabaseClientImpl client = - (DatabaseClientImpl) spanner.getDatabaseClient(DatabaseId.of("p", "i", "d")); - assertEquals(minSessions, mockSpanner.getSessions().size()); - - // Remove all sessions from the backend. - mockSpanner.getSessions().clear(); - - // Sessions have been removed from the backend, but this will still succeed, as Session not - // found errors are retried by the client. - try (ResultSet resultSet = client.singleUse().executeQuery(Statement.of("SELECT 1"))) { - assertTrue(resultSet.next()); - assertEquals(1L, resultSet.getLong(0)); - assertFalse(resultSet.next()); - } - - int numRequests = mockSpanner.countRequestsOfType(ExecuteSqlRequest.class); - assertTrue( - String.format("Number of requests should be larger than 1, but was %d", numRequests), - numRequests > 1); - } - - @Test - public void testMaintainerReplenishesPoolIfAllAreInvalid() { - int minSessions = spanner.getOptions().getSessionPoolOptions().getMinSessions(); - DatabaseClientImpl client = - (DatabaseClientImpl) spanner.getDatabaseClient(DatabaseId.of("p", "i", "d")); - assertEquals(minSessions, mockSpanner.getSessions().size()); - - // Remove all sessions from the backend. - mockSpanner.getSessions().clear(); - // Advance the clock of the maintainer to mark all sessions are eligible for maintenance. - clock.currentTimeMillis.addAndGet(Duration.ofMinutes(35).toMillis()); - // Run the maintainer. This will ping one session, which again will cause it to be replaced. - client.pool.poolMaintainer.maintainPool(); - assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)); - - // The session will be replaced using a single BatchCreateSessions call. - Stopwatch watch = Stopwatch.createStarted(); - //noinspection StatementWithEmptyBody - while (client.pool.totalSessions() < minSessions - && watch.elapsed(TimeUnit.MILLISECONDS) - < spanner.getOptions().getSessionPoolOptions().getWaitForMinSessions().toMillis()) { - // wait for the pool to be replenished. - } - assertEquals(minSessions, client.pool.totalSessions()); - assertEquals( - spanner.getOptions().getNumChannels() + 1, - mockSpanner.countRequestsOfType(BatchCreateSessionsRequest.class)); - } -} diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolMaintainerTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolMaintainerTest.java deleted file mode 100644 index 276a3ac813f..00000000000 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolMaintainerTest.java +++ /dev/null @@ -1,398 +0,0 @@ -/* - * Copyright 2020 Google LLC - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package com.google.cloud.spanner; - -import static com.google.common.truth.Truth.assertThat; -import static org.junit.Assert.assertEquals; -import static org.junit.Assert.assertFalse; -import static org.junit.Assert.assertTrue; -import static org.mockito.Mockito.any; -import static org.mockito.Mockito.doAnswer; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.when; -import static org.mockito.MockitoAnnotations.initMocks; - -import com.google.cloud.spanner.SessionClient.SessionConsumer; -import com.google.cloud.spanner.SessionPool.PooledSession; -import com.google.cloud.spanner.SessionPool.PooledSessionFuture; -import com.google.cloud.spanner.SessionPool.Position; -import com.google.cloud.spanner.SessionPool.SessionConsumerImpl; -import com.google.common.base.Preconditions; -import io.opencensus.trace.Tracing; -import io.opentelemetry.api.OpenTelemetry; -import java.util.ArrayList; -import java.util.HashMap; -import java.util.List; -import java.util.Map; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.Executors; -import java.util.concurrent.TimeUnit; -import org.junit.Before; -import org.junit.Test; -import org.junit.runner.RunWith; -import org.junit.runners.JUnit4; -import org.mockito.Mock; -import org.mockito.Mockito; - -@RunWith(JUnit4.class) -public class SessionPoolMaintainerTest extends BaseSessionPoolTest { - private ExecutorService executor = Executors.newSingleThreadExecutor(); - private @Mock SpannerImpl client; - private @Mock SessionClient sessionClient; - private @Mock SpannerOptions spannerOptions; - private DatabaseId db = DatabaseId.of("projects/p/instances/i/databases/unused"); - private SessionPoolOptions options; - private FakeClock clock = new FakeClock(); - private List idledSessions = new ArrayList<>(); - private Map pingedSessions = new HashMap<>(); - - @Before - public void setUp() { - initMocks(this); - when(client.getOptions()).thenReturn(spannerOptions); - when(client.getSessionClient(db)).thenReturn(sessionClient); - when(sessionClient.getSpanner()).thenReturn(client); - when(spannerOptions.getNumChannels()).thenReturn(4); - when(spannerOptions.getDatabaseRole()).thenReturn("role"); - setupMockSessionCreation(); - options = - SessionPoolOptions.newBuilder() - .setMinSessions(1) - .setMaxIdleSessions(1) - .setMaxSessions(5) - .setIncStep(1) - .setKeepAliveIntervalMinutes(2) - .setPoolMaintainerClock(clock) - .build(); - when(spannerOptions.getSessionPoolOptions()).thenReturn(options); - idledSessions.clear(); - pingedSessions.clear(); - } - - private void setupMockSessionCreation() { - doAnswer( - invocation -> { - executor.submit( - () -> { - int sessionCount = invocation.getArgument(0, Integer.class); - SessionConsumerImpl consumer = - invocation.getArgument(2, SessionConsumerImpl.class); - for (int i = 0; i < sessionCount; i++) { - ReadContext mockContext = mock(ReadContext.class); - consumer.onSessionReady( - setupMockSession(buildMockSession(client, mockContext), mockContext)); - } - }); - return null; - }) - .when(sessionClient) - .asyncBatchCreateSessions( - Mockito.anyInt(), Mockito.anyBoolean(), any(SessionConsumer.class)); - } - - private SessionImpl setupMockSession(final SessionImpl session, final ReadContext mockContext) { - final ResultSet mockResult = mock(ResultSet.class); - when(mockContext.executeQuery(any(Statement.class))) - .thenAnswer( - invocation -> { - Integer currentValue = pingedSessions.get(session.getName()); - if (currentValue == null) { - currentValue = 0; - } - pingedSessions.put(session.getName(), ++currentValue); - return mockResult; - }); - when(mockResult.next()).thenReturn(true); - return session; - } - - private SessionPool createPool() throws Exception { - return createPool(this.options); - } - - private SessionPool createPool(SessionPoolOptions options) throws Exception { - // Allow sessions to be added to the head of the pool in all cases in this test, as it is - // otherwise impossible to know which session exactly is getting pinged at what point in time. - SessionPool pool = - SessionPool.createPool( - options, - new TestExecutorFactory(), - client.getSessionClient(db), - clock, - Position.FIRST, - new TraceWrapper(Tracing.getTracer(), OpenTelemetry.noop().getTracer(""), false), - OpenTelemetry.noop()); - pool.idleSessionRemovedListener = - input -> { - idledSessions.add(input); - return null; - }; - // Wait until pool has initialized. - while (pool.totalSessions() < options.getMinSessions()) { - Thread.sleep(1L); - } - return pool; - } - - @Test - public void testKeepAlive() throws Exception { - SessionPool pool = createPool(); - assertThat(pingedSessions).isEmpty(); - // Run one maintenance loop. No sessions should get a keep-alive ping. - runMaintenanceLoop(clock, pool, 1); - assertThat(pingedSessions).isEmpty(); - - // Checkout two sessions and do a maintenance loop. Still no sessions should be getting any - // pings. - Session session1 = pool.getSession(); - Session session2 = pool.getSession(); - runMaintenanceLoop(clock, pool, 1); - assertThat(pingedSessions).isEmpty(); - - // Check the sessions back into the pool and do a maintenance loop. - session2.close(); - session1.close(); - runMaintenanceLoop(clock, pool, 1); - assertThat(pingedSessions).isEmpty(); - - // Now advance the time enough for both sessions in the pool to be idled. Then do one - // maintenance loop. This should cause the last session to have been checked back into the pool - // to get a ping, but not the second session. - clock.currentTimeMillis.addAndGet( - TimeUnit.MINUTES.toMillis(options.getKeepAliveIntervalMinutes()) + 1); - runMaintenanceLoop(clock, pool, 1); - assertThat(pingedSessions).containsExactly(session1.getName(), 1); - // Do another maintenance loop. This should cause the other session to also get a ping. - runMaintenanceLoop(clock, pool, 1); - assertThat(pingedSessions).containsExactly(session1.getName(), 1, session2.getName(), 1); - - // Now check out three sessions so the pool will create an additional session. The pool will - // only keep 2 sessions alive, as that is the setting for MinSessions. - Session session3 = pool.getSession(); - Session session4 = pool.getSession(); - Session session5 = pool.getSession(); - // Pinging a session will put it at the back of the pool. A session that needed a ping to be - // kept alive is not one that should be preferred for use. This means that session2 is the last - // session in the pool, and session1 the second-to-last. - assertEquals(session1.getName(), session3.getName()); - assertEquals(session2.getName(), session4.getName()); - session5.close(); - session4.close(); - session3.close(); - // Advance the clock to force pings for the sessions in the pool and do three maintenance loops. - // This should ping the sessions in the following order: - // 1. session3 (=session1) - // 2. session4 (=session2) - // The pinged sessions already contains: {session1: 1, session2: 1} - // Note that the pool only pings up to MinSessions sessions. - clock.currentTimeMillis.addAndGet( - TimeUnit.MINUTES.toMillis(options.getKeepAliveIntervalMinutes()) + 1); - runMaintenanceLoop(clock, pool, 3); - assertThat(pingedSessions).containsExactly(session1.getName(), 2, session2.getName(), 2); - - // Advance the clock to idle all sessions in the pool again and then check out one session. This - // should cause only one session to get a ping. - clock.currentTimeMillis.addAndGet( - TimeUnit.MINUTES.toMillis(options.getKeepAliveIntervalMinutes()) + 1); - // This will be session1, as all sessions were pinged in the previous 3 maintenance loops, and - // this will have brought session1 back to the front of the pool. - Session session6 = pool.getSession(); - // The session that was first in the pool now is equal to the initial first session as each full - // round of pings will swap the order of the first MinSessions sessions in the pool. - assertThat(session6.getName()).isEqualTo(session1.getName()); - runMaintenanceLoop(clock, pool, 3); - // Running 3 cycles will only ping the 2 sessions in the pool once. - assertThat(pool.totalSessions()).isEqualTo(3); - assertThat(pingedSessions).containsExactly(session1.getName(), 2, session2.getName(), 3); - // Update the last use date and release the session to the pool and do another maintenance - // cycle. This should not ping any sessions. - ((PooledSessionFuture) session6).get().markUsed(); - session6.close(); - runMaintenanceLoop(clock, pool, 3); - assertThat(pingedSessions).containsExactly(session1.getName(), 2, session2.getName(), 3); - - // Now check out 3 sessions again and make sure the 'extra' session is checked in last. That - // will make it eligible for pings. - Session session7 = pool.getSession(); - Session session8 = pool.getSession(); - Session session9 = pool.getSession(); - - assertThat(session7.getName()).isEqualTo(session1.getName()); - assertThat(session8.getName()).isEqualTo(session2.getName()); - assertThat(session9.getName()).isEqualTo(session5.getName()); - - session7.close(); - session8.close(); - session9.close(); - - clock.currentTimeMillis.addAndGet( - TimeUnit.MINUTES.toMillis(options.getKeepAliveIntervalMinutes()) + 1); - runMaintenanceLoop(clock, pool, 3); - // session1 will not get a ping this time, as it was checked in first and is now the last - // session in the pool. - assertThat(pingedSessions) - .containsExactly(session1.getName(), 2, session2.getName(), 4, session5.getName(), 1); - } - - @Test - public void testIdleSessions() throws Exception { - SessionPool pool = createPool(); - long loopsToIdleSessions = - Double.valueOf( - Math.ceil( - (double) options.getRemoveInactiveSessionAfter().toMillis() - / pool.poolMaintainer.loopFrequency)) - .longValue() - + 2L; - assertThat(idledSessions).isEmpty(); - // Run one maintenance loop. No sessions should be removed from the pool. - runMaintenanceLoop(clock, pool, 1); - assertThat(idledSessions).isEmpty(); - - // Checkout two sessions and do a maintenance loop. Still no sessions should be removed. - Session session1 = pool.getSession(); - Session session2 = pool.getSession(); - runMaintenanceLoop(clock, pool, 1); - assertThat(idledSessions).isEmpty(); - - // Check the sessions back into the pool and do a maintenance loop. - session2.close(); - session1.close(); - runMaintenanceLoop(clock, pool, 1); - assertThat(idledSessions).isEmpty(); - - // Now advance the time enough for both sessions in the pool to be idled. Both sessions should - // be kept alive by the maintainer and remain in the pool. - runMaintenanceLoop(clock, pool, loopsToIdleSessions); - assertThat(idledSessions).isEmpty(); - - // Now check out three sessions so the pool will create an additional session. The pool will - // only keep 2 sessions alive, as that is the setting for MinSessions. - Session session3 = pool.getSession().get(); - Session session4 = pool.getSession().get(); - Session session5 = pool.getSession().get(); - // Note that pinging sessions does not change the order of the pool. This means that session2 - // is still the last session in the pool. - assertThat(session3.getName()).isEqualTo(session1.getName()); - assertThat(session4.getName()).isEqualTo(session2.getName()); - session5.close(); - session4.close(); - session3.close(); - // Advance the clock to idle sessions. The pool will keep session4 and session3 alive, session5 - // will be idled and removed. - runMaintenanceLoop(clock, pool, loopsToIdleSessions); - assertThat(idledSessions).containsExactly(session5); - assertThat(pool.totalSessions()).isEqualTo(2); - - // Check out three sessions again and keep one session checked out. - Session session6 = pool.getSession().get(); - Session session7 = pool.getSession().get(); - Session session8 = pool.getSession().get(); - session8.close(); - session7.close(); - // Now advance the clock to idle sessions. This should remove session8 from the pool. - runMaintenanceLoop(clock, pool, loopsToIdleSessions); - assertThat(idledSessions).containsExactly(session5, session8); - assertThat(pool.totalSessions()).isEqualTo(2); - ((PooledSession) session6).markUsed(); - session6.close(); - - // Check out three sessions and keep them all checked out. No sessions should be removed from - // the pool. - Session session9 = pool.getSession().get(); - Session session10 = pool.getSession().get(); - Session session11 = pool.getSession().get(); - runMaintenanceLoop(clock, pool, loopsToIdleSessions); - assertThat(idledSessions).containsExactly(session5, session8); - assertThat(pool.totalSessions()).isEqualTo(3); - // Return the sessions to the pool. As they have not been used, they are all into idle time. - // Running the maintainer will now remove all the sessions from the pool and then start the - // replenish method. - session9.close(); - session10.close(); - session11.close(); - runMaintenanceLoop(clock, pool, 1); - assertThat(idledSessions).containsExactly(session5, session8, session9, session10, session11); - // Check that the pool is replenished. - while (pool.totalSessions() < options.getMinSessions()) { - Thread.sleep(1L); - } - assertThat(pool.totalSessions()).isEqualTo(options.getMinSessions()); - } - - @Test - public void testRandomizeThreshold() throws Exception { - SessionPool pool = - createPool( - this.options.toBuilder() - .setMaxSessions(400) - .setLoopFrequency(1000L) - .setRandomizePositionQPSThreshold(4) - .build()); - List sessions; - - // Run a maintenance loop. No sessions have been checked out so far, so the TPS should be 0. - runMaintenanceLoop(clock, pool, 1); - assertFalse(pool.shouldRandomize()); - - // Get and return one session. This means TPS == 1. - returnSessions(1, useSessions(1, pool)); - runMaintenanceLoop(clock, pool, 1); - assertFalse(pool.shouldRandomize()); - - // Get and return four sessions. This means TPS == 4, and that no sessions are checked out. - returnSessions(4, useSessions(4, pool)); - runMaintenanceLoop(clock, pool, 1); - assertFalse(pool.shouldRandomize()); - - // Get four sessions without returning them. - // This means TPS == 4 and that they are all still checked out. - sessions = useSessions(4, pool); - runMaintenanceLoop(clock, pool, 1); - assertTrue(pool.shouldRandomize()); - // Returning one of the sessions reduces the number of checked out sessions enough to stop the - // randomizing. - returnSessions(1, sessions); - runMaintenanceLoop(clock, pool, 1); - assertFalse(pool.shouldRandomize()); - - // Get three more session and run the maintenance loop. - // The TPS is then 3, as we've only gotten 3 sessions since the last maintenance run. - // That means that we should not randomize. - sessions.addAll(useSessions(3, pool)); - runMaintenanceLoop(clock, pool, 1); - assertFalse(pool.shouldRandomize()); - - returnSessions(sessions.size(), sessions); - } - - private List useSessions(int numSessions, SessionPool pool) { - List sessions = new ArrayList<>(numSessions); - for (int i = 0; i < numSessions; i++) { - sessions.add(pool.getSession()); - sessions.get(sessions.size() - 1).singleUse().executeQuery(Statement.of("SELECT 1")).next(); - } - return sessions; - } - - private void returnSessions(int numSessions, List sessions) { - Preconditions.checkArgument(numSessions <= sessions.size()); - for (int i = 0; i < numSessions; i++) { - sessions.remove(0).close(); - } - } -} diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolStressTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolStressTest.java deleted file mode 100644 index 33771962828..00000000000 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolStressTest.java +++ /dev/null @@ -1,296 +0,0 @@ -/* - * Copyright 2017 Google LLC - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package com.google.cloud.spanner; - -import static com.google.common.truth.Truth.assertThat; -import static org.mockito.Mockito.any; -import static org.mockito.Mockito.doAnswer; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.when; - -import com.google.api.core.ApiFuture; -import com.google.api.core.ApiFutures; -import com.google.cloud.spanner.SessionClient.SessionConsumer; -import com.google.cloud.spanner.SessionPool.PooledSessionFuture; -import com.google.cloud.spanner.SessionPool.Position; -import com.google.cloud.spanner.SessionPool.SessionConsumerImpl; -import com.google.cloud.spanner.SessionPoolOptions.ActionOnInactiveTransaction; -import com.google.cloud.spanner.SessionPoolOptions.InactiveTransactionRemovalOptions; -import com.google.cloud.spanner.spi.v1.SpannerRpc.Option; -import com.google.common.util.concurrent.Uninterruptibles; -import com.google.protobuf.Empty; -import io.opencensus.trace.Tracing; -import io.opentelemetry.api.OpenTelemetry; -import java.util.ArrayList; -import java.util.Collection; -import java.util.HashMap; -import java.util.HashSet; -import java.util.List; -import java.util.Map; -import java.util.Random; -import java.util.Set; -import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.CountDownLatch; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.Executors; -import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicBoolean; -import org.junit.Test; -import org.junit.runner.RunWith; -import org.junit.runners.Parameterized; -import org.junit.runners.Parameterized.Parameter; -import org.junit.runners.Parameterized.Parameters; -import org.mockito.Mockito; - -/** - * Stress test for {@code SessionPool} which does multiple operations on the pool, making some of - * them fail and asserts that all the invariants are maintained. - */ -@RunWith(Parameterized.class) -public class SessionPoolStressTest extends BaseSessionPoolTest { - - @Parameter(0) - public boolean shouldBlock; - - DatabaseId db = DatabaseId.of("projects/p/instances/i/databases/unused"); - SessionPool pool; - ExecutorService createExecutor = Executors.newSingleThreadExecutor(); - final Object lock = new Object(); - Random random = new Random(); - FakeClock clock = new FakeClock(); - final Map sessions = new ConcurrentHashMap<>(); - // Exception keeps track of where the session was closed at. - Map closedSessions = new HashMap<>(); - Set expiredSessions = new HashSet<>(); - SpannerImpl mockSpanner; - SpannerOptions spannerOptions; - int maxAliveSessions; - int minSessionsWhenSessionClosed = Integer.MAX_VALUE; - Exception e; - - @Parameters(name = "should block = {0}") - public static Collection data() { - List params = new ArrayList<>(); - params.add(new Object[] {true}); - params.add(new Object[] {false}); - return params; - } - - private void setupSpanner(DatabaseId db) { - ReadContext context = mock(ReadContext.class); - mockSpanner = mock(SpannerImpl.class); - spannerOptions = mock(SpannerOptions.class); - when(spannerOptions.getNumChannels()).thenReturn(4); - when(spannerOptions.getDatabaseRole()).thenReturn("role"); - SessionClient sessionClient = mock(SessionClient.class); - when(sessionClient.getSpanner()).thenReturn(mockSpanner); - when(mockSpanner.getSessionClient(db)).thenReturn(sessionClient); - when(mockSpanner.getOptions()).thenReturn(spannerOptions); - doAnswer( - invocation -> { - createExecutor.submit( - () -> { - int sessionCount = invocation.getArgument(0, Integer.class); - for (int s = 0; s < sessionCount; s++) { - SessionImpl session; - synchronized (lock) { - session = getMockedSession(mockSpanner, context); - setupSession(session, context); - sessions.put(session.getName(), false); - if (sessions.size() > maxAliveSessions) { - maxAliveSessions = sessions.size(); - } - } - SessionConsumerImpl consumer = - invocation.getArgument(2, SessionConsumerImpl.class); - consumer.onSessionReady(session); - } - }); - return null; - }) - .when(sessionClient) - .asyncBatchCreateSessions( - Mockito.anyInt(), Mockito.anyBoolean(), Mockito.any(SessionConsumer.class)); - } - - SessionImpl getMockedSession(SpannerImpl spanner, ReadContext context) { - Map options = new HashMap<>(); - options.put(Option.CHANNEL_HINT, channelHint.getAndIncrement()); - final SessionImpl session = - new SessionImpl( - spanner, - new SessionReference( - "projects/dummy/instances/dummy/databases/dummy/sessions/session" + sessionIndex, - options)) { - @Override - public ReadContext singleUse(TimestampBound bound) { - // The below stubs are added so that we can mock keep-alive. - return context; - } - - @Override - public ApiFuture asyncClose() { - synchronized (lock) { - if (expiredSessions.contains(this.getName())) { - return ApiFutures.immediateFailedFuture( - SpannerExceptionFactoryTest.newSessionNotFoundException(this.getName())); - } - if (sessions.remove(this.getName()) == null) { - setFailed(closedSessions.get(this.getName())); - } - closedSessions.put(this.getName(), new Exception("Session closed at:")); - if (sessions.size() < minSessionsWhenSessionClosed) { - minSessionsWhenSessionClosed = sessions.size(); - } - } - return ApiFutures.immediateFuture(Empty.getDefaultInstance()); - } - }; - sessionIndex++; - return session; - } - - private void setupSession(final SessionImpl session, final ReadContext mockContext) { - final ResultSet mockResult = mock(ResultSet.class); - when(mockContext.executeQuery(any(Statement.class))) - .thenAnswer( - invocation -> { - resetTransaction(session); - return mockResult; - }); - when(mockResult.next()).thenReturn(true); - } - - private void resetTransaction(SessionImpl session) { - String name = session.getName(); - synchronized (lock) { - sessions.put(name, false); - } - } - - private void setFailed(Exception cause) { - e = new Exception(cause); - } - - private void setFailed() { - e = new Exception(); - } - - private Exception getFailedError() { - synchronized (lock) { - return e; - } - } - - @Test - public void stressTest() throws Exception { - int concurrentThreads = 10; - final int numOperationsPerThread = 1000; - final CountDownLatch releaseThreads = new CountDownLatch(1); - final CountDownLatch threadsDone = new CountDownLatch(concurrentThreads); - setupSpanner(db); - int minSessions = 2; - int maxSessions = concurrentThreads / 2; - SessionPoolOptions.Builder builder = - SessionPoolOptions.newBuilder() - .setPoolMaintainerClock(clock) - .setMinSessions(minSessions) - .setMaxSessions(maxSessions) - .setInactiveTransactionRemovalOptions( - InactiveTransactionRemovalOptions.newBuilder() - .setActionOnInactiveTransaction(ActionOnInactiveTransaction.CLOSE) - .build()); - if (shouldBlock) { - builder.setBlockIfPoolExhausted(); - } else { - builder.setFailIfPoolExhausted(); - } - SessionPoolOptions sessionPoolOptions = builder.build(); - when(spannerOptions.getSessionPoolOptions()).thenReturn(sessionPoolOptions); - pool = - SessionPool.createPool( - sessionPoolOptions, - new TestExecutorFactory(), - mockSpanner.getSessionClient(db), - clock, - Position.RANDOM, - new TraceWrapper(Tracing.getTracer(), OpenTelemetry.noop().getTracer(""), false), - OpenTelemetry.noop()); - pool.idleSessionRemovedListener = - pooled -> { - String name = pooled.getName(); - // We do not take the test lock here, as we already hold the session pool lock. Taking the - // test lock as well here can cause a deadlock. - sessions.remove(name); - return null; - }; - pool.longRunningSessionRemovedListener = - pooled -> { - String name = pooled.getName(); - // We do not take the test lock here, as we already hold the session pool lock. Taking the - // test lock as well here can cause a deadlock. - sessions.remove(name); - return null; - }; - for (int i = 0; i < concurrentThreads; i++) { - new Thread( - () -> { - Uninterruptibles.awaitUninterruptibly(releaseThreads); - for (int j = 0; j < numOperationsPerThread; j++) { - try { - PooledSessionFuture session = pool.getSession(); - session.get(); - Uninterruptibles.sleepUninterruptibly(random.nextInt(2), TimeUnit.MILLISECONDS); - resetTransaction(session.get().delegate); - session.close(); - } catch (SpannerException e) { - if (e.getErrorCode() != ErrorCode.RESOURCE_EXHAUSTED || shouldBlock) { - setFailed(e); - } - } catch (Exception e) { - setFailed(e); - } - } - threadsDone.countDown(); - }) - .start(); - } - // Start maintenance threads in tight loop - final AtomicBoolean stopMaintenance = new AtomicBoolean(false); - new Thread( - () -> { - while (!stopMaintenance.get()) { - runMaintenanceLoop(clock, pool, 1); - // Sleep 1ms between maintenance loops to prevent the long-running session remover - // from stealing all sessions before they can be used. - Uninterruptibles.sleepUninterruptibly(1L, TimeUnit.MILLISECONDS); - } - }) - .start(); - releaseThreads.countDown(); - threadsDone.await(); - synchronized (lock) { - assertThat(pool.totalSessions()).isAtMost(maxSessions); - } - stopMaintenance.set(true); - pool.closeAsync(new SpannerImpl.ClosedException()).get(); - Exception e = getFailedError(); - if (e != null) { - throw e; - } - } -} diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java deleted file mode 100644 index 8d00f0889b8..00000000000 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java +++ /dev/null @@ -1,2336 +0,0 @@ -/* - * Copyright 2017 Google LLC - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package com.google.cloud.spanner; - -import static com.google.cloud.spanner.MetricRegistryConstants.GET_SESSION_TIMEOUTS; -import static com.google.cloud.spanner.MetricRegistryConstants.IS_MULTIPLEXED_KEY; -import static com.google.cloud.spanner.MetricRegistryConstants.MAX_ALLOWED_SESSIONS; -import static com.google.cloud.spanner.MetricRegistryConstants.MAX_IN_USE_SESSIONS; -import static com.google.cloud.spanner.MetricRegistryConstants.METRIC_PREFIX; -import static com.google.cloud.spanner.MetricRegistryConstants.NUM_ACQUIRED_SESSIONS; -import static com.google.cloud.spanner.MetricRegistryConstants.NUM_IN_USE_SESSIONS; -import static com.google.cloud.spanner.MetricRegistryConstants.NUM_READ_SESSIONS; -import static com.google.cloud.spanner.MetricRegistryConstants.NUM_RELEASED_SESSIONS; -import static com.google.cloud.spanner.MetricRegistryConstants.NUM_SESSIONS_AVAILABLE; -import static com.google.cloud.spanner.MetricRegistryConstants.NUM_SESSIONS_BEING_PREPARED; -import static com.google.cloud.spanner.MetricRegistryConstants.NUM_SESSIONS_IN_POOL; -import static com.google.cloud.spanner.MetricRegistryConstants.NUM_SESSIONS_IN_USE; -import static com.google.cloud.spanner.MetricRegistryConstants.NUM_WRITE_SESSIONS; -import static com.google.cloud.spanner.MetricRegistryConstants.SPANNER_DEFAULT_LABEL_VALUES; -import static com.google.cloud.spanner.MetricRegistryConstants.SPANNER_LABEL_KEYS; -import static com.google.cloud.spanner.MetricRegistryConstants.SPANNER_LABEL_KEYS_WITH_MULTIPLEXED_SESSIONS; -import static com.google.cloud.spanner.MetricRegistryConstants.SPANNER_LABEL_KEYS_WITH_TYPE; -import static com.google.cloud.spanner.SpannerOptionsTest.runWithSystemProperty; -import static com.google.common.truth.Truth.assertThat; -import static org.junit.Assert.assertEquals; -import static org.junit.Assert.assertNotEquals; -import static org.junit.Assert.assertNotNull; -import static org.junit.Assert.assertThrows; -import static org.junit.Assert.assertTrue; -import static org.junit.Assert.fail; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.anyInt; -import static org.mockito.ArgumentMatchers.eq; -import static org.mockito.Mockito.any; -import static org.mockito.Mockito.atMost; -import static org.mockito.Mockito.doAnswer; -import static org.mockito.Mockito.doThrow; -import static org.mockito.Mockito.eq; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.never; -import static org.mockito.Mockito.times; -import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; -import static org.mockito.MockitoAnnotations.initMocks; - -import com.google.api.core.ApiFutures; -import com.google.cloud.Timestamp; -import com.google.cloud.spanner.ErrorHandler.DefaultErrorHandler; -import com.google.cloud.spanner.MetricRegistryTestUtils.FakeMetricRegistry; -import com.google.cloud.spanner.MetricRegistryTestUtils.MetricsRecord; -import com.google.cloud.spanner.MetricRegistryTestUtils.PointWithFunction; -import com.google.cloud.spanner.ReadContext.QueryAnalyzeMode; -import com.google.cloud.spanner.SessionClient.SessionConsumer; -import com.google.cloud.spanner.SessionPool.PooledSession; -import com.google.cloud.spanner.SessionPool.PooledSessionFuture; -import com.google.cloud.spanner.SessionPool.Position; -import com.google.cloud.spanner.SessionPool.SessionConsumerImpl; -import com.google.cloud.spanner.SpannerImpl.ClosedException; -import com.google.cloud.spanner.TransactionRunner.TransactionCallable; -import com.google.cloud.spanner.TransactionRunnerImpl.TransactionContextImpl; -import com.google.cloud.spanner.spi.v1.SpannerRpc; -import com.google.cloud.spanner.spi.v1.SpannerRpc.ResultStreamConsumer; -import com.google.cloud.spanner.v1.stub.SpannerStubSettings; -import com.google.common.base.Stopwatch; -import com.google.common.collect.Lists; -import com.google.common.util.concurrent.ListenableFuture; -import com.google.common.util.concurrent.Uninterruptibles; -import com.google.protobuf.ByteString; -import com.google.protobuf.Empty; -import com.google.spanner.v1.CommitRequest; -import com.google.spanner.v1.CommitResponse; -import com.google.spanner.v1.ExecuteBatchDmlRequest; -import com.google.spanner.v1.ExecuteSqlRequest; -import com.google.spanner.v1.ResultSetStats; -import com.google.spanner.v1.RollbackRequest; -import com.google.spanner.v1.Transaction; -import com.google.spanner.v1.TransactionOptions; -import io.opencensus.metrics.LabelValue; -import io.opencensus.metrics.MetricRegistry; -import io.opencensus.metrics.Metrics; -import io.opencensus.trace.Tracing; -import io.opentelemetry.api.OpenTelemetry; -import io.opentelemetry.api.common.AttributeKey; -import io.opentelemetry.api.common.Attributes; -import io.opentelemetry.api.common.AttributesBuilder; -import io.opentelemetry.api.trace.Span; -import io.opentelemetry.context.Scope; -import io.opentelemetry.sdk.OpenTelemetrySdk; -import io.opentelemetry.sdk.metrics.SdkMeterProvider; -import io.opentelemetry.sdk.metrics.data.LongPointData; -import io.opentelemetry.sdk.metrics.data.MetricData; -import io.opentelemetry.sdk.testing.exporter.InMemoryMetricReader; -import java.io.PrintWriter; -import java.io.StringWriter; -import java.time.Duration; -import java.time.Instant; -import java.time.temporal.ChronoUnit; -import java.util.ArrayList; -import java.util.Arrays; -import java.util.Collection; -import java.util.Collections; -import java.util.LinkedList; -import java.util.List; -import java.util.Map; -import java.util.Objects; -import java.util.concurrent.CountDownLatch; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.Executors; -import java.util.concurrent.Future; -import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicBoolean; -import java.util.concurrent.atomic.AtomicLong; -import java.util.logging.Level; -import java.util.logging.Logger; -import java.util.stream.Collectors; -import org.junit.AfterClass; -import org.junit.Before; -import org.junit.BeforeClass; -import org.junit.Test; -import org.junit.runner.RunWith; -import org.junit.runners.Parameterized; -import org.junit.runners.Parameterized.Parameter; -import org.junit.runners.Parameterized.Parameters; -import org.mockito.Mock; -import org.mockito.Mockito; - -/** Tests for SessionPool that mock out the underlying stub. */ -@RunWith(Parameterized.class) -public class SessionPoolTest extends BaseSessionPoolTest { - private static Level originalLogLevel; - - private final ExecutorService executor = Executors.newSingleThreadExecutor(); - @Parameter public int minSessions; - - @Mock SpannerImpl client; - @Mock SessionClient sessionClient; - @Mock SpannerOptions spannerOptions; - DatabaseId db = DatabaseId.of("projects/p/instances/i/databases/unused"); - SessionPool pool; - SessionPoolOptions options; - private String sessionName = String.format("%s/sessions/s", db.getName()); - private String TEST_DATABASE_ROLE = "my-role"; - - private final TraceWrapper tracer = - new TraceWrapper(Tracing.getTracer(), OpenTelemetry.noop().getTracer(""), false); - - @Parameters(name = "min sessions = {0}") - public static Collection data() { - return Arrays.asList(new Object[][] {{0}, {1}}); - } - - private SessionPool createPool() { - return SessionPool.createPool( - options, - new TestExecutorFactory(), - client.getSessionClient(db), - tracer, - OpenTelemetry.noop()); - } - - private SessionPool createPool(Clock clock) { - return SessionPool.createPool( - options, - new TestExecutorFactory(), - client.getSessionClient(db), - clock, - Position.RANDOM, - tracer, - OpenTelemetry.noop()); - } - - private SessionPool createPool( - Clock clock, MetricRegistry metricRegistry, List labelValues) { - return SessionPool.createPool( - options, - TEST_DATABASE_ROLE, - new TestExecutorFactory(), - client.getSessionClient(db), - clock, - Position.RANDOM, - metricRegistry, - tracer, - labelValues, - OpenTelemetry.noop(), - null, - new AtomicLong(), - new AtomicLong()); - } - - private SessionPool createPool( - Clock clock, - MetricRegistry metricRegistry, - List labelValues, - OpenTelemetry openTelemetry, - Attributes attributes) { - return SessionPool.createPool( - options, - TEST_DATABASE_ROLE, - new TestExecutorFactory(), - client.getSessionClient(db), - clock, - Position.RANDOM, - metricRegistry, - tracer, - labelValues, - openTelemetry, - attributes, - new AtomicLong(), - new AtomicLong()); - } - - @BeforeClass - public static void disableLogging() { - Logger logger = Logger.getLogger(""); - originalLogLevel = logger.getLevel(); - logger.setLevel(Level.OFF); - } - - @AfterClass - public static void resetLogging() { - Logger logger = Logger.getLogger(""); - logger.setLevel(originalLogLevel); - } - - @Before - public void setUp() { - initMocks(this); - SpannerOptions.resetActiveTracingFramework(); - SpannerOptions.enableOpenTelemetryTraces(); - when(client.getOptions()).thenReturn(spannerOptions); - when(client.getSessionClient(db)).thenReturn(sessionClient); - when(sessionClient.getSpanner()).thenReturn(client); - when(spannerOptions.getNumChannels()).thenReturn(4); - when(spannerOptions.getDatabaseRole()).thenReturn("role"); - options = - SessionPoolOptions.newBuilder() - .setMinSessions(minSessions) - .setMaxSessions(2) - .setIncStep(1) - .setBlockIfPoolExhausted() - .build(); - } - - private void setupMockSessionCreation() { - doAnswer( - invocation -> { - executor.submit( - () -> { - int sessionCount = invocation.getArgument(0, Integer.class); - SessionConsumerImpl consumer = - invocation.getArgument(2, SessionConsumerImpl.class); - for (int i = 0; i < sessionCount; i++) { - consumer.onSessionReady(mockSession()); - } - }); - return null; - }) - .when(sessionClient) - .asyncBatchCreateSessions( - Mockito.anyInt(), Mockito.anyBoolean(), any(SessionConsumer.class)); - doAnswer( - invocation -> - executor.submit( - () -> { - SessionConsumer consumer = invocation.getArgument(0, SessionConsumer.class); - consumer.onSessionReady(mockMultiplexedSession()); - })) - .when(sessionClient) - .asyncCreateMultiplexedSession(any(SessionConsumer.class)); - } - - @Test - public void testClosedPoolIncludesClosedException() { - pool = createPool(); - assertTrue(pool.isValid()); - closePoolWithStacktrace(); - IllegalStateException e = assertThrows(IllegalStateException.class, () -> pool.getSession()); - assertThat(e.getCause()).isInstanceOf(ClosedException.class); - StringWriter sw = new StringWriter(); - e.getCause().printStackTrace(new PrintWriter(sw)); - assertThat(sw.toString()).contains("closePoolWithStacktrace"); - } - - private void closePoolWithStacktrace() { - pool.closeAsync(new SpannerImpl.ClosedException()); - } - - @Test - public void sessionCreation() { - setupMockSessionCreation(); - pool = createPool(); - try (Session session = pool.getSession()) { - assertThat(session).isNotNull(); - } - } - - @Test - public void poolLifo() { - setupMockSessionCreation(); - options = - options.toBuilder() - .setMinSessions(2) - .setWaitForMinSessionsDuration(Duration.ofSeconds(10L)) - .build(); - pool = createPool(); - pool.maybeWaitOnMinSessions(); - Session session1 = pool.getSession().get(); - Session session2 = pool.getSession().get(); - assertThat(session1).isNotEqualTo(session2); - - session2.close(); - session1.close(); - - // Check the session out and back in once more to finalize their positions. - session1 = pool.getSession().get(); - session2 = pool.getSession().get(); - session2.close(); - session1.close(); - - Session session3 = pool.getSession().get(); - Session session4 = pool.getSession().get(); - assertThat(session3).isEqualTo(session1); - assertThat(session4).isEqualTo(session2); - session3.close(); - session4.close(); - } - - @Test - public void poolFifo() throws Exception { - setupMockSessionCreation(); - runWithSystemProperty( - "com.google.cloud.spanner.session_pool_release_to_position", - "LAST", - () -> { - options = - options.toBuilder() - .setMinSessions(2) - .setWaitForMinSessionsDuration(Duration.ofSeconds(10L)) - .build(); - pool = createPool(); - pool.maybeWaitOnMinSessions(); - Session session1 = pool.getSession().get(); - Session session2 = pool.getSession().get(); - assertNotEquals(session1, session2); - - session2.close(); - session1.close(); - - // Check the session out and back in once more to finalize their positions. - session1 = pool.getSession().get(); - session2 = pool.getSession().get(); - session2.close(); - session1.close(); - - // Verify that we get the sessions in FIFO order, so in this order: - // 1. session2 - // 2. session1 - Session session3 = pool.getSession().get(); - Session session4 = pool.getSession().get(); - assertEquals(session2, session3); - assertEquals(session1, session4); - session3.close(); - session4.close(); - - return null; - }); - } - - @Test - public void poolAllPositions() throws Exception { - int maxAttempts = 100; - setupMockSessionCreation(); - for (Position position : Position.values()) { - runWithSystemProperty( - "com.google.cloud.spanner.session_pool_release_to_position", - position.name(), - () -> { - int attempt = 0; - while (attempt < maxAttempts) { - int numSessions = 5; - options = - options.toBuilder() - .setMinSessions(numSessions) - .setMaxSessions(numSessions) - .setWaitForMinSessionsDuration(Duration.ofSeconds(10L)) - .build(); - pool = createPool(); - pool.maybeWaitOnMinSessions(); - // First check out and release the sessions twice to the pool, so we know that we have - // finalized the position of them. - for (int n = 0; n < 2; n++) { - checkoutAndReleaseAllSessions(); - } - - // Now verify that if we get all sessions twice, they will be in random order. - List> allSessions = new ArrayList<>(2); - for (int n = 0; n < 2; n++) { - allSessions.add(checkoutAndReleaseAllSessions()); - } - List firstTime = - allSessions.get(0).stream() - .map(PooledSessionFuture::get) - .collect(Collectors.toList()); - List secondTime = - allSessions.get(1).stream() - .map(PooledSessionFuture::get) - .collect(Collectors.toList()); - switch (position) { - case FIRST: - // LIFO: - // First check out all sessions, so we have 1, 2, 3, 4, ..., N - // Then release them all back into the pool in the same order (1, 2, 3, 4, ..., N) - // That will give us the list N, ..., 4, 3, 2, 1 because each session is added at - // the front of the pool. - assertEquals(firstTime, Lists.reverse(secondTime)); - break; - case LAST: - // FIFO: - // First check out all sessions, so we have 1, 2, 3, 4, ..., N - // Then release them all back into the pool in the same order (1, 2, 3, 4, ..., N) - // That will give us the list 1, 2, 3, 4, ..., N because each session is added at - // the end of the pool. - assertEquals(firstTime, secondTime); - break; - case RANDOM: - // Random means that we should not get the same order twice (unless the randomizer - // got lucky, and then we retry). - if (attempt < (maxAttempts - 1)) { - if (Objects.equals(firstTime, secondTime)) { - attempt++; - continue; - } - } - assertNotEquals(firstTime, secondTime); - } - break; - } - return null; - }); - } - } - - private List checkoutAndReleaseAllSessions() { - List sessions = new ArrayList<>(pool.totalSessions()); - for (int i = 0; i < pool.totalSessions(); i++) { - sessions.add(pool.getSession()); - } - for (Session session : sessions) { - session.close(); - } - return sessions; - } - - @Test - public void poolClosure() throws Exception { - setupMockSessionCreation(); - pool = createPool(); - pool.closeAsync(new SpannerImpl.ClosedException()).get(5L, TimeUnit.SECONDS); - } - - @Test - public void poolClosureClosesLeakedSessions() throws Exception { - SessionImpl mockSession1 = mockSession(); - SessionImpl mockSession2 = mockSession(); - final LinkedList sessions = - new LinkedList<>(Arrays.asList(mockSession1, mockSession2)); - doAnswer( - invocation -> { - executor.submit( - () -> { - SessionConsumerImpl consumer = - invocation.getArgument(2, SessionConsumerImpl.class); - consumer.onSessionReady(sessions.pop()); - }); - return null; - }) - .when(sessionClient) - .asyncBatchCreateSessions(Mockito.eq(1), Mockito.anyBoolean(), any(SessionConsumer.class)); - pool = createPool(); - Session session1 = pool.getSession(); - // Leaked sessions - PooledSessionFuture leakedSession = pool.getSession(); - // Clear the leaked exception to suppress logging of expected exceptions. - leakedSession.clearLeakedException(); - session1.close(); - pool.closeAsync(new SpannerImpl.ClosedException()).get(5L, TimeUnit.SECONDS); - verify(mockSession1).asyncClose(); - verify(mockSession2).asyncClose(); - } - - @Test - public void poolClosesWhenMaintenanceLoopIsRunning() throws Exception { - setupMockSessionCreation(); - final FakeClock clock = new FakeClock(); - pool = createPool(clock); - final AtomicBoolean stop = new AtomicBoolean(false); - new Thread( - () -> { - // Run in a tight loop. - while (!stop.get()) { - runMaintenanceLoop(clock, pool, 1); - } - }) - .start(); - pool.closeAsync(new SpannerImpl.ClosedException()).get(5L, TimeUnit.SECONDS); - stop.set(true); - } - - @Test - public void poolClosureFailsPendingReadWaiters() throws Exception { - final CountDownLatch insideCreation = new CountDownLatch(1); - final CountDownLatch releaseCreation = new CountDownLatch(1); - final SessionImpl session1 = mockSession(); - final SessionImpl session2 = mockSession(); - doAnswer( - invocation -> { - executor.submit( - () -> { - SessionConsumerImpl consumer = - invocation.getArgument(2, SessionConsumerImpl.class); - consumer.onSessionReady(session1); - }); - return null; - }) - .doAnswer( - invocation -> { - executor.submit( - () -> { - insideCreation.countDown(); - releaseCreation.await(); - SessionConsumerImpl consumer = - invocation.getArgument(2, SessionConsumerImpl.class); - consumer.onSessionReady(session2); - return null; - }); - return null; - }) - .when(sessionClient) - .asyncBatchCreateSessions(Mockito.eq(1), Mockito.anyBoolean(), any(SessionConsumer.class)); - - pool = createPool(); - PooledSessionFuture leakedSession = pool.getSession(); - // Suppress expected leakedSession warning. - leakedSession.clearLeakedException(); - AtomicBoolean failed = new AtomicBoolean(false); - CountDownLatch latch = new CountDownLatch(1); - getSessionAsync(latch, failed); - insideCreation.await(); - pool.closeAsync(new SpannerImpl.ClosedException()); - releaseCreation.countDown(); - latch.await(5L, TimeUnit.SECONDS); - assertThat(failed.get()).isTrue(); - } - - @Test - public void poolClosureFailsPendingWriteWaiters() throws Exception { - final CountDownLatch insideCreation = new CountDownLatch(1); - final CountDownLatch releaseCreation = new CountDownLatch(1); - final SessionImpl session1 = mockSession(); - final SessionImpl session2 = mockSession(); - doAnswer( - invocation -> { - executor.submit( - () -> { - SessionConsumerImpl consumer = - invocation.getArgument(2, SessionConsumerImpl.class); - consumer.onSessionReady(session1); - }); - return null; - }) - .doAnswer( - invocation -> { - executor.submit( - () -> { - insideCreation.countDown(); - releaseCreation.await(); - SessionConsumerImpl consumer = - invocation.getArgument(2, SessionConsumerImpl.class); - consumer.onSessionReady(session2); - return null; - }); - return null; - }) - .when(sessionClient) - .asyncBatchCreateSessions(Mockito.eq(1), Mockito.anyBoolean(), any(SessionConsumer.class)); - - pool = createPool(); - PooledSessionFuture leakedSession = pool.getSession(); - // Suppress expected leakedSession warning. - leakedSession.clearLeakedException(); - AtomicBoolean failed = new AtomicBoolean(false); - CountDownLatch latch = new CountDownLatch(1); - getSessionAsync(latch, failed); - insideCreation.await(); - pool.closeAsync(new SpannerImpl.ClosedException()); - releaseCreation.countDown(); - latch.await(); - assertThat(failed.get()).isTrue(); - } - - @Test - public void poolClosesEvenIfCreationFails() throws Exception { - final CountDownLatch insideCreation = new CountDownLatch(1); - final CountDownLatch releaseCreation = new CountDownLatch(1); - doAnswer( - invocation -> { - executor.submit( - () -> { - insideCreation.countDown(); - releaseCreation.await(); - SessionConsumerImpl consumer = - invocation.getArgument(2, SessionConsumerImpl.class); - consumer.onSessionCreateFailure( - SpannerExceptionFactory.newSpannerException(new RuntimeException()), 1); - return null; - }); - return null; - }) - .when(sessionClient) - .asyncBatchCreateSessions(Mockito.eq(1), Mockito.anyBoolean(), any(SessionConsumer.class)); - pool = createPool(); - AtomicBoolean failed = new AtomicBoolean(false); - CountDownLatch latch = new CountDownLatch(1); - getSessionAsync(latch, failed); - insideCreation.await(); - ListenableFuture f = pool.closeAsync(new SpannerImpl.ClosedException()); - releaseCreation.countDown(); - f.get(); - assertThat(f.isDone()).isTrue(); - } - - @Test - public void poolClosureFailsNewRequests() { - final SessionImpl session = mockSession(); - doAnswer( - invocation -> { - executor.submit( - () -> { - SessionConsumerImpl consumer = - invocation.getArgument(2, SessionConsumerImpl.class); - consumer.onSessionReady(session); - }); - return null; - }) - .when(sessionClient) - .asyncBatchCreateSessions(Mockito.eq(1), Mockito.anyBoolean(), any(SessionConsumer.class)); - pool = createPool(); - PooledSessionFuture leakedSession = pool.getSession(); - leakedSession.get(); - // Suppress expected leakedSession warning. - leakedSession.clearLeakedException(); - pool.closeAsync(new SpannerImpl.ClosedException()); - IllegalStateException e = assertThrows(IllegalStateException.class, () -> pool.getSession()); - assertNotNull(e.getMessage()); - } - - @Test - public void atMostMaxSessionsCreated() { - setupMockSessionCreation(); - AtomicBoolean failed = new AtomicBoolean(false); - pool = createPool(); - int numSessions = 10; - final CountDownLatch latch = new CountDownLatch(numSessions); - for (int i = 0; i < numSessions; i++) { - getSessionAsync(latch, failed); - } - Uninterruptibles.awaitUninterruptibly(latch); - verify(sessionClient, atMost(options.getMaxSessions())) - .asyncBatchCreateSessions(eq(1), Mockito.anyBoolean(), any(SessionConsumer.class)); - assertThat(failed.get()).isFalse(); - } - - @Test - public void creationExceptionPropagatesToReadSession() { - doAnswer( - invocation -> { - executor.submit( - () -> { - SessionConsumerImpl consumer = - invocation.getArgument(2, SessionConsumerImpl.class); - consumer.onSessionCreateFailure( - SpannerExceptionFactory.newSpannerException(ErrorCode.INTERNAL, ""), 1); - return null; - }); - return null; - }) - .when(sessionClient) - .asyncBatchCreateSessions(Mockito.eq(1), Mockito.anyBoolean(), any(SessionConsumer.class)); - pool = createPool(); - SpannerException e = assertThrows(SpannerException.class, () -> pool.getSession().get()); - assertEquals(ErrorCode.INTERNAL, e.getErrorCode()); - } - - @Test - public void failOnPoolExhaustion() { - options = - SessionPoolOptions.newBuilder() - .setMinSessions(1) - .setMaxSessions(1) - .setFailIfPoolExhausted() - .build(); - doAnswer( - invocation -> { - executor.submit( - () -> { - SessionConsumerImpl consumer = - invocation.getArgument(2, SessionConsumerImpl.class); - consumer.onSessionReady(mockSession()); - }); - return null; - }) - .when(sessionClient) - .asyncBatchCreateSessions(Mockito.eq(1), Mockito.anyBoolean(), any(SessionConsumer.class)); - pool = createPool(); - Session session1 = pool.getSession(); - SpannerException e = assertThrows(SpannerException.class, () -> pool.getSession()); - assertEquals(ErrorCode.RESOURCE_EXHAUSTED, e.getErrorCode()); - session1.close(); - session1 = pool.getSession(); - assertThat(session1).isNotNull(); - session1.close(); - } - - @Test - public void idleSessionCleanup() throws Exception { - ReadContext context = mock(ReadContext.class); - - FakeClock clock = new FakeClock(); - clock.currentTimeMillis.set(System.currentTimeMillis()); - options = - SessionPoolOptions.newBuilder() - .setMinSessions(1) - .setMaxSessions(3) - .setIncStep(1) - .setMaxIdleSessions(0) - .setPoolMaintainerClock(clock) - .build(); - SpannerImpl spanner = mock(SpannerImpl.class); - SpannerOptions spannerOptions = mock(SpannerOptions.class); - when(spanner.getOptions()).thenReturn(spannerOptions); - when(spannerOptions.getSessionPoolOptions()).thenReturn(options); - SessionImpl session1 = buildMockSession(spanner, context); - SessionImpl session2 = buildMockSession(spanner, context); - SessionImpl session3 = buildMockSession(spanner, context); - final LinkedList sessions = - new LinkedList<>(Arrays.asList(session1, session2, session3)); - doAnswer( - invocation -> { - executor.submit( - () -> { - SessionConsumerImpl consumer = - invocation.getArgument(2, SessionConsumerImpl.class); - consumer.onSessionReady(sessions.pop()); - }); - return null; - }) - .when(sessionClient) - .asyncBatchCreateSessions(Mockito.eq(1), Mockito.anyBoolean(), any(SessionConsumer.class)); - - mockKeepAlive(context); - - pool = createPool(clock); - // Make sure pool has been initialized - pool.getSession().close(); - runMaintenanceLoop(clock, pool, pool.poolMaintainer.numClosureCycles); - assertThat(pool.numIdleSessionsRemoved()).isEqualTo(0L); - PooledSessionFuture readSession1 = pool.getSession(); - PooledSessionFuture readSession2 = pool.getSession(); - PooledSessionFuture readSession3 = pool.getSession(); - // Wait until the sessions have actually been gotten in order to make sure they are in use in - // parallel. - readSession1.get(); - readSession2.get(); - readSession3.get(); - readSession1.close(); - readSession2.close(); - readSession3.close(); - // Now there are 3 sessions in the pool but since none of them has timed out, they will all be - // kept in the pool. - runMaintenanceLoop(clock, pool, pool.poolMaintainer.numClosureCycles); - assertThat(pool.numIdleSessionsRemoved()).isEqualTo(0L); - // Counters have now been reset - // Use all 3 sessions sequentially - pool.getSession().close(); - pool.getSession().close(); - pool.getSession().close(); - // Advance the time by running the maintainer. This should cause - // one session to be kept alive and two sessions to be removed. - long cycles = - options.getRemoveInactiveSessionAfter().toMillis() / pool.poolMaintainer.loopFrequency; - runMaintenanceLoop(clock, pool, cycles); - // We will still close 2 sessions since at any point in time only 1 session was in use. - assertThat(pool.numIdleSessionsRemoved()).isEqualTo(2L); - pool.closeAsync(new SpannerImpl.ClosedException()).get(5L, TimeUnit.SECONDS); - } - - @Test - public void longRunningTransactionsCleanup_whenActionSetToClose_verifyInactiveSessionsClosed() - throws Exception { - Clock clock = mock(Clock.class); - when(clock.instant()).thenReturn(Instant.now()); - options = - SessionPoolOptions.newBuilder() - .setMinSessions(1) - .setMaxSessions(3) - .setIncStep(1) - .setMaxIdleSessions(0) - .setPoolMaintainerClock(clock) - .setCloseIfInactiveTransactions() // set option to close inactive transactions - .build(); - setupForLongRunningTransactionsCleanup(options); - - pool = createPool(clock); - // Make sure pool has been initialized - pool.getSession().close(); - - // All 3 sessions used. 100% of pool utilised. - PooledSessionFuture readSession1 = pool.getSession(); - PooledSessionFuture readSession2 = pool.getSession(); - PooledSessionFuture readSession3 = pool.getSession(); - - // complete the async tasks - readSession1.get().setEligibleForLongRunning(false); - readSession2.get().setEligibleForLongRunning(false); - readSession3.get().setEligibleForLongRunning(true); - - assertEquals(3, pool.totalSessions()); - assertEquals(3, pool.checkedOutSessions.size()); - - // ensure that the sessions are in use for > 60 minutes - pool.poolMaintainer.lastExecutionTime = Instant.now(); - when(clock.instant()).thenReturn(Instant.now().plus(61, ChronoUnit.MINUTES)); - - pool.poolMaintainer.maintainPool(); - - // the two session that were un-expectedly long-running were removed from the pool. - // verify that only 1 session that is unexpected to be long-running remains in the pool. - assertEquals(1, pool.totalSessions()); - assertEquals(2, pool.numLeakedSessionsRemoved()); - pool.closeAsync(new SpannerImpl.ClosedException()).get(5L, TimeUnit.SECONDS); - } - - @Test - public void longRunningTransactionsCleanup_whenActionSetToWarn_verifyInactiveSessionsOpen() - throws Exception { - Clock clock = mock(Clock.class); - when(clock.instant()).thenReturn(Instant.now()); - options = - SessionPoolOptions.newBuilder() - .setMinSessions(1) - .setMaxSessions(3) - .setIncStep(1) - .setPoolMaintainerClock(clock) - .setWarnIfInactiveTransactions() // set option to warn (via logs) inactive transactions - .build(); - setupForLongRunningTransactionsCleanup(options); - - pool = createPool(clock); - // Make sure pool has been initialized - pool.getSession().close(); - - // All 3 sessions used. 100% of pool utilised. - PooledSessionFuture readSession1 = pool.getSession(); - PooledSessionFuture readSession2 = pool.getSession(); - PooledSessionFuture readSession3 = pool.getSession(); - - // complete the async tasks - readSession1.get().setEligibleForLongRunning(false); - readSession2.get().setEligibleForLongRunning(false); - readSession3.get().setEligibleForLongRunning(true); - - assertEquals(3, pool.totalSessions()); - assertEquals(3, pool.checkedOutSessions.size()); - - // ensure that the sessions are in use for > 60 minutes - pool.poolMaintainer.lastExecutionTime = Instant.now(); - when(clock.instant()).thenReturn(Instant.now().plus(61, ChronoUnit.MINUTES)); - - pool.poolMaintainer.maintainPool(); - - assertEquals(3, pool.totalSessions()); - assertEquals(3, pool.checkedOutSessions.size()); - assertEquals(0, pool.numLeakedSessionsRemoved()); - - readSession1.close(); - readSession2.close(); - readSession3.close(); - pool.closeAsync(new SpannerImpl.ClosedException()).get(5L, TimeUnit.SECONDS); - } - - @Test - public void - longRunningTransactionsCleanup_whenUtilisationBelowThreshold_verifyInactiveSessionsOpen() - throws Exception { - Clock clock = mock(Clock.class); - when(clock.instant()).thenReturn(Instant.now()); - options = - SessionPoolOptions.newBuilder() - .setMinSessions(1) - .setMaxSessions(3) - .setIncStep(1) - .setMaxIdleSessions(0) - .setPoolMaintainerClock(clock) - .setCloseIfInactiveTransactions() // set option to close inactive transactions - .build(); - setupForLongRunningTransactionsCleanup(options); - - pool = createPool(clock); - pool.getSession().close(); - - // 2/3 sessions are used. Hence utilisation < 95% - PooledSessionFuture readSession1 = pool.getSession(); - PooledSessionFuture readSession2 = pool.getSession(); - - // complete the async tasks and mark sessions as checked out - readSession1.get().setEligibleForLongRunning(false); - readSession2.get().setEligibleForLongRunning(false); - - assertEquals(2, pool.totalSessions()); - assertEquals(2, pool.checkedOutSessions.size()); - - // ensure that the sessions are in use for > 60 minutes - pool.poolMaintainer.lastExecutionTime = Instant.now(); - when(clock.instant()).thenReturn(Instant.now().plus(61, ChronoUnit.MINUTES)); - - pool.poolMaintainer.maintainPool(); - - assertEquals(2, pool.totalSessions()); - assertEquals(2, pool.checkedOutSessions.size()); - assertEquals(0, pool.numLeakedSessionsRemoved()); - pool.closeAsync(new SpannerImpl.ClosedException()).get(5L, TimeUnit.SECONDS); - } - - @Test - public void - longRunningTransactionsCleanup_whenAllAreExpectedlyLongRunning_verifyInactiveSessionsOpen() - throws Exception { - SessionImpl session1 = mockSession(); - SessionImpl session2 = mockSession(); - SessionImpl session3 = mockSession(); - - final LinkedList sessions = - new LinkedList<>(Arrays.asList(session1, session2, session3)); - doAnswer( - invocation -> { - executor.submit( - () -> { - SessionConsumerImpl consumer = - invocation.getArgument(2, SessionConsumerImpl.class); - consumer.onSessionReady(sessions.pop()); - }); - return null; - }) - .when(sessionClient) - .asyncBatchCreateSessions(Mockito.eq(1), Mockito.anyBoolean(), any(SessionConsumer.class)); - - for (SessionImpl session : sessions) { - mockKeepAlive(session); - } - options = - SessionPoolOptions.newBuilder() - .setMinSessions(1) - .setMaxSessions(3) - .setIncStep(1) - .setMaxIdleSessions(0) - .setCloseIfInactiveTransactions() // set option to close inactive transactions - .build(); - Clock clock = mock(Clock.class); - when(clock.instant()).thenReturn(Instant.now()); - - pool = createPool(clock); - // Make sure pool has been initialized - pool.getSession().close(); - - // All 3 sessions used. 100% of pool utilised. - PooledSessionFuture readSession1 = pool.getSession(); - PooledSessionFuture readSession2 = pool.getSession(); - PooledSessionFuture readSession3 = pool.getSession(); - - // complete the async tasks - readSession1.get().setEligibleForLongRunning(true); - readSession2.get().setEligibleForLongRunning(true); - readSession3.get().setEligibleForLongRunning(true); - - assertEquals(3, pool.totalSessions()); - assertEquals(3, pool.checkedOutSessions.size()); - - // ensure that the sessions are in use for > 60 minutes - pool.poolMaintainer.lastExecutionTime = Instant.now(); - when(clock.instant()).thenReturn(Instant.now().plus(61, ChronoUnit.MINUTES)); - - pool.poolMaintainer.maintainPool(); - - assertEquals(3, pool.totalSessions()); - assertEquals(3, pool.checkedOutSessions.size()); - assertEquals(0, pool.numLeakedSessionsRemoved()); - pool.closeAsync(new SpannerImpl.ClosedException()).get(5L, TimeUnit.SECONDS); - } - - @Test - public void longRunningTransactionsCleanup_whenBelowDurationThreshold_verifyInactiveSessionsOpen() - throws Exception { - Clock clock = mock(Clock.class); - when(clock.instant()).thenReturn(Instant.now()); - options = - SessionPoolOptions.newBuilder() - .setMinSessions(1) - .setMaxSessions(3) - .setIncStep(1) - .setMaxIdleSessions(0) - .setPoolMaintainerClock(clock) - .setCloseIfInactiveTransactions() // set option to close inactive transactions - .build(); - setupForLongRunningTransactionsCleanup(options); - - pool = createPool(clock); - // Make sure pool has been initialized - pool.getSession().close(); - - // All 3 sessions used. 100% of pool utilised. - PooledSessionFuture readSession1 = pool.getSession(); - PooledSessionFuture readSession2 = pool.getSession(); - PooledSessionFuture readSession3 = pool.getSession(); - - // complete the async tasks - readSession1.get().setEligibleForLongRunning(false); - readSession2.get().setEligibleForLongRunning(false); - readSession3.get().setEligibleForLongRunning(true); - - assertEquals(3, pool.totalSessions()); - assertEquals(3, pool.checkedOutSessions.size()); - - // ensure that the sessions are in use for < 60 minutes - pool.poolMaintainer.lastExecutionTime = Instant.now(); - when(clock.instant()).thenReturn(Instant.now().plus(50, ChronoUnit.MINUTES)); - - pool.poolMaintainer.maintainPool(); - - assertEquals(3, pool.totalSessions()); - assertEquals(3, pool.checkedOutSessions.size()); - assertEquals(0, pool.numLeakedSessionsRemoved()); - pool.closeAsync(new SpannerImpl.ClosedException()).get(5L, TimeUnit.SECONDS); - } - - @Test - public void longRunningTransactionsCleanup_whenException_doNothing() throws Exception { - Clock clock = mock(Clock.class); - when(clock.instant()).thenReturn(Instant.now()); - options = - SessionPoolOptions.newBuilder() - .setMinSessions(1) - .setMaxSessions(3) - .setIncStep(1) - .setMaxIdleSessions(0) - .setPoolMaintainerClock(clock) - .setCloseIfInactiveTransactions() // set option to close inactive transactions - .build(); - setupForLongRunningTransactionsCleanup(options); - - pool = createPool(clock); - // Make sure pool has been initialized - pool.getSession().close(); - - // All 3 sessions used. 100% of pool utilised. - PooledSessionFuture readSession1 = pool.getSession(); - PooledSessionFuture readSession2 = pool.getSession(); - PooledSessionFuture readSession3 = pool.getSession(); - - // complete the async tasks - readSession1.get().setEligibleForLongRunning(false); - readSession2.get().setEligibleForLongRunning(false); - readSession3.get().setEligibleForLongRunning(true); - - assertEquals(3, pool.totalSessions()); - assertEquals(3, pool.checkedOutSessions.size()); - - when(clock.instant()).thenReturn(Instant.now().plus(50, ChronoUnit.MINUTES)); - - pool.poolMaintainer.lastExecutionTime = null; // setting null to throw exception - pool.poolMaintainer.maintainPool(); - - assertEquals(3, pool.totalSessions()); - assertEquals(3, pool.checkedOutSessions.size()); - assertEquals(0, pool.numLeakedSessionsRemoved()); - pool.closeAsync(new SpannerImpl.ClosedException()).get(5L, TimeUnit.SECONDS); - } - - @Test - public void - longRunningTransactionsCleanup_whenTaskRecurrenceBelowThreshold_verifyInactiveSessionsOpen() - throws Exception { - Clock clock = mock(Clock.class); - when(clock.instant()).thenReturn(Instant.now()); - options = - SessionPoolOptions.newBuilder() - .setMinSessions(1) - .setMaxSessions(3) - .setIncStep(1) - .setMaxIdleSessions(0) - .setPoolMaintainerClock(clock) - .setCloseIfInactiveTransactions() // set option to close inactive transactions - .build(); - setupForLongRunningTransactionsCleanup(options); - - pool = createPool(clock); - // Make sure pool has been initialized - pool.getSession().close(); - - // All 3 sessions used. 100% of pool utilised. - PooledSessionFuture readSession1 = pool.getSession(); - PooledSessionFuture readSession2 = pool.getSession(); - PooledSessionFuture readSession3 = pool.getSession(); - - // complete the async tasks - readSession1.get(); - readSession2.get(); - readSession3.get(); - - assertEquals(3, pool.totalSessions()); - assertEquals(3, pool.checkedOutSessions.size()); - - pool.poolMaintainer.lastExecutionTime = Instant.now(); - when(clock.instant()).thenReturn(Instant.now().plus(10, ChronoUnit.SECONDS)); - - pool.poolMaintainer.maintainPool(); - - assertEquals(3, pool.totalSessions()); - assertEquals(3, pool.checkedOutSessions.size()); - assertEquals(0, pool.numLeakedSessionsRemoved()); - - readSession1.close(); - readSession2.close(); - readSession3.close(); - pool.closeAsync(new SpannerImpl.ClosedException()).get(5L, TimeUnit.SECONDS); - } - - private void setupForLongRunningTransactionsCleanup(SessionPoolOptions sessionPoolOptions) { - ReadContext context = mock(ReadContext.class); - SpannerImpl spanner = mock(SpannerImpl.class); - SpannerOptions options = mock(SpannerOptions.class); - when(spanner.getOptions()).thenReturn(options); - when(options.getSessionPoolOptions()).thenReturn(sessionPoolOptions); - SessionImpl session1 = buildMockSession(spanner, context); - SessionImpl session2 = buildMockSession(spanner, context); - SessionImpl session3 = buildMockSession(spanner, context); - - final LinkedList sessions = - new LinkedList<>(Arrays.asList(session1, session2, session3)); - doAnswer( - invocation -> { - executor.submit( - () -> { - SessionConsumerImpl consumer = - invocation.getArgument(2, SessionConsumerImpl.class); - consumer.onSessionReady(sessions.pop()); - }); - return null; - }) - .when(sessionClient) - .asyncBatchCreateSessions(Mockito.eq(1), Mockito.anyBoolean(), any(SessionConsumer.class)); - - mockKeepAlive(context); - } - - @Test - public void keepAlive() throws Exception { - ReadContext context = mock(ReadContext.class); - FakeClock clock = new FakeClock(); - clock.currentTimeMillis.set(System.currentTimeMillis()); - options = - SessionPoolOptions.newBuilder() - .setMinSessions(2) - .setMaxSessions(3) - .setPoolMaintainerClock(clock) - .build(); - SpannerImpl spanner = mock(SpannerImpl.class); - SpannerOptions spannerOptions = mock(SpannerOptions.class); - when(spanner.getOptions()).thenReturn(spannerOptions); - when(spannerOptions.getSessionPoolOptions()).thenReturn(options); - final SessionImpl mockSession1 = buildMockSession(spanner, context); - final SessionImpl mockSession2 = buildMockSession(spanner, context); - final SessionImpl mockSession3 = buildMockSession(spanner, context); - final LinkedList sessions = - new LinkedList<>(Arrays.asList(mockSession1, mockSession2, mockSession3)); - - mockKeepAlive(context); - // This is cheating as we are returning the same session each but it makes the verification - // easier. - doAnswer( - invocation -> { - executor.submit( - () -> { - int sessionCount = invocation.getArgument(0, Integer.class); - SessionConsumerImpl consumer = - invocation.getArgument(2, SessionConsumerImpl.class); - for (int i = 0; i < sessionCount; i++) { - consumer.onSessionReady(sessions.pop()); - } - }); - return null; - }) - .when(sessionClient) - .asyncBatchCreateSessions(anyInt(), Mockito.anyBoolean(), any(SessionConsumer.class)); - pool = createPool(clock); - PooledSessionFuture session1 = pool.getSession(); - PooledSessionFuture session2 = pool.getSession(); - session1.get(); - session2.get(); - session1.close(); - session2.close(); - runMaintenanceLoop(clock, pool, pool.poolMaintainer.numKeepAliveCycles); - verify(context, never()).executeQuery(any(Statement.class)); - runMaintenanceLoop(clock, pool, pool.poolMaintainer.numKeepAliveCycles); - verify(context, times(2)).executeQuery(Statement.newBuilder("SELECT 1").build()); - clock.currentTimeMillis.addAndGet( - clock.currentTimeMillis.get() + (options.getKeepAliveIntervalMinutes() + 5L) * 60L * 1000L); - session1 = pool.getSession(); - session1.writeAtLeastOnceWithOptions(new ArrayList<>()); - session1.close(); - runMaintenanceLoop(clock, pool, pool.poolMaintainer.numKeepAliveCycles); - // The session pool only keeps MinSessions + MaxIdleSessions alive. - verify(context, times(options.getMinSessions() + options.getMaxIdleSessions())) - .executeQuery(Statement.newBuilder("SELECT 1").build()); - pool.closeAsync(new SpannerImpl.ClosedException()).get(5L, TimeUnit.SECONDS); - } - - @Test - public void blockAndTimeoutOnPoolExhaustion() throws Exception { - // Create a session pool with max 1 session and a low timeout for waiting for a session. - options = - SessionPoolOptions.newBuilder() - .setMinSessions(minSessions) - .setMaxSessions(1) - .setInitialWaitForSessionTimeoutMillis(20L) - .setAcquireSessionTimeout(null) - .build(); - setupMockSessionCreation(); - pool = createPool(); - // Take the only session that can be in the pool. - PooledSessionFuture checkedOutSession = pool.getSession(); - checkedOutSession.get(); - ExecutorService executor = Executors.newFixedThreadPool(1); - final CountDownLatch latch = new CountDownLatch(1); - // Then try asynchronously to take another session. This attempt should time out. - Future fut = - executor.submit( - () -> { - latch.countDown(); - PooledSessionFuture session = pool.getSession(); - session.close(); - return null; - }); - // Wait until the background thread is actually waiting for a session. - latch.await(); - // Wait until the request has timed out. - int waitCount = 0; - while (pool.getNumWaiterTimeouts() == 0L && waitCount < 5000) { - Thread.sleep(1L); - waitCount++; - } - // Return the checked out session to the pool so the async request will get a session and - // finish. - checkedOutSession.close(); - // Verify that the async request also succeeds. - fut.get(10L, TimeUnit.SECONDS); - executor.shutdown(); - - // Verify that the session was returned to the pool and that we can get it again. - Session session = pool.getSession(); - assertThat(session).isNotNull(); - session.close(); - assertThat(pool.getNumWaiterTimeouts()).isAtLeast(1L); - } - - @Test - public void blockAndTimeoutOnPoolExhaustion_withAcquireSessionTimeout() throws Exception { - // Create a session pool with max 1 session and a low timeout for waiting for a session. - options = - SessionPoolOptions.newBuilder() - .setMinSessions(minSessions) - .setMaxSessions(1) - .setInitialWaitForSessionTimeoutMillis(20L) - .setAcquireSessionTimeout(null) - .build(); - setupMockSessionCreation(); - pool = createPool(); - // Take the only session that can be in the pool. - PooledSessionFuture checkedOutSession = pool.getSession(); - checkedOutSession.get(); - ExecutorService executor = Executors.newFixedThreadPool(1); - final CountDownLatch latch = new CountDownLatch(1); - // Then try asynchronously to take another session. This attempt should time out. - Future fut = - executor.submit( - () -> { - PooledSessionFuture session = pool.getSession(); - latch.countDown(); - session.get(); - session.close(); - return null; - }); - // Wait until the background thread is actually waiting for a session. - latch.await(); - // Wait until the request has timed out. - Stopwatch watch = Stopwatch.createStarted(); - while (pool.getNumWaiterTimeouts() == 0L && watch.elapsed(TimeUnit.MILLISECONDS) < 1000) { - Thread.yield(); - } - // Return the checked out session to the pool so the async request will get a session and - // finish. - checkedOutSession.close(); - // Verify that the async request also succeeds. - fut.get(10L, TimeUnit.SECONDS); - executor.shutdown(); - assertTrue(executor.awaitTermination(10L, TimeUnit.SECONDS)); - - // Verify that the session was returned to the pool and that we can get it again. - PooledSessionFuture session = pool.getSession(); - assertThat(session.get()).isNotNull(); - session.close(); - assertThat(pool.getNumWaiterTimeouts()).isAtLeast(1L); - } - - @Test - public void testSessionNotFoundSingleUse() { - Statement statement = Statement.of("SELECT 1"); - final SessionImpl closedSession = mockSession(); - ReadContext closedContext = mock(ReadContext.class); - ResultSet closedResultSet = mock(ResultSet.class); - when(closedResultSet.next()) - .thenThrow(SpannerExceptionFactoryTest.newSessionNotFoundException(sessionName)); - when(closedContext.executeQuery(statement)).thenReturn(closedResultSet); - when(closedSession.singleUse()).thenReturn(closedContext); - - final SessionImpl openSession = mockSession(); - ReadContext openContext = mock(ReadContext.class); - ResultSet openResultSet = mock(ResultSet.class); - when(openResultSet.next()).thenReturn(true, false); - when(openContext.executeQuery(statement)).thenReturn(openResultSet); - when(openSession.singleUse()).thenReturn(openContext); - - doAnswer( - invocation -> { - executor.submit( - () -> { - SessionConsumerImpl consumer = - invocation.getArgument(2, SessionConsumerImpl.class); - consumer.onSessionReady(closedSession); - }); - return null; - }) - .doAnswer( - invocation -> { - executor.submit( - () -> { - SessionConsumerImpl consumer = - invocation.getArgument(2, SessionConsumerImpl.class); - consumer.onSessionReady(openSession); - }); - return null; - }) - .when(sessionClient) - .asyncBatchCreateSessions(Mockito.eq(1), Mockito.anyBoolean(), any(SessionConsumer.class)); - FakeClock clock = new FakeClock(); - clock.currentTimeMillis.set(System.currentTimeMillis()); - pool = createPool(clock); - ReadContext context = pool.getSession().singleUse(); - ResultSet resultSet = context.executeQuery(statement); - assertThat(resultSet.next()).isTrue(); - } - - @Test - public void testSessionNotFoundReadOnlyTransaction() { - Statement statement = Statement.of("SELECT 1"); - final SessionImpl closedSession = mockSession(); - when(closedSession.readOnlyTransaction()) - .thenThrow(SpannerExceptionFactoryTest.newSessionNotFoundException(sessionName)); - - final SessionImpl openSession = mockSession(); - ReadOnlyTransaction openTransaction = mock(ReadOnlyTransaction.class); - ResultSet openResultSet = mock(ResultSet.class); - when(openResultSet.next()).thenReturn(true, false); - when(openTransaction.executeQuery(statement)).thenReturn(openResultSet); - when(openSession.readOnlyTransaction()).thenReturn(openTransaction); - - doAnswer( - invocation -> { - executor.submit( - () -> { - SessionConsumerImpl consumer = - invocation.getArgument(2, SessionConsumerImpl.class); - consumer.onSessionReady(closedSession); - }); - return null; - }) - .doAnswer( - invocation -> { - executor.submit( - () -> { - SessionConsumerImpl consumer = - invocation.getArgument(2, SessionConsumerImpl.class); - consumer.onSessionReady(openSession); - }); - return null; - }) - .when(sessionClient) - .asyncBatchCreateSessions(Mockito.eq(1), Mockito.anyBoolean(), any(SessionConsumer.class)); - FakeClock clock = new FakeClock(); - clock.currentTimeMillis.set(System.currentTimeMillis()); - pool = createPool(clock); - ReadOnlyTransaction transaction = pool.getSession().readOnlyTransaction(); - ResultSet resultSet = transaction.executeQuery(statement); - assertThat(resultSet.next()).isTrue(); - } - - private enum ReadWriteTransactionTestStatementType { - QUERY, - ANALYZE, - UPDATE, - BATCH_UPDATE, - WRITE, - EXCEPTION - } - - @SuppressWarnings("unchecked") - @Test - public void testSessionNotFoundReadWriteTransaction() { - final Statement queryStatement = Statement.of("SELECT 1"); - final Statement updateStatement = Statement.of("UPDATE FOO SET BAR=1 WHERE ID=2"); - final SpannerException sessionNotFound = - SpannerExceptionFactoryTest.newSessionNotFoundException(sessionName); - for (ReadWriteTransactionTestStatementType statementType : - ReadWriteTransactionTestStatementType.values()) { - final ReadWriteTransactionTestStatementType executeStatementType = statementType; - SpannerRpc.StreamingCall closedStreamingCall = mock(SpannerRpc.StreamingCall.class); - doThrow(sessionNotFound).when(closedStreamingCall).request(Mockito.anyInt()); - SpannerRpc rpc = mock(SpannerRpc.class); - when(rpc.asyncDeleteSession(Mockito.anyString(), Mockito.anyMap())) - .thenReturn(ApiFutures.immediateFuture(Empty.getDefaultInstance())); - when(rpc.executeQuery( - any(ExecuteSqlRequest.class), - any(ResultStreamConsumer.class), - any(Map.class), - eq(true))) - .thenReturn(closedStreamingCall); - when(rpc.executeQuery(any(ExecuteSqlRequest.class), any(Map.class), eq(true))) - .thenThrow(sessionNotFound); - when(rpc.executeBatchDml(any(ExecuteBatchDmlRequest.class), any(Map.class))) - .thenThrow(sessionNotFound); - when(rpc.commitAsync(any(CommitRequest.class), any(Map.class))) - .thenReturn(ApiFutures.immediateFailedFuture(sessionNotFound)); - when(rpc.rollbackAsync(any(RollbackRequest.class), any(Map.class))) - .thenReturn(ApiFutures.immediateFailedFuture(sessionNotFound)); - when(rpc.getReadRetrySettings()) - .thenReturn(SpannerStubSettings.newBuilder().streamingReadSettings().getRetrySettings()); - when(rpc.getReadRetryableCodes()) - .thenReturn(SpannerStubSettings.newBuilder().streamingReadSettings().getRetryableCodes()); - when(rpc.getExecuteQueryRetrySettings()) - .thenReturn( - SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetrySettings()); - when(rpc.getExecuteQueryRetryableCodes()) - .thenReturn( - SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetryableCodes()); - final SessionImpl closedSession = mock(SessionImpl.class); - when(closedSession.defaultTransactionOptions()) - .thenReturn(TransactionOptions.getDefaultInstance()); - when(closedSession.getName()) - .thenReturn("projects/dummy/instances/dummy/database/dummy/sessions/session-closed"); - when(closedSession.getErrorHandler()).thenReturn(DefaultErrorHandler.INSTANCE); - when(closedSession.getRequestIdCreator()) - .thenReturn(new XGoogSpannerRequestId.NoopRequestIdCreator()); - - Span oTspan = mock(Span.class); - ISpan span = new OpenTelemetrySpan(oTspan); - when(oTspan.makeCurrent()).thenReturn(mock(Scope.class)); - - final TransactionContextImpl closedTransactionContext = - TransactionContextImpl.newBuilder() - .setSession(closedSession) - .setOptions(Options.fromTransactionOptions()) - .setRpc(rpc) - .setTracer(tracer) - .setSpan(span) - .build(); - when(closedSession.asyncClose()) - .thenReturn(ApiFutures.immediateFuture(Empty.getDefaultInstance())); - when(closedSession.newTransaction(eq(Options.fromTransactionOptions()), any())) - .thenReturn(closedTransactionContext); - when(closedSession.beginTransactionAsync(any(), eq(true), any(), any(), any())) - .thenThrow(sessionNotFound); - when(closedSession.getTracer()).thenReturn(tracer); - TransactionRunnerImpl closedTransactionRunner = new TransactionRunnerImpl(closedSession); - closedTransactionRunner.setSpan(span); - when(closedSession.readWriteTransaction()).thenReturn(closedTransactionRunner); - - final SessionImpl openSession = mock(SessionImpl.class); - when(openSession.getErrorHandler()).thenReturn(DefaultErrorHandler.INSTANCE); - when(openSession.asyncClose()) - .thenReturn(ApiFutures.immediateFuture(Empty.getDefaultInstance())); - when(openSession.getName()) - .thenReturn("projects/dummy/instances/dummy/database/dummy/sessions/session-open"); - final TransactionContextImpl openTransactionContext = mock(TransactionContextImpl.class); - when(openSession.newTransaction(eq(Options.fromTransactionOptions()), any())) - .thenReturn(openTransactionContext); - Transaction txn = Transaction.newBuilder().setId(ByteString.copyFromUtf8("open-txn")).build(); - when(openSession.beginTransactionAsync(any(), eq(true), any(), any(), any())) - .thenReturn(ApiFutures.immediateFuture(txn)); - when(openSession.getTracer()).thenReturn(tracer); - TransactionRunnerImpl openTransactionRunner = new TransactionRunnerImpl(openSession); - openTransactionRunner.setSpan(span); - when(openSession.readWriteTransaction()).thenReturn(openTransactionRunner); - when(openSession.getRequestIdCreator()) - .thenReturn(new XGoogSpannerRequestId.NoopRequestIdCreator()); - - ResultSet openResultSet = mock(ResultSet.class); - when(openResultSet.next()).thenReturn(true, false); - ResultSet planResultSet = mock(ResultSet.class); - when(planResultSet.getStats()).thenReturn(ResultSetStats.getDefaultInstance()); - when(openTransactionContext.executeQuery(queryStatement)).thenReturn(openResultSet); - when(openTransactionContext.analyzeQuery(queryStatement, QueryAnalyzeMode.PLAN)) - .thenReturn(planResultSet); - when(openTransactionContext.executeUpdate(updateStatement)).thenReturn(1L); - when(openTransactionContext.batchUpdate(Arrays.asList(updateStatement, updateStatement))) - .thenReturn(new long[] {1L, 1L}); - SpannerImpl spanner = mock(SpannerImpl.class); - SessionClient sessionClient = mock(SessionClient.class); - when(spanner.getSessionClient(db)).thenReturn(sessionClient); - when(sessionClient.getSpanner()).thenReturn(spanner); - - doAnswer( - invocation -> { - executor.submit( - () -> { - SessionConsumerImpl consumer = - invocation.getArgument(2, SessionConsumerImpl.class); - consumer.onSessionReady(closedSession); - }); - return null; - }) - .doAnswer( - invocation -> { - executor.submit( - () -> { - SessionConsumerImpl consumer = - invocation.getArgument(2, SessionConsumerImpl.class); - consumer.onSessionReady(openSession); - }); - return null; - }) - .when(sessionClient) - .asyncBatchCreateSessions( - Mockito.eq(1), Mockito.anyBoolean(), any(SessionConsumer.class)); - SessionPoolOptions options = - SessionPoolOptions.newBuilder() - .setMinSessions(0) // The pool should not auto-create any sessions - .setMaxSessions(2) - .setIncStep(1) - .setBlockIfPoolExhausted() - .build(); - SpannerOptions spannerOptions = mock(SpannerOptions.class); - when(spannerOptions.getSessionPoolOptions()).thenReturn(options); - when(spannerOptions.getNumChannels()).thenReturn(4); - when(spannerOptions.getDatabaseRole()).thenReturn("role"); - when(spanner.getOptions()).thenReturn(spannerOptions); - SessionPool pool = - SessionPool.createPool( - options, - new TestExecutorFactory(), - spanner.getSessionClient(db), - tracer, - OpenTelemetry.noop()); - try (PooledSessionFuture readWriteSession = pool.getSession()) { - TransactionRunner runner = readWriteSession.readWriteTransaction(); - try { - runner.run( - new TransactionCallable() { - private int callNumber = 0; - - @Override - public Integer run(TransactionContext transaction) { - callNumber++; - if (callNumber == 1) { - assertThat(transaction).isEqualTo(closedTransactionContext); - } else { - assertThat(transaction).isEqualTo(openTransactionContext); - } - switch (executeStatementType) { - case QUERY: - ResultSet resultSet = transaction.executeQuery(queryStatement); - assertThat(resultSet.next()).isTrue(); - break; - case ANALYZE: - ResultSet planResultSet = - transaction.analyzeQuery(queryStatement, QueryAnalyzeMode.PLAN); - assertThat(planResultSet.next()).isFalse(); - assertThat(planResultSet.getStats()).isNotNull(); - break; - case UPDATE: - long updateCount = transaction.executeUpdate(updateStatement); - assertThat(updateCount).isEqualTo(1L); - break; - case BATCH_UPDATE: - long[] updateCounts = - transaction.batchUpdate(Arrays.asList(updateStatement, updateStatement)); - assertThat(updateCounts).isEqualTo(new long[] {1L, 1L}); - break; - case WRITE: - transaction.buffer(Mutation.delete("FOO", Key.of(1L))); - break; - case EXCEPTION: - throw new RuntimeException("rollback at call " + callNumber); - default: - fail("Unknown statement type: " + executeStatementType); - } - return callNumber; - } - }); - } catch (Exception e) { - // The rollback will also cause a SessionNotFoundException, but this is caught, logged - // and further ignored by the library, meaning that the session will not be re-created - // for retry. Hence rollback at call 1. - assertThat(executeStatementType) - .isEqualTo(ReadWriteTransactionTestStatementType.EXCEPTION); - assertThat(e.getMessage()).contains("rollback at call 1"); - } - } - pool.closeAsync(new SpannerImpl.ClosedException()); - } - } - - @Test - public void testSessionNotFoundWrite() { - SpannerException sessionNotFound = - SpannerExceptionFactoryTest.newSessionNotFoundException(sessionName); - List mutations = Collections.singletonList(Mutation.newInsertBuilder("FOO").build()); - final SessionImpl closedSession = mockSession(); - closedSession.setRequestIdCreator(new XGoogSpannerRequestId.NoopRequestIdCreator()); - when(closedSession.writeWithOptions(eq(mutations), any())).thenThrow(sessionNotFound); - - final SessionImpl openSession = mockSession(); - com.google.cloud.spanner.CommitResponse response = - mock(com.google.cloud.spanner.CommitResponse.class); - when(response.getCommitTimestamp()).thenReturn(Timestamp.now()); - openSession.setRequestIdCreator(new XGoogSpannerRequestId.NoopRequestIdCreator()); - when(openSession.writeWithOptions(eq(mutations), any())).thenReturn(response); - doAnswer( - invocation -> { - executor.submit( - () -> { - SessionConsumerImpl consumer = - invocation.getArgument(2, SessionConsumerImpl.class); - consumer.onSessionReady(closedSession); - }); - return null; - }) - .doAnswer( - invocation -> { - executor.submit( - () -> { - SessionConsumerImpl consumer = - invocation.getArgument(2, SessionConsumerImpl.class); - consumer.onSessionReady(openSession); - }); - return null; - }) - .when(sessionClient) - .asyncBatchCreateSessions(Mockito.eq(1), Mockito.anyBoolean(), any(SessionConsumer.class)); - - FakeClock clock = new FakeClock(); - clock.currentTimeMillis.set(System.currentTimeMillis()); - pool = createPool(clock); - DatabaseClientImpl impl = new DatabaseClientImpl(pool, tracer); - assertThat(impl.write(mutations)).isNotNull(); - } - - @Test - public void testSessionNotFoundWriteAtLeastOnce() { - SpannerException sessionNotFound = - SpannerExceptionFactoryTest.newSessionNotFoundException(sessionName); - List mutations = Collections.singletonList(Mutation.newInsertBuilder("FOO").build()); - final SessionImpl closedSession = mockSession(); - closedSession.setRequestIdCreator(new XGoogSpannerRequestId.NoopRequestIdCreator()); - when(closedSession.writeAtLeastOnceWithOptions(eq(mutations), any())) - .thenThrow(sessionNotFound); - - final SessionImpl openSession = mockSession(); - com.google.cloud.spanner.CommitResponse response = - mock(com.google.cloud.spanner.CommitResponse.class); - when(response.getCommitTimestamp()).thenReturn(Timestamp.now()); - openSession.setRequestIdCreator(new XGoogSpannerRequestId.NoopRequestIdCreator()); - when(openSession.writeAtLeastOnceWithOptions(eq(mutations), any())).thenReturn(response); - doAnswer( - invocation -> { - executor.submit( - () -> { - SessionConsumerImpl consumer = - invocation.getArgument(2, SessionConsumerImpl.class); - consumer.onSessionReady(closedSession); - }); - return null; - }) - .doAnswer( - invocation -> { - executor.submit( - () -> { - SessionConsumerImpl consumer = - invocation.getArgument(2, SessionConsumerImpl.class); - consumer.onSessionReady(openSession); - }); - return null; - }) - .when(sessionClient) - .asyncBatchCreateSessions(Mockito.eq(1), Mockito.anyBoolean(), any(SessionConsumer.class)); - FakeClock clock = new FakeClock(); - clock.currentTimeMillis.set(System.currentTimeMillis()); - pool = createPool(clock); - DatabaseClientImpl impl = new DatabaseClientImpl(pool, tracer); - assertThat(impl.writeAtLeastOnce(mutations)).isNotNull(); - } - - @Test - public void testSessionNotFoundPartitionedUpdate() { - SpannerException sessionNotFound = - SpannerExceptionFactoryTest.newSessionNotFoundException(sessionName); - Statement statement = Statement.of("UPDATE FOO SET BAR=1 WHERE 1=1"); - final SessionImpl closedSession = mockSession(); - when(closedSession.executePartitionedUpdate(eq(statement), any())).thenThrow(sessionNotFound); - - final SessionImpl openSession = mockSession(); - when(openSession.executePartitionedUpdate(eq(statement), any())).thenReturn(1L); - doAnswer( - invocation -> { - executor.submit( - () -> { - SessionConsumerImpl consumer = - invocation.getArgument(2, SessionConsumerImpl.class); - consumer.onSessionReady(closedSession); - }); - return null; - }) - .doAnswer( - invocation -> { - executor.submit( - () -> { - SessionConsumerImpl consumer = - invocation.getArgument(2, SessionConsumerImpl.class); - consumer.onSessionReady(openSession); - }); - return null; - }) - .when(sessionClient) - .asyncBatchCreateSessions(Mockito.eq(1), Mockito.anyBoolean(), any(SessionConsumer.class)); - FakeClock clock = new FakeClock(); - clock.currentTimeMillis.set(System.currentTimeMillis()); - pool = createPool(clock); - DatabaseClientImpl impl = new DatabaseClientImpl(pool, mock(TraceWrapper.class)); - assertThat(impl.executePartitionedUpdate(statement)).isEqualTo(1L); - } - - @SuppressWarnings("rawtypes") - @Test - public void testOpenCensusSessionMetrics() throws Exception { - // Create a session pool with max 2 session and a low timeout for waiting for a session. - options = - SessionPoolOptions.newBuilder() - .setMinSessions(1) - .setMaxSessions(2) - .setInitialWaitForSessionTimeoutMillis(50L) - .setAcquireSessionTimeout(null) - .build(); - FakeClock clock = new FakeClock(); - clock.currentTimeMillis.set(System.currentTimeMillis()); - FakeMetricRegistry metricRegistry = new FakeMetricRegistry(); - List labelValues = - Arrays.asList( - LabelValue.create("client1"), - LabelValue.create("database1"), - LabelValue.create("instance1"), - LabelValue.create("1.0.0")); - - setupMockSessionCreation(); - pool = createPool(clock, metricRegistry, labelValues); - PooledSessionFuture session1 = pool.getSession(); - PooledSessionFuture session2 = pool.getSession(); - session1.get(); - session2.get(); - - MetricsRecord record = metricRegistry.pollRecord(); - assertThat(record.getMetrics().size()).isEqualTo(6); - - List maxInUseSessions = - record.getMetrics().get(METRIC_PREFIX + MAX_IN_USE_SESSIONS); - assertThat(maxInUseSessions.size()).isEqualTo(1); - assertThat(maxInUseSessions.get(0).value()).isEqualTo(2L); - assertThat(maxInUseSessions.get(0).keys()).isEqualTo(SPANNER_LABEL_KEYS); - assertThat(maxInUseSessions.get(0).values()).isEqualTo(labelValues); - - List getSessionsTimeouts = - record.getMetrics().get(METRIC_PREFIX + GET_SESSION_TIMEOUTS); - assertThat(getSessionsTimeouts.size()).isEqualTo(1); - assertThat(getSessionsTimeouts.get(0).value()).isAtMost(1L); - assertThat(getSessionsTimeouts.get(0).keys()).isEqualTo(SPANNER_LABEL_KEYS); - assertThat(getSessionsTimeouts.get(0).values()).isEqualTo(labelValues); - - List labelValuesWithRegularSessions = new ArrayList<>(labelValues); - labelValuesWithRegularSessions.add(LabelValue.create("false")); - List labelValuesWithMultiplexedSessions = new ArrayList<>(labelValues); - labelValuesWithMultiplexedSessions.add(LabelValue.create("true")); - List numAcquiredSessions = - record.getMetrics().get(METRIC_PREFIX + NUM_ACQUIRED_SESSIONS); - assertThat(numAcquiredSessions.size()).isEqualTo(2); - PointWithFunction regularSessionMetric = - numAcquiredSessions.stream() - .filter( - x -> - x.keys().contains(IS_MULTIPLEXED_KEY) - && x.values().contains(LabelValue.create("false"))) - .findFirst() - .get(); - PointWithFunction multiplexedSessionMetric = - numAcquiredSessions.stream() - .filter( - x -> - x.keys().contains(IS_MULTIPLEXED_KEY) - && x.values().contains(LabelValue.create("true"))) - .findFirst() - .get(); - // verify metrics for regular sessions - assertThat(regularSessionMetric.value()).isEqualTo(2L); - assertThat(regularSessionMetric.keys()).isEqualTo(SPANNER_LABEL_KEYS_WITH_MULTIPLEXED_SESSIONS); - assertThat(regularSessionMetric.values()).isEqualTo(labelValuesWithRegularSessions); - - // verify metrics for multiplexed sessions - assertThat(multiplexedSessionMetric.value()).isEqualTo(0L); - assertThat(multiplexedSessionMetric.keys()) - .isEqualTo(SPANNER_LABEL_KEYS_WITH_MULTIPLEXED_SESSIONS); - assertThat(multiplexedSessionMetric.values()).isEqualTo(labelValuesWithMultiplexedSessions); - - List numReleasedSessions = - record.getMetrics().get(METRIC_PREFIX + NUM_RELEASED_SESSIONS); - assertThat(numReleasedSessions.size()).isEqualTo(2); - - regularSessionMetric = - numReleasedSessions.stream() - .filter( - x -> - x.keys().contains(IS_MULTIPLEXED_KEY) - && x.values().contains(LabelValue.create("false"))) - .findFirst() - .get(); - multiplexedSessionMetric = - numReleasedSessions.stream() - .filter( - x -> - x.keys().contains(IS_MULTIPLEXED_KEY) - && x.values().contains(LabelValue.create("true"))) - .findFirst() - .get(); - // verify metrics for regular sessions - assertThat(regularSessionMetric.value()).isEqualTo(0L); - assertThat(regularSessionMetric.keys()).isEqualTo(SPANNER_LABEL_KEYS_WITH_MULTIPLEXED_SESSIONS); - assertThat(regularSessionMetric.values()).isEqualTo(labelValuesWithRegularSessions); - - // verify metrics for multiplexed sessions - assertThat(multiplexedSessionMetric.value()).isEqualTo(0L); - assertThat(multiplexedSessionMetric.keys()) - .isEqualTo(SPANNER_LABEL_KEYS_WITH_MULTIPLEXED_SESSIONS); - assertThat(multiplexedSessionMetric.values()).isEqualTo(labelValuesWithMultiplexedSessions); - - List maxAllowedSessions = - record.getMetrics().get(METRIC_PREFIX + MAX_ALLOWED_SESSIONS); - assertThat(maxAllowedSessions.size()).isEqualTo(1); - assertThat(maxAllowedSessions.get(0).value()).isEqualTo(options.getMaxSessions()); - assertThat(maxAllowedSessions.get(0).keys()).isEqualTo(SPANNER_LABEL_KEYS); - assertThat(maxAllowedSessions.get(0).values()).isEqualTo(labelValues); - - List numSessionsInPool = - record.getMetrics().get(METRIC_PREFIX + NUM_SESSIONS_IN_POOL); - assertThat(numSessionsInPool.size()).isEqualTo(4); - PointWithFunction beingPrepared = numSessionsInPool.get(0); - List labelValuesWithBeingPreparedType = new ArrayList<>(labelValues); - labelValuesWithBeingPreparedType.add(NUM_SESSIONS_BEING_PREPARED); - assertThat(beingPrepared.value()).isEqualTo(0L); - assertThat(beingPrepared.keys()).isEqualTo(SPANNER_LABEL_KEYS_WITH_TYPE); - assertThat(beingPrepared.values()).isEqualTo(labelValuesWithBeingPreparedType); - PointWithFunction numSessionsInUse = numSessionsInPool.get(1); - List labelValuesWithInUseType = new ArrayList<>(labelValues); - labelValuesWithInUseType.add(NUM_IN_USE_SESSIONS); - assertThat(numSessionsInUse.value()).isEqualTo(2L); - assertThat(numSessionsInUse.keys()).isEqualTo(SPANNER_LABEL_KEYS_WITH_TYPE); - assertThat(numSessionsInUse.values()).isEqualTo(labelValuesWithInUseType); - PointWithFunction readSessions = numSessionsInPool.get(2); - List labelValuesWithReadType = new ArrayList<>(labelValues); - labelValuesWithReadType.add(NUM_READ_SESSIONS); - assertThat(readSessions.value()).isEqualTo(0L); - assertThat(readSessions.keys()).isEqualTo(SPANNER_LABEL_KEYS_WITH_TYPE); - assertThat(readSessions.values()).isEqualTo(labelValuesWithReadType); - PointWithFunction writePreparedSessions = numSessionsInPool.get(3); - List labelValuesWithWriteType = new ArrayList<>(labelValues); - labelValuesWithWriteType.add(NUM_WRITE_SESSIONS); - assertThat(writePreparedSessions.value()).isEqualTo(0L); - assertThat(writePreparedSessions.keys()).isEqualTo(SPANNER_LABEL_KEYS_WITH_TYPE); - assertThat(writePreparedSessions.values()).isEqualTo(labelValuesWithWriteType); - - final CountDownLatch latch = new CountDownLatch(1); - // Try asynchronously to take another session. This attempt should time out. - Future fut = - executor.submit( - () -> { - latch.countDown(); - Session session = pool.getSession(); - session.close(); - return null; - }); - // Wait until the background thread is actually waiting for a session. - latch.await(); - // Wait until the request has timed out. - int waitCount = 0; - while (pool.getNumWaiterTimeouts() == 0L && waitCount < 5000) { - //noinspection BusyWait - Thread.sleep(1L); - waitCount++; - } - assertTrue(pool.getNumWaiterTimeouts() > 0L); - // Return the checked out session to the pool so the async request will get a session and - // finish. - session2.close(); - // Verify that the async request also succeeds. - fut.get(10L, TimeUnit.SECONDS); - executor.shutdown(); - - session1.close(); - numAcquiredSessions = record.getMetrics().get(METRIC_PREFIX + NUM_ACQUIRED_SESSIONS); - assertThat(numAcquiredSessions.size()).isEqualTo(2); - regularSessionMetric = - numAcquiredSessions.stream() - .filter( - x -> - x.keys().contains(IS_MULTIPLEXED_KEY) - && x.values().contains(LabelValue.create("false"))) - .findFirst() - .get(); - multiplexedSessionMetric = - numAcquiredSessions.stream() - .filter( - x -> - x.keys().contains(IS_MULTIPLEXED_KEY) - && x.values().contains(LabelValue.create("true"))) - .findFirst() - .get(); - assertThat(regularSessionMetric.value()).isEqualTo(3L); - assertThat(multiplexedSessionMetric.value()).isEqualTo(0L); - - numReleasedSessions = record.getMetrics().get(METRIC_PREFIX + NUM_RELEASED_SESSIONS); - assertThat(numReleasedSessions.size()).isEqualTo(2); - regularSessionMetric = - numReleasedSessions.stream() - .filter( - x -> - x.keys().contains(IS_MULTIPLEXED_KEY) - && x.values().contains(LabelValue.create("false"))) - .findFirst() - .get(); - multiplexedSessionMetric = - numReleasedSessions.stream() - .filter( - x -> - x.keys().contains(IS_MULTIPLEXED_KEY) - && x.values().contains(LabelValue.create("true"))) - .findFirst() - .get(); - assertThat(regularSessionMetric.value()).isEqualTo(3L); - assertThat(multiplexedSessionMetric.value()).isEqualTo(0L); - - maxInUseSessions = record.getMetrics().get(METRIC_PREFIX + MAX_IN_USE_SESSIONS); - assertThat(maxInUseSessions.size()).isEqualTo(1); - assertThat(maxInUseSessions.get(0).value()).isEqualTo(2L); - - numSessionsInPool = record.getMetrics().get(METRIC_PREFIX + NUM_SESSIONS_IN_POOL); - assertThat(numSessionsInPool.size()).isEqualTo(4); - beingPrepared = numSessionsInPool.get(0); - assertThat(beingPrepared.value()).isEqualTo(0L); - numSessionsInUse = numSessionsInPool.get(1); - assertThat(numSessionsInUse.value()).isEqualTo(0L); - readSessions = numSessionsInPool.get(2); - assertThat(readSessions.value()).isEqualTo(2L); - writePreparedSessions = numSessionsInPool.get(3); - assertThat(writePreparedSessions.value()).isEqualTo(0L); - } - - @Test - public void testOpenCensusMetricsDisable() { - SpannerOptions.disableOpenCensusMetrics(); - // Create a session pool with max 2 session and a low timeout for waiting for a session. - options = - SessionPoolOptions.newBuilder() - .setMinSessions(1) - .setMaxSessions(2) - .setMaxIdleSessions(0) - .setInitialWaitForSessionTimeoutMillis(50L) - .build(); - FakeClock clock = new FakeClock(); - clock.currentTimeMillis.set(System.currentTimeMillis()); - FakeMetricRegistry metricRegistry = new FakeMetricRegistry(); - List labelValues = - Arrays.asList( - LabelValue.create("client1"), - LabelValue.create("database1"), - LabelValue.create("instance1"), - LabelValue.create("1.0.0")); - - setupMockSessionCreation(); - pool = createPool(clock, metricRegistry, labelValues); - PooledSessionFuture session1 = pool.getSession(); - PooledSessionFuture session2 = pool.getSession(); - session1.get(); - session2.get(); - - MetricsRecord record = metricRegistry.pollRecord(); - assertThat(record.getMetrics().size()).isEqualTo(0); - SpannerOptions.enableOpenCensusMetrics(); - } - - @Test - public void testOpenTelemetrySessionMetrics() throws Exception { - SpannerOptions.resetActiveTracingFramework(); - SpannerOptions.enableOpenTelemetryMetrics(); - // Create a session pool with max 3 session and a low timeout for waiting for a session. - if (minSessions == 1) { - options = - SessionPoolOptions.newBuilder() - .setMinSessions(1) - .setMaxSessions(3) - // This must be set to null for the setInitialWaitForSessionTimeoutMillis call to have - // any effect. - .setAcquireSessionTimeout(null) - .setInitialWaitForSessionTimeoutMillis(1L) - .build(); - FakeClock clock = new FakeClock(); - clock.currentTimeMillis.set(System.currentTimeMillis()); - - InMemoryMetricReader inMemoryMetricReader = InMemoryMetricReader.create(); - SdkMeterProvider sdkMeterProvider = - SdkMeterProvider.builder().registerMetricReader(inMemoryMetricReader).build(); - OpenTelemetry openTelemetry = - OpenTelemetrySdk.builder().setMeterProvider(sdkMeterProvider).build(); - - setupMockSessionCreation(); - - AttributesBuilder attributesBuilder = Attributes.builder(); - attributesBuilder.put("client_id", "testClient"); - attributesBuilder.put("database", "testDb"); - attributesBuilder.put("instance_id", "test_instance"); - attributesBuilder.put("library_version", "test_version"); - - pool = - createPool( - clock, - Metrics.getMetricRegistry(), - SPANNER_DEFAULT_LABEL_VALUES, - openTelemetry, - attributesBuilder.build()); - PooledSessionFuture session1 = pool.getSession(); - PooledSessionFuture session2 = pool.getSession(); - session1.get(); - session2.get(); - - Collection metricDataCollection = inMemoryMetricReader.collectAllMetrics(); - // Acquired sessions are 2. - verifyMetricData(metricDataCollection, NUM_ACQUIRED_SESSIONS, 1, 2L); - // Max in use session are 2. - verifyMetricData(metricDataCollection, MAX_IN_USE_SESSIONS, 1, 2D); - // Max Allowed sessions should be 3 - verifyMetricData(metricDataCollection, MAX_ALLOWED_SESSIONS, 1, 3D); - // Released sessions should be 0 - verifyMetricData(metricDataCollection, NUM_RELEASED_SESSIONS, 1, 0L); - // Num sessions in pool - verifyMetricData(metricDataCollection, NUM_SESSIONS_IN_POOL, 1, NUM_SESSIONS_IN_USE, 2); - - PooledSessionFuture session3 = pool.getSession(); - session3.get(); - - final CountDownLatch latch = new CountDownLatch(1); - // Try asynchronously to take another session. This attempt should time out. - Future fut = - executor.submit( - () -> { - PooledSessionFuture session = pool.getSession(); - latch.countDown(); - session.get(); - session.close(); - return null; - }); - // Wait until the background thread is actually waiting for a session. - latch.await(); - // Wait until the request has timed out. - Stopwatch watch = Stopwatch.createStarted(); - while (pool.getNumWaiterTimeouts() == 0L && watch.elapsed(TimeUnit.MILLISECONDS) < 100) { - Thread.yield(); - } - assertTrue(pool.getNumWaiterTimeouts() > 0); - // Return the checked out session to the pool so the async request will get a session and - // finish. - session2.close(); - // Verify that the async request also succeeds. - fut.get(10L, TimeUnit.SECONDS); - executor.shutdown(); - assertTrue(executor.awaitTermination(10L, TimeUnit.SECONDS)); - - inMemoryMetricReader.forceFlush(); - metricDataCollection = inMemoryMetricReader.collectAllMetrics(); - - // Max Allowed sessions should be 3 - verifyMetricData(metricDataCollection, MAX_ALLOWED_SESSIONS, 1, 3D); - // Session timeouts 1 - // verifyMetricData(metricDataCollection, GET_SESSION_TIMEOUTS, 1, 1L); - // Max in use session are 2. - verifyMetricData(metricDataCollection, MAX_IN_USE_SESSIONS, 1, 3D); - // Session released 2 - verifyMetricData(metricDataCollection, NUM_RELEASED_SESSIONS, 1, 2L); - // Acquired sessions are 4. - verifyMetricData(metricDataCollection, NUM_ACQUIRED_SESSIONS, 1, 4L); - // Num sessions in pool - verifyMetricData(metricDataCollection, NUM_SESSIONS_IN_POOL, 1, NUM_SESSIONS_IN_USE, 2); - verifyMetricData(metricDataCollection, NUM_SESSIONS_IN_POOL, 1, NUM_SESSIONS_AVAILABLE, 1); - } - } - - private static void verifyMetricData( - Collection metricDataCollection, String metricName, int size, long value) { - Collection metricDataFiltered = - metricDataCollection.stream() - .filter(x -> x.getName().equals(metricName)) - .collect(Collectors.toList()); - - assertEquals(metricDataFiltered.size(), size); - MetricData metricData = metricDataFiltered.stream().findFirst().get(); - LongPointData regularSessionMetric = - metricData.getLongSumData().getPoints().stream() - .filter( - x -> - Boolean.FALSE.equals( - x.getAttributes().get(AttributeKey.booleanKey("is_multiplexed")))) - .findFirst() - .get(); - LongPointData multiplexedSessionMetric = - metricData.getLongSumData().getPoints().stream() - .filter( - x -> - Boolean.TRUE.equals( - x.getAttributes().get(AttributeKey.booleanKey("is_multiplexed")))) - .findFirst() - .get(); - assertEquals(value, regularSessionMetric.getValue()); - assertEquals(0, multiplexedSessionMetric.getValue()); - } - - private static void verifyMetricData( - Collection metricDataCollection, String metricName, int size, double value) { - Collection metricDataFiltered = - metricDataCollection.stream() - .filter(x -> x.getName().equals(metricName)) - .collect(Collectors.toList()); - - assertEquals(metricDataFiltered.size(), size); - MetricData metricData = metricDataFiltered.stream().findFirst().get(); - assertEquals( - metricData.getDoubleGaugeData().getPoints().stream().findFirst().get().getValue(), - value, - 0.0); - } - - private static void verifyMetricData( - Collection metricDataCollection, - String metricName, - int size, - String labelName, - long value) { - Collection metricDataFiltered = - metricDataCollection.stream() - .filter(x -> x.getName().equals(metricName)) - .collect(Collectors.toList()); - - assertEquals(metricDataFiltered.size(), size); - - MetricData metricData = metricDataFiltered.stream().findFirst().get(); - - assertEquals( - metricData.getLongSumData().getPoints().stream() - .filter(x -> x.getAttributes().asMap().containsValue(labelName)) - .findFirst() - .get() - .getValue(), - value); - } - - @Test - public void testGetDatabaseRole() throws Exception { - setupMockSessionCreation(); - pool = createPool(new FakeClock(), new FakeMetricRegistry(), SPANNER_DEFAULT_LABEL_VALUES); - assertEquals(TEST_DATABASE_ROLE, pool.getDatabaseRole()); - } - - @Test - public void testWaitOnMinSessionsWhenSessionsAreCreatedBeforeTimeout() { - options = - SessionPoolOptions.newBuilder() - .setMinSessions(minSessions) - .setMaxSessions(minSessions + 1) - .setWaitForMinSessionsDuration(Duration.ofSeconds(5)) - .build(); - doAnswer( - invocation -> - executor.submit( - () -> { - SessionConsumerImpl consumer = - invocation.getArgument(2, SessionConsumerImpl.class); - consumer.onSessionReady(mockSession()); - })) - .when(sessionClient) - .asyncBatchCreateSessions(Mockito.eq(1), Mockito.anyBoolean(), any(SessionConsumer.class)); - - pool = createPool(new FakeClock(), new FakeMetricRegistry(), SPANNER_DEFAULT_LABEL_VALUES); - pool.maybeWaitOnMinSessions(); - assertTrue(pool.getNumberOfSessionsInPool() >= minSessions); - } - - @Test(expected = SpannerException.class) - public void testWaitOnMinSessionsThrowsExceptionWhenTimeoutIsReached() { - // Does not call onSessionReady, so session pool is never populated - doAnswer(invocation -> null) - .when(sessionClient) - .asyncBatchCreateSessions(Mockito.eq(1), Mockito.anyBoolean(), any(SessionConsumer.class)); - - options = - SessionPoolOptions.newBuilder() - .setMinSessions(minSessions + 1) - .setMaxSessions(minSessions + 1) - .setWaitForMinSessionsDuration(Duration.ofMillis(100)) - .build(); - pool = createPool(new FakeClock(), new FakeMetricRegistry(), SPANNER_DEFAULT_LABEL_VALUES); - pool.maybeWaitOnMinSessions(); - } - - @Test - public void reset_maxSessionsInUse() { - Clock clock = mock(Clock.class); - when(clock.instant()).thenReturn(Instant.now()); - options = - SessionPoolOptions.newBuilder() - .setMinSessions(1) - .setMaxSessions(3) - .setIncStep(1) - .setMaxIdleSessions(0) - .setPoolMaintainerClock(clock) - .build(); - setupForLongRunningTransactionsCleanup(options); - - pool = createPool(clock); - // Make sure pool has been initialized - pool.getSession().close(); - - // All 3 sessions used. 100% of pool utilised. - PooledSessionFuture readSession1 = pool.getSession(); - PooledSessionFuture readSession2 = pool.getSession(); - PooledSessionFuture readSession3 = pool.getSession(); - - // complete the async tasks - readSession1.get().setEligibleForLongRunning(false); - readSession2.get().setEligibleForLongRunning(false); - readSession3.get().setEligibleForLongRunning(true); - - assertEquals(3, pool.getMaxSessionsInUse()); - assertEquals(3, pool.getNumberOfSessionsInUse()); - - // Release 1 session - readSession1.get().close(); - - // Verify that numSessionsInUse reduces to 2 while maxSessionsInUse remain 3 - assertEquals(3, pool.getMaxSessionsInUse()); - assertEquals(2, pool.getNumberOfSessionsInUse()); - - // ensure that the lastResetTime for maxSessionsInUse > 10 minutes - when(clock.instant()).thenReturn(Instant.now().plus(11, ChronoUnit.MINUTES)); - - pool.poolMaintainer.maintainPool(); - - // Verify that maxSessionsInUse is reset to numSessionsInUse - assertEquals(2, pool.getMaxSessionsInUse()); - assertEquals(2, pool.getNumberOfSessionsInUse()); - } - - private void mockKeepAlive(ReadContext context) { - ResultSet resultSet = mock(ResultSet.class); - when(resultSet.next()).thenReturn(true, false); - when(context.executeQuery(any(Statement.class))).thenReturn(resultSet); - } - - private void mockKeepAlive(Session session) { - ReadContext context = mock(ReadContext.class); - ResultSet resultSet = mock(ResultSet.class); - when(resultSet.next()).thenReturn(true, false); - when(session.singleUse(any(TimestampBound.class))).thenReturn(context); - when(context.executeQuery(any(Statement.class))).thenReturn(resultSet); - } - - private void getSessionAsync(final CountDownLatch latch, final AtomicBoolean failed) { - new Thread( - () -> { - try (PooledSessionFuture future = pool.getSession()) { - PooledSession session = future.get(); - failed.compareAndSet(false, session == null); - Uninterruptibles.sleepUninterruptibly(10, TimeUnit.MILLISECONDS); - } catch (Throwable e) { - failed.compareAndSet(false, true); - } finally { - latch.countDown(); - } - }) - .start(); - } -} diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolUnbalancedTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolUnbalancedTest.java deleted file mode 100644 index 5a9365eaed9..00000000000 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolUnbalancedTest.java +++ /dev/null @@ -1,241 +0,0 @@ -/* - * Copyright 2023 Google LLC - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package com.google.cloud.spanner; - -import static com.google.cloud.spanner.SessionPool.isUnbalanced; -import static org.junit.Assert.assertFalse; -import static org.junit.Assert.assertTrue; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.when; - -import com.google.cloud.spanner.SessionPool.PooledSession; -import com.google.cloud.spanner.SessionPool.PooledSessionFuture; -import java.util.Arrays; -import java.util.List; -import java.util.Set; -import java.util.stream.Collectors; -import org.junit.Test; -import org.junit.runner.RunWith; -import org.junit.runners.JUnit4; - -@RunWith(JUnit4.class) -public class SessionPoolUnbalancedTest { - - static PooledSession mockedSession(int channel) { - PooledSession session = mock(PooledSession.class); - when(session.getChannel()).thenReturn(channel); - return session; - } - - static List mockedSessions(int... channels) { - return Arrays.stream(channels) - .mapToObj(SessionPoolUnbalancedTest::mockedSession) - .collect(Collectors.toList()); - } - - static PooledSessionFuture mockedCheckedOutSession(int channel) { - PooledSession session = mockedSession(channel); - PooledSessionFuture future = mock(PooledSessionFuture.class); - when(future.get()).thenReturn(session); - when(future.isDone()).thenReturn(true); - return future; - } - - static Set mockedCheckedOutSessions(int... channels) { - return Arrays.stream(channels) - .mapToObj(SessionPoolUnbalancedTest::mockedCheckedOutSession) - .collect(Collectors.toSet()); - } - - @Test - public void testIsUnbalancedBasics() { - // An empty session pool is never unbalanced. - assertFalse(isUnbalanced(1, mockedSessions(), mockedCheckedOutSessions(1, 1, 1), 1)); - assertFalse(isUnbalanced(1, mockedSessions(), mockedCheckedOutSessions(1, 1, 1), 2)); - assertFalse(isUnbalanced(1, mockedSessions(), mockedCheckedOutSessions(1, 1, 1), 4)); - assertFalse(isUnbalanced(1, mockedSessions(), mockedCheckedOutSessions(1, 1, 1, 1), 1)); - assertFalse(isUnbalanced(1, mockedSessions(), mockedCheckedOutSessions(1, 1, 1, 1), 2)); - assertFalse(isUnbalanced(1, mockedSessions(), mockedCheckedOutSessions(1, 1, 1, 1), 4)); - assertFalse(isUnbalanced(1, mockedSessions(), mockedCheckedOutSessions(1, 1, 1, 1, 1), 1)); - assertFalse(isUnbalanced(1, mockedSessions(), mockedCheckedOutSessions(1, 1, 1, 1, 1), 2)); - assertFalse(isUnbalanced(1, mockedSessions(), mockedCheckedOutSessions(1, 1, 1, 1, 1), 4)); - - // A session pool that has 2 or fewer sessions checked out is never unbalanced. - // This prevents low-QPS scenarios from re-balancing the pool. - assertFalse(isUnbalanced(1, mockedSessions(1, 1, 1), mockedCheckedOutSessions(), 1)); - assertFalse(isUnbalanced(1, mockedSessions(1, 1, 1), mockedCheckedOutSessions(), 2)); - assertFalse(isUnbalanced(1, mockedSessions(1, 1, 1), mockedCheckedOutSessions(), 4)); - assertFalse(isUnbalanced(1, mockedSessions(1, 1, 1, 1), mockedCheckedOutSessions(1), 1)); - assertFalse(isUnbalanced(1, mockedSessions(1, 1, 1, 1), mockedCheckedOutSessions(1), 2)); - assertFalse(isUnbalanced(1, mockedSessions(1, 1, 1, 1), mockedCheckedOutSessions(1), 4)); - assertFalse(isUnbalanced(1, mockedSessions(1, 1, 1, 1, 1), mockedCheckedOutSessions(1, 1), 1)); - assertFalse(isUnbalanced(1, mockedSessions(1, 1, 1, 1, 1), mockedCheckedOutSessions(1, 1), 2)); - assertFalse(isUnbalanced(1, mockedSessions(1, 1, 1, 1, 1), mockedCheckedOutSessions(1, 1), 4)); - - // A session pool that uses only 1 channel is never unbalanced. - assertFalse(isUnbalanced(1, mockedSessions(1, 1, 1), mockedCheckedOutSessions(), 1)); - assertFalse(isUnbalanced(1, mockedSessions(1, 1, 1, 1), mockedCheckedOutSessions(), 1)); - assertFalse(isUnbalanced(1, mockedSessions(1, 1, 1, 1, 1), mockedCheckedOutSessions(), 1)); - assertFalse(isUnbalanced(1, mockedSessions(1, 1, 1, 1, 1, 1), mockedCheckedOutSessions(), 1)); - assertFalse(isUnbalanced(1, mockedSessions(1, 1, 1), mockedCheckedOutSessions(1, 1, 1), 1)); - assertFalse( - isUnbalanced(1, mockedSessions(1, 1, 1, 1), mockedCheckedOutSessions(1, 1, 1, 1), 1)); - assertFalse( - isUnbalanced(1, mockedSessions(1, 1, 1, 1, 1), mockedCheckedOutSessions(1, 1, 1, 1, 1), 1)); - assertFalse( - isUnbalanced( - 1, mockedSessions(1, 1, 1, 1, 1, 1), mockedCheckedOutSessions(1, 1, 1, 1, 1, 1), 1)); - } - - @Test - public void testIsUnbalanced_returnsFalseForBalancedPool() { - assertFalse( - isUnbalanced(1, mockedSessions(1, 2, 3, 4), mockedCheckedOutSessions(1, 2, 3, 4), 4)); - assertFalse( - isUnbalanced(2, mockedSessions(1, 2, 3, 4), mockedCheckedOutSessions(1, 2, 3, 4), 4)); - assertFalse( - isUnbalanced(3, mockedSessions(1, 2, 3, 4), mockedCheckedOutSessions(1, 2, 3, 4), 4)); - assertFalse( - isUnbalanced(4, mockedSessions(1, 2, 3, 4), mockedCheckedOutSessions(1, 2, 3, 4), 4)); - - assertFalse( - isUnbalanced(1, mockedSessions(1, 2, 3, 4), mockedCheckedOutSessions(4, 3, 2, 1), 4)); - assertFalse( - isUnbalanced(2, mockedSessions(1, 2, 3, 4), mockedCheckedOutSessions(4, 3, 2, 1), 4)); - assertFalse( - isUnbalanced(3, mockedSessions(1, 2, 3, 4), mockedCheckedOutSessions(4, 3, 2, 1), 4)); - assertFalse( - isUnbalanced(4, mockedSessions(1, 2, 3, 4), mockedCheckedOutSessions(4, 3, 2, 1), 4)); - - assertFalse( - isUnbalanced( - 1, - mockedSessions(1, 2, 3, 4, 1, 2, 3, 4), - mockedCheckedOutSessions(1, 2, 3, 4, 1, 2, 3, 4), - 4)); - - // We only check the first numChannels sessions that are in the pool, so the fact that the end - // of the pool is unbalanced is not a reason to re-balance. - assertFalse( - isUnbalanced( - 1, mockedSessions(1, 2, 3, 4, 1, 1, 1, 1), mockedCheckedOutSessions(1, 2, 3, 4), 4)); - assertFalse( - isUnbalanced(1, mockedSessions(1, 2, 1, 1, 1, 1), mockedCheckedOutSessions(1, 2), 2)); - assertFalse( - isUnbalanced( - 1, - mockedSessions(1, 2, 3, 4, 1, 2, 3, 4, 1, 1, 1, 1), - mockedCheckedOutSessions(1, 2, 3, 4), - 8)); - assertFalse( - isUnbalanced( - 1, - mockedSessions(1, 1, 2, 2, 3, 3, 4, 4, 1, 1, 1, 1), - mockedCheckedOutSessions(1, 2, 3, 4), - 8)); - - // The list of checked out sessions is allowed to contain up to twice the number of sessions - // with a given channel than it should for a perfect distribution (perfect means - // num_sessions_with_a_channel == num_channels). - assertFalse( - isUnbalanced(1, mockedSessions(1, 2, 3, 4), mockedCheckedOutSessions(1, 1, 2, 3), 4)); - assertFalse( - isUnbalanced( - 1, - mockedSessions(1, 2, 3, 4), - mockedCheckedOutSessions(1, 1, 1, 1, 2, 3, 4, 5, 6, 7, 8, 2, 3, 4, 5, 6), - 8)); - // We're only checking the list of checked out sessions against the channel that is being added - // to the pool. - assertFalse( - isUnbalanced(1, mockedSessions(1, 2, 3, 4), mockedCheckedOutSessions(2, 2, 2, 2), 4)); - - // We do not consider a pool unbalanced if the list of checked out sessions only contains 2 of - // the same channel, even if that would still be 'more than twice the ideal number'. This - // prevents that a small number of checked out sessions that happen to use the same channel - // causes the pool to be considered unbalanced. - assertFalse( - isUnbalanced( - 1, mockedSessions(1, 2, 3, 4, 5, 6, 7, 8), mockedCheckedOutSessions(1, 1, 2), 8)); - - // A larger number of checked out sessions means that we can also have a 'large' number of the - // same channels in that list, as long as it does not exceed twice the number that it should be - // for an ideal distribution. - assertFalse( - isUnbalanced( - 1, - mockedSessions(1, 2, 3, 4, 5, 6, 7, 8), - mockedCheckedOutSessions(1, 1, 1, 1, 1, 2, 3, 4, 5, 6, 7, 8, 2, 4, 5, 5, 3, 4, 8, 8), - 8)); - } - - @Test - public void testIsUnbalanced_returnsTrueForUnbalancedPool() { - // The pool is considered unbalanced if the first numChannel sessions contain 3 or more of the - // same sessions as the one that is being added. Also; if the pool uses only 2 channels, then it - // is also considered unbalanced if the two first sessions in the pool already use the same - // channel as the one being added. - assertTrue(isUnbalanced(1, mockedSessions(1, 1), mockedCheckedOutSessions(1, 2, 1, 2), 2)); - assertTrue(isUnbalanced(2, mockedSessions(2, 2), mockedCheckedOutSessions(1, 2, 1, 2), 2)); - - assertTrue( - isUnbalanced(1, mockedSessions(1, 1, 1, 4), mockedCheckedOutSessions(1, 2, 3, 4), 4)); - assertTrue( - isUnbalanced(2, mockedSessions(2, 2, 2, 4), mockedCheckedOutSessions(1, 2, 3, 4), 4)); - assertTrue( - isUnbalanced(3, mockedSessions(1, 3, 3, 3), mockedCheckedOutSessions(1, 2, 3, 4), 4)); - assertTrue( - isUnbalanced(4, mockedSessions(1, 4, 4, 4), mockedCheckedOutSessions(1, 2, 3, 4), 4)); - - assertTrue( - isUnbalanced( - 1, mockedSessions(1, 2, 3, 4, 5, 6, 1, 1), mockedCheckedOutSessions(1, 2, 3, 4), 8)); - assertTrue( - isUnbalanced( - 2, mockedSessions(1, 3, 4, 5, 6, 2, 2, 2), mockedCheckedOutSessions(1, 2, 3, 4), 8)); - assertTrue( - isUnbalanced( - 3, mockedSessions(1, 2, 3, 3, 4, 5, 3, 6), mockedCheckedOutSessions(1, 2, 3, 4), 8)); - assertTrue( - isUnbalanced( - 4, mockedSessions(1, 2, 3, 4, 5, 4, 5, 4), mockedCheckedOutSessions(1, 2, 3, 4), 8)); - - // The pool is also considered unbalanced if the list of checked out sessions contain more than - // 2 times as many sessions of the one being returned as it should. - assertTrue( - isUnbalanced(1, mockedSessions(1, 2, 3, 4), mockedCheckedOutSessions(1, 1, 2, 1), 4)); - assertTrue( - isUnbalanced(2, mockedSessions(1, 2, 3, 4), mockedCheckedOutSessions(2, 2, 2, 4), 4)); - assertTrue( - isUnbalanced(3, mockedSessions(1, 2, 3, 4), mockedCheckedOutSessions(1, 3, 3, 3), 4)); - assertTrue( - isUnbalanced(4, mockedSessions(1, 2, 3, 4), mockedCheckedOutSessions(4, 2, 4, 4), 4)); - assertTrue( - isUnbalanced( - 1, mockedSessions(1, 2, 3, 4), mockedCheckedOutSessions(1, 1, 2, 1, 1, 2, 3, 1), 4)); - - assertTrue( - isUnbalanced( - 1, mockedSessions(1, 2, 3, 4, 5, 6, 7, 8), mockedCheckedOutSessions(1, 1, 1, 3), 8)); - assertTrue( - isUnbalanced( - 1, - mockedSessions(1, 2, 3, 4, 5, 6, 7, 8), - mockedCheckedOutSessions(1, 1, 1, 2, 3, 4, 5, 6, 7, 8, 1, 1), - 8)); - } -} diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerImplTest.java index 3cf13dc58d3..a675605e768 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerImplTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerImplTest.java @@ -28,7 +28,6 @@ import com.google.cloud.NoCredentials; import com.google.cloud.ServiceRpc; import com.google.cloud.grpc.GrpcTransportOptions; -import com.google.cloud.spanner.SpannerException.DoNotConstructDirectly; import com.google.cloud.spanner.SpannerImpl.ClosedException; import com.google.cloud.spanner.admin.database.v1.DatabaseAdminClient; import com.google.cloud.spanner.admin.database.v1.stub.DatabaseAdminStub; @@ -45,7 +44,6 @@ import java.util.Collections; import java.util.HashMap; import java.util.Map; -import java.util.UUID; import org.junit.After; import org.junit.Before; import org.junit.BeforeClass; @@ -222,63 +220,6 @@ public void testSpannerClosed() { spanner4.close(); } - @Test - public void testClientId() { - // Create a unique database id to be sure it has not yet been used in the lifetime of this JVM. - String dbName = - String.format("projects/p1/instances/i1/databases/%s", UUID.randomUUID().toString()); - DatabaseId db = DatabaseId.of(dbName); - - Mockito.when(spannerOptions.getTransportOptions()) - .thenReturn(GrpcTransportOptions.newBuilder().build()); - Mockito.when(spannerOptions.getSessionPoolOptions()) - .thenReturn(SessionPoolOptions.newBuilder().setMinSessions(0).build()); - Mockito.when(spannerOptions.getDatabaseRole()).thenReturn("role"); - - DatabaseClientImpl databaseClient = (DatabaseClientImpl) impl.getDatabaseClient(db); - assertThat(databaseClient.clientId).isEqualTo("client-1"); - - // Get same db client again. - DatabaseClientImpl databaseClient1 = (DatabaseClientImpl) impl.getDatabaseClient(db); - assertThat(databaseClient1.clientId).isEqualTo(databaseClient.clientId); - - // Get a db client for a different database. - String dbName2 = - String.format("projects/p1/instances/i1/databases/%s", UUID.randomUUID().toString()); - DatabaseId db2 = DatabaseId.of(dbName2); - DatabaseClientImpl databaseClient2 = (DatabaseClientImpl) impl.getDatabaseClient(db2); - assertThat(databaseClient2.clientId).isEqualTo("client-1"); - - // Getting a new database client for an invalidated database should use the same client id. - databaseClient.pool.setResourceNotFoundException( - new DatabaseNotFoundException(DoNotConstructDirectly.ALLOWED, "not found", null, null)); - DatabaseClientImpl revalidated = (DatabaseClientImpl) impl.getDatabaseClient(db); - assertThat(revalidated).isNotSameInstanceAs(databaseClient); - assertThat(revalidated.clientId).isEqualTo(databaseClient.clientId); - - // Now invalidate the second client and request a new one. - revalidated.pool.setResourceNotFoundException( - new DatabaseNotFoundException(DoNotConstructDirectly.ALLOWED, "not found", null, null)); - DatabaseClientImpl revalidated2 = (DatabaseClientImpl) impl.getDatabaseClient(db); - assertThat(revalidated2).isNotSameInstanceAs(revalidated); - assertThat(revalidated2.clientId).isEqualTo(revalidated.clientId); - - // Create a new Spanner instance. This will generate new database clients with new ids. - try (Spanner spanner = - SpannerOptions.newBuilder() - .setProjectId("p1") - .setCredentials(NoCredentials.getInstance()) - .build() - .getService()) { - - // Get a database client for the same database as the first database. As this goes through a - // different Spanner instance with potentially different options, it will get a different - // client id. - DatabaseClientImpl databaseClient3 = (DatabaseClientImpl) spanner.getDatabaseClient(db); - assertThat(databaseClient3.clientId).isEqualTo("client-2"); - } - } - @Test public void testClosedException() { Spanner spanner = new SpannerImpl(rpc, spannerOptions); diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java index 2325f2ac40f..b0a79fe5645 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java @@ -335,7 +335,7 @@ public void inlineBegin() { new SessionImpl( spanner, new SessionReference( - "projects/p/instances/i/databases/d/sessions/s", Collections.EMPTY_MAP)) {}; + "projects/p/instances/i/databases/d/sessions/s", null, Collections.EMPTY_MAP)) {}; session.setRequestIdCreator(new XGoogSpannerRequestId.NoopRequestIdCreator()); session.setCurrentSpan(new OpenTelemetrySpan(mock(io.opentelemetry.api.trace.Span.class))); TransactionRunnerImpl runner = new TransactionRunnerImpl(session); diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionTest.java index 72ef90c9cf6..c8469ae08a9 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionTest.java @@ -23,8 +23,6 @@ import static org.junit.Assert.assertTrue; import static org.junit.Assert.fail; -import com.google.api.core.ApiFuture; -import com.google.api.core.ApiFutures; import com.google.cloud.spanner.AbortedException; import com.google.cloud.spanner.Dialect; import com.google.cloud.spanner.ErrorCode; @@ -36,10 +34,7 @@ import com.google.cloud.spanner.SpannerOptions; import com.google.cloud.spanner.Statement; import com.google.cloud.spanner.TimestampBound; -import com.google.cloud.spanner.connection.ConnectionOptions.Builder; -import com.google.cloud.spanner.connection.StatementExecutor.StatementExecutorType; import com.google.common.collect.ImmutableList; -import com.google.spanner.v1.BatchCreateSessionsRequest; import com.google.spanner.v1.CommitRequest; import com.google.spanner.v1.DirectedReadOptions; import com.google.spanner.v1.DirectedReadOptions.ExcludeReplicas; @@ -50,11 +45,8 @@ import com.google.spanner.v1.RequestOptions; import java.nio.charset.StandardCharsets; import java.time.Duration; -import java.util.Arrays; import java.util.Collections; -import java.util.concurrent.ExecutionException; import java.util.concurrent.TimeUnit; -import java.util.concurrent.TimeoutException; import java.util.function.Consumer; import java.util.function.Supplier; import javax.annotation.Nonnull; @@ -385,82 +377,6 @@ private void assertResetProperty( } } - public static class ConnectionMinSessionsTest extends AbstractMockServerTest { - - @AfterClass - public static void reset() { - mockSpanner.reset(); - } - - protected String getBaseUrl() { - return super.getBaseUrl() + ";minSessions=1"; - } - - @Test - public void testMinSessions() throws InterruptedException, TimeoutException { - try (Connection connection = createConnection()) { - mockSpanner.waitForRequestsToContain( - input -> - input instanceof BatchCreateSessionsRequest - && ((BatchCreateSessionsRequest) input).getSessionCount() == 1, - 5000L); - } - } - } - - public static class ConnectionMaxSessionsTest extends AbstractMockServerTest { - - @AfterClass - public static void reset() { - mockSpanner.reset(); - } - - protected String getBaseUrl() { - return super.getBaseUrl() + ";maxSessions=1"; - } - - @Override - protected Builder configureConnectionOptions(Builder builder) { - return builder.setStatementExecutorType(StatementExecutorType.PLATFORM_THREAD); - } - - @Test - public void testMaxSessions() - throws InterruptedException, TimeoutException, ExecutionException { - try (Connection connection1 = createConnection(); - Connection connection2 = createConnection()) { - connection1.beginTransactionAsync(); - connection2.beginTransactionAsync(); - - ApiFuture count1 = connection1.executeUpdateAsync(INSERT_STATEMENT); - ApiFuture count2 = connection2.executeUpdateAsync(INSERT_STATEMENT); - - // Commit the transactions. Both should be able to finish, but both used the same session. - ApiFuture commit1 = connection1.commitAsync(); - ApiFuture commit2 = connection2.commitAsync(); - - // At least one transaction must wait until the other has finished before it can get a - // session. - assertThat(count1.isDone() && count2.isDone()).isFalse(); - assertThat(commit1.isDone() && commit2.isDone()).isFalse(); - - // Wait until both finishes. - ApiFutures.allAsList(Arrays.asList(commit1, commit2)).get(5L, TimeUnit.SECONDS); - - assertThat(count1.isDone()).isTrue(); - assertThat(count2.isDone()).isTrue(); - if (isMultiplexedSessionsEnabled(connection1.getSpanner())) { - // We don't use the multiplexed session, so we don't know whether the server had time to - // create it or not. That means that we have between 1 and 2 sessions on the server. - assertThat(mockSpanner.numSessionsCreated()).isAtLeast(1); - assertThat(mockSpanner.numSessionsCreated()).isAtMost(2); - } else { - assertThat(mockSpanner.numSessionsCreated()).isEqualTo(1); - } - } - } - } - public static class ConnectionRPCPriorityTest extends AbstractMockServerTest { @AfterClass diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITClosedSessionTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITClosedSessionTest.java deleted file mode 100644 index 6ffb0e1ca68..00000000000 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITClosedSessionTest.java +++ /dev/null @@ -1,286 +0,0 @@ -/* - * Copyright 2019 Google LLC - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package com.google.cloud.spanner.it; - -import static com.google.common.truth.Truth.assertThat; -import static org.junit.Assert.assertNotNull; -import static org.junit.Assert.fail; -import static org.junit.Assume.assumeFalse; - -import com.google.cloud.spanner.AbortedException; -import com.google.cloud.spanner.Database; -import com.google.cloud.spanner.IntegrationTestWithClosedSessionsEnv; -import com.google.cloud.spanner.IntegrationTestWithClosedSessionsEnv.DatabaseClientWithClosedSessionImpl; -import com.google.cloud.spanner.ParallelIntegrationTest; -import com.google.cloud.spanner.ReadOnlyTransaction; -import com.google.cloud.spanner.ResultSet; -import com.google.cloud.spanner.SessionNotFoundException; -import com.google.cloud.spanner.Statement; -import com.google.cloud.spanner.TimestampBound; -import com.google.cloud.spanner.TransactionContext; -import com.google.cloud.spanner.TransactionManager; -import com.google.cloud.spanner.TransactionRunner; -import java.util.concurrent.TimeUnit; -import org.junit.Before; -import org.junit.BeforeClass; -import org.junit.ClassRule; -import org.junit.Test; -import org.junit.experimental.categories.Category; -import org.junit.runner.RunWith; -import org.junit.runners.JUnit4; - -/** Test the automatic re-creation of sessions that have been invalidated by the server. */ -@Category(ParallelIntegrationTest.class) -@RunWith(JUnit4.class) -public class ITClosedSessionTest { - // Run each test case twice to ensure that a retried session does not affect subsequent - // transactions. - private static final int RUNS_PER_TEST_CASE = 2; - - @ClassRule - public static IntegrationTestWithClosedSessionsEnv env = - new IntegrationTestWithClosedSessionsEnv(); - - private static Database db; - private static DatabaseClientWithClosedSessionImpl client; - - @BeforeClass - public static void setUpDatabase() { - // For multiplexed sessions, it will never be invalidated by the server and hence the client - // will never receive an exception with code NOT_FOUND and the text 'Session not found'. - assumeFalse( - env.getTestHelper().getOptions().getSessionPoolOptions().getUseMultiplexedSession()); - - // Empty database. - db = env.getTestHelper().createTestDatabase(); - client = (DatabaseClientWithClosedSessionImpl) env.getTestHelper().getDatabaseClient(db); - } - - @Before - public void setup() { - client.setAllowSessionReplacing(true); - } - - @Test - public void testSingleUse() { - // This should trigger an exception with code NOT_FOUND and the text 'Session not found'. - client.invalidateNextSession(); - for (int run = 0; run < RUNS_PER_TEST_CASE; run++) { - try (ResultSet rs = Statement.of("SELECT 1").executeQuery(client.singleUse())) { - assertThat(rs.next()).isTrue(); - assertThat(rs.getLong(0)).isEqualTo(1L); - assertThat(rs.next()).isFalse(); - } - } - } - - @Test - public void testSingleUseNoRecreation() { - // This should trigger an exception with code NOT_FOUND and the text 'Session not found'. - client.setAllowSessionReplacing(false); - client.invalidateNextSession(); - try (ResultSet rs = Statement.of("SELECT 1").executeQuery(client.singleUse())) { - rs.next(); - fail("Expected exception"); - } catch (SessionNotFoundException ex) { - assertNotNull(ex.getMessage()); - } - } - - @Test - public void testSingleUseBound() { - // This should trigger an exception with code NOT_FOUND and the text 'Session not found'. - client.invalidateNextSession(); - for (int run = 0; run < RUNS_PER_TEST_CASE; run++) { - try (ResultSet rs = - Statement.of("SELECT 1") - .executeQuery( - client.singleUse(TimestampBound.ofExactStaleness(10L, TimeUnit.SECONDS)))) { - assertThat(rs.next()).isTrue(); - assertThat(rs.getLong(0)).isEqualTo(1L); - assertThat(rs.next()).isFalse(); - } - } - } - - @Test - public void testSingleUseReadOnlyTransaction() { - client.invalidateNextSession(); - for (int run = 0; run < RUNS_PER_TEST_CASE; run++) { - try (ReadOnlyTransaction txn = client.singleUseReadOnlyTransaction()) { - try (ResultSet rs = txn.executeQuery(Statement.of("SELECT 1"))) { - assertThat(rs.next()).isTrue(); - assertThat(rs.getLong(0)).isEqualTo(1L); - assertThat(rs.next()).isFalse(); - } - assertThat(txn.getReadTimestamp()).isNotNull(); - } - } - } - - @Test - public void testSingleUseReadOnlyTransactionBound() { - client.invalidateNextSession(); - for (int run = 0; run < RUNS_PER_TEST_CASE; run++) { - try (ReadOnlyTransaction txn = - client.singleUseReadOnlyTransaction( - TimestampBound.ofMaxStaleness(10L, TimeUnit.SECONDS))) { - try (ResultSet rs = txn.executeQuery(Statement.of("SELECT 1"))) { - assertThat(rs.next()).isTrue(); - assertThat(rs.getLong(0)).isEqualTo(1L); - assertThat(rs.next()).isFalse(); - } - assertThat(txn.getReadTimestamp()).isNotNull(); - } - } - } - - @Test - public void testReadOnlyTransaction() { - client.invalidateNextSession(); - for (int run = 0; run < RUNS_PER_TEST_CASE; run++) { - try (ReadOnlyTransaction txn = client.readOnlyTransaction()) { - for (int i = 0; i < 2; i++) { - try (ResultSet rs = txn.executeQuery(Statement.of("SELECT 1"))) { - assertThat(rs.next()).isTrue(); - assertThat(rs.getLong(0)).isEqualTo(1L); - assertThat(rs.next()).isFalse(); - } - } - assertThat(txn.getReadTimestamp()).isNotNull(); - } - } - } - - @Test - public void testReadOnlyTransactionNoRecreation() { - client.setAllowSessionReplacing(false); - client.invalidateNextSession(); - try (ReadOnlyTransaction txn = client.readOnlyTransaction()) { - try (ResultSet rs = txn.executeQuery(Statement.of("SELECT 1"))) { - rs.next(); - fail("Expected exception"); - } - fail("Expected exception"); - } catch (SessionNotFoundException ex) { - assertNotNull(ex.getMessage()); - } - } - - @Test - public void testReadOnlyTransactionBound() { - client.invalidateNextSession(); - for (int run = 0; run < RUNS_PER_TEST_CASE; run++) { - try (ReadOnlyTransaction txn = - client.readOnlyTransaction(TimestampBound.ofExactStaleness(10L, TimeUnit.SECONDS))) { - for (int i = 0; i < 2; i++) { - try (ResultSet rs = txn.executeQuery(Statement.of("SELECT 1"))) { - assertThat(rs.next()).isTrue(); - assertThat(rs.getLong(0)).isEqualTo(1L); - assertThat(rs.next()).isFalse(); - } - } - assertThat(txn.getReadTimestamp()).isNotNull(); - } - } - } - - @Test - public void testReadWriteTransaction() { - client.invalidateNextSession(); - for (int run = 0; run < RUNS_PER_TEST_CASE; run++) { - TransactionRunner txn = client.readWriteTransaction(); - txn.run( - transaction -> { - for (int i = 0; i < 2; i++) { - try (ResultSet rs = transaction.executeQuery(Statement.of("SELECT 1"))) { - assertThat(rs.next()).isTrue(); - assertThat(rs.getLong(0)).isEqualTo(1L); - assertThat(rs.next()).isFalse(); - } - } - return null; - }); - } - } - - @Test - public void testReadWriteTransactionNoRecreation() { - client.setAllowSessionReplacing(false); - client.invalidateNextSession(); - try { - TransactionRunner txn = client.readWriteTransaction(); - txn.run( - transaction -> { - try (ResultSet rs = transaction.executeQuery(Statement.of("SELECT 1"))) { - rs.next(); - fail("Expected exception"); - } - return null; - }); - fail("Expected exception"); - } catch (SessionNotFoundException ex) { - assertNotNull(ex.getMessage()); - } - } - - @Test - public void testTransactionManager() throws InterruptedException { - client.invalidateNextSession(); - for (int run = 0; run < 2; run++) { - try (TransactionManager manager = client.transactionManager()) { - TransactionContext txn = manager.begin(); - try { - while (true) { - for (int i = 0; i < 2; i++) { - try (ResultSet rs = txn.executeQuery(Statement.of("SELECT 1"))) { - assertThat(rs.next()).isTrue(); - assertThat(rs.getLong(0)).isEqualTo(1L); - assertThat(rs.next()).isFalse(); - } - } - manager.commit(); - break; - } - } catch (AbortedException e) { - long retryDelayInMillis = e.getRetryDelayInMillis(); - if (retryDelayInMillis > 0) { - Thread.sleep(retryDelayInMillis); - } - txn = manager.resetForRetry(); - } - } - } - } - - @Test - public void testTransactionManagerNoRecreation() { - client.setAllowSessionReplacing(false); - client.invalidateNextSession(); - try (TransactionManager manager = client.transactionManager()) { - TransactionContext txn = manager.begin(); - while (true) { - try (ResultSet rs = txn.executeQuery(Statement.of("SELECT 1"))) { - rs.next(); - fail("Expected exception"); - } - } - } catch (SessionNotFoundException ex) { - assertNotNull(ex.getMessage()); - } - } -}