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 27253bf1e13..ebfb0e0a774 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 @@ -40,17 +40,6 @@ public String getDatabaseRole() { throw new UnsupportedOperationException(); } - @Override - public Timestamp write(Iterable mutations) throws SpannerException { - throw new UnsupportedOperationException(); - } - - @Override - public CommitResponse writeWithOptions(Iterable mutations, TransactionOption... options) - throws SpannerException { - throw new UnsupportedOperationException(); - } - @Override public Timestamp writeAtLeastOnce(Iterable mutations) throws SpannerException { return writeAtLeastOnceWithOptions(mutations).getCommitTimestamp(); @@ -63,26 +52,6 @@ public ServerStream batchWriteAtLeastOnce( throw new UnsupportedOperationException(); } - @Override - public TransactionRunner readWriteTransaction(TransactionOption... options) { - throw new UnsupportedOperationException(); - } - - @Override - public TransactionManager transactionManager(TransactionOption... options) { - throw new UnsupportedOperationException(); - } - - @Override - public AsyncRunner runAsync(TransactionOption... options) { - throw new UnsupportedOperationException(); - } - - @Override - public AsyncTransactionManager transactionManagerAsync(TransactionOption... options) { - throw new UnsupportedOperationException(); - } - @Override public long executePartitionedUpdate(Statement stmt, UpdateOption... options) { throw new UnsupportedOperationException(); diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManagerImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManagerImpl.java index 8b20dd824a0..882f709966b 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManagerImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManagerImpl.java @@ -66,6 +66,9 @@ public ApiFuture closeAsync() { if (txn != null) { txn.close(); } + if (session != null) { + session.onTransactionDone(); + } return MoreObjects.firstNonNull(res, ApiFutures.immediateFuture(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 909d731818f..b9d1ce054d7 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 @@ -36,6 +36,7 @@ class DatabaseClientImpl implements DatabaseClient { @VisibleForTesting final String clientId; @VisibleForTesting final SessionPool pool; @VisibleForTesting final MultiplexedSessionDatabaseClient multiplexedSessionDatabaseClient; + @VisibleForTesting final boolean useMultiplexedSessionForRW; final boolean useMultiplexedSessionBlindWrite; @@ -46,7 +47,8 @@ class DatabaseClientImpl implements DatabaseClient { pool, /* useMultiplexedSessionBlindWrite = */ false, /* multiplexedSessionDatabaseClient = */ null, - tracer); + tracer, + /* useMultiplexedSessionForRW = */ false); } @VisibleForTesting @@ -56,7 +58,8 @@ class DatabaseClientImpl implements DatabaseClient { pool, /* useMultiplexedSessionBlindWrite = */ false, /* multiplexedSessionDatabaseClient = */ null, - tracer); + tracer, + false); } DatabaseClientImpl( @@ -64,12 +67,14 @@ class DatabaseClientImpl implements DatabaseClient { SessionPool pool, boolean useMultiplexedSessionBlindWrite, @Nullable MultiplexedSessionDatabaseClient multiplexedSessionDatabaseClient, - TraceWrapper tracer) { + TraceWrapper tracer, + boolean useMultiplexedSessionForRW) { this.clientId = clientId; this.pool = pool; this.useMultiplexedSessionBlindWrite = useMultiplexedSessionBlindWrite; this.multiplexedSessionDatabaseClient = multiplexedSessionDatabaseClient; this.tracer = tracer; + this.useMultiplexedSessionForRW = useMultiplexedSessionForRW; } @VisibleForTesting @@ -85,6 +90,14 @@ DatabaseClient getMultiplexedSession() { return pool.getMultiplexedSessionWithFallback(); } + @VisibleForTesting + DatabaseClient getMultiplexedSessionForRW() { + if (this.useMultiplexedSessionForRW) { + return getMultiplexedSession(); + } + return getSession(); + } + private MultiplexedSessionDatabaseClient getMultiplexedSessionDatabaseClient() { return canUseMultiplexedSessions() ? this.multiplexedSessionDatabaseClient : null; } @@ -116,6 +129,9 @@ public CommitResponse writeWithOptions( throws SpannerException { ISpan span = tracer.spanBuilder(READ_WRITE_TRANSACTION, options); try (IScope s = tracer.withSpan(span)) { + if (this.useMultiplexedSessionForRW && getMultiplexedSessionDatabaseClient() != null) { + return getMultiplexedSessionDatabaseClient().writeWithOptions(mutations, options); + } return runWithSessionRetry(session -> session.writeWithOptions(mutations, options)); } catch (RuntimeException e) { span.setStatus(e); @@ -241,7 +257,7 @@ public ReadOnlyTransaction readOnlyTransaction(TimestampBound bound) { public TransactionRunner readWriteTransaction(TransactionOption... options) { ISpan span = tracer.spanBuilder(READ_WRITE_TRANSACTION, options); try (IScope s = tracer.withSpan(span)) { - return getSession().readWriteTransaction(options); + return getMultiplexedSessionForRW().readWriteTransaction(options); } catch (RuntimeException e) { span.setStatus(e); span.end(); @@ -253,7 +269,7 @@ public TransactionRunner readWriteTransaction(TransactionOption... options) { public TransactionManager transactionManager(TransactionOption... options) { ISpan span = tracer.spanBuilder(READ_WRITE_TRANSACTION, options); try (IScope s = tracer.withSpan(span)) { - return getSession().transactionManager(options); + return getMultiplexedSessionForRW().transactionManager(options); } catch (RuntimeException e) { span.setStatus(e); span.end(); @@ -265,7 +281,7 @@ public TransactionManager transactionManager(TransactionOption... options) { public AsyncRunner runAsync(TransactionOption... options) { ISpan span = tracer.spanBuilder(READ_WRITE_TRANSACTION, options); try (IScope s = tracer.withSpan(span)) { - return getSession().runAsync(options); + return getMultiplexedSessionForRW().runAsync(options); } catch (RuntimeException e) { span.setStatus(e); span.end(); @@ -277,7 +293,7 @@ public AsyncRunner runAsync(TransactionOption... options) { public AsyncTransactionManager transactionManagerAsync(TransactionOption... options) { ISpan span = tracer.spanBuilder(READ_WRITE_TRANSACTION, options); try (IScope s = tracer.withSpan(span)) { - return getSession().transactionManagerAsync(options); + return getMultiplexedSessionForRW().transactionManagerAsync(options); } catch (RuntimeException e) { span.setStatus(e); span.end(); diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DelayedAsyncRunner.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DelayedAsyncRunner.java new file mode 100644 index 00000000000..3783a84903f --- /dev/null +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DelayedAsyncRunner.java @@ -0,0 +1,74 @@ +/* + * 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 com.google.api.core.ApiFuture; +import com.google.api.core.ApiFutures; +import com.google.cloud.Timestamp; +import com.google.common.util.concurrent.MoreExecutors; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executor; + +/** + * Represents a {@link AsyncRunner} using a multiplexed session that is not yet ready. The execution + * will be delayed until the multiplexed session has been created and is ready. This class is only + * used during the startup of the client and the multiplexed session has not yet been created. + */ +public class DelayedAsyncRunner implements AsyncRunner { + + private final ApiFuture asyncRunnerFuture; + + public DelayedAsyncRunner(ApiFuture asyncRunnerFuture) { + this.asyncRunnerFuture = asyncRunnerFuture; + } + + ApiFuture getAsyncRunner() { + return ApiFutures.catchingAsync( + asyncRunnerFuture, + Exception.class, + exception -> { + if (exception instanceof InterruptedException) { + throw SpannerExceptionFactory.propagateInterrupt((InterruptedException) exception); + } + if (exception instanceof ExecutionException) { + throw SpannerExceptionFactory.causeAsRunTimeException((ExecutionException) exception); + } + throw exception; + }, + MoreExecutors.directExecutor()); + } + + @Override + public ApiFuture runAsync(AsyncWork work, Executor executor) { + return ApiFutures.transformAsync( + getAsyncRunner(), + asyncRunner -> asyncRunner.runAsync(work, executor), + MoreExecutors.directExecutor()); + } + + @Override + public ApiFuture getCommitTimestamp() { + return ApiFutures.transformAsync( + getAsyncRunner(), AsyncRunner::getCommitTimestamp, MoreExecutors.directExecutor()); + } + + @Override + public ApiFuture getCommitResponse() { + return ApiFutures.transformAsync( + getAsyncRunner(), AsyncRunner::getCommitResponse, MoreExecutors.directExecutor()); + } +} diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DelayedAsyncTransactionManager.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DelayedAsyncTransactionManager.java new file mode 100644 index 00000000000..56b874e4a87 --- /dev/null +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DelayedAsyncTransactionManager.java @@ -0,0 +1,82 @@ +/* + * 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 com.google.api.core.ApiFuture; +import com.google.cloud.spanner.TransactionManager.TransactionState; +import java.util.concurrent.ExecutionException; + +/** + * Represents a {@link AsyncTransactionManager} using a multiplexed session that is not yet ready. + * The execution will be delayed until the multiplexed session has been created and is ready. This + * class is only used during the startup of the client and the multiplexed session has not yet been + * created. + */ +public class DelayedAsyncTransactionManager implements AsyncTransactionManager { + + private final ApiFuture asyncTransactionManagerApiFuture; + + DelayedAsyncTransactionManager( + ApiFuture asyncTransactionManagerApiFuture) { + this.asyncTransactionManagerApiFuture = asyncTransactionManagerApiFuture; + } + + AsyncTransactionManager getAsyncTransactionManager() { + try { + return this.asyncTransactionManagerApiFuture.get(); + } catch (ExecutionException executionException) { + throw SpannerExceptionFactory.causeAsRunTimeException(executionException); + } catch (InterruptedException interruptedException) { + throw SpannerExceptionFactory.propagateInterrupt(interruptedException); + } + } + + @Override + public TransactionContextFuture beginAsync() { + return getAsyncTransactionManager().beginAsync(); + } + + @Override + public ApiFuture rollbackAsync() { + return getAsyncTransactionManager().rollbackAsync(); + } + + @Override + public TransactionContextFuture resetForRetryAsync() { + return getAsyncTransactionManager().resetForRetryAsync(); + } + + @Override + public TransactionState getState() { + return getAsyncTransactionManager().getState(); + } + + @Override + public ApiFuture getCommitResponse() { + return getAsyncTransactionManager().getCommitResponse(); + } + + @Override + public void close() { + getAsyncTransactionManager().close(); + } + + @Override + public ApiFuture closeAsync() { + return getAsyncTransactionManager().closeAsync(); + } +} 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 36750eaccd1..ad3e6b0cf70 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 @@ -20,6 +20,7 @@ import com.google.api.core.ApiFuture; import com.google.api.core.ApiFutures; +import com.google.cloud.Timestamp; import com.google.cloud.spanner.DelayedReadContext.DelayedReadOnlyTransaction; import com.google.cloud.spanner.MultiplexedSessionDatabaseClient.MultiplexedSessionTransaction; import com.google.cloud.spanner.Options.TransactionOption; @@ -57,7 +58,7 @@ public ReadContext singleUse() { this.sessionFuture, sessionReference -> new MultiplexedSessionTransaction( - client, span, sessionReference, NO_CHANNEL_HINT, true) + client, span, sessionReference, NO_CHANNEL_HINT, /* singleUse = */ true) .singleUse(), MoreExecutors.directExecutor())); } @@ -69,7 +70,7 @@ public ReadContext singleUse(TimestampBound bound) { this.sessionFuture, sessionReference -> new MultiplexedSessionTransaction( - client, span, sessionReference, NO_CHANNEL_HINT, true) + client, span, sessionReference, NO_CHANNEL_HINT, /* singleUse = */ true) .singleUse(bound), MoreExecutors.directExecutor())); } @@ -81,7 +82,7 @@ public ReadOnlyTransaction singleUseReadOnlyTransaction() { this.sessionFuture, sessionReference -> new MultiplexedSessionTransaction( - client, span, sessionReference, NO_CHANNEL_HINT, true) + client, span, sessionReference, NO_CHANNEL_HINT, /* singleUse = */ true) .singleUseReadOnlyTransaction(), MoreExecutors.directExecutor())); } @@ -93,7 +94,7 @@ public ReadOnlyTransaction singleUseReadOnlyTransaction(TimestampBound bound) { this.sessionFuture, sessionReference -> new MultiplexedSessionTransaction( - client, span, sessionReference, NO_CHANNEL_HINT, true) + client, span, sessionReference, NO_CHANNEL_HINT, /* singleUse = */ true) .singleUseReadOnlyTransaction(bound), MoreExecutors.directExecutor())); } @@ -105,7 +106,7 @@ public ReadOnlyTransaction readOnlyTransaction() { this.sessionFuture, sessionReference -> new MultiplexedSessionTransaction( - client, span, sessionReference, NO_CHANNEL_HINT, false) + client, span, sessionReference, NO_CHANNEL_HINT, /* singleUse = */ false) .readOnlyTransaction(), MoreExecutors.directExecutor())); } @@ -117,7 +118,7 @@ public ReadOnlyTransaction readOnlyTransaction(TimestampBound bound) { this.sessionFuture, sessionReference -> new MultiplexedSessionTransaction( - client, span, sessionReference, NO_CHANNEL_HINT, false) + client, span, sessionReference, NO_CHANNEL_HINT, /* singleUse = */ false) .readOnlyTransaction(bound), MoreExecutors.directExecutor())); } @@ -131,11 +132,85 @@ public CommitResponse writeAtLeastOnceWithOptions( Iterable mutations, TransactionOption... options) throws SpannerException { SessionReference sessionReference = getSessionReference(); try (MultiplexedSessionTransaction transaction = - new MultiplexedSessionTransaction(client, span, sessionReference, NO_CHANNEL_HINT, true)) { + new MultiplexedSessionTransaction( + client, span, sessionReference, NO_CHANNEL_HINT, /* singleUse = */ true)) { return transaction.writeAtLeastOnceWithOptions(mutations, options); } } + // This is a blocking method, as the interface that it implements is also defined as a blocking + // method. + @Override + public Timestamp write(Iterable mutations) throws SpannerException { + SessionReference sessionReference = getSessionReference(); + try (MultiplexedSessionTransaction transaction = + new MultiplexedSessionTransaction( + client, span, sessionReference, NO_CHANNEL_HINT, /* singleUse = */ false)) { + return transaction.write(mutations); + } + } + + // This is a blocking method, as the interface that it implements is also defined as a blocking + // method. + @Override + public CommitResponse writeWithOptions(Iterable mutations, TransactionOption... options) + throws SpannerException { + SessionReference sessionReference = getSessionReference(); + try (MultiplexedSessionTransaction transaction = + new MultiplexedSessionTransaction( + client, span, sessionReference, NO_CHANNEL_HINT, /* singleUse = */ false)) { + return transaction.writeWithOptions(mutations, options); + } + } + + @Override + public TransactionRunner readWriteTransaction(TransactionOption... options) { + return new DelayedTransactionRunner( + ApiFutures.transform( + this.sessionFuture, + sessionReference -> + new MultiplexedSessionTransaction( + client, span, sessionReference, NO_CHANNEL_HINT, /* singleUse = */ false) + .readWriteTransaction(options), + MoreExecutors.directExecutor())); + } + + @Override + public TransactionManager transactionManager(TransactionOption... options) { + return new DelayedTransactionManager( + ApiFutures.transform( + this.sessionFuture, + sessionReference -> + new MultiplexedSessionTransaction( + client, span, sessionReference, NO_CHANNEL_HINT, /* singleUse = */ false) + .transactionManager(options), + MoreExecutors.directExecutor())); + } + + @Override + public AsyncRunner runAsync(TransactionOption... options) { + return new DelayedAsyncRunner( + ApiFutures.transform( + this.sessionFuture, + sessionReference -> + new MultiplexedSessionTransaction( + client, span, sessionReference, NO_CHANNEL_HINT, /* singleUse = */ false) + .runAsync(options), + MoreExecutors.directExecutor())); + } + + @Override + public AsyncTransactionManager transactionManagerAsync(TransactionOption... options) { + return new DelayedAsyncTransactionManager( + ApiFutures.transform( + this.sessionFuture, + sessionReference -> + new MultiplexedSessionTransaction( + client, span, sessionReference, NO_CHANNEL_HINT, /* singleUse = */ false) + .transactionManagerAsync(options), + MoreExecutors.directExecutor())); + } + /** * Gets the session reference that this delayed transaction is waiting for. This method should * only be called by methods that are allowed to be blocking. @@ -144,12 +219,7 @@ private SessionReference getSessionReference() { try { return this.sessionFuture.get(); } catch (ExecutionException executionException) { - // Propagate the underlying exception as a RuntimeException (SpannerException is also a - // RuntimeException). - if (executionException.getCause() instanceof RuntimeException) { - throw (RuntimeException) executionException.getCause(); - } - throw SpannerExceptionFactory.asSpannerException(executionException.getCause()); + throw SpannerExceptionFactory.causeAsRunTimeException(executionException); } catch (InterruptedException interruptedException) { throw SpannerExceptionFactory.propagateInterrupt(interruptedException); } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DelayedReadContext.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DelayedReadContext.java index 62bc5711852..86d6ae079d3 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DelayedReadContext.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DelayedReadContext.java @@ -44,12 +44,7 @@ T getReadContext() { try { return this.readContextFuture.get(); } catch (ExecutionException executionException) { - // Propagate the underlying exception as a RuntimeException (SpannerException is also a - // RuntimeException). - if (executionException.getCause() instanceof RuntimeException) { - throw (RuntimeException) executionException.getCause(); - } - throw SpannerExceptionFactory.asSpannerException(executionException.getCause()); + throw SpannerExceptionFactory.causeAsRunTimeException(executionException); } catch (InterruptedException interruptedException) { throw SpannerExceptionFactory.propagateInterrupt(interruptedException); } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DelayedTransactionManager.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DelayedTransactionManager.java new file mode 100644 index 00000000000..29eae6477fc --- /dev/null +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DelayedTransactionManager.java @@ -0,0 +1,86 @@ +/* + * 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 com.google.api.core.ApiFuture; +import com.google.cloud.Timestamp; +import java.util.concurrent.ExecutionException; + +/** + * Represents a {@link TransactionManager} using a multiplexed session that is not yet ready. The + * execution will be delayed until the multiplexed session has been created and is ready. This class + * is only used during the startup of the client and the multiplexed session has not yet been + * created. + */ +class DelayedTransactionManager implements TransactionManager { + + private final ApiFuture transactionManagerFuture; + + DelayedTransactionManager(ApiFuture transactionManagerFuture) { + this.transactionManagerFuture = transactionManagerFuture; + } + + TransactionManager getTransactionManager() { + try { + return this.transactionManagerFuture.get(); + } catch (ExecutionException executionException) { + throw SpannerExceptionFactory.causeAsRunTimeException(executionException); + } catch (InterruptedException interruptedException) { + throw SpannerExceptionFactory.propagateInterrupt(interruptedException); + } + } + + @Override + public TransactionContext begin() { + return getTransactionManager().begin(); + } + + @Override + public void commit() { + getTransactionManager().commit(); + } + + @Override + public void rollback() { + getTransactionManager().rollback(); + } + + @Override + public TransactionContext resetForRetry() { + return getTransactionManager().resetForRetry(); + } + + @Override + public Timestamp getCommitTimestamp() { + return getTransactionManager().getCommitTimestamp(); + } + + @Override + public CommitResponse getCommitResponse() { + return getTransactionManager().getCommitResponse(); + } + + @Override + public TransactionState getState() { + return getTransactionManager().getState(); + } + + @Override + public void close() { + getTransactionManager().close(); + } +} diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DelayedTransactionRunner.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DelayedTransactionRunner.java new file mode 100644 index 00000000000..bf0ed1b880a --- /dev/null +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DelayedTransactionRunner.java @@ -0,0 +1,67 @@ +/* + * 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 com.google.api.core.ApiFuture; +import com.google.cloud.Timestamp; +import java.util.concurrent.ExecutionException; +import javax.annotation.Nullable; + +/** + * Represents a {@link TransactionRunner} using a multiplexed session that is not yet ready. The + * execution will be delayed until the multiplexed session has been created and is ready. This class + * is only used during the startup of the client and the multiplexed session has not yet been + * created. + */ +class DelayedTransactionRunner implements TransactionRunner { + private final ApiFuture transactionRunnerFuture; + + DelayedTransactionRunner(ApiFuture transactionRunnerFuture) { + this.transactionRunnerFuture = transactionRunnerFuture; + } + + TransactionRunner getTransactionRunner() { + try { + return this.transactionRunnerFuture.get(); + } catch (ExecutionException executionException) { + throw SpannerExceptionFactory.causeAsRunTimeException(executionException); + } catch (InterruptedException interruptedException) { + throw SpannerExceptionFactory.propagateInterrupt(interruptedException); + } + } + + @Nullable + @Override + public T run(TransactionCallable callable) { + return getTransactionRunner().run(callable); + } + + @Override + public Timestamp getCommitTimestamp() { + return getTransactionRunner().getCommitTimestamp(); + } + + @Override + public CommitResponse getCommitResponse() { + return getTransactionRunner().getCommitResponse(); + } + + @Override + public TransactionRunner allowNestedTransaction() { + return getTransactionRunner().allowNestedTransaction(); + } +} 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 81415e80d25..0c20c4cbc76 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 @@ -21,6 +21,7 @@ import com.google.api.core.ApiFuture; 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.SessionClient.SessionConsumer; import com.google.cloud.spanner.SpannerException.ResourceNotFoundException; @@ -367,41 +368,78 @@ private int getSingleUseChannelHint() { } } + @Override + public Timestamp write(Iterable mutations) throws SpannerException { + return createMultiplexedSessionTransaction(/* singleUse = */ false).write(mutations); + } + + @Override + public CommitResponse writeWithOptions( + final Iterable mutations, final TransactionOption... options) + throws SpannerException { + return createMultiplexedSessionTransaction(/* singleUse = */ false) + .writeWithOptions(mutations, options); + } + @Override public CommitResponse writeAtLeastOnceWithOptions( Iterable mutations, TransactionOption... options) throws SpannerException { - return createMultiplexedSessionTransaction(true) + return createMultiplexedSessionTransaction(/* singleUse = */ true) .writeAtLeastOnceWithOptions(mutations, options); } @Override public ReadContext singleUse() { - return createMultiplexedSessionTransaction(true).singleUse(); + return createMultiplexedSessionTransaction(/* singleUse = */ true).singleUse(); } @Override public ReadContext singleUse(TimestampBound bound) { - return createMultiplexedSessionTransaction(true).singleUse(bound); + return createMultiplexedSessionTransaction(/* singleUse = */ true).singleUse(bound); } @Override public ReadOnlyTransaction singleUseReadOnlyTransaction() { - return createMultiplexedSessionTransaction(true).singleUseReadOnlyTransaction(); + return createMultiplexedSessionTransaction(/* singleUse = */ true) + .singleUseReadOnlyTransaction(); } @Override public ReadOnlyTransaction singleUseReadOnlyTransaction(TimestampBound bound) { - return createMultiplexedSessionTransaction(true).singleUseReadOnlyTransaction(bound); + return createMultiplexedSessionTransaction(/* singleUse = */ true) + .singleUseReadOnlyTransaction(bound); } @Override public ReadOnlyTransaction readOnlyTransaction() { - return createMultiplexedSessionTransaction(false).readOnlyTransaction(); + return createMultiplexedSessionTransaction(/* singleUse = */ false).readOnlyTransaction(); } @Override public ReadOnlyTransaction readOnlyTransaction(TimestampBound bound) { - return createMultiplexedSessionTransaction(false).readOnlyTransaction(bound); + return createMultiplexedSessionTransaction(/* singleUse = */ false).readOnlyTransaction(bound); + } + + @Override + public TransactionRunner readWriteTransaction(TransactionOption... options) { + return createMultiplexedSessionTransaction(/* singleUse = */ false) + .readWriteTransaction(options); + } + + @Override + public TransactionManager transactionManager(TransactionOption... options) { + return createMultiplexedSessionTransaction(/* singleUse = */ false).transactionManager(options); + } + + @Override + public AsyncRunner runAsync(TransactionOption... options) { + return createMultiplexedSessionTransaction(/* singleUse = */ false).runAsync(options); + } + + @Override + public AsyncTransactionManager transactionManagerAsync(TransactionOption... options) { + return createMultiplexedSessionTransaction(/* singleUse = */ false) + .transactionManagerAsync(options); } /** 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 ba335cf8f9f..89de8df3ca9 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 @@ -822,6 +822,8 @@ Builder setUseMultiplexedSessionBlindWrite(boolean useMultiplexedSessionBlindWri * Sets whether the client should use multiplexed session for R/W operations or not. This method * is intentionally package-private and intended for internal use. */ + @InternalApi + @VisibleForTesting Builder setUseMultiplexedSessionForRW(boolean useMultiplexedSessionForRW) { this.useMultiplexedSessionForRW = useMultiplexedSessionForRW; return this; diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerExceptionFactory.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerExceptionFactory.java index 39b254fe997..6af5dee8e74 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerExceptionFactory.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerExceptionFactory.java @@ -31,6 +31,7 @@ import io.grpc.StatusRuntimeException; import io.grpc.protobuf.ProtoUtils; import java.util.concurrent.CancellationException; +import java.util.concurrent.ExecutionException; import java.util.concurrent.TimeoutException; import javax.annotation.Nullable; @@ -181,6 +182,15 @@ public static SpannerException newSpannerException(@Nullable Context context, Th return newSpannerException(ErrorCode.fromGrpcStatus(status), cause.getMessage(), cause); } + public static RuntimeException causeAsRunTimeException(ExecutionException executionException) { + // Propagate the underlying exception as a RuntimeException (SpannerException is also a + // RuntimeException). + if (executionException.getCause() instanceof RuntimeException) { + throw (RuntimeException) executionException.getCause(); + } + throw asSpannerException(executionException.getCause()); + } + /** * Creates a new SpannerException that indicates that the RPC or transaction should be retried on * a different gRPC channel. This is an experimental feature that can be removed in the future. 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 e5982cba0c8..dac1fc2c82b 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 @@ -307,7 +307,8 @@ public DatabaseClient getDatabaseClient(DatabaseId db) { clientId, pool, getOptions().getSessionPoolOptions().getUseMultiplexedSessionBlindWrite(), - multiplexedSessionDatabaseClient); + multiplexedSessionDatabaseClient, + useMultiplexedSessionForRW); dbClients.put(db, dbClient); return dbClient; } @@ -319,9 +320,15 @@ DatabaseClientImpl createDatabaseClient( String clientId, SessionPool pool, boolean useMultiplexedSessionBlindWrite, - @Nullable MultiplexedSessionDatabaseClient multiplexedSessionClient) { + @Nullable MultiplexedSessionDatabaseClient multiplexedSessionClient, + boolean useMultiplexedSessionForRW) { return new DatabaseClientImpl( - clientId, pool, useMultiplexedSessionBlindWrite, multiplexedSessionClient, tracer); + clientId, + pool, + useMultiplexedSessionBlindWrite, + multiplexedSessionClient, + tracer, + useMultiplexedSessionForRW); } @Override diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionManagerImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionManagerImpl.java index 95ffd1168b2..bbbc9aeb447 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionManagerImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionManagerImpl.java @@ -136,6 +136,7 @@ public void close() { } } finally { span.end(); + session.onTransactionDone(); } } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java index c8bf6dc833b..a7250e1ef79 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java @@ -1105,6 +1105,7 @@ public T run(TransactionCallable callable) { // was running. SessionImpl.hasPendingTransaction.remove(); span.end(); + session.onTransactionDone(); } } 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 index 7627ed54883..6deca476fc3 100644 --- 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 @@ -50,7 +50,8 @@ DatabaseClientImpl createDatabaseClient( String clientId, SessionPool pool, boolean useMultiplexedSessionBlindWriteIgnore, - MultiplexedSessionDatabaseClient ignore) { + MultiplexedSessionDatabaseClient ignore, + boolean useMultiplexedSessionForRWIgnore) { return new DatabaseClientWithClosedSessionImpl(clientId, pool, tracer); } } 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 b6dff424079..2e412537882 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 @@ -16,6 +16,8 @@ package com.google.cloud.spanner; +import static com.google.cloud.spanner.MockSpannerTestUtil.UPDATE_COUNT; +import static com.google.cloud.spanner.MockSpannerTestUtil.UPDATE_STATEMENT; import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; @@ -24,14 +26,20 @@ import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; +import com.google.api.core.ApiFuture; +import com.google.api.core.ApiFutures; import com.google.cloud.NoCredentials; import com.google.cloud.Timestamp; +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.SimulatedExecutionTime; import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult; import com.google.cloud.spanner.Options.RpcPriority; import com.google.cloud.spanner.connection.RandomResultSetGenerator; import com.google.common.base.Stopwatch; import com.google.common.collect.ImmutableList; +import com.google.common.util.concurrent.MoreExecutors; import com.google.protobuf.ByteString; import com.google.spanner.v1.CommitRequest; import com.google.spanner.v1.ExecuteSqlRequest; @@ -43,6 +51,9 @@ import java.util.List; import java.util.Set; import java.util.UUID; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; import java.util.stream.Collectors; import org.junit.Before; import org.junit.BeforeClass; @@ -58,6 +69,7 @@ public class MultiplexedSessionDatabaseClientMockServerTest extends AbstractMock public static void setupResults() { mockSpanner.putStatementResults( StatementResult.query(STATEMENT, new RandomResultSetGenerator(1).generate())); + mockSpanner.putStatementResult(StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT)); } @Before @@ -71,6 +83,7 @@ public void createSpannerInstance() { SessionPoolOptions.newBuilder() .setUseMultiplexedSession(true) .setUseMultiplexedSessionBlindWrite(true) + .setUseMultiplexedSessionForRW(true) // Set the maintainer to loop once every 1ms .setMultiplexedSessionMaintenanceLoopFrequency(Duration.ofMillis(1L)) // Set multiplexed sessions to be replaced once every 1ms @@ -467,6 +480,271 @@ public void testWriteAtLeastOnceWithExcludeTxnFromChangeStreams() { assertEquals(1L, client.multiplexedSessionDatabaseClient.getNumSessionsReleased().get()); } + @Test + public void testReadWriteTransactionUsingTransactionRunner() { + // Queries executed within a R/W transaction via TransactionRunner should use a multiplexed + // session. + // During a retry (due to an ABORTED error), the transaction should use the same multiplexed + // session as before, assuming the maintainer hasn't run in the meantime. + DatabaseClientImpl client = + (DatabaseClientImpl) spanner.getDatabaseClient(DatabaseId.of("p", "i", "d")); + // Force the Commit RPC to return Aborted the first time it is called. The exception is cleared + // after the first call, so the retry should succeed. + mockSpanner.setCommitExecutionTime( + SimulatedExecutionTime.ofException( + mockSpanner.createAbortedException(ByteString.copyFromUtf8("test")))); + + client + .readWriteTransaction() + .run( + transaction -> { + try (ResultSet resultSet = transaction.executeQuery(STATEMENT)) { + //noinspection StatementWithEmptyBody + while (resultSet.next()) { + // ignore + } + } + return null; + }); + + List executeSqlRequests = + mockSpanner.getRequestsOfType(ExecuteSqlRequest.class); + assertEquals(2, executeSqlRequests.size()); + assertEquals(executeSqlRequests.get(0).getSession(), executeSqlRequests.get(1).getSession()); + + // Verify the requests are executed using multiplexed sessions + for (ExecuteSqlRequest request : executeSqlRequests) { + assertTrue(mockSpanner.getSession(request.getSession()).getMultiplexed()); + } + + assertNotNull(client.multiplexedSessionDatabaseClient); + assertEquals(1L, client.multiplexedSessionDatabaseClient.getNumSessionsAcquired().get()); + assertEquals(1L, client.multiplexedSessionDatabaseClient.getNumSessionsReleased().get()); + } + + @Test + public void testReadWriteTransactionUsingTransactionManager() { + // Queries executed within a R/W transaction via TransactionManager should use a multiplexed + // session. + // During a retry (due to an ABORTED error), the transaction should use the same multiplexed + // session as before, assuming the maintainer hasn't run in the meantime. + DatabaseClientImpl client = + (DatabaseClientImpl) spanner.getDatabaseClient(DatabaseId.of("p", "i", "d")); + // Force the Commit RPC to return Aborted the first time it is called. The exception is cleared + // after the first call, so the retry should succeed. + mockSpanner.setCommitExecutionTime( + SimulatedExecutionTime.ofException( + mockSpanner.createAbortedException(ByteString.copyFromUtf8("test")))); + + try (TransactionManager manager = client.transactionManager()) { + TransactionContext transaction = manager.begin(); + while (true) { + try { + try (ResultSet resultSet = transaction.executeQuery(STATEMENT)) { + //noinspection StatementWithEmptyBody + while (resultSet.next()) { + // ignore + } + } + manager.commit(); + assertNotNull(manager.getCommitTimestamp()); + break; + } catch (AbortedException e) { + transaction = manager.resetForRetry(); + } + } + } + + List executeSqlRequests = + mockSpanner.getRequestsOfType(ExecuteSqlRequest.class); + assertEquals(2, executeSqlRequests.size()); + assertEquals(executeSqlRequests.get(0).getSession(), executeSqlRequests.get(1).getSession()); + + // Verify the requests are executed using multiplexed sessions + for (ExecuteSqlRequest request : executeSqlRequests) { + assertTrue(mockSpanner.getSession(request.getSession()).getMultiplexed()); + } + + assertNotNull(client.multiplexedSessionDatabaseClient); + assertEquals(1L, client.multiplexedSessionDatabaseClient.getNumSessionsAcquired().get()); + assertEquals(1L, client.multiplexedSessionDatabaseClient.getNumSessionsReleased().get()); + } + + @Test + public void testMutationUsingWrite() { + DatabaseClientImpl client = + (DatabaseClientImpl) spanner.getDatabaseClient(DatabaseId.of("p", "i", "d")); + // Force the Commit RPC to return Aborted the first time it is called. The exception is cleared + // after the first call, so the retry should succeed. + mockSpanner.setCommitExecutionTime( + SimulatedExecutionTime.ofException( + mockSpanner.createAbortedException(ByteString.copyFromUtf8("test")))); + Timestamp timestamp = + client.write( + Collections.singletonList( + Mutation.newInsertBuilder("FOO").set("ID").to(1L).set("NAME").to("Bar").build())); + assertNotNull(timestamp); + + List commitRequests = mockSpanner.getRequestsOfType(CommitRequest.class); + assertEquals(2, commitRequests.size()); + for (CommitRequest request : commitRequests) { + assertTrue(mockSpanner.getSession(request.getSession()).getMultiplexed()); + } + + assertNotNull(client.multiplexedSessionDatabaseClient); + assertEquals(1L, client.multiplexedSessionDatabaseClient.getNumSessionsAcquired().get()); + assertEquals(1L, client.multiplexedSessionDatabaseClient.getNumSessionsReleased().get()); + } + + @Test + public void testMutationUsingWriteWithOptions() { + DatabaseClientImpl client = + (DatabaseClientImpl) spanner.getDatabaseClient(DatabaseId.of("p", "i", "d")); + CommitResponse response = + client.writeWithOptions( + Collections.singletonList( + Mutation.newInsertBuilder("FOO").set("ID").to(1L).set("NAME").to("Bar").build()), + Options.tag("app=spanner,env=test")); + assertNotNull(response); + assertNotNull(response.getCommitTimestamp()); + + List commitRequests = mockSpanner.getRequestsOfType(CommitRequest.class); + assertEquals(1L, commitRequests.size()); + CommitRequest commit = commitRequests.get(0); + assertNotNull(commit.getRequestOptions()); + assertEquals("app=spanner,env=test", commit.getRequestOptions().getTransactionTag()); + assertTrue(mockSpanner.getSession(commit.getSession()).getMultiplexed()); + + assertNotNull(client.multiplexedSessionDatabaseClient); + assertEquals(1L, client.multiplexedSessionDatabaseClient.getNumSessionsAcquired().get()); + assertEquals(1L, client.multiplexedSessionDatabaseClient.getNumSessionsReleased().get()); + } + + @Test + public void testReadWriteTransactionUsingAsyncTransactionManager() throws Exception { + // Updates executed within a R/W transaction via AsyncTransactionManager should use a + // multiplexed session. + // During a retry (due to an ABORTED error), the transaction should use the same multiplexed + // session as before, assuming the maintainer hasn't run in the meantime. + final AtomicInteger attempt = new AtomicInteger(); + CountDownLatch abortedLatch = new CountDownLatch(1); + DatabaseClientImpl client = + (DatabaseClientImpl) spanner.getDatabaseClient(DatabaseId.of("p", "i", "d")); + try (AsyncTransactionManager manager = client.transactionManagerAsync()) { + TransactionContextFuture transactionContextFuture = manager.beginAsync(); + while (true) { + try { + attempt.incrementAndGet(); + AsyncTransactionStep updateCount = + transactionContextFuture.then( + (transaction, ignored) -> transaction.executeUpdateAsync(UPDATE_STATEMENT), + MoreExecutors.directExecutor()); + updateCount.then( + (transaction, ignored) -> { + if (attempt.get() == 1) { + mockSpanner.abortTransaction(transaction); + abortedLatch.countDown(); + } + return ApiFutures.immediateFuture(null); + }, + MoreExecutors.directExecutor()); + abortedLatch.await(10L, TimeUnit.SECONDS); + CommitTimestampFuture commitTimestamp = updateCount.commitAsync(); + assertEquals(UPDATE_COUNT, updateCount.get().longValue()); + assertNotNull(commitTimestamp.get()); + assertEquals(2L, attempt.get()); + break; + } catch (AbortedException e) { + transactionContextFuture = manager.resetForRetryAsync(); + } + } + } + + List executeSqlRequests = + mockSpanner.getRequestsOfType(ExecuteSqlRequest.class); + assertEquals(2, executeSqlRequests.size()); + assertEquals(executeSqlRequests.get(0).getSession(), executeSqlRequests.get(1).getSession()); + + // Verify the requests are executed using multiplexed sessions + for (ExecuteSqlRequest request : executeSqlRequests) { + assertTrue(mockSpanner.getSession(request.getSession()).getMultiplexed()); + } + + assertNotNull(client.multiplexedSessionDatabaseClient); + assertEquals(1L, client.multiplexedSessionDatabaseClient.getNumSessionsAcquired().get()); + assertEquals(1L, client.multiplexedSessionDatabaseClient.getNumSessionsReleased().get()); + } + + @Test + public void testReadWriteTransactionUsingAsyncRunner() throws Exception { + // Updates executed within a R/W transaction via AsyncRunner should use a multiplexed + // session. + // During a retry (due to an ABORTED error), the transaction should use the same multiplexed + // session as before, assuming the maintainer hasn't run in the meantime. + final AtomicInteger attempt = new AtomicInteger(); + DatabaseClientImpl client = + (DatabaseClientImpl) spanner.getDatabaseClient(DatabaseId.of("p", "i", "d")); + AsyncRunner runner = client.runAsync(); + ApiFuture updateCount = + runner.runAsync( + txn -> { + ApiFuture updateCount1 = txn.executeUpdateAsync(UPDATE_STATEMENT); + if (attempt.incrementAndGet() == 1) { + mockSpanner.abortTransaction(txn); + } + return updateCount1; + }, + MoreExecutors.directExecutor()); + assertEquals(UPDATE_COUNT, updateCount.get().longValue()); + assertEquals(2L, attempt.get()); + + List executeSqlRequests = + mockSpanner.getRequestsOfType(ExecuteSqlRequest.class); + assertEquals(2L, executeSqlRequests.size()); + assertEquals(executeSqlRequests.get(0).getSession(), executeSqlRequests.get(1).getSession()); + + // Verify the requests are executed using multiplexed sessions + for (ExecuteSqlRequest request : executeSqlRequests) { + assertTrue(mockSpanner.getSession(request.getSession()).getMultiplexed()); + } + + assertNotNull(client.multiplexedSessionDatabaseClient); + assertEquals(1L, client.multiplexedSessionDatabaseClient.getNumSessionsAcquired().get()); + assertEquals(1L, client.multiplexedSessionDatabaseClient.getNumSessionsReleased().get()); + } + + @Test + public void testAsyncRunnerIsNonBlockingWithMultiplexedSession() throws Exception { + mockSpanner.freeze(); + DatabaseClientImpl client = + (DatabaseClientImpl) spanner.getDatabaseClient(DatabaseId.of("p", "i", "d")); + AsyncRunner runner = client.runAsync(); + ApiFuture res = + runner.runAsync( + txn -> { + txn.executeUpdateAsync(UPDATE_STATEMENT); + return ApiFutures.immediateFuture(null); + }, + MoreExecutors.directExecutor()); + ApiFuture ts = runner.getCommitTimestamp(); + mockSpanner.unfreeze(); + assertThat(res.get()).isNull(); + assertThat(ts.get()).isNotNull(); + + List executeSqlRequests = + mockSpanner.getRequestsOfType(ExecuteSqlRequest.class); + assertEquals(1L, executeSqlRequests.size()); + + // Verify the requests are executed using multiplexed sessions + for (ExecuteSqlRequest request : executeSqlRequests) { + assertTrue(mockSpanner.getSession(request.getSession()).getMultiplexed()); + } + + assertNotNull(client.multiplexedSessionDatabaseClient); + assertEquals(1L, client.multiplexedSessionDatabaseClient.getNumSessionsAcquired().get()); + assertEquals(1L, client.multiplexedSessionDatabaseClient.getNumSessionsReleased().get()); + } + private void waitForSessionToBeReplaced(DatabaseClientImpl client) { assertNotNull(client.multiplexedSessionDatabaseClient); SessionReference sessionReference =