From 3dd0675d2d7882d40a6af1e12fda3b4617019870 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= Date: Thu, 8 Oct 2020 16:08:50 +0200 Subject: [PATCH] feat!: async connection API (#392) * feat: support setting timeout per RPC The Spanner client allows a user to set custom timeouts while creating a SpannerOptions instance, but these timeouts are static and are applied to all invocations of the RPCs. This change introduces the possibility to set custom timeouts and other call options on a per-RPC basis. Fixes #378 * fix: change grpc deps from test to compile scope * feat: add async api for connection * fix: fix test failures * fix: move state handling from callback to callable * fix: fix integration tests with emulator * fix: fix timeout integration test on emulator * fix: prevent flakiness in DDL tests * fix: fix clirr build failures * fix: do not set transaction state for Aborted err * fix: set transaction state after retry * cleanup: remove sync methods and use async instead * cleanup: remove unused code * feat: make ddl async * fix: reduce timeout and remove debug info * feat: make runBatch async * test: set forkCount to 1 to investigate test failure * fix: linting + clirr * fix: prevent deadlock in DmlBatch * fix: fix DMLBatch state handling * tests: add tests for aborted async transactions * test: add aborted tests * fix: add change to clirr + more tests * fix: require a rollback after a tx has aborted * docs: add javadoc for new methods * tests: add integration tests * fix: wait for commit before select * fix: fix handling aborted commit * docs: document behavior -Async methods * fix: iterating without callback could cause exception * fix: remove todos and commented code * feat: keep track of caller to include in stacktrace * docs: explain why Aborted is active * fix: use ticker for better testability * test: increase coverage and remove unused code * test: add additional tests * docs: add missing @override * docs: fix comment --- .../clirr-ignored-differences.xml | 52 + .../cloud/spanner/AsyncResultSetImpl.java | 37 +- .../com/google/cloud/spanner/ErrorCode.java | 2 +- .../com/google/cloud/spanner/ResultSets.java | 40 +- .../cloud/spanner/SpannerApiFutures.java | 43 + .../spanner/SpannerExceptionFactory.java | 26 + .../cloud/spanner/TransactionRunnerImpl.java | 10 +- .../connection/AbstractBaseUnitOfWork.java | 144 +- .../AbstractMultiUseTransaction.java | 15 +- .../connection/AsyncStatementResult.java | 47 + .../connection/AsyncStatementResultImpl.java | 130 ++ .../spanner/connection/ChecksumResultSet.java | 18 +- .../cloud/spanner/connection/Connection.java | 247 +++- .../spanner/connection/ConnectionImpl.java | 276 ++-- .../spanner/connection/ConnectionOptions.java | 22 + .../cloud/spanner/connection/DdlBatch.java | 118 +- .../cloud/spanner/connection/DmlBatch.java | 68 +- .../connection/ReadOnlyTransaction.java | 27 +- .../connection/ReadWriteTransaction.java | 586 ++++---- .../connection/SingleUseTransaction.java | 391 +++--- .../cloud/spanner/connection/SpannerPool.java | 28 +- .../spanner/connection/StatementExecutor.java | 73 +- .../connection/StatementResultImpl.java | 23 + .../cloud/spanner/connection/UnitOfWork.java | 60 +- .../cloud/spanner/MockSpannerServiceImpl.java | 83 +- .../google/cloud/spanner/ResultSetsTest.java | 137 ++ .../cloud/spanner/SpannerApiFuturesTest.java | 118 ++ .../connection/AbstractMockServerTest.java | 63 +- .../AsyncStatementResultImplTest.java | 99 ++ .../ConnectionAsyncApiAbortedTest.java | 688 ++++++++++ .../connection/ConnectionAsyncApiTest.java | 833 ++++++++++++ .../connection/ConnectionImplTest.java | 24 +- .../spanner/connection/ConnectionTest.java | 77 ++ .../spanner/connection/DdlBatchTest.java | 141 +- .../spanner/connection/DmlBatchTest.java | 41 +- .../connection/ITAbstractSpannerTest.java | 16 + .../connection/ReadOnlyTransactionTest.java | 50 +- .../connection/ReadWriteTransactionTest.java | 41 +- .../connection/SingleUseTransactionTest.java | 165 +-- .../spanner/connection/SpannerPoolTest.java | 145 +- .../connection/StatementTimeoutTest.java | 1194 ++++++++--------- .../it/ITAsyncTransactionRetryTest.java | 1015 ++++++++++++++ .../connection/it/ITReadOnlySpannerTest.java | 17 +- .../connection/it/ITSqlMusicScriptTest.java | 2 + .../connection/it/ITTransactionRetryTest.java | 4 +- .../ITSqlScriptTest_TestStatementTimeout.sql | 16 +- 46 files changed, 5690 insertions(+), 1762 deletions(-) create mode 100644 google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerApiFutures.java create mode 100644 google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AsyncStatementResult.java create mode 100644 google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AsyncStatementResultImpl.java create mode 100644 google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerApiFuturesTest.java create mode 100644 google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AsyncStatementResultImplTest.java create mode 100644 google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionAsyncApiAbortedTest.java create mode 100644 google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionAsyncApiTest.java create mode 100644 google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/it/ITAsyncTransactionRetryTest.java diff --git a/google-cloud-spanner/clirr-ignored-differences.xml b/google-cloud-spanner/clirr-ignored-differences.xml index bc33de3bbb..cfbcb88f85 100644 --- a/google-cloud-spanner/clirr-ignored-differences.xml +++ b/google-cloud-spanner/clirr-ignored-differences.xml @@ -319,4 +319,56 @@ com/google/cloud/spanner/Value java.util.List getNumericArray() + + + + 7012 + com/google/cloud/spanner/connection/Connection + com.google.api.core.ApiFuture beginTransactionAsync() + + + 7012 + com/google/cloud/spanner/connection/Connection + com.google.api.core.ApiFuture commitAsync() + + + 7012 + com/google/cloud/spanner/connection/Connection + com.google.cloud.spanner.connection.AsyncStatementResult executeAsync(com.google.cloud.spanner.Statement) + + + 7012 + com/google/cloud/spanner/connection/Connection + com.google.api.core.ApiFuture executeBatchUpdateAsync(java.lang.Iterable) + + + 7012 + com/google/cloud/spanner/connection/Connection + com.google.api.core.ApiFuture executeUpdateAsync(com.google.cloud.spanner.Statement) + + + 7012 + com/google/cloud/spanner/connection/Connection + com.google.api.core.ApiFuture rollbackAsync() + + + 7012 + com/google/cloud/spanner/connection/Connection + com.google.api.core.ApiFuture runBatchAsync() + + + 7012 + com/google/cloud/spanner/connection/Connection + com.google.api.core.ApiFuture writeAsync(com.google.cloud.spanner.Mutation) + + + 7012 + com/google/cloud/spanner/connection/Connection + com.google.api.core.ApiFuture writeAsync(java.lang.Iterable) + + + 7004 + com/google/cloud/spanner/ResultSets + com.google.cloud.spanner.AsyncResultSet toAsyncResultSet(com.google.cloud.spanner.ResultSet, com.google.api.gax.core.ExecutorProvider) + diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncResultSetImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncResultSetImpl.java index 1cb768ea85..fd172e96f9 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncResultSetImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncResultSetImpl.java @@ -25,6 +25,8 @@ import com.google.cloud.spanner.AbstractReadContext.ListenableAsyncResultSet; import com.google.common.base.Function; import com.google.common.base.Preconditions; +import com.google.common.base.Supplier; +import com.google.common.base.Suppliers; import com.google.common.collect.ImmutableList; import com.google.common.util.concurrent.ListeningScheduledExecutorService; import com.google.common.util.concurrent.MoreExecutors; @@ -88,8 +90,8 @@ private State(boolean shouldStop) { private final BlockingDeque buffer; private Struct currentRow; - /** The underlying synchronous {@link ResultSet} that is producing the rows. */ - private final ResultSet delegateResultSet; + /** Supplies the underlying synchronous {@link ResultSet} that will be producing the rows. */ + private final Supplier delegateResultSet; /** * Any exception that occurs while executing the query and iterating over the result set will be @@ -144,6 +146,11 @@ private State(boolean shouldStop) { private volatile CountDownLatch consumingLatch = new CountDownLatch(0); AsyncResultSetImpl(ExecutorProvider executorProvider, ResultSet delegate, int bufferSize) { + this(executorProvider, Suppliers.ofInstance(Preconditions.checkNotNull(delegate)), bufferSize); + } + + AsyncResultSetImpl( + ExecutorProvider executorProvider, Supplier delegate, int bufferSize) { super(delegate); this.executorProvider = Preconditions.checkNotNull(executorProvider); this.delegateResultSet = Preconditions.checkNotNull(delegate); @@ -165,7 +172,7 @@ public void close() { return; } if (state == State.INITIALIZED || state == State.SYNC) { - delegateResultSet.close(); + delegateResultSet.get().close(); } this.closed = true; } @@ -228,7 +235,7 @@ public CursorState tryNext() throws SpannerException { private void closeDelegateResultSet() { try { - delegateResultSet.close(); + delegateResultSet.get().close(); } catch (Throwable t) { log.log(Level.FINE, "Ignoring error from closing delegate result set", t); } @@ -261,7 +268,7 @@ public void run() { // we'll keep the cancelled state. return; } - executionException = SpannerExceptionFactory.newSpannerException(e); + executionException = SpannerExceptionFactory.asSpannerException(e); cursorReturnedDoneOrException = true; } return; @@ -325,10 +332,10 @@ public Void call() throws Exception { boolean stop = false; boolean hasNext = false; try { - hasNext = delegateResultSet.next(); + hasNext = delegateResultSet.get().next(); } catch (Throwable e) { synchronized (monitor) { - executionException = SpannerExceptionFactory.newSpannerException(e); + executionException = SpannerExceptionFactory.asSpannerException(e); } } try { @@ -357,13 +364,13 @@ public Void call() throws Exception { } } if (!stop) { - buffer.put(delegateResultSet.getCurrentRowAsStruct()); + buffer.put(delegateResultSet.get().getCurrentRowAsStruct()); startCallbackIfNecessary(); - hasNext = delegateResultSet.next(); + hasNext = delegateResultSet.get().next(); } } catch (Throwable e) { synchronized (monitor) { - executionException = SpannerExceptionFactory.newSpannerException(e); + executionException = SpannerExceptionFactory.asSpannerException(e); stop = true; } } @@ -544,9 +551,9 @@ public List toList(Function transformer) throws SpannerE try { return future.get(); } catch (ExecutionException e) { - throw SpannerExceptionFactory.newSpannerException(e.getCause()); + throw SpannerExceptionFactory.asSpannerException(e.getCause()); } catch (Throwable e) { - throw SpannerExceptionFactory.newSpannerException(e); + throw SpannerExceptionFactory.asSpannerException(e); } } @@ -558,14 +565,14 @@ public boolean next() throws SpannerException { "Cannot call next() on a result set with a callback."); this.state = State.SYNC; } - boolean res = delegateResultSet.next(); - currentRow = res ? delegateResultSet.getCurrentRowAsStruct() : null; + boolean res = delegateResultSet.get().next(); + currentRow = res ? delegateResultSet.get().getCurrentRowAsStruct() : null; return res; } @Override public ResultSetStats getStats() { - return delegateResultSet.getStats(); + return delegateResultSet.get().getStats(); } @Override diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ErrorCode.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ErrorCode.java index a9df5ab59b..9896cc8aec 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ErrorCode.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ErrorCode.java @@ -89,7 +89,7 @@ static ErrorCode valueOf(String name, ErrorCode defaultValue) { /** * Returns the error code corresponding to a gRPC status, or {@code UNKNOWN} if not recognized. */ - static ErrorCode fromGrpcStatus(Status status) { + public static ErrorCode fromGrpcStatus(Status status) { ErrorCode code = errorByRpcCode.get(status.getCode().value()); return code == null ? UNKNOWN : code; } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ResultSets.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ResultSets.java index ee9e715a25..5ec54960e6 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ResultSets.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ResultSets.java @@ -16,14 +16,17 @@ package com.google.cloud.spanner; +import com.google.api.core.ApiFuture; import com.google.api.gax.core.ExecutorProvider; import com.google.api.gax.core.InstantiatingExecutorProvider; import com.google.cloud.ByteArray; import com.google.cloud.Date; import com.google.cloud.Timestamp; +import com.google.cloud.spanner.Options.QueryOption; import com.google.cloud.spanner.Type.Code; import com.google.cloud.spanner.Type.StructField; import com.google.common.base.Preconditions; +import com.google.common.base.Supplier; import com.google.common.collect.Lists; import com.google.common.util.concurrent.ThreadFactoryBuilder; import com.google.spanner.v1.ResultSetStats; @@ -65,8 +68,41 @@ public static AsyncResultSet toAsyncResultSet(ResultSet delegate) { * ExecutorProvider}. */ public static AsyncResultSet toAsyncResultSet( - ResultSet delegate, ExecutorProvider executorProvider) { - return new AsyncResultSetImpl(executorProvider, delegate, 100); + ResultSet delegate, ExecutorProvider executorProvider, QueryOption... options) { + Options readOptions = Options.fromQueryOptions(options); + final int bufferRows = + readOptions.hasBufferRows() + ? readOptions.bufferRows() + : AsyncResultSetImpl.DEFAULT_BUFFER_SIZE; + return new AsyncResultSetImpl(executorProvider, delegate, bufferRows); + } + + /** + * Converts the {@link ResultSet} that will be returned by the given {@link ApiFuture} to an + * {@link AsyncResultSet} using the given {@link ExecutorProvider}. + */ + public static AsyncResultSet toAsyncResultSet( + ApiFuture delegate, ExecutorProvider executorProvider, QueryOption... options) { + Options readOptions = Options.fromQueryOptions(options); + final int bufferRows = + readOptions.hasBufferRows() + ? readOptions.bufferRows() + : AsyncResultSetImpl.DEFAULT_BUFFER_SIZE; + return new AsyncResultSetImpl( + executorProvider, new FutureResultSetSupplier(delegate), bufferRows); + } + + private static class FutureResultSetSupplier implements Supplier { + final ApiFuture delegate; + + FutureResultSetSupplier(ApiFuture delegate) { + this.delegate = Preconditions.checkNotNull(delegate); + } + + @Override + public ResultSet get() { + return SpannerApiFutures.get(delegate); + } } private static class PrePopulatedResultSet implements ResultSet { diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerApiFutures.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerApiFutures.java new file mode 100644 index 0000000000..39afc1b81a --- /dev/null +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerApiFutures.java @@ -0,0 +1,43 @@ +/* + * 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.common.base.Preconditions; +import java.util.concurrent.CancellationException; +import java.util.concurrent.ExecutionException; + +public class SpannerApiFutures { + public static T get(ApiFuture future) throws SpannerException { + return getOrNull(Preconditions.checkNotNull(future)); + } + + public static T getOrNull(ApiFuture future) throws SpannerException { + try { + return future == null ? null : future.get(); + } catch (ExecutionException e) { + if (e.getCause() instanceof SpannerException) { + throw (SpannerException) e.getCause(); + } + throw SpannerExceptionFactory.newSpannerException(e.getCause()); + } catch (InterruptedException e) { + throw SpannerExceptionFactory.propagateInterrupt(e); + } catch (CancellationException e) { + throw SpannerExceptionFactory.newSpannerExceptionForCancellation(null, e); + } + } +} 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 3fa756875b..774aaf472e 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 @@ -83,6 +83,18 @@ public static SpannerException propagateTimeout(TimeoutException e) { ErrorCode.DEADLINE_EXCEEDED, "Operation did not complete in the given time", e); } + /** + * Converts the given {@link Throwable} to a {@link SpannerException}. If t is + * already a (subclass of a) {@link SpannerException}, t is returned unaltered. + * Otherwise, a new {@link SpannerException} is created with t as its cause. + */ + public static SpannerException asSpannerException(Throwable t) { + if (t instanceof SpannerException) { + return (SpannerException) t; + } + return newSpannerException(t); + } + /** * Creates a new exception based on {@code cause}. * @@ -126,6 +138,20 @@ public static SpannerBatchUpdateException newSpannerBatchUpdateException( databaseError); } + /** + * Constructs a new {@link AbortedDueToConcurrentModificationException} that can be re-thrown for + * a transaction that had already been aborted, but that the client application tried to use for + * additional statements. + */ + public static AbortedDueToConcurrentModificationException + newAbortedDueToConcurrentModificationException( + AbortedDueToConcurrentModificationException cause) { + return new AbortedDueToConcurrentModificationException( + DoNotConstructDirectly.ALLOWED, + "This transaction has already been aborted and could not be retried due to a concurrent modification. Rollback this transaction to start a new one.", + cause); + } + /** * Creates a new exception based on {@code cause}. If {@code cause} indicates cancellation, {@code * context} will be inspected to establish the type of cancellation. 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 dc6cb56f30..ab4a80b340 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 @@ -150,7 +150,7 @@ public void removeListener(Runnable listener) { @GuardedBy("lock") private long retryDelayInMillis = -1L; - private ByteString transactionId; + private volatile ByteString transactionId; private Timestamp commitTimestamp; private TransactionContextImpl(Builder builder) { @@ -238,12 +238,17 @@ void commit() { try { commitTimestamp = commitAsync().get(); } catch (InterruptedException e) { + if (commitFuture != null) { + commitFuture.cancel(true); + } throw SpannerExceptionFactory.propagateInterrupt(e); } catch (ExecutionException e) { throw SpannerExceptionFactory.newSpannerException(e.getCause() == null ? e : e.getCause()); } } + volatile ApiFuture commitFuture; + ApiFuture commitAsync() { final SettableApiFuture res = SettableApiFuture.create(); final SettableApiFuture latch; @@ -273,8 +278,7 @@ public void run() { span.addAnnotation("Starting Commit"); final Span opSpan = tracer.spanBuilderWithExplicitParent(SpannerImpl.COMMIT, span).startSpan(); - final ApiFuture commitFuture = - rpc.commitAsync(commitRequest, session.getOptions()); + commitFuture = rpc.commitAsync(commitRequest, session.getOptions()); commitFuture.addListener( tracer.withSpan( opSpan, diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AbstractBaseUnitOfWork.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AbstractBaseUnitOfWork.java index 3fcffce0ac..9ba86b3ec5 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AbstractBaseUnitOfWork.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AbstractBaseUnitOfWork.java @@ -16,12 +16,27 @@ package com.google.cloud.spanner.connection; +import com.google.api.core.ApiFunction; +import com.google.api.core.ApiFuture; +import com.google.api.core.ApiFutures; +import com.google.api.gax.grpc.GrpcCallContext; +import com.google.api.gax.longrunning.OperationFuture; +import com.google.api.gax.rpc.ApiCallContext; import com.google.cloud.spanner.ErrorCode; import com.google.cloud.spanner.SpannerException; import com.google.cloud.spanner.SpannerExceptionFactory; +import com.google.cloud.spanner.SpannerOptions; +import com.google.cloud.spanner.Statement; import com.google.cloud.spanner.connection.StatementExecutor.StatementTimeout; import com.google.cloud.spanner.connection.StatementParser.ParsedStatement; import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableList; +import com.google.common.util.concurrent.MoreExecutors; +import io.grpc.Context; +import io.grpc.MethodDescriptor; +import io.grpc.Status; +import java.util.Collection; +import java.util.Collections; import java.util.HashSet; import java.util.Set; import java.util.concurrent.Callable; @@ -30,6 +45,7 @@ import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; +import javax.annotation.Nullable; import javax.annotation.concurrent.GuardedBy; /** Base for all {@link Connection}-based transactions and batches. */ @@ -37,12 +53,27 @@ abstract class AbstractBaseUnitOfWork implements UnitOfWork { private final StatementExecutor statementExecutor; private final StatementTimeout statementTimeout; + /** Class for keeping track of the stacktrace of the caller of an async statement. */ + static final class SpannerAsyncExecutionException extends RuntimeException { + final Statement statement; + + SpannerAsyncExecutionException(Statement statement) { + this.statement = statement; + } + + public String getMessage() { + // We only include the SQL of the statement and not the parameter values to prevent + // potentially sensitive data to escape into an error message. + return String.format("Execution failed for statement: %s", statement.getSql()); + } + } + /** * The {@link Future} that monitors the result of the statement currently being executed for this * unit of work. */ @GuardedBy("this") - private Future currentlyRunningStatementFuture = null; + private volatile Future currentlyRunningStatementFuture = null; enum InterceptorsUsage { INVOKE_INTERCEPTORS, @@ -100,34 +131,38 @@ public void cancel() { } } - T asyncExecuteStatement(ParsedStatement statement, Callable callable) { - return asyncExecuteStatement(statement, callable, InterceptorsUsage.INVOKE_INTERCEPTORS); + ApiFuture executeStatementAsync( + ParsedStatement statement, + Callable callable, + @Nullable MethodDescriptor applyStatementTimeoutToMethod) { + return executeStatementAsync( + statement, + callable, + InterceptorsUsage.INVOKE_INTERCEPTORS, + applyStatementTimeoutToMethod == null + ? Collections.>emptySet() + : ImmutableList.>of(applyStatementTimeoutToMethod)); } - T asyncExecuteStatement( - ParsedStatement statement, Callable callable, InterceptorsUsage interceptorUsage) { - Preconditions.checkNotNull(statement); - Preconditions.checkNotNull(callable); + ApiFuture executeStatementAsync( + ParsedStatement statement, + Callable callable, + Collection> applyStatementTimeoutToMethods) { + return executeStatementAsync( + statement, callable, InterceptorsUsage.INVOKE_INTERCEPTORS, applyStatementTimeoutToMethods); + } - if (interceptorUsage == InterceptorsUsage.INVOKE_INTERCEPTORS) { - statementExecutor.invokeInterceptors( - statement, StatementExecutionStep.EXECUTE_STATEMENT, this); - } - Future future = statementExecutor.submit(callable); - synchronized (this) { - this.currentlyRunningStatementFuture = future; - } - T res; + ResponseT getWithStatementTimeout( + OperationFuture operation, ParsedStatement statement) { + ResponseT res; try { if (statementTimeout.hasTimeout()) { TimeUnit unit = statementTimeout.getAppropriateTimeUnit(); - res = future.get(statementTimeout.getTimeoutValue(unit), unit); + res = operation.get(statementTimeout.getTimeoutValue(unit), unit); } else { - res = future.get(); + res = operation.get(); } } catch (TimeoutException e) { - // statement timed out, cancel the execution - future.cancel(true); throw SpannerExceptionFactory.newSpannerException( ErrorCode.DEADLINE_EXCEEDED, "Statement execution timeout occurred for " + statement.getSqlWithoutComments(), @@ -143,7 +178,7 @@ T asyncExecuteStatement( cause = cause.getCause(); } throw SpannerExceptionFactory.newSpannerException( - ErrorCode.UNKNOWN, + ErrorCode.fromGrpcStatus(Status.fromThrowable(e)), "Statement execution failed for " + statement.getSqlWithoutComments(), e); } catch (InterruptedException e) { @@ -152,11 +187,70 @@ T asyncExecuteStatement( } catch (CancellationException e) { throw SpannerExceptionFactory.newSpannerException( ErrorCode.CANCELLED, "Statement execution was cancelled", e); - } finally { - synchronized (this) { - this.currentlyRunningStatementFuture = null; - } } return res; } + + ApiFuture executeStatementAsync( + ParsedStatement statement, + Callable callable, + InterceptorsUsage interceptorUsage, + final Collection> applyStatementTimeoutToMethods) { + Preconditions.checkNotNull(statement); + Preconditions.checkNotNull(callable); + + if (interceptorUsage == InterceptorsUsage.INVOKE_INTERCEPTORS) { + statementExecutor.invokeInterceptors( + statement, StatementExecutionStep.EXECUTE_STATEMENT, this); + } + Context context = Context.current(); + if (statementTimeout.hasTimeout() && !applyStatementTimeoutToMethods.isEmpty()) { + context = + context.withValue( + SpannerOptions.CALL_CONTEXT_CONFIGURATOR_KEY, + new SpannerOptions.CallContextConfigurator() { + @Override + public ApiCallContext configure( + ApiCallContext context, ReqT request, MethodDescriptor method) { + if (statementTimeout.hasTimeout() + && applyStatementTimeoutToMethods.contains(method)) { + return GrpcCallContext.createDefault() + .withTimeout(statementTimeout.asDuration()); + } + return null; + } + }); + } + ApiFuture f = statementExecutor.submit(context.wrap(callable)); + final SpannerAsyncExecutionException caller = + new SpannerAsyncExecutionException(statement.getStatement()); + final ApiFuture future = + ApiFutures.catching( + f, + Throwable.class, + new ApiFunction() { + @Override + public T apply(Throwable input) { + input.addSuppressed(caller); + throw SpannerExceptionFactory.asSpannerException(input); + } + }, + MoreExecutors.directExecutor()); + synchronized (this) { + this.currentlyRunningStatementFuture = future; + } + future.addListener( + new Runnable() { + @Override + public void run() { + synchronized (this) { + if (currentlyRunningStatementFuture == future) { + currentlyRunningStatementFuture = null; + } + } + } + }, + MoreExecutors.directExecutor()); + return future; + } } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AbstractMultiUseTransaction.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AbstractMultiUseTransaction.java index cb8cf3bc55..33cef1fedb 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AbstractMultiUseTransaction.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AbstractMultiUseTransaction.java @@ -16,6 +16,7 @@ package com.google.cloud.spanner.connection; +import com.google.api.core.ApiFuture; import com.google.cloud.spanner.ErrorCode; import com.google.cloud.spanner.Options.QueryOption; import com.google.cloud.spanner.ReadContext; @@ -24,6 +25,7 @@ import com.google.cloud.spanner.SpannerExceptionFactory; import com.google.cloud.spanner.connection.StatementParser.ParsedStatement; import com.google.common.base.Preconditions; +import com.google.spanner.v1.SpannerGrpc; import java.util.concurrent.Callable; /** @@ -46,6 +48,8 @@ public boolean isActive() { return getState().isActive(); } + abstract void checkAborted(); + /** * Check that the current transaction actually has a valid underlying transaction. If not, the * method will throw a {@link SpannerException}. @@ -55,22 +59,23 @@ public boolean isActive() { /** Returns the {@link ReadContext} that can be used for queries on this transaction. */ abstract ReadContext getReadContext(); - @Override - public ResultSet executeQuery( + public ApiFuture executeQueryAsync( final ParsedStatement statement, final AnalyzeMode analyzeMode, final QueryOption... options) { Preconditions.checkArgument(statement.isQuery(), "Statement is not a query"); checkValidTransaction(); - return asyncExecuteStatement( + return executeStatementAsync( statement, new Callable() { @Override public ResultSet call() throws Exception { + checkAborted(); return DirectExecuteResultSet.ofResultSet( internalExecuteQuery(statement, analyzeMode, options)); } - }); + }, + SpannerGrpc.getExecuteStreamingSqlMethod()); } ResultSet internalExecuteQuery( @@ -83,7 +88,7 @@ ResultSet internalExecuteQuery( } @Override - public long[] runBatch() { + public ApiFuture runBatchAsync() { throw SpannerExceptionFactory.newSpannerException( ErrorCode.FAILED_PRECONDITION, "Run batch is not supported for transactions"); } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AsyncStatementResult.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AsyncStatementResult.java new file mode 100644 index 0000000000..fef96ab456 --- /dev/null +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AsyncStatementResult.java @@ -0,0 +1,47 @@ +/* + * 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.connection; + +import com.google.api.core.ApiFuture; +import com.google.api.core.InternalApi; +import com.google.cloud.spanner.AsyncResultSet; +import com.google.cloud.spanner.connection.StatementResult.ResultType; + +@InternalApi +public interface AsyncStatementResult extends StatementResult { + /** + * Returns the {@link AsyncResultSet} held by this result. May only be called if the type of this + * result is {@link ResultType#RESULT_SET}. + * + * @return the {@link AsyncResultSet} held by this result. + */ + AsyncResultSet getResultSetAsync(); + + /** + * Returns the update count held by this result. May only be called if the type of this result is + * {@link ResultType#UPDATE_COUNT}. + * + * @return the update count held by this result. + */ + ApiFuture getUpdateCountAsync(); + + /** + * Returns a future that tracks the progress of a statement that returns no result. This could be + * a DDL statement or a client side statement that does not return a result. + */ + ApiFuture getNoResultAsync(); +} diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AsyncStatementResultImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AsyncStatementResultImpl.java new file mode 100644 index 0000000000..7d0b0fc3b5 --- /dev/null +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AsyncStatementResultImpl.java @@ -0,0 +1,130 @@ +/* + * 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.connection; + +import static com.google.cloud.spanner.SpannerApiFutures.get; + +import com.google.api.core.ApiFuture; +import com.google.api.core.ApiFutures; +import com.google.api.gax.core.ExecutorProvider; +import com.google.cloud.spanner.AsyncResultSet; +import com.google.cloud.spanner.ResultSet; +import com.google.cloud.spanner.ResultSets; +import com.google.common.base.Preconditions; + +class AsyncStatementResultImpl implements AsyncStatementResult { + + static AsyncStatementResult of(AsyncResultSet resultSet) { + return new AsyncStatementResultImpl(Preconditions.checkNotNull(resultSet), null); + } + + static AsyncStatementResult of(ApiFuture updateCount) { + return new AsyncStatementResultImpl(Preconditions.checkNotNull(updateCount)); + } + + static AsyncStatementResult of( + StatementResult clientSideStatementResult, ExecutorProvider executorProvider) { + Preconditions.checkNotNull(clientSideStatementResult.getClientSideStatementType()); + Preconditions.checkNotNull(executorProvider); + if (clientSideStatementResult.getResultType() == ResultType.RESULT_SET) { + return new AsyncStatementResultImpl( + ResultSets.toAsyncResultSet(clientSideStatementResult.getResultSet(), executorProvider), + clientSideStatementResult.getClientSideStatementType()); + } else { + return new AsyncStatementResultImpl( + clientSideStatementResult.getClientSideStatementType(), + ApiFutures.immediateFuture(null)); + } + } + + static AsyncStatementResult noResult(ApiFuture result) { + return new AsyncStatementResultImpl(null, Preconditions.checkNotNull(result)); + } + + private final ResultType type; + private final ClientSideStatementType clientSideStatementType; + private final AsyncResultSet resultSet; + private final ApiFuture updateCount; + private final ApiFuture noResult; + + private AsyncStatementResultImpl( + AsyncResultSet resultSet, ClientSideStatementType clientSideStatementType) { + this.type = ResultType.RESULT_SET; + this.clientSideStatementType = clientSideStatementType; + this.resultSet = resultSet; + this.updateCount = null; + this.noResult = null; + } + + private AsyncStatementResultImpl(ApiFuture updateCount) { + this.type = ResultType.UPDATE_COUNT; + this.clientSideStatementType = null; + this.resultSet = null; + this.updateCount = updateCount; + this.noResult = null; + } + + private AsyncStatementResultImpl( + ClientSideStatementType clientSideStatementType, ApiFuture result) { + this.type = ResultType.NO_RESULT; + this.clientSideStatementType = clientSideStatementType; + this.resultSet = null; + this.updateCount = null; + this.noResult = result; + } + + @Override + public ResultType getResultType() { + return type; + } + + @Override + public ClientSideStatementType getClientSideStatementType() { + return clientSideStatementType; + } + + @Override + public ResultSet getResultSet() { + return getResultSetAsync(); + } + + @Override + public Long getUpdateCount() { + return get(getUpdateCountAsync()); + } + + @Override + public AsyncResultSet getResultSetAsync() { + ConnectionPreconditions.checkState( + resultSet != null, "This result does not contain a ResultSet"); + return resultSet; + } + + @Override + public ApiFuture getUpdateCountAsync() { + ConnectionPreconditions.checkState( + updateCount != null, "This result does not contain an update count"); + return updateCount; + } + + @Override + public ApiFuture getNoResultAsync() { + ConnectionPreconditions.checkState( + type == ResultType.NO_RESULT, "This result does not contain a 'no-result' result"); + return noResult; + } +} diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ChecksumResultSet.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ChecksumResultSet.java index 649d6c51fd..f2d1ba548e 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ChecksumResultSet.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ChecksumResultSet.java @@ -66,7 +66,7 @@ @VisibleForTesting class ChecksumResultSet extends ReplaceableForwardingResultSet implements RetriableStatement { private final ReadWriteTransaction transaction; - private long numberOfNextCalls; + private volatile long numberOfNextCalls; private final ParsedStatement statement; private final AnalyzeMode analyzeMode; private final QueryOption[] options; @@ -98,7 +98,13 @@ public Boolean call() throws Exception { .getStatementExecutor() .invokeInterceptors( statement, StatementExecutionStep.CALL_NEXT_ON_RESULT_SET, transaction); - return ChecksumResultSet.super.next(); + boolean res = ChecksumResultSet.super.next(); + // Only update the checksum if there was another row to be consumed. + if (res) { + checksumCalculator.calculateNextChecksum(getCurrentRowAsStruct()); + } + numberOfNextCalls++; + return res; } } @@ -107,13 +113,7 @@ public Boolean call() throws Exception { @Override public boolean next() { // Call next() with retry. - boolean res = transaction.runWithRetry(nextCallable); - // Only update the checksum if there was another row to be consumed. - if (res) { - checksumCalculator.calculateNextChecksum(getCurrentRowAsStruct()); - } - numberOfNextCalls++; - return res; + return transaction.runWithRetry(nextCallable); } @VisibleForTesting diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/Connection.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/Connection.java index 5247ce2c13..71b03e2e0b 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/Connection.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/Connection.java @@ -16,10 +16,12 @@ package com.google.cloud.spanner.connection; +import com.google.api.core.ApiFuture; import com.google.api.core.InternalApi; import com.google.cloud.Timestamp; import com.google.cloud.spanner.AbortedDueToConcurrentModificationException; import com.google.cloud.spanner.AbortedException; +import com.google.cloud.spanner.AsyncResultSet; import com.google.cloud.spanner.ErrorCode; import com.google.cloud.spanner.Mutation; import com.google.cloud.spanner.Options.QueryOption; @@ -31,6 +33,7 @@ import com.google.cloud.spanner.TimestampBound; import com.google.cloud.spanner.connection.StatementResult.ResultType; import java.util.Iterator; +import java.util.concurrent.ExecutionException; import java.util.concurrent.TimeUnit; /** @@ -41,6 +44,10 @@ * only exception is the {@link Connection#cancel()} method that may be called by any other thread * to stop the execution of the current statement on the connection. * + *

All -Async methods on {@link Connection} are guaranteed to be executed in the order that they + * are issued on the {@link Connection}. Mixing synchronous and asynchronous method calls is also + * supported, and these are also guaranteed to be executed in the order that they are issued. + * *

Connections accept a number of additional SQL statements for setting or changing the state of * a {@link Connection}. These statements can only be executed using the {@link * Connection#execute(Statement)} method: @@ -259,6 +266,25 @@ public interface Connection extends AutoCloseable { */ void beginTransaction(); + /** + * Begins a new transaction for this connection. This method is guaranteed to be non-blocking. The + * returned {@link ApiFuture} will be done when the transaction has been initialized. + * + *

    + *
  • Calling this method on a connection that has no transaction and that is + * not in autocommit mode, will register a new transaction that has not yet + * started on this connection + *
  • Calling this method on a connection that has no transaction and that is + * in autocommit mode, will register a new transaction that has not yet started on this + * connection, and temporarily turn off autocommit mode until the next commit/rollback + *
  • Calling this method on a connection that already has a transaction that has not yet + * started, will cause a {@link SpannerException} + *
  • Calling this method on a connection that already has a transaction that has started, will + * cause a {@link SpannerException} (no nested transactions) + *
+ */ + ApiFuture beginTransactionAsync(); + /** * Sets the transaction mode to use for current transaction. This method may only be called when * in a transaction, and before the transaction is actually started, i.e. before any statements @@ -450,6 +476,53 @@ public interface Connection extends AutoCloseable { */ void commit(); + /** + * Commits the current transaction of this connection. All mutations that have been buffered + * during the current transaction will be written to the database. + * + *

This method is guaranteed to be non-blocking. The returned {@link ApiFuture} will be done + * when the transaction has committed or the commit has failed. + * + *

Calling this method will always end the current transaction and start a new transaction when + * the next statement is executed, regardless whether this commit call succeeded or failed. If the + * next statement(s) rely on the results of the transaction that is being committed, it is + * recommended to check the status of this commit by inspecting the value of the returned {@link + * ApiFuture} before executing the next statement, to ensure that the commit actually succeeded. + * + *

If the connection is in autocommit mode, and there is a temporary transaction active on this + * connection, calling this method will cause the connection to go back to autocommit mode after + * calling this method. + * + *

This method will throw a {@link SpannerException} with code {@link + * ErrorCode#DEADLINE_EXCEEDED} if a statement timeout has been set on this connection, and the + * commit operation takes longer than this timeout. + * + *

    + *
  • Calling this method on a connection in autocommit mode and with no temporary transaction, + * will cause an exception + *
  • Calling this method while a DDL batch is active will cause an exception + *
  • Calling this method on a connection with a transaction that has not yet started, will end + * that transaction and any properties that might have been set on that transaction, and + * return the connection to its previous state. This means that if a transaction is created + * and set to read-only, and then committed before any statements have been executed, the + * read-only transaction is ended and any subsequent statements will be executed in a new + * transaction. If the connection is in read-write mode, the default for new transactions + * will be {@link TransactionMode#READ_WRITE_TRANSACTION}. Committing an empty transaction + * also does not generate a read timestamp or a commit timestamp, and calling one of the + * methods {@link Connection#getReadTimestamp()} or {@link Connection#getCommitTimestamp()} + * will cause an exception. + *
  • Calling this method on a connection with a {@link TransactionMode#READ_ONLY_TRANSACTION} + * transaction will end that transaction. If the connection is in read-write mode, any + * subsequent transaction will by default be a {@link + * TransactionMode#READ_WRITE_TRANSACTION} transaction, unless any following transaction is + * explicitly set to {@link TransactionMode#READ_ONLY_TRANSACTION} + *
  • Calling this method on a connection with a {@link TransactionMode#READ_WRITE_TRANSACTION} + * transaction will send all buffered mutations to the database, commit any DML statements + * that have been executed during this transaction and end the transaction. + *
+ */ + ApiFuture commitAsync(); + /** * Rollbacks the current transaction of this connection. All mutations or DDL statements that have * been buffered during the current transaction will be removed from the buffer. @@ -481,6 +554,40 @@ public interface Connection extends AutoCloseable { */ void rollback(); + /** + * Rollbacks the current transaction of this connection. All mutations or DDL statements that have + * been buffered during the current transaction will be removed from the buffer. + * + *

This method is guaranteed to be non-blocking. The returned {@link ApiFuture} will be done + * when the transaction has been rolled back. + * + *

If the connection is in autocommit mode, and there is a temporary transaction active on this + * connection, calling this method will cause the connection to go back to autocommit mode after + * calling this method. + * + *

    + *
  • Calling this method on a connection in autocommit mode and with no temporary transaction + * will cause an exception + *
  • Calling this method while a DDL batch is active will cause an exception + *
  • Calling this method on a connection with a transaction that has not yet started, will end + * that transaction and any properties that might have been set on that transaction, and + * return the connection to its previous state. This means that if a transaction is created + * and set to read-only, and then rolled back before any statements have been executed, the + * read-only transaction is ended and any subsequent statements will be executed in a new + * transaction. If the connection is in read-write mode, the default for new transactions + * will be {@link TransactionMode#READ_WRITE_TRANSACTION}. + *
  • Calling this method on a connection with a {@link TransactionMode#READ_ONLY_TRANSACTION} + * transaction will end that transaction. If the connection is in read-write mode, any + * subsequent transaction will by default be a {@link + * TransactionMode#READ_WRITE_TRANSACTION} transaction, unless any following transaction is + * explicitly set to {@link TransactionMode#READ_ONLY_TRANSACTION} + *
  • Calling this method on a connection with a {@link TransactionMode#READ_WRITE_TRANSACTION} + * transaction will clear all buffered mutations, rollback any DML statements that have been + * executed during this transaction and end the transaction. + *
+ */ + ApiFuture rollbackAsync(); + /** * @return true if this connection has a transaction (that has not necessarily * started). This method will only return false when the {@link Connection} is in autocommit @@ -572,11 +679,30 @@ public interface Connection extends AutoCloseable { *

This method may only be called when a (possibly empty) batch is active. * * @return the update counts in case of a DML batch. Returns an array containing 1 for each - * successful statement and 0 for each failed statement or statement that was not executed DDL - * in case of a DDL batch. + * successful statement and 0 for each failed statement or statement that was not executed in + * case of a DDL batch. */ long[] runBatch(); + /** + * Sends all buffered DML or DDL statements of the current batch to the database, waits for these + * to be executed and ends the current batch. The method will throw an exception for the first + * statement that cannot be executed, or return successfully if all statements could be executed. + * If an exception is thrown for a statement in the batch, the preceding statements in the same + * batch may still have been applied to the database. + * + *

This method is guaranteed to be non-blocking. The returned {@link ApiFuture} will be done + * when the batch has been successfully applied, or when one or more of the statements in the + * batch has failed and the further execution of the batch has been halted. + * + *

This method may only be called when a (possibly empty) batch is active. + * + * @return an {@link ApiFuture} containing the update counts in case of a DML batch. The {@link + * ApiFuture} contains an array containing 1 for each successful statement and 0 for each + * failed statement or statement that was not executed in case of a DDL batch. + */ + ApiFuture runBatchAsync(); + /** * Clears all buffered statements in the current batch and ends the batch. * @@ -608,6 +734,30 @@ public interface Connection extends AutoCloseable { */ StatementResult execute(Statement statement); + /** + * Executes the given statement if allowed in the current {@link TransactionMode} and connection + * state asynchronously. The returned value depends on the type of statement: + * + *

    + *
  • Queries will return an {@link AsyncResultSet} + *
  • DML statements will return an {@link ApiFuture} with an update count that is done when + * the DML statement has been applied successfully, or that throws an {@link + * ExecutionException} if the DML statement failed. + *
  • DDL statements will return an {@link ApiFuture} containing a {@link Void} that is done + * when the DDL statement has been applied successfully, or that throws an {@link + * ExecutionException} if the DDL statement failed. + *
  • Connection and transaction statements (SET AUTOCOMMIT=TRUE|FALSE, SHOW AUTOCOMMIT, SET + * TRANSACTION READ ONLY, etc) will return either a {@link ResultSet} or {@link + * ResultType#NO_RESULT}, depending on the type of statement (SHOW or SET) + *
+ * + * This method is guaranteed to be non-blocking. + * + * @param statement The statement to execute + * @return the result of the statement + */ + AsyncStatementResult executeAsync(Statement statement); + /** * Executes the given statement as a query and returns the result as a {@link ResultSet}. This * method blocks and waits for a response from Spanner. If the statement does not contain a valid @@ -619,6 +769,31 @@ public interface Connection extends AutoCloseable { */ ResultSet executeQuery(Statement query, QueryOption... options); + /** + * Same as {@link #executeQuery(Statement, QueryOption...)}, but is guaranteed to be non-blocking + * and returns the query result as an {@link AsyncResultSet}. See {@link + * AsyncResultSet#setCallback(java.util.concurrent.Executor, + * com.google.cloud.spanner.AsyncResultSet.ReadyCallback)} for more information on how to consume + * the results of the query asynchronously. + */ + /** + * Executes the given statement asynchronously as a query and returns the result as an {@link + * AsyncResultSet}. This method is guaranteed to be non-blocking. If the statement does not + * contain a valid query, the method will throw a {@link SpannerException}. + * + *

See {@link AsyncResultSet#setCallback(java.util.concurrent.Executor, + * com.google.cloud.spanner.AsyncResultSet.ReadyCallback)} for more information on how to consume + * the results of the query asynchronously. + * + *

It is also possible to consume the returned {@link AsyncResultSet} in the same way as a + * normal {@link ResultSet}, i.e. in a while-loop calling {@link AsyncResultSet#next()}. + * + * @param query The query statement to execute + * @param options the options to configure the query + * @return an {@link AsyncResultSet} with the results of the query + */ + AsyncResultSet executeQueryAsync(Statement query, QueryOption... options); + /** * Analyzes a query and returns query plan and/or query execution statistics information. * @@ -655,6 +830,18 @@ public interface Connection extends AutoCloseable { */ long executeUpdate(Statement update); + /** + * Executes the given statement asynchronously as a DML statement. If the statement does not + * contain a valid DML statement, the method will throw a {@link SpannerException}. + * + *

This method is guaranteed to be non-blocking. + * + * @param update The update statement to execute + * @return an {@link ApiFuture} containing the number of records that were + * inserted/updated/deleted by this statement + */ + ApiFuture executeUpdateAsync(Statement update); + /** * Executes a list of DML statements in a single request. The statements will be executed in order * and the semantics is the same as if each statement is executed by {@link @@ -677,6 +864,31 @@ public interface Connection extends AutoCloseable { */ long[] executeBatchUpdate(Iterable updates); + /** + * Executes a list of DML statements in a single request. The statements will be executed in order + * and the semantics is the same as if each statement is executed by {@link + * Connection#executeUpdate(Statement)} in a loop. This method returns an {@link ApiFuture} that + * contains an array of long integers, each representing the number of rows modified by each + * statement. + * + *

This method is guaranteed to be non-blocking. + * + *

If an individual statement fails, execution stops and a {@code SpannerBatchUpdateException} + * is returned, which includes the error and the number of rows affected by the statements that + * are run prior to the error. + * + *

For example, if statements contains 3 statements, and the 2nd one is not a valid DML. This + * method throws a {@code SpannerBatchUpdateException} that contains the error message from the + * 2nd statement, and an array of length 1 that contains the number of rows modified by the 1st + * statement. The 3rd statement will not run. Executes the given statements as DML statements in + * one batch. If one of the statements does not contain a valid DML statement, the method will + * throw a {@link SpannerException}. + * + * @param updates The update statements that will be executed as one batch. + * @return an {@link ApiFuture} containing an array with the update counts per statement. + */ + ApiFuture executeBatchUpdateAsync(Iterable updates); + /** * Writes the specified mutation directly to the database and commits the change. The value is * readable after the successful completion of this method. Writing multiple mutations to a @@ -692,6 +904,23 @@ public interface Connection extends AutoCloseable { */ void write(Mutation mutation); + /** + * Writes the specified mutation directly to the database and commits the change. The value is + * readable after the successful completion of the returned {@link ApiFuture}. Writing multiple + * mutations to a database by calling this method multiple times mode is inefficient, as each call + * will need a round trip to the database. Instead, you should consider writing the mutations + * together by calling {@link Connection#writeAsync(Iterable)}. + * + *

This method is guaranteed to be non-blocking. + * + *

Calling this method is only allowed in autocommit mode. See {@link + * Connection#bufferedWrite(Iterable)} for writing mutations in transactions. + * + * @param mutation The {@link Mutation} to write to the database + * @throws SpannerException if the {@link Connection} is not in autocommit mode + */ + ApiFuture writeAsync(Mutation mutation); + /** * Writes the specified mutations directly to the database and commits the changes. The values are * readable after the successful completion of this method. @@ -704,6 +933,20 @@ public interface Connection extends AutoCloseable { */ void write(Iterable mutations); + /** + * Writes the specified mutations directly to the database and commits the changes. The values are + * readable after the successful completion of the returned {@link ApiFuture}. + * + *

This method is guaranteed to be non-blocking. + * + *

Calling this method is only allowed in autocommit mode. See {@link + * Connection#bufferedWrite(Iterable)} for writing mutations in transactions. + * + * @param mutations The {@link Mutation}s to write to the database + * @throws SpannerException if the {@link Connection} is not in autocommit mode + */ + ApiFuture writeAsync(Iterable mutations); + /** * Buffers the given mutation locally on the current transaction of this {@link Connection}. The * mutation will be written to the database at the next call to {@link Connection#commit()}. The diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionImpl.java index ce24791859..b49adbf124 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionImpl.java @@ -16,13 +16,19 @@ package com.google.cloud.spanner.connection; +import static com.google.cloud.spanner.SpannerApiFutures.get; + +import com.google.api.core.ApiFuture; +import com.google.api.core.ApiFutures; import com.google.cloud.Timestamp; +import com.google.cloud.spanner.AsyncResultSet; import com.google.cloud.spanner.DatabaseClient; import com.google.cloud.spanner.ErrorCode; import com.google.cloud.spanner.Mutation; import com.google.cloud.spanner.Options.QueryOption; import com.google.cloud.spanner.ReadContext.QueryAnalyzeMode; import com.google.cloud.spanner.ResultSet; +import com.google.cloud.spanner.ResultSets; import com.google.cloud.spanner.Spanner; import com.google.cloud.spanner.SpannerException; import com.google.cloud.spanner.SpannerExceptionFactory; @@ -583,6 +589,11 @@ private void setDefaultTransactionOptions() { @Override public void beginTransaction() { + get(beginTransactionAsync()); + } + + @Override + public ApiFuture beginTransactionAsync() { ConnectionPreconditions.checkState(!isClosed(), CLOSED_ERROR_MSG); ConnectionPreconditions.checkState( !isBatchActive(), "This connection has an active batch and cannot begin a transaction"); @@ -596,17 +607,18 @@ public void beginTransaction() { if (isAutocommit()) { inTransaction = true; } + return ApiFutures.immediateFuture(null); } /** Internal interface for ending a transaction (commit/rollback). */ private static interface EndTransactionMethod { - public void end(UnitOfWork t); + public ApiFuture endAsync(UnitOfWork t); } private static final class Commit implements EndTransactionMethod { @Override - public void end(UnitOfWork t) { - t.commit(); + public ApiFuture endAsync(UnitOfWork t) { + return t.commitAsync(); } } @@ -614,14 +626,18 @@ public void end(UnitOfWork t) { @Override public void commit() { + get(commitAsync()); + } + + public ApiFuture commitAsync() { ConnectionPreconditions.checkState(!isClosed(), CLOSED_ERROR_MSG); - endCurrentTransaction(commit); + return endCurrentTransactionAsync(commit); } private static final class Rollback implements EndTransactionMethod { @Override - public void end(UnitOfWork t) { - t.rollback(); + public ApiFuture endAsync(UnitOfWork t) { + return t.rollbackAsync(); } } @@ -629,18 +645,24 @@ public void end(UnitOfWork t) { @Override public void rollback() { + get(rollbackAsync()); + } + + public ApiFuture rollbackAsync() { ConnectionPreconditions.checkState(!isClosed(), CLOSED_ERROR_MSG); - endCurrentTransaction(rollback); + return endCurrentTransactionAsync(rollback); } - private void endCurrentTransaction(EndTransactionMethod endTransactionMethod) { + private ApiFuture endCurrentTransactionAsync(EndTransactionMethod endTransactionMethod) { ConnectionPreconditions.checkState(!isBatchActive(), "This connection has an active batch"); ConnectionPreconditions.checkState(isInTransaction(), "This connection has no transaction"); + ApiFuture res; try { if (isTransactionStarted()) { - endTransactionMethod.end(getCurrentUnitOfWorkOrStartNewUnitOfWork()); + res = endTransactionMethod.endAsync(getCurrentUnitOfWorkOrStartNewUnitOfWork()); } else { this.currentUnitOfWork = null; + res = ApiFutures.immediateFuture(null); } } finally { transactionBeginMarked = false; @@ -649,6 +671,7 @@ private void endCurrentTransaction(EndTransactionMethod endTransactionMethod) { } setDefaultTransactionOptions(); } + return res; } @Override @@ -664,9 +687,9 @@ public StatementResult execute(Statement statement) { case QUERY: return StatementResultImpl.of(internalExecuteQuery(parsedStatement, AnalyzeMode.NONE)); case UPDATE: - return StatementResultImpl.of(internalExecuteUpdate(parsedStatement)); + return StatementResultImpl.of(get(internalExecuteUpdateAsync(parsedStatement))); case DDL: - executeDdl(parsedStatement); + get(executeDdlAsync(parsedStatement)); return StatementResultImpl.noResult(); case UNKNOWN: default: @@ -676,11 +699,43 @@ public StatementResult execute(Statement statement) { "Unknown statement: " + parsedStatement.getSqlWithoutComments()); } + @Override + public AsyncStatementResult executeAsync(Statement statement) { + Preconditions.checkNotNull(statement); + ConnectionPreconditions.checkState(!isClosed(), CLOSED_ERROR_MSG); + ParsedStatement parsedStatement = parser.parse(statement, this.queryOptions); + switch (parsedStatement.getType()) { + case CLIENT_SIDE: + return AsyncStatementResultImpl.of( + parsedStatement + .getClientSideStatement() + .execute(connectionStatementExecutor, parsedStatement.getSqlWithoutComments()), + spanner.getAsyncExecutorProvider()); + case QUERY: + return AsyncStatementResultImpl.of( + internalExecuteQueryAsync(parsedStatement, AnalyzeMode.NONE)); + case UPDATE: + return AsyncStatementResultImpl.of(internalExecuteUpdateAsync(parsedStatement)); + case DDL: + return AsyncStatementResultImpl.noResult(executeDdlAsync(parsedStatement)); + case UNKNOWN: + default: + } + throw SpannerExceptionFactory.newSpannerException( + ErrorCode.INVALID_ARGUMENT, + "Unknown statement: " + parsedStatement.getSqlWithoutComments()); + } + @Override public ResultSet executeQuery(Statement query, QueryOption... options) { return parseAndExecuteQuery(query, AnalyzeMode.NONE, options); } + @Override + public AsyncResultSet executeQueryAsync(Statement query, QueryOption... options) { + return parseAndExecuteQueryAsync(query, AnalyzeMode.NONE, options); + } + @Override public ResultSet analyzeQuery(Statement query, QueryAnalyzeMode queryMode) { Preconditions.checkNotNull(queryMode); @@ -717,6 +772,34 @@ private ResultSet parseAndExecuteQuery( "Statement is not a query: " + parsedStatement.getSqlWithoutComments()); } + private AsyncResultSet parseAndExecuteQueryAsync( + Statement query, AnalyzeMode analyzeMode, QueryOption... options) { + Preconditions.checkNotNull(query); + ConnectionPreconditions.checkState(!isClosed(), CLOSED_ERROR_MSG); + ParsedStatement parsedStatement = parser.parse(query, this.queryOptions); + if (parsedStatement.isQuery()) { + switch (parsedStatement.getType()) { + case CLIENT_SIDE: + return ResultSets.toAsyncResultSet( + parsedStatement + .getClientSideStatement() + .execute(connectionStatementExecutor, parsedStatement.getSqlWithoutComments()) + .getResultSet(), + spanner.getAsyncExecutorProvider(), + options); + case QUERY: + return internalExecuteQueryAsync(parsedStatement, analyzeMode, options); + case UPDATE: + case DDL: + case UNKNOWN: + default: + } + } + throw SpannerExceptionFactory.newSpannerException( + ErrorCode.INVALID_ARGUMENT, + "Statement is not a query: " + parsedStatement.getSqlWithoutComments()); + } + @Override public long executeUpdate(Statement update) { Preconditions.checkNotNull(update); @@ -725,7 +808,27 @@ public long executeUpdate(Statement update) { if (parsedStatement.isUpdate()) { switch (parsedStatement.getType()) { case UPDATE: - return internalExecuteUpdate(parsedStatement); + return get(internalExecuteUpdateAsync(parsedStatement)); + case CLIENT_SIDE: + case QUERY: + case DDL: + case UNKNOWN: + default: + } + } + throw SpannerExceptionFactory.newSpannerException( + ErrorCode.INVALID_ARGUMENT, + "Statement is not an update statement: " + parsedStatement.getSqlWithoutComments()); + } + + public ApiFuture executeUpdateAsync(Statement update) { + Preconditions.checkNotNull(update); + ConnectionPreconditions.checkState(!isClosed(), CLOSED_ERROR_MSG); + ParsedStatement parsedStatement = parser.parse(update); + if (parsedStatement.isUpdate()) { + switch (parsedStatement.getType()) { + case UPDATE: + return internalExecuteUpdateAsync(parsedStatement); case CLIENT_SIDE: case QUERY: case DDL: @@ -746,24 +849,48 @@ public long[] executeBatchUpdate(Iterable updates) { List parsedStatements = new LinkedList<>(); for (Statement update : updates) { ParsedStatement parsedStatement = parser.parse(update); - if (parsedStatement.isUpdate()) { - switch (parsedStatement.getType()) { - case UPDATE: - parsedStatements.add(parsedStatement); - break; - case CLIENT_SIDE: - case QUERY: - case DDL: - case UNKNOWN: - default: - throw SpannerExceptionFactory.newSpannerException( - ErrorCode.INVALID_ARGUMENT, - "The batch update list contains a statement that is not an update statement: " - + parsedStatement.getSqlWithoutComments()); - } + switch (parsedStatement.getType()) { + case UPDATE: + parsedStatements.add(parsedStatement); + break; + case CLIENT_SIDE: + case QUERY: + case DDL: + case UNKNOWN: + default: + throw SpannerExceptionFactory.newSpannerException( + ErrorCode.INVALID_ARGUMENT, + "The batch update list contains a statement that is not an update statement: " + + parsedStatement.getSqlWithoutComments()); + } + } + return get(internalExecuteBatchUpdateAsync(parsedStatements)); + } + + @Override + public ApiFuture executeBatchUpdateAsync(Iterable updates) { + Preconditions.checkNotNull(updates); + ConnectionPreconditions.checkState(!isClosed(), CLOSED_ERROR_MSG); + // Check that there are only DML statements in the input. + List parsedStatements = new LinkedList<>(); + for (Statement update : updates) { + ParsedStatement parsedStatement = parser.parse(update); + switch (parsedStatement.getType()) { + case UPDATE: + parsedStatements.add(parsedStatement); + break; + case CLIENT_SIDE: + case QUERY: + case DDL: + case UNKNOWN: + default: + throw SpannerExceptionFactory.newSpannerException( + ErrorCode.INVALID_ARGUMENT, + "The batch update list contains a statement that is not an update statement: " + + parsedStatement.getSqlWithoutComments()); } } - return internalExecuteBatchUpdate(parsedStatements); + return internalExecuteBatchUpdateAsync(parsedStatements); } private ResultSet internalExecuteQuery( @@ -773,52 +900,32 @@ private ResultSet internalExecuteQuery( Preconditions.checkArgument( statement.getType() == StatementType.QUERY, "Statement must be a query"); UnitOfWork transaction = getCurrentUnitOfWorkOrStartNewUnitOfWork(); - try { - return transaction.executeQuery(statement, analyzeMode, options); - } catch (SpannerException e) { - // In case of a timed out or cancelled query we need to replace the executor to ensure that we - // have an executor that is not busy executing a statement. Although we try to cancel the - // current statement, it is not guaranteed to actually stop the execution directly. - if (e.getErrorCode() == ErrorCode.DEADLINE_EXCEEDED - || e.getErrorCode() == ErrorCode.CANCELLED) { - this.statementExecutor.recreate(); - } - throw e; - } + return get(transaction.executeQueryAsync(statement, analyzeMode, options)); } - private long internalExecuteUpdate(final ParsedStatement update) { + private AsyncResultSet internalExecuteQueryAsync( + final ParsedStatement statement, + final AnalyzeMode analyzeMode, + final QueryOption... options) { + Preconditions.checkArgument( + statement.getType() == StatementType.QUERY, "Statement must be a query"); + UnitOfWork transaction = getCurrentUnitOfWorkOrStartNewUnitOfWork(); + return ResultSets.toAsyncResultSet( + transaction.executeQueryAsync(statement, analyzeMode, options), + spanner.getAsyncExecutorProvider(), + options); + } + + private ApiFuture internalExecuteUpdateAsync(final ParsedStatement update) { Preconditions.checkArgument( update.getType() == StatementType.UPDATE, "Statement must be an update"); UnitOfWork transaction = getCurrentUnitOfWorkOrStartNewUnitOfWork(); - try { - return transaction.executeUpdate(update); - } catch (SpannerException e) { - // In case of a timed out or cancelled query we need to replace the executor to ensure that we - // have an executor that is not busy executing a statement. Although we try to cancel the - // current statement, it is not guaranteed to actually stop the execution directly. - if (e.getErrorCode() == ErrorCode.DEADLINE_EXCEEDED - || e.getErrorCode() == ErrorCode.CANCELLED) { - this.statementExecutor.recreate(); - } - throw e; - } + return transaction.executeUpdateAsync(update); } - private long[] internalExecuteBatchUpdate(final List updates) { + private ApiFuture internalExecuteBatchUpdateAsync(List updates) { UnitOfWork transaction = getCurrentUnitOfWorkOrStartNewUnitOfWork(); - try { - return transaction.executeBatchUpdate(updates); - } catch (SpannerException e) { - // In case of a timed out or cancelled query we need to replace the executor to ensure that we - // have an executor that is not busy executing a statement. Although we try to cancel the - // current statement, it is not guaranteed to actually stop the execution directly. - if (e.getErrorCode() == ErrorCode.DEADLINE_EXCEEDED - || e.getErrorCode() == ErrorCode.CANCELLED) { - this.statementExecutor.recreate(); - } - throw e; - } + return transaction.executeBatchUpdateAsync(updates); } /** @@ -898,32 +1005,36 @@ private void popUnitOfWorkFromTransactionStack() { this.currentUnitOfWork = transactionStack.pop(); } - private void executeDdl(ParsedStatement ddl) { - getCurrentUnitOfWorkOrStartNewUnitOfWork().executeDdl(ddl); + private ApiFuture executeDdlAsync(ParsedStatement ddl) { + return getCurrentUnitOfWorkOrStartNewUnitOfWork().executeDdlAsync(ddl); } @Override public void write(Mutation mutation) { - Preconditions.checkNotNull(mutation); - ConnectionPreconditions.checkState(!isClosed(), CLOSED_ERROR_MSG); - ConnectionPreconditions.checkState(isAutocommit(), ONLY_ALLOWED_IN_AUTOCOMMIT); - getCurrentUnitOfWorkOrStartNewUnitOfWork().write(mutation); + get(writeAsync(Collections.singleton(Preconditions.checkNotNull(mutation)))); + } + + @Override + public ApiFuture writeAsync(Mutation mutation) { + return writeAsync(Collections.singleton(Preconditions.checkNotNull(mutation))); } @Override public void write(Iterable mutations) { + get(writeAsync(Preconditions.checkNotNull(mutations))); + } + + @Override + public ApiFuture writeAsync(Iterable mutations) { Preconditions.checkNotNull(mutations); ConnectionPreconditions.checkState(!isClosed(), CLOSED_ERROR_MSG); ConnectionPreconditions.checkState(isAutocommit(), ONLY_ALLOWED_IN_AUTOCOMMIT); - getCurrentUnitOfWorkOrStartNewUnitOfWork().write(mutations); + return getCurrentUnitOfWorkOrStartNewUnitOfWork().writeAsync(mutations); } @Override public void bufferedWrite(Mutation mutation) { - Preconditions.checkNotNull(mutation); - ConnectionPreconditions.checkState(!isClosed(), CLOSED_ERROR_MSG); - ConnectionPreconditions.checkState(!isAutocommit(), NOT_ALLOWED_IN_AUTOCOMMIT); - getCurrentUnitOfWorkOrStartNewUnitOfWork().write(mutation); + bufferedWrite(Preconditions.checkNotNull(Collections.singleton(mutation))); } @Override @@ -931,7 +1042,7 @@ public void bufferedWrite(Iterable mutations) { Preconditions.checkNotNull(mutations); ConnectionPreconditions.checkState(!isClosed(), CLOSED_ERROR_MSG); ConnectionPreconditions.checkState(!isAutocommit(), NOT_ALLOWED_IN_AUTOCOMMIT); - getCurrentUnitOfWorkOrStartNewUnitOfWork().write(mutations); + get(getCurrentUnitOfWorkOrStartNewUnitOfWork().writeAsync(mutations)); } @Override @@ -973,13 +1084,18 @@ public void startBatchDml() { @Override public long[] runBatch() { + return get(runBatchAsync()); + } + + @Override + public ApiFuture runBatchAsync() { ConnectionPreconditions.checkState(!isClosed(), CLOSED_ERROR_MSG); ConnectionPreconditions.checkState(isBatchActive(), "This connection has no active batch"); try { if (this.currentUnitOfWork != null) { - return this.currentUnitOfWork.runBatch(); + return this.currentUnitOfWork.runBatchAsync(); } - return new long[0]; + return ApiFutures.immediateFuture(new long[0]); } finally { this.batchMode = BatchMode.NONE; setDefaultTransactionOptions(); diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionOptions.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionOptions.java index 379459884c..d2a341430e 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionOptions.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionOptions.java @@ -237,6 +237,15 @@ public static void closeSpanner() { SpannerPool.INSTANCE.checkAndCloseSpanners(); } + /** + * {@link SpannerOptionsConfigurator} can be used to add additional configuration for a {@link + * Spanner} instance. Intended for tests. + */ + @VisibleForTesting + interface SpannerOptionsConfigurator { + void configure(SpannerOptions.Builder options); + } + /** Builder for {@link ConnectionOptions} instances. */ public static class Builder { private String uri; @@ -246,6 +255,7 @@ public static class Builder { private SessionPoolOptions sessionPoolOptions; private List statementExecutionInterceptors = Collections.emptyList(); + private SpannerOptionsConfigurator configurator; private Builder() {} @@ -358,6 +368,12 @@ Builder setStatementExecutionInterceptors(List in return this; } + @VisibleForTesting + Builder setConfigurator(SpannerOptionsConfigurator configurator) { + this.configurator = Preconditions.checkNotNull(configurator); + return this; + } + @VisibleForTesting Builder setCredentials(Credentials credentials) { this.credentials = credentials; @@ -401,6 +417,7 @@ public static Builder newBuilder() { private final boolean readOnly; private final boolean retryAbortsInternally; private final List statementExecutionInterceptors; + private final SpannerOptionsConfigurator configurator; private ConnectionOptions(Builder builder) { Matcher matcher = Builder.SPANNER_URI_PATTERN.matcher(builder.uri); @@ -473,6 +490,11 @@ private ConnectionOptions(Builder builder) { this.retryAbortsInternally = parseRetryAbortsInternally(this.uri); this.statementExecutionInterceptors = Collections.unmodifiableList(builder.statementExecutionInterceptors); + this.configurator = builder.configurator; + } + + SpannerOptionsConfigurator getConfigurator() { + return configurator; } @VisibleForTesting diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/DdlBatch.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/DdlBatch.java index a80e93dfc0..7d4f18c4db 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/DdlBatch.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/DdlBatch.java @@ -16,6 +16,8 @@ package com.google.cloud.spanner.connection; +import com.google.api.core.ApiFuture; +import com.google.api.core.ApiFutures; import com.google.api.gax.longrunning.OperationFuture; import com.google.cloud.Timestamp; import com.google.cloud.spanner.DatabaseClient; @@ -31,15 +33,14 @@ import com.google.cloud.spanner.connection.StatementParser.StatementType; import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Preconditions; +import com.google.spanner.admin.database.v1.DatabaseAdminGrpc; import com.google.spanner.admin.database.v1.UpdateDatabaseDdlMetadata; +import com.google.spanner.v1.SpannerGrpc; import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; -import java.util.HashSet; import java.util.List; -import java.util.Set; import java.util.concurrent.Callable; -import java.util.concurrent.ExecutionException; /** * {@link UnitOfWork} that is used when a DDL batch is started. These batches only accept DDL @@ -111,8 +112,7 @@ public boolean isReadOnly() { return false; } - @Override - public ResultSet executeQuery( + public ApiFuture executeQueryAsync( final ParsedStatement statement, AnalyzeMode analyzeMode, QueryOption... options) { if (options != null) { for (int i = 0; i < options.length; i++) { @@ -136,7 +136,8 @@ public ResultSet call() throws Exception { dbClient.singleUse().executeQuery(statement.getStatement(), internalOptions)); } }; - return asyncExecuteStatement(statement, callable); + return executeStatementAsync( + statement, callable, SpannerGrpc.getExecuteStreamingSqlMethod()); } } } @@ -168,7 +169,7 @@ public Timestamp getCommitTimestampOrNull() { } @Override - public void executeDdl(ParsedStatement ddl) { + public ApiFuture executeDdlAsync(ParsedStatement ddl) { ConnectionPreconditions.checkState( state == UnitOfWorkState.STARTED, "The batch is no longer active and cannot be used for further statements"); @@ -178,28 +179,23 @@ public void executeDdl(ParsedStatement ddl) { + ddl.getSqlWithoutComments() + "\" is not a DDL-statement."); statements.add(ddl.getSqlWithoutComments()); + return ApiFutures.immediateFuture(null); } @Override - public long executeUpdate(ParsedStatement update) { + public ApiFuture executeUpdateAsync(ParsedStatement update) { throw SpannerExceptionFactory.newSpannerException( ErrorCode.FAILED_PRECONDITION, "Executing updates is not allowed for DDL batches."); } @Override - public long[] executeBatchUpdate(Iterable updates) { + public ApiFuture executeBatchUpdateAsync(Iterable updates) { throw SpannerExceptionFactory.newSpannerException( ErrorCode.FAILED_PRECONDITION, "Executing batch updates is not allowed for DDL batches."); } @Override - public void write(Mutation mutation) { - throw SpannerExceptionFactory.newSpannerException( - ErrorCode.FAILED_PRECONDITION, "Writing mutations is not allowed for DDL batches."); - } - - @Override - public void write(Iterable mutations) { + public ApiFuture writeAsync(Iterable mutations) { throw SpannerExceptionFactory.newSpannerException( ErrorCode.FAILED_PRECONDITION, "Writing mutations is not allowed for DDL batches."); } @@ -214,62 +210,50 @@ public void write(Iterable mutations) { StatementParser.INSTANCE.parse(Statement.of("RUN BATCH")); @Override - public long[] runBatch() { + public ApiFuture runBatchAsync() { ConnectionPreconditions.checkState( state == UnitOfWorkState.STARTED, "The batch is no longer active and cannot be ran"); - try { - if (!statements.isEmpty()) { - // create a statement that can be passed in to the execute method - Callable callable = - new Callable() { - @Override - public UpdateDatabaseDdlMetadata call() throws Exception { - OperationFuture operation = - ddlClient.executeDdl(statements); - try { - // Wait until the operation has finished. - operation.get(); - // Return metadata. - return operation.getMetadata().get(); - } catch (ExecutionException e) { - SpannerException spannerException = extractSpannerCause(e); - long[] updateCounts = extractUpdateCounts(operation.getMetadata().get()); - throw SpannerExceptionFactory.newSpannerBatchUpdateException( - spannerException == null - ? ErrorCode.UNKNOWN - : spannerException.getErrorCode(), - e.getMessage(), - updateCounts); - } catch (InterruptedException e) { - long[] updateCounts = extractUpdateCounts(operation.getMetadata().get()); - throw SpannerExceptionFactory.newSpannerBatchUpdateException( - ErrorCode.CANCELLED, e.getMessage(), updateCounts); - } - } - }; - asyncExecuteStatement(RUN_BATCH, callable); - } + if (statements.isEmpty()) { this.state = UnitOfWorkState.RAN; - long[] updateCounts = new long[statements.size()]; - Arrays.fill(updateCounts, 1L); - return updateCounts; - } catch (SpannerException e) { - this.state = UnitOfWorkState.RUN_FAILED; - throw e; + return ApiFutures.immediateFuture(new long[0]); } + // create a statement that can be passed in to the execute method + Callable callable = + new Callable() { + @Override + public long[] call() throws Exception { + try { + OperationFuture operation = + ddlClient.executeDdl(statements); + try { + // Wait until the operation has finished. + getWithStatementTimeout(operation, RUN_BATCH); + long[] updateCounts = new long[statements.size()]; + Arrays.fill(updateCounts, 1L); + state = UnitOfWorkState.RAN; + return updateCounts; + } catch (SpannerException e) { + long[] updateCounts = extractUpdateCounts(operation); + throw SpannerExceptionFactory.newSpannerBatchUpdateException( + e.getErrorCode(), e.getMessage(), updateCounts); + } + } catch (Throwable t) { + state = UnitOfWorkState.RUN_FAILED; + throw t; + } + } + }; + this.state = UnitOfWorkState.RUNNING; + return executeStatementAsync( + RUN_BATCH, callable, DatabaseAdminGrpc.getUpdateDatabaseDdlMethod()); } - private SpannerException extractSpannerCause(ExecutionException e) { - Throwable cause = e.getCause(); - Set causes = new HashSet<>(); - while (cause != null && !causes.contains(cause)) { - if (cause instanceof SpannerException) { - return (SpannerException) cause; - } - causes.add(cause); - cause = cause.getCause(); + long[] extractUpdateCounts(OperationFuture operation) { + try { + return extractUpdateCounts(operation.getMetadata().get()); + } catch (Throwable t) { + return new long[0]; } - return null; } @VisibleForTesting @@ -293,13 +277,13 @@ public void abortBatch() { } @Override - public void commit() { + public ApiFuture commitAsync() { throw SpannerExceptionFactory.newSpannerException( ErrorCode.FAILED_PRECONDITION, "Commit is not allowed for DDL batches."); } @Override - public void rollback() { + public ApiFuture rollbackAsync() { throw SpannerExceptionFactory.newSpannerException( ErrorCode.FAILED_PRECONDITION, "Rollback is not allowed for DDL batches."); } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/DmlBatch.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/DmlBatch.java index ff38338d62..b5b80e46cf 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/DmlBatch.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/DmlBatch.java @@ -16,16 +16,20 @@ package com.google.cloud.spanner.connection; +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.ErrorCode; import com.google.cloud.spanner.Mutation; import com.google.cloud.spanner.Options.QueryOption; import com.google.cloud.spanner.ResultSet; -import com.google.cloud.spanner.SpannerException; import com.google.cloud.spanner.SpannerExceptionFactory; import com.google.cloud.spanner.connection.StatementParser.ParsedStatement; import com.google.cloud.spanner.connection.StatementParser.StatementType; import com.google.common.base.Preconditions; +import com.google.common.util.concurrent.MoreExecutors; import java.util.ArrayList; import java.util.List; @@ -87,7 +91,7 @@ public boolean isReadOnly() { } @Override - public ResultSet executeQuery( + public ApiFuture executeQueryAsync( ParsedStatement statement, AnalyzeMode analyzeMode, QueryOption... options) { throw SpannerExceptionFactory.newSpannerException( ErrorCode.FAILED_PRECONDITION, "Executing queries is not allowed for DML batches."); @@ -116,13 +120,13 @@ public Timestamp getCommitTimestampOrNull() { } @Override - public void executeDdl(ParsedStatement ddl) { + public ApiFuture executeDdlAsync(ParsedStatement ddl) { throw SpannerExceptionFactory.newSpannerException( ErrorCode.FAILED_PRECONDITION, "Executing DDL statements is not allowed for DML batches."); } @Override - public long executeUpdate(ParsedStatement update) { + public ApiFuture executeUpdateAsync(ParsedStatement update) { ConnectionPreconditions.checkState( state == UnitOfWorkState.STARTED, "The batch is no longer active and cannot be used for further statements"); @@ -132,44 +136,54 @@ public long executeUpdate(ParsedStatement update) { + update.getSqlWithoutComments() + "\" is not a DML-statement."); statements.add(update); - return -1L; + return ApiFutures.immediateFuture(-1L); } @Override - public long[] executeBatchUpdate(Iterable updates) { + public ApiFuture executeBatchUpdateAsync(Iterable updates) { throw SpannerExceptionFactory.newSpannerException( ErrorCode.FAILED_PRECONDITION, "Executing batch updates is not allowed for DML batches."); } @Override - public void write(Mutation mutation) { + public ApiFuture writeAsync(Iterable mutations) { throw SpannerExceptionFactory.newSpannerException( ErrorCode.FAILED_PRECONDITION, "Writing mutations is not allowed for DML batches."); } @Override - public void write(Iterable mutations) { - throw SpannerExceptionFactory.newSpannerException( - ErrorCode.FAILED_PRECONDITION, "Writing mutations is not allowed for DML batches."); - } - - @Override - public long[] runBatch() { + public ApiFuture runBatchAsync() { ConnectionPreconditions.checkState( state == UnitOfWorkState.STARTED, "The batch is no longer active and cannot be ran"); - try { - long[] res; - if (statements.isEmpty()) { - res = new long[0]; - } else { - res = transaction.executeBatchUpdate(statements); - } + if (statements.isEmpty()) { this.state = UnitOfWorkState.RAN; - return res; - } catch (SpannerException e) { - this.state = UnitOfWorkState.RUN_FAILED; - throw e; + return ApiFutures.immediateFuture(new long[0]); } + this.state = UnitOfWorkState.RUNNING; + // Use a SettableApiFuture to return the result, instead of directly returning the future that + // is returned by the executeBatchUpdateAsync method. This is needed because the state of the + // batch is set after the update has finished, and this happens in a listener. A listener is + // executed AFTER a Future is done, which means that a user could read the state of the Batch + // before it has been changed. + final SettableApiFuture res = SettableApiFuture.create(); + ApiFuture updateCounts = transaction.executeBatchUpdateAsync(statements); + ApiFutures.addCallback( + updateCounts, + new ApiFutureCallback() { + @Override + public void onFailure(Throwable t) { + state = UnitOfWorkState.RUN_FAILED; + res.setException(t); + } + + @Override + public void onSuccess(long[] result) { + state = UnitOfWorkState.RAN; + res.set(result); + } + }, + MoreExecutors.directExecutor()); + return res; } @Override @@ -180,13 +194,13 @@ public void abortBatch() { } @Override - public void commit() { + public ApiFuture commitAsync() { throw SpannerExceptionFactory.newSpannerException( ErrorCode.FAILED_PRECONDITION, "Commit is not allowed for DML batches."); } @Override - public void rollback() { + public ApiFuture rollbackAsync() { throw SpannerExceptionFactory.newSpannerException( ErrorCode.FAILED_PRECONDITION, "Rollback is not allowed for DML batches."); } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ReadOnlyTransaction.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ReadOnlyTransaction.java index c9435886c0..09f3efc6d5 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ReadOnlyTransaction.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ReadOnlyTransaction.java @@ -16,6 +16,8 @@ package com.google.cloud.spanner.connection; +import com.google.api.core.ApiFuture; +import com.google.api.core.ApiFutures; import com.google.cloud.Timestamp; import com.google.cloud.spanner.DatabaseClient; import com.google.cloud.spanner.ErrorCode; @@ -83,6 +85,11 @@ public boolean isReadOnly() { return true; } + @Override + void checkAborted() { + // No-op for read-only transactions as they cannot abort. + } + @Override void checkValidTransaction() { if (transaction == null) { @@ -130,49 +137,45 @@ public Timestamp getCommitTimestampOrNull() { } @Override - public void executeDdl(ParsedStatement ddl) { + public ApiFuture executeDdlAsync(ParsedStatement ddl) { throw SpannerExceptionFactory.newSpannerException( ErrorCode.FAILED_PRECONDITION, "DDL statements are not allowed for read-only transactions"); } @Override - public long executeUpdate(ParsedStatement update) { + public ApiFuture executeUpdateAsync(ParsedStatement update) { throw SpannerExceptionFactory.newSpannerException( ErrorCode.FAILED_PRECONDITION, "Update statements are not allowed for read-only transactions"); } @Override - public long[] executeBatchUpdate(Iterable updates) { + public ApiFuture executeBatchUpdateAsync(Iterable updates) { throw SpannerExceptionFactory.newSpannerException( ErrorCode.FAILED_PRECONDITION, "Batch updates are not allowed for read-only transactions."); } @Override - public void write(Mutation mutation) { - throw SpannerExceptionFactory.newSpannerException( - ErrorCode.FAILED_PRECONDITION, "Mutations are not allowed for read-only transactions"); - } - - @Override - public void write(Iterable mutations) { + public ApiFuture writeAsync(Iterable mutations) { throw SpannerExceptionFactory.newSpannerException( ErrorCode.FAILED_PRECONDITION, "Mutations are not allowed for read-only transactions"); } @Override - public void commit() { + public ApiFuture commitAsync() { if (this.transaction != null) { this.transaction.close(); } this.state = UnitOfWorkState.COMMITTED; + return ApiFutures.immediateFuture(null); } @Override - public void rollback() { + public ApiFuture rollbackAsync() { if (this.transaction != null) { this.transaction.close(); } this.state = UnitOfWorkState.ROLLED_BACK; + return ApiFutures.immediateFuture(null); } } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ReadWriteTransaction.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ReadWriteTransaction.java index 7a0155cbfb..0a8e322e79 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ReadWriteTransaction.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ReadWriteTransaction.java @@ -16,8 +16,13 @@ package com.google.cloud.spanner.connection; +import static com.google.cloud.spanner.SpannerApiFutures.get; import static com.google.common.base.Preconditions.checkNotNull; +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.AbortedDueToConcurrentModificationException; import com.google.cloud.spanner.AbortedException; @@ -35,6 +40,10 @@ import com.google.cloud.spanner.connection.TransactionRetryListener.RetryResult; import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableList; +import com.google.common.util.concurrent.MoreExecutors; +import com.google.spanner.v1.SpannerGrpc; +import io.grpc.MethodDescriptor; import java.util.ArrayList; import java.util.LinkedList; import java.util.List; @@ -65,12 +74,15 @@ class ReadWriteTransaction extends AbstractMultiUseTransaction { private int transactionRetryAttempts; private int successfulRetries; private final List transactionRetryListeners; - private volatile TransactionContext txContext; + private volatile ApiFuture txContextFuture; + private volatile SettableApiFuture commitTimestampFuture; private volatile UnitOfWorkState state = UnitOfWorkState.STARTED; + private volatile AbortedException abortedException; private boolean timedOutOrCancelled = false; private final List statements = new ArrayList<>(); private final List mutations = new ArrayList<>(); private Timestamp transactionStarted; + final Object abortedLock = new Object(); static class Builder extends AbstractMultiUseTransaction.Builder { private DatabaseClient dbClient; @@ -154,36 +166,80 @@ public boolean isReadOnly() { return false; } + private static final ParsedStatement BEGIN_STATEMENT = + StatementParser.INSTANCE.parse(Statement.of("BEGIN")); + @Override void checkValidTransaction() { + checkValidState(); + if (txContextFuture == null) { + transactionStarted = Timestamp.now(); + txContextFuture = + executeStatementAsync( + BEGIN_STATEMENT, + new Callable() { + @Override + public TransactionContext call() throws Exception { + return txManager.begin(); + } + }, + SpannerGrpc.getBeginTransactionMethod()); + } + } + + private void checkValidState() { ConnectionPreconditions.checkState( - state == UnitOfWorkState.STARTED, + this.state == UnitOfWorkState.STARTED || this.state == UnitOfWorkState.ABORTED, "This transaction has status " - + state.name() + + this.state.name() + ", only " + UnitOfWorkState.STARTED + + "or " + + UnitOfWorkState.ABORTED + " is allowed."); ConnectionPreconditions.checkState( !timedOutOrCancelled, "The last statement of this transaction timed out or was cancelled. " + "The transaction is no longer usable. " + "Rollback the transaction and start a new one."); - if (txManager.getState() == null) { - transactionStarted = Timestamp.now(); - txContext = txManager.begin(); - } - if (txManager.getState() - != com.google.cloud.spanner.TransactionManager.TransactionState.STARTED) { - throw SpannerExceptionFactory.newSpannerException( - ErrorCode.FAILED_PRECONDITION, - String.format("Invalid transaction state: %s", txManager.getState())); + } + + @Override + public boolean isActive() { + // Consider ABORTED an active state, as it is something that is automatically set if the + // transaction is aborted by the backend. That means that we should not automatically create a + // new transaction for the following statement after a transaction has aborted, and instead we + // should wait until the application has rolled back the current transaction. + // + // Othwerwise the following list of statements could show unexpected behavior: + + // connection.executeUpdateAsync("UPDATE FOO SET BAR=1 ..."); + // connection.executeUpdateAsync("UPDATE BAR SET FOO=2 ..."); + // connection.commitAsync(); + // + // If the first update statement fails with an aborted exception, the second update statement + // should not be executed in a new transaction, but should also abort. + return getState().isActive() || state == UnitOfWorkState.ABORTED; + } + + void checkAborted() { + if (this.state == UnitOfWorkState.ABORTED && this.abortedException != null) { + if (this.abortedException instanceof AbortedDueToConcurrentModificationException) { + throw SpannerExceptionFactory.newAbortedDueToConcurrentModificationException( + (AbortedDueToConcurrentModificationException) this.abortedException); + } else { + throw SpannerExceptionFactory.newSpannerException( + ErrorCode.ABORTED, + "This transaction has already been aborted. Rollback this transaction to start a new one.", + this.abortedException); + } } } @Override TransactionContext getReadContext() { - ConnectionPreconditions.checkState(txContext != null, "Missing transaction context"); - return txContext; + ConnectionPreconditions.checkState(txContextFuture != null, "Missing transaction context"); + return get(txContextFuture); } @Override @@ -199,23 +255,22 @@ public Timestamp getReadTimestampOrNull() { } private boolean hasCommitTimestamp() { - return txManager.getState() - == com.google.cloud.spanner.TransactionManager.TransactionState.COMMITTED; + return commitTimestampFuture != null; } @Override public Timestamp getCommitTimestamp() { ConnectionPreconditions.checkState(hasCommitTimestamp(), "This transaction has not committed."); - return txManager.getCommitTimestamp(); + return get(commitTimestampFuture); } @Override public Timestamp getCommitTimestampOrNull() { - return hasCommitTimestamp() ? txManager.getCommitTimestamp() : null; + return hasCommitTimestamp() ? get(commitTimestampFuture) : null; } @Override - public void executeDdl(ParsedStatement ddl) { + public ApiFuture executeDdlAsync(ParsedStatement ddl) { throw SpannerExceptionFactory.newSpannerException( ErrorCode.FAILED_PRECONDITION, "DDL-statements are not allowed inside a read/write transaction."); @@ -229,108 +284,138 @@ private void handlePossibleInvalidatingException(SpannerException e) { } @Override - public ResultSet executeQuery( + public ApiFuture executeQueryAsync( final ParsedStatement statement, final AnalyzeMode analyzeMode, final QueryOption... options) { Preconditions.checkArgument(statement.isQuery(), "Statement is not a query"); checkValidTransaction(); - try { - if (retryAbortsInternally) { - return asyncExecuteStatement( - statement, - new Callable() { - @Override - public ResultSet call() throws Exception { - return runWithRetry( - new Callable() { - @Override - public ResultSet call() throws Exception { - try { - getStatementExecutor() - .invokeInterceptors( - statement, - StatementExecutionStep.EXECUTE_STATEMENT, - ReadWriteTransaction.this); - ResultSet delegate = - DirectExecuteResultSet.ofResultSet( - internalExecuteQuery(statement, analyzeMode, options)); - return createAndAddRetryResultSet( - delegate, statement, analyzeMode, options); - } catch (AbortedException e) { - throw e; - } catch (SpannerException e) { - createAndAddFailedQuery(e, statement, analyzeMode, options); - throw e; + + ApiFuture res; + if (retryAbortsInternally) { + res = + executeStatementAsync( + statement, + new Callable() { + @Override + public ResultSet call() throws Exception { + return runWithRetry( + new Callable() { + @Override + public ResultSet call() throws Exception { + try { + getStatementExecutor() + .invokeInterceptors( + statement, + StatementExecutionStep.EXECUTE_STATEMENT, + ReadWriteTransaction.this); + ResultSet delegate = + DirectExecuteResultSet.ofResultSet( + internalExecuteQuery(statement, analyzeMode, options)); + return createAndAddRetryResultSet( + delegate, statement, analyzeMode, options); + } catch (AbortedException e) { + throw e; + } catch (SpannerException e) { + createAndAddFailedQuery(e, statement, analyzeMode, options); + throw e; + } } - } - }); - } - }, - InterceptorsUsage - .IGNORE_INTERCEPTORS); // ignore interceptors here as they are invoked in the - // Callable. - } else { - return super.executeQuery(statement, analyzeMode, options); - } - } catch (SpannerException e) { - handlePossibleInvalidatingException(e); - throw e; + }); + } + }, + // ignore interceptors here as they are invoked in the Callable. + InterceptorsUsage.IGNORE_INTERCEPTORS, + ImmutableList.>of(SpannerGrpc.getExecuteStreamingSqlMethod())); + } else { + res = super.executeQueryAsync(statement, analyzeMode, options); } + + ApiFutures.addCallback( + res, + new ApiFutureCallback() { + @Override + public void onFailure(Throwable t) { + if (t instanceof SpannerException) { + handlePossibleInvalidatingException((SpannerException) t); + } + } + + @Override + public void onSuccess(ResultSet result) {} + }, + MoreExecutors.directExecutor()); + return res; } @Override - public long executeUpdate(final ParsedStatement update) { + public ApiFuture executeUpdateAsync(final ParsedStatement update) { Preconditions.checkNotNull(update); Preconditions.checkArgument(update.isUpdate(), "The statement is not an update statement"); checkValidTransaction(); - try { - if (retryAbortsInternally) { - return asyncExecuteStatement( - update, - new Callable() { - @Override - public Long call() throws Exception { - return runWithRetry( - new Callable() { - @Override - public Long call() throws Exception { - try { - getStatementExecutor() - .invokeInterceptors( - update, - StatementExecutionStep.EXECUTE_STATEMENT, - ReadWriteTransaction.this); - long updateCount = txContext.executeUpdate(update.getStatement()); - createAndAddRetriableUpdate(update, updateCount); - return updateCount; - } catch (AbortedException e) { - throw e; - } catch (SpannerException e) { - createAndAddFailedUpdate(e, update); - throw e; + ApiFuture res; + if (retryAbortsInternally) { + res = + executeStatementAsync( + update, + new Callable() { + @Override + public Long call() throws Exception { + return runWithRetry( + new Callable() { + @Override + public Long call() throws Exception { + try { + getStatementExecutor() + .invokeInterceptors( + update, + StatementExecutionStep.EXECUTE_STATEMENT, + ReadWriteTransaction.this); + long updateCount = + get(txContextFuture).executeUpdate(update.getStatement()); + createAndAddRetriableUpdate(update, updateCount); + return updateCount; + } catch (AbortedException e) { + throw e; + } catch (SpannerException e) { + createAndAddFailedUpdate(e, update); + throw e; + } } - } - }); - } - }, - InterceptorsUsage - .IGNORE_INTERCEPTORS); // ignore interceptors here as they are invoked in the - // Callable. - } else { - return asyncExecuteStatement( - update, - new Callable() { - @Override - public Long call() throws Exception { - return txContext.executeUpdate(update.getStatement()); - } - }); - } - } catch (SpannerException e) { - handlePossibleInvalidatingException(e); - throw e; + }); + } + }, + // ignore interceptors here as they are invoked in the Callable. + InterceptorsUsage.IGNORE_INTERCEPTORS, + ImmutableList.>of(SpannerGrpc.getExecuteSqlMethod())); + } else { + res = + executeStatementAsync( + update, + new Callable() { + @Override + public Long call() throws Exception { + checkAborted(); + return get(txContextFuture).executeUpdate(update.getStatement()); + } + }, + SpannerGrpc.getExecuteSqlMethod()); } + ApiFutures.addCallback( + res, + new ApiFutureCallback() { + @Override + public void onFailure(Throwable t) { + if (t instanceof SpannerException) { + handlePossibleInvalidatingException((SpannerException) t); + } + } + + @Override + public void onSuccess(Long result) {} + }, + MoreExecutors.directExecutor()); + return res; } /** @@ -348,7 +433,7 @@ public Long call() throws Exception { StatementParser.INSTANCE.parse(Statement.of("RUN BATCH")); @Override - public long[] executeBatchUpdate(final Iterable updates) { + public ApiFuture executeBatchUpdateAsync(Iterable updates) { Preconditions.checkNotNull(updates); final List updateStatements = new LinkedList<>(); for (ParsedStatement update : updates) { @@ -358,69 +443,81 @@ public long[] executeBatchUpdate(final Iterable updates) { updateStatements.add(update.getStatement()); } checkValidTransaction(); - try { - if (retryAbortsInternally) { - return asyncExecuteStatement( - EXECUTE_BATCH_UPDATE_STATEMENT, - new Callable() { - @Override - public long[] call() throws Exception { - return runWithRetry( - new Callable() { - @Override - public long[] call() throws Exception { - try { - getStatementExecutor() - .invokeInterceptors( - EXECUTE_BATCH_UPDATE_STATEMENT, - StatementExecutionStep.EXECUTE_STATEMENT, - ReadWriteTransaction.this); - long[] updateCounts = txContext.batchUpdate(updateStatements); - createAndAddRetriableBatchUpdate(updateStatements, updateCounts); - return updateCounts; - } catch (AbortedException e) { - throw e; - } catch (SpannerException e) { - createAndAddFailedBatchUpdate(e, updateStatements); - throw e; + + ApiFuture res; + if (retryAbortsInternally) { + res = + executeStatementAsync( + EXECUTE_BATCH_UPDATE_STATEMENT, + new Callable() { + @Override + public long[] call() throws Exception { + return runWithRetry( + new Callable() { + @Override + public long[] call() throws Exception { + try { + getStatementExecutor() + .invokeInterceptors( + EXECUTE_BATCH_UPDATE_STATEMENT, + StatementExecutionStep.EXECUTE_STATEMENT, + ReadWriteTransaction.this); + long[] updateCounts = + get(txContextFuture).batchUpdate(updateStatements); + createAndAddRetriableBatchUpdate(updateStatements, updateCounts); + return updateCounts; + } catch (AbortedException e) { + throw e; + } catch (SpannerException e) { + createAndAddFailedBatchUpdate(e, updateStatements); + throw e; + } } - } - }); - } - }, - InterceptorsUsage - .IGNORE_INTERCEPTORS); // ignore interceptors here as they are invoked in the - // Callable. - } else { - return asyncExecuteStatement( - EXECUTE_BATCH_UPDATE_STATEMENT, - new Callable() { - @Override - public long[] call() throws Exception { - return txContext.batchUpdate(updateStatements); - } - }); - } - } catch (SpannerException e) { - handlePossibleInvalidatingException(e); - throw e; + }); + } + }, + // ignore interceptors here as they are invoked in the Callable. + InterceptorsUsage.IGNORE_INTERCEPTORS, + ImmutableList.>of(SpannerGrpc.getExecuteBatchDmlMethod())); + } else { + res = + executeStatementAsync( + EXECUTE_BATCH_UPDATE_STATEMENT, + new Callable() { + @Override + public long[] call() throws Exception { + checkAborted(); + return get(txContextFuture).batchUpdate(updateStatements); + } + }, + SpannerGrpc.getExecuteBatchDmlMethod()); } - } - @Override - public void write(Mutation mutation) { - Preconditions.checkNotNull(mutation); - checkValidTransaction(); - mutations.add(mutation); + ApiFutures.addCallback( + res, + new ApiFutureCallback() { + @Override + public void onFailure(Throwable t) { + if (t instanceof SpannerException) { + handlePossibleInvalidatingException((SpannerException) t); + } + } + + @Override + public void onSuccess(long[] result) {} + }, + MoreExecutors.directExecutor()); + return res; } @Override - public void write(Iterable mutations) { + public ApiFuture writeAsync(Iterable mutations) { Preconditions.checkNotNull(mutations); checkValidTransaction(); for (Mutation mutation : mutations) { this.mutations.add(checkNotNull(mutation)); } + return ApiFutures.immediateFuture(null); } /** @@ -440,51 +537,79 @@ public void write(Iterable mutations) { new Callable() { @Override public Void call() throws Exception { - txContext.buffer(mutations); + checkAborted(); + get(txContextFuture).buffer(mutations); txManager.commit(); + commitTimestampFuture.set(txManager.getCommitTimestamp()); + state = UnitOfWorkState.COMMITTED; return null; } }; @Override - public void commit() { + public ApiFuture commitAsync() { checkValidTransaction(); - try { - if (retryAbortsInternally) { - asyncExecuteStatement( - COMMIT_STATEMENT, - new Callable() { - @Override - public Void call() throws Exception { - return runWithRetry( - new Callable() { - @Override - public Void call() throws Exception { - getStatementExecutor() - .invokeInterceptors( - COMMIT_STATEMENT, - StatementExecutionStep.EXECUTE_STATEMENT, - ReadWriteTransaction.this); - commitCallable.call(); - return null; - } - }); - } - }, - InterceptorsUsage.IGNORE_INTERCEPTORS); - } else { - asyncExecuteStatement(COMMIT_STATEMENT, commitCallable); - } - ReadWriteTransaction.this.state = UnitOfWorkState.COMMITTED; - } catch (SpannerException e) { - try { - txManager.close(); - } catch (Throwable t) { - // ignore - } - this.state = UnitOfWorkState.COMMIT_FAILED; - throw e; + state = UnitOfWorkState.COMMITTING; + commitTimestampFuture = SettableApiFuture.create(); + ApiFuture res; + if (retryAbortsInternally) { + res = + executeStatementAsync( + COMMIT_STATEMENT, + new Callable() { + @Override + public Void call() throws Exception { + try { + return runWithRetry( + new Callable() { + @Override + public Void call() throws Exception { + getStatementExecutor() + .invokeInterceptors( + COMMIT_STATEMENT, + StatementExecutionStep.EXECUTE_STATEMENT, + ReadWriteTransaction.this); + return commitCallable.call(); + } + }); + } catch (Throwable t) { + commitTimestampFuture.setException(t); + state = UnitOfWorkState.COMMIT_FAILED; + try { + txManager.close(); + } catch (Throwable t2) { + // Ignore. + } + throw t; + } + } + }, + InterceptorsUsage.IGNORE_INTERCEPTORS, + ImmutableList.>of(SpannerGrpc.getCommitMethod())); + } else { + res = + executeStatementAsync( + COMMIT_STATEMENT, + new Callable() { + @Override + public Void call() throws Exception { + try { + return commitCallable.call(); + } catch (Throwable t) { + commitTimestampFuture.setException(t); + state = UnitOfWorkState.COMMIT_FAILED; + try { + txManager.close(); + } catch (Throwable t2) { + // Ignore. + } + throw t; + } + } + }, + SpannerGrpc.getCommitMethod()); } + return res; } /** @@ -508,18 +633,17 @@ public Void call() throws Exception { */ T runWithRetry(Callable callable) throws SpannerException { while (true) { - try { - return callable.call(); - } catch (final AbortedException aborted) { - if (retryAbortsInternally) { + synchronized (abortedLock) { + checkAborted(); + try { + return callable.call(); + } catch (final AbortedException aborted) { handleAborted(aborted); - } else { - throw aborted; + } catch (SpannerException e) { + throw e; + } catch (Exception e) { + throw SpannerExceptionFactory.asSpannerException(e); } - } catch (SpannerException e) { - throw e; - } catch (Exception e) { - throw SpannerExceptionFactory.newSpannerException(ErrorCode.UNKNOWN, e.getMessage(), e); } } } @@ -609,7 +733,7 @@ private void handleAborted(AbortedException aborted) { ErrorCode.CANCELLED, "The statement was cancelled"); } try { - txContext = txManager.resetForRetry(); + txContextFuture = ApiFutures.immediateFuture(txManager.resetForRetry()); // Inform listeners about the transaction retry that is about to start. invokeTransactionRetryListenersOnStart(); // Then retry all transaction statements. @@ -630,13 +754,14 @@ private void handleAborted(AbortedException aborted) { RetryResult.RETRY_ABORTED_DUE_TO_CONCURRENT_MODIFICATION); logger.fine( toString() + ": Internal transaction retry aborted due to a concurrent modification"); - // Try to rollback the new transaction and ignore any exceptions. + // Do a shoot and forget rollback. try { txManager.rollback(); } catch (Throwable t) { // ignore } this.state = UnitOfWorkState.ABORTED; + this.abortedException = e; throw e; } catch (AbortedException e) { // Retry aborted, do another retry of the transaction. @@ -651,7 +776,7 @@ private void handleAborted(AbortedException aborted) { Level.FINE, toString() + ": Internal transaction retry failed due to an unexpected exception", e); - // Try to rollback the new transaction and ignore any exceptions. + // Do a shoot and forget rollback. try { txManager.rollback(); } catch (Throwable t) { @@ -659,6 +784,7 @@ private void handleAborted(AbortedException aborted) { } // Set transaction state to aborted as the retry failed. this.state = UnitOfWorkState.ABORTED; + this.abortedException = aborted; // Re-throw underlying exception. throw e; } @@ -671,6 +797,7 @@ private void handleAborted(AbortedException aborted) { } // Internal retry is not enabled. this.state = UnitOfWorkState.ABORTED; + this.abortedException = aborted; throw aborted; } } @@ -689,8 +816,11 @@ private void throwAbortWithRetryAttemptsExceeded() throws SpannerException { // ignore } this.state = UnitOfWorkState.ABORTED; - throw SpannerExceptionFactory.newSpannerException( - ErrorCode.ABORTED, MAX_INTERNAL_RETRIES_EXCEEDED); + this.abortedException = + (AbortedException) + SpannerExceptionFactory.newSpannerException( + ErrorCode.ABORTED, MAX_INTERNAL_RETRIES_EXCEEDED); + throw this.abortedException; } private void invokeTransactionRetryListenersOnStart() { @@ -713,26 +843,30 @@ private void invokeTransactionRetryListenersOnFinish(RetryResult result) { new Callable() { @Override public Void call() throws Exception { - txManager.rollback(); - return null; + try { + if (state != UnitOfWorkState.ABORTED) { + // Make sure the transaction has actually started before we try to rollback. + get(txContextFuture); + txManager.rollback(); + } + return null; + } finally { + txManager.close(); + } } }; @Override - public void rollback() { + public ApiFuture rollbackAsync() { ConnectionPreconditions.checkState( - state == UnitOfWorkState.STARTED, "This transaction has status " + state.name()); - try { - asyncExecuteStatement(rollbackStatement, rollbackCallable); - } finally { - // Whatever happens, we should always call close in order to return the underlying session to - // the session pool to avoid any session leaks. - try { - txManager.close(); - } catch (Throwable e) { - // ignore - } - this.state = UnitOfWorkState.ROLLED_BACK; + state == UnitOfWorkState.STARTED || state == UnitOfWorkState.ABORTED, + "This transaction has status " + state.name()); + state = UnitOfWorkState.ROLLED_BACK; + if (txContextFuture != null && state != UnitOfWorkState.ABORTED) { + return executeStatementAsync( + rollbackStatement, rollbackCallable, SpannerGrpc.getRollbackMethod()); + } else { + return ApiFutures.immediateFuture(null); } } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/SingleUseTransaction.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/SingleUseTransaction.java index 614d0c61e5..52011eb910 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/SingleUseTransaction.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/SingleUseTransaction.java @@ -16,15 +16,17 @@ package com.google.cloud.spanner.connection; +import com.google.api.core.ApiFuture; +import com.google.api.core.SettableApiFuture; import com.google.api.gax.longrunning.OperationFuture; import com.google.cloud.Timestamp; -import com.google.cloud.spanner.AbortedException; import com.google.cloud.spanner.DatabaseClient; import com.google.cloud.spanner.ErrorCode; import com.google.cloud.spanner.Mutation; import com.google.cloud.spanner.Options.QueryOption; import com.google.cloud.spanner.ReadOnlyTransaction; import com.google.cloud.spanner.ResultSet; +import com.google.cloud.spanner.SpannerApiFutures; import com.google.cloud.spanner.SpannerBatchUpdateException; import com.google.cloud.spanner.SpannerException; import com.google.cloud.spanner.SpannerExceptionFactory; @@ -36,13 +38,15 @@ import com.google.cloud.spanner.TransactionRunner.TransactionCallable; import com.google.cloud.spanner.connection.StatementParser.ParsedStatement; import com.google.cloud.spanner.connection.StatementParser.StatementType; +import com.google.common.base.Function; import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.Iterables; +import com.google.spanner.admin.database.v1.DatabaseAdminGrpc; import com.google.spanner.admin.database.v1.UpdateDatabaseDdlMetadata; -import java.util.Arrays; -import java.util.LinkedList; -import java.util.List; +import com.google.spanner.v1.SpannerGrpc; +import io.grpc.MethodDescriptor; import java.util.concurrent.Callable; -import java.util.concurrent.TimeUnit; /** * Transaction that is used when a {@link Connection} is in autocommit mode. Each method on this @@ -66,11 +70,11 @@ class SingleUseTransaction extends AbstractBaseUnitOfWork { private final DatabaseClient dbClient; private final TimestampBound readOnlyStaleness; private final AutocommitDmlMode autocommitDmlMode; - private Timestamp readTimestamp = null; + private volatile SettableApiFuture readTimestamp = null; private volatile TransactionManager txManager; - private TransactionRunner writeTransaction; + private volatile TransactionRunner writeTransaction; private boolean used = false; - private UnitOfWorkState state = UnitOfWorkState.STARTED; + private volatile UnitOfWorkState state = UnitOfWorkState.STARTED; static class Builder extends AbstractBaseUnitOfWork.Builder { private DdlClient ddlClient; @@ -160,7 +164,7 @@ private void checkAndMarkUsed() { } @Override - public ResultSet executeQuery( + public ApiFuture executeQueryAsync( final ParsedStatement statement, final AnalyzeMode analyzeMode, final QueryOption... options) { @@ -185,42 +189,43 @@ public ResultSet call() throws Exception { } // Return a DirectExecuteResultSet, which will directly do a next() call in order to // ensure that the query is actually sent to Spanner. - return DirectExecuteResultSet.ofResultSet(rs); - } finally { + ResultSet directRs = DirectExecuteResultSet.ofResultSet(rs); + state = UnitOfWorkState.COMMITTED; + readTimestamp.set(currentTransaction.getReadTimestamp()); + return directRs; + } catch (Throwable t) { + state = UnitOfWorkState.COMMIT_FAILED; + readTimestamp.set(null); currentTransaction.close(); + throw t; } } }; - try { - ResultSet res = asyncExecuteStatement(statement, callable); - readTimestamp = currentTransaction.getReadTimestamp(); - state = UnitOfWorkState.COMMITTED; - return res; - } catch (Throwable e) { - state = UnitOfWorkState.COMMIT_FAILED; - throw e; - } finally { - currentTransaction.close(); - } + readTimestamp = SettableApiFuture.create(); + ApiFuture res = + executeStatementAsync(statement, callable, SpannerGrpc.getExecuteStreamingSqlMethod()); + return res; } @Override public Timestamp getReadTimestamp() { ConnectionPreconditions.checkState( - readTimestamp != null, "There is no read timestamp available for this transaction."); - return readTimestamp; + SpannerApiFutures.getOrNull(readTimestamp) != null, + "There is no read timestamp available for this transaction."); + return SpannerApiFutures.get(readTimestamp); } @Override public Timestamp getReadTimestampOrNull() { - return readTimestamp; + return SpannerApiFutures.getOrNull(readTimestamp); } private boolean hasCommitTimestamp() { - return writeTransaction != null - || (txManager != null - && txManager.getState() - == com.google.cloud.spanner.TransactionManager.TransactionState.COMMITTED); + return state == UnitOfWorkState.COMMITTED + && (writeTransaction != null + || (txManager != null + && txManager.getState() + == com.google.cloud.spanner.TransactionManager.TransactionState.COMMITTED)); } @Override @@ -247,7 +252,7 @@ public Timestamp getCommitTimestampOrNull() { } @Override - public void executeDdl(final ParsedStatement ddl) { + public ApiFuture executeDdlAsync(final ParsedStatement ddl) { Preconditions.checkNotNull(ddl); Preconditions.checkArgument( ddl.getType() == StatementType.DDL, "Statement is not a ddl statement"); @@ -255,70 +260,53 @@ public void executeDdl(final ParsedStatement ddl) { !isReadOnly(), "DDL statements are not allowed in read-only mode"); checkAndMarkUsed(); - try { - Callable callable = - new Callable() { - @Override - public Void call() throws Exception { + Callable callable = + new Callable() { + @Override + public Void call() throws Exception { + try { OperationFuture operation = ddlClient.executeDdl(ddl.getSqlWithoutComments()); - return operation.get(); + Void res = getWithStatementTimeout(operation, ddl); + state = UnitOfWorkState.COMMITTED; + return res; + } catch (Throwable t) { + state = UnitOfWorkState.COMMIT_FAILED; + throw t; } - }; - asyncExecuteStatement(ddl, callable); - state = UnitOfWorkState.COMMITTED; - } catch (Throwable e) { - state = UnitOfWorkState.COMMIT_FAILED; - throw e; - } + } + }; + return executeStatementAsync(ddl, callable, DatabaseAdminGrpc.getUpdateDatabaseDdlMethod()); } @Override - public long executeUpdate(final ParsedStatement update) { + public ApiFuture executeUpdateAsync(ParsedStatement update) { Preconditions.checkNotNull(update); Preconditions.checkArgument(update.isUpdate(), "Statement is not an update statement"); ConnectionPreconditions.checkState( !isReadOnly(), "Update statements are not allowed in read-only mode"); checkAndMarkUsed(); - long res; - try { - switch (autocommitDmlMode) { - case TRANSACTIONAL: - res = executeAsyncTransactionalUpdate(update, new TransactionalUpdateCallable(update)); - break; - case PARTITIONED_NON_ATOMIC: - res = executeAsyncPartitionedUpdate(update); - break; - default: - throw SpannerExceptionFactory.newSpannerException( - ErrorCode.FAILED_PRECONDITION, "Unknown dml mode: " + autocommitDmlMode); - } - } catch (Throwable e) { - state = UnitOfWorkState.COMMIT_FAILED; - throw e; + ApiFuture res; + switch (autocommitDmlMode) { + case TRANSACTIONAL: + res = executeTransactionalUpdateAsync(update); + break; + case PARTITIONED_NON_ATOMIC: + res = executePartitionedUpdateAsync(update); + break; + default: + throw SpannerExceptionFactory.newSpannerException( + ErrorCode.FAILED_PRECONDITION, "Unknown dml mode: " + autocommitDmlMode); } - state = UnitOfWorkState.COMMITTED; return res; } - /** Execute an update statement as a partitioned DML statement. */ - private long executeAsyncPartitionedUpdate(final ParsedStatement update) { - Callable callable = - new Callable() { - @Override - public Long call() throws Exception { - return dbClient.executePartitionedUpdate(update.getStatement()); - } - }; - return asyncExecuteStatement(update, callable); - } - private final ParsedStatement executeBatchUpdateStatement = StatementParser.INSTANCE.parse(Statement.of("RUN BATCH")); @Override - public long[] executeBatchUpdate(Iterable updates) { + public ApiFuture executeBatchUpdateAsync(Iterable updates) { Preconditions.checkNotNull(updates); for (ParsedStatement update : updates) { Preconditions.checkArgument( @@ -329,170 +317,157 @@ public long[] executeBatchUpdate(Iterable updates) { !isReadOnly(), "Batch update statements are not allowed in read-only mode"); checkAndMarkUsed(); - long[] res; - try { - switch (autocommitDmlMode) { - case TRANSACTIONAL: - res = - executeAsyncTransactionalUpdate( - executeBatchUpdateStatement, new TransactionalBatchUpdateCallable(updates)); - break; - case PARTITIONED_NON_ATOMIC: - throw SpannerExceptionFactory.newSpannerException( - ErrorCode.FAILED_PRECONDITION, - "Batch updates are not allowed in " + autocommitDmlMode); - default: - throw SpannerExceptionFactory.newSpannerException( - ErrorCode.FAILED_PRECONDITION, "Unknown dml mode: " + autocommitDmlMode); - } - } catch (SpannerBatchUpdateException e) { - // Batch update exceptions does not cause a rollback. - state = UnitOfWorkState.COMMITTED; - throw e; - } catch (Throwable e) { - state = UnitOfWorkState.COMMIT_FAILED; - throw e; + switch (autocommitDmlMode) { + case TRANSACTIONAL: + return executeTransactionalBatchUpdateAsync(updates); + case PARTITIONED_NON_ATOMIC: + throw SpannerExceptionFactory.newSpannerException( + ErrorCode.FAILED_PRECONDITION, "Batch updates are not allowed in " + autocommitDmlMode); + default: + throw SpannerExceptionFactory.newSpannerException( + ErrorCode.FAILED_PRECONDITION, "Unknown dml mode: " + autocommitDmlMode); } - state = UnitOfWorkState.COMMITTED; - return res; } - /** Base class for executing DML updates (both single statements and batches). */ - private abstract class AbstractUpdateCallable implements Callable { - abstract T executeUpdate(TransactionContext txContext); - - @Override - public T call() throws Exception { - try { - txManager = dbClient.transactionManager(); - // Check the interrupted state after each (possible) round-trip to the db to allow the - // statement to be cancelled. - checkInterrupted(); - try (TransactionContext txContext = txManager.begin()) { - checkInterrupted(); - T res = executeUpdate(txContext); - checkInterrupted(); - txManager.commit(); - checkInterrupted(); - return res; - } - } finally { - if (txManager != null) { - // Calling txManager.close() will rollback the transaction if it is still active, i.e. if - // an error occurred before the commit() call returned successfully. - txManager.close(); - } - } - } - } - - /** {@link Callable} for a single update statement. */ - private final class TransactionalUpdateCallable extends AbstractUpdateCallable { - private final ParsedStatement update; - - private TransactionalUpdateCallable(ParsedStatement update) { - this.update = update; - } - - @Override - Long executeUpdate(TransactionContext txContext) { - return txContext.executeUpdate(update.getStatement()); - } - } - - /** {@link Callable} for a batch update. */ - private final class TransactionalBatchUpdateCallable extends AbstractUpdateCallable { - private final List updates; - - private TransactionalBatchUpdateCallable(Iterable updates) { - this.updates = new LinkedList<>(); - for (ParsedStatement update : updates) { - this.updates.add(update.getStatement()); - } - } - - @Override - long[] executeUpdate(TransactionContext txContext) { - return txContext.batchUpdate(updates); - } + private ApiFuture executeTransactionalUpdateAsync(final ParsedStatement update) { + Callable callable = + new Callable() { + @Override + public Long call() throws Exception { + try { + writeTransaction = dbClient.readWriteTransaction(); + Long res = + writeTransaction.run( + new TransactionCallable() { + @Override + public Long run(TransactionContext transaction) throws Exception { + return transaction.executeUpdate(update.getStatement()); + } + }); + state = UnitOfWorkState.COMMITTED; + return res; + } catch (Throwable t) { + state = UnitOfWorkState.COMMIT_FAILED; + throw t; + } + } + }; + return executeStatementAsync( + update, + callable, + ImmutableList.>of( + SpannerGrpc.getExecuteSqlMethod(), SpannerGrpc.getCommitMethod())); } - private T executeAsyncTransactionalUpdate( - final ParsedStatement update, final AbstractUpdateCallable callable) { - long startedTime = System.currentTimeMillis(); - // This method uses a TransactionManager instead of the TransactionRunner in order to be able to - // handle timeouts and canceling of a statement. - while (true) { - try { - return asyncExecuteStatement(update, callable); - } catch (AbortedException e) { - try { - Thread.sleep(e.getRetryDelayInMillis() / 1000); - } catch (InterruptedException e1) { - throw SpannerExceptionFactory.newSpannerException( - ErrorCode.CANCELLED, "Statement execution was interrupted", e1); - } - // Check whether the timeout time has been exceeded. - long executionTime = System.currentTimeMillis() - startedTime; - if (getStatementTimeout().hasTimeout() - && executionTime > getStatementTimeout().getTimeoutValue(TimeUnit.MILLISECONDS)) { - throw SpannerExceptionFactory.newSpannerException( - ErrorCode.DEADLINE_EXCEEDED, - "Statement execution timeout occurred for " + update.getSqlWithoutComments()); - } - } - } + private ApiFuture executePartitionedUpdateAsync(final ParsedStatement update) { + Callable callable = + new Callable() { + @Override + public Long call() throws Exception { + try { + Long res = dbClient.executePartitionedUpdate(update.getStatement()); + state = UnitOfWorkState.COMMITTED; + return res; + } catch (Throwable t) { + state = UnitOfWorkState.COMMIT_FAILED; + throw t; + } + } + }; + return executeStatementAsync(update, callable, SpannerGrpc.getExecuteStreamingSqlMethod()); } - private void checkInterrupted() throws InterruptedException { - if (Thread.currentThread().isInterrupted()) { - throw new InterruptedException(); - } + private ApiFuture executeTransactionalBatchUpdateAsync( + final Iterable updates) { + Callable callable = + new Callable() { + @Override + public long[] call() throws Exception { + writeTransaction = dbClient.readWriteTransaction(); + return writeTransaction.run( + new TransactionCallable() { + @Override + public long[] run(TransactionContext transaction) throws Exception { + try { + long[] res = + transaction.batchUpdate( + Iterables.transform( + updates, + new Function() { + @Override + public Statement apply(ParsedStatement input) { + return input.getStatement(); + } + })); + state = UnitOfWorkState.COMMITTED; + return res; + } catch (Throwable t) { + if (t instanceof SpannerBatchUpdateException) { + // Batch update exceptions does not cause a rollback. + state = UnitOfWorkState.COMMITTED; + } else { + state = UnitOfWorkState.COMMIT_FAILED; + } + throw t; + } + } + }); + } + }; + return executeStatementAsync( + executeBatchUpdateStatement, callable, SpannerGrpc.getExecuteBatchDmlMethod()); } - @Override - public void write(final Mutation mutation) { - write(Arrays.asList(mutation)); - } + private final ParsedStatement commitStatement = + StatementParser.INSTANCE.parse(Statement.of("COMMIT")); @Override - public void write(final Iterable mutations) { + public ApiFuture writeAsync(final Iterable mutations) { Preconditions.checkNotNull(mutations); ConnectionPreconditions.checkState( !isReadOnly(), "Update statements are not allowed in read-only mode"); checkAndMarkUsed(); - writeTransaction = dbClient.readWriteTransaction(); - try { - writeTransaction.run( - new TransactionCallable() { - @Override - public Void run(TransactionContext transaction) throws Exception { - transaction.buffer(mutations); - return null; + Callable callable = + new Callable() { + @Override + public Void call() throws Exception { + try { + writeTransaction = dbClient.readWriteTransaction(); + Void res = + writeTransaction.run( + new TransactionCallable() { + @Override + public Void run(TransactionContext transaction) throws Exception { + transaction.buffer(mutations); + return null; + } + }); + state = UnitOfWorkState.COMMITTED; + return res; + } catch (Throwable t) { + state = UnitOfWorkState.COMMIT_FAILED; + throw t; } - }); - } catch (Throwable e) { - state = UnitOfWorkState.COMMIT_FAILED; - throw e; - } - state = UnitOfWorkState.COMMITTED; + } + }; + return executeStatementAsync(commitStatement, callable, SpannerGrpc.getCommitMethod()); } @Override - public void commit() { + public ApiFuture commitAsync() { throw SpannerExceptionFactory.newSpannerException( ErrorCode.FAILED_PRECONDITION, "Commit is not supported for single-use transactions"); } @Override - public void rollback() { + public ApiFuture rollbackAsync() { throw SpannerExceptionFactory.newSpannerException( ErrorCode.FAILED_PRECONDITION, "Rollback is not supported for single-use transactions"); } @Override - public long[] runBatch() { + public ApiFuture runBatchAsync() { throw SpannerExceptionFactory.newSpannerException( ErrorCode.FAILED_PRECONDITION, "Run batch is not supported for single-use transactions"); } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/SpannerPool.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/SpannerPool.java index ecf13cd399..350cf61394 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/SpannerPool.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/SpannerPool.java @@ -28,6 +28,7 @@ import com.google.common.base.MoreObjects; import com.google.common.base.Preconditions; import com.google.common.base.Predicates; +import com.google.common.base.Ticker; import com.google.common.collect.Iterables; import io.grpc.ManagedChannelBuilder; import java.util.ArrayList; @@ -80,7 +81,7 @@ public static void closeSpannerPool() { private static final long DEFAULT_CLOSE_SPANNER_AFTER_MILLISECONDS_UNUSED = 60000L; static final SpannerPool INSTANCE = - new SpannerPool(DEFAULT_CLOSE_SPANNER_AFTER_MILLISECONDS_UNUSED); + new SpannerPool(DEFAULT_CLOSE_SPANNER_AFTER_MILLISECONDS_UNUSED, Ticker.systemTicker()); @VisibleForTesting enum CheckAndCloseSpannersMode { @@ -236,14 +237,17 @@ public int hashCode() { @GuardedBy("this") private final Map lastConnectionClosedAt = new HashMap<>(); + private final Ticker ticker; + @VisibleForTesting - SpannerPool() { - this(0L); + SpannerPool(Ticker ticker) { + this(0L, ticker); } @VisibleForTesting - SpannerPool(long closeSpannerAfterMillisecondsUnused) { + SpannerPool(long closeSpannerAfterMillisecondsUnused, Ticker ticker) { this.closeSpannerAfterMillisecondsUnused = closeSpannerAfterMillisecondsUnused; + this.ticker = ticker; } /** @@ -333,6 +337,9 @@ public ManagedChannelBuilder apply(ManagedChannelBuilder input) { } }); } + if (options.getConfigurator() != null) { + options.getConfigurator().configure(builder); + } return builder.build().getService(); } @@ -360,7 +367,8 @@ void removeConnection(ConnectionOptions options, ConnectionImpl connection) { if (registeredConnections.isEmpty()) { // Register the moment the last connection for this Spanner key was removed, so we know // which Spanner objects we could close. - lastConnectionClosedAt.put(key, System.currentTimeMillis()); + lastConnectionClosedAt.put( + key, TimeUnit.MILLISECONDS.convert(ticker.read(), TimeUnit.NANOSECONDS)); } } } else { @@ -443,7 +451,8 @@ void closeUnusedSpanners(long closeSpannerAfterMillisecondsUnused) { // Check whether the last connection was closed more than // closeSpannerAfterMillisecondsUnused milliseconds ago. if (closedAt != null - && ((System.currentTimeMillis() - closedAt.longValue())) + && ((TimeUnit.MILLISECONDS.convert(ticker.read(), TimeUnit.NANOSECONDS) + - closedAt.longValue())) > closeSpannerAfterMillisecondsUnused) { Spanner spanner = spanners.get(entry.getKey()); if (spanner != null) { @@ -463,4 +472,11 @@ void closeUnusedSpanners(long closeSpannerAfterMillisecondsUnused) { } } } + + @VisibleForTesting + int getCurrentSpannerCount() { + synchronized (this) { + return spanners.size(); + } + } } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/StatementExecutor.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/StatementExecutor.java index bb1fa28126..baaadbe167 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/StatementExecutor.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/StatementExecutor.java @@ -16,10 +16,13 @@ package com.google.cloud.spanner.connection; +import com.google.api.core.ApiFuture; +import com.google.api.core.ListenableFutureToApiFuture; import com.google.cloud.spanner.connection.ReadOnlyStalenessUtil.DurationValueGetter; import com.google.cloud.spanner.connection.StatementParser.ParsedStatement; import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Preconditions; +import com.google.common.util.concurrent.ListeningExecutorService; import com.google.common.util.concurrent.MoreExecutors; import com.google.common.util.concurrent.ThreadFactoryBuilder; import com.google.protobuf.Duration; @@ -27,11 +30,11 @@ import java.util.List; import java.util.concurrent.Callable; import java.util.concurrent.ExecutorService; -import java.util.concurrent.Future; import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.ThreadFactory; import java.util.concurrent.ThreadPoolExecutor; import java.util.concurrent.TimeUnit; +import org.threeten.bp.temporal.ChronoUnit; /** * {@link StatementExecutor} is responsible for executing statements on a {@link Connection}. @@ -55,23 +58,7 @@ static boolean isValidTimeoutUnit(TimeUnit unit) { } /** The statement timeout. */ - private Duration duration = null; - - /** Creates a {@link StatementTimeout} that will never timeout. */ - @VisibleForTesting - static StatementTimeout nullTimeout() { - return new StatementTimeout(); - } - - /** Creates a {@link StatementTimeout} with the given duration. */ - @VisibleForTesting - static StatementTimeout of(long timeout, TimeUnit unit) { - Preconditions.checkArgument(timeout > 0L); - Preconditions.checkArgument(isValidTimeoutUnit(unit)); - StatementTimeout res = new StatementTimeout(); - res.duration = ReadOnlyStalenessUtil.createDuration(timeout, unit); - return res; - } + private volatile Duration duration = null; /** * Does this {@link StatementTimeout} have an actual timeout (i.e. it will eventually timeout). @@ -115,6 +102,31 @@ public boolean hasDuration() { } }); } + + org.threeten.bp.Duration asDuration() { + if (!hasTimeout()) { + return org.threeten.bp.Duration.ZERO; + } + TimeUnit unit = getAppropriateTimeUnit(); + switch (unit) { + case DAYS: + return org.threeten.bp.Duration.ofDays(getTimeoutValue(unit)); + case HOURS: + return org.threeten.bp.Duration.ofHours(getTimeoutValue(unit)); + case MICROSECONDS: + return org.threeten.bp.Duration.of(getTimeoutValue(unit), ChronoUnit.MICROS); + case MILLISECONDS: + return org.threeten.bp.Duration.ofMillis(getTimeoutValue(unit)); + case MINUTES: + return org.threeten.bp.Duration.ofMinutes(getTimeoutValue(unit)); + case NANOSECONDS: + return org.threeten.bp.Duration.ofNanos(getTimeoutValue(unit)); + case SECONDS: + return org.threeten.bp.Duration.ofSeconds(getTimeoutValue(unit)); + default: + throw new IllegalStateException("invalid time unit: " + unit); + } + } } /** @@ -129,12 +141,13 @@ public boolean hasDuration() { .build(); /** Creates an {@link ExecutorService} for a {@link StatementExecutor}. */ - private static ExecutorService createExecutorService() { - return new ThreadPoolExecutor( - 1, 1, 0L, TimeUnit.MILLISECONDS, new LinkedBlockingQueue(), THREAD_FACTORY); + private static ListeningExecutorService createExecutorService() { + return MoreExecutors.listeningDecorator( + new ThreadPoolExecutor( + 1, 1, 0L, TimeUnit.MILLISECONDS, new LinkedBlockingQueue(), THREAD_FACTORY)); } - private ExecutorService executor = createExecutorService(); + private ListeningExecutorService executor = createExecutorService(); /** * Interceptors that should be invoked before or after a statement is executed can be registered @@ -151,18 +164,6 @@ private static ExecutorService createExecutorService() { this.interceptors = Collections.unmodifiableList(interceptors); } - /** - * Recreates this {@link StatementExecutor} and its {@link ExecutorService}. This can be necessary - * if a statement times out or is cancelled, and it cannot be guaranteed that the statement - * execution can be terminated. In order to prevent the single threaded {@link ExecutorService} to - * continue to block on the timed out/cancelled statement, a new {@link ExecutorService} is - * created. - */ - void recreate() { - executor.shutdown(); - executor = createExecutorService(); - } - /** * Shutdown this executor now and do not wait for any statement that is being executed to finish. */ @@ -171,8 +172,8 @@ List shutdownNow() { } /** Execute a statement on this {@link StatementExecutor}. */ - Future submit(Callable callable) { - return executor.submit(callable); + ApiFuture submit(Callable callable) { + return new ListenableFutureToApiFuture<>(executor.submit(callable)); } /** diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/StatementResultImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/StatementResultImpl.java index 6221cc447b..37e8d7e5a0 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/StatementResultImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/StatementResultImpl.java @@ -16,6 +16,8 @@ package com.google.cloud.spanner.connection; +import static com.google.cloud.spanner.SpannerApiFutures.get; + import com.google.cloud.Timestamp; import com.google.cloud.spanner.ResultSet; import com.google.cloud.spanner.ResultSets; @@ -27,6 +29,27 @@ /** Implementation of {@link StatementResult} */ class StatementResultImpl implements StatementResult { + /** + * Returns the {@link AsyncStatementResult} as a {@link StatementResult} with the guarantee that + * the underlying result is available. + */ + static StatementResult of(AsyncStatementResult delegate) { + switch (delegate.getResultType()) { + case NO_RESULT: + get(delegate.getNoResultAsync()); + break; + case RESULT_SET: + delegate.getResultSet(); + break; + case UPDATE_COUNT: + delegate.getUpdateCount(); + break; + default: + throw new IllegalStateException("Unknown result type: " + delegate.getResultType()); + } + return delegate; + } + /** {@link StatementResult} containing a {@link ResultSet} returned by Cloud Spanner. */ static StatementResult of(ResultSet resultSet) { return new StatementResultImpl(resultSet, null); diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/UnitOfWork.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/UnitOfWork.java index e372229c64..eb3c47d4bf 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/UnitOfWork.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/UnitOfWork.java @@ -16,6 +16,7 @@ package com.google.cloud.spanner.connection; +import com.google.api.core.ApiFuture; import com.google.api.core.InternalApi; import com.google.cloud.Timestamp; import com.google.cloud.spanner.Mutation; @@ -26,6 +27,7 @@ import com.google.cloud.spanner.TransactionContext; import com.google.cloud.spanner.connection.StatementParser.ParsedStatement; import com.google.spanner.v1.ResultSetStats; +import java.util.concurrent.ExecutionException; /** Internal interface for transactions and batches on {@link Connection}s. */ @InternalApi @@ -39,9 +41,11 @@ enum Type { enum UnitOfWorkState { STARTED, + COMMITTING, COMMITTED, COMMIT_FAILED, ROLLED_BACK, + RUNNING, RAN, RUN_FAILED, ABORTED; @@ -67,30 +71,35 @@ public boolean isActive() { * Commits the changes in this unit of work to the database. For read-only transactions, this only * closes the {@link ReadContext}. This method will throw a {@link SpannerException} if called for * a {@link Type#BATCH}. + * + * @return An {@link ApiFuture} that is done when the commit has finished. */ - void commit(); + ApiFuture commitAsync(); /** * Rollbacks any changes in this unit of work. For read-only transactions, this only closes the * {@link ReadContext}. This method will throw a {@link SpannerException} if called for a {@link * Type#BATCH}. + * + * @return An {@link ApiFuture} that is done when the rollback has finished. */ - void rollback(); + ApiFuture rollbackAsync(); /** * Sends the currently buffered statements in this unit of work to the database and ends the * batch. This method will throw a {@link SpannerException} if called for a {@link * Type#TRANSACTION}. * - * @return the update counts in case of a DML batch. Returns an array containing 1 for each - * successful statement and 0 for each failed statement or statement that was not executed DDL - * in case of a DDL batch. + * @return an {@link ApiFuture} containing the update counts in case of a DML batch. Returns an + * array containing 1 for each successful statement and 0 for each failed statement or + * statement that was not executed in case of a DDL batch. */ - long[] runBatch(); + ApiFuture runBatchAsync(); /** * Clears the currently buffered statements in this unit of work and ends the batch. This method - * will throw a {@link SpannerException} if called for a {@link Type#TRANSACTION}. + * will throw a {@link SpannerException} if called for a {@link Type#TRANSACTION}. This method is + * always non-blocking. */ void abortBatch(); @@ -107,11 +116,12 @@ public boolean isActive() { * ResultSet} or not. Cannot be used in combination with {@link QueryOption}s. * @param options the options to configure the query. May only be set if analyzeMode is set to * {@link AnalyzeMode#NONE}. - * @return a {@link ResultSet} with the results of the query. - * @throws SpannerException if the query is not allowed on this {@link UnitOfWork}, or if a - * database error occurs. + * @return an {@link ApiFuture} containing a {@link ResultSet} with the results of the query. + * @throws SpannerException if the query is not allowed on this {@link UnitOfWork}. The {@link + * ApiFuture} will return a {@link SpannerException} wrapped in an {@link ExecutionException} + * if a database error occurs. */ - ResultSet executeQuery( + ApiFuture executeQueryAsync( ParsedStatement statement, AnalyzeMode analyzeMode, QueryOption... options); /** @@ -139,36 +149,28 @@ ResultSet executeQuery( * statement directly on Spanner. * * @param ddl The DDL statement to execute. + * @return an {@link ApiFuture} that is done when the DDL operation has finished. */ - void executeDdl(ParsedStatement ddl); + ApiFuture executeDdlAsync(ParsedStatement ddl); /** * Execute a DML statement on Spanner. * * @param update The DML statement to execute. - * @return the number of records that were inserted/updated/deleted by this statement. + * @return an {@link ApiFuture} containing the number of records that were + * inserted/updated/deleted by this statement. */ - long executeUpdate(ParsedStatement update); + ApiFuture executeUpdateAsync(ParsedStatement update); /** * Execute a batch of DML statements on Spanner. * * @param updates The DML statements to execute. - * @return an array containing the number of records that were inserted/updated/deleted per - * statement. + * @return an {@link ApiFuture} containing an array with the number of records that were + * inserted/updated/deleted per statement. * @see TransactionContext#batchUpdate(Iterable) */ - long[] executeBatchUpdate(Iterable updates); - - /** - * Writes a {@link Mutation} to Spanner. For {@link ReadWriteTransaction}s, this means buffering - * the {@link Mutation} locally and writing the {@link Mutation} to Spanner upon {@link - * UnitOfWork#commit()}. For {@link SingleUseTransaction}s, the {@link Mutation} will be sent - * directly to Spanner. - * - * @param mutation The mutation to write. - */ - void write(Mutation mutation); + ApiFuture executeBatchUpdateAsync(Iterable updates); /** * Writes a batch of {@link Mutation}s to Spanner. For {@link ReadWriteTransaction}s, this means @@ -177,6 +179,8 @@ ResultSet executeQuery( * sent directly to Spanner. * * @param mutations The mutations to write. + * @return an {@link ApiFuture} that is done when the {@link Mutation}s have been successfully + * buffered or written to Cloud Spanner. */ - void write(Iterable mutations); + ApiFuture writeAsync(Iterable mutations); } 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 cae41510fc..4f55cd5ebd 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 @@ -212,6 +212,7 @@ public PartialResultSet next() { int recordCount = 0; while (recordCount < MAX_ROWS_IN_CHUNK && currentRow < resultSet.getRowsCount()) { builder.addAllValues(resultSet.getRows(currentRow).getValuesList()); + builder.setResumeToken(ByteString.copyFromUtf8(String.format("%010d", currentRow))); recordCount++; currentRow++; } @@ -408,6 +409,7 @@ public static class SimulatedExecutionTime { private final int randomExecutionTime; private final Queue exceptions; private final boolean stickyException; + private final Queue streamIndices; /** * Creates a simulated execution time that will always be somewhere between @@ -430,11 +432,18 @@ public static SimulatedExecutionTime none() { } public static SimulatedExecutionTime ofException(Exception exception) { - return new SimulatedExecutionTime(0, 0, Arrays.asList(exception), false); + return new SimulatedExecutionTime( + 0, 0, Arrays.asList(exception), false, Collections.emptySet()); } public static SimulatedExecutionTime ofStickyException(Exception exception) { - return new SimulatedExecutionTime(0, 0, Arrays.asList(exception), true); + return new SimulatedExecutionTime( + 0, 0, Arrays.asList(exception), true, Collections.emptySet()); + } + + public static SimulatedExecutionTime ofStreamException(Exception exception, long streamIndex) { + return new SimulatedExecutionTime( + 0, 0, Arrays.asList(exception), false, Collections.singleton(streamIndex)); } public static SimulatedExecutionTime stickyDatabaseNotFoundException(String name) { @@ -443,27 +452,37 @@ public static SimulatedExecutionTime stickyDatabaseNotFoundException(String name } public static SimulatedExecutionTime ofExceptions(Collection exceptions) { - return new SimulatedExecutionTime(0, 0, exceptions, false); + return new SimulatedExecutionTime(0, 0, exceptions, false, Collections.emptySet()); } public static SimulatedExecutionTime ofMinimumAndRandomTimeAndExceptions( int minimumExecutionTime, int randomExecutionTime, Collection exceptions) { return new SimulatedExecutionTime( - minimumExecutionTime, randomExecutionTime, exceptions, false); + minimumExecutionTime, + randomExecutionTime, + exceptions, + false, + Collections.emptySet()); } private SimulatedExecutionTime(int minimum, int random) { - this(minimum, random, Collections.emptyList(), false); + this( + minimum, random, Collections.emptyList(), false, Collections.emptySet()); } private SimulatedExecutionTime( - int minimum, int random, Collection exceptions, boolean stickyException) { + int minimum, + int random, + Collection exceptions, + boolean stickyException, + Collection streamIndices) { Preconditions.checkArgument(minimum >= 0, "Minimum execution time must be >= 0"); Preconditions.checkArgument(random >= 0, "Random execution time must be >= 0"); this.minimumExecutionTime = minimum; this.randomExecutionTime = random; this.exceptions = new LinkedList<>(exceptions); this.stickyException = stickyException; + this.streamIndices = new LinkedList<>(streamIndices); } void simulateExecutionTime( @@ -472,7 +491,9 @@ void simulateExecutionTime( CountDownLatch freezeLock) { Uninterruptibles.awaitUninterruptibly(freezeLock); checkException(globalExceptions, stickyGlobalExceptions); - checkException(this.exceptions, stickyException); + if (streamIndices.isEmpty()) { + checkException(this.exceptions, stickyException); + } if (minimumExecutionTime > 0 || randomExecutionTime > 0) { Uninterruptibles.sleepUninterruptibly( (randomExecutionTime == 0 ? 0 : RANDOM.nextInt(randomExecutionTime)) @@ -488,6 +509,18 @@ private static void checkException(Queue exceptions, boolean keepExce throw Status.INTERNAL.withDescription(e.getMessage()).withCause(e).asRuntimeException(); } } + + private static void checkStreamException( + long streamIndex, Queue exceptions, Queue streamIndices) { + Exception e = exceptions.peek(); + Long index = streamIndices.peek(); + if (e != null && index != null && index == streamIndex) { + exceptions.poll(); + streamIndices.poll(); + Throwables.throwIfUnchecked(e); + throw Status.INTERNAL.withDescription(e.getMessage()).withCause(e).asRuntimeException(); + } + } } public static final SimulatedExecutionTime NO_EXECUTION_TIME = SimulatedExecutionTime.none(); @@ -1096,7 +1129,11 @@ public void executeStreamingSql( throw res.getException(); case RESULT_SET: returnPartialResultSet( - res.getResultSet(), transactionId, request.getTransaction(), responseObserver); + res.getResultSet(), + transactionId, + request.getTransaction(), + responseObserver, + getExecuteStreamingSqlExecutionTime()); break; case UPDATE_COUNT: if (isPartitioned) { @@ -1425,7 +1462,11 @@ public Iterator iterator() { .asRuntimeException(); } returnPartialResultSet( - res.getResultSet(), transactionId, request.getTransaction(), responseObserver); + res.getResultSet(), + transactionId, + request.getTransaction(), + responseObserver, + getStreamingReadExecutionTime()); } catch (StatusRuntimeException e) { responseObserver.onError(e); } catch (Throwable t) { @@ -1437,7 +1478,8 @@ private void returnPartialResultSet( ResultSet resultSet, ByteString transactionId, TransactionSelector transactionSelector, - StreamObserver responseObserver) { + StreamObserver responseObserver, + SimulatedExecutionTime executionTime) { ResultSetMetadata metadata = resultSet.getMetadata(); if (transactionId == null) { Transaction transaction = getTemporaryTransactionOrNull(transactionSelector); @@ -1451,8 +1493,12 @@ private void returnPartialResultSet( } resultSet = resultSet.toBuilder().setMetadata(metadata).build(); PartialResultSetsIterator iterator = new PartialResultSetsIterator(resultSet); + long index = 0L; while (iterator.hasNext()) { + SimulatedExecutionTime.checkStreamException( + index, executionTime.exceptions, executionTime.streamIndices); responseObserver.onNext(iterator.next()); + index++; } responseObserver.onCompleted(); } @@ -1699,6 +1745,11 @@ public void commit(CommitRequest request, StreamObserver respons .build()); } else if (request.getTransactionId() != null) { transaction = transactions.get(request.getTransactionId()); + Optional aborted = + Optional.fromNullable(abortedTransactions.get(request.getTransactionId())); + if (aborted.or(Boolean.FALSE)) { + throwTransactionAborted(request.getTransactionId()); + } } else { // No transaction mode specified responseObserver.onError( @@ -1864,6 +1915,18 @@ public void waitForLastRequestToBe(Class type, long t } } + public void waitForRequestsToContain(Class type, long timeoutMillis) + throws InterruptedException, TimeoutException { + Stopwatch watch = Stopwatch.createStarted(); + while (countRequestsOfType(type) == 0) { + Thread.sleep(10L); + if (watch.elapsed(TimeUnit.MILLISECONDS) > timeoutMillis) { + throw new TimeoutException( + "Timeout while waiting for requests to contain " + type.getName()); + } + } + } + @Override public void addResponse(AbstractMessage response) { throw new UnsupportedOperationException(); diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ResultSetsTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ResultSetsTest.java index bd3c0c9c52..2d7d695ea2 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ResultSetsTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ResultSetsTest.java @@ -20,15 +20,24 @@ import static org.junit.Assert.assertNotNull; import static org.junit.Assert.fail; +import com.google.api.core.ApiFuture; +import com.google.api.core.ApiFutures; +import com.google.api.gax.core.ExecutorProvider; import com.google.cloud.ByteArray; import com.google.cloud.Date; import com.google.cloud.Timestamp; +import com.google.cloud.spanner.AsyncResultSet.CallbackResponse; +import com.google.cloud.spanner.AsyncResultSet.ReadyCallback; import com.google.common.primitives.Booleans; import com.google.common.primitives.Doubles; import com.google.common.primitives.Longs; +import com.google.common.util.concurrent.MoreExecutors; import java.math.BigDecimal; import java.util.Arrays; import java.util.List; +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.atomic.AtomicInteger; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -350,4 +359,132 @@ public void exceptionIfNextIsNotCalled() { assertNotNull(ex.getMessage()); } } + + @Test + public void testToAsyncResultSet() { + ResultSet delegate = + ResultSets.forRows( + Type.struct(Type.StructField.of("f1", Type.string())), + Arrays.asList(Struct.newBuilder().set("f1").to("x").build())); + + final AtomicInteger count = new AtomicInteger(); + AsyncResultSet rs = ResultSets.toAsyncResultSet(delegate); + ApiFuture fut = + rs.setCallback( + MoreExecutors.directExecutor(), + new ReadyCallback() { + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + while (true) { + switch (resultSet.tryNext()) { + case DONE: + return CallbackResponse.DONE; + case NOT_READY: + return CallbackResponse.CONTINUE; + case OK: + count.incrementAndGet(); + assertThat(resultSet.getString("f1")).isEqualTo("x"); + } + } + } + }); + SpannerApiFutures.get(fut); + assertThat(count.get()).isEqualTo(1); + } + + @Test + public void testToAsyncResultSetWithExecProvider() { + ResultSet delegate = + ResultSets.forRows( + Type.struct(Type.StructField.of("f1", Type.string())), + Arrays.asList(Struct.newBuilder().set("f1").to("x").build())); + + ExecutorProvider provider = + new ExecutorProvider() { + final ScheduledExecutorService executor = Executors.newSingleThreadScheduledExecutor(); + + @Override + public boolean shouldAutoClose() { + return true; + } + + @Override + public ScheduledExecutorService getExecutor() { + return executor; + } + }; + final AtomicInteger count = new AtomicInteger(); + AsyncResultSet rs = ResultSets.toAsyncResultSet(delegate, provider); + ApiFuture fut = + rs.setCallback( + MoreExecutors.directExecutor(), + new ReadyCallback() { + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + while (true) { + switch (resultSet.tryNext()) { + case DONE: + return CallbackResponse.DONE; + case NOT_READY: + return CallbackResponse.CONTINUE; + case OK: + count.incrementAndGet(); + assertThat(resultSet.getString("f1")).isEqualTo("x"); + } + } + } + }); + SpannerApiFutures.get(fut); + assertThat(count.get()).isEqualTo(1); + assertThat(provider.getExecutor().isShutdown()).isTrue(); + } + + @Test + public void testToAsyncResultSetWithFuture() { + ApiFuture delegateFuture = + ApiFutures.immediateFuture( + ResultSets.forRows( + Type.struct(Type.StructField.of("f1", Type.string())), + Arrays.asList(Struct.newBuilder().set("f1").to("x").build()))); + + ExecutorProvider provider = + new ExecutorProvider() { + final ScheduledExecutorService executor = Executors.newSingleThreadScheduledExecutor(); + + @Override + public boolean shouldAutoClose() { + return false; + } + + @Override + public ScheduledExecutorService getExecutor() { + return executor; + } + }; + final AtomicInteger count = new AtomicInteger(); + AsyncResultSet rs = ResultSets.toAsyncResultSet(delegateFuture, provider); + ApiFuture fut = + rs.setCallback( + MoreExecutors.directExecutor(), + new ReadyCallback() { + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + while (true) { + switch (resultSet.tryNext()) { + case DONE: + return CallbackResponse.DONE; + case NOT_READY: + return CallbackResponse.CONTINUE; + case OK: + count.incrementAndGet(); + assertThat(resultSet.getString("f1")).isEqualTo("x"); + } + } + } + }); + SpannerApiFutures.get(fut); + assertThat(count.get()).isEqualTo(1); + assertThat(provider.getExecutor().isShutdown()).isFalse(); + provider.getExecutor().shutdown(); + } } diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerApiFuturesTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerApiFuturesTest.java new file mode 100644 index 0000000000..8b0d03717a --- /dev/null +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerApiFuturesTest.java @@ -0,0 +1,118 @@ +/* + * 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.SpannerApiFutures.get; +import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.fail; + +import com.google.api.core.ApiFuture; +import com.google.api.core.ApiFutures; +import com.google.api.core.ForwardingApiFuture; +import java.util.concurrent.CancellationException; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class SpannerApiFuturesTest { + + @Test + public void testGet() { + ApiFuture fut = ApiFutures.immediateFuture(1L); + assertThat(get(fut)).isEqualTo(1L); + } + + @Test + public void testGetNull() { + try { + get(null); + fail("Missing expected exception"); + } catch (NullPointerException e) { + // Ignore, this is the expected exception. + } + } + + @Test + public void testGetOrNull() { + assertThat(SpannerApiFutures.getOrNull(null)).isNull(); + } + + @Test + public void testGetSpannerException() { + ApiFuture fut = + ApiFutures.immediateFailedFuture( + SpannerExceptionFactory.newSpannerException( + ErrorCode.FAILED_PRECONDITION, "test exception")); + try { + get(fut); + fail("Missing expected exception"); + } catch (SpannerException e) { + assertThat(e.getErrorCode()).isEqualTo(ErrorCode.FAILED_PRECONDITION); + assertThat(e.getMessage()).contains("test exception"); + } + } + + @Test + public void testGetOtherException() { + ApiFuture fut = + ApiFutures.immediateFailedFuture(new RuntimeException("test runtime exception")); + try { + get(fut); + fail("Missing expected exception"); + } catch (SpannerException e) { + assertThat(e.getErrorCode()).isEqualTo(ErrorCode.UNKNOWN); + assertThat(e.getMessage()).contains("test runtime exception"); + } + } + + @Test + public void testGetInterruptedException() { + ApiFuture fut = + new ForwardingApiFuture(ApiFutures.immediateFuture(null)) { + public Void get() throws InterruptedException { + throw new InterruptedException("test interrupted exception"); + } + }; + try { + get(fut); + fail("Missing expected exception"); + } catch (SpannerException e) { + assertThat(e.getErrorCode()).isEqualTo(ErrorCode.CANCELLED); + // The message of an interrupted exception is not included in the SpannerException. + assertThat(e.getMessage()).doesNotContain("test interrupted exception"); + } + } + + @Test + public void testGetCancellationException() { + ApiFuture fut = + new ForwardingApiFuture(ApiFutures.immediateFuture(null)) { + public Void get() throws InterruptedException { + throw new CancellationException("test cancellation exception"); + } + }; + try { + get(fut); + fail("Missing expected exception"); + } catch (SpannerException e) { + assertThat(e.getErrorCode()).isEqualTo(ErrorCode.CANCELLED); + // The message of an cancellation exception is included in the SpannerException. + assertThat(e.getMessage()).contains("test cancellation exception"); + } + } +} diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AbstractMockServerTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AbstractMockServerTest.java index a54a5b848a..b2ebd82661 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AbstractMockServerTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AbstractMockServerTest.java @@ -23,7 +23,13 @@ import com.google.cloud.spanner.admin.instance.v1.MockInstanceAdminImpl; import com.google.cloud.spanner.connection.ITAbstractSpannerTest.AbortInterceptor; import com.google.cloud.spanner.connection.ITAbstractSpannerTest.ITConnection; +import com.google.common.util.concurrent.AbstractFuture; +import com.google.longrunning.GetOperationRequest; +import com.google.longrunning.Operation; +import com.google.longrunning.OperationsGrpc.OperationsImplBase; import com.google.protobuf.AbstractMessage; +import com.google.protobuf.Any; +import com.google.protobuf.Empty; import com.google.protobuf.ListValue; import com.google.protobuf.Value; import com.google.spanner.v1.ExecuteSqlRequest; @@ -33,7 +39,9 @@ import com.google.spanner.v1.Type; import com.google.spanner.v1.TypeCode; import io.grpc.Server; +import io.grpc.internal.LogExceptionRunnable; import io.grpc.netty.shaded.io.grpc.netty.NettyServerBuilder; +import io.grpc.stub.StreamObserver; import java.io.IOException; import java.net.InetSocketAddress; import java.sql.DriverManager; @@ -41,6 +49,7 @@ import java.util.Arrays; import java.util.Collections; import java.util.List; +import java.util.logging.Logger; import org.junit.After; import org.junit.AfterClass; import org.junit.Before; @@ -87,31 +96,58 @@ public abstract class AbstractMockServerTest { .build(); public static final Statement INSERT_STATEMENT = Statement.of("INSERT INTO TEST (ID, NAME) VALUES (1, 'test aborted')"); - public static final int UPDATE_COUNT = 1; + public static final long UPDATE_COUNT = 1L; + + public static final int RANDOM_RESULT_SET_ROW_COUNT = 100; + public static final Statement SELECT_RANDOM_STATEMENT = Statement.of("SELECT * FROM RANDOM"); + public static final com.google.spanner.v1.ResultSet RANDOM_RESULT_SET = + new RandomResultSetGenerator(RANDOM_RESULT_SET_ROW_COUNT).generate(); public static MockSpannerServiceImpl mockSpanner; public static MockInstanceAdminImpl mockInstanceAdmin; public static MockDatabaseAdminImpl mockDatabaseAdmin; + public static OperationsImplBase mockOperations; private static Server server; private static InetSocketAddress address; + private boolean futureParentHandlers; + private boolean exceptionRunnableParentHandlers; + private boolean nettyServerParentHandlers; + @BeforeClass public static void startStaticServer() throws IOException { mockSpanner = new MockSpannerServiceImpl(); mockSpanner.setAbortProbability(0.0D); // We don't want any unpredictable aborted transactions. mockInstanceAdmin = new MockInstanceAdminImpl(); mockDatabaseAdmin = new MockDatabaseAdminImpl(); + mockOperations = + new OperationsImplBase() { + @Override + public void getOperation( + GetOperationRequest request, StreamObserver responseObserver) { + responseObserver.onNext( + Operation.newBuilder() + .setDone(false) + .setName(request.getName()) + .setMetadata(Any.pack(Empty.getDefaultInstance())) + .build()); + responseObserver.onCompleted(); + } + }; address = new InetSocketAddress("localhost", 0); server = NettyServerBuilder.forAddress(address) .addService(mockSpanner) .addService(mockInstanceAdmin) .addService(mockDatabaseAdmin) + .addService(mockOperations) .build() .start(); mockSpanner.putStatementResult( StatementResult.query(SELECT_COUNT_STATEMENT, SELECT_COUNT_RESULTSET_BEFORE_INSERT)); mockSpanner.putStatementResult(StatementResult.update(INSERT_STATEMENT, UPDATE_COUNT)); + mockSpanner.putStatementResult( + StatementResult.query(SELECT_RANDOM_STATEMENT, RANDOM_RESULT_SET)); } @AfterClass @@ -124,11 +160,32 @@ public static void stopServer() throws Exception { @Before public void setupResults() { mockSpanner.reset(); + mockDatabaseAdmin.reset(); + mockInstanceAdmin.reset(); + + futureParentHandlers = Logger.getLogger(AbstractFuture.class.getName()).getUseParentHandlers(); + exceptionRunnableParentHandlers = + Logger.getLogger(LogExceptionRunnable.class.getName()).getUseParentHandlers(); + nettyServerParentHandlers = + Logger.getLogger("io.grpc.netty.shaded.io.grpc.netty.NettyServerHandler") + .getUseParentHandlers(); + Logger.getLogger(AbstractFuture.class.getName()).setUseParentHandlers(false); + Logger.getLogger(LogExceptionRunnable.class.getName()).setUseParentHandlers(false); + Logger.getLogger("io.grpc.netty.shaded.io.grpc.netty.NettyServerHandler") + .setUseParentHandlers(false); } @After public void closeSpannerPool() { - SpannerPool.closeSpannerPool(); + try { + SpannerPool.closeSpannerPool(); + } finally { + Logger.getLogger(AbstractFuture.class.getName()).setUseParentHandlers(futureParentHandlers); + Logger.getLogger(LogExceptionRunnable.class.getName()) + .setUseParentHandlers(exceptionRunnableParentHandlers); + Logger.getLogger("io.grpc.netty.shaded.io.grpc.netty.NettyServerHandler") + .setUseParentHandlers(nettyServerParentHandlers); + } } protected java.sql.Connection createJdbcConnection() throws SQLException { @@ -184,7 +241,7 @@ protected ExecuteSqlRequest getLastExecuteSqlRequest() { throw new IllegalStateException("No ExecuteSqlRequest found in requests"); } - private ITConnection createITConnection(ConnectionOptions options) { + ITConnection createITConnection(ConnectionOptions options) { return new ITConnectionImpl(options); } } diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AsyncStatementResultImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AsyncStatementResultImplTest.java new file mode 100644 index 0000000000..53c3e1a1fc --- /dev/null +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AsyncStatementResultImplTest.java @@ -0,0 +1,99 @@ +/* + * 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.connection; + +import static com.google.cloud.spanner.SpannerApiFutures.get; +import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.fail; +import static org.mockito.Mockito.mock; + +import com.google.api.core.ApiFutures; +import com.google.cloud.spanner.AsyncResultSet; +import com.google.cloud.spanner.ErrorCode; +import com.google.cloud.spanner.SpannerException; +import com.google.cloud.spanner.connection.StatementResult.ResultType; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class AsyncStatementResultImplTest { + + @Test + public void testNoResultGetResultSetAsync() { + AsyncStatementResult subject = + AsyncStatementResultImpl.noResult(ApiFutures.immediateFuture(null)); + assertThat(subject.getResultType()).isEqualTo(ResultType.NO_RESULT); + try { + subject.getResultSetAsync(); + fail("Expected exception"); + } catch (SpannerException ex) { + assertThat(ex.getErrorCode()).isEqualTo(ErrorCode.FAILED_PRECONDITION); + } + } + + @Test + public void testNoResultGetUpdateCountAsync() { + AsyncStatementResult subject = + AsyncStatementResultImpl.noResult(ApiFutures.immediateFuture(null)); + assertThat(subject.getResultType()).isEqualTo(ResultType.NO_RESULT); + try { + subject.getUpdateCountAsync(); + fail("Expected exception"); + } catch (SpannerException ex) { + assertThat(ex.getErrorCode()).isEqualTo(ErrorCode.FAILED_PRECONDITION); + } + } + + @Test + public void testResultSetGetResultSetAsync() { + AsyncStatementResult subject = AsyncStatementResultImpl.of(mock(AsyncResultSet.class)); + assertThat(subject.getResultType()).isEqualTo(ResultType.RESULT_SET); + assertThat(subject.getResultSetAsync()).isNotNull(); + } + + @Test + public void testResultSetGetUpdateCountAsync() { + AsyncStatementResult subject = AsyncStatementResultImpl.of(mock(AsyncResultSet.class)); + assertThat(subject.getResultType()).isEqualTo(ResultType.RESULT_SET); + try { + subject.getUpdateCountAsync(); + fail("Expected exception"); + } catch (SpannerException ex) { + assertThat(ex.getErrorCode()).isEqualTo(ErrorCode.FAILED_PRECONDITION); + } + } + + @Test + public void testUpdateCountGetResultSetAsync() { + AsyncStatementResult subject = AsyncStatementResultImpl.of(ApiFutures.immediateFuture(1L)); + assertThat(subject.getResultType()).isEqualTo(ResultType.UPDATE_COUNT); + try { + subject.getResultSetAsync(); + fail("Expected exception"); + } catch (SpannerException ex) { + assertThat(ex.getErrorCode()).isEqualTo(ErrorCode.FAILED_PRECONDITION); + } + } + + @Test + public void testUpdateCountGetUpdateCountAsync() { + AsyncStatementResult subject = AsyncStatementResultImpl.of(ApiFutures.immediateFuture(1L)); + assertThat(subject.getResultType()).isEqualTo(ResultType.UPDATE_COUNT); + assertThat(get(subject.getUpdateCountAsync())).isEqualTo(1L); + } +} diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionAsyncApiAbortedTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionAsyncApiAbortedTest.java new file mode 100644 index 0000000000..a209bfa312 --- /dev/null +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionAsyncApiAbortedTest.java @@ -0,0 +1,688 @@ +/* + * 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.connection; + +import static com.google.cloud.spanner.SpannerApiFutures.get; +import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.fail; + +import com.google.api.core.ApiFuture; +import com.google.api.core.SettableApiFuture; +import com.google.cloud.Timestamp; +import com.google.cloud.spanner.AbortedDueToConcurrentModificationException; +import com.google.cloud.spanner.AsyncResultSet; +import com.google.cloud.spanner.AsyncResultSet.CallbackResponse; +import com.google.cloud.spanner.AsyncResultSet.ReadyCallback; +import com.google.cloud.spanner.MockSpannerServiceImpl.SimulatedExecutionTime; +import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult; +import com.google.cloud.spanner.Options; +import com.google.cloud.spanner.SpannerExceptionFactory; +import com.google.cloud.spanner.Statement; +import com.google.cloud.spanner.connection.ITAbstractSpannerTest.ITConnection; +import com.google.common.base.Predicate; +import com.google.common.collect.Collections2; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.Lists; +import com.google.common.util.concurrent.MoreExecutors; +import com.google.protobuf.AbstractMessage; +import com.google.spanner.v1.ExecuteSqlRequest; +import io.grpc.Status; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.Executor; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.After; +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.Test; + +/** Tests retry handling of read/write transactions using the Async Connection API. */ +public class ConnectionAsyncApiAbortedTest extends AbstractMockServerTest { + private static final class QueryResult { + final ApiFuture finished; + final AtomicInteger rowCount; + + QueryResult(ApiFuture finished, AtomicInteger rowCount) { + this.finished = finished; + this.rowCount = rowCount; + } + } + + private static final class RetryCounter implements TransactionRetryListener { + final CountDownLatch latch; + int retryCount = 0; + + RetryCounter() { + this(0); + } + + RetryCounter(int countDown) { + latch = new CountDownLatch(countDown); + } + + @Override + public void retryStarting(Timestamp transactionStarted, long transactionId, int retryAttempt) { + retryCount++; + latch.countDown(); + } + + @Override + public void retryFinished( + Timestamp transactionStarted, long transactionId, int retryAttempt, RetryResult result) {} + } + + private static final ExecutorService singleThreadedExecutor = Executors.newSingleThreadExecutor(); + private static final ExecutorService multiThreadedExecutor = Executors.newFixedThreadPool(8); + public static final int RANDOM_RESULT_SET_ROW_COUNT_2 = 50; + public static final Statement SELECT_RANDOM_STATEMENT_2 = Statement.of("SELECT * FROM RANDOM2"); + public static final com.google.spanner.v1.ResultSet RANDOM_RESULT_SET_2 = + new RandomResultSetGenerator(RANDOM_RESULT_SET_ROW_COUNT_2).generate(); + + @BeforeClass + public static void setupAdditionalResults() { + mockSpanner.putStatementResult( + StatementResult.query(SELECT_RANDOM_STATEMENT_2, RANDOM_RESULT_SET_2)); + } + + @AfterClass + public static void stopExecutor() { + singleThreadedExecutor.shutdown(); + multiThreadedExecutor.shutdown(); + } + + @After + public void reset() { + mockSpanner.removeAllExecutionTimes(); + } + + ITConnection createConnection(TransactionRetryListener listener) { + ITConnection connection = + super.createConnection( + ImmutableList.of(), ImmutableList.of(listener)); + connection.setAutocommit(false); + return connection; + } + + @Test + public void testSingleQueryAborted() { + RetryCounter counter = new RetryCounter(); + try (Connection connection = createConnection(counter)) { + assertThat(counter.retryCount).isEqualTo(0); + mockSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.ofException(Status.ABORTED.asRuntimeException())); + QueryResult res = executeQueryAsync(connection, SELECT_RANDOM_STATEMENT); + + assertThat(get(res.finished)).isNull(); + assertThat(res.rowCount.get()).isEqualTo(RANDOM_RESULT_SET_ROW_COUNT); + assertThat(counter.retryCount).isEqualTo(1); + } + } + + @Test + public void testTwoQueriesSecondAborted() { + RetryCounter counter = new RetryCounter(); + try (Connection connection = createConnection(counter)) { + assertThat(counter.retryCount).isEqualTo(0); + QueryResult res1 = executeQueryAsync(connection, SELECT_RANDOM_STATEMENT); + mockSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.ofException(Status.ABORTED.asRuntimeException())); + QueryResult res2 = executeQueryAsync(connection, SELECT_RANDOM_STATEMENT_2); + + assertThat(get(res1.finished)).isNull(); + assertThat(res1.rowCount.get()).isEqualTo(RANDOM_RESULT_SET_ROW_COUNT); + assertThat(get(res2.finished)).isNull(); + assertThat(res2.rowCount.get()).isEqualTo(RANDOM_RESULT_SET_ROW_COUNT_2); + assertThat(counter.retryCount).isEqualTo(1); + } + } + + @Test + public void testTwoQueriesBothAborted() throws InterruptedException { + RetryCounter counter = new RetryCounter(1); + try (Connection connection = createConnection(counter)) { + assertThat(counter.retryCount).isEqualTo(0); + mockSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.ofException(Status.ABORTED.asRuntimeException())); + QueryResult res1 = executeQueryAsync(connection, SELECT_RANDOM_STATEMENT); + // Wait until the first query aborted. + assertThat(counter.latch.await(10L, TimeUnit.SECONDS)).isTrue(); + mockSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.ofException(Status.ABORTED.asRuntimeException())); + QueryResult res2 = executeQueryAsync(connection, SELECT_RANDOM_STATEMENT_2); + + assertThat(get(res1.finished)).isNull(); + assertThat(res1.rowCount.get()).isEqualTo(RANDOM_RESULT_SET_ROW_COUNT); + assertThat(get(res2.finished)).isNull(); + assertThat(res2.rowCount.get()).isEqualTo(RANDOM_RESULT_SET_ROW_COUNT_2); + assertThat(counter.retryCount).isEqualTo(2); + } + } + + @Test + public void testSingleQueryAbortedMidway() { + mockSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.ofStreamException( + Status.ABORTED.asRuntimeException(), RANDOM_RESULT_SET_ROW_COUNT / 2)); + RetryCounter counter = new RetryCounter(); + try (Connection connection = createConnection(counter)) { + assertThat(counter.retryCount).isEqualTo(0); + QueryResult res = executeQueryAsync(connection, SELECT_RANDOM_STATEMENT); + + assertThat(get(res.finished)).isNull(); + assertThat(res.rowCount.get()).isEqualTo(RANDOM_RESULT_SET_ROW_COUNT); + assertThat(counter.retryCount).isEqualTo(1); + } + } + + @Test + public void testTwoQueriesSecondAbortedMidway() { + RetryCounter counter = new RetryCounter(); + try (Connection connection = createConnection(counter)) { + assertThat(counter.retryCount).isEqualTo(0); + QueryResult res1 = executeQueryAsync(connection, SELECT_RANDOM_STATEMENT); + mockSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.ofStreamException( + Status.ABORTED.asRuntimeException(), RANDOM_RESULT_SET_ROW_COUNT_2 / 2)); + QueryResult res2 = executeQueryAsync(connection, SELECT_RANDOM_STATEMENT_2); + + assertThat(get(res1.finished)).isNull(); + assertThat(res1.rowCount.get()).isEqualTo(RANDOM_RESULT_SET_ROW_COUNT); + assertThat(get(res2.finished)).isNull(); + assertThat(res2.rowCount.get()).isEqualTo(RANDOM_RESULT_SET_ROW_COUNT_2); + assertThat(counter.retryCount).isEqualTo(1); + } + } + + @Test + public void testTwoQueriesOneAbortedMidway() { + mockSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.ofStreamException( + Status.ABORTED.asRuntimeException(), + Math.min(RANDOM_RESULT_SET_ROW_COUNT / 2, RANDOM_RESULT_SET_ROW_COUNT_2 / 2))); + RetryCounter counter = new RetryCounter(); + try (Connection connection = createConnection(counter)) { + assertThat(counter.retryCount).isEqualTo(0); + // These AsyncResultSets will be consumed in parallel. One of them will (at random) abort + // halfway. + QueryResult res1 = + executeQueryAsync(connection, SELECT_RANDOM_STATEMENT, multiThreadedExecutor); + QueryResult res2 = + executeQueryAsync(connection, SELECT_RANDOM_STATEMENT_2, multiThreadedExecutor); + + assertThat(get(res1.finished)).isNull(); + assertThat(res1.rowCount.get()).isEqualTo(RANDOM_RESULT_SET_ROW_COUNT); + assertThat(get(res2.finished)).isNull(); + assertThat(res2.rowCount.get()).isEqualTo(RANDOM_RESULT_SET_ROW_COUNT_2); + assertThat(counter.retryCount).isEqualTo(1); + } + } + + @Test + public void testUpdateAndQueryAbortedMidway() throws InterruptedException { + mockSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.ofStreamException( + Status.ABORTED.asRuntimeException(), RANDOM_RESULT_SET_ROW_COUNT / 2)); + final RetryCounter counter = new RetryCounter(); + try (Connection connection = createConnection(counter)) { + assertThat(counter.retryCount).isEqualTo(0); + final SettableApiFuture rowCount = SettableApiFuture.create(); + final CountDownLatch updateLatch = new CountDownLatch(1); + final CountDownLatch queryLatch = new CountDownLatch(1); + ApiFuture finished; + try (AsyncResultSet rs = + connection.executeQueryAsync( + SELECT_RANDOM_STATEMENT, Options.bufferRows(RANDOM_RESULT_SET_ROW_COUNT / 2 - 1))) { + finished = + rs.setCallback( + singleThreadedExecutor, + new ReadyCallback() { + long count; + + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + // Indicate that the query has been executed. + queryLatch.countDown(); + try { + // Wait until the update is on its way. + updateLatch.await(10L, TimeUnit.SECONDS); + while (true) { + switch (resultSet.tryNext()) { + case OK: + count++; + break; + case DONE: + rowCount.set(count); + return CallbackResponse.DONE; + case NOT_READY: + return CallbackResponse.CONTINUE; + } + } + } catch (InterruptedException e) { + throw SpannerExceptionFactory.propagateInterrupt(e); + } + } + }); + } + // Wait until the query has actually executed. + queryLatch.await(10L, TimeUnit.SECONDS); + ApiFuture updateCount = connection.executeUpdateAsync(INSERT_STATEMENT); + updateCount.addListener( + new Runnable() { + @Override + public void run() { + updateLatch.countDown(); + } + }, + MoreExecutors.directExecutor()); + + // We should not commit before the AsyncResultSet has finished. + assertThat(get(finished)).isNull(); + ApiFuture commit = connection.commitAsync(); + + assertThat(get(rowCount)).isEqualTo(RANDOM_RESULT_SET_ROW_COUNT); + assertThat(get(updateCount)).isEqualTo(UPDATE_COUNT); + assertThat(get(commit)).isNull(); + assertThat(counter.retryCount).isEqualTo(1); + + // Verify the order of the statements on the server. + List requests = + Lists.newArrayList( + Collections2.filter( + mockSpanner.getRequests(), + new Predicate() { + @Override + public boolean apply(AbstractMessage input) { + return input instanceof ExecuteSqlRequest; + } + })); + // The entire transaction should be retried. + assertThat(requests).hasSize(4); + assertThat(((ExecuteSqlRequest) requests.get(0)).getSeqno()).isEqualTo(1L); + assertThat(((ExecuteSqlRequest) requests.get(0)).getSql()) + .isEqualTo(SELECT_RANDOM_STATEMENT.getSql()); + assertThat(((ExecuteSqlRequest) requests.get(1)).getSeqno()).isEqualTo(2L); + assertThat(((ExecuteSqlRequest) requests.get(1)).getSql()) + .isEqualTo(INSERT_STATEMENT.getSql()); + assertThat(((ExecuteSqlRequest) requests.get(2)).getSeqno()).isEqualTo(1L); + assertThat(((ExecuteSqlRequest) requests.get(2)).getSql()) + .isEqualTo(SELECT_RANDOM_STATEMENT.getSql()); + assertThat(((ExecuteSqlRequest) requests.get(3)).getSeqno()).isEqualTo(2L); + assertThat(((ExecuteSqlRequest) requests.get(3)).getSql()) + .isEqualTo(INSERT_STATEMENT.getSql()); + } + } + + @Test + public void testUpdateAndQueryAbortedMidway_UpdateCountChanged() throws InterruptedException { + mockSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.ofStreamException( + Status.ABORTED.asRuntimeException(), RANDOM_RESULT_SET_ROW_COUNT / 2)); + final RetryCounter counter = new RetryCounter(); + try (Connection connection = createConnection(counter)) { + assertThat(counter.retryCount).isEqualTo(0); + final CountDownLatch updateLatch = new CountDownLatch(1); + final CountDownLatch queryLatch = new CountDownLatch(1); + ApiFuture finished; + try (AsyncResultSet rs = + connection.executeQueryAsync( + SELECT_RANDOM_STATEMENT, Options.bufferRows(RANDOM_RESULT_SET_ROW_COUNT / 2 - 1))) { + finished = + rs.setCallback( + singleThreadedExecutor, + new ReadyCallback() { + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + // Indicate that the query has been executed. + queryLatch.countDown(); + try { + // Wait until the update is on its way. + updateLatch.await(10L, TimeUnit.SECONDS); + while (true) { + switch (resultSet.tryNext()) { + case OK: + break; + case DONE: + return CallbackResponse.DONE; + case NOT_READY: + return CallbackResponse.CONTINUE; + } + } + } catch (InterruptedException e) { + throw SpannerExceptionFactory.propagateInterrupt(e); + } + } + }); + } + // Wait until the query has actually executed. + queryLatch.await(10L, TimeUnit.SECONDS); + // Execute an update statement and wait until it has finished before allowing the + // AsyncResultSet to continue processing. Also change the result of the update statement after + // it has finished. The AsyncResultSet will see an aborted transaction halfway, and then + // during the retry, it will get a different result for this update statement. That will cause + // the retry to be aborted. + get(connection.executeUpdateAsync(INSERT_STATEMENT)); + try { + mockSpanner.putStatementResult(StatementResult.update(INSERT_STATEMENT, UPDATE_COUNT + 1)); + updateLatch.countDown(); + get(finished); + fail("Missing expected exception"); + } catch (AbortedDueToConcurrentModificationException e) { + assertThat(counter.retryCount).isEqualTo(1); + } finally { + mockSpanner.putStatementResult(StatementResult.update(INSERT_STATEMENT, UPDATE_COUNT)); + } + + // Verify the order of the statements on the server. + List requests = + Lists.newArrayList( + Collections2.filter( + mockSpanner.getRequests(), + new Predicate() { + @Override + public boolean apply(AbstractMessage input) { + return input instanceof ExecuteSqlRequest; + } + })); + // The entire transaction should be retried, but will not succeed as the result of the update + // statement was different during the retry. + assertThat(requests).hasSize(4); + assertThat(((ExecuteSqlRequest) requests.get(0)).getSeqno()).isEqualTo(1L); + assertThat(((ExecuteSqlRequest) requests.get(0)).getSql()) + .isEqualTo(SELECT_RANDOM_STATEMENT.getSql()); + assertThat(((ExecuteSqlRequest) requests.get(1)).getSeqno()).isEqualTo(2L); + assertThat(((ExecuteSqlRequest) requests.get(1)).getSql()) + .isEqualTo(INSERT_STATEMENT.getSql()); + assertThat(((ExecuteSqlRequest) requests.get(2)).getSeqno()).isEqualTo(1L); + assertThat(((ExecuteSqlRequest) requests.get(2)).getSql()) + .isEqualTo(SELECT_RANDOM_STATEMENT.getSql()); + assertThat(((ExecuteSqlRequest) requests.get(3)).getSeqno()).isEqualTo(2L); + assertThat(((ExecuteSqlRequest) requests.get(3)).getSql()) + .isEqualTo(INSERT_STATEMENT.getSql()); + } + } + + @Test + public void testQueriesAbortedMidway_ResultsChanged() throws InterruptedException { + mockSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.ofStreamException( + Status.ABORTED.asRuntimeException(), RANDOM_RESULT_SET_ROW_COUNT - 1)); + final Statement statement = Statement.of("SELECT * FROM TEST_TABLE"); + final RandomResultSetGenerator generator = + new RandomResultSetGenerator(RANDOM_RESULT_SET_ROW_COUNT - 10); + mockSpanner.putStatementResult(StatementResult.query(statement, generator.generate())); + + final CountDownLatch latch = new CountDownLatch(1); + final RetryCounter counter = new RetryCounter(); + try (Connection connection = createConnection(counter)) { + ApiFuture res1; + try (AsyncResultSet rs = + connection.executeQueryAsync(SELECT_RANDOM_STATEMENT, Options.bufferRows(5))) { + res1 = + rs.setCallback( + multiThreadedExecutor, + new ReadyCallback() { + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + try { + latch.await(10L, TimeUnit.SECONDS); + while (true) { + switch (resultSet.tryNext()) { + case OK: + break; + case DONE: + return CallbackResponse.DONE; + case NOT_READY: + return CallbackResponse.CONTINUE; + } + } + } catch (Throwable t) { + throw SpannerExceptionFactory.asSpannerException(t); + } + } + }); + } + try (AsyncResultSet rs = connection.executeQueryAsync(statement, Options.bufferRows(5))) { + rs.setCallback( + multiThreadedExecutor, + new ReadyCallback() { + boolean replaced; + + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + if (!replaced) { + // Replace the result of the query on the server after the first execution. + mockSpanner.putStatementResult( + StatementResult.query(statement, generator.generate())); + replaced = true; + } + while (true) { + switch (resultSet.tryNext()) { + case OK: + break; + case DONE: + latch.countDown(); + return CallbackResponse.DONE; + case NOT_READY: + return CallbackResponse.CONTINUE; + } + } + } + }); + } + try { + get(res1); + fail("Missing expected exception"); + } catch (AbortedDueToConcurrentModificationException e) { + assertThat(counter.retryCount).isEqualTo(1); + } + } + } + + @Test + public void testBlindUpdateAborted() { + RetryCounter counter = new RetryCounter(); + try (Connection connection = createConnection(counter)) { + mockSpanner.abortNextStatement(); + ApiFuture updateCount = connection.executeUpdateAsync(INSERT_STATEMENT); + get(connection.commitAsync()); + + assertThat(get(updateCount)).isEqualTo(UPDATE_COUNT); + assertThat(counter.retryCount).isEqualTo(1); + } + } + + @Test + public void testBlindUpdateAborted_WithConcurrentModification() { + Statement update1 = Statement.of("UPDATE FOO SET BAR=1 WHERE BAZ=100"); + mockSpanner.putStatementResult(StatementResult.update(update1, 100)); + + RetryCounter counter = new RetryCounter(); + try (Connection connection = createConnection(counter)) { + // Execute an update statement and then change the result for the next time it is executed. + get(connection.executeUpdateAsync(update1)); + mockSpanner.putStatementResult(StatementResult.update(update1, 200)); + + // Abort on the next statement. The retry should now fail because of the changed result of the + // first update. + mockSpanner.abortNextStatement(); + connection.executeUpdateAsync(INSERT_STATEMENT); + + try { + get(connection.commitAsync()); + fail("Missing expected exception"); + } catch (AbortedDueToConcurrentModificationException e) { + assertThat(counter.retryCount).isEqualTo(1); + } + } + } + + @Test + public void testMultipleBlindUpdatesAborted_WithConcurrentModification() { + Statement update1 = Statement.of("UPDATE FOO SET BAR=1 WHERE BAZ=100"); + mockSpanner.putStatementResult(StatementResult.update(update1, 100)); + + RetryCounter counter = new RetryCounter(); + try (Connection connection = createConnection(counter)) { + // Execute an update statement and then change the result for the next time it is executed. + get(connection.executeUpdateAsync(update1)); + mockSpanner.putStatementResult(StatementResult.update(update1, 200)); + + // Abort the transaction on the next statement. The retry should now fail because of the + // changed result of the first update. + mockSpanner.abortNextStatement(); + + // Continue to (try to) execute blind updates. This should not cause any exceptions, although + // all of the returned futures will fail. + List> futures = new ArrayList<>(); + for (int i = 0; i < 3; i++) { + futures.add(connection.executeUpdateAsync(INSERT_STATEMENT)); + } + + for (ApiFuture fut : futures) { + try { + get(fut); + fail("Missing expected exception"); + } catch (AbortedDueToConcurrentModificationException e) { + assertThat(counter.retryCount).isEqualTo(1); + } + } + } + } + + @Test + public void testBlindUpdateAborted_ThenAsyncQuery_WithConcurrentModification() { + Statement update1 = Statement.of("UPDATE FOO SET BAR=1 WHERE BAZ=100"); + mockSpanner.putStatementResult(StatementResult.update(update1, 100)); + + RetryCounter counter = new RetryCounter(); + try (Connection connection = createConnection(counter)) { + // Execute an update statement and then change the result for the next time it is executed. + get(connection.executeUpdateAsync(update1)); + mockSpanner.putStatementResult(StatementResult.update(update1, 200)); + + // Abort on the next statement. The retry should now fail because of the changed result of the + // first update. + mockSpanner.abortNextStatement(); + connection.executeUpdateAsync(INSERT_STATEMENT); + + // Try to execute an async query. The callback should also receive the + // AbortedDueToConcurrentModificationException. + try (AsyncResultSet rs = connection.executeQueryAsync(SELECT_RANDOM_STATEMENT)) { + ApiFuture fut = + rs.setCallback( + singleThreadedExecutor, + new ReadyCallback() { + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + // The following line should throw AbortedDueToConcurrentModificationException. + resultSet.tryNext(); + return CallbackResponse.DONE; + } + }); + try { + assertThat(get(fut)).isNull(); + fail("Missing expected exception"); + } catch (AbortedDueToConcurrentModificationException e) { + assertThat(counter.retryCount).isEqualTo(1); + } + } + + // Ensure that a rollback and then a new statement does succeed. + connection.rollbackAsync(); + try (AsyncResultSet rs = connection.executeQueryAsync(SELECT_RANDOM_STATEMENT)) { + ApiFuture fut = + rs.setCallback( + singleThreadedExecutor, + new ReadyCallback() { + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + resultSet.tryNext(); + return CallbackResponse.DONE; + } + }); + assertThat(get(fut)).isNull(); + } + get(connection.commitAsync()); + } + } + + @Test + public void testBlindUpdateAborted_SelectResults() { + final Statement update1 = Statement.of("UPDATE FOO SET BAR=1 WHERE BAZ=100"); + mockSpanner.putStatementResult(StatementResult.update(update1, 100)); + + RetryCounter counter = new RetryCounter(); + try (Connection connection = createConnection(counter)) { + // Execute an update statement and then change the result for the next time it is executed. + connection.executeUpdate(update1); + // Abort on the next statement. The retry should now fail because of the changed result of the + // first update. + mockSpanner.abortNextStatement(); + mockSpanner.putStatementResult(StatementResult.update(update1, 200)); + connection.executeUpdateAsync(INSERT_STATEMENT); + ApiFuture commit = connection.commitAsync(); + + try (AsyncResultSet rs = connection.executeQueryAsync(SELECT_RANDOM_STATEMENT)) { + while (rs.next()) {} + } + get(connection.commitAsync()); + + try { + get(commit); + fail("Missing expected exception"); + } catch (AbortedDueToConcurrentModificationException e) { + assertThat(counter.retryCount).isEqualTo(1); + } + } + } + + private QueryResult executeQueryAsync(Connection connection, Statement statement) { + return executeQueryAsync(connection, statement, singleThreadedExecutor); + } + + private QueryResult executeQueryAsync( + Connection connection, Statement statement, Executor executor) { + ApiFuture res; + final AtomicInteger rowCount = new AtomicInteger(); + try (AsyncResultSet rs = connection.executeQueryAsync(statement, Options.bufferRows(5))) { + res = + rs.setCallback( + executor, + new ReadyCallback() { + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + while (true) { + switch (resultSet.tryNext()) { + case OK: + rowCount.incrementAndGet(); + break; + case DONE: + return CallbackResponse.DONE; + case NOT_READY: + return CallbackResponse.CONTINUE; + } + } + } + }); + return new QueryResult(res, rowCount); + } + } +} diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionAsyncApiTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionAsyncApiTest.java new file mode 100644 index 0000000000..39d33ae1ca --- /dev/null +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionAsyncApiTest.java @@ -0,0 +1,833 @@ +/* + * 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.connection; + +import static com.google.cloud.spanner.SpannerApiFutures.get; +import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.fail; + +import com.google.api.core.ApiFuture; +import com.google.api.core.SettableApiFuture; +import com.google.cloud.spanner.AsyncResultSet; +import com.google.cloud.spanner.AsyncResultSet.CallbackResponse; +import com.google.cloud.spanner.AsyncResultSet.ReadyCallback; +import com.google.cloud.spanner.ErrorCode; +import com.google.cloud.spanner.MockSpannerServiceImpl.SimulatedExecutionTime; +import com.google.cloud.spanner.Mutation; +import com.google.cloud.spanner.ResultSet; +import com.google.cloud.spanner.SpannerApiFutures; +import com.google.cloud.spanner.SpannerException; +import com.google.cloud.spanner.Statement; +import com.google.cloud.spanner.connection.StatementResult.ResultType; +import com.google.common.base.Function; +import com.google.common.base.Predicate; +import com.google.common.collect.Collections2; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.Lists; +import com.google.protobuf.AbstractMessage; +import com.google.spanner.v1.ExecuteBatchDmlRequest; +import com.google.spanner.v1.ExecuteSqlRequest; +import java.util.List; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.After; +import org.junit.AfterClass; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class ConnectionAsyncApiTest extends AbstractMockServerTest { + private static final ExecutorService executor = Executors.newSingleThreadExecutor(); + private static final Function AUTOCOMMIT = + new Function() { + @Override + public Void apply(Connection input) { + input.setAutocommit(true); + return null; + } + }; + private static final Function READ_ONLY = + new Function() { + @Override + public Void apply(Connection input) { + input.setReadOnly(true); + return null; + } + }; + private static final Function READ_WRITE = + new Function() { + @Override + public Void apply(Connection input) { + return null; + } + }; + + @AfterClass + public static void stopExecutor() { + executor.shutdown(); + } + + @After + public void reset() { + mockSpanner.removeAllExecutionTimes(); + } + + @Test + public void testExecuteQueryAsyncAutocommit() { + testExecuteQueryAsync(AUTOCOMMIT); + } + + @Test + public void testExecuteQueryAsyncAutocommitIsNonBlocking() { + testExecuteQueryAsyncIsNonBlocking(AUTOCOMMIT); + } + + @Test + public void testExecuteQueryAsStatementAsyncAutocommit() { + testExecuteQueryAsync(AUTOCOMMIT, true); + } + + @Test + public void testExecuteQueryAutocommit() { + testExecuteQuery(AUTOCOMMIT); + } + + @Test + public void testExecuteUpdateAsyncAutocommit() { + testExecuteUpdateAsync(AUTOCOMMIT); + } + + @Test + public void testExecuteUpdateAsyncAutocommitIsNonBlocking() { + testExecuteUpdateAsyncIsNonBlocking(AUTOCOMMIT); + } + + @Test + public void testExecuteUpdateAsStatementAsyncAutocommit() { + testExecuteUpdateAsync(AUTOCOMMIT, true); + } + + @Test + public void testExecuteUpdateAutocommit() { + testExecuteUpdate(AUTOCOMMIT); + } + + @Test + public void testExecuteBatchUpdateAsyncAutocommit() { + testExecuteBatchUpdateAsync(AUTOCOMMIT); + } + + @Test + public void testExecuteBatchUpdateAsyncAutocommitIsNonBlocking() { + testExecuteBatchUpdateAsyncIsNonBlocking(AUTOCOMMIT); + } + + @Test + public void testExecuteBatchUpdateAutocommit() { + testExecuteBatchUpdate(AUTOCOMMIT); + } + + @Test + public void testWriteAsyncAutocommit() { + testWriteAsync(AUTOCOMMIT); + } + + @Test + public void testWriteAutocommit() { + testWrite(AUTOCOMMIT); + } + + @Test + public void testExecuteQueryAsyncReadOnly() { + testExecuteQueryAsync(READ_ONLY); + } + + @Test + public void testExecuteQueryAsyncReadOnlyIsNonBlocking() { + testExecuteQueryAsyncIsNonBlocking(READ_ONLY); + } + + @Test + public void testExecuteQueryAsStatementAsyncReadOnly() { + testExecuteQueryAsync(READ_ONLY, true); + } + + @Test + public void testExecuteQueryReadOnly() { + testExecuteQuery(READ_ONLY); + } + + @Test + public void testExecuteQueryAsyncReadWrite() { + testExecuteQueryAsync(READ_WRITE); + } + + @Test + public void testExecuteQueryAsyncReadWriteIsNonBlocking() { + testExecuteQueryAsyncIsNonBlocking(READ_WRITE); + } + + @Test + public void testExecuteQueryAsStatementAsyncReadWrite() { + testExecuteQueryAsync(READ_WRITE, true); + } + + @Test + public void testExecuteQueryReadWrite() { + testExecuteQuery(READ_WRITE); + } + + @Test + public void testExecuteUpdateAsyncReadWrite() { + testExecuteUpdateAsync(READ_WRITE); + } + + @Test + public void testExecuteUpdateAsyncReadWriteIsNonBlocking() { + testExecuteUpdateAsyncIsNonBlocking(READ_WRITE); + } + + @Test + public void testExecuteUpdateAsStatementAsyncReadWrite() { + testExecuteUpdateAsync(READ_WRITE, true); + } + + @Test + public void testExecuteUpdateReadWrite() { + testExecuteUpdate(READ_WRITE); + } + + @Test + public void testExecuteBatchUpdateAsyncReadWrite() { + testExecuteBatchUpdateAsync(READ_WRITE); + } + + @Test + public void testExecuteBatchUpdateAsyncReadWriteIsNonBlocking() { + testExecuteBatchUpdateAsyncIsNonBlocking(READ_WRITE); + } + + @Test + public void testExecuteBatchUpdateReadWrite() { + testExecuteBatchUpdate(READ_WRITE); + } + + @Test + public void testBufferedWriteReadWrite() { + testBufferedWrite(READ_WRITE); + } + + @Test + public void testReadWriteMultipleAsyncStatements() { + try (Connection connection = createConnection()) { + assertThat(connection.isAutocommit()).isFalse(); + ApiFuture update1 = connection.executeUpdateAsync(INSERT_STATEMENT); + ApiFuture update2 = connection.executeUpdateAsync(INSERT_STATEMENT); + ApiFuture batch = + connection.executeBatchUpdateAsync(ImmutableList.of(INSERT_STATEMENT, INSERT_STATEMENT)); + final SettableApiFuture rowCount = SettableApiFuture.create(); + try (AsyncResultSet rs = connection.executeQueryAsync(SELECT_RANDOM_STATEMENT)) { + rs.setCallback( + executor, + new ReadyCallback() { + int count = 0; + + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + try { + while (true) { + switch (resultSet.tryNext()) { + case DONE: + rowCount.set(count); + return CallbackResponse.DONE; + case NOT_READY: + return CallbackResponse.CONTINUE; + case OK: + count++; + } + } + } catch (SpannerException e) { + rowCount.setException(e); + return CallbackResponse.DONE; + } + } + }); + } + connection.commitAsync(); + assertThat(get(update1)).isEqualTo(UPDATE_COUNT); + assertThat(get(update2)).isEqualTo(UPDATE_COUNT); + assertThat(get(batch)).asList().containsExactly(1L, 1L); + assertThat(get(rowCount)).isEqualTo(RANDOM_RESULT_SET_ROW_COUNT); + + // Verify the order of the statements on the server. + List requests = + Lists.newArrayList( + Collections2.filter( + mockSpanner.getRequests(), + new Predicate() { + @Override + public boolean apply(AbstractMessage input) { + return input instanceof ExecuteSqlRequest + || input instanceof ExecuteBatchDmlRequest; + } + })); + assertThat(requests).hasSize(4); + assertThat(requests.get(0)).isInstanceOf(ExecuteSqlRequest.class); + assertThat(((ExecuteSqlRequest) requests.get(0)).getSeqno()).isEqualTo(1L); + assertThat(requests.get(1)).isInstanceOf(ExecuteSqlRequest.class); + assertThat(((ExecuteSqlRequest) requests.get(1)).getSeqno()).isEqualTo(2L); + assertThat(requests.get(2)).isInstanceOf(ExecuteBatchDmlRequest.class); + assertThat(((ExecuteBatchDmlRequest) requests.get(2)).getSeqno()).isEqualTo(3L); + assertThat(requests.get(3)).isInstanceOf(ExecuteSqlRequest.class); + assertThat(((ExecuteSqlRequest) requests.get(3)).getSeqno()).isEqualTo(4L); + } + } + + @Test + public void testAutocommitRunBatch() { + try (Connection connection = createConnection()) { + connection.setAutocommit(true); + connection.execute(Statement.of("START BATCH DML")); + connection.execute(INSERT_STATEMENT); + connection.execute(INSERT_STATEMENT); + StatementResult res = connection.execute(Statement.of("RUN BATCH")); + assertThat(res.getResultType()).isEqualTo(ResultType.RESULT_SET); + try (ResultSet rs = res.getResultSet()) { + assertThat(rs.next()).isTrue(); + assertThat(rs.getLongList(0)).containsExactly(1L, 1L); + assertThat(rs.next()).isFalse(); + } + } + } + + @Test + public void testAutocommitRunBatchAsync() { + try (Connection connection = createConnection()) { + connection.executeAsync(Statement.of("SET AUTOCOMMIT = TRUE")); + connection.executeAsync(Statement.of("START BATCH DML")); + connection.executeAsync(INSERT_STATEMENT); + connection.executeAsync(INSERT_STATEMENT); + ApiFuture res = connection.runBatchAsync(); + assertThat(get(res)).asList().containsExactly(1L, 1L); + } + } + + @Test + public void testExecuteDdlAsync() { + try (Connection connection = createConnection()) { + connection.executeAsync(Statement.of("SET AUTOCOMMIT = TRUE")); + connection.executeAsync(Statement.of("START BATCH DDL")); + connection.executeAsync(Statement.of("CREATE TABLE FOO (ID INT64) PRIMARY KEY (ID)")); + connection.executeAsync(Statement.of("ABORT BATCH")); + } + } + + @Test + public void testExecuteInvalidStatementAsync() { + try (Connection connection = createConnection()) { + try { + connection.executeAsync(Statement.of("UPSERT INTO FOO (ID, VAL) VALUES (1, 'foo')")); + fail("Missing expected exception"); + } catch (SpannerException e) { + assertThat(e.getErrorCode()).isEqualTo(ErrorCode.INVALID_ARGUMENT); + } + } + } + + @Test + public void testExecuteClientSideQueryAsync() { + try (Connection connection = createConnection()) { + connection.executeAsync(Statement.of("SET AUTOCOMMIT = TRUE")); + final SettableApiFuture autocommit = SettableApiFuture.create(); + try (AsyncResultSet rs = + connection.executeQueryAsync(Statement.of("SHOW VARIABLE AUTOCOMMIT"))) { + rs.setCallback( + executor, + new ReadyCallback() { + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + while (true) { + switch (resultSet.tryNext()) { + case DONE: + return CallbackResponse.DONE; + case NOT_READY: + return CallbackResponse.CONTINUE; + case OK: + autocommit.set(resultSet.getBoolean("AUTOCOMMIT")); + } + } + } + }); + } + assertThat(get(autocommit)).isTrue(); + } + } + + @Test + public void testExecuteInvalidQueryAsync() { + try (Connection connection = createConnection()) { + try { + connection.executeQueryAsync(INSERT_STATEMENT); + fail("Missing expected exception"); + } catch (SpannerException e) { + assertThat(e.getErrorCode()).isEqualTo(ErrorCode.INVALID_ARGUMENT); + } + } + } + + @Test + public void testExecuteInvalidUpdateAsync() { + try (Connection connection = createConnection()) { + try { + connection.executeUpdateAsync(SELECT_RANDOM_STATEMENT); + fail("Missing expected exception"); + } catch (SpannerException e) { + assertThat(e.getErrorCode()).isEqualTo(ErrorCode.INVALID_ARGUMENT); + } + } + } + + @Test + public void testExecuteInvalidBatchUpdateAsync() { + try (Connection connection = createConnection()) { + try { + connection.executeBatchUpdateAsync( + ImmutableList.of(INSERT_STATEMENT, SELECT_RANDOM_STATEMENT)); + fail("Missing expected exception"); + } catch (SpannerException e) { + assertThat(e.getErrorCode()).isEqualTo(ErrorCode.INVALID_ARGUMENT); + } + } + } + + @Test + public void testRunEmptyBatchAsync() { + try (Connection connection = createConnection()) { + connection.startBatchDml(); + assertThat(get(connection.runBatchAsync())).isEqualTo(new long[0]); + } + } + + private void testExecuteQueryAsync(Function connectionConfigurator) { + testExecuteQueryAsync(connectionConfigurator, false); + } + + private void testExecuteQueryAsync( + Function connectionConfigurator, boolean executeAsStatement) { + ApiFuture res; + try (Connection connection = createConnection()) { + connectionConfigurator.apply(connection); + for (boolean timeout : new boolean[] {true, false}) { + final AtomicInteger rowCount = new AtomicInteger(); + final AtomicBoolean receivedTimeout = new AtomicBoolean(); + if (timeout) { + mockSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.ofMinimumAndRandomTime(10, 0)); + connection.setStatementTimeout(1L, TimeUnit.NANOSECONDS); + } else { + mockSpanner.removeAllExecutionTimes(); + connection.clearStatementTimeout(); + } + try (AsyncResultSet rs = + executeAsStatement + ? connection.executeAsync(SELECT_RANDOM_STATEMENT).getResultSetAsync() + : connection.executeQueryAsync(SELECT_RANDOM_STATEMENT)) { + res = + rs.setCallback( + executor, + new ReadyCallback() { + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + try { + while (true) { + switch (resultSet.tryNext()) { + case OK: + rowCount.incrementAndGet(); + break; + case DONE: + return CallbackResponse.DONE; + case NOT_READY: + return CallbackResponse.CONTINUE; + } + } + } catch (SpannerException e) { + receivedTimeout.set(e.getErrorCode() == ErrorCode.DEADLINE_EXCEEDED); + throw e; + } + } + }); + } + try { + SpannerApiFutures.get(res); + assertThat(rowCount.get()).isEqualTo(RANDOM_RESULT_SET_ROW_COUNT); + if (connection.isReadOnly() || !connection.isInTransaction()) { + assertThat(connection.getReadTimestamp()).isNotNull(); + } + assertThat(timeout).isFalse(); + } catch (SpannerException e) { + assertThat(e.getSuppressed()).hasLength(1); + assertThat(e.getSuppressed()[0].getMessage()).contains(SELECT_RANDOM_STATEMENT.getSql()); + assertThat(e.getErrorCode()).isEqualTo(ErrorCode.DEADLINE_EXCEEDED); + assertThat(timeout).isTrue(); + assertThat(receivedTimeout.get()).isTrue(); + // Start a new transaction if a timeout occurred on a read/write transaction, as that will + // invalidate that transaction. + if (!connection.isReadOnly() && connection.isInTransaction()) { + connection.clearStatementTimeout(); + connection.rollback(); + } + } + } + } + } + + private void testExecuteQuery(Function connectionConfigurator) { + long rowCount = 0L; + try (Connection connection = createConnection()) { + connectionConfigurator.apply(connection); + for (boolean timeout : new boolean[] {true, false}) { + if (timeout) { + mockSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.ofMinimumAndRandomTime(10, 0)); + connection.setStatementTimeout(1L, TimeUnit.NANOSECONDS); + } else { + mockSpanner.removeAllExecutionTimes(); + connection.clearStatementTimeout(); + } + try (ResultSet rs = connection.executeQuery(SELECT_RANDOM_STATEMENT)) { + while (rs.next()) { + rowCount++; + } + assertThat(rowCount).isEqualTo(RANDOM_RESULT_SET_ROW_COUNT); + if (connection.isReadOnly() || !connection.isInTransaction()) { + assertThat(connection.getReadTimestamp()).isNotNull(); + } + assertThat(timeout).isFalse(); + } catch (SpannerException e) { + assertThat(timeout).isTrue(); + assertThat(e.getErrorCode()).isEqualTo(ErrorCode.DEADLINE_EXCEEDED); + // Start a new transaction if a timeout occurred on a read/write transaction, as that will + // invalidate that transaction. + if (!connection.isReadOnly() && connection.isInTransaction()) { + connection.clearStatementTimeout(); + connection.rollback(); + } + } + } + } + } + + private void testExecuteUpdateAsync(Function connectionConfigurator) { + testExecuteUpdateAsync(connectionConfigurator, false); + } + + private void testExecuteUpdateAsync( + Function connectionConfigurator, boolean executeAsStatement) { + try (Connection connection = createConnection()) { + connectionConfigurator.apply(connection); + for (boolean timeout : new boolean[] {true, false}) { + if (timeout) { + mockSpanner.setExecuteSqlExecutionTime( + SimulatedExecutionTime.ofMinimumAndRandomTime(10, 0)); + connection.setStatementTimeout(1L, TimeUnit.NANOSECONDS); + } else { + mockSpanner.removeAllExecutionTimes(); + connection.clearStatementTimeout(); + } + ApiFuture updateCount = + executeAsStatement + ? connection.executeAsync(INSERT_STATEMENT).getUpdateCountAsync() + : connection.executeUpdateAsync(INSERT_STATEMENT); + try { + assertThat(get(updateCount)).isEqualTo(1L); + if (connection.isInTransaction()) { + connection.commitAsync(); + } + assertThat(connection.getCommitTimestamp()).isNotNull(); + assertThat(timeout).isFalse(); + } catch (SpannerException e) { + assertThat(timeout).isTrue(); + assertThat(e.getErrorCode()).isEqualTo(ErrorCode.DEADLINE_EXCEEDED); + // Start a new transaction if a timeout occurred on a read/write transaction, as that will + // invalidate that transaction. + if (!connection.isReadOnly() && connection.isInTransaction()) { + connection.clearStatementTimeout(); + connection.rollback(); + } + } + } + } + } + + private void testExecuteUpdate(Function connectionConfigurator) { + try (Connection connection = createConnection()) { + connectionConfigurator.apply(connection); + for (boolean timeout : new boolean[] {true, false}) { + if (timeout) { + mockSpanner.setExecuteSqlExecutionTime( + SimulatedExecutionTime.ofMinimumAndRandomTime(10, 0)); + connection.setStatementTimeout(1L, TimeUnit.NANOSECONDS); + } else { + mockSpanner.removeAllExecutionTimes(); + connection.clearStatementTimeout(); + } + try { + long updateCount = connection.executeUpdate(INSERT_STATEMENT); + assertThat(updateCount).isEqualTo(1L); + if (connection.isInTransaction()) { + connection.commit(); + } + assertThat(connection.getCommitTimestamp()).isNotNull(); + assertThat(timeout).isFalse(); + } catch (SpannerException e) { + assertThat(timeout).isTrue(); + assertThat(e.getErrorCode()).isEqualTo(ErrorCode.DEADLINE_EXCEEDED); + // Start a new transaction if a timeout occurred on a read/write transaction, as that will + // invalidate that transaction. + if (!connection.isReadOnly() && connection.isInTransaction()) { + connection.clearStatementTimeout(); + connection.rollback(); + } + } + } + } + } + + private void testExecuteBatchUpdateAsync(Function connectionConfigurator) { + try (Connection connection = createConnection()) { + connectionConfigurator.apply(connection); + for (boolean timeout : new boolean[] {true, false}) { + if (timeout) { + mockSpanner.setExecuteBatchDmlExecutionTime( + SimulatedExecutionTime.ofMinimumAndRandomTime(10, 0)); + connection.setStatementTimeout(1L, TimeUnit.NANOSECONDS); + } else { + mockSpanner.removeAllExecutionTimes(); + connection.clearStatementTimeout(); + } + ApiFuture updateCounts = + connection.executeBatchUpdateAsync( + ImmutableList.of(INSERT_STATEMENT, INSERT_STATEMENT)); + try { + assertThat(get(updateCounts)).asList().containsExactly(1L, 1L); + if (connection.isInTransaction()) { + connection.commitAsync(); + } + assertThat(connection.getCommitTimestamp()).isNotNull(); + assertThat(timeout).isFalse(); + } catch (SpannerException e) { + assertThat(timeout).isTrue(); + assertThat(e.getErrorCode()).isEqualTo(ErrorCode.DEADLINE_EXCEEDED); + // Start a new transaction if a timeout occurred on a read/write transaction, as that will + // invalidate that transaction. + if (!connection.isReadOnly() && connection.isInTransaction()) { + connection.clearStatementTimeout(); + connection.rollback(); + } + } + } + } + } + + private void testExecuteBatchUpdate(Function connectionConfigurator) { + try (Connection connection = createConnection()) { + connectionConfigurator.apply(connection); + for (boolean timeout : new boolean[] {true, false}) { + if (timeout) { + mockSpanner.setExecuteBatchDmlExecutionTime( + SimulatedExecutionTime.ofMinimumAndRandomTime(10, 0)); + connection.setStatementTimeout(1L, TimeUnit.NANOSECONDS); + } else { + mockSpanner.removeAllExecutionTimes(); + connection.clearStatementTimeout(); + } + try { + long[] updateCounts = + connection.executeBatchUpdate(ImmutableList.of(INSERT_STATEMENT, INSERT_STATEMENT)); + assertThat(updateCounts).asList().containsExactly(1L, 1L); + if (connection.isInTransaction()) { + connection.commit(); + } + assertThat(connection.getCommitTimestamp()).isNotNull(); + assertThat(timeout).isFalse(); + } catch (SpannerException e) { + assertThat(e.getErrorCode()).isEqualTo(ErrorCode.DEADLINE_EXCEEDED); + assertThat(timeout).isTrue(); + // Start a new transaction if a timeout occurred on a read/write transaction, as that will + // invalidate that transaction. + if (!connection.isReadOnly() && connection.isInTransaction()) { + connection.clearStatementTimeout(); + connection.rollback(); + } + } + } + } + } + + private void testWriteAsync(Function connectionConfigurator) { + try (Connection connection = createConnection()) { + connectionConfigurator.apply(connection); + for (boolean timeout : new boolean[] {true, false}) { + if (timeout) { + mockSpanner.setCommitExecutionTime(SimulatedExecutionTime.ofMinimumAndRandomTime(10, 0)); + connection.setStatementTimeout(1L, TimeUnit.NANOSECONDS); + } else { + mockSpanner.removeAllExecutionTimes(); + connection.clearStatementTimeout(); + } + ApiFuture fut = + connection.writeAsync( + ImmutableList.of( + Mutation.newInsertBuilder("foo").build(), + Mutation.newInsertBuilder("bar").build())); + try { + assertThat(get(fut)).isNull(); + assertThat(connection.getCommitTimestamp()).isNotNull(); + assertThat(timeout).isFalse(); + } catch (SpannerException e) { + assertThat(timeout).isTrue(); + assertThat(e.getErrorCode()).isEqualTo(ErrorCode.DEADLINE_EXCEEDED); + } + } + } + } + + private void testWrite(Function connectionConfigurator) { + try (Connection connection = createConnection()) { + connectionConfigurator.apply(connection); + for (boolean timeout : new boolean[] {true, false}) { + if (timeout) { + mockSpanner.setCommitExecutionTime(SimulatedExecutionTime.ofMinimumAndRandomTime(10, 0)); + connection.setStatementTimeout(1L, TimeUnit.NANOSECONDS); + } else { + mockSpanner.removeAllExecutionTimes(); + connection.clearStatementTimeout(); + } + try { + connection.write( + ImmutableList.of( + Mutation.newInsertBuilder("foo").build(), + Mutation.newInsertBuilder("bar").build())); + assertThat(connection.getCommitTimestamp()).isNotNull(); + assertThat(timeout).isFalse(); + } catch (SpannerException e) { + assertThat(timeout).isTrue(); + assertThat(e.getErrorCode()).isEqualTo(ErrorCode.DEADLINE_EXCEEDED); + } + } + } + } + + private void testBufferedWrite(Function connectionConfigurator) { + try (Connection connection = createConnection()) { + connectionConfigurator.apply(connection); + for (boolean timeout : new boolean[] {true, false}) { + if (timeout) { + mockSpanner.setCommitExecutionTime( + SimulatedExecutionTime.ofMinimumAndRandomTime(1000, 0)); + connection.setStatementTimeout(1L, TimeUnit.NANOSECONDS); + } else { + mockSpanner.removeAllExecutionTimes(); + connection.clearStatementTimeout(); + } + try { + connection.bufferedWrite( + ImmutableList.of( + Mutation.newInsertBuilder("foo").build(), + Mutation.newInsertBuilder("bar").build())); + connection.commitAsync(); + assertThat(connection.getCommitTimestamp()).isNotNull(); + assertThat(timeout).isFalse(); + } catch (SpannerException e) { + assertThat(timeout).isTrue(); + assertThat(e.getErrorCode()).isEqualTo(ErrorCode.DEADLINE_EXCEEDED); + connection.clearStatementTimeout(); + connection.rollbackAsync(); + } + } + } + } + + private void testExecuteQueryAsyncIsNonBlocking( + Function connectionConfigurator) { + ApiFuture res; + final AtomicInteger rowCount = new AtomicInteger(); + mockSpanner.freeze(); + try (Connection connection = createConnection()) { + connectionConfigurator.apply(connection); + try (AsyncResultSet rs = connection.executeQueryAsync(SELECT_RANDOM_STATEMENT)) { + res = + rs.setCallback( + executor, + new ReadyCallback() { + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + while (true) { + switch (resultSet.tryNext()) { + case OK: + rowCount.incrementAndGet(); + break; + case DONE: + return CallbackResponse.DONE; + case NOT_READY: + return CallbackResponse.CONTINUE; + } + } + } + }); + mockSpanner.unfreeze(); + } + SpannerApiFutures.get(res); + assertThat(rowCount.get()).isEqualTo(RANDOM_RESULT_SET_ROW_COUNT); + } + } + + private void testExecuteUpdateAsyncIsNonBlocking( + Function connectionConfigurator) { + mockSpanner.freeze(); + try (Connection connection = createConnection()) { + connectionConfigurator.apply(connection); + ApiFuture updateCount = connection.executeUpdateAsync(INSERT_STATEMENT); + if (connection.isInTransaction()) { + connection.commitAsync(); + } + mockSpanner.unfreeze(); + assertThat(get(updateCount)).isEqualTo(1L); + assertThat(connection.getCommitTimestamp()).isNotNull(); + } + } + + private void testExecuteBatchUpdateAsyncIsNonBlocking( + Function connectionConfigurator) { + mockSpanner.freeze(); + try (Connection connection = createConnection()) { + connectionConfigurator.apply(connection); + ApiFuture updateCounts = + connection.executeBatchUpdateAsync(ImmutableList.of(INSERT_STATEMENT, INSERT_STATEMENT)); + if (connection.isInTransaction()) { + connection.commitAsync(); + } + mockSpanner.unfreeze(); + assertThat(get(updateCounts)).asList().containsExactly(1L, 1L); + assertThat(connection.getCommitTimestamp()).isNotNull(); + } + } +} diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionImplTest.java index f5295ef96b..88f942122a 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionImplTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionImplTest.java @@ -53,14 +53,13 @@ import com.google.cloud.spanner.TimestampBound.Mode; import com.google.cloud.spanner.TransactionContext; import com.google.cloud.spanner.TransactionManager; -import com.google.cloud.spanner.TransactionManager.TransactionState; import com.google.cloud.spanner.TransactionRunner; -import com.google.cloud.spanner.TransactionRunner.TransactionCallable; import com.google.cloud.spanner.Type; import com.google.cloud.spanner.connection.AbstractConnectionImplTest.ConnectionConsumer; import com.google.cloud.spanner.connection.ConnectionImpl.UnitOfWorkType; import com.google.cloud.spanner.connection.ConnectionStatementExecutorImpl.StatementTimeoutGetter; import com.google.cloud.spanner.connection.ReadOnlyStalenessUtil.GetExactStaleness; +import com.google.cloud.spanner.connection.StatementParser.ParsedStatement; import com.google.cloud.spanner.connection.StatementResult.ResultType; import com.google.spanner.admin.database.v1.UpdateDatabaseDdlMetadata; import com.google.spanner.v1.ExecuteSqlRequest.QueryOptions; @@ -74,6 +73,7 @@ import org.junit.runner.RunWith; import org.junit.runners.JUnit4; import org.mockito.Matchers; +import org.mockito.Mockito; import org.mockito.invocation.InvocationOnMock; import org.mockito.stubbing.Answer; @@ -319,11 +319,16 @@ public TransactionRunner answer(InvocationOnMock invocation) { new TransactionRunner() { private Timestamp commitTimestamp; - @SuppressWarnings("unchecked") @Override public T run(TransactionCallable callable) { this.commitTimestamp = Timestamp.now(); - return (T) Long.valueOf(1L); + TransactionContext tx = mock(TransactionContext.class); + when(tx.executeUpdate(Statement.of(UPDATE))).thenReturn(1L); + try { + return callable.run(tx); + } catch (Exception e) { + throw SpannerExceptionFactory.newSpannerException(e); + } } @Override @@ -1199,6 +1204,9 @@ public void testMergeQueryOptions() { DdlClient ddlClient = mock(DdlClient.class); DatabaseClient dbClient = mock(DatabaseClient.class); final UnitOfWork unitOfWork = mock(UnitOfWork.class); + when(unitOfWork.executeQueryAsync( + any(ParsedStatement.class), any(AnalyzeMode.class), Mockito.anyVararg())) + .thenReturn(ApiFutures.immediateFuture(mock(ResultSet.class))); try (ConnectionImpl impl = new ConnectionImpl(connectionOptions, spannerPool, ddlClient, dbClient) { @Override @@ -1210,7 +1218,7 @@ UnitOfWork getCurrentUnitOfWorkOrStartNewUnitOfWork() { impl.setOptimizerVersion("1"); impl.executeQuery(Statement.of("SELECT FOO FROM BAR")); verify(unitOfWork) - .executeQuery( + .executeQueryAsync( StatementParser.INSTANCE.parse( Statement.newBuilder("SELECT FOO FROM BAR") .withQueryOptions(QueryOptions.newBuilder().setOptimizerVersion("1").build()) @@ -1221,7 +1229,7 @@ UnitOfWork getCurrentUnitOfWorkOrStartNewUnitOfWork() { impl.setOptimizerVersion("2"); impl.executeQuery(Statement.of("SELECT FOO FROM BAR")); verify(unitOfWork) - .executeQuery( + .executeQueryAsync( StatementParser.INSTANCE.parse( Statement.newBuilder("SELECT FOO FROM BAR") .withQueryOptions(QueryOptions.newBuilder().setOptimizerVersion("2").build()) @@ -1234,7 +1242,7 @@ UnitOfWork getCurrentUnitOfWorkOrStartNewUnitOfWork() { impl.setOptimizerVersion("3"); impl.executeQuery(Statement.of("SELECT FOO FROM BAR"), prefetchOption); verify(unitOfWork) - .executeQuery( + .executeQueryAsync( StatementParser.INSTANCE.parse( Statement.newBuilder("SELECT FOO FROM BAR") .withQueryOptions(QueryOptions.newBuilder().setOptimizerVersion("3").build()) @@ -1251,7 +1259,7 @@ UnitOfWork getCurrentUnitOfWorkOrStartNewUnitOfWork() { .build(), prefetchOption); verify(unitOfWork) - .executeQuery( + .executeQueryAsync( StatementParser.INSTANCE.parse( Statement.newBuilder("SELECT FOO FROM BAR") .withQueryOptions(QueryOptions.newBuilder().setOptimizerVersion("5").build()) 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 5838e7778d..de820ccbcc 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 @@ -17,10 +17,15 @@ package com.google.cloud.spanner.connection; import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.fail; +import com.google.cloud.spanner.AbortedException; +import com.google.cloud.spanner.ErrorCode; import com.google.cloud.spanner.ResultSet; +import com.google.cloud.spanner.SpannerException; import com.google.cloud.spanner.SpannerOptions; import com.google.cloud.spanner.Statement; +import com.google.common.collect.ImmutableList; import com.google.spanner.v1.ExecuteSqlRequest; import com.google.spanner.v1.ExecuteSqlRequest.QueryOptions; import org.junit.Test; @@ -94,4 +99,76 @@ public String getOptimizerVersion() { SpannerOptions.useDefaultEnvironment(); } } + + @Test + public void testExecuteInvalidBatchUpdate() { + try (Connection connection = createConnection()) { + try { + connection.executeBatchUpdate(ImmutableList.of(INSERT_STATEMENT, SELECT_RANDOM_STATEMENT)); + fail("Missing expected exception"); + } catch (SpannerException e) { + assertThat(e.getErrorCode()).isEqualTo(ErrorCode.INVALID_ARGUMENT); + } + } + } + + @Test + public void testQueryAborted() { + try (Connection connection = createConnection()) { + connection.setRetryAbortsInternally(false); + for (boolean abort : new Boolean[] {true, false}) { + try { + if (abort) { + mockSpanner.abortNextStatement(); + } + connection.executeQuery(SELECT_RANDOM_STATEMENT); + assertThat(abort).isFalse(); + connection.commit(); + } catch (AbortedException e) { + assertThat(abort).isTrue(); + connection.rollback(); + } + } + } + } + + @Test + public void testUpdateAborted() { + try (Connection connection = createConnection()) { + connection.setRetryAbortsInternally(false); + for (boolean abort : new Boolean[] {true, false}) { + try { + if (abort) { + mockSpanner.abortNextStatement(); + } + connection.executeUpdate(INSERT_STATEMENT); + assertThat(abort).isFalse(); + connection.commit(); + } catch (AbortedException e) { + assertThat(abort).isTrue(); + connection.rollback(); + } + } + } + } + + @Test + public void testBatchUpdateAborted() { + try (Connection connection = createConnection()) { + connection.setRetryAbortsInternally(false); + for (boolean abort : new Boolean[] {true, false}) { + try { + if (abort) { + mockSpanner.abortNextStatement(); + } + connection.executeBatchUpdate(ImmutableList.of(INSERT_STATEMENT, INSERT_STATEMENT)); + assertThat(abort).isFalse(); + connection.commit(); + } catch (AbortedException e) { + assertThat(abort).isTrue(); + connection.rollback(); + } + } + } + } } diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/DdlBatchTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/DdlBatchTest.java index 6194b7f73b..1e09eb70f5 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/DdlBatchTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/DdlBatchTest.java @@ -16,6 +16,7 @@ package com.google.cloud.spanner.connection; +import static com.google.cloud.spanner.SpannerApiFutures.get; import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.CoreMatchers.is; import static org.hamcrest.MatcherAssert.assertThat; @@ -50,6 +51,7 @@ import com.google.spanner.admin.database.v1.UpdateDatabaseDdlMetadata; import io.grpc.Status; import java.util.Arrays; +import java.util.Collections; import java.util.List; import java.util.concurrent.ExecutionException; import java.util.concurrent.Executors; @@ -137,7 +139,7 @@ private DdlBatch createSubject(DdlClient ddlClient, DatabaseClient dbClient) { public void testExecuteQuery() { DdlBatch batch = createSubject(); try { - batch.executeQuery(mock(ParsedStatement.class), AnalyzeMode.NONE); + batch.executeQueryAsync(mock(ParsedStatement.class), AnalyzeMode.NONE); fail("expected FAILED_PRECONDITION"); } catch (SpannerException e) { assertEquals(ErrorCode.FAILED_PRECONDITION, e.getErrorCode()); @@ -156,16 +158,18 @@ public void testExecuteMetadataQuery() { when(singleUse.executeQuery(statement)).thenReturn(resultSet); when(dbClient.singleUse()).thenReturn(singleUse); DdlBatch batch = createSubject(createDefaultMockDdlClient(), dbClient); - ResultSet result = - batch.executeQuery(parsedStatement, AnalyzeMode.NONE, InternalMetadataQuery.INSTANCE); - assertThat(result.hashCode(), is(equalTo(resultSet.hashCode()))); + assertThat( + get(batch.executeQueryAsync( + parsedStatement, AnalyzeMode.NONE, InternalMetadataQuery.INSTANCE)) + .hashCode(), + is(equalTo(resultSet.hashCode()))); } @Test public void testExecuteUpdate() { DdlBatch batch = createSubject(); try { - batch.executeUpdate(mock(ParsedStatement.class)); + batch.executeUpdateAsync(mock(ParsedStatement.class)); fail("expected FAILED_PRECONDITION"); } catch (SpannerException e) { assertEquals(ErrorCode.FAILED_PRECONDITION, e.getErrorCode()); @@ -173,11 +177,10 @@ public void testExecuteUpdate() { } @Test - public void testGetCommitTimestamp() { + public void testExecuteBatchUpdate() { DdlBatch batch = createSubject(); - batch.runBatch(); try { - batch.getCommitTimestamp(); + batch.executeBatchUpdateAsync(Collections.singleton(mock(ParsedStatement.class))); fail("expected FAILED_PRECONDITION"); } catch (SpannerException e) { assertEquals(ErrorCode.FAILED_PRECONDITION, e.getErrorCode()); @@ -185,11 +188,11 @@ public void testGetCommitTimestamp() { } @Test - public void testGetReadTimestamp() { + public void testGetCommitTimestamp() { DdlBatch batch = createSubject(); - batch.runBatch(); + get(batch.runBatchAsync()); try { - batch.getReadTimestamp(); + batch.getCommitTimestamp(); fail("expected FAILED_PRECONDITION"); } catch (SpannerException e) { assertEquals(ErrorCode.FAILED_PRECONDITION, e.getErrorCode()); @@ -197,10 +200,11 @@ public void testGetReadTimestamp() { } @Test - public void testWrite() { + public void testGetReadTimestamp() { DdlBatch batch = createSubject(); + get(batch.runBatchAsync()); try { - batch.write(Mutation.newInsertBuilder("foo").build()); + batch.getReadTimestamp(); fail("expected FAILED_PRECONDITION"); } catch (SpannerException e) { assertEquals(ErrorCode.FAILED_PRECONDITION, e.getErrorCode()); @@ -211,7 +215,7 @@ public void testWrite() { public void testWriteIterable() { DdlBatch batch = createSubject(); try { - batch.write(Arrays.asList(Mutation.newInsertBuilder("foo").build())); + batch.writeAsync(Arrays.asList(Mutation.newInsertBuilder("foo").build())); fail("expected FAILED_PRECONDITION"); } catch (SpannerException e) { assertEquals(ErrorCode.FAILED_PRECONDITION, e.getErrorCode()); @@ -229,7 +233,7 @@ public void testGetStateAndIsActive() { DdlBatch batch = createSubject(); assertThat(batch.getState(), is(UnitOfWorkState.STARTED)); assertThat(batch.isActive(), is(true)); - batch.runBatch(); + get(batch.runBatchAsync()); assertThat(batch.getState(), is(UnitOfWorkState.RAN)); assertThat(batch.isActive(), is(false)); @@ -241,7 +245,9 @@ public void testGetStateAndIsActive() { assertThat(batch.isActive(), is(false)); DdlClient client = mock(DdlClient.class); - doThrow(SpannerException.class).when(client).executeDdl(anyListOf(String.class)); + SpannerException exception = mock(SpannerException.class); + when(exception.getErrorCode()).thenReturn(ErrorCode.FAILED_PRECONDITION); + doThrow(exception).when(client).executeDdl(anyListOf(String.class)); batch = createSubject(client); assertThat(batch.getState(), is(UnitOfWorkState.STARTED)); assertThat(batch.isActive(), is(true)); @@ -249,14 +255,13 @@ public void testGetStateAndIsActive() { when(statement.getStatement()).thenReturn(Statement.of("CREATE TABLE FOO")); when(statement.getSqlWithoutComments()).thenReturn("CREATE TABLE FOO"); when(statement.getType()).thenReturn(StatementType.DDL); - batch.executeDdl(statement); - boolean exception = false; + batch.executeDdlAsync(statement); try { - batch.runBatch(); + get(batch.runBatchAsync()); + fail("Missing expected exception"); } catch (SpannerException e) { - exception = true; + assertThat(e.getErrorCode(), is(equalTo(ErrorCode.FAILED_PRECONDITION))); } - assertThat(exception, is(true)); assertThat(batch.getState(), is(UnitOfWorkState.RUN_FAILED)); assertThat(batch.isActive(), is(false)); } @@ -287,7 +292,7 @@ public boolean matches(Object list) { public void testRunBatch() { DdlClient client = createDefaultMockDdlClient(); DdlBatch batch = createSubject(client); - batch.runBatch(); + get(batch.runBatchAsync()); assertThat(batch.getState(), is(UnitOfWorkState.RAN)); verify(client, never()).executeDdl(anyString()); verify(client, never()).executeDdl(argThat(isEmptyListOfStrings())); @@ -299,20 +304,20 @@ public void testRunBatch() { client = createDefaultMockDdlClient(); batch = createSubject(client); - batch.executeDdl(statement); - batch.runBatch(); + batch.executeDdlAsync(statement); + get(batch.runBatchAsync()); verify(client).executeDdl(argThat(isListOfStringsWithSize(1))); client = createDefaultMockDdlClient(); batch = createSubject(client); - batch.executeDdl(statement); - batch.executeDdl(statement); - batch.runBatch(); + batch.executeDdlAsync(statement); + batch.executeDdlAsync(statement); + get(batch.runBatchAsync()); verify(client).executeDdl(argThat(isListOfStringsWithSize(2))); assertThat(batch.getState(), is(UnitOfWorkState.RAN)); boolean exception = false; try { - batch.runBatch(); + get(batch.runBatchAsync()); } catch (SpannerException e) { if (e.getErrorCode() != ErrorCode.FAILED_PRECONDITION) { throw e; @@ -323,7 +328,7 @@ public void testRunBatch() { assertThat(batch.getState(), is(UnitOfWorkState.RAN)); exception = false; try { - batch.executeDdl(statement); + batch.executeDdlAsync(statement); } catch (SpannerException e) { if (e.getErrorCode() != ErrorCode.FAILED_PRECONDITION) { throw e; @@ -333,7 +338,7 @@ public void testRunBatch() { assertThat(exception, is(true)); exception = false; try { - batch.executeDdl(statement); + batch.executeDdlAsync(statement); } catch (SpannerException e) { if (e.getErrorCode() != ErrorCode.FAILED_PRECONDITION) { throw e; @@ -344,11 +349,11 @@ public void testRunBatch() { client = createDefaultMockDdlClient(true); batch = createSubject(client); - batch.executeDdl(statement); - batch.executeDdl(statement); + batch.executeDdlAsync(statement); + batch.executeDdlAsync(statement); exception = false; try { - batch.runBatch(); + get(batch.runBatchAsync()); } catch (SpannerException e) { exception = true; } @@ -380,9 +385,9 @@ public void testUpdateCount() throws InterruptedException, ExecutionException { .setDdlClient(client) .setDatabaseClient(mock(DatabaseClient.class)) .build(); - batch.executeDdl(StatementParser.INSTANCE.parse(Statement.of("CREATE TABLE FOO"))); - batch.executeDdl(StatementParser.INSTANCE.parse(Statement.of("CREATE TABLE BAR"))); - long[] updateCounts = batch.runBatch(); + batch.executeDdlAsync(StatementParser.INSTANCE.parse(Statement.of("CREATE TABLE FOO"))); + batch.executeDdlAsync(StatementParser.INSTANCE.parse(Statement.of("CREATE TABLE BAR"))); + long[] updateCounts = get(batch.runBatchAsync()); assertThat(updateCounts.length, is(equalTo(2))); assertThat(updateCounts[0], is(equalTo(1L))); assertThat(updateCounts[1], is(equalTo(1L))); @@ -412,10 +417,48 @@ public void testFailedUpdateCount() throws InterruptedException, ExecutionExcept .setDdlClient(client) .setDatabaseClient(mock(DatabaseClient.class)) .build(); - batch.executeDdl(StatementParser.INSTANCE.parse(Statement.of("CREATE TABLE FOO"))); - batch.executeDdl(StatementParser.INSTANCE.parse(Statement.of("CREATE TABLE INVALID_TABLE"))); + batch.executeDdlAsync(StatementParser.INSTANCE.parse(Statement.of("CREATE TABLE FOO"))); + batch.executeDdlAsync( + StatementParser.INSTANCE.parse(Statement.of("CREATE TABLE INVALID_TABLE"))); + try { + get(batch.runBatchAsync()); + fail("missing expected exception"); + } catch (SpannerBatchUpdateException e) { + assertThat(e.getUpdateCounts().length, is(equalTo(2))); + assertThat(e.getUpdateCounts()[0], is(equalTo(1L))); + assertThat(e.getUpdateCounts()[1], is(equalTo(0L))); + } + } + + @Test + public void testFailedAfterFirstStatement() throws InterruptedException, ExecutionException { + DdlClient client = mock(DdlClient.class); + UpdateDatabaseDdlMetadata metadata = + UpdateDatabaseDdlMetadata.newBuilder() + .addCommitTimestamps( + Timestamp.newBuilder().setSeconds(System.currentTimeMillis() * 1000L - 1L)) + .addAllStatements(Arrays.asList("CREATE TABLE FOO", "CREATE TABLE INVALID_TABLE")) + .build(); + ApiFuture metadataFuture = ApiFutures.immediateFuture(metadata); + @SuppressWarnings("unchecked") + OperationFuture operationFuture = mock(OperationFuture.class); + when(operationFuture.get()) + .thenThrow( + new ExecutionException( + "ddl statement failed", Status.INVALID_ARGUMENT.asRuntimeException())); + when(operationFuture.getMetadata()).thenReturn(metadataFuture); + when(client.executeDdl(argThat(isListOfStringsWithSize(2)))).thenReturn(operationFuture); + DdlBatch batch = + DdlBatch.newBuilder() + .withStatementExecutor(new StatementExecutor()) + .setDdlClient(client) + .setDatabaseClient(mock(DatabaseClient.class)) + .build(); + batch.executeDdlAsync(StatementParser.INSTANCE.parse(Statement.of("CREATE TABLE FOO"))); + batch.executeDdlAsync( + StatementParser.INSTANCE.parse(Statement.of("CREATE TABLE INVALID_TABLE"))); try { - batch.runBatch(); + get(batch.runBatchAsync()); fail("missing expected exception"); } catch (SpannerBatchUpdateException e) { assertThat(e.getUpdateCounts().length, is(equalTo(2))); @@ -440,26 +483,26 @@ public void testAbort() { client = createDefaultMockDdlClient(); batch = createSubject(client); - batch.executeDdl(statement); + batch.executeDdlAsync(statement); batch.abortBatch(); verify(client, never()).executeDdl(anyListOf(String.class)); client = createDefaultMockDdlClient(); batch = createSubject(client); - batch.executeDdl(statement); - batch.executeDdl(statement); + batch.executeDdlAsync(statement); + batch.executeDdlAsync(statement); batch.abortBatch(); verify(client, never()).executeDdl(anyListOf(String.class)); client = createDefaultMockDdlClient(); batch = createSubject(client); - batch.executeDdl(statement); - batch.executeDdl(statement); + batch.executeDdlAsync(statement); + batch.executeDdlAsync(statement); batch.abortBatch(); verify(client, never()).executeDdl(anyListOf(String.class)); boolean exception = false; try { - batch.runBatch(); + get(batch.runBatchAsync()); } catch (SpannerException e) { if (e.getErrorCode() != ErrorCode.FAILED_PRECONDITION) { throw e; @@ -479,7 +522,7 @@ public void testCancel() { DdlClient client = createDefaultMockDdlClient(10000L); final DdlBatch batch = createSubject(client); - batch.executeDdl(statement); + batch.executeDdlAsync(statement); Executors.newSingleThreadScheduledExecutor() .schedule( new Runnable() { @@ -491,7 +534,7 @@ public void run() { 100, TimeUnit.MILLISECONDS); try { - batch.runBatch(); + get(batch.runBatchAsync()); fail("expected CANCELLED"); } catch (SpannerException e) { assertEquals(ErrorCode.CANCELLED, e.getErrorCode()); @@ -502,7 +545,7 @@ public void run() { public void testCommit() { DdlBatch batch = createSubject(); try { - batch.commit(); + batch.commitAsync(); fail("expected FAILED_PRECONDITION"); } catch (SpannerException e) { assertEquals(ErrorCode.FAILED_PRECONDITION, e.getErrorCode()); @@ -513,7 +556,7 @@ public void testCommit() { public void testRollback() { DdlBatch batch = createSubject(); try { - batch.rollback(); + batch.rollbackAsync(); fail("expected FAILED_PRECONDITION"); } catch (SpannerException e) { assertEquals(ErrorCode.FAILED_PRECONDITION, e.getErrorCode()); diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/DmlBatchTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/DmlBatchTest.java index e841601db7..0f1ca38cd7 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/DmlBatchTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/DmlBatchTest.java @@ -16,15 +16,16 @@ package com.google.cloud.spanner.connection; +import static com.google.cloud.spanner.SpannerApiFutures.get; import static org.hamcrest.CoreMatchers.is; import static org.hamcrest.MatcherAssert.assertThat; import static org.junit.Assert.assertEquals; import static org.junit.Assert.fail; import static org.mockito.Matchers.anyListOf; -import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; +import com.google.api.core.ApiFutures; import com.google.cloud.spanner.ErrorCode; import com.google.cloud.spanner.Mutation; import com.google.cloud.spanner.SpannerException; @@ -47,8 +48,8 @@ public class DmlBatchTest { private DmlBatch createSubject() { UnitOfWork transaction = mock(UnitOfWork.class); - when(transaction.executeBatchUpdate(Arrays.asList(statement1, statement2))) - .thenReturn(new long[] {3L, 5L}); + when(transaction.executeBatchUpdateAsync(Arrays.asList(statement1, statement2))) + .thenReturn(ApiFutures.immediateFuture(new long[] {3L, 5L})); return createSubject(transaction); } @@ -63,7 +64,7 @@ private DmlBatch createSubject(UnitOfWork transaction) { public void testExecuteQuery() { DmlBatch batch = createSubject(); try { - batch.executeQuery(mock(ParsedStatement.class), AnalyzeMode.NONE); + batch.executeQueryAsync(mock(ParsedStatement.class), AnalyzeMode.NONE); fail("Expected exception"); } catch (SpannerException e) { assertEquals(ErrorCode.FAILED_PRECONDITION, e.getErrorCode()); @@ -74,7 +75,7 @@ public void testExecuteQuery() { public void testExecuteDdl() { DmlBatch batch = createSubject(); try { - batch.executeDdl(mock(ParsedStatement.class)); + batch.executeDdlAsync(mock(ParsedStatement.class)); fail("Expected exception"); } catch (SpannerException e) { assertEquals(ErrorCode.FAILED_PRECONDITION, e.getErrorCode()); @@ -84,7 +85,7 @@ public void testExecuteDdl() { @Test public void testGetReadTimestamp() { DmlBatch batch = createSubject(); - batch.runBatch(); + get(batch.runBatchAsync()); try { batch.getReadTimestamp(); fail("Expected exception"); @@ -102,7 +103,7 @@ public void testIsReadOnly() { @Test public void testGetCommitTimestamp() { DmlBatch batch = createSubject(); - batch.runBatch(); + get(batch.runBatchAsync()); try { batch.getCommitTimestamp(); fail("Expected exception"); @@ -111,22 +112,11 @@ public void testGetCommitTimestamp() { } } - @Test - public void testWrite() { - DmlBatch batch = createSubject(); - try { - batch.write(Mutation.newInsertBuilder("foo").build()); - fail("Expected exception"); - } catch (SpannerException e) { - assertEquals(ErrorCode.FAILED_PRECONDITION, e.getErrorCode()); - } - } - @Test public void testWriteIterable() { DmlBatch batch = createSubject(); try { - batch.write(Arrays.asList(Mutation.newInsertBuilder("foo").build())); + batch.writeAsync(Arrays.asList(Mutation.newInsertBuilder("foo").build())); fail("Expected exception"); } catch (SpannerException e) { assertEquals(ErrorCode.FAILED_PRECONDITION, e.getErrorCode()); @@ -138,7 +128,7 @@ public void testGetStateAndIsActive() { DmlBatch batch = createSubject(); assertThat(batch.getState(), is(UnitOfWorkState.STARTED)); assertThat(batch.isActive(), is(true)); - batch.runBatch(); + get(batch.runBatchAsync()); assertThat(batch.getState(), is(UnitOfWorkState.RAN)); assertThat(batch.isActive(), is(false)); @@ -150,7 +140,8 @@ public void testGetStateAndIsActive() { assertThat(batch.isActive(), is(false)); UnitOfWork tx = mock(UnitOfWork.class); - doThrow(SpannerException.class).when(tx).executeBatchUpdate(anyListOf(ParsedStatement.class)); + when(tx.executeBatchUpdateAsync(anyListOf(ParsedStatement.class))) + .thenReturn(ApiFutures.immediateFailedFuture(mock(SpannerException.class))); batch = createSubject(tx); assertThat(batch.getState(), is(UnitOfWorkState.STARTED)); assertThat(batch.isActive(), is(true)); @@ -158,10 +149,10 @@ public void testGetStateAndIsActive() { when(statement.getStatement()).thenReturn(Statement.of("UPDATE TEST SET COL1=2")); when(statement.getSqlWithoutComments()).thenReturn("UPDATE TEST SET COL1=2"); when(statement.getType()).thenReturn(StatementType.UPDATE); - batch.executeUpdate(statement); + get(batch.executeUpdateAsync(statement)); boolean exception = false; try { - batch.runBatch(); + get(batch.runBatchAsync()); } catch (SpannerException e) { exception = true; } @@ -174,7 +165,7 @@ public void testGetStateAndIsActive() { public void testCommit() { DmlBatch batch = createSubject(); try { - batch.commit(); + batch.commitAsync(); fail("Expected exception"); } catch (SpannerException e) { assertEquals(ErrorCode.FAILED_PRECONDITION, e.getErrorCode()); @@ -185,7 +176,7 @@ public void testCommit() { public void testRollback() { DmlBatch batch = createSubject(); try { - batch.rollback(); + batch.rollbackAsync(); fail("Expected exception"); } catch (SpannerException e) { assertEquals(ErrorCode.FAILED_PRECONDITION, e.getErrorCode()); diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ITAbstractSpannerTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ITAbstractSpannerTest.java index fae463ceb1..88cbcc108b 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ITAbstractSpannerTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ITAbstractSpannerTest.java @@ -32,6 +32,7 @@ import com.google.cloud.spanner.connection.SqlScriptVerifier.SpannerGenericConnection; import com.google.cloud.spanner.connection.StatementParser.ParsedStatement; import com.google.common.base.Preconditions; +import com.google.common.base.Stopwatch; import com.google.common.base.Strings; import java.lang.reflect.Field; import java.nio.file.Files; @@ -40,6 +41,7 @@ import java.util.Collections; import java.util.List; import java.util.Random; +import java.util.concurrent.TimeUnit; import org.junit.AfterClass; import org.junit.Before; import org.junit.BeforeClass; @@ -119,13 +121,27 @@ public void intercept( try { Field field = ReadWriteTransaction.class.getDeclaredField("txManager"); field.setAccessible(true); + Stopwatch watch = Stopwatch.createStarted(); + while (field.get(transaction) == null && watch.elapsed(TimeUnit.MILLISECONDS) < 100) { + Thread.sleep(1L); + } TransactionManager tx = (TransactionManager) field.get(transaction); + if (tx == null) { + return; + } Class cls = Class.forName("com.google.cloud.spanner.TransactionManagerImpl"); Class cls2 = Class.forName("com.google.cloud.spanner.SessionPool$AutoClosingTransactionManager"); Field delegateField = cls2.getDeclaredField("delegate"); delegateField.setAccessible(true); + watch = watch.reset().start(); + while (delegateField.get(tx) == null && watch.elapsed(TimeUnit.MILLISECONDS) < 100) { + Thread.sleep(1L); + } TransactionManager delegate = (TransactionManager) delegateField.get(tx); + if (delegate == null) { + return; + } Field stateField = cls.getDeclaredField("txnState"); stateField.setAccessible(true); diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ReadOnlyTransactionTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ReadOnlyTransactionTest.java index 118f596c86..ad8b2849a1 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ReadOnlyTransactionTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ReadOnlyTransactionTest.java @@ -16,6 +16,7 @@ package com.google.cloud.spanner.connection; +import static com.google.cloud.spanner.SpannerApiFutures.get; import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.CoreMatchers.is; import static org.hamcrest.CoreMatchers.notNullValue; @@ -183,7 +184,7 @@ public void testExecuteDdl() { ParsedStatement ddl = mock(ParsedStatement.class); when(ddl.getType()).thenReturn(StatementType.DDL); try { - createSubject().executeDdl(ddl); + createSubject().executeDdlAsync(ddl); fail("Expected exception"); } catch (SpannerException ex) { assertEquals(ErrorCode.FAILED_PRECONDITION, ex.getErrorCode()); @@ -195,18 +196,7 @@ public void testExecuteUpdate() { ParsedStatement update = mock(ParsedStatement.class); when(update.getType()).thenReturn(StatementType.UPDATE); try { - createSubject().executeUpdate(update); - fail("Expected exception"); - } catch (SpannerException ex) { - assertEquals(ErrorCode.FAILED_PRECONDITION, ex.getErrorCode()); - } - } - - @Test - public void testWrite() { - Mutation mutation = Mutation.newInsertBuilder("foo").build(); - try { - createSubject().write(mutation); + createSubject().executeUpdateAsync(update); fail("Expected exception"); } catch (SpannerException ex) { assertEquals(ErrorCode.FAILED_PRECONDITION, ex.getErrorCode()); @@ -217,7 +207,7 @@ public void testWrite() { public void testWriteIterable() { Mutation mutation = Mutation.newInsertBuilder("foo").build(); try { - createSubject().write(Arrays.asList(mutation, mutation)); + createSubject().writeAsync(Arrays.asList(mutation, mutation)); fail("Expected exception"); } catch (SpannerException ex) { assertEquals(ErrorCode.FAILED_PRECONDITION, ex.getErrorCode()); @@ -228,7 +218,7 @@ public void testWriteIterable() { public void testRunBatch() { ReadOnlyTransaction subject = createSubject(); try { - subject.runBatch(); + subject.runBatchAsync(); fail("Expected exception"); } catch (SpannerException ex) { assertEquals(ErrorCode.FAILED_PRECONDITION, ex.getErrorCode()); @@ -249,7 +239,7 @@ public void testAbortBatch() { @Test public void testGetCommitTimestamp() { ReadOnlyTransaction transaction = createSubject(); - transaction.commit(); + get(transaction.commitAsync()); assertThat(transaction.getState(), is(UnitOfWorkState.COMMITTED)); try { transaction.getCommitTimestamp(); @@ -275,7 +265,7 @@ public void testExecuteQuery() { when(parsedStatement.getSqlWithoutComments()).thenReturn(statement.getSql()); ReadOnlyTransaction transaction = createSubject(staleness); - ResultSet rs = transaction.executeQuery(parsedStatement, AnalyzeMode.NONE); + ResultSet rs = get(transaction.executeQueryAsync(parsedStatement, AnalyzeMode.NONE)); assertThat(rs, is(notNullValue())); assertThat(rs.getStats(), is(nullValue())); } @@ -308,11 +298,11 @@ public void testExecuteQueryWithOptionsTest() { .build(); ResultSet expectedWithOptions = DirectExecuteResultSet.ofResultSet(resWithOptions); assertThat( - transaction.executeQuery(parsedStatement, AnalyzeMode.NONE, option), + get(transaction.executeQueryAsync(parsedStatement, AnalyzeMode.NONE, option)), is(equalTo(expectedWithOptions))); ResultSet expectedWithoutOptions = DirectExecuteResultSet.ofResultSet(resWithoutOptions); assertThat( - transaction.executeQuery(parsedStatement, AnalyzeMode.NONE), + get(transaction.executeQueryAsync(parsedStatement, AnalyzeMode.NONE)), is(equalTo(expectedWithoutOptions))); } @@ -327,7 +317,7 @@ public void testPlanQuery() { when(parsedStatement.getSqlWithoutComments()).thenReturn(statement.getSql()); ReadOnlyTransaction transaction = createSubject(staleness); - ResultSet rs = transaction.executeQuery(parsedStatement, AnalyzeMode.PLAN); + ResultSet rs = get(transaction.executeQueryAsync(parsedStatement, AnalyzeMode.PLAN)); assertThat(rs, is(notNullValue())); // get all results and then get the stats while (rs.next()) { @@ -348,7 +338,7 @@ public void testProfileQuery() { when(parsedStatement.getSqlWithoutComments()).thenReturn(statement.getSql()); ReadOnlyTransaction transaction = createSubject(staleness); - ResultSet rs = transaction.executeQuery(parsedStatement, AnalyzeMode.PROFILE); + ResultSet rs = get(transaction.executeQueryAsync(parsedStatement, AnalyzeMode.PROFILE)); assertThat(rs, is(notNullValue())); // get all results and then get the stats while (rs.next()) { @@ -378,7 +368,9 @@ public void testGetReadTimestamp() { } } assertThat(expectedException, is(true)); - assertThat(transaction.executeQuery(parsedStatement, AnalyzeMode.NONE), is(notNullValue())); + assertThat( + get(transaction.executeQueryAsync(parsedStatement, AnalyzeMode.NONE)), + is(notNullValue())); assertThat(transaction.getReadTimestamp(), is(notNullValue())); } } @@ -406,7 +398,7 @@ public void testState() { transaction.getState(), is(equalTo(com.google.cloud.spanner.connection.UnitOfWork.UnitOfWorkState.STARTED))); assertThat(transaction.isActive(), is(true)); - transaction.commit(); + get(transaction.commitAsync()); assertThat( transaction.getState(), is(equalTo(com.google.cloud.spanner.connection.UnitOfWork.UnitOfWorkState.COMMITTED))); @@ -417,13 +409,14 @@ public void testState() { transaction.getState(), is(equalTo(com.google.cloud.spanner.connection.UnitOfWork.UnitOfWorkState.STARTED))); assertThat(transaction.isActive(), is(true)); - assertThat(transaction.executeQuery(parsedStatement, AnalyzeMode.NONE), is(notNullValue())); + assertThat( + get(transaction.executeQueryAsync(parsedStatement, AnalyzeMode.NONE)), is(notNullValue())); assertThat( transaction.getState(), is(equalTo(com.google.cloud.spanner.connection.UnitOfWork.UnitOfWorkState.STARTED))); assertThat(transaction.isActive(), is(true)); - transaction.commit(); + get(transaction.commitAsync()); assertThat( transaction.getState(), is(equalTo(com.google.cloud.spanner.connection.UnitOfWork.UnitOfWorkState.COMMITTED))); @@ -435,7 +428,7 @@ public void testState() { transaction.getState(), is(equalTo(com.google.cloud.spanner.connection.UnitOfWork.UnitOfWorkState.STARTED))); assertThat(transaction.isActive(), is(true)); - transaction.rollback(); + get(transaction.rollbackAsync()); assertThat( transaction.getState(), is(equalTo(com.google.cloud.spanner.connection.UnitOfWork.UnitOfWorkState.ROLLED_BACK))); @@ -446,12 +439,13 @@ public void testState() { transaction.getState(), is(equalTo(com.google.cloud.spanner.connection.UnitOfWork.UnitOfWorkState.STARTED))); assertThat(transaction.isActive(), is(true)); - assertThat(transaction.executeQuery(parsedStatement, AnalyzeMode.NONE), is(notNullValue())); + assertThat( + get(transaction.executeQueryAsync(parsedStatement, AnalyzeMode.NONE)), is(notNullValue())); assertThat( transaction.getState(), is(equalTo(com.google.cloud.spanner.connection.UnitOfWork.UnitOfWorkState.STARTED))); assertThat(transaction.isActive(), is(true)); - transaction.rollback(); + get(transaction.rollbackAsync()); assertThat( transaction.getState(), is(equalTo(com.google.cloud.spanner.connection.UnitOfWork.UnitOfWorkState.ROLLED_BACK))); diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ReadWriteTransactionTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ReadWriteTransactionTest.java index 1a332ab438..1e094eaeb6 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ReadWriteTransactionTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ReadWriteTransactionTest.java @@ -16,6 +16,7 @@ package com.google.cloud.spanner.connection; +import static com.google.cloud.spanner.SpannerApiFutures.get; import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.CoreMatchers.is; import static org.hamcrest.CoreMatchers.not; @@ -172,7 +173,7 @@ public void testExecuteDdl() { ReadWriteTransaction transaction = createSubject(); try { - transaction.executeDdl(statement); + transaction.executeDdlAsync(statement); fail("Expected exception"); } catch (SpannerException ex) { assertEquals(ErrorCode.FAILED_PRECONDITION, ex.getErrorCode()); @@ -183,7 +184,7 @@ public void testExecuteDdl() { public void testRunBatch() { ReadWriteTransaction subject = createSubject(); try { - subject.runBatch(); + subject.runBatchAsync(); fail("Expected exception"); } catch (SpannerException ex) { assertEquals(ErrorCode.FAILED_PRECONDITION, ex.getErrorCode()); @@ -210,7 +211,7 @@ public void testExecuteQuery() { when(parsedStatement.getStatement()).thenReturn(statement); ReadWriteTransaction transaction = createSubject(); - ResultSet rs = transaction.executeQuery(parsedStatement, AnalyzeMode.NONE); + ResultSet rs = get(transaction.executeQueryAsync(parsedStatement, AnalyzeMode.NONE)); assertThat(rs, is(notNullValue())); assertThat(rs.getStats(), is(nullValue())); } @@ -224,7 +225,7 @@ public void testPlanQuery() { when(parsedStatement.getStatement()).thenReturn(statement); ReadWriteTransaction transaction = createSubject(); - ResultSet rs = transaction.executeQuery(parsedStatement, AnalyzeMode.PLAN); + ResultSet rs = get(transaction.executeQueryAsync(parsedStatement, AnalyzeMode.PLAN)); assertThat(rs, is(notNullValue())); while (rs.next()) { // do nothing @@ -241,7 +242,7 @@ public void testProfileQuery() { when(parsedStatement.getStatement()).thenReturn(statement); ReadWriteTransaction transaction = createSubject(); - ResultSet rs = transaction.executeQuery(parsedStatement, AnalyzeMode.PROFILE); + ResultSet rs = get(transaction.executeQueryAsync(parsedStatement, AnalyzeMode.PROFILE)); assertThat(rs, is(notNullValue())); while (rs.next()) { // do nothing @@ -258,7 +259,7 @@ public void testExecuteUpdate() { when(parsedStatement.getStatement()).thenReturn(statement); ReadWriteTransaction transaction = createSubject(); - assertThat(transaction.executeUpdate(parsedStatement), is(1L)); + assertThat(get(transaction.executeUpdateAsync(parsedStatement)), is(1L)); } @Test @@ -270,7 +271,7 @@ public void testGetCommitTimestampBeforeCommit() { when(parsedStatement.getStatement()).thenReturn(statement); ReadWriteTransaction transaction = createSubject(); - assertThat(transaction.executeUpdate(parsedStatement), is(1L)); + assertThat(get(transaction.executeUpdateAsync(parsedStatement)), is(1L)); try { transaction.getCommitTimestamp(); fail("Expected exception"); @@ -288,8 +289,8 @@ public void testGetCommitTimestampAfterCommit() { when(parsedStatement.getStatement()).thenReturn(statement); ReadWriteTransaction transaction = createSubject(); - assertThat(transaction.executeUpdate(parsedStatement), is(1L)); - transaction.commit(); + assertThat(get(transaction.executeUpdateAsync(parsedStatement)), is(1L)); + get(transaction.commitAsync()); assertThat(transaction.getCommitTimestamp(), is(notNullValue())); } @@ -303,7 +304,8 @@ public void testGetReadTimestamp() { when(parsedStatement.getStatement()).thenReturn(statement); ReadWriteTransaction transaction = createSubject(); - assertThat(transaction.executeQuery(parsedStatement, AnalyzeMode.NONE), is(notNullValue())); + assertThat( + get(transaction.executeQueryAsync(parsedStatement, AnalyzeMode.NONE)), is(notNullValue())); try { transaction.getReadTimestamp(); fail("Expected exception"); @@ -325,13 +327,14 @@ public void testState() { transaction.getState(), is(equalTo(com.google.cloud.spanner.connection.UnitOfWork.UnitOfWorkState.STARTED))); assertThat(transaction.isActive(), is(true)); - assertThat(transaction.executeQuery(parsedStatement, AnalyzeMode.NONE), is(notNullValue())); + assertThat( + get(transaction.executeQueryAsync(parsedStatement, AnalyzeMode.NONE)), is(notNullValue())); assertThat( transaction.getState(), is(equalTo(com.google.cloud.spanner.connection.UnitOfWork.UnitOfWorkState.STARTED))); assertThat(transaction.isActive(), is(true)); - transaction.commit(); + get(transaction.commitAsync()); assertThat( transaction.getState(), is(equalTo(com.google.cloud.spanner.connection.UnitOfWork.UnitOfWorkState.COMMITTED))); @@ -343,7 +346,7 @@ public void testState() { transaction.getState(), is(equalTo(com.google.cloud.spanner.connection.UnitOfWork.UnitOfWorkState.STARTED))); assertThat(transaction.isActive(), is(true)); - transaction.rollback(); + get(transaction.rollbackAsync()); assertThat( transaction.getState(), is(equalTo(com.google.cloud.spanner.connection.UnitOfWork.UnitOfWorkState.ROLLED_BACK))); @@ -356,7 +359,7 @@ public void testState() { is(equalTo(com.google.cloud.spanner.connection.UnitOfWork.UnitOfWorkState.STARTED))); assertThat(transaction.isActive(), is(true)); try { - transaction.commit(); + get(transaction.commitAsync()); } catch (SpannerException e) { // ignore } @@ -372,7 +375,7 @@ public void testState() { is(equalTo(com.google.cloud.spanner.connection.UnitOfWork.UnitOfWorkState.STARTED))); assertThat(transaction.isActive(), is(true)); try { - transaction.commit(); + get(transaction.commitAsync()); } catch (AbortedException e) { // ignore } @@ -388,7 +391,7 @@ public void testState() { transaction.getState(), is(equalTo(com.google.cloud.spanner.connection.UnitOfWork.UnitOfWorkState.STARTED))); assertThat(transaction.isActive(), is(true)); - transaction.commit(); + get(transaction.commitAsync()); assertThat( transaction.getState(), is(equalTo(com.google.cloud.spanner.connection.UnitOfWork.UnitOfWorkState.COMMITTED))); @@ -452,11 +455,11 @@ public void testRetry() { .setDatabaseClient(client) .withStatementExecutor(new StatementExecutor()) .build(); - subject.executeUpdate(update1); - subject.executeUpdate(update2); + subject.executeUpdateAsync(update1); + subject.executeUpdateAsync(update2); boolean expectedException = false; try { - subject.commit(); + get(subject.commitAsync()); } catch (SpannerException e) { if (results == RetryResults.DIFFERENT && e.getErrorCode() == ErrorCode.ABORTED) { // expected diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/SingleUseTransactionTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/SingleUseTransactionTest.java index e73eb8e0b2..76ef62a21a 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/SingleUseTransactionTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/SingleUseTransactionTest.java @@ -16,6 +16,7 @@ package com.google.cloud.spanner.connection; +import static com.google.cloud.spanner.SpannerApiFutures.get; import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.fail; import static org.mockito.Matchers.anyListOf; @@ -45,12 +46,15 @@ import com.google.cloud.spanner.TransactionContext; import com.google.cloud.spanner.TransactionManager; import com.google.cloud.spanner.TransactionRunner; +import com.google.cloud.spanner.connection.StatementExecutor.StatementTimeout; import com.google.cloud.spanner.connection.StatementParser.ParsedStatement; import com.google.cloud.spanner.connection.StatementParser.StatementType; +import com.google.common.base.Preconditions; import com.google.spanner.admin.database.v1.UpdateDatabaseDdlMetadata; import com.google.spanner.v1.ResultSetStats; import java.util.Arrays; import java.util.Calendar; +import java.util.Collections; import java.util.List; import java.util.concurrent.TimeUnit; import org.junit.Test; @@ -77,6 +81,20 @@ private enum CommitBehavior { ABORT; } + /** Creates a {@link StatementTimeout} that will never timeout. */ + static StatementTimeout nullTimeout() { + return new StatementTimeout(); + } + + /** Creates a {@link StatementTimeout} with the given duration. */ + static StatementTimeout timeout(long timeout, TimeUnit unit) { + Preconditions.checkArgument(timeout > 0L); + Preconditions.checkArgument(StatementTimeout.isValidTimeoutUnit(unit)); + StatementTimeout res = new StatementTimeout(); + res.setTimeoutValue(timeout, unit); + return res; + } + private static class SimpleTransactionManager implements TransactionManager { private TransactionState state; private Timestamp commitTimestamp; @@ -287,16 +305,6 @@ private SingleUseTransaction createSubject() { 0L); } - private SingleUseTransaction createSubjectWithTimeout(long timeout) { - return createSubject( - createDefaultMockDdlClient(), - false, - TimestampBound.strong(), - AutocommitDmlMode.TRANSACTIONAL, - CommitBehavior.SUCCEED, - timeout); - } - private SingleUseTransaction createSubject(AutocommitDmlMode dmlMode) { return createSubject( createDefaultMockDdlClient(), @@ -349,7 +357,7 @@ private SingleUseTransaction createSubject( new SimpleReadOnlyTransaction(staleness); when(dbClient.singleUseReadOnlyTransaction(staleness)).thenReturn(singleUse); - TransactionContext txContext = mock(TransactionContext.class); + final TransactionContext txContext = mock(TransactionContext.class); when(txContext.executeUpdate(Statement.of(VALID_UPDATE))).thenReturn(VALID_UPDATE_COUNT); when(txContext.executeUpdate(Statement.of(SLOW_UPDATE))) .thenAnswer( @@ -381,12 +389,17 @@ public TransactionRunner answer(InvocationOnMock invocation) { new TransactionRunner() { private Timestamp commitTimestamp; - @SuppressWarnings("unchecked") @Override public T run(TransactionCallable callable) { if (commitBehavior == CommitBehavior.SUCCEED) { + T res; + try { + res = callable.run(txContext); + } catch (Exception e) { + throw SpannerExceptionFactory.newSpannerException(e); + } this.commitTimestamp = Timestamp.now(); - return (T) Long.valueOf(1L); + return res; } else if (commitBehavior == CommitBehavior.FAIL) { throw SpannerExceptionFactory.newSpannerException( ErrorCode.UNKNOWN, "commit failed"); @@ -420,9 +433,7 @@ public TransactionRunner allowNestedTransaction() { .setReadOnly(readOnly) .setReadOnlyStaleness(staleness) .setStatementTimeout( - timeout == 0L - ? StatementExecutor.StatementTimeout.nullTimeout() - : StatementExecutor.StatementTimeout.of(timeout, TimeUnit.MILLISECONDS)) + timeout == 0L ? nullTimeout() : timeout(timeout, TimeUnit.MILLISECONDS)) .withStatementExecutor(executor) .build(); } @@ -464,7 +475,7 @@ private List getTestTimestampBounds() { public void testCommit() { SingleUseTransaction subject = createSubject(); try { - subject.commit(); + subject.commitAsync(); fail("missing expected exception"); } catch (SpannerException e) { assertThat(e.getErrorCode()).isEqualTo(ErrorCode.FAILED_PRECONDITION); @@ -475,7 +486,7 @@ public void testCommit() { public void testRollback() { SingleUseTransaction subject = createSubject(); try { - subject.rollback(); + subject.rollbackAsync(); fail("missing expected exception"); } catch (SpannerException e) { assertThat(e.getErrorCode()).isEqualTo(ErrorCode.FAILED_PRECONDITION); @@ -486,7 +497,7 @@ public void testRollback() { public void testRunBatch() { SingleUseTransaction subject = createSubject(); try { - subject.runBatch(); + subject.runBatchAsync(); fail("missing expected exception"); } catch (SpannerException e) { assertThat(e.getErrorCode()).isEqualTo(ErrorCode.FAILED_PRECONDITION); @@ -510,7 +521,7 @@ public void testExecuteDdl() { ParsedStatement ddl = createParsedDdl(sql); DdlClient ddlClient = createDefaultMockDdlClient(); SingleUseTransaction subject = createDdlSubject(ddlClient); - subject.executeDdl(ddl); + get(subject.executeDdlAsync(ddl)); verify(ddlClient).executeDdl(sql); } @@ -519,7 +530,7 @@ public void testExecuteQuery() { for (TimestampBound staleness : getTestTimestampBounds()) { for (AnalyzeMode analyzeMode : AnalyzeMode.values()) { SingleUseTransaction subject = createReadOnlySubject(staleness); - ResultSet rs = subject.executeQuery(createParsedQuery(VALID_QUERY), analyzeMode); + ResultSet rs = get(subject.executeQueryAsync(createParsedQuery(VALID_QUERY), analyzeMode)); assertThat(rs).isNotNull(); assertThat(subject.getReadTimestamp()).isNotNull(); assertThat(subject.getState()) @@ -537,7 +548,7 @@ public void testExecuteQuery() { for (TimestampBound staleness : getTestTimestampBounds()) { SingleUseTransaction subject = createReadOnlySubject(staleness); try { - subject.executeQuery(createParsedQuery(INVALID_QUERY), AnalyzeMode.NONE); + get(subject.executeQueryAsync(createParsedQuery(INVALID_QUERY), AnalyzeMode.NONE)); fail("missing expected exception"); } catch (SpannerException e) { assertThat(e.getErrorCode()).isEqualTo(ErrorCode.UNKNOWN); @@ -570,14 +581,15 @@ public void testExecuteQueryWithOptionsTest() { .withStatementExecutor(executor) .setReadOnlyStaleness(TimestampBound.strong()) .build(); - assertThat(transaction.executeQuery(parsedStatement, AnalyzeMode.NONE, option)).isNotNull(); + assertThat(get(transaction.executeQueryAsync(parsedStatement, AnalyzeMode.NONE, option))) + .isNotNull(); } @Test public void testExecuteUpdate_Transactional_Valid() { ParsedStatement update = createParsedUpdate(VALID_UPDATE); SingleUseTransaction subject = createSubject(); - long updateCount = subject.executeUpdate(update); + long updateCount = get(subject.executeUpdateAsync(update)); assertThat(updateCount).isEqualTo(VALID_UPDATE_COUNT); assertThat(subject.getCommitTimestamp()).isNotNull(); assertThat(subject.getState()) @@ -589,7 +601,7 @@ public void testExecuteUpdate_Transactional_Invalid() { ParsedStatement update = createParsedUpdate(INVALID_UPDATE); SingleUseTransaction subject = createSubject(); try { - subject.executeUpdate(update); + get(subject.executeUpdateAsync(update)); fail("missing expected exception"); } catch (SpannerException e) { assertThat(e.getErrorCode()).isEqualTo(ErrorCode.UNKNOWN); @@ -602,7 +614,7 @@ public void testExecuteUpdate_Transactional_Valid_FailedCommit() { ParsedStatement update = createParsedUpdate(VALID_UPDATE); SingleUseTransaction subject = createSubject(CommitBehavior.FAIL); try { - subject.executeUpdate(update); + get(subject.executeUpdateAsync(update)); fail("missing expected exception"); } catch (SpannerException e) { assertThat(e.getErrorCode()).isEqualTo(ErrorCode.UNKNOWN); @@ -610,23 +622,11 @@ public void testExecuteUpdate_Transactional_Valid_FailedCommit() { } } - @Test - public void testExecuteUpdate_Transactional_Valid_AbortedCommit() { - ParsedStatement update = createParsedUpdate(VALID_UPDATE); - SingleUseTransaction subject = createSubject(CommitBehavior.ABORT); - // even though the transaction aborts at first, it will be retried and eventually succeed - long updateCount = subject.executeUpdate(update); - assertThat(updateCount).isEqualTo(VALID_UPDATE_COUNT); - assertThat(subject.getCommitTimestamp()).isNotNull(); - assertThat(subject.getState()) - .isEqualTo(com.google.cloud.spanner.connection.UnitOfWork.UnitOfWorkState.COMMITTED); - } - @Test public void testExecuteUpdate_Partitioned_Valid() { ParsedStatement update = createParsedUpdate(VALID_UPDATE); SingleUseTransaction subject = createSubject(AutocommitDmlMode.PARTITIONED_NON_ATOMIC); - long updateCount = subject.executeUpdate(update); + long updateCount = get(subject.executeUpdateAsync(update)); assertThat(updateCount).isEqualTo(VALID_UPDATE_COUNT); assertThat(subject.getState()) .isEqualTo(com.google.cloud.spanner.connection.UnitOfWork.UnitOfWorkState.COMMITTED); @@ -637,7 +637,7 @@ public void testExecuteUpdate_Partitioned_Invalid() { ParsedStatement update = createParsedUpdate(INVALID_UPDATE); SingleUseTransaction subject = createSubject(AutocommitDmlMode.PARTITIONED_NON_ATOMIC); try { - subject.executeUpdate(update); + get(subject.executeUpdateAsync(update)); fail("missing expected exception"); } catch (SpannerException e) { assertThat(e.getErrorCode()).isEqualTo(ErrorCode.UNKNOWN); @@ -645,32 +645,11 @@ public void testExecuteUpdate_Partitioned_Invalid() { } } - @Test - public void testWrite() { - SingleUseTransaction subject = createSubject(); - subject.write(Mutation.newInsertBuilder("FOO").build()); - assertThat(subject.getCommitTimestamp()).isNotNull(); - assertThat(subject.getState()) - .isEqualTo(com.google.cloud.spanner.connection.UnitOfWork.UnitOfWorkState.COMMITTED); - } - - @Test - public void testWriteFail() { - SingleUseTransaction subject = createSubject(CommitBehavior.FAIL); - try { - subject.write(Mutation.newInsertBuilder("FOO").build()); - fail("missing expected exception"); - } catch (SpannerException e) { - assertThat(e.getErrorCode()).isEqualTo(ErrorCode.UNKNOWN); - assertThat(e.getMessage()).contains("commit failed"); - } - } - @Test public void testWriteIterable() { SingleUseTransaction subject = createSubject(); Mutation mutation = Mutation.newInsertBuilder("FOO").build(); - subject.write(Arrays.asList(mutation, mutation)); + get(subject.writeAsync(Arrays.asList(mutation, mutation))); assertThat(subject.getCommitTimestamp()).isNotNull(); assertThat(subject.getState()) .isEqualTo(com.google.cloud.spanner.connection.UnitOfWork.UnitOfWorkState.COMMITTED); @@ -681,7 +660,7 @@ public void testWriteIterableFail() { SingleUseTransaction subject = createSubject(CommitBehavior.FAIL); Mutation mutation = Mutation.newInsertBuilder("FOO").build(); try { - subject.write(Arrays.asList(mutation, mutation)); + get(subject.writeAsync(Arrays.asList(mutation, mutation))); fail("missing expected exception"); } catch (SpannerException e) { assertThat(e.getErrorCode()).isEqualTo(ErrorCode.UNKNOWN); @@ -693,11 +672,12 @@ public void testWriteIterableFail() { public void testMultiUse() { for (TimestampBound staleness : getTestTimestampBounds()) { SingleUseTransaction subject = createReadOnlySubject(staleness); - ResultSet rs = subject.executeQuery(createParsedQuery(VALID_QUERY), AnalyzeMode.NONE); + ResultSet rs = + get(subject.executeQueryAsync(createParsedQuery(VALID_QUERY), AnalyzeMode.NONE)); assertThat(rs).isNotNull(); assertThat(subject.getReadTimestamp()).isNotNull(); try { - subject.executeQuery(createParsedQuery(VALID_QUERY), AnalyzeMode.NONE); + get(subject.executeQueryAsync(createParsedQuery(VALID_QUERY), AnalyzeMode.NONE)); fail("missing expected exception"); } catch (IllegalStateException e) { } @@ -707,81 +687,42 @@ public void testMultiUse() { ParsedStatement ddl = createParsedDdl(sql); DdlClient ddlClient = createDefaultMockDdlClient(); SingleUseTransaction subject = createDdlSubject(ddlClient); - subject.executeDdl(ddl); + get(subject.executeDdlAsync(ddl)); verify(ddlClient).executeDdl(sql); try { - subject.executeDdl(ddl); + get(subject.executeDdlAsync(ddl)); fail("missing expected exception"); } catch (IllegalStateException e) { } ParsedStatement update = createParsedUpdate(VALID_UPDATE); subject = createSubject(); - long updateCount = subject.executeUpdate(update); + long updateCount = get(subject.executeUpdateAsync(update)); assertThat(updateCount).isEqualTo(VALID_UPDATE_COUNT); assertThat(subject.getCommitTimestamp()).isNotNull(); try { - subject.executeUpdate(update); + get(subject.executeUpdateAsync(update)); fail("missing expected exception"); } catch (IllegalStateException e) { } subject = createSubject(); - subject.write(Mutation.newInsertBuilder("FOO").build()); + get(subject.writeAsync(Collections.singleton(Mutation.newInsertBuilder("FOO").build()))); assertThat(subject.getCommitTimestamp()).isNotNull(); try { - subject.write(Mutation.newInsertBuilder("FOO").build()); + get(subject.writeAsync(Collections.singleton(Mutation.newInsertBuilder("FOO").build()))); fail("missing expected exception"); } catch (IllegalStateException e) { } subject = createSubject(); Mutation mutation = Mutation.newInsertBuilder("FOO").build(); - subject.write(Arrays.asList(mutation, mutation)); + get(subject.writeAsync(Arrays.asList(mutation, mutation))); assertThat(subject.getCommitTimestamp()).isNotNull(); try { - subject.write(Arrays.asList(mutation, mutation)); + get(subject.writeAsync(Arrays.asList(mutation, mutation))); fail("missing expected exception"); } catch (IllegalStateException e) { } } - - @Test - public void testExecuteQueryWithTimeout() { - SingleUseTransaction subject = createSubjectWithTimeout(1L); - try { - subject.executeQuery(createParsedQuery(SLOW_QUERY), AnalyzeMode.NONE); - } catch (SpannerException e) { - if (e.getErrorCode() != ErrorCode.DEADLINE_EXCEEDED) { - throw e; - } - } - assertThat(subject.getState()) - .isEqualTo(com.google.cloud.spanner.connection.UnitOfWork.UnitOfWorkState.COMMIT_FAILED); - try { - subject.getReadTimestamp(); - fail("missing expected exception"); - } catch (SpannerException e) { - assertThat(e.getErrorCode()).isEqualTo(ErrorCode.FAILED_PRECONDITION); - } - } - - @Test - public void testExecuteUpdateWithTimeout() { - SingleUseTransaction subject = createSubjectWithTimeout(1L); - try { - subject.executeUpdate(createParsedUpdate(SLOW_UPDATE)); - fail("missing expected exception"); - } catch (SpannerException e) { - assertThat(e.getErrorCode()).isEqualTo(ErrorCode.DEADLINE_EXCEEDED); - } - assertThat(subject.getState()) - .isEqualTo(com.google.cloud.spanner.connection.UnitOfWork.UnitOfWorkState.COMMIT_FAILED); - try { - subject.getCommitTimestamp(); - fail("missing expected exception"); - } catch (SpannerException e) { - assertThat(e.getErrorCode()).isEqualTo(ErrorCode.FAILED_PRECONDITION); - } - } } diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/SpannerPoolTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/SpannerPoolTest.java index 3c0e9cf160..afc0512b4e 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/SpannerPoolTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/SpannerPoolTest.java @@ -33,9 +33,11 @@ import com.google.cloud.spanner.SpannerException; import com.google.cloud.spanner.connection.ConnectionImpl.LeakedConnectionException; import com.google.cloud.spanner.connection.SpannerPool.CheckAndCloseSpannersMode; -import com.google.cloud.spanner.connection.SpannerPool.SpannerPoolKey; +import com.google.common.base.Ticker; +import com.google.common.testing.FakeTicker; import java.io.ByteArrayOutputStream; import java.io.OutputStream; +import java.util.concurrent.TimeUnit; import java.util.logging.Handler; import java.util.logging.Logger; import java.util.logging.StreamHandler; @@ -62,12 +64,13 @@ public class SpannerPoolTest { private ConnectionOptions options6 = mock(ConnectionOptions.class); private SpannerPool createSubjectAndMocks() { - return createSubjectAndMocks(0L); + return createSubjectAndMocks(0L, Ticker.systemTicker()); } - private SpannerPool createSubjectAndMocks(long closeSpannerAfterMillisecondsUnused) { + private SpannerPool createSubjectAndMocks( + long closeSpannerAfterMillisecondsUnused, Ticker ticker) { SpannerPool pool = - new SpannerPool(closeSpannerAfterMillisecondsUnused) { + new SpannerPool(closeSpannerAfterMillisecondsUnused, ticker) { @Override Spanner createSpanner(SpannerPoolKey key, ConnectionOptions options) { return mock(Spanner.class); @@ -340,73 +343,77 @@ public void testCloseUnusedSpanners() { verify(spanner3).close(); } - /** Allow the automatic close test to be run multiple times to ensure it is stable */ - private static final int NUMBER_OF_AUTOMATIC_CLOSE_TEST_RUNS = 1; - - private static final long TEST_AUTOMATIC_CLOSE_TIMEOUT = 2L; - private static final long SLEEP_BEFORE_VERIFICATION = 100L; + private static final long TEST_AUTOMATIC_CLOSE_TIMEOUT_MILLIS = 60_000L; + private static final long TEST_AUTOMATIC_CLOSE_TIMEOUT_NANOS = + TimeUnit.NANOSECONDS.convert(TEST_AUTOMATIC_CLOSE_TIMEOUT_MILLIS, TimeUnit.MILLISECONDS); + private static final long MILLISECOND = TimeUnit.NANOSECONDS.convert(1L, TimeUnit.MILLISECONDS); @Test public void testAutomaticCloser() throws InterruptedException { - for (int testRun = 0; testRun < NUMBER_OF_AUTOMATIC_CLOSE_TEST_RUNS; testRun++) { - // create a pool that will close unused spanners after 5 milliseconds - SpannerPool pool = createSubjectAndMocks(TEST_AUTOMATIC_CLOSE_TIMEOUT); - Spanner spanner1; - Spanner spanner2; - Spanner spanner3; - - // create two connections that use the same Spanner - spanner1 = pool.getSpanner(options1, connection1); - spanner2 = pool.getSpanner(options1, connection2); - assertThat(spanner1, is(equalTo(spanner2))); - - // all spanners are in use, this should have no effect - Thread.sleep(SLEEP_BEFORE_VERIFICATION); - verify(spanner1, never()).close(); - - // close one connection. This should also have no effect. - pool.removeConnection(options1, connection1); - Thread.sleep(SLEEP_BEFORE_VERIFICATION); - verify(spanner1, never()).close(); - - // close the other connection as well, the Spanner object should now be closed. - pool.removeConnection(options1, connection2); - Thread.sleep(SLEEP_BEFORE_VERIFICATION); - verify(spanner1).close(); - - // create three connections that use two different Spanners - spanner1 = pool.getSpanner(options1, connection1); - spanner2 = pool.getSpanner(options2, connection2); - spanner3 = pool.getSpanner(options2, connection3); - assertThat(spanner1, not(equalTo(spanner2))); - assertThat(spanner2, is(equalTo(spanner3))); - - // all spanners are in use, this should have no effect - Thread.sleep(SLEEP_BEFORE_VERIFICATION); - verify(spanner1, never()).close(); - verify(spanner2, never()).close(); - verify(spanner3, never()).close(); - - // close connection1. That should also mark spanner1 as no longer in use - pool.removeConnection(options1, connection1); - Thread.sleep(SLEEP_BEFORE_VERIFICATION); - verify(spanner1).close(); - verify(spanner2, never()).close(); - verify(spanner3, never()).close(); - - // close connection2. That should have no effect, as connection3 is still using spanner2 - pool.removeConnection(options2, connection2); - Thread.sleep(SLEEP_BEFORE_VERIFICATION); - verify(spanner1).close(); - verify(spanner2, never()).close(); - verify(spanner3, never()).close(); - - // close connection3. Now all should be closed. - pool.removeConnection(options2, connection3); - Thread.sleep(SLEEP_BEFORE_VERIFICATION); - verify(spanner1).close(); - verify(spanner2).close(); - verify(spanner3).close(); - } + FakeTicker ticker = new FakeTicker(); + SpannerPool pool = createSubjectAndMocks(TEST_AUTOMATIC_CLOSE_TIMEOUT_MILLIS, ticker); + Spanner spanner1; + Spanner spanner2; + Spanner spanner3; + + // create two connections that use the same Spanner + spanner1 = pool.getSpanner(options1, connection1); + spanner2 = pool.getSpanner(options1, connection2); + assertThat(spanner1, is(equalTo(spanner2))); + + // all spanners are in use, this should have no effect + ticker.advance(TEST_AUTOMATIC_CLOSE_TIMEOUT_NANOS + MILLISECOND); + pool.closeUnusedSpanners(TEST_AUTOMATIC_CLOSE_TIMEOUT_MILLIS); + verify(spanner1, never()).close(); + + // close one connection. This should also have no effect. + pool.removeConnection(options1, connection1); + ticker.advance(TEST_AUTOMATIC_CLOSE_TIMEOUT_NANOS + MILLISECOND); + pool.closeUnusedSpanners(TEST_AUTOMATIC_CLOSE_TIMEOUT_MILLIS); + verify(spanner1, never()).close(); + + // close the other connection as well, the Spanner object should now be closed. + pool.removeConnection(options1, connection2); + ticker.advance(TEST_AUTOMATIC_CLOSE_TIMEOUT_NANOS + MILLISECOND); + pool.closeUnusedSpanners(TEST_AUTOMATIC_CLOSE_TIMEOUT_MILLIS); + verify(spanner1).close(); + + // create three connections that use two different Spanners + spanner1 = pool.getSpanner(options1, connection1); + spanner2 = pool.getSpanner(options2, connection2); + spanner3 = pool.getSpanner(options2, connection3); + assertThat(spanner1, not(equalTo(spanner2))); + assertThat(spanner2, is(equalTo(spanner3))); + + // all spanners are in use, this should have no effect + ticker.advance(TEST_AUTOMATIC_CLOSE_TIMEOUT_NANOS + MILLISECOND); + pool.closeUnusedSpanners(TEST_AUTOMATIC_CLOSE_TIMEOUT_MILLIS); + verify(spanner1, never()).close(); + verify(spanner2, never()).close(); + verify(spanner3, never()).close(); + + // close connection1. That should also mark spanner1 as no longer in use + pool.removeConnection(options1, connection1); + ticker.advance(TEST_AUTOMATIC_CLOSE_TIMEOUT_NANOS + MILLISECOND); + pool.closeUnusedSpanners(TEST_AUTOMATIC_CLOSE_TIMEOUT_MILLIS); + verify(spanner1).close(); + verify(spanner2, never()).close(); + verify(spanner3, never()).close(); + + // close connection2. That should have no effect, as connection3 is still using spanner2 + pool.removeConnection(options2, connection2); + ticker.advance(TEST_AUTOMATIC_CLOSE_TIMEOUT_NANOS + MILLISECOND); + pool.closeUnusedSpanners(TEST_AUTOMATIC_CLOSE_TIMEOUT_MILLIS); + verify(spanner1).close(); + verify(spanner2, never()).close(); + verify(spanner3, never()).close(); + + // close connection3. Now all should be closed. + pool.removeConnection(options2, connection3); + ticker.advance(TEST_AUTOMATIC_CLOSE_TIMEOUT_NANOS + MILLISECOND); + pool.closeUnusedSpanners(TEST_AUTOMATIC_CLOSE_TIMEOUT_MILLIS); + verify(spanner1).close(); + verify(spanner2).close(); + verify(spanner3).close(); } } diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/StatementTimeoutTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/StatementTimeoutTest.java index e483a50279..eac4c38d17 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/StatementTimeoutTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/StatementTimeoutTest.java @@ -22,59 +22,53 @@ import static org.hamcrest.MatcherAssert.assertThat; import static org.junit.Assert.assertEquals; import static org.junit.Assert.fail; -import static org.mockito.Matchers.any; -import static org.mockito.Matchers.anyListOf; -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.api.gax.longrunning.OperationFuture; -import com.google.cloud.NoCredentials; -import com.google.cloud.spanner.DatabaseClient; + +import com.google.api.gax.longrunning.OperationTimedPollAlgorithm; +import com.google.api.gax.retrying.RetrySettings; import com.google.cloud.spanner.ErrorCode; -import com.google.cloud.spanner.ReadOnlyTransaction; +import com.google.cloud.spanner.MockSpannerServiceImpl.SimulatedExecutionTime; import com.google.cloud.spanner.ResultSet; -import com.google.cloud.spanner.Spanner; import com.google.cloud.spanner.SpannerException; import com.google.cloud.spanner.SpannerExceptionFactory; +import com.google.cloud.spanner.SpannerOptions.Builder; 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.TransactionManager.TransactionState; import com.google.cloud.spanner.connection.AbstractConnectionImplTest.ConnectionConsumer; +import com.google.cloud.spanner.connection.ConnectionOptions.SpannerOptionsConfigurator; +import com.google.cloud.spanner.connection.ITAbstractSpannerTest.ITConnection; +import com.google.common.util.concurrent.Uninterruptibles; +import com.google.longrunning.Operation; +import com.google.protobuf.AbstractMessage; +import com.google.protobuf.Any; +import com.google.protobuf.Empty; import com.google.spanner.admin.database.v1.UpdateDatabaseDdlMetadata; -import java.util.Arrays; +import com.google.spanner.v1.CommitRequest; +import com.google.spanner.v1.ExecuteSqlRequest; +import io.grpc.Status; import java.util.concurrent.Callable; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import org.junit.After; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; -import org.mockito.Matchers; -import org.mockito.invocation.InvocationOnMock; -import org.mockito.stubbing.Answer; +import org.threeten.bp.Duration; @RunWith(JUnit4.class) -public class StatementTimeoutTest { +public class StatementTimeoutTest extends AbstractMockServerTest { - private static final String URI = - "cloudspanner:/projects/test-project-123/instances/test-instance/databases/test-database"; private static final String SLOW_SELECT = "SELECT foo FROM bar"; private static final String INVALID_SELECT = "SELECT FROM bar"; // missing columns / * - private static final String FAST_SELECT = "SELECT fast_column FROM fast_table"; private static final String SLOW_DDL = "CREATE TABLE foo"; private static final String FAST_DDL = "CREATE TABLE fast_table"; private static final String SLOW_UPDATE = "UPDATE foo SET col1=1 WHERE id=2"; - private static final String FAST_UPDATE = "UPDATE fast_table SET foo=1 WHERE bar=2"; /** Execution time for statements that have been defined as slow. */ - private static final long EXECUTION_TIME_SLOW_STATEMENT = 10_000L; + private static final int EXECUTION_TIME_SLOW_STATEMENT = 10_000; /** * This timeout should be high enough that it will never be exceeded, even on a slow build * environment, but still significantly lower than the expected execution time of the slow @@ -87,193 +81,51 @@ public class StatementTimeoutTest { * still high enough that it would normally not be exceeded for a statement that is executed * directly. */ - private static final long TIMEOUT_FOR_SLOW_STATEMENTS = 20L; - /** - * The number of milliseconds to wait before cancelling a query should be high enough to not cause - * flakiness on a slow environment, but at the same time low enough that it does not slow down the - * test case unnecessarily. - */ - private static final int WAIT_BEFORE_CANCEL = 100; - - private enum CommitRollbackBehavior { - FAST, - SLOW_COMMIT, - SLOW_ROLLBACK; - } - - private static final class DelayedQueryExecution implements Answer { - @Override - public ResultSet answer(InvocationOnMock invocation) throws Throwable { - Thread.sleep(EXECUTION_TIME_SLOW_STATEMENT); - return mock(ResultSet.class); - } - } - - private DdlClient createDefaultMockDdlClient(final long waitForMillis) { - try { - DdlClient ddlClient = mock(DdlClient.class); - UpdateDatabaseDdlMetadata metadata = UpdateDatabaseDdlMetadata.getDefaultInstance(); - ApiFuture futureMetadata = ApiFutures.immediateFuture(metadata); - @SuppressWarnings("unchecked") - final OperationFuture operation = - mock(OperationFuture.class); - if (waitForMillis > 0L) { - when(operation.get()) - .thenAnswer( - new Answer() { + private static final int TIMEOUT_FOR_SLOW_STATEMENTS = 20; + + ITConnection createConnection() { + StringBuilder url = new StringBuilder(getBaseUrl()); + ConnectionOptions options = + ConnectionOptions.newBuilder() + .setUri(url.toString()) + .setConfigurator( + new SpannerOptionsConfigurator() { @Override - public Void answer(InvocationOnMock invocation) throws Throwable { - Thread.sleep(waitForMillis); - return null; + public void configure(Builder options) { + options + .getDatabaseAdminStubSettingsBuilder() + .updateDatabaseDdlOperationSettings() + .setPollingAlgorithm( + OperationTimedPollAlgorithm.create( + RetrySettings.newBuilder() + .setInitialRetryDelay(Duration.ofMillis(1L)) + .setMaxRetryDelay(Duration.ofMillis(1L)) + .setRetryDelayMultiplier(1.0) + .setTotalTimeout(Duration.ofMinutes(10L)) + .build())); } - }); - } else { - when(operation.get()).thenReturn(null); - } - when(operation.getMetadata()).thenReturn(futureMetadata); - when(ddlClient.executeDdl(SLOW_DDL)).thenCallRealMethod(); - when(ddlClient.executeDdl(anyListOf(String.class))).thenReturn(operation); - - @SuppressWarnings("unchecked") - final OperationFuture fastOperation = - mock(OperationFuture.class); - when(fastOperation.isDone()).thenReturn(true); - when(fastOperation.get()).thenReturn(null); - when(fastOperation.getMetadata()).thenReturn(futureMetadata); - when(ddlClient.executeDdl(FAST_DDL)).thenReturn(fastOperation); - when(ddlClient.executeDdl(Arrays.asList(FAST_DDL))).thenReturn(fastOperation); - return ddlClient; - } catch (Exception e) { - throw new RuntimeException(e); - } - } - - private ConnectionImpl createConnection(ConnectionOptions options) { - return createConnection(options, CommitRollbackBehavior.FAST); + }) + .build(); + return createITConnection(options); } - /** - * Creates a connection on which the statements {@link StatementTimeoutTest#SLOW_SELECT} and - * {@link StatementTimeoutTest#SLOW_DDL} will take at least 10,000 milliseconds - */ - private ConnectionImpl createConnection( - ConnectionOptions options, final CommitRollbackBehavior commitRollbackBehavior) { - DatabaseClient dbClient = mock(DatabaseClient.class); - Spanner spanner = mock(Spanner.class); - SpannerPool spannerPool = mock(SpannerPool.class); - when(spannerPool.getSpanner(any(ConnectionOptions.class), any(ConnectionImpl.class))) - .thenReturn(spanner); - DdlClient ddlClient = createDefaultMockDdlClient(EXECUTION_TIME_SLOW_STATEMENT); - final ResultSet invalidResultSet = mock(ResultSet.class); - when(invalidResultSet.next()) - .thenThrow( - SpannerExceptionFactory.newSpannerException( - ErrorCode.INVALID_ARGUMENT, "invalid query")); - - ReadOnlyTransaction singleUseReadOnlyTx = mock(ReadOnlyTransaction.class); - when(singleUseReadOnlyTx.executeQuery(Statement.of(SLOW_SELECT))) - .thenAnswer(new DelayedQueryExecution()); - when(singleUseReadOnlyTx.executeQuery(Statement.of(FAST_SELECT))) - .thenReturn(mock(ResultSet.class)); - when(singleUseReadOnlyTx.executeQuery(Statement.of(INVALID_SELECT))) - .thenReturn(invalidResultSet); - when(dbClient.singleUseReadOnlyTransaction(Matchers.any(TimestampBound.class))) - .thenReturn(singleUseReadOnlyTx); - - ReadOnlyTransaction readOnlyTx = mock(ReadOnlyTransaction.class); - when(readOnlyTx.executeQuery(Statement.of(SLOW_SELECT))) - .thenAnswer(new DelayedQueryExecution()); - when(readOnlyTx.executeQuery(Statement.of(FAST_SELECT))).thenReturn(mock(ResultSet.class)); - when(readOnlyTx.executeQuery(Statement.of(INVALID_SELECT))).thenReturn(invalidResultSet); - when(dbClient.readOnlyTransaction(Matchers.any(TimestampBound.class))).thenReturn(readOnlyTx); - - when(dbClient.transactionManager()) - .thenAnswer( - new Answer() { - @Override - public TransactionManager answer(InvocationOnMock invocation) { - TransactionManager txManager = mock(TransactionManager.class); - when(txManager.getState()).thenReturn(null, TransactionState.STARTED); - when(txManager.begin()) - .thenAnswer( - new Answer() { - @Override - public TransactionContext answer(InvocationOnMock invocation) { - TransactionContext txContext = mock(TransactionContext.class); - when(txContext.executeQuery(Statement.of(SLOW_SELECT))) - .thenAnswer(new DelayedQueryExecution()); - when(txContext.executeQuery(Statement.of(FAST_SELECT))) - .thenReturn(mock(ResultSet.class)); - when(txContext.executeQuery(Statement.of(INVALID_SELECT))) - .thenReturn(invalidResultSet); - when(txContext.executeUpdate(Statement.of(SLOW_UPDATE))) - .thenAnswer( - new Answer() { - @Override - public Long answer(InvocationOnMock invocation) - throws Throwable { - Thread.sleep(EXECUTION_TIME_SLOW_STATEMENT); - return 1L; - } - }); - when(txContext.executeUpdate(Statement.of(FAST_UPDATE))).thenReturn(1L); - return txContext; - } - }); - if (commitRollbackBehavior == CommitRollbackBehavior.SLOW_COMMIT) { - doAnswer( - new Answer() { - @Override - public Void answer(InvocationOnMock invocation) throws Throwable { - Thread.sleep(EXECUTION_TIME_SLOW_STATEMENT); - return null; - } - }) - .when(txManager) - .commit(); - } - if (commitRollbackBehavior == CommitRollbackBehavior.SLOW_ROLLBACK) { - doAnswer( - new Answer() { - @Override - public Void answer(InvocationOnMock invocation) throws Throwable { - Thread.sleep(EXECUTION_TIME_SLOW_STATEMENT); - return null; - } - }) - .when(txManager) - .rollback(); - } - - return txManager; - } - }); - when(dbClient.executePartitionedUpdate(Statement.of(FAST_UPDATE))).thenReturn(1L); - when(dbClient.executePartitionedUpdate(Statement.of(SLOW_UPDATE))) - .thenAnswer( - new Answer() { - @Override - public Long answer(InvocationOnMock invocation) throws Throwable { - Thread.sleep(EXECUTION_TIME_SLOW_STATEMENT); - return 1L; - } - }); - return new ConnectionImpl(options, spannerPool, ddlClient, dbClient); + @After + public void clearExecutionTimes() { + mockSpanner.removeAllExecutionTimes(); } @Test public void testTimeoutExceptionReadOnlyAutocommit() { - try (Connection connection = - createConnection( - ConnectionOptions.newBuilder() - .setCredentials(NoCredentials.getInstance()) - .setUri(URI) - .build())) { + mockSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.ofMinimumAndRandomTime(EXECUTION_TIME_SLOW_STATEMENT, 0)); + + try (Connection connection = createConnection()) { + connection.setAutocommit(true); connection.setReadOnly(true); connection.setStatementTimeout(TIMEOUT_FOR_SLOW_STATEMENTS, TimeUnit.MILLISECONDS); try { - connection.executeQuery(Statement.of(SLOW_SELECT)); - fail("Expected exception"); + connection.executeQuery(SELECT_RANDOM_STATEMENT); + fail("missing expected exception"); } catch (SpannerException ex) { assertEquals(ErrorCode.DEADLINE_EXCEEDED, ex.getErrorCode()); } @@ -282,44 +134,43 @@ public void testTimeoutExceptionReadOnlyAutocommit() { @Test public void testTimeoutExceptionReadOnlyAutocommitMultipleStatements() { - try (Connection connection = - createConnection( - ConnectionOptions.newBuilder() - .setCredentials(NoCredentials.getInstance()) - .setUri(URI) - .build())) { + mockSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.ofMinimumAndRandomTime(EXECUTION_TIME_SLOW_STATEMENT, 0)); + + try (Connection connection = createConnection()) { + connection.setAutocommit(true); connection.setReadOnly(true); connection.setStatementTimeout(TIMEOUT_FOR_SLOW_STATEMENTS, TimeUnit.MILLISECONDS); // assert that multiple statements after each other also time out for (int i = 0; i < 2; i++) { - boolean timedOut = false; try { - connection.executeQuery(Statement.of(SLOW_SELECT)); + connection.executeQuery(SELECT_RANDOM_STATEMENT); + fail("missing expected exception"); } catch (SpannerException e) { - timedOut = e.getErrorCode() == ErrorCode.DEADLINE_EXCEEDED; + assertEquals(ErrorCode.DEADLINE_EXCEEDED, e.getErrorCode()); } - assertThat(timedOut, is(true)); } // try to do a new query that is fast. + mockSpanner.removeAllExecutionTimes(); connection.setStatementTimeout(TIMEOUT_FOR_FAST_STATEMENTS, TimeUnit.MILLISECONDS); - assertThat(connection.executeQuery(Statement.of(FAST_SELECT)), is(notNullValue())); + try (ResultSet rs = connection.executeQuery(SELECT_RANDOM_STATEMENT)) { + assertThat(rs, is(notNullValue())); + } } } @Test public void testTimeoutExceptionReadOnlyTransactional() { - try (Connection connection = - createConnection( - ConnectionOptions.newBuilder() - .setCredentials(NoCredentials.getInstance()) - .setUri(URI) - .build())) { + mockSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.ofMinimumAndRandomTime(EXECUTION_TIME_SLOW_STATEMENT, 0)); + + try (Connection connection = createConnection()) { connection.setReadOnly(true); connection.setAutocommit(false); connection.setStatementTimeout(TIMEOUT_FOR_SLOW_STATEMENTS, TimeUnit.MILLISECONDS); try { - connection.executeQuery(Statement.of(SLOW_SELECT)); - fail("Expected exception"); + connection.executeQuery(SELECT_RANDOM_STATEMENT); + fail("missing expected exception"); } catch (SpannerException ex) { assertEquals(ErrorCode.DEADLINE_EXCEEDED, ex.getErrorCode()); } @@ -328,46 +179,45 @@ public void testTimeoutExceptionReadOnlyTransactional() { @Test public void testTimeoutExceptionReadOnlyTransactionMultipleStatements() { - try (Connection connection = - createConnection( - ConnectionOptions.newBuilder() - .setCredentials(NoCredentials.getInstance()) - .setUri(URI) - .build())) { + mockSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.ofMinimumAndRandomTime(EXECUTION_TIME_SLOW_STATEMENT, 0)); + + try (Connection connection = createConnection()) { connection.setReadOnly(true); connection.setAutocommit(false); connection.setStatementTimeout(TIMEOUT_FOR_SLOW_STATEMENTS, TimeUnit.MILLISECONDS); // assert that multiple statements after each other also time out for (int i = 0; i < 2; i++) { - boolean timedOut = false; try { - connection.executeQuery(Statement.of(SLOW_SELECT)); + connection.executeQuery(SELECT_RANDOM_STATEMENT); + fail("missing expected exception"); } catch (SpannerException e) { - timedOut = e.getErrorCode() == ErrorCode.DEADLINE_EXCEEDED; + assertEquals(ErrorCode.DEADLINE_EXCEEDED, e.getErrorCode()); } - assertThat(timedOut, is(true)); } // do a rollback without any chance of a timeout connection.clearStatementTimeout(); connection.rollback(); // try to do a new query that is fast. + mockSpanner.removeAllExecutionTimes(); connection.setStatementTimeout(TIMEOUT_FOR_FAST_STATEMENTS, TimeUnit.MILLISECONDS); - assertThat(connection.executeQuery(Statement.of(FAST_SELECT)), is(notNullValue())); + try (ResultSet rs = connection.executeQuery(SELECT_RANDOM_STATEMENT)) { + assertThat(rs, is(notNullValue())); + } } } @Test public void testTimeoutExceptionReadWriteAutocommit() { - try (Connection connection = - createConnection( - ConnectionOptions.newBuilder() - .setCredentials(NoCredentials.getInstance()) - .setUri(URI) - .build())) { + mockSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.ofMinimumAndRandomTime(EXECUTION_TIME_SLOW_STATEMENT, 0)); + + try (Connection connection = createConnection()) { + connection.setAutocommit(true); connection.setStatementTimeout(TIMEOUT_FOR_SLOW_STATEMENTS, TimeUnit.MILLISECONDS); try { - connection.executeQuery(Statement.of(SLOW_SELECT)); - fail("Expected exception"); + connection.executeQuery(SELECT_RANDOM_STATEMENT); + fail("missing expected exception"); } catch (SpannerException ex) { assertEquals(ErrorCode.DEADLINE_EXCEEDED, ex.getErrorCode()); } @@ -376,41 +226,41 @@ public void testTimeoutExceptionReadWriteAutocommit() { @Test public void testTimeoutExceptionReadWriteAutocommitMultipleStatements() { - try (Connection connection = - createConnection( - ConnectionOptions.newBuilder() - .setCredentials(NoCredentials.getInstance()) - .setUri(URI) - .build())) { + mockSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.ofMinimumAndRandomTime(EXECUTION_TIME_SLOW_STATEMENT, 0)); + + try (Connection connection = createConnection()) { + connection.setAutocommit(true); connection.setStatementTimeout(TIMEOUT_FOR_SLOW_STATEMENTS, TimeUnit.MILLISECONDS); // assert that multiple statements after each other also time out for (int i = 0; i < 2; i++) { - boolean timedOut = false; try { - connection.executeQuery(Statement.of(SLOW_SELECT)); + connection.executeQuery(SELECT_RANDOM_STATEMENT); + fail("missing expected exception"); } catch (SpannerException e) { - timedOut = e.getErrorCode() == ErrorCode.DEADLINE_EXCEEDED; + assertEquals(ErrorCode.DEADLINE_EXCEEDED, e.getErrorCode()); } - assertThat(timedOut, is(true)); } // try to do a new query that is fast. + mockSpanner.removeAllExecutionTimes(); connection.setStatementTimeout(TIMEOUT_FOR_FAST_STATEMENTS, TimeUnit.MILLISECONDS); - assertThat(connection.executeQuery(Statement.of(FAST_SELECT)), is(notNullValue())); + try (ResultSet rs = connection.executeQuery(SELECT_RANDOM_STATEMENT)) { + assertThat(rs, is(notNullValue())); + } } } @Test public void testTimeoutExceptionReadWriteAutocommitSlowUpdate() { - try (Connection connection = - createConnection( - ConnectionOptions.newBuilder() - .setCredentials(NoCredentials.getInstance()) - .setUri(URI) - .build())) { + mockSpanner.setExecuteSqlExecutionTime( + SimulatedExecutionTime.ofMinimumAndRandomTime(EXECUTION_TIME_SLOW_STATEMENT, 0)); + + try (Connection connection = createConnection()) { + connection.setAutocommit(true); connection.setStatementTimeout(TIMEOUT_FOR_SLOW_STATEMENTS, TimeUnit.MILLISECONDS); try { - connection.execute(Statement.of(SLOW_UPDATE)); - fail("Expected exception"); + connection.execute(INSERT_STATEMENT); + fail("Missing expected exception"); } catch (SpannerException ex) { assertEquals(ErrorCode.DEADLINE_EXCEEDED, ex.getErrorCode()); } @@ -419,44 +269,40 @@ public void testTimeoutExceptionReadWriteAutocommitSlowUpdate() { @Test public void testTimeoutExceptionReadWriteAutocommitSlowUpdateMultipleStatements() { - try (Connection connection = - createConnection( - ConnectionOptions.newBuilder() - .setCredentials(NoCredentials.getInstance()) - .setUri(URI) - .build())) { + mockSpanner.setExecuteSqlExecutionTime( + SimulatedExecutionTime.ofMinimumAndRandomTime(EXECUTION_TIME_SLOW_STATEMENT, 0)); + + try (Connection connection = createConnection()) { + connection.setAutocommit(true); connection.setStatementTimeout(TIMEOUT_FOR_SLOW_STATEMENTS, TimeUnit.MILLISECONDS); // assert that multiple statements after each other also time out for (int i = 0; i < 2; i++) { - boolean timedOut = false; try { connection.execute(Statement.of(SLOW_UPDATE)); + fail("missing expected exception"); } catch (SpannerException e) { - timedOut = e.getErrorCode() == ErrorCode.DEADLINE_EXCEEDED; + assertEquals(ErrorCode.DEADLINE_EXCEEDED, e.getErrorCode()); } - assertThat(timedOut, is(true)); } // try to do a new update that is fast. + mockSpanner.removeAllExecutionTimes(); connection.setStatementTimeout(TIMEOUT_FOR_FAST_STATEMENTS, TimeUnit.MILLISECONDS); - assertThat(connection.execute(Statement.of(FAST_UPDATE)).getUpdateCount(), is(equalTo(1L))); + assertThat(connection.execute(INSERT_STATEMENT).getUpdateCount(), is(equalTo(UPDATE_COUNT))); } } @Test public void testTimeoutExceptionReadWriteAutocommitSlowCommit() { - try (Connection connection = - createConnection( - ConnectionOptions.newBuilder() - .setCredentials(NoCredentials.getInstance()) - .setUri(URI) - .build(), - CommitRollbackBehavior.SLOW_COMMIT)) { + mockSpanner.setCommitExecutionTime( + SimulatedExecutionTime.ofMinimumAndRandomTime(EXECUTION_TIME_SLOW_STATEMENT, 0)); + + try (Connection connection = createConnection()) { connection.setStatementTimeout(TIMEOUT_FOR_FAST_STATEMENTS, TimeUnit.MILLISECONDS); // First verify that the fast update does not timeout when in transactional mode (as it is the // commit that is slow). connection.setAutocommit(false); - connection.execute(Statement.of(FAST_UPDATE)); + connection.execute(INSERT_STATEMENT); connection.rollback(); // Then verify that the update does timeout when executed in autocommit mode, as the commit @@ -464,8 +310,8 @@ public void testTimeoutExceptionReadWriteAutocommitSlowCommit() { connection.setStatementTimeout(TIMEOUT_FOR_SLOW_STATEMENTS, TimeUnit.MILLISECONDS); connection.setAutocommit(true); try { - connection.execute(Statement.of(FAST_UPDATE)); - fail("Expected exception"); + connection.execute(INSERT_STATEMENT); + fail("missing expected exception"); } catch (SpannerException ex) { assertEquals(ErrorCode.DEADLINE_EXCEEDED, ex.getErrorCode()); } @@ -474,47 +320,47 @@ public void testTimeoutExceptionReadWriteAutocommitSlowCommit() { @Test public void testTimeoutExceptionReadWriteAutocommitSlowCommitMultipleStatements() { - try (Connection connection = - createConnection( - ConnectionOptions.newBuilder() - .setCredentials(NoCredentials.getInstance()) - .setUri(URI) - .build(), - CommitRollbackBehavior.SLOW_COMMIT)) { + mockSpanner.setCommitExecutionTime( + SimulatedExecutionTime.ofMinimumAndRandomTime(EXECUTION_TIME_SLOW_STATEMENT, 0)); + + try (Connection connection = createConnection()) { + connection.setAutocommit(true); connection.setStatementTimeout(TIMEOUT_FOR_SLOW_STATEMENTS, TimeUnit.MILLISECONDS); // assert that multiple statements after each other also time out for (int i = 0; i < 2; i++) { - boolean timedOut = false; try { - connection.execute(Statement.of(FAST_UPDATE)); + connection.execute(INSERT_STATEMENT); + fail("Missing expected exception"); } catch (SpannerException e) { - timedOut = e.getErrorCode() == ErrorCode.DEADLINE_EXCEEDED; + assertThat(e.getErrorCode(), is(equalTo(ErrorCode.DEADLINE_EXCEEDED))); } - assertThat(timedOut, is(true)); } - // try to do a new query that is fast. + // try to do a query in autocommit mode. This will use a single-use read-only transaction that + // does not need to commit, i.e. it should succeed. connection.setStatementTimeout(TIMEOUT_FOR_FAST_STATEMENTS, TimeUnit.MILLISECONDS); - assertThat(connection.executeQuery(Statement.of(FAST_SELECT)), is(notNullValue())); + try (ResultSet rs = connection.executeQuery(SELECT_RANDOM_STATEMENT)) { + assertThat(rs, is(notNullValue())); + } } } @Test public void testTimeoutExceptionReadWriteAutocommitPartitioned() { - try (Connection connection = - createConnection( - ConnectionOptions.newBuilder() - .setCredentials(NoCredentials.getInstance()) - .setUri(URI) - .build())) { + try (Connection connection = createConnection()) { + connection.setAutocommit(true); connection.setAutocommitDmlMode(AutocommitDmlMode.PARTITIONED_NON_ATOMIC); - // first verify that the fast update does not timeout + // First verify that the statement will not timeout by default. connection.setStatementTimeout(TIMEOUT_FOR_FAST_STATEMENTS, TimeUnit.MILLISECONDS); - connection.execute(Statement.of(FAST_UPDATE)); + connection.execute(INSERT_STATEMENT); + // Now slow down the execution and verify that it times out. PDML uses the ExecuteStreamingSql + // RPC. + mockSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.ofMinimumAndRandomTime(EXECUTION_TIME_SLOW_STATEMENT, 0)); connection.setStatementTimeout(TIMEOUT_FOR_SLOW_STATEMENTS, TimeUnit.MILLISECONDS); try { - connection.execute(Statement.of(SLOW_UPDATE)); - fail("Expected exception"); + connection.execute(INSERT_STATEMENT); + fail("Missing expected exception"); } catch (SpannerException ex) { assertEquals(ErrorCode.DEADLINE_EXCEEDED, ex.getErrorCode()); } @@ -523,17 +369,15 @@ public void testTimeoutExceptionReadWriteAutocommitPartitioned() { @Test public void testTimeoutExceptionReadWriteTransactional() { - try (Connection connection = - createConnection( - ConnectionOptions.newBuilder() - .setCredentials(NoCredentials.getInstance()) - .setUri(URI) - .build())) { + mockSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.ofMinimumAndRandomTime(EXECUTION_TIME_SLOW_STATEMENT, 0)); + + try (Connection connection = createConnection()) { connection.setAutocommit(false); connection.setStatementTimeout(TIMEOUT_FOR_SLOW_STATEMENTS, TimeUnit.MILLISECONDS); try { - connection.executeQuery(Statement.of(SLOW_SELECT)); - fail("Expected exception"); + connection.executeQuery(SELECT_RANDOM_STATEMENT); + fail("Missing expected exception"); } catch (SpannerException ex) { assertEquals(ErrorCode.DEADLINE_EXCEEDED, ex.getErrorCode()); } @@ -542,57 +386,55 @@ public void testTimeoutExceptionReadWriteTransactional() { @Test public void testTimeoutExceptionReadWriteTransactionMultipleStatements() { - try (Connection connection = - createConnection( - ConnectionOptions.newBuilder() - .setCredentials(NoCredentials.getInstance()) - .setUri(URI) - .build())) { + mockSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.ofMinimumAndRandomTime(EXECUTION_TIME_SLOW_STATEMENT, 0)); + + try (Connection connection = createConnection()) { connection.setAutocommit(false); connection.setStatementTimeout(TIMEOUT_FOR_SLOW_STATEMENTS, TimeUnit.MILLISECONDS); // Assert that multiple statements after each other will timeout the first time, and then // throw a SpannerException with code FAILED_PRECONDITION. - boolean timedOut = false; for (int i = 0; i < 2; i++) { try { - connection.executeQuery(Statement.of(SLOW_SELECT)); + connection.executeQuery(SELECT_RANDOM_STATEMENT); + fail("Missing expected exception"); } catch (SpannerException e) { if (i == 0) { assertThat(e.getErrorCode(), is(equalTo(ErrorCode.DEADLINE_EXCEEDED))); - timedOut = true; } else { assertThat(e.getErrorCode(), is(equalTo(ErrorCode.FAILED_PRECONDITION))); } } } - assertThat(timedOut, is(true)); // do a rollback without any chance of a timeout connection.clearStatementTimeout(); connection.rollback(); // try to do a new query that is fast. + mockSpanner.removeAllExecutionTimes(); connection.setStatementTimeout(TIMEOUT_FOR_FAST_STATEMENTS, TimeUnit.MILLISECONDS); - assertThat(connection.executeQuery(Statement.of(FAST_SELECT)), is(notNullValue())); + try (ResultSet rs = connection.executeQuery(SELECT_RANDOM_STATEMENT)) { + assertThat(rs, is(notNullValue())); + } } } @Test public void testTimeoutExceptionReadWriteTransactionalSlowCommit() { - try (Connection connection = - createConnection( - ConnectionOptions.newBuilder() - .setCredentials(NoCredentials.getInstance()) - .setUri(URI) - .build(), - CommitRollbackBehavior.SLOW_COMMIT)) { + mockSpanner.setCommitExecutionTime( + SimulatedExecutionTime.ofMinimumAndRandomTime(EXECUTION_TIME_SLOW_STATEMENT, 0)); + + try (Connection connection = createConnection()) { connection.setAutocommit(false); connection.setStatementTimeout(TIMEOUT_FOR_FAST_STATEMENTS, TimeUnit.MILLISECONDS); - connection.executeQuery(Statement.of(FAST_SELECT)); + try (ResultSet rs = connection.executeQuery(SELECT_RANDOM_STATEMENT)) { + assertThat(rs, is(notNullValue())); + } connection.setStatementTimeout(TIMEOUT_FOR_SLOW_STATEMENTS, TimeUnit.MILLISECONDS); try { connection.commit(); - fail("Expected exception"); + fail("Missing expected exception"); } catch (SpannerException ex) { assertEquals(ErrorCode.DEADLINE_EXCEEDED, ex.getErrorCode()); } @@ -601,30 +443,27 @@ public void testTimeoutExceptionReadWriteTransactionalSlowCommit() { @Test public void testTimeoutExceptionReadWriteTransactionalSlowRollback() { - try (Connection connection = - createConnection( - ConnectionOptions.newBuilder() - .setCredentials(NoCredentials.getInstance()) - .setUri(URI) - .build(), - CommitRollbackBehavior.SLOW_ROLLBACK)) { + mockSpanner.setRollbackExecutionTime( + SimulatedExecutionTime.ofMinimumAndRandomTime(EXECUTION_TIME_SLOW_STATEMENT, 0)); + + try (Connection connection = createConnection()) { connection.setAutocommit(false); connection.setStatementTimeout(TIMEOUT_FOR_FAST_STATEMENTS, TimeUnit.MILLISECONDS); - connection.executeQuery(Statement.of(FAST_SELECT)); - connection.setStatementTimeout(TIMEOUT_FOR_SLOW_STATEMENTS, TimeUnit.MILLISECONDS); - try { - connection.rollback(); - fail("Expected exception"); - } catch (SpannerException ex) { - assertEquals(ErrorCode.DEADLINE_EXCEEDED, ex.getErrorCode()); + try (ResultSet rs = connection.executeQuery(SELECT_RANDOM_STATEMENT)) { + assertThat(rs, is(notNullValue())); } + connection.setStatementTimeout(TIMEOUT_FOR_SLOW_STATEMENTS, TimeUnit.MILLISECONDS); + // Rollback timeouts are not propagated as exceptions, as all errors during a Rollback RPC are + // ignored by the client library. + connection.rollback(); } } private static final class ConnectionReadOnlyAutocommit implements ConnectionConsumer { @Override public void accept(Connection t) { + t.setAutocommit(true); t.setReadOnly(true); } } @@ -651,7 +490,10 @@ public void testInterruptedExceptionReadOnlyTransactional() private static final class ConnectionReadWriteAutocommit implements ConnectionConsumer { @Override - public void accept(Connection t) {} + public void accept(Connection t) { + t.setAutocommit(true); + t.setReadOnly(false); + } } @Test @@ -664,6 +506,7 @@ private static final class ConnectionReadWriteTransactional implements Connectio @Override public void accept(Connection t) { t.setAutocommit(false); + t.setReadOnly(false); } } @@ -675,51 +518,45 @@ public void testInterruptedExceptionReadWriteTransactional() private void testInterruptedException(final ConnectionConsumer consumer) throws InterruptedException, ExecutionException { + mockSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.ofMinimumAndRandomTime(EXECUTION_TIME_SLOW_STATEMENT, 0)); + + final CountDownLatch latch = new CountDownLatch(1); ExecutorService executor = Executors.newSingleThreadExecutor(); Future future = executor.submit( new Callable() { @Override public Boolean call() { - try (Connection connection = - createConnection( - ConnectionOptions.newBuilder() - .setCredentials(NoCredentials.getInstance()) - .setUri(URI) - .build())) { + try (Connection connection = createConnection()) { consumer.accept(connection); connection.setStatementTimeout(10000L, TimeUnit.MILLISECONDS); - connection.executeQuery(Statement.of(SLOW_SELECT)); + latch.countDown(); + try (ResultSet rs = connection.executeQuery(SELECT_RANDOM_STATEMENT)) {} + return false; } catch (SpannerException e) { - if (e.getErrorCode() == ErrorCode.CANCELLED) { - return Boolean.TRUE; - } else { - return Boolean.FALSE; - } + return e.getErrorCode() == ErrorCode.CANCELLED; } - return Boolean.FALSE; } }); - // wait a little bit to ensure that the task has started - Thread.sleep(10L); + latch.await(10L, TimeUnit.SECONDS); executor.shutdownNow(); assertThat(future.get(), is(true)); } @Test public void testInvalidQueryReadOnlyAutocommit() { - try (Connection connection = - createConnection( - ConnectionOptions.newBuilder() - .setUri(URI) - .setCredentials(NoCredentials.getInstance()) - .build())) { + mockSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.ofException(Status.INVALID_ARGUMENT.asRuntimeException())); + + try (Connection connection = createConnection()) { + connection.setAutocommit(true); connection.setReadOnly(true); connection.setStatementTimeout(TIMEOUT_FOR_FAST_STATEMENTS, TimeUnit.MILLISECONDS); try { connection.executeQuery(Statement.of(INVALID_SELECT)); - fail("Expected exception"); + fail("Missing expected exception"); } catch (SpannerException ex) { assertEquals(ErrorCode.INVALID_ARGUMENT, ex.getErrorCode()); } @@ -728,18 +565,16 @@ public void testInvalidQueryReadOnlyAutocommit() { @Test public void testInvalidQueryReadOnlyTransactional() { - try (Connection connection = - createConnection( - ConnectionOptions.newBuilder() - .setCredentials(NoCredentials.getInstance()) - .setUri(URI) - .build())) { + mockSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.ofException(Status.INVALID_ARGUMENT.asRuntimeException())); + + try (Connection connection = createConnection()) { connection.setReadOnly(true); connection.setAutocommit(false); connection.setStatementTimeout(TIMEOUT_FOR_FAST_STATEMENTS, TimeUnit.MILLISECONDS); try { connection.executeQuery(Statement.of(INVALID_SELECT)); - fail("Expected exception"); + fail("Missing expected exception"); } catch (SpannerException ex) { assertEquals(ErrorCode.INVALID_ARGUMENT, ex.getErrorCode()); } @@ -748,16 +583,15 @@ public void testInvalidQueryReadOnlyTransactional() { @Test public void testInvalidQueryReadWriteAutocommit() { - try (Connection connection = - createConnection( - ConnectionOptions.newBuilder() - .setCredentials(NoCredentials.getInstance()) - .setUri(URI) - .build())) { + mockSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.ofException(Status.INVALID_ARGUMENT.asRuntimeException())); + + try (Connection connection = createConnection()) { + connection.setAutocommit(true); connection.setStatementTimeout(TIMEOUT_FOR_FAST_STATEMENTS, TimeUnit.MILLISECONDS); try { connection.executeQuery(Statement.of(INVALID_SELECT)); - fail("Expected exception"); + fail("Missing expected exception"); } catch (SpannerException ex) { assertEquals(ErrorCode.INVALID_ARGUMENT, ex.getErrorCode()); } @@ -766,394 +600,427 @@ public void testInvalidQueryReadWriteAutocommit() { @Test public void testInvalidQueryReadWriteTransactional() { - try (Connection connection = - createConnection( - ConnectionOptions.newBuilder() - .setCredentials(NoCredentials.getInstance()) - .setUri(URI) - .build())) { + mockSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.ofException(Status.INVALID_ARGUMENT.asRuntimeException())); + + try (Connection connection = createConnection()) { connection.setAutocommit(false); connection.setStatementTimeout(TIMEOUT_FOR_FAST_STATEMENTS, TimeUnit.MILLISECONDS); try { connection.executeQuery(Statement.of(INVALID_SELECT)); - fail("Expected exception"); + fail("Missing expected exception"); } catch (SpannerException ex) { assertEquals(ErrorCode.INVALID_ARGUMENT, ex.getErrorCode()); } } } + static void waitForRequestsToContain(Class request) { + try { + mockSpanner.waitForRequestsToContain(request, EXECUTION_TIME_SLOW_STATEMENT); + } catch (InterruptedException e) { + throw SpannerExceptionFactory.propagateInterrupt(e); + } catch (TimeoutException e) { + throw SpannerExceptionFactory.propagateTimeout(e); + } + } + @Test public void testCancelReadOnlyAutocommit() { - try (Connection connection = - createConnection( - ConnectionOptions.newBuilder() - .setCredentials(NoCredentials.getInstance()) - .setUri(URI) - .build())) { + mockSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.ofMinimumAndRandomTime(EXECUTION_TIME_SLOW_STATEMENT, 0)); + + ExecutorService executor = Executors.newSingleThreadExecutor(); + try (Connection connection = createConnection()) { + connection.setAutocommit(true); connection.setReadOnly(true); - Executors.newSingleThreadScheduledExecutor() - .schedule( - new Runnable() { - @Override - public void run() { - connection.cancel(); - } - }, - WAIT_BEFORE_CANCEL, - TimeUnit.MILLISECONDS); + executor.execute( + new Runnable() { + @Override + public void run() { + waitForRequestsToContain(ExecuteSqlRequest.class); + connection.cancel(); + } + }); try { - connection.executeQuery(Statement.of(SLOW_SELECT)); - fail("Expected exception"); + connection.executeQuery(SELECT_RANDOM_STATEMENT); + fail("Missing expected exception"); } catch (SpannerException ex) { assertEquals(ErrorCode.CANCELLED, ex.getErrorCode()); } + } finally { + executor.shutdown(); } } @Test public void testCancelReadOnlyAutocommitMultipleStatements() { - try (Connection connection = - createConnection( - ConnectionOptions.newBuilder() - .setCredentials(NoCredentials.getInstance()) - .setUri(URI) - .build())) { - connection.setReadOnly(true); - Executors.newSingleThreadScheduledExecutor() - .schedule( - new Runnable() { - @Override - public void run() { - connection.cancel(); - } - }, - WAIT_BEFORE_CANCEL, - TimeUnit.MILLISECONDS); + mockSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.ofMinimumAndRandomTime(EXECUTION_TIME_SLOW_STATEMENT, 0)); - boolean cancelled = false; - try { - connection.executeQuery(Statement.of(SLOW_SELECT)); + ExecutorService executor = Executors.newSingleThreadExecutor(); + try (Connection connection = createConnection()) { + connection.setAutocommit(true); + connection.setReadOnly(true); + executor.execute( + new Runnable() { + @Override + public void run() { + waitForRequestsToContain(ExecuteSqlRequest.class); + connection.cancel(); + } + }); + + try (ResultSet rs = connection.executeQuery(SELECT_RANDOM_STATEMENT)) { + fail("Missing expected exception"); } catch (SpannerException e) { - cancelled = e.getErrorCode() == ErrorCode.CANCELLED; + assertThat(e.getErrorCode(), is(equalTo(ErrorCode.CANCELLED))); } - assertThat(cancelled, is(true)); - // try to do a new query that is fast. + mockSpanner.removeAllExecutionTimes(); connection.setStatementTimeout(TIMEOUT_FOR_FAST_STATEMENTS, TimeUnit.MILLISECONDS); - assertThat(connection.executeQuery(Statement.of(FAST_SELECT)), is(notNullValue())); + try (ResultSet rs = connection.executeQuery(SELECT_RANDOM_STATEMENT)) { + assertThat(rs, is(notNullValue())); + } + } finally { + executor.shutdown(); } } @Test public void testCancelReadOnlyTransactional() { - try (Connection connection = - createConnection( - ConnectionOptions.newBuilder() - .setCredentials(NoCredentials.getInstance()) - .setUri(URI) - .build())) { + mockSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.ofMinimumAndRandomTime(EXECUTION_TIME_SLOW_STATEMENT, 0)); + + ExecutorService executor = Executors.newSingleThreadExecutor(); + try (Connection connection = createConnection()) { connection.setReadOnly(true); connection.setAutocommit(false); - Executors.newSingleThreadScheduledExecutor() - .schedule( - new Runnable() { - @Override - public void run() { - connection.cancel(); - } - }, - WAIT_BEFORE_CANCEL, - TimeUnit.MILLISECONDS); + executor.execute( + new Runnable() { + @Override + public void run() { + waitForRequestsToContain(ExecuteSqlRequest.class); + connection.cancel(); + } + }); try { - connection.executeQuery(Statement.of(SLOW_SELECT)); - fail("Expected exception"); + connection.executeQuery(SELECT_RANDOM_STATEMENT); + fail("Missing expected exception"); } catch (SpannerException ex) { assertEquals(ErrorCode.CANCELLED, ex.getErrorCode()); } + } finally { + executor.shutdown(); } } @Test public void testCancelReadOnlyTransactionalMultipleStatements() { - try (Connection connection = - createConnection( - ConnectionOptions.newBuilder() - .setCredentials(NoCredentials.getInstance()) - .setUri(URI) - .build())) { + mockSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.ofMinimumAndRandomTime(EXECUTION_TIME_SLOW_STATEMENT, 0)); + + ExecutorService executor = Executors.newSingleThreadExecutor(); + try (Connection connection = createConnection()) { connection.setReadOnly(true); connection.setAutocommit(false); - Executors.newSingleThreadScheduledExecutor() - .schedule( - new Runnable() { - @Override - public void run() { - connection.cancel(); - } - }, - WAIT_BEFORE_CANCEL, - TimeUnit.MILLISECONDS); - - boolean cancelled = false; + executor.execute( + new Runnable() { + @Override + public void run() { + waitForRequestsToContain(ExecuteSqlRequest.class); + connection.cancel(); + } + }); try { connection.executeQuery(Statement.of(SLOW_SELECT)); + fail("Missing expected exception"); } catch (SpannerException e) { - cancelled = e.getErrorCode() == ErrorCode.CANCELLED; + assertEquals(ErrorCode.CANCELLED, e.getErrorCode()); } - assertThat(cancelled, is(true)); // try to do a new query that is fast. + mockSpanner.removeAllExecutionTimes(); connection.setStatementTimeout(TIMEOUT_FOR_FAST_STATEMENTS, TimeUnit.MILLISECONDS); - assertThat(connection.executeQuery(Statement.of(FAST_SELECT)), is(notNullValue())); + try (ResultSet rs = connection.executeQuery(SELECT_RANDOM_STATEMENT)) { + assertThat(rs, is(notNullValue())); + } // rollback and do another fast query connection.rollback(); - assertThat(connection.executeQuery(Statement.of(FAST_SELECT)), is(notNullValue())); + try (ResultSet rs = connection.executeQuery(SELECT_RANDOM_STATEMENT)) { + assertThat(rs, is(notNullValue())); + } + } finally { + executor.shutdown(); } } @Test public void testCancelReadWriteAutocommit() { - try (Connection connection = - createConnection( - ConnectionOptions.newBuilder() - .setCredentials(NoCredentials.getInstance()) - .setUri(URI) - .build())) { - Executors.newSingleThreadScheduledExecutor() - .schedule( - new Runnable() { - @Override - public void run() { - connection.cancel(); - } - }, - WAIT_BEFORE_CANCEL, - TimeUnit.MILLISECONDS); + mockSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.ofMinimumAndRandomTime(EXECUTION_TIME_SLOW_STATEMENT, 0)); + + ExecutorService executor = Executors.newSingleThreadExecutor(); + try (Connection connection = createConnection()) { + connection.setAutocommit(true); + executor.execute( + new Runnable() { + @Override + public void run() { + waitForRequestsToContain(ExecuteSqlRequest.class); + connection.cancel(); + } + }); try { - connection.executeQuery(Statement.of(SLOW_SELECT)); - fail("Expected exception"); + connection.executeQuery(SELECT_RANDOM_STATEMENT); + fail("Missing expected exception"); } catch (SpannerException ex) { assertEquals(ErrorCode.CANCELLED, ex.getErrorCode()); } + } finally { + executor.shutdown(); } } @Test public void testCancelReadWriteAutocommitMultipleStatements() { - try (Connection connection = - createConnection( - ConnectionOptions.newBuilder() - .setCredentials(NoCredentials.getInstance()) - .setUri(URI) - .build())) { - Executors.newSingleThreadScheduledExecutor() - .schedule( - new Runnable() { - @Override - public void run() { - connection.cancel(); - } - }, - WAIT_BEFORE_CANCEL, - TimeUnit.MILLISECONDS); + mockSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.ofMinimumAndRandomTime(EXECUTION_TIME_SLOW_STATEMENT, 0)); - boolean cancelled = false; + ExecutorService executor = Executors.newSingleThreadExecutor(); + try (Connection connection = createConnection()) { + connection.setAutocommit(true); + executor.execute( + new Runnable() { + @Override + public void run() { + waitForRequestsToContain(ExecuteSqlRequest.class); + connection.cancel(); + } + }); try { - connection.executeQuery(Statement.of(SLOW_SELECT)); - } catch (SpannerException e) { - cancelled = e.getErrorCode() == ErrorCode.CANCELLED; + connection.executeQuery(SELECT_RANDOM_STATEMENT); + fail("Missing expected exception"); + } catch (SpannerException ex) { + assertEquals(ErrorCode.CANCELLED, ex.getErrorCode()); } - assertThat(cancelled, is(true)); // try to do a new query that is fast. + mockSpanner.removeAllExecutionTimes(); connection.setStatementTimeout(TIMEOUT_FOR_FAST_STATEMENTS, TimeUnit.MILLISECONDS); - assertThat(connection.executeQuery(Statement.of(FAST_SELECT)), is(notNullValue())); + try (ResultSet rs = connection.executeQuery(SELECT_RANDOM_STATEMENT)) { + assertThat(rs, is(notNullValue())); + } + } finally { + executor.shutdown(); } } @Test public void testCancelReadWriteAutocommitSlowUpdate() { - try (Connection connection = - createConnection( - ConnectionOptions.newBuilder() - .setCredentials(NoCredentials.getInstance()) - .setUri(URI) - .build())) { - Executors.newSingleThreadScheduledExecutor() - .schedule( - new Runnable() { - @Override - public void run() { - connection.cancel(); - } - }, - WAIT_BEFORE_CANCEL, - TimeUnit.MILLISECONDS); + mockSpanner.setExecuteSqlExecutionTime( + SimulatedExecutionTime.ofMinimumAndRandomTime(EXECUTION_TIME_SLOW_STATEMENT, 0)); + + ExecutorService executor = Executors.newSingleThreadExecutor(); + try (Connection connection = createConnection()) { + connection.setAutocommit(true); + executor.execute( + new Runnable() { + @Override + public void run() { + waitForRequestsToContain(ExecuteSqlRequest.class); + connection.cancel(); + } + }); try { - connection.execute(Statement.of(SLOW_UPDATE)); - fail("Expected exception"); + connection.execute(INSERT_STATEMENT); + fail("Missing expected exception"); } catch (SpannerException ex) { assertEquals(ErrorCode.CANCELLED, ex.getErrorCode()); } + } finally { + executor.shutdown(); } } @Test public void testCancelReadWriteAutocommitSlowCommit() { - try (Connection connection = - createConnection( - ConnectionOptions.newBuilder() - .setCredentials(NoCredentials.getInstance()) - .setUri(URI) - .build(), - CommitRollbackBehavior.SLOW_COMMIT)) { - Executors.newSingleThreadScheduledExecutor() - .schedule( - new Runnable() { - @Override - public void run() { - connection.cancel(); - } - }, - WAIT_BEFORE_CANCEL, - TimeUnit.MILLISECONDS); - connection.execute(Statement.of(FAST_UPDATE)); - fail("Expected exception"); + mockSpanner.setCommitExecutionTime( + SimulatedExecutionTime.ofMinimumAndRandomTime(EXECUTION_TIME_SLOW_STATEMENT, 0)); + + ExecutorService executor = Executors.newSingleThreadExecutor(); + try (Connection connection = createConnection()) { + connection.setAutocommit(true); + executor.execute( + new Runnable() { + @Override + public void run() { + waitForRequestsToContain(CommitRequest.class); + connection.cancel(); + } + }); + connection.execute(INSERT_STATEMENT); + fail("Missing expected exception"); } catch (SpannerException ex) { assertEquals(ErrorCode.CANCELLED, ex.getErrorCode()); + } finally { + executor.shutdown(); } } @Test public void testCancelReadWriteTransactional() { - try (Connection connection = - createConnection( - ConnectionOptions.newBuilder() - .setCredentials(NoCredentials.getInstance()) - .setUri(URI) - .build())) { - connection.setAutocommit(false); - Executors.newSingleThreadScheduledExecutor() - .schedule( - new Runnable() { - @Override - public void run() { - connection.cancel(); - } - }, - WAIT_BEFORE_CANCEL, - TimeUnit.MILLISECONDS); + mockSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.ofMinimumAndRandomTime(EXECUTION_TIME_SLOW_STATEMENT, 0)); - connection.executeQuery(Statement.of(SLOW_SELECT)); - fail("Expected exception"); + ExecutorService executor = Executors.newSingleThreadExecutor(); + try (Connection connection = createConnection()) { + connection.setAutocommit(false); + executor.execute( + new Runnable() { + @Override + public void run() { + waitForRequestsToContain(ExecuteSqlRequest.class); + connection.cancel(); + } + }); + connection.executeQuery(SELECT_RANDOM_STATEMENT); + fail("Missing expected exception"); } catch (SpannerException ex) { assertEquals(ErrorCode.CANCELLED, ex.getErrorCode()); + } finally { + executor.shutdown(); } } @Test public void testCancelReadWriteTransactionalMultipleStatements() { - try (Connection connection = - createConnection( - ConnectionOptions.newBuilder() - .setCredentials(NoCredentials.getInstance()) - .setUri(URI) - .build())) { - connection.setAutocommit(false); - Executors.newSingleThreadScheduledExecutor() - .schedule( - new Runnable() { - @Override - public void run() { - connection.cancel(); - } - }, - WAIT_BEFORE_CANCEL, - TimeUnit.MILLISECONDS); + mockSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.ofMinimumAndRandomTime(EXECUTION_TIME_SLOW_STATEMENT, 0)); - boolean cancelled = false; + ExecutorService executor = Executors.newSingleThreadExecutor(); + try (Connection connection = createConnection()) { + connection.setAutocommit(false); + executor.execute( + new Runnable() { + @Override + public void run() { + waitForRequestsToContain(ExecuteSqlRequest.class); + connection.cancel(); + } + }); try { - connection.executeQuery(Statement.of(SLOW_SELECT)); - fail("Expected exception"); + connection.executeQuery(SELECT_RANDOM_STATEMENT); + fail("Missing expected exception"); } catch (SpannerException e) { - cancelled = e.getErrorCode() == ErrorCode.CANCELLED; + assertEquals(ErrorCode.CANCELLED, e.getErrorCode()); } - assertThat(cancelled, is(true)); // Rollback the transaction as it is no longer usable. connection.rollback(); // Try to do a new query that is fast. + mockSpanner.removeAllExecutionTimes(); connection.setStatementTimeout(TIMEOUT_FOR_FAST_STATEMENTS, TimeUnit.MILLISECONDS); - assertThat(connection.executeQuery(Statement.of(FAST_SELECT)), is(notNullValue())); + try (ResultSet rs = connection.executeQuery(SELECT_RANDOM_STATEMENT)) { + assertThat(rs, is(notNullValue())); + } + } finally { + executor.shutdown(); + } + } + + static void addSlowMockDdlOperation() { + addSlowMockDdlOperations(1); + } + + static void addSlowMockDdlOperations(int count) { + addMockDdlOperations(count, false); + } + + static void addFastMockDdlOperation() { + addFastMockDdlOperations(1); + } + + static void addFastMockDdlOperations(int count) { + addMockDdlOperations(count, true); + } + + static void addMockDdlOperations(int count, boolean done) { + for (int i = 0; i < count; i++) { + mockDatabaseAdmin.addResponse( + Operation.newBuilder() + .setMetadata( + Any.pack( + UpdateDatabaseDdlMetadata.newBuilder() + .addStatements(SLOW_DDL) + .setDatabase("projects/proj/instances/inst/databases/db") + .build())) + .setName("projects/proj/instances/inst/databases/db/operations/1") + .setDone(done) + .setResponse(Any.pack(Empty.getDefaultInstance())) + .build()); } } @Test public void testCancelDdlBatch() { - try (Connection connection = - createConnection( - ConnectionOptions.newBuilder() - .setCredentials(NoCredentials.getInstance()) - .setUri(URI) - .build())) { + addSlowMockDdlOperation(); + + ExecutorService executor = Executors.newSingleThreadExecutor(); + try (Connection connection = createConnection()) { connection.setAutocommit(false); connection.startBatchDdl(); connection.execute(Statement.of(SLOW_DDL)); - Executors.newSingleThreadScheduledExecutor() - .schedule( - new Runnable() { - @Override - public void run() { - connection.cancel(); - } - }, - WAIT_BEFORE_CANCEL, - TimeUnit.MILLISECONDS); + executor.execute( + new Runnable() { + @Override + public void run() { + Uninterruptibles.sleepUninterruptibly(100L, TimeUnit.MILLISECONDS); + connection.cancel(); + } + }); connection.runBatch(); - fail("Expected exception"); + fail("Missing expected exception"); } catch (SpannerException ex) { assertEquals(ErrorCode.CANCELLED, ex.getErrorCode()); + } finally { + executor.shutdown(); } } @Test public void testCancelDdlAutocommit() { - try (Connection connection = - createConnection( - ConnectionOptions.newBuilder() - .setCredentials(NoCredentials.getInstance()) - .setUri(URI) - .build())) { - Executors.newSingleThreadScheduledExecutor() - .schedule( - new Runnable() { - @Override - public void run() { - connection.cancel(); - } - }, - WAIT_BEFORE_CANCEL, - TimeUnit.MILLISECONDS); + addSlowMockDdlOperation(); + ExecutorService executor = Executors.newSingleThreadExecutor(); + try (Connection connection = createConnection()) { + connection.setAutocommit(true); + executor.execute( + new Runnable() { + @Override + public void run() { + Uninterruptibles.sleepUninterruptibly(100L, TimeUnit.MILLISECONDS); + connection.cancel(); + } + }); connection.execute(Statement.of(SLOW_DDL)); - fail("Expected exception"); + fail("Missing expected exception"); } catch (SpannerException ex) { assertEquals(ErrorCode.CANCELLED, ex.getErrorCode()); + } finally { + executor.shutdown(); } } @Test public void testTimeoutExceptionDdlAutocommit() { - try (Connection connection = - createConnection( - ConnectionOptions.newBuilder() - .setCredentials(NoCredentials.getInstance()) - .setUri(URI) - .build())) { + addSlowMockDdlOperations(10); + + try (Connection connection = createConnection()) { + connection.setAutocommit(true); connection.setStatementTimeout(TIMEOUT_FOR_SLOW_STATEMENTS, TimeUnit.MILLISECONDS); connection.execute(Statement.of(SLOW_DDL)); - fail("Expected exception"); + fail("Missing expected exception"); } catch (SpannerException ex) { assertEquals(ErrorCode.DEADLINE_EXCEEDED, ex.getErrorCode()); } @@ -1161,25 +1028,24 @@ public void testTimeoutExceptionDdlAutocommit() { @Test public void testTimeoutExceptionDdlAutocommitMultipleStatements() { - try (Connection connection = - createConnection( - ConnectionOptions.newBuilder() - .setCredentials(NoCredentials.getInstance()) - .setUri(URI) - .build())) { + addSlowMockDdlOperations(20); + + try (Connection connection = createConnection()) { + connection.setAutocommit(true); connection.setStatementTimeout(TIMEOUT_FOR_SLOW_STATEMENTS, TimeUnit.MILLISECONDS); // assert that multiple statements after each other also time out for (int i = 0; i < 2; i++) { - boolean timedOut = false; try { connection.execute(Statement.of(SLOW_DDL)); + fail("Missing expected exception"); } catch (SpannerException e) { - timedOut = e.getErrorCode() == ErrorCode.DEADLINE_EXCEEDED; + assertEquals(ErrorCode.DEADLINE_EXCEEDED, e.getErrorCode()); } - assertThat(timedOut, is(true)); } // try to do a new DDL statement that is fast. + mockDatabaseAdmin.reset(); + addFastMockDdlOperation(); connection.setStatementTimeout(TIMEOUT_FOR_FAST_STATEMENTS, TimeUnit.MILLISECONDS); assertThat(connection.execute(Statement.of(FAST_DDL)), is(notNullValue())); } @@ -1187,21 +1053,18 @@ public void testTimeoutExceptionDdlAutocommitMultipleStatements() { @Test public void testTimeoutExceptionDdlBatch() { - try (Connection connection = - createConnection( - ConnectionOptions.newBuilder() - .setCredentials(NoCredentials.getInstance()) - .setUri(URI) - .build())) { + addSlowMockDdlOperations(10); + + try (Connection connection = createConnection()) { connection.setAutocommit(false); connection.startBatchDdl(); connection.setStatementTimeout(TIMEOUT_FOR_SLOW_STATEMENTS, TimeUnit.MILLISECONDS); // the following statement will NOT timeout as the statement is only buffered locally connection.execute(Statement.of(SLOW_DDL)); - // the commit sends the statement to the server and should timeout + // the runBatch() statement sends the statement to the server and should timeout connection.runBatch(); - fail("Expected exception"); + fail("Missing expected exception"); } catch (SpannerException ex) { assertEquals(ErrorCode.DEADLINE_EXCEEDED, ex.getErrorCode()); } @@ -1209,28 +1072,27 @@ public void testTimeoutExceptionDdlBatch() { @Test public void testTimeoutExceptionDdlBatchMultipleStatements() { - try (Connection connection = - createConnection( - ConnectionOptions.newBuilder() - .setCredentials(NoCredentials.getInstance()) - .setUri(URI) - .build())) { + addSlowMockDdlOperations(20); + + try (Connection connection = createConnection()) { connection.setAutocommit(false); connection.setStatementTimeout(TIMEOUT_FOR_SLOW_STATEMENTS, TimeUnit.MILLISECONDS); // assert that multiple statements after each other also time out for (int i = 0; i < 2; i++) { - boolean timedOut = false; + connection.startBatchDdl(); connection.execute(Statement.of(SLOW_DDL)); try { connection.runBatch(); + fail("Missing expected exception"); } catch (SpannerException e) { - timedOut = e.getErrorCode() == ErrorCode.DEADLINE_EXCEEDED; + assertEquals(ErrorCode.DEADLINE_EXCEEDED, e.getErrorCode()); } - assertThat(timedOut, is(true)); } // try to do a new DDL statement that is fast. + mockDatabaseAdmin.reset(); + addFastMockDdlOperation(); connection.setStatementTimeout(TIMEOUT_FOR_FAST_STATEMENTS, TimeUnit.MILLISECONDS); connection.startBatchDdl(); assertThat(connection.execute(Statement.of(FAST_DDL)), is(notNullValue())); @@ -1240,21 +1102,19 @@ public void testTimeoutExceptionDdlBatchMultipleStatements() { @Test public void testTimeoutDifferentTimeUnits() { - try (Connection connection = - createConnection( - ConnectionOptions.newBuilder() - .setCredentials(NoCredentials.getInstance()) - .setUri(URI) - .build())) { + mockSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.ofMinimumAndRandomTime(EXECUTION_TIME_SLOW_STATEMENT, 0)); + + try (Connection connection = createConnection()) { + connection.setAutocommit(true); for (TimeUnit unit : ReadOnlyStalenessUtil.SUPPORTED_UNITS) { connection.setStatementTimeout(1L, unit); - boolean timedOut = false; try { - connection.execute(Statement.of(SLOW_SELECT)); + connection.execute(SELECT_RANDOM_STATEMENT); + fail("Missing expected exception"); } catch (SpannerException e) { - timedOut = e.getErrorCode() == ErrorCode.DEADLINE_EXCEEDED; + assertEquals(ErrorCode.DEADLINE_EXCEEDED, e.getErrorCode()); } - assertThat(timedOut, is(true)); } } } diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/it/ITAsyncTransactionRetryTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/it/ITAsyncTransactionRetryTest.java new file mode 100644 index 0000000000..721dccc651 --- /dev/null +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/it/ITAsyncTransactionRetryTest.java @@ -0,0 +1,1015 @@ +/* + * 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.connection.it; + +import static com.google.cloud.spanner.SpannerApiFutures.get; +import static com.google.cloud.spanner.testing.EmulatorSpannerHelper.isUsingEmulator; +import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.fail; +import static org.junit.Assume.assumeFalse; + +import com.google.api.core.ApiFuture; +import com.google.api.core.SettableApiFuture; +import com.google.cloud.Timestamp; +import com.google.cloud.spanner.AbortedDueToConcurrentModificationException; +import com.google.cloud.spanner.AbortedException; +import com.google.cloud.spanner.AsyncResultSet; +import com.google.cloud.spanner.AsyncResultSet.CallbackResponse; +import com.google.cloud.spanner.AsyncResultSet.ReadyCallback; +import com.google.cloud.spanner.ErrorCode; +import com.google.cloud.spanner.KeySet; +import com.google.cloud.spanner.Mutation; +import com.google.cloud.spanner.Options; +import com.google.cloud.spanner.ParallelIntegrationTest; +import com.google.cloud.spanner.ResultSet; +import com.google.cloud.spanner.SpannerExceptionFactory; +import com.google.cloud.spanner.Statement; +import com.google.cloud.spanner.Struct; +import com.google.cloud.spanner.connection.Connection; +import com.google.cloud.spanner.connection.ITAbstractSpannerTest; +import com.google.cloud.spanner.connection.TransactionRetryListener; +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.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; +import java.util.logging.Level; +import java.util.logging.Logger; +import org.junit.After; +import org.junit.AfterClass; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.rules.TestName; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** + * This integration test tests the different scenarios for automatically retrying read/write + * transactions, both when possible and when the transaction must abort because of a concurrent + * update. + */ +@Category(ParallelIntegrationTest.class) +@RunWith(JUnit4.class) +public class ITAsyncTransactionRetryTest extends ITAbstractSpannerTest { + private static final Logger logger = + Logger.getLogger(ITAsyncTransactionRetryTest.class.getName()); + + @Rule public TestName testName = new TestName(); + + private static final ExecutorService executor = Executors.newFixedThreadPool(4); + + @AfterClass + public static void shutdownExecutor() { + executor.shutdown(); + } + + @Override + protected void appendConnectionUri(StringBuilder uri) { + uri.append(";autocommit=false;retryAbortsInternally=true"); + } + + @Override + public boolean doCreateDefaultTestTable() { + return true; + } + + /** Clear the test table before each test run */ + @Before + public void clearTable() { + try (ITConnection connection = createConnection()) { + connection.bufferedWrite(Mutation.delete("TEST", KeySet.all())); + get(connection.commitAsync()); + } + } + + @Before + public void clearStatistics() { + RETRY_STATISTICS.clear(); + } + + @Before + public void logStart() { + logger.fine( + "--------------------------------------------------------------\n" + + testName.getMethodName() + + " started"); + } + + @After + public void logFinished() { + logger.fine( + "--------------------------------------------------------------\n" + + testName.getMethodName() + + " finished"); + } + + /** Simple data structure to keep track of retry statistics */ + private static class RetryStatistics { + private int totalRetryAttemptsStarted; + private int totalRetryAttemptsFinished; + private int totalSuccessfulRetries; + private int totalErroredRetries; + private int totalNestedAborts; + private int totalMaxAttemptsExceeded; + private int totalConcurrentModifications; + + private void clear() { + totalRetryAttemptsStarted = 0; + totalRetryAttemptsFinished = 0; + totalSuccessfulRetries = 0; + totalErroredRetries = 0; + totalNestedAborts = 0; + totalMaxAttemptsExceeded = 0; + totalConcurrentModifications = 0; + } + } + + /** + * Static to allow access from the {@link CountTransactionRetryListener}. Statistics are + * automatically cleared before each test case. + */ + public static final RetryStatistics RETRY_STATISTICS = new RetryStatistics(); + + /** + * Simple {@link TransactionRetryListener} that keeps track of the total count of the different + * transaction retry events of a {@link Connection}. Note that as {@link + * TransactionRetryListener}s are instantiated once per connection, the listener keeps track of + * the total statistics of a connection and not only of the last transaction. + */ + public static class CountTransactionRetryListener implements TransactionRetryListener { + + @Override + public void retryStarting(Timestamp transactionStarted, long transactionId, int retryAttempt) { + RETRY_STATISTICS.totalRetryAttemptsStarted++; + } + + @Override + public void retryFinished( + Timestamp transactionStarted, long transactionId, int retryAttempt, RetryResult result) { + RETRY_STATISTICS.totalRetryAttemptsFinished++; + switch (result) { + case RETRY_ABORTED_AND_MAX_ATTEMPTS_EXCEEDED: + RETRY_STATISTICS.totalMaxAttemptsExceeded++; + break; + case RETRY_ABORTED_AND_RESTARTING: + RETRY_STATISTICS.totalNestedAborts++; + break; + case RETRY_ABORTED_DUE_TO_CONCURRENT_MODIFICATION: + RETRY_STATISTICS.totalConcurrentModifications++; + break; + case RETRY_ERROR: + RETRY_STATISTICS.totalErroredRetries++; + break; + case RETRY_SUCCESSFUL: + RETRY_STATISTICS.totalSuccessfulRetries++; + break; + default: + break; + } + } + } + + private ApiFuture getTestRecordCountAsync(Connection connection) { + final SettableApiFuture count = SettableApiFuture.create(); + try (AsyncResultSet rs = + connection.executeQueryAsync(Statement.of("SELECT COUNT(*) AS C FROM TEST WHERE ID=1"))) { + rs.setCallback( + executor, + new ReadyCallback() { + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + while (true) { + switch (resultSet.tryNext()) { + case DONE: + return CallbackResponse.DONE; + case NOT_READY: + return CallbackResponse.CONTINUE; + case OK: + count.set(resultSet.getLong("C")); + break; + } + } + } + }); + } + return count; + } + + private void verifyRecordCount(Connection connection, long expected) { + try (AsyncResultSet rs = + connection.executeQueryAsync(Statement.of("SELECT COUNT(*) AS C FROM TEST"))) { + assertThat(rs.next()).isTrue(); + assertThat(rs.getLong("C")).isEqualTo(expected); + assertThat(rs.next()).isFalse(); + } + } + + /** Test successful retry when the commit aborts */ + @Test + public void testCommitAborted() { + AbortInterceptor interceptor = new AbortInterceptor(0); + try (ITConnection connection = + createConnection(interceptor, new CountTransactionRetryListener())) { + ApiFuture count = getTestRecordCountAsync(connection); + // do an insert + ApiFuture updateCount = + connection.executeUpdateAsync( + Statement.of("INSERT INTO TEST (ID, NAME) VALUES (1, 'test aborted')")); + // indicate that the next statement should abort + interceptor.setProbability(1.0); + interceptor.setOnlyInjectOnce(true); + // do a commit that will first abort, and then on retry will succeed + ApiFuture commit = connection.commitAsync(); + + assertThat(get(count)).isEqualTo(0L); + // Wait until the commit has finished before checking retry stats. + assertThat(get(commit)).isNull(); + assertThat(get(updateCount)).isEqualTo(1L); + assertThat(RETRY_STATISTICS.totalRetryAttemptsStarted >= 1).isTrue(); + assertThat(RETRY_STATISTICS.totalRetryAttemptsFinished >= 1).isTrue(); + assertThat(RETRY_STATISTICS.totalSuccessfulRetries >= 1).isTrue(); + assertThat(RETRY_STATISTICS.totalErroredRetries).isEqualTo(0); + assertThat(RETRY_STATISTICS.totalConcurrentModifications).isEqualTo(0); + assertThat(RETRY_STATISTICS.totalMaxAttemptsExceeded).isEqualTo(0); + // verify that the insert succeeded + verifyRecordCount(connection, 1L); + } + } + + /** Test successful retry when an insert statement aborts */ + @Test + public void testInsertAborted() { + AbortInterceptor interceptor = new AbortInterceptor(0); + try (ITConnection connection = + createConnection(interceptor, new CountTransactionRetryListener())) { + ApiFuture count = getTestRecordCountAsync(connection); + // indicate that the next statement should abort + interceptor.setProbability(1.0); + interceptor.setOnlyInjectOnce(true); + // do an insert that will abort + connection.executeUpdateAsync( + Statement.of("INSERT INTO TEST (ID, NAME) VALUES (1, 'test aborted')")); + // do a commit + ApiFuture commit = connection.commitAsync(); + assertThat(get(count)).isEqualTo(0L); + assertThat(get(commit)).isNull(); + assertThat(RETRY_STATISTICS.totalSuccessfulRetries >= 1).isTrue(); + // verify that the insert succeeded + verifyRecordCount(connection, 1L); + } + } + + /** Test successful retry when an update statement aborts */ + @Test + public void testUpdateAborted() { + AbortInterceptor interceptor = new AbortInterceptor(0); + try (ITConnection connection = + createConnection(interceptor, new CountTransactionRetryListener())) { + ApiFuture count = getTestRecordCountAsync(connection); + // insert a test record + connection.executeUpdateAsync( + Statement.of("INSERT INTO TEST (ID, NAME) VALUES (1, 'test aborted')")); + // indicate that the next statement should abort + interceptor.setProbability(1.0); + interceptor.setOnlyInjectOnce(true); + // do an update that will abort + connection.executeUpdateAsync( + Statement.of("UPDATE TEST SET NAME='update aborted' WHERE ID=1")); + // do a commit + ApiFuture commit = connection.commitAsync(); + assertThat(get(count)).isEqualTo(0L); + assertThat(get(commit)).isNull(); + assertThat(RETRY_STATISTICS.totalSuccessfulRetries >= 1).isTrue(); + // verify that the update succeeded + try (AsyncResultSet rs = + connection.executeQueryAsync( + Statement.of( + "SELECT COUNT(*) AS C FROM TEST WHERE ID=1 AND NAME='update aborted'"))) { + assertThat(rs.next()).isTrue(); + assertThat(rs.getLong("C")).isEqualTo(1L); + assertThat(rs.next()).isFalse(); + } + } + } + + /** Test successful retry when a query aborts */ + @Test + public void testQueryAborted() { + AbortInterceptor interceptor = new AbortInterceptor(0); + try (ITConnection connection = + createConnection(interceptor, new CountTransactionRetryListener())) { + // insert a test record + connection.executeUpdateAsync( + Statement.of("INSERT INTO TEST (ID, NAME) VALUES (1, 'test aborted')")); + // indicate that the next statement should abort + interceptor.setProbability(1.0); + interceptor.setOnlyInjectOnce(true); + // do a query that will abort + final SettableApiFuture countAfterInsert = SettableApiFuture.create(); + try (AsyncResultSet rs = + connection.executeQueryAsync(Statement.of("SELECT COUNT(*) AS C FROM TEST WHERE ID=1"))) { + rs.setCallback( + executor, + new ReadyCallback() { + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + try { + while (true) { + switch (resultSet.tryNext()) { + case DONE: + return CallbackResponse.DONE; + case NOT_READY: + return CallbackResponse.CONTINUE; + case OK: + countAfterInsert.set(resultSet.getLong("C")); + break; + } + } + } catch (Throwable t) { + countAfterInsert.setException(t); + return CallbackResponse.DONE; + } + } + }); + } + connection.commitAsync(); + assertThat(get(countAfterInsert)).isEqualTo(1L); + assertThat(RETRY_STATISTICS.totalSuccessfulRetries >= 1).isTrue(); + // verify that the update succeeded + try (ResultSet rs = + connection.executeQueryAsync(Statement.of("SELECT COUNT(*) AS C FROM TEST WHERE ID=1"))) { + assertThat(rs.next()).isTrue(); + assertThat(rs.getLong("C")).isEqualTo(1L); + assertThat(rs.next()).isFalse(); + } + } + } + + /** Test successful retry when a call to {@link ResultSet#next()} aborts */ + @Test + public void testNextCallAborted() { + AbortInterceptor interceptor = new AbortInterceptor(0); + try (ITConnection connection = + createConnection(interceptor, new CountTransactionRetryListener())) { + // insert two test records + connection.executeUpdateAsync( + Statement.of("INSERT INTO TEST (ID, NAME) VALUES (1, 'test 1')")); + connection.executeUpdateAsync( + Statement.of("INSERT INTO TEST (ID, NAME) VALUES (2, 'test 2')")); + // do a query + try (AsyncResultSet rs = + connection.executeQueryAsync(Statement.of("SELECT * FROM TEST ORDER BY ID"))) { + // the first record should be accessible without any problems + assertThat(rs.next()).isTrue(); + assertThat(rs.getLong("ID")).isEqualTo(1L); + + // indicate that the next statement should abort + interceptor.setProbability(1.0); + interceptor.setOnlyInjectOnce(true); + assertThat(rs.next()).isTrue(); + assertThat(rs.getLong("ID")).isEqualTo(2L); + assertThat(RETRY_STATISTICS.totalSuccessfulRetries >= 1).isTrue(); + // there should be only two records + assertThat(rs.next()).isFalse(); + } + connection.commitAsync(); + // verify that the transaction succeeded + verifyRecordCount(connection, 2L); + } + } + + /** Test successful retry after multiple aborts */ + @Test + public void testMultipleAborts() { + AbortInterceptor interceptor = new AbortInterceptor(0); + try (ITConnection connection = + createConnection(interceptor, new CountTransactionRetryListener())) { + ApiFuture count = getTestRecordCountAsync(connection); + // do three inserts which all will abort and retry + interceptor.setProbability(1.0); + interceptor.setOnlyInjectOnce(true); + get( + connection.executeUpdateAsync( + Statement.of("INSERT INTO TEST (ID, NAME) VALUES (1, 'test 1')"))); + interceptor.setProbability(1.0); + interceptor.setOnlyInjectOnce(true); + get( + connection.executeUpdateAsync( + Statement.of("INSERT INTO TEST (ID, NAME) VALUES (2, 'test 2')"))); + interceptor.setProbability(1.0); + interceptor.setOnlyInjectOnce(true); + get( + connection.executeUpdateAsync( + Statement.of("INSERT INTO TEST (ID, NAME) VALUES (3, 'test 3')"))); + + ApiFuture commit = connection.commitAsync(); + assertThat(get(count)).isEqualTo(0L); + assertThat(get(commit)).isNull(); + assertThat(RETRY_STATISTICS.totalSuccessfulRetries).isAtLeast(3); + // verify that the inserts succeeded + verifyRecordCount(connection, 3L); + } + } + + /** + * Tests that a transaction retry can be successful after a select, as long as the select returns + * the same results during the retry + */ + @Test + public void testAbortAfterSelect() { + AbortInterceptor interceptor = new AbortInterceptor(0); + try (ITConnection connection = + createConnection(interceptor, new CountTransactionRetryListener())) { + ApiFuture count = getTestRecordCountAsync(connection); + // insert a test record + connection.executeUpdateAsync( + Statement.of("INSERT INTO TEST (ID, NAME) VALUES (1, 'test 1')")); + // select the test record + final SettableApiFuture initialRecord = SettableApiFuture.create(); + try (AsyncResultSet rs = + connection.executeQueryAsync(Statement.of("SELECT * FROM TEST WHERE ID=1"))) { + rs.setCallback( + executor, + new ReadyCallback() { + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + try { + while (true) { + switch (resultSet.tryNext()) { + case DONE: + return CallbackResponse.DONE; + case NOT_READY: + return CallbackResponse.CONTINUE; + case OK: + initialRecord.set(resultSet.getCurrentRowAsStruct()); + } + } + } catch (Throwable t) { + initialRecord.setException(t); + return CallbackResponse.DONE; + } + } + }); + } + // do another insert that will abort and retry + interceptor.setProbability(1.0); + interceptor.setOnlyInjectOnce(true); + connection.executeUpdateAsync( + Statement.of("INSERT INTO TEST (ID, NAME) VALUES (2, 'test 2')")); + + // select the first test record again + final SettableApiFuture secondRecord = SettableApiFuture.create(); + try (AsyncResultSet rs = + connection.executeQueryAsync(Statement.of("SELECT * FROM TEST WHERE ID=1"))) { + rs.setCallback( + executor, + new ReadyCallback() { + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + try { + while (true) { + switch (resultSet.tryNext()) { + case DONE: + return CallbackResponse.DONE; + case NOT_READY: + return CallbackResponse.CONTINUE; + case OK: + secondRecord.set(resultSet.getCurrentRowAsStruct()); + } + } + } catch (Throwable t) { + secondRecord.setException(t); + return CallbackResponse.DONE; + } + } + }); + } + ApiFuture commit = connection.commitAsync(); + assertThat(get(count)).isEqualTo(0L); + assertThat(get(initialRecord)).isEqualTo(get(secondRecord)); + assertThat(get(commit)).isNull(); + assertThat(RETRY_STATISTICS.totalSuccessfulRetries >= 1).isTrue(); + } + } + + /** + * Test a successful retry when a {@link ResultSet} has been consumed half way. The {@link + * ResultSet} should still be at the same position and still behave as if the original transaction + * did not abort. + */ + @Test + public void testAbortWithResultSetHalfway() { + AbortInterceptor interceptor = new AbortInterceptor(0); + try (ITConnection connection = + createConnection(interceptor, new CountTransactionRetryListener())) { + // insert two test records + connection.executeUpdateAsync( + Statement.of("INSERT INTO TEST (ID, NAME) VALUES (1, 'test 1')")); + connection.executeUpdateAsync( + Statement.of("INSERT INTO TEST (ID, NAME) VALUES (2, 'test 2')")); + // select the test records + try (AsyncResultSet rs = + connection.executeQueryAsync(Statement.of("SELECT * FROM TEST ORDER BY ID"))) { + // iterate one step + assertThat(rs.next()).isTrue(); + assertThat(rs.getLong("ID")).isEqualTo(1L); + // do another insert that will abort and retry + interceptor.setProbability(1.0); + interceptor.setOnlyInjectOnce(true); + connection.executeUpdateAsync( + Statement.of("INSERT INTO TEST (ID, NAME) VALUES (3, 'test 3')")); + // iterate another step + assertThat(rs.next()).isTrue(); + assertThat(rs.getLong("ID")).isEqualTo(2L); + // ensure we are at the end of the result set + assertThat(rs.next()).isFalse(); + } + get(connection.commitAsync()); + assertThat(RETRY_STATISTICS.totalSuccessfulRetries).isAtLeast(1); + // verify that all the inserts succeeded + verifyRecordCount(connection, 3L); + } + } + + /** Test successful retry after a {@link ResultSet} has been fully consumed. */ + @Test + public void testAbortWithResultSetFullyConsumed() { + AbortInterceptor interceptor = new AbortInterceptor(0); + try (ITConnection connection = + createConnection(interceptor, new CountTransactionRetryListener())) { + // insert two test records + connection.executeUpdateAsync( + Statement.of("INSERT INTO TEST (ID, NAME) VALUES (1, 'test 1')")); + connection.executeUpdateAsync( + Statement.of("INSERT INTO TEST (ID, NAME) VALUES (2, 'test 2')")); + // select the test records and iterate over them + try (AsyncResultSet rs = + connection.executeQueryAsync(Statement.of("SELECT * FROM TEST ORDER BY ID"))) { + // do nothing, just consume the result set + rs.setCallback( + executor, + new ReadyCallback() { + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + while (true) { + switch (resultSet.tryNext()) { + case DONE: + return CallbackResponse.DONE; + case NOT_READY: + return CallbackResponse.CONTINUE; + case OK: + break; + } + } + } + }); + } + // do another insert that will abort and retry + interceptor.setProbability(1.0); + interceptor.setOnlyInjectOnce(true); + connection.executeUpdateAsync( + Statement.of("INSERT INTO TEST (ID, NAME) VALUES (3, 'test 3')")); + get(connection.commitAsync()); + + assertThat(RETRY_STATISTICS.totalSuccessfulRetries).isAtLeast(1); + // verify that all the inserts succeeded + verifyRecordCount(connection, 3L); + } + } + + @Test + public void testAbortWithConcurrentInsert() { + assumeFalse("concurrent transactions are not supported on the emulator", isUsingEmulator()); + AbortInterceptor interceptor = new AbortInterceptor(0); + try (ITConnection connection = + createConnection(interceptor, new CountTransactionRetryListener())) { + // insert two test records + connection.executeUpdateAsync( + Statement.of("INSERT INTO TEST (ID, NAME) VALUES (1, 'test 1')")); + connection.executeUpdateAsync( + Statement.of("INSERT INTO TEST (ID, NAME) VALUES (2, 'test 2')")); + // select the test records and consume the entire result set + try (AsyncResultSet rs = + connection.executeQueryAsync(Statement.of("SELECT * FROM TEST ORDER BY ID"))) { + get( + rs.setCallback( + executor, + new ReadyCallback() { + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + while (true) { + switch (resultSet.tryNext()) { + case DONE: + return CallbackResponse.DONE; + case NOT_READY: + return CallbackResponse.CONTINUE; + case OK: + break; + } + } + } + })); + } + // open a new connection and transaction and do an additional insert + try (ITConnection connection2 = createConnection()) { + connection2.executeUpdateAsync( + Statement.of("INSERT INTO TEST (ID, NAME) VALUES (3, 'test 3')")); + get(connection2.commitAsync()); + } + // now try to do an insert that will abort. The retry should now fail as there has been a + // concurrent modification + interceptor.setProbability(1.0); + interceptor.setOnlyInjectOnce(true); + ApiFuture updateCount = + connection.executeUpdateAsync( + Statement.of("INSERT INTO TEST (ID, NAME) VALUES (4, 'test 4')")); + try { + get(updateCount); + fail("Missing expected exception"); + } catch (AbortedDueToConcurrentModificationException e) { + assertRetryStatistics(1, 1, 0); + } + } + } + + @Test + public void testAbortWithConcurrentDelete() { + assumeFalse("concurrent transactions are not supported on the emulator", isUsingEmulator()); + AbortInterceptor interceptor = new AbortInterceptor(0); + // first insert two test records + try (ITConnection connection = createConnection()) { + connection.executeUpdateAsync( + Statement.of("INSERT INTO TEST (ID, NAME) VALUES (1, 'test 1')")); + connection.executeUpdateAsync( + Statement.of("INSERT INTO TEST (ID, NAME) VALUES (2, 'test 2')")); + get(connection.commitAsync()); + } + // open a new connection and select the two test records + try (ITConnection connection = + createConnection(interceptor, new CountTransactionRetryListener())) { + // select the test records and consume the entire result set + try (AsyncResultSet rs = + connection.executeQueryAsync(Statement.of("SELECT * FROM TEST ORDER BY ID"))) { + get( + rs.setCallback( + executor, + new ReadyCallback() { + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + while (true) { + switch (resultSet.tryNext()) { + case DONE: + return CallbackResponse.DONE; + case NOT_READY: + return CallbackResponse.CONTINUE; + case OK: + break; + } + } + } + })); + } + // open a new connection and transaction and remove one of the test records + try (ITConnection connection2 = createConnection()) { + connection2.executeUpdateAsync(Statement.of("DELETE FROM TEST WHERE ID=1")); + get(connection2.commitAsync()); + } + // now try to do an insert that will abort. The retry should now fail as there has been a + // concurrent modification + interceptor.setProbability(1.0); + interceptor.setOnlyInjectOnce(true); + try { + get( + connection.executeUpdateAsync( + Statement.of("INSERT INTO TEST (ID, NAME) VALUES (3, 'test 3')"))); + fail("Missing expected exception"); + } catch (AbortedDueToConcurrentModificationException e) { + assertRetryStatistics(1, 1, 0); + } + } + } + + @Test + public void testAbortWithConcurrentUpdate() { + assumeFalse("concurrent transactions are not supported on the emulator", isUsingEmulator()); + AbortInterceptor interceptor = new AbortInterceptor(0); + // first insert two test records + try (ITConnection connection = createConnection()) { + connection.executeUpdateAsync( + Statement.of("INSERT INTO TEST (ID, NAME) VALUES (1, 'test 1')")); + connection.executeUpdateAsync( + Statement.of("INSERT INTO TEST (ID, NAME) VALUES (2, 'test 2')")); + get(connection.commitAsync()); + } + // open a new connection and select the two test records + try (ITConnection connection = + createConnection(interceptor, new CountTransactionRetryListener())) { + // select the test records and consume the entire result set + try (AsyncResultSet rs = + connection.executeQueryAsync(Statement.of("SELECT * FROM TEST ORDER BY ID"))) { + get( + rs.setCallback( + executor, + new ReadyCallback() { + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + while (true) { + switch (resultSet.tryNext()) { + case DONE: + return CallbackResponse.DONE; + case NOT_READY: + return CallbackResponse.CONTINUE; + case OK: + break; + } + } + } + })); + } + // open a new connection and transaction and update one of the test records + try (ITConnection connection2 = createConnection()) { + connection2.executeUpdateAsync( + Statement.of("UPDATE TEST SET NAME='test updated' WHERE ID=2")); + get(connection2.commitAsync()); + } + // now try to do an insert that will abort. The retry should now fail as there has been a + // concurrent modification + interceptor.setProbability(1.0); + interceptor.setOnlyInjectOnce(true); + try { + get( + connection.executeUpdateAsync( + Statement.of("INSERT INTO TEST (ID, NAME) VALUES (3, 'test 3')"))); + fail("Missing expected exception"); + } catch (AbortedDueToConcurrentModificationException e) { + assertRetryStatistics(1, 1, 0); + } + } + } + + /** + * Test that shows that a transaction retry is possible even when there is a concurrent insert + * that has an impact on a query that has been executed, as long as the user hasn't actually seen + * the relevant part of the result of the query + */ + @Test + public void testAbortWithUnseenConcurrentInsert() throws InterruptedException { + assumeFalse("concurrent transactions are not supported on the emulator", isUsingEmulator()); + AbortInterceptor interceptor = new AbortInterceptor(0); + try (ITConnection connection = + createConnection(interceptor, new CountTransactionRetryListener())) { + // insert three test records + connection.executeUpdateAsync( + Statement.of("INSERT INTO TEST (ID, NAME) VALUES (1, 'test 1')")); + connection.executeUpdateAsync( + Statement.of("INSERT INTO TEST (ID, NAME) VALUES (2, 'test 2')")); + connection.executeUpdateAsync( + Statement.of("INSERT INTO TEST (ID, NAME) VALUES (3, 'test 3')")); + // select the test records and consume part of the result set + final AtomicInteger count = new AtomicInteger(); + final AtomicLong lastSeenId = new AtomicLong(); + final CountDownLatch latch1 = new CountDownLatch(1); + final CountDownLatch latch2 = new CountDownLatch(1); + // Use buffer size 1. This means that the underlying result set will see 2 records (1 in the + // buffer and 1 waiting to be put in the buffer). + try (AsyncResultSet rs = + connection.executeQueryAsync( + Statement.of("SELECT * FROM TEST ORDER BY ID"), Options.bufferRows(1))) { + ApiFuture finished = + rs.setCallback( + executor, + new ReadyCallback() { + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + try { + while (true) { + switch (resultSet.tryNext()) { + case DONE: + return CallbackResponse.DONE; + case NOT_READY: + return CallbackResponse.CONTINUE; + case OK: + count.incrementAndGet(); + lastSeenId.set(resultSet.getLong("ID")); + break; + } + if (count.get() == 1) { + // Let the other transaction proceed. + latch1.countDown(); + // Wait until the transaction has been aborted and retried. + if (!latch2.await(120L, TimeUnit.SECONDS)) { + throw SpannerExceptionFactory.newSpannerException( + ErrorCode.DEADLINE_EXCEEDED, "Timeout while waiting for latch2"); + } + } + } + } catch (Throwable t) { + throw SpannerExceptionFactory.asSpannerException(t); + } + } + }); + // Open a new connection and transaction and do an additional insert. This insert will be + // included in a retry of the above query, but this has not yet been 'seen' by the user, + // hence is not a problem for retrying the transaction. + try (ITConnection connection2 = createConnection()) { + assertThat(latch1.await(60L, TimeUnit.SECONDS)).isTrue(); + connection2.executeUpdateAsync( + Statement.of("INSERT INTO TEST (ID, NAME) VALUES (4, 'test 4')")); + get(connection2.commitAsync()); + } + // now try to do an insert that will abort. The retry should still succeed. + interceptor.setProbability(1.0); + interceptor.setOnlyInjectOnce(true); + int currentRetryCount = RETRY_STATISTICS.totalRetryAttemptsStarted; + get( + connection.executeUpdateAsync( + Statement.of("INSERT INTO TEST (ID, NAME) VALUES (5, 'test 5')"))); + assertThat(RETRY_STATISTICS.totalRetryAttemptsStarted).isAtLeast(currentRetryCount + 1); + // Consume the rest of the result set. The insert by the other transaction should now be + // included in the result set as the transaction retried. Although this means that the + // result + // is different after a retry, it is not different as seen by the user, as the user didn't + // know that the result set did not have any more results before the transaction retry. + latch2.countDown(); + get(finished); + // record with id 5 should not be visible, as it was added to the transaction after the + // query + // was executed + assertThat(count.get()).isEqualTo(4); + assertThat(lastSeenId.get()).isEqualTo(4L); + } + get(connection.commitAsync()); + assertThat(RETRY_STATISTICS.totalSuccessfulRetries).isAtLeast(1); + } + } + + /** Test the successful retry of a transaction with a large {@link ResultSet} */ + @Test + public void testRetryLargeResultSet() { + final int NUMBER_OF_TEST_RECORDS = 100000; + final long UPDATED_RECORDS = 1000L; + AbortInterceptor interceptor = new AbortInterceptor(0); + try (ITConnection connection = createConnection()) { + // insert test records + for (int i = 0; i < NUMBER_OF_TEST_RECORDS; i++) { + connection.bufferedWrite( + Mutation.newInsertBuilder("TEST").set("ID").to(i).set("NAME").to("test " + i).build()); + if (i % 1000 == 0) { + connection.commitAsync(); + } + } + get(connection.commitAsync()); + } + try (ITConnection connection = + createConnection(interceptor, new CountTransactionRetryListener())) { + // select the test records and iterate over them + try (AsyncResultSet rs = + connection.executeQueryAsync(Statement.of("SELECT * FROM TEST ORDER BY ID"))) { + ApiFuture finished = + rs.setCallback( + executor, + new ReadyCallback() { + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + // do nothing, just consume the result set + while (true) { + switch (resultSet.tryNext()) { + case DONE: + return CallbackResponse.DONE; + case NOT_READY: + return CallbackResponse.CONTINUE; + case OK: + break; + } + } + } + }); + // Wait until the entire result set has been consumed. + get(finished); + } + // Do an update that will abort and retry. + interceptor.setProbability(1.0); + interceptor.setOnlyInjectOnce(true); + connection.executeUpdateAsync( + Statement.newBuilder("UPDATE TEST SET NAME='updated' WHERE ID<@max_id") + .bind("max_id") + .to(UPDATED_RECORDS) + .build()); + connection.commitAsync(); + // verify that the update succeeded + try (AsyncResultSet rs = + connection.executeQueryAsync( + Statement.of("SELECT COUNT(*) AS C FROM TEST WHERE NAME='updated'"))) { + assertThat(rs.next()).isTrue(); + assertThat(rs.getLong("C")).isEqualTo(UPDATED_RECORDS); + assertThat(rs.next()).isFalse(); + } + // Verify that the transaction retried. + assertRetryStatistics(1, 0, 1); + } + } + + /** Test the successful retry of a transaction with a high chance of multiple aborts */ + @Test + public void testRetryHighAbortRate() { + final int NUMBER_OF_TEST_RECORDS = 10000; + final long UPDATED_RECORDS = 1000L; + // abort on 25% of all statements + AbortInterceptor interceptor = new AbortInterceptor(0.25D); + try (ITConnection connection = + createConnection(interceptor, new CountTransactionRetryListener())) { + // insert test records + for (int i = 0; i < NUMBER_OF_TEST_RECORDS; i++) { + connection.bufferedWrite( + Mutation.newInsertBuilder("TEST").set("ID").to(i).set("NAME").to("test " + i).build()); + if (i % 1000 == 0) { + connection.commitAsync(); + } + } + connection.commitAsync(); + // select the test records and iterate over them + // reduce the abort rate to 0.01% as each next() call could abort + interceptor.setProbability(0.0001D); + try (AsyncResultSet rs = + connection.executeQueryAsync(Statement.of("SELECT * FROM TEST ORDER BY ID"))) { + ApiFuture finished = + rs.setCallback( + executor, + new ReadyCallback() { + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + // do nothing, just consume the result set + while (true) { + switch (resultSet.tryNext()) { + case DONE: + return CallbackResponse.DONE; + case NOT_READY: + return CallbackResponse.CONTINUE; + case OK: + break; + } + } + } + }); + // Wait until the entire result set has been consumed. + get(finished); + } + // increase the abort rate to 50% + interceptor.setProbability(0.50D); + connection.executeUpdateAsync( + Statement.newBuilder("UPDATE TEST SET NAME='updated' WHERE ID<@max_id") + .bind("max_id") + .to(UPDATED_RECORDS) + .build()); + // Wait for the commit to finish, as it could be that the transaction is aborted so many times + // that the last update does not succeed. + get(connection.commitAsync()); + // verify that the update succeeded + try (AsyncResultSet rs = + connection.executeQueryAsync( + Statement.of("SELECT COUNT(*) AS C FROM TEST WHERE NAME='updated'"))) { + assertThat(rs.next()).isTrue(); + assertThat(rs.getLong("C")).isEqualTo(UPDATED_RECORDS); + assertThat(rs.next()).isFalse(); + } + get(connection.commitAsync()); + } catch (AbortedException e) { + // This could happen if the number of aborts exceeds the max number of retries. + logger.log(Level.FINE, "testRetryHighAbortRate aborted because of too many retries", e); + } + logger.fine("Total number of retries started: " + RETRY_STATISTICS.totalRetryAttemptsStarted); + logger.fine("Total number of retries finished: " + RETRY_STATISTICS.totalRetryAttemptsFinished); + logger.fine("Total number of retries successful: " + RETRY_STATISTICS.totalSuccessfulRetries); + logger.fine("Total number of retries aborted: " + RETRY_STATISTICS.totalNestedAborts); + logger.fine( + "Total number of times the max retry count was exceeded: " + + RETRY_STATISTICS.totalMaxAttemptsExceeded); + } + + private void assertRetryStatistics( + int minAttemptsStartedExpected, + int concurrentModificationsExpected, + int successfulRetriesExpected) { + assertThat(RETRY_STATISTICS.totalRetryAttemptsStarted).isAtLeast(minAttemptsStartedExpected); + assertThat(RETRY_STATISTICS.totalConcurrentModifications) + .isEqualTo(concurrentModificationsExpected); + assertThat(RETRY_STATISTICS.totalSuccessfulRetries).isAtLeast(successfulRetriesExpected); + } +} diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/it/ITReadOnlySpannerTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/it/ITReadOnlySpannerTest.java index d6c89c65d9..899771b9e5 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/it/ITReadOnlySpannerTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/it/ITReadOnlySpannerTest.java @@ -39,7 +39,6 @@ import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.TimeUnit; -import java.util.logging.Logger; import org.junit.Before; import org.junit.Test; import org.junit.experimental.categories.Category; @@ -53,7 +52,6 @@ @Category(ParallelIntegrationTest.class) @RunWith(JUnit4.class) public class ITReadOnlySpannerTest extends ITAbstractSpannerTest { - private static final Logger logger = Logger.getLogger(ITReadOnlySpannerTest.class.getName()); private static final long TEST_ROWS_COUNT = 1000L; @Override @@ -126,30 +124,21 @@ public void testStatementTimeoutTransactional() { @Test public void testStatementTimeoutTransactionalMultipleStatements() { - long startTime = System.currentTimeMillis(); try (ITConnection connection = createConnection()) { connection.beginTransaction(); for (int i = 0; i < 3; i++) { - boolean timedOut = false; - connection.setStatementTimeout(1L, TimeUnit.MILLISECONDS); + connection.setStatementTimeout(1L, TimeUnit.MICROSECONDS); try (ResultSet rs = connection.executeQuery( Statement.of( "SELECT (SELECT COUNT(*) FROM PRIME_NUMBERS)/(SELECT COUNT(*) FROM NUMBERS) AS PRIME_NUMBER_RATIO"))) { + fail("Missing expected exception"); } catch (SpannerException e) { - timedOut = e.getErrorCode() == ErrorCode.DEADLINE_EXCEEDED; + assertThat(e.getErrorCode(), is(ErrorCode.DEADLINE_EXCEEDED)); } - assertThat(timedOut, is(true)); } connection.commit(); } - long endTime = System.currentTimeMillis(); - long executionTime = endTime - startTime; - if (executionTime > 25L) { - logger.warning("Total test execution time exceeded 25 milliseconds: " + executionTime); - } else { - logger.info("Total test execution time: " + executionTime); - } } @Test diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/it/ITSqlMusicScriptTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/it/ITSqlMusicScriptTest.java index a6b4fc8873..e8a479c6d6 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/it/ITSqlMusicScriptTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/it/ITSqlMusicScriptTest.java @@ -186,6 +186,8 @@ public void test02_RunAbortedTest() { // verify that the commit aborted, an internal retry was started and then aborted because of // the concurrent modification assertThat(expectedException, is(true)); + // Rollback the transaction to start a new one. + connection.rollback(); // verify that the prices were changed try (ResultSet rs = connection.executeQuery( diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/it/ITTransactionRetryTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/it/ITTransactionRetryTest.java index c1567496bc..1d7de23cb4 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/it/ITTransactionRetryTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/it/ITTransactionRetryTest.java @@ -767,7 +767,8 @@ public void testAbortWithConcurrentInsertAndContinue() { } assertThat(expectedException, is(true)); assertRetryStatistics(1, 1, 0); - // the next statement should be in a new transaction as the previous transaction rolled back + // Rollback the aborted transaction to start a new one. + connection.rollback(); try (ResultSet rs = connection.executeQuery(Statement.of("SELECT * FROM TEST"))) { // there should be one record from the transaction on connection2 assertThat(rs.next(), is(true)); @@ -1519,6 +1520,7 @@ public void testRetryHighAbortRate() { @Test public void testAbortWithConcurrentInsertOnEmptyTable() { assumeFalse("concurrent transactions are not supported on the emulator", isUsingEmulator()); + AbortInterceptor interceptor = new AbortInterceptor(0); try (ITConnection connection = createConnection(interceptor, new CountTransactionRetryListener())) { diff --git a/google-cloud-spanner/src/test/resources/com/google/cloud/spanner/connection/ITSqlScriptTest_TestStatementTimeout.sql b/google-cloud-spanner/src/test/resources/com/google/cloud/spanner/connection/ITSqlScriptTest_TestStatementTimeout.sql index 9a9894fafa..7e8d907b95 100644 --- a/google-cloud-spanner/src/test/resources/com/google/cloud/spanner/connection/ITSqlScriptTest_TestStatementTimeout.sql +++ b/google-cloud-spanner/src/test/resources/com/google/cloud/spanner/connection/ITSqlScriptTest_TestStatementTimeout.sql @@ -70,7 +70,7 @@ SET STATEMENT_TIMEOUT='1ns'; SHOW VARIABLE STATEMENT_TIMEOUT; -- Do a somewhat complex query that should now timeout -@EXPECT EXCEPTION DEADLINE_EXCEEDED 'DEADLINE_EXCEEDED: Statement execution timeout occurred' +@EXPECT EXCEPTION DEADLINE_EXCEEDED 'DEADLINE_EXCEEDED:' SELECT COUNT(*) AS ACTUAL, 0 AS EXPECTED FROM ( SELECT * @@ -97,7 +97,7 @@ FROM ( ; -- Try to execute an update that should also timeout -@EXPECT EXCEPTION DEADLINE_EXCEEDED 'DEADLINE_EXCEEDED: Statement execution timeout occurred' +@EXPECT EXCEPTION DEADLINE_EXCEEDED 'DEADLINE_EXCEEDED:' UPDATE Singers SET LastName='Some Other Last Name' /* It used to be 'Last 1' */ WHERE SingerId=1 OR LastName IN ( @@ -176,7 +176,7 @@ SET STATEMENT_TIMEOUT='1ns'; SHOW VARIABLE STATEMENT_TIMEOUT; -- Do a somewhat complex query that should now timeout -@EXPECT EXCEPTION DEADLINE_EXCEEDED 'DEADLINE_EXCEEDED: Statement execution timeout occurred' +@EXPECT EXCEPTION DEADLINE_EXCEEDED 'DEADLINE_EXCEEDED:' SELECT COUNT(*) AS ACTUAL, 0 AS EXPECTED FROM ( SELECT * @@ -202,11 +202,17 @@ FROM ( ) RES ; -- We need to rollback the transaction as it is no longer usable. -@EXPECT EXCEPTION DEADLINE_EXCEEDED 'DEADLINE_EXCEEDED: Statement execution timeout occurred' +-- A timeout during a rollback is ignored, and also not rolling back +-- a transaction on the emulator will make the transaction remain the +-- current transaction. We therefore remove the timeout before the +-- rollback call. +SET STATEMENT_TIMEOUT=null; ROLLBACK; +SET STATEMENT_TIMEOUT='1ns'; + -- Try to execute an update that should also timeout -@EXPECT EXCEPTION DEADLINE_EXCEEDED 'DEADLINE_EXCEEDED: Statement execution timeout occurred' +@EXPECT EXCEPTION DEADLINE_EXCEEDED 'DEADLINE_EXCEEDED:' UPDATE Singers SET LastName='Some Other Last Name' /* It used to be 'Last 1' */ WHERE SingerId=1 OR LastName IN (