diff --git a/.github/sync-repo-settings.yaml b/.github/sync-repo-settings.yaml index 39ea467f10..e26b598a47 100644 --- a/.github/sync-repo-settings.yaml +++ b/.github/sync-repo-settings.yaml @@ -14,13 +14,10 @@ branchProtectionRules: - units (8) - units (11) - 'Kokoro - Test: Integration' - - 'Kokoro - Test: Integration with Multiplexed Sessions' - cla/google - checkstyle - compile (8) - compile (11) - - units-with-multiplexed-session (8) - - units-with-multiplexed-session (11) - unmanaged_dependency_check - library_generation - pattern: 3.3.x @@ -154,7 +151,6 @@ branchProtectionRules: - units (8) - units (11) - 'Kokoro - Test: Integration' - - 'Kokoro - Test: Integration with Multiplexed Sessions' - cla/google - checkstyle - compile (8) @@ -173,7 +169,6 @@ branchProtectionRules: - units (8) - units (11) - 'Kokoro - Test: Integration' - - 'Kokoro - Test: Integration with Multiplexed Sessions' - cla/google - checkstyle - compile (8) @@ -194,7 +189,6 @@ branchProtectionRules: - units (8) - units (11) - 'Kokoro - Test: Integration' - - 'Kokoro - Test: Integration with Multiplexed Sessions' - cla/google - checkstyle - compile (8) @@ -215,7 +209,6 @@ branchProtectionRules: - units (8) - units (11) - 'Kokoro - Test: Integration' - - 'Kokoro - Test: Integration with Multiplexed Sessions' - cla/google - checkstyle - compile (8) diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index e2e6881c78..ae7ec53f5d 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -36,41 +36,6 @@ jobs: - run: .kokoro/build.sh env: JOB_TYPE: test - units-with-multiplexed-session: - runs-on: ubuntu-latest - strategy: - fail-fast: false - matrix: - java: [ 11, 17, 21 ] - steps: - - uses: actions/checkout@v4 - - uses: actions/setup-java@v3 - with: - distribution: temurin - java-version: ${{matrix.java}} - - run: java -version - - run: .kokoro/build.sh - env: - JOB_TYPE: test - GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_PARTITIONED_OPS: true - GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_FOR_RW: true - units-with-regular-session: - runs-on: ubuntu-latest - strategy: - fail-fast: false - matrix: - java: [ 11, 17, 21 ] - steps: - - uses: actions/checkout@v4 - - uses: actions/setup-java@v3 - with: - distribution: temurin - java-version: ${{matrix.java}} - - run: java -version - - run: .kokoro/build.sh - env: - JOB_TYPE: test - GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS: false units-java8: # Building using Java 17 and run the tests with Java 8 runtime name: "units (8)" @@ -90,48 +55,6 @@ jobs: - run: .kokoro/build.sh env: JOB_TYPE: test - units-with-multiplexed-session8: - # Building using Java 17 and run the tests with Java 8 runtime - name: "units-with-multiplexed-session (8)" - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v4 - - uses: actions/setup-java@v3 - with: - java-version: 8 - distribution: temurin - - run: echo "SUREFIRE_JVM_OPT=-Djvm=${JAVA_HOME}/bin/java" >> $GITHUB_ENV - shell: bash - - uses: actions/setup-java@v3 - with: - java-version: 17 - distribution: temurin - - run: .kokoro/build.sh - env: - JOB_TYPE: test - GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS: true - GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_PARTITIONED_OPS: true - GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_FOR_RW: true - units-with-regular-session8: - # Building using Java 17 and run the tests with Java 8 runtime - name: "units-with-regular-session (8)" - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v4 - - uses: actions/setup-java@v3 - with: - java-version: 8 - distribution: temurin - - run: echo "SUREFIRE_JVM_OPT=-Djvm=${JAVA_HOME}/bin/java" >> $GITHUB_ENV - shell: bash - - uses: actions/setup-java@v3 - with: - java-version: 17 - distribution: temurin - - run: .kokoro/build.sh - env: - JOB_TYPE: test - GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS: false windows: runs-on: windows-latest steps: diff --git a/.github/workflows/integration-tests-against-emulator-with-regular-session.yaml b/.github/workflows/integration-tests-against-emulator-with-regular-session.yaml deleted file mode 100644 index 371620cf5a..0000000000 --- a/.github/workflows/integration-tests-against-emulator-with-regular-session.yaml +++ /dev/null @@ -1,42 +0,0 @@ -on: - push: - branches: - - main - pull_request: -name: integration-tests-against-emulator-with-multiplexed-session -jobs: - units: - runs-on: ubuntu-latest - - services: - emulator: - image: gcr.io/cloud-spanner-emulator/emulator:latest - ports: - - 9010:9010 - - 9020:9020 - - steps: - - uses: actions/checkout@v5 - - uses: stCarolas/setup-maven@v5 - with: - maven-version: 3.8.1 - # Build with JDK 11 and run tests with JDK 8 - - uses: actions/setup-java@v5 - with: - java-version: 11 - distribution: temurin - - name: Compiling main library - run: .kokoro/build.sh - - uses: actions/setup-java@v5 - with: - java-version: 8 - distribution: temurin - - name: Running tests - run: | - mvn -V -B -Dspanner.testenv.instance="" -Penable-integration-tests \ - -DtrimStackTrace=false -Dclirr.skip=true -Denforcer.skip=true \ - -Dmaven.main.skip=true -fae verify - env: - JOB_TYPE: test - SPANNER_EMULATOR_HOST: localhost:9010 - GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS: false \ No newline at end of file diff --git a/.kokoro/build.sh b/.kokoro/build.sh index eb66097a4f..caeff7a8b9 100755 --- a/.kokoro/build.sh +++ b/.kokoro/build.sh @@ -104,21 +104,6 @@ integration) verify RETURN_CODE=$? ;; -integration-regular-sessions) - mvn -B ${INTEGRATION_TEST_ARGS} \ - -ntp \ - -Penable-integration-tests \ - -Djava.net.preferIPv4Stack=true \ - -DtrimStackTrace=false \ - -Dclirr.skip=true \ - -Denforcer.skip=true \ - -Dmaven.main.skip=true \ - -Dspanner.gce.config.project_id=gcloud-devel \ - -Dspanner.testenv.instance=projects/gcloud-devel/instances/java-client-integration-tests-regular-sessions \ - -fae \ - verify - RETURN_CODE=$? - ;; integration-directpath-enabled) mvn -B ${INTEGRATION_TEST_ARGS} \ -ntp \ diff --git a/.kokoro/presubmit/integration-regular-sessions-enabled.cfg b/.kokoro/presubmit/integration-regular-sessions-enabled.cfg deleted file mode 100644 index b454868ebf..0000000000 --- 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-regular-sessions" -} - -# 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 f9b04136ec..7d083db211 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/BatchClientImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/BatchClientImpl.java index 6e8340784b..bc0ccb316d 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/BatchClientImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/BatchClientImpl.java @@ -270,10 +270,6 @@ private List partitionReadUsingIndex( } return partitions.build(); } catch (SpannerException e) { - if (!isFallback && maybeMarkUnimplementedForPartitionedOps(e)) { - return partitionReadUsingIndex( - partitionOptions, table, index, keys, columns, true, option); - } e.setRequestId(reqId); throw e; } @@ -330,32 +326,11 @@ private List partitionQuery( } return partitions.build(); } catch (SpannerException e) { - if (!isFallback && maybeMarkUnimplementedForPartitionedOps(e)) { - return partitionQuery(partitionOptions, statement, true, option); - } e.setRequestId(reqId); throw e; } } - boolean maybeMarkUnimplementedForPartitionedOps(SpannerException spannerException) { - if (MultiplexedSessionDatabaseClient.verifyErrorMessage( - spannerException, "Partitioned operations are not supported with multiplexed sessions")) { - synchronized (fallbackInitiated) { - if (!fallbackInitiated.get()) { - session.setFallbackSessionReference( - sessionClient.createSession().getSessionReference()); - sessionName = session.getName(); - initFallbackTransaction(); - unimplementedForPartitionedOps.set(true); - fallbackInitiated.set(true); - } - return true; - } - } - return false; - } - @Override public ResultSet execute(Partition partition) throws SpannerException { if (partition.getStatement() != null) { 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 40dbd710bd..a8a8e8edf7 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 debb07d7af..81e29cfda4 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 ef1f4a8894..d8baa71253 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,16 @@ 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 +259,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 +278,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 +290,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 +315,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 +344,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 +354,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 +379,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 +406,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 +561,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 20c86bdf25..5823f87e9f 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 68c37561c9..e9a68bfcee 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 42a67a6629..0000000000 --- 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 5e48d1b78b..0000000000 --- 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 605f96da74..51d0ca3d47 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 e96be9effa..1fd6c303ed 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 34fad2a69c..1bd3146225 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 @@ -342,25 +342,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; } @@ -370,27 +354,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 56d33c5487..5e34fcc7e7 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 59657026e5..0e06f9554c 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 76cc20cb4d..0000000000 --- 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 939114a7f6..0000000000 --- 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 4ff0675549..0000000000 --- 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 a06eeb9166..8177ee2527 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 7988e53f84..0000000000 --- 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 9ca5c17330..827476d57c 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 35712cd5b4..28580f5336 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 ab6dfea4d6..0000000000 --- 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 1448ebbc96..c3063f4d6c 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 f852fc2903..0000000000 --- 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 58eb423a5d..0000000000 --- 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 a47aecdccc..e6928dbcb3 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 0448656475..0ad4c6b82b 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 c6f7e22f28..f71fdfe37a 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/OpenTelemetrySpanTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/OpenTelemetrySpanTest.java index 732709a731..3b6c63ca35 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/OpenTelemetrySpanTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/OpenTelemetrySpanTest.java @@ -117,14 +117,6 @@ public class OpenTelemetrySpanTest { ImmutableList.of("Request for 1 multiplexed session returned 1 session"); private int expectedCreateMultiplexedSessionsRequestEventsCount = 1; - private List expectedBatchCreateSessionsRequestEvents = - ImmutableList.of("Requesting 2 sessions", "Request for 2 sessions returned 2 sessions"); - - private int expectedBatchCreateSessionsRequestEventsCount = 2; - - private List expectedBatchCreateSessionsEvents = ImmutableList.of("Creating 2 sessions"); - - private int expectedBatchCreateSessionsEventsCount = 1; private List expectedExecuteStreamingQueryEvents = ImmutableList.of("Starting/Resuming stream"); @@ -276,18 +268,10 @@ public void singleUse() { List expectedReadOnlyTransactionSingleUseEvents = getExpectedReadOnlyTransactionSingleUseEvents(); List expectedReadOnlyTransactionSpans = - isMultiplexedSessionsEnabled() - ? ImmutableList.of( - "CloudSpannerOperation.CreateMultiplexedSession", - "CloudSpannerOperation.BatchCreateSessionsRequest", - "CloudSpannerOperation.ExecuteStreamingQuery", - "CloudSpannerOperation.BatchCreateSessions", - "CloudSpanner.ReadOnlyTransaction") - : ImmutableList.of( - "CloudSpannerOperation.BatchCreateSessionsRequest", - "CloudSpannerOperation.ExecuteStreamingQuery", - "CloudSpannerOperation.BatchCreateSessions", - "CloudSpanner.ReadOnlyTransaction"); + ImmutableList.of( + "CloudSpannerOperation.CreateMultiplexedSession", + "CloudSpannerOperation.ExecuteStreamingQuery", + "CloudSpanner.ReadOnlyTransaction"); int expectedReadOnlyTransactionSingleUseEventsCount = expectedReadOnlyTransactionSingleUseEvents.size(); @@ -314,18 +298,6 @@ public void singleUse() { expectedCreateMultiplexedSessionsRequestEvents, expectedCreateMultiplexedSessionsRequestEventsCount); break; - case "CloudSpannerOperation.BatchCreateSessionsRequest": - verifyRequestEvents( - spanItem, - expectedBatchCreateSessionsRequestEvents, - expectedBatchCreateSessionsRequestEventsCount); - break; - case "CloudSpannerOperation.BatchCreateSessions": - verifyRequestEvents( - spanItem, - expectedBatchCreateSessionsEvents, - expectedBatchCreateSessionsEventsCount); - break; case "CloudSpannerOperation.ExecuteStreamingQuery": verifyRequestEvents( spanItem, @@ -361,31 +333,12 @@ private List getExpectedReadOnlyTransactionSingleUseEvents() { @Test public void multiUse() { List expectedReadOnlyTransactionSpans = - isMultiplexedSessionsEnabled() - ? ImmutableList.of( - "CloudSpannerOperation.CreateMultiplexedSession", - "CloudSpannerOperation.BatchCreateSessionsRequest", - "CloudSpannerOperation.ExecuteStreamingQuery", - "CloudSpannerOperation.BatchCreateSessions", - "CloudSpanner.ReadOnlyTransaction") - : ImmutableList.of( - "CloudSpannerOperation.BatchCreateSessionsRequest", - "CloudSpannerOperation.ExecuteStreamingQuery", - "CloudSpannerOperation.BatchCreateSessions", - "CloudSpanner.ReadOnlyTransaction"); - List expectedReadOnlyTransactionMultiUseEvents; - if (isMultiplexedSessionsEnabled()) { - expectedReadOnlyTransactionMultiUseEvents = - ImmutableList.of("Creating Transaction", "Transaction Creation Done"); - } else { - expectedReadOnlyTransactionMultiUseEvents = - ImmutableList.of( - "Acquiring session", - "Acquired session", - "Using Session", - "Creating Transaction", - "Transaction Creation Done"); - } + ImmutableList.of( + "CloudSpannerOperation.CreateMultiplexedSession", + "CloudSpannerOperation.ExecuteStreamingQuery", + "CloudSpanner.ReadOnlyTransaction"); + List expectedReadOnlyTransactionMultiUseEvents = + ImmutableList.of("Creating Transaction", "Transaction Creation Done"); int expectedReadOnlyTransactionMultiUseEventsCount = expectedReadOnlyTransactionMultiUseEvents.size(); @@ -411,18 +364,6 @@ public void multiUse() { expectedCreateMultiplexedSessionsRequestEvents, expectedCreateMultiplexedSessionsRequestEventsCount); break; - case "CloudSpannerOperation.BatchCreateSessionsRequest": - verifyRequestEvents( - spanItem, - expectedBatchCreateSessionsRequestEvents, - expectedBatchCreateSessionsRequestEventsCount); - break; - case "CloudSpannerOperation.BatchCreateSessions": - verifyRequestEvents( - spanItem, - expectedBatchCreateSessionsEvents, - expectedBatchCreateSessionsEventsCount); - break; case "CloudSpannerOperation.ExecuteStreamingQuery": verifyRequestEvents( spanItem, @@ -447,38 +388,27 @@ public void multiUse() { @Test public void transactionRunner() { List expectedReadWriteTransactionWithCommitSpans = - isMultiplexedSessionsEnabled() - ? ImmutableList.of( - "CloudSpannerOperation.CreateMultiplexedSession", - "CloudSpannerOperation.BatchCreateSessionsRequest", - "CloudSpannerOperation.ExecuteUpdate", - "CloudSpannerOperation.Commit", - "CloudSpannerOperation.BatchCreateSessions", - "CloudSpanner.ReadWriteTransaction") - : ImmutableList.of( - "CloudSpannerOperation.BatchCreateSessionsRequest", - "CloudSpannerOperation.ExecuteUpdate", - "CloudSpannerOperation.Commit", - "CloudSpannerOperation.BatchCreateSessions", - "CloudSpanner.ReadWriteTransaction"); - - if (isMultiplexedSessionsEnabledForRW()) { - expectedReadWriteTransactionEvents = - ImmutableList.of( - "Starting Transaction Attempt", - "Starting Commit", - "Commit Done", - "Transaction Attempt Succeeded"); - expectedReadWriteTransactionEventsCount = 4; - } + ImmutableList.of( + "CloudSpannerOperation.CreateMultiplexedSession", + "CloudSpannerOperation.ExecuteUpdate", + "CloudSpannerOperation.Commit", + "CloudSpanner.ReadWriteTransaction"); + + expectedReadWriteTransactionEvents = + ImmutableList.of( + "Starting Transaction Attempt", + "Starting Commit", + "Commit Done", + "Transaction Attempt Succeeded"); + expectedReadWriteTransactionEventsCount = 4; DatabaseClient client = getClient(); TransactionRunner runner = client.readWriteTransaction(); runner.run(transaction -> transaction.executeUpdate(UPDATE_STATEMENT)); - // Wait until the list of spans contains "CloudSpannerOperation.BatchCreateSessions", as this is + // Wait until the list of spans contains "CloudSpannerOperation.CreateSession", as this is // an async operation. Stopwatch stopwatch = Stopwatch.createStarted(); while (spanExporter.getFinishedSpanItems().stream() - .noneMatch(span -> span.getName().equals("CloudSpannerOperation.BatchCreateSessions")) + .noneMatch(span -> span.getName().equals("CloudSpannerOperation.CreateSession")) && stopwatch.elapsed(TimeUnit.MILLISECONDS) < 100) { Thread.yield(); } @@ -495,18 +425,6 @@ public void transactionRunner() { expectedCreateMultiplexedSessionsRequestEvents, expectedCreateMultiplexedSessionsRequestEventsCount); break; - case "CloudSpannerOperation.BatchCreateSessionsRequest": - verifyRequestEvents( - spanItem, - expectedBatchCreateSessionsRequestEvents, - expectedBatchCreateSessionsRequestEventsCount); - break; - case "CloudSpannerOperation.BatchCreateSessions": - verifyRequestEvents( - spanItem, - expectedBatchCreateSessionsEvents, - expectedBatchCreateSessionsEventsCount); - break; case "CloudSpannerOperation.Commit": case "CloudSpannerOperation.ExecuteUpdate": assertEquals(0, spanItem.getEvents().size()); @@ -529,26 +447,16 @@ public void transactionRunner() { @Test public void transactionRunnerWithError() { List expectedReadWriteTransactionSpans = - isMultiplexedSessionsEnabled() - ? ImmutableList.of( - "CloudSpannerOperation.CreateMultiplexedSession", - "CloudSpannerOperation.BatchCreateSessionsRequest", - "CloudSpannerOperation.BatchCreateSessions", - "CloudSpannerOperation.ExecuteUpdate", - "CloudSpanner.ReadWriteTransaction") - : ImmutableList.of( - "CloudSpannerOperation.BatchCreateSessionsRequest", - "CloudSpannerOperation.BatchCreateSessions", - "CloudSpannerOperation.ExecuteUpdate", - "CloudSpanner.ReadWriteTransaction"); - if (isMultiplexedSessionsEnabledForRW()) { - expectedReadWriteTransactionErrorEvents = - ImmutableList.of( - "Starting Transaction Attempt", - "Transaction Attempt Failed in user operation", - "exception"); - expectedReadWriteTransactionErrorEventsCount = 3; - } + ImmutableList.of( + "CloudSpannerOperation.CreateMultiplexedSession", + "CloudSpannerOperation.ExecuteUpdate", + "CloudSpanner.ReadWriteTransaction"); + expectedReadWriteTransactionErrorEvents = + ImmutableList.of( + "Starting Transaction Attempt", + "Transaction Attempt Failed in user operation", + "exception"); + expectedReadWriteTransactionErrorEventsCount = 3; DatabaseClient client = getClient(); TransactionRunner runner = client.readWriteTransaction(); SpannerException e = @@ -570,18 +478,6 @@ public void transactionRunnerWithError() { expectedCreateMultiplexedSessionsRequestEvents, expectedCreateMultiplexedSessionsRequestEventsCount); break; - case "CloudSpannerOperation.BatchCreateSessionsRequest": - verifyRequestEvents( - spanItem, - expectedBatchCreateSessionsRequestEvents, - expectedBatchCreateSessionsRequestEventsCount); - break; - case "CloudSpannerOperation.BatchCreateSessions": - verifyRequestEvents( - spanItem, - expectedBatchCreateSessionsEvents, - expectedBatchCreateSessionsEventsCount); - break; case "CloudSpanner.ReadWriteTransaction": verifyRequestEvents( spanItem, @@ -605,23 +501,19 @@ public void transactionRunnerWithFailedAndBeginTransaction() { List expectedReadWriteTransactionWithCommitAndBeginTransactionSpans = ImmutableList.of( "CloudSpannerOperation.BeginTransaction", - "CloudSpannerOperation.BatchCreateSessionsRequest", "CloudSpannerOperation.ExecuteUpdate", "CloudSpannerOperation.Commit", - "CloudSpannerOperation.BatchCreateSessions", "CloudSpanner.ReadWriteTransaction"); - if (isMultiplexedSessionsEnabledForRW()) { - expectedReadWriteTransactionErrorWithBeginTransactionEvents = - ImmutableList.of( - "Starting Transaction Attempt", - "Transaction Attempt Aborted in user operation. Retrying", - "Creating Transaction", - "Transaction Creation Done", - "Starting Commit", - "Commit Done", - "Transaction Attempt Succeeded"); - expectedReadWriteTransactionErrorWithBeginTransactionEventsCount = 8; - } + expectedReadWriteTransactionErrorWithBeginTransactionEvents = + ImmutableList.of( + "Starting Transaction Attempt", + "Transaction Attempt Aborted in user operation. Retrying", + "Creating Transaction", + "Transaction Creation Done", + "Starting Commit", + "Commit Done", + "Transaction Attempt Succeeded"); + expectedReadWriteTransactionErrorWithBeginTransactionEventsCount = 8; DatabaseClient client = getClient(); assertEquals( Long.valueOf(1L), @@ -641,7 +533,7 @@ public void transactionRunnerWithFailedAndBeginTransaction() { return transaction.executeUpdate(UPDATE_STATEMENT); })); // Wait for all spans to finish. Failing to do so can cause the test to miss the - // BatchCreateSessions span, as that span is executed asynchronously in the SessionClient, and + // CreateSession span, as that span is executed asynchronously in the SessionClient, and // the SessionClient returns the session to the pool before the span has finished fully. Stopwatch stopwatch = Stopwatch.createStarted(); while (spanExporter.getFinishedSpanItems().size() @@ -667,18 +559,6 @@ public void transactionRunnerWithFailedAndBeginTransaction() { expectedCreateMultiplexedSessionsRequestEvents, expectedCreateMultiplexedSessionsRequestEventsCount); break; - case "CloudSpannerOperation.BatchCreateSessionsRequest": - verifyRequestEvents( - spanItem, - expectedBatchCreateSessionsRequestEvents, - expectedBatchCreateSessionsRequestEventsCount); - break; - case "CloudSpannerOperation.BatchCreateSessions": - verifyRequestEvents( - spanItem, - expectedBatchCreateSessionsEvents, - expectedBatchCreateSessionsEventsCount); - break; case "CloudSpannerOperation.Commit": case "CloudSpannerOperation.BeginTransaction": case "CloudSpannerOperation.ExecuteUpdate": @@ -718,7 +598,7 @@ public void testTransactionRunnerWithRetryOnBeginTransaction() { }); assertEquals(2, mockSpanner.countRequestsOfType(BeginTransactionRequest.class)); - int numExpectedSpans = isMultiplexedSessionsEnabled() ? 10 : 8; + int numExpectedSpans = 7; waitForFinishedSpans(numExpectedSpans); List finishedSpans = spanExporter.getFinishedSpanItems(); List finishedSpanNames = @@ -731,13 +611,7 @@ public void testTransactionRunnerWithRetryOnBeginTransaction() { assertTrue( actualSpanNames, finishedSpanNames.contains("CloudSpannerOperation.BeginTransaction")); assertTrue(actualSpanNames, finishedSpanNames.contains("CloudSpannerOperation.Commit")); - assertTrue( - actualSpanNames, finishedSpanNames.contains("CloudSpannerOperation.BatchCreateSessions")); - assertTrue( - actualSpanNames, - finishedSpanNames.contains("CloudSpannerOperation.BatchCreateSessionsRequest")); - assertTrue(actualSpanNames, finishedSpanNames.contains("Spanner.BatchCreateSessions")); assertTrue(actualSpanNames, finishedSpanNames.contains("Spanner.BeginTransaction")); assertTrue(actualSpanNames, finishedSpanNames.contains("Spanner.Commit")); @@ -768,7 +642,7 @@ public void testSingleUseRetryOnExecuteStreamingSql() { } assertEquals(2, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)); - int numExpectedSpans = isMultiplexedSessionsEnabled() ? 9 : 7; + int numExpectedSpans = 6; waitForFinishedSpans(numExpectedSpans); List finishedSpans = spanExporter.getFinishedSpanItems(); List finishedSpanNames = @@ -780,13 +654,7 @@ public void testSingleUseRetryOnExecuteStreamingSql() { assertTrue(actualSpanNames, finishedSpanNames.contains("CloudSpanner.ReadOnlyTransaction")); assertTrue( actualSpanNames, finishedSpanNames.contains("CloudSpannerOperation.ExecuteStreamingQuery")); - assertTrue( - actualSpanNames, finishedSpanNames.contains("CloudSpannerOperation.BatchCreateSessions")); - assertTrue( - actualSpanNames, - finishedSpanNames.contains("CloudSpannerOperation.BatchCreateSessionsRequest")); - assertTrue(actualSpanNames, finishedSpanNames.contains("Spanner.BatchCreateSessions")); assertTrue(actualSpanNames, finishedSpanNames.contains("Spanner.ExecuteStreamingSql")); // UNAVAILABLE errors on ExecuteStreamingSql are handled manually in the client library, which @@ -817,7 +685,7 @@ public void testRetryOnExecuteSql() { .run(transaction -> transaction.executeUpdate(UPDATE_STATEMENT)); assertEquals(2, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)); - int numExpectedSpans = isMultiplexedSessionsEnabled() ? 10 : 8; + int numExpectedSpans = 7; waitForFinishedSpans(numExpectedSpans); List finishedSpans = spanExporter.getFinishedSpanItems(); List finishedSpanNames = @@ -828,13 +696,6 @@ public void testRetryOnExecuteSql() { assertTrue(actualSpanNames, finishedSpanNames.contains("CloudSpanner.ReadWriteTransaction")); assertTrue(actualSpanNames, finishedSpanNames.contains("CloudSpannerOperation.Commit")); - assertTrue( - actualSpanNames, finishedSpanNames.contains("CloudSpannerOperation.BatchCreateSessions")); - assertTrue( - actualSpanNames, - finishedSpanNames.contains("CloudSpannerOperation.BatchCreateSessionsRequest")); - - assertTrue(actualSpanNames, finishedSpanNames.contains("Spanner.BatchCreateSessions")); assertTrue(actualSpanNames, finishedSpanNames.contains("Spanner.ExecuteSql")); assertTrue(actualSpanNames, finishedSpanNames.contains("Spanner.Commit")); 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 3704b11890..ede202ced3 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 6a0f16c02b..0000000000 --- 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 3106bd1652..2e9d4185cb 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 e18cddd3bf..e1f93d334d 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 4415ba7d70..0000000000 --- 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 080091e661..0000000000 --- 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 0370f5420e..0000000000 --- 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 99a773eeb0..0000000000 --- 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 276a3ac813..0000000000 --- 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 3377196282..0000000000 --- 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 8d00f0889b..0000000000 --- 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 5a9365eaed..0000000000 --- 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/SpanTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpanTest.java index 38aa31ad0b..3957697a19 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpanTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpanTest.java @@ -306,39 +306,16 @@ private void verifySingleUseSpans() { // OpenCensus spans and events verification Map spans = failOnOverkillTraceComponent.getSpans(); assertThat(spans).containsEntry("CloudSpanner.ReadOnlyTransaction", true); - assertThat(spans).containsEntry("CloudSpannerOperation.BatchCreateSessions", true); - assertThat(spans).containsEntry("CloudSpannerOperation.BatchCreateSessionsRequest", true); assertThat(spans).containsEntry("CloudSpannerOperation.ExecuteStreamingQuery", true); - List expectedAnnotations = - ImmutableList.of( - "Requesting 2 sessions", - "Request for 2 sessions returned 2 sessions", - "Creating 2 sessions", - "Acquiring session", - "Acquired session", - "Using Session", - "Starting/Resuming stream"); List expectedAnnotationsForMultiplexedSession = ImmutableList.of( - "Requesting 2 sessions", - "Request for 2 sessions returned 2 sessions", - "Request for 1 multiplexed session returned 1 session", - "Creating 2 sessions", - "Starting/Resuming stream"); - if (spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSession()) { - verifyAnnotations( - failOnOverkillTraceComponent.getAnnotations().stream() - .distinct() - .collect(Collectors.toList()), - expectedAnnotationsForMultiplexedSession); - } else { - verifyAnnotations( - failOnOverkillTraceComponent.getAnnotations().stream() - .distinct() - .collect(Collectors.toList()), - expectedAnnotations); - } + "Request for 1 multiplexed session returned 1 session", "Starting/Resuming stream"); + verifyAnnotations( + failOnOverkillTraceComponent.getAnnotations().stream() + .distinct() + .collect(Collectors.toList()), + expectedAnnotationsForMultiplexedSession); } @Test @@ -359,41 +336,18 @@ public void singleUseWithError() { // OpenCensus spans and events verification Map spans = failOnOverkillTraceComponent.getSpans(); assertThat(spans).containsEntry("CloudSpanner.ReadOnlyTransaction", true); - assertThat(spans).containsEntry("CloudSpannerOperation.BatchCreateSessions", true); - assertThat(spans).containsEntry("CloudSpannerOperation.BatchCreateSessionsRequest", true); assertThat(spans).containsEntry("CloudSpannerOperation.ExecuteStreamingQuery", true); - List expectedAnnotations = - ImmutableList.of( - "Requesting 2 sessions", - "Request for 2 sessions returned 2 sessions", - "Creating 2 sessions", - "Acquiring session", - "Acquired session", - "Using Session", - "Starting/Resuming stream", - "Stream broken. Not safe to retry"); List expectedAnnotationsForMultiplexedSession = ImmutableList.of( - "Requesting 2 sessions", - "Request for 2 sessions returned 2 sessions", "Request for 1 multiplexed session returned 1 session", - "Creating 2 sessions", "Starting/Resuming stream", "Stream broken. Not safe to retry"); - if (spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSession()) { - verifyAnnotations( - failOnOverkillTraceComponent.getAnnotations().stream() - .distinct() - .collect(Collectors.toList()), - expectedAnnotationsForMultiplexedSession); - } else { - verifyAnnotations( - failOnOverkillTraceComponent.getAnnotations().stream() - .distinct() - .collect(Collectors.toList()), - expectedAnnotations); - } + verifyAnnotations( + failOnOverkillTraceComponent.getAnnotations().stream() + .distinct() + .collect(Collectors.toList()), + expectedAnnotationsForMultiplexedSession); } @Test @@ -408,43 +362,19 @@ public void multiUse() { Map spans = failOnOverkillTraceComponent.getSpans(); assertThat(spans).containsEntry("CloudSpanner.ReadOnlyTransaction", true); - assertThat(spans).containsEntry("CloudSpannerOperation.BatchCreateSessions", true); - assertThat(spans).containsEntry("CloudSpannerOperation.BatchCreateSessionsRequest", true); assertThat(spans).containsEntry("CloudSpannerOperation.ExecuteStreamingQuery", true); - List expectedAnnotations = - ImmutableList.of( - "Requesting 2 sessions", - "Request for 2 sessions returned 2 sessions", - "Creating 2 sessions", - "Acquiring session", - "Acquired session", - "Using Session", - "Starting/Resuming stream", - "Creating Transaction", - "Transaction Creation Done"); List expectedAnnotationsForMultiplexedSession = ImmutableList.of( - "Requesting 2 sessions", - "Request for 2 sessions returned 2 sessions", "Request for 1 multiplexed session returned 1 session", - "Creating 2 sessions", "Starting/Resuming stream", "Creating Transaction", "Transaction Creation Done"); - if (isMultiplexedSessionsEnabled()) { - verifyAnnotations( - failOnOverkillTraceComponent.getAnnotations().stream() - .distinct() - .collect(Collectors.toList()), - expectedAnnotationsForMultiplexedSession); - } else { - verifyAnnotations( - failOnOverkillTraceComponent.getAnnotations().stream() - .distinct() - .collect(Collectors.toList()), - expectedAnnotations); - } + verifyAnnotations( + failOnOverkillTraceComponent.getAnnotations().stream() + .distinct() + .collect(Collectors.toList()), + expectedAnnotationsForMultiplexedSession); } @Test @@ -453,64 +383,20 @@ public void transactionRunner() { runner.run(transaction -> transaction.executeUpdate(UPDATE_STATEMENT)); Map spans = failOnOverkillTraceComponent.getSpans(); assertThat(spans).containsEntry("CloudSpanner.ReadWriteTransaction", true); - assertThat(spans).containsEntry("CloudSpannerOperation.BatchCreateSessions", true); - assertThat(spans).containsEntry("CloudSpannerOperation.BatchCreateSessionsRequest", true); assertThat(spans).containsEntry("CloudSpannerOperation.Commit", true); - List expectedAnnotations = - ImmutableList.of( - "Acquiring session", - "Acquired session", - "Using Session", - "Starting Transaction Attempt", - "Starting Commit", - "Commit Done", - "Transaction Attempt Succeeded", - "Requesting 2 sessions", - "Request for 2 sessions returned 2 sessions", - "Creating 2 sessions"); - List expectedAnnotationsForMultiplexedSession = - ImmutableList.of( - "Acquiring session", - "Acquired session", - "Using Session", - "Starting Transaction Attempt", - "Starting Commit", - "Commit Done", - "Transaction Attempt Succeeded", - "Requesting 2 sessions", - "Request for 2 sessions returned 2 sessions", - "Request for 1 multiplexed session returned 1 session", - "Creating 2 sessions"); List expectedAnnotationsForMultiplexedSessionsRW = ImmutableList.of( "Starting Transaction Attempt", "Starting Commit", "Commit Done", "Transaction Attempt Succeeded", - "Requesting 2 sessions", - "Request for 2 sessions returned 2 sessions", - "Request for 1 multiplexed session returned 1 session", - "Creating 2 sessions"); - if (isMultiplexedSessionsEnabledForRW()) { - verifyAnnotations( - failOnOverkillTraceComponent.getAnnotations().stream() - .distinct() - .collect(Collectors.toList()), - expectedAnnotationsForMultiplexedSessionsRW); - } else if (isMultiplexedSessionsEnabled()) { - verifyAnnotations( - failOnOverkillTraceComponent.getAnnotations().stream() - .distinct() - .collect(Collectors.toList()), - expectedAnnotationsForMultiplexedSession); - } else { - verifyAnnotations( - failOnOverkillTraceComponent.getAnnotations().stream() - .distinct() - .collect(Collectors.toList()), - expectedAnnotations); - } + "Request for 1 multiplexed session returned 1 session"); + verifyAnnotations( + failOnOverkillTraceComponent.getAnnotations().stream() + .distinct() + .collect(Collectors.toList()), + expectedAnnotationsForMultiplexedSessionsRW); } @Test @@ -524,65 +410,21 @@ public void transactionRunnerWithError() { Map spans = failOnOverkillTraceComponent.getSpans(); - if (isMultiplexedSessionsEnabled()) { - assertEquals(spans.toString(), 5, spans.size()); - assertThat(spans).containsEntry("CloudSpannerOperation.CreateMultiplexedSession", true); - } else { - assertThat(spans.size()).isEqualTo(4); - } + assertEquals(spans.toString(), 3, spans.size()); + assertThat(spans).containsEntry("CloudSpannerOperation.CreateMultiplexedSession", true); assertThat(spans).containsEntry("CloudSpanner.ReadWriteTransaction", true); assertThat(spans).containsEntry("CloudSpannerOperation.ExecuteUpdate", true); - assertThat(spans).containsEntry("CloudSpannerOperation.BatchCreateSessions", true); - assertThat(spans).containsEntry("CloudSpannerOperation.BatchCreateSessionsRequest", true); - List expectedAnnotations = - ImmutableList.of( - "Acquiring session", - "Acquired session", - "Using Session", - "Starting Transaction Attempt", - "Transaction Attempt Failed in user operation", - "Requesting 2 sessions", - "Request for 2 sessions returned 2 sessions", - "Creating 2 sessions"); - List expectedAnnotationsForMultiplexedSession = - ImmutableList.of( - "Acquiring session", - "Acquired session", - "Using Session", - "Starting Transaction Attempt", - "Transaction Attempt Failed in user operation", - "Requesting 2 sessions", - "Request for 1 multiplexed session returned 1 session", - "Request for 2 sessions returned 2 sessions", - "Creating 2 sessions"); List expectedAnnotationsForMultiplexedSessionsRW = ImmutableList.of( "Starting Transaction Attempt", "Transaction Attempt Failed in user operation", - "Requesting 2 sessions", - "Request for 1 multiplexed session returned 1 session", - "Request for 2 sessions returned 2 sessions", - "Creating 2 sessions"); - if (isMultiplexedSessionsEnabledForRW()) { - verifyAnnotations( - failOnOverkillTraceComponent.getAnnotations().stream() - .distinct() - .collect(Collectors.toList()), - expectedAnnotationsForMultiplexedSessionsRW); - } else if (isMultiplexedSessionsEnabled()) { - verifyAnnotations( - failOnOverkillTraceComponent.getAnnotations().stream() - .distinct() - .collect(Collectors.toList()), - expectedAnnotationsForMultiplexedSession); - } else { - verifyAnnotations( - failOnOverkillTraceComponent.getAnnotations().stream() - .distinct() - .collect(Collectors.toList()), - expectedAnnotations); - } + "Request for 1 multiplexed session returned 1 session"); + verifyAnnotations( + failOnOverkillTraceComponent.getAnnotations().stream() + .distinct() + .collect(Collectors.toList()), + expectedAnnotationsForMultiplexedSessionsRW); } private void verifyAnnotations(List actualAnnotations, List expectedAnnotations) { 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 3cf13dc58d..a675605e76 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 2325f2ac40..b0a79fe564 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 72ef90c9cf..c8469ae08a 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 6ffb0e1ca6..0000000000 --- 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()); - } - } -}