From e5e8df8c8b3cc17f5abe60d89994a5074e1f5db5 Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Thu, 20 Feb 2020 08:29:48 +0100 Subject: [PATCH 01/49] feat: add async api --- .../cloud/spanner/AbstractReadContext.java | 11 + .../google/cloud/spanner/AsyncResultSet.java | 232 +++++++++ .../cloud/spanner/AsyncResultSetImpl.java | 485 ++++++++++++++++++ .../cloud/spanner/DatabaseClientImpl.java | 8 +- .../com/google/cloud/spanner/ReadContext.java | 2 + .../com/google/cloud/spanner/SessionPool.java | 448 +++++++++++----- .../com/google/cloud/spanner/SpannerImpl.java | 6 + .../cloud/spanner/TransactionRunnerImpl.java | 2 + .../spanner/AsyncResultSetImplStressTest.java | 462 +++++++++++++++++ .../cloud/spanner/AsyncResultSetImplTest.java | 440 ++++++++++++++++ .../cloud/spanner/DatabaseClientImplTest.java | 176 +++++++ .../IntegrationTestWithClosedSessionsEnv.java | 29 +- .../spanner/RandomResultSetGenerator.java | 166 ++++++ .../cloud/spanner/SessionPoolStressTest.java | 13 +- .../google/cloud/spanner/SessionPoolTest.java | 80 +-- .../spanner/TransactionContextImplTest.java | 1 + 16 files changed, 2384 insertions(+), 177 deletions(-) create mode 100644 google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncResultSet.java create mode 100644 google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncResultSetImpl.java create mode 100644 google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncResultSetImplStressTest.java create mode 100644 google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncResultSetImplTest.java create mode 100644 google-cloud-spanner/src/test/java/com/google/cloud/spanner/RandomResultSetGenerator.java diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java index 685e9a1e31b..7f7d02f95d5 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java @@ -22,6 +22,7 @@ import static com.google.common.base.Preconditions.checkState; import com.google.cloud.Timestamp; +import com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory; import com.google.cloud.spanner.AbstractResultSet.CloseableIterator; import com.google.cloud.spanner.AbstractResultSet.GrpcResultSet; import com.google.cloud.spanner.AbstractResultSet.GrpcStreamIterator; @@ -45,6 +46,7 @@ import io.opencensus.trace.Span; import io.opencensus.trace.Tracing; import java.util.Map; +import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.atomic.AtomicLong; import javax.annotation.Nullable; import javax.annotation.concurrent.GuardedBy; @@ -350,6 +352,7 @@ void initTransaction() { final Object lock = new Object(); final SessionImpl session; final SpannerRpc rpc; + final ExecutorFactory executorFactory; final Span span; private final int defaultPrefetchChunks; private final QueryOptions defaultQueryOptions; @@ -415,6 +418,14 @@ public final ResultSet executeQuery(Statement statement, QueryOption... options) statement, com.google.spanner.v1.ExecuteSqlRequest.QueryMode.NORMAL, options); } + @Override + public final AsyncResultSet executeQueryAsync(Statement statement, QueryOption... options) { + return new AsyncResultSetImpl( + executorFactory, + executeQueryInternal( + statement, com.google.spanner.v1.ExecuteSqlRequest.QueryMode.NORMAL, options)); + } + @Override public final ResultSet analyzeQuery(Statement statement, QueryAnalyzeMode readContextQueryMode) { switch (readContextQueryMode) { diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncResultSet.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncResultSet.java new file mode 100644 index 00000000000..cb052042257 --- /dev/null +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncResultSet.java @@ -0,0 +1,232 @@ +/* + * 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.Function; +import com.google.common.collect.ImmutableList; +import java.util.concurrent.Executor; + +public interface AsyncResultSet extends AutoCloseable, StructReader { + public static interface ReadyCallback { + CallbackResponse cursorReady(AsyncResultSet resultSet); + } + /** Response code from {@code tryNext()}. */ + public enum CursorState { + /** Cursor has been moved to a new row. */ + OK, + /** Read is complete, all rows have been consumed, and there are no more. */ + DONE, + /** No further information known at this time, thus current row not available. */ + NOT_READY + } + + @Override + void close(); + + /** + * Creates an immutable version of the row that the result set is positioned over. This may + * involve copying internal data structures, and so converting all rows to {@code Struct} objects + * is generally more expensive than processing the {@code ResultSet} directly. + */ + Struct getCurrentRowAsStruct(); + + /** + * Non-blocking call that attempts to step the cursor to the next position in the stream. The + * cursor may be inspected only if the cursor returns {@code CursorState.OK}. + * + *

A caller will typically call {@link #tryNext()} in a loop inside the ReadyCallback, + * consuming all results available. For more information see {@link #setCallback(Executor, + * ReadyCallback)}. + * + *

Currently this method may only be called if a ReadyCallback has been registered. This is for + * safety purposes only, and may be relaxed in future. + * + * @return current cursor readiness state + * @throws SpannerException When an unrecoverable problem downstream occurs. Once this occurs you + * will get no further callbacks. You should return CallbackResponse.DONE back from callback. + */ + CursorState tryNext() throws SpannerException; + + /** + * Register a callback with the ResultSet to be made aware when more data is available, changing + * the usage pattern from sync to async. Details: + * + *

+ * + *

Flow Control

+ * + * If no flow control is needed (say because result sizes are known in advance to be finite in + * size) then async processing is simple. The following is a code example that transfers work from + * the cursor to an upstream sink: + * + * 
{@code
+   * @Override
+   * public CallbackResponse cursorReady(ResultSet cursor) {
+   *   try {
+   *     while (true) {
+   *       switch (cursor.tryNext()) {
+   *         case OK:    upstream.emit(cursor.getRow()); break;
+   *         case DONE:  upstream.done(); return CallbackResponse.DONE;
+   *         case NOT_READY:  return CallbackResponse.CONTINUE;
+   *       }
+   *     }
+   *   } catch (SpannerException e) {
+   *     upstream.doneWithError(e);
+   *     return CallbackResponse.DONE;
+   *   }
+   * }
+   * }
+ * + * Flow control may be needed if for example the upstream system may not always be ready to handle + * more data. In this case the app developer has two main options: + * + * + * + * Note that it would have been equivalent to have the app be responsible for draining the cursor + * instead of calling {@code resume()} (which has basically the same effect, namely running the + * application callback.) The explicit pause and resume was chosen to make the flow control + * behavior more explicit in application code. + * + * @param exec executor on which to run all callbacks. Typically use a threadpool. If the executor + * is one that runs the work on the submitting thread, you must be very careful not to throw + * RuntimeException up the stack, lest you do damage to calling components. For example, it + * may cause an event dispatcher thread to crash. + * @param cb ready callback + */ + void setCallback(Executor exec, ReadyCallback cb); + + /** + * Attempt to cancel this operation and free all resources. Non-blocking. This is a no-op for + * child row cursors and does not cancel the parent cursor. + */ + void cancel(); + + public enum CallbackResponse { + /** + * Tell the cursor to continue issuing callbacks when data is available. This is the standard + * "I'm ready for more" response. If cursor is not completely drained of all ready results the + * callback will be called again immediately. + */ + CONTINUE, + + /** + * Tell the cursor to suspend all callbacks until application calls {@link RowCursor#resume()}. + */ + PAUSE, + + /** + * Tell the cursor you are done receiving results, even if there are more results sitting in the + * buffer. Once you return DONE, you will receive no further callbacks. + * + *

Approximately equivalent to calling {@link RowCursor#cancel()}, and then returning {@code + * PAUSE}, but more clear, immediate, and idiomatic. + * + *

It is legal to commit a transaction that owns this read before actually returning {@code + * DONE}. + */ + DONE, + } + + /** + * Resume callbacks from the cursor. If there is more data available, a callback will be + * dispatched immediately. This can be called from any thread. + */ + void resume(); + + /** + * Transforms the row cursor into an immutable list using the given transformer function. {@code + * transformer} will be called once per row, thus the returned list will contain one entry per + * row. The returned future will throw a {@link SpannerException} if the row cursor encountered + * any error or if the transformer threw an exception on any row. + * + *

The transformer will be run on the supplied executor. The implementation may batch multiple + * transformer invocations together into a single {@code Runnable} when possible to increase + * efficiency. At any point in time, there will be at most one invocation of the transformer in + * progress. + * + *

WARNING: This will result in materializing the entire list so this should be used + * judiciously after considering the memory requirements of the returned list. + * + *

WARNING: The {@code RowBase} object passed to transformer function is not immutable and is + * not guaranteed to remain valid after the transformer function returns. The same {@code RowBase} + * object might be passed multiple times to the transformer with different underlying data each + * time. So *NEVER* keep a reference to the {@code RowBase} outside of the transformer. + * Specifically do not use {@link com.google.common.base.Functions#identity()} function. + * + * @param transformer function which will be used to transform the row. It should not return null. + * @param executor executor on which the transformer will be run. This should ideally not be an + * inline executor such as {@code MoreExecutors.directExecutor()}; using such an executor may + * degrade the performance of the Spanner library. + */ + ApiFuture> toListAsync( + Function transformer, Executor executor); + + /** + * Transforms the row cursor into an immutable list using the given transformer function. {@code + * transformer} will be called once per row, thus the returned list will contain one entry per + * row. This method will block until all the rows have been yielded by the cursor. + * + *

WARNING: This will result in consuming the entire list so this should be used judiciously + * after considering the memory requirements of the returned list. + * + *

WARNING: The {@code RowBase} object passed to transformer function is not immutable and is + * not guaranteed to remain valid after the transformer function returns. The same {@code RowBase} + * object might be passed multiple times to the transformer with different underlying data each + * time. So *NEVER* keep a reference to the {@code RowBase} outside of the transformer. + * Specifically do not use {@link com.google.common.base.Functions#identity()} function. + * + * @param transformer function which will be used to transform the row. It should not return null. + */ + ImmutableList toList(Function transformer) throws SpannerException; +} 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 new file mode 100644 index 00000000000..9b7609d4dca --- /dev/null +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncResultSetImpl.java @@ -0,0 +1,485 @@ +/* + * Copyright 2020 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.spanner; + +import com.google.api.core.ApiFuture; +import com.google.api.core.SettableApiFuture; +import com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory; +import com.google.common.base.Function; +import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableList; +import com.google.common.util.concurrent.MoreExecutors; +import java.util.concurrent.BlockingDeque; +import java.util.concurrent.Callable; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executor; +import java.util.concurrent.Future; +import java.util.concurrent.LinkedBlockingDeque; +import java.util.concurrent.ScheduledExecutorService; + +/** Default implementation for {@link AsyncResultSet}. */ +class AsyncResultSetImpl extends ForwardingStructReader implements AsyncResultSet { + /** State of an {@link AsyncResultSetImpl}. */ + private enum State { + INITIALIZED, + CONSUMING, + RUNNING, + PAUSED, + CANCELLED(true), + DONE(true); + + /** Does this state mean that the result set should permanently stop producing rows. */ + private final boolean shouldStop; + + private State() { + shouldStop = false; + } + + private State(boolean shouldStop) { + this.shouldStop = shouldStop; + } + } + + private static final int DEFAULT_BUFFER_SIZE = 10; + private static final int MAX_WAIT_FOR_BUFFER_CONSUMPTION = 10; + + private final Object monitor = new Object(); + private boolean closed; + /** + * {@link ExecutorFactory} produces executors that are used to fetch data from the backend and put + * these into the buffer for further consumpation by the callback. + */ + private final ExecutorFactory executorFactory; + + private final ScheduledExecutorService service; + private final BlockingDeque buffer; + private Struct currentRow; + /** The underlying synchronous {@link ResultSet} that is producing the rows. */ + private final ResultSet delegateResultSet; + /** + * Any exception that occurs while executing the query and iterating over the result set will be + * stored in this variable and propagated to the user through {@link #tryNext()}. + */ + private volatile SpannerException executionException; + + /** + * Executor for callbacks. Regardless of the type of executor that is provided, the {@link + * AsyncResultSetImpl} will ensure that at most 1 callback call will be active at any one time. + */ + private Executor executor; + + private ReadyCallback callback; + + private State state = State.INITIALIZED; + + /** + * {@link #finished} indicates whether all the results from the underlying result set have been + * read. + */ + private volatile boolean finished; + + private final Future result; + + /** + * {@link #cursorReturnedDoneOrException} indicates whether {@link #tryNext()} has returned {@link + * CursorState#DONE} or a {@link SpannerException}. + */ + private volatile boolean cursorReturnedDoneOrException; + + /** + * {@link #pausedLatch} is used to pause the producer when the {@link AsyncResultSet} is paused. + * The production of rows that are put into the buffer is only paused once the buffer is full. + */ + private volatile CountDownLatch pausedLatch = new CountDownLatch(1); + /** + * {@link #bufferConsumptionLatch} is used to pause the producer when the buffer is full and the + * consumer needs some time to catch up. + */ + private volatile CountDownLatch bufferConsumptionLatch = new CountDownLatch(0); + /** + * {@link #consumingLatch} is used to pause the producer when all rows have been put into the + * buffer, but the consumer (the callback) has not yet received and processed all rows. + */ + private volatile CountDownLatch consumingLatch = new CountDownLatch(0); + + AsyncResultSetImpl( + ExecutorFactory executorFactory, ResultSet delegate) { + this(executorFactory, delegate, DEFAULT_BUFFER_SIZE); + } + + AsyncResultSetImpl( + ExecutorFactory executorFactory, + ResultSet delegate, + int bufferSize) { + super(delegate); + this.buffer = new LinkedBlockingDeque<>(bufferSize); + this.executorFactory = executorFactory; + this.service = executorFactory.get(); + this.delegateResultSet = delegate; + // Eagerly start to fetch data and buffer these. + this.result = this.service.submit(new ProduceRowsCallable()); + } + + /** + * Closes the {@link AsyncResultSet}. {@link #close()} is non-blocking and may be called multiple + * times without side effects. An {@link AsyncResultSet} may be closed before all rows have been + * returned to the callback, and calling {@link #tryNext()} on a closed {@link AsyncResultSet} is + * allowed as long as this is done from within a {@link ReadyCallback}. Calling {@link #resume()} + * on a closed {@link AsyncResultSet} is also allowed. + */ + @Override + public void close() { + synchronized (monitor) { + if (this.closed) { + return; + } + this.closed = true; + } + } + + public Struct getCurrentRowAsStruct() { + return currentRow; + } + + /** + * Tries to advance this {@link AsyncResultSet} to the next row. This method may only be called + * from within a {@link ReadyCallback}. + */ + @Override + public CursorState tryNext() throws SpannerException { + synchronized (monitor) { + if (state == State.CANCELLED) { + cursorReturnedDoneOrException = true; + throw SpannerExceptionFactory.newSpannerException( + ErrorCode.CANCELLED, "This AsyncResultSet has been cancelled"); + } + if (buffer.isEmpty() && executionException != null) { + cursorReturnedDoneOrException = true; + throw executionException; + } + Preconditions.checkState( + this.callback != null, "tryNext may only be called after a callback has been set."); + Preconditions.checkState( + this.state == State.CONSUMING, + "tryNext may only be called from a DataReady callback. Current state: " + + this.state.name()); + + if (finished && buffer.isEmpty()) { + cursorReturnedDoneOrException = true; + return CursorState.DONE; + } + } + if (!buffer.isEmpty()) { + // Set the next row from the buffer as the current row of the StructReader. + replaceDelegate(currentRow = buffer.pop()); + synchronized (monitor) { + bufferConsumptionLatch.countDown(); + } + return CursorState.OK; + } + return CursorState.NOT_READY; + } + + /** + * {@link CallbackRunnable} calls the {@link ReadyCallback} registered for this {@link + * AsyncResultSet}. + */ + private class CallbackRunnable implements Runnable { + @Override + public void run() { + try { + while (true) { + synchronized (monitor) { + if (cursorReturnedDoneOrException) { + break; + } + } + CallbackResponse response; + try { + response = callback.cursorReady(AsyncResultSetImpl.this); + } catch (Throwable e) { + synchronized (monitor) { + if (cursorReturnedDoneOrException + && state == State.CANCELLED + && e instanceof SpannerException + && ((SpannerException) e).getErrorCode() == ErrorCode.CANCELLED) { + // The callback did not catch the cancelled exception (which it should have), but + // we'll keep the cancelled state. + return; + } + executionException = SpannerExceptionFactory.newSpannerException(e); + cursorReturnedDoneOrException = true; + } + return; + } + synchronized (monitor) { + if (state == State.CANCELLED) { + if (cursorReturnedDoneOrException) { + return; + } + } else { + switch (response) { + case DONE: + state = State.DONE; + return; + case PAUSE: + state = State.PAUSED; + // Make sure no-one else is waiting on the current pause latch and create a new + // one. + pausedLatch.countDown(); + pausedLatch = new CountDownLatch(1); + return; + case CONTINUE: + if (buffer.isEmpty()) { + // Call the callback once more if the entire result set has been processed but + // the + // callback has not yet received a CursorState.DONE or a CANCELLED error. + if (finished && !cursorReturnedDoneOrException) { + break; + } + state = State.RUNNING; + return; + } + break; + default: + throw new IllegalStateException("Unknown response: " + response); + } + } + } + } + } finally { + synchronized (monitor) { + // Count down all latches that the producer might be waiting on. + consumingLatch.countDown(); + while (bufferConsumptionLatch.getCount() > 0L) { + bufferConsumptionLatch.countDown(); + } + } + } + } + } + + private final CallbackRunnable callbackRunnable = new CallbackRunnable(); + + /** + * {@link ProduceRowsCallable} reads data from the underlying {@link ResultSet}, places these in + * the buffer and dispatches the {@link CallbackRunnable} when data is ready to be consumed. + */ + private class ProduceRowsCallable implements Callable { + @Override + public Void call() throws Exception { + boolean stop = false; + boolean hasNext = false; + try { + hasNext = delegateResultSet.next(); + } catch (Throwable e) { + synchronized (monitor) { + executionException = SpannerExceptionFactory.newSpannerException(e); + } + } + try { + while (!stop && hasNext) { + try { + synchronized (monitor) { + stop = state.shouldStop; + } + if (!stop) { + while (buffer.remainingCapacity() == 0 && !stop) { + waitIfPaused(); + // The buffer is full and we should let the callback consume a number of rows before + // we proceed with producing any more rows to prevent us from potentially waiting on + // a full buffer repeatedly. + // Wait until at least half of the buffer is available, or if it's a bigger buffer, + // wait until at least 10 rows can be placed in it. + // TODO: Make this more dynamic / configurable? + startCallbackWithBufferLatchIfNecessary( + Math.min( + Math.min(buffer.size() / 2 + 1, buffer.size()), + MAX_WAIT_FOR_BUFFER_CONSUMPTION)); + bufferConsumptionLatch.await(); + synchronized (monitor) { + stop = state.shouldStop; + } + } + } + if (!stop) { + buffer.put(delegateResultSet.getCurrentRowAsStruct()); + startCallbackIfNecessary(); + } + hasNext = delegateResultSet.next(); + } catch (Throwable e) { + synchronized (monitor) { + executionException = SpannerExceptionFactory.newSpannerException(e); + stop = true; + } + } + } + // Ensure that the callback has been called at least once, even if the result set was + // cancelled. + synchronized (monitor) { + finished = true; + stop = cursorReturnedDoneOrException; + } + // Call the callback if there are still rows in the buffer that need to be processed. + while (!stop) { + waitIfPaused(); + startCallbackIfNecessary(); + synchronized (monitor) { + stop = state.shouldStop || cursorReturnedDoneOrException; + } + // Make sure we wait until the callback runner has actually finished. + consumingLatch.await(); + } + } finally { + delegateResultSet.close(); + executorFactory.release(service); + synchronized (monitor) { + if (executionException != null) { + throw executionException; + } + } + } + return null; + } + + private void waitIfPaused() throws InterruptedException { + CountDownLatch pause; + synchronized (monitor) { + pause = pausedLatch; + } + pause.await(); + } + + private void startCallbackIfNecessary() { + startCallbackWithBufferLatchIfNecessary(0); + } + + private void startCallbackWithBufferLatchIfNecessary(int bufferLatch) { + synchronized (monitor) { + if ((state == State.RUNNING || state == State.CANCELLED) + && !cursorReturnedDoneOrException) { + consumingLatch = new CountDownLatch(1); + if (bufferLatch > 0) { + bufferConsumptionLatch = new CountDownLatch(bufferLatch); + } + if (state == State.RUNNING) { + state = State.CONSUMING; + } + executor.execute(callbackRunnable); + } + } + } + } + + /** Sets the callback for this {@link AsyncResultSet}. */ + @Override + public void setCallback(Executor exec, ReadyCallback cb) { + synchronized (monitor) { + Preconditions.checkState(!closed, "This AsyncResultSet has been closed"); + Preconditions.checkState( + this.state == State.INITIALIZED, "callback may not be set multiple times"); + this.executor = MoreExecutors.newSequentialExecutor(Preconditions.checkNotNull(exec)); + this.callback = Preconditions.checkNotNull(cb); + this.state = State.RUNNING; + pausedLatch.countDown(); + } + } + + Future getResult() { + return result; + } + + @Override + public void cancel() { + synchronized (monitor) { + Preconditions.checkState( + state != State.INITIALIZED, "cannot cancel a result set without a callback"); + state = State.CANCELLED; + pausedLatch.countDown(); + } + } + + @Override + public void resume() { + synchronized (monitor) { + Preconditions.checkState( + state != State.INITIALIZED, "cannot resume a result set without a callback"); + if (state == State.PAUSED) { + state = State.RUNNING; + pausedLatch.countDown(); + } + } + } + + private static class CreateListCallback implements ReadyCallback { + private final SettableApiFuture> future; + private final Function transformer; + private final ImmutableList.Builder builder = ImmutableList.builder(); + + private CreateListCallback( + SettableApiFuture> future, Function transformer) { + this.future = future; + this.transformer = transformer; + } + + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + CursorState state; + try { + while ((state = resultSet.tryNext()) == CursorState.OK) { + builder.add(transformer.apply(resultSet)); + } + } catch (SpannerException e) { + future.setException(e); + return CallbackResponse.DONE; + } + if (state == CursorState.DONE) { + future.set(builder.build()); + return CallbackResponse.DONE; + } + return CallbackResponse.CONTINUE; + } + } + + @Override + public ApiFuture> toListAsync( + Function transformer, Executor executor) { + synchronized (monitor) { + Preconditions.checkState(!closed, "This AsyncResultSet has been closed"); + Preconditions.checkState( + this.state == State.INITIALIZED, "This AsyncResultSet has already been used."); + SettableApiFuture> res = SettableApiFuture.>create(); + CreateListCallback callback = new CreateListCallback(res, transformer); + setCallback(executor, callback); + return res; + } + } + + @Override + public ImmutableList toList(Function transformer) + throws SpannerException { + ApiFuture> future = toListAsync(transformer, MoreExecutors.directExecutor()); + try { + return future.get(); + } catch (ExecutionException e) { + throw SpannerExceptionFactory.newSpannerException(e.getCause()); + } catch (Throwable e) { + throw SpannerExceptionFactory.newSpannerException(e); + } + } +} diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClientImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClientImpl.java index 129924fbbc4..607684611c4 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClientImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClientImpl.java @@ -17,7 +17,7 @@ package com.google.cloud.spanner; import com.google.cloud.Timestamp; -import com.google.cloud.spanner.SessionPool.PooledSession; +import com.google.cloud.spanner.SessionPool.PooledSessionFuture; import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Function; import com.google.common.util.concurrent.ListenableFuture; @@ -51,12 +51,12 @@ private enum SessionMode { } @VisibleForTesting - PooledSession getReadSession() { + PooledSessionFuture getReadSession() { return pool.getReadSession(); } @VisibleForTesting - PooledSession getReadWriteSession() { + PooledSessionFuture getReadWriteSession() { return pool.getReadWriteSession(); } @@ -211,7 +211,7 @@ public Long apply(Session session) { } private T runWithSessionRetry(SessionMode mode, Function callable) { - PooledSession session = + PooledSessionFuture session = mode == SessionMode.READ_WRITE ? getReadWriteSession() : getReadSession(); while (true) { try { diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ReadContext.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ReadContext.java index 16f40769fae..542c3da4771 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ReadContext.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ReadContext.java @@ -160,6 +160,8 @@ ResultSet readUsingIndex( */ ResultSet executeQuery(Statement statement, QueryOption... options); + AsyncResultSet executeQueryAsync(Statement statement, QueryOption... options); + /** * Analyzes a query and returns query plan and/or query execution statistics information. * diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java index 5023f0fafb9..4880cfb370d 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java @@ -49,6 +49,8 @@ import com.google.common.base.Supplier; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; +import com.google.common.util.concurrent.ForwardingFuture.SimpleForwardingFuture; +import com.google.common.util.concurrent.Futures; import com.google.common.util.concurrent.ListenableFuture; import com.google.common.util.concurrent.MoreExecutors; import com.google.common.util.concurrent.SettableFuture; @@ -81,7 +83,16 @@ import java.util.concurrent.ScheduledFuture; import java.util.concurrent.SynchronousQueue; import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.Callable; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Future; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.ScheduledFuture; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicLong; import java.util.logging.Level; import java.util.logging.Logger; @@ -117,18 +128,18 @@ Instant instant() { * finished, if it is a single use context. */ private static class AutoClosingReadContext implements ReadContext { - private final Function readContextDelegateSupplier; + private final Function readContextDelegateSupplier; private T readContextDelegate; private final SessionPool sessionPool; - private PooledSession session; + private PooledSessionFuture session; private final boolean isSingleUse; private boolean closed; private boolean sessionUsedForQuery = false; private AutoClosingReadContext( - Function delegateSupplier, + Function delegateSupplier, SessionPool sessionPool, - PooledSession session, + PooledSessionFuture session, boolean isSingleUse) { this.readContextDelegateSupplier = delegateSupplier; this.sessionPool = sessionPool; @@ -177,7 +188,7 @@ private boolean internalNext() { try { boolean ret = super.next(); if (beforeFirst) { - session.markUsed(); + session.get().markUsed(); beforeFirst = false; sessionUsedForQuery = true; } @@ -189,7 +200,7 @@ private boolean internalNext() { throw e; } catch (SpannerException e) { if (!closed && isSingleUse) { - session.lastException = e; + session.get().lastException = e; AutoClosingReadContext.this.close(); } throw e; @@ -206,14 +217,14 @@ public void close() { }; } - private void replaceSessionIfPossible(SessionNotFoundException e) { + private void replaceSessionIfPossible(SessionNotFoundException notFound) { if (isSingleUse || !sessionUsedForQuery) { // This class is only used by read-only transactions, so we know that we only need a // read-only session. - session = sessionPool.replaceReadSession(e, session); + session = sessionPool.replaceReadSession(notFound, session); readContextDelegate = readContextDelegateSupplier.apply(session); } else { - throw e; + throw notFound; } } @@ -254,7 +265,7 @@ public Struct readRow(String table, Key key, Iterable columns) { try { while (true) { try { - session.markUsed(); + session.get().markUsed(); return readContextDelegate.readRow(table, key, columns); } catch (SessionNotFoundException e) { replaceSessionIfPossible(e); @@ -274,7 +285,7 @@ public Struct readRowUsingIndex(String table, String index, Key key, Iterable() { + @Override + public ResultSet get() { + return readContextDelegate.executeQuery(statement, options); + } + })); + } + @Override public ResultSet analyzeQuery(final Statement statement, final QueryAnalyzeMode queryMode) { return wrap( @@ -325,9 +350,9 @@ private static class AutoClosingReadTransaction extends AutoClosingReadContext implements ReadOnlyTransaction { AutoClosingReadTransaction( - Function txnSupplier, + Function txnSupplier, SessionPool sessionPool, - PooledSession session, + PooledSessionFuture session, boolean isSingleUse) { super(txnSupplier, sessionPool, session, isSingleUse); } @@ -437,6 +462,11 @@ public ResultSet executeQuery(Statement statement, QueryOption... options) { return new SessionPoolResultSet(delegate.executeQuery(statement, options)); } + @Override + public AsyncResultSet executeQueryAsync(Statement statement, QueryOption... options) { + return null; + } + @Override public ResultSet analyzeQuery(Statement statement, QueryAnalyzeMode queryMode) { return new SessionPoolResultSet(delegate.analyzeQuery(statement, queryMode)); @@ -450,39 +480,41 @@ public void close() { private TransactionManager delegate; private final SessionPool sessionPool; - private PooledSession session; + private PooledSessionFuture session; private boolean closed; private boolean restartedAfterSessionNotFound; - AutoClosingTransactionManager(SessionPool sessionPool, PooledSession session) { + AutoClosingTransactionManager(SessionPool sessionPool, PooledSessionFuture session) { this.sessionPool = sessionPool; this.session = session; - this.delegate = session.delegate.transactionManager(); + // this.delegate = session.delegate.transactionManager(); } @Override public TransactionContext begin() { + this.delegate = session.get().transactionManager(); while (true) { try { return internalBegin(); } catch (SessionNotFoundException e) { session = sessionPool.replaceReadWriteSession(e, session); - delegate = session.delegate.transactionManager(); + delegate = session.get().delegate.transactionManager(); } } } private TransactionContext internalBegin() { TransactionContext res = new SessionPoolTransactionContext(delegate.begin()); - session.markUsed(); + session.get().markUsed(); return res; } - private SpannerException handleSessionNotFound(SessionNotFoundException e) { - session = sessionPool.replaceReadWriteSession(e, session); - delegate = session.delegate.transactionManager(); + private SpannerException handleSessionNotFound(SessionNotFoundException notFound) { + session = sessionPool.replaceReadWriteSession(notFound, session); + delegate = session.get().delegate.transactionManager(); restartedAfterSessionNotFound = true; - return SpannerExceptionFactory.newSpannerException(ErrorCode.ABORTED, e.getMessage(), e); + return SpannerExceptionFactory.newSpannerException( + ErrorCode.ABORTED, notFound.getMessage(), notFound); } @Override @@ -520,7 +552,7 @@ public TransactionContext resetForRetry() { } } catch (SessionNotFoundException e) { session = sessionPool.replaceReadWriteSession(e, session); - delegate = session.delegate.transactionManager(); + delegate = session.get().delegate.transactionManager(); restartedAfterSessionNotFound = true; } } @@ -560,13 +592,19 @@ public TransactionState getState() { */ private static final class SessionPoolTransactionRunner implements TransactionRunner { private final SessionPool sessionPool; - private PooledSession session; + private PooledSessionFuture session; private TransactionRunner runner; - private SessionPoolTransactionRunner(SessionPool sessionPool, PooledSession session) { + private SessionPoolTransactionRunner(SessionPool sessionPool, PooledSessionFuture session) { this.sessionPool = sessionPool; this.session = session; - this.runner = session.delegate.readWriteTransaction(); + } + + private TransactionRunner getRunner() { + if (this.runner == null) { + this.runner = session.get().delegate.readWriteTransaction(); + } + return runner; } @Override @@ -576,17 +614,17 @@ public T run(TransactionCallable callable) { T result; while (true) { try { - result = runner.run(callable); + result = getRunner().run(callable); break; } catch (SessionNotFoundException e) { session = sessionPool.replaceReadWriteSession(e, session); - runner = session.delegate.readWriteTransaction(); + runner = session.get().delegate.readWriteTransaction(); } } - session.markUsed(); + session.get().markUsed(); return result; } catch (SpannerException e) { - throw session.lastException = e; + throw session.get().lastException = e; } finally { session.close(); } @@ -594,12 +632,12 @@ public T run(TransactionCallable callable) { @Override public Timestamp getCommitTimestamp() { - return runner.getCommitTimestamp(); + return getRunner().getCommitTimestamp(); } @Override public TransactionRunner allowNestedTransaction() { - runner.allowNestedTransaction(); + getRunner().allowNestedTransaction(); return this; } } @@ -620,25 +658,23 @@ private enum SessionState { CLOSING, } - final class PooledSession implements Session { - @VisibleForTesting SessionImpl delegate; - private volatile Instant lastUseTime; - private volatile SpannerException lastException; - private volatile LeakedSessionException leakedException; - private volatile boolean allowReplacing = true; + private PooledSessionFuture createPooledSessionFuture(Future future, Span span) { + return new PooledSessionFuture(future, span); + } - @GuardedBy("lock") - private SessionState state; + private PooledSessionFuture createPooledSessionFuture(PooledSession session, Span span) { + return new PooledSessionFuture(Futures.immediateFuture(session), span); + } - private PooledSession(SessionImpl delegate) { - this.delegate = delegate; - this.state = SessionState.AVAILABLE; - this.lastUseTime = clock.instant(); - } + final class PooledSessionFuture extends SimpleForwardingFuture implements Session { + private volatile LeakedSessionException leakedException; + private volatile AtomicBoolean inUse = new AtomicBoolean(); + private volatile CountDownLatch initialized = new CountDownLatch(1); + private final Span span; - @VisibleForTesting - void setAllowReplacing(boolean allowReplacing) { - this.allowReplacing = allowReplacing; + private PooledSessionFuture(Future delegate, Span span) { + super(delegate); + this.span = span; } @VisibleForTesting @@ -646,34 +682,14 @@ void clearLeakedException() { this.leakedException = null; } - private void markBusy() { - this.state = SessionState.BUSY; + private void markCheckedOut() { this.leakedException = new LeakedSessionException(); } - private void markClosing() { - this.state = SessionState.CLOSING; - } - @Override public Timestamp write(Iterable mutations) throws SpannerException { try { - markUsed(); - return delegate.write(mutations); - } catch (SpannerException e) { - throw lastException = e; - } finally { - close(); - } - } - - @Override - public long executePartitionedUpdate(Statement stmt) throws SpannerException { - try { - markUsed(); - return delegate.executePartitionedUpdate(stmt); - } catch (SpannerException e) { - throw lastException = e; + return get().write(mutations); } finally { close(); } @@ -682,10 +698,7 @@ public long executePartitionedUpdate(Statement stmt) throws SpannerException { @Override public Timestamp writeAtLeastOnce(Iterable mutations) throws SpannerException { try { - markUsed(); - return delegate.writeAtLeastOnce(mutations); - } catch (SpannerException e) { - throw lastException = e; + return get().writeAtLeastOnce(mutations); } finally { close(); } @@ -695,10 +708,10 @@ public Timestamp writeAtLeastOnce(Iterable mutations) throws SpannerEx public ReadContext singleUse() { try { return new AutoClosingReadContext<>( - new Function() { + new Function() { @Override - public ReadContext apply(PooledSession session) { - return session.delegate.singleUse(); + public ReadContext apply(PooledSessionFuture session) { + return session.get().delegate.singleUse(); } }, SessionPool.this, @@ -714,10 +727,10 @@ public ReadContext apply(PooledSession session) { public ReadContext singleUse(final TimestampBound bound) { try { return new AutoClosingReadContext<>( - new Function() { + new Function() { @Override - public ReadContext apply(PooledSession session) { - return session.delegate.singleUse(bound); + public ReadContext apply(PooledSessionFuture session) { + return session.get().delegate.singleUse(bound); } }, SessionPool.this, @@ -732,10 +745,10 @@ public ReadContext apply(PooledSession session) { @Override public ReadOnlyTransaction singleUseReadOnlyTransaction() { return internalReadOnlyTransaction( - new Function() { + new Function() { @Override - public ReadOnlyTransaction apply(PooledSession session) { - return session.delegate.singleUseReadOnlyTransaction(); + public ReadOnlyTransaction apply(PooledSessionFuture session) { + return session.get().delegate.singleUseReadOnlyTransaction(); } }, true); @@ -744,10 +757,10 @@ public ReadOnlyTransaction apply(PooledSession session) { @Override public ReadOnlyTransaction singleUseReadOnlyTransaction(final TimestampBound bound) { return internalReadOnlyTransaction( - new Function() { + new Function() { @Override - public ReadOnlyTransaction apply(PooledSession session) { - return session.delegate.singleUseReadOnlyTransaction(bound); + public ReadOnlyTransaction apply(PooledSessionFuture session) { + return session.get().delegate.singleUseReadOnlyTransaction(bound); } }, true); @@ -756,10 +769,10 @@ public ReadOnlyTransaction apply(PooledSession session) { @Override public ReadOnlyTransaction readOnlyTransaction() { return internalReadOnlyTransaction( - new Function() { + new Function() { @Override - public ReadOnlyTransaction apply(PooledSession session) { - return session.delegate.readOnlyTransaction(); + public ReadOnlyTransaction apply(PooledSessionFuture session) { + return session.get().delegate.readOnlyTransaction(); } }, false); @@ -768,17 +781,18 @@ public ReadOnlyTransaction apply(PooledSession session) { @Override public ReadOnlyTransaction readOnlyTransaction(final TimestampBound bound) { return internalReadOnlyTransaction( - new Function() { + new Function() { @Override - public ReadOnlyTransaction apply(PooledSession session) { - return session.delegate.readOnlyTransaction(bound); + public ReadOnlyTransaction apply(PooledSessionFuture session) { + return session.get().delegate.readOnlyTransaction(bound); } }, false); } private ReadOnlyTransaction internalReadOnlyTransaction( - Function transactionSupplier, boolean isSingleUse) { + Function transactionSupplier, + boolean isSingleUse) { try { return new AutoClosingReadTransaction( transactionSupplier, SessionPool.this, this, isSingleUse); @@ -793,6 +807,161 @@ public TransactionRunner readWriteTransaction() { return new SessionPoolTransactionRunner(SessionPool.this, this); } + @Override + public TransactionManager transactionManager() { + return new AutoClosingTransactionManager(SessionPool.this, this); + } + + @Override + public long executePartitionedUpdate(Statement stmt) { + try { + return get().executePartitionedUpdate(stmt); + } finally { + close(); + } + } + + @Override + public String getName() { + return get().getName(); + } + + @Override + public void prepareReadWriteTransaction() { + get().prepareReadWriteTransaction(); + } + + @Override + public void close() { + synchronized (lock) { + leakedException = null; + checkedOutSessions.remove(this); + } + get().close(); + } + + @Override + public ApiFuture asyncClose() { + synchronized (lock) { + leakedException = null; + checkedOutSessions.remove(this); + } + return get().asyncClose(); + } + + @Override + public PooledSession get() { + if (inUse.compareAndSet(false, true)) { + try { + PooledSession res = super.get(); + synchronized (lock) { + res.markBusy(); + span.addAnnotation(sessionAnnotation(res)); + incrementNumSessionsInUse(); + checkedOutSessions.add(this); + } + initialized.countDown(); + } catch (Throwable e) { + initialized.countDown(); + // ignore and fallthrough. + } + } + try { + initialized.await(); + return super.get(); + } catch (ExecutionException e) { + throw SpannerExceptionFactory.newSpannerException(e.getCause()); + } catch (InterruptedException e) { + throw SpannerExceptionFactory.propagateInterrupt(e); + } + } + } + + final class PooledSession implements Session { + @VisibleForTesting SessionImpl delegate; + private volatile Instant lastUseTime; + private volatile SpannerException lastException; + private volatile boolean allowReplacing = true; + + @GuardedBy("lock") + private SessionState state; + + private PooledSession(SessionImpl delegate) { + this.delegate = delegate; + this.state = SessionState.AVAILABLE; + this.lastUseTime = clock.instant(); + } + + @VisibleForTesting + void setAllowReplacing(boolean allowReplacing) { + this.allowReplacing = allowReplacing; + } + + @Override + public Timestamp write(Iterable mutations) throws SpannerException { + try { + markUsed(); + return delegate.write(mutations); + } catch (SpannerException e) { + throw lastException = e; + } + } + + @Override + public Timestamp writeAtLeastOnce(Iterable mutations) throws SpannerException { + try { + markUsed(); + return delegate.writeAtLeastOnce(mutations); + } catch (SpannerException e) { + throw lastException = e; + } + } + + @Override + public long executePartitionedUpdate(Statement stmt) throws SpannerException { + try { + markUsed(); + return delegate.executePartitionedUpdate(stmt); + } catch (SpannerException e) { + throw lastException = e; + } + } + + @Override + public ReadContext singleUse() { + return delegate.singleUse(); + } + + @Override + public ReadContext singleUse(TimestampBound bound) { + return delegate.singleUse(bound); + } + + @Override + public ReadOnlyTransaction singleUseReadOnlyTransaction() { + return delegate.singleUseReadOnlyTransaction(); + } + + @Override + public ReadOnlyTransaction singleUseReadOnlyTransaction(TimestampBound bound) { + return delegate.singleUseReadOnlyTransaction(bound); + } + + @Override + public ReadOnlyTransaction readOnlyTransaction() { + return delegate.readOnlyTransaction(); + } + + @Override + public ReadOnlyTransaction readOnlyTransaction(TimestampBound bound) { + return delegate.readOnlyTransaction(bound); + } + + @Override + public TransactionRunner readWriteTransaction() { + return delegate.readWriteTransaction(); + } + @Override public ApiFuture asyncClose() { close(); @@ -805,7 +974,6 @@ public void close() { numSessionsInUse--; numSessionsReleased++; } - leakedException = null; if (lastException != null && isSessionNotFound(lastException)) { invalidateSession(this); } else { @@ -848,13 +1016,21 @@ private void keepAlive() { } } + private void markBusy() { + this.state = SessionState.BUSY; + } + + private void markClosing() { + this.state = SessionState.CLOSING; + } + private void markUsed() { lastUseTime = clock.instant(); } @Override public TransactionManager transactionManager() { - return new AutoClosingTransactionManager(SessionPool.this, this); + return delegate.transactionManager(); } } @@ -875,7 +1051,8 @@ private static final class SessionOrError { private final class Waiter { private static final long MAX_SESSION_WAIT_TIMEOUT = 240_000L; - private final SynchronousQueue waiter = new SynchronousQueue<>(); + private final BlockingQueue waiter = new LinkedBlockingQueue<>(1); + // private final SynchronousQueue waiter = new SynchronousQueue<>(); private void put(PooledSession session) { Uninterruptibles.putUninterruptibly(waiter, new SessionOrError(session)); @@ -1093,6 +1270,18 @@ private static enum Position { private final ScheduledExecutorService executor; private final ExecutorFactory executorFactory; private final ScheduledExecutorService prepareExecutor; + private final ScheduledExecutorService readWaiterExecutor = + Executors.newSingleThreadScheduledExecutor( + new ThreadFactoryBuilder() + .setDaemon(true) + .setNameFormat("session-pool-read-waiter-%d") + .build()); + private final ScheduledExecutorService writeWaiterExecutor = + Executors.newSingleThreadScheduledExecutor( + new ThreadFactoryBuilder() + .setDaemon(true) + .setNameFormat("session-pool-write-waiter-%d") + .build()); final PoolMaintainer poolMaintainer; private final Clock clock; private final Object lock = new Object(); @@ -1142,6 +1331,9 @@ private static enum Position { @GuardedBy("lock") private final Set allSessions = new HashSet<>(); + @GuardedBy("lock") + private final Set checkedOutSessions = new HashSet<>(); + private final SessionConsumer sessionConsumer = new SessionConsumerImpl(); /** @@ -1348,7 +1540,7 @@ boolean isValid() { * session being returned to the pool or a new session being created. * */ - PooledSession getReadSession() throws SpannerException { + PooledSessionFuture getReadSession() throws SpannerException { Span span = Tracing.getTracer().getCurrentSpan(); span.addAnnotation("Acquiring session"); Waiter waiter = null; @@ -1381,18 +1573,8 @@ PooledSession getReadSession() throws SpannerException { } else { span.addAnnotation("Acquired read only session"); } + return checkoutSession(span, sess, waiter, false); } - if (waiter != null) { - logger.log( - Level.FINE, - "No session available in the pool. Blocking for one to become available/created"); - span.addAnnotation("Waiting for read only session to be available"); - sess = waiter.take(); - } - sess.markBusy(); - incrementNumSessionsInUse(); - span.addAnnotation(sessionAnnotation(sess)); - return sess; } /** @@ -1413,7 +1595,7 @@ PooledSession getReadSession() throws SpannerException { * to the pool which is then write prepared. * */ - PooledSession getReadWriteSession() { + PooledSessionFuture getReadWriteSession() { Span span = Tracing.getTracer().getCurrentSpan(); span.addAnnotation("Acquiring read write session"); Waiter waiter = null; @@ -1449,37 +1631,57 @@ PooledSession getReadWriteSession() { } else { span.addAnnotation("Acquired read write session"); } + return checkoutSession(span, sess, waiter, true); } + } + + private PooledSessionFuture checkoutSession( + final Span span, final PooledSession sess, final Waiter waiter, boolean write) { + final PooledSessionFuture res; if (waiter != null) { logger.log( Level.FINE, "No session available in the pool. Blocking for one to become available/created"); - span.addAnnotation("Waiting for read write session to be available"); - sess = waiter.take(); + span.addAnnotation( + String.format( + "Waiting for %s session to be available", write ? "read write" : "read only")); + ScheduledExecutorService executor = write ? writeWaiterExecutor : readWaiterExecutor; + res = + createPooledSessionFuture( + executor.submit( + new Callable() { + @Override + public PooledSession call() throws Exception { + return waiter.take(); + } + }), + span); + } else { + res = createPooledSessionFuture(sess, span); } - sess.markBusy(); - incrementNumSessionsInUse(); - span.addAnnotation(sessionAnnotation(sess)); - return sess; + res.markCheckedOut(); + return res; } - PooledSession replaceReadSession(SessionNotFoundException e, PooledSession session) { + PooledSessionFuture replaceReadSession(SessionNotFoundException e, PooledSessionFuture session) { return replaceSession(e, session, false); } - PooledSession replaceReadWriteSession(SessionNotFoundException e, PooledSession session) { + PooledSessionFuture replaceReadWriteSession( + SessionNotFoundException e, PooledSessionFuture session) { return replaceSession(e, session, true); } - private PooledSession replaceSession( - SessionNotFoundException e, PooledSession session, boolean write) { - if (!options.isFailIfSessionNotFound() && session.allowReplacing) { + private PooledSessionFuture replaceSession( + SessionNotFoundException e, PooledSessionFuture session, boolean write) { + if (!options.isFailIfSessionNotFound() && session.get().allowReplacing) { synchronized (lock) { numSessionsInUse--; numSessionsReleased++; + checkedOutSessions.remove(session); } session.leakedException = null; - invalidateSession(session); + invalidateSession(session.get()); return write ? getReadWriteSession() : getReadSession(); } else { throw e; @@ -1668,10 +1870,12 @@ public void run() { } } }); - for (final PooledSession session : ImmutableList.copyOf(allSessions)) { + for (PooledSessionFuture session : checkedOutSessions) { if (session.leakedException != null) { logger.log(Level.WARNING, "Leaked session", session.leakedException); } + } + for (final PooledSession session : ImmutableList.copyOf(allSessions)) { if (session.state != SessionState.CLOSING) { closeSessionAsync(session); } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java index bf0a47222bc..344c46f82bc 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java @@ -22,6 +22,7 @@ import com.google.cloud.PageImpl; import com.google.cloud.PageImpl.NextPageFetcher; import com.google.cloud.grpc.GrpcTransportOptions; +import com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory; import com.google.cloud.spanner.SessionClient.SessionId; import com.google.cloud.spanner.spi.v1.SpannerRpc; import com.google.cloud.spanner.spi.v1.SpannerRpc.Paginated; @@ -42,6 +43,7 @@ import java.util.concurrent.ExecutionException; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; +import java.util.concurrent.ScheduledExecutorService; import java.util.logging.Level; import java.util.logging.Logger; import javax.annotation.Nullable; @@ -125,6 +127,10 @@ QueryOptions getDefaultQueryOptions(DatabaseId databaseId) { return getOptions().getDefaultQueryOptions(databaseId); } + ExecutorFactory getExecutorFactory() { + return ((GrpcTransportOptions) getOptions().getTransportOptions()).getExecutorFactory(); + } + SessionImpl sessionWithId(String name) { Preconditions.checkArgument(!Strings.isNullOrEmpty(name), "name is null or empty"); SessionId id = SessionId.of(name); 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 cfa8b73c4a4..fc72793e860 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 @@ -22,6 +22,7 @@ import static com.google.common.base.Preconditions.checkState; import com.google.cloud.Timestamp; +import com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory; import com.google.cloud.spanner.SessionImpl.SessionTransaction; import com.google.cloud.spanner.spi.v1.SpannerRpc; import com.google.common.annotations.VisibleForTesting; @@ -43,6 +44,7 @@ import java.util.ArrayList; import java.util.List; import java.util.concurrent.Callable; +import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.atomic.AtomicInteger; import java.util.logging.Level; import java.util.logging.Logger; diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncResultSetImplStressTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncResultSetImplStressTest.java new file mode 100644 index 00000000000..ea8396b7ed0 --- /dev/null +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncResultSetImplStressTest.java @@ -0,0 +1,462 @@ +/* + * Copyright 2020 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.spanner; + +import static com.google.common.truth.Truth.assertThat; + +import com.google.api.core.ApiFuture; +import com.google.api.core.ApiFutures; +import com.google.api.core.SettableApiFuture; +import com.google.cloud.grpc.GrpcTransportOptions; +import com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory; +import com.google.cloud.spanner.AsyncResultSet.CallbackResponse; +import com.google.cloud.spanner.AsyncResultSet.CursorState; +import com.google.cloud.spanner.AsyncResultSet.ReadyCallback; +import com.google.cloud.spanner.Type.StructField; +import com.google.common.base.Function; +import com.google.common.collect.ImmutableList; +import com.google.common.util.concurrent.MoreExecutors; +import com.google.common.util.concurrent.ThreadFactoryBuilder; +import java.util.ArrayList; +import java.util.Collection; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Objects; +import java.util.Random; +import java.util.Set; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executor; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.Parameterized; +import org.junit.runners.Parameterized.Parameter; +import org.junit.runners.Parameterized.Parameters; + +@RunWith(Parameterized.class) +public class AsyncResultSetImplStressTest { + private static final int TEST_RUNS = 1000; + + @Parameter(0) + public int resultSetSize; + + @Parameters(name = "rows = {0}") + public static Collection data() { + List params = new ArrayList<>(); + for (int rows : new int[] {0, 1, 5, 10}) { + params.add(new Object[] {rows}); + } + return params; + } + + /** POJO representing a row in the test {@link ResultSet}. */ + private static final class Row { + private final Long id; + private final String name; + + static Row create(StructReader reader) { + return new Row(reader.getLong("ID"), reader.getString("NAME")); + } + + private Row(Long id, String name) { + this.id = id; + this.name = name; + } + + @Override + public boolean equals(Object o) { + if (!(o instanceof Row)) { + return false; + } + Row other = (Row) o; + return Objects.equals(this.id, other.id) && Objects.equals(this.name, other.name); + } + + @Override + public int hashCode() { + return Objects.hash(this.id, this.name); + } + + @Override + public String toString() { + return String.format("ID: %d, NAME: %s", id, name); + } + } + + private static final class ResultSetWithRandomErrors extends ForwardingResultSet { + private final Random random = new Random(); + private final double errorFraction; + + private ResultSetWithRandomErrors(ResultSet delegate, double errorFraction) { + super(delegate); + this.errorFraction = errorFraction; + } + + @Override + public boolean next() { + if (random.nextDouble() < errorFraction) { + throw SpannerExceptionFactory.newSpannerException( + ErrorCode.INVALID_ARGUMENT, "random error"); + } + return super.next(); + } + } + + /** Creates a simple in-mem {@link ResultSet}. */ + private ResultSet createResultSet() { + List rows = new ArrayList<>(resultSetSize); + for (int i = 0; i < resultSetSize; i++) { + rows.add( + Struct.newBuilder() + .set("ID") + .to(i + 1) + .set("NAME") + .to(String.format("Row %d", (i + 1))) + .build()); + } + return ResultSets.forRows( + Type.struct(StructField.of("ID", Type.int64()), StructField.of("NAME", Type.string())), + rows); + } + + private ResultSet createResultSetWithErrors(double errorFraction) { + return new ResultSetWithRandomErrors(createResultSet(), errorFraction); + } + + /** + * Generates a list of {@link Row} instances that correspond with the rows in {@link + * #createResultSet()}. + */ + private List createExpectedRows() { + List rows = new ArrayList<>(resultSetSize); + for (int i = 0; i < resultSetSize; i++) { + rows.add(new Row(Long.valueOf(i + 1), String.format("Row %d", (i + 1)))); + } + return rows; + } + + /** Creates a single-threaded {@link ExecutorService}. */ + private static ScheduledExecutorService createExecService() { + return createExecService(1); + } + + /** Creates an {@link ExecutorService} using a bounded pool of threadCount threads. */ + private static ScheduledExecutorService createExecService(int threadCount) { + return Executors.newScheduledThreadPool( + threadCount, new ThreadFactoryBuilder().setDaemon(true).build()); + } + + @Test + public void toList() throws Exception { + ExecutorFactory executorFactory = + new GrpcTransportOptions.DefaultExecutorFactory(); + for (int bufferSize = 1; bufferSize < resultSetSize * 2; bufferSize *= 2) { + for (int i = 0; i < TEST_RUNS; i++) { + try (AsyncResultSetImpl impl = + new AsyncResultSetImpl(executorFactory, createResultSet(), bufferSize)) { + ImmutableList list = + impl.toList( + new Function() { + @Override + public Row apply(StructReader input) { + return Row.create(input); + } + }); + assertThat(list).containsExactlyElementsIn(createExpectedRows()); + } + } + } + } + + @Test + public void toListWithErrors() throws Exception { + ExecutorFactory executorFactory = + new GrpcTransportOptions.DefaultExecutorFactory(); + for (int bufferSize = 1; bufferSize < resultSetSize * 2; bufferSize *= 2) { + for (int i = 0; i < TEST_RUNS; i++) { + try (AsyncResultSetImpl impl = + new AsyncResultSetImpl( + executorFactory, createResultSetWithErrors(1.0 / resultSetSize), bufferSize)) { + ImmutableList list = + impl.toList( + new Function() { + @Override + public Row apply(StructReader input) { + return Row.create(input); + } + }); + assertThat(list).containsExactlyElementsIn(createExpectedRows()); + } catch (SpannerException e) { + assertThat(e.getErrorCode()).isEqualTo(ErrorCode.INVALID_ARGUMENT); + assertThat(e.getMessage()).contains("random error"); + } + } + } + } + + @Test + public void asyncToList() throws Exception { + ExecutorFactory executorFactory = + new GrpcTransportOptions.DefaultExecutorFactory(); + for (int bufferSize = 1; bufferSize < resultSetSize * 2; bufferSize *= 2) { + List>> futures = new ArrayList<>(TEST_RUNS); + ExecutorService executor = createExecService(32); + for (int i = 0; i < TEST_RUNS; i++) { + try (AsyncResultSet impl = + new AsyncResultSetImpl(executorFactory, createResultSet(), bufferSize)) { + futures.add( + impl.toListAsync( + new Function() { + @Override + public Row apply(StructReader input) { + return Row.create(input); + } + }, + executor)); + } + } + List> lists = ApiFutures.allAsList(futures).get(); + for (ImmutableList list : lists) { + assertThat(list).containsExactlyElementsIn(createExpectedRows()); + } + executor.shutdown(); + } + } + + @Test + public void consume() throws Exception { + ExecutorFactory executorFactory = + new GrpcTransportOptions.DefaultExecutorFactory(); + final Random random = new Random(); + for (Executor executor : + new Executor[] { + MoreExecutors.directExecutor(), createExecService(), createExecService(32) + }) { + for (int bufferSize = 1; bufferSize < resultSetSize * 2; bufferSize *= 2) { + for (int i = 0; i < TEST_RUNS; i++) { + final SettableApiFuture> future = SettableApiFuture.create(); + try (AsyncResultSetImpl impl = + new AsyncResultSetImpl(executorFactory, createResultSet(), bufferSize)) { + final ImmutableList.Builder builder = ImmutableList.builder(); + impl.setCallback( + executor, + new ReadyCallback() { + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + // Randomly do something with the received data or not. Not calling tryNext() in + // the onDataReady is not 'normal', but users may do it, and the result set + // should be able to handle that. + if (random.nextBoolean()) { + CursorState state; + while ((state = resultSet.tryNext()) == CursorState.OK) { + builder.add(Row.create(resultSet)); + } + if (state == CursorState.DONE) { + future.set(builder.build()); + } + } + return CallbackResponse.CONTINUE; + } + }); + assertThat(future.get()).containsExactlyElementsIn(createExpectedRows()); + } + } + } + } + } + + @Test + public void pauseResume() throws Exception { + ExecutorFactory executorFactory = + new GrpcTransportOptions.DefaultExecutorFactory(); + final Random random = new Random(); + List>> futures = new ArrayList<>(); + for (Executor executor : + new Executor[] { + MoreExecutors.directExecutor(), createExecService(), createExecService(32) + }) { + final List resultSets = + Collections.synchronizedList(new ArrayList()); + for (int bufferSize = 1; bufferSize < resultSetSize * 2; bufferSize *= 2) { + for (int i = 0; i < TEST_RUNS; i++) { + final SettableApiFuture> future = SettableApiFuture.create(); + futures.add(future); + try (AsyncResultSetImpl impl = + new AsyncResultSetImpl(executorFactory, createResultSet(), bufferSize)) { + resultSets.add(impl); + final ImmutableList.Builder builder = ImmutableList.builder(); + impl.setCallback( + executor, + new ReadyCallback() { + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + CursorState state; + while ((state = resultSet.tryNext()) == CursorState.OK) { + builder.add(Row.create(resultSet)); + // Randomly request the iterator to pause. + if (random.nextBoolean()) { + return CallbackResponse.PAUSE; + } + } + if (state == CursorState.DONE) { + future.set(builder.build()); + } + return CallbackResponse.CONTINUE; + } + }); + } + } + } + final AtomicBoolean finished = new AtomicBoolean(false); + ExecutorService resumeService = createExecService(); + resumeService.execute( + new Runnable() { + @Override + public void run() { + while (!finished.get()) { + // Randomly resume result sets. + resultSets.get(random.nextInt(resultSets.size())).resume(); + } + } + }); + List> lists = ApiFutures.allAsList(futures).get(); + for (ImmutableList list : lists) { + assertThat(list).containsExactlyElementsIn(createExpectedRows()); + } + if (executor instanceof ExecutorService) { + ((ExecutorService) executor).shutdown(); + } + finished.set(true); + resumeService.shutdown(); + } + } + + @Test + public void cancel() throws Exception { + ExecutorFactory executorFactory = + new GrpcTransportOptions.DefaultExecutorFactory(); + final Random random = new Random(); + for (Executor executor : + new Executor[] { + MoreExecutors.directExecutor(), createExecService(), createExecService(32) + }) { + List>> futures = new ArrayList<>(); + final List resultSets = + Collections.synchronizedList(new ArrayList()); + final Set cancelledIndexes = new HashSet<>(); + for (int bufferSize = 1; bufferSize < resultSetSize * 2; bufferSize *= 2) { + for (int i = 0; i < TEST_RUNS; i++) { + final SettableApiFuture> future = SettableApiFuture.create(); + futures.add(future); + try (AsyncResultSetImpl impl = + new AsyncResultSetImpl(executorFactory, createResultSet(), bufferSize)) { + resultSets.add(impl); + final ImmutableList.Builder builder = ImmutableList.builder(); + impl.setCallback( + executor, + new ReadyCallback() { + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + try { + CursorState state; + while ((state = resultSet.tryNext()) == CursorState.OK) { + builder.add(Row.create(resultSet)); + // Randomly request the iterator to pause. + if (random.nextBoolean()) { + return CallbackResponse.PAUSE; + } + } + if (state == CursorState.DONE) { + future.set(builder.build()); + } + return CallbackResponse.CONTINUE; + } catch (SpannerException e) { + future.setException(e); + throw e; + } + } + }); + } + } + } + final AtomicBoolean finished = new AtomicBoolean(false); + // Both resume and cancel result sets randomly. + ExecutorService resumeService = createExecService(); + resumeService.execute( + new Runnable() { + @Override + public void run() { + while (!finished.get()) { + // Randomly resume result sets. + resultSets.get(random.nextInt(resultSets.size())).resume(); + } + } + }); + ExecutorService cancelService = createExecService(); + cancelService.execute( + new Runnable() { + @Override + public void run() { + while (!finished.get()) { + // Randomly cancel result sets. + int index = random.nextInt(resultSets.size()); + resultSets.get(index).cancel(); + cancelledIndexes.add(index); + } + } + }); + + // First wait until all result sets have finished. + for (ApiFuture> future : futures) { + try { + future.get(); + } catch (Throwable e) { + // ignore for now. + } + } + finished.set(true); + cancelService.shutdown(); + cancelService.awaitTermination(10L, TimeUnit.SECONDS); + + int index = 0; + for (ApiFuture> future : futures) { + try { + ImmutableList list = future.get(); + // Note that the fact that the call succeeded for for this result set, does not + // necessarily mean that the result set was not cancelled. Cancelling a result set is a + // best-effort operation, and the entire result set may still be produced and returned to + // the user. + assertThat(list).containsExactlyElementsIn(createExpectedRows()); + } catch (ExecutionException e) { + assertThat(e.getCause()).isInstanceOf(SpannerException.class); + SpannerException se = (SpannerException) e.getCause(); + assertThat(se.getErrorCode()).isEqualTo(ErrorCode.CANCELLED); + assertThat(cancelledIndexes).contains(index); + } + index++; + } + if (executor instanceof ExecutorService) { + ((ExecutorService) executor).shutdown(); + } + } + } +} diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncResultSetImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncResultSetImplTest.java new file mode 100644 index 00000000000..f9df08563da --- /dev/null +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncResultSetImplTest.java @@ -0,0 +1,440 @@ +/* + * Copyright 2020 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.spanner; + +import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.fail; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import com.google.api.core.ApiFuture; +import com.google.cloud.grpc.GrpcTransportOptions; +import com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory; +import com.google.cloud.spanner.AsyncResultSet.CallbackResponse; +import com.google.cloud.spanner.AsyncResultSet.CursorState; +import com.google.cloud.spanner.AsyncResultSet.ReadyCallback; +import com.google.common.base.Function; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.Range; +import java.util.concurrent.BlockingDeque; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executor; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.LinkedBlockingDeque; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class AsyncResultSetImplTest { + private ExecutorFactory mockedFactory; + private ExecutorFactory simpleFactory; + + @SuppressWarnings("unchecked") + @Before + public void setup() { + mockedFactory = mock(ExecutorFactory.class); + when(mockedFactory.get()).thenReturn(mock(ScheduledExecutorService.class)); + simpleFactory = + new GrpcTransportOptions.ExecutorFactory() { + @Override + public ScheduledExecutorService get() { + return Executors.newScheduledThreadPool(1); + } + + @Override + public void release(ScheduledExecutorService executor) { + executor.shutdown(); + } + }; + } + + @SuppressWarnings("unchecked") + @Test + public void close() { + AsyncResultSetImpl rs = new AsyncResultSetImpl(mockedFactory, mock(ResultSet.class)); + rs.close(); + // Closing a second time should be a no-op. + rs.close(); + + // The following methods are not allowed to call after closing the result set. + try { + rs.setCallback(mock(Executor.class), mock(ReadyCallback.class)); + fail("missing expected exception"); + } catch (IllegalStateException e) { + } + try { + rs.toList(mock(Function.class)); + fail("missing expected exception"); + } catch (IllegalStateException e) { + } + try { + rs.toListAsync(mock(Function.class), mock(Executor.class)); + fail("missing expected exception"); + } catch (IllegalStateException e) { + } + + // The following methods are allowed on a closed result set. + AsyncResultSetImpl rs2 = new AsyncResultSetImpl(mockedFactory, mock(ResultSet.class)); + rs2.setCallback(mock(Executor.class), mock(ReadyCallback.class)); + rs2.close(); + rs2.cancel(); + rs2.resume(); + } + + @Test + public void tryNextNotAllowed() { + try (AsyncResultSetImpl rs = new AsyncResultSetImpl(mockedFactory, mock(ResultSet.class))) { + rs.setCallback(mock(Executor.class), mock(ReadyCallback.class)); + try { + rs.tryNext(); + fail("missing expected exception"); + } catch (IllegalStateException e) { + assertThat(e.getMessage()) + .contains("tryNext may only be called from a DataReady callback."); + } + } + } + + @Test + public void toList() { + ResultSet delegate = mock(ResultSet.class); + when(delegate.next()).thenReturn(true, true, true, false); + when(delegate.getCurrentRowAsStruct()).thenReturn(mock(Struct.class)); + try (AsyncResultSetImpl rs = new AsyncResultSetImpl(simpleFactory, delegate)) { + ImmutableList list = + rs.toList( + new Function() { + @Override + public Object apply(StructReader input) { + return new Object(); + } + }); + assertThat(list).hasSize(3); + } + } + + @Test + public void toListPropagatesError() { + ResultSet delegate = mock(ResultSet.class); + when(delegate.next()) + .thenThrow( + SpannerExceptionFactory.newSpannerException( + ErrorCode.INVALID_ARGUMENT, "invalid query")); + try (AsyncResultSetImpl rs = new AsyncResultSetImpl(simpleFactory, delegate)) { + rs.toList( + new Function() { + @Override + public Object apply(StructReader input) { + return new Object(); + } + }); + fail("missing expected exception"); + } catch (SpannerException e) { + assertThat(e.getErrorCode()).isEqualTo(ErrorCode.INVALID_ARGUMENT); + assertThat(e.getMessage()).contains("invalid query"); + } + } + + @Test + public void toListAsync() throws InterruptedException, ExecutionException { + ExecutorService executor = Executors.newFixedThreadPool(1); + ResultSet delegate = mock(ResultSet.class); + when(delegate.next()).thenReturn(true, true, true, false); + when(delegate.getCurrentRowAsStruct()).thenReturn(mock(Struct.class)); + try (AsyncResultSetImpl rs = new AsyncResultSetImpl(simpleFactory, delegate)) { + ApiFuture> future = + rs.toListAsync( + new Function() { + @Override + public Object apply(StructReader input) { + return new Object(); + } + }, + executor); + assertThat(future.get()).hasSize(3); + } + executor.shutdown(); + } + + @Test + public void toListAsyncPropagatesError() throws InterruptedException { + ExecutorService executor = Executors.newFixedThreadPool(1); + ResultSet delegate = mock(ResultSet.class); + when(delegate.next()) + .thenThrow( + SpannerExceptionFactory.newSpannerException( + ErrorCode.INVALID_ARGUMENT, "invalid query")); + try (AsyncResultSetImpl rs = new AsyncResultSetImpl(simpleFactory, delegate)) { + rs.toListAsync( + new Function() { + @Override + public Object apply(StructReader input) { + return new Object(); + } + }, + executor) + .get(); + fail("missing expected exception"); + } catch (ExecutionException e) { + assertThat(e.getCause()).isInstanceOf(SpannerException.class); + SpannerException se = (SpannerException) e.getCause(); + assertThat(se.getErrorCode()).isEqualTo(ErrorCode.INVALID_ARGUMENT); + assertThat(se.getMessage()).contains("invalid query"); + } + executor.shutdown(); + } + + @Test + public void withCallback() throws InterruptedException { + Executor executor = Executors.newSingleThreadExecutor(); + ResultSet delegate = mock(ResultSet.class); + when(delegate.next()).thenReturn(true, true, true, false); + when(delegate.getCurrentRowAsStruct()).thenReturn(mock(Struct.class)); + final AtomicInteger callbackCounter = new AtomicInteger(); + final AtomicInteger rowCounter = new AtomicInteger(); + final CountDownLatch finishedLatch = new CountDownLatch(1); + try (AsyncResultSetImpl rs = new AsyncResultSetImpl(simpleFactory, delegate)) { + rs.setCallback( + executor, + new ReadyCallback() { + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + callbackCounter.incrementAndGet(); + CursorState state; + while ((state = resultSet.tryNext()) == CursorState.OK) { + rowCounter.incrementAndGet(); + } + if (state == CursorState.DONE) { + finishedLatch.countDown(); + } + return CallbackResponse.CONTINUE; + } + }); + } + finishedLatch.await(); + // There should be between 1 and 4 callbacks, depending on the timing of the threads. + // Normally, there should be just 1 callback. + assertThat(callbackCounter.get()).isIn(Range.closed(1, 4)); + assertThat(rowCounter.get()).isEqualTo(3); + } + + @Test + public void callbackReceivesError() throws InterruptedException { + Executor executor = Executors.newSingleThreadExecutor(); + ResultSet delegate = mock(ResultSet.class); + when(delegate.next()) + .thenThrow( + SpannerExceptionFactory.newSpannerException( + ErrorCode.INVALID_ARGUMENT, "invalid query")); + final BlockingDeque receivedErr = new LinkedBlockingDeque<>(1); + try (AsyncResultSetImpl rs = new AsyncResultSetImpl(simpleFactory, delegate)) { + rs.setCallback( + executor, + new ReadyCallback() { + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + try { + resultSet.tryNext(); + receivedErr.push(new Exception("missing expected exception")); + } catch (SpannerException e) { + receivedErr.push(e); + } + return CallbackResponse.DONE; + } + }); + } + Exception e = receivedErr.take(); + assertThat(e).isInstanceOf(SpannerException.class); + SpannerException se = (SpannerException) e; + assertThat(se.getErrorCode()).isEqualTo(ErrorCode.INVALID_ARGUMENT); + assertThat(se.getMessage()).contains("invalid query"); + } + + @Test + public void callbackReceivesErrorHalfwayThrough() throws InterruptedException { + Executor executor = Executors.newSingleThreadExecutor(); + ResultSet delegate = mock(ResultSet.class); + when(delegate.next()) + .thenReturn(true) + .thenThrow( + SpannerExceptionFactory.newSpannerException( + ErrorCode.INVALID_ARGUMENT, "invalid query")); + when(delegate.getCurrentRowAsStruct()).thenReturn(mock(Struct.class)); + final AtomicInteger rowCount = new AtomicInteger(); + final BlockingDeque receivedErr = new LinkedBlockingDeque<>(1); + try (AsyncResultSetImpl rs = new AsyncResultSetImpl(simpleFactory, delegate)) { + rs.setCallback( + executor, + new ReadyCallback() { + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + try { + if (resultSet.tryNext() != CursorState.DONE) { + rowCount.incrementAndGet(); + return CallbackResponse.CONTINUE; + } + } catch (SpannerException e) { + receivedErr.push(e); + } + return CallbackResponse.DONE; + } + }); + } + Exception e = receivedErr.take(); + assertThat(e).isInstanceOf(SpannerException.class); + SpannerException se = (SpannerException) e; + assertThat(se.getErrorCode()).isEqualTo(ErrorCode.INVALID_ARGUMENT); + assertThat(se.getMessage()).contains("invalid query"); + assertThat(rowCount.get()).isEqualTo(1); + } + + @Test + public void pauseResume() throws InterruptedException { + Executor executor = Executors.newSingleThreadExecutor(); + ResultSet delegate = mock(ResultSet.class); + when(delegate.next()).thenReturn(true, true, true, false); + when(delegate.getCurrentRowAsStruct()).thenReturn(mock(Struct.class)); + final AtomicInteger callbackCounter = new AtomicInteger(); + final BlockingDeque queue = new LinkedBlockingDeque<>(1); + final AtomicBoolean finished = new AtomicBoolean(false); + try (AsyncResultSetImpl rs = new AsyncResultSetImpl(simpleFactory, delegate)) { + rs.setCallback( + executor, + new ReadyCallback() { + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + callbackCounter.incrementAndGet(); + CursorState state = resultSet.tryNext(); + if (state == CursorState.OK) { + try { + queue.put(new Object()); + } catch (InterruptedException e) { + // Finish early if an error occurs. + return CallbackResponse.DONE; + } + return CallbackResponse.PAUSE; + } + finished.set(true); + return CallbackResponse.DONE; + } + }); + int rowCounter = 0; + while (!finished.get()) { + Object o = queue.poll(1L, TimeUnit.MILLISECONDS); + if (o != null) { + rowCounter++; + } + rs.resume(); + } + // There should be exactly 4 callbacks as we only consume one row per callback. + assertThat(callbackCounter.get()).isEqualTo(4); + assertThat(rowCounter).isEqualTo(3); + } + } + + @Test + public void cancel() throws InterruptedException { + Executor executor = Executors.newSingleThreadExecutor(); + ResultSet delegate = mock(ResultSet.class); + when(delegate.next()).thenReturn(true, true, true, false); + when(delegate.getCurrentRowAsStruct()).thenReturn(mock(Struct.class)); + final AtomicInteger callbackCounter = new AtomicInteger(); + final BlockingDeque queue = new LinkedBlockingDeque<>(1); + final AtomicBoolean finished = new AtomicBoolean(false); + try (AsyncResultSetImpl rs = new AsyncResultSetImpl(simpleFactory, delegate)) { + rs.setCallback( + executor, + new ReadyCallback() { + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + callbackCounter.incrementAndGet(); + try { + CursorState state = resultSet.tryNext(); + if (state == CursorState.OK) { + try { + queue.put(new Object()); + } catch (InterruptedException e) { + // Finish early if an error occurs. + return CallbackResponse.DONE; + } + } + // Pause after 2 rows to make sure that no more data is consumed until the cancel + // call has been received. + return callbackCounter.get() == 2 + ? CallbackResponse.PAUSE + : CallbackResponse.CONTINUE; + } catch (SpannerException e) { + if (e.getErrorCode() == ErrorCode.CANCELLED) { + finished.set(true); + } + } + return CallbackResponse.DONE; + } + }); + int rowCounter = 0; + while (!finished.get()) { + Object o = queue.poll(1L, TimeUnit.MILLISECONDS); + if (o != null) { + rowCounter++; + } + if (rowCounter == 2) { + // Cancel the result set and then resume it to get the cancelled error. + rs.cancel(); + rs.resume(); + } + } + assertThat(callbackCounter.get()).isIn(Range.closed(2, 4)); + assertThat(rowCounter).isIn(Range.closed(2, 3)); + } + } + + @Test + public void callbackReturnsError() throws InterruptedException { + Executor executor = Executors.newSingleThreadExecutor(); + ResultSet delegate = mock(ResultSet.class); + when(delegate.next()).thenReturn(true, true, true, false); + when(delegate.getCurrentRowAsStruct()).thenReturn(mock(Struct.class)); + final AtomicInteger callbackCounter = new AtomicInteger(); + try (AsyncResultSetImpl rs = new AsyncResultSetImpl(simpleFactory, delegate)) { + rs.setCallback( + executor, + new ReadyCallback() { + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + callbackCounter.incrementAndGet(); + throw new RuntimeException("async test"); + } + }); + rs.getResult().get(); + fail("missing expected exception"); + } catch (ExecutionException e) { + assertThat(e.getCause()).isInstanceOf(SpannerException.class); + SpannerException se = (SpannerException) e.getCause(); + assertThat(se.getErrorCode()).isEqualTo(ErrorCode.UNKNOWN); + assertThat(se.getMessage()).contains("async test"); + assertThat(callbackCounter.get()).isEqualTo(1); + } + } +} diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java index f665d66adab..066cf5123ad 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java @@ -25,6 +25,8 @@ import com.google.api.gax.grpc.testing.LocalChannelProvider; import com.google.api.gax.retrying.RetrySettings; import com.google.cloud.NoCredentials; +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.ReadContext.QueryAnalyzeMode; @@ -45,6 +47,11 @@ import io.grpc.inprocess.InProcessServerBuilder; import java.io.IOException; import java.util.List; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; import java.util.concurrent.ScheduledThreadPoolExecutor; import java.util.concurrent.TimeUnit; import org.junit.After; @@ -145,6 +152,136 @@ public void tearDown() throws Exception { mockSpanner.removeAllExecutionTimes(); } + @Test + public void write() { + DatabaseClient client = + spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); + client.write( + Arrays.asList( + Mutation.newInsertBuilder("FOO").set("ID").to(1L).set("NAME").to("Bar").build())); + } + + @Test + public void writeAtLeastOnce() { + DatabaseClient client = + spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); + client.writeAtLeastOnce( + Arrays.asList( + Mutation.newInsertBuilder("FOO").set("ID").to(1L).set("NAME").to("Bar").build())); + } + + @Test + public void singleUse() { + DatabaseClient client = + spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); + try (ResultSet rs = client.singleUse().executeQuery(SELECT1)) { + assertThat(rs.next()).isTrue(); + assertThat(rs.getLong(0)).isEqualTo(1L); + assertThat(rs.next()).isFalse(); + } + } + + @Test + public void singleUseBound() { + DatabaseClient client = + spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); + try (ResultSet rs = + client + .singleUse(TimestampBound.ofExactStaleness(15L, TimeUnit.SECONDS)) + .executeQuery(SELECT1)) { + assertThat(rs.next()).isTrue(); + assertThat(rs.getLong(0)).isEqualTo(1L); + assertThat(rs.next()).isFalse(); + } + } + + @Test + public void singleUseTransaction() { + DatabaseClient client = + spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); + try (ResultSet rs = client.singleUseReadOnlyTransaction().executeQuery(SELECT1)) { + assertThat(rs.next()).isTrue(); + assertThat(rs.getLong(0)).isEqualTo(1L); + assertThat(rs.next()).isFalse(); + } + } + + @Test + public void singleUseTransactionBound() { + DatabaseClient client = + spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); + try (ResultSet rs = + client + .singleUseReadOnlyTransaction(TimestampBound.ofExactStaleness(15L, TimeUnit.SECONDS)) + .executeQuery(SELECT1)) { + assertThat(rs.next()).isTrue(); + assertThat(rs.getLong(0)).isEqualTo(1L); + assertThat(rs.next()).isFalse(); + } + } + + @Test + public void readOnlyTransaction() { + DatabaseClient client = + spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); + try (ReadOnlyTransaction tx = + client.readOnlyTransaction(TimestampBound.ofExactStaleness(15L, TimeUnit.SECONDS))) { + try (ResultSet rs = tx.executeQuery(SELECT1)) { + assertThat(rs.next()).isTrue(); + assertThat(rs.getLong(0)).isEqualTo(1L); + assertThat(rs.next()).isFalse(); + } + } + } + + @Test + public void readOnlyTransactionBound() { + DatabaseClient client = + spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); + try (ReadOnlyTransaction tx = + client.readOnlyTransaction(TimestampBound.ofExactStaleness(15L, TimeUnit.SECONDS))) { + try (ResultSet rs = tx.executeQuery(SELECT1)) { + assertThat(rs.next()).isTrue(); + assertThat(rs.getLong(0)).isEqualTo(1L); + assertThat(rs.next()).isFalse(); + } + } + } + + @Test + public void readWriteTransaction() { + DatabaseClient client = + spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); + TransactionRunner runner = client.readWriteTransaction(); + runner.run( + new TransactionCallable() { + @Override + public Void run(TransactionContext transaction) throws Exception { + transaction.executeUpdate(UPDATE_STATEMENT); + return null; + } + }); + } + + @Test + public void transactionManager() throws Exception { + DatabaseClient client = + spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); + try (TransactionManager txManager = client.transactionManager()) { + while (true) { + TransactionContext tx = txManager.begin(); + try { + tx.executeUpdate(UPDATE_STATEMENT); + txManager.commit(); + break; + } catch (AbortedException e) { + Thread.sleep(e.getRetryDelayInMillis() / 1000); + tx = txManager.resetForRetry(); + } + } + } + } + /** * Test that the update statement can be executed as a partitioned transaction that returns a * lower bound update count. @@ -825,4 +962,43 @@ public void testBackendPartitionQueryOptions() { assertThat(request.getQueryOptions().getOptimizerVersion()).isEqualTo("1"); } } + + public void testAsyncQuery() throws InterruptedException { + final int EXPECTED_ROW_COUNT = 10; + RandomResultSetGenerator generator = new RandomResultSetGenerator(EXPECTED_ROW_COUNT); + com.google.spanner.v1.ResultSet resultSet = generator.generate(); + mockSpanner.putStatementResult( + StatementResult.query(Statement.of("SELECT * FROM RANDOM"), resultSet)); + DatabaseClient client = + spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); + ExecutorService executor = Executors.newSingleThreadExecutor(); + final CountDownLatch finished = new CountDownLatch(1); + final List receivedResults = new ArrayList<>(); + try (AsyncResultSet rs = + client.singleUse().executeQueryAsync(Statement.of("SELECT * FROM RANDOM"))) { + rs.setCallback( + executor, + new ReadyCallback() { + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + while (true) { + switch (rs.tryNext()) { + case DONE: + finished.countDown(); + return CallbackResponse.DONE; + case NOT_READY: + return CallbackResponse.CONTINUE; + case OK: + receivedResults.add(resultSet.getCurrentRowAsStruct()); + break; + default: + throw new IllegalStateException("Unknown cursor state"); + } + } + } + }); + } + finished.await(); + assertThat(receivedResults.size()).isEqualTo(EXPECTED_ROW_COUNT); + } } diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/IntegrationTestWithClosedSessionsEnv.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/IntegrationTestWithClosedSessionsEnv.java index 6b22ba77c33..edbc7976c07 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/IntegrationTestWithClosedSessionsEnv.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/IntegrationTestWithClosedSessionsEnv.java @@ -17,6 +17,7 @@ package com.google.cloud.spanner; import com.google.cloud.spanner.SessionPool.PooledSession; +import com.google.cloud.spanner.SessionPool.PooledSessionFuture; import com.google.cloud.spanner.testing.RemoteSpannerHelper; /** @@ -73,30 +74,30 @@ public void setAllowSessionReplacing(boolean allow) { } @Override - PooledSession getReadSession() { - PooledSession session = super.getReadSession(); + PooledSessionFuture getReadSession() { + PooledSessionFuture session = super.getReadSession(); if (invalidateNextSession) { - session.delegate.close(); - session.setAllowReplacing(false); - awaitDeleted(session.delegate); - session.setAllowReplacing(allowReplacing); + session.get().delegate.close(); + session.get().setAllowReplacing(false); + awaitDeleted(session.get().delegate); + session.get().setAllowReplacing(allowReplacing); invalidateNextSession = false; } - session.setAllowReplacing(allowReplacing); + session.get().setAllowReplacing(allowReplacing); return session; } @Override - PooledSession getReadWriteSession() { - PooledSession session = super.getReadWriteSession(); + PooledSessionFuture getReadWriteSession() { + PooledSessionFuture session = super.getReadWriteSession(); if (invalidateNextSession) { - session.delegate.close(); - session.setAllowReplacing(false); - awaitDeleted(session.delegate); - session.setAllowReplacing(allowReplacing); + session.get().delegate.close(); + session.get().setAllowReplacing(false); + awaitDeleted(session.get().delegate); + session.get().setAllowReplacing(allowReplacing); invalidateNextSession = false; } - session.setAllowReplacing(allowReplacing); + session.get().setAllowReplacing(allowReplacing); return session; } diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/RandomResultSetGenerator.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/RandomResultSetGenerator.java new file mode 100644 index 00000000000..63bc234a417 --- /dev/null +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/RandomResultSetGenerator.java @@ -0,0 +1,166 @@ +/* + * Copyright 2019 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.spanner; + +import com.google.api.client.util.Base64; +import com.google.cloud.Date; +import com.google.cloud.Timestamp; +import com.google.protobuf.ListValue; +import com.google.protobuf.NullValue; +import com.google.protobuf.Value; +import com.google.protobuf.util.Timestamps; +import com.google.spanner.v1.ResultSet; +import com.google.spanner.v1.ResultSetMetadata; +import com.google.spanner.v1.StructType; +import com.google.spanner.v1.StructType.Field; +import com.google.spanner.v1.Type; +import com.google.spanner.v1.TypeCode; +import java.util.Random; + +public class RandomResultSetGenerator { + private static final Type TYPES[] = + new Type[] { + Type.newBuilder().setCode(TypeCode.BOOL).build(), + Type.newBuilder().setCode(TypeCode.INT64).build(), + Type.newBuilder().setCode(TypeCode.FLOAT64).build(), + Type.newBuilder().setCode(TypeCode.STRING).build(), + Type.newBuilder().setCode(TypeCode.BYTES).build(), + Type.newBuilder().setCode(TypeCode.DATE).build(), + Type.newBuilder().setCode(TypeCode.TIMESTAMP).build(), + Type.newBuilder() + .setCode(TypeCode.ARRAY) + .setArrayElementType(Type.newBuilder().setCode(TypeCode.BOOL)) + .build(), + Type.newBuilder() + .setCode(TypeCode.ARRAY) + .setArrayElementType(Type.newBuilder().setCode(TypeCode.INT64)) + .build(), + Type.newBuilder() + .setCode(TypeCode.ARRAY) + .setArrayElementType(Type.newBuilder().setCode(TypeCode.FLOAT64)) + .build(), + Type.newBuilder() + .setCode(TypeCode.ARRAY) + .setArrayElementType(Type.newBuilder().setCode(TypeCode.STRING)) + .build(), + Type.newBuilder() + .setCode(TypeCode.ARRAY) + .setArrayElementType(Type.newBuilder().setCode(TypeCode.BYTES)) + .build(), + Type.newBuilder() + .setCode(TypeCode.ARRAY) + .setArrayElementType(Type.newBuilder().setCode(TypeCode.DATE)) + .build(), + Type.newBuilder() + .setCode(TypeCode.ARRAY) + .setArrayElementType(Type.newBuilder().setCode(TypeCode.TIMESTAMP)) + .build(), + }; + + private static final ResultSetMetadata generateMetadata() { + StructType.Builder rowTypeBuilder = StructType.newBuilder(); + for (int col = 0; col < TYPES.length; col++) { + rowTypeBuilder.addFields(Field.newBuilder().setName("COL" + col).setType(TYPES[col])).build(); + } + ResultSetMetadata.Builder builder = ResultSetMetadata.newBuilder(); + builder.setRowType(rowTypeBuilder.build()); + return builder.build(); + } + + private static final ResultSetMetadata METADATA = generateMetadata(); + + private final int rowCount; + private final Random random = new Random(); + + public RandomResultSetGenerator(int rowCount) { + this.rowCount = rowCount; + } + + public ResultSet generate() { + ResultSet.Builder builder = ResultSet.newBuilder(); + for (int row = 0; row < rowCount; row++) { + ListValue.Builder rowBuilder = ListValue.newBuilder(); + for (int col = 0; col < TYPES.length; col++) { + Value.Builder valueBuilder = Value.newBuilder(); + setRandomValue(valueBuilder, TYPES[col]); + rowBuilder.addValues(valueBuilder.build()); + } + builder.addRows(rowBuilder.build()); + } + builder.setMetadata(METADATA); + return builder.build(); + } + + private void setRandomValue(Value.Builder builder, Type type) { + if (randomNull()) { + builder.setNullValue(NullValue.NULL_VALUE); + } else { + switch (type.getCode()) { + case ARRAY: + int length = random.nextInt(20) + 1; + ListValue.Builder arrayBuilder = ListValue.newBuilder(); + for (int i = 0; i < length; i++) { + Value.Builder valueBuilder = Value.newBuilder(); + setRandomValue(valueBuilder, type.getArrayElementType()); + arrayBuilder.addValues(valueBuilder.build()); + } + builder.setListValue(arrayBuilder.build()); + break; + case BOOL: + builder.setBoolValue(random.nextBoolean()); + break; + case STRING: + case BYTES: + byte[] bytes = new byte[random.nextInt(200)]; + random.nextBytes(bytes); + builder.setStringValue(Base64.encodeBase64String(bytes)); + break; + case DATE: + Date date = + Date.fromYearMonthDay( + random.nextInt(2019) + 1, random.nextInt(11) + 1, random.nextInt(28) + 1); + builder.setStringValue(date.toString()); + break; + case FLOAT64: + builder.setNumberValue(random.nextDouble()); + break; + case INT64: + builder.setStringValue(String.valueOf(random.nextLong())); + break; + case TIMESTAMP: + com.google.protobuf.Timestamp ts = + Timestamps.add( + Timestamps.EPOCH, + com.google.protobuf.Duration.newBuilder() + .setSeconds(random.nextInt(100_000_000)) + .setNanos(random.nextInt(1000_000_000)) + .build()); + builder.setStringValue(Timestamp.fromProto(ts).toString()); + break; + case STRUCT: + case TYPE_CODE_UNSPECIFIED: + case UNRECOGNIZED: + default: + throw new IllegalArgumentException("Unknown or unsupported type: " + type.getCode()); + } + } + } + + private boolean randomNull() { + return random.nextInt(10) == 0; + } +} diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolStressTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolStressTest.java index 60eb151f6e4..9f03833e720 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolStressTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolStressTest.java @@ -25,6 +25,7 @@ import com.google.api.core.ApiFuture; import com.google.api.core.ApiFutures; import com.google.cloud.spanner.SessionClient.SessionConsumer; +import com.google.cloud.spanner.SessionPool.PooledSessionFuture; import com.google.cloud.spanner.SessionPool.SessionConsumerImpl; import com.google.common.util.concurrent.Uninterruptibles; import com.google.protobuf.Empty; @@ -97,12 +98,12 @@ private void setupSpanner(DatabaseId db) { when(mockSpanner.getOptions()).thenReturn(spannerOptions); when(sessionClient.createSession()) .thenAnswer( - new Answer() { + new Answer() { @Override - public Session answer(InvocationOnMock invocation) throws Throwable { + public SessionImpl answer(InvocationOnMock invocation) throws Throwable { synchronized (lock) { - Session session = mockSession(); + SessionImpl session = mockSession(); setupSession(session); sessions.put(session.getName(), false); @@ -139,7 +140,7 @@ public Void answer(InvocationOnMock invocation) throws Throwable { .asyncBatchCreateSessions(Mockito.anyInt(), Mockito.any(SessionConsumer.class)); } - private void setupSession(final Session session) { + private void setupSession(final SessionImpl session) { ReadContext mockContext = mock(ReadContext.class); final ResultSet mockResult = mock(ResultSet.class); when(session.singleUse(any(TimestampBound.class))).thenReturn(mockContext); @@ -266,12 +267,14 @@ public void run() { Uninterruptibles.awaitUninterruptibly(releaseThreads); for (int j = 0; j < numOperationsPerThread; j++) { try { - Session session = null; + PooledSessionFuture session = null; if (random.nextInt(10) < writeOperationFraction) { session = pool.getReadWriteSession(); + session.get(); assertWritePrepared(session); } else { session = pool.getReadSession(); + session.get(); } Uninterruptibles.sleepUninterruptibly( random.nextInt(5), TimeUnit.MILLISECONDS); diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java index 6daa36f1d6c..c1acecdbc2f 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java @@ -36,12 +36,14 @@ import com.google.api.core.ApiFuture; import com.google.api.core.ApiFutures; import com.google.cloud.Timestamp; +import com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory; import com.google.cloud.spanner.MetricRegistryTestUtils.FakeMetricRegistry; import com.google.cloud.spanner.MetricRegistryTestUtils.MetricsRecord; import com.google.cloud.spanner.ReadContext.QueryAnalyzeMode; import com.google.cloud.spanner.SessionClient.SessionConsumer; import com.google.cloud.spanner.SessionPool.Clock; import com.google.cloud.spanner.SessionPool.PooledSession; +import com.google.cloud.spanner.SessionPool.PooledSessionFuture; import com.google.cloud.spanner.SessionPool.SessionConsumerImpl; import com.google.cloud.spanner.TransactionRunner.TransactionCallable; import com.google.cloud.spanner.TransactionRunnerImpl.TransactionContextImpl; @@ -177,21 +179,21 @@ public void sessionCreation() { public void poolLifo() { setupMockSessionCreation(); pool = createPool(); - Session session1 = pool.getReadSession(); - Session session2 = pool.getReadSession(); + Session session1 = pool.getReadSession().get(); + Session session2 = pool.getReadSession().get(); assertThat(session1).isNotEqualTo(session2); session2.close(); session1.close(); - Session session3 = pool.getReadSession(); - Session session4 = pool.getReadSession(); + Session session3 = pool.getReadSession().get(); + Session session4 = pool.getReadSession().get(); assertThat(session3).isEqualTo(session1); assertThat(session4).isEqualTo(session2); session3.close(); session4.close(); - Session session5 = pool.getReadWriteSession(); - Session session6 = pool.getReadWriteSession(); + Session session5 = pool.getReadWriteSession().get(); + Session session6 = pool.getReadWriteSession().get(); assertThat(session5).isEqualTo(session4); assertThat(session6).isEqualTo(session3); session6.close(); @@ -232,7 +234,7 @@ public void run() { pool = createPool(); Session session1 = pool.getReadSession(); // Leaked sessions - PooledSession leakedSession = pool.getReadSession(); + PooledSessionFuture leakedSession = pool.getReadSession(); // Clear the leaked exception to suppress logging of expected exceptions. leakedSession.clearLeakedException(); session1.close(); @@ -308,7 +310,7 @@ public Void call() throws Exception { .asyncBatchCreateSessions(Mockito.eq(1), any(SessionConsumer.class)); pool = createPool(); - PooledSession leakedSession = pool.getReadSession(); + PooledSessionFuture leakedSession = pool.getReadSession(); // Suppress expected leakedSession warning. leakedSession.clearLeakedException(); AtomicBoolean failed = new AtomicBoolean(false); @@ -366,7 +368,7 @@ public Void call() throws Exception { .asyncBatchCreateSessions(Mockito.eq(1), any(SessionConsumer.class)); pool = createPool(); - PooledSession leakedSession = pool.getReadSession(); + PooledSessionFuture leakedSession = pool.getReadSession(); // Suppress expected leakedSession warning. leakedSession.clearLeakedException(); AtomicBoolean failed = new AtomicBoolean(false); @@ -483,7 +485,8 @@ public void run() { .when(sessionClient) .asyncBatchCreateSessions(Mockito.eq(1), any(SessionConsumer.class)); pool = createPool(); - PooledSession leakedSession = pool.getReadSession(); + PooledSessionFuture leakedSession = pool.getReadSession(); + leakedSession.get(); // Suppress expected leakedSession warning. leakedSession.clearLeakedException(); pool.closeAsync(); @@ -531,7 +534,7 @@ public Void call() throws Exception { .asyncBatchCreateSessions(Mockito.eq(1), any(SessionConsumer.class)); pool = createPool(); expectedException.expect(isSpannerException(ErrorCode.INTERNAL)); - pool.getReadSession(); + pool.getReadSession().get(); } @Test @@ -558,7 +561,7 @@ public Void call() throws Exception { .asyncBatchCreateSessions(Mockito.eq(1), any(SessionConsumer.class)); pool = createPool(); expectedException.expect(isSpannerException(ErrorCode.INTERNAL)); - pool.getReadWriteSession(); + pool.getReadWriteSession().get(); } @Test @@ -587,7 +590,7 @@ public void run() { .prepareReadWriteTransaction(); pool = createPool(); expectedException.expect(isSpannerException(ErrorCode.INTERNAL)); - pool.getReadWriteSession(); + pool.getReadWriteSession().get(); } @Test @@ -612,14 +615,15 @@ public void run() { .when(sessionClient) .asyncBatchCreateSessions(Mockito.eq(1), any(SessionConsumer.class)); pool = createPool(); - try (Session session = pool.getReadWriteSession()) { + try (PooledSessionFuture session = pool.getReadWriteSession()) { assertThat(session).isNotNull(); + session.get(); verify(mockSession).prepareReadWriteTransaction(); } } @Test - public void getMultipleReadWriteSessions() { + public void getMultipleReadWriteSessions() throws Exception { SessionImpl mockSession1 = mockSession(); SessionImpl mockSession2 = mockSession(); final LinkedList sessions = @@ -643,8 +647,10 @@ public void run() { .when(sessionClient) .asyncBatchCreateSessions(Mockito.eq(1), any(SessionConsumer.class)); pool = createPool(); - Session session1 = pool.getReadWriteSession(); - Session session2 = pool.getReadWriteSession(); + PooledSessionFuture session1 = pool.getReadWriteSession(); + PooledSessionFuture session2 = pool.getReadWriteSession(); + session1.get(); + session2.get(); verify(mockSession1).prepareReadWriteTransaction(); verify(mockSession2).prepareReadWriteTransaction(); session1.close(); @@ -739,8 +745,8 @@ public void run() { pool = createPool(); // One of the sessions would be pre prepared. Uninterruptibles.awaitUninterruptibly(prepareLatch); - PooledSession readSession = pool.getReadSession(); - PooledSession writeSession = pool.getReadWriteSession(); + PooledSession readSession = pool.getReadSession().get(); + PooledSession writeSession = pool.getReadWriteSession().get(); verify(writeSession.delegate, times(1)).prepareReadWriteTransaction(); verify(readSession.delegate, never()).prepareReadWriteTransaction(); readSession.close(); @@ -789,7 +795,7 @@ public void run() { pool.getReadWriteSession().close(); prepareLatch.await(); // This session should also be write prepared. - PooledSession readSession = pool.getReadSession(); + PooledSession readSession = pool.getReadSession().get(); verify(readSession.delegate, times(2)).prepareReadWriteTransaction(); } @@ -857,7 +863,7 @@ public void run() { .when(sessionClient) .asyncBatchCreateSessions(Mockito.eq(1), any(SessionConsumer.class)); pool = createPool(); - assertThat(pool.getReadWriteSession().delegate).isEqualTo(mockSession2); + assertThat(pool.getReadWriteSession().get().delegate).isEqualTo(mockSession2); } @Test @@ -912,9 +918,14 @@ public ApiFuture answer(InvocationOnMock invocation) throws Throwable { pool.getReadSession().close(); runMaintainanceLoop(clock, pool, pool.poolMaintainer.numClosureCycles); assertThat(numSessionClosed.get()).isEqualTo(0); - Session readSession1 = pool.getReadSession(); - Session readSession2 = pool.getReadSession(); - Session readSession3 = pool.getReadSession(); + PooledSessionFuture readSession1 = pool.getReadSession(); + PooledSessionFuture readSession2 = pool.getReadSession(); + PooledSessionFuture readSession3 = pool.getReadSession(); + // Wait until the sessions have actually been gotten in order to make sure they are in use in + // parallel. + readSession1.get(); + readSession2.get(); + readSession3.get(); readSession1.close(); readSession2.close(); readSession3.close(); @@ -995,7 +1006,8 @@ public void blockAndTimeoutOnPoolExhaustion() throws Exception { setupMockSessionCreation(); pool = createPool(); // Take the only session that can be in the pool. - Session checkedOutSession = pool.getReadSession(); + PooledSessionFuture checkedOutSession = pool.getReadSession(); + checkedOutSession.get(); final Boolean finWrite = write; ExecutorService executor = Executors.newFixedThreadPool(1); final CountDownLatch latch = new CountDownLatch(1); @@ -1005,7 +1017,7 @@ public void blockAndTimeoutOnPoolExhaustion() throws Exception { new Callable() { @Override public Void call() throws Exception { - Session session; + PooledSessionFuture session; latch.countDown(); if (finWrite) { session = pool.getReadWriteSession(); @@ -1282,7 +1294,7 @@ public void run() { SessionPool pool = SessionPool.createPool( options, new TestExecutorFactory(), spanner.getSessionClient(db)); - try (PooledSession readWriteSession = pool.getReadWriteSession()) { + try (PooledSessionFuture readWriteSession = pool.getReadWriteSession()) { TransactionRunner runner = readWriteSession.readWriteTransaction(); try { runner.run( @@ -1406,7 +1418,7 @@ public void run() { FakeClock clock = new FakeClock(); clock.currentTimeMillis = System.currentTimeMillis(); pool = createPool(clock); - PooledSession session = pool.getReadWriteSession(); + PooledSession session = pool.getReadWriteSession().get(); assertThat(session.delegate).isEqualTo(openSession); } @@ -1586,8 +1598,10 @@ public void testSessionMetrics() throws Exception { setupMockSessionCreation(); pool = createPool(clock, metricRegistry, labelValues); - Session session1 = pool.getReadSession(); - Session session2 = pool.getReadSession(); + PooledSessionFuture session1 = pool.getReadSession(); + PooledSessionFuture session2 = pool.getReadSession(); + session1.get(); + session2.get(); MetricsRecord record = metricRegistry.pollRecord(); assertThat(record.getMetrics().size()).isEqualTo(6); @@ -1654,7 +1668,8 @@ private void getSessionAsync(final CountDownLatch latch, final AtomicBoolean fai new Runnable() { @Override public void run() { - try (Session session = pool.getReadSession()) { + try (PooledSessionFuture future = pool.getReadSession()) { + PooledSession session = future.get(); failed.compareAndSet(false, session == null); Uninterruptibles.sleepUninterruptibly(10, TimeUnit.MILLISECONDS); } catch (Throwable e) { @@ -1672,7 +1687,8 @@ private void getReadWriteSessionAsync(final CountDownLatch latch, final AtomicBo new Runnable() { @Override public void run() { - try (Session session = pool.getReadWriteSession()) { + try (PooledSessionFuture future = pool.getReadWriteSession()) { + PooledSession session = future.get(); failed.compareAndSet(false, session == null); Uninterruptibles.sleepUninterruptibly(2, TimeUnit.MILLISECONDS); } catch (SpannerException e) { diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionContextImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionContextImplTest.java index 077b6605766..d319d13de2b 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionContextImplTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionContextImplTest.java @@ -19,6 +19,7 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; +import com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory; import com.google.cloud.spanner.TransactionRunnerImpl.TransactionContextImpl; import com.google.cloud.spanner.spi.v1.SpannerRpc; import com.google.protobuf.ByteString; From c7db649ed353ac7812ee6dc418978101152e8d5e Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Sun, 23 Feb 2020 22:18:18 +0100 Subject: [PATCH 02/49] feat: session pool is non-blocking --- .../cloud/spanner/AbstractReadContext.java | 7 +- .../google/cloud/spanner/AsyncResultSet.java | 32 +-- .../cloud/spanner/AsyncResultSetImpl.java | 85 +++++-- .../google/cloud/spanner/DatabaseClient.java | 33 +++ .../cloud/spanner/DatabaseClientImpl.java | 13 + .../cloud/spanner/ForwardingResultSet.java | 21 +- .../cloud/spanner/ForwardingStructReader.java | 146 +++++++---- .../com/google/cloud/spanner/SessionImpl.java | 41 +++ .../com/google/cloud/spanner/SessionPool.java | 192 ++++++++++---- .../cloud/spanner/SessionPoolOptions.java | 19 ++ .../com/google/cloud/spanner/SpannerImpl.java | 9 +- .../google/cloud/spanner/SpannerOptions.java | 43 ++++ .../cloud/spanner/TransactionRunnerImpl.java | 3 +- .../spanner/AsyncResultSetImplStressTest.java | 33 +-- .../cloud/spanner/AsyncResultSetImplTest.java | 51 ++-- .../cloud/spanner/DatabaseClientImplTest.java | 234 +++++++++++++++++- .../cloud/spanner/MockSpannerServiceImpl.java | 36 +-- .../google/cloud/spanner/SessionPoolTest.java | 2 +- .../spanner/TransactionContextImplTest.java | 2 +- .../spanner/TransactionRunnerImplTest.java | 1 + 20 files changed, 780 insertions(+), 223 deletions(-) diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java index 7f7d02f95d5..1b10f3ba34d 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java @@ -21,8 +21,8 @@ import static com.google.common.base.Preconditions.checkNotNull; import static com.google.common.base.Preconditions.checkState; +import com.google.api.gax.core.ExecutorProvider; import com.google.cloud.Timestamp; -import com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory; import com.google.cloud.spanner.AbstractResultSet.CloseableIterator; import com.google.cloud.spanner.AbstractResultSet.GrpcResultSet; import com.google.cloud.spanner.AbstractResultSet.GrpcStreamIterator; @@ -46,7 +46,6 @@ import io.opencensus.trace.Span; import io.opencensus.trace.Tracing; import java.util.Map; -import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.atomic.AtomicLong; import javax.annotation.Nullable; import javax.annotation.concurrent.GuardedBy; @@ -352,7 +351,7 @@ void initTransaction() { final Object lock = new Object(); final SessionImpl session; final SpannerRpc rpc; - final ExecutorFactory executorFactory; + final ExecutorProvider executorProvider; final Span span; private final int defaultPrefetchChunks; private final QueryOptions defaultQueryOptions; @@ -421,7 +420,7 @@ public final ResultSet executeQuery(Statement statement, QueryOption... options) @Override public final AsyncResultSet executeQueryAsync(Statement statement, QueryOption... options) { return new AsyncResultSetImpl( - executorFactory, + executorProvider, executeQueryInternal( statement, com.google.spanner.v1.ExecuteSqlRequest.QueryMode.NORMAL, options)); } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncResultSet.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncResultSet.java index cb052042257..79c4b3b7686 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncResultSet.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncResultSet.java @@ -21,10 +21,17 @@ import com.google.common.collect.ImmutableList; import java.util.concurrent.Executor; -public interface AsyncResultSet extends AutoCloseable, StructReader { +/** Interface for result sets returned by async query methods. */ +public interface AsyncResultSet extends ResultSet { + + /** + * Interface for receiving asynchronous callbacks when new data is ready. See {@link + * AsyncResultSet#setCallback(Executor, ReadyCallback)}. + */ public static interface ReadyCallback { CallbackResponse cursorReady(AsyncResultSet resultSet); } + /** Response code from {@code tryNext()}. */ public enum CursorState { /** Cursor has been moved to a new row. */ @@ -35,16 +42,6 @@ public enum CursorState { NOT_READY } - @Override - void close(); - - /** - * Creates an immutable version of the row that the result set is positioned over. This may - * involve copying internal data structures, and so converting all rows to {@code Struct} objects - * is generally more expensive than processing the {@code ResultSet} directly. - */ - Struct getCurrentRowAsStruct(); - /** * Non-blocking call that attempts to step the cursor to the next position in the stream. The * cursor may be inspected only if the cursor returns {@code CursorState.OK}. @@ -87,13 +84,9 @@ public enum CursorState { * *
  • Callback may possibly be invoked after a call to {@link ResultSet#cancel()} call, but the * subsequent call to {@link #tryNext()} will yield a SpannerException. - *
  • Spurious callbacks are possible where cursors is not actually ready. Typically callback + *
  • Spurious callbacks are possible where cursors are not actually ready. Typically callback * should return {@link CallbackResponse#CONTINUE} any time it sees {@link - * CursorState#NOT_READY}. This is similar to pthreads "Spurious Wakeups", - * http://en.wikipedia.org/wiki/Spurious_wakeup TODO: consider squelching spurious wakeups - * by adding a "lookahead & store result" buffer of at most 1 item. Reasons to is to - * simplify this explanation, but user code is unlikely to change either way... its just - * weird. + * CursorState#NOT_READY}. * * *

    Flow Control

    @@ -133,11 +126,6 @@ public enum CursorState { * callback again. * * - * Note that it would have been equivalent to have the app be responsible for draining the cursor - * instead of calling {@code resume()} (which has basically the same effect, namely running the - * application callback.) The explicit pause and resume was chosen to make the flow control - * behavior more explicit in application code. - * * @param exec executor on which to run all callbacks. Typically use a threadpool. If the executor * is one that runs the work on the submitting thread, you must be very careful not to throw * RuntimeException up the stack, lest you do damage to calling components. For example, it 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 9b7609d4dca..fcd0bc770ae 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 @@ -18,11 +18,12 @@ import com.google.api.core.ApiFuture; import com.google.api.core.SettableApiFuture; -import com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory; +import com.google.api.gax.core.ExecutorProvider; import com.google.common.base.Function; import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableList; import com.google.common.util.concurrent.MoreExecutors; +import com.google.spanner.v1.ResultSetStats; import java.util.concurrent.BlockingDeque; import java.util.concurrent.Callable; import java.util.concurrent.CountDownLatch; @@ -37,6 +38,8 @@ class AsyncResultSetImpl extends ForwardingStructReader implements AsyncResultSe /** State of an {@link AsyncResultSetImpl}. */ private enum State { INITIALIZED, + /** SYNC indicates that the {@link ResultSet} is used in sync pattern. */ + SYNC, CONSUMING, RUNNING, PAUSED, @@ -60,13 +63,15 @@ private State(boolean shouldStop) { private final Object monitor = new Object(); private boolean closed; + /** - * {@link ExecutorFactory} produces executors that are used to fetch data from the backend and put - * these into the buffer for further consumpation by the callback. + * {@link ExecutorProvider} provides executor services that are used to fetch data from the + * backend and put these into the buffer for further consumption by the callback. */ - private final ExecutorFactory executorFactory; + private final ExecutorProvider executorProvider; private final ScheduledExecutorService service; + private final BlockingDeque buffer; private Struct currentRow; /** The underlying synchronous {@link ResultSet} that is producing the rows. */ @@ -93,7 +98,7 @@ private State(boolean shouldStop) { */ private volatile boolean finished; - private final Future result; + private volatile Future result; /** * {@link #cursorReturnedDoneOrException} indicates whether {@link #tryNext()} has returned {@link @@ -117,22 +122,16 @@ private State(boolean shouldStop) { */ private volatile CountDownLatch consumingLatch = new CountDownLatch(0); - AsyncResultSetImpl( - ExecutorFactory executorFactory, ResultSet delegate) { - this(executorFactory, delegate, DEFAULT_BUFFER_SIZE); + AsyncResultSetImpl(ExecutorProvider executorProvider, ResultSet delegate) { + this(executorProvider, delegate, DEFAULT_BUFFER_SIZE); } - AsyncResultSetImpl( - ExecutorFactory executorFactory, - ResultSet delegate, - int bufferSize) { + AsyncResultSetImpl(ExecutorProvider executorProvider, ResultSet delegate, int bufferSize) { super(delegate); this.buffer = new LinkedBlockingDeque<>(bufferSize); - this.executorFactory = executorFactory; - this.service = executorFactory.get(); + this.executorProvider = executorProvider; + this.service = executorProvider.getExecutor(); this.delegateResultSet = delegate; - // Eagerly start to fetch data and buffer these. - this.result = this.service.submit(new ProduceRowsCallable()); } /** @@ -148,14 +147,13 @@ public void close() { if (this.closed) { return; } + if (state == State.INITIALIZED || state == State.SYNC) { + delegateResultSet.close(); + } this.closed = true; } } - public Struct getCurrentRowAsStruct() { - return currentRow; - } - /** * Tries to advance this {@link AsyncResultSet} to the next row. This method may only be called * from within a {@link ReadyCallback}. @@ -347,7 +345,9 @@ public Void call() throws Exception { } } finally { delegateResultSet.close(); - executorFactory.release(service); + if (executorProvider.shouldAutoClose()) { + service.shutdown(); + } synchronized (monitor) { if (executionException != null) { throw executionException; @@ -393,6 +393,9 @@ public void setCallback(Executor exec, ReadyCallback cb) { Preconditions.checkState(!closed, "This AsyncResultSet has been closed"); Preconditions.checkState( this.state == State.INITIALIZED, "callback may not be set multiple times"); + + // Start to fetch data and buffer these. + this.result = this.service.submit(new ProduceRowsCallable()); this.executor = MoreExecutors.newSequentialExecutor(Preconditions.checkNotNull(exec)); this.callback = Preconditions.checkNotNull(cb); this.state = State.RUNNING; @@ -408,7 +411,8 @@ Future getResult() { public void cancel() { synchronized (monitor) { Preconditions.checkState( - state != State.INITIALIZED, "cannot cancel a result set without a callback"); + state != State.INITIALIZED && state != State.SYNC, + "cannot cancel a result set without a callback"); state = State.CANCELLED; pausedLatch.countDown(); } @@ -418,7 +422,8 @@ public void cancel() { public void resume() { synchronized (monitor) { Preconditions.checkState( - state != State.INITIALIZED, "cannot resume a result set without a callback"); + state != State.INITIALIZED && state != State.SYNC, + "cannot resume a result set without a callback"); if (state == State.PAUSED) { state = State.RUNNING; pausedLatch.countDown(); @@ -482,4 +487,38 @@ public ImmutableList toList(Function transformer) throw SpannerExceptionFactory.newSpannerException(e); } } + + @Override + public boolean next() throws SpannerException { + synchronized (monitor) { + Preconditions.checkState( + this.state == State.INITIALIZED || this.state == State.SYNC, + "Cannot call next() on a result set with a callback."); + this.state = State.SYNC; + } + boolean res = delegateResultSet.next(); + currentRow = delegateResultSet.getCurrentRowAsStruct(); + return res; + } + + @Override + public ResultSetStats getStats() { + return delegateResultSet.getStats(); + } + + @Override + protected void checkValidState() { + synchronized (monitor) { + Preconditions.checkState( + state == State.SYNC || state == State.CONSUMING || state == State.CANCELLED, + "only allowed after a next() call or from within a ReadyCallback#cursorReady callback"); + Preconditions.checkState(state != State.SYNC || !closed, "ResultSet is closed"); + } + } + + @Override + public Struct getCurrentRowAsStruct() { + checkValidState(); + return currentRow; + } } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClient.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClient.java index ac29ba2b374..dd9e9769066 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClient.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClient.java @@ -16,7 +16,9 @@ package com.google.cloud.spanner; +import com.google.api.core.ApiFuture; import com.google.cloud.Timestamp; +import java.util.concurrent.Executor; /** * Interface for all the APIs that are used to read/write data into a Cloud Spanner database. An @@ -278,6 +280,37 @@ public interface DatabaseClient { */ TransactionManager transactionManager(); + public static interface AsyncWork { + /** + * Performs a single transaction attempt. All reads/writes should be performed using {@code + * txn}. + * + *

    Implementations of this method should not attempt to commit the transaction directly: + * returning normally will result in the runner attempting to commit the transaction once the + * returned future completes, retrying on abort. + * + *

    In most cases, the implementation will not need to catch {@code SpannerException}s from + * Spanner operations, instead letting these propagate to the framework. The transaction runner + * + *

    will take appropriate action based on the type of exception. In particular, + * implementations should never catch an exception of type {@link SpannerErrors#isAborted}: + * these indicate that some reads may have returned inconsistent data and the transaction + * attempt must be aborted. + * + *

    If any exception is thrown, the runner will validate the reads performed in the current + * transaction attempt using {@link Transaction#commitReadsOnly}: if validation succeeds, the + * exception is propagated to the caller; if validation aborts, the exception is thrown away and + * the work is retried; if the commit fails for some other reason, the corresponding {@code + * SpannerException} is returned to the caller. Any buffered mutations will be ignored. + * + * @param txn the transaction + * @return future over the result of the work + */ + ApiFuture doWorkAsync(TransactionContext txn); + } + + ApiFuture runAsync(AsyncWork work, Executor executor); + /** * Returns the lower bound of rows modified by this DML statement. * diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClientImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClientImpl.java index 607684611c4..f1bd75b3ccb 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClientImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClientImpl.java @@ -16,6 +16,7 @@ package com.google.cloud.spanner; +import com.google.api.core.ApiFuture; import com.google.cloud.Timestamp; import com.google.cloud.spanner.SessionPool.PooledSessionFuture; import com.google.common.annotations.VisibleForTesting; @@ -25,6 +26,7 @@ import io.opencensus.trace.Span; import io.opencensus.trace.Tracer; import io.opencensus.trace.Tracing; +import java.util.concurrent.Executor; class DatabaseClientImpl implements DatabaseClient { private static final String READ_WRITE_TRANSACTION = "CloudSpanner.ReadWriteTransaction"; @@ -190,6 +192,17 @@ public TransactionManager transactionManager() { } } + @Override + public ApiFuture runAsync(AsyncWork work, Executor executor) { + Span span = tracer.spanBuilder(READ_WRITE_TRANSACTION).startSpan(); + try (Scope s = tracer.withSpan(span)) { + return getReadWriteSession().runAsync(work, executor); + } catch (RuntimeException e) { + TraceUtil.endSpanWithFailure(span, e); + throw e; + } + } + @Override public long executePartitionedUpdate(final Statement stmt) { Span span = tracer.spanBuilder(PARTITION_DML_TRANSACTION).startSpan(); diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ForwardingResultSet.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ForwardingResultSet.java index 753c3f6f390..4cc0ab9b9e7 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ForwardingResultSet.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ForwardingResultSet.java @@ -17,16 +17,23 @@ package com.google.cloud.spanner; import com.google.common.base.Preconditions; +import com.google.common.base.Supplier; +import com.google.common.base.Suppliers; import com.google.spanner.v1.ResultSetStats; /** Forwarding implementation of ResultSet that forwards all calls to a delegate. */ public class ForwardingResultSet extends ForwardingStructReader implements ResultSet { - private ResultSet delegate; + private Supplier delegate; public ForwardingResultSet(ResultSet delegate) { super(delegate); - this.delegate = Preconditions.checkNotNull(delegate); + this.delegate = Suppliers.ofInstance(Preconditions.checkNotNull(delegate)); + } + + public ForwardingResultSet(Supplier supplier) { + super(supplier); + this.delegate = supplier; } /** @@ -39,26 +46,26 @@ public ForwardingResultSet(ResultSet delegate) { void replaceDelegate(ResultSet newDelegate) { Preconditions.checkNotNull(newDelegate); super.replaceDelegate(newDelegate); - this.delegate = newDelegate; + this.delegate = Suppliers.ofInstance(Preconditions.checkNotNull(newDelegate)); } @Override public boolean next() throws SpannerException { - return delegate.next(); + return delegate.get().next(); } @Override public Struct getCurrentRowAsStruct() { - return delegate.getCurrentRowAsStruct(); + return delegate.get().getCurrentRowAsStruct(); } @Override public void close() { - delegate.close(); + delegate.get().close(); } @Override public ResultSetStats getStats() { - return delegate.getStats(); + return delegate.get().getStats(); } } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ForwardingStructReader.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ForwardingStructReader.java index 9b30b899852..67e546ad5a6 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ForwardingStructReader.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ForwardingStructReader.java @@ -20,14 +20,20 @@ import com.google.cloud.Date; import com.google.cloud.Timestamp; import com.google.common.base.Preconditions; +import com.google.common.base.Supplier; +import com.google.common.base.Suppliers; import java.util.List; /** Forwarding implements of StructReader */ public class ForwardingStructReader implements StructReader { - private StructReader delegate; + private Supplier delegate; public ForwardingStructReader(StructReader delegate) { + this.delegate = Suppliers.ofInstance(Preconditions.checkNotNull(delegate)); + } + + public ForwardingStructReader(Supplier delegate) { this.delegate = Preconditions.checkNotNull(delegate); } @@ -39,221 +45,271 @@ public ForwardingStructReader(StructReader delegate) { * returned to the user. */ void replaceDelegate(StructReader newDelegate) { - this.delegate = Preconditions.checkNotNull(newDelegate); + this.delegate = Suppliers.ofInstance(Preconditions.checkNotNull(newDelegate)); } + /** + * Called before each forwarding call to allow sub classes to do additional state checking. Sub + * classes should throw an {@link Exception} if the current state is not valid for reading data + * from this {@link ForwardingStructReader}. The default implementation does nothing. + */ + protected void checkValidState() {} + @Override public Type getType() { - return delegate.getType(); + checkValidState(); + return delegate.get().getType(); } @Override public int getColumnCount() { - return delegate.getColumnCount(); + checkValidState(); + return delegate.get().getColumnCount(); } @Override public int getColumnIndex(String columnName) { - return delegate.getColumnIndex(columnName); + checkValidState(); + return delegate.get().getColumnIndex(columnName); } @Override public Type getColumnType(int columnIndex) { - return delegate.getColumnType(columnIndex); + checkValidState(); + return delegate.get().getColumnType(columnIndex); } @Override public Type getColumnType(String columnName) { - return delegate.getColumnType(columnName); + checkValidState(); + return delegate.get().getColumnType(columnName); } @Override public boolean isNull(int columnIndex) { - return delegate.isNull(columnIndex); + checkValidState(); + return delegate.get().isNull(columnIndex); } @Override public boolean isNull(String columnName) { - return delegate.isNull(columnName); + checkValidState(); + return delegate.get().isNull(columnName); } @Override public boolean getBoolean(int columnIndex) { - return delegate.getBoolean(columnIndex); + checkValidState(); + return delegate.get().getBoolean(columnIndex); } @Override public boolean getBoolean(String columnName) { - return delegate.getBoolean(columnName); + checkValidState(); + return delegate.get().getBoolean(columnName); } @Override public long getLong(int columnIndex) { - return delegate.getLong(columnIndex); + checkValidState(); + return delegate.get().getLong(columnIndex); } @Override public long getLong(String columnName) { - return delegate.getLong(columnName); + checkValidState(); + return delegate.get().getLong(columnName); } @Override public double getDouble(int columnIndex) { - return delegate.getDouble(columnIndex); + checkValidState(); + return delegate.get().getDouble(columnIndex); } @Override public double getDouble(String columnName) { - return delegate.getDouble(columnName); + checkValidState(); + return delegate.get().getDouble(columnName); } @Override public String getString(int columnIndex) { - return delegate.getString(columnIndex); + checkValidState(); + return delegate.get().getString(columnIndex); } @Override public String getString(String columnName) { - return delegate.getString(columnName); + checkValidState(); + return delegate.get().getString(columnName); } @Override public ByteArray getBytes(int columnIndex) { - return delegate.getBytes(columnIndex); + checkValidState(); + return delegate.get().getBytes(columnIndex); } @Override public ByteArray getBytes(String columnName) { - return delegate.getBytes(columnName); + checkValidState(); + return delegate.get().getBytes(columnName); } @Override public Timestamp getTimestamp(int columnIndex) { - return delegate.getTimestamp(columnIndex); + checkValidState(); + return delegate.get().getTimestamp(columnIndex); } @Override public Timestamp getTimestamp(String columnName) { - return delegate.getTimestamp(columnName); + checkValidState(); + return delegate.get().getTimestamp(columnName); } @Override public Date getDate(int columnIndex) { - return delegate.getDate(columnIndex); + checkValidState(); + return delegate.get().getDate(columnIndex); } @Override public Date getDate(String columnName) { - return delegate.getDate(columnName); + checkValidState(); + return delegate.get().getDate(columnName); } @Override public boolean[] getBooleanArray(int columnIndex) { - return delegate.getBooleanArray(columnIndex); + checkValidState(); + return delegate.get().getBooleanArray(columnIndex); } @Override public boolean[] getBooleanArray(String columnName) { - return delegate.getBooleanArray(columnName); + checkValidState(); + return delegate.get().getBooleanArray(columnName); } @Override public List getBooleanList(int columnIndex) { - return delegate.getBooleanList(columnIndex); + checkValidState(); + return delegate.get().getBooleanList(columnIndex); } @Override public List getBooleanList(String columnName) { - return delegate.getBooleanList(columnName); + checkValidState(); + return delegate.get().getBooleanList(columnName); } @Override public long[] getLongArray(int columnIndex) { - return delegate.getLongArray(columnIndex); + checkValidState(); + return delegate.get().getLongArray(columnIndex); } @Override public long[] getLongArray(String columnName) { - return delegate.getLongArray(columnName); + checkValidState(); + return delegate.get().getLongArray(columnName); } @Override public List getLongList(int columnIndex) { - return delegate.getLongList(columnIndex); + checkValidState(); + return delegate.get().getLongList(columnIndex); } @Override public List getLongList(String columnName) { - return delegate.getLongList(columnName); + checkValidState(); + return delegate.get().getLongList(columnName); } @Override public double[] getDoubleArray(int columnIndex) { - return delegate.getDoubleArray(columnIndex); + checkValidState(); + return delegate.get().getDoubleArray(columnIndex); } @Override public double[] getDoubleArray(String columnName) { - return delegate.getDoubleArray(columnName); + checkValidState(); + return delegate.get().getDoubleArray(columnName); } @Override public List getDoubleList(int columnIndex) { - return delegate.getDoubleList(columnIndex); + checkValidState(); + return delegate.get().getDoubleList(columnIndex); } @Override public List getDoubleList(String columnName) { - return delegate.getDoubleList(columnName); + checkValidState(); + return delegate.get().getDoubleList(columnName); } @Override public List getStringList(int columnIndex) { - return delegate.getStringList(columnIndex); + checkValidState(); + return delegate.get().getStringList(columnIndex); } @Override public List getStringList(String columnName) { - return delegate.getStringList(columnName); + checkValidState(); + return delegate.get().getStringList(columnName); } @Override public List getBytesList(int columnIndex) { - return delegate.getBytesList(columnIndex); + checkValidState(); + return delegate.get().getBytesList(columnIndex); } @Override public List getBytesList(String columnName) { - return delegate.getBytesList(columnName); + checkValidState(); + return delegate.get().getBytesList(columnName); } @Override public List getTimestampList(int columnIndex) { - return delegate.getTimestampList(columnIndex); + checkValidState(); + return delegate.get().getTimestampList(columnIndex); } @Override public List getTimestampList(String columnName) { - return delegate.getTimestampList(columnName); + checkValidState(); + return delegate.get().getTimestampList(columnName); } @Override public List getDateList(int columnIndex) { - return delegate.getDateList(columnIndex); + checkValidState(); + return delegate.get().getDateList(columnIndex); } @Override public List getDateList(String columnName) { - return delegate.getDateList(columnName); + checkValidState(); + return delegate.get().getDateList(columnName); } @Override public List getStructList(int columnIndex) { - return delegate.getStructList(columnIndex); + checkValidState(); + return delegate.get().getStructList(columnIndex); } @Override public List getStructList(String columnName) { - return delegate.getStructList(columnName); + checkValidState(); + return delegate.get().getStructList(columnName); } } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java index 015e1862d6f..0a8bfd2316d 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java @@ -20,11 +20,16 @@ import static com.google.common.base.Preconditions.checkNotNull; import com.google.api.core.ApiFuture; +import com.google.api.core.SettableApiFuture; import com.google.cloud.Timestamp; import com.google.cloud.spanner.AbstractReadContext.MultiUseReadOnlyTransaction; import com.google.cloud.spanner.AbstractReadContext.SingleReadContext; import com.google.cloud.spanner.AbstractReadContext.SingleUseReadOnlyTransaction; +<<<<<<< HEAD import com.google.cloud.spanner.SessionClient.SessionId; +======= +import com.google.cloud.spanner.TransactionRunner.TransactionCallable; +>>>>>>> feat: session pool is non-blocking import com.google.cloud.spanner.TransactionRunnerImpl.TransactionContextImpl; import com.google.cloud.spanner.spi.v1.SpannerRpc; import com.google.common.collect.Lists; @@ -43,6 +48,8 @@ import java.util.Collection; import java.util.List; import java.util.Map; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executor; import javax.annotation.Nullable; /** @@ -212,6 +219,40 @@ public TransactionRunner readWriteTransaction() { new TransactionRunnerImpl(this, spanner.getRpc(), spanner.getDefaultPrefetchChunks())); } + @Override + public ApiFuture runAsync(final AsyncWork work, Executor executor) { + final SettableApiFuture res = SettableApiFuture.create(); + final TransactionRunner runner = + setActive( + new TransactionRunnerImpl(this, spanner.getRpc(), spanner.getDefaultPrefetchChunks())); + executor.execute( + new Runnable() { + @Override + public void run() { + try { + R r = + runner.run( + new TransactionCallable() { + @Override + public R run(TransactionContext transaction) throws Exception { + try { + return work.doWorkAsync(transaction).get(); + } catch (ExecutionException e) { + throw SpannerExceptionFactory.newSpannerException(e.getCause()); + } catch (InterruptedException e) { + throw SpannerExceptionFactory.propagateInterrupt(e); + } + } + }); + res.set(r); + } catch (Throwable t) { + res.setException(t); + } + } + }); + return res; + } + @Override public void prepareReadWriteTransaction() { setActive(null); diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java index 4880cfb370d..facab374cd0 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java @@ -35,9 +35,11 @@ import com.google.api.core.ApiFuture; import com.google.api.core.ApiFutures; +import com.google.api.core.SettableApiFuture; import com.google.cloud.Timestamp; import com.google.cloud.grpc.GrpcTransportOptions; import com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory; +import com.google.cloud.spanner.DatabaseClient.AsyncWork; import com.google.cloud.spanner.Options.QueryOption; import com.google.cloud.spanner.Options.ReadOption; import com.google.cloud.spanner.SessionClient.SessionConsumer; @@ -87,6 +89,7 @@ import java.util.concurrent.Callable; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executor; import java.util.concurrent.Future; import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.ScheduledExecutorService; @@ -123,6 +126,24 @@ Instant instant() { } } + private abstract static class CachedResultSetSupplier implements Supplier { + private ResultSet cached; + + abstract ResultSet load(); + + ResultSet reload() { + return cached = load(); + } + + @Override + public ResultSet get() { + if (cached == null) { + cached = load(); + } + return cached; + } + } + /** * Wrapper around {@code ReadContext} that releases the session to the pool once the call is * finished, if it is a single use context. @@ -145,31 +166,33 @@ private AutoClosingReadContext( this.sessionPool = sessionPool; this.session = session; this.isSingleUse = isSingleUse; - while (true) { - try { - this.readContextDelegate = readContextDelegateSupplier.apply(this.session); - break; - } catch (SessionNotFoundException e) { - replaceSessionIfPossible(e); - } - } } T getReadContextDelegate() { + if (readContextDelegate == null) { + while (true) { + try { + this.readContextDelegate = readContextDelegateSupplier.apply(this.session); + break; + } catch (SessionNotFoundException e) { + replaceSessionIfPossible(e); + } + } + } return readContextDelegate; } - private ResultSet wrap(final Supplier resultSetSupplier) { - ResultSet res; - while (true) { - try { - res = resultSetSupplier.get(); - break; - } catch (SessionNotFoundException e) { - replaceSessionIfPossible(e); - } - } - return new ForwardingResultSet(res) { + private ResultSet wrap(final CachedResultSetSupplier resultSetSupplier) { + // ResultSet res; + // while (true) { + // try { + // res = resultSetSupplier.get(); + // break; + // } catch (SessionNotFoundException e) { + // replaceSessionIfPossible(e); + // } + // } + return new ForwardingResultSet(resultSetSupplier) { private boolean beforeFirst = true; @Override @@ -178,8 +201,18 @@ public boolean next() throws SpannerException { try { return internalNext(); } catch (SessionNotFoundException e) { - replaceSessionIfPossible(e); - replaceDelegate(resultSetSupplier.get()); + while (true) { + // Keep the replace-if-possible outside the try-block to let the exception bubble up + // if it's too late to replace the session. + replaceSessionIfPossible(e); + try { + replaceDelegate(resultSetSupplier.reload()); + break; + } catch (SessionNotFoundException snfe) { + e = snfe; + // retry on yet another session. + } + } } } } @@ -235,10 +268,10 @@ public ResultSet read( final Iterable columns, final ReadOption... options) { return wrap( - new Supplier() { + new CachedResultSetSupplier() { @Override - public ResultSet get() { - return readContextDelegate.read(table, keys, columns, options); + ResultSet load() { + return getReadContextDelegate().read(table, keys, columns, options); } }); } @@ -251,10 +284,10 @@ public ResultSet readUsingIndex( final Iterable columns, final ReadOption... options) { return wrap( - new Supplier() { + new CachedResultSetSupplier() { @Override - public ResultSet get() { - return readContextDelegate.readUsingIndex(table, index, keys, columns, options); + ResultSet load() { + return getReadContextDelegate().readUsingIndex(table, index, keys, columns, options); } }); } @@ -266,7 +299,7 @@ public Struct readRow(String table, Key key, Iterable columns) { while (true) { try { session.get().markUsed(); - return readContextDelegate.readRow(table, key, columns); + return getReadContextDelegate().readRow(table, key, columns); } catch (SessionNotFoundException e) { replaceSessionIfPossible(e); } @@ -286,7 +319,7 @@ public Struct readRowUsingIndex(String table, String index, Key key, Iterable() { + new CachedResultSetSupplier() { @Override - public ResultSet get() { - return readContextDelegate.executeQuery(statement, options); + ResultSet load() { + return getReadContextDelegate().executeQuery(statement, options); } }); } @@ -314,12 +347,12 @@ public ResultSet get() { public AsyncResultSet executeQueryAsync( final Statement statement, final QueryOption... options) { return new AsyncResultSetImpl( - sessionPool.executorFactory, + sessionPool.sessionClient.getSpanner().getAsyncExecutorProvider(), wrap( - new Supplier() { + new CachedResultSetSupplier() { @Override - public ResultSet get() { - return readContextDelegate.executeQuery(statement, options); + ResultSet load() { + return getReadContextDelegate().executeQuery(statement, options); } })); } @@ -327,10 +360,10 @@ public ResultSet get() { @Override public ResultSet analyzeQuery(final Statement statement, final QueryAnalyzeMode queryMode) { return wrap( - new Supplier() { + new CachedResultSetSupplier() { @Override - public ResultSet get() { - return readContextDelegate.analyzeQuery(statement, queryMode); + ResultSet load() { + return getReadContextDelegate().analyzeQuery(statement, queryMode); } }); } @@ -341,7 +374,9 @@ public void close() { return; } closed = true; - readContextDelegate.close(); + if (readContextDelegate != null) { + readContextDelegate.close(); + } session.close(); } } @@ -602,7 +637,7 @@ private SessionPoolTransactionRunner(SessionPool sessionPool, PooledSessionFutur private TransactionRunner getRunner() { if (this.runner == null) { - this.runner = session.get().delegate.readWriteTransaction(); + this.runner = session.get().readWriteTransaction(); } return runner; } @@ -642,9 +677,66 @@ public TransactionRunner allowNestedTransaction() { } } + private static class SessionPoolAsyncRunner { + private final SessionPool sessionPool; + private volatile PooledSessionFuture session; + private final AsyncWork work; + private final Executor executor; + + private SessionPoolAsyncRunner( + SessionPool sessionPool, + PooledSessionFuture session, + AsyncWork work, + Executor executor) { + this.sessionPool = sessionPool; + this.session = session; + this.work = work; + this.executor = executor; + } + + private ApiFuture runAsync() { + final SettableApiFuture res = SettableApiFuture.create(); + executor.execute( + new Runnable() { + @Override + public void run() { + SpannerException se = null; + R r = null; + while (true) { + try { + r = session.get().runAsync(work, MoreExecutors.directExecutor()).get(); + break; + } catch (ExecutionException e) { + se = SpannerExceptionFactory.newSpannerException(e.getCause()); + } catch (InterruptedException e) { + se = SpannerExceptionFactory.propagateInterrupt(e); + } catch (Throwable t) { + se = SpannerExceptionFactory.newSpannerException(t); + } finally { + if (se != null && se instanceof SessionNotFoundException) { + session = + sessionPool.replaceReadWriteSession((SessionNotFoundException) se, session); + } else { + break; + } + } + } + session.get().markUsed(); + session.close(); + if (se != null) { + res.setException(se); + } else { + res.set(r); + } + } + }); + return res; + } + } + // Exception class used just to track the stack trace at the point when a session was handed out // from the pool. - private final class LeakedSessionException extends RuntimeException { + final class LeakedSessionException extends RuntimeException { private static final long serialVersionUID = 1451131180314064914L; private LeakedSessionException() { @@ -812,6 +904,11 @@ public TransactionManager transactionManager() { return new AutoClosingTransactionManager(SessionPool.this, this); } + @Override + public ApiFuture runAsync(AsyncWork work, Executor executor) { + return new SessionPoolAsyncRunner<>(SessionPool.this, this, work, executor).runAsync(); + } + @Override public long executePartitionedUpdate(Statement stmt) { try { @@ -962,6 +1059,11 @@ public TransactionRunner readWriteTransaction() { return delegate.readWriteTransaction(); } + @Override + public ApiFuture runAsync(AsyncWork work, Executor executor) { + return delegate.runAsync(work, executor); + } + @Override public ApiFuture asyncClose() { close(); @@ -1872,7 +1974,11 @@ public void run() { }); for (PooledSessionFuture session : checkedOutSessions) { if (session.leakedException != null) { - logger.log(Level.WARNING, "Leaked session", session.leakedException); + if (options.isFailOnSessionLeak()) { + throw session.leakedException; + } else { + logger.log(Level.WARNING, "Leaked session", session.leakedException); + } } } for (final PooledSession session : ImmutableList.copyOf(allSessions)) { diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPoolOptions.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPoolOptions.java index 45289fb3cd2..27257bc65e9 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPoolOptions.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPoolOptions.java @@ -32,6 +32,7 @@ public class SessionPoolOptions { private final ActionOnExhaustion actionOnExhaustion; private final int keepAliveIntervalMinutes; private final ActionOnSessionNotFound actionOnSessionNotFound; + private final ActionOnSessionLeak actionOnSessionLeak; private final long initialWaitForSessionTimeoutMillis; private SessionPoolOptions(Builder builder) { @@ -44,6 +45,7 @@ private SessionPoolOptions(Builder builder) { this.writeSessionsFraction = builder.writeSessionsFraction; this.actionOnExhaustion = builder.actionOnExhaustion; this.actionOnSessionNotFound = builder.actionOnSessionNotFound; + this.actionOnSessionLeak = builder.actionOnSessionLeak; this.initialWaitForSessionTimeoutMillis = builder.initialWaitForSessionTimeoutMillis; this.keepAliveIntervalMinutes = builder.keepAliveIntervalMinutes; } @@ -86,6 +88,11 @@ boolean isFailIfSessionNotFound() { return actionOnSessionNotFound == ActionOnSessionNotFound.FAIL; } + @VisibleForTesting + boolean isFailOnSessionLeak() { + return actionOnSessionLeak == ActionOnSessionLeak.FAIL; + } + public static Builder newBuilder() { return new Builder(); } @@ -100,6 +107,11 @@ private static enum ActionOnSessionNotFound { FAIL; } + private static enum ActionOnSessionLeak { + WARN, + FAIL; + } + /** Builder for creating SessionPoolOptions. */ public static class Builder { private boolean minSessionsSet = false; @@ -110,6 +122,7 @@ public static class Builder { private ActionOnExhaustion actionOnExhaustion = DEFAULT_ACTION; private long initialWaitForSessionTimeoutMillis = 30_000L; private ActionOnSessionNotFound actionOnSessionNotFound = ActionOnSessionNotFound.RETRY; + private ActionOnSessionLeak actionOnSessionLeak = ActionOnSessionLeak.WARN; private int keepAliveIntervalMinutes = 30; /** @@ -197,6 +210,12 @@ Builder setFailIfSessionNotFound() { return this; } + @VisibleForTesting + Builder setFailOnSessionLeak() { + this.actionOnSessionLeak = ActionOnSessionLeak.FAIL; + return this; + } + /** * Fraction of sessions to be kept prepared for write transactions. This is an optimisation to * avoid the cost of sending a BeginTransaction() rpc. If all such sessions are in use and a diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java index 344c46f82bc..ef8e08ef2a7 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java @@ -16,13 +16,13 @@ package com.google.cloud.spanner; +import com.google.api.gax.core.ExecutorProvider; import com.google.api.gax.core.GaxProperties; import com.google.api.gax.paging.Page; import com.google.cloud.BaseService; import com.google.cloud.PageImpl; import com.google.cloud.PageImpl.NextPageFetcher; import com.google.cloud.grpc.GrpcTransportOptions; -import com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory; import com.google.cloud.spanner.SessionClient.SessionId; import com.google.cloud.spanner.spi.v1.SpannerRpc; import com.google.cloud.spanner.spi.v1.SpannerRpc.Paginated; @@ -127,8 +127,11 @@ QueryOptions getDefaultQueryOptions(DatabaseId databaseId) { return getOptions().getDefaultQueryOptions(databaseId); } - ExecutorFactory getExecutorFactory() { - return ((GrpcTransportOptions) getOptions().getTransportOptions()).getExecutorFactory(); + /** + * Returns the {@link ExecutorProvider} to use for async methods that need a background executor. + */ + ExecutorProvider getAsyncExecutorProvider() { + return getOptions().getAsyncExecutorProvider(); } SessionImpl sessionWithId(String name) { diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java index 32dc3b71572..efe9b381ade 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java @@ -17,6 +17,8 @@ package com.google.cloud.spanner; import com.google.api.core.ApiFunction; +import com.google.api.gax.core.ExecutorProvider; +import com.google.api.gax.core.FixedExecutorProvider; import com.google.api.gax.grpc.GrpcInterceptorProvider; import com.google.api.gax.longrunning.OperationSnapshot; import com.google.api.gax.longrunning.OperationTimedPollAlgorithm; @@ -39,6 +41,7 @@ import com.google.cloud.spanner.spi.v1.SpannerRpc; import com.google.cloud.spanner.v1.SpannerSettings; import com.google.cloud.spanner.v1.stub.SpannerStubSettings; +import com.google.common.annotations.VisibleForTesting; import com.google.common.base.MoreObjects; import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableMap; @@ -47,6 +50,7 @@ import com.google.spanner.admin.database.v1.CreateDatabaseRequest; import com.google.spanner.admin.database.v1.RestoreDatabaseRequest; import com.google.spanner.v1.ExecuteSqlRequest.QueryOptions; +import com.google.common.util.concurrent.ThreadFactoryBuilder; import io.grpc.CallCredentials; import io.grpc.ManagedChannelBuilder; import java.io.IOException; @@ -57,6 +61,10 @@ import java.util.Map.Entry; import java.util.Set; import javax.annotation.Nonnull; +import java.util.concurrent.ScheduledThreadPoolExecutor; +import java.util.concurrent.ThreadFactory; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; import org.threeten.bp.Duration; /** Options for the Cloud Spanner service. */ @@ -103,6 +111,7 @@ public class SpannerOptions extends ServiceOptions { private final Map mergedQueryOptions; private final CallCredentialsProvider callCredentialsProvider; + private final ExecutorProvider asyncExecutorProvider; /** * Interface that can be used to provide {@link CallCredentials} instead of {@link Credentials} to @@ -133,6 +142,31 @@ public ServiceRpc create(SpannerOptions options) { } } + private static final AtomicInteger DEFAULT_POOL_COUNT = new AtomicInteger(); + + /** + * Default {@link ExecutorProvider} for high-level async calls that need an executor. The default + * uses a cached thread pool containing a max of 8 threads. The pool is lazily initialized and + * will not create any threads if the user application does not use any async methods. It will + * also scale down the thread usage if the async load allows for that. + */ + @VisibleForTesting + static ExecutorProvider createDefaultAsyncExecutorProvider() { + return createAsyncExecutorProvider(8, 60L, TimeUnit.SECONDS); + } + + @VisibleForTesting + static ExecutorProvider createAsyncExecutorProvider( + int poolSize, long keepAliveTime, TimeUnit unit) { + String format = String.format("async-pool-%d-thread-%%d", DEFAULT_POOL_COUNT.incrementAndGet()); + ThreadFactory threadFactory = + new ThreadFactoryBuilder().setDaemon(true).setNameFormat(format).build(); + ScheduledThreadPoolExecutor executor = new ScheduledThreadPoolExecutor(poolSize, threadFactory); + executor.setKeepAliveTime(keepAliveTime, unit); + executor.allowCoreThreadTimeOut(true); + return FixedExecutorProvider.create(executor); + } + private SpannerOptions(Builder builder) { super(SpannerFactory.class, SpannerRpcFactory.class, builder, new SpannerDefaults()); numChannels = builder.numChannels; @@ -173,6 +207,9 @@ private SpannerOptions(Builder builder) { this.mergedQueryOptions = ImmutableMap.copyOf(merged); } callCredentialsProvider = builder.callCredentialsProvider; + asyncExecutorProvider = + MoreObjects.firstNonNull( + builder.asyncExecutorProvider, createDefaultAsyncExecutorProvider()); } /** @@ -237,6 +274,7 @@ public static class Builder private boolean autoThrottleAdministrativeRequests = false; private Map defaultQueryOptions = new HashMap<>(); private CallCredentialsProvider callCredentialsProvider; + private ExecutorProvider asyncExecutorProvider; private String emulatorHost = System.getenv("SPANNER_EMULATOR_HOST"); private Builder() { @@ -307,6 +345,7 @@ private Builder() { this.autoThrottleAdministrativeRequests = options.autoThrottleAdministrativeRequests; this.defaultQueryOptions = options.defaultQueryOptions; this.callCredentialsProvider = options.callCredentialsProvider; + this.asyncExecutorProvider = options.asyncExecutorProvider; this.channelProvider = options.channelProvider; this.channelConfigurator = options.channelConfigurator; this.interceptorProvider = options.interceptorProvider; @@ -692,6 +731,10 @@ public QueryOptions getDefaultQueryOptions(DatabaseId databaseId) { return options; } + public ExecutorProvider getAsyncExecutorProvider() { + return asyncExecutorProvider; + } + public int getPrefetchChunks() { return prefetchChunks; } 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 fc72793e860..ecd04f26e07 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 @@ -21,8 +21,8 @@ import static com.google.common.base.Preconditions.checkNotNull; import static com.google.common.base.Preconditions.checkState; +import com.google.api.gax.core.ExecutorProvider; import com.google.cloud.Timestamp; -import com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory; import com.google.cloud.spanner.SessionImpl.SessionTransaction; import com.google.cloud.spanner.spi.v1.SpannerRpc; import com.google.common.annotations.VisibleForTesting; @@ -44,7 +44,6 @@ import java.util.ArrayList; import java.util.List; import java.util.concurrent.Callable; -import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.atomic.AtomicInteger; import java.util.logging.Level; import java.util.logging.Logger; diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncResultSetImplStressTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncResultSetImplStressTest.java index ea8396b7ed0..0bf7a1a2c96 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncResultSetImplStressTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncResultSetImplStressTest.java @@ -21,8 +21,7 @@ import com.google.api.core.ApiFuture; import com.google.api.core.ApiFutures; import com.google.api.core.SettableApiFuture; -import com.google.cloud.grpc.GrpcTransportOptions; -import com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory; +import com.google.api.gax.core.ExecutorProvider; import com.google.cloud.spanner.AsyncResultSet.CallbackResponse; import com.google.cloud.spanner.AsyncResultSet.CursorState; import com.google.cloud.spanner.AsyncResultSet.ReadyCallback; @@ -167,12 +166,11 @@ private static ScheduledExecutorService createExecService(int threadCount) { @Test public void toList() throws Exception { - ExecutorFactory executorFactory = - new GrpcTransportOptions.DefaultExecutorFactory(); + ExecutorProvider executorProvider = SpannerOptions.createDefaultAsyncExecutorProvider(); for (int bufferSize = 1; bufferSize < resultSetSize * 2; bufferSize *= 2) { for (int i = 0; i < TEST_RUNS; i++) { try (AsyncResultSetImpl impl = - new AsyncResultSetImpl(executorFactory, createResultSet(), bufferSize)) { + new AsyncResultSetImpl(executorProvider, createResultSet(), bufferSize)) { ImmutableList list = impl.toList( new Function() { @@ -189,13 +187,12 @@ public Row apply(StructReader input) { @Test public void toListWithErrors() throws Exception { - ExecutorFactory executorFactory = - new GrpcTransportOptions.DefaultExecutorFactory(); + ExecutorProvider executorProvider = SpannerOptions.createDefaultAsyncExecutorProvider(); for (int bufferSize = 1; bufferSize < resultSetSize * 2; bufferSize *= 2) { for (int i = 0; i < TEST_RUNS; i++) { try (AsyncResultSetImpl impl = new AsyncResultSetImpl( - executorFactory, createResultSetWithErrors(1.0 / resultSetSize), bufferSize)) { + executorProvider, createResultSetWithErrors(1.0 / resultSetSize), bufferSize)) { ImmutableList list = impl.toList( new Function() { @@ -215,14 +212,13 @@ public Row apply(StructReader input) { @Test public void asyncToList() throws Exception { - ExecutorFactory executorFactory = - new GrpcTransportOptions.DefaultExecutorFactory(); + ExecutorProvider executorProvider = SpannerOptions.createDefaultAsyncExecutorProvider(); for (int bufferSize = 1; bufferSize < resultSetSize * 2; bufferSize *= 2) { List>> futures = new ArrayList<>(TEST_RUNS); ExecutorService executor = createExecService(32); for (int i = 0; i < TEST_RUNS; i++) { try (AsyncResultSet impl = - new AsyncResultSetImpl(executorFactory, createResultSet(), bufferSize)) { + new AsyncResultSetImpl(executorProvider, createResultSet(), bufferSize)) { futures.add( impl.toListAsync( new Function() { @@ -244,8 +240,7 @@ public Row apply(StructReader input) { @Test public void consume() throws Exception { - ExecutorFactory executorFactory = - new GrpcTransportOptions.DefaultExecutorFactory(); + ExecutorProvider executorProvider = SpannerOptions.createDefaultAsyncExecutorProvider(); final Random random = new Random(); for (Executor executor : new Executor[] { @@ -255,7 +250,7 @@ public void consume() throws Exception { for (int i = 0; i < TEST_RUNS; i++) { final SettableApiFuture> future = SettableApiFuture.create(); try (AsyncResultSetImpl impl = - new AsyncResultSetImpl(executorFactory, createResultSet(), bufferSize)) { + new AsyncResultSetImpl(executorProvider, createResultSet(), bufferSize)) { final ImmutableList.Builder builder = ImmutableList.builder(); impl.setCallback( executor, @@ -286,8 +281,7 @@ public CallbackResponse cursorReady(AsyncResultSet resultSet) { @Test public void pauseResume() throws Exception { - ExecutorFactory executorFactory = - new GrpcTransportOptions.DefaultExecutorFactory(); + ExecutorProvider executorProvider = SpannerOptions.createDefaultAsyncExecutorProvider(); final Random random = new Random(); List>> futures = new ArrayList<>(); for (Executor executor : @@ -301,7 +295,7 @@ public void pauseResume() throws Exception { final SettableApiFuture> future = SettableApiFuture.create(); futures.add(future); try (AsyncResultSetImpl impl = - new AsyncResultSetImpl(executorFactory, createResultSet(), bufferSize)) { + new AsyncResultSetImpl(executorProvider, createResultSet(), bufferSize)) { resultSets.add(impl); final ImmutableList.Builder builder = ImmutableList.builder(); impl.setCallback( @@ -352,8 +346,7 @@ public void run() { @Test public void cancel() throws Exception { - ExecutorFactory executorFactory = - new GrpcTransportOptions.DefaultExecutorFactory(); + ExecutorProvider executorProvider = SpannerOptions.createDefaultAsyncExecutorProvider(); final Random random = new Random(); for (Executor executor : new Executor[] { @@ -368,7 +361,7 @@ public void cancel() throws Exception { final SettableApiFuture> future = SettableApiFuture.create(); futures.add(future); try (AsyncResultSetImpl impl = - new AsyncResultSetImpl(executorFactory, createResultSet(), bufferSize)) { + new AsyncResultSetImpl(executorProvider, createResultSet(), bufferSize)) { resultSets.add(impl); final ImmutableList.Builder builder = ImmutableList.builder(); impl.setCallback( diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncResultSetImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncResultSetImplTest.java index f9df08563da..cd5588187e6 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncResultSetImplTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncResultSetImplTest.java @@ -22,8 +22,7 @@ import static org.mockito.Mockito.when; import com.google.api.core.ApiFuture; -import com.google.cloud.grpc.GrpcTransportOptions; -import com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory; +import com.google.api.gax.core.ExecutorProvider; import com.google.cloud.spanner.AsyncResultSet.CallbackResponse; import com.google.cloud.spanner.AsyncResultSet.CursorState; import com.google.cloud.spanner.AsyncResultSet.ReadyCallback; @@ -48,32 +47,20 @@ @RunWith(JUnit4.class) public class AsyncResultSetImplTest { - private ExecutorFactory mockedFactory; - private ExecutorFactory simpleFactory; + private ExecutorProvider mockedProvider; + private ExecutorProvider simpleProvider; - @SuppressWarnings("unchecked") @Before public void setup() { - mockedFactory = mock(ExecutorFactory.class); - when(mockedFactory.get()).thenReturn(mock(ScheduledExecutorService.class)); - simpleFactory = - new GrpcTransportOptions.ExecutorFactory() { - @Override - public ScheduledExecutorService get() { - return Executors.newScheduledThreadPool(1); - } - - @Override - public void release(ScheduledExecutorService executor) { - executor.shutdown(); - } - }; + mockedProvider = mock(ExecutorProvider.class); + when(mockedProvider.getExecutor()).thenReturn(mock(ScheduledExecutorService.class)); + simpleProvider = SpannerOptions.createAsyncExecutorProvider(1, 1L, TimeUnit.SECONDS); } @SuppressWarnings("unchecked") @Test public void close() { - AsyncResultSetImpl rs = new AsyncResultSetImpl(mockedFactory, mock(ResultSet.class)); + AsyncResultSetImpl rs = new AsyncResultSetImpl(mockedProvider, mock(ResultSet.class)); rs.close(); // Closing a second time should be a no-op. rs.close(); @@ -96,7 +83,7 @@ public void close() { } // The following methods are allowed on a closed result set. - AsyncResultSetImpl rs2 = new AsyncResultSetImpl(mockedFactory, mock(ResultSet.class)); + AsyncResultSetImpl rs2 = new AsyncResultSetImpl(mockedProvider, mock(ResultSet.class)); rs2.setCallback(mock(Executor.class), mock(ReadyCallback.class)); rs2.close(); rs2.cancel(); @@ -105,7 +92,7 @@ public void close() { @Test public void tryNextNotAllowed() { - try (AsyncResultSetImpl rs = new AsyncResultSetImpl(mockedFactory, mock(ResultSet.class))) { + try (AsyncResultSetImpl rs = new AsyncResultSetImpl(mockedProvider, mock(ResultSet.class))) { rs.setCallback(mock(Executor.class), mock(ReadyCallback.class)); try { rs.tryNext(); @@ -122,7 +109,7 @@ public void toList() { ResultSet delegate = mock(ResultSet.class); when(delegate.next()).thenReturn(true, true, true, false); when(delegate.getCurrentRowAsStruct()).thenReturn(mock(Struct.class)); - try (AsyncResultSetImpl rs = new AsyncResultSetImpl(simpleFactory, delegate)) { + try (AsyncResultSetImpl rs = new AsyncResultSetImpl(simpleProvider, delegate)) { ImmutableList list = rs.toList( new Function() { @@ -142,7 +129,7 @@ public void toListPropagatesError() { .thenThrow( SpannerExceptionFactory.newSpannerException( ErrorCode.INVALID_ARGUMENT, "invalid query")); - try (AsyncResultSetImpl rs = new AsyncResultSetImpl(simpleFactory, delegate)) { + try (AsyncResultSetImpl rs = new AsyncResultSetImpl(simpleProvider, delegate)) { rs.toList( new Function() { @Override @@ -163,7 +150,7 @@ public void toListAsync() throws InterruptedException, ExecutionException { ResultSet delegate = mock(ResultSet.class); when(delegate.next()).thenReturn(true, true, true, false); when(delegate.getCurrentRowAsStruct()).thenReturn(mock(Struct.class)); - try (AsyncResultSetImpl rs = new AsyncResultSetImpl(simpleFactory, delegate)) { + try (AsyncResultSetImpl rs = new AsyncResultSetImpl(simpleProvider, delegate)) { ApiFuture> future = rs.toListAsync( new Function() { @@ -186,7 +173,7 @@ public void toListAsyncPropagatesError() throws InterruptedException { .thenThrow( SpannerExceptionFactory.newSpannerException( ErrorCode.INVALID_ARGUMENT, "invalid query")); - try (AsyncResultSetImpl rs = new AsyncResultSetImpl(simpleFactory, delegate)) { + try (AsyncResultSetImpl rs = new AsyncResultSetImpl(simpleProvider, delegate)) { rs.toListAsync( new Function() { @Override @@ -215,7 +202,7 @@ public void withCallback() throws InterruptedException { final AtomicInteger callbackCounter = new AtomicInteger(); final AtomicInteger rowCounter = new AtomicInteger(); final CountDownLatch finishedLatch = new CountDownLatch(1); - try (AsyncResultSetImpl rs = new AsyncResultSetImpl(simpleFactory, delegate)) { + try (AsyncResultSetImpl rs = new AsyncResultSetImpl(simpleProvider, delegate)) { rs.setCallback( executor, new ReadyCallback() { @@ -249,7 +236,7 @@ public void callbackReceivesError() throws InterruptedException { SpannerExceptionFactory.newSpannerException( ErrorCode.INVALID_ARGUMENT, "invalid query")); final BlockingDeque receivedErr = new LinkedBlockingDeque<>(1); - try (AsyncResultSetImpl rs = new AsyncResultSetImpl(simpleFactory, delegate)) { + try (AsyncResultSetImpl rs = new AsyncResultSetImpl(simpleProvider, delegate)) { rs.setCallback( executor, new ReadyCallback() { @@ -284,7 +271,7 @@ public void callbackReceivesErrorHalfwayThrough() throws InterruptedException { when(delegate.getCurrentRowAsStruct()).thenReturn(mock(Struct.class)); final AtomicInteger rowCount = new AtomicInteger(); final BlockingDeque receivedErr = new LinkedBlockingDeque<>(1); - try (AsyncResultSetImpl rs = new AsyncResultSetImpl(simpleFactory, delegate)) { + try (AsyncResultSetImpl rs = new AsyncResultSetImpl(simpleProvider, delegate)) { rs.setCallback( executor, new ReadyCallback() { @@ -319,7 +306,7 @@ public void pauseResume() throws InterruptedException { final AtomicInteger callbackCounter = new AtomicInteger(); final BlockingDeque queue = new LinkedBlockingDeque<>(1); final AtomicBoolean finished = new AtomicBoolean(false); - try (AsyncResultSetImpl rs = new AsyncResultSetImpl(simpleFactory, delegate)) { + try (AsyncResultSetImpl rs = new AsyncResultSetImpl(simpleProvider, delegate)) { rs.setCallback( executor, new ReadyCallback() { @@ -363,7 +350,7 @@ public void cancel() throws InterruptedException { final AtomicInteger callbackCounter = new AtomicInteger(); final BlockingDeque queue = new LinkedBlockingDeque<>(1); final AtomicBoolean finished = new AtomicBoolean(false); - try (AsyncResultSetImpl rs = new AsyncResultSetImpl(simpleFactory, delegate)) { + try (AsyncResultSetImpl rs = new AsyncResultSetImpl(simpleProvider, delegate)) { rs.setCallback( executor, new ReadyCallback() { @@ -417,7 +404,7 @@ public void callbackReturnsError() throws InterruptedException { when(delegate.next()).thenReturn(true, true, true, false); when(delegate.getCurrentRowAsStruct()).thenReturn(mock(Struct.class)); final AtomicInteger callbackCounter = new AtomicInteger(); - try (AsyncResultSetImpl rs = new AsyncResultSetImpl(simpleFactory, delegate)) { + try (AsyncResultSetImpl rs = new AsyncResultSetImpl(simpleProvider, delegate)) { rs.setCallback( executor, new ReadyCallback() { diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java index 066cf5123ad..aa514499704 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java @@ -22,11 +22,14 @@ import static org.hamcrest.MatcherAssert.assertThat; import static org.junit.Assert.fail; +import com.google.api.core.ApiFuture; +import com.google.api.core.SettableApiFuture; import com.google.api.gax.grpc.testing.LocalChannelProvider; import com.google.api.gax.retrying.RetrySettings; import com.google.cloud.NoCredentials; import com.google.cloud.spanner.AsyncResultSet.CallbackResponse; import com.google.cloud.spanner.AsyncResultSet.ReadyCallback; +import com.google.cloud.spanner.DatabaseClient.AsyncWork; import com.google.cloud.spanner.MockSpannerServiceImpl.SimulatedExecutionTime; import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult; import com.google.cloud.spanner.ReadContext.QueryAnalyzeMode; @@ -50,6 +53,7 @@ import java.util.ArrayList; import java.util.Arrays; import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.ScheduledThreadPoolExecutor; @@ -58,7 +62,9 @@ import org.junit.AfterClass; import org.junit.Before; import org.junit.BeforeClass; +import org.junit.Rule; import org.junit.Test; +import org.junit.rules.Timeout; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; import org.threeten.bp.Duration; @@ -105,6 +111,9 @@ public class DatabaseClientImplTest { .setMetadata(SELECT1_METADATA) .build(); private Spanner spanner; + private Spanner spannerWithEmptySessionPool; + + @Rule public Timeout globalTimeout = new Timeout(5L, TimeUnit.SECONDS); @BeforeClass public static void startStaticServer() throws IOException { @@ -141,13 +150,24 @@ public void setUp() throws IOException { .setProjectId(TEST_PROJECT) .setChannelProvider(channelProvider) .setCredentials(NoCredentials.getInstance()) + .setSessionPoolOption(SessionPoolOptions.newBuilder().setFailOnSessionLeak().build()) + .build() + .getService(); + spannerWithEmptySessionPool = + spanner + .getOptions() + .toBuilder() + .setSessionPoolOption( + SessionPoolOptions.newBuilder().setMinSessions(0).setFailOnSessionLeak().build()) .build() .getService(); } @After public void tearDown() throws Exception { + mockSpanner.unfreeze(); spanner.close(); + spannerWithEmptySessionPool.close(); mockSpanner.reset(); mockSpanner.removeAllExecutionTimes(); } @@ -181,6 +201,23 @@ public void singleUse() { } } + @Test + public void singleUseIsNonBlocking() { + mockSpanner.freeze(); + // Use a Spanner instance with no initial sessions in the pool to show that getting a session + // from the pool and then preparing a query is non-blocking (i.e. does not wait on a reply from + // the server). + DatabaseClient client = + spannerWithEmptySessionPool.getDatabaseClient( + DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); + try (ResultSet rs = client.singleUse().executeQuery(SELECT1)) { + mockSpanner.unfreeze(); + assertThat(rs.next()).isTrue(); + assertThat(rs.getLong(0)).isEqualTo(1L); + assertThat(rs.next()).isFalse(); + } + } + @Test public void singleUseBound() { DatabaseClient client = @@ -195,6 +232,23 @@ public void singleUseBound() { } } + @Test + public void singleUseBoundIsNonBlocking() { + mockSpanner.freeze(); + DatabaseClient client = + spannerWithEmptySessionPool.getDatabaseClient( + DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); + try (ResultSet rs = + client + .singleUse(TimestampBound.ofExactStaleness(15L, TimeUnit.SECONDS)) + .executeQuery(SELECT1)) { + mockSpanner.unfreeze(); + assertThat(rs.next()).isTrue(); + assertThat(rs.getLong(0)).isEqualTo(1L); + assertThat(rs.next()).isFalse(); + } + } + @Test public void singleUseTransaction() { DatabaseClient client = @@ -206,6 +260,20 @@ public void singleUseTransaction() { } } + @Test + public void singleUseTransactionIsNonBlocking() { + mockSpanner.freeze(); + DatabaseClient client = + spannerWithEmptySessionPool.getDatabaseClient( + DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); + try (ResultSet rs = client.singleUseReadOnlyTransaction().executeQuery(SELECT1)) { + mockSpanner.unfreeze(); + assertThat(rs.next()).isTrue(); + assertThat(rs.getLong(0)).isEqualTo(1L); + assertThat(rs.next()).isFalse(); + } + } + @Test public void singleUseTransactionBound() { DatabaseClient client = @@ -220,12 +288,28 @@ public void singleUseTransactionBound() { } } + @Test + public void singleUseTransactionBoundIsNonBlocking() { + mockSpanner.freeze(); + DatabaseClient client = + spannerWithEmptySessionPool.getDatabaseClient( + DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); + try (ResultSet rs = + client + .singleUseReadOnlyTransaction(TimestampBound.ofExactStaleness(15L, TimeUnit.SECONDS)) + .executeQuery(SELECT1)) { + mockSpanner.unfreeze(); + assertThat(rs.next()).isTrue(); + assertThat(rs.getLong(0)).isEqualTo(1L); + assertThat(rs.next()).isFalse(); + } + } + @Test public void readOnlyTransaction() { DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); - try (ReadOnlyTransaction tx = - client.readOnlyTransaction(TimestampBound.ofExactStaleness(15L, TimeUnit.SECONDS))) { + try (ReadOnlyTransaction tx = client.readOnlyTransaction()) { try (ResultSet rs = tx.executeQuery(SELECT1)) { assertThat(rs.next()).isTrue(); assertThat(rs.getLong(0)).isEqualTo(1L); @@ -234,6 +318,22 @@ public void readOnlyTransaction() { } } + @Test + public void readOnlyTransactionIsNonBlocking() { + mockSpanner.freeze(); + DatabaseClient client = + spannerWithEmptySessionPool.getDatabaseClient( + DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); + try (ReadOnlyTransaction tx = client.readOnlyTransaction()) { + try (ResultSet rs = tx.executeQuery(SELECT1)) { + mockSpanner.unfreeze(); + assertThat(rs.next()).isTrue(); + assertThat(rs.getLong(0)).isEqualTo(1L); + assertThat(rs.next()).isFalse(); + } + } + } + @Test public void readOnlyTransactionBound() { DatabaseClient client = @@ -248,6 +348,23 @@ public void readOnlyTransactionBound() { } } + @Test + public void readOnlyTransactionBoundIsNonBlocking() { + mockSpanner.freeze(); + DatabaseClient client = + spannerWithEmptySessionPool.getDatabaseClient( + DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); + try (ReadOnlyTransaction tx = + client.readOnlyTransaction(TimestampBound.ofExactStaleness(15L, TimeUnit.SECONDS))) { + try (ResultSet rs = tx.executeQuery(SELECT1)) { + mockSpanner.unfreeze(); + assertThat(rs.next()).isTrue(); + assertThat(rs.getLong(0)).isEqualTo(1L); + assertThat(rs.next()).isFalse(); + } + } + } + @Test public void readWriteTransaction() { DatabaseClient client = @@ -263,6 +380,96 @@ public Void run(TransactionContext transaction) throws Exception { }); } + @Test + public void readWriteTransactionIsNonBlocking() { + mockSpanner.freeze(); + DatabaseClient client = + spannerWithEmptySessionPool.getDatabaseClient( + DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); + TransactionRunner runner = client.readWriteTransaction(); + // The runner.run(...) method cannot be made non-blocking, as it returns the result of the + // transaction. + mockSpanner.unfreeze(); + runner.run( + new TransactionCallable() { + @Override + public Void run(TransactionContext transaction) throws Exception { + transaction.executeUpdate(UPDATE_STATEMENT); + return null; + } + }); + } + + @Test + public void runAsync() throws Exception { + DatabaseClient client = + spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); + ExecutorService executor = Executors.newSingleThreadExecutor(); + ApiFuture fut = + client.runAsync( + new AsyncWork() { + @Override + public ApiFuture doWorkAsync(TransactionContext txn) { + SettableApiFuture res = SettableApiFuture.create(); + res.set(txn.executeUpdate(UPDATE_STATEMENT)); + return res; + } + }, + executor); + assertThat(fut.get()).isEqualTo(UPDATE_COUNT); + executor.shutdown(); + } + + @Test + public void runAsyncIsNonBlocking() throws Exception { + mockSpanner.freeze(); + DatabaseClient client = + spannerWithEmptySessionPool.getDatabaseClient( + DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); + ExecutorService executor = Executors.newSingleThreadExecutor(); + ApiFuture fut = + client.runAsync( + new AsyncWork() { + @Override + public ApiFuture doWorkAsync(TransactionContext txn) { + SettableApiFuture res = SettableApiFuture.create(); + res.set(txn.executeUpdate(UPDATE_STATEMENT)); + return res; + } + }, + executor); + mockSpanner.unfreeze(); + assertThat(fut.get()).isEqualTo(UPDATE_COUNT); + executor.shutdown(); + } + + @Test + public void runAsyncWithException() throws Exception { + DatabaseClient client = + spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); + ExecutorService executor = Executors.newSingleThreadExecutor(); + ApiFuture fut = + client.runAsync( + new AsyncWork() { + @Override + public ApiFuture doWorkAsync(TransactionContext txn) { + SettableApiFuture res = SettableApiFuture.create(); + res.set(txn.executeUpdate(INVALID_UPDATE_STATEMENT)); + return res; + } + }, + executor); + try { + fut.get(); + fail("missing expected exception"); + } catch (ExecutionException e) { + assertThat(e.getCause()).isInstanceOf(SpannerException.class); + SpannerException se = (SpannerException) e.getCause(); + assertThat(se.getErrorCode()).isEqualTo(ErrorCode.INVALID_ARGUMENT); + } + executor.shutdown(); + } + @Test public void transactionManager() throws Exception { DatabaseClient client = @@ -282,6 +489,28 @@ public void transactionManager() throws Exception { } } + @Test + public void transactionManagerIsNonBlocking() throws Exception { + mockSpanner.freeze(); + DatabaseClient client = + spannerWithEmptySessionPool.getDatabaseClient( + DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); + try (TransactionManager txManager = client.transactionManager()) { + while (true) { + mockSpanner.unfreeze(); + TransactionContext tx = txManager.begin(); + try { + tx.executeUpdate(UPDATE_STATEMENT); + txManager.commit(); + break; + } catch (AbortedException e) { + Thread.sleep(e.getRetryDelayInMillis() / 1000); + tx = txManager.resetForRetry(); + } + } + } + } + /** * Test that the update statement can be executed as a partitioned transaction that returns a * lower bound update count. @@ -537,6 +766,7 @@ public void testDatabaseOrInstanceDoesNotExistOnCreate() throws Exception { DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); // The create session failure should propagate to the client and not retry. try (ResultSet rs = dbClient.singleUse().executeQuery(SELECT1)) { + rs.next(); fail("missing expected exception"); } catch (DatabaseNotFoundException | InstanceNotFoundException e) { // The server should only receive one BatchCreateSessions request. 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 118b2c57fe0..54ae39bbc71 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 @@ -90,11 +90,10 @@ import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentLinkedQueue; import java.util.concurrent.ConcurrentMap; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicLong; -import java.util.concurrent.locks.ReadWriteLock; -import java.util.concurrent.locks.ReentrantReadWriteLock; import org.threeten.bp.Instant; /** @@ -420,19 +419,15 @@ private SimulatedExecutionTime( private void simulateExecutionTime( Queue globalExceptions, boolean stickyGlobalExceptions, - ReadWriteLock freezeLock) { - try { - freezeLock.readLock().lock(); - checkException(globalExceptions, stickyGlobalExceptions); - checkException(this.exceptions, stickyException); - if (minimumExecutionTime > 0 || randomExecutionTime > 0) { - Uninterruptibles.sleepUninterruptibly( - (randomExecutionTime == 0 ? 0 : RANDOM.nextInt(randomExecutionTime)) - + minimumExecutionTime, - TimeUnit.MILLISECONDS); - } - } finally { - freezeLock.readLock().unlock(); + CountDownLatch freezeLock) { + Uninterruptibles.awaitUninterruptibly(freezeLock); + checkException(globalExceptions, stickyGlobalExceptions); + checkException(this.exceptions, stickyException); + if (minimumExecutionTime > 0 || randomExecutionTime > 0) { + Uninterruptibles.sleepUninterruptibly( + (randomExecutionTime == 0 ? 0 : RANDOM.nextInt(randomExecutionTime)) + + minimumExecutionTime, + TimeUnit.MILLISECONDS); } } @@ -451,7 +446,7 @@ private static void checkException(Queue exceptions, boolean keepExce private double abortProbability = 0.0010D; private final Queue requests = new ConcurrentLinkedQueue<>(); - private final ReadWriteLock freezeLock = new ReentrantReadWriteLock(); + private volatile CountDownLatch freezeLock = new CountDownLatch(0); private final Queue exceptions = new ConcurrentLinkedQueue<>(); private boolean stickyGlobalExceptions = false; private final ConcurrentMap statementResults = @@ -591,11 +586,11 @@ public void abortAllTransactions() { } public void freeze() { - freezeLock.writeLock().lock(); + freezeLock = new CountDownLatch(1); } public void unfreeze() { - freezeLock.writeLock().unlock(); + freezeLock.countDown(); } public void setMaxSessionsInOneBatch(int max) { @@ -1638,6 +1633,10 @@ public void addException(Exception exception) { exceptions.add(exception); } + public void clearExceptions() { + exceptions.clear(); + } + public void setStickyGlobalExceptions(boolean sticky) { this.stickyGlobalExceptions = sticky; } @@ -1661,6 +1660,7 @@ public void reset() { transactionLastUsed.clear(); exceptions.clear(); stickyGlobalExceptions = false; + freezeLock.countDown(); } public void removeAllExecutionTimes() { diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java index c1acecdbc2f..8735a2eece5 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java @@ -35,8 +35,8 @@ import com.google.api.core.ApiFuture; import com.google.api.core.ApiFutures; +import com.google.api.gax.core.ExecutorProvider; import com.google.cloud.Timestamp; -import com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory; import com.google.cloud.spanner.MetricRegistryTestUtils.FakeMetricRegistry; import com.google.cloud.spanner.MetricRegistryTestUtils.MetricsRecord; import com.google.cloud.spanner.ReadContext.QueryAnalyzeMode; diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionContextImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionContextImplTest.java index d319d13de2b..061187696a0 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionContextImplTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionContextImplTest.java @@ -19,7 +19,7 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; -import com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory; +import com.google.api.gax.core.ExecutorProvider; import com.google.cloud.spanner.TransactionRunnerImpl.TransactionContextImpl; import com.google.cloud.spanner.spi.v1.SpannerRpc; import com.google.protobuf.ByteString; diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java index 3dcd523c13a..1f2df00e057 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java @@ -25,6 +25,7 @@ import static org.mockito.Mockito.when; import com.google.api.core.ApiFutures; +import com.google.api.gax.core.ExecutorProvider; import com.google.cloud.grpc.GrpcTransportOptions; import com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory; import com.google.cloud.spanner.SessionClient.SessionId; From e485709e46a8dbc231f851cfead5aaa3c48056ef Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Mon, 24 Feb 2020 15:27:56 +0100 Subject: [PATCH 03/49] tests: fix integration tests that assumed tx was blocking Some integration tests started transactions without executing a query, and expected these transactions to fail. However, as the client is now non-blocking up until the first call to ResultSet#next(), no exception would occur. --- .../google/cloud/spanner/it/ITDatabaseTest.java | 1 + .../google/cloud/spanner/it/ITReadOnlyTxnTest.java | 14 ++++++++++++-- .../google/cloud/spanner/it/ITTransactionTest.java | 7 ++++++- 3 files changed, 19 insertions(+), 3 deletions(-) diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITDatabaseTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITDatabaseTest.java index 3a7125312bd..b83be137a13 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITDatabaseTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITDatabaseTest.java @@ -133,6 +133,7 @@ public void instanceNotFound() { .getClient() .getDatabaseClient(DatabaseId.of(nonExistingInstanceId, "some-db")); try (ResultSet rs = client.singleUse().executeQuery(Statement.of("SELECT 1"))) { + rs.next(); fail("missing expected exception"); } catch (InstanceNotFoundException e) { assertThat(e.getResourceName()).isEqualTo(nonExistingInstanceId.getName()); diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITReadOnlyTxnTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITReadOnlyTxnTest.java index e6e473779d4..f38809615b6 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITReadOnlyTxnTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITReadOnlyTxnTest.java @@ -312,7 +312,12 @@ public void multiReadTimestamp() { public void multiMinReadTimestamp() { // Cannot use bounded modes with multi-read transactions. expectedException.expect(IllegalArgumentException.class); - client.readOnlyTransaction(TimestampBound.ofMinReadTimestamp(history.get(2).timestamp)); + try (ReadOnlyTransaction tx = + client.readOnlyTransaction(TimestampBound.ofMinReadTimestamp(history.get(2).timestamp))) { + try (ResultSet rs = tx.executeQuery(Statement.of("SELECT 1"))) { + rs.next(); + } + } } @Test @@ -339,6 +344,11 @@ public void multiExactStaleness() { public void multiMaxStaleness() { // Cannot use bounded modes with multi-read transactions. expectedException.expect(IllegalArgumentException.class); - client.readOnlyTransaction(TimestampBound.ofMaxStaleness(1, TimeUnit.SECONDS)); + try (ReadOnlyTransaction tx = + client.readOnlyTransaction(TimestampBound.ofMaxStaleness(1, TimeUnit.SECONDS))) { + try (ResultSet rs = tx.executeQuery(Statement.of("SELECT 1"))) { + rs.next(); + } + } } } diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITTransactionTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITTransactionTest.java index 4e95f8efe81..5a163d4f6d5 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITTransactionTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITTransactionTest.java @@ -446,7 +446,12 @@ public void nestedSingleUseReadTxnThrows() { new TransactionCallable() { @Override public Void run(TransactionContext transaction) throws SpannerException { - client.singleUseReadOnlyTransaction(); + try (ResultSet rs = + client + .singleUseReadOnlyTransaction() + .executeQuery(Statement.of("SELECT 1"))) { + rs.next(); + } return null; } From e3ebeb313fadfb37016998f5b03e261d58306f6f Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Tue, 25 Feb 2020 12:04:22 +0100 Subject: [PATCH 04/49] feat: add read methods support --- .../cloud/spanner/AbstractReadContext.java | 83 ++++++ .../com/google/cloud/spanner/ReadContext.java | 12 + .../com/google/cloud/spanner/SessionPool.java | 106 ++++++- .../cloud/spanner/DatabaseClientImplTest.java | 39 ++- .../cloud/spanner/MockSpannerServiceImpl.java | 23 +- .../google/cloud/spanner/ReadAsyncTest.java | 250 ++++++++++++++++ .../spanner/SpannerExceptionFactoryTest.java | 5 + .../google/cloud/spanner/SpannerMatchers.java | 37 +++ .../cloud/spanner/it/ITAsyncAPITest.java | 274 ++++++++++++++++++ .../google/cloud/spanner/it/ITReadTest.java | 3 + 10 files changed, 801 insertions(+), 31 deletions(-) create mode 100644 google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadAsyncTest.java create mode 100644 google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITAsyncAPITest.java diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java index 1b10f3ba34d..d395345ace7 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java @@ -21,17 +21,22 @@ import static com.google.common.base.Preconditions.checkNotNull; import static com.google.common.base.Preconditions.checkState; +import com.google.api.core.ApiFuture; +import com.google.api.core.SettableApiFuture; import com.google.api.gax.core.ExecutorProvider; import com.google.cloud.Timestamp; import com.google.cloud.spanner.AbstractResultSet.CloseableIterator; import com.google.cloud.spanner.AbstractResultSet.GrpcResultSet; import com.google.cloud.spanner.AbstractResultSet.GrpcStreamIterator; import com.google.cloud.spanner.AbstractResultSet.ResumableStreamIterator; +import com.google.cloud.spanner.AsyncResultSet.CallbackResponse; +import com.google.cloud.spanner.AsyncResultSet.ReadyCallback; import com.google.cloud.spanner.Options.QueryOption; import com.google.cloud.spanner.Options.ReadOption; import com.google.cloud.spanner.SessionImpl.SessionTransaction; import com.google.cloud.spanner.spi.v1.SpannerRpc; import com.google.common.annotations.VisibleForTesting; +import com.google.common.util.concurrent.MoreExecutors; import com.google.protobuf.ByteString; import com.google.spanner.v1.BeginTransactionRequest; import com.google.spanner.v1.ExecuteBatchDmlRequest; @@ -388,12 +393,26 @@ public final ResultSet read( return readInternal(table, null, keys, columns, options); } + @Override + public final AsyncResultSet readAsync( + String table, KeySet keys, Iterable columns, ReadOption... options) { + return new AsyncResultSetImpl( + executorProvider, readInternal(table, null, keys, columns, options)); + } + @Override public final ResultSet readUsingIndex( String table, String index, KeySet keys, Iterable columns, ReadOption... options) { return readInternal(table, checkNotNull(index), keys, columns, options); } + @Override + public final AsyncResultSet readUsingIndexAsync( + String table, String index, KeySet keys, Iterable columns, ReadOption... options) { + return new AsyncResultSetImpl( + executorProvider, readInternal(table, checkNotNull(index), keys, columns, options)); + } + @Nullable @Override public final Struct readRow(String table, Key key, Iterable columns) { @@ -402,6 +421,13 @@ public final Struct readRow(String table, Key key, Iterable columns) { } } + @Override + public final ApiFuture readRowAsync(String table, Key key, Iterable columns) { + try (AsyncResultSet resultSet = readAsync(table, KeySet.singleKey(key), columns)) { + return consumeSingleRowAsync(resultSet); + } + } + @Nullable @Override public final Struct readRowUsingIndex( @@ -411,6 +437,15 @@ public final Struct readRowUsingIndex( } } + @Override + public final ApiFuture readRowUsingIndexAsync( + String table, String index, Key key, Iterable columns) { + try (AsyncResultSet resultSet = + readUsingIndexAsync(table, index, KeySet.singleKey(key), columns)) { + return consumeSingleRowAsync(resultSet); + } + } + @Override public final ResultSet executeQuery(Statement statement, QueryOption... options) { return executeQueryInternal( @@ -676,4 +711,52 @@ private Struct consumeSingleRow(ResultSet resultSet) { } return row; } + + private ApiFuture consumeSingleRowAsync(AsyncResultSet resultSet) { + SettableApiFuture result = SettableApiFuture.create(); + // We can safely use a directExecutor here, as we will only be consuming one row, and we will + // not be doing any blocking stuff in the handler. + resultSet.setCallback(MoreExecutors.directExecutor(), ConsumeSingleRowCallback.create(result)); + return result; + } + + /** + * {@link ReadyCallback} for returning the first row in a result set as a future {@link Struct}. + */ + static class ConsumeSingleRowCallback implements ReadyCallback { + private final SettableApiFuture result; + private Struct row; + + static ConsumeSingleRowCallback create(SettableApiFuture result) { + return new ConsumeSingleRowCallback(result); + } + + private ConsumeSingleRowCallback(SettableApiFuture result) { + this.result = result; + } + + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + try { + while (true) { + switch (resultSet.tryNext()) { + case DONE: + result.set(row); + return CallbackResponse.DONE; + case NOT_READY: + return CallbackResponse.CONTINUE; + case OK: + if (row != null) { + throw newSpannerException( + ErrorCode.INTERNAL, "Multiple rows returned for single key"); + } + row = resultSet.getCurrentRowAsStruct(); + } + } + } catch (Throwable t) { + result.setException(t); + return CallbackResponse.DONE; + } + } + } } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ReadContext.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ReadContext.java index 542c3da4771..904fa4176b1 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ReadContext.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ReadContext.java @@ -16,6 +16,7 @@ package com.google.cloud.spanner; +import com.google.api.core.ApiFuture; import com.google.cloud.spanner.Options.QueryOption; import com.google.cloud.spanner.Options.ReadOption; import javax.annotation.Nullable; @@ -65,6 +66,9 @@ enum QueryAnalyzeMode { */ ResultSet read(String table, KeySet keys, Iterable columns, ReadOption... options); + AsyncResultSet readAsync( + String table, KeySet keys, Iterable columns, ReadOption... options); + /** * Reads zero or more rows from a database using an index. * @@ -93,6 +97,9 @@ enum QueryAnalyzeMode { ResultSet readUsingIndex( String table, String index, KeySet keys, Iterable columns, ReadOption... options); + AsyncResultSet readUsingIndexAsync( + String table, String index, KeySet keys, Iterable columns, ReadOption... options); + /** * Reads a single row from a database, returning {@code null} if the row does not exist. * @@ -112,6 +119,8 @@ ResultSet readUsingIndex( @Nullable Struct readRow(String table, Key key, Iterable columns); + ApiFuture readRowAsync(String table, Key key, Iterable columns); + /** * Reads a single row from a database using an index, returning {@code null} if the row does not * exist. @@ -134,6 +143,9 @@ ResultSet readUsingIndex( @Nullable Struct readRowUsingIndex(String table, String index, Key key, Iterable columns); + ApiFuture readRowUsingIndexAsync( + String table, String index, Key key, Iterable columns); + /** * Executes a query against the database. * diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java index facab374cd0..0cff2f0687c 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java @@ -39,6 +39,7 @@ import com.google.cloud.Timestamp; import com.google.cloud.grpc.GrpcTransportOptions; import com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory; +import com.google.cloud.spanner.AbstractReadContext.ConsumeSingleRowCallback; import com.google.cloud.spanner.DatabaseClient.AsyncWork; import com.google.cloud.spanner.Options.QueryOption; import com.google.cloud.spanner.Options.ReadOption; @@ -183,15 +184,6 @@ T getReadContextDelegate() { } private ResultSet wrap(final CachedResultSetSupplier resultSetSupplier) { - // ResultSet res; - // while (true) { - // try { - // res = resultSetSupplier.get(); - // break; - // } catch (SessionNotFoundException e) { - // replaceSessionIfPossible(e); - // } - // } return new ForwardingResultSet(resultSetSupplier) { private boolean beforeFirst = true; @@ -276,6 +268,23 @@ ResultSet load() { }); } + @Override + public AsyncResultSet readAsync( + final String table, + final KeySet keys, + final Iterable columns, + final ReadOption... options) { + return new AsyncResultSetImpl( + sessionPool.sessionClient.getSpanner().getAsyncExecutorProvider(), + wrap( + new CachedResultSetSupplier() { + @Override + ResultSet load() { + return getReadContextDelegate().read(table, keys, columns, options); + } + })); + } + @Override public ResultSet readUsingIndex( final String table, @@ -292,6 +301,25 @@ ResultSet load() { }); } + @Override + public AsyncResultSet readUsingIndexAsync( + final String table, + final String index, + final KeySet keys, + final Iterable columns, + final ReadOption... options) { + return new AsyncResultSetImpl( + sessionPool.sessionClient.getSpanner().getAsyncExecutorProvider(), + wrap( + new CachedResultSetSupplier() { + @Override + ResultSet load() { + return getReadContextDelegate() + .readUsingIndex(table, index, keys, columns, options); + } + })); + } + @Override @Nullable public Struct readRow(String table, Key key, Iterable columns) { @@ -312,6 +340,15 @@ public Struct readRow(String table, Key key, Iterable columns) { } } + @Override + public ApiFuture readRowAsync(String table, Key key, Iterable columns) { + SettableApiFuture result = SettableApiFuture.create(); + try (AsyncResultSet rs = readAsync(table, KeySet.singleKey(key), columns)) { + rs.setCallback(MoreExecutors.directExecutor(), ConsumeSingleRowCallback.create(result)); + } + return result; + } + @Override @Nullable public Struct readRowUsingIndex(String table, String index, Key key, Iterable columns) { @@ -332,6 +369,16 @@ public Struct readRowUsingIndex(String table, String index, Key key, Iterable readRowUsingIndexAsync( + String table, String index, Key key, Iterable columns) { + SettableApiFuture result = SettableApiFuture.create(); + try (AsyncResultSet rs = readUsingIndexAsync(table, index, KeySet.singleKey(key), columns)) { + rs.setCallback(MoreExecutors.directExecutor(), ConsumeSingleRowCallback.create(result)); + } + return result; + } + @Override public ResultSet executeQuery(final Statement statement, final QueryOption... options) { return wrap( @@ -434,6 +481,13 @@ public ResultSet read( return new SessionPoolResultSet(delegate.read(table, keys, columns, options)); } + @Override + public AsyncResultSet readAsync( + String table, KeySet keys, Iterable columns, ReadOption... options) { + throw SpannerExceptionFactory.newSpannerException( + ErrorCode.UNIMPLEMENTED, "not yet implemented"); + } + @Override public ResultSet readUsingIndex( String table, @@ -445,6 +499,17 @@ public ResultSet readUsingIndex( delegate.readUsingIndex(table, index, keys, columns, options)); } + @Override + public AsyncResultSet readUsingIndexAsync( + String table, + String index, + KeySet keys, + Iterable columns, + ReadOption... options) { + throw SpannerExceptionFactory.newSpannerException( + ErrorCode.UNIMPLEMENTED, "not yet implemented"); + } + @Override public Struct readRow(String table, Key key, Iterable columns) { try { @@ -454,6 +519,15 @@ public Struct readRow(String table, Key key, Iterable columns) { } } + @Override + public ApiFuture readRowAsync(String table, Key key, Iterable columns) { + SettableApiFuture result = SettableApiFuture.create(); + try (AsyncResultSet rs = readAsync(table, KeySet.singleKey(key), columns)) { + rs.setCallback(MoreExecutors.directExecutor(), ConsumeSingleRowCallback.create(result)); + } + return result; + } + @Override public void buffer(Mutation mutation) { delegate.buffer(mutation); @@ -469,6 +543,17 @@ public Struct readRowUsingIndex( } } + @Override + public ApiFuture readRowUsingIndexAsync( + String table, String index, Key key, Iterable columns) { + SettableApiFuture result = SettableApiFuture.create(); + try (AsyncResultSet rs = + readUsingIndexAsync(table, index, KeySet.singleKey(key), columns)) { + rs.setCallback(MoreExecutors.directExecutor(), ConsumeSingleRowCallback.create(result)); + } + return result; + } + @Override public void buffer(Iterable mutations) { delegate.buffer(mutations); @@ -499,7 +584,8 @@ public ResultSet executeQuery(Statement statement, QueryOption... options) { @Override public AsyncResultSet executeQueryAsync(Statement statement, QueryOption... options) { - return null; + throw SpannerExceptionFactory.newSpannerException( + ErrorCode.UNIMPLEMENTED, "not yet implemented"); } @Override diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java index aa514499704..83913f1bf8b 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java @@ -36,6 +36,7 @@ import com.google.cloud.spanner.TransactionRunner.TransactionCallable; import com.google.common.base.Stopwatch; import com.google.protobuf.AbstractMessage; +import com.google.common.util.concurrent.SettableFuture; import com.google.protobuf.ListValue; import com.google.spanner.v1.ExecuteSqlRequest; import com.google.spanner.v1.ExecuteSqlRequest.QueryMode; @@ -88,6 +89,7 @@ public class DatabaseClientImplTest { Statement.of("UPDATE NON_EXISTENT_TABLE SET BAR=1 WHERE BAZ=2"); private static final long UPDATE_COUNT = 1L; private static final Statement SELECT1 = Statement.of("SELECT 1 AS COL1"); + private static final Statement READ1_STATEMENT = Statement.of("SELECT COL1 FROM FOO WHERE ID=1"); private static final ResultSetMetadata SELECT1_METADATA = ResultSetMetadata.newBuilder() .setRowType( @@ -121,6 +123,7 @@ public static void startStaticServer() throws IOException { mockSpanner.setAbortProbability(0.0D); // We don't want any unpredictable aborted transactions. mockSpanner.putStatementResult(StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT)); mockSpanner.putStatementResult(StatementResult.query(SELECT1, SELECT1_RESULTSET)); + mockSpanner.putStatementResult(StatementResult.query(READ1_STATEMENT, SELECT1_RESULTSET)); mockSpanner.putStatementResult( StatementResult.exception( INVALID_UPDATE_STATEMENT, @@ -1193,7 +1196,8 @@ public void testBackendPartitionQueryOptions() { } } - public void testAsyncQuery() throws InterruptedException { + @Test + public void testAsyncQuery() throws Exception { final int EXPECTED_ROW_COUNT = 10; RandomResultSetGenerator generator = new RandomResultSetGenerator(EXPECTED_ROW_COUNT); com.google.spanner.v1.ResultSet resultSet = generator.generate(); @@ -1202,7 +1206,7 @@ public void testAsyncQuery() throws InterruptedException { DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); ExecutorService executor = Executors.newSingleThreadExecutor(); - final CountDownLatch finished = new CountDownLatch(1); + final SettableFuture finished = SettableFuture.create(); final List receivedResults = new ArrayList<>(); try (AsyncResultSet rs = client.singleUse().executeQueryAsync(Statement.of("SELECT * FROM RANDOM"))) { @@ -1211,24 +1215,29 @@ public void testAsyncQuery() throws InterruptedException { new ReadyCallback() { @Override public CallbackResponse cursorReady(AsyncResultSet resultSet) { - while (true) { - switch (rs.tryNext()) { - case DONE: - finished.countDown(); - return CallbackResponse.DONE; - case NOT_READY: - return CallbackResponse.CONTINUE; - case OK: - receivedResults.add(resultSet.getCurrentRowAsStruct()); - break; - default: - throw new IllegalStateException("Unknown cursor state"); + try { + while (true) { + switch (rs.tryNext()) { + case DONE: + finished.set(true); + return CallbackResponse.DONE; + case NOT_READY: + return CallbackResponse.CONTINUE; + case OK: + receivedResults.add(resultSet.getCurrentRowAsStruct()); + break; + default: + throw new IllegalStateException("Unknown cursor state"); + } } + } catch (Throwable t) { + finished.setException(t); + return CallbackResponse.DONE; } } }); } - finished.await(); + assertThat(finished.get()).isTrue(); assertThat(receivedResults.size()).isEqualTo(EXPECTED_ROW_COUNT); } } 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 54ae39bbc71..8e0375bd279 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 @@ -273,6 +273,7 @@ public static Statement createReadStatement( builder.append(", "); } builder.append(col); + first = false; } builder.append(" FROM ").append(table); if (keySet.isAll()) { @@ -392,6 +393,11 @@ public static SimulatedExecutionTime ofStickyException(Exception exception) { return new SimulatedExecutionTime(0, 0, Arrays.asList(exception), true); } + public static SimulatedExecutionTime stickyDatabaseNotFoundException(String name) { + return ofStickyException( + SpannerExceptionFactoryTest.newStatusDatabaseNotFoundException(name)); + } + public static SimulatedExecutionTime ofExceptions(Collection exceptions) { return new SimulatedExecutionTime(0, 0, exceptions, false); } @@ -1227,12 +1233,17 @@ public Iterator iterator() { return request.getColumnsList().iterator(); } }; - StatementResult res = - statementResults.get( - StatementResult.createReadStatement( - request.getTable(), - request.getKeySet().getAll() ? KeySet.all() : KeySet.singleKey(Key.of()), - cols)); + Statement statement = + StatementResult.createReadStatement( + request.getTable(), + request.getKeySet().getAll() ? KeySet.all() : KeySet.singleKey(Key.of()), + cols); + StatementResult res = statementResults.get(statement); + if (res == null) { + throw Status.NOT_FOUND + .withDescription("No result found for " + statement.toString()) + .asRuntimeException(); + } returnPartialResultSet( res.getResultSet(), transactionId, request.getTransaction(), responseObserver); } catch (StatusRuntimeException e) { diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadAsyncTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadAsyncTest.java new file mode 100644 index 00000000000..112da3948b0 --- /dev/null +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadAsyncTest.java @@ -0,0 +1,250 @@ +/* + * Copyright 2020 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.spanner; + +import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.fail; + +import com.google.api.core.ApiFuture; +import com.google.api.gax.grpc.testing.LocalChannelProvider; +import com.google.cloud.NoCredentials; +import com.google.cloud.spanner.AsyncResultSet.CallbackResponse; +import com.google.cloud.spanner.AsyncResultSet.ReadyCallback; +import com.google.cloud.spanner.MockSpannerServiceImpl.SimulatedExecutionTime; +import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult; +import com.google.cloud.spanner.Type.StructField; +import com.google.common.util.concurrent.SettableFuture; +import com.google.protobuf.ListValue; +import com.google.spanner.v1.ResultSetMetadata; +import com.google.spanner.v1.StructType; +import com.google.spanner.v1.StructType.Field; +import com.google.spanner.v1.TypeCode; +import io.grpc.Server; +import io.grpc.Status; +import io.grpc.inprocess.InProcessServerBuilder; +import java.util.Arrays; +import java.util.List; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledThreadPoolExecutor; +import org.junit.After; +import org.junit.AfterClass; +import org.junit.Before; +import org.junit.BeforeClass; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class ReadAsyncTest { + private static final String EMPTY_TABLE_NAME = "EmptyTestTable"; + private static final String TABLE_NAME = "TestTable"; + private static final List ALL_COLUMNS = Arrays.asList("Key", "StringValue"); + private static final Type TABLE_TYPE = + Type.struct( + StructField.of("Key", Type.string()), StructField.of("StringValue", Type.string())); + private static final String TEST_PROJECT = "my-project"; + private static final String TEST_INSTANCE = "my-instance"; + private static final String TEST_DATABASE = "my-database"; + private static MockSpannerServiceImpl mockSpanner; + private static Server server; + private static LocalChannelProvider channelProvider; + private static final Statement EMPTY_READ_STATEMENT = + Statement.of("SELECT Key, StringValue FROM EmptyTestTable WHERE ID=1"); + private static final Statement READ1_STATEMENT = + Statement.of("SELECT Key, StringValue FROM TestTable WHERE ID=1"); + private static final ResultSetMetadata READ1_METADATA = + ResultSetMetadata.newBuilder() + .setRowType( + StructType.newBuilder() + .addFields( + Field.newBuilder() + .setName("Key") + .setType( + com.google.spanner.v1.Type.newBuilder() + .setCode(TypeCode.STRING) + .build()) + .build()) + .addFields( + Field.newBuilder() + .setName("StringValue") + .setType( + com.google.spanner.v1.Type.newBuilder() + .setCode(TypeCode.STRING) + .build()) + .build()) + .build()) + .build(); + private static final com.google.spanner.v1.ResultSet EMPTY_RESULTSET = + com.google.spanner.v1.ResultSet.newBuilder() + .addRows(ListValue.newBuilder().build()) + .setMetadata(READ1_METADATA) + .build(); + private static final com.google.spanner.v1.ResultSet READ1_RESULTSET = + com.google.spanner.v1.ResultSet.newBuilder() + .addRows( + ListValue.newBuilder() + .addValues(com.google.protobuf.Value.newBuilder().setStringValue("k1").build()) + .addValues(com.google.protobuf.Value.newBuilder().setStringValue("v1").build()) + .build()) + .setMetadata(READ1_METADATA) + .build(); + + private static ExecutorService executor; + private Spanner spanner; + private DatabaseClient client; + + @BeforeClass + public static void setup() throws Exception { + mockSpanner = new MockSpannerServiceImpl(); + mockSpanner.putStatementResult(StatementResult.query(READ1_STATEMENT, READ1_RESULTSET)); + mockSpanner.putStatementResult(StatementResult.query(EMPTY_READ_STATEMENT, EMPTY_RESULTSET)); + + String uniqueName = InProcessServerBuilder.generateName(); + server = + InProcessServerBuilder.forName(uniqueName) + .scheduledExecutorService(new ScheduledThreadPoolExecutor(1)) + .addService(mockSpanner) + .build() + .start(); + channelProvider = LocalChannelProvider.create(uniqueName); + executor = Executors.newSingleThreadExecutor(); + } + + @AfterClass + public static void teardown() throws Exception { + executor.shutdown(); + server.shutdown(); + server.awaitTermination(); + } + + @Before + public void before() { + spanner = + SpannerOptions.newBuilder() + .setProjectId(TEST_PROJECT) + .setChannelProvider(channelProvider) + .setCredentials(NoCredentials.getInstance()) + .setSessionPoolOption(SessionPoolOptions.newBuilder().setFailOnSessionLeak().build()) + .build() + .getService(); + client = spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); + } + + @After + public void after() { + spanner.close(); + mockSpanner.removeAllExecutionTimes(); + } + + @Test + public void emptyReadAsync() throws Exception { + final SettableFuture result = SettableFuture.create(); + AsyncResultSet resultSet = + client + .singleUse(TimestampBound.strong()) + .readAsync(EMPTY_TABLE_NAME, KeySet.singleKey(Key.of("k99")), ALL_COLUMNS); + resultSet.setCallback( + executor, + new ReadyCallback() { + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + try { + while (true) { + switch (resultSet.tryNext()) { + case OK: + fail("received unexpected data"); + case NOT_READY: + return CallbackResponse.CONTINUE; + case DONE: + assertThat(resultSet.getType()).isEqualTo(TABLE_TYPE); + result.set(true); + return CallbackResponse.DONE; + } + } + } catch (Throwable t) { + result.setException(t); + return CallbackResponse.DONE; + } + } + }); + assertThat(result.get()).isTrue(); + } + + @Test + public void pointReadAsync() throws Exception { + ApiFuture row = + client + .singleUse(TimestampBound.strong()) + .readRowAsync(TABLE_NAME, Key.of("k1"), ALL_COLUMNS); + assertThat(row.get()).isNotNull(); + assertThat(row.get().getString(0)).isEqualTo("k1"); + assertThat(row.get().getString(1)).isEqualTo("v1"); + assertThat(row.get()) + .isEqualTo(Struct.newBuilder().set("Key").to("k1").set("StringValue").to("v1").build()); + } + + @Test + public void pointReadNotFound() throws Exception { + ApiFuture row = + client + .singleUse(TimestampBound.strong()) + .readRowAsync(EMPTY_TABLE_NAME, Key.of("k999"), ALL_COLUMNS); + assertThat(row.get()).isNull(); + } + + @Test + public void invalidDatabase() throws Exception { + mockSpanner.setBatchCreateSessionsExecutionTime( + SimulatedExecutionTime.stickyDatabaseNotFoundException("invalid-database")); + DatabaseClient invalidClient = + spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, "invalid-database")); + ApiFuture row = + invalidClient + .singleUse(TimestampBound.strong()) + .readRowAsync(TABLE_NAME, Key.of("k99"), ALL_COLUMNS); + try { + row.get(); + fail("missing expected exception"); + } catch (ExecutionException e) { + assertThat(e.getCause()).isInstanceOf(DatabaseNotFoundException.class); + } + } + + @Test + public void tableNotFound() throws Exception { + mockSpanner.setStreamingReadExecutionTime( + SimulatedExecutionTime.ofException( + Status.NOT_FOUND + .withDescription("Table not found: BadTableName") + .asRuntimeException())); + ApiFuture row = + client + .singleUse(TimestampBound.strong()) + .readRowAsync("BadTableName", Key.of("k1"), ALL_COLUMNS); + try { + row.get(); + fail("missing expected exception"); + } catch (ExecutionException e) { + assertThat(e.getCause()).isInstanceOf(SpannerException.class); + SpannerException se = (SpannerException) e.getCause(); + assertThat(se.getErrorCode()).isEqualTo(ErrorCode.NOT_FOUND); + assertThat(se.getMessage()).contains("BadTableName"); + } + } +} diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerExceptionFactoryTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerExceptionFactoryTest.java index bc7dd5498de..49cbfb905d2 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerExceptionFactoryTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerExceptionFactoryTest.java @@ -52,6 +52,11 @@ static DatabaseNotFoundException newDatabaseNotFoundException(String name) { "Database", SpannerExceptionFactory.DATABASE_RESOURCE_TYPE, name); } + static StatusRuntimeException newStatusDatabaseNotFoundException(String name) { + return newStatusResourceNotFoundException( + "Database", SpannerExceptionFactory.DATABASE_RESOURCE_TYPE, name); + } + static InstanceNotFoundException newInstanceNotFoundException(String name) { return (InstanceNotFoundException) newResourceNotFoundException( diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerMatchers.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerMatchers.java index 9662047867a..4723497a472 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerMatchers.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerMatchers.java @@ -21,6 +21,7 @@ import com.google.protobuf.Message; import com.google.protobuf.TextFormat; import java.lang.reflect.InvocationTargetException; +import java.util.concurrent.ExecutionException; import org.hamcrest.BaseMatcher; import org.hamcrest.Description; import org.hamcrest.Matcher; @@ -47,6 +48,15 @@ public static Matcher isSpannerException(ErrorCode code return new SpannerExceptionMatcher<>(code); } + /** + * Returns a method that checks that a {@link Throwable} is an {@link ExecutionException} where + * the cause is a {@link SpannerException} with an error code to {@code code}. + */ + public static Matcher isExecutionExceptionWithSpannerCause( + ErrorCode code) { + return new ExecutionExceptionWithSpannerCauseMatcher<>(code); + } + private static class ProtoTextMatcher extends BaseMatcher { private final T expected; @@ -110,4 +120,31 @@ public void describeTo(Description description) { description.appendText("SpannerException[" + expectedCode + "]"); } } + + private static class ExecutionExceptionWithSpannerCauseMatcher + extends BaseMatcher { + private final ErrorCode expectedCode; + + ExecutionExceptionWithSpannerCauseMatcher(ErrorCode expectedCode) { + this.expectedCode = checkNotNull(expectedCode); + } + + @Override + public boolean matches(Object item) { + if (!(item instanceof ExecutionException)) { + return false; + } + ExecutionException ee = (ExecutionException) item; + if (!(ee.getCause() instanceof SpannerException)) { + return false; + } + SpannerException e = (SpannerException) ee.getCause(); + return e.getErrorCode() == expectedCode; + } + + @Override + public void describeTo(Description description) { + description.appendText("ExecutionException[SpannerException[" + expectedCode + "]]"); + } + } } diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITAsyncAPITest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITAsyncAPITest.java new file mode 100644 index 00000000000..bc46ff8b0f2 --- /dev/null +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITAsyncAPITest.java @@ -0,0 +1,274 @@ +/* + * 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.it; + +import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.fail; + +import com.google.api.core.ApiFuture; +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.Database; +import com.google.cloud.spanner.DatabaseClient; +import com.google.cloud.spanner.DatabaseId; +import com.google.cloud.spanner.ErrorCode; +import com.google.cloud.spanner.IntegrationTest; +import com.google.cloud.spanner.IntegrationTestEnv; +import com.google.cloud.spanner.Key; +import com.google.cloud.spanner.KeyRange; +import com.google.cloud.spanner.KeySet; +import com.google.cloud.spanner.Mutation; +import com.google.cloud.spanner.SpannerException; +import com.google.cloud.spanner.Struct; +import com.google.cloud.spanner.TimestampBound; +import com.google.cloud.spanner.Type; +import com.google.cloud.spanner.Type.StructField; +import com.google.cloud.spanner.testing.RemoteSpannerHelper; +import com.google.common.util.concurrent.SettableFuture; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.ClassRule; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Integration tests for asynchronous APIs. */ +@Category(IntegrationTest.class) +@RunWith(JUnit4.class) +public class ITAsyncAPITest { + @ClassRule public static IntegrationTestEnv env = new IntegrationTestEnv(); + private static final String TABLE_NAME = "TestTable"; + private static final String INDEX_NAME = "TestTableByValue"; + private static final List ALL_COLUMNS = Arrays.asList("Key", "StringValue"); + private static final Type TABLE_TYPE = + Type.struct( + StructField.of("Key", Type.string()), StructField.of("StringValue", Type.string())); + + private static Database db; + private static DatabaseClient client; + private static ExecutorService executor; + + @BeforeClass + public static void setUpDatabase() { + db = + env.getTestHelper() + .createTestDatabase( + "CREATE TABLE TestTable (" + + " Key STRING(MAX) NOT NULL," + + " StringValue STRING(MAX)," + + ") PRIMARY KEY (Key)", + "CREATE INDEX TestTableByValue ON TestTable(StringValue)", + "CREATE INDEX TestTableByValueDesc ON TestTable(StringValue DESC)"); + client = env.getTestHelper().getDatabaseClient(db); + + // Includes k0..k14. Note that strings k{10,14} sort between k1 and k2. + List mutations = new ArrayList<>(); + for (int i = 0; i < 15; ++i) { + mutations.add( + Mutation.newInsertOrUpdateBuilder(TABLE_NAME) + .set("Key") + .to("k" + i) + .set("StringValue") + .to("v" + i) + .build()); + } + client.write(mutations); + executor = Executors.newSingleThreadExecutor(); + } + + @AfterClass + public static void cleanup() { + executor.shutdown(); + } + + @Test + public void emptyReadAsync() throws Exception { + final SettableFuture result = SettableFuture.create(); + AsyncResultSet resultSet = + client + .singleUse(TimestampBound.strong()) + .readAsync( + TABLE_NAME, + KeySet.range(KeyRange.closedOpen(Key.of("k99"), Key.of("z"))), + ALL_COLUMNS); + resultSet.setCallback( + executor, + new ReadyCallback() { + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + try { + while (true) { + switch (resultSet.tryNext()) { + case OK: + fail("received unexpected data"); + case NOT_READY: + return CallbackResponse.CONTINUE; + case DONE: + assertThat(resultSet.getType()).isEqualTo(TABLE_TYPE); + result.set(true); + return CallbackResponse.DONE; + } + } + } catch (Throwable t) { + result.setException(t); + return CallbackResponse.DONE; + } + } + }); + assertThat(result.get()).isTrue(); + } + + @Test + public void indexEmptyReadAsync() throws Exception { + final SettableFuture result = SettableFuture.create(); + AsyncResultSet resultSet = + client + .singleUse(TimestampBound.strong()) + .readUsingIndexAsync( + TABLE_NAME, + INDEX_NAME, + KeySet.range(KeyRange.closedOpen(Key.of("v99"), Key.of("z"))), + ALL_COLUMNS); + resultSet.setCallback( + executor, + new ReadyCallback() { + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + try { + while (true) { + switch (resultSet.tryNext()) { + case OK: + fail("received unexpected data"); + case NOT_READY: + return CallbackResponse.CONTINUE; + case DONE: + assertThat(resultSet.getType()).isEqualTo(TABLE_TYPE); + result.set(true); + return CallbackResponse.DONE; + } + } + } catch (Throwable t) { + result.setException(t); + return CallbackResponse.DONE; + } + } + }); + assertThat(result.get()).isTrue(); + } + + @Test + public void pointReadAsync() throws Exception { + ApiFuture row = + client + .singleUse(TimestampBound.strong()) + .readRowAsync(TABLE_NAME, Key.of("k1"), ALL_COLUMNS); + assertThat(row.get()).isNotNull(); + assertThat(row.get().getString(0)).isEqualTo("k1"); + assertThat(row.get().getString(1)).isEqualTo("v1"); + // Ensure that the Struct implementation supports equality properly. + assertThat(row.get()) + .isEqualTo(Struct.newBuilder().set("Key").to("k1").set("StringValue").to("v1").build()); + } + + @Test + public void indexPointReadAsync() throws Exception { + ApiFuture row = + client + .singleUse(TimestampBound.strong()) + .readRowUsingIndexAsync(TABLE_NAME, INDEX_NAME, Key.of("v1"), ALL_COLUMNS); + assertThat(row.get()).isNotNull(); + assertThat(row.get().getString(0)).isEqualTo("k1"); + assertThat(row.get().getString(1)).isEqualTo("v1"); + } + + @Test + public void pointReadNotFound() throws Exception { + ApiFuture row = + client + .singleUse(TimestampBound.strong()) + .readRowAsync(TABLE_NAME, Key.of("k999"), ALL_COLUMNS); + assertThat(row.get()).isNull(); + } + + @Test + public void indexPointReadNotFound() throws Exception { + ApiFuture row = + client + .singleUse(TimestampBound.strong()) + .readRowUsingIndexAsync(TABLE_NAME, INDEX_NAME, Key.of("v999"), ALL_COLUMNS); + assertThat(row.get()).isNull(); + } + + @Test + public void invalidDatabase() throws Exception { + RemoteSpannerHelper helper = env.getTestHelper(); + DatabaseClient invalidClient = + helper.getClient().getDatabaseClient(DatabaseId.of(helper.getInstanceId(), "invalid")); + ApiFuture row = + invalidClient + .singleUse(TimestampBound.strong()) + .readRowAsync(TABLE_NAME, Key.of("k99"), ALL_COLUMNS); + try { + row.get(); + fail("missing expected exception"); + } catch (ExecutionException e) { + assertThat(e.getCause()).isInstanceOf(SpannerException.class); + SpannerException se = (SpannerException) e.getCause(); + assertThat(se.getErrorCode()).isEqualTo(ErrorCode.NOT_FOUND); + } + } + + @Test + public void tableNotFound() throws Exception { + ApiFuture row = + client + .singleUse(TimestampBound.strong()) + .readRowAsync("BadTableName", Key.of("k1"), ALL_COLUMNS); + try { + row.get(); + } catch (ExecutionException e) { + assertThat(e.getCause()).isInstanceOf(SpannerException.class); + SpannerException se = (SpannerException) e.getCause(); + assertThat(se.getErrorCode()).isEqualTo(ErrorCode.NOT_FOUND); + assertThat(se.getMessage()).contains("BadTableName"); + } + } + + @Test + public void columnNotFound() throws Exception { + ApiFuture row = + client + .singleUse(TimestampBound.strong()) + .readRowAsync(TABLE_NAME, Key.of("k1"), Arrays.asList("Key", "BadColumnName")); + try { + row.get(); + } catch (ExecutionException e) { + assertThat(e.getCause()).isInstanceOf(SpannerException.class); + SpannerException se = (SpannerException) e.getCause(); + assertThat(se.getErrorCode()).isEqualTo(ErrorCode.NOT_FOUND); + assertThat(se.getMessage()).contains("BadColumnName"); + } + } +} diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITReadTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITReadTest.java index b06d3ae1524..ac73cc4f1a0 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITReadTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITReadTest.java @@ -20,6 +20,7 @@ import static com.google.cloud.spanner.Type.StructField; import static com.google.common.truth.Truth.assertThat; import static com.google.common.truth.Truth.assertWithMessage; +import static org.junit.Assert.fail; import com.google.cloud.spanner.Database; import com.google.cloud.spanner.DatabaseClient; @@ -345,6 +346,7 @@ public void run() { try { work.run(); + fail("missing expected exception"); } catch (SpannerException e) { MatcherAssert.assertThat(e, isSpannerException(ErrorCode.CANCELLED)); } @@ -368,6 +370,7 @@ public void run() { try { work.run(); + fail("missing expected exception"); } catch (SpannerException e) { MatcherAssert.assertThat(e, isSpannerException(ErrorCode.DEADLINE_EXCEEDED)); } finally { From 54629ad2724bc7d643c1141d0a08dc00db9708bd Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Tue, 25 Feb 2020 16:48:05 +0100 Subject: [PATCH 05/49] tests: test async runner --- .../com/google/cloud/spanner/SessionPool.java | 9 + .../cloud/spanner/TransactionContext.java | 4 + .../cloud/spanner/TransactionRunnerImpl.java | 49 ++++ .../cloud/spanner/spi/v1/GapicSpannerRpc.java | 8 +- .../cloud/spanner/spi/v1/SpannerRpc.java | 3 + .../google/cloud/spanner/AsyncRunnerTest.java | 258 ++++++++++++++++++ .../cloud/spanner/DatabaseClientImplTest.java | 2 - .../google/cloud/spanner/ReadAsyncTest.java | 29 ++ 8 files changed, 359 insertions(+), 3 deletions(-) create mode 100644 google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncRunnerTest.java diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java index 0cff2f0687c..9366488791e 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java @@ -568,6 +568,15 @@ public long executeUpdate(Statement statement) { } } + @Override + public ApiFuture executeUpdateAsync(Statement statement) { + try { + return delegate.executeUpdateAsync(statement); + } catch (SessionNotFoundException e) { + throw handleSessionNotFound(e); + } + } + @Override public long[] batchUpdate(Iterable statements) { try { diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionContext.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionContext.java index a529c4c492b..7e09da901c2 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionContext.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionContext.java @@ -16,6 +16,8 @@ package com.google.cloud.spanner; +import com.google.api.core.ApiFuture; + /** * Context for a single attempt of a locking read-write transaction. This type of transaction is the * only way to write data into Cloud Spanner; {@link Session#write(Iterable)} and {@link @@ -102,6 +104,8 @@ public interface TransactionContext extends ReadContext { */ long executeUpdate(Statement statement); + ApiFuture executeUpdateAsync(Statement statement); + /** * 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 {@code executeUpdate} in a 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 ecd04f26e07..6aa15134112 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 @@ -21,12 +21,16 @@ import static com.google.common.base.Preconditions.checkNotNull; import static com.google.common.base.Preconditions.checkState; +import com.google.api.core.ApiFunction; +import com.google.api.core.ApiFuture; +import com.google.api.core.ApiFutures; import com.google.api.gax.core.ExecutorProvider; import com.google.cloud.Timestamp; import com.google.cloud.spanner.SessionImpl.SessionTransaction; import com.google.cloud.spanner.spi.v1.SpannerRpc; import com.google.common.annotations.VisibleForTesting; import com.google.common.collect.ImmutableMap; +import com.google.common.util.concurrent.MoreExecutors; import com.google.protobuf.ByteString; import com.google.rpc.Code; import com.google.spanner.v1.CommitRequest; @@ -34,6 +38,7 @@ import com.google.spanner.v1.ExecuteBatchDmlRequest; import com.google.spanner.v1.ExecuteSqlRequest; import com.google.spanner.v1.ExecuteSqlRequest.QueryMode; +import com.google.spanner.v1.ResultSet; import com.google.spanner.v1.RollbackRequest; import com.google.spanner.v1.TransactionSelector; import io.opencensus.common.Scope; @@ -44,6 +49,7 @@ import java.util.ArrayList; import java.util.List; import java.util.concurrent.Callable; +import java.util.concurrent.ExecutionException; import java.util.concurrent.atomic.AtomicInteger; import java.util.logging.Level; import java.util.logging.Logger; @@ -253,6 +259,49 @@ public long executeUpdate(Statement statement) { } } + @Override + public ApiFuture executeUpdateAsync(Statement statement) { + beforeReadOrQuery(); + final ExecuteSqlRequest.Builder builder = + getExecuteSqlRequestBuilder(statement, QueryMode.NORMAL); + ApiFuture resultSet = + rpc.executeQueryAsync(builder.build(), session.getOptions()); + final ApiFuture updateCount = + ApiFutures.transform( + resultSet, + new ApiFunction() { + @Override + public Long apply(ResultSet input) { + if (!input.hasStats()) { + SpannerException e = + SpannerExceptionFactory.newSpannerException( + ErrorCode.INVALID_ARGUMENT, + "DML response missing stats possibly due to non-DML statement as input"); + onError(e); + throw e; + } + // For standard DML, using the exact row count. + return input.getStats().getRowCountExact(); + } + }, + MoreExecutors.directExecutor()); + updateCount.addListener( + new Runnable() { + @Override + public void run() { + try { + updateCount.get(); + } catch (ExecutionException e) { + onError(SpannerExceptionFactory.newSpannerException(e.getCause())); + } catch (InterruptedException e) { + onError(SpannerExceptionFactory.propagateInterrupt(e)); + } + } + }, + MoreExecutors.directExecutor()); + return updateCount; + } + @Override public long[] batchUpdate(Iterable statements) { beforeReadOrQuery(); diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java index de1a09158cb..c696459feee 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java @@ -729,8 +729,14 @@ public void cancel(String message) { @Override public ResultSet executeQuery(ExecuteSqlRequest request, @Nullable Map options) { + return get(executeQueryAsync(request, options)); + } + + @Override + public ApiFuture executeQueryAsync( + ExecuteSqlRequest request, @Nullable Map options) { GrpcCallContext context = newCallContext(options, request.getSession()); - return get(spannerStub.executeSqlCallable().futureCall(request, context)); + return spannerStub.executeSqlCallable().futureCall(request, context); } @Override diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerRpc.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerRpc.java index 497be948cc8..a9114bc22b5 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerRpc.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerRpc.java @@ -282,6 +282,9 @@ StreamingCall read( ResultSet executeQuery(ExecuteSqlRequest request, @Nullable Map options); + ApiFuture executeQueryAsync( + ExecuteSqlRequest request, @Nullable Map options); + ResultSet executePartitionedDml( ExecuteSqlRequest request, @Nullable Map options, Duration timeout); diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncRunnerTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncRunnerTest.java new file mode 100644 index 00000000000..c3a14e7bc52 --- /dev/null +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncRunnerTest.java @@ -0,0 +1,258 @@ +/* + * Copyright 2020 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.spanner; + +import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.fail; + +import com.google.api.core.ApiFuture; +import com.google.api.core.SettableApiFuture; +import com.google.api.gax.grpc.testing.LocalChannelProvider; +import com.google.cloud.NoCredentials; +import com.google.cloud.spanner.DatabaseClient.AsyncWork; +import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult; +import io.grpc.Server; +import io.grpc.Status; +import io.grpc.inprocess.InProcessServerBuilder; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledThreadPoolExecutor; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.After; +import org.junit.AfterClass; +import org.junit.Before; +import org.junit.BeforeClass; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class AsyncRunnerTest { + private static final String TEST_PROJECT = "my-project"; + private static final String TEST_INSTANCE = "my-instance"; + private static final String TEST_DATABASE = "my-database"; + + private static final Statement UPDATE_STATEMENT = + Statement.of("UPDATE FOO SET BAR=1 WHERE BAZ=2"); + private static final Statement INVALID_UPDATE_STATEMENT = + Statement.of("UPDATE NON_EXISTENT_TABLE SET BAR=1 WHERE BAZ=2"); + private static final long UPDATE_COUNT = 1L; + + private static MockSpannerServiceImpl mockSpanner; + private static Server server; + private static LocalChannelProvider channelProvider; + private static ExecutorService executor; + + private Spanner spanner; + private Spanner spannerWithEmptySessionPool; + private DatabaseClient client; + private DatabaseClient clientWithEmptySessionPool; + + @BeforeClass + public static void setup() throws Exception { + mockSpanner = new MockSpannerServiceImpl(); + mockSpanner.setAbortProbability(0.0D); + mockSpanner.putStatementResult(StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT)); + mockSpanner.putStatementResult( + StatementResult.exception( + INVALID_UPDATE_STATEMENT, + Status.INVALID_ARGUMENT.withDescription("invalid statement").asRuntimeException())); + String uniqueName = InProcessServerBuilder.generateName(); + server = + InProcessServerBuilder.forName(uniqueName) + .scheduledExecutorService(new ScheduledThreadPoolExecutor(1)) + .addService(mockSpanner) + .build() + .start(); + channelProvider = LocalChannelProvider.create(uniqueName); + executor = Executors.newSingleThreadExecutor(); + } + + @AfterClass + public static void teardown() throws Exception { + executor.shutdown(); + server.shutdown(); + server.awaitTermination(); + } + + @Before + public void before() { + spanner = + SpannerOptions.newBuilder() + .setProjectId(TEST_PROJECT) + .setChannelProvider(channelProvider) + .setCredentials(NoCredentials.getInstance()) + .setSessionPoolOption(SessionPoolOptions.newBuilder().setFailOnSessionLeak().build()) + .build() + .getService(); + client = spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); + spannerWithEmptySessionPool = + spanner + .getOptions() + .toBuilder() + .setSessionPoolOption( + SessionPoolOptions.newBuilder().setFailOnSessionLeak().setMinSessions(0).build()) + .build() + .getService(); + clientWithEmptySessionPool = + spannerWithEmptySessionPool.getDatabaseClient( + DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); + } + + @After + public void after() { + spanner.close(); + spannerWithEmptySessionPool.close(); + mockSpanner.removeAllExecutionTimes(); + } + + @Test + public void asyncRunnerUpdate() throws Exception { + ApiFuture updateCount = + client.runAsync( + new AsyncWork() { + @Override + public ApiFuture doWorkAsync(TransactionContext txn) { + return txn.executeUpdateAsync(UPDATE_STATEMENT); + } + }, + executor); + assertThat(updateCount.get()).isEqualTo(UPDATE_COUNT); + } + + @Test + public void asyncRunnerIsNonBlocking() throws Exception { + mockSpanner.freeze(); + final SettableApiFuture finished = SettableApiFuture.create(); + ApiFuture res = + clientWithEmptySessionPool.runAsync( + new AsyncWork() { + @Override + public ApiFuture doWorkAsync(TransactionContext txn) { + finished.set(null); + return finished; + } + }, + executor); + mockSpanner.unfreeze(); + assertThat(res.get()).isNull(); + } + + @Test + public void asyncRunnerInvalidUpdate() throws Exception { + ApiFuture updateCount = + client.runAsync( + new AsyncWork() { + @Override + public ApiFuture doWorkAsync(TransactionContext txn) { + return txn.executeUpdateAsync(INVALID_UPDATE_STATEMENT); + } + }, + executor); + try { + updateCount.get(); + fail("missing expected exception"); + } catch (ExecutionException e) { + assertThat(e.getCause()).isInstanceOf(SpannerException.class); + SpannerException se = (SpannerException) e.getCause(); + assertThat(se.getErrorCode()).isEqualTo(ErrorCode.INVALID_ARGUMENT); + assertThat(se.getMessage()).contains("invalid statement"); + } + } + + @Test + public void asyncRunnerUpdateAborted() throws Exception { + try { + // Temporarily set the result of the update to 2 rows. + mockSpanner.putStatementResult(StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT + 1L)); + final AtomicInteger attempt = new AtomicInteger(); + ApiFuture updateCount = + client.runAsync( + new AsyncWork() { + @Override + public ApiFuture doWorkAsync(TransactionContext txn) { + if (attempt.incrementAndGet() == 1) { + mockSpanner.abortTransaction(txn); + } else { + // Set the result of the update statement back to 1 row. + mockSpanner.putStatementResult( + StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT)); + } + return txn.executeUpdateAsync(UPDATE_STATEMENT); + } + }, + executor); + assertThat(updateCount.get()).isEqualTo(UPDATE_COUNT); + assertThat(attempt.get()).isEqualTo(2); + } finally { + mockSpanner.putStatementResult(StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT)); + } + } + + @Test + public void asyncRunnerCommitAborted() throws Exception { + try { + // Temporarily set the result of the update to 2 rows. + mockSpanner.putStatementResult(StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT + 1L)); + final AtomicInteger attempt = new AtomicInteger(); + ApiFuture updateCount = + client.runAsync( + new AsyncWork() { + @Override + public ApiFuture doWorkAsync(TransactionContext txn) { + ApiFuture updateCount = txn.executeUpdateAsync(UPDATE_STATEMENT); + if (attempt.incrementAndGet() == 1) { + mockSpanner.abortTransaction(txn); + } else { + // Set the result of the update statement back to 1 row. + mockSpanner.putStatementResult( + StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT)); + } + return updateCount; + } + }, + executor); + assertThat(updateCount.get()).isEqualTo(UPDATE_COUNT); + assertThat(attempt.get()).isEqualTo(2); + } finally { + mockSpanner.putStatementResult(StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT)); + } + } + + @Test + public void asyncRunnerUpdateAbortedWithoutGettingResult() throws Exception { + final SettableApiFuture finished = SettableApiFuture.create(); + final AtomicInteger attempt = new AtomicInteger(); + ApiFuture result = + client.runAsync( + new AsyncWork() { + @Override + public ApiFuture doWorkAsync(TransactionContext txn) { + if (attempt.incrementAndGet() == 1) { + mockSpanner.abortTransaction(txn); + } + txn.executeUpdateAsync(UPDATE_STATEMENT); + finished.set(null); + return finished; + } + }, + executor); + assertThat(result.get()).isNull(); + assertThat(attempt.get()).isEqualTo(2); + } +} diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java index 83913f1bf8b..6c90bba5111 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java @@ -89,7 +89,6 @@ public class DatabaseClientImplTest { Statement.of("UPDATE NON_EXISTENT_TABLE SET BAR=1 WHERE BAZ=2"); private static final long UPDATE_COUNT = 1L; private static final Statement SELECT1 = Statement.of("SELECT 1 AS COL1"); - private static final Statement READ1_STATEMENT = Statement.of("SELECT COL1 FROM FOO WHERE ID=1"); private static final ResultSetMetadata SELECT1_METADATA = ResultSetMetadata.newBuilder() .setRowType( @@ -123,7 +122,6 @@ public static void startStaticServer() throws IOException { mockSpanner.setAbortProbability(0.0D); // We don't want any unpredictable aborted transactions. mockSpanner.putStatementResult(StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT)); mockSpanner.putStatementResult(StatementResult.query(SELECT1, SELECT1_RESULTSET)); - mockSpanner.putStatementResult(StatementResult.query(READ1_STATEMENT, SELECT1_RESULTSET)); mockSpanner.putStatementResult( StatementResult.exception( INVALID_UPDATE_STATEMENT, diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadAsyncTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadAsyncTest.java index 112da3948b0..21442d3f2e5 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadAsyncTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadAsyncTest.java @@ -61,6 +61,35 @@ public class ReadAsyncTest { private static final String TEST_PROJECT = "my-project"; private static final String TEST_INSTANCE = "my-instance"; private static final String TEST_DATABASE = "my-database"; + private static final Statement UPDATE_STATEMENT = + Statement.of("UPDATE FOO SET BAR=1 WHERE BAZ=2"); + private static final Statement INVALID_UPDATE_STATEMENT = + Statement.of("UPDATE NON_EXISTENT_TABLE SET BAR=1 WHERE BAZ=2"); + private static final long UPDATE_COUNT = 1L; + private static final Statement SELECT1 = Statement.of("SELECT 1 AS COL1"); + private static final ResultSetMetadata SELECT1_METADATA = + ResultSetMetadata.newBuilder() + .setRowType( + StructType.newBuilder() + .addFields( + Field.newBuilder() + .setName("COL1") + .setType( + com.google.spanner.v1.Type.newBuilder() + .setCode(TypeCode.INT64) + .build()) + .build()) + .build()) + .build(); + private static final com.google.spanner.v1.ResultSet SELECT1_RESULTSET = + com.google.spanner.v1.ResultSet.newBuilder() + .addRows( + ListValue.newBuilder() + .addValues(com.google.protobuf.Value.newBuilder().setStringValue("1").build()) + .build()) + .setMetadata(SELECT1_METADATA) + .build(); + private static MockSpannerServiceImpl mockSpanner; private static Server server; private static LocalChannelProvider channelProvider; From 8a10b646b416ab61c079f00294e20529ca9f7b9c Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Tue, 25 Feb 2020 21:07:58 +0100 Subject: [PATCH 06/49] feat: create async runner --- .../cloud/spanner/AsyncResultSetImpl.java | 11 +- .../com/google/cloud/spanner/AsyncRunner.java | 66 +++++++++ .../google/cloud/spanner/AsyncRunnerImpl.java | 79 ++++++++++ .../google/cloud/spanner/DatabaseClient.java | 33 +---- .../cloud/spanner/DatabaseClientImpl.java | 6 +- .../com/google/cloud/spanner/SessionImpl.java | 39 +---- .../com/google/cloud/spanner/SessionPool.java | 43 +++--- .../google/cloud/spanner/AsyncRunnerTest.java | 138 +++++++++++++++--- .../cloud/spanner/DatabaseClientImplTest.java | 52 ++----- .../cloud/spanner/MockSpannerTestUtil.java | 124 ++++++++++++++++ .../google/cloud/spanner/ReadAsyncTest.java | 28 ---- .../RetryOnInvalidatedSessionTest.java | 37 ++++- 12 files changed, 474 insertions(+), 182 deletions(-) create mode 100644 google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncRunner.java create mode 100644 google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncRunnerImpl.java create mode 100644 google-cloud-spanner/src/test/java/com/google/cloud/spanner/MockSpannerTestUtil.java 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 fcd0bc770ae..87891f33010 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 @@ -245,8 +245,7 @@ public void run() { case CONTINUE: if (buffer.isEmpty()) { // Call the callback once more if the entire result set has been processed but - // the - // callback has not yet received a CursorState.DONE or a CANCELLED error. + // the callback has not yet received a CursorState.DONE or a CANCELLED error. if (finished && !cursorReturnedDoneOrException) { break; } @@ -327,6 +326,13 @@ public Void call() throws Exception { } } } + // We don't need any more data from the underlying result set, so we close it as soon as + // possible. Any error that might occur during this will be ignored. + try { + delegateResultSet.close(); + } catch (Throwable t) { + } + // Ensure that the callback has been called at least once, even if the result set was // cancelled. synchronized (monitor) { @@ -344,7 +350,6 @@ public Void call() throws Exception { consumingLatch.await(); } } finally { - delegateResultSet.close(); if (executorProvider.shouldAutoClose()) { service.shutdown(); } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncRunner.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncRunner.java new file mode 100644 index 00000000000..432d6a8645d --- /dev/null +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncRunner.java @@ -0,0 +1,66 @@ +/* + * 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.cloud.Timestamp; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executor; + +public interface AsyncRunner { + + interface AsyncWork { + /** + * Performs a single transaction attempt. All reads/writes should be performed using {@code + * txn}. + * + *

    Implementations of this method should not attempt to commit the transaction directly: + * returning normally will result in the runner attempting to commit the transaction once the + * returned future completes, retrying on abort. + * + *

    In most cases, the implementation will not need to catch {@code SpannerException}s from + * Spanner operations, instead letting these propagate to the framework. The transaction runner + * + *

    will take appropriate action based on the type of exception. In particular, + * implementations should never catch an exception of type {@link SpannerErrors#isAborted}: + * these indicate that some reads may have returned inconsistent data and the transaction + * attempt must be aborted. + * + *

    If any exception is thrown, the runner will validate the reads performed in the current + * transaction attempt using {@link Transaction#commitReadsOnly}: if validation succeeds, the + * exception is propagated to the caller; if validation aborts, the exception is thrown away and + * the work is retried; if the commit fails for some other reason, the corresponding {@code + * SpannerException} is returned to the caller. Any buffered mutations will be ignored. + * + * @param txn the transaction + * @return future over the result of the work + *

    TODO(loite): It's probably better to let this method return `R` instead of + * `ApiFuture`, as we need to wait until the result of the work has actually finished + * before we can commit the transaction. Returning an ApiFuture here just means that the + * underlying framework code still has to call {@link ApiFuture#get()} before committing. + */ + ApiFuture doWorkAsync(TransactionContext txn); + } + + ApiFuture runAsync(AsyncWork work, Executor executor); + + /** + * Returns the timestamp at which the transaction committed. {@link ApiFuture#get()} will throw an + * {@link ExecutionException} if the transaction did not commit. + */ + ApiFuture getCommitTimestamp(); +} diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncRunnerImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncRunnerImpl.java new file mode 100644 index 00000000000..6ffb3214907 --- /dev/null +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncRunnerImpl.java @@ -0,0 +1,79 @@ +/* + * Copyright 2020 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.spanner; + +import com.google.api.core.ApiFuture; +import com.google.api.core.SettableApiFuture; +import com.google.cloud.Timestamp; +import com.google.cloud.spanner.TransactionRunner.TransactionCallable; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executor; + +class AsyncRunnerImpl implements AsyncRunner { + private final TransactionRunnerImpl delegate; + private final SettableApiFuture commitTimestamp = SettableApiFuture.create(); + + AsyncRunnerImpl(TransactionRunnerImpl delegate) { + this.delegate = delegate; + } + + @Override + public ApiFuture runAsync(final AsyncWork work, Executor executor) { + final SettableApiFuture res = SettableApiFuture.create(); + executor.execute( + new Runnable() { + @Override + public void run() { + try { + R r = + delegate.run( + new TransactionCallable() { + @Override + public R run(TransactionContext transaction) throws Exception { + try { + return work.doWorkAsync(transaction).get(); + } catch (ExecutionException e) { + throw SpannerExceptionFactory.newSpannerException(e.getCause()); + } catch (InterruptedException e) { + throw SpannerExceptionFactory.propagateInterrupt(e); + } + } + }); + res.set(r); + } catch (Throwable t) { + res.setException(t); + } finally { + setCommitTimestamp(); + } + } + }); + return res; + } + + private void setCommitTimestamp() { + try { + commitTimestamp.set(delegate.getCommitTimestamp()); + } catch (Throwable t) { + commitTimestamp.setException(t); + } + } + + @Override + public ApiFuture getCommitTimestamp() { + return commitTimestamp; + } +} diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClient.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClient.java index dd9e9769066..18cbc3ca858 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClient.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClient.java @@ -16,9 +16,7 @@ package com.google.cloud.spanner; -import com.google.api.core.ApiFuture; import com.google.cloud.Timestamp; -import java.util.concurrent.Executor; /** * Interface for all the APIs that are used to read/write data into a Cloud Spanner database. An @@ -280,36 +278,7 @@ public interface DatabaseClient { */ TransactionManager transactionManager(); - public static interface AsyncWork { - /** - * Performs a single transaction attempt. All reads/writes should be performed using {@code - * txn}. - * - *

    Implementations of this method should not attempt to commit the transaction directly: - * returning normally will result in the runner attempting to commit the transaction once the - * returned future completes, retrying on abort. - * - *

    In most cases, the implementation will not need to catch {@code SpannerException}s from - * Spanner operations, instead letting these propagate to the framework. The transaction runner - * - *

    will take appropriate action based on the type of exception. In particular, - * implementations should never catch an exception of type {@link SpannerErrors#isAborted}: - * these indicate that some reads may have returned inconsistent data and the transaction - * attempt must be aborted. - * - *

    If any exception is thrown, the runner will validate the reads performed in the current - * transaction attempt using {@link Transaction#commitReadsOnly}: if validation succeeds, the - * exception is propagated to the caller; if validation aborts, the exception is thrown away and - * the work is retried; if the commit fails for some other reason, the corresponding {@code - * SpannerException} is returned to the caller. Any buffered mutations will be ignored. - * - * @param txn the transaction - * @return future over the result of the work - */ - ApiFuture doWorkAsync(TransactionContext txn); - } - - ApiFuture runAsync(AsyncWork work, Executor executor); + AsyncRunner runAsync(); /** * Returns the lower bound of rows modified by this DML statement. diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClientImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClientImpl.java index f1bd75b3ccb..44f386a7273 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClientImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClientImpl.java @@ -16,7 +16,6 @@ package com.google.cloud.spanner; -import com.google.api.core.ApiFuture; import com.google.cloud.Timestamp; import com.google.cloud.spanner.SessionPool.PooledSessionFuture; import com.google.common.annotations.VisibleForTesting; @@ -26,7 +25,6 @@ import io.opencensus.trace.Span; import io.opencensus.trace.Tracer; import io.opencensus.trace.Tracing; -import java.util.concurrent.Executor; class DatabaseClientImpl implements DatabaseClient { private static final String READ_WRITE_TRANSACTION = "CloudSpanner.ReadWriteTransaction"; @@ -193,10 +191,10 @@ public TransactionManager transactionManager() { } @Override - public ApiFuture runAsync(AsyncWork work, Executor executor) { + public AsyncRunner runAsync() { Span span = tracer.spanBuilder(READ_WRITE_TRANSACTION).startSpan(); try (Scope s = tracer.withSpan(span)) { - return getReadWriteSession().runAsync(work, executor); + return getReadWriteSession().runAsync(); } catch (RuntimeException e) { TraceUtil.endSpanWithFailure(span, e); throw e; diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java index 0a8bfd2316d..672fa519905 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java @@ -20,16 +20,12 @@ import static com.google.common.base.Preconditions.checkNotNull; import com.google.api.core.ApiFuture; -import com.google.api.core.SettableApiFuture; import com.google.cloud.Timestamp; import com.google.cloud.spanner.AbstractReadContext.MultiUseReadOnlyTransaction; import com.google.cloud.spanner.AbstractReadContext.SingleReadContext; import com.google.cloud.spanner.AbstractReadContext.SingleUseReadOnlyTransaction; -<<<<<<< HEAD import com.google.cloud.spanner.SessionClient.SessionId; -======= import com.google.cloud.spanner.TransactionRunner.TransactionCallable; ->>>>>>> feat: session pool is non-blocking import com.google.cloud.spanner.TransactionRunnerImpl.TransactionContextImpl; import com.google.cloud.spanner.spi.v1.SpannerRpc; import com.google.common.collect.Lists; @@ -48,8 +44,6 @@ import java.util.Collection; import java.util.List; import java.util.Map; -import java.util.concurrent.ExecutionException; -import java.util.concurrent.Executor; import javax.annotation.Nullable; /** @@ -220,37 +214,10 @@ public TransactionRunner readWriteTransaction() { } @Override - public ApiFuture runAsync(final AsyncWork work, Executor executor) { - final SettableApiFuture res = SettableApiFuture.create(); - final TransactionRunner runner = + public AsyncRunner runAsync() { + return new AsyncRunnerImpl( setActive( - new TransactionRunnerImpl(this, spanner.getRpc(), spanner.getDefaultPrefetchChunks())); - executor.execute( - new Runnable() { - @Override - public void run() { - try { - R r = - runner.run( - new TransactionCallable() { - @Override - public R run(TransactionContext transaction) throws Exception { - try { - return work.doWorkAsync(transaction).get(); - } catch (ExecutionException e) { - throw SpannerExceptionFactory.newSpannerException(e.getCause()); - } catch (InterruptedException e) { - throw SpannerExceptionFactory.propagateInterrupt(e); - } - } - }); - res.set(r); - } catch (Throwable t) { - res.setException(t); - } - } - }); - return res; + new TransactionRunnerImpl(this, spanner.getRpc(), spanner.getDefaultPrefetchChunks()))); } @Override diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java index 9366488791e..b8e15ad4aba 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java @@ -40,7 +40,6 @@ import com.google.cloud.grpc.GrpcTransportOptions; import com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory; import com.google.cloud.spanner.AbstractReadContext.ConsumeSingleRowCallback; -import com.google.cloud.spanner.DatabaseClient.AsyncWork; import com.google.cloud.spanner.Options.QueryOption; import com.google.cloud.spanner.Options.ReadOption; import com.google.cloud.spanner.SessionClient.SessionConsumer; @@ -772,24 +771,18 @@ public TransactionRunner allowNestedTransaction() { } } - private static class SessionPoolAsyncRunner { + private static class SessionPoolAsyncRunner implements AsyncRunner { private final SessionPool sessionPool; private volatile PooledSessionFuture session; - private final AsyncWork work; - private final Executor executor; + private final SettableApiFuture commitTimestamp = SettableApiFuture.create(); - private SessionPoolAsyncRunner( - SessionPool sessionPool, - PooledSessionFuture session, - AsyncWork work, - Executor executor) { + private SessionPoolAsyncRunner(SessionPool sessionPool, PooledSessionFuture session) { this.sessionPool = sessionPool; this.session = session; - this.work = work; - this.executor = executor; } - private ApiFuture runAsync() { + @Override + public ApiFuture runAsync(final AsyncWork work, Executor executor) { final SettableApiFuture res = SettableApiFuture.create(); executor.execute( new Runnable() { @@ -797,9 +790,11 @@ private ApiFuture runAsync() { public void run() { SpannerException se = null; R r = null; + AsyncRunner runner = null; while (true) { try { - r = session.get().runAsync(work, MoreExecutors.directExecutor()).get(); + runner = session.get().runAsync(); + r = runner.runAsync(work, MoreExecutors.directExecutor()).get(); break; } catch (ExecutionException e) { se = SpannerExceptionFactory.newSpannerException(e.getCause()); @@ -818,6 +813,7 @@ public void run() { } session.get().markUsed(); session.close(); + setCommitTimestamp(runner); if (se != null) { res.setException(se); } else { @@ -827,6 +823,19 @@ public void run() { }); return res; } + + private void setCommitTimestamp(AsyncRunner delegate) { + try { + commitTimestamp.set(delegate.getCommitTimestamp().get()); + } catch (Throwable t) { + commitTimestamp.setException(t); + } + } + + @Override + public ApiFuture getCommitTimestamp() { + return commitTimestamp; + } } // Exception class used just to track the stack trace at the point when a session was handed out @@ -1000,8 +1009,8 @@ public TransactionManager transactionManager() { } @Override - public ApiFuture runAsync(AsyncWork work, Executor executor) { - return new SessionPoolAsyncRunner<>(SessionPool.this, this, work, executor).runAsync(); + public AsyncRunner runAsync() { + return new SessionPoolAsyncRunner(SessionPool.this, this); } @Override @@ -1155,8 +1164,8 @@ public TransactionRunner readWriteTransaction() { } @Override - public ApiFuture runAsync(AsyncWork work, Executor executor) { - return delegate.runAsync(work, executor); + public AsyncRunner runAsync() { + return delegate.runAsync(); } @Override diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncRunnerTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncRunnerTest.java index c3a14e7bc52..fbc1f2177df 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncRunnerTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncRunnerTest.java @@ -16,15 +16,22 @@ package com.google.cloud.spanner; +import static com.google.cloud.spanner.MockSpannerTestUtil.*; import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.fail; +import com.google.api.core.ApiFunction; import com.google.api.core.ApiFuture; -import com.google.api.core.SettableApiFuture; +import com.google.api.core.ApiFutures; import com.google.api.gax.grpc.testing.LocalChannelProvider; import com.google.cloud.NoCredentials; -import com.google.cloud.spanner.DatabaseClient.AsyncWork; +import com.google.cloud.Timestamp; +import com.google.cloud.spanner.AsyncRunner.AsyncWork; +import com.google.cloud.spanner.MockSpannerServiceImpl.SimulatedExecutionTime; import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult; +import com.google.common.base.Function; +import com.google.common.collect.ImmutableList; +import com.google.common.util.concurrent.MoreExecutors; import io.grpc.Server; import io.grpc.Status; import io.grpc.inprocess.InProcessServerBuilder; @@ -43,15 +50,6 @@ @RunWith(JUnit4.class) public class AsyncRunnerTest { - private static final String TEST_PROJECT = "my-project"; - private static final String TEST_INSTANCE = "my-instance"; - private static final String TEST_DATABASE = "my-database"; - - private static final Statement UPDATE_STATEMENT = - Statement.of("UPDATE FOO SET BAR=1 WHERE BAZ=2"); - private static final Statement INVALID_UPDATE_STATEMENT = - Statement.of("UPDATE NON_EXISTENT_TABLE SET BAR=1 WHERE BAZ=2"); - private static final long UPDATE_COUNT = 1L; private static MockSpannerServiceImpl mockSpanner; private static Server server; @@ -67,6 +65,13 @@ public class AsyncRunnerTest { public static void setup() throws Exception { mockSpanner = new MockSpannerServiceImpl(); mockSpanner.setAbortProbability(0.0D); + mockSpanner.putStatementResult( + StatementResult.query(READ_EMPTY_KEY_VALUE_STATEMENT, EMPTY_KEY_VALUE_RESULTSET)); + mockSpanner.putStatementResult( + StatementResult.query(READ_ONE_KEY_VALUE_STATEMENT, READ_ONE_KEY_VALUE_RESULTSET)); + mockSpanner.putStatementResult( + StatementResult.query( + READ_MULTIPLE_KEY_VALUE_STATEMENT, READ_MULTIPLE_KEY_VALUE_RESULTSET)); mockSpanner.putStatementResult(StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT)); mockSpanner.putStatementResult( StatementResult.exception( @@ -123,8 +128,9 @@ public void after() { @Test public void asyncRunnerUpdate() throws Exception { + AsyncRunner runner = client.runAsync(); ApiFuture updateCount = - client.runAsync( + runner.runAsync( new AsyncWork() { @Override public ApiFuture doWorkAsync(TransactionContext txn) { @@ -138,25 +144,27 @@ public ApiFuture doWorkAsync(TransactionContext txn) { @Test public void asyncRunnerIsNonBlocking() throws Exception { mockSpanner.freeze(); - final SettableApiFuture finished = SettableApiFuture.create(); + AsyncRunner runner = clientWithEmptySessionPool.runAsync(); ApiFuture res = - clientWithEmptySessionPool.runAsync( + runner.runAsync( new AsyncWork() { @Override public ApiFuture doWorkAsync(TransactionContext txn) { - finished.set(null); - return finished; + return ApiFutures.immediateFuture(null); } }, executor); + ApiFuture ts = runner.getCommitTimestamp(); mockSpanner.unfreeze(); assertThat(res.get()).isNull(); + assertThat(ts.get()).isNotNull(); } @Test public void asyncRunnerInvalidUpdate() throws Exception { + AsyncRunner runner = client.runAsync(); ApiFuture updateCount = - client.runAsync( + runner.runAsync( new AsyncWork() { @Override public ApiFuture doWorkAsync(TransactionContext txn) { @@ -181,8 +189,9 @@ public void asyncRunnerUpdateAborted() throws Exception { // Temporarily set the result of the update to 2 rows. mockSpanner.putStatementResult(StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT + 1L)); final AtomicInteger attempt = new AtomicInteger(); + AsyncRunner runner = client.runAsync(); ApiFuture updateCount = - client.runAsync( + runner.runAsync( new AsyncWork() { @Override public ApiFuture doWorkAsync(TransactionContext txn) { @@ -210,8 +219,9 @@ public void asyncRunnerCommitAborted() throws Exception { // Temporarily set the result of the update to 2 rows. mockSpanner.putStatementResult(StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT + 1L)); final AtomicInteger attempt = new AtomicInteger(); + AsyncRunner runner = client.runAsync(); ApiFuture updateCount = - client.runAsync( + runner.runAsync( new AsyncWork() { @Override public ApiFuture doWorkAsync(TransactionContext txn) { @@ -236,23 +246,105 @@ public ApiFuture doWorkAsync(TransactionContext txn) { @Test public void asyncRunnerUpdateAbortedWithoutGettingResult() throws Exception { - final SettableApiFuture finished = SettableApiFuture.create(); final AtomicInteger attempt = new AtomicInteger(); + AsyncRunner runner = client.runAsync(); ApiFuture result = - client.runAsync( + runner.runAsync( new AsyncWork() { @Override public ApiFuture doWorkAsync(TransactionContext txn) { if (attempt.incrementAndGet() == 1) { mockSpanner.abortTransaction(txn); } + // This update statement will be aborted, but the error will not propagated to the + // transaction runner and cause the transaction to retry. Instead, the commit call + // will do that. txn.executeUpdateAsync(UPDATE_STATEMENT); - finished.set(null); - return finished; + // Resolving this future will not resolve the result of the entire transaction. The + // transaction result will be resolved when the commit has actually finished + // successfully. + return ApiFutures.immediateFuture(null); } }, executor); assertThat(result.get()).isNull(); assertThat(attempt.get()).isEqualTo(2); } + + @Test + public void asyncRunnerCommitFails() throws Exception { + mockSpanner.setCommitExecutionTime( + SimulatedExecutionTime.ofException( + Status.RESOURCE_EXHAUSTED + .withDescription("mutation limit exceeded") + .asRuntimeException())); + AsyncRunner runner = client.runAsync(); + ApiFuture updateCount = + runner.runAsync( + new AsyncWork() { + @Override + public ApiFuture doWorkAsync(TransactionContext txn) { + // This statement will succeed, but the commit will fail. The error from the commit + // will bubble up to the future that is returned by the transaction, and the update + // count returned here will never reach the user application. + return txn.executeUpdateAsync(UPDATE_STATEMENT); + } + }, + executor); + try { + updateCount.get(); + fail("missing expected exception"); + } catch (ExecutionException e) { + assertThat(e.getCause()).isInstanceOf(SpannerException.class); + SpannerException se = (SpannerException) e.getCause(); + assertThat(se.getErrorCode()).isEqualTo(ErrorCode.RESOURCE_EXHAUSTED); + assertThat(se.getMessage()).contains("mutation limit exceeded"); + } + } + + @Test + public void asyncRunnerReadRow() throws Exception { + AsyncRunner runner = client.runAsync(); + ApiFuture val = + runner.runAsync( + new AsyncWork() { + @Override + public ApiFuture doWorkAsync(TransactionContext txn) { + return ApiFutures.transform( + txn.readRowAsync(READ_TABLE_NAME, Key.of(1L), READ_COLUMN_NAMES), + new ApiFunction() { + @Override + public String apply(Struct input) { + return input.getString("Value"); + } + }, + MoreExecutors.directExecutor()); + } + }, + executor); + assertThat(val.get()).isEqualTo("v1"); + } + + @Test + public void asyncRunnerRead() throws Exception { + AsyncRunner runner = client.runAsync(); + ApiFuture> val = + runner.runAsync( + new AsyncWork>() { + @Override + public ApiFuture> doWorkAsync(TransactionContext txn) { + return txn.readAsync(READ_TABLE_NAME, KeySet.all(), READ_COLUMN_NAMES) + .toListAsync( + new Function() { + @Override + public String apply(StructReader input) { + return input.getString("Value"); + } + }, + MoreExecutors.directExecutor()); + } + }, + executor); + assertThat(val.get()).containsExactly("v1", "v2", "v3"); + } } diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java index 6c90bba5111..7d7b919fc18 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java @@ -16,6 +16,7 @@ package com.google.cloud.spanner; +import static com.google.cloud.spanner.MockSpannerTestUtil.SELECT1; import static com.google.common.truth.Truth.assertThat; import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.CoreMatchers.is; @@ -23,13 +24,13 @@ import static org.junit.Assert.fail; import com.google.api.core.ApiFuture; -import com.google.api.core.SettableApiFuture; +import com.google.api.core.ApiFutures; import com.google.api.gax.grpc.testing.LocalChannelProvider; import com.google.api.gax.retrying.RetrySettings; import com.google.cloud.NoCredentials; import com.google.cloud.spanner.AsyncResultSet.CallbackResponse; import com.google.cloud.spanner.AsyncResultSet.ReadyCallback; -import com.google.cloud.spanner.DatabaseClient.AsyncWork; +import com.google.cloud.spanner.AsyncRunner.AsyncWork; import com.google.cloud.spanner.MockSpannerServiceImpl.SimulatedExecutionTime; import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult; import com.google.cloud.spanner.ReadContext.QueryAnalyzeMode; @@ -88,29 +89,6 @@ public class DatabaseClientImplTest { private static final Statement INVALID_UPDATE_STATEMENT = Statement.of("UPDATE NON_EXISTENT_TABLE SET BAR=1 WHERE BAZ=2"); private static final long UPDATE_COUNT = 1L; - private static final Statement SELECT1 = Statement.of("SELECT 1 AS COL1"); - private static final ResultSetMetadata SELECT1_METADATA = - ResultSetMetadata.newBuilder() - .setRowType( - StructType.newBuilder() - .addFields( - Field.newBuilder() - .setName("COL1") - .setType( - com.google.spanner.v1.Type.newBuilder() - .setCode(TypeCode.INT64) - .build()) - .build()) - .build()) - .build(); - private static final com.google.spanner.v1.ResultSet SELECT1_RESULTSET = - com.google.spanner.v1.ResultSet.newBuilder() - .addRows( - ListValue.newBuilder() - .addValues(com.google.protobuf.Value.newBuilder().setStringValue("1").build()) - .build()) - .setMetadata(SELECT1_METADATA) - .build(); private Spanner spanner; private Spanner spannerWithEmptySessionPool; @@ -121,7 +99,8 @@ public static void startStaticServer() throws IOException { mockSpanner = new MockSpannerServiceImpl(); mockSpanner.setAbortProbability(0.0D); // We don't want any unpredictable aborted transactions. mockSpanner.putStatementResult(StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT)); - mockSpanner.putStatementResult(StatementResult.query(SELECT1, SELECT1_RESULTSET)); + mockSpanner.putStatementResult( + StatementResult.query(SELECT1, MockSpannerTestUtil.SELECT1_RESULTSET)); mockSpanner.putStatementResult( StatementResult.exception( INVALID_UPDATE_STATEMENT, @@ -406,14 +385,13 @@ public void runAsync() throws Exception { DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); ExecutorService executor = Executors.newSingleThreadExecutor(); + AsyncRunner runner = client.runAsync(); ApiFuture fut = - client.runAsync( + runner.runAsync( new AsyncWork() { @Override public ApiFuture doWorkAsync(TransactionContext txn) { - SettableApiFuture res = SettableApiFuture.create(); - res.set(txn.executeUpdate(UPDATE_STATEMENT)); - return res; + return ApiFutures.immediateFuture(txn.executeUpdate(UPDATE_STATEMENT)); } }, executor); @@ -428,14 +406,13 @@ public void runAsyncIsNonBlocking() throws Exception { spannerWithEmptySessionPool.getDatabaseClient( DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); ExecutorService executor = Executors.newSingleThreadExecutor(); + AsyncRunner runner = client.runAsync(); ApiFuture fut = - client.runAsync( + runner.runAsync( new AsyncWork() { @Override public ApiFuture doWorkAsync(TransactionContext txn) { - SettableApiFuture res = SettableApiFuture.create(); - res.set(txn.executeUpdate(UPDATE_STATEMENT)); - return res; + return ApiFutures.immediateFuture(txn.executeUpdate(UPDATE_STATEMENT)); } }, executor); @@ -449,14 +426,13 @@ public void runAsyncWithException() throws Exception { DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); ExecutorService executor = Executors.newSingleThreadExecutor(); + AsyncRunner runner = client.runAsync(); ApiFuture fut = - client.runAsync( + runner.runAsync( new AsyncWork() { @Override public ApiFuture doWorkAsync(TransactionContext txn) { - SettableApiFuture res = SettableApiFuture.create(); - res.set(txn.executeUpdate(INVALID_UPDATE_STATEMENT)); - return res; + return ApiFutures.immediateFuture(txn.executeUpdate(INVALID_UPDATE_STATEMENT)); } }, executor); diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MockSpannerTestUtil.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MockSpannerTestUtil.java new file mode 100644 index 00000000000..00c296391a2 --- /dev/null +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MockSpannerTestUtil.java @@ -0,0 +1,124 @@ +/* + * 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.protobuf.ListValue; +import com.google.spanner.v1.ResultSetMetadata; +import com.google.spanner.v1.StructType; +import com.google.spanner.v1.StructType.Field; +import com.google.spanner.v1.TypeCode; +import java.util.Arrays; + +public class MockSpannerTestUtil { + static final Statement SELECT1 = Statement.of("SELECT 1 AS COL1"); + private static final ResultSetMetadata SELECT1_METADATA = + ResultSetMetadata.newBuilder() + .setRowType( + StructType.newBuilder() + .addFields( + Field.newBuilder() + .setName("COL1") + .setType( + com.google.spanner.v1.Type.newBuilder() + .setCode(TypeCode.INT64) + .build()) + .build()) + .build()) + .build(); + static final com.google.spanner.v1.ResultSet SELECT1_RESULTSET = + com.google.spanner.v1.ResultSet.newBuilder() + .addRows( + ListValue.newBuilder() + .addValues(com.google.protobuf.Value.newBuilder().setStringValue("1").build()) + .build()) + .setMetadata(SELECT1_METADATA) + .build(); + + static final String TEST_PROJECT = "my-project"; + static final String TEST_INSTANCE = "my-instance"; + static final String TEST_DATABASE = "my-database"; + + static final Statement UPDATE_STATEMENT = Statement.of("UPDATE FOO SET BAR=1 WHERE BAZ=2"); + static final Statement INVALID_UPDATE_STATEMENT = + Statement.of("UPDATE NON_EXISTENT_TABLE SET BAR=1 WHERE BAZ=2"); + static final long UPDATE_COUNT = 1L; + + static final String READ_TABLE_NAME = "TestTable"; + static final String EMPTY_READ_TABLE_NAME = "EmptyTestTable"; + static final Iterable READ_COLUMN_NAMES = Arrays.asList("Key", "Value"); + static final Statement READ_ONE_KEY_VALUE_STATEMENT = + Statement.of("SELECT Key, Value FROM TestTable WHERE ID=1"); + static final Statement READ_MULTIPLE_KEY_VALUE_STATEMENT = + Statement.of("SELECT Key, Value FROM TestTable WHERE 1=1"); + static final Statement READ_EMPTY_KEY_VALUE_STATEMENT = + Statement.of("SELECT Key, Value From EmptyTestTable"); + static final ResultSetMetadata READ_KEY_VALUE_METADATA = + ResultSetMetadata.newBuilder() + .setRowType( + StructType.newBuilder() + .addFields( + Field.newBuilder() + .setName("Key") + .setType( + com.google.spanner.v1.Type.newBuilder() + .setCode(TypeCode.STRING) + .build()) + .build()) + .addFields( + Field.newBuilder() + .setName("Value") + .setType( + com.google.spanner.v1.Type.newBuilder() + .setCode(TypeCode.STRING) + .build()) + .build()) + .build()) + .build(); + static final com.google.spanner.v1.ResultSet EMPTY_KEY_VALUE_RESULTSET = + com.google.spanner.v1.ResultSet.newBuilder() + .addRows(ListValue.newBuilder().build()) + .setMetadata(READ_KEY_VALUE_METADATA) + .build(); + static final com.google.spanner.v1.ResultSet READ_ONE_KEY_VALUE_RESULTSET = + com.google.spanner.v1.ResultSet.newBuilder() + .addRows( + ListValue.newBuilder() + .addValues(com.google.protobuf.Value.newBuilder().setStringValue("k1").build()) + .addValues(com.google.protobuf.Value.newBuilder().setStringValue("v1").build()) + .build()) + .setMetadata(READ_KEY_VALUE_METADATA) + .build(); + static final com.google.spanner.v1.ResultSet READ_MULTIPLE_KEY_VALUE_RESULTSET = + com.google.spanner.v1.ResultSet.newBuilder() + .addRows( + ListValue.newBuilder() + .addValues(com.google.protobuf.Value.newBuilder().setStringValue("k1").build()) + .addValues(com.google.protobuf.Value.newBuilder().setStringValue("v1").build()) + .build()) + .addRows( + ListValue.newBuilder() + .addValues(com.google.protobuf.Value.newBuilder().setStringValue("k2").build()) + .addValues(com.google.protobuf.Value.newBuilder().setStringValue("v2").build()) + .build()) + .addRows( + ListValue.newBuilder() + .addValues(com.google.protobuf.Value.newBuilder().setStringValue("k3").build()) + .addValues(com.google.protobuf.Value.newBuilder().setStringValue("v3").build()) + .build()) + .setMetadata(READ_KEY_VALUE_METADATA) + .build(); +} diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadAsyncTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadAsyncTest.java index 21442d3f2e5..06413ae51e9 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadAsyncTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadAsyncTest.java @@ -61,34 +61,6 @@ public class ReadAsyncTest { private static final String TEST_PROJECT = "my-project"; private static final String TEST_INSTANCE = "my-instance"; private static final String TEST_DATABASE = "my-database"; - private static final Statement UPDATE_STATEMENT = - Statement.of("UPDATE FOO SET BAR=1 WHERE BAZ=2"); - private static final Statement INVALID_UPDATE_STATEMENT = - Statement.of("UPDATE NON_EXISTENT_TABLE SET BAR=1 WHERE BAZ=2"); - private static final long UPDATE_COUNT = 1L; - private static final Statement SELECT1 = Statement.of("SELECT 1 AS COL1"); - private static final ResultSetMetadata SELECT1_METADATA = - ResultSetMetadata.newBuilder() - .setRowType( - StructType.newBuilder() - .addFields( - Field.newBuilder() - .setName("COL1") - .setType( - com.google.spanner.v1.Type.newBuilder() - .setCode(TypeCode.INT64) - .build()) - .build()) - .build()) - .build(); - private static final com.google.spanner.v1.ResultSet SELECT1_RESULTSET = - com.google.spanner.v1.ResultSet.newBuilder() - .addRows( - ListValue.newBuilder() - .addValues(com.google.protobuf.Value.newBuilder().setStringValue("1").build()) - .build()) - .setMetadata(SELECT1_METADATA) - .build(); private static MockSpannerServiceImpl mockSpanner; private static Server server; diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/RetryOnInvalidatedSessionTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/RetryOnInvalidatedSessionTest.java index 72d537f11ab..e386a463f27 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/RetryOnInvalidatedSessionTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/RetryOnInvalidatedSessionTest.java @@ -19,6 +19,7 @@ import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.fail; +import com.google.api.core.ApiFuture; import com.google.api.gax.core.NoCredentialsProvider; import com.google.api.gax.grpc.testing.LocalChannelProvider; import com.google.cloud.NoCredentials; @@ -28,7 +29,9 @@ import com.google.cloud.spanner.v1.SpannerClient; import com.google.cloud.spanner.v1.SpannerClient.ListSessionsPagedResponse; import com.google.cloud.spanner.v1.SpannerSettings; +import com.google.common.base.Function; import com.google.common.base.Stopwatch; +import com.google.common.collect.ImmutableList; import com.google.protobuf.ListValue; import com.google.spanner.v1.ResultSetMetadata; import com.google.spanner.v1.StructType; @@ -41,7 +44,11 @@ import java.util.Arrays; import java.util.Collection; import java.util.List; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; import java.util.concurrent.TimeUnit; +import org.hamcrest.Matchers; import org.junit.After; import org.junit.AfterClass; import org.junit.Before; @@ -56,6 +63,15 @@ @RunWith(Parameterized.class) public class RetryOnInvalidatedSessionTest { + private static final class ToLongTransformer implements Function { + @Override + public Long apply(StructReader input) { + return input.getLong(0); + } + } + + private static final ToLongTransformer TO_LONG = new ToLongTransformer(); + @Rule public ExpectedException expected = ExpectedException.none(); @Parameter(0) @@ -141,6 +157,7 @@ public static Collection data() { private static SpannerClient spannerClient; private static Spanner spanner; private static DatabaseClient client; + private static ExecutorService executor; @BeforeClass public static void startStaticServer() throws IOException { @@ -169,6 +186,7 @@ public static void startStaticServer() throws IOException { .setCredentialsProvider(NoCredentialsProvider.create()) .build(); spannerClient = SpannerClient.create(settings); + executor = Executors.newSingleThreadExecutor(); } @AfterClass @@ -176,13 +194,16 @@ public static void stopServer() throws InterruptedException { spannerClient.close(); server.shutdown(); server.awaitTermination(); + executor.shutdown(); } @Before public void setUp() throws IOException { mockSpanner.reset(); SessionPoolOptions.Builder builder = - SessionPoolOptions.newBuilder().setWriteSessionsFraction(WRITE_SESSIONS_FRACTION); + SessionPoolOptions.newBuilder() + .setWriteSessionsFraction(WRITE_SESSIONS_FRACTION) + .setFailOnSessionLeak(); if (failOnInvalidatedSession) { builder.setFailIfSessionNotFound(); } @@ -254,6 +275,20 @@ public void singleUseSelect() throws InterruptedException { assertThat(count).isEqualTo(2); } + @Test + public void singleUseSelectAsync() throws Exception { + if (failOnInvalidatedSession) { + expected.expect(ExecutionException.class); + expected.expectCause(Matchers.instanceOf(SessionNotFoundException.class)); + } + invalidateSessionPool(); + ApiFuture> list; + try (AsyncResultSet rs = client.singleUse().executeQueryAsync(SELECT1AND2)) { + list = rs.toListAsync(TO_LONG, executor); + } + assertThat(list.get()).containsExactly(1L, 2L); + } + @Test public void singleUseRead() throws InterruptedException { if (failOnInvalidatedSession) { From 91253cf7dbd943a5162fab209bdb54607b321bdf Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Wed, 26 Feb 2020 09:48:09 +0100 Subject: [PATCH 07/49] tests: centralize some commonly used test objects --- .../google/cloud/spanner/AsyncRunnerTest.java | 2 +- .../cloud/spanner/MockSpannerTestUtil.java | 10 ++- .../google/cloud/spanner/ReadAsyncTest.java | 77 +++---------------- 3 files changed, 18 insertions(+), 71 deletions(-) diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncRunnerTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncRunnerTest.java index fbc1f2177df..5782ed8e663 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncRunnerTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncRunnerTest.java @@ -66,7 +66,7 @@ public static void setup() throws Exception { mockSpanner = new MockSpannerServiceImpl(); mockSpanner.setAbortProbability(0.0D); mockSpanner.putStatementResult( - StatementResult.query(READ_EMPTY_KEY_VALUE_STATEMENT, EMPTY_KEY_VALUE_RESULTSET)); + StatementResult.query(READ_ONE_EMPTY_KEY_VALUE_STATEMENT, EMPTY_KEY_VALUE_RESULTSET)); mockSpanner.putStatementResult( StatementResult.query(READ_ONE_KEY_VALUE_STATEMENT, READ_ONE_KEY_VALUE_RESULTSET)); mockSpanner.putStatementResult( diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MockSpannerTestUtil.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MockSpannerTestUtil.java index 00c296391a2..336ae9d70e0 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MockSpannerTestUtil.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MockSpannerTestUtil.java @@ -16,6 +16,7 @@ package com.google.cloud.spanner; +import com.google.cloud.spanner.Type.StructField; import com.google.protobuf.ListValue; import com.google.spanner.v1.ResultSetMetadata; import com.google.spanner.v1.StructType; @@ -64,8 +65,10 @@ public class MockSpannerTestUtil { Statement.of("SELECT Key, Value FROM TestTable WHERE ID=1"); static final Statement READ_MULTIPLE_KEY_VALUE_STATEMENT = Statement.of("SELECT Key, Value FROM TestTable WHERE 1=1"); - static final Statement READ_EMPTY_KEY_VALUE_STATEMENT = - Statement.of("SELECT Key, Value From EmptyTestTable"); + static final Statement READ_ONE_EMPTY_KEY_VALUE_STATEMENT = + Statement.of("SELECT Key, Value FROM EmptyTestTable WHERE ID=1"); + static final Statement READ_ALL_EMPTY_KEY_VALUE_STATEMENT = + Statement.of("SELECT Key, Value FROM EmptyTestTable WHERE 1=1"); static final ResultSetMetadata READ_KEY_VALUE_METADATA = ResultSetMetadata.newBuilder() .setRowType( @@ -88,6 +91,9 @@ public class MockSpannerTestUtil { .build()) .build()) .build(); + static final Type READ_TABLE_TYPE = + Type.struct( + StructField.of("Key", Type.string()), StructField.of("Value", Type.string())); static final com.google.spanner.v1.ResultSet EMPTY_KEY_VALUE_RESULTSET = com.google.spanner.v1.ResultSet.newBuilder() .addRows(ListValue.newBuilder().build()) diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadAsyncTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadAsyncTest.java index 06413ae51e9..a8410410b48 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadAsyncTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadAsyncTest.java @@ -16,6 +16,7 @@ package com.google.cloud.spanner; +import static com.google.cloud.spanner.MockSpannerTestUtil.*; import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.fail; @@ -26,18 +27,10 @@ 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.Type.StructField; import com.google.common.util.concurrent.SettableFuture; -import com.google.protobuf.ListValue; -import com.google.spanner.v1.ResultSetMetadata; -import com.google.spanner.v1.StructType; -import com.google.spanner.v1.StructType.Field; -import com.google.spanner.v1.TypeCode; import io.grpc.Server; import io.grpc.Status; import io.grpc.inprocess.InProcessServerBuilder; -import java.util.Arrays; -import java.util.List; import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; @@ -52,59 +45,9 @@ @RunWith(JUnit4.class) public class ReadAsyncTest { - private static final String EMPTY_TABLE_NAME = "EmptyTestTable"; - private static final String TABLE_NAME = "TestTable"; - private static final List ALL_COLUMNS = Arrays.asList("Key", "StringValue"); - private static final Type TABLE_TYPE = - Type.struct( - StructField.of("Key", Type.string()), StructField.of("StringValue", Type.string())); - private static final String TEST_PROJECT = "my-project"; - private static final String TEST_INSTANCE = "my-instance"; - private static final String TEST_DATABASE = "my-database"; - private static MockSpannerServiceImpl mockSpanner; private static Server server; private static LocalChannelProvider channelProvider; - private static final Statement EMPTY_READ_STATEMENT = - Statement.of("SELECT Key, StringValue FROM EmptyTestTable WHERE ID=1"); - private static final Statement READ1_STATEMENT = - Statement.of("SELECT Key, StringValue FROM TestTable WHERE ID=1"); - private static final ResultSetMetadata READ1_METADATA = - ResultSetMetadata.newBuilder() - .setRowType( - StructType.newBuilder() - .addFields( - Field.newBuilder() - .setName("Key") - .setType( - com.google.spanner.v1.Type.newBuilder() - .setCode(TypeCode.STRING) - .build()) - .build()) - .addFields( - Field.newBuilder() - .setName("StringValue") - .setType( - com.google.spanner.v1.Type.newBuilder() - .setCode(TypeCode.STRING) - .build()) - .build()) - .build()) - .build(); - private static final com.google.spanner.v1.ResultSet EMPTY_RESULTSET = - com.google.spanner.v1.ResultSet.newBuilder() - .addRows(ListValue.newBuilder().build()) - .setMetadata(READ1_METADATA) - .build(); - private static final com.google.spanner.v1.ResultSet READ1_RESULTSET = - com.google.spanner.v1.ResultSet.newBuilder() - .addRows( - ListValue.newBuilder() - .addValues(com.google.protobuf.Value.newBuilder().setStringValue("k1").build()) - .addValues(com.google.protobuf.Value.newBuilder().setStringValue("v1").build()) - .build()) - .setMetadata(READ1_METADATA) - .build(); private static ExecutorService executor; private Spanner spanner; @@ -113,8 +56,8 @@ public class ReadAsyncTest { @BeforeClass public static void setup() throws Exception { mockSpanner = new MockSpannerServiceImpl(); - mockSpanner.putStatementResult(StatementResult.query(READ1_STATEMENT, READ1_RESULTSET)); - mockSpanner.putStatementResult(StatementResult.query(EMPTY_READ_STATEMENT, EMPTY_RESULTSET)); + mockSpanner.putStatementResult(StatementResult.query(READ_ONE_KEY_VALUE_STATEMENT, READ_ONE_KEY_VALUE_RESULTSET)); + mockSpanner.putStatementResult(StatementResult.query(READ_ONE_EMPTY_KEY_VALUE_STATEMENT, EMPTY_KEY_VALUE_RESULTSET)); String uniqueName = InProcessServerBuilder.generateName(); server = @@ -159,7 +102,7 @@ public void emptyReadAsync() throws Exception { AsyncResultSet resultSet = client .singleUse(TimestampBound.strong()) - .readAsync(EMPTY_TABLE_NAME, KeySet.singleKey(Key.of("k99")), ALL_COLUMNS); + .readAsync(EMPTY_READ_TABLE_NAME, KeySet.singleKey(Key.of("k99")), READ_COLUMN_NAMES); resultSet.setCallback( executor, new ReadyCallback() { @@ -173,7 +116,7 @@ public CallbackResponse cursorReady(AsyncResultSet resultSet) { case NOT_READY: return CallbackResponse.CONTINUE; case DONE: - assertThat(resultSet.getType()).isEqualTo(TABLE_TYPE); + assertThat(resultSet.getType()).isEqualTo(READ_TABLE_TYPE); result.set(true); return CallbackResponse.DONE; } @@ -192,12 +135,10 @@ public void pointReadAsync() throws Exception { ApiFuture row = client .singleUse(TimestampBound.strong()) - .readRowAsync(TABLE_NAME, Key.of("k1"), ALL_COLUMNS); + .readRowAsync(READ_TABLE_NAME, Key.of("k1"), READ_COLUMN_NAMES); assertThat(row.get()).isNotNull(); assertThat(row.get().getString(0)).isEqualTo("k1"); assertThat(row.get().getString(1)).isEqualTo("v1"); - assertThat(row.get()) - .isEqualTo(Struct.newBuilder().set("Key").to("k1").set("StringValue").to("v1").build()); } @Test @@ -205,7 +146,7 @@ public void pointReadNotFound() throws Exception { ApiFuture row = client .singleUse(TimestampBound.strong()) - .readRowAsync(EMPTY_TABLE_NAME, Key.of("k999"), ALL_COLUMNS); + .readRowAsync(EMPTY_READ_TABLE_NAME, Key.of("k999"), READ_COLUMN_NAMES); assertThat(row.get()).isNull(); } @@ -218,7 +159,7 @@ public void invalidDatabase() throws Exception { ApiFuture row = invalidClient .singleUse(TimestampBound.strong()) - .readRowAsync(TABLE_NAME, Key.of("k99"), ALL_COLUMNS); + .readRowAsync(READ_TABLE_NAME, Key.of("k99"), READ_COLUMN_NAMES); try { row.get(); fail("missing expected exception"); @@ -237,7 +178,7 @@ public void tableNotFound() throws Exception { ApiFuture row = client .singleUse(TimestampBound.strong()) - .readRowAsync("BadTableName", Key.of("k1"), ALL_COLUMNS); + .readRowAsync("BadTableName", Key.of("k1"), READ_COLUMN_NAMES); try { row.get(); fail("missing expected exception"); From a2d28cd6601706fd3f5b32a35dc4d7745eb417b8 Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Wed, 26 Feb 2020 16:47:52 +0100 Subject: [PATCH 08/49] feat: keep session checked out until async finishes --- .../cloud/spanner/AsyncResultSetImpl.java | 9 ++ .../com/google/cloud/spanner/Options.java | 34 ++++ .../com/google/cloud/spanner/SessionPool.java | 152 +++++++++++++----- .../cloud/spanner/MockSpannerTestUtil.java | 3 +- .../google/cloud/spanner/ReadAsyncTest.java | 82 +++++++++- 5 files changed, 238 insertions(+), 42 deletions(-) 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 87891f33010..21f9094b24d 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 @@ -154,6 +154,13 @@ public void close() { } } + /** + * Called when no more rows will be read from the underlying {@link ResultSet}, either because all + * rows have been read, or because {@link ReadyCallback#cursorReady(AsyncResultSet)} returned + * {@link CallbackResponse#DONE}. + */ + void onFinished() {} + /** * Tries to advance this {@link AsyncResultSet} to the next row. This method may only be called * from within a {@link ReadyCallback}. @@ -331,6 +338,8 @@ public Void call() throws Exception { try { delegateResultSet.close(); } catch (Throwable t) { + } finally { + onFinished(); } // Ensure that the callback has been called at least once, even if the result set was diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/Options.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/Options.java index d193ad1c75c..879b632d175 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/Options.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/Options.java @@ -59,6 +59,11 @@ public static ReadAndQueryOption prefetchChunks(int prefetchChunks) { return new FlowControlOption(prefetchChunks); } + public static ReadAndQueryOption bufferRows(int bufferRows) { + Preconditions.checkArgument(bufferRows > 0, "bufferRows should be greater than 0"); + return new BufferRowsOption(bufferRows); + } + /** * Specifying this will cause the list operations to fetch at most this many records in a page. */ @@ -115,8 +120,22 @@ void appendToOptions(Options options) { } } + static final class BufferRowsOption extends InternalOption implements ReadAndQueryOption { + final int bufferRows; + + BufferRowsOption(int bufferRows) { + this.bufferRows = bufferRows; + } + + @Override + void appendToOptions(Options options) { + options.bufferRows = bufferRows; + } + } + private Long limit; private Integer prefetchChunks; + private Integer bufferRows; private Integer pageSize; private String pageToken; private String filter; @@ -140,6 +159,14 @@ int prefetchChunks() { return prefetchChunks; } + boolean hasBufferRows() { + return bufferRows != null; + } + + int bufferRows() { + return bufferRows; + } + boolean hasPageSize() { return pageSize != null; } @@ -203,6 +230,10 @@ public boolean equals(Object o) { || hasPrefetchChunks() && that.hasPrefetchChunks() && Objects.equals(prefetchChunks(), that.prefetchChunks())) + && (!hasBufferRows() && !that.hasBufferRows() + || hasBufferRows() + && that.hasBufferRows() + && Objects.equals(bufferRows(), that.bufferRows())) && (!hasPageSize() && !that.hasPageSize() || hasPageSize() && that.hasPageSize() && Objects.equals(pageSize(), that.pageSize())) && Objects.equals(pageToken(), that.pageToken()) @@ -218,6 +249,9 @@ public int hashCode() { if (prefetchChunks != null) { result = 31 * result + prefetchChunks.hashCode(); } + if (bufferRows != null) { + result = 31 * result + bufferRows.hashCode(); + } if (pageSize != null) { result = 31 * result + pageSize.hashCode(); } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java index b8e15ad4aba..c921eea51fe 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java @@ -36,6 +36,7 @@ import com.google.api.core.ApiFuture; import com.google.api.core.ApiFutures; import com.google.api.core.SettableApiFuture; +import com.google.api.gax.core.ExecutorProvider; import com.google.cloud.Timestamp; import com.google.cloud.grpc.GrpcTransportOptions; import com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory; @@ -96,6 +97,7 @@ import java.util.concurrent.ScheduledFuture; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicLong; import java.util.logging.Level; import java.util.logging.Logger; @@ -149,14 +151,51 @@ public ResultSet get() { * finished, if it is a single use context. */ private static class AutoClosingReadContext implements ReadContext { + private class AutoClosingReadContextAsyncResultSetImpl extends AsyncResultSetImpl { + private AutoClosingReadContextAsyncResultSetImpl( + ExecutorProvider executorProvider, ResultSet delegate, int bufferRows) { + super(executorProvider, delegate, bufferRows); + } + + @Override + public void setCallback(Executor exec, ReadyCallback cb) { + asyncOperationsCount.incrementAndGet(); + super.setCallback(exec, cb); + } + + @Override + void onFinished() { + synchronized (lock) { + if (asyncOperationsCount.decrementAndGet() == 0) { + if (closed) { + // All async operations for this read context have finished. + AutoClosingReadContext.this.close(); + } + } + } + } + } + private final Function readContextDelegateSupplier; private T readContextDelegate; private final SessionPool sessionPool; - private PooledSessionFuture session; private final boolean isSingleUse; - private boolean closed; + private final AtomicInteger asyncOperationsCount = new AtomicInteger(); + + private Object lock = new Object(); + + @GuardedBy("lock") private boolean sessionUsedForQuery = false; + @GuardedBy("lock") + private PooledSessionFuture session; + + @GuardedBy("lock") + private boolean closed; + + @GuardedBy("lock") + private boolean delegateClosed; + private AutoClosingReadContext( Function delegateSupplier, SessionPool sessionPool, @@ -170,12 +209,14 @@ private AutoClosingReadContext( T getReadContextDelegate() { if (readContextDelegate == null) { - while (true) { - try { - this.readContextDelegate = readContextDelegateSupplier.apply(this.session); - break; - } catch (SessionNotFoundException e) { - replaceSessionIfPossible(e); + synchronized (lock) { + while (true) { + try { + this.readContextDelegate = readContextDelegateSupplier.apply(this.session); + break; + } catch (SessionNotFoundException e) { + replaceSessionIfPossible(e); + } } } } @@ -212,9 +253,11 @@ private boolean internalNext() { try { boolean ret = super.next(); if (beforeFirst) { - session.get().markUsed(); - beforeFirst = false; - sessionUsedForQuery = true; + synchronized (lock) { + session.get().markUsed(); + beforeFirst = false; + sessionUsedForQuery = true; + } } if (!ret && isSingleUse) { close(); @@ -223,9 +266,11 @@ private boolean internalNext() { } catch (SessionNotFoundException e) { throw e; } catch (SpannerException e) { - if (!closed && isSingleUse) { - session.get().lastException = e; - AutoClosingReadContext.this.close(); + synchronized (lock) { + if (!closed && isSingleUse) { + session.get().lastException = e; + AutoClosingReadContext.this.close(); + } } throw e; } @@ -242,13 +287,15 @@ public void close() { } private void replaceSessionIfPossible(SessionNotFoundException notFound) { - if (isSingleUse || !sessionUsedForQuery) { - // This class is only used by read-only transactions, so we know that we only need a - // read-only session. - session = sessionPool.replaceReadSession(notFound, session); - readContextDelegate = readContextDelegateSupplier.apply(session); - } else { - throw notFound; + synchronized (lock) { + if (isSingleUse || !sessionUsedForQuery) { + // This class is only used by read-only transactions, so we know that we only need a + // read-only session. + session = sessionPool.replaceReadSession(notFound, session); + readContextDelegate = readContextDelegateSupplier.apply(session); + } else { + throw notFound; + } } } @@ -273,7 +320,9 @@ public AsyncResultSet readAsync( final KeySet keys, final Iterable columns, final ReadOption... options) { - return new AsyncResultSetImpl( + Options readOptions = Options.fromReadOptions(options); + final int bufferRows = readOptions.hasBufferRows() ? readOptions.bufferRows() : 10; + return new AutoClosingReadContextAsyncResultSetImpl( sessionPool.sessionClient.getSpanner().getAsyncExecutorProvider(), wrap( new CachedResultSetSupplier() { @@ -281,7 +330,8 @@ public AsyncResultSet readAsync( ResultSet load() { return getReadContextDelegate().read(table, keys, columns, options); } - })); + }), + bufferRows); } @Override @@ -307,7 +357,9 @@ public AsyncResultSet readUsingIndexAsync( final KeySet keys, final Iterable columns, final ReadOption... options) { - return new AsyncResultSetImpl( + Options readOptions = Options.fromReadOptions(options); + final int bufferRows = readOptions.hasBufferRows() ? readOptions.bufferRows() : 10; + return new AutoClosingReadContextAsyncResultSetImpl( sessionPool.sessionClient.getSpanner().getAsyncExecutorProvider(), wrap( new CachedResultSetSupplier() { @@ -316,7 +368,8 @@ ResultSet load() { return getReadContextDelegate() .readUsingIndex(table, index, keys, columns, options); } - })); + }), + bufferRows); } @Override @@ -325,14 +378,18 @@ public Struct readRow(String table, Key key, Iterable columns) { try { while (true) { try { - session.get().markUsed(); + synchronized (lock) { + session.get().markUsed(); + } return getReadContextDelegate().readRow(table, key, columns); } catch (SessionNotFoundException e) { replaceSessionIfPossible(e); } } } finally { - sessionUsedForQuery = true; + synchronized (lock) { + sessionUsedForQuery = true; + } if (isSingleUse) { close(); } @@ -354,14 +411,18 @@ public Struct readRowUsingIndex(String table, String index, Key key, Iterable results = new SynchronousQueue<>(); + final SettableApiFuture finished = SettableApiFuture.create(); + DatabaseClientImpl clientImpl = (DatabaseClientImpl) client; + + // There should currently not be any sessions checked out of the pool. + assertThat(clientImpl.pool.getNumberOfSessionsInUse()).isEqualTo(0); + + final CountDownLatch dataReceived = new CountDownLatch(1); + try (ReadOnlyTransaction tx = client.readOnlyTransaction()) { + try (AsyncResultSet rs = + tx.readAsync(READ_TABLE_NAME, KeySet.all(), READ_COLUMN_NAMES, Options.bufferRows(1))) { + rs.setCallback( + executor, + new ReadyCallback() { + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + try { + while (true) { + switch (resultSet.tryNext()) { + case DONE: + finished.set(true); + return CallbackResponse.DONE; + case NOT_READY: + return CallbackResponse.CONTINUE; + case OK: + dataReceived.countDown(); + results.put(resultSet.getString(0)); + } + } + } catch (Throwable t) { + finished.setException(t); + return CallbackResponse.DONE; + } + } + }); + } + // Wait until at least one row has been fetched. At that moment there should be one session + // checked out. + dataReceived.await(); + assertThat(clientImpl.pool.getNumberOfSessionsInUse()).isEqualTo(1); + } + // The read-only transaction is now closed, but the ready callback will continue to receive + // data. As it tries to put the data into a synchronous queue and the underlying buffer can also + // only hold 1 row, the async result set has not yet finished. The read-only transaction will + // release the session back into the pool when all async statements have finished. The number of + // sessions in use is therefore still 1. + assertThat(clientImpl.pool.getNumberOfSessionsInUse()).isEqualTo(1); + List resultList = new ArrayList<>(); + do { + results.drainTo(resultList); + } while (!finished.isDone() || results.size() > 0); + assertThat(finished.get()).isTrue(); + assertThat(resultList).containsExactly("k1", "k2", "k3"); + // The session will be released back into the pool by the asynchronous result set when it has + // returned all rows. As this is done in the background, it could take a couple of milliseconds. + Thread.sleep(10L); + assertThat(clientImpl.pool.getNumberOfSessionsInUse()).isEqualTo(0); + } } From 2a63e62560a1c7b266fc896f1133b03925b69875 Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Wed, 26 Feb 2020 19:01:07 +0100 Subject: [PATCH 09/49] fix: fix span test cases after rebase --- .../clirr-ignored-differences.xml | 38 +++++++++++++++- .../cloud/spanner/AbstractReadContext.java | 8 +++- .../google/cloud/spanner/BatchClientImpl.java | 3 ++ .../spanner/PartitionedDMLTransaction.java | 4 ++ .../com/google/cloud/spanner/SessionImpl.java | 13 +++++- .../com/google/cloud/spanner/SessionPool.java | 5 ++- .../cloud/spanner/TransactionManagerImpl.java | 11 +++-- .../cloud/spanner/TransactionRunnerImpl.java | 8 +++- .../google/cloud/spanner/SessionImplTest.java | 2 + .../google/cloud/spanner/SessionPoolTest.java | 3 ++ .../com/google/cloud/spanner/SpanTest.java | 43 ++++++++++--------- .../spanner/TransactionContextImplTest.java | 1 + .../spanner/TransactionManagerImplTest.java | 3 +- .../spanner/TransactionRunnerImplTest.java | 3 ++ 14 files changed, 113 insertions(+), 32 deletions(-) diff --git a/google-cloud-spanner/clirr-ignored-differences.xml b/google-cloud-spanner/clirr-ignored-differences.xml index a8afa4b642f..d1096aaa6a1 100644 --- a/google-cloud-spanner/clirr-ignored-differences.xml +++ b/google-cloud-spanner/clirr-ignored-differences.xml @@ -93,7 +93,6 @@ com/google/cloud/spanner/DatabaseAdminClient com.google.cloud.spanner.Backup updateBackup(java.lang.String, java.lang.String, com.google.cloud.Timestamp) - 7012 com/google/cloud/spanner/spi/v1/SpannerRpc @@ -147,4 +146,41 @@ com.google.api.gax.paging.Page listDatabases() + + + 7012 + com/google/cloud/spanner/DatabaseClient + * runAsync(*) + + + 7012 + com/google/cloud/spanner/ReadContext + * executeQueryAsync(*) + + + 7012 + com/google/cloud/spanner/ReadContext + * readAsync(*) + + + 7012 + com/google/cloud/spanner/ReadContext + * readRowAsync(*) + + + 7012 + com/google/cloud/spanner/ReadContext + * readUsingIndexAsync(*) + + + 7012 + com/google/cloud/spanner/ReadContext + * readRowUsingIndexAsync(*) + + + 7012 + com/google/cloud/spanner/TransactionContext + * executeUpdateAsync(*) + + diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java index d395345ace7..f4a12970358 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java @@ -49,7 +49,6 @@ import com.google.spanner.v1.TransactionOptions; import com.google.spanner.v1.TransactionSelector; import io.opencensus.trace.Span; -import io.opencensus.trace.Tracing; import java.util.Map; import java.util.concurrent.atomic.AtomicLong; import javax.annotation.Nullable; @@ -357,7 +356,7 @@ void initTransaction() { final SessionImpl session; final SpannerRpc rpc; final ExecutorProvider executorProvider; - final Span span; + Span span; private final int defaultPrefetchChunks; private final QueryOptions defaultQueryOptions; @@ -383,6 +382,11 @@ void initTransaction() { this.span = builder.span; } + @Override + public void setSpan(Span span) { + this.span = span; + } + long getSeqNo() { return seqNo.incrementAndGet(); } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/BatchClientImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/BatchClientImpl.java index 43de2be092f..39827647b6b 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/BatchClientImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/BatchClientImpl.java @@ -30,6 +30,7 @@ import com.google.spanner.v1.PartitionReadRequest; import com.google.spanner.v1.PartitionResponse; import com.google.spanner.v1.TransactionSelector; +import io.opencensus.trace.Tracing; import java.util.List; import java.util.Map; @@ -81,6 +82,7 @@ private static class BatchReadOnlyTransactionImpl extends MultiUseReadOnlyTransa super(builder.setTimestampBound(bound)); this.sessionName = session.getName(); this.options = session.getOptions(); + setSpan(Tracing.getTracer().getCurrentSpan()); initTransaction(); } @@ -89,6 +91,7 @@ private static class BatchReadOnlyTransactionImpl extends MultiUseReadOnlyTransa super(builder.setTransactionId(batchTransactionId.getTransactionId())); this.sessionName = session.getName(); this.options = session.getOptions(); + setSpan(Tracing.getTracer().getCurrentSpan()); } @Override diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/PartitionedDMLTransaction.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/PartitionedDMLTransaction.java index ded74ce85a1..351b7596287 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/PartitionedDMLTransaction.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/PartitionedDMLTransaction.java @@ -27,6 +27,7 @@ import com.google.spanner.v1.Transaction; import com.google.spanner.v1.TransactionOptions; import com.google.spanner.v1.TransactionSelector; +import io.opencensus.trace.Span; import java.util.Map; import java.util.concurrent.Callable; import org.threeten.bp.Duration; @@ -101,4 +102,7 @@ public com.google.spanner.v1.ResultSet call() throws Exception { public void invalidate() { isValid = false; } + + @Override + public void setSpan(Span span) {} } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java index 672fa519905..1d9c99eeb59 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java @@ -77,6 +77,8 @@ static void throwIfTransactionsPending() { static interface SessionTransaction { /** Invalidates the transaction, generally because a new one has been started on the session. */ void invalidate(); + /** Registers the current span on the transaction. */ + void setSpan(Span span); } private final SpannerImpl spanner; @@ -85,6 +87,7 @@ static interface SessionTransaction { private SessionTransaction activeTransaction; private ByteString readyTransactionId; private final Map options; + private Span currentSpan; SessionImpl(SpannerImpl spanner, String name, Map options) { this.spanner = spanner; @@ -102,6 +105,10 @@ public String getName() { return options; } + void setCurrentSpan(Span span) { + currentSpan = span; + } + @Override public long executePartitionedUpdate(Statement stmt) { setActive(null); @@ -273,6 +280,7 @@ TransactionContextImpl newTransaction() { .setRpc(spanner.getRpc()) .setDefaultQueryOptions(spanner.getDefaultQueryOptions(databaseId)) .setDefaultPrefetchChunks(spanner.getDefaultPrefetchChunks()) + .setSpan(currentSpan) .build(); } @@ -284,11 +292,14 @@ T setActive(@Nullable T ctx) { } activeTransaction = ctx; readyTransactionId = null; + if (activeTransaction != null) { + activeTransaction.setSpan(currentSpan); + } return ctx; } @Override public TransactionManager transactionManager() { - return new TransactionManagerImpl(this); + return new TransactionManagerImpl(this, currentSpan); } } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java index c921eea51fe..c3b32351ad0 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java @@ -1125,7 +1125,7 @@ public PooledSession get() { try { PooledSession res = super.get(); synchronized (lock) { - res.markBusy(); + res.markBusy(span); span.addAnnotation(sessionAnnotation(res)); incrementNumSessionsInUse(); checkedOutSessions.add(this); @@ -1291,7 +1291,8 @@ private void keepAlive() { } } - private void markBusy() { + private void markBusy(Span span) { + this.delegate.setCurrentSpan(span); this.state = SessionState.BUSY; } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionManagerImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionManagerImpl.java index bdf7ec954f0..35184cdf9c9 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionManagerImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionManagerImpl.java @@ -29,14 +29,19 @@ final class TransactionManagerImpl implements TransactionManager, SessionTransac private static final Tracer tracer = Tracing.getTracer(); private final SessionImpl session; - private final Span span; + private Span span; private TransactionRunnerImpl.TransactionContextImpl txn; private TransactionState txnState; - TransactionManagerImpl(SessionImpl session) { + TransactionManagerImpl(SessionImpl session, Span span) { this.session = session; - this.span = Tracing.getTracer().getCurrentSpan(); + this.span = span; + } + + @Override + public void setSpan(Span span) { + this.span = span; } @Override 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 6aa15134112..d292ad52281 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 @@ -335,7 +335,7 @@ public long[] batchUpdate(Iterable statements) { private boolean blockNestedTxn = true; private final SessionImpl session; - private final Span span; + private Span span; private TransactionContextImpl txn; private volatile boolean isValid = true; @@ -347,10 +347,14 @@ public TransactionRunner allowNestedTransaction() { TransactionRunnerImpl(SessionImpl session, SpannerRpc rpc, int defaultPrefetchChunks) { this.session = session; - this.span = Tracing.getTracer().getCurrentSpan(); this.txn = session.newTransaction(); } + @Override + public void setSpan(Span span) { + this.span = span; + } + @Nullable @Override public T run(TransactionCallable callable) { diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionImplTest.java index f3f205a5294..cc21774daea 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionImplTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionImplTest.java @@ -40,6 +40,7 @@ import com.google.spanner.v1.ResultSetMetadata; import com.google.spanner.v1.Session; import com.google.spanner.v1.Transaction; +import io.opencensus.trace.Span; import java.text.ParseException; import java.util.Arrays; import java.util.Calendar; @@ -111,6 +112,7 @@ public void setUp() { Mockito.when(rpc.commit(Mockito.any(CommitRequest.class), Mockito.any(Map.class))) .thenReturn(commitResponse); session = spanner.getSessionClient(db).createSession(); + ((SessionImpl) session).setCurrentSpan(mock(Span.class)); // We expect the same options, "options", on all calls on "session". options = optionsCaptor.getValue(); } diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java index 8735a2eece5..54250420833 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java @@ -60,6 +60,7 @@ import com.google.spanner.v1.RollbackRequest; import io.opencensus.metrics.LabelValue; import io.opencensus.metrics.MetricRegistry; +import io.opencensus.trace.Span; import java.util.ArrayList; import java.util.Arrays; import java.util.Collection; @@ -1219,6 +1220,7 @@ public void testSessionNotFoundReadWriteTransaction() { when(closedSession.beginTransaction()).thenThrow(sessionNotFound); TransactionRunnerImpl closedTransactionRunner = new TransactionRunnerImpl(closedSession, rpc, 10); + closedTransactionRunner.setSpan(mock(Span.class)); when(closedSession.readWriteTransaction()).thenReturn(closedTransactionRunner); final SessionImpl openSession = mock(SessionImpl.class); @@ -1231,6 +1233,7 @@ public void testSessionNotFoundReadWriteTransaction() { when(openSession.beginTransaction()).thenReturn(ByteString.copyFromUtf8("open-txn")); TransactionRunnerImpl openTransactionRunner = new TransactionRunnerImpl(openSession, mock(SpannerRpc.class), 10); + openTransactionRunner.setSpan(mock(Span.class)); when(openSession.readWriteTransaction()).thenReturn(openTransactionRunner); ResultSet openResultSet = mock(ResultSet.class); diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpanTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpanTest.java index e60522dc1d8..0caab4f574e 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpanTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpanTest.java @@ -302,26 +302,29 @@ public Void run(TransactionContext transaction) throws Exception { @Test public void transactionRunnerWithError() { - TransactionRunner runner = client.readWriteTransaction(); - try { - runner.run( - new TransactionCallable() { - @Override - public Void run(TransactionContext transaction) throws Exception { - transaction.executeUpdate(INVALID_UPDATE_STATEMENT); - return null; - } - }); - fail("missing expected exception"); - } catch (SpannerException e) { - assertThat(e.getErrorCode()).isEqualTo(ErrorCode.INVALID_ARGUMENT); - } + for (int i = 0; i < 1000; i++) { + TransactionRunner runner = client.readWriteTransaction(); + try { + runner.run( + new TransactionCallable() { + @Override + public Void run(TransactionContext transaction) throws Exception { + transaction.executeUpdate(INVALID_UPDATE_STATEMENT); + return null; + } + }); + fail("missing expected exception"); + } catch (SpannerException e) { + assertThat(e.getErrorCode()).isEqualTo(ErrorCode.INVALID_ARGUMENT); + } - Map spans = failOnOverkillTraceComponent.getSpans(); - assertThat(spans).containsEntry("CloudSpanner.ReadWriteTransaction", true); - assertThat(spans).containsEntry("CloudSpannerOperation.BatchCreateSessions", true); - assertThat(spans).containsEntry("SessionPool.WaitForSession", true); - assertThat(spans).containsEntry("CloudSpannerOperation.BatchCreateSessionsRequest", true); - assertThat(spans).containsEntry("CloudSpannerOperation.BeginTransaction", true); + Map spans = failOnOverkillTraceComponent.getSpans(); + assertThat(spans.size()).isEqualTo(5); + assertThat(spans).containsEntry("CloudSpanner.ReadWriteTransaction", true); + assertThat(spans).containsEntry("CloudSpannerOperation.BatchCreateSessions", true); + assertThat(spans).containsEntry("SessionPool.WaitForSession", true); + assertThat(spans).containsEntry("CloudSpannerOperation.BatchCreateSessionsRequest", true); + assertThat(spans).containsEntry("CloudSpannerOperation.BeginTransaction", true); + } } } diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionContextImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionContextImplTest.java index 061187696a0..3446ba9fc84 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionContextImplTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionContextImplTest.java @@ -27,6 +27,7 @@ import com.google.rpc.Status; import com.google.spanner.v1.ExecuteBatchDmlRequest; import com.google.spanner.v1.ExecuteBatchDmlResponse; +import io.opencensus.trace.Span; import java.util.Arrays; import org.junit.Test; import org.junit.runner.RunWith; diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionManagerImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionManagerImplTest.java index cc522f3f457..ad56c3cc09a 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionManagerImplTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionManagerImplTest.java @@ -37,6 +37,7 @@ import com.google.spanner.v1.CommitRequest; import com.google.spanner.v1.CommitResponse; import com.google.spanner.v1.Transaction; +import io.opencensus.trace.Span; import java.util.Arrays; import java.util.Collections; import java.util.List; @@ -78,7 +79,7 @@ public void release(ScheduledExecutorService exec) { @Before public void setUp() { initMocks(this); - manager = new TransactionManagerImpl(session); + manager = new TransactionManagerImpl(session, mock(Span.class)); } @Test diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java index 1f2df00e057..8bcc05aa001 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java @@ -51,6 +51,7 @@ import io.grpc.Status; import io.grpc.StatusRuntimeException; import io.grpc.protobuf.ProtoUtils; +import io.opencensus.trace.Span; import java.util.Arrays; import java.util.Collections; import java.util.List; @@ -96,6 +97,7 @@ public void setUp() throws Exception { firstRun = true; when(session.newTransaction()).thenReturn(txn); transactionRunner = new TransactionRunnerImpl(session, rpc, 1); + transactionRunner.setSpan(mock(Span.class)); } @SuppressWarnings("unchecked") @@ -278,6 +280,7 @@ private long[] batchDmlException(int status) { .thenReturn(ByteString.copyFromUtf8(UUID.randomUUID().toString())); when(session.getName()).thenReturn(SessionId.of("p", "i", "d", "test").getName()); TransactionRunnerImpl runner = new TransactionRunnerImpl(session, rpc, 10); + runner.setSpan(mock(Span.class)); ExecuteBatchDmlResponse response1 = ExecuteBatchDmlResponse.newBuilder() .addResultSets( From 9d58bc366dc3c9317a5b6419c2fb4a616e38b39a Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Thu, 27 Feb 2020 08:25:13 +0100 Subject: [PATCH 10/49] fix: fix async runner tests --- .../com/google/cloud/spanner/AsyncRunner.java | 4 - .../google/cloud/spanner/AsyncRunnerTest.java | 84 +++++++++++++++++++ 2 files changed, 84 insertions(+), 4 deletions(-) diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncRunner.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncRunner.java index 432d6a8645d..de15d79c7ae 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncRunner.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncRunner.java @@ -48,10 +48,6 @@ interface AsyncWork { * * @param txn the transaction * @return future over the result of the work - *

    TODO(loite): It's probably better to let this method return `R` instead of - * `ApiFuture`, as we need to wait until the result of the work has actually finished - * before we can commit the transaction. Returning an ApiFuture here just means that the - * underlying framework code still has to call {@link ApiFuture#get()} before committing. */ ApiFuture doWorkAsync(TransactionContext txn); } diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncRunnerTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncRunnerTest.java index 5782ed8e663..5dbdd1092f9 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncRunnerTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncRunnerTest.java @@ -23,9 +23,12 @@ import com.google.api.core.ApiFunction; import com.google.api.core.ApiFuture; import com.google.api.core.ApiFutures; +import com.google.api.core.SettableApiFuture; import com.google.api.gax.grpc.testing.LocalChannelProvider; import com.google.cloud.NoCredentials; import com.google.cloud.Timestamp; +import com.google.cloud.spanner.AsyncResultSet.CallbackResponse; +import com.google.cloud.spanner.AsyncResultSet.ReadyCallback; import com.google.cloud.spanner.AsyncRunner.AsyncWork; import com.google.cloud.spanner.MockSpannerServiceImpl.SimulatedExecutionTime; import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult; @@ -35,10 +38,15 @@ import io.grpc.Server; import io.grpc.Status; import io.grpc.inprocess.InProcessServerBuilder; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.ScheduledThreadPoolExecutor; +import java.util.concurrent.SynchronousQueue; import java.util.concurrent.atomic.AtomicInteger; import org.junit.After; import org.junit.AfterClass; @@ -302,6 +310,82 @@ public ApiFuture doWorkAsync(TransactionContext txn) { } } + @Test + public void asyncRunnerWaitsUntilAsyncUpdateHasFinished() { + AsyncRunner runner = client.runAsync(); + runner.runAsync( + new AsyncWork() { + @Override + public ApiFuture doWorkAsync(TransactionContext txn) { + txn.executeUpdateAsync(UPDATE_STATEMENT); + return ApiFutures.immediateFuture(null); + } + }, + executor); + } + @Test + public void closeTransactionBeforeEndOfAsyncQuery() throws Exception { + final BlockingQueue results = new SynchronousQueue<>(); + final SettableApiFuture finished = SettableApiFuture.create(); + DatabaseClientImpl clientImpl = (DatabaseClientImpl) client; + + // There should currently not be any sessions checked out of the pool. + assertThat(clientImpl.pool.getNumberOfSessionsInUse()).isEqualTo(0); + + AsyncRunner runner = client.runAsync(); + final CountDownLatch dataReceived = new CountDownLatch(1); + ApiFuture res = runner.runAsync( + new AsyncWork() { + @Override + public ApiFuture doWorkAsync(TransactionContext txn) { + try (AsyncResultSet rs = + txn.readAsync(READ_TABLE_NAME, KeySet.all(), READ_COLUMN_NAMES, Options.bufferRows(1))) { + rs.setCallback( + Executors.newSingleThreadExecutor(), + new ReadyCallback() { + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + try { + while (true) { + switch (resultSet.tryNext()) { + case DONE: + finished.set(true); + return CallbackResponse.DONE; + case NOT_READY: + return CallbackResponse.CONTINUE; + case OK: + dataReceived.countDown(); + results.put(resultSet.getString(0)); + } + } + } catch (Throwable t) { + finished.setException(t); + dataReceived.countDown(); + return CallbackResponse.DONE; + } + } + }); + } + return ApiFutures.immediateFuture(null); + } + }, + executor); + // Wait until at least one row has been fetched. At that moment there should be one session + // checked out. + dataReceived.await(); + assertThat(clientImpl.pool.getNumberOfSessionsInUse()).isEqualTo(1); + assertThat(res.isDone()).isFalse(); + // Get the data from the transaction. + List resultList = new ArrayList<>(); + do { + results.drainTo(resultList); + } while (!finished.isDone() || results.size() > 0); + assertThat(finished.get()).isTrue(); + assertThat(resultList).containsExactly("k1", "k2", "k3"); + assertThat(res.get()).isNull(); + assertThat(clientImpl.pool.getNumberOfSessionsInUse()).isEqualTo(0); + } + @Test public void asyncRunnerReadRow() throws Exception { AsyncRunner runner = client.runAsync(); From 4f796325b2c9f38bd2afe839604abc262a9ae48b Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Fri, 28 Feb 2020 14:15:50 +0100 Subject: [PATCH 11/49] fix: make async runner wait for async operations --- .../cloud/spanner/AbstractReadContext.java | 41 ++++- .../cloud/spanner/AsyncResultSetImpl.java | 40 +++-- .../spanner/ForwardingAsyncResultSet.java | 66 +++++++ .../com/google/cloud/spanner/SessionPool.java | 66 ++++--- .../com/google/cloud/spanner/SpannerImpl.java | 1 - .../cloud/spanner/TransactionRunnerImpl.java | 103 ++++++++++- .../cloud/spanner/AsyncResultSetImplTest.java | 42 +++-- .../google/cloud/spanner/AsyncRunnerTest.java | 169 +++++++++++------- .../cloud/spanner/MockSpannerServiceImpl.java | 18 ++ .../com/google/cloud/spanner/SpanTest.java | 44 +++-- .../cloud/spanner/it/ITAsyncAPITest.java | 31 ++++ 11 files changed, 482 insertions(+), 139 deletions(-) create mode 100644 google-cloud-spanner/src/main/java/com/google/cloud/spanner/ForwardingAsyncResultSet.java diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java index f4a12970358..f191a858a11 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java @@ -103,6 +103,17 @@ B setDefaultQueryOptions(QueryOptions defaultQueryOptions) { abstract T build(); } + /** + * {@link AsyncResultSet} that supports adding listeners that are called when all rows from the + * underlying result stream have been fetched. + */ + interface ListenableAsyncResultSet extends AsyncResultSet { + /** Adds a listener to this {@link AsyncResultSet}. */ + void addListener(Runnable listener); + + void removeListener(Runnable listener); + } + /** * A {@code ReadContext} for standalone reads. This can only be used for a single operation, since * each standalone read may see a different timestamp of Cloud Spanner data. @@ -398,10 +409,15 @@ public final ResultSet read( } @Override - public final AsyncResultSet readAsync( + public ListenableAsyncResultSet readAsync( String table, KeySet keys, Iterable columns, ReadOption... options) { + Options readOptions = Options.fromReadOptions(options); + final int bufferRows = + readOptions.hasBufferRows() + ? readOptions.bufferRows() + : AsyncResultSetImpl.DEFAULT_BUFFER_SIZE; return new AsyncResultSetImpl( - executorProvider, readInternal(table, null, keys, columns, options)); + executorProvider, readInternal(table, null, keys, columns, options), bufferRows); } @Override @@ -411,10 +427,17 @@ public final ResultSet readUsingIndex( } @Override - public final AsyncResultSet readUsingIndexAsync( + public ListenableAsyncResultSet readUsingIndexAsync( String table, String index, KeySet keys, Iterable columns, ReadOption... options) { + Options readOptions = Options.fromReadOptions(options); + final int bufferRows = + readOptions.hasBufferRows() + ? readOptions.bufferRows() + : AsyncResultSetImpl.DEFAULT_BUFFER_SIZE; return new AsyncResultSetImpl( - executorProvider, readInternal(table, checkNotNull(index), keys, columns, options)); + executorProvider, + readInternal(table, checkNotNull(index), keys, columns, options), + bufferRows); } @Nullable @@ -457,11 +480,17 @@ public final ResultSet executeQuery(Statement statement, QueryOption... options) } @Override - public final AsyncResultSet executeQueryAsync(Statement statement, QueryOption... options) { + public ListenableAsyncResultSet executeQueryAsync(Statement statement, QueryOption... options) { + Options readOptions = Options.fromQueryOptions(options); + final int bufferRows = + readOptions.hasBufferRows() + ? readOptions.bufferRows() + : AsyncResultSetImpl.DEFAULT_BUFFER_SIZE; return new AsyncResultSetImpl( executorProvider, executeQueryInternal( - statement, com.google.spanner.v1.ExecuteSqlRequest.QueryMode.NORMAL, options)); + statement, com.google.spanner.v1.ExecuteSqlRequest.QueryMode.NORMAL, options), + bufferRows); } @Override 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 21f9094b24d..82ed5aab989 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 @@ -19,11 +19,14 @@ import com.google.api.core.ApiFuture; import com.google.api.core.SettableApiFuture; import com.google.api.gax.core.ExecutorProvider; +import com.google.cloud.spanner.AbstractReadContext.ListenableAsyncResultSet; import com.google.common.base.Function; import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableList; import com.google.common.util.concurrent.MoreExecutors; import com.google.spanner.v1.ResultSetStats; +import java.util.Collection; +import java.util.LinkedList; import java.util.concurrent.BlockingDeque; import java.util.concurrent.Callable; import java.util.concurrent.CountDownLatch; @@ -34,7 +37,8 @@ import java.util.concurrent.ScheduledExecutorService; /** Default implementation for {@link AsyncResultSet}. */ -class AsyncResultSetImpl extends ForwardingStructReader implements AsyncResultSet { +class AsyncResultSetImpl extends ForwardingStructReader implements ListenableAsyncResultSet { + /** State of an {@link AsyncResultSetImpl}. */ private enum State { INITIALIZED, @@ -58,7 +62,7 @@ private State(boolean shouldStop) { } } - private static final int DEFAULT_BUFFER_SIZE = 10; + static final int DEFAULT_BUFFER_SIZE = 10; private static final int MAX_WAIT_FOR_BUFFER_CONSUMPTION = 10; private final Object monitor = new Object(); @@ -90,6 +94,12 @@ private State(boolean shouldStop) { private ReadyCallback callback; + /** + * Listeners that will be called when the {@link AsyncResultSetImpl} has finished fetching all + * rows and any underlying transaction or session can be closed. + */ + private Collection listeners = new LinkedList<>(); + private State state = State.INITIALIZED; /** @@ -122,10 +132,6 @@ private State(boolean shouldStop) { */ private volatile CountDownLatch consumingLatch = new CountDownLatch(0); - AsyncResultSetImpl(ExecutorProvider executorProvider, ResultSet delegate) { - this(executorProvider, delegate, DEFAULT_BUFFER_SIZE); - } - AsyncResultSetImpl(ExecutorProvider executorProvider, ResultSet delegate, int bufferSize) { super(delegate); this.buffer = new LinkedBlockingDeque<>(bufferSize); @@ -155,11 +161,21 @@ public void close() { } /** - * Called when no more rows will be read from the underlying {@link ResultSet}, either because all - * rows have been read, or because {@link ReadyCallback#cursorReady(AsyncResultSet)} returned - * {@link CallbackResponse#DONE}. + * Adds a listener that will be called when no more rows will be read from the underlying {@link + * ResultSet}, either because all rows have been read, or because {@link + * ReadyCallback#cursorReady(AsyncResultSet)} returned {@link CallbackResponse#DONE}. */ - void onFinished() {} + @Override + public void addListener(Runnable listener) { + Preconditions.checkState(state == State.INITIALIZED); + listeners.add(listener); + } + + @Override + public void removeListener(Runnable listener) { + Preconditions.checkState(state == State.INITIALIZED); + listeners.remove(listener); + } /** * Tries to advance this {@link AsyncResultSet} to the next row. This method may only be called @@ -339,7 +355,9 @@ public Void call() throws Exception { delegateResultSet.close(); } catch (Throwable t) { } finally { - onFinished(); + for (Runnable listener : listeners) { + listener.run(); + } } // Ensure that the callback has been called at least once, even if the result set was diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ForwardingAsyncResultSet.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ForwardingAsyncResultSet.java new file mode 100644 index 00000000000..c5535bc4490 --- /dev/null +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ForwardingAsyncResultSet.java @@ -0,0 +1,66 @@ +/* + * 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.Function; +import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableList; +import java.util.concurrent.Executor; + +/** Forwarding implementation of {@link AsyncResultSet} that forwards all calls to a delegate. */ +public class ForwardingAsyncResultSet extends ForwardingResultSet implements AsyncResultSet { + final AsyncResultSet delegate; + + public ForwardingAsyncResultSet(AsyncResultSet delegate) { + super(Preconditions.checkNotNull(delegate)); + this.delegate = delegate; + } + + @Override + public CursorState tryNext() throws SpannerException { + return delegate.tryNext(); + } + + @Override + public void setCallback(Executor exec, ReadyCallback cb) { + delegate.setCallback(exec, cb); + ; + } + + @Override + public void cancel() { + delegate.cancel(); + } + + @Override + public void resume() { + delegate.resume(); + } + + @Override + public ApiFuture> toListAsync( + Function transformer, Executor executor) { + return delegate.toListAsync(transformer, executor); + } + + @Override + public ImmutableList toList(Function transformer) + throws SpannerException { + return delegate.toList(transformer); + } +} diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java index c3b32351ad0..a0942cc19f5 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java @@ -81,20 +81,17 @@ import java.util.Queue; import java.util.Random; import java.util.Set; -import java.util.concurrent.Executors; -import java.util.concurrent.ScheduledExecutorService; -import java.util.concurrent.ScheduledFuture; -import java.util.concurrent.SynchronousQueue; -import java.util.concurrent.ThreadPoolExecutor; import java.util.concurrent.BlockingQueue; import java.util.concurrent.Callable; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; import java.util.concurrent.Executor; +import java.util.concurrent.Executors; import java.util.concurrent.Future; import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.ScheduledFuture; +import java.util.concurrent.ThreadPoolExecutor; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; @@ -151,6 +148,11 @@ public ResultSet get() { * finished, if it is a single use context. */ private static class AutoClosingReadContext implements ReadContext { + /** + * {@link AsyncResultSet} implementation that keeps track of the async operations that are still + * running for this {@link ReadContext} and that should finish before the {@link ReadContext} + * releases its session back into the pool. + */ private class AutoClosingReadContextAsyncResultSetImpl extends AsyncResultSetImpl { private AutoClosingReadContextAsyncResultSetImpl( ExecutorProvider executorProvider, ResultSet delegate, int bufferRows) { @@ -159,19 +161,28 @@ private AutoClosingReadContextAsyncResultSetImpl( @Override public void setCallback(Executor exec, ReadyCallback cb) { - asyncOperationsCount.incrementAndGet(); - super.setCallback(exec, cb); - } - - @Override - void onFinished() { - synchronized (lock) { - if (asyncOperationsCount.decrementAndGet() == 0) { - if (closed) { - // All async operations for this read context have finished. - AutoClosingReadContext.this.close(); - } - } + Runnable listener = + new Runnable() { + @Override + public void run() { + synchronized (lock) { + if (asyncOperationsCount.decrementAndGet() == 0) { + if (closed) { + // All async operations for this read context have finished. + AutoClosingReadContext.this.close(); + } + } + } + } + }; + try { + asyncOperationsCount.incrementAndGet(); + addListener(listener); + super.setCallback(exec, cb); + } catch (Throwable t) { + removeListener(listener); + asyncOperationsCount.decrementAndGet(); + throw t; } } } @@ -321,7 +332,10 @@ public AsyncResultSet readAsync( final Iterable columns, final ReadOption... options) { Options readOptions = Options.fromReadOptions(options); - final int bufferRows = readOptions.hasBufferRows() ? readOptions.bufferRows() : 10; + final int bufferRows = + readOptions.hasBufferRows() + ? readOptions.bufferRows() + : AsyncResultSetImpl.DEFAULT_BUFFER_SIZE; return new AutoClosingReadContextAsyncResultSetImpl( sessionPool.sessionClient.getSpanner().getAsyncExecutorProvider(), wrap( @@ -358,7 +372,10 @@ public AsyncResultSet readUsingIndexAsync( final Iterable columns, final ReadOption... options) { Options readOptions = Options.fromReadOptions(options); - final int bufferRows = readOptions.hasBufferRows() ? readOptions.bufferRows() : 10; + final int bufferRows = + readOptions.hasBufferRows() + ? readOptions.bufferRows() + : AsyncResultSetImpl.DEFAULT_BUFFER_SIZE; return new AutoClosingReadContextAsyncResultSetImpl( sessionPool.sessionClient.getSpanner().getAsyncExecutorProvider(), wrap( @@ -454,7 +471,10 @@ ResultSet load() { public AsyncResultSet executeQueryAsync( final Statement statement, final QueryOption... options) { Options queryOptions = Options.fromQueryOptions(options); - final int bufferRows = queryOptions.hasBufferRows() ? queryOptions.bufferRows() : 10; + final int bufferRows = + queryOptions.hasBufferRows() + ? queryOptions.bufferRows() + : AsyncResultSetImpl.DEFAULT_BUFFER_SIZE; return new AutoClosingReadContextAsyncResultSetImpl( sessionPool.sessionClient.getSpanner().getAsyncExecutorProvider(), wrap( @@ -1546,6 +1566,9 @@ private static enum Position { private final ScheduledExecutorService executor; private final ExecutorFactory executorFactory; private final ScheduledExecutorService prepareExecutor; + + // TODO(loite): Refactor Waiter to use a SettableFuture that can be set when a session is released + // into the pool, instead of using a thread waiting on a synchronous queue. private final ScheduledExecutorService readWaiterExecutor = Executors.newSingleThreadScheduledExecutor( new ThreadFactoryBuilder() @@ -1558,6 +1581,7 @@ private static enum Position { .setDaemon(true) .setNameFormat("session-pool-write-waiter-%d") .build()); + final PoolMaintainer poolMaintainer; private final Clock clock; private final Object lock = new Object(); diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java index ef8e08ef2a7..4aaa8ae97bb 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java @@ -43,7 +43,6 @@ import java.util.concurrent.ExecutionException; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; -import java.util.concurrent.ScheduledExecutorService; import java.util.logging.Level; import java.util.logging.Logger; import javax.annotation.Nullable; 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 d292ad52281..e4fae3e3a68 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 @@ -26,6 +26,8 @@ import com.google.api.core.ApiFutures; import com.google.api.gax.core.ExecutorProvider; import com.google.cloud.Timestamp; +import com.google.cloud.spanner.Options.QueryOption; +import com.google.cloud.spanner.Options.ReadOption; import com.google.cloud.spanner.SessionImpl.SessionTransaction; import com.google.cloud.spanner.spi.v1.SpannerRpc; import com.google.common.annotations.VisibleForTesting; @@ -49,7 +51,9 @@ import java.util.ArrayList; import java.util.List; import java.util.concurrent.Callable; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executor; import java.util.concurrent.atomic.AtomicInteger; import java.util.logging.Level; import java.util.logging.Logger; @@ -83,6 +87,54 @@ static Builder newBuilder() { return new Builder(); } + /** + * {@link AsyncResultSet} implementation that keeps track of the async operations that are still + * running for this {@link TransactionContext} and that should finish before the {@link + * TransactionContext} can commit and release its session back into the pool. + */ + private class TransactionContextAsyncResultSetImpl extends ForwardingAsyncResultSet + implements ListenableAsyncResultSet { + private TransactionContextAsyncResultSetImpl(ListenableAsyncResultSet delegate) { + super(delegate); + } + + @Override + public void setCallback(Executor exec, ReadyCallback cb) { + Runnable listener = + new Runnable() { + @Override + public void run() { + finishedAsyncOperations.countDown(); + } + }; + try { + increaseAsynOperations(); + addListener(listener); + super.setCallback(exec, cb); + } catch (Throwable t) { + removeListener(listener); + finishedAsyncOperations.countDown(); + throw t; + } + } + + @Override + public void addListener(Runnable listener) { + ((ListenableAsyncResultSet) this.delegate).addListener(listener); + } + + @Override + public void removeListener(Runnable listener) { + ((ListenableAsyncResultSet) this.delegate).removeListener(listener); + } + } + + @GuardedBy("lock") + private volatile boolean committing; + + @GuardedBy("lock") + private volatile CountDownLatch finishedAsyncOperations = new CountDownLatch(0); + @GuardedBy("lock") private List mutations = new ArrayList<>(); @@ -101,6 +153,12 @@ private TransactionContextImpl(Builder builder) { this.transactionId = builder.transactionId; } + private void increaseAsynOperations() { + synchronized (lock) { + finishedAsyncOperations = new CountDownLatch((int) finishedAsyncOperations.getCount() + 1); + } + } + void ensureTxn() { if (transactionId == null || isAborted()) { span.addAnnotation("Creating Transaction"); @@ -131,6 +189,15 @@ void ensureTxn() { } void commit() { + CountDownLatch latch; + synchronized (lock) { + latch = finishedAsyncOperations; + } + try { + latch.await(); + } catch (InterruptedException e) { + throw SpannerExceptionFactory.propagateInterrupt(e); + } span.addAnnotation("Starting Commit"); CommitRequest.Builder builder = CommitRequest.newBuilder().setSession(session.getName()).setTransactionId(transactionId); @@ -264,8 +331,16 @@ public ApiFuture executeUpdateAsync(Statement statement) { beforeReadOrQuery(); final ExecuteSqlRequest.Builder builder = getExecuteSqlRequestBuilder(statement, QueryMode.NORMAL); - ApiFuture resultSet = - rpc.executeQueryAsync(builder.build(), session.getOptions()); + ApiFuture resultSet; + try { + // Register the update as an async operation that must finish before the transaction may + // commit. + increaseAsynOperations(); + resultSet = rpc.executeQueryAsync(builder.build(), session.getOptions()); + } catch (Throwable t) { + finishedAsyncOperations.countDown(); + throw t; + } final ApiFuture updateCount = ApiFutures.transform( resultSet, @@ -295,6 +370,8 @@ public void run() { onError(SpannerExceptionFactory.newSpannerException(e.getCause())); } catch (InterruptedException e) { onError(SpannerExceptionFactory.propagateInterrupt(e)); + } finally { + finishedAsyncOperations.countDown(); } } }, @@ -331,6 +408,28 @@ public long[] batchUpdate(Iterable statements) { throw e; } } + + private ListenableAsyncResultSet wrap(ListenableAsyncResultSet delegate) { + return new TransactionContextAsyncResultSetImpl(delegate); + } + + @Override + public ListenableAsyncResultSet readAsync( + String table, KeySet keys, Iterable columns, ReadOption... options) { + return wrap(super.readAsync(table, keys, columns, options)); + } + + @Override + public ListenableAsyncResultSet readUsingIndexAsync( + String table, String index, KeySet keys, Iterable columns, ReadOption... options) { + return wrap(super.readUsingIndexAsync(table, index, keys, columns, options)); + } + + @Override + public ListenableAsyncResultSet executeQueryAsync( + final Statement statement, final QueryOption... options) { + return wrap(super.executeQueryAsync(statement, options)); + } } private boolean blockNestedTxn = true; diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncResultSetImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncResultSetImplTest.java index cd5588187e6..9359dc66946 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncResultSetImplTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncResultSetImplTest.java @@ -60,7 +60,9 @@ public void setup() { @SuppressWarnings("unchecked") @Test public void close() { - AsyncResultSetImpl rs = new AsyncResultSetImpl(mockedProvider, mock(ResultSet.class)); + AsyncResultSetImpl rs = + new AsyncResultSetImpl( + mockedProvider, mock(ResultSet.class), AsyncResultSetImpl.DEFAULT_BUFFER_SIZE); rs.close(); // Closing a second time should be a no-op. rs.close(); @@ -83,7 +85,9 @@ public void close() { } // The following methods are allowed on a closed result set. - AsyncResultSetImpl rs2 = new AsyncResultSetImpl(mockedProvider, mock(ResultSet.class)); + AsyncResultSetImpl rs2 = + new AsyncResultSetImpl( + mockedProvider, mock(ResultSet.class), AsyncResultSetImpl.DEFAULT_BUFFER_SIZE); rs2.setCallback(mock(Executor.class), mock(ReadyCallback.class)); rs2.close(); rs2.cancel(); @@ -92,7 +96,9 @@ public void close() { @Test public void tryNextNotAllowed() { - try (AsyncResultSetImpl rs = new AsyncResultSetImpl(mockedProvider, mock(ResultSet.class))) { + try (AsyncResultSetImpl rs = + new AsyncResultSetImpl( + mockedProvider, mock(ResultSet.class), AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { rs.setCallback(mock(Executor.class), mock(ReadyCallback.class)); try { rs.tryNext(); @@ -109,7 +115,8 @@ public void toList() { ResultSet delegate = mock(ResultSet.class); when(delegate.next()).thenReturn(true, true, true, false); when(delegate.getCurrentRowAsStruct()).thenReturn(mock(Struct.class)); - try (AsyncResultSetImpl rs = new AsyncResultSetImpl(simpleProvider, delegate)) { + try (AsyncResultSetImpl rs = + new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { ImmutableList list = rs.toList( new Function() { @@ -129,7 +136,8 @@ public void toListPropagatesError() { .thenThrow( SpannerExceptionFactory.newSpannerException( ErrorCode.INVALID_ARGUMENT, "invalid query")); - try (AsyncResultSetImpl rs = new AsyncResultSetImpl(simpleProvider, delegate)) { + try (AsyncResultSetImpl rs = + new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { rs.toList( new Function() { @Override @@ -150,7 +158,8 @@ public void toListAsync() throws InterruptedException, ExecutionException { ResultSet delegate = mock(ResultSet.class); when(delegate.next()).thenReturn(true, true, true, false); when(delegate.getCurrentRowAsStruct()).thenReturn(mock(Struct.class)); - try (AsyncResultSetImpl rs = new AsyncResultSetImpl(simpleProvider, delegate)) { + try (AsyncResultSetImpl rs = + new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { ApiFuture> future = rs.toListAsync( new Function() { @@ -173,7 +182,8 @@ public void toListAsyncPropagatesError() throws InterruptedException { .thenThrow( SpannerExceptionFactory.newSpannerException( ErrorCode.INVALID_ARGUMENT, "invalid query")); - try (AsyncResultSetImpl rs = new AsyncResultSetImpl(simpleProvider, delegate)) { + try (AsyncResultSetImpl rs = + new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { rs.toListAsync( new Function() { @Override @@ -202,7 +212,8 @@ public void withCallback() throws InterruptedException { final AtomicInteger callbackCounter = new AtomicInteger(); final AtomicInteger rowCounter = new AtomicInteger(); final CountDownLatch finishedLatch = new CountDownLatch(1); - try (AsyncResultSetImpl rs = new AsyncResultSetImpl(simpleProvider, delegate)) { + try (AsyncResultSetImpl rs = + new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { rs.setCallback( executor, new ReadyCallback() { @@ -236,7 +247,8 @@ public void callbackReceivesError() throws InterruptedException { SpannerExceptionFactory.newSpannerException( ErrorCode.INVALID_ARGUMENT, "invalid query")); final BlockingDeque receivedErr = new LinkedBlockingDeque<>(1); - try (AsyncResultSetImpl rs = new AsyncResultSetImpl(simpleProvider, delegate)) { + try (AsyncResultSetImpl rs = + new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { rs.setCallback( executor, new ReadyCallback() { @@ -271,7 +283,8 @@ public void callbackReceivesErrorHalfwayThrough() throws InterruptedException { when(delegate.getCurrentRowAsStruct()).thenReturn(mock(Struct.class)); final AtomicInteger rowCount = new AtomicInteger(); final BlockingDeque receivedErr = new LinkedBlockingDeque<>(1); - try (AsyncResultSetImpl rs = new AsyncResultSetImpl(simpleProvider, delegate)) { + try (AsyncResultSetImpl rs = + new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { rs.setCallback( executor, new ReadyCallback() { @@ -306,7 +319,8 @@ public void pauseResume() throws InterruptedException { final AtomicInteger callbackCounter = new AtomicInteger(); final BlockingDeque queue = new LinkedBlockingDeque<>(1); final AtomicBoolean finished = new AtomicBoolean(false); - try (AsyncResultSetImpl rs = new AsyncResultSetImpl(simpleProvider, delegate)) { + try (AsyncResultSetImpl rs = + new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { rs.setCallback( executor, new ReadyCallback() { @@ -350,7 +364,8 @@ public void cancel() throws InterruptedException { final AtomicInteger callbackCounter = new AtomicInteger(); final BlockingDeque queue = new LinkedBlockingDeque<>(1); final AtomicBoolean finished = new AtomicBoolean(false); - try (AsyncResultSetImpl rs = new AsyncResultSetImpl(simpleProvider, delegate)) { + try (AsyncResultSetImpl rs = + new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { rs.setCallback( executor, new ReadyCallback() { @@ -404,7 +419,8 @@ public void callbackReturnsError() throws InterruptedException { when(delegate.next()).thenReturn(true, true, true, false); when(delegate.getCurrentRowAsStruct()).thenReturn(mock(Struct.class)); final AtomicInteger callbackCounter = new AtomicInteger(); - try (AsyncResultSetImpl rs = new AsyncResultSetImpl(simpleProvider, delegate)) { + try (AsyncResultSetImpl rs = + new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { rs.setCallback( executor, new ReadyCallback() { diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncRunnerTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncRunnerTest.java index 5dbdd1092f9..d8267dcb4e3 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncRunnerTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncRunnerTest.java @@ -35,6 +35,10 @@ import com.google.common.base.Function; import com.google.common.collect.ImmutableList; import com.google.common.util.concurrent.MoreExecutors; +import com.google.spanner.v1.BatchCreateSessionsRequest; +import com.google.spanner.v1.BeginTransactionRequest; +import com.google.spanner.v1.CommitRequest; +import com.google.spanner.v1.ExecuteSqlRequest; import io.grpc.Server; import io.grpc.Status; import io.grpc.inprocess.InProcessServerBuilder; @@ -66,8 +70,6 @@ public class AsyncRunnerTest { private Spanner spanner; private Spanner spannerWithEmptySessionPool; - private DatabaseClient client; - private DatabaseClient clientWithEmptySessionPool; @BeforeClass public static void setup() throws Exception { @@ -113,7 +115,6 @@ public void before() { .setSessionPoolOption(SessionPoolOptions.newBuilder().setFailOnSessionLeak().build()) .build() .getService(); - client = spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); spannerWithEmptySessionPool = spanner .getOptions() @@ -122,9 +123,15 @@ public void before() { SessionPoolOptions.newBuilder().setFailOnSessionLeak().setMinSessions(0).build()) .build() .getService(); - clientWithEmptySessionPool = - spannerWithEmptySessionPool.getDatabaseClient( - DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); + } + + private DatabaseClient client() { + return spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); + } + + private DatabaseClient clientWithEmptySessionPool() { + return spannerWithEmptySessionPool.getDatabaseClient( + DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); } @After @@ -132,11 +139,12 @@ public void after() { spanner.close(); spannerWithEmptySessionPool.close(); mockSpanner.removeAllExecutionTimes(); + mockSpanner.reset(); } @Test public void asyncRunnerUpdate() throws Exception { - AsyncRunner runner = client.runAsync(); + AsyncRunner runner = client().runAsync(); ApiFuture updateCount = runner.runAsync( new AsyncWork() { @@ -152,12 +160,13 @@ public ApiFuture doWorkAsync(TransactionContext txn) { @Test public void asyncRunnerIsNonBlocking() throws Exception { mockSpanner.freeze(); - AsyncRunner runner = clientWithEmptySessionPool.runAsync(); + AsyncRunner runner = clientWithEmptySessionPool().runAsync(); ApiFuture res = runner.runAsync( new AsyncWork() { @Override public ApiFuture doWorkAsync(TransactionContext txn) { + txn.executeUpdateAsync(UPDATE_STATEMENT); return ApiFutures.immediateFuture(null); } }, @@ -170,7 +179,7 @@ public ApiFuture doWorkAsync(TransactionContext txn) { @Test public void asyncRunnerInvalidUpdate() throws Exception { - AsyncRunner runner = client.runAsync(); + AsyncRunner runner = client().runAsync(); ApiFuture updateCount = runner.runAsync( new AsyncWork() { @@ -191,13 +200,29 @@ public ApiFuture doWorkAsync(TransactionContext txn) { } } + @Test + public void asyncRunnerFireAndForgetInvalidUpdate() throws Exception { + AsyncRunner runner = client().runAsync(); + ApiFuture res = + runner.runAsync( + new AsyncWork() { + @Override + public ApiFuture doWorkAsync(TransactionContext txn) { + txn.executeUpdateAsync(INVALID_UPDATE_STATEMENT); + return txn.executeUpdateAsync(UPDATE_STATEMENT); + } + }, + executor); + assertThat(res.get()).isEqualTo(UPDATE_COUNT); + } + @Test public void asyncRunnerUpdateAborted() throws Exception { try { // Temporarily set the result of the update to 2 rows. mockSpanner.putStatementResult(StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT + 1L)); final AtomicInteger attempt = new AtomicInteger(); - AsyncRunner runner = client.runAsync(); + AsyncRunner runner = client().runAsync(); ApiFuture updateCount = runner.runAsync( new AsyncWork() { @@ -227,7 +252,7 @@ public void asyncRunnerCommitAborted() throws Exception { // Temporarily set the result of the update to 2 rows. mockSpanner.putStatementResult(StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT + 1L)); final AtomicInteger attempt = new AtomicInteger(); - AsyncRunner runner = client.runAsync(); + AsyncRunner runner = client().runAsync(); ApiFuture updateCount = runner.runAsync( new AsyncWork() { @@ -255,7 +280,7 @@ public ApiFuture doWorkAsync(TransactionContext txn) { @Test public void asyncRunnerUpdateAbortedWithoutGettingResult() throws Exception { final AtomicInteger attempt = new AtomicInteger(); - AsyncRunner runner = client.runAsync(); + AsyncRunner runner = clientWithEmptySessionPool().runAsync(); ApiFuture result = runner.runAsync( new AsyncWork() { @@ -277,6 +302,15 @@ public ApiFuture doWorkAsync(TransactionContext txn) { executor); assertThat(result.get()).isNull(); assertThat(attempt.get()).isEqualTo(2); + assertThat(mockSpanner.getRequestTypes()) + .containsExactly( + BatchCreateSessionsRequest.class, + BeginTransactionRequest.class, + ExecuteSqlRequest.class, + CommitRequest.class, + BeginTransactionRequest.class, + ExecuteSqlRequest.class, + CommitRequest.class); } @Test @@ -286,7 +320,7 @@ public void asyncRunnerCommitFails() throws Exception { Status.RESOURCE_EXHAUSTED .withDescription("mutation limit exceeded") .asRuntimeException())); - AsyncRunner runner = client.runAsync(); + AsyncRunner runner = client().runAsync(); ApiFuture updateCount = runner.runAsync( new AsyncWork() { @@ -311,65 +345,76 @@ public ApiFuture doWorkAsync(TransactionContext txn) { } @Test - public void asyncRunnerWaitsUntilAsyncUpdateHasFinished() { - AsyncRunner runner = client.runAsync(); - runner.runAsync( - new AsyncWork() { - @Override - public ApiFuture doWorkAsync(TransactionContext txn) { - txn.executeUpdateAsync(UPDATE_STATEMENT); - return ApiFutures.immediateFuture(null); - } - }, - executor); + public void asyncRunnerWaitsUntilAsyncUpdateHasFinished() throws Exception { + AsyncRunner runner = clientWithEmptySessionPool().runAsync(); + ApiFuture res = + runner.runAsync( + new AsyncWork() { + @Override + public ApiFuture doWorkAsync(TransactionContext txn) { + txn.executeUpdateAsync(UPDATE_STATEMENT); + return ApiFutures.immediateFuture(null); + } + }, + executor); + res.get(); + assertThat(mockSpanner.getRequestTypes()) + .containsExactly( + BatchCreateSessionsRequest.class, + BeginTransactionRequest.class, + ExecuteSqlRequest.class, + CommitRequest.class); } + @Test public void closeTransactionBeforeEndOfAsyncQuery() throws Exception { final BlockingQueue results = new SynchronousQueue<>(); final SettableApiFuture finished = SettableApiFuture.create(); - DatabaseClientImpl clientImpl = (DatabaseClientImpl) client; + DatabaseClientImpl clientImpl = (DatabaseClientImpl) client(); // There should currently not be any sessions checked out of the pool. assertThat(clientImpl.pool.getNumberOfSessionsInUse()).isEqualTo(0); - AsyncRunner runner = client.runAsync(); + AsyncRunner runner = clientImpl.runAsync(); final CountDownLatch dataReceived = new CountDownLatch(1); - ApiFuture res = runner.runAsync( - new AsyncWork() { - @Override - public ApiFuture doWorkAsync(TransactionContext txn) { - try (AsyncResultSet rs = - txn.readAsync(READ_TABLE_NAME, KeySet.all(), READ_COLUMN_NAMES, Options.bufferRows(1))) { - rs.setCallback( - Executors.newSingleThreadExecutor(), - new ReadyCallback() { - @Override - public CallbackResponse cursorReady(AsyncResultSet resultSet) { - try { - while (true) { - switch (resultSet.tryNext()) { - case DONE: - finished.set(true); - return CallbackResponse.DONE; - case NOT_READY: - return CallbackResponse.CONTINUE; - case OK: - dataReceived.countDown(); - results.put(resultSet.getString(0)); + ApiFuture res = + runner.runAsync( + new AsyncWork() { + @Override + public ApiFuture doWorkAsync(TransactionContext txn) { + try (AsyncResultSet rs = + txn.readAsync( + READ_TABLE_NAME, KeySet.all(), READ_COLUMN_NAMES, Options.bufferRows(1))) { + rs.setCallback( + Executors.newSingleThreadExecutor(), + new ReadyCallback() { + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + try { + while (true) { + switch (resultSet.tryNext()) { + case DONE: + finished.set(true); + return CallbackResponse.DONE; + case NOT_READY: + return CallbackResponse.CONTINUE; + case OK: + dataReceived.countDown(); + results.put(resultSet.getString(0)); + } + } + } catch (Throwable t) { + finished.setException(t); + dataReceived.countDown(); + return CallbackResponse.DONE; } } - } catch (Throwable t) { - finished.setException(t); - dataReceived.countDown(); - return CallbackResponse.DONE; - } - } - }); - } - return ApiFutures.immediateFuture(null); - } - }, - executor); + }); + } + return ApiFutures.immediateFuture(null); + } + }, + executor); // Wait until at least one row has been fetched. At that moment there should be one session // checked out. dataReceived.await(); @@ -388,7 +433,7 @@ public CallbackResponse cursorReady(AsyncResultSet resultSet) { @Test public void asyncRunnerReadRow() throws Exception { - AsyncRunner runner = client.runAsync(); + AsyncRunner runner = client().runAsync(); ApiFuture val = runner.runAsync( new AsyncWork() { @@ -411,7 +456,7 @@ public String apply(Struct input) { @Test public void asyncRunnerRead() throws Exception { - AsyncRunner runner = client.runAsync(); + AsyncRunner runner = client().runAsync(); ApiFuture> val = runner.runAsync( new AsyncWork>() { 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 8e0375bd279..7cb75673213 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 @@ -1634,6 +1634,24 @@ public List getRequests() { return new ArrayList<>(this.requests); } + public Iterable> getRequestTypes() { + List> res = new LinkedList<>(); + for (AbstractMessage m : this.requests) { + res.add(m.getClass()); + } + return res; + } + + public int countRequestsOfType(Class type) { + int c = 0; + for (AbstractMessage m : this.requests) { + if (m.getClass().equals(type)) { + c++; + } + } + return c; + } + @Override public void addResponse(AbstractMessage response) { throw new UnsupportedOperationException(); diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpanTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpanTest.java index 0caab4f574e..ff211c4e831 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpanTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpanTest.java @@ -302,29 +302,27 @@ public Void run(TransactionContext transaction) throws Exception { @Test public void transactionRunnerWithError() { - for (int i = 0; i < 1000; i++) { - TransactionRunner runner = client.readWriteTransaction(); - try { - runner.run( - new TransactionCallable() { - @Override - public Void run(TransactionContext transaction) throws Exception { - transaction.executeUpdate(INVALID_UPDATE_STATEMENT); - return null; - } - }); - fail("missing expected exception"); - } catch (SpannerException e) { - assertThat(e.getErrorCode()).isEqualTo(ErrorCode.INVALID_ARGUMENT); - } - - Map spans = failOnOverkillTraceComponent.getSpans(); - assertThat(spans.size()).isEqualTo(5); - assertThat(spans).containsEntry("CloudSpanner.ReadWriteTransaction", true); - assertThat(spans).containsEntry("CloudSpannerOperation.BatchCreateSessions", true); - assertThat(spans).containsEntry("SessionPool.WaitForSession", true); - assertThat(spans).containsEntry("CloudSpannerOperation.BatchCreateSessionsRequest", true); - assertThat(spans).containsEntry("CloudSpannerOperation.BeginTransaction", true); + TransactionRunner runner = client.readWriteTransaction(); + try { + runner.run( + new TransactionCallable() { + @Override + public Void run(TransactionContext transaction) throws Exception { + transaction.executeUpdate(INVALID_UPDATE_STATEMENT); + return null; + } + }); + fail("missing expected exception"); + } catch (SpannerException e) { + assertThat(e.getErrorCode()).isEqualTo(ErrorCode.INVALID_ARGUMENT); } + + Map spans = failOnOverkillTraceComponent.getSpans(); + assertThat(spans.size()).isEqualTo(5); + assertThat(spans).containsEntry("CloudSpanner.ReadWriteTransaction", true); + assertThat(spans).containsEntry("CloudSpannerOperation.BatchCreateSessions", true); + assertThat(spans).containsEntry("SessionPool.WaitForSession", true); + assertThat(spans).containsEntry("CloudSpannerOperation.BatchCreateSessionsRequest", true); + assertThat(spans).containsEntry("CloudSpannerOperation.BeginTransaction", true); } } diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITAsyncAPITest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITAsyncAPITest.java index bc46ff8b0f2..721536cb6b3 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITAsyncAPITest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITAsyncAPITest.java @@ -23,6 +23,8 @@ 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.AsyncRunner; +import com.google.cloud.spanner.AsyncRunner.AsyncWork; import com.google.cloud.spanner.Database; import com.google.cloud.spanner.DatabaseClient; import com.google.cloud.spanner.DatabaseId; @@ -34,8 +36,10 @@ import com.google.cloud.spanner.KeySet; import com.google.cloud.spanner.Mutation; import com.google.cloud.spanner.SpannerException; +import com.google.cloud.spanner.Statement; import com.google.cloud.spanner.Struct; import com.google.cloud.spanner.TimestampBound; +import com.google.cloud.spanner.TransactionContext; import com.google.cloud.spanner.Type; import com.google.cloud.spanner.Type.StructField; import com.google.cloud.spanner.testing.RemoteSpannerHelper; @@ -271,4 +275,31 @@ public void columnNotFound() throws Exception { assertThat(se.getMessage()).contains("BadColumnName"); } } + + @Test + public void asyncRunnerFireAndForgetInvalidUpdate() throws Exception { + try { + assertThat(client.singleUse().readRow("TestTable", Key.of("k999"), ALL_COLUMNS)).isNull(); + AsyncRunner runner = client.runAsync(); + ApiFuture res = + runner.runAsync( + new AsyncWork() { + @Override + public ApiFuture doWorkAsync(TransactionContext txn) { + // The error returned by this update statement will not bubble up and fail the + // transaction. + txn.executeUpdateAsync(Statement.of("UPDATE BadTableName SET FOO=1 WHERE ID=2")); + return txn.executeUpdateAsync( + Statement.of( + "INSERT INTO TestTable (Key, StringValue) VALUES ('k999', 'v999')")); + } + }, + executor); + assertThat(res.get()).isEqualTo(1L); + assertThat(client.singleUse().readRow("TestTable", Key.of("k999"), ALL_COLUMNS)).isNotNull(); + } finally { + client.writeAtLeastOnce(Arrays.asList(Mutation.delete("TestTable", Key.of("k999")))); + assertThat(client.singleUse().readRow("TestTable", Key.of("k999"), ALL_COLUMNS)).isNull(); + } + } } From cfd1802183bfe7d16d00b9acc163fb1283922148 Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Fri, 28 Feb 2020 17:37:30 +0100 Subject: [PATCH 12/49] examples: add example integration test --- .../cloud/spanner/AsyncResultSetImpl.java | 23 +- .../google/cloud/spanner/DatabaseClient.java | 31 ++ .../com/google/cloud/spanner/SessionPool.java | 4 +- .../cloud/spanner/MockSpannerTestUtil.java | 31 +- .../google/cloud/spanner/ReadAsyncTest.java | 68 +++- .../cloud/spanner/it/ITAsyncExamplesTest.java | 331 ++++++++++++++++++ 6 files changed, 457 insertions(+), 31 deletions(-) create mode 100644 google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITAsyncExamplesTest.java 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 82ed5aab989..fcccf05aacd 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 @@ -476,20 +476,23 @@ private CreateListCallback( @Override public CallbackResponse cursorReady(AsyncResultSet resultSet) { - CursorState state; try { - while ((state = resultSet.tryNext()) == CursorState.OK) { - builder.add(transformer.apply(resultSet)); + while (true) { + switch (resultSet.tryNext()) { + case DONE: + future.set(builder.build()); + return CallbackResponse.DONE; + case NOT_READY: + return CallbackResponse.CONTINUE; + case OK: + builder.add(transformer.apply(resultSet)); + break; + } } - } catch (SpannerException e) { - future.setException(e); - return CallbackResponse.DONE; - } - if (state == CursorState.DONE) { - future.set(builder.build()); + } catch (Throwable t) { + future.setException(t); return CallbackResponse.DONE; } - return CallbackResponse.CONTINUE; } } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClient.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClient.java index 18cbc3ca858..3298e6a2ab0 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClient.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClient.java @@ -278,6 +278,37 @@ public interface DatabaseClient { */ TransactionManager transactionManager(); + /** + * Returns an asynchronous transaction runner for executing a single logical transaction with + * retries. The returned runner can only be used once. + * + *

    Example of a read write transaction. + * + * 

     
+   * Executor executor = Executors.newSingleThreadExecutor();
+   * final long singerId = my_singer_id;
+   * AsyncRunner runner = client.runAsync();
+   * ApiFuture rowCount =
+   *     runner.runAsync(
+   *         new AsyncWork() {
+   *           @Override
+   *           public ApiFuture doWorkAsync(TransactionContext txn) {
+   *             String column = "FirstName";
+   *             Struct row =
+   *                 txn.readRow("Singers", Key.of(singerId), Collections.singleton("Name"));
+   *             String name = row.getString("Name");
+   *             return txn.executeUpdateAsync(
+   *                 Statement.newBuilder("UPDATE Singers SET Name=@name WHERE SingerId=@id")
+   *                     .bind("id")
+   *                     .to(singerId)
+   *                     .bind("name")
+   *                     .to(name.toUpperCase())
+   *                     .build());
+   *           }
+   *         },
+   *         executor);
+   *
    + */ AsyncRunner runAsync(); /** diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java index a0942cc19f5..ee09f28341d 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java @@ -219,8 +219,8 @@ private AutoClosingReadContext( } T getReadContextDelegate() { - if (readContextDelegate == null) { - synchronized (lock) { + synchronized (lock) { + if (readContextDelegate == null) { while (true) { try { this.readContextDelegate = readContextDelegateSupplier.apply(this.session); diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MockSpannerTestUtil.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MockSpannerTestUtil.java index 981f9244484..0a8f680f998 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MockSpannerTestUtil.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MockSpannerTestUtil.java @@ -108,22 +108,17 @@ public class MockSpannerTestUtil { .setMetadata(READ_KEY_VALUE_METADATA) .build(); static final com.google.spanner.v1.ResultSet READ_MULTIPLE_KEY_VALUE_RESULTSET = - com.google.spanner.v1.ResultSet.newBuilder() - .addRows( - ListValue.newBuilder() - .addValues(com.google.protobuf.Value.newBuilder().setStringValue("k1").build()) - .addValues(com.google.protobuf.Value.newBuilder().setStringValue("v1").build()) - .build()) - .addRows( - ListValue.newBuilder() - .addValues(com.google.protobuf.Value.newBuilder().setStringValue("k2").build()) - .addValues(com.google.protobuf.Value.newBuilder().setStringValue("v2").build()) - .build()) - .addRows( - ListValue.newBuilder() - .addValues(com.google.protobuf.Value.newBuilder().setStringValue("k3").build()) - .addValues(com.google.protobuf.Value.newBuilder().setStringValue("v3").build()) - .build()) - .setMetadata(READ_KEY_VALUE_METADATA) - .build(); + generateKeyValueResultSet(1, 3); + + static com.google.spanner.v1.ResultSet generateKeyValueResultSet(int beginRow, int endRow) { + com.google.spanner.v1.ResultSet.Builder builder = com.google.spanner.v1.ResultSet.newBuilder(); + for (int row = beginRow; row <= endRow; row++) { + builder.addRows( + ListValue.newBuilder() + .addValues(com.google.protobuf.Value.newBuilder().setStringValue("k" + row).build()) + .addValues(com.google.protobuf.Value.newBuilder().setStringValue("v" + row).build()) + .build()); + } + return builder.setMetadata(READ_KEY_VALUE_METADATA).build(); + } } diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadAsyncTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadAsyncTest.java index f73549e465c..f324905ef54 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadAsyncTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadAsyncTest.java @@ -20,7 +20,9 @@ import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.fail; +import com.google.api.core.ApiFunction; import com.google.api.core.ApiFuture; +import com.google.api.core.ApiFutures; import com.google.api.core.SettableApiFuture; import com.google.api.gax.grpc.testing.LocalChannelProvider; import com.google.cloud.NoCredentials; @@ -28,11 +30,16 @@ import com.google.cloud.spanner.AsyncResultSet.ReadyCallback; import com.google.cloud.spanner.MockSpannerServiceImpl.SimulatedExecutionTime; import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult; +import com.google.common.base.Function; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.Iterables; import com.google.common.util.concurrent.SettableFuture; import io.grpc.Server; import io.grpc.Status; import io.grpc.inprocess.InProcessServerBuilder; import java.util.ArrayList; +import java.util.Arrays; +import java.util.Comparator; import java.util.List; import java.util.concurrent.BlockingQueue; import java.util.concurrent.CountDownLatch; @@ -95,7 +102,8 @@ public void before() { .setProjectId(TEST_PROJECT) .setChannelProvider(channelProvider) .setCredentials(NoCredentials.getInstance()) - .setSessionPoolOption(SessionPoolOptions.newBuilder().setFailOnSessionLeak().build()) + .setSessionPoolOption( + SessionPoolOptions.newBuilder().setFailOnSessionLeak().setMinSessions(0).build()) .build() .getService(); client = spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); @@ -267,4 +275,62 @@ public CallbackResponse cursorReady(AsyncResultSet resultSet) { Thread.sleep(10L); assertThat(clientImpl.pool.getNumberOfSessionsInUse()).isEqualTo(0); } + + @Test + public void readOnlyTransaction() throws Exception { + Statement statement1 = + Statement.of("SELECT * FROM TestTable WHERE Key IN ('k10', 'k11', 'k12')"); + Statement statement2 = Statement.of("SELECT * FROM TestTable WHERE Key IN ('k1', 'k2', 'k3"); + mockSpanner.putStatementResult( + StatementResult.query(statement1, generateKeyValueResultSet(10, 12))); + mockSpanner.putStatementResult( + StatementResult.query(statement2, generateKeyValueResultSet(1, 3))); + + ApiFuture> values1; + ApiFuture> values2; + try (ReadOnlyTransaction tx = client.readOnlyTransaction()) { + try (AsyncResultSet rs = tx.executeQueryAsync(statement1)) { + values1 = + rs.toListAsync( + new Function() { + @Override + public String apply(StructReader input) { + return input.getString("Value"); + } + }, + executor); + } + try (AsyncResultSet rs = tx.executeQueryAsync(statement2)) { + values2 = + rs.toListAsync( + new Function() { + @Override + public String apply(StructReader input) { + return input.getString("Value"); + } + }, + executor); + } + } + ApiFuture> allValues = + ApiFutures.transform( + ApiFutures.allAsList(Arrays.asList(values1, values2)), + new ApiFunction>, Iterable>() { + @Override + public Iterable apply(List> input) { + return Iterables.mergeSorted( + input, + new Comparator() { + @Override + public int compare(String o1, String o2) { + // Return in numerical order (i.e. without the preceding 'v'). + return Integer.valueOf(o1.substring(1)) + .compareTo(Integer.valueOf(o2.substring(1))); + } + }); + } + }, + executor); + assertThat(allValues.get()).containsExactly("v1", "v2", "v3", "v10", "v11", "v12"); + } } diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITAsyncExamplesTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITAsyncExamplesTest.java new file mode 100644 index 00000000000..7c80632ed26 --- /dev/null +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITAsyncExamplesTest.java @@ -0,0 +1,331 @@ +/* + * 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.it; + +import static com.google.common.truth.Truth.assertThat; + +import com.google.api.core.ApiFunction; +import com.google.api.core.ApiFuture; +import com.google.api.core.ApiFutures; +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.AsyncRunner; +import com.google.cloud.spanner.AsyncRunner.AsyncWork; +import com.google.cloud.spanner.Database; +import com.google.cloud.spanner.DatabaseClient; +import com.google.cloud.spanner.IntegrationTest; +import com.google.cloud.spanner.IntegrationTestEnv; +import com.google.cloud.spanner.Key; +import com.google.cloud.spanner.KeySet; +import com.google.cloud.spanner.Mutation; +import com.google.cloud.spanner.ReadOnlyTransaction; +import com.google.cloud.spanner.Statement; +import com.google.cloud.spanner.Struct; +import com.google.cloud.spanner.StructReader; +import com.google.cloud.spanner.TransactionContext; +import com.google.common.base.Function; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.Iterables; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Comparator; +import java.util.LinkedList; +import java.util.List; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.ClassRule; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Integration tests for asynchronous APIs. */ +@Category(IntegrationTest.class) +@RunWith(JUnit4.class) +public class ITAsyncExamplesTest { + @ClassRule public static IntegrationTestEnv env = new IntegrationTestEnv(); + private static final String TABLE_NAME = "TestTable"; + private static final String INDEX_NAME = "TestTableByValue"; + private static final List ALL_COLUMNS = Arrays.asList("Key", "StringValue"); + private static final ImmutableList ALL_VALUES_IN_PK_ORDER = + ImmutableList.of( + "v0", "v1", "v10", "v11", "v12", "v13", "v14", "v2", "v3", "v4", "v5", "v6", "v7", "v8", + "v9"); + + private static Database db; + private static DatabaseClient client; + private static ExecutorService executor; + + @BeforeClass + public static void setUpDatabase() { + db = + env.getTestHelper() + .createTestDatabase( + "CREATE TABLE TestTable (" + + " Key STRING(MAX) NOT NULL," + + " StringValue STRING(MAX)," + + ") PRIMARY KEY (Key)", + "CREATE INDEX TestTableByValue ON TestTable(StringValue)", + "CREATE INDEX TestTableByValueDesc ON TestTable(StringValue DESC)"); + client = env.getTestHelper().getDatabaseClient(db); + + // Includes k0..k14. Note that strings k{10,14} sort between k1 and k2. + List mutations = new ArrayList<>(); + for (int i = 0; i < 15; ++i) { + mutations.add( + Mutation.newInsertOrUpdateBuilder(TABLE_NAME) + .set("Key") + .to("k" + i) + .set("StringValue") + .to("v" + i) + .build()); + } + client.write(mutations); + executor = Executors.newScheduledThreadPool(8); + } + + @AfterClass + public static void cleanup() { + executor.shutdown(); + } + + @Test + public void readAsync() throws Exception { + final SettableApiFuture> future = SettableApiFuture.create(); + try (AsyncResultSet rs = client.singleUse().readAsync(TABLE_NAME, KeySet.all(), ALL_COLUMNS)) { + rs.setCallback( + executor, + new ReadyCallback() { + final List values = new LinkedList<>(); + + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + try { + while (true) { + switch (resultSet.tryNext()) { + case DONE: + future.set(values); + return CallbackResponse.DONE; + case NOT_READY: + return CallbackResponse.CONTINUE; + case OK: + values.add(resultSet.getString("StringValue")); + break; + } + } + } catch (Throwable t) { + future.setException(t); + return CallbackResponse.DONE; + } + } + }); + } + assertThat(future.get()).containsExactlyElementsIn(ALL_VALUES_IN_PK_ORDER); + } + + @Test + public void readUsingIndexAsync() throws Exception { + final SettableApiFuture> future = SettableApiFuture.create(); + try (AsyncResultSet rs = + client.singleUse().readUsingIndexAsync(TABLE_NAME, INDEX_NAME, KeySet.all(), ALL_COLUMNS)) { + rs.setCallback( + executor, + new ReadyCallback() { + final List values = new LinkedList<>(); + + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + try { + while (true) { + switch (resultSet.tryNext()) { + case DONE: + future.set(values); + return CallbackResponse.DONE; + case NOT_READY: + return CallbackResponse.CONTINUE; + case OK: + values.add(resultSet.getString("StringValue")); + break; + } + } + } catch (Throwable t) { + future.setException(t); + return CallbackResponse.DONE; + } + } + }); + } + assertThat(future.get()).containsExactlyElementsIn(ALL_VALUES_IN_PK_ORDER); + } + + @Test + public void readRowAsync() throws Exception { + ApiFuture row = client.singleUse().readRowAsync(TABLE_NAME, Key.of("k1"), ALL_COLUMNS); + assertThat(row.get().getString("StringValue")).isEqualTo("v1"); + } + + @Test + public void readRowUsingIndexAsync() throws Exception { + ApiFuture row = + client + .singleUse() + .readRowUsingIndexAsync(TABLE_NAME, INDEX_NAME, Key.of("v2"), ALL_COLUMNS); + assertThat(row.get().getString("Key")).isEqualTo("k2"); + } + + @Test + public void executeQueryAsync() throws Exception { + final ImmutableList keys = ImmutableList.of("k3", "k4"); + final SettableApiFuture> future = SettableApiFuture.create(); + try (AsyncResultSet rs = + client + .singleUse() + .executeQueryAsync( + Statement.newBuilder("SELECT StringValue FROM TestTable WHERE Key IN UNNEST(@keys)") + .bind("keys") + .toStringArray(keys) + .build())) { + rs.setCallback( + executor, + new ReadyCallback() { + final List values = new LinkedList<>(); + + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + try { + while (true) { + switch (resultSet.tryNext()) { + case DONE: + future.set(values); + return CallbackResponse.DONE; + case NOT_READY: + return CallbackResponse.CONTINUE; + case OK: + values.add(resultSet.getString("StringValue")); + break; + } + } + } catch (Throwable t) { + future.setException(t); + return CallbackResponse.DONE; + } + } + }); + } + assertThat(future.get()).containsExactly("v3", "v4"); + } + + @Test + public void runAsync() throws Exception { + AsyncRunner runner = client.runAsync(); + ApiFuture deleteCount = + runner.runAsync( + new AsyncWork() { + @Override + public ApiFuture doWorkAsync(TransactionContext txn) { + // Even though this is a shoot-and-forget asynchronous DML statement, it is + // guaranteed to be executed within the transaction before the commit is executed. + txn.executeUpdateAsync( + Statement.newBuilder( + "INSERT INTO TestTable (Key, StringValue) VALUES (@key, @value)") + .bind("key") + .to("k999") + .bind("value") + .to("v999") + .build()); + // Note that even though both DML statements are executed asynchronously, they are + // guaranteed to be executed in the order they are submitted to the transaction, as + // they receive a monotonically increasing sequence number at the moment that they + // are submitted. If they arrive out of order on the backend, the backend may abort + // the transaction and the transaction will be retried. + return txn.executeUpdateAsync( + Statement.newBuilder("DELETE FROM TestTable WHERE Key=@key") + .bind("key") + .to("k999") + .build()); + } + }, + executor); + assertThat(deleteCount.get()).isEqualTo(1L); + } + + @Test + public void readOnlyTransaction() throws Exception { + ImmutableList keys1 = ImmutableList.of("k10", "k11", "k12"); + ImmutableList keys2 = ImmutableList.of("k1", "k2", "k3"); + ApiFuture> values1; + ApiFuture> values2; + try (ReadOnlyTransaction tx = client.readOnlyTransaction()) { + try (AsyncResultSet rs = + tx.executeQueryAsync( + Statement.newBuilder("SELECT * FROM TestTable WHERE Key IN UNNEST(@keys)") + .bind("keys") + .toStringArray(keys1) + .build())) { + values1 = + rs.toListAsync( + new Function() { + @Override + public String apply(StructReader input) { + return input.getString("StringValue"); + } + }, + executor); + } + try (AsyncResultSet rs = + tx.executeQueryAsync( + Statement.newBuilder("SELECT * FROM TestTable WHERE Key IN UNNEST(@keys)") + .bind("keys") + .toStringArray(keys2) + .build())) { + values2 = + rs.toListAsync( + new Function() { + @Override + public String apply(StructReader input) { + return input.getString("StringValue"); + } + }, + executor); + } + } + ApiFuture> allValues = + ApiFutures.transform( + ApiFutures.allAsList(Arrays.asList(values1, values2)), + new ApiFunction>, Iterable>() { + @Override + public Iterable apply(List> input) { + return Iterables.mergeSorted( + input, + new Comparator() { + @Override + public int compare(String o1, String o2) { + // Compare based on numerical order (i.e. without the preceding 'v'). + return Integer.valueOf(o1.substring(1)) + .compareTo(Integer.valueOf(o2.substring(1))); + } + }); + } + }, + executor); + assertThat(allValues.get()).containsExactly("v1", "v2", "v3", "v10", "v11", "v12"); + } +} From a2e28a577aa007133c723ea2d7fc785d66053c3a Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Fri, 28 Feb 2020 19:34:58 +0100 Subject: [PATCH 13/49] examples: add more examples --- .../cloud/spanner/MockSpannerTestUtil.java | 7 +- .../google/cloud/spanner/ReadAsyncTest.java | 143 +++++++++++++++++- .../cloud/spanner/it/ITAsyncExamplesTest.java | 138 +++++++++++++++++ 3 files changed, 283 insertions(+), 5 deletions(-) diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MockSpannerTestUtil.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MockSpannerTestUtil.java index 0a8f680f998..fe85ef7c908 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MockSpannerTestUtil.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MockSpannerTestUtil.java @@ -17,6 +17,7 @@ package com.google.cloud.spanner; import com.google.cloud.spanner.Type.StructField; +import com.google.common.collect.ContiguousSet; import com.google.protobuf.ListValue; import com.google.spanner.v1.ResultSetMetadata; import com.google.spanner.v1.StructType; @@ -108,11 +109,11 @@ public class MockSpannerTestUtil { .setMetadata(READ_KEY_VALUE_METADATA) .build(); static final com.google.spanner.v1.ResultSet READ_MULTIPLE_KEY_VALUE_RESULTSET = - generateKeyValueResultSet(1, 3); + generateKeyValueResultSet(ContiguousSet.closed(1, 3)); - static com.google.spanner.v1.ResultSet generateKeyValueResultSet(int beginRow, int endRow) { + static com.google.spanner.v1.ResultSet generateKeyValueResultSet(Iterable rows) { com.google.spanner.v1.ResultSet.Builder builder = com.google.spanner.v1.ResultSet.newBuilder(); - for (int row = beginRow; row <= endRow; row++) { + for (Integer row : rows) { builder.addRows( ListValue.newBuilder() .addValues(com.google.protobuf.Value.newBuilder().setStringValue("k" + row).build()) diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadAsyncTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadAsyncTest.java index f324905ef54..53e5041891a 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadAsyncTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadAsyncTest.java @@ -31,7 +31,9 @@ import com.google.cloud.spanner.MockSpannerServiceImpl.SimulatedExecutionTime; import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult; import com.google.common.base.Function; +import com.google.common.collect.ContiguousSet; import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableSet; import com.google.common.collect.Iterables; import com.google.common.util.concurrent.SettableFuture; import io.grpc.Server; @@ -40,6 +42,7 @@ import java.util.ArrayList; import java.util.Arrays; import java.util.Comparator; +import java.util.LinkedList; import java.util.List; import java.util.concurrent.BlockingQueue; import java.util.concurrent.CountDownLatch; @@ -282,9 +285,9 @@ public void readOnlyTransaction() throws Exception { Statement.of("SELECT * FROM TestTable WHERE Key IN ('k10', 'k11', 'k12')"); Statement statement2 = Statement.of("SELECT * FROM TestTable WHERE Key IN ('k1', 'k2', 'k3"); mockSpanner.putStatementResult( - StatementResult.query(statement1, generateKeyValueResultSet(10, 12))); + StatementResult.query(statement1, generateKeyValueResultSet(ContiguousSet.closed(10, 12)))); mockSpanner.putStatementResult( - StatementResult.query(statement2, generateKeyValueResultSet(1, 3))); + StatementResult.query(statement2, generateKeyValueResultSet(ContiguousSet.closed(1, 3)))); ApiFuture> values1; ApiFuture> values2; @@ -333,4 +336,140 @@ public int compare(String o1, String o2) { executor); assertThat(allValues.get()).containsExactly("v1", "v2", "v3", "v10", "v11", "v12"); } + + @Test + public void pauseResume() throws Exception { + Statement unevenStatement = + Statement.of("SELECT * FROM TestTable WHERE MOD(CAST(SUBSTR(Key, 2) AS INT64), 2) = 1"); + Statement evenStatement = + Statement.of("SELECT * FROM TestTable WHERE MOD(CAST(SUBSTR(Key, 2) AS INT64), 2) = 0"); + mockSpanner.putStatementResult( + StatementResult.query( + unevenStatement, generateKeyValueResultSet(ImmutableSet.of(1, 3, 5, 7, 9)))); + mockSpanner.putStatementResult( + StatementResult.query( + evenStatement, generateKeyValueResultSet(ImmutableSet.of(2, 4, 6, 8, 10)))); + + final SettableApiFuture evenFinished = SettableApiFuture.create(); + final SettableApiFuture unevenFinished = SettableApiFuture.create(); + final List allValues = new LinkedList<>(); + try (ReadOnlyTransaction tx = client.readOnlyTransaction()) { + try (AsyncResultSet evenRs = tx.executeQueryAsync(evenStatement); + AsyncResultSet unevenRs = tx.executeQueryAsync(unevenStatement)) { + evenRs.setCallback( + executor, + new ReadyCallback() { + private boolean firstRow = true; + + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + if (firstRow) { + // Make sure the uneven result set returns the first result. + firstRow = false; + return CallbackResponse.PAUSE; + } + try { + while (true) { + switch (resultSet.tryNext()) { + case DONE: + evenFinished.set(true); + return CallbackResponse.DONE; + case NOT_READY: + return CallbackResponse.CONTINUE; + case OK: + allValues.add(resultSet.getString("Value")); + return CallbackResponse.PAUSE; + } + } + } catch (Throwable t) { + evenFinished.setException(t); + return CallbackResponse.DONE; + } finally { + unevenRs.resume(); + } + } + }); + + unevenRs.setCallback( + executor, + new ReadyCallback() { + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + try { + while (true) { + switch (resultSet.tryNext()) { + case DONE: + unevenFinished.set(true); + return CallbackResponse.DONE; + case NOT_READY: + return CallbackResponse.CONTINUE; + case OK: + allValues.add(resultSet.getString("Value")); + return CallbackResponse.PAUSE; + } + } + } catch (Throwable t) { + unevenFinished.setException(t); + return CallbackResponse.DONE; + } finally { + evenRs.resume(); + } + } + }); + } + } + assertThat(ApiFutures.allAsList(Arrays.asList(evenFinished, unevenFinished)).get()) + .containsExactly(Boolean.TRUE, Boolean.TRUE); + assertThat(allValues) + .containsExactly("v1", "v2", "v3", "v4", "v5", "v6", "v7", "v8", "v9", "v10"); + } + + @Test + public void cancel() throws Exception { + final List values = new LinkedList<>(); + final SettableApiFuture finished = SettableApiFuture.create(); + final CountDownLatch receivedFirstRow = new CountDownLatch(1); + final CountDownLatch cancelled = new CountDownLatch(1); + try (AsyncResultSet rs = + client.singleUse().readAsync(READ_TABLE_NAME, KeySet.all(), READ_COLUMN_NAMES)) { + rs.setCallback( + executor, + new ReadyCallback() { + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + try { + while (true) { + switch (resultSet.tryNext()) { + case DONE: + finished.set(true); + return CallbackResponse.DONE; + case NOT_READY: + return CallbackResponse.CONTINUE; + case OK: + values.add(resultSet.getString("Value")); + receivedFirstRow.countDown(); + cancelled.await(); + break; + } + } + } catch (Throwable t) { + finished.setException(t); + return CallbackResponse.DONE; + } + } + }); + receivedFirstRow.await(); + rs.cancel(); + } + cancelled.countDown(); + try { + finished.get(); + fail("missing expected exception"); + } catch (ExecutionException e) { + assertThat(e.getCause()).isInstanceOf(SpannerException.class); + SpannerException se = (SpannerException) e.getCause(); + assertThat(se.getErrorCode()).isEqualTo(ErrorCode.CANCELLED); + assertThat(values).containsExactly("v1"); + } + } } diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITAsyncExamplesTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITAsyncExamplesTest.java index 7c80632ed26..1f9b0dd81d3 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITAsyncExamplesTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITAsyncExamplesTest.java @@ -17,6 +17,7 @@ package com.google.cloud.spanner.it; import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.fail; import com.google.api.core.ApiFunction; import com.google.api.core.ApiFuture; @@ -29,12 +30,14 @@ import com.google.cloud.spanner.AsyncRunner.AsyncWork; import com.google.cloud.spanner.Database; import com.google.cloud.spanner.DatabaseClient; +import com.google.cloud.spanner.ErrorCode; import com.google.cloud.spanner.IntegrationTest; import com.google.cloud.spanner.IntegrationTestEnv; import com.google.cloud.spanner.Key; import com.google.cloud.spanner.KeySet; import com.google.cloud.spanner.Mutation; import com.google.cloud.spanner.ReadOnlyTransaction; +import com.google.cloud.spanner.SpannerException; import com.google.cloud.spanner.Statement; import com.google.cloud.spanner.Struct; import com.google.cloud.spanner.StructReader; @@ -47,6 +50,8 @@ import java.util.Comparator; import java.util.LinkedList; import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import org.junit.AfterClass; @@ -328,4 +333,137 @@ public int compare(String o1, String o2) { executor); assertThat(allValues.get()).containsExactly("v1", "v2", "v3", "v10", "v11", "v12"); } + + @Test + public void pauseResume() throws Exception { + Statement unevenStatement = + Statement.of( + "SELECT * FROM TestTable WHERE MOD(CAST(SUBSTR(Key, 2) AS INT64), 2) = 1 ORDER BY CAST(SUBSTR(Key, 2) AS INT64)"); + Statement evenStatement = + Statement.of( + "SELECT * FROM TestTable WHERE MOD(CAST(SUBSTR(Key, 2) AS INT64), 2) = 0 ORDER BY CAST(SUBSTR(Key, 2) AS INT64)"); + + final SettableApiFuture evenFinished = SettableApiFuture.create(); + final SettableApiFuture unevenFinished = SettableApiFuture.create(); + final List allValues = new LinkedList<>(); + try (ReadOnlyTransaction tx = client.readOnlyTransaction()) { + try (AsyncResultSet evenRs = tx.executeQueryAsync(evenStatement); + AsyncResultSet unevenRs = tx.executeQueryAsync(unevenStatement)) { + evenRs.setCallback( + executor, + new ReadyCallback() { + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + try { + while (true) { + switch (resultSet.tryNext()) { + case DONE: + evenFinished.set(true); + return CallbackResponse.DONE; + case NOT_READY: + return CallbackResponse.CONTINUE; + case OK: + allValues.add(resultSet.getString("StringValue")); + return CallbackResponse.PAUSE; + } + } + } catch (Throwable t) { + evenFinished.setException(t); + return CallbackResponse.DONE; + } finally { + unevenRs.resume(); + } + } + }); + + unevenRs.setCallback( + executor, + new ReadyCallback() { + private boolean firstRow = true; + + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + if (firstRow) { + // Make sure the even result set returns the first result. + firstRow = false; + return CallbackResponse.PAUSE; + } + try { + while (true) { + switch (resultSet.tryNext()) { + case DONE: + unevenFinished.set(true); + return CallbackResponse.DONE; + case NOT_READY: + return CallbackResponse.CONTINUE; + case OK: + allValues.add(resultSet.getString("StringValue")); + return CallbackResponse.PAUSE; + } + } + } catch (Throwable t) { + unevenFinished.setException(t); + return CallbackResponse.DONE; + } finally { + evenRs.resume(); + } + } + }); + } + } + assertThat(ApiFutures.allAsList(Arrays.asList(evenFinished, unevenFinished)).get()) + .containsExactly(Boolean.TRUE, Boolean.TRUE); + assertThat(allValues) + .containsExactly( + "v0", "v1", "v2", "v3", "v4", "v5", "v6", "v7", "v8", "v9", "v10", "v11", "v12", "v13", + "v14"); + } + + @Test + public void cancel() throws Exception { + final List values = new LinkedList<>(); + final SettableApiFuture finished = SettableApiFuture.create(); + final CountDownLatch receivedFirstRow = new CountDownLatch(1); + final CountDownLatch cancelled = new CountDownLatch(1); + try (AsyncResultSet rs = client.singleUse().readAsync(TABLE_NAME, KeySet.all(), ALL_COLUMNS)) { + rs.setCallback( + executor, + new ReadyCallback() { + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + try { + while (true) { + switch (resultSet.tryNext()) { + case DONE: + finished.set(true); + return CallbackResponse.DONE; + case NOT_READY: + return CallbackResponse.CONTINUE; + case OK: + values.add(resultSet.getString("StringValue")); + receivedFirstRow.countDown(); + cancelled.await(); + break; + } + } + } catch (Throwable t) { + finished.setException(t); + return CallbackResponse.DONE; + } + } + }); + receivedFirstRow.await(); + rs.cancel(); + } + cancelled.countDown(); + try { + finished.get(); + fail("missing expected exception"); + } catch (ExecutionException e) { + assertThat(e.getCause()).isInstanceOf(SpannerException.class); + SpannerException se = (SpannerException) e.getCause(); + assertThat(se.getErrorCode()).isEqualTo(ErrorCode.CANCELLED); + assertThat(values).containsExactly("v0"); + } + } } From fa61e7d14a3cd2a26a396107ce2b5a551a8587ba Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Fri, 28 Feb 2020 21:54:49 +0100 Subject: [PATCH 14/49] tests: fix flaky tests --- .../google/cloud/spanner/AsyncRunnerTest.java | 9 +-- .../google/cloud/spanner/ReadAsyncTest.java | 59 ++++++++++++------- .../cloud/spanner/it/ITAsyncExamplesTest.java | 43 +++++++++----- 3 files changed, 71 insertions(+), 40 deletions(-) diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncRunnerTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncRunnerTest.java index d8267dcb4e3..b20cd3b7a8a 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncRunnerTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncRunnerTest.java @@ -258,14 +258,15 @@ public void asyncRunnerCommitAborted() throws Exception { new AsyncWork() { @Override public ApiFuture doWorkAsync(TransactionContext txn) { - ApiFuture updateCount = txn.executeUpdateAsync(UPDATE_STATEMENT); - if (attempt.incrementAndGet() == 1) { - mockSpanner.abortTransaction(txn); - } else { + if (attempt.get() > 0) { // Set the result of the update statement back to 1 row. mockSpanner.putStatementResult( StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT)); } + ApiFuture updateCount = txn.executeUpdateAsync(UPDATE_STATEMENT); + if (attempt.incrementAndGet() == 1) { + mockSpanner.abortTransaction(txn); + } return updateCount; } }, diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadAsyncTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadAsyncTest.java index 53e5041891a..944a7d70f51 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadAsyncTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadAsyncTest.java @@ -42,9 +42,11 @@ import java.util.ArrayList; import java.util.Arrays; import java.util.Comparator; +import java.util.Deque; import java.util.LinkedList; import java.util.List; import java.util.concurrent.BlockingQueue; +import java.util.concurrent.ConcurrentLinkedDeque; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; @@ -88,7 +90,7 @@ public static void setup() throws Exception { .build() .start(); channelProvider = LocalChannelProvider.create(uniqueName); - executor = Executors.newSingleThreadExecutor(); + executor = Executors.newScheduledThreadPool(8); } @AfterClass @@ -350,72 +352,85 @@ public void pauseResume() throws Exception { StatementResult.query( evenStatement, generateKeyValueResultSet(ImmutableSet.of(2, 4, 6, 8, 10)))); + final Object lock = new Object(); final SettableApiFuture evenFinished = SettableApiFuture.create(); final SettableApiFuture unevenFinished = SettableApiFuture.create(); - final List allValues = new LinkedList<>(); + final CountDownLatch unevenReturnedFirstRow = new CountDownLatch(1); + final Deque allValues = new ConcurrentLinkedDeque<>(); try (ReadOnlyTransaction tx = client.readOnlyTransaction()) { try (AsyncResultSet evenRs = tx.executeQueryAsync(evenStatement); AsyncResultSet unevenRs = tx.executeQueryAsync(unevenStatement)) { - evenRs.setCallback( + unevenRs.setCallback( executor, new ReadyCallback() { - private boolean firstRow = true; - @Override public CallbackResponse cursorReady(AsyncResultSet resultSet) { - if (firstRow) { - // Make sure the uneven result set returns the first result. - firstRow = false; - return CallbackResponse.PAUSE; - } try { while (true) { switch (resultSet.tryNext()) { case DONE: - evenFinished.set(true); + unevenFinished.set(true); return CallbackResponse.DONE; case NOT_READY: return CallbackResponse.CONTINUE; case OK: - allValues.add(resultSet.getString("Value")); + synchronized (lock) { + allValues.add(resultSet.getString("Value")); + } + unevenReturnedFirstRow.countDown(); return CallbackResponse.PAUSE; } } } catch (Throwable t) { - evenFinished.setException(t); + unevenFinished.setException(t); return CallbackResponse.DONE; - } finally { - unevenRs.resume(); } } }); - - unevenRs.setCallback( + evenRs.setCallback( executor, new ReadyCallback() { @Override public CallbackResponse cursorReady(AsyncResultSet resultSet) { try { + // Make sure the uneven result set has returned the first before we start the even + // results. + unevenReturnedFirstRow.await(); while (true) { switch (resultSet.tryNext()) { case DONE: - unevenFinished.set(true); + evenFinished.set(true); return CallbackResponse.DONE; case NOT_READY: return CallbackResponse.CONTINUE; case OK: - allValues.add(resultSet.getString("Value")); + synchronized (lock) { + allValues.add(resultSet.getString("Value")); + } return CallbackResponse.PAUSE; } } } catch (Throwable t) { - unevenFinished.setException(t); + evenFinished.setException(t); return CallbackResponse.DONE; - } finally { - evenRs.resume(); } } }); + while (!(evenFinished.isDone() && unevenFinished.isDone())) { + synchronized (lock) { + if (allValues.peekLast() != null) { + if (Integer.valueOf(allValues.peekLast().substring(1)) % 2 == 1) { + evenRs.resume(); + } else { + unevenRs.resume(); + } + } + if (allValues.size() == 10) { + unevenRs.resume(); + evenRs.resume(); + } + } + } } } assertThat(ApiFutures.allAsList(Arrays.asList(evenFinished, unevenFinished)).get()) diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITAsyncExamplesTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITAsyncExamplesTest.java index 1f9b0dd81d3..c328f16ef9c 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITAsyncExamplesTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITAsyncExamplesTest.java @@ -48,6 +48,7 @@ import java.util.ArrayList; import java.util.Arrays; import java.util.Comparator; +import java.util.Deque; import java.util.LinkedList; import java.util.List; import java.util.concurrent.CountDownLatch; @@ -343,9 +344,11 @@ public void pauseResume() throws Exception { Statement.of( "SELECT * FROM TestTable WHERE MOD(CAST(SUBSTR(Key, 2) AS INT64), 2) = 0 ORDER BY CAST(SUBSTR(Key, 2) AS INT64)"); + final Object lock = new Object(); final SettableApiFuture evenFinished = SettableApiFuture.create(); final SettableApiFuture unevenFinished = SettableApiFuture.create(); - final List allValues = new LinkedList<>(); + final CountDownLatch evenReturnedFirstRow = new CountDownLatch(1); + final Deque allValues = new LinkedList<>(); try (ReadOnlyTransaction tx = client.readOnlyTransaction()) { try (AsyncResultSet evenRs = tx.executeQueryAsync(evenStatement); AsyncResultSet unevenRs = tx.executeQueryAsync(unevenStatement)) { @@ -363,15 +366,16 @@ public CallbackResponse cursorReady(AsyncResultSet resultSet) { case NOT_READY: return CallbackResponse.CONTINUE; case OK: - allValues.add(resultSet.getString("StringValue")); + synchronized (lock) { + allValues.add(resultSet.getString("StringValue")); + } + evenReturnedFirstRow.countDown(); return CallbackResponse.PAUSE; } } } catch (Throwable t) { evenFinished.setException(t); return CallbackResponse.DONE; - } finally { - unevenRs.resume(); } } }); @@ -379,16 +383,12 @@ public CallbackResponse cursorReady(AsyncResultSet resultSet) { unevenRs.setCallback( executor, new ReadyCallback() { - private boolean firstRow = true; - @Override public CallbackResponse cursorReady(AsyncResultSet resultSet) { - if (firstRow) { - // Make sure the even result set returns the first result. - firstRow = false; - return CallbackResponse.PAUSE; - } try { + // Make sure the even result set has returned the first before we start the uneven + // results. + evenReturnedFirstRow.await(); while (true) { switch (resultSet.tryNext()) { case DONE: @@ -397,18 +397,33 @@ public CallbackResponse cursorReady(AsyncResultSet resultSet) { case NOT_READY: return CallbackResponse.CONTINUE; case OK: - allValues.add(resultSet.getString("StringValue")); + synchronized (lock) { + allValues.add(resultSet.getString("StringValue")); + } return CallbackResponse.PAUSE; } } } catch (Throwable t) { unevenFinished.setException(t); return CallbackResponse.DONE; - } finally { - evenRs.resume(); } } }); + while (!(evenFinished.isDone() && unevenFinished.isDone())) { + synchronized (lock) { + if (allValues.peekLast() != null) { + if (Integer.valueOf(allValues.peekLast().substring(1)) % 2 == 1) { + evenRs.resume(); + } else { + unevenRs.resume(); + } + } + if (allValues.size() == 15) { + unevenRs.resume(); + evenRs.resume(); + } + } + } } } assertThat(ApiFutures.allAsList(Arrays.asList(evenFinished, unevenFinished)).get()) From fc53dbf2a62f199466da10655c2c355b36745144 Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Mon, 16 Mar 2020 14:54:46 +0100 Subject: [PATCH 15/49] rebase: rebase on current master --- .../com/google/cloud/spanner/AbstractReadContext.java | 8 ++++++++ .../java/com/google/cloud/spanner/BatchClientImpl.java | 2 ++ .../java/com/google/cloud/spanner/SessionImpl.java | 8 +++++++- .../java/com/google/cloud/spanner/SpannerOptions.java | 3 ++- .../google/cloud/spanner/TransactionRunnerImpl.java | 1 - .../google/cloud/spanner/AbstractReadContextTest.java | 2 ++ .../google/cloud/spanner/DatabaseClientImplTest.java | 10 ++-------- .../java/com/google/cloud/spanner/SessionPoolTest.java | 1 - .../cloud/spanner/TransactionContextImplTest.java | 2 -- .../cloud/spanner/TransactionRunnerImplTest.java | 1 - 10 files changed, 23 insertions(+), 15 deletions(-) diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java index f191a858a11..9fff15ed62f 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java @@ -49,6 +49,7 @@ import com.google.spanner.v1.TransactionOptions; import com.google.spanner.v1.TransactionSelector; import io.opencensus.trace.Span; +import io.opencensus.trace.Tracing; import java.util.Map; import java.util.concurrent.atomic.AtomicLong; import javax.annotation.Nullable; @@ -67,6 +68,7 @@ abstract static class Builder, T extends AbstractReadCon private Span span = Tracing.getTracer().getCurrentSpan(); private int defaultPrefetchChunks = SpannerOptions.Builder.DEFAULT_PREFETCH_CHUNKS; private QueryOptions defaultQueryOptions = SpannerOptions.Builder.DEFAULT_QUERY_OPTIONS; + private ExecutorProvider executorProvider; Builder() {} @@ -100,6 +102,11 @@ B setDefaultQueryOptions(QueryOptions defaultQueryOptions) { return self(); } + B setExecutorProvider(ExecutorProvider executorProvider) { + this.executorProvider = executorProvider; + return self(); + } + abstract T build(); } @@ -391,6 +398,7 @@ void initTransaction() { this.defaultPrefetchChunks = builder.defaultPrefetchChunks; this.defaultQueryOptions = builder.defaultQueryOptions; this.span = builder.span; + this.executorProvider = builder.executorProvider; } @Override diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/BatchClientImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/BatchClientImpl.java index 39827647b6b..c84bef77cf8 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/BatchClientImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/BatchClientImpl.java @@ -52,6 +52,7 @@ public BatchReadOnlyTransaction batchReadOnlyTransaction(TimestampBound bound) { .setTimestampBound(bound) .setDefaultQueryOptions( sessionClient.getSpanner().getDefaultQueryOptions(sessionClient.getDatabaseId())) + .setExecutorProvider(sessionClient.getSpanner().getAsyncExecutorProvider()) .setDefaultPrefetchChunks(sessionClient.getSpanner().getDefaultPrefetchChunks()), checkNotNull(bound)); } @@ -68,6 +69,7 @@ public BatchReadOnlyTransaction batchReadOnlyTransaction(BatchTransactionId batc .setTimestamp(batchTransactionId.getTimestamp()) .setDefaultQueryOptions( sessionClient.getSpanner().getDefaultQueryOptions(sessionClient.getDatabaseId())) + .setExecutorProvider(sessionClient.getSpanner().getAsyncExecutorProvider()) .setDefaultPrefetchChunks(sessionClient.getSpanner().getDefaultPrefetchChunks()), batchTransactionId); } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java index 1d9c99eeb59..77f8946bb12 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java @@ -25,7 +25,6 @@ import com.google.cloud.spanner.AbstractReadContext.SingleReadContext; import com.google.cloud.spanner.AbstractReadContext.SingleUseReadOnlyTransaction; import com.google.cloud.spanner.SessionClient.SessionId; -import com.google.cloud.spanner.TransactionRunner.TransactionCallable; import com.google.cloud.spanner.TransactionRunnerImpl.TransactionContextImpl; import com.google.cloud.spanner.spi.v1.SpannerRpc; import com.google.common.collect.Lists; @@ -177,6 +176,8 @@ public ReadContext singleUse(TimestampBound bound) { .setRpc(spanner.getRpc()) .setDefaultQueryOptions(spanner.getDefaultQueryOptions(databaseId)) .setDefaultPrefetchChunks(spanner.getDefaultPrefetchChunks()) + .setSpan(currentSpan) + .setExecutorProvider(spanner.getAsyncExecutorProvider()) .build()); } @@ -194,6 +195,8 @@ public ReadOnlyTransaction singleUseReadOnlyTransaction(TimestampBound bound) { .setRpc(spanner.getRpc()) .setDefaultQueryOptions(spanner.getDefaultQueryOptions(databaseId)) .setDefaultPrefetchChunks(spanner.getDefaultPrefetchChunks()) + .setSpan(currentSpan) + .setExecutorProvider(spanner.getAsyncExecutorProvider()) .buildSingleUseReadOnlyTransaction()); } @@ -211,6 +214,8 @@ public ReadOnlyTransaction readOnlyTransaction(TimestampBound bound) { .setRpc(spanner.getRpc()) .setDefaultQueryOptions(spanner.getDefaultQueryOptions(databaseId)) .setDefaultPrefetchChunks(spanner.getDefaultPrefetchChunks()) + .setSpan(currentSpan) + .setExecutorProvider(spanner.getAsyncExecutorProvider()) .build()); } @@ -281,6 +286,7 @@ TransactionContextImpl newTransaction() { .setDefaultQueryOptions(spanner.getDefaultQueryOptions(databaseId)) .setDefaultPrefetchChunks(spanner.getDefaultPrefetchChunks()) .setSpan(currentSpan) + .setExecutorProvider(spanner.getAsyncExecutorProvider()) .build(); } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java index efe9b381ade..3b73158bac7 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java @@ -51,6 +51,7 @@ import com.google.spanner.admin.database.v1.RestoreDatabaseRequest; import com.google.spanner.v1.ExecuteSqlRequest.QueryOptions; import com.google.common.util.concurrent.ThreadFactoryBuilder; +import com.google.spanner.v1.ExecuteSqlRequest.QueryOptions; import io.grpc.CallCredentials; import io.grpc.ManagedChannelBuilder; import java.io.IOException; @@ -60,11 +61,11 @@ import java.util.Map; import java.util.Map.Entry; import java.util.Set; -import javax.annotation.Nonnull; import java.util.concurrent.ScheduledThreadPoolExecutor; import java.util.concurrent.ThreadFactory; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; +import javax.annotation.Nonnull; import org.threeten.bp.Duration; /** Options for the Cloud Spanner service. */ 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 e4fae3e3a68..0f098f4f6ef 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 @@ -24,7 +24,6 @@ import com.google.api.core.ApiFunction; import com.google.api.core.ApiFuture; import com.google.api.core.ApiFutures; -import com.google.api.gax.core.ExecutorProvider; import com.google.cloud.Timestamp; import com.google.cloud.spanner.Options.QueryOption; import com.google.cloud.spanner.Options.ReadOption; diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AbstractReadContextTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AbstractReadContextTest.java index bfd739d5530..f9cee1c4889 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AbstractReadContextTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AbstractReadContextTest.java @@ -20,6 +20,7 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; +import com.google.api.gax.core.ExecutorProvider; import com.google.cloud.spanner.spi.v1.SpannerRpc; import com.google.spanner.v1.ExecuteSqlRequest; import com.google.spanner.v1.ExecuteSqlRequest.QueryMode; @@ -80,6 +81,7 @@ public void setup() { .setSession(session) .setRpc(mock(SpannerRpc.class)) .setDefaultQueryOptions(defaultQueryOptions) + .setExecutorProvider(mock(ExecutorProvider.class)) .build(); } diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java index 7d7b919fc18..f6814af8dc3 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java @@ -36,25 +36,19 @@ import com.google.cloud.spanner.ReadContext.QueryAnalyzeMode; import com.google.cloud.spanner.TransactionRunner.TransactionCallable; import com.google.common.base.Stopwatch; -import com.google.protobuf.AbstractMessage; import com.google.common.util.concurrent.SettableFuture; -import com.google.protobuf.ListValue; +import com.google.protobuf.AbstractMessage; import com.google.spanner.v1.ExecuteSqlRequest; import com.google.spanner.v1.ExecuteSqlRequest.QueryMode; import com.google.spanner.v1.ExecuteSqlRequest.QueryOptions; -import com.google.spanner.v1.ResultSetMetadata; -import com.google.spanner.v1.StructType; -import com.google.spanner.v1.StructType.Field; -import com.google.spanner.v1.TypeCode; import io.grpc.Server; import io.grpc.Status; import io.grpc.StatusRuntimeException; import io.grpc.inprocess.InProcessServerBuilder; import java.io.IOException; -import java.util.List; import java.util.ArrayList; import java.util.Arrays; -import java.util.concurrent.CountDownLatch; +import java.util.List; import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java index 54250420833..fc85384453c 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java @@ -35,7 +35,6 @@ import com.google.api.core.ApiFuture; import com.google.api.core.ApiFutures; -import com.google.api.gax.core.ExecutorProvider; import com.google.cloud.Timestamp; import com.google.cloud.spanner.MetricRegistryTestUtils.FakeMetricRegistry; import com.google.cloud.spanner.MetricRegistryTestUtils.MetricsRecord; diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionContextImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionContextImplTest.java index 3446ba9fc84..077b6605766 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionContextImplTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionContextImplTest.java @@ -19,7 +19,6 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; -import com.google.api.gax.core.ExecutorProvider; import com.google.cloud.spanner.TransactionRunnerImpl.TransactionContextImpl; import com.google.cloud.spanner.spi.v1.SpannerRpc; import com.google.protobuf.ByteString; @@ -27,7 +26,6 @@ import com.google.rpc.Status; import com.google.spanner.v1.ExecuteBatchDmlRequest; import com.google.spanner.v1.ExecuteBatchDmlResponse; -import io.opencensus.trace.Span; import java.util.Arrays; import org.junit.Test; import org.junit.runner.RunWith; diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java index 8bcc05aa001..b790ed93b2e 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java @@ -25,7 +25,6 @@ import static org.mockito.Mockito.when; import com.google.api.core.ApiFutures; -import com.google.api.gax.core.ExecutorProvider; import com.google.cloud.grpc.GrpcTransportOptions; import com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory; import com.google.cloud.spanner.SessionClient.SessionId; From 0eee1f64f3767574608237f31d7a8cb1d7df6601 Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Fri, 20 Mar 2020 17:25:57 +0100 Subject: [PATCH 16/49] fix: run code formatter --- .../src/main/java/com/google/cloud/spanner/SpannerOptions.java | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java index 3b73158bac7..34b5e728e1e 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java @@ -46,12 +46,11 @@ import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; +import com.google.common.util.concurrent.ThreadFactoryBuilder; import com.google.spanner.admin.database.v1.CreateBackupRequest; import com.google.spanner.admin.database.v1.CreateDatabaseRequest; import com.google.spanner.admin.database.v1.RestoreDatabaseRequest; import com.google.spanner.v1.ExecuteSqlRequest.QueryOptions; -import com.google.common.util.concurrent.ThreadFactoryBuilder; -import com.google.spanner.v1.ExecuteSqlRequest.QueryOptions; import io.grpc.CallCredentials; import io.grpc.ManagedChannelBuilder; import java.io.IOException; From d3d2ffc49f2c6a73b08af77d420b12a3ff2065ab Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Mon, 23 Mar 2020 18:52:39 +0100 Subject: [PATCH 17/49] feat: add support for poller --- .../clirr-ignored-differences.xml | 5 ++ .../cloud/spanner/MockSpannerServiceImpl.java | 46 ++++++++++++++++--- 2 files changed, 45 insertions(+), 6 deletions(-) diff --git a/google-cloud-spanner/clirr-ignored-differences.xml b/google-cloud-spanner/clirr-ignored-differences.xml index d1096aaa6a1..4695ca92700 100644 --- a/google-cloud-spanner/clirr-ignored-differences.xml +++ b/google-cloud-spanner/clirr-ignored-differences.xml @@ -147,6 +147,11 @@ + + 7012 + com/google/cloud/spanner/spi/v1/SpannerRpc + com.google.api.core.ApiFuture executeQueryAsync(com.google.spanner.v1.ExecuteSqlRequest, java.util.Map) + 7012 com/google/cloud/spanner/DatabaseClient 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 7cb75673213..1fe8ab4170a 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 @@ -79,6 +79,7 @@ import java.util.Collection; import java.util.Collections; import java.util.Comparator; +import java.util.Deque; import java.util.Iterator; import java.util.LinkedList; import java.util.List; @@ -230,7 +231,7 @@ private enum StatementResultType { private final StatementResultType type; private final Statement statement; private final Long updateCount; - private final ResultSet resultSet; + private final Deque resultSets; private final StatusRuntimeException exception; /** Creates a {@link StatementResult} for a query that returns a {@link ResultSet}. */ @@ -238,6 +239,11 @@ public static StatementResult query(Statement statement, ResultSet resultSet) { return new StatementResult(statement, resultSet); } + /** Creates a {@link StatementResult} for a query that returns a {@link ResultSet} the first time, and a different {@link ResultSet} for all subsequent calls. */ + public static StatementResult queryAndThen(Statement statement, ResultSet resultSet, ResultSet next) { + return new StatementResult(statement, resultSet); + } + /** Creates a {@link StatementResult} for a read request. */ public static StatementResult read( String table, KeySet keySet, Iterable columns, ResultSet resultSet) { @@ -254,6 +260,25 @@ public static StatementResult exception(Statement statement, StatusRuntimeExcept return new StatementResult(statement, exception); } + private static class KeepLastElementDeque extends LinkedList { + private static KeepLastElementDeque singleton(E item) { + return new KeepLastElementDeque(Collections.singleton(item)); + } + + private static KeepLastElementDeque of(E first, E second) { + return new KeepLastElementDeque(Arrays.asList(first, second)); + } + + private KeepLastElementDeque(Collection coll) { + super(coll); + } + + @Override + public E pop() { + return this.size() == 1 ? super.peek() : super.pop(); + } + } + /** * Creates a {@link Statement} for a read statement. This {@link Statement} can be used to mock * a result for a read request. @@ -301,14 +326,22 @@ private static boolean isValidKeySet(KeySet keySet) { private StatementResult(Statement statement, Long updateCount) { this.statement = Preconditions.checkNotNull(statement); this.updateCount = Preconditions.checkNotNull(updateCount); - this.resultSet = null; + this.resultSets = null; this.exception = null; this.type = StatementResultType.UPDATE_COUNT; } private StatementResult(Statement statement, ResultSet resultSet) { this.statement = Preconditions.checkNotNull(statement); - this.resultSet = Preconditions.checkNotNull(resultSet); + this.resultSets = KeepLastElementDeque.singleton(Preconditions.checkNotNull(resultSet)); + this.updateCount = null; + this.exception = null; + this.type = StatementResultType.RESULT_SET; + } + + private StatementResult(Statement statement, ResultSet resultSet, ResultSet andThen) { + this.statement = Preconditions.checkNotNull(statement); + this.resultSets = KeepLastElementDeque.of(Preconditions.checkNotNull(resultSet), Preconditions.checkNotNull(andThen)); this.updateCount = null; this.exception = null; this.type = StatementResultType.RESULT_SET; @@ -317,7 +350,7 @@ private StatementResult(Statement statement, ResultSet resultSet) { private StatementResult( String table, KeySet keySet, Iterable columns, ResultSet resultSet) { this.statement = createReadStatement(table, keySet, columns); - this.resultSet = Preconditions.checkNotNull(resultSet); + this.resultSets = KeepLastElementDeque.singleton(Preconditions.checkNotNull(resultSet)); this.updateCount = null; this.exception = null; this.type = StatementResultType.RESULT_SET; @@ -326,7 +359,7 @@ private StatementResult( private StatementResult(Statement statement, StatusRuntimeException exception) { this.statement = Preconditions.checkNotNull(statement); this.exception = Preconditions.checkNotNull(exception); - this.resultSet = null; + this.resultSets = null; this.updateCount = null; this.type = StatementResultType.EXCEPTION; } @@ -339,7 +372,7 @@ private ResultSet getResultSet() { Preconditions.checkState( type == StatementResultType.RESULT_SET, "This statement result does not contain a result set"); - return resultSet; + return resultSets.pop(); } private Long getUpdateCount() { @@ -1102,6 +1135,7 @@ private Statement buildStatement( case STRUCT: throw new IllegalArgumentException("Struct parameters not (yet) supported"); case TIMESTAMP: + builder.bind(entry.getKey()).to(com.google.cloud.Timestamp.parseTimestamp(value.getStringValue())); break; case TYPE_CODE_UNSPECIFIED: case UNRECOGNIZED: From 2e01ca71854f6d7e39ad0e9cf6bb184c2699a7b9 Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Thu, 2 Apr 2020 08:39:18 +0200 Subject: [PATCH 18/49] tests: support more param types --- .../cloud/spanner/AbstractResultSet.java | 12 +- .../cloud/spanner/MockSpannerServiceImpl.java | 117 +++++++++++++++++- 2 files changed, 117 insertions(+), 12 deletions(-) diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractResultSet.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractResultSet.java index 7b248bfb9d4..6b0681b5889 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractResultSet.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractResultSet.java @@ -495,7 +495,7 @@ private static Struct decodeStructValue(Type structType, ListValue structValue) return new GrpcStruct(structType, fields); } - private static Object decodeArrayValue(Type elementType, ListValue listValue) { + static Object decodeArrayValue(Type elementType, ListValue listValue) { switch (elementType.getCode()) { case BOOL: // Use a view: element conversion is virtually free. @@ -1009,7 +1009,7 @@ protected PartialResultSet computeNext() { } } - private static double valueProtoToFloat64(com.google.protobuf.Value proto) { + static double valueProtoToFloat64(com.google.protobuf.Value proto) { if (proto.getKindCase() == KindCase.STRING_VALUE) { switch (proto.getStringValue()) { case "-Infinity": @@ -1037,7 +1037,7 @@ private static double valueProtoToFloat64(com.google.protobuf.Value proto) { return proto.getNumberValue(); } - private static NullPointerException throwNotNull(int columnIndex) { + static NullPointerException throwNotNull(int columnIndex) { throw new NullPointerException( "Cannot call array getter for column " + columnIndex + " with null elements"); } @@ -1048,7 +1048,7 @@ private static NullPointerException throwNotNull(int columnIndex) { * {@code BigDecimal} respectively. Rather than construct new wrapper objects for each array * element, we use primitive arrays and a {@code BitSet} to track nulls. */ - private abstract static class PrimitiveArray extends AbstractList { + abstract static class PrimitiveArray extends AbstractList { private final A data; private final BitSet nulls; private final int size; @@ -1103,7 +1103,7 @@ A toPrimitiveArray(int columnIndex) { } } - private static class Int64Array extends PrimitiveArray { + static class Int64Array extends PrimitiveArray { Int64Array(ListValue protoList) { super(protoList); } @@ -1128,7 +1128,7 @@ Long get(long[] array, int i) { } } - private static class Float64Array extends PrimitiveArray { + static class Float64Array extends PrimitiveArray { Float64Array(ListValue protoList) { super(protoList); } 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 1fe8ab4170a..fb48bb8cd02 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 @@ -19,6 +19,7 @@ import com.google.api.gax.grpc.testing.MockGrpcService; import com.google.cloud.ByteArray; import com.google.cloud.Date; +import com.google.cloud.spanner.AbstractResultSet.GrpcStruct; import com.google.cloud.spanner.TransactionRunnerImpl.TransactionContextImpl; import com.google.common.base.Optional; import com.google.common.base.Preconditions; @@ -239,8 +240,12 @@ public static StatementResult query(Statement statement, ResultSet resultSet) { return new StatementResult(statement, resultSet); } - /** Creates a {@link StatementResult} for a query that returns a {@link ResultSet} the first time, and a different {@link ResultSet} for all subsequent calls. */ - public static StatementResult queryAndThen(Statement statement, ResultSet resultSet, ResultSet next) { + /** + * Creates a {@link StatementResult} for a query that returns a {@link ResultSet} the first + * time, and a different {@link ResultSet} for all subsequent calls. + */ + public static StatementResult queryAndThen( + Statement statement, ResultSet resultSet, ResultSet next) { return new StatementResult(statement, resultSet); } @@ -341,7 +346,9 @@ private StatementResult(Statement statement, ResultSet resultSet) { private StatementResult(Statement statement, ResultSet resultSet, ResultSet andThen) { this.statement = Preconditions.checkNotNull(statement); - this.resultSets = KeepLastElementDeque.of(Preconditions.checkNotNull(resultSet), Preconditions.checkNotNull(andThen)); + this.resultSets = + KeepLastElementDeque.of( + Preconditions.checkNotNull(resultSet), Preconditions.checkNotNull(andThen)); this.updateCount = null; this.exception = null; this.type = StatementResultType.RESULT_SET; @@ -1071,6 +1078,7 @@ public void executeStreamingSql( } } + @SuppressWarnings("unchecked") private Statement buildStatement( String sql, Map paramTypes, com.google.protobuf.Struct params) { Statement.Builder builder = Statement.newBuilder(sql); @@ -1079,7 +1087,37 @@ private Statement buildStatement( if (value.getKindCase() == KindCase.NULL_VALUE) { switch (entry.getValue().getCode()) { case ARRAY: - throw new IllegalArgumentException("Array parameters not (yet) supported"); + switch (entry.getValue().getArrayElementType().getCode()) { + case BOOL: + builder.bind(entry.getKey()).toBoolArray((Iterable) null); + break; + case BYTES: + builder.bind(entry.getKey()).toBytesArray(null); + break; + case DATE: + builder.bind(entry.getKey()).toDateArray(null); + break; + case FLOAT64: + builder.bind(entry.getKey()).toFloat64Array((Iterable) null); + break; + case INT64: + builder.bind(entry.getKey()).toInt64Array((Iterable) null); + break; + case STRING: + builder.bind(entry.getKey()).toStringArray(null); + break; + case TIMESTAMP: + builder.bind(entry.getKey()).toTimestampArray(null); + break; + case STRUCT: + case TYPE_CODE_UNSPECIFIED: + case UNRECOGNIZED: + default: + throw new IllegalArgumentException( + "Unknown or invalid array parameter type: " + + entry.getValue().getArrayElementType().getCode()); + } + break; case BOOL: builder.bind(entry.getKey()).to((Boolean) null); break; @@ -1113,7 +1151,72 @@ private Statement buildStatement( } else { switch (entry.getValue().getCode()) { case ARRAY: - throw new IllegalArgumentException("Array parameters not (yet) supported"); + switch (entry.getValue().getArrayElementType().getCode()) { + case BOOL: + builder + .bind(entry.getKey()) + .toBoolArray( + (Iterable) + GrpcStruct.decodeArrayValue( + com.google.cloud.spanner.Type.bool(), value.getListValue())); + break; + case BYTES: + builder + .bind(entry.getKey()) + .toBytesArray( + (Iterable) + GrpcStruct.decodeArrayValue( + com.google.cloud.spanner.Type.bytes(), value.getListValue())); + break; + case DATE: + builder + .bind(entry.getKey()) + .toDateArray( + (Iterable) + GrpcStruct.decodeArrayValue( + com.google.cloud.spanner.Type.date(), value.getListValue())); + break; + case FLOAT64: + builder + .bind(entry.getKey()) + .toFloat64Array( + (Iterable) + GrpcStruct.decodeArrayValue( + com.google.cloud.spanner.Type.float64(), value.getListValue())); + break; + case INT64: + builder + .bind(entry.getKey()) + .toInt64Array( + (Iterable) + GrpcStruct.decodeArrayValue( + com.google.cloud.spanner.Type.int64(), value.getListValue())); + break; + case STRING: + builder + .bind(entry.getKey()) + .toStringArray( + (Iterable) + GrpcStruct.decodeArrayValue( + com.google.cloud.spanner.Type.string(), value.getListValue())); + break; + case TIMESTAMP: + builder + .bind(entry.getKey()) + .toTimestampArray( + (Iterable) + GrpcStruct.decodeArrayValue( + com.google.cloud.spanner.Type.timestamp(), value.getListValue())); + break; + case STRUCT: + case TYPE_CODE_UNSPECIFIED: + case UNRECOGNIZED: + default: + throw new IllegalArgumentException( + "Unknown or invalid array parameter type: " + + entry.getValue().getArrayElementType().getCode()); + } + break; case BOOL: builder.bind(entry.getKey()).to(value.getBoolValue()); break; @@ -1135,7 +1238,9 @@ private Statement buildStatement( case STRUCT: throw new IllegalArgumentException("Struct parameters not (yet) supported"); case TIMESTAMP: - builder.bind(entry.getKey()).to(com.google.cloud.Timestamp.parseTimestamp(value.getStringValue())); + builder + .bind(entry.getKey()) + .to(com.google.cloud.Timestamp.parseTimestamp(value.getStringValue())); break; case TYPE_CODE_UNSPECIFIED: case UNRECOGNIZED: From 5e63f2b02bd487a9b74e851f5af82e2ac85ec1cb Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Wed, 8 Apr 2020 07:02:36 +0200 Subject: [PATCH 19/49] fix: fix race conditions --- .../cloud/spanner/MockSpannerServiceImpl.java | 32 +++++++++++++------ 1 file changed, 23 insertions(+), 9 deletions(-) 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 fb48bb8cd02..d55808ece0d 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 @@ -491,6 +491,7 @@ private static void checkException(Queue exceptions, boolean keepExce private final Random random = new Random(); private double abortProbability = 0.0010D; + private final Object lock = new Object(); private final Queue requests = new ConcurrentLinkedQueue<>(); private volatile CountDownLatch freezeLock = new CountDownLatch(0); private final Queue exceptions = new ConcurrentLinkedQueue<>(); @@ -571,11 +572,24 @@ private Timestamp getCurrentGoogleTimestamp() { */ public void putStatementResult(StatementResult result) { Preconditions.checkNotNull(result); - statementResults.put(result.statement, result); + synchronized (lock) { + statementResults.put(result.statement, result); + } + } + + public void putStatementResults(StatementResult... results) { + synchronized (lock) { + for (StatementResult result : results) { + statementResults.put(result.statement, result); + } + } } private StatementResult getResult(Statement statement) { - StatementResult res = statementResults.get(statement); + StatementResult res; + synchronized (lock) { + res = statementResults.get(statement); + } if (res == null) { throw Status.INTERNAL .withDescription( @@ -1322,12 +1336,12 @@ public Iterator iterator() { return request.getColumnsList().iterator(); } }; - StatementResult res = - statementResults.get( - StatementResult.createReadStatement( - request.getTable(), - request.getKeySet().getAll() ? KeySet.all() : KeySet.singleKey(Key.of()), - cols)); + Statement statement = + StatementResult.createReadStatement( + request.getTable(), + request.getKeySet().getAll() ? KeySet.all() : KeySet.singleKey(Key.of()), + cols); + StatementResult res = getResult(statement); returnResultSet( res.getResultSet(), transactionId, request.getTransaction(), responseObserver); responseObserver.onCompleted(); @@ -1377,7 +1391,7 @@ public Iterator iterator() { request.getTable(), request.getKeySet().getAll() ? KeySet.all() : KeySet.singleKey(Key.of()), cols); - StatementResult res = statementResults.get(statement); + StatementResult res = getResult(statement); if (res == null) { throw Status.NOT_FOUND .withDescription("No result found for " + statement.toString()) From abe455e7e8bd2483ff419419f0ebe1b152c73ae8 Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Thu, 9 Apr 2020 18:20:10 +0200 Subject: [PATCH 20/49] feat: return ApiFuture to monitor end of AsyncResultSet --- .../google/cloud/spanner/AsyncResultSet.java | 8 ++- .../cloud/spanner/AsyncResultSetImpl.java | 15 +++--- .../spanner/ForwardingAsyncResultSet.java | 5 +- .../com/google/cloud/spanner/SessionPool.java | 6 ++- .../com/google/cloud/spanner/SpannerImpl.java | 11 +++- .../google/cloud/spanner/SpannerOptions.java | 53 +++++++++++++++---- .../cloud/spanner/TransactionRunnerImpl.java | 4 +- .../cloud/spanner/MockSpannerServiceImpl.java | 6 +++ 8 files changed, 83 insertions(+), 25 deletions(-) diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncResultSet.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncResultSet.java index 79c4b3b7686..0dcc379e091 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncResultSet.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncResultSet.java @@ -19,6 +19,7 @@ import com.google.api.core.ApiFuture; import com.google.common.base.Function; import com.google.common.collect.ImmutableList; +import java.util.concurrent.ExecutionException; import java.util.concurrent.Executor; /** Interface for result sets returned by async query methods. */ @@ -131,8 +132,13 @@ public enum CursorState { * RuntimeException up the stack, lest you do damage to calling components. For example, it * may cause an event dispatcher thread to crash. * @param cb ready callback + * @return An {@link ApiFuture} that returns null when the consumption of the {@link + * AsyncResultSet} has finished successfully. No more calls to the {@link ReadyCallback} will + * follow and all resources used by the {@link AsyncResultSet} have been cleaned up. The + * {@link ApiFuture} throws an {@link ExecutionException} if the consumption of the {@link + * AsyncResultSet} finished with an error. */ - void setCallback(Executor exec, ReadyCallback cb); + ApiFuture setCallback(Executor exec, ReadyCallback cb); /** * Attempt to cancel this operation and free all resources. Non-blocking. This is a no-op for 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 fcccf05aacd..a6739688157 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 @@ -17,12 +17,14 @@ package com.google.cloud.spanner; import com.google.api.core.ApiFuture; +import com.google.api.core.ListenableFutureToApiFuture; import com.google.api.core.SettableApiFuture; import com.google.api.gax.core.ExecutorProvider; import com.google.cloud.spanner.AbstractReadContext.ListenableAsyncResultSet; import com.google.common.base.Function; import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableList; +import com.google.common.util.concurrent.ListeningScheduledExecutorService; import com.google.common.util.concurrent.MoreExecutors; import com.google.spanner.v1.ResultSetStats; import java.util.Collection; @@ -34,7 +36,6 @@ import java.util.concurrent.Executor; import java.util.concurrent.Future; import java.util.concurrent.LinkedBlockingDeque; -import java.util.concurrent.ScheduledExecutorService; /** Default implementation for {@link AsyncResultSet}. */ class AsyncResultSetImpl extends ForwardingStructReader implements ListenableAsyncResultSet { @@ -74,7 +75,7 @@ private State(boolean shouldStop) { */ private final ExecutorProvider executorProvider; - private final ScheduledExecutorService service; + private final ListeningScheduledExecutorService service; private final BlockingDeque buffer; private Struct currentRow; @@ -108,7 +109,7 @@ private State(boolean shouldStop) { */ private volatile boolean finished; - private volatile Future result; + private volatile ApiFuture result; /** * {@link #cursorReturnedDoneOrException} indicates whether {@link #tryNext()} has returned {@link @@ -136,7 +137,7 @@ private State(boolean shouldStop) { super(delegate); this.buffer = new LinkedBlockingDeque<>(bufferSize); this.executorProvider = executorProvider; - this.service = executorProvider.getExecutor(); + this.service = MoreExecutors.listeningDecorator(executorProvider.getExecutor()); this.delegateResultSet = delegate; } @@ -420,18 +421,20 @@ private void startCallbackWithBufferLatchIfNecessary(int bufferLatch) { /** Sets the callback for this {@link AsyncResultSet}. */ @Override - public void setCallback(Executor exec, ReadyCallback cb) { + public ApiFuture setCallback(Executor exec, ReadyCallback cb) { synchronized (monitor) { Preconditions.checkState(!closed, "This AsyncResultSet has been closed"); Preconditions.checkState( this.state == State.INITIALIZED, "callback may not be set multiple times"); // Start to fetch data and buffer these. - this.result = this.service.submit(new ProduceRowsCallable()); + this.result = + new ListenableFutureToApiFuture<>(this.service.submit(new ProduceRowsCallable())); this.executor = MoreExecutors.newSequentialExecutor(Preconditions.checkNotNull(exec)); this.callback = Preconditions.checkNotNull(cb); this.state = State.RUNNING; pausedLatch.countDown(); + return result; } } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ForwardingAsyncResultSet.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ForwardingAsyncResultSet.java index c5535bc4490..78e35059988 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ForwardingAsyncResultSet.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ForwardingAsyncResultSet.java @@ -37,9 +37,8 @@ public CursorState tryNext() throws SpannerException { } @Override - public void setCallback(Executor exec, ReadyCallback cb) { - delegate.setCallback(exec, cb); - ; + public ApiFuture setCallback(Executor exec, ReadyCallback cb) { + return delegate.setCallback(exec, cb); } @Override diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java index ee09f28341d..a8fad6c4d0f 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java @@ -160,7 +160,7 @@ private AutoClosingReadContextAsyncResultSetImpl( } @Override - public void setCallback(Executor exec, ReadyCallback cb) { + public ApiFuture setCallback(Executor exec, ReadyCallback cb) { Runnable listener = new Runnable() { @Override @@ -178,7 +178,7 @@ public void run() { try { asyncOperationsCount.incrementAndGet(); addListener(listener); - super.setCallback(exec, cb); + return super.setCallback(exec, cb); } catch (Throwable t) { removeListener(listener); asyncOperationsCount.decrementAndGet(); @@ -2164,6 +2164,8 @@ ListenableFuture closeAsync() { readSessions.clear(); writePreparedSessions.clear(); prepareExecutor.shutdown(); + readWaiterExecutor.shutdown(); + writeWaiterExecutor.shutdown(); executor.submit( new Runnable() { @Override diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java index 4aaa8ae97bb..162f7495259 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java @@ -24,9 +24,11 @@ import com.google.cloud.PageImpl.NextPageFetcher; import com.google.cloud.grpc.GrpcTransportOptions; import com.google.cloud.spanner.SessionClient.SessionId; +import com.google.cloud.spanner.SpannerOptions.CloseableExecutorProvider; import com.google.cloud.spanner.spi.v1.SpannerRpc; import com.google.cloud.spanner.spi.v1.SpannerRpc.Paginated; import com.google.common.annotations.VisibleForTesting; +import com.google.common.base.MoreObjects; import com.google.common.base.Preconditions; import com.google.common.base.Strings; import com.google.common.collect.ImmutableList; @@ -86,6 +88,8 @@ private static String nextDatabaseClientId(DatabaseId databaseId) { @GuardedBy("this") private final Map dbClients = new HashMap<>(); + private final CloseableExecutorProvider asyncExecutorProvider; + @GuardedBy("this") private final List invalidatedDbClients = new ArrayList<>(); @@ -102,6 +106,10 @@ private static String nextDatabaseClientId(DatabaseId databaseId) { SpannerImpl(SpannerRpc gapicRpc, SpannerOptions options) { super(options); this.gapicRpc = gapicRpc; + this.asyncExecutorProvider = + MoreObjects.firstNonNull( + options.getAsyncExecutorProvider(), + SpannerOptions.createDefaultAsyncExecutorProvider()); this.dbAdminClient = new DatabaseAdminClientImpl(options.getProjectId(), gapicRpc); this.instanceClient = new InstanceAdminClientImpl(options.getProjectId(), gapicRpc, dbAdminClient); @@ -130,7 +138,7 @@ QueryOptions getDefaultQueryOptions(DatabaseId databaseId) { * Returns the {@link ExecutorProvider} to use for async methods that need a background executor. */ ExecutorProvider getAsyncExecutorProvider() { - return getOptions().getAsyncExecutorProvider(); + return asyncExecutorProvider; } SessionImpl sessionWithId(String name) { @@ -232,6 +240,7 @@ void close(long timeout, TimeUnit unit) { sessionClient.close(); } sessionClients.clear(); + asyncExecutorProvider.close(); try { gapicRpc.shutdown(); } catch (RuntimeException e) { diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java index 34b5e728e1e..dd917683620 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java @@ -18,7 +18,6 @@ import com.google.api.core.ApiFunction; import com.google.api.gax.core.ExecutorProvider; -import com.google.api.gax.core.FixedExecutorProvider; import com.google.api.gax.grpc.GrpcInterceptorProvider; import com.google.api.gax.longrunning.OperationSnapshot; import com.google.api.gax.longrunning.OperationTimedPollAlgorithm; @@ -60,6 +59,7 @@ import java.util.Map; import java.util.Map.Entry; import java.util.Set; +import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.ScheduledThreadPoolExecutor; import java.util.concurrent.ThreadFactory; import java.util.concurrent.TimeUnit; @@ -111,7 +111,7 @@ public class SpannerOptions extends ServiceOptions { private final Map mergedQueryOptions; private final CallCredentialsProvider callCredentialsProvider; - private final ExecutorProvider asyncExecutorProvider; + private final CloseableExecutorProvider asyncExecutorProvider; /** * Interface that can be used to provide {@link CallCredentials} instead of {@link Credentials} to @@ -144,6 +144,41 @@ public ServiceRpc create(SpannerOptions options) { private static final AtomicInteger DEFAULT_POOL_COUNT = new AtomicInteger(); + /** {@link ExecutorProvider} that is used for {@link AsyncResultSet}. */ + interface CloseableExecutorProvider extends ExecutorProvider, AutoCloseable { + /** Overridden to suppress the throws declaration of the super interface. */ + @Override + public void close(); + } + + static class FixedCloseableExecutorProvider implements CloseableExecutorProvider { + private final ScheduledExecutorService executor; + + private FixedCloseableExecutorProvider(ScheduledExecutorService executor) { + this.executor = Preconditions.checkNotNull(executor); + } + + @Override + public void close() { + executor.shutdown(); + } + + @Override + public ScheduledExecutorService getExecutor() { + return executor; + } + + @Override + public boolean shouldAutoClose() { + return false; + } + + /** Creates a FixedCloseableExecutorProvider. */ + static FixedCloseableExecutorProvider create(ScheduledExecutorService executor) { + return new FixedCloseableExecutorProvider(executor); + } + } + /** * Default {@link ExecutorProvider} for high-level async calls that need an executor. The default * uses a cached thread pool containing a max of 8 threads. The pool is lazily initialized and @@ -151,12 +186,12 @@ public ServiceRpc create(SpannerOptions options) { * also scale down the thread usage if the async load allows for that. */ @VisibleForTesting - static ExecutorProvider createDefaultAsyncExecutorProvider() { + static CloseableExecutorProvider createDefaultAsyncExecutorProvider() { return createAsyncExecutorProvider(8, 60L, TimeUnit.SECONDS); } @VisibleForTesting - static ExecutorProvider createAsyncExecutorProvider( + static CloseableExecutorProvider createAsyncExecutorProvider( int poolSize, long keepAliveTime, TimeUnit unit) { String format = String.format("async-pool-%d-thread-%%d", DEFAULT_POOL_COUNT.incrementAndGet()); ThreadFactory threadFactory = @@ -164,7 +199,7 @@ static ExecutorProvider createAsyncExecutorProvider( ScheduledThreadPoolExecutor executor = new ScheduledThreadPoolExecutor(poolSize, threadFactory); executor.setKeepAliveTime(keepAliveTime, unit); executor.allowCoreThreadTimeOut(true); - return FixedExecutorProvider.create(executor); + return FixedCloseableExecutorProvider.create(executor); } private SpannerOptions(Builder builder) { @@ -207,9 +242,7 @@ private SpannerOptions(Builder builder) { this.mergedQueryOptions = ImmutableMap.copyOf(merged); } callCredentialsProvider = builder.callCredentialsProvider; - asyncExecutorProvider = - MoreObjects.firstNonNull( - builder.asyncExecutorProvider, createDefaultAsyncExecutorProvider()); + asyncExecutorProvider = builder.asyncExecutorProvider; } /** @@ -274,7 +307,7 @@ public static class Builder private boolean autoThrottleAdministrativeRequests = false; private Map defaultQueryOptions = new HashMap<>(); private CallCredentialsProvider callCredentialsProvider; - private ExecutorProvider asyncExecutorProvider; + private CloseableExecutorProvider asyncExecutorProvider; private String emulatorHost = System.getenv("SPANNER_EMULATOR_HOST"); private Builder() { @@ -731,7 +764,7 @@ public QueryOptions getDefaultQueryOptions(DatabaseId databaseId) { return options; } - public ExecutorProvider getAsyncExecutorProvider() { + CloseableExecutorProvider getAsyncExecutorProvider() { return asyncExecutorProvider; } 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 0f098f4f6ef..259364e7049 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 @@ -98,7 +98,7 @@ private TransactionContextAsyncResultSetImpl(ListenableAsyncResultSet delegate) } @Override - public void setCallback(Executor exec, ReadyCallback cb) { + public ApiFuture setCallback(Executor exec, ReadyCallback cb) { Runnable listener = new Runnable() { @Override @@ -109,7 +109,7 @@ public void run() { try { increaseAsynOperations(); addListener(listener); - super.setCallback(exec, cb); + return super.setCallback(exec, cb); } catch (Throwable t) { removeListener(listener); finishedAsyncOperations.countDown(); 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 d55808ece0d..9caf7e2cd5a 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 @@ -498,6 +498,7 @@ private static void checkException(Queue exceptions, boolean keepExce private boolean stickyGlobalExceptions = false; private final ConcurrentMap statementResults = new ConcurrentHashMap<>(); + private final ConcurrentMap statementGetCounts = new ConcurrentHashMap<>(); private final ConcurrentMap sessions = new ConcurrentHashMap<>(); private ConcurrentMap sessionLastUsed = new ConcurrentHashMap<>(); private final ConcurrentMap transactions = new ConcurrentHashMap<>(); @@ -589,6 +590,11 @@ private StatementResult getResult(Statement statement) { StatementResult res; synchronized (lock) { res = statementResults.get(statement); + if (statementGetCounts.containsKey(statement)) { + statementGetCounts.put(statement, statementGetCounts.get(statement) + 1L); + } else { + statementGetCounts.put(statement, 1L); + } } if (res == null) { throw Status.INTERNAL From cc091e8d97ae2b01aea73571164ae746d81d8597 Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Sun, 19 Apr 2020 20:21:33 +0200 Subject: [PATCH 21/49] feat: add helper method for create test result sets --- .../src/main/java/com/google/cloud/spanner/ResultSets.java | 7 +++++++ 1 file changed, 7 insertions(+) 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 29c3e52c6ac..26c4832697a 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,6 +16,7 @@ package com.google.cloud.spanner; +import com.google.api.gax.core.InstantiatingExecutorProvider; import com.google.cloud.ByteArray; import com.google.cloud.Date; import com.google.cloud.Timestamp; @@ -23,6 +24,7 @@ import com.google.cloud.spanner.Type.StructField; import com.google.common.base.Preconditions; import com.google.common.collect.Lists; +import com.google.common.util.concurrent.ThreadFactoryBuilder; import com.google.spanner.v1.ResultSetStats; import java.util.List; @@ -41,6 +43,11 @@ public static ResultSet forRows(Type type, Iterable rows) { return new PrePopulatedResultSet(type, rows); } + /** Converts the given {@link ResultSet} to an {@link AsyncResultSet}. */ + public static AsyncResultSet toAsyncResultSet(ResultSet delegate) { + return new AsyncResultSetImpl(InstantiatingExecutorProvider.newBuilder().setExecutorThreadCount(1).setThreadFactory(new ThreadFactoryBuilder().setDaemon(true).setNameFormat("test-async-resultset-%d").build()).build(), delegate, 100); + } + private static class PrePopulatedResultSet implements ResultSet { private final List rows; private final Type type; From d75f9795e2782bd7d628da289e3c2f5103aa8136 Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Tue, 21 Apr 2020 21:43:29 +0200 Subject: [PATCH 22/49] feat: add batchUpdateAsync --- .../com/google/cloud/spanner/ReadContext.java | 17 ++ .../com/google/cloud/spanner/SessionPool.java | 9 + .../cloud/spanner/TransactionContext.java | 20 ++ .../cloud/spanner/TransactionRunnerImpl.java | 68 +++++- .../cloud/spanner/spi/v1/GapicSpannerRpc.java | 17 +- .../cloud/spanner/spi/v1/SpannerRpc.java | 6 + .../google/cloud/spanner/AsyncRunnerTest.java | 224 ++++++++++++++++++ .../cloud/spanner/MockSpannerServiceImpl.java | 1 + .../cloud/spanner/MockSpannerTestUtil.java | 2 + .../cloud/spanner/it/ITAsyncExamplesTest.java | 80 ++++++- 10 files changed, 428 insertions(+), 16 deletions(-) diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ReadContext.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ReadContext.java index 904fa4176b1..e87d40fb207 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ReadContext.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ReadContext.java @@ -66,6 +66,10 @@ enum QueryAnalyzeMode { */ ResultSet read(String table, KeySet keys, Iterable columns, ReadOption... options); + /** + * Same as {@link #read(String, KeySet, Iterable, ReadOption...)}, but is guaranteed to be + * non-blocking and will return the results as an {@link AsyncResultSet}. + */ AsyncResultSet readAsync( String table, KeySet keys, Iterable columns, ReadOption... options); @@ -97,6 +101,10 @@ AsyncResultSet readAsync( ResultSet readUsingIndex( String table, String index, KeySet keys, Iterable columns, ReadOption... options); + /** + * Same as {@link #readUsingIndex(String, String, KeySet, Iterable, ReadOption...)}, but is + * guaranteed to be non-blocking and will return its results as an {@link AsyncResultSet}. + */ AsyncResultSet readUsingIndexAsync( String table, String index, KeySet keys, Iterable columns, ReadOption... options); @@ -119,6 +127,7 @@ AsyncResultSet readUsingIndexAsync( @Nullable Struct readRow(String table, Key key, Iterable columns); + /** Same as {@link #readRow(String, Key, Iterable)}, but is guaranteed to be non-blocking. */ ApiFuture readRowAsync(String table, Key key, Iterable columns); /** @@ -143,6 +152,10 @@ AsyncResultSet readUsingIndexAsync( @Nullable Struct readRowUsingIndex(String table, String index, Key key, Iterable columns); + /** + * Same as {@link #readRowUsingIndex(String, String, Key, Iterable)}, but is guaranteed to be + * non-blocking. + */ ApiFuture readRowUsingIndexAsync( String table, String index, Key key, Iterable columns); @@ -172,6 +185,10 @@ ApiFuture readRowUsingIndexAsync( */ ResultSet executeQuery(Statement statement, QueryOption... options); + /** + * Same as {@link #executeQuery(Statement, QueryOption...)}, but is guaranteed to be non-blocking + * and returns its results as an {@link AsyncResultSet}. + */ AsyncResultSet executeQueryAsync(Statement statement, QueryOption... options); /** diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java index 382bbb3097b..b6fc3fe0653 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java @@ -674,6 +674,15 @@ public long[] batchUpdate(Iterable statements) { } } + @Override + public ApiFuture batchUpdateAsync(Iterable statements) { + try { + return delegate.batchUpdateAsync(statements); + } catch (SessionNotFoundException e) { + throw handleSessionNotFound(e); + } + } + @Override public ResultSet executeQuery(Statement statement, QueryOption... options) { return new SessionPoolResultSet(delegate.executeQuery(statement, options)); diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionContext.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionContext.java index 7e09da901c2..0b4a92f989e 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionContext.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionContext.java @@ -104,6 +104,15 @@ public interface TransactionContext extends ReadContext { */ long executeUpdate(Statement statement); + /** + * Same as {@link #executeUpdate(Statement)}, but is guaranteed to be non-blocking. If multiple + * asynchronous update statements are submitted to the same read/write transaction, the statements + * are guaranteed to be submitted to Cloud Spanner in the order that they were submitted in the + * client. This does however not guarantee that an asynchronous update statement will see the + * results of all previously submitted statements, as the execution of the statements can be + * parallel. If you rely on the results of a previous statement, you should block until the result + * of that statement is known and has been returned to the client. + */ ApiFuture executeUpdateAsync(Statement statement); /** @@ -122,4 +131,15 @@ public interface TransactionContext extends ReadContext { * statement. The 3rd statement will not run. */ long[] batchUpdate(Iterable statements); + + /** + * Same as {@link #batchUpdate(Iterable)}, but is guaranteed to be non-blocking. If multiple + * asynchronous update statements are submitted to the same read/write transaction, the statements + * are guaranteed to be submitted to Cloud Spanner in the order that they were submitted in the + * client. This does however not guarantee that an asynchronous update statement will see the + * results of all previously submitted statements, as the execution of the statements can be + * parallel. If you rely on the results of a previous statement, you should block until the result + * of that statement is known and has been returned to the client. + */ + ApiFuture batchUpdateAsync(Iterable statements); } 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 259364e7049..c3e13f2a87f 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 @@ -37,6 +37,7 @@ import com.google.spanner.v1.CommitRequest; import com.google.spanner.v1.CommitResponse; import com.google.spanner.v1.ExecuteBatchDmlRequest; +import com.google.spanner.v1.ExecuteBatchDmlResponse; import com.google.spanner.v1.ExecuteSqlRequest; import com.google.spanner.v1.ExecuteSqlRequest.QueryMode; import com.google.spanner.v1.ResultSet; @@ -347,12 +348,9 @@ public ApiFuture executeUpdateAsync(Statement statement) { @Override public Long apply(ResultSet input) { if (!input.hasStats()) { - SpannerException e = - SpannerExceptionFactory.newSpannerException( - ErrorCode.INVALID_ARGUMENT, - "DML response missing stats possibly due to non-DML statement as input"); - onError(e); - throw e; + throw SpannerExceptionFactory.newSpannerException( + ErrorCode.INVALID_ARGUMENT, + "DML response missing stats possibly due to non-DML statement as input"); } // For standard DML, using the exact row count. return input.getStats().getRowCountExact(); @@ -408,6 +406,64 @@ public long[] batchUpdate(Iterable statements) { } } + @Override + public ApiFuture batchUpdateAsync(Iterable statements) { + beforeReadOrQuery(); + final ExecuteBatchDmlRequest.Builder builder = getExecuteBatchDmlRequestBuilder(statements); + ApiFuture response; + try { + // Register the update as an async operation that must finish before the transaction may + // commit. + increaseAsynOperations(); + response = rpc.executeBatchDmlAsync(builder.build(), session.getOptions()); + } catch (Throwable t) { + finishedAsyncOperations.countDown(); + throw t; + } + final ApiFuture updateCounts = + ApiFutures.transform( + response, + new ApiFunction() { + @Override + public long[] apply(ExecuteBatchDmlResponse input) { + long[] results = new long[input.getResultSetsCount()]; + for (int i = 0; i < input.getResultSetsCount(); ++i) { + results[i] = input.getResultSets(i).getStats().getRowCountExact(); + } + // If one of the DML statements was aborted, we should throw an aborted exception. + // In all other cases, we should throw a BatchUpdateException. + if (input.getStatus().getCode() == Code.ABORTED_VALUE) { + throw newSpannerException( + ErrorCode.fromRpcStatus(input.getStatus()), input.getStatus().getMessage()); + } else if (input.getStatus().getCode() != 0) { + throw newSpannerBatchUpdateException( + ErrorCode.fromRpcStatus(input.getStatus()), + input.getStatus().getMessage(), + results); + } + return results; + } + }, + MoreExecutors.directExecutor()); + updateCounts.addListener( + new Runnable() { + @Override + public void run() { + try { + updateCounts.get(); + } catch (ExecutionException e) { + onError(SpannerExceptionFactory.newSpannerException(e.getCause())); + } catch (InterruptedException e) { + onError(SpannerExceptionFactory.propagateInterrupt(e)); + } finally { + finishedAsyncOperations.countDown(); + } + } + }, + MoreExecutors.directExecutor()); + return updateCounts; + } + private ListenableAsyncResultSet wrap(ListenableAsyncResultSet delegate) { return new TransactionContextAsyncResultSetImpl(delegate); } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java index a6e589429f5..0fc704f5f7c 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java @@ -1065,16 +1065,27 @@ public void cancel(String message) { @Override public ExecuteBatchDmlResponse executeBatchDml( ExecuteBatchDmlRequest request, @Nullable Map options) { + return get(executeBatchDmlAsync(request, options)); + } + + @Override + public ApiFuture executeBatchDmlAsync( + ExecuteBatchDmlRequest request, @Nullable Map options) { + GrpcCallContext context = newCallContext(options, request.getSession()); + return spannerStub.executeBatchDmlCallable().futureCall(request, context); + } + @Override + public ApiFuture beginTransactionAsync( + BeginTransactionRequest request, @Nullable Map options) { GrpcCallContext context = newCallContext(options, request.getSession()); - return get(spannerStub.executeBatchDmlCallable().futureCall(request, context)); + return spannerStub.beginTransactionCallable().futureCall(request, context); } @Override public Transaction beginTransaction( BeginTransactionRequest request, @Nullable Map options) throws SpannerException { - GrpcCallContext context = newCallContext(options, request.getSession()); - return get(spannerStub.beginTransactionCallable().futureCall(request, context)); + return get(beginTransactionAsync(request, options)); } @Override diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerRpc.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerRpc.java index a9114bc22b5..0334ca4e403 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerRpc.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerRpc.java @@ -293,9 +293,15 @@ StreamingCall executeQuery( ExecuteBatchDmlResponse executeBatchDml(ExecuteBatchDmlRequest build, Map options); + ApiFuture executeBatchDmlAsync( + ExecuteBatchDmlRequest build, Map options); + Transaction beginTransaction(BeginTransactionRequest request, @Nullable Map options) throws SpannerException; + ApiFuture beginTransactionAsync( + BeginTransactionRequest request, @Nullable Map options); + CommitResponse commit(CommitRequest commitRequest, @Nullable Map options) throws SpannerException; diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncRunnerTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncRunnerTest.java index 5d67de3fc25..cbe78fe0d01 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncRunnerTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncRunnerTest.java @@ -38,6 +38,7 @@ import com.google.spanner.v1.BatchCreateSessionsRequest; import com.google.spanner.v1.BeginTransactionRequest; import com.google.spanner.v1.CommitRequest; +import com.google.spanner.v1.ExecuteBatchDmlRequest; import com.google.spanner.v1.ExecuteSqlRequest; import io.grpc.Server; import io.grpc.Status; @@ -87,6 +88,10 @@ public static void setup() throws Exception { StatementResult.exception( INVALID_UPDATE_STATEMENT, Status.INVALID_ARGUMENT.withDescription("invalid statement").asRuntimeException())); + mockSpanner.putStatementResult( + StatementResult.exception( + UPDATE_ABORTED_STATEMENT, + Status.ABORTED.withDescription("Transaction was aborted").asRuntimeException())); String uniqueName = InProcessServerBuilder.generateName(); server = InProcessServerBuilder.forName(uniqueName) @@ -371,6 +376,225 @@ public ApiFuture doWorkAsync(TransactionContext txn) { CommitRequest.class); } + @Test + public void asyncRunnerBatchUpdate() throws Exception { + AsyncRunner runner = client().runAsync(); + ApiFuture updateCount = + runner.runAsync( + new AsyncWork() { + @Override + public ApiFuture doWorkAsync(TransactionContext txn) { + return txn.batchUpdateAsync(ImmutableList.of(UPDATE_STATEMENT, UPDATE_STATEMENT)); + } + }, + executor); + assertThat(updateCount.get()).asList().containsExactly(UPDATE_COUNT, UPDATE_COUNT); + } + + @Test + public void asyncRunnerIsNonBlockingWithBatchUpdate() throws Exception { + mockSpanner.freeze(); + AsyncRunner runner = clientWithEmptySessionPool().runAsync(); + ApiFuture res = + runner.runAsync( + new AsyncWork() { + @Override + public ApiFuture doWorkAsync(TransactionContext txn) { + txn.batchUpdateAsync(ImmutableList.of(UPDATE_STATEMENT)); + return ApiFutures.immediateFuture(null); + } + }, + executor); + ApiFuture ts = runner.getCommitTimestamp(); + mockSpanner.unfreeze(); + assertThat(res.get()).isNull(); + assertThat(ts.get()).isNotNull(); + } + + @Test + public void asyncRunnerInvalidBatchUpdate() throws Exception { + AsyncRunner runner = client().runAsync(); + ApiFuture updateCount = + runner.runAsync( + new AsyncWork() { + @Override + public ApiFuture doWorkAsync(TransactionContext txn) { + return txn.batchUpdateAsync( + ImmutableList.of(UPDATE_STATEMENT, INVALID_UPDATE_STATEMENT)); + } + }, + executor); + try { + updateCount.get(); + fail("missing expected exception"); + } catch (ExecutionException e) { + assertThat(e.getCause()).isInstanceOf(SpannerException.class); + SpannerException se = (SpannerException) e.getCause(); + assertThat(se.getErrorCode()).isEqualTo(ErrorCode.INVALID_ARGUMENT); + assertThat(se.getMessage()).contains("invalid statement"); + } + } + + @Test + public void asyncRunnerFireAndForgetInvalidBatchUpdate() throws Exception { + AsyncRunner runner = client().runAsync(); + ApiFuture res = + runner.runAsync( + new AsyncWork() { + @Override + public ApiFuture doWorkAsync(TransactionContext txn) { + txn.batchUpdateAsync(ImmutableList.of(UPDATE_STATEMENT, INVALID_UPDATE_STATEMENT)); + return txn.batchUpdateAsync(ImmutableList.of(UPDATE_STATEMENT, UPDATE_STATEMENT)); + } + }, + executor); + assertThat(res.get()).asList().containsExactly(UPDATE_COUNT, UPDATE_COUNT); + } + + @Test + public void asyncRunnerBatchUpdateAborted() throws Exception { + final AtomicInteger attempt = new AtomicInteger(); + AsyncRunner runner = client().runAsync(); + ApiFuture updateCount = + runner.runAsync( + new AsyncWork() { + @Override + public ApiFuture doWorkAsync(TransactionContext txn) { + if (attempt.incrementAndGet() == 1) { + return txn.batchUpdateAsync( + ImmutableList.of(UPDATE_STATEMENT, UPDATE_ABORTED_STATEMENT)); + } else { + return txn.batchUpdateAsync(ImmutableList.of(UPDATE_STATEMENT, UPDATE_STATEMENT)); + } + } + }, + executor); + assertThat(updateCount.get()).asList().containsExactly(UPDATE_COUNT, UPDATE_COUNT); + assertThat(attempt.get()).isEqualTo(2); + } + + @Test + public void asyncRunnerWithBatchUpdateCommitAborted() throws Exception { + try { + // Temporarily set the result of the update to 2 rows. + mockSpanner.putStatementResult(StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT + 1L)); + final AtomicInteger attempt = new AtomicInteger(); + AsyncRunner runner = client().runAsync(); + ApiFuture updateCount = + runner.runAsync( + new AsyncWork() { + @Override + public ApiFuture doWorkAsync(TransactionContext txn) { + if (attempt.get() > 0) { + // Set the result of the update statement back to 1 row. + mockSpanner.putStatementResult( + StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT)); + } + ApiFuture updateCount = + txn.batchUpdateAsync(ImmutableList.of(UPDATE_STATEMENT, UPDATE_STATEMENT)); + if (attempt.incrementAndGet() == 1) { + mockSpanner.abortTransaction(txn); + } + return updateCount; + } + }, + executor); + assertThat(updateCount.get()).asList().containsExactly(UPDATE_COUNT, UPDATE_COUNT); + assertThat(attempt.get()).isEqualTo(2); + } finally { + mockSpanner.putStatementResult(StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT)); + } + } + + @Test + public void asyncRunnerBatchUpdateAbortedWithoutGettingResult() throws Exception { + final AtomicInteger attempt = new AtomicInteger(); + AsyncRunner runner = clientWithEmptySessionPool().runAsync(); + ApiFuture result = + runner.runAsync( + new AsyncWork() { + @Override + public ApiFuture doWorkAsync(TransactionContext txn) { + if (attempt.incrementAndGet() == 1) { + mockSpanner.abortTransaction(txn); + } + // This update statement will be aborted, but the error will not propagated to the + // transaction runner and cause the transaction to retry. Instead, the commit call + // will do that. + txn.batchUpdateAsync(ImmutableList.of(UPDATE_STATEMENT, UPDATE_STATEMENT)); + // Resolving this future will not resolve the result of the entire transaction. The + // transaction result will be resolved when the commit has actually finished + // successfully. + return ApiFutures.immediateFuture(null); + } + }, + executor); + assertThat(result.get()).isNull(); + assertThat(attempt.get()).isEqualTo(2); + assertThat(mockSpanner.getRequestTypes()) + .containsExactly( + BatchCreateSessionsRequest.class, + BeginTransactionRequest.class, + ExecuteBatchDmlRequest.class, + CommitRequest.class, + BeginTransactionRequest.class, + ExecuteBatchDmlRequest.class, + CommitRequest.class); + } + + @Test + public void asyncRunnerWithBatchUpdateCommitFails() throws Exception { + mockSpanner.setCommitExecutionTime( + SimulatedExecutionTime.ofException( + Status.RESOURCE_EXHAUSTED + .withDescription("mutation limit exceeded") + .asRuntimeException())); + AsyncRunner runner = client().runAsync(); + ApiFuture updateCount = + runner.runAsync( + new AsyncWork() { + @Override + public ApiFuture doWorkAsync(TransactionContext txn) { + // This statement will succeed, but the commit will fail. The error from the commit + // will bubble up to the future that is returned by the transaction, and the update + // count returned here will never reach the user application. + return txn.batchUpdateAsync(ImmutableList.of(UPDATE_STATEMENT, UPDATE_STATEMENT)); + } + }, + executor); + try { + updateCount.get(); + fail("missing expected exception"); + } catch (ExecutionException e) { + assertThat(e.getCause()).isInstanceOf(SpannerException.class); + SpannerException se = (SpannerException) e.getCause(); + assertThat(se.getErrorCode()).isEqualTo(ErrorCode.RESOURCE_EXHAUSTED); + assertThat(se.getMessage()).contains("mutation limit exceeded"); + } + } + + @Test + public void asyncRunnerWaitsUntilAsyncBatchUpdateHasFinished() throws Exception { + AsyncRunner runner = clientWithEmptySessionPool().runAsync(); + ApiFuture res = + runner.runAsync( + new AsyncWork() { + @Override + public ApiFuture doWorkAsync(TransactionContext txn) { + txn.batchUpdateAsync(ImmutableList.of(UPDATE_STATEMENT)); + return ApiFutures.immediateFuture(null); + } + }, + executor); + res.get(); + assertThat(mockSpanner.getRequestTypes()) + .containsExactly( + BatchCreateSessionsRequest.class, + BeginTransactionRequest.class, + ExecuteBatchDmlRequest.class, + CommitRequest.class); + } + @Test public void closeTransactionBeforeEndOfAsyncQuery() throws Exception { final BlockingQueue results = new SynchronousQueue<>(); 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 f2589ea6d73..73d903de5b9 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 @@ -996,6 +996,7 @@ public void executeBatchDml( status = com.google.rpc.Status.newBuilder() .setCode(res.getException().getStatus().getCode().value()) + .setMessage(res.getException().getMessage()) .build(); break resultLoop; case RESULT_SET: diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MockSpannerTestUtil.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MockSpannerTestUtil.java index fe85ef7c908..c0d3bdc7d1c 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MockSpannerTestUtil.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MockSpannerTestUtil.java @@ -57,6 +57,8 @@ public class MockSpannerTestUtil { static final Statement UPDATE_STATEMENT = Statement.of("UPDATE FOO SET BAR=1 WHERE BAZ=2"); static final Statement INVALID_UPDATE_STATEMENT = Statement.of("UPDATE NON_EXISTENT_TABLE SET BAR=1 WHERE BAZ=2"); + static final Statement UPDATE_ABORTED_STATEMENT = + Statement.of("UPDATE FOO SET BAR=1 WHERE BAZ=2 AND THIS_WILL_ABORT=TRUE"); static final long UPDATE_COUNT = 1L; static final String READ_TABLE_NAME = "TestTable"; diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITAsyncExamplesTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITAsyncExamplesTest.java index c328f16ef9c..c5e2419ba6b 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITAsyncExamplesTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITAsyncExamplesTest.java @@ -242,14 +242,14 @@ public CallbackResponse cursorReady(AsyncResultSet resultSet) { @Test public void runAsync() throws Exception { AsyncRunner runner = client.runAsync(); - ApiFuture deleteCount = + ApiFuture insertCount = runner.runAsync( new AsyncWork() { @Override public ApiFuture doWorkAsync(TransactionContext txn) { // Even though this is a shoot-and-forget asynchronous DML statement, it is // guaranteed to be executed within the transaction before the commit is executed. - txn.executeUpdateAsync( + return txn.executeUpdateAsync( Statement.newBuilder( "INSERT INTO TestTable (Key, StringValue) VALUES (@key, @value)") .bind("key") @@ -257,11 +257,15 @@ public ApiFuture doWorkAsync(TransactionContext txn) { .bind("value") .to("v999") .build()); - // Note that even though both DML statements are executed asynchronously, they are - // guaranteed to be executed in the order they are submitted to the transaction, as - // they receive a monotonically increasing sequence number at the moment that they - // are submitted. If they arrive out of order on the backend, the backend may abort - // the transaction and the transaction will be retried. + } + }, + executor); + assertThat(insertCount.get()).isEqualTo(1L); + ApiFuture deleteCount = + runner.runAsync( + new AsyncWork() { + @Override + public ApiFuture doWorkAsync(TransactionContext txn) { return txn.executeUpdateAsync( Statement.newBuilder("DELETE FROM TestTable WHERE Key=@key") .bind("key") @@ -273,6 +277,68 @@ public ApiFuture doWorkAsync(TransactionContext txn) { assertThat(deleteCount.get()).isEqualTo(1L); } + @Test + public void runAsyncBatchUpdate() throws Exception { + AsyncRunner runner = client.runAsync(); + ApiFuture insertCount = + runner.runAsync( + new AsyncWork() { + @Override + public ApiFuture doWorkAsync(TransactionContext txn) { + // Even though this is a shoot-and-forget asynchronous DML statement, it is + // guaranteed to be executed within the transaction before the commit is executed. + return txn.batchUpdateAsync( + ImmutableList.of( + Statement.newBuilder( + "INSERT INTO TestTable (Key, StringValue) VALUES (@key, @value)") + .bind("key") + .to("k997") + .bind("value") + .to("v997") + .build(), + Statement.newBuilder( + "INSERT INTO TestTable (Key, StringValue) VALUES (@key, @value)") + .bind("key") + .to("k998") + .bind("value") + .to("v998") + .build(), + Statement.newBuilder( + "INSERT INTO TestTable (Key, StringValue) VALUES (@key, @value)") + .bind("key") + .to("k999") + .bind("value") + .to("v999") + .build())); + } + }, + executor); + assertThat(insertCount.get()).asList().containsExactly(1L, 1L, 1L); + ApiFuture deleteCount = + runner.runAsync( + new AsyncWork() { + @Override + public ApiFuture doWorkAsync(TransactionContext txn) { + return txn.batchUpdateAsync( + ImmutableList.of( + Statement.newBuilder("DELETE FROM TestTable WHERE Key=@key") + .bind("key") + .to("k997") + .build(), + Statement.newBuilder("DELETE FROM TestTable WHERE Key=@key") + .bind("key") + .to("k998") + .build(), + Statement.newBuilder("DELETE FROM TestTable WHERE Key=@key") + .bind("key") + .to("k999") + .build())); + } + }, + executor); + assertThat(deleteCount.get()).asList().containsExactly(1L, 1L, 1L); + } + @Test public void readOnlyTransaction() throws Exception { ImmutableList keys1 = ImmutableList.of("k10", "k11", "k12"); From 7f06e845181ac240dbd5ec664f702946c321bd84 Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Sun, 26 Apr 2020 11:17:32 +0200 Subject: [PATCH 23/49] fix: add ignored interface differences --- .../clirr-ignored-differences.xml | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/google-cloud-spanner/clirr-ignored-differences.xml b/google-cloud-spanner/clirr-ignored-differences.xml index 4695ca92700..2931407b31b 100644 --- a/google-cloud-spanner/clirr-ignored-differences.xml +++ b/google-cloud-spanner/clirr-ignored-differences.xml @@ -182,10 +182,25 @@ com/google/cloud/spanner/ReadContext * readRowUsingIndexAsync(*) + + 7012 + com/google/cloud/spanner/TransactionContext + * batchUpdateAsync(*) + 7012 com/google/cloud/spanner/TransactionContext * executeUpdateAsync(*) + + 7012 + com/google/cloud/spanner/spi/v1/SpannerRpc + * beginTransactionAsync(*) + + + 7012 + com/google/cloud/spanner/spi/v1/SpannerRpc + * executeBatchDmlAsync(*) + From 00b83d2ef88fa5ccee369809033a55cfbe9e298f Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Sun, 26 Apr 2020 22:33:24 +0200 Subject: [PATCH 24/49] refactor: use future as waiter in SessionPool --- .../com/google/cloud/spanner/SessionImpl.java | 6 +- .../com/google/cloud/spanner/SessionPool.java | 355 ++++++++++-------- .../cloud/spanner/BaseSessionPoolTest.java | 2 +- .../spanner/SessionPoolMaintainerTest.java | 21 +- .../cloud/spanner/SessionPoolStressTest.java | 91 ++--- .../google/cloud/spanner/SessionPoolTest.java | 2 +- 6 files changed, 254 insertions(+), 223 deletions(-) diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java index 77f8946bb12..a425ea56118 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java @@ -84,7 +84,7 @@ static interface SessionTransaction { private final String name; private final DatabaseId databaseId; private SessionTransaction activeTransaction; - private ByteString readyTransactionId; + ByteString readyTransactionId; private final Map options; private Span currentSpan; @@ -304,6 +304,10 @@ T setActive(@Nullable T ctx) { return ctx; } + boolean hasReadyTransaction() { + return readyTransactionId != null; + } + @Override public TransactionManager transactionManager() { return new TransactionManagerImpl(this, currentSpan); diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java index 36b61a920aa..e89f03178db 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java @@ -52,13 +52,12 @@ import com.google.common.base.Supplier; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; +import com.google.common.util.concurrent.ForwardingFuture; import com.google.common.util.concurrent.ForwardingFuture.SimpleForwardingFuture; -import com.google.common.util.concurrent.Futures; import com.google.common.util.concurrent.ListenableFuture; import com.google.common.util.concurrent.MoreExecutors; import com.google.common.util.concurrent.SettableFuture; import com.google.common.util.concurrent.ThreadFactoryBuilder; -import com.google.common.util.concurrent.Uninterruptibles; import com.google.protobuf.Empty; import io.opencensus.common.Scope; import io.opencensus.common.ToLongFunction; @@ -82,18 +81,16 @@ import java.util.Queue; import java.util.Random; import java.util.Set; -import java.util.concurrent.BlockingQueue; -import java.util.concurrent.Callable; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; import java.util.concurrent.Executor; import java.util.concurrent.Executors; import java.util.concurrent.Future; -import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.ScheduledFuture; import java.util.concurrent.ThreadPoolExecutor; import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicLong; @@ -957,10 +954,6 @@ private PooledSessionFuture createPooledSessionFuture(Future futu return new PooledSessionFuture(future, span); } - private PooledSessionFuture createPooledSessionFuture(PooledSession session, Span span) { - return new PooledSessionFuture(Futures.immediateFuture(session), span); - } - final class PooledSessionFuture extends SimpleForwardingFuture implements Session { private volatile LeakedSessionException leakedException; private volatile AtomicBoolean inUse = new AtomicBoolean(); @@ -1192,6 +1185,11 @@ private PooledSession(SessionImpl delegate) { this.lastUseTime = clock.instant(); } + @Override + public String toString() { + return getName(); + } + @VisibleForTesting void setAllowReplacing(boolean allowReplacing) { this.allowReplacing = allowReplacing; @@ -1340,50 +1338,37 @@ public TransactionManager transactionManager() { } } - private static final class SessionOrError { - private final PooledSession session; - private final SpannerException e; - - SessionOrError(PooledSession session) { - this.session = session; - this.e = null; - } + private final class WaiterFuture extends ForwardingFuture { + private static final long MAX_SESSION_WAIT_TIMEOUT = 240_000L; + private final SettableApiFuture waiter = SettableApiFuture.create(); - SessionOrError(SpannerException e) { - this.session = null; - this.e = e; + @Override + protected Future delegate() { + return waiter; } - } - - private final class Waiter { - private static final long MAX_SESSION_WAIT_TIMEOUT = 240_000L; - private final BlockingQueue waiter = new LinkedBlockingQueue<>(1); - // private final SynchronousQueue waiter = new SynchronousQueue<>(); private void put(PooledSession session) { - Uninterruptibles.putUninterruptibly(waiter, new SessionOrError(session)); + waiter.set(session); } private void put(SpannerException e) { - Uninterruptibles.putUninterruptibly(waiter, new SessionOrError(e)); + waiter.setException(e); } - private PooledSession take() throws SpannerException { + @Override + public PooledSession get() { long currentTimeout = options.getInitialWaitForSessionTimeoutMillis(); while (true) { Span span = tracer.spanBuilder(WAIT_FOR_SESSION).startSpan(); try (Scope waitScope = tracer.withSpan(span)) { - SessionOrError s = pollUninterruptiblyWithTimeout(currentTimeout); + PooledSession s = pollUninterruptiblyWithTimeout(currentTimeout); if (s == null) { // Set the status to DEADLINE_EXCEEDED and retry. numWaiterTimeouts.incrementAndGet(); tracer.getCurrentSpan().setStatus(Status.DEADLINE_EXCEEDED); currentTimeout = Math.min(currentTimeout * 2, MAX_SESSION_WAIT_TIMEOUT); } else { - if (s.e != null) { - throw newSpannerException(s.e); - } - return s.session; + return s; } } catch (Exception e) { TraceUtil.setWithFailure(span, e); @@ -1394,14 +1379,18 @@ private PooledSession take() throws SpannerException { } } - private SessionOrError pollUninterruptiblyWithTimeout(long timeoutMillis) { + private PooledSession pollUninterruptiblyWithTimeout(long timeoutMillis) { boolean interrupted = false; try { while (true) { try { - return waiter.poll(timeoutMillis, TimeUnit.MILLISECONDS); + return waiter.get(timeoutMillis, TimeUnit.MILLISECONDS); } catch (InterruptedException e) { interrupted = true; + } catch (TimeoutException e) { + return null; + } catch (ExecutionException e) { + throw SpannerExceptionFactory.newSpannerException(e.getCause()); } } } finally { @@ -1582,21 +1571,6 @@ private static enum Position { private final ExecutorFactory executorFactory; private final ScheduledExecutorService prepareExecutor; - // TODO(loite): Refactor Waiter to use a SettableFuture that can be set when a session is released - // into the pool, instead of using a thread waiting on a synchronous queue. - private final ScheduledExecutorService readWaiterExecutor = - Executors.newSingleThreadScheduledExecutor( - new ThreadFactoryBuilder() - .setDaemon(true) - .setNameFormat("session-pool-read-waiter-%d") - .build()); - private final ScheduledExecutorService writeWaiterExecutor = - Executors.newSingleThreadScheduledExecutor( - new ThreadFactoryBuilder() - .setDaemon(true) - .setNameFormat("session-pool-write-waiter-%d") - .build()); - private final int prepareThreadPoolSize; final PoolMaintainer poolMaintainer; private final Clock clock; @@ -1619,10 +1593,10 @@ private static enum Position { private final LinkedList writePreparedSessions = new LinkedList<>(); @GuardedBy("lock") - private final Queue readWaiters = new LinkedList<>(); + private final Queue readWaiters = new LinkedList<>(); @GuardedBy("lock") - private final Queue readWriteWaiters = new LinkedList<>(); + private final Queue readWriteWaiters = new LinkedList<>(); @GuardedBy("lock") private int numSessionsBeingPrepared = 0; @@ -1779,9 +1753,9 @@ void removeFromPool(PooledSession session) { session.markClosing(); allSessions.remove(session); numIdleSessionsRemoved++; - if (idleSessionRemovedListener != null) { - idleSessionRemovedListener.apply(session); - } + } + if (idleSessionRemovedListener != null) { + idleSessionRemovedListener.apply(session); } } @@ -1914,7 +1888,7 @@ boolean isValid() { PooledSessionFuture getReadSession() throws SpannerException { Span span = Tracing.getTracer().getCurrentSpan(); span.addAnnotation("Acquiring session"); - Waiter waiter = null; + WaiterFuture waiter = null; PooledSession sess = null; synchronized (lock) { if (closureFuture != null) { @@ -1936,7 +1910,7 @@ PooledSessionFuture getReadSession() throws SpannerException { if (sess == null) { span.addAnnotation("No session available"); maybeCreateSession(); - waiter = new Waiter(); + waiter = new WaiterFuture(); readWaiters.add(waiter); } else { span.addAnnotation("Acquired read write session"); @@ -1944,7 +1918,7 @@ PooledSessionFuture getReadSession() throws SpannerException { } else { span.addAnnotation("Acquired read only session"); } - return checkoutSession(span, sess, waiter, false); + return checkoutSession(span, sess, waiter, false, false); } } @@ -1970,103 +1944,80 @@ PooledSessionFuture getReadWriteSession() { Span span = Tracing.getTracer().getCurrentSpan(); span.addAnnotation("Acquiring read write session"); PooledSession sess = null; - // Loop to retry SessionNotFoundExceptions that might occur during in-process prepare of a - // session. - while (true) { - Waiter waiter = null; - boolean inProcessPrepare = false; - synchronized (lock) { - if (closureFuture != null) { - span.addAnnotation("Pool has been closed"); - throw new IllegalStateException("Pool has been closed"); - } - if (resourceNotFoundException != null) { - span.addAnnotation("Database has been deleted"); - throw SpannerExceptionFactory.newSpannerException( - ErrorCode.NOT_FOUND, - String.format( - "The session pool has been invalidated because a previous RPC returned 'Database not found': %s", - resourceNotFoundException.getMessage()), - resourceNotFoundException); - } - sess = writePreparedSessions.poll(); - if (sess == null) { - if (numSessionsBeingPrepared <= prepareThreadPoolSize) { - if (numSessionsBeingPrepared <= readWriteWaiters.size()) { - PooledSession readSession = readSessions.poll(); - if (readSession != null) { - span.addAnnotation( - "Acquired read only session. Preparing for read write transaction"); - prepareSession(readSession); - } else { - span.addAnnotation("No session available"); - maybeCreateSession(); - } - } - } else { - inProcessPrepare = true; - numSessionsInProcessPrepared++; + WaiterFuture waiter = null; + boolean inProcessPrepare = false; + synchronized (lock) { + if (closureFuture != null) { + span.addAnnotation("Pool has been closed"); + throw new IllegalStateException("Pool has been closed"); + } + if (resourceNotFoundException != null) { + span.addAnnotation("Database has been deleted"); + throw SpannerExceptionFactory.newSpannerException( + ErrorCode.NOT_FOUND, + String.format( + "The session pool has been invalidated because a previous RPC returned 'Database not found': %s", + resourceNotFoundException.getMessage()), + resourceNotFoundException); + } + sess = writePreparedSessions.poll(); + if (sess == null) { + if (numSessionsBeingPrepared <= prepareThreadPoolSize) { + if (numSessionsBeingPrepared <= readWriteWaiters.size()) { PooledSession readSession = readSessions.poll(); if (readSession != null) { - // Create a read/write transaction in-process if there is already a queue for prepared - // sessions. This is more efficient than doing it asynchronously, as it scales with - // the number of user threads. The thread pool for asynchronously preparing sessions - // is fixed. span.addAnnotation( - "Acquired read only session. Preparing in-process for read write transaction"); - sess = readSession; + "Acquired read only session. Preparing for read write transaction"); + prepareSession(readSession); } else { span.addAnnotation("No session available"); maybeCreateSession(); } } - if (sess == null) { - waiter = new Waiter(); - if (inProcessPrepare) { - // inProcessPrepare=true means that we have already determined that the queue for - // preparing read/write sessions is larger than the number of threads in the prepare - // thread pool, and that it's more efficient to do the prepare in-process. We will - // therefore create a waiter for a read-only session, even though a read/write session - // has been requested. - readWaiters.add(waiter); - } else { - readWriteWaiters.add(waiter); - } - } } else { - span.addAnnotation("Acquired read write session"); - } - } - if (waiter != null) { - logger.log( - Level.FINE, - "No session available in the pool. Blocking for one to become available/created"); - span.addAnnotation("Waiting for read write session to be available"); - sess = waiter.take(); - } - if (inProcessPrepare) { - try { - sess.prepareReadWriteTransaction(); - } catch (Throwable t) { - sess = null; - SpannerException e = newSpannerException(t); - if (!isClosed()) { - handlePrepareSessionFailure(e, sess, false); + inProcessPrepare = true; + numSessionsInProcessPrepared++; + PooledSession readSession = readSessions.poll(); + if (readSession != null) { + // Create a read/write transaction in-process if there is already a queue for prepared + // sessions. This is more efficient than doing it asynchronously, as it scales with + // the number of user threads. The thread pool for asynchronously preparing sessions + // is fixed. + span.addAnnotation( + "Acquired read only session. Preparing in-process for read write transaction"); + sess = readSession; + } else { + span.addAnnotation("No session available"); + maybeCreateSession(); } - if (!isSessionNotFound(e)) { - throw e; + } + if (sess == null) { + waiter = new WaiterFuture(); + if (inProcessPrepare) { + // inProcessPrepare=true means that we have already determined that the queue for + // preparing read/write sessions is larger than the number of threads in the prepare + // thread pool, and that it's more efficient to do the prepare in-process. We will + // therefore create a waiter for a read-only session, even though a read/write session + // has been requested. + readWaiters.add(waiter); + } else { + readWriteWaiters.add(waiter); } } + } else { + span.addAnnotation("Acquired read write session"); } - if (sess != null) { - return checkoutSession(span, sess, waiter, true); - } + return checkoutSession(span, sess, waiter, true, inProcessPrepare); } } private PooledSessionFuture checkoutSession( - final Span span, final PooledSession sess, final Waiter waiter, boolean write) { - final PooledSessionFuture res; + final Span span, + final PooledSession readySession, + WaiterFuture waiter, + boolean write, + final boolean inProcessPrepare) { + Future sessionFuture; if (waiter != null) { logger.log( Level.FINE, @@ -2074,20 +2025,95 @@ private PooledSessionFuture checkoutSession( span.addAnnotation( String.format( "Waiting for %s session to be available", write ? "read write" : "read only")); - ScheduledExecutorService executor = write ? writeWaiterExecutor : readWaiterExecutor; - res = - createPooledSessionFuture( - executor.submit( - new Callable() { - @Override - public PooledSession call() throws Exception { - return waiter.take(); - } - }), - span); + + // sessionFuture = ApiFutures.transform(finalWaiter.waiter, new + // ApiFunction(){ + // @Override + // public PooledSession apply(SessionOrError input) { + // if (input.session != null) { + // return input.session; + // } + // throw input.e; + // } + // }, MoreExecutors.directExecutor()); + sessionFuture = waiter; } else { - res = createPooledSessionFuture(sess, span); + sessionFuture = ApiFutures.immediateFuture(readySession); } + SimpleForwardingFuture forwardingFuture = + new SimpleForwardingFuture(sessionFuture) { + private volatile boolean initialized = false; + private final Object prepareLock = new Object(); + private volatile PooledSession result; + + @Override + public PooledSession get() throws InterruptedException, ExecutionException { + try { + return initialize(super.get()); + } catch (ExecutionException e) { + throw SpannerExceptionFactory.newSpannerException(e.getCause()); + } catch (InterruptedException e) { + throw SpannerExceptionFactory.propagateInterrupt(e); + } + } + + @Override + public PooledSession get(long timeout, TimeUnit unit) + throws InterruptedException, ExecutionException, TimeoutException { + try { + return initialize(super.get(timeout, unit)); + } catch (ExecutionException e) { + throw SpannerExceptionFactory.newSpannerException(e.getCause()); + } catch (InterruptedException e) { + throw SpannerExceptionFactory.propagateInterrupt(e); + } catch (TimeoutException e) { + throw SpannerExceptionFactory.propagateTimeout(e); + } + } + + private PooledSession initialize(PooledSession sess) { + if (!initialized) { + synchronized (prepareLock) { + if (!initialized) { + result = prepare(sess); + initialized = true; + } + } + } + return result; + } + + private PooledSession prepare(PooledSession sess) { + if (inProcessPrepare && !sess.delegate.hasReadyTransaction()) { + while (true) { + try { + sess.prepareReadWriteTransaction(); + break; + } catch (Throwable t) { + if (isClosed()) { + span.addAnnotation("Pool has been closed"); + throw new IllegalStateException("Pool has been closed"); + } + SpannerException e = newSpannerException(t); + WaiterFuture waiter = new WaiterFuture(); + synchronized (lock) { + handlePrepareSessionFailure(e, sess, false); + if (!isSessionNotFound(e)) { + throw e; + } + readWaiters.add(waiter); + } + sess = waiter.get(); + if (sess.delegate.hasReadyTransaction()) { + break; + } + } + } + } + return sess; + } + }; + PooledSessionFuture res = createPooledSessionFuture(forwardingFuture, span); res.markCheckedOut(); return res; } @@ -2207,10 +2233,12 @@ private void handleCreateSessionsFailure(SpannerException e, int count) { break; } } - this.resourceNotFoundException = - MoreObjects.firstNonNull( - this.resourceNotFoundException, - isDatabaseOrInstanceNotFound(e) ? (ResourceNotFoundException) e : null); + if (isDatabaseOrInstanceNotFound(e)) { + this.resourceNotFoundException = + MoreObjects.firstNonNull( + this.resourceNotFoundException, + isDatabaseOrInstanceNotFound(e) ? (ResourceNotFoundException) e : null); + } } } @@ -2230,14 +2258,13 @@ private void handlePrepareSessionFailure( readWaiters.poll().put(e); } // Remove the session from the pool. - allSessions.remove(session); - if (isClosed()) { - decrementPendingClosures(1); + removeFromPool(session); + if (isDatabaseOrInstanceNotFound(e)) { + this.resourceNotFoundException = + MoreObjects.firstNonNull( + this.resourceNotFoundException, + isDatabaseOrInstanceNotFound(e) ? (ResourceNotFoundException) e : null); } - this.resourceNotFoundException = - MoreObjects.firstNonNull( - this.resourceNotFoundException, - isDatabaseOrInstanceNotFound(e) ? (ResourceNotFoundException) e : null); } else if (informFirstWaiter && readWriteWaiters.size() > 0) { releaseSession(session, Position.FIRST); readWriteWaiters.poll().put(e); @@ -2266,7 +2293,7 @@ ListenableFuture closeAsync() { throw new IllegalStateException("Close has already been invoked"); } // Fail all pending waiters. - Waiter waiter = readWaiters.poll(); + WaiterFuture waiter = readWaiters.poll(); while (waiter != null) { waiter.put(newSpannerException(ErrorCode.INTERNAL, "Client has been closed")); waiter = readWaiters.poll(); @@ -2287,8 +2314,6 @@ ListenableFuture closeAsync() { readSessions.clear(); writePreparedSessions.clear(); prepareExecutor.shutdown(); - readWaiterExecutor.shutdown(); - writeWaiterExecutor.shutdown(); executor.submit( new Runnable() { @Override diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/BaseSessionPoolTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/BaseSessionPoolTest.java index 26bbef4535b..1bcb303f72b 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/BaseSessionPoolTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/BaseSessionPoolTest.java @@ -59,7 +59,7 @@ public void release(ScheduledExecutorService executor) { } SessionImpl mockSession() { - SessionImpl session = mock(SessionImpl.class); + final SessionImpl session = mock(SessionImpl.class); when(session.getName()) .thenReturn( "projects/dummy/instances/dummy/database/dummy/sessions/session" + sessionIndex); diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolMaintainerTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolMaintainerTest.java index 8d1b7804327..d71a03a31c8 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolMaintainerTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolMaintainerTest.java @@ -25,6 +25,7 @@ import com.google.cloud.spanner.SessionClient.SessionConsumer; import com.google.cloud.spanner.SessionPool.PooledSession; +import com.google.cloud.spanner.SessionPool.PooledSessionFuture; import com.google.cloud.spanner.SessionPool.SessionConsumerImpl; import com.google.common.base.Function; import java.util.ArrayList; @@ -200,7 +201,7 @@ public void testKeepAlive() throws Exception { assertThat(pingedSessions).containsExactly(session1.getName(), 2, session2.getName(), 3); // Update the last use date and release the session to the pool and do another maintenance // cycle. - ((PooledSession) session6).markUsed(); + ((PooledSessionFuture) session6).get().markUsed(); session6.close(); runMaintainanceLoop(clock, pool, 3); assertThat(pingedSessions).containsExactly(session1.getName(), 2, session2.getName(), 3); @@ -261,9 +262,9 @@ public void testIdleSessions() throws Exception { // Now check out three sessions so the pool will create an additional session. The pool will // only keep 2 sessions alive, as that is the setting for MinSessions. - Session session3 = pool.getReadSession(); - Session session4 = pool.getReadSession(); - Session session5 = pool.getReadSession(); + Session session3 = pool.getReadSession().get(); + Session session4 = pool.getReadSession().get(); + Session session5 = pool.getReadSession().get(); // Note that session2 was now the first session in the pool as it was the last to receive a // ping. assertThat(session3.getName()).isEqualTo(session2.getName()); @@ -278,9 +279,9 @@ public void testIdleSessions() throws Exception { assertThat(pool.totalSessions()).isEqualTo(2); // Check out three sessions again and keep one session checked out. - Session session6 = pool.getReadSession(); - Session session7 = pool.getReadSession(); - Session session8 = pool.getReadSession(); + Session session6 = pool.getReadSession().get(); + Session session7 = pool.getReadSession().get(); + Session session8 = pool.getReadSession().get(); session8.close(); session7.close(); // Now advance the clock to idle sessions. This should remove session8 from the pool. @@ -292,9 +293,9 @@ public void testIdleSessions() throws Exception { // Check out three sessions and keep them all checked out. No sessions should be removed from // the pool. - Session session9 = pool.getReadSession(); - Session session10 = pool.getReadSession(); - Session session11 = pool.getReadSession(); + Session session9 = pool.getReadSession().get(); + Session session10 = pool.getReadSession().get(); + Session session11 = pool.getReadSession().get(); runMaintainanceLoop(clock, pool, loopsToIdleSessions); assertThat(idledSessions).containsExactly(session5, session8); assertThat(pool.totalSessions()).isEqualTo(3); diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolStressTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolStressTest.java index 9bd989b98f4..24edff9e6bd 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolStressTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolStressTest.java @@ -17,7 +17,7 @@ package com.google.cloud.spanner; import static com.google.common.truth.Truth.assertThat; -import static org.mockito.Mockito.any; +import static org.mockito.Matchers.any; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; @@ -25,11 +25,12 @@ import com.google.api.core.ApiFuture; import com.google.api.core.ApiFutures; import com.google.cloud.spanner.SessionClient.SessionConsumer; -import com.google.cloud.spanner.SessionPool.PooledSessionFuture; import com.google.cloud.spanner.SessionPool.PooledSession; +import com.google.cloud.spanner.SessionPool.PooledSessionFuture; import com.google.cloud.spanner.SessionPool.SessionConsumerImpl; import com.google.common.base.Function; import com.google.common.util.concurrent.Uninterruptibles; +import com.google.protobuf.ByteString; import com.google.protobuf.Empty; import java.util.ArrayList; import java.util.Collection; @@ -40,6 +41,8 @@ import java.util.Random; import java.util.Set; import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import org.junit.Test; @@ -67,6 +70,7 @@ public class SessionPoolStressTest extends BaseSessionPoolTest { DatabaseId db = DatabaseId.of("projects/p/instances/i/databases/unused"); SessionPool pool; SessionPoolOptions options; + ExecutorService createExecutor = Executors.newSingleThreadExecutor(); Object lock = new Object(); Random random = new Random(); FakeClock clock = new FakeClock(); @@ -98,43 +102,32 @@ private void setupSpanner(DatabaseId db) { SessionClient sessionClient = mock(SessionClient.class); when(mockSpanner.getSessionClient(db)).thenReturn(sessionClient); when(mockSpanner.getOptions()).thenReturn(spannerOptions); - when(sessionClient.createSession()) - .thenAnswer( - new Answer() { - - @Override - public SessionImpl answer(InvocationOnMock invocation) throws Throwable { - synchronized (lock) { - SessionImpl session = mockSession(); - setupSession(session); - - sessions.put(session.getName(), false); - if (sessions.size() > maxAliveSessions) { - maxAliveSessions = sessions.size(); - } - return session; - } - } - }); doAnswer( new Answer() { @Override - public Void answer(InvocationOnMock invocation) throws Throwable { - int sessionCount = invocation.getArgumentAt(0, Integer.class); - for (int s = 0; s < sessionCount; s++) { - synchronized (lock) { - SessionImpl session = mockSession(); - setupSession(session); + public Void answer(final InvocationOnMock invocation) throws Throwable { + createExecutor.submit( + new Runnable() { + @Override + public void run() { + int sessionCount = invocation.getArgumentAt(0, Integer.class); + for (int s = 0; s < sessionCount; s++) { + SessionImpl session; + synchronized (lock) { + session = mockSession(); + setupSession(session); - sessions.put(session.getName(), false); - if (sessions.size() > maxAliveSessions) { - maxAliveSessions = sessions.size(); - } - SessionConsumerImpl consumer = - invocation.getArgumentAt(2, SessionConsumerImpl.class); - consumer.onSessionReady(session); - } - } + sessions.put(session.getName(), false); + if (sessions.size() > maxAliveSessions) { + maxAliveSessions = sessions.size(); + } + } + SessionConsumerImpl consumer = + invocation.getArgumentAt(2, SessionConsumerImpl.class); + consumer.onSessionReady(session); + } + } + }); return null; } }) @@ -190,36 +183,43 @@ public Void answer(InvocationOnMock invocation) throws Throwable { expireSession(session); throw SpannerExceptionFactoryTest.newSessionNotFoundException(session.getName()); } + String name = session.getName(); synchronized (lock) { - if (sessions.put(session.getName(), true)) { + if (sessions.put(name, true)) { setFailed(); } + session.readyTransactionId = ByteString.copyFromUtf8("foo"); } return null; } }) .when(session) .prepareReadWriteTransaction(); + when(session.hasReadyTransaction()).thenCallRealMethod(); } private void expireSession(Session session) { + String name = session.getName(); synchronized (lock) { - sessions.remove(session.getName()); - expiredSessions.add(session.getName()); + sessions.remove(name); + expiredSessions.add(name); } } private void assertWritePrepared(Session session) { + String name = session.getName(); synchronized (lock) { - if (!sessions.get(session.getName())) { + if (!sessions.containsKey(name) || !sessions.get(name)) { setFailed(); } } } - private void resetTransaction(Session session) { + private void resetTransaction(SessionImpl session) { + String name = session.getName(); synchronized (lock) { - sessions.put(session.getName(), false); + session.readyTransactionId = null; + sessions.put(name, false); } } @@ -265,8 +265,9 @@ public void stressTest() throws Exception { new Function() { @Override public Void apply(PooledSession pooled) { + String name = pooled.getName(); synchronized (lock) { - sessions.remove(pooled.getName()); + sessions.remove(name); return null; } } @@ -283,15 +284,15 @@ public void run() { PooledSessionFuture session = null; if (random.nextInt(10) < writeOperationFraction) { session = pool.getReadWriteSession(); - session.get(); - assertWritePrepared(session); + PooledSession sess = session.get(); + assertWritePrepared(sess); } else { session = pool.getReadSession(); session.get(); } Uninterruptibles.sleepUninterruptibly( random.nextInt(5), TimeUnit.MILLISECONDS); - resetTransaction(session); + resetTransaction(session.get().delegate); session.close(); } catch (SpannerException e) { if (e.getErrorCode() != ErrorCode.RESOURCE_EXHAUSTED || shouldBlock) { diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java index acf48c4510e..18f45494d82 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java @@ -320,7 +320,7 @@ public Void call() throws Exception { insideCreation.await(); pool.closeAsync(); releaseCreation.countDown(); - latch.await(); + latch.await(5L, TimeUnit.SECONDS); assertThat(failed.get()).isTrue(); } From 51b511349e094ae5bac9225a9d22d6d9cad1c212 Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Wed, 29 Apr 2020 09:25:39 +0200 Subject: [PATCH 25/49] format: run code formatter --- .../src/main/java/com/google/cloud/spanner/SessionPool.java | 2 +- .../java/com/google/cloud/spanner/DatabaseClientImplTest.java | 4 +--- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java index 7ad5293259b..267f8a363f3 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java @@ -52,9 +52,9 @@ import com.google.common.base.Supplier; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; import com.google.common.util.concurrent.ForwardingFuture; import com.google.common.util.concurrent.ForwardingFuture.SimpleForwardingFuture; -import com.google.common.collect.ImmutableSet; import com.google.common.util.concurrent.ListenableFuture; import com.google.common.util.concurrent.MoreExecutors; import com.google.common.util.concurrent.SettableFuture; diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java index b370fd9e6ac..93f91f3004d 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java @@ -58,9 +58,7 @@ import org.junit.AfterClass; import org.junit.Before; import org.junit.BeforeClass; -import org.junit.Rule; import org.junit.Test; -import org.junit.rules.Timeout; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; import org.threeten.bp.Duration; @@ -86,7 +84,7 @@ public class DatabaseClientImplTest { private Spanner spanner; private Spanner spannerWithEmptySessionPool; -// @Rule public Timeout globalTimeout = new Timeout(5L, TimeUnit.SECONDS); + // @Rule public Timeout globalTimeout = new Timeout(5L, TimeUnit.SECONDS); @BeforeClass public static void startStaticServer() throws IOException { From a03380be0c2a720131615f793a4b807adc593b9d Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Wed, 29 Apr 2020 10:23:42 +0200 Subject: [PATCH 26/49] tests: fix test case + remove commented code --- .../java/com/google/cloud/spanner/SessionPool.java | 11 ----------- .../com/google/cloud/spanner/SessionPoolTest.java | 6 ++++-- 2 files changed, 4 insertions(+), 13 deletions(-) diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java index 267f8a363f3..451063403e2 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java @@ -2041,17 +2041,6 @@ private PooledSessionFuture checkoutSession( span.addAnnotation( String.format( "Waiting for %s session to be available", write ? "read write" : "read only")); - - // sessionFuture = ApiFutures.transform(finalWaiter.waiter, new - // ApiFunction(){ - // @Override - // public PooledSession apply(SessionOrError input) { - // if (input.session != null) { - // return input.session; - // } - // throw input.e; - // } - // }, MoreExecutors.directExecutor()); sessionFuture = waiter; } else { sessionFuture = ApiFutures.immediateFuture(readySession); diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java index 18f45494d82..c57545bad1b 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java @@ -970,8 +970,10 @@ public void run() { FakeClock clock = new FakeClock(); clock.currentTimeMillis = System.currentTimeMillis(); pool = createPool(clock); - Session session1 = pool.getReadSession(); - Session session2 = pool.getReadSession(); + PooledSessionFuture session1 = pool.getReadSession(); + PooledSessionFuture session2 = pool.getReadSession(); + session1.get(); + session2.get(); session1.close(); session2.close(); runMaintainanceLoop(clock, pool, pool.poolMaintainer.numKeepAliveCycles); From 8bc3a135382517e3e08097bb5e669946496d3028 Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Wed, 29 Apr 2020 15:01:58 +0200 Subject: [PATCH 27/49] fix: AsyncResultSet should throw Cancelled --- .../cloud/spanner/AsyncResultSetImpl.java | 9 +++- .../google/cloud/spanner/ReadAsyncTest.java | 51 +++++++++---------- 2 files changed, 32 insertions(+), 28 deletions(-) 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 a6739688157..931e87033a1 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 @@ -65,6 +65,9 @@ private State(boolean shouldStop) { static final int DEFAULT_BUFFER_SIZE = 10; private static final int MAX_WAIT_FOR_BUFFER_CONSUMPTION = 10; + private static final SpannerException CANCELLED_EXCEPTION = + SpannerExceptionFactory.newSpannerException( + ErrorCode.CANCELLED, "This AsyncResultSet has been cancelled"); private final Object monitor = new Object(); private boolean closed; @@ -187,8 +190,7 @@ public CursorState tryNext() throws SpannerException { synchronized (monitor) { if (state == State.CANCELLED) { cursorReturnedDoneOrException = true; - throw SpannerExceptionFactory.newSpannerException( - ErrorCode.CANCELLED, "This AsyncResultSet has been cancelled"); + throw CANCELLED_EXCEPTION; } if (buffer.isEmpty() && executionException != null) { cursorReturnedDoneOrException = true; @@ -385,6 +387,9 @@ public Void call() throws Exception { if (executionException != null) { throw executionException; } + if (state == State.CANCELLED) { + throw CANCELLED_EXCEPTION; + } } } return null; diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadAsyncTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadAsyncTest.java index 944a7d70f51..50350fd2a14 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadAsyncTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadAsyncTest.java @@ -442,43 +442,42 @@ public CallbackResponse cursorReady(AsyncResultSet resultSet) { @Test public void cancel() throws Exception { final List values = new LinkedList<>(); - final SettableApiFuture finished = SettableApiFuture.create(); final CountDownLatch receivedFirstRow = new CountDownLatch(1); final CountDownLatch cancelled = new CountDownLatch(1); + final ApiFuture res; try (AsyncResultSet rs = client.singleUse().readAsync(READ_TABLE_NAME, KeySet.all(), READ_COLUMN_NAMES)) { - rs.setCallback( - executor, - new ReadyCallback() { - @Override - public CallbackResponse cursorReady(AsyncResultSet resultSet) { - try { - while (true) { - switch (resultSet.tryNext()) { - case DONE: - finished.set(true); - return CallbackResponse.DONE; - case NOT_READY: - return CallbackResponse.CONTINUE; - case OK: - values.add(resultSet.getString("Value")); - receivedFirstRow.countDown(); - cancelled.await(); - break; + res = + 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: + values.add(resultSet.getString("Value")); + receivedFirstRow.countDown(); + cancelled.await(); + break; + } + } + } catch (Throwable t) { + return CallbackResponse.DONE; } } - } catch (Throwable t) { - finished.setException(t); - return CallbackResponse.DONE; - } - } - }); + }); receivedFirstRow.await(); rs.cancel(); } cancelled.countDown(); try { - finished.get(); + res.get(); fail("missing expected exception"); } catch (ExecutionException e) { assertThat(e.getCause()).isInstanceOf(SpannerException.class); From e486c54a913797d73e9f83c00491812f5deb7d6d Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Mon, 11 May 2020 15:05:54 +0200 Subject: [PATCH 28/49] feat: expose DatabaseId.of(String name) --- .../src/main/java/com/google/cloud/spanner/DatabaseId.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseId.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseId.java index d2c732750e9..dd13df65e81 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseId.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseId.java @@ -81,7 +81,7 @@ public String toString() { * projects/PROJECT_ID/instances/INSTANCE_ID/databases/DATABASE_ID} * @throws IllegalArgumentException if {@code name} does not conform to the expected pattern */ - static DatabaseId of(String name) { + public static DatabaseId of(String name) { Preconditions.checkNotNull(name); Map parts = NAME_TEMPLATE.match(name); Preconditions.checkArgument( From a7dd1dd3376784cf24e8aeb64a7155e26b9192d5 Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Tue, 12 May 2020 13:37:08 +0200 Subject: [PATCH 29/49] deps: set version to 1.53 to match bom --- google-cloud-spanner-bom/pom.xml | 16 ++++++++-------- google-cloud-spanner/pom.xml | 4 ++-- .../pom.xml | 4 ++-- .../pom.xml | 4 ++-- grpc-google-cloud-spanner-v1/pom.xml | 4 ++-- pom.xml | 16 ++++++++-------- .../pom.xml | 4 ++-- .../pom.xml | 4 ++-- proto-google-cloud-spanner-v1/pom.xml | 4 ++-- versions.txt | 14 +++++++------- 10 files changed, 37 insertions(+), 37 deletions(-) diff --git a/google-cloud-spanner-bom/pom.xml b/google-cloud-spanner-bom/pom.xml index 3c9e80b8a98..b67e3fb9436 100644 --- a/google-cloud-spanner-bom/pom.xml +++ b/google-cloud-spanner-bom/pom.xml @@ -3,7 +3,7 @@ 4.0.0 com.google.cloud google-cloud-spanner-bom - 1.54.0 + 1.53.0 pom com.google.cloud @@ -64,37 +64,37 @@ com.google.api.grpc proto-google-cloud-spanner-admin-instance-v1 - 1.54.0 + 1.53.0 com.google.api.grpc grpc-google-cloud-spanner-v1 - 1.54.0 + 1.53.0 com.google.api.grpc proto-google-cloud-spanner-v1 - 1.54.0 + 1.53.0 com.google.api.grpc proto-google-cloud-spanner-admin-database-v1 - 1.54.0 + 1.53.0 com.google.cloud google-cloud-spanner - 1.54.0 + 1.53.0 com.google.api.grpc grpc-google-cloud-spanner-admin-instance-v1 - 1.54.0 + 1.53.0 com.google.api.grpc grpc-google-cloud-spanner-admin-database-v1 - 1.54.0 + 1.53.0 diff --git a/google-cloud-spanner/pom.xml b/google-cloud-spanner/pom.xml index abc4c0a687d..ff76c61b3f6 100644 --- a/google-cloud-spanner/pom.xml +++ b/google-cloud-spanner/pom.xml @@ -3,7 +3,7 @@ 4.0.0 com.google.cloud google-cloud-spanner - 1.54.0 + 1.53.0 jar Google Cloud Spanner https://github.com/googleapis/java-spanner @@ -11,7 +11,7 @@ com.google.cloud google-cloud-spanner-parent - 1.54.0 + 1.53.0 google-cloud-spanner diff --git a/grpc-google-cloud-spanner-admin-database-v1/pom.xml b/grpc-google-cloud-spanner-admin-database-v1/pom.xml index 76e3cdc389d..eee25be70cd 100644 --- a/grpc-google-cloud-spanner-admin-database-v1/pom.xml +++ b/grpc-google-cloud-spanner-admin-database-v1/pom.xml @@ -4,13 +4,13 @@ 4.0.0 com.google.api.grpc grpc-google-cloud-spanner-admin-database-v1 - 1.54.0 + 1.53.0 grpc-google-cloud-spanner-admin-database-v1 GRPC library for grpc-google-cloud-spanner-admin-database-v1 com.google.cloud google-cloud-spanner-parent - 1.54.0 + 1.53.0 diff --git a/grpc-google-cloud-spanner-admin-instance-v1/pom.xml b/grpc-google-cloud-spanner-admin-instance-v1/pom.xml index 85ae91dd144..7cd460a9588 100644 --- a/grpc-google-cloud-spanner-admin-instance-v1/pom.xml +++ b/grpc-google-cloud-spanner-admin-instance-v1/pom.xml @@ -4,13 +4,13 @@ 4.0.0 com.google.api.grpc grpc-google-cloud-spanner-admin-instance-v1 - 1.54.0 + 1.53.0 grpc-google-cloud-spanner-admin-instance-v1 GRPC library for grpc-google-cloud-spanner-admin-instance-v1 com.google.cloud google-cloud-spanner-parent - 1.54.0 + 1.53.0 diff --git a/grpc-google-cloud-spanner-v1/pom.xml b/grpc-google-cloud-spanner-v1/pom.xml index 7a3ddb3d90d..784dbbcd9fa 100644 --- a/grpc-google-cloud-spanner-v1/pom.xml +++ b/grpc-google-cloud-spanner-v1/pom.xml @@ -4,13 +4,13 @@ 4.0.0 com.google.api.grpc grpc-google-cloud-spanner-v1 - 1.54.0 + 1.53.0 grpc-google-cloud-spanner-v1 GRPC library for grpc-google-cloud-spanner-v1 com.google.cloud google-cloud-spanner-parent - 1.54.0 + 1.53.0 diff --git a/pom.xml b/pom.xml index 9f6bcebc299..3e53f9b4841 100644 --- a/pom.xml +++ b/pom.xml @@ -4,7 +4,7 @@ com.google.cloud google-cloud-spanner-parent pom - 1.54.0 + 1.53.0 Google Cloud Spanner Parent https://github.com/googleapis/java-spanner @@ -70,37 +70,37 @@ com.google.api.grpc proto-google-cloud-spanner-admin-instance-v1 - 1.54.0 + 1.53.0 com.google.api.grpc proto-google-cloud-spanner-v1 - 1.54.0 + 1.53.0 com.google.api.grpc proto-google-cloud-spanner-admin-database-v1 - 1.54.0 + 1.53.0 com.google.api.grpc grpc-google-cloud-spanner-v1 - 1.54.0 + 1.53.0 com.google.api.grpc grpc-google-cloud-spanner-admin-instance-v1 - 1.54.0 + 1.53.0 com.google.api.grpc grpc-google-cloud-spanner-admin-database-v1 - 1.54.0 + 1.53.0 com.google.cloud google-cloud-spanner - 1.54.0 + 1.53.0 diff --git a/proto-google-cloud-spanner-admin-database-v1/pom.xml b/proto-google-cloud-spanner-admin-database-v1/pom.xml index 3c69ba5cd78..77eb238bf2a 100644 --- a/proto-google-cloud-spanner-admin-database-v1/pom.xml +++ b/proto-google-cloud-spanner-admin-database-v1/pom.xml @@ -4,13 +4,13 @@ 4.0.0 com.google.api.grpc proto-google-cloud-spanner-admin-database-v1 - 1.54.0 + 1.53.0 proto-google-cloud-spanner-admin-database-v1 PROTO library for proto-google-cloud-spanner-admin-database-v1 com.google.cloud google-cloud-spanner-parent - 1.54.0 + 1.53.0 diff --git a/proto-google-cloud-spanner-admin-instance-v1/pom.xml b/proto-google-cloud-spanner-admin-instance-v1/pom.xml index bd1f1b48128..738ccd42c58 100644 --- a/proto-google-cloud-spanner-admin-instance-v1/pom.xml +++ b/proto-google-cloud-spanner-admin-instance-v1/pom.xml @@ -4,13 +4,13 @@ 4.0.0 com.google.api.grpc proto-google-cloud-spanner-admin-instance-v1 - 1.54.0 + 1.53.0 proto-google-cloud-spanner-admin-instance-v1 PROTO library for proto-google-cloud-spanner-admin-instance-v1 com.google.cloud google-cloud-spanner-parent - 1.54.0 + 1.53.0 diff --git a/proto-google-cloud-spanner-v1/pom.xml b/proto-google-cloud-spanner-v1/pom.xml index 25af3e2c2fb..8efb3a4677d 100644 --- a/proto-google-cloud-spanner-v1/pom.xml +++ b/proto-google-cloud-spanner-v1/pom.xml @@ -4,13 +4,13 @@ 4.0.0 com.google.api.grpc proto-google-cloud-spanner-v1 - 1.54.0 + 1.53.0 proto-google-cloud-spanner-v1 PROTO library for proto-google-cloud-spanner-v1 com.google.cloud google-cloud-spanner-parent - 1.54.0 + 1.53.0 diff --git a/versions.txt b/versions.txt index 907ad4e3708..79fe4e3607d 100644 --- a/versions.txt +++ b/versions.txt @@ -1,10 +1,10 @@ # Format: # module:released-version:current-version -proto-google-cloud-spanner-admin-instance-v1:1.54.0:1.54.0 -proto-google-cloud-spanner-v1:1.54.0:1.54.0 -proto-google-cloud-spanner-admin-database-v1:1.54.0:1.54.0 -grpc-google-cloud-spanner-v1:1.54.0:1.54.0 -grpc-google-cloud-spanner-admin-instance-v1:1.54.0:1.54.0 -grpc-google-cloud-spanner-admin-database-v1:1.54.0:1.54.0 -google-cloud-spanner:1.54.0:1.54.0 \ No newline at end of file +proto-google-cloud-spanner-admin-instance-v1:1.53.0:1.53.0 +proto-google-cloud-spanner-v1:1.53.0:1.53.0 +proto-google-cloud-spanner-admin-database-v1:1.53.0:1.53.0 +grpc-google-cloud-spanner-v1:1.53.0:1.53.0 +grpc-google-cloud-spanner-admin-instance-v1:1.53.0:1.53.0 +grpc-google-cloud-spanner-admin-database-v1:1.53.0:1.53.0 +google-cloud-spanner:1.53.0:1.53.0 \ No newline at end of file From c662ed75db1d9e4630ef0d7289f95fb86b081925 Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Tue, 9 Jun 2020 16:33:28 +0200 Subject: [PATCH 30/49] feat: steps to add async support for tx manager --- .../clirr-ignored-differences.xml | 10 ++ .../cloud/spanner/AsyncResultSetImpl.java | 6 +- .../spanner/AsyncTransactionManager.java | 84 +++++++++ .../spanner/AsyncTransactionManagerImpl.java | 158 ++++++++++++++++ .../google/cloud/spanner/DatabaseClient.java | 2 + .../cloud/spanner/DatabaseClientImpl.java | 11 ++ .../com/google/cloud/spanner/ResultSets.java | 10 ++ .../com/google/cloud/spanner/SessionImpl.java | 82 ++++++--- .../com/google/cloud/spanner/SessionPool.java | 129 +++++++++++-- .../com/google/cloud/spanner/Spanner.java | 4 + .../com/google/cloud/spanner/SpannerImpl.java | 2 +- .../com/google/cloud/spanner/TraceUtil.java | 2 +- .../cloud/spanner/TransactionManagerImpl.java | 4 + .../cloud/spanner/TransactionRunnerImpl.java | 170 ++++++++++++++++-- .../AbstractMultiUseTransaction.java | 11 ++ .../connection/AsyncChecksumResultSet.java | 73 ++++++++ .../cloud/spanner/connection/Connection.java | 3 + .../spanner/connection/ConnectionImpl.java | 49 +++++ .../cloud/spanner/connection/DdlBatch.java | 33 ++-- .../cloud/spanner/connection/DmlBatch.java | 8 + .../connection/ReadWriteTransaction.java | 44 +++++ .../connection/SingleUseTransaction.java | 55 ++++-- .../connection/StatementResultImpl.java | 1 - .../cloud/spanner/connection/UnitOfWork.java | 4 + .../cloud/spanner/spi/v1/GapicSpannerRpc.java | 18 +- .../cloud/spanner/spi/v1/SpannerRpc.java | 4 + .../google/cloud/spanner/AsyncRunnerTest.java | 20 ++- .../cloud/spanner/DatabaseClientImplTest.java | 108 ++++++++++- .../google/cloud/spanner/SessionImplTest.java | 10 +- .../google/cloud/spanner/SessionPoolTest.java | 8 +- .../spanner/TransactionManagerImplTest.java | 24 +-- .../spanner/TransactionRunnerImplTest.java | 31 ++-- .../connection/AbstractMockServerTest.java | 7 + .../connection/AsyncConnectionApiTest.java | 112 ++++++++++++ .../connection/SingleUseTransactionTest.java | 1 + 35 files changed, 1171 insertions(+), 127 deletions(-) create mode 100644 google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManager.java create mode 100644 google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManagerImpl.java create mode 100644 google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AsyncChecksumResultSet.java create mode 100644 google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AsyncConnectionApiTest.java diff --git a/google-cloud-spanner/clirr-ignored-differences.xml b/google-cloud-spanner/clirr-ignored-differences.xml index 5c8b59d72b3..42875a9a8bf 100644 --- a/google-cloud-spanner/clirr-ignored-differences.xml +++ b/google-cloud-spanner/clirr-ignored-differences.xml @@ -157,6 +157,16 @@ com/google/cloud/spanner/DatabaseClient * runAsync(*) + + 7012 + com/google/cloud/spanner/DatabaseClient + * transactionManagerAsync(*) + + + 7012 + com/google/cloud/spanner/Spanner + * getAsyncExecutorProvider(*) + 7012 com/google/cloud/spanner/ReadContext 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 931e87033a1..d8f35c31e22 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 @@ -138,10 +138,10 @@ private State(boolean shouldStop) { AsyncResultSetImpl(ExecutorProvider executorProvider, ResultSet delegate, int bufferSize) { super(delegate); - this.buffer = new LinkedBlockingDeque<>(bufferSize); - this.executorProvider = executorProvider; + this.executorProvider = Preconditions.checkNotNull(executorProvider); + this.delegateResultSet = Preconditions.checkNotNull(delegate); this.service = MoreExecutors.listeningDecorator(executorProvider.getExecutor()); - this.delegateResultSet = delegate; + this.buffer = new LinkedBlockingDeque<>(bufferSize); } /** diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManager.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManager.java new file mode 100644 index 00000000000..a86c5a776c4 --- /dev/null +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManager.java @@ -0,0 +1,84 @@ +/* + * Copyright 2017 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.spanner; + +import com.google.api.core.ApiFuture; +import com.google.cloud.Timestamp; +import com.google.cloud.spanner.TransactionManager.TransactionState; + +/** + * An interface for managing the life cycle of a read write transaction including all its retries. + * See {@link TransactionContext} for a description of transaction semantics. + * + *

    At any point in time there can be at most one active transaction in this manager. When that + * transaction is committed, if it fails with an {@code ABORTED} error, calling {@link + * #resetForRetry()} would create a new {@link TransactionContext}. The newly created transaction + * would use the same session thus increasing its lock priority. If the transaction is committed + * successfully, or is rolled back or commit fails with any error other than {@code ABORTED}, the + * manager is considered complete and no further transactions are allowed to be created in it. + * + *

    Every {@code AsyncTransactionManager} should either be committed or rolled back. Failure to do so + * can cause resources to be leaked and deadlocks. Easiest way to guarantee this is by calling + * {@link #close()} in a finally block. + * + * @see DatabaseClient#transactionManager() + */ +public interface AsyncTransactionManager extends AutoCloseable { + /** + * Creates a new read write transaction. This must be called before doing any other operation and + * can only be called once. To create a new transaction for subsequent retries, see {@link + * #resetForRetry()}. + */ + ApiFuture beginAsync(); + + /** + * Commits the currently active transaction. If the transaction was already aborted, then this + * would throw an {@link AbortedException}. + */ + ApiFuture commitAsync(); + + /** + * Rolls back the currently active transaction. In most cases there should be no need to call this + * explicitly since {@link #close()} would automatically roll back any active transaction. + */ + ApiFuture rollbackAsync(); + + /** + * Creates a new transaction for retry. This should only be called if the previous transaction + * failed with {@code ABORTED}. In all other cases, this will throw an {@link + * IllegalStateException}. Users should backoff before calling this method. Backoff delay is + * specified by {@link SpannerException#getRetryDelayInMillis()} on the {@code SpannerException} + * throw by the previous commit call. + */ + ApiFuture resetForRetryAsync(); + + /** + * Returns the commit timestamp if the transaction committed successfully otherwise it will throw + * {@code IllegalStateException}. + */ + ApiFuture getCommitTimestampAsync(); + + /** Returns the state of the transaction. */ + TransactionState getState(); + + /** + * Closes the manager. If there is an active transaction, it will be rolled back. Underlying + * session will be released back to the session pool. + */ + @Override + void close(); +} diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManagerImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManagerImpl.java new file mode 100644 index 00000000000..945498407dd --- /dev/null +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManagerImpl.java @@ -0,0 +1,158 @@ +/* + * Copyright 2017 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.spanner; + +import java.util.concurrent.ExecutionException; +import com.google.api.core.ApiFuture; +import com.google.api.core.ApiFutures; +import com.google.api.core.SettableApiFuture; +import com.google.cloud.Timestamp; +import com.google.cloud.spanner.SessionImpl.SessionTransaction; +import com.google.cloud.spanner.TransactionManager.TransactionState; +import com.google.common.base.Preconditions; +import com.google.common.util.concurrent.MoreExecutors; +import com.google.common.util.concurrent.SettableFuture; +import io.opencensus.common.Scope; +import io.opencensus.trace.Span; +import io.opencensus.trace.Tracer; +import io.opencensus.trace.Tracing; + +/** Implementation of {@link AsyncTransactionManager}. */ +final class AsyncTransactionManagerImpl implements AsyncTransactionManager, SessionTransaction { + private static final Tracer tracer = Tracing.getTracer(); + + private final SessionImpl session; + private Span span; + + private TransactionRunnerImpl.TransactionContextImpl txn; + private TransactionState txnState; + + AsyncTransactionManagerImpl(SessionImpl session, Span span) { + this.session = session; + this.span = span; + } + + @Override + public void setSpan(Span span) { + this.span = span; + } + + @Override + public ApiFuture beginAsync() { + Preconditions.checkState(txn == null, "begin can only be called once"); + txnState = TransactionState.STARTED; + txn = session.newTransaction(); + session.setActive(this); + final SettableApiFuture res = SettableApiFuture.create(); + final ApiFuture fut = txn.ensureTxnAsync(); + fut.addListener(tracer.withSpan(span, new Runnable(){ + @Override + public void run() { + try { + fut.get(); + res.set(txn); + } catch (ExecutionException e) { + res.setException(e.getCause() == null ? e : e.getCause()); + } catch (InterruptedException e) { + res.setException(SpannerExceptionFactory.propagateInterrupt(e)); + } + } + }), MoreExecutors.directExecutor()); + return res; + } + + @Override + public ApiFuture commitAsync() { + Preconditions.checkState( + txnState == TransactionState.STARTED, + "commit can only be invoked if the transaction is in progress"); + SettableApiFuture res = SettableApiFuture.create(); + if (txn.isAborted()) { + txnState = TransactionState.ABORTED; + res.setException(SpannerExceptionFactory.newSpannerException( + ErrorCode.ABORTED, "Transaction already aborted")); + } + try { + txn.commit(); + txnState = TransactionState.COMMITTED; + return ApiFutures.immediateFuture(null); + } catch (AbortedException e1) { + txnState = TransactionState.ABORTED; + return ApiFutures.immediateFailedFuture(e1); + } catch (SpannerException e2) { + txnState = TransactionState.COMMIT_FAILED; + return ApiFutures.immediateFailedFuture(e2); + } + } + + @Override + public ApiFuture rollbackAsync() { + Preconditions.checkState( + txnState == TransactionState.STARTED, + "rollback can only be called if the transaction is in progress"); + try { + txn.rollback(); + } finally { + txnState = TransactionState.ROLLED_BACK; + } + return ApiFutures.immediateFuture(null); + } + + @Override + public ApiFuture resetForRetryAsync() { + if (txn == null || !txn.isAborted() && txnState != TransactionState.ABORTED) { + throw new IllegalStateException( + "resetForRetry can only be called if the previous attempt" + " aborted"); + } + try (Scope s = tracer.withSpan(span)) { + txn = session.newTransaction(); + txn.ensureTxn(); + txnState = TransactionState.STARTED; + return ApiFutures.immediateFuture(txn); + } + } + + @Override + public ApiFuture getCommitTimestampAsync() { + Preconditions.checkState( + txnState == TransactionState.COMMITTED, + "getCommitTimestamp can only be invoked if the transaction committed successfully"); + return ApiFutures.immediateFuture(txn.commitTimestamp()); + } + + @Override + public void close() { + try { + if (txnState == TransactionState.STARTED && !txn.isAborted()) { + txn.rollback(); + txnState = TransactionState.ROLLED_BACK; + } + } finally { + span.end(TraceUtil.END_SPAN_OPTIONS); + } + } + + @Override + public TransactionState getState() { + return txnState; + } + + @Override + public void invalidate() { + close(); + } +} diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClient.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClient.java index 3298e6a2ab0..2792fd0c866 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClient.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClient.java @@ -311,6 +311,8 @@ public interface DatabaseClient { */ AsyncRunner runAsync(); + AsyncTransactionManager transactionManagerAsync(); + /** * Returns the lower bound of rows modified by this DML statement. * diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClientImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClientImpl.java index 44f386a7273..5299d47d10b 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClientImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClientImpl.java @@ -201,6 +201,17 @@ public AsyncRunner runAsync() { } } + @Override + public AsyncTransactionManager transactionManagerAsync() { + Span span = tracer.spanBuilder(READ_WRITE_TRANSACTION).startSpan(); + try (Scope s = tracer.withSpan(span)) { + return getReadWriteSession().transactionManagerAsync(); + } catch (RuntimeException e) { + TraceUtil.endSpanWithFailure(span, e); + throw e; + } + } + @Override public long executePartitionedUpdate(final Statement stmt) { Span span = tracer.spanBuilder(PARTITION_DML_TRANSACTION).startSpan(); 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 d440162213d..278b15d967b 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,6 +16,7 @@ package com.google.cloud.spanner; +import com.google.api.gax.core.ExecutorProvider; import com.google.api.gax.core.InstantiatingExecutorProvider; import com.google.cloud.ByteArray; import com.google.cloud.Date; @@ -58,6 +59,15 @@ public static AsyncResultSet toAsyncResultSet(ResultSet delegate) { 100); } + /** + * Converts the given {@link ResultSet} to an {@link AsyncResultSet} using the given {@link + * ExecutorProvider}. + */ + public static AsyncResultSet toAsyncResultSet( + ResultSet delegate, ExecutorProvider executorProvider) { + return new AsyncResultSetImpl(executorProvider, delegate, 100); + } + private static class PrePopulatedResultSet implements ResultSet { private final List rows; private final Type type; diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java index a425ea56118..01fd8d758de 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java @@ -18,8 +18,10 @@ import static com.google.cloud.spanner.SpannerExceptionFactory.newSpannerException; import static com.google.common.base.Preconditions.checkNotNull; - +import com.google.api.core.ApiFunction; import com.google.api.core.ApiFuture; +import com.google.api.core.ApiFutures; +import com.google.api.core.SettableApiFuture; import com.google.cloud.Timestamp; import com.google.cloud.spanner.AbstractReadContext.MultiUseReadOnlyTransaction; import com.google.cloud.spanner.AbstractReadContext.SingleReadContext; @@ -28,6 +30,7 @@ import com.google.cloud.spanner.TransactionRunnerImpl.TransactionContextImpl; import com.google.cloud.spanner.spi.v1.SpannerRpc; import com.google.common.collect.Lists; +import com.google.common.util.concurrent.MoreExecutors; import com.google.protobuf.ByteString; import com.google.protobuf.Empty; import com.google.spanner.v1.BeginTransactionRequest; @@ -35,6 +38,7 @@ import com.google.spanner.v1.CommitResponse; import com.google.spanner.v1.Transaction; import com.google.spanner.v1.TransactionOptions; +import io.grpc.Context; import io.opencensus.common.Scope; import io.opencensus.trace.Span; import io.opencensus.trace.Tracer; @@ -43,6 +47,7 @@ import java.util.Collection; import java.util.List; import java.util.Map; +import java.util.concurrent.ExecutionException; import javax.annotation.Nullable; /** @@ -232,6 +237,16 @@ public AsyncRunner runAsync() { new TransactionRunnerImpl(this, spanner.getRpc(), spanner.getDefaultPrefetchChunks()))); } + @Override + public TransactionManager transactionManager() { + return new TransactionManagerImpl(this, currentSpan); + } + + @Override + public AsyncTransactionManager transactionManagerAsync() { + return new AsyncTransactionManagerImpl(this, currentSpan); + } + @Override public void prepareReadWriteTransaction() { setActive(null); @@ -257,27 +272,51 @@ public void close() { } ByteString beginTransaction() { - Span span = tracer.spanBuilder(SpannerImpl.BEGIN_TRANSACTION).startSpan(); - try (Scope s = tracer.withSpan(span)) { - final BeginTransactionRequest request = - BeginTransactionRequest.newBuilder() - .setSession(name) - .setOptions( - TransactionOptions.newBuilder() - .setReadWrite(TransactionOptions.ReadWrite.getDefaultInstance())) - .build(); - Transaction txn = spanner.getRpc().beginTransaction(request, options); - if (txn.getId().isEmpty()) { - throw newSpannerException(ErrorCode.INTERNAL, "Missing id in transaction\n" + getName()); - } - span.end(TraceUtil.END_SPAN_OPTIONS); - return txn.getId(); - } catch (RuntimeException e) { - TraceUtil.endSpanWithFailure(span, e); - throw e; + try { + return beginTransactionAsync().get(); + } catch (ExecutionException e) { + throw SpannerExceptionFactory.newSpannerException(e.getCause() == null ? e : e.getCause()); + } catch (InterruptedException e) { + throw SpannerExceptionFactory.propagateInterrupt(e); } } + ApiFuture beginTransactionAsync() { + final SettableApiFuture res = SettableApiFuture.create(); + final Span span = tracer.spanBuilder(SpannerImpl.BEGIN_TRANSACTION).startSpan(); + final BeginTransactionRequest request = + BeginTransactionRequest.newBuilder() + .setSession(name) + .setOptions( + TransactionOptions.newBuilder() + .setReadWrite(TransactionOptions.ReadWrite.getDefaultInstance())) + .build(); + final ApiFuture requestFuture = spanner.getRpc().beginTransactionAsync(request, options); + requestFuture.addListener(tracer.withSpan(span, new Runnable(){ + @Override + public void run() { + try { + Transaction txn = requestFuture.get(); + if (txn.getId().isEmpty()) { + throw newSpannerException(ErrorCode.INTERNAL, "Missing id in transaction\n" + getName()); + } + span.end(TraceUtil.END_SPAN_OPTIONS); + res.set(txn.getId()); + } catch (ExecutionException e) { + TraceUtil.endSpanWithFailure(span, e); + res.setException(SpannerExceptionFactory.newSpannerException(e.getCause() == null ? e : e.getCause())); + } catch (InterruptedException e) { + TraceUtil.endSpanWithFailure(span, e); + res.setException(SpannerExceptionFactory.propagateInterrupt(e)); + } catch (Exception e) { + TraceUtil.endSpanWithFailure(span, e); + res.setException(e); + } + } + }), MoreExecutors.directExecutor()); + return res; + } + TransactionContextImpl newTransaction() { return TransactionContextImpl.newBuilder() .setSession(this) @@ -307,9 +346,4 @@ T setActive(@Nullable T ctx) { boolean hasReadyTransaction() { return readyTransactionId != null; } - - @Override - public TransactionManager transactionManager() { - return new TransactionManagerImpl(this, currentSpan); - } } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java index 451063403e2..9b90ff64776 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java @@ -41,10 +41,14 @@ import com.google.cloud.grpc.GrpcTransportOptions; import com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory; import com.google.cloud.spanner.AbstractReadContext.ConsumeSingleRowCallback; +import com.google.cloud.spanner.AbstractReadContext.ListenableAsyncResultSet; +import com.google.cloud.spanner.AsyncResultSet.ReadyCallback; import com.google.cloud.spanner.Options.QueryOption; import com.google.cloud.spanner.Options.ReadOption; import com.google.cloud.spanner.SessionClient.SessionConsumer; import com.google.cloud.spanner.SpannerException.ResourceNotFoundException; +import com.google.cloud.spanner.TransactionManager.TransactionState; +import com.google.cloud.spanner.TransactionRunnerImpl.TransactionContextImpl; import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Function; import com.google.common.base.MoreObjects; @@ -55,6 +59,8 @@ import com.google.common.collect.ImmutableSet; import com.google.common.util.concurrent.ForwardingFuture; import com.google.common.util.concurrent.ForwardingFuture.SimpleForwardingFuture; +import com.google.common.util.concurrent.ForwardingListenableFuture; +import com.google.common.util.concurrent.ForwardingListenableFuture.SimpleForwardingListenableFuture; import com.google.common.util.concurrent.ListenableFuture; import com.google.common.util.concurrent.MoreExecutors; import com.google.common.util.concurrent.SettableFuture; @@ -700,8 +706,11 @@ public ResultSet executeQuery(Statement statement, QueryOption... options) { @Override public AsyncResultSet executeQueryAsync(Statement statement, QueryOption... options) { - throw SpannerExceptionFactory.newSpannerException( - ErrorCode.UNIMPLEMENTED, "not yet implemented"); + try { + return delegate.executeQueryAsync(statement, options); + } catch (SessionNotFoundException e) { + throw handleSessionNotFound(e); + } } @Override @@ -724,7 +733,6 @@ public void close() { AutoClosingTransactionManager(SessionPool sessionPool, PooledSessionFuture session) { this.sessionPool = sessionPool; this.session = session; - // this.delegate = session.delegate.transactionManager(); } @Override @@ -946,6 +954,87 @@ public ApiFuture getCommitTimestamp() { } } + private static class SessionPoolAsyncTransactionManager implements AsyncTransactionManager { + private final SessionPool sessionPool; + private volatile PooledSessionFuture session; + private final SettableApiFuture commitTimestamp = SettableApiFuture.create(); + private final SettableApiFuture delegate = SettableApiFuture.create(); + + private SessionPoolAsyncTransactionManager(SessionPool sessionPool, PooledSessionFuture session) { + this.sessionPool = sessionPool; + this.session = session; + this.session.addListener(new Runnable(){ + @Override + public void run() { + try { + delegate.set(SessionPoolAsyncTransactionManager.this.session.get().transactionManagerAsync()); + } catch (Throwable t) { + delegate.setException(t); + } + } + }, MoreExecutors.directExecutor()); + } + + @Override + public ApiFuture beginAsync() { + final SettableApiFuture res = SettableApiFuture.create(); + delegate.addListener(new Runnable(){ + @Override + public void run() { + try { + res.set(delegate.get().beginAsync().get()); + } catch (Throwable t) { + res.setException(t); + } + } + }, MoreExecutors.directExecutor()); + return res; + } + + @Override + public ApiFuture commitAsync() { + // TODO Auto-generated method stub + return null; + } + + @Override + public ApiFuture rollbackAsync() { + // TODO Auto-generated method stub + return null; + } + + @Override + public ApiFuture resetForRetryAsync() { + // TODO Auto-generated method stub + return null; + } + + @Override + public TransactionState getState() { + // TODO Auto-generated method stub + return null; + } + + @Override + public void close() { + // TODO Auto-generated method stub + + } + + private void setCommitTimestamp(AsyncTransactionManager delegate) { + try { + commitTimestamp.set(delegate.getCommitTimestampAsync().get()); + } catch (Throwable t) { + commitTimestamp.setException(t); + } + } + + @Override + public ApiFuture getCommitTimestampAsync() { + return commitTimestamp; + } + } + // Exception class used just to track the stack trace at the point when a session was handed out // from the pool. final class LeakedSessionException extends RuntimeException { @@ -962,17 +1051,17 @@ private enum SessionState { CLOSING, } - private PooledSessionFuture createPooledSessionFuture(Future future, Span span) { + private PooledSessionFuture createPooledSessionFuture(ListenableFuture future, Span span) { return new PooledSessionFuture(future, span); } - final class PooledSessionFuture extends SimpleForwardingFuture implements Session { + final class PooledSessionFuture extends SimpleForwardingListenableFuture implements Session { private volatile LeakedSessionException leakedException; private volatile AtomicBoolean inUse = new AtomicBoolean(); private volatile CountDownLatch initialized = new CountDownLatch(1); private final Span span; - private PooledSessionFuture(Future delegate, Span span) { + private PooledSessionFuture(ListenableFuture delegate, Span span) { super(delegate); this.span = span; } @@ -1117,6 +1206,11 @@ public AsyncRunner runAsync() { return new SessionPoolAsyncRunner(SessionPool.this, this); } + @Override + public AsyncTransactionManager transactionManagerAsync() { + return new SessionPoolAsyncTransactionManager(SessionPool.this, this); + } + @Override public long executePartitionedUpdate(Statement stmt) { try { @@ -1277,6 +1371,11 @@ public AsyncRunner runAsync() { return delegate.runAsync(); } + @Override + public AsyncTransactionManager transactionManagerAsync() { + return delegate.transactionManagerAsync(); + } + @Override public ApiFuture asyncClose() { close(); @@ -1350,12 +1449,12 @@ public TransactionManager transactionManager() { } } - private final class WaiterFuture extends ForwardingFuture { + private final class WaiterFuture extends ForwardingListenableFuture { private static final long MAX_SESSION_WAIT_TIMEOUT = 240_000L; - private final SettableApiFuture waiter = SettableApiFuture.create(); + private final SettableFuture waiter = SettableFuture.create(); @Override - protected Future delegate() { + protected ListenableFuture delegate() { return waiter; } @@ -2033,7 +2132,7 @@ private PooledSessionFuture checkoutSession( WaiterFuture waiter, boolean write, final boolean inProcessPrepare) { - Future sessionFuture; + ListenableFuture sessionFuture; if (waiter != null) { logger.log( Level.FINE, @@ -2043,10 +2142,12 @@ private PooledSessionFuture checkoutSession( "Waiting for %s session to be available", write ? "read write" : "read only")); sessionFuture = waiter; } else { - sessionFuture = ApiFutures.immediateFuture(readySession); + SettableFuture fut = SettableFuture.create(); + fut.set(readySession); + sessionFuture = fut; } - SimpleForwardingFuture forwardingFuture = - new SimpleForwardingFuture(sessionFuture) { + SimpleForwardingListenableFuture forwardingFuture = + new SimpleForwardingListenableFuture(sessionFuture) { private volatile boolean initialized = false; private final Object prepareLock = new Object(); private volatile PooledSession result; @@ -2426,7 +2527,7 @@ public void run() { } } }, - executor); + MoreExecutors.directExecutor()); return res; } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/Spanner.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/Spanner.java index 0c6bec4ea83..52c35cb7130 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/Spanner.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/Spanner.java @@ -16,6 +16,7 @@ package com.google.cloud.spanner; +import com.google.api.gax.core.ExecutorProvider; import com.google.cloud.Service; /** @@ -108,4 +109,7 @@ public interface Spanner extends Service, AutoCloseable { /** @return true if this {@link Spanner} object is closed. */ boolean isClosed(); + + /** @return the {@link ExecutorProvider} that is used for asynchronous queries and operations. */ + ExecutorProvider getAsyncExecutorProvider(); } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java index 162f7495259..753adc1e862 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java @@ -137,7 +137,7 @@ QueryOptions getDefaultQueryOptions(DatabaseId databaseId) { /** * Returns the {@link ExecutorProvider} to use for async methods that need a background executor. */ - ExecutorProvider getAsyncExecutorProvider() { + public ExecutorProvider getAsyncExecutorProvider() { return asyncExecutorProvider; } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TraceUtil.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TraceUtil.java index c5488ac55d3..0d429661ad2 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TraceUtil.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TraceUtil.java @@ -40,7 +40,7 @@ static Map getTransactionAnnotations(Transaction t) { AttributeValue.stringAttributeValue(Timestamp.fromProto(t.getReadTimestamp()).toString())); } - static ImmutableMap getExceptionAnnotations(RuntimeException e) { + static ImmutableMap getExceptionAnnotations(Throwable e) { if (e instanceof SpannerException) { return ImmutableMap.of( "Status", diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionManagerImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionManagerImpl.java index 35184cdf9c9..8dbab883140 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionManagerImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionManagerImpl.java @@ -39,6 +39,10 @@ final class TransactionManagerImpl implements TransactionManager, SessionTransac this.span = span; } + Span getSpan() { + return span; + } + @Override public void setSpan(Span span) { this.span = span; 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 c3e13f2a87f..a91c6a90ef1 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 @@ -24,6 +24,7 @@ import com.google.api.core.ApiFunction; import com.google.api.core.ApiFuture; import com.google.api.core.ApiFutures; +import com.google.api.core.SettableApiFuture; import com.google.cloud.Timestamp; import com.google.cloud.spanner.Options.QueryOption; import com.google.cloud.spanner.Options.ReadOption; @@ -104,7 +105,7 @@ public ApiFuture setCallback(Executor exec, ReadyCallback cb) { new Runnable() { @Override public void run() { - finishedAsyncOperations.countDown(); + decreaseAsyncOperations(); } }; try { @@ -113,7 +114,7 @@ public void run() { return super.setCallback(exec, cb); } catch (Throwable t) { removeListener(listener); - finishedAsyncOperations.countDown(); + decreaseAsyncOperations(); throw t; } } @@ -133,7 +134,12 @@ public void removeListener(Runnable listener) { private volatile boolean committing; @GuardedBy("lock") - private volatile CountDownLatch finishedAsyncOperations = new CountDownLatch(0); + private volatile SettableApiFuture finishedAsyncOperations = SettableApiFuture.create(); + @GuardedBy("lock") + private volatile int runningAsyncOperations; + +// @GuardedBy("lock") +// private volatile CountDownLatch finishedAsyncOperations = new CountDownLatch(0); @GuardedBy("lock") private List mutations = new ArrayList<>(); @@ -151,15 +157,28 @@ public void removeListener(Runnable listener) { private TransactionContextImpl(Builder builder) { super(builder); this.transactionId = builder.transactionId; + this.finishedAsyncOperations.set(null); } private void increaseAsynOperations() { synchronized (lock) { - finishedAsyncOperations = new CountDownLatch((int) finishedAsyncOperations.getCount() + 1); + if (runningAsyncOperations == 0) { + finishedAsyncOperations = SettableApiFuture.create(); + } + runningAsyncOperations++; } } - void ensureTxn() { + private void decreaseAsyncOperations() { + synchronized (lock) { + runningAsyncOperations--; + if (runningAsyncOperations == 0) { + finishedAsyncOperations.set(null); + } + } + } + + void ensureTxn_old() { if (transactionId == null || isAborted()) { span.addAnnotation("Creating Transaction"); try { @@ -188,15 +207,68 @@ void ensureTxn() { } } - void commit() { - CountDownLatch latch; + void ensureTxn() { + try { + ensureTxnAsync().get(); + } catch (ExecutionException e) { + throw SpannerExceptionFactory.newSpannerException(e.getCause() == null ? e : e.getCause()); + } catch (InterruptedException e) { + throw SpannerExceptionFactory.propagateInterrupt(e); + } + } + + ApiFuture ensureTxnAsync() { + final SettableApiFuture res = SettableApiFuture.create(); + if (transactionId == null || isAborted()) { + span.addAnnotation("Creating Transaction"); + final ApiFuture fut = session.beginTransactionAsync(); + fut.addListener(new Runnable(){ + @Override + public void run() { + try { + transactionId = fut.get(); + span.addAnnotation( + "Transaction Creation Done", + ImmutableMap.of( + "Id", AttributeValue.stringAttributeValue(transactionId.toStringUtf8()))); + txnLogger.log( + Level.FINER, + "Started transaction {0}", + txnLogger.isLoggable(Level.FINER) ? transactionId.asReadOnlyByteBuffer() : null); + res.set(null); + } catch (ExecutionException e) { + span.addAnnotation("Transaction Creation Failed", TraceUtil.getExceptionAnnotations(e.getCause() == null ? e : e.getCause())); + res.setException(e.getCause() == null ? e : e.getCause()); + } catch (InterruptedException e) { + res.setException(SpannerExceptionFactory.propagateInterrupt(e)); + } + } + }, MoreExecutors.directExecutor()); + } else { + span.addAnnotation( + "Transaction Initialized", + ImmutableMap.of( + "Id", AttributeValue.stringAttributeValue(transactionId.toStringUtf8()))); + txnLogger.log( + Level.FINER, + "Using prepared transaction {0}", + txnLogger.isLoggable(Level.FINER) ? transactionId.asReadOnlyByteBuffer() : null); + res.set(null); + } + return res; + } + + void commit_old() { + SettableApiFuture latch; synchronized (lock) { latch = finishedAsyncOperations; } try { - latch.await(); + latch.get(); } catch (InterruptedException e) { throw SpannerExceptionFactory.propagateInterrupt(e); + } catch (ExecutionException e) { + throw SpannerExceptionFactory.newSpannerException(e.getCause() == null ? e : e.getCause()); } span.addAnnotation("Starting Commit"); CommitRequest.Builder builder = @@ -231,6 +303,80 @@ void commit() { span.addAnnotation("Commit Done"); } + void commit() { + try { + commitAsync().get(); + } catch (InterruptedException e) { + throw SpannerExceptionFactory.propagateInterrupt(e); + } catch (ExecutionException e) { + throw SpannerExceptionFactory.newSpannerException(e.getCause() == null ? e : e.getCause()); + } + } + + ApiFuture commitAsync() { + final SettableApiFuture res = SettableApiFuture.create(); + CommitRequest.Builder builder = + CommitRequest.newBuilder().setSession(session.getName()).setTransactionId(transactionId); + synchronized (lock) { + if (!mutations.isEmpty()) { + List mutationsProto = new ArrayList<>(); + Mutation.toProto(mutations, mutationsProto); + builder.addAllMutations(mutationsProto); + } + // Ensure that no call to buffer mutations that would be lost can succeed. + mutations = null; + } + final CommitRequest commitRequest = builder.build(); + final SettableApiFuture latch; + synchronized (lock) { + latch = finishedAsyncOperations; + } + latch.addListener(new Runnable(){ + @Override + public void run() { + try { + latch.get(); + span.addAnnotation("Starting Commit"); + final Span opSpan = tracer.spanBuilderWithExplicitParent(SpannerImpl.COMMIT, span).startSpan(); + final ApiFuture commitFuture = rpc.commitAsync(commitRequest, session.getOptions()); + commitFuture.addListener(tracer.withSpan(opSpan, new Runnable(){ + @Override + public void run() { + try { + CommitResponse commitResponse = commitFuture.get(); + if (!commitResponse.hasCommitTimestamp()) { + throw newSpannerException( + ErrorCode.INTERNAL, "Missing commitTimestamp:\n" + session.getName()); + } + commitTimestamp = Timestamp.fromProto(commitResponse.getCommitTimestamp()); + span.addAnnotation("Commit Done"); + opSpan.end(TraceUtil.END_SPAN_OPTIONS); + res.set(null); + } catch (Throwable e) { + if (e instanceof ExecutionException) { + e = SpannerExceptionFactory.newSpannerException(e.getCause() == null ? e : e.getCause()); + } else if (e instanceof InterruptedException) { + e = SpannerExceptionFactory.propagateInterrupt((InterruptedException) e); + } else { + e = SpannerExceptionFactory.newSpannerException(e); + } + span.addAnnotation("Commit Failed", TraceUtil.getExceptionAnnotations(e)); + TraceUtil.endSpanWithFailure(opSpan, e); + onError((SpannerException) e); + res.setException(e); + } + } + }), MoreExecutors.directExecutor()); + } catch (InterruptedException e) { + res.setException(SpannerExceptionFactory.propagateInterrupt(e)); + } catch (ExecutionException e) { + res.setException(SpannerExceptionFactory.newSpannerException(e.getCause() == null ? e : e.getCause())); + } + } + }, MoreExecutors.directExecutor()); + return res; + } + Timestamp commitTimestamp() { checkState(commitTimestamp != null, "run() has not yet returned normally"); return commitTimestamp; @@ -338,7 +484,7 @@ public ApiFuture executeUpdateAsync(Statement statement) { increaseAsynOperations(); resultSet = rpc.executeQueryAsync(builder.build(), session.getOptions()); } catch (Throwable t) { - finishedAsyncOperations.countDown(); + decreaseAsyncOperations(); throw t; } final ApiFuture updateCount = @@ -368,7 +514,7 @@ public void run() { } catch (InterruptedException e) { onError(SpannerExceptionFactory.propagateInterrupt(e)); } finally { - finishedAsyncOperations.countDown(); + decreaseAsyncOperations(); } } }, @@ -417,7 +563,7 @@ public ApiFuture batchUpdateAsync(Iterable statements) { increaseAsynOperations(); response = rpc.executeBatchDmlAsync(builder.build(), session.getOptions()); } catch (Throwable t) { - finishedAsyncOperations.countDown(); + decreaseAsyncOperations(); throw t; } final ApiFuture updateCounts = @@ -456,7 +602,7 @@ public void run() { } catch (InterruptedException e) { onError(SpannerExceptionFactory.propagateInterrupt(e)); } finally { - finishedAsyncOperations.countDown(); + decreaseAsyncOperations(); } } }, 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 cb8cf3bc557..9f278fb11db 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.cloud.spanner.AsyncResultSet; import com.google.cloud.spanner.ErrorCode; import com.google.cloud.spanner.Options.QueryOption; import com.google.cloud.spanner.ReadContext; @@ -73,6 +74,16 @@ public ResultSet call() throws Exception { }); } + @Override + public AsyncResultSet executeQueryAsync( + final ParsedStatement statement, + final AnalyzeMode analyzeMode, + final QueryOption... options) { + Preconditions.checkArgument(statement.isQuery(), "Statement is not a query"); + checkValidTransaction(); + return getReadContext().executeQueryAsync(statement.getStatement(), options); + } + ResultSet internalExecuteQuery( final ParsedStatement statement, AnalyzeMode analyzeMode, QueryOption... options) { if (analyzeMode == AnalyzeMode.NONE) { diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AsyncChecksumResultSet.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AsyncChecksumResultSet.java new file mode 100644 index 00000000000..95c1077fa60 --- /dev/null +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AsyncChecksumResultSet.java @@ -0,0 +1,73 @@ +/* + * 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.cloud.spanner.AsyncResultSet; +import com.google.cloud.spanner.Options.QueryOption; +import com.google.cloud.spanner.SpannerException; +import com.google.cloud.spanner.StructReader; +import com.google.cloud.spanner.connection.StatementParser.ParsedStatement; +import com.google.common.base.Function; +import com.google.common.collect.ImmutableList; +import java.util.concurrent.Executor; + +class AsyncChecksumResultSet extends ChecksumResultSet implements AsyncResultSet { + private AsyncResultSet delegate; + + AsyncChecksumResultSet( + ReadWriteTransaction transaction, + AsyncResultSet delegate, + ParsedStatement statement, + AnalyzeMode analyzeMode, + QueryOption... options) { + super(transaction, delegate, statement, analyzeMode, options); + this.delegate = delegate; + } + + @Override + public CursorState tryNext() throws SpannerException { + return delegate.tryNext(); + } + + @Override + public ApiFuture setCallback(Executor exec, ReadyCallback cb) { + return delegate.setCallback(exec, cb); + } + + @Override + public void cancel() { + delegate.cancel(); + } + + @Override + public void resume() { + delegate.resume(); + } + + @Override + public ApiFuture> toListAsync( + Function transformer, Executor executor) { + return delegate.toListAsync(transformer, executor); + } + + @Override + public ImmutableList toList(Function transformer) + throws SpannerException { + return delegate.toList(transformer); + } +} 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 5247ce2c130..9a1dc69a0cd 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 @@ -20,6 +20,7 @@ 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; @@ -619,6 +620,8 @@ public interface Connection extends AutoCloseable { */ ResultSet executeQuery(Statement query, QueryOption... options); + AsyncResultSet executeQueryAsync(Statement query, QueryOption... options); + /** * Analyzes a query and returns query plan and/or query execution statistics information. * 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 ce24791859e..1e37f3927e9 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 @@ -17,12 +17,14 @@ package com.google.cloud.spanner.connection; 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; @@ -681,6 +683,11 @@ 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 +724,38 @@ private ResultSet parseAndExecuteQuery( "Statement is not a query: " + parsedStatement.getSqlWithoutComments()); } + /** + * Parses the given statement as a query and executes it asynchronously. Throws a {@link + * SpannerException} if the statement is not a query. + */ + private AsyncResultSet parseAndExecuteQueryAsync( + Statement query, AnalyzeMode analyzeMode, QueryOption... options) { + Preconditions.checkNotNull(query); + Preconditions.checkNotNull(analyzeMode); + 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()); + 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); @@ -787,6 +826,16 @@ private ResultSet internalExecuteQuery( } } + 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 transaction.executeQueryAsync(statement, analyzeMode, options); + } + private long internalExecuteUpdate(final ParsedStatement update) { Preconditions.checkArgument( update.getType() == StatementType.UPDATE, "Statement must be an update"); 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 b18f3fa891c..b49f443227b 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 @@ -18,6 +18,7 @@ import com.google.api.gax.longrunning.OperationFuture; 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; @@ -114,6 +115,27 @@ public boolean isReadOnly() { @Override public ResultSet executeQuery( final ParsedStatement statement, AnalyzeMode analyzeMode, QueryOption... options) { + final QueryOption[] internalOptions = verifyQueryForDdlBatch(statement, analyzeMode, options); + Callable callable = + new Callable() { + @Override + public ResultSet call() throws Exception { + return DirectExecuteResultSet.ofResultSet( + dbClient.singleUse().executeQuery(statement.getStatement(), internalOptions)); + } + }; + return asyncExecuteStatement(statement, callable); + } + + @Override + public AsyncResultSet executeQueryAsync( + final ParsedStatement statement, AnalyzeMode analyzeMode, QueryOption... options) { + final QueryOption[] internalOptions = verifyQueryForDdlBatch(statement, analyzeMode, options); + return dbClient.singleUse().executeQueryAsync(statement.getStatement(), internalOptions); + } + + private QueryOption[] verifyQueryForDdlBatch( + ParsedStatement statement, AnalyzeMode analyzeMode, QueryOption... options) { if (options != null) { for (int i = 0; i < options.length; i++) { if (options[i] instanceof InternalMetadataQuery) { @@ -124,16 +146,7 @@ public ResultSet executeQuery( // Queries marked with internal metadata queries are allowed during a DDL batch. // These can only be generated by library internal methods and may be used to check // whether a database object such as table or an index exists. - final QueryOption[] internalOptions = ArrayUtils.remove(options, i); - Callable callable = - new Callable() { - @Override - public ResultSet call() throws Exception { - return DirectExecuteResultSet.ofResultSet( - dbClient.singleUse().executeQuery(statement.getStatement(), internalOptions)); - } - }; - return asyncExecuteStatement(statement, callable); + return ArrayUtils.remove(options, i); } } } 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 ff38338d623..250e7a1cc72 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 @@ -17,6 +17,7 @@ package com.google.cloud.spanner.connection; import com.google.cloud.Timestamp; +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; @@ -93,6 +94,13 @@ public ResultSet executeQuery( ErrorCode.FAILED_PRECONDITION, "Executing queries is not allowed for DML batches."); } + @Override + public AsyncResultSet executeQueryAsync( + ParsedStatement statement, AnalyzeMode analyzeMode, QueryOption... options) { + throw SpannerExceptionFactory.newSpannerException( + ErrorCode.FAILED_PRECONDITION, "Executing queries is not allowed for DML batches."); + } + @Override public Timestamp getReadTimestamp() { throw SpannerExceptionFactory.newSpannerException( 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 7a0155cbfb8..3689d4b8d95 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 @@ -21,6 +21,7 @@ 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.DatabaseClient; import com.google.cloud.spanner.ErrorCode; import com.google.cloud.spanner.Mutation; @@ -279,6 +280,21 @@ public ResultSet call() throws Exception { } } + @Override + public AsyncResultSet executeQueryAsync( + final ParsedStatement statement, + final AnalyzeMode analyzeMode, + final QueryOption... options) { + Preconditions.checkArgument(statement.isQuery(), "Statement is not a query"); + checkValidTransaction(); + if (retryAbortsInternally) { + AsyncResultSet delegate = super.executeQueryAsync(statement, analyzeMode, options); + return createAndAddAsyncRetryResultSet(delegate, statement, analyzeMode, options); + } else { + return super.executeQueryAsync(statement, analyzeMode, options); + } + } + @Override public long executeUpdate(final ParsedStatement update) { Preconditions.checkNotNull(update); @@ -542,6 +558,24 @@ private ResultSet createAndAddRetryResultSet( return resultSet; } + /** + * Registers a {@link AsyncResultSet} on this transaction that must be checked during a retry, and + * returns a retryable {@link AsyncResultSet}. + */ + private AsyncResultSet createAndAddAsyncRetryResultSet( + AsyncResultSet resultSet, + ParsedStatement statement, + AnalyzeMode analyzeMode, + QueryOption... options) { + if (retryAbortsInternally) { + AsyncChecksumResultSet checksumResultSet = + createAsyncChecksumResultSet(resultSet, statement, analyzeMode, options); + addRetryStatement(checksumResultSet); + return checksumResultSet; + } + return resultSet; + } + /** Registers the statement as a query that should return an error during a retry. */ private void createAndAddFailedQuery( SpannerException e, @@ -759,4 +793,14 @@ ChecksumResultSet createChecksumResultSet( QueryOption... options) { return new ChecksumResultSet(this, delegate, statement, analyzeMode, options); } + + /** Creates a {@link AsyncChecksumResultSet} for this {@link ReadWriteTransaction}. */ + @VisibleForTesting + AsyncChecksumResultSet createAsyncChecksumResultSet( + AsyncResultSet delegate, + ParsedStatement statement, + AnalyzeMode analyzeMode, + QueryOption... options) { + return new AsyncChecksumResultSet(this, delegate, statement, analyzeMode, options); + } } 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 614d0c61e52..edf4a9a5289 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 @@ -19,6 +19,7 @@ import com.google.api.gax.longrunning.OperationFuture; import com.google.cloud.Timestamp; import com.google.cloud.spanner.AbortedException; +import com.google.cloud.spanner.AsyncResultSet; import com.google.cloud.spanner.DatabaseClient; import com.google.cloud.spanner.ErrorCode; import com.google.cloud.spanner.Mutation; @@ -66,7 +67,7 @@ class SingleUseTransaction extends AbstractBaseUnitOfWork { private final DatabaseClient dbClient; private final TimestampBound readOnlyStaleness; private final AutocommitDmlMode autocommitDmlMode; - private Timestamp readTimestamp = null; + private ReadOnlyTransaction readOnlyTransaction; private volatile TransactionManager txManager; private TransactionRunner writeTransaction; private boolean used = false; @@ -168,52 +169,78 @@ public ResultSet executeQuery( Preconditions.checkArgument(statement.isQuery(), "Statement is not a query"); checkAndMarkUsed(); - final ReadOnlyTransaction currentTransaction = - dbClient.singleUseReadOnlyTransaction(readOnlyStaleness); + readOnlyTransaction = dbClient.singleUseReadOnlyTransaction(readOnlyStaleness); Callable callable = new Callable() { @Override public ResultSet call() throws Exception { - try { +// try { ResultSet rs; if (analyzeMode == AnalyzeMode.NONE) { - rs = currentTransaction.executeQuery(statement.getStatement(), options); + rs = readOnlyTransaction.executeQuery(statement.getStatement(), options); } else { rs = - currentTransaction.analyzeQuery( + readOnlyTransaction.analyzeQuery( statement.getStatement(), analyzeMode.getQueryAnalyzeMode()); } // 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 { - currentTransaction.close(); - } +// } catch (Exception e) { +// readOnlyTransaction.close(); +// throw e; +// } finally { +// readOnlyTransaction.close(); +// currentTransaction.close(); +// } } }; try { ResultSet res = asyncExecuteStatement(statement, callable); - readTimestamp = currentTransaction.getReadTimestamp(); + // readTimestamp = currentTransaction.getReadTimestamp(); state = UnitOfWorkState.COMMITTED; return res; } catch (Throwable e) { state = UnitOfWorkState.COMMIT_FAILED; throw e; } finally { - currentTransaction.close(); + readOnlyTransaction.close(); + } + } + + @Override + public AsyncResultSet executeQueryAsync( + final ParsedStatement statement, + final AnalyzeMode analyzeMode, + final QueryOption... options) { + Preconditions.checkNotNull(statement); + Preconditions.checkArgument(statement.isQuery(), "Statement is not a query"); + checkAndMarkUsed(); + + readOnlyTransaction = dbClient.singleUseReadOnlyTransaction(readOnlyStaleness); + try { + AsyncResultSet res = readOnlyTransaction.executeQueryAsync(statement.getStatement(), options); + state = UnitOfWorkState.COMMITTED; + return res; + } catch (Throwable e) { + readOnlyTransaction.close(); + state = UnitOfWorkState.COMMIT_FAILED; + throw e; + // } finally { + // currentTransaction.close(); } } @Override public Timestamp getReadTimestamp() { ConnectionPreconditions.checkState( - readTimestamp != null, "There is no read timestamp available for this transaction."); - return readTimestamp; + readOnlyTransaction != null && state == UnitOfWorkState.COMMITTED, "There is no read timestamp available for this transaction."); + return readOnlyTransaction.getReadTimestamp(); } @Override public Timestamp getReadTimestampOrNull() { - return readTimestamp; + return readOnlyTransaction == null || state != UnitOfWorkState.COMMITTED ? null : readOnlyTransaction.getReadTimestamp(); } private boolean hasCommitTimestamp() { 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 6221cc447b6..ab5610d0723 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 @@ -26,7 +26,6 @@ /** Implementation of {@link StatementResult} */ class StatementResultImpl implements StatementResult { - /** {@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 e372229c64c..49001cd8d8f 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 @@ -18,6 +18,7 @@ import com.google.api.core.InternalApi; import com.google.cloud.Timestamp; +import com.google.cloud.spanner.AsyncResultSet; import com.google.cloud.spanner.Mutation; import com.google.cloud.spanner.Options.QueryOption; import com.google.cloud.spanner.ReadContext; @@ -114,6 +115,9 @@ public boolean isActive() { ResultSet executeQuery( ParsedStatement statement, AnalyzeMode analyzeMode, QueryOption... options); + AsyncResultSet executeQueryAsync( + ParsedStatement statement, AnalyzeMode analyzeMode, QueryOption... options); + /** * @return the read timestamp of this transaction. Will throw a {@link SpannerException} if there * is no read timestamp. diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java index fc027576c11..57ad55480c4 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java @@ -1088,18 +1088,28 @@ public Transaction beginTransaction( return get(beginTransactionAsync(request, options)); } + @Override + public ApiFuture commitAsync(CommitRequest commitRequest, @Nullable Map options) { + GrpcCallContext context = newCallContext(options, commitRequest.getSession()); + return spannerStub.commitCallable().futureCall(commitRequest, context); + } + @Override public CommitResponse commit(CommitRequest commitRequest, @Nullable Map options) throws SpannerException { - GrpcCallContext context = newCallContext(options, commitRequest.getSession()); - return get(spannerStub.commitCallable().futureCall(commitRequest, context)); + return get(commitAsync(commitRequest, options)); + } + + @Override + public ApiFuture rollbackAsync(RollbackRequest request, @Nullable Map options) { + GrpcCallContext context = newCallContext(options, request.getSession()); + return spannerStub.rollbackCallable().futureCall(request, context); } @Override public void rollback(RollbackRequest request, @Nullable Map options) throws SpannerException { - GrpcCallContext context = newCallContext(options, request.getSession()); - get(spannerStub.rollbackCallable().futureCall(request, context)); + get(rollbackAsync(request, options)); } @Override diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerRpc.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerRpc.java index 0334ca4e403..31347ae547d 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerRpc.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerRpc.java @@ -305,8 +305,12 @@ ApiFuture beginTransactionAsync( CommitResponse commit(CommitRequest commitRequest, @Nullable Map options) throws SpannerException; + ApiFuture commitAsync(CommitRequest commitRequest, @Nullable Map options); + void rollback(RollbackRequest request, @Nullable Map options) throws SpannerException; + ApiFuture rollbackAsync(RollbackRequest request, @Nullable Map options); + PartitionResponse partitionQuery(PartitionQueryRequest request, @Nullable Map options) throws SpannerException; diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncRunnerTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncRunnerTest.java index cbe78fe0d01..a020997fd7d 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncRunnerTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncRunnerTest.java @@ -32,6 +32,7 @@ import com.google.cloud.spanner.AsyncRunner.AsyncWork; import com.google.cloud.spanner.MockSpannerServiceImpl.SimulatedExecutionTime; import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult; +import com.google.cloud.spanner.TransactionRunnerImpl.TransactionContextImpl; import com.google.common.base.Function; import com.google.common.collect.ImmutableList; import com.google.common.util.concurrent.MoreExecutors; @@ -266,7 +267,7 @@ public void asyncRunnerCommitAborted() throws Exception { runner.runAsync( new AsyncWork() { @Override - public ApiFuture doWorkAsync(TransactionContext txn) { + public ApiFuture doWorkAsync(final TransactionContext txn) { if (attempt.get() > 0) { // Set the result of the update statement back to 1 row. mockSpanner.putStatementResult( @@ -484,7 +485,7 @@ public void asyncRunnerWithBatchUpdateCommitAborted() throws Exception { runner.runAsync( new AsyncWork() { @Override - public ApiFuture doWorkAsync(TransactionContext txn) { + public ApiFuture doWorkAsync(final TransactionContext txn) { if (attempt.get() > 0) { // Set the result of the update statement back to 1 row. mockSpanner.putStatementResult( @@ -606,11 +607,12 @@ public void closeTransactionBeforeEndOfAsyncQuery() throws Exception { AsyncRunner runner = clientImpl.runAsync(); final CountDownLatch dataReceived = new CountDownLatch(1); + final CountDownLatch dataChecked = new CountDownLatch(1); ApiFuture res = runner.runAsync( new AsyncWork() { @Override - public ApiFuture doWorkAsync(TransactionContext txn) { + public ApiFuture doWorkAsync(TransactionContext txn) { try (AsyncResultSet rs = txn.readAsync( READ_TABLE_NAME, KeySet.all(), READ_COLUMN_NAMES, Options.bufferRows(1))) { @@ -619,6 +621,7 @@ public ApiFuture doWorkAsync(TransactionContext txn) { new ReadyCallback() { @Override public CallbackResponse cursorReady(AsyncResultSet resultSet) { + dataReceived.countDown(); try { while (true) { switch (resultSet.tryNext()) { @@ -628,19 +631,23 @@ public CallbackResponse cursorReady(AsyncResultSet resultSet) { case NOT_READY: return CallbackResponse.CONTINUE; case OK: - dataReceived.countDown(); + dataChecked.await(); results.put(resultSet.getString(0)); } } } catch (Throwable t) { finished.setException(t); - dataReceived.countDown(); return CallbackResponse.DONE; } } }); } - return ApiFutures.immediateFuture(null); + try { + dataReceived.await(); + return ApiFutures.immediateFuture(null); + } catch (InterruptedException e) { + return ApiFutures.immediateFailedFuture(SpannerExceptionFactory.propagateInterrupt(e)); + } } }, executor); @@ -649,6 +656,7 @@ public CallbackResponse cursorReady(AsyncResultSet resultSet) { dataReceived.await(); assertThat(clientImpl.pool.getNumberOfSessionsInUse()).isEqualTo(1); assertThat(res.isDone()).isFalse(); + dataChecked.countDown(); // Get the data from the transaction. List resultList = new ArrayList<>(); do { diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java index 93f91f3004d..5100b2723ee 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java @@ -54,6 +54,7 @@ import java.util.concurrent.Executors; import java.util.concurrent.ScheduledThreadPoolExecutor; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; import org.junit.After; import org.junit.AfterClass; import org.junit.Before; @@ -83,8 +84,7 @@ public class DatabaseClientImplTest { private static final long UPDATE_COUNT = 1L; private Spanner spanner; private Spanner spannerWithEmptySessionPool; - - // @Rule public Timeout globalTimeout = new Timeout(5L, TimeUnit.SECONDS); + private static final ExecutorService executor = Executors.newSingleThreadExecutor(); @BeforeClass public static void startStaticServer() throws IOException { @@ -113,6 +113,7 @@ public static void startStaticServer() throws IOException { public static void stopServer() throws InterruptedException { server.shutdown(); server.awaitTermination(); + executor.shutdown(); } @Before @@ -190,6 +191,37 @@ public void singleUseIsNonBlocking() { } } + @Test + public void singleUseAsync() throws Exception { + DatabaseClient client = + spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); + final AtomicInteger rowCount = new AtomicInteger(); + ApiFuture res; + try (AsyncResultSet rs = client.singleUse().executeQueryAsync(SELECT1)) { + 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; + } + } + } + }); + } + res.get(); + assertThat(rowCount.get()).isEqualTo(1); + } + @Test public void singleUseBound() { DatabaseClient client = @@ -221,6 +253,40 @@ public void singleUseBoundIsNonBlocking() { } } + @Test + public void singleUseBoundAsync() throws Exception { + DatabaseClient client = + spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); + final AtomicInteger rowCount = new AtomicInteger(); + ApiFuture res; + try (AsyncResultSet rs = + client + .singleUse(TimestampBound.ofExactStaleness(15L, TimeUnit.SECONDS)) + .executeQueryAsync(SELECT1)) { + 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; + } + } + } + }); + } + res.get(); + assertThat(rowCount.get()).isEqualTo(1); + } + @Test public void singleUseTransaction() { DatabaseClient client = @@ -480,6 +546,44 @@ public void transactionManagerIsNonBlocking() throws Exception { } } + @Test + public void transactionManagerExecuteQueryAsync() throws Exception { + DatabaseClient client = + spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); + final AtomicInteger rowCount = new AtomicInteger(); + try (TransactionManager txManager = client.transactionManager()) { + while (true) { + TransactionContext tx = txManager.begin(); + try { + try(AsyncResultSet rs = tx.executeQueryAsync(SELECT1)) { + 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; + } + } + } + }); + } + txManager.commit(); + break; + } catch (AbortedException e) { + Thread.sleep(e.getRetryDelayInMillis() / 1000); + tx = txManager.resetForRetry(); + } + } + } + assertThat(rowCount.get()).isEqualTo(1); + } + /** * Test that the update statement can be executed as a partitioned transaction that returns a * lower bound update count. diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionImplTest.java index cc21774daea..cbcb10d25bc 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionImplTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionImplTest.java @@ -20,7 +20,7 @@ import static org.junit.Assert.fail; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; - +import com.google.api.core.ApiFutures; import com.google.api.core.NanoClock; import com.google.api.gax.retrying.RetrySettings; import com.google.cloud.Timestamp; @@ -102,15 +102,15 @@ public void setUp() { .thenReturn(sessionProto); Transaction txn = Transaction.newBuilder().setId(ByteString.copyFromUtf8("TEST")).build(); Mockito.when( - rpc.beginTransaction( + rpc.beginTransactionAsync( Mockito.any(BeginTransactionRequest.class), Mockito.any(Map.class))) - .thenReturn(txn); + .thenReturn(ApiFutures.immediateFuture(txn)); CommitResponse commitResponse = CommitResponse.newBuilder() .setCommitTimestamp(com.google.protobuf.Timestamp.getDefaultInstance()) .build(); - Mockito.when(rpc.commit(Mockito.any(CommitRequest.class), Mockito.any(Map.class))) - .thenReturn(commitResponse); + Mockito.when(rpc.commitAsync(Mockito.any(CommitRequest.class), Mockito.any(Map.class))) + .thenReturn(ApiFutures.immediateFuture(commitResponse)); session = spanner.getSessionClient(db).createSession(); ((SessionImpl) session).setCurrentSpan(mock(Span.class)); // We expect the same options, "options", on all calls on "session". diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java index c57545bad1b..538075fe4f9 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java @@ -53,6 +53,7 @@ import com.google.protobuf.ByteString; import com.google.protobuf.Empty; import com.google.spanner.v1.CommitRequest; +import com.google.spanner.v1.CommitResponse; import com.google.spanner.v1.ExecuteBatchDmlRequest; import com.google.spanner.v1.ExecuteSqlRequest; import com.google.spanner.v1.ResultSetStats; @@ -68,6 +69,7 @@ import java.util.Map; 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; @@ -1296,7 +1298,7 @@ public void testSessionNotFoundReadWriteTransaction() { .thenThrow(sessionNotFound); when(rpc.executeBatchDml(any(ExecuteBatchDmlRequest.class), any(Map.class))) .thenThrow(sessionNotFound); - when(rpc.commit(any(CommitRequest.class), any(Map.class))).thenThrow(sessionNotFound); + when(rpc.commitAsync(any(CommitRequest.class), any(Map.class))).thenReturn(ApiFutures.immediateFailedFuture(sessionNotFound)); doThrow(sessionNotFound).when(rpc).rollback(any(RollbackRequest.class), any(Map.class)); final SessionImpl closedSession = mock(SessionImpl.class); when(closedSession.getName()) @@ -1312,7 +1314,7 @@ public void testSessionNotFoundReadWriteTransaction() { when(closedSession.asyncClose()) .thenReturn(ApiFutures.immediateFuture(Empty.getDefaultInstance())); when(closedSession.newTransaction()).thenReturn(closedTransactionContext); - when(closedSession.beginTransaction()).thenThrow(sessionNotFound); + when(closedSession.beginTransactionAsync()).thenThrow(sessionNotFound); TransactionRunnerImpl closedTransactionRunner = new TransactionRunnerImpl(closedSession, rpc, 10); closedTransactionRunner.setSpan(mock(Span.class)); @@ -1325,7 +1327,7 @@ public void testSessionNotFoundReadWriteTransaction() { .thenReturn("projects/dummy/instances/dummy/database/dummy/sessions/session-open"); final TransactionContextImpl openTransactionContext = mock(TransactionContextImpl.class); when(openSession.newTransaction()).thenReturn(openTransactionContext); - when(openSession.beginTransaction()).thenReturn(ByteString.copyFromUtf8("open-txn")); + when(openSession.beginTransactionAsync()).thenReturn(ApiFutures.immediateFuture(ByteString.copyFromUtf8("open-txn"))); TransactionRunnerImpl openTransactionRunner = new TransactionRunnerImpl(openSession, mock(SpannerRpc.class), 10); openTransactionRunner.setSpan(mock(Span.class)); diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionManagerImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionManagerImplTest.java index b270c227f8d..ead4a6bb159 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionManagerImplTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionManagerImplTest.java @@ -24,7 +24,7 @@ import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import static org.mockito.MockitoAnnotations.initMocks; - +import com.google.api.core.ApiFuture; import com.google.api.core.ApiFutures; import com.google.cloud.Timestamp; import com.google.cloud.grpc.GrpcTransportOptions; @@ -219,26 +219,26 @@ public List answer(InvocationOnMock invocation) .build()); } }); - when(rpc.beginTransaction(Mockito.any(BeginTransactionRequest.class), Mockito.anyMap())) + when(rpc.beginTransactionAsync(Mockito.any(BeginTransactionRequest.class), Mockito.anyMap())) .thenAnswer( - new Answer() { + new Answer>() { @Override - public Transaction answer(InvocationOnMock invocation) throws Throwable { - return Transaction.newBuilder() + public ApiFuture answer(InvocationOnMock invocation) throws Throwable { + return ApiFutures.immediateFuture(Transaction.newBuilder() .setId(ByteString.copyFromUtf8(UUID.randomUUID().toString())) - .build(); + .build()); } }); - when(rpc.commit(Mockito.any(CommitRequest.class), Mockito.anyMap())) + when(rpc.commitAsync(Mockito.any(CommitRequest.class), Mockito.anyMap())) .thenAnswer( - new Answer() { + new Answer>() { @Override - public CommitResponse answer(InvocationOnMock invocation) throws Throwable { - return CommitResponse.newBuilder() + public ApiFuture answer(InvocationOnMock invocation) throws Throwable { + return ApiFutures.immediateFuture(CommitResponse.newBuilder() .setCommitTimestamp( com.google.protobuf.Timestamp.newBuilder() .setSeconds(System.currentTimeMillis() * 1000)) - .build(); + .build()); } }); DatabaseId db = DatabaseId.of("test", "test", "test"); @@ -249,7 +249,7 @@ public CommitResponse answer(InvocationOnMock invocation) throws Throwable { mgr.commit(); } verify(rpc, times(1)) - .beginTransaction(Mockito.any(BeginTransactionRequest.class), Mockito.anyMap()); + .beginTransactionAsync(Mockito.any(BeginTransactionRequest.class), Mockito.anyMap()); } } } diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java index 10924933317..74d8b1c9054 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java @@ -23,7 +23,7 @@ import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; - +import com.google.api.core.ApiFuture; import com.google.api.core.ApiFutures; import com.google.cloud.grpc.GrpcTransportOptions; import com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory; @@ -96,6 +96,7 @@ public void setUp() throws Exception { firstRun = true; when(session.newTransaction()).thenReturn(txn); transactionRunner = new TransactionRunnerImpl(session, rpc, 1); + when(rpc.commitAsync(Mockito.any(CommitRequest.class), Mockito.anyMap())).thenReturn(ApiFutures.immediateFuture(CommitResponse.newBuilder().setCommitTimestamp(Timestamp.getDefaultInstance()).build())); transactionRunner.setSpan(mock(Span.class)); } @@ -129,25 +130,25 @@ public List answer(InvocationOnMock invocation) .build()); } }); - when(rpc.beginTransaction(Mockito.any(BeginTransactionRequest.class), Mockito.anyMap())) + when(rpc.beginTransactionAsync(Mockito.any(BeginTransactionRequest.class), Mockito.anyMap())) .thenAnswer( - new Answer() { + new Answer< ApiFuture>() { @Override - public Transaction answer(InvocationOnMock invocation) throws Throwable { - return Transaction.newBuilder() + public ApiFuture answer(InvocationOnMock invocation) throws Throwable { + return ApiFutures.immediateFuture(Transaction.newBuilder() .setId(ByteString.copyFromUtf8(UUID.randomUUID().toString())) - .build(); + .build()); } }); - when(rpc.commit(Mockito.any(CommitRequest.class), Mockito.anyMap())) + when(rpc.commitAsync(Mockito.any(CommitRequest.class), Mockito.anyMap())) .thenAnswer( - new Answer() { + new Answer>() { @Override - public CommitResponse answer(InvocationOnMock invocation) throws Throwable { - return CommitResponse.newBuilder() + public ApiFuture answer(InvocationOnMock invocation) throws Throwable { + return ApiFutures.immediateFuture(CommitResponse.newBuilder() .setCommitTimestamp( Timestamp.newBuilder().setSeconds(System.currentTimeMillis() * 1000)) - .build(); + .build()); } }); DatabaseId db = DatabaseId.of("test", "test", "test"); @@ -163,7 +164,7 @@ public Void run(TransactionContext transaction) throws Exception { } }); verify(rpc, times(1)) - .beginTransaction(Mockito.any(BeginTransactionRequest.class), Mockito.anyMap()); + .beginTransactionAsync(Mockito.any(BeginTransactionRequest.class), Mockito.anyMap()); } } @@ -275,8 +276,8 @@ private long[] batchDmlException(int status) { .setRpc(rpc) .build(); when(session.newTransaction()).thenReturn(transaction); - when(session.beginTransaction()) - .thenReturn(ByteString.copyFromUtf8(UUID.randomUUID().toString())); + when(session.beginTransactionAsync()) + .thenReturn(ApiFutures.immediateFuture(ByteString.copyFromUtf8(UUID.randomUUID().toString()))); when(session.getName()).thenReturn(SessionId.of("p", "i", "d", "test").getName()); TransactionRunnerImpl runner = new TransactionRunnerImpl(session, rpc, 10); runner.setSpan(mock(Span.class)); @@ -304,7 +305,7 @@ private long[] batchDmlException(int status) { .thenReturn(response1, response2); CommitResponse commitResponse = CommitResponse.newBuilder().setCommitTimestamp(Timestamp.getDefaultInstance()).build(); - when(rpc.commit(Mockito.any(CommitRequest.class), Mockito.anyMap())).thenReturn(commitResponse); + when(rpc.commitAsync(Mockito.any(CommitRequest.class), Mockito.anyMap())).thenReturn(ApiFutures.immediateFuture(commitResponse)); final Statement statement = Statement.of("UPDATE FOO SET BAR=1"); final AtomicInteger numCalls = new AtomicInteger(0); long updateCount[] = 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 3497b42bc7d..1e05dd6e12e 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 @@ -18,6 +18,7 @@ import com.google.cloud.spanner.MockSpannerServiceImpl; import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult; +import com.google.cloud.spanner.ResultSet; import com.google.cloud.spanner.Statement; import com.google.cloud.spanner.admin.database.v1.MockDatabaseAdminImpl; import com.google.cloud.spanner.admin.instance.v1.MockInstanceAdminImpl; @@ -89,6 +90,11 @@ public abstract class AbstractMockServerTest { Statement.of("INSERT INTO TEST (ID, NAME) VALUES (1, 'test aborted')"); public static final int UPDATE_COUNT = 1; + 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; @@ -112,6 +118,7 @@ public static void startStaticServer() throws IOException { 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 diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AsyncConnectionApiTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AsyncConnectionApiTest.java new file mode 100644 index 00000000000..f5e2e9ad318 --- /dev/null +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AsyncConnectionApiTest.java @@ -0,0 +1,112 @@ +/* + * 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.common.truth.Truth.assertThat; + +import com.google.api.core.ApiFuture; +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.StatementResult; +import com.google.cloud.spanner.Statement; +import com.google.cloud.spanner.connection.ITAbstractSpannerTest.ITConnection; +import com.google.common.base.Function; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.AfterClass; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class AsyncConnectionApiTest extends AbstractMockServerTest { + private static final ExecutorService executor = Executors.newSingleThreadExecutor(); + + @AfterClass + public static void stopExecutor() { + executor.shutdown(); + } + + @Test + public void testSimpleSelectAutocommit() throws Exception { + testSimpleSelect(new Function(){ + @Override + public Void apply(Connection input) { + input.setAutocommit(true); + return null; + } + }); + } + + @Test + public void testSimpleSelectReadOnly() throws Exception { + testSimpleSelect(new Function(){ + @Override + public Void apply(Connection input) { + input.setReadOnly(true); + return null; + } + }); + } + + @Test + public void testSimpleSelectReadWrite() throws Exception { + testSimpleSelect(new Function(){ + @Override + public Void apply(Connection input) { + return null; + } + }); + } + + private void testSimpleSelect(Function connectionConfigurator) throws Exception { + final AtomicInteger rowCount = new AtomicInteger(); + ApiFuture res; + try (ITConnection connection = createConnection()) { + connectionConfigurator.apply(connection); + // Verify that the call is non-blocking. +// mockSpanner.freeze(); + try (AsyncResultSet rs = + connection.executeQueryAsync(SELECT_RANDOM_STATEMENT)) { +// mockSpanner.unfreeze(); + 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; + } + } + } + }); + } + res.get(); + assertThat(rowCount.get()).isEqualTo(RANDOM_RESULT_SET_ROW_COUNT); + } + } +} 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 28bc4564762..e7ac02c49c8 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 @@ -751,6 +751,7 @@ public void testExecuteQueryWithTimeout() { SingleUseTransaction subject = createSubjectWithTimeout(1L); try { subject.executeQuery(createParsedQuery(SLOW_QUERY), AnalyzeMode.NONE); + fail("missing expected exception"); } catch (SpannerException e) { if (e.getErrorCode() != ErrorCode.DEADLINE_EXCEEDED) { throw e; From 226f91bcaf0178c62d49156e6d973c26e66d3676 Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Thu, 11 Jun 2020 18:30:44 +0200 Subject: [PATCH 31/49] review: process review comments --- .../cloud/spanner/AbstractReadContext.java | 29 +- .../google/cloud/spanner/AsyncResultSet.java | 70 ++--- .../cloud/spanner/AsyncResultSetImpl.java | 5 + .../com/google/cloud/spanner/AsyncRunner.java | 19 +- .../google/cloud/spanner/AsyncRunnerImpl.java | 32 +- .../spanner/AsyncTransactionManager.java | 4 +- .../spanner/AsyncTransactionManagerImpl.java | 38 +-- .../spanner/PartitionedDMLTransaction.java | 1 + .../com/google/cloud/spanner/SessionImpl.java | 58 ++-- .../com/google/cloud/spanner/SessionPool.java | 273 ++++++++++-------- .../google/cloud/spanner/SpannerOptions.java | 3 +- .../cloud/spanner/TransactionRunnerImpl.java | 157 +++++----- .../connection/SingleUseTransaction.java | 45 +-- .../cloud/spanner/spi/v1/GapicSpannerRpc.java | 3 +- .../cloud/spanner/spi/v1/SpannerRpc.java | 3 +- .../google/cloud/spanner/AsyncRunnerTest.java | 6 +- .../cloud/spanner/DatabaseClientImplTest.java | 34 ++- .../google/cloud/spanner/SessionImplTest.java | 1 + .../google/cloud/spanner/SessionPoolTest.java | 7 +- .../spanner/TransactionManagerImplTest.java | 22 +- .../spanner/TransactionRunnerImplTest.java | 35 ++- .../connection/AbstractMockServerTest.java | 10 +- .../connection/AsyncConnectionApiTest.java | 55 ++-- 23 files changed, 501 insertions(+), 409 deletions(-) diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java index 9fff15ed62f..b1d752e6f41 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java @@ -779,20 +779,21 @@ private ConsumeSingleRowCallback(SettableApiFuture result) { @Override public CallbackResponse cursorReady(AsyncResultSet resultSet) { try { - while (true) { - switch (resultSet.tryNext()) { - case DONE: - result.set(row); - return CallbackResponse.DONE; - case NOT_READY: - return CallbackResponse.CONTINUE; - case OK: - if (row != null) { - throw newSpannerException( - ErrorCode.INTERNAL, "Multiple rows returned for single key"); - } - row = resultSet.getCurrentRowAsStruct(); - } + switch (resultSet.tryNext()) { + case DONE: + result.set(row); + return CallbackResponse.DONE; + case NOT_READY: + return CallbackResponse.CONTINUE; + case OK: + if (row != null) { + throw newSpannerException( + ErrorCode.INTERNAL, "Multiple rows returned for single key"); + } + row = resultSet.getCurrentRowAsStruct(); + return CallbackResponse.CONTINUE; + default: + throw new IllegalStateException(); } } catch (Throwable t) { result.setException(t); diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncResultSet.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncResultSet.java index 0dcc379e091..c44a42994ed 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncResultSet.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncResultSet.java @@ -25,16 +25,8 @@ /** Interface for result sets returned by async query methods. */ public interface AsyncResultSet extends ResultSet { - /** - * Interface for receiving asynchronous callbacks when new data is ready. See {@link - * AsyncResultSet#setCallback(Executor, ReadyCallback)}. - */ - public static interface ReadyCallback { - CallbackResponse cursorReady(AsyncResultSet resultSet); - } - /** Response code from {@code tryNext()}. */ - public enum CursorState { + enum CursorState { /** Cursor has been moved to a new row. */ OK, /** Read is complete, all rows have been consumed, and there are no more. */ @@ -60,6 +52,40 @@ public enum CursorState { */ CursorState tryNext() throws SpannerException; + enum CallbackResponse { + /** + * Tell the cursor to continue issuing callbacks when data is available. This is the standard + * "I'm ready for more" response. If cursor is not completely drained of all ready results the + * callback will be called again immediately. + */ + CONTINUE, + + /** + * Tell the cursor to suspend all callbacks until application calls {@link RowCursor#resume()}. + */ + PAUSE, + + /** + * Tell the cursor you are done receiving results, even if there are more results sitting in the + * buffer. Once you return DONE, you will receive no further callbacks. + * + *

    Approximately equivalent to calling {@link RowCursor#cancel()}, and then returning {@code + * PAUSE}, but more clear, immediate, and idiomatic. + * + *

    It is legal to commit a transaction that owns this read before actually returning {@code + * DONE}. + */ + DONE, + } + + /** + * Interface for receiving asynchronous callbacks when new data is ready. See {@link + * AsyncResultSet#setCallback(Executor, ReadyCallback)}. + */ + interface ReadyCallback { + CallbackResponse cursorReady(AsyncResultSet resultSet); + } + /** * Register a callback with the ResultSet to be made aware when more data is available, changing * the usage pattern from sync to async. Details: @@ -146,32 +172,6 @@ public enum CursorState { */ void cancel(); - public enum CallbackResponse { - /** - * Tell the cursor to continue issuing callbacks when data is available. This is the standard - * "I'm ready for more" response. If cursor is not completely drained of all ready results the - * callback will be called again immediately. - */ - CONTINUE, - - /** - * Tell the cursor to suspend all callbacks until application calls {@link RowCursor#resume()}. - */ - PAUSE, - - /** - * Tell the cursor you are done receiving results, even if there are more results sitting in the - * buffer. Once you return DONE, you will receive no further callbacks. - * - *

    Approximately equivalent to calling {@link RowCursor#cancel()}, and then returning {@code - * PAUSE}, but more clear, immediate, and idiomatic. - * - *

    It is legal to commit a transaction that owns this read before actually returning {@code - * DONE}. - */ - DONE, - } - /** * Resume callbacks from the cursor. If there is more data available, a callback will be * dispatched immediately. This can be called from any thread. 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 d8f35c31e22..a92026536b6 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 @@ -36,9 +36,12 @@ import java.util.concurrent.Executor; import java.util.concurrent.Future; import java.util.concurrent.LinkedBlockingDeque; +import java.util.logging.Level; +import java.util.logging.Logger; /** Default implementation for {@link AsyncResultSet}. */ class AsyncResultSetImpl extends ForwardingStructReader implements ListenableAsyncResultSet { + private static final Logger log = Logger.getLogger(AsyncResultSetImpl.class.getName()); /** State of an {@link AsyncResultSetImpl}. */ private enum State { @@ -84,6 +87,7 @@ private State(boolean shouldStop) { private Struct currentRow; /** The underlying synchronous {@link ResultSet} that is producing the rows. */ private final ResultSet delegateResultSet; + /** * Any exception that occurs while executing the query and iterating over the result set will be * stored in this variable and propagated to the user through {@link #tryNext()}. @@ -357,6 +361,7 @@ public Void call() throws Exception { try { delegateResultSet.close(); } catch (Throwable t) { + log.log(Level.INFO, "Ignoring error from closing delegate result set", t); } finally { for (Runnable listener : listeners) { listener.run(); diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncRunner.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncRunner.java index de15d79c7ae..3cae49e65b4 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncRunner.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncRunner.java @@ -23,6 +23,10 @@ public interface AsyncRunner { + /** + * Functional interface for executing a read/write transaction asynchronously that returns a + * result of type R. + */ interface AsyncWork { /** * Performs a single transaction attempt. All reads/writes should be performed using {@code @@ -34,17 +38,9 @@ interface AsyncWork { * *

    In most cases, the implementation will not need to catch {@code SpannerException}s from * Spanner operations, instead letting these propagate to the framework. The transaction runner - * - *

    will take appropriate action based on the type of exception. In particular, - * implementations should never catch an exception of type {@link SpannerErrors#isAborted}: - * these indicate that some reads may have returned inconsistent data and the transaction - * attempt must be aborted. - * - *

    If any exception is thrown, the runner will validate the reads performed in the current - * transaction attempt using {@link Transaction#commitReadsOnly}: if validation succeeds, the - * exception is propagated to the caller; if validation aborts, the exception is thrown away and - * the work is retried; if the commit fails for some other reason, the corresponding {@code - * SpannerException} is returned to the caller. Any buffered mutations will be ignored. + * will take appropriate action based on the type of exception. In particular, implementations + * should never catch an exception of type {@link SpannerErrors#isAborted}: these indicate that + * some reads may have returned inconsistent data and the transaction attempt must be aborted. * * @param txn the transaction * @return future over the result of the work @@ -52,6 +48,7 @@ interface AsyncWork { ApiFuture doWorkAsync(TransactionContext txn); } + /** Executes a read/write transaction asynchronously using the given executor. */ ApiFuture runAsync(AsyncWork work, Executor executor); /** diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncRunnerImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncRunnerImpl.java index 6ffb3214907..5b83402919e 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncRunnerImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncRunnerImpl.java @@ -39,21 +39,7 @@ public ApiFuture runAsync(final AsyncWork work, Executor executor) { @Override public void run() { try { - R r = - delegate.run( - new TransactionCallable() { - @Override - public R run(TransactionContext transaction) throws Exception { - try { - return work.doWorkAsync(transaction).get(); - } catch (ExecutionException e) { - throw SpannerExceptionFactory.newSpannerException(e.getCause()); - } catch (InterruptedException e) { - throw SpannerExceptionFactory.propagateInterrupt(e); - } - } - }); - res.set(r); + res.set(runTransaction(work)); } catch (Throwable t) { res.setException(t); } finally { @@ -64,6 +50,22 @@ public R run(TransactionContext transaction) throws Exception { return res; } + private R runTransaction(final AsyncWork work) { + return delegate.run( + new TransactionCallable() { + @Override + public R run(TransactionContext transaction) throws Exception { + try { + return work.doWorkAsync(transaction).get(); + } catch (ExecutionException e) { + throw SpannerExceptionFactory.newSpannerException(e.getCause()); + } catch (InterruptedException e) { + throw SpannerExceptionFactory.propagateInterrupt(e); + } + } + }); + } + private void setCommitTimestamp() { try { commitTimestamp.set(delegate.getCommitTimestamp()); diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManager.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManager.java index a86c5a776c4..e5b5b3bb036 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManager.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManager.java @@ -31,8 +31,8 @@ * successfully, or is rolled back or commit fails with any error other than {@code ABORTED}, the * manager is considered complete and no further transactions are allowed to be created in it. * - *

    Every {@code AsyncTransactionManager} should either be committed or rolled back. Failure to do so - * can cause resources to be leaked and deadlocks. Easiest way to guarantee this is by calling + *

    Every {@code AsyncTransactionManager} should either be committed or rolled back. Failure to do + * so can cause resources to be leaked and deadlocks. Easiest way to guarantee this is by calling * {@link #close()} in a finally block. * * @see DatabaseClient#transactionManager() diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManagerImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManagerImpl.java index 945498407dd..7c7f8504c74 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManagerImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManagerImpl.java @@ -16,7 +16,6 @@ package com.google.cloud.spanner; -import java.util.concurrent.ExecutionException; import com.google.api.core.ApiFuture; import com.google.api.core.ApiFutures; import com.google.api.core.SettableApiFuture; @@ -25,11 +24,11 @@ import com.google.cloud.spanner.TransactionManager.TransactionState; import com.google.common.base.Preconditions; import com.google.common.util.concurrent.MoreExecutors; -import com.google.common.util.concurrent.SettableFuture; import io.opencensus.common.Scope; import io.opencensus.trace.Span; import io.opencensus.trace.Tracer; import io.opencensus.trace.Tracing; +import java.util.concurrent.ExecutionException; /** Implementation of {@link AsyncTransactionManager}. */ final class AsyncTransactionManagerImpl implements AsyncTransactionManager, SessionTransaction { @@ -59,19 +58,23 @@ public ApiFuture beginAsync() { session.setActive(this); final SettableApiFuture res = SettableApiFuture.create(); final ApiFuture fut = txn.ensureTxnAsync(); - fut.addListener(tracer.withSpan(span, new Runnable(){ - @Override - public void run() { - try { - fut.get(); - res.set(txn); - } catch (ExecutionException e) { - res.setException(e.getCause() == null ? e : e.getCause()); - } catch (InterruptedException e) { - res.setException(SpannerExceptionFactory.propagateInterrupt(e)); - } - } - }), MoreExecutors.directExecutor()); + fut.addListener( + tracer.withSpan( + span, + new Runnable() { + @Override + public void run() { + try { + fut.get(); + res.set(txn); + } catch (ExecutionException e) { + res.setException(e.getCause() == null ? e : e.getCause()); + } catch (InterruptedException e) { + res.setException(SpannerExceptionFactory.propagateInterrupt(e)); + } + } + }), + MoreExecutors.directExecutor()); return res; } @@ -83,8 +86,9 @@ public ApiFuture commitAsync() { SettableApiFuture res = SettableApiFuture.create(); if (txn.isAborted()) { txnState = TransactionState.ABORTED; - res.setException(SpannerExceptionFactory.newSpannerException( - ErrorCode.ABORTED, "Transaction already aborted")); + res.setException( + SpannerExceptionFactory.newSpannerException( + ErrorCode.ABORTED, "Transaction already aborted")); } try { txn.commit(); diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/PartitionedDMLTransaction.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/PartitionedDMLTransaction.java index 351b7596287..3892a018f1c 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/PartitionedDMLTransaction.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/PartitionedDMLTransaction.java @@ -103,6 +103,7 @@ public void invalidate() { isValid = false; } + // No-op method needed to implement SessionTransaction interface. @Override public void setSpan(Span span) {} } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java index 01fd8d758de..f0e01cf986a 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java @@ -18,9 +18,8 @@ import static com.google.cloud.spanner.SpannerExceptionFactory.newSpannerException; import static com.google.common.base.Preconditions.checkNotNull; -import com.google.api.core.ApiFunction; + import com.google.api.core.ApiFuture; -import com.google.api.core.ApiFutures; import com.google.api.core.SettableApiFuture; import com.google.cloud.Timestamp; import com.google.cloud.spanner.AbstractReadContext.MultiUseReadOnlyTransaction; @@ -38,7 +37,6 @@ import com.google.spanner.v1.CommitResponse; import com.google.spanner.v1.Transaction; import com.google.spanner.v1.TransactionOptions; -import io.grpc.Context; import io.opencensus.common.Scope; import io.opencensus.trace.Span; import io.opencensus.trace.Tracer; @@ -291,29 +289,37 @@ ApiFuture beginTransactionAsync() { TransactionOptions.newBuilder() .setReadWrite(TransactionOptions.ReadWrite.getDefaultInstance())) .build(); - final ApiFuture requestFuture = spanner.getRpc().beginTransactionAsync(request, options); - requestFuture.addListener(tracer.withSpan(span, new Runnable(){ - @Override - public void run() { - try { - Transaction txn = requestFuture.get(); - if (txn.getId().isEmpty()) { - throw newSpannerException(ErrorCode.INTERNAL, "Missing id in transaction\n" + getName()); - } - span.end(TraceUtil.END_SPAN_OPTIONS); - res.set(txn.getId()); - } catch (ExecutionException e) { - TraceUtil.endSpanWithFailure(span, e); - res.setException(SpannerExceptionFactory.newSpannerException(e.getCause() == null ? e : e.getCause())); - } catch (InterruptedException e) { - TraceUtil.endSpanWithFailure(span, e); - res.setException(SpannerExceptionFactory.propagateInterrupt(e)); - } catch (Exception e) { - TraceUtil.endSpanWithFailure(span, e); - res.setException(e); - } - } - }), MoreExecutors.directExecutor()); + final ApiFuture requestFuture = + spanner.getRpc().beginTransactionAsync(request, options); + requestFuture.addListener( + tracer.withSpan( + span, + new Runnable() { + @Override + public void run() { + try { + Transaction txn = requestFuture.get(); + if (txn.getId().isEmpty()) { + throw newSpannerException( + ErrorCode.INTERNAL, "Missing id in transaction\n" + getName()); + } + span.end(TraceUtil.END_SPAN_OPTIONS); + res.set(txn.getId()); + } catch (ExecutionException e) { + TraceUtil.endSpanWithFailure(span, e); + res.setException( + SpannerExceptionFactory.newSpannerException( + e.getCause() == null ? e : e.getCause())); + } catch (InterruptedException e) { + TraceUtil.endSpanWithFailure(span, e); + res.setException(SpannerExceptionFactory.propagateInterrupt(e)); + } catch (Exception e) { + TraceUtil.endSpanWithFailure(span, e); + res.setException(e); + } + } + }), + MoreExecutors.directExecutor()); return res; } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java index 9b90ff64776..60282aa74b3 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java @@ -41,14 +41,13 @@ import com.google.cloud.grpc.GrpcTransportOptions; import com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory; import com.google.cloud.spanner.AbstractReadContext.ConsumeSingleRowCallback; -import com.google.cloud.spanner.AbstractReadContext.ListenableAsyncResultSet; import com.google.cloud.spanner.AsyncResultSet.ReadyCallback; import com.google.cloud.spanner.Options.QueryOption; import com.google.cloud.spanner.Options.ReadOption; import com.google.cloud.spanner.SessionClient.SessionConsumer; +import com.google.cloud.spanner.SessionPool.PooledSession; import com.google.cloud.spanner.SpannerException.ResourceNotFoundException; import com.google.cloud.spanner.TransactionManager.TransactionState; -import com.google.cloud.spanner.TransactionRunnerImpl.TransactionContextImpl; import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Function; import com.google.common.base.MoreObjects; @@ -57,8 +56,6 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; -import com.google.common.util.concurrent.ForwardingFuture; -import com.google.common.util.concurrent.ForwardingFuture.SimpleForwardingFuture; import com.google.common.util.concurrent.ForwardingListenableFuture; import com.google.common.util.concurrent.ForwardingListenableFuture.SimpleForwardingListenableFuture; import com.google.common.util.concurrent.ListenableFuture; @@ -92,7 +89,6 @@ import java.util.concurrent.ExecutionException; import java.util.concurrent.Executor; import java.util.concurrent.Executors; -import java.util.concurrent.Future; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.ScheduledFuture; import java.util.concurrent.ThreadPoolExecutor; @@ -182,11 +178,9 @@ public ApiFuture setCallback(Executor exec, ReadyCallback cb) { @Override public void run() { synchronized (lock) { - if (asyncOperationsCount.decrementAndGet() == 0) { - if (closed) { - // All async operations for this read context have finished. - AutoClosingReadContext.this.close(); - } + if (asyncOperationsCount.decrementAndGet() == 0 && closed) { + // All async operations for this read context have finished. + AutoClosingReadContext.this.close(); } } } @@ -960,34 +954,43 @@ private static class SessionPoolAsyncTransactionManager implements AsyncTransact private final SettableApiFuture commitTimestamp = SettableApiFuture.create(); private final SettableApiFuture delegate = SettableApiFuture.create(); - private SessionPoolAsyncTransactionManager(SessionPool sessionPool, PooledSessionFuture session) { + private SessionPoolAsyncTransactionManager( + SessionPool sessionPool, PooledSessionFuture session) { this.sessionPool = sessionPool; this.session = session; - this.session.addListener(new Runnable(){ - @Override - public void run() { - try { - delegate.set(SessionPoolAsyncTransactionManager.this.session.get().transactionManagerAsync()); - } catch (Throwable t) { - delegate.setException(t); - } - } - }, MoreExecutors.directExecutor()); + this.session.addListener( + new Runnable() { + @Override + public void run() { + try { + delegate.set( + SessionPoolAsyncTransactionManager.this + .session + .get() + .transactionManagerAsync()); + } catch (Throwable t) { + delegate.setException(t); + } + } + }, + MoreExecutors.directExecutor()); } @Override public ApiFuture beginAsync() { final SettableApiFuture res = SettableApiFuture.create(); - delegate.addListener(new Runnable(){ - @Override - public void run() { - try { - res.set(delegate.get().beginAsync().get()); - } catch (Throwable t) { - res.setException(t); - } - } - }, MoreExecutors.directExecutor()); + delegate.addListener( + new Runnable() { + @Override + public void run() { + try { + res.set(delegate.get().beginAsync().get()); + } catch (Throwable t) { + res.setException(t); + } + } + }, + MoreExecutors.directExecutor()); return res; } @@ -1051,11 +1054,116 @@ private enum SessionState { CLOSING, } - private PooledSessionFuture createPooledSessionFuture(ListenableFuture future, Span span) { + /** + * Forwarding future that will return a {@link PooledSession}. If {@link #inProcessPrepare} has + * been set to true, the returned session will be prepared with a read/write session using the + * thread of the caller to {@link #get()}. This ensures that the executor that is responsible for + * background preparing of read/write transactions is not overwhelmed by requests in case of a + * large burst of write requests. Instead of filling up the queue of the background executor, the + * caller threads will be used for the BeginTransaction call. + */ + private final class ForwardingListenablePooledSessionFuture + extends SimpleForwardingListenableFuture { + private final boolean inProcessPrepare; + private final Span span; + private volatile boolean initialized = false; + private final Object prepareLock = new Object(); + private volatile PooledSession result; + private volatile SpannerException error; + + private ForwardingListenablePooledSessionFuture( + ListenableFuture delegate, boolean inProcessPrepare, Span span) { + super(delegate); + this.inProcessPrepare = inProcessPrepare; + this.span = span; + } + + @Override + public PooledSession get() throws InterruptedException, ExecutionException { + try { + return initialize(super.get()); + } catch (ExecutionException e) { + throw SpannerExceptionFactory.newSpannerException(e.getCause()); + } catch (InterruptedException e) { + throw SpannerExceptionFactory.propagateInterrupt(e); + } + } + + @Override + public PooledSession get(long timeout, TimeUnit unit) + throws InterruptedException, ExecutionException, TimeoutException { + try { + return initialize(super.get(timeout, unit)); + } catch (ExecutionException e) { + throw SpannerExceptionFactory.newSpannerException(e.getCause()); + } catch (InterruptedException e) { + throw SpannerExceptionFactory.propagateInterrupt(e); + } catch (TimeoutException e) { + throw SpannerExceptionFactory.propagateTimeout(e); + } + } + + private PooledSession initialize(PooledSession sess) { + if (!initialized) { + synchronized (prepareLock) { + if (!initialized) { + try { + result = prepare(sess); + } catch (Throwable t) { + error = SpannerExceptionFactory.newSpannerException(t); + } finally { + initialized = true; + } + } + } + } + if (error != null) { + throw error; + } + return result; + } + + private PooledSession prepare(PooledSession sess) { + if (inProcessPrepare && !sess.delegate.hasReadyTransaction()) { + while (true) { + try { + sess.prepareReadWriteTransaction(); + synchronized (lock) { + stopAutomaticPrepare = false; + } + break; + } catch (Throwable t) { + if (isClosed()) { + span.addAnnotation("Pool has been closed"); + throw new IllegalStateException("Pool has been closed"); + } + SpannerException e = newSpannerException(t); + WaiterFuture waiter = new WaiterFuture(); + synchronized (lock) { + handlePrepareSessionFailure(e, sess, false); + if (!isSessionNotFound(e)) { + throw e; + } + readWaiters.add(waiter); + } + sess = waiter.get(); + if (sess.delegate.hasReadyTransaction()) { + break; + } + } + } + } + return sess; + } + } + + private PooledSessionFuture createPooledSessionFuture( + ListenableFuture future, Span span) { return new PooledSessionFuture(future, span); } - final class PooledSessionFuture extends SimpleForwardingListenableFuture implements Session { + final class PooledSessionFuture extends SimpleForwardingListenableFuture + implements Session { private volatile LeakedSessionException leakedException; private volatile AtomicBoolean inUse = new AtomicBoolean(); private volatile CountDownLatch initialized = new CountDownLatch(1); @@ -1251,19 +1359,21 @@ public ApiFuture asyncClose() { @Override public PooledSession get() { if (inUse.compareAndSet(false, true)) { + PooledSession res = null; try { - PooledSession res = super.get(); + res = super.get(); + } catch (Throwable e) { + // ignore the exception as it will be handled by the call to super.get() below. + } + if (res != null) { synchronized (lock) { res.markBusy(span); span.addAnnotation(sessionAnnotation(res)); incrementNumSessionsInUse(); checkedOutSessions.add(this); } - initialized.countDown(); - } catch (Throwable e) { - initialized.countDown(); - // ignore and fallthrough. } + initialized.countDown(); } try { initialized.await(); @@ -2146,91 +2256,8 @@ private PooledSessionFuture checkoutSession( fut.set(readySession); sessionFuture = fut; } - SimpleForwardingListenableFuture forwardingFuture = - new SimpleForwardingListenableFuture(sessionFuture) { - private volatile boolean initialized = false; - private final Object prepareLock = new Object(); - private volatile PooledSession result; - private volatile SpannerException error; - - @Override - public PooledSession get() throws InterruptedException, ExecutionException { - try { - return initialize(super.get()); - } catch (ExecutionException e) { - throw SpannerExceptionFactory.newSpannerException(e.getCause()); - } catch (InterruptedException e) { - throw SpannerExceptionFactory.propagateInterrupt(e); - } - } - - @Override - public PooledSession get(long timeout, TimeUnit unit) - throws InterruptedException, ExecutionException, TimeoutException { - try { - return initialize(super.get(timeout, unit)); - } catch (ExecutionException e) { - throw SpannerExceptionFactory.newSpannerException(e.getCause()); - } catch (InterruptedException e) { - throw SpannerExceptionFactory.propagateInterrupt(e); - } catch (TimeoutException e) { - throw SpannerExceptionFactory.propagateTimeout(e); - } - } - - private PooledSession initialize(PooledSession sess) { - if (!initialized) { - synchronized (prepareLock) { - if (!initialized) { - try { - result = prepare(sess); - } catch (Throwable t) { - error = SpannerExceptionFactory.newSpannerException(t); - } finally { - initialized = true; - } - } - } - } - if (error != null) { - throw error; - } - return result; - } - - private PooledSession prepare(PooledSession sess) { - if (inProcessPrepare && !sess.delegate.hasReadyTransaction()) { - while (true) { - try { - sess.prepareReadWriteTransaction(); - synchronized (lock) { - stopAutomaticPrepare = false; - } - break; - } catch (Throwable t) { - if (isClosed()) { - span.addAnnotation("Pool has been closed"); - throw new IllegalStateException("Pool has been closed"); - } - SpannerException e = newSpannerException(t); - WaiterFuture waiter = new WaiterFuture(); - synchronized (lock) { - handlePrepareSessionFailure(e, sess, false); - if (!isSessionNotFound(e)) { - throw e; - } - readWaiters.add(waiter); - } - sess = waiter.get(); - if (sess.delegate.hasReadyTransaction()) { - break; - } - } - } - } - return sess; - } - }; + ForwardingListenablePooledSessionFuture forwardingFuture = + new ForwardingListenablePooledSessionFuture(sessionFuture, inProcessPrepare, span); PooledSessionFuture res = createPooledSessionFuture(forwardingFuture, span); res.markCheckedOut(); return res; diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java index d0bad0e6529..d8145a5a25f 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java @@ -194,7 +194,8 @@ static CloseableExecutorProvider createDefaultAsyncExecutorProvider() { @VisibleForTesting static CloseableExecutorProvider createAsyncExecutorProvider( int poolSize, long keepAliveTime, TimeUnit unit) { - String format = String.format("async-pool-%d-thread-%%d", DEFAULT_POOL_COUNT.incrementAndGet()); + String format = + String.format("spanner-async-pool-%d-thread-%%d", DEFAULT_POOL_COUNT.incrementAndGet()); ThreadFactory threadFactory = new ThreadFactoryBuilder().setDaemon(true).setNameFormat(format).build(); ScheduledThreadPoolExecutor executor = new ScheduledThreadPoolExecutor(poolSize, threadFactory); 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 a91c6a90ef1..9373edde2aa 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 @@ -52,7 +52,6 @@ import java.util.ArrayList; import java.util.List; import java.util.concurrent.Callable; -import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; import java.util.concurrent.Executor; import java.util.concurrent.atomic.AtomicInteger; @@ -135,11 +134,12 @@ public void removeListener(Runnable listener) { @GuardedBy("lock") private volatile SettableApiFuture finishedAsyncOperations = SettableApiFuture.create(); + @GuardedBy("lock") private volatile int runningAsyncOperations; -// @GuardedBy("lock") -// private volatile CountDownLatch finishedAsyncOperations = new CountDownLatch(0); + // @GuardedBy("lock") + // private volatile CountDownLatch finishedAsyncOperations = new CountDownLatch(0); @GuardedBy("lock") private List mutations = new ArrayList<>(); @@ -222,28 +222,34 @@ ApiFuture ensureTxnAsync() { if (transactionId == null || isAborted()) { span.addAnnotation("Creating Transaction"); final ApiFuture fut = session.beginTransactionAsync(); - fut.addListener(new Runnable(){ - @Override - public void run() { - try { - transactionId = fut.get(); - span.addAnnotation( - "Transaction Creation Done", - ImmutableMap.of( - "Id", AttributeValue.stringAttributeValue(transactionId.toStringUtf8()))); - txnLogger.log( - Level.FINER, - "Started transaction {0}", - txnLogger.isLoggable(Level.FINER) ? transactionId.asReadOnlyByteBuffer() : null); - res.set(null); - } catch (ExecutionException e) { - span.addAnnotation("Transaction Creation Failed", TraceUtil.getExceptionAnnotations(e.getCause() == null ? e : e.getCause())); - res.setException(e.getCause() == null ? e : e.getCause()); - } catch (InterruptedException e) { - res.setException(SpannerExceptionFactory.propagateInterrupt(e)); - } - } - }, MoreExecutors.directExecutor()); + fut.addListener( + new Runnable() { + @Override + public void run() { + try { + transactionId = fut.get(); + span.addAnnotation( + "Transaction Creation Done", + ImmutableMap.of( + "Id", AttributeValue.stringAttributeValue(transactionId.toStringUtf8()))); + txnLogger.log( + Level.FINER, + "Started transaction {0}", + txnLogger.isLoggable(Level.FINER) + ? transactionId.asReadOnlyByteBuffer() + : null); + res.set(null); + } catch (ExecutionException e) { + span.addAnnotation( + "Transaction Creation Failed", + TraceUtil.getExceptionAnnotations(e.getCause() == null ? e : e.getCause())); + res.setException(e.getCause() == null ? e : e.getCause()); + } catch (InterruptedException e) { + res.setException(SpannerExceptionFactory.propagateInterrupt(e)); + } + } + }, + MoreExecutors.directExecutor()); } else { span.addAnnotation( "Transaction Initialized", @@ -331,49 +337,66 @@ ApiFuture commitAsync() { synchronized (lock) { latch = finishedAsyncOperations; } - latch.addListener(new Runnable(){ - @Override - public void run() { - try { - latch.get(); - span.addAnnotation("Starting Commit"); - final Span opSpan = tracer.spanBuilderWithExplicitParent(SpannerImpl.COMMIT, span).startSpan(); - final ApiFuture commitFuture = rpc.commitAsync(commitRequest, session.getOptions()); - commitFuture.addListener(tracer.withSpan(opSpan, new Runnable(){ - @Override - public void run() { - try { - CommitResponse commitResponse = commitFuture.get(); - if (!commitResponse.hasCommitTimestamp()) { - throw newSpannerException( - ErrorCode.INTERNAL, "Missing commitTimestamp:\n" + session.getName()); - } - commitTimestamp = Timestamp.fromProto(commitResponse.getCommitTimestamp()); - span.addAnnotation("Commit Done"); - opSpan.end(TraceUtil.END_SPAN_OPTIONS); - res.set(null); - } catch (Throwable e) { - if (e instanceof ExecutionException) { - e = SpannerExceptionFactory.newSpannerException(e.getCause() == null ? e : e.getCause()); - } else if (e instanceof InterruptedException) { - e = SpannerExceptionFactory.propagateInterrupt((InterruptedException) e); - } else { - e = SpannerExceptionFactory.newSpannerException(e); - } - span.addAnnotation("Commit Failed", TraceUtil.getExceptionAnnotations(e)); - TraceUtil.endSpanWithFailure(opSpan, e); - onError((SpannerException) e); - res.setException(e); - } + latch.addListener( + new Runnable() { + @Override + public void run() { + try { + latch.get(); + span.addAnnotation("Starting Commit"); + final Span opSpan = + tracer.spanBuilderWithExplicitParent(SpannerImpl.COMMIT, span).startSpan(); + final ApiFuture commitFuture = + rpc.commitAsync(commitRequest, session.getOptions()); + commitFuture.addListener( + tracer.withSpan( + opSpan, + new Runnable() { + @Override + public void run() { + try { + CommitResponse commitResponse = commitFuture.get(); + if (!commitResponse.hasCommitTimestamp()) { + throw newSpannerException( + ErrorCode.INTERNAL, + "Missing commitTimestamp:\n" + session.getName()); + } + commitTimestamp = + Timestamp.fromProto(commitResponse.getCommitTimestamp()); + span.addAnnotation("Commit Done"); + opSpan.end(TraceUtil.END_SPAN_OPTIONS); + res.set(null); + } catch (Throwable e) { + if (e instanceof ExecutionException) { + e = + SpannerExceptionFactory.newSpannerException( + e.getCause() == null ? e : e.getCause()); + } else if (e instanceof InterruptedException) { + e = + SpannerExceptionFactory.propagateInterrupt( + (InterruptedException) e); + } else { + e = SpannerExceptionFactory.newSpannerException(e); + } + span.addAnnotation( + "Commit Failed", TraceUtil.getExceptionAnnotations(e)); + TraceUtil.endSpanWithFailure(opSpan, e); + onError((SpannerException) e); + res.setException(e); + } + } + }), + MoreExecutors.directExecutor()); + } catch (InterruptedException e) { + res.setException(SpannerExceptionFactory.propagateInterrupt(e)); + } catch (ExecutionException e) { + res.setException( + SpannerExceptionFactory.newSpannerException( + e.getCause() == null ? e : e.getCause())); } - }), MoreExecutors.directExecutor()); - } catch (InterruptedException e) { - res.setException(SpannerExceptionFactory.propagateInterrupt(e)); - } catch (ExecutionException e) { - res.setException(SpannerExceptionFactory.newSpannerException(e.getCause() == null ? e : e.getCause())); - } - } - }, MoreExecutors.directExecutor()); + } + }, + MoreExecutors.directExecutor()); return res; } 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 edf4a9a5289..3da17cf26d4 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 @@ -174,25 +174,25 @@ public ResultSet executeQuery( new Callable() { @Override public ResultSet call() throws Exception { -// try { - ResultSet rs; - if (analyzeMode == AnalyzeMode.NONE) { - rs = readOnlyTransaction.executeQuery(statement.getStatement(), options); - } else { - rs = - readOnlyTransaction.analyzeQuery( - statement.getStatement(), analyzeMode.getQueryAnalyzeMode()); - } - // 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); -// } catch (Exception e) { -// readOnlyTransaction.close(); -// throw e; -// } finally { -// readOnlyTransaction.close(); -// currentTransaction.close(); -// } + // try { + ResultSet rs; + if (analyzeMode == AnalyzeMode.NONE) { + rs = readOnlyTransaction.executeQuery(statement.getStatement(), options); + } else { + rs = + readOnlyTransaction.analyzeQuery( + statement.getStatement(), analyzeMode.getQueryAnalyzeMode()); + } + // 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); + // } catch (Exception e) { + // readOnlyTransaction.close(); + // throw e; + // } finally { + // readOnlyTransaction.close(); + // currentTransaction.close(); + // } } }; try { @@ -234,13 +234,16 @@ public AsyncResultSet executeQueryAsync( @Override public Timestamp getReadTimestamp() { ConnectionPreconditions.checkState( - readOnlyTransaction != null && state == UnitOfWorkState.COMMITTED, "There is no read timestamp available for this transaction."); + readOnlyTransaction != null && state == UnitOfWorkState.COMMITTED, + "There is no read timestamp available for this transaction."); return readOnlyTransaction.getReadTimestamp(); } @Override public Timestamp getReadTimestampOrNull() { - return readOnlyTransaction == null || state != UnitOfWorkState.COMMITTED ? null : readOnlyTransaction.getReadTimestamp(); + return readOnlyTransaction == null || state != UnitOfWorkState.COMMITTED + ? null + : readOnlyTransaction.getReadTimestamp(); } private boolean hasCommitTimestamp() { diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java index 57ad55480c4..fb9dc1da598 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java @@ -1089,7 +1089,8 @@ public Transaction beginTransaction( } @Override - public ApiFuture commitAsync(CommitRequest commitRequest, @Nullable Map options) { + public ApiFuture commitAsync( + CommitRequest commitRequest, @Nullable Map options) { GrpcCallContext context = newCallContext(options, commitRequest.getSession()); return spannerStub.commitCallable().futureCall(commitRequest, context); } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerRpc.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerRpc.java index 31347ae547d..84d8b1b6be6 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerRpc.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerRpc.java @@ -305,7 +305,8 @@ ApiFuture beginTransactionAsync( CommitResponse commit(CommitRequest commitRequest, @Nullable Map options) throws SpannerException; - ApiFuture commitAsync(CommitRequest commitRequest, @Nullable Map options); + ApiFuture commitAsync( + CommitRequest commitRequest, @Nullable Map options); void rollback(RollbackRequest request, @Nullable Map options) throws SpannerException; diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncRunnerTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncRunnerTest.java index a020997fd7d..a3ba51d0004 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncRunnerTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncRunnerTest.java @@ -32,7 +32,6 @@ import com.google.cloud.spanner.AsyncRunner.AsyncWork; import com.google.cloud.spanner.MockSpannerServiceImpl.SimulatedExecutionTime; import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult; -import com.google.cloud.spanner.TransactionRunnerImpl.TransactionContextImpl; import com.google.common.base.Function; import com.google.common.collect.ImmutableList; import com.google.common.util.concurrent.MoreExecutors; @@ -612,7 +611,7 @@ public void closeTransactionBeforeEndOfAsyncQuery() throws Exception { runner.runAsync( new AsyncWork() { @Override - public ApiFuture doWorkAsync(TransactionContext txn) { + public ApiFuture doWorkAsync(TransactionContext txn) { try (AsyncResultSet rs = txn.readAsync( READ_TABLE_NAME, KeySet.all(), READ_COLUMN_NAMES, Options.bufferRows(1))) { @@ -646,7 +645,8 @@ public CallbackResponse cursorReady(AsyncResultSet resultSet) { dataReceived.await(); return ApiFutures.immediateFuture(null); } catch (InterruptedException e) { - return ApiFutures.immediateFailedFuture(SpannerExceptionFactory.propagateInterrupt(e)); + return ApiFutures.immediateFailedFuture( + SpannerExceptionFactory.propagateInterrupt(e)); } } }, diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java index 5100b2723ee..e2655ccc722 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java @@ -555,23 +555,25 @@ public void transactionManagerExecuteQueryAsync() throws Exception { while (true) { TransactionContext tx = txManager.begin(); try { - try(AsyncResultSet rs = tx.executeQueryAsync(SELECT1)) { - 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; + try (AsyncResultSet rs = tx.executeQueryAsync(SELECT1)) { + 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; + } + } } - } - } - }); + }); } txManager.commit(); break; diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionImplTest.java index cbcb10d25bc..06eff8ebfb9 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionImplTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionImplTest.java @@ -20,6 +20,7 @@ import static org.junit.Assert.fail; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; + import com.google.api.core.ApiFutures; import com.google.api.core.NanoClock; import com.google.api.gax.retrying.RetrySettings; diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java index 538075fe4f9..8e635380793 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java @@ -69,7 +69,6 @@ import java.util.Map; 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; @@ -1298,7 +1297,8 @@ public void testSessionNotFoundReadWriteTransaction() { .thenThrow(sessionNotFound); when(rpc.executeBatchDml(any(ExecuteBatchDmlRequest.class), any(Map.class))) .thenThrow(sessionNotFound); - when(rpc.commitAsync(any(CommitRequest.class), any(Map.class))).thenReturn(ApiFutures.immediateFailedFuture(sessionNotFound)); + when(rpc.commitAsync(any(CommitRequest.class), any(Map.class))) + .thenReturn(ApiFutures.immediateFailedFuture(sessionNotFound)); doThrow(sessionNotFound).when(rpc).rollback(any(RollbackRequest.class), any(Map.class)); final SessionImpl closedSession = mock(SessionImpl.class); when(closedSession.getName()) @@ -1327,7 +1327,8 @@ public void testSessionNotFoundReadWriteTransaction() { .thenReturn("projects/dummy/instances/dummy/database/dummy/sessions/session-open"); final TransactionContextImpl openTransactionContext = mock(TransactionContextImpl.class); when(openSession.newTransaction()).thenReturn(openTransactionContext); - when(openSession.beginTransactionAsync()).thenReturn(ApiFutures.immediateFuture(ByteString.copyFromUtf8("open-txn"))); + when(openSession.beginTransactionAsync()) + .thenReturn(ApiFutures.immediateFuture(ByteString.copyFromUtf8("open-txn"))); TransactionRunnerImpl openTransactionRunner = new TransactionRunnerImpl(openSession, mock(SpannerRpc.class), 10); openTransactionRunner.setSpan(mock(Span.class)); diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionManagerImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionManagerImplTest.java index ead4a6bb159..ad32b20586c 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionManagerImplTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionManagerImplTest.java @@ -24,6 +24,7 @@ import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import static org.mockito.MockitoAnnotations.initMocks; + import com.google.api.core.ApiFuture; import com.google.api.core.ApiFutures; import com.google.cloud.Timestamp; @@ -224,21 +225,24 @@ public List answer(InvocationOnMock invocation) new Answer>() { @Override public ApiFuture answer(InvocationOnMock invocation) throws Throwable { - return ApiFutures.immediateFuture(Transaction.newBuilder() - .setId(ByteString.copyFromUtf8(UUID.randomUUID().toString())) - .build()); + return ApiFutures.immediateFuture( + Transaction.newBuilder() + .setId(ByteString.copyFromUtf8(UUID.randomUUID().toString())) + .build()); } }); when(rpc.commitAsync(Mockito.any(CommitRequest.class), Mockito.anyMap())) .thenAnswer( new Answer>() { @Override - public ApiFuture answer(InvocationOnMock invocation) throws Throwable { - return ApiFutures.immediateFuture(CommitResponse.newBuilder() - .setCommitTimestamp( - com.google.protobuf.Timestamp.newBuilder() - .setSeconds(System.currentTimeMillis() * 1000)) - .build()); + public ApiFuture answer(InvocationOnMock invocation) + throws Throwable { + return ApiFutures.immediateFuture( + CommitResponse.newBuilder() + .setCommitTimestamp( + com.google.protobuf.Timestamp.newBuilder() + .setSeconds(System.currentTimeMillis() * 1000)) + .build()); } }); DatabaseId db = DatabaseId.of("test", "test", "test"); diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java index 74d8b1c9054..a9abea938d8 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java @@ -23,6 +23,7 @@ import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; + import com.google.api.core.ApiFuture; import com.google.api.core.ApiFutures; import com.google.cloud.grpc.GrpcTransportOptions; @@ -96,7 +97,12 @@ public void setUp() throws Exception { firstRun = true; when(session.newTransaction()).thenReturn(txn); transactionRunner = new TransactionRunnerImpl(session, rpc, 1); - when(rpc.commitAsync(Mockito.any(CommitRequest.class), Mockito.anyMap())).thenReturn(ApiFutures.immediateFuture(CommitResponse.newBuilder().setCommitTimestamp(Timestamp.getDefaultInstance()).build())); + when(rpc.commitAsync(Mockito.any(CommitRequest.class), Mockito.anyMap())) + .thenReturn( + ApiFutures.immediateFuture( + CommitResponse.newBuilder() + .setCommitTimestamp(Timestamp.getDefaultInstance()) + .build())); transactionRunner.setSpan(mock(Span.class)); } @@ -132,23 +138,26 @@ public List answer(InvocationOnMock invocation) }); when(rpc.beginTransactionAsync(Mockito.any(BeginTransactionRequest.class), Mockito.anyMap())) .thenAnswer( - new Answer< ApiFuture>() { + new Answer>() { @Override public ApiFuture answer(InvocationOnMock invocation) throws Throwable { - return ApiFutures.immediateFuture(Transaction.newBuilder() - .setId(ByteString.copyFromUtf8(UUID.randomUUID().toString())) - .build()); + return ApiFutures.immediateFuture( + Transaction.newBuilder() + .setId(ByteString.copyFromUtf8(UUID.randomUUID().toString())) + .build()); } }); when(rpc.commitAsync(Mockito.any(CommitRequest.class), Mockito.anyMap())) .thenAnswer( new Answer>() { @Override - public ApiFuture answer(InvocationOnMock invocation) throws Throwable { - return ApiFutures.immediateFuture(CommitResponse.newBuilder() - .setCommitTimestamp( - Timestamp.newBuilder().setSeconds(System.currentTimeMillis() * 1000)) - .build()); + public ApiFuture answer(InvocationOnMock invocation) + throws Throwable { + return ApiFutures.immediateFuture( + CommitResponse.newBuilder() + .setCommitTimestamp( + Timestamp.newBuilder().setSeconds(System.currentTimeMillis() * 1000)) + .build()); } }); DatabaseId db = DatabaseId.of("test", "test", "test"); @@ -277,7 +286,8 @@ private long[] batchDmlException(int status) { .build(); when(session.newTransaction()).thenReturn(transaction); when(session.beginTransactionAsync()) - .thenReturn(ApiFutures.immediateFuture(ByteString.copyFromUtf8(UUID.randomUUID().toString()))); + .thenReturn( + ApiFutures.immediateFuture(ByteString.copyFromUtf8(UUID.randomUUID().toString()))); when(session.getName()).thenReturn(SessionId.of("p", "i", "d", "test").getName()); TransactionRunnerImpl runner = new TransactionRunnerImpl(session, rpc, 10); runner.setSpan(mock(Span.class)); @@ -305,7 +315,8 @@ private long[] batchDmlException(int status) { .thenReturn(response1, response2); CommitResponse commitResponse = CommitResponse.newBuilder().setCommitTimestamp(Timestamp.getDefaultInstance()).build(); - when(rpc.commitAsync(Mockito.any(CommitRequest.class), Mockito.anyMap())).thenReturn(ApiFutures.immediateFuture(commitResponse)); + when(rpc.commitAsync(Mockito.any(CommitRequest.class), Mockito.anyMap())) + .thenReturn(ApiFutures.immediateFuture(commitResponse)); final Statement statement = Statement.of("UPDATE FOO SET BAR=1"); final AtomicInteger numCalls = new AtomicInteger(0); long updateCount[] = 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 1e05dd6e12e..29532c54830 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 @@ -18,7 +18,6 @@ import com.google.cloud.spanner.MockSpannerServiceImpl; import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult; -import com.google.cloud.spanner.ResultSet; import com.google.cloud.spanner.Statement; import com.google.cloud.spanner.admin.database.v1.MockDatabaseAdminImpl; import com.google.cloud.spanner.admin.instance.v1.MockInstanceAdminImpl; @@ -91,9 +90,9 @@ public abstract class AbstractMockServerTest { public static final int UPDATE_COUNT = 1; 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 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; @@ -118,7 +117,8 @@ public static void startStaticServer() throws IOException { 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)); + mockSpanner.putStatementResult( + StatementResult.query(SELECT_RANDOM_STATEMENT, RANDOM_RESULT_SET)); } @AfterClass diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AsyncConnectionApiTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AsyncConnectionApiTest.java index f5e2e9ad318..c1381f1b241 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AsyncConnectionApiTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AsyncConnectionApiTest.java @@ -22,8 +22,6 @@ 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.StatementResult; -import com.google.cloud.spanner.Statement; import com.google.cloud.spanner.connection.ITAbstractSpannerTest.ITConnection; import com.google.common.base.Function; import java.util.concurrent.ExecutorService; @@ -45,46 +43,49 @@ public static void stopExecutor() { @Test public void testSimpleSelectAutocommit() throws Exception { - testSimpleSelect(new Function(){ - @Override - public Void apply(Connection input) { - input.setAutocommit(true); - return null; - } - }); + testSimpleSelect( + new Function() { + @Override + public Void apply(Connection input) { + input.setAutocommit(true); + return null; + } + }); } @Test public void testSimpleSelectReadOnly() throws Exception { - testSimpleSelect(new Function(){ - @Override - public Void apply(Connection input) { - input.setReadOnly(true); - return null; - } - }); + testSimpleSelect( + new Function() { + @Override + public Void apply(Connection input) { + input.setReadOnly(true); + return null; + } + }); } @Test public void testSimpleSelectReadWrite() throws Exception { - testSimpleSelect(new Function(){ - @Override - public Void apply(Connection input) { - return null; - } - }); + testSimpleSelect( + new Function() { + @Override + public Void apply(Connection input) { + return null; + } + }); } - private void testSimpleSelect(Function connectionConfigurator) throws Exception { + private void testSimpleSelect(Function connectionConfigurator) + throws Exception { final AtomicInteger rowCount = new AtomicInteger(); ApiFuture res; try (ITConnection connection = createConnection()) { connectionConfigurator.apply(connection); // Verify that the call is non-blocking. -// mockSpanner.freeze(); - try (AsyncResultSet rs = - connection.executeQueryAsync(SELECT_RANDOM_STATEMENT)) { -// mockSpanner.unfreeze(); + // mockSpanner.freeze(); + try (AsyncResultSet rs = connection.executeQueryAsync(SELECT_RANDOM_STATEMENT)) { + // mockSpanner.unfreeze(); res = rs.setCallback( executor, From 64b8a349a85a9a16404939f3e9384eeb9f9b8eda Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Thu, 11 Jun 2020 19:02:19 +0200 Subject: [PATCH 32/49] chore: remove unused code --- .../cloud/spanner/TransactionRunnerImpl.java | 77 ------------------- 1 file changed, 77 deletions(-) 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 9373edde2aa..ed53fba6d20 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 @@ -138,9 +138,6 @@ public void removeListener(Runnable listener) { @GuardedBy("lock") private volatile int runningAsyncOperations; - // @GuardedBy("lock") - // private volatile CountDownLatch finishedAsyncOperations = new CountDownLatch(0); - @GuardedBy("lock") private List mutations = new ArrayList<>(); @@ -178,35 +175,6 @@ private void decreaseAsyncOperations() { } } - void ensureTxn_old() { - if (transactionId == null || isAborted()) { - span.addAnnotation("Creating Transaction"); - try { - transactionId = session.beginTransaction(); - span.addAnnotation( - "Transaction Creation Done", - ImmutableMap.of( - "Id", AttributeValue.stringAttributeValue(transactionId.toStringUtf8()))); - txnLogger.log( - Level.FINER, - "Started transaction {0}", - txnLogger.isLoggable(Level.FINER) ? transactionId.asReadOnlyByteBuffer() : null); - } catch (SpannerException e) { - span.addAnnotation("Transaction Creation Failed", TraceUtil.getExceptionAnnotations(e)); - throw e; - } - } else { - span.addAnnotation( - "Transaction Initialized", - ImmutableMap.of( - "Id", AttributeValue.stringAttributeValue(transactionId.toStringUtf8()))); - txnLogger.log( - Level.FINER, - "Using prepared transaction {0}", - txnLogger.isLoggable(Level.FINER) ? transactionId.asReadOnlyByteBuffer() : null); - } - } - void ensureTxn() { try { ensureTxnAsync().get(); @@ -264,51 +232,6 @@ public void run() { return res; } - void commit_old() { - SettableApiFuture latch; - synchronized (lock) { - latch = finishedAsyncOperations; - } - try { - latch.get(); - } catch (InterruptedException e) { - throw SpannerExceptionFactory.propagateInterrupt(e); - } catch (ExecutionException e) { - throw SpannerExceptionFactory.newSpannerException(e.getCause() == null ? e : e.getCause()); - } - span.addAnnotation("Starting Commit"); - CommitRequest.Builder builder = - CommitRequest.newBuilder().setSession(session.getName()).setTransactionId(transactionId); - synchronized (lock) { - if (!mutations.isEmpty()) { - List mutationsProto = new ArrayList<>(); - Mutation.toProto(mutations, mutationsProto); - builder.addAllMutations(mutationsProto); - } - // Ensure that no call to buffer mutations that would be lost can succeed. - mutations = null; - } - final CommitRequest commitRequest = builder.build(); - Span opSpan = tracer.spanBuilderWithExplicitParent(SpannerImpl.COMMIT, span).startSpan(); - try (Scope s = tracer.withSpan(opSpan)) { - CommitResponse commitResponse = rpc.commit(commitRequest, session.getOptions()); - if (!commitResponse.hasCommitTimestamp()) { - throw newSpannerException( - ErrorCode.INTERNAL, "Missing commitTimestamp:\n" + session.getName()); - } - commitTimestamp = Timestamp.fromProto(commitResponse.getCommitTimestamp()); - opSpan.end(TraceUtil.END_SPAN_OPTIONS); - } catch (RuntimeException e) { - span.addAnnotation("Commit Failed", TraceUtil.getExceptionAnnotations(e)); - TraceUtil.endSpanWithFailure(opSpan, e); - if (e instanceof SpannerException) { - onError((SpannerException) e); - } - throw e; - } - span.addAnnotation("Commit Done"); - } - void commit() { try { commitAsync().get(); From 3354344a4d1ce961e076b0ecddc036323da5e85b Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Thu, 11 Jun 2020 20:47:36 +0200 Subject: [PATCH 33/49] clirr: add ignored differences to clirr --- .../clirr-ignored-differences.xml | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/google-cloud-spanner/clirr-ignored-differences.xml b/google-cloud-spanner/clirr-ignored-differences.xml index baf7670e501..f0c67e02c1a 100644 --- a/google-cloud-spanner/clirr-ignored-differences.xml +++ b/google-cloud-spanner/clirr-ignored-differences.xml @@ -207,11 +207,26 @@ com/google/cloud/spanner/spi/v1/SpannerRpc * beginTransactionAsync(*) + + 7012 + com/google/cloud/spanner/spi/v1/SpannerRpc + * commitAsync(*) + + + 7012 + com/google/cloud/spanner/spi/v1/SpannerRpc + * rollbackAsync(*) + 7012 com/google/cloud/spanner/spi/v1/SpannerRpc * executeBatchDmlAsync(*) + + 7012 + com/google/cloud/spanner/connection/Connection + * executeQueryAsync(*) + From 1fada1c780452a251507048b6580d168cb84b9a0 Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Fri, 12 Jun 2020 11:17:21 +0200 Subject: [PATCH 34/49] fix: call listeners after all rows have been consumed --- .../cloud/spanner/AbstractReadContext.java | 27 ++++++++++++++++--- .../cloud/spanner/AsyncResultSetImpl.java | 9 +++---- .../com/google/cloud/spanner/SessionPool.java | 19 +++---------- .../cloud/spanner/DatabaseClientImplTest.java | 22 ++++++++------- 4 files changed, 44 insertions(+), 33 deletions(-) diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java index b1d752e6f41..f4e6596fdf5 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java @@ -51,6 +51,7 @@ import io.opencensus.trace.Span; import io.opencensus.trace.Tracing; import java.util.Map; +import java.util.concurrent.ExecutionException; import java.util.concurrent.atomic.AtomicLong; import javax.annotation.Nullable; import javax.annotation.concurrent.GuardedBy; @@ -753,18 +754,36 @@ private Struct consumeSingleRow(ResultSet resultSet) { return row; } - private ApiFuture consumeSingleRowAsync(AsyncResultSet resultSet) { - SettableApiFuture result = SettableApiFuture.create(); + static ApiFuture consumeSingleRowAsync(AsyncResultSet resultSet) { + final SettableApiFuture result = SettableApiFuture.create(); // We can safely use a directExecutor here, as we will only be consuming one row, and we will // not be doing any blocking stuff in the handler. - resultSet.setCallback(MoreExecutors.directExecutor(), ConsumeSingleRowCallback.create(result)); + final SettableApiFuture row = SettableApiFuture.create(); + resultSet + .setCallback(MoreExecutors.directExecutor(), ConsumeSingleRowCallback.create(row)) + .addListener( + new Runnable() { + @Override + public void run() { + try { + result.set(row.get()); + } catch (ExecutionException e) { + result.setException( + SpannerExceptionFactory.newSpannerException( + e.getCause() == null ? e : e.getCause())); + } catch (InterruptedException e) { + result.setException(SpannerExceptionFactory.propagateInterrupt(e)); + } + } + }, + MoreExecutors.directExecutor()); return result; } /** * {@link ReadyCallback} for returning the first row in a result set as a future {@link Struct}. */ - static class ConsumeSingleRowCallback implements ReadyCallback { + private static class ConsumeSingleRowCallback implements ReadyCallback { private final SettableApiFuture result; private Struct row; 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 a92026536b6..a86af434342 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 @@ -361,11 +361,7 @@ public Void call() throws Exception { try { delegateResultSet.close(); } catch (Throwable t) { - log.log(Level.INFO, "Ignoring error from closing delegate result set", t); - } finally { - for (Runnable listener : listeners) { - listener.run(); - } + log.log(Level.FINE, "Ignoring error from closing delegate result set", t); } // Ensure that the callback has been called at least once, even if the result set was @@ -388,6 +384,9 @@ public Void call() throws Exception { if (executorProvider.shouldAutoClose()) { service.shutdown(); } + for (Runnable listener : listeners) { + listener.run(); + } synchronized (monitor) { if (executionException != null) { throw executionException; diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java index 276915a4b78..ed854b964f6 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java @@ -45,12 +45,9 @@ import com.google.cloud.Timestamp; import com.google.cloud.grpc.GrpcTransportOptions; import com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory; -import com.google.cloud.spanner.AbstractReadContext.ConsumeSingleRowCallback; -import com.google.cloud.spanner.AsyncResultSet.ReadyCallback; import com.google.cloud.spanner.Options.QueryOption; import com.google.cloud.spanner.Options.ReadOption; import com.google.cloud.spanner.SessionClient.SessionConsumer; -import com.google.cloud.spanner.SessionPool.PooledSession; import com.google.cloud.spanner.SpannerException.ResourceNotFoundException; import com.google.cloud.spanner.SpannerImpl.ClosedException; import com.google.cloud.spanner.TransactionManager.TransactionState; @@ -432,11 +429,9 @@ public Struct readRow(String table, Key key, Iterable columns) { @Override public ApiFuture readRowAsync(String table, Key key, Iterable columns) { - SettableApiFuture result = SettableApiFuture.create(); try (AsyncResultSet rs = readAsync(table, KeySet.singleKey(key), columns)) { - rs.setCallback(MoreExecutors.directExecutor(), ConsumeSingleRowCallback.create(result)); + return AbstractReadContext.consumeSingleRowAsync(rs); } - return result; } @Override @@ -466,11 +461,9 @@ public Struct readRowUsingIndex(String table, String index, Key key, Iterable readRowUsingIndexAsync( String table, String index, Key key, Iterable columns) { - SettableApiFuture result = SettableApiFuture.create(); try (AsyncResultSet rs = readUsingIndexAsync(table, index, KeySet.singleKey(key), columns)) { - rs.setCallback(MoreExecutors.directExecutor(), ConsumeSingleRowCallback.create(result)); + return AbstractReadContext.consumeSingleRowAsync(rs); } - return result; } @Override @@ -626,11 +619,9 @@ public Struct readRow(String table, Key key, Iterable columns) { @Override public ApiFuture readRowAsync(String table, Key key, Iterable columns) { - SettableApiFuture result = SettableApiFuture.create(); try (AsyncResultSet rs = readAsync(table, KeySet.singleKey(key), columns)) { - rs.setCallback(MoreExecutors.directExecutor(), ConsumeSingleRowCallback.create(result)); + return AbstractReadContext.consumeSingleRowAsync(rs); } - return result; } @Override @@ -651,12 +642,10 @@ public Struct readRowUsingIndex( @Override public ApiFuture readRowUsingIndexAsync( String table, String index, Key key, Iterable columns) { - SettableApiFuture result = SettableApiFuture.create(); try (AsyncResultSet rs = readUsingIndexAsync(table, index, KeySet.singleKey(key), columns)) { - rs.setCallback(MoreExecutors.directExecutor(), ConsumeSingleRowCallback.create(result)); + return AbstractReadContext.consumeSingleRowAsync(rs); } - return result; } @Override diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java index cd5e15f5fc9..26e86c289a6 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java @@ -558,16 +558,20 @@ public void transactionManagerExecuteQueryAsync() throws Exception { 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; + try { + while (true) { + switch (resultSet.tryNext()) { + case OK: + rowCount.incrementAndGet(); + break; + case DONE: + return CallbackResponse.DONE; + case NOT_READY: + return CallbackResponse.CONTINUE; + } } + } catch (Throwable t) { + return CallbackResponse.DONE; } } }); From 35743e6a8dc2949c9fd2d37a57233aa7a9b05759 Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Fri, 12 Jun 2020 22:43:51 +0200 Subject: [PATCH 35/49] feat: towards AsyncTransactionManager --- .../spanner/AsyncTransactionManager.java | 18 +- .../spanner/AsyncTransactionManagerImpl.java | 73 +- .../com/google/cloud/spanner/SessionPool.java | 117 +-- .../cloud/spanner/TransactionRunnerImpl.java | 56 +- .../spanner/AbstractAsyncTransactionTest.java | 135 ++++ .../google/cloud/spanner/AsyncRunnerTest.java | 100 +-- .../spanner/AsyncTransactionManagerTest.java | 724 ++++++++++++++++++ .../google/cloud/spanner/ReadAsyncTest.java | 47 +- 8 files changed, 1034 insertions(+), 236 deletions(-) create mode 100644 google-cloud-spanner/src/test/java/com/google/cloud/spanner/AbstractAsyncTransactionTest.java create mode 100644 google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncTransactionManagerTest.java diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManager.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManager.java index e5b5b3bb036..8e8f808d0b8 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManager.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManager.java @@ -37,19 +37,19 @@ * * @see DatabaseClient#transactionManager() */ -public interface AsyncTransactionManager extends AutoCloseable { +public interface AsyncTransactionManager { /** * Creates a new read write transaction. This must be called before doing any other operation and * can only be called once. To create a new transaction for subsequent retries, see {@link * #resetForRetry()}. */ - ApiFuture beginAsync(); + ApiFuture beginAsync(); /** * Commits the currently active transaction. If the transaction was already aborted, then this * would throw an {@link AbortedException}. */ - ApiFuture commitAsync(); + ApiFuture commitAsync(); /** * Rolls back the currently active transaction. In most cases there should be no need to call this @@ -64,13 +64,7 @@ public interface AsyncTransactionManager extends AutoCloseable { * specified by {@link SpannerException#getRetryDelayInMillis()} on the {@code SpannerException} * throw by the previous commit call. */ - ApiFuture resetForRetryAsync(); - - /** - * Returns the commit timestamp if the transaction committed successfully otherwise it will throw - * {@code IllegalStateException}. - */ - ApiFuture getCommitTimestampAsync(); + ApiFuture resetForRetryAsync(); /** Returns the state of the transaction. */ TransactionState getState(); @@ -79,6 +73,6 @@ public interface AsyncTransactionManager extends AutoCloseable { * Closes the manager. If there is an active transaction, it will be rolled back. Underlying * session will be released back to the session pool. */ - @Override - void close(); + // @Override + // void close(); } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManagerImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManagerImpl.java index 7c7f8504c74..b830a8a61b3 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManagerImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManagerImpl.java @@ -17,6 +17,7 @@ package com.google.cloud.spanner; import com.google.api.core.ApiFuture; +import com.google.api.core.ApiFutureCallback; import com.google.api.core.ApiFutures; import com.google.api.core.SettableApiFuture; import com.google.cloud.Timestamp; @@ -39,6 +40,7 @@ final class AsyncTransactionManagerImpl implements AsyncTransactionManager, Sess private TransactionRunnerImpl.TransactionContextImpl txn; private TransactionState txnState; + private final SettableApiFuture commitTimestamp = SettableApiFuture.create(); AsyncTransactionManagerImpl(SessionImpl session, Span span) { this.session = session; @@ -51,7 +53,7 @@ public void setSpan(Span span) { } @Override - public ApiFuture beginAsync() { + public ApiFuture beginAsync() { Preconditions.checkState(txn == null, "begin can only be called once"); txnState = TransactionState.STARTED; txn = session.newTransaction(); @@ -79,28 +81,38 @@ public void run() { } @Override - public ApiFuture commitAsync() { + public ApiFuture commitAsync() { Preconditions.checkState( txnState == TransactionState.STARTED, "commit can only be invoked if the transaction is in progress"); - SettableApiFuture res = SettableApiFuture.create(); if (txn.isAborted()) { txnState = TransactionState.ABORTED; - res.setException( + return ApiFutures.immediateFailedFuture( SpannerExceptionFactory.newSpannerException( ErrorCode.ABORTED, "Transaction already aborted")); } - try { - txn.commit(); - txnState = TransactionState.COMMITTED; - return ApiFutures.immediateFuture(null); - } catch (AbortedException e1) { - txnState = TransactionState.ABORTED; - return ApiFutures.immediateFailedFuture(e1); - } catch (SpannerException e2) { - txnState = TransactionState.COMMIT_FAILED; - return ApiFutures.immediateFailedFuture(e2); - } + ApiFuture res = txn.commitAsync(); + txnState = TransactionState.COMMITTED; + ApiFutures.addCallback( + res, + new ApiFutureCallback() { + @Override + public void onFailure(Throwable t) { + if (t instanceof AbortedException) { + txnState = TransactionState.ABORTED; + } else { + txnState = TransactionState.COMMIT_FAILED; + commitTimestamp.setException(t); + } + } + + @Override + public void onSuccess(Timestamp result) { + commitTimestamp.set(result); + } + }, + MoreExecutors.directExecutor()); + return res; } @Override @@ -109,15 +121,14 @@ public ApiFuture rollbackAsync() { txnState == TransactionState.STARTED, "rollback can only be called if the transaction is in progress"); try { - txn.rollback(); + return txn.rollbackAsync(); } finally { txnState = TransactionState.ROLLED_BACK; } - return ApiFutures.immediateFuture(null); } @Override - public ApiFuture resetForRetryAsync() { + public ApiFuture resetForRetryAsync() { if (txn == null || !txn.isAborted() && txnState != TransactionState.ABORTED) { throw new IllegalStateException( "resetForRetry can only be called if the previous attempt" + " aborted"); @@ -126,27 +137,7 @@ public ApiFuture resetForRetryAsync() { txn = session.newTransaction(); txn.ensureTxn(); txnState = TransactionState.STARTED; - return ApiFutures.immediateFuture(txn); - } - } - - @Override - public ApiFuture getCommitTimestampAsync() { - Preconditions.checkState( - txnState == TransactionState.COMMITTED, - "getCommitTimestamp can only be invoked if the transaction committed successfully"); - return ApiFutures.immediateFuture(txn.commitTimestamp()); - } - - @Override - public void close() { - try { - if (txnState == TransactionState.STARTED && !txn.isAborted()) { - txn.rollback(); - txnState = TransactionState.ROLLED_BACK; - } - } finally { - span.end(TraceUtil.END_SPAN_OPTIONS); + return ApiFutures.immediateFuture((TransactionContext) txn); } } @@ -157,6 +148,8 @@ public TransactionState getState() { @Override public void invalidate() { - close(); + if (txnState == TransactionState.STARTED || txnState == null) { + txnState = TransactionState.ROLLED_BACK; + } } } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java index ed854b964f6..6126a077bfc 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java @@ -38,6 +38,7 @@ import static com.google.cloud.spanner.MetricRegistryConstants.SPANNER_LABEL_KEYS_WITH_TYPE; import static com.google.cloud.spanner.SpannerExceptionFactory.newSpannerException; +import com.google.api.core.ApiAsyncFunction; import com.google.api.core.ApiFuture; import com.google.api.core.ApiFutures; import com.google.api.core.SettableApiFuture; @@ -303,9 +304,12 @@ private boolean internalNext() { @Override public void close() { - super.close(); - if (isSingleUse) { - AutoClosingReadContext.this.close(); + try { + super.close(); + } finally { + if (isSingleUse) { + AutoClosingReadContext.this.close(); + } } } }; @@ -945,14 +949,11 @@ public ApiFuture getCommitTimestamp() { } private static class SessionPoolAsyncTransactionManager implements AsyncTransactionManager { - private final SessionPool sessionPool; private volatile PooledSessionFuture session; - private final SettableApiFuture commitTimestamp = SettableApiFuture.create(); private final SettableApiFuture delegate = SettableApiFuture.create(); + private volatile ApiFuture commitTimestamp; - private SessionPoolAsyncTransactionManager( - SessionPool sessionPool, PooledSessionFuture session) { - this.sessionPool = sessionPool; + private SessionPoolAsyncTransactionManager(PooledSessionFuture session) { this.session = session; this.session.addListener( new Runnable() { @@ -974,64 +975,86 @@ public void run() { @Override public ApiFuture beginAsync() { - final SettableApiFuture res = SettableApiFuture.create(); - delegate.addListener( - new Runnable() { + return ApiFutures.transformAsync( + delegate, + new ApiAsyncFunction() { @Override - public void run() { - try { - res.set(delegate.get().beginAsync().get()); - } catch (Throwable t) { - res.setException(t); - } + public ApiFuture apply(AsyncTransactionManager input) + throws Exception { + return input.beginAsync(); } }, MoreExecutors.directExecutor()); - return res; } @Override - public ApiFuture commitAsync() { - // TODO Auto-generated method stub - return null; + public ApiFuture commitAsync() { + return ApiFutures.transformAsync( + delegate, + new ApiAsyncFunction() { + @Override + public ApiFuture apply(AsyncTransactionManager input) throws Exception { + ApiFuture res = input.commitAsync(); + res.addListener( + new Runnable() { + @Override + public void run() { + session.close(); + } + }, + MoreExecutors.directExecutor()); + return res; + } + }, + MoreExecutors.directExecutor()); } @Override public ApiFuture rollbackAsync() { - // TODO Auto-generated method stub - return null; + return ApiFutures.transformAsync( + delegate, + new ApiAsyncFunction() { + @Override + public ApiFuture apply(AsyncTransactionManager input) throws Exception { + ApiFuture res = input.rollbackAsync(); + res.addListener( + new Runnable() { + @Override + public void run() { + session.close(); + } + }, + MoreExecutors.directExecutor()); + return res; + } + }, + MoreExecutors.directExecutor()); } @Override public ApiFuture resetForRetryAsync() { - // TODO Auto-generated method stub - return null; + return ApiFutures.transformAsync( + delegate, + new ApiAsyncFunction() { + @Override + public ApiFuture apply(AsyncTransactionManager input) + throws Exception { + return input.resetForRetryAsync(); + } + }, + MoreExecutors.directExecutor()); } @Override public TransactionState getState() { - // TODO Auto-generated method stub - return null; - } - - @Override - public void close() { - // TODO Auto-generated method stub - - } - - private void setCommitTimestamp(AsyncTransactionManager delegate) { try { - commitTimestamp.set(delegate.getCommitTimestampAsync().get()); - } catch (Throwable t) { - commitTimestamp.setException(t); + return delegate.get().getState(); + } catch (InterruptedException e) { + throw SpannerExceptionFactory.propagateInterrupt(e); + } catch (ExecutionException e) { + throw SpannerExceptionFactory.newSpannerException(e.getCause() == null ? e : e.getCause()); } } - - @Override - public ApiFuture getCommitTimestampAsync() { - return commitTimestamp; - } } // Exception class used just to track the stack trace at the point when a session was handed out @@ -1312,7 +1335,7 @@ public AsyncRunner runAsync() { @Override public AsyncTransactionManager transactionManagerAsync() { - return new SessionPoolAsyncTransactionManager(SessionPool.this, this); + return new SessionPoolAsyncTransactionManager(this); } @Override @@ -1362,9 +1385,9 @@ public PooledSession get() { // ignore the exception as it will be handled by the call to super.get() below. } if (res != null) { + res.markBusy(span); + span.addAnnotation(sessionAnnotation(res)); synchronized (lock) { - res.markBusy(span); - span.addAnnotation(sessionAnnotation(res)); incrementNumSessionsInUse(); checkedOutSessions.add(this); } 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 ed53fba6d20..36565fb0f8d 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 @@ -21,6 +21,7 @@ import static com.google.common.base.Preconditions.checkNotNull; import static com.google.common.base.Preconditions.checkState; +import com.google.api.core.ApiAsyncFunction; import com.google.api.core.ApiFunction; import com.google.api.core.ApiFuture; import com.google.api.core.ApiFutures; @@ -34,6 +35,7 @@ import com.google.common.collect.ImmutableMap; import com.google.common.util.concurrent.MoreExecutors; import com.google.protobuf.ByteString; +import com.google.protobuf.Empty; import com.google.rpc.Code; import com.google.spanner.v1.CommitRequest; import com.google.spanner.v1.CommitResponse; @@ -234,7 +236,7 @@ public void run() { void commit() { try { - commitAsync().get(); + commitTimestamp = commitAsync().get(); } catch (InterruptedException e) { throw SpannerExceptionFactory.propagateInterrupt(e); } catch (ExecutionException e) { @@ -242,8 +244,8 @@ void commit() { } } - ApiFuture commitAsync() { - final SettableApiFuture res = SettableApiFuture.create(); + ApiFuture commitAsync() { + final SettableApiFuture res = SettableApiFuture.create(); CommitRequest.Builder builder = CommitRequest.newBuilder().setSession(session.getName()).setTransactionId(transactionId); synchronized (lock) { @@ -284,11 +286,11 @@ public void run() { ErrorCode.INTERNAL, "Missing commitTimestamp:\n" + session.getName()); } - commitTimestamp = + Timestamp ts = Timestamp.fromProto(commitResponse.getCommitTimestamp()); span.addAnnotation("Commit Done"); opSpan.end(TraceUtil.END_SPAN_OPTIONS); - res.set(null); + res.set(ts); } catch (Throwable e) { if (e instanceof ExecutionException) { e = @@ -356,6 +358,25 @@ void rollback() { } } + ApiFuture rollbackAsync() { + span.addAnnotation("Starting Rollback"); + return ApiFutures.transformAsync( + rpc.rollbackAsync( + RollbackRequest.newBuilder() + .setSession(session.getName()) + .setTransactionId(transactionId) + .build(), + session.getOptions()), + new ApiAsyncFunction() { + @Override + public ApiFuture apply(Empty input) throws Exception { + span.addAnnotation("Rollback Done"); + return ApiFutures.immediateFuture(null); + } + }, + MoreExecutors.directExecutor()); + } + @Nullable @Override TransactionSelector getTransactionSelector() { @@ -433,7 +454,7 @@ public ApiFuture executeUpdateAsync(Statement statement) { decreaseAsyncOperations(); throw t; } - final ApiFuture updateCount = + ApiFuture updateCount = ApiFutures.transform( resultSet, new ApiFunction() { @@ -449,19 +470,24 @@ public Long apply(ResultSet input) { } }, MoreExecutors.directExecutor()); + updateCount = + ApiFutures.catching( + updateCount, + Throwable.class, + new ApiFunction() { + @Override + public Long apply(Throwable input) { + SpannerException e = SpannerExceptionFactory.newSpannerException(input); + onError(e); + throw e; + } + }, + MoreExecutors.directExecutor()); updateCount.addListener( new Runnable() { @Override public void run() { - try { - updateCount.get(); - } catch (ExecutionException e) { - onError(SpannerExceptionFactory.newSpannerException(e.getCause())); - } catch (InterruptedException e) { - onError(SpannerExceptionFactory.propagateInterrupt(e)); - } finally { - decreaseAsyncOperations(); - } + decreaseAsyncOperations(); } }, MoreExecutors.directExecutor()); diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AbstractAsyncTransactionTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AbstractAsyncTransactionTest.java new file mode 100644 index 00000000000..fbd8e44ffef --- /dev/null +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AbstractAsyncTransactionTest.java @@ -0,0 +1,135 @@ +/* + * 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.MockSpannerTestUtil.EMPTY_KEY_VALUE_RESULTSET; +import static com.google.cloud.spanner.MockSpannerTestUtil.INVALID_UPDATE_STATEMENT; +import static com.google.cloud.spanner.MockSpannerTestUtil.READ_MULTIPLE_KEY_VALUE_RESULTSET; +import static com.google.cloud.spanner.MockSpannerTestUtil.READ_MULTIPLE_KEY_VALUE_STATEMENT; +import static com.google.cloud.spanner.MockSpannerTestUtil.READ_ONE_EMPTY_KEY_VALUE_STATEMENT; +import static com.google.cloud.spanner.MockSpannerTestUtil.READ_ONE_KEY_VALUE_RESULTSET; +import static com.google.cloud.spanner.MockSpannerTestUtil.READ_ONE_KEY_VALUE_STATEMENT; +import static com.google.cloud.spanner.MockSpannerTestUtil.TEST_DATABASE; +import static com.google.cloud.spanner.MockSpannerTestUtil.TEST_INSTANCE; +import static com.google.cloud.spanner.MockSpannerTestUtil.TEST_PROJECT; +import static com.google.cloud.spanner.MockSpannerTestUtil.UPDATE_ABORTED_STATEMENT; +import static com.google.cloud.spanner.MockSpannerTestUtil.UPDATE_COUNT; +import static com.google.cloud.spanner.MockSpannerTestUtil.UPDATE_STATEMENT; + +import com.google.api.gax.grpc.testing.LocalChannelProvider; +import com.google.cloud.NoCredentials; +import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult; +import io.grpc.Server; +import io.grpc.Status; +import io.grpc.inprocess.InProcessServerBuilder; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledThreadPoolExecutor; +import org.junit.After; +import org.junit.AfterClass; +import org.junit.Before; +import org.junit.BeforeClass; + +/** Base class for {@link AsyncRunnerTest} and {@link AsyncTransactionManagerTest}. */ +public abstract class AbstractAsyncTransactionTest { + static MockSpannerServiceImpl mockSpanner; + private static Server server; + private static LocalChannelProvider channelProvider; + static ExecutorService executor; + + Spanner spanner; + Spanner spannerWithEmptySessionPool; + + @BeforeClass + public static void setup() throws Exception { + mockSpanner = new MockSpannerServiceImpl(); + mockSpanner.setAbortProbability(0.0D); + mockSpanner.putStatementResult( + StatementResult.query(READ_ONE_EMPTY_KEY_VALUE_STATEMENT, EMPTY_KEY_VALUE_RESULTSET)); + mockSpanner.putStatementResult( + StatementResult.query(READ_ONE_KEY_VALUE_STATEMENT, READ_ONE_KEY_VALUE_RESULTSET)); + mockSpanner.putStatementResult( + StatementResult.query( + READ_MULTIPLE_KEY_VALUE_STATEMENT, READ_MULTIPLE_KEY_VALUE_RESULTSET)); + mockSpanner.putStatementResult(StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT)); + mockSpanner.putStatementResult( + StatementResult.exception( + INVALID_UPDATE_STATEMENT, + Status.INVALID_ARGUMENT.withDescription("invalid statement").asRuntimeException())); + mockSpanner.putStatementResult( + StatementResult.exception( + UPDATE_ABORTED_STATEMENT, + Status.ABORTED.withDescription("Transaction was aborted").asRuntimeException())); + String uniqueName = InProcessServerBuilder.generateName(); + server = + InProcessServerBuilder.forName(uniqueName) + .scheduledExecutorService(new ScheduledThreadPoolExecutor(1)) + .addService(mockSpanner) + .build() + .start(); + channelProvider = LocalChannelProvider.create(uniqueName); + executor = Executors.newSingleThreadExecutor(); + } + + @AfterClass + public static void teardown() throws Exception { + executor.shutdown(); + server.shutdown(); + server.awaitTermination(); + } + + @Before + public void before() { + spanner = + SpannerOptions.newBuilder() + .setProjectId(TEST_PROJECT) + .setChannelProvider(channelProvider) + .setCredentials(NoCredentials.getInstance()) + .setSessionPoolOption(SessionPoolOptions.newBuilder().setFailOnSessionLeak().build()) + .build() + .getService(); + spannerWithEmptySessionPool = + spanner + .getOptions() + .toBuilder() + .setSessionPoolOption( + SessionPoolOptions.newBuilder() + .setFailOnSessionLeak() + .setMinSessions(0) + .setIncStep(1) + .build()) + .build() + .getService(); + } + + DatabaseClient client() { + return spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); + } + + DatabaseClient clientWithEmptySessionPool() { + return spannerWithEmptySessionPool.getDatabaseClient( + DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); + } + + @After + public void after() { + spanner.close(); + spannerWithEmptySessionPool.close(); + mockSpanner.removeAllExecutionTimes(); + mockSpanner.reset(); + } +} diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncRunnerTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncRunnerTest.java index a3ba51d0004..eb00047ca47 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncRunnerTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncRunnerTest.java @@ -24,8 +24,6 @@ import com.google.api.core.ApiFuture; import com.google.api.core.ApiFutures; import com.google.api.core.SettableApiFuture; -import com.google.api.gax.grpc.testing.LocalChannelProvider; -import com.google.cloud.NoCredentials; import com.google.cloud.Timestamp; import com.google.cloud.spanner.AsyncResultSet.CallbackResponse; import com.google.cloud.spanner.AsyncResultSet.ReadyCallback; @@ -40,117 +38,21 @@ import com.google.spanner.v1.CommitRequest; import com.google.spanner.v1.ExecuteBatchDmlRequest; import com.google.spanner.v1.ExecuteSqlRequest; -import io.grpc.Server; import io.grpc.Status; -import io.grpc.inprocess.InProcessServerBuilder; import java.util.ArrayList; import java.util.List; import java.util.concurrent.BlockingQueue; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; -import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; -import java.util.concurrent.ScheduledThreadPoolExecutor; import java.util.concurrent.SynchronousQueue; import java.util.concurrent.atomic.AtomicInteger; -import org.junit.After; -import org.junit.AfterClass; -import org.junit.Before; -import org.junit.BeforeClass; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @RunWith(JUnit4.class) -public class AsyncRunnerTest { - - private static MockSpannerServiceImpl mockSpanner; - private static Server server; - private static LocalChannelProvider channelProvider; - private static ExecutorService executor; - - private Spanner spanner; - private Spanner spannerWithEmptySessionPool; - - @BeforeClass - public static void setup() throws Exception { - mockSpanner = new MockSpannerServiceImpl(); - mockSpanner.setAbortProbability(0.0D); - mockSpanner.putStatementResult( - StatementResult.query(READ_ONE_EMPTY_KEY_VALUE_STATEMENT, EMPTY_KEY_VALUE_RESULTSET)); - mockSpanner.putStatementResult( - StatementResult.query(READ_ONE_KEY_VALUE_STATEMENT, READ_ONE_KEY_VALUE_RESULTSET)); - mockSpanner.putStatementResult( - StatementResult.query( - READ_MULTIPLE_KEY_VALUE_STATEMENT, READ_MULTIPLE_KEY_VALUE_RESULTSET)); - mockSpanner.putStatementResult(StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT)); - mockSpanner.putStatementResult( - StatementResult.exception( - INVALID_UPDATE_STATEMENT, - Status.INVALID_ARGUMENT.withDescription("invalid statement").asRuntimeException())); - mockSpanner.putStatementResult( - StatementResult.exception( - UPDATE_ABORTED_STATEMENT, - Status.ABORTED.withDescription("Transaction was aborted").asRuntimeException())); - String uniqueName = InProcessServerBuilder.generateName(); - server = - InProcessServerBuilder.forName(uniqueName) - .scheduledExecutorService(new ScheduledThreadPoolExecutor(1)) - .addService(mockSpanner) - .build() - .start(); - channelProvider = LocalChannelProvider.create(uniqueName); - executor = Executors.newSingleThreadExecutor(); - } - - @AfterClass - public static void teardown() throws Exception { - executor.shutdown(); - server.shutdown(); - server.awaitTermination(); - } - - @Before - public void before() { - spanner = - SpannerOptions.newBuilder() - .setProjectId(TEST_PROJECT) - .setChannelProvider(channelProvider) - .setCredentials(NoCredentials.getInstance()) - .setSessionPoolOption(SessionPoolOptions.newBuilder().setFailOnSessionLeak().build()) - .build() - .getService(); - spannerWithEmptySessionPool = - spanner - .getOptions() - .toBuilder() - .setSessionPoolOption( - SessionPoolOptions.newBuilder() - .setFailOnSessionLeak() - .setMinSessions(0) - .setIncStep(1) - .build()) - .build() - .getService(); - } - - private DatabaseClient client() { - return spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); - } - - private DatabaseClient clientWithEmptySessionPool() { - return spannerWithEmptySessionPool.getDatabaseClient( - DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); - } - - @After - public void after() { - spanner.close(); - spannerWithEmptySessionPool.close(); - mockSpanner.removeAllExecutionTimes(); - mockSpanner.reset(); - } - +public class AsyncRunnerTest extends AbstractAsyncTransactionTest { @Test public void asyncRunnerUpdate() throws Exception { AsyncRunner runner = client().runAsync(); diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncTransactionManagerTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncTransactionManagerTest.java new file mode 100644 index 00000000000..8b6772e43b9 --- /dev/null +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncTransactionManagerTest.java @@ -0,0 +1,724 @@ +/* + * 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.MockSpannerTestUtil.*; +import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.fail; + +import com.google.api.core.ApiAsyncFunction; +import com.google.api.core.ApiFunction; +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.AsyncResultSet.CallbackResponse; +import com.google.cloud.spanner.AsyncResultSet.ReadyCallback; +import com.google.cloud.spanner.AsyncRunner.AsyncWork; +import com.google.cloud.spanner.MockSpannerServiceImpl.SimulatedExecutionTime; +import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult; +import com.google.common.base.Function; +import com.google.common.collect.ImmutableList; +import com.google.common.util.concurrent.MoreExecutors; +import com.google.spanner.v1.BatchCreateSessionsRequest; +import com.google.spanner.v1.BeginTransactionRequest; +import com.google.spanner.v1.CommitRequest; +import com.google.spanner.v1.ExecuteBatchDmlRequest; +import com.google.spanner.v1.ExecuteSqlRequest; +import io.grpc.Status; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executors; +import java.util.concurrent.SynchronousQueue; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class AsyncTransactionManagerTest extends AbstractAsyncTransactionTest { + + @Test + public void asyncTransactionManagerUpdate() throws Exception { + final SettableApiFuture> commitTimestamp = SettableApiFuture.create(); + final AsyncTransactionManager manager = client().transactionManagerAsync(); + ApiFuture txn = manager.beginAsync(); + ApiFuture updateCount = + ApiFutures.transformAsync( + txn, + new ApiAsyncFunction() { + @Override + public ApiFuture apply(TransactionContext input) throws Exception { + ApiFuture res = input.executeUpdateAsync(UPDATE_STATEMENT); + commitTimestamp.set( + ApiFutures.transformAsync( + res, + new ApiAsyncFunction() { + @Override + public ApiFuture apply(Long input) throws Exception { + return manager.commitAsync(); + } + }, + executor)); + return res; + } + }, + executor); + assertThat(updateCount.get()).isEqualTo(UPDATE_COUNT); + assertThat(commitTimestamp.get().get()).isNotNull(); + } + + @Test + public void asyncTransactionManagerIsNonBlocking() throws Exception { + mockSpanner.freeze(); + final SettableApiFuture> commitTimestamp = SettableApiFuture.create(); + final AsyncTransactionManager manager = client().transactionManagerAsync(); + ApiFuture txn = manager.beginAsync(); + ApiFuture updateCount = + ApiFutures.transformAsync( + txn, + new ApiAsyncFunction() { + @Override + public ApiFuture apply(TransactionContext input) throws Exception { + ApiFuture res = input.executeUpdateAsync(UPDATE_STATEMENT); + commitTimestamp.set( + ApiFutures.transformAsync( + res, + new ApiAsyncFunction() { + @Override + public ApiFuture apply(Long input) throws Exception { + return manager.commitAsync(); + } + }, + executor)); + return res; + } + }, + executor); + mockSpanner.unfreeze(); + assertThat(updateCount.get(10L, TimeUnit.SECONDS)).isEqualTo(UPDATE_COUNT); + assertThat(commitTimestamp.get().get(10L, TimeUnit.SECONDS)).isNotNull(); + } + + @Test + public void asyncTransactionManagerInvalidUpdate() throws Exception { + final AsyncTransactionManager manager = client().transactionManagerAsync(); + final SettableApiFuture rollbacked = SettableApiFuture.create(); + ApiFuture txnFut = manager.beginAsync(); + ApiFuture updateCount = + ApiFutures.transformAsync( + txnFut, + new ApiAsyncFunction() { + @Override + public ApiFuture apply(TransactionContext txn) throws Exception { + ApiFuture res = txn.executeUpdateAsync(INVALID_UPDATE_STATEMENT); + ApiFutures.addCallback( + res, + new ApiFutureCallback() { + @Override + public void onFailure(Throwable t) { + manager + .rollbackAsync() + .addListener( + new Runnable() { + @Override + public void run() { + rollbacked.set(null); + } + }, + executor); + ; + } + + @Override + public void onSuccess(Long result) { + fail("update should not succeed"); + } + }, + MoreExecutors.directExecutor()); + return res; + } + }, + executor); + try { + updateCount.get(); + fail("missing expected exception"); + } catch (ExecutionException e) { + assertThat(e.getCause()).isInstanceOf(SpannerException.class); + SpannerException se = (SpannerException) e.getCause(); + assertThat(se.getErrorCode()).isEqualTo(ErrorCode.INVALID_ARGUMENT); + assertThat(se.getMessage()).contains("invalid statement"); + } + // This is to ensure that the test case does not end before the session has been returned to the + // pool. + rollbacked.get(); + } + + @Test + public void asyncTransactionManagerCommitAborted() throws Exception { + final AtomicInteger attempt = new AtomicInteger(); + final AsyncTransactionManager manager = clientWithEmptySessionPool().transactionManagerAsync(); + ApiFuture txn = manager.beginAsync(); + while (true) { + try { + final SettableApiFuture> commitTimestamp = SettableApiFuture.create(); + ApiFuture updateCount = + ApiFutures.transformAsync( + txn, + new ApiAsyncFunction() { + @Override + public ApiFuture apply(TransactionContext input) throws Exception { + ApiFuture res = input.executeUpdateAsync(UPDATE_STATEMENT); + commitTimestamp.set( + ApiFutures.transformAsync( + res, + new ApiAsyncFunction() { + @Override + public ApiFuture apply(Long input) throws Exception { + if (attempt.incrementAndGet() == 1) { + mockSpanner.abortAllTransactions(); + } + return manager.commitAsync(); + } + }, + executor)); + return res; + } + }, + executor); + assertThat(updateCount.get()).isEqualTo(UPDATE_COUNT); + assertThat(commitTimestamp.get().get()).isNotNull(); + assertThat(attempt.get()).isEqualTo(2); + break; + } catch (ExecutionException e) { + assertThat(e.getCause()).isInstanceOf(AbortedException.class); + assertThat(attempt.get()).isEqualTo(1); + txn = manager.resetForRetryAsync(); + } + } + } + + @Test + public void asyncRunnerFireAndForgetInvalidUpdate() throws Exception { + AsyncRunner runner = client().runAsync(); + ApiFuture res = + runner.runAsync( + new AsyncWork() { + @Override + public ApiFuture doWorkAsync(TransactionContext txn) { + txn.executeUpdateAsync(INVALID_UPDATE_STATEMENT); + return txn.executeUpdateAsync(UPDATE_STATEMENT); + } + }, + executor); + assertThat(res.get()).isEqualTo(UPDATE_COUNT); + } + + @Test + public void asyncRunnerUpdateAborted() throws Exception { + try { + // Temporarily set the result of the update to 2 rows. + mockSpanner.putStatementResult(StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT + 1L)); + final AtomicInteger attempt = new AtomicInteger(); + AsyncRunner runner = client().runAsync(); + ApiFuture updateCount = + runner.runAsync( + new AsyncWork() { + @Override + public ApiFuture doWorkAsync(TransactionContext txn) { + if (attempt.incrementAndGet() == 1) { + mockSpanner.abortTransaction(txn); + } else { + // Set the result of the update statement back to 1 row. + mockSpanner.putStatementResult( + StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT)); + } + return txn.executeUpdateAsync(UPDATE_STATEMENT); + } + }, + executor); + assertThat(updateCount.get()).isEqualTo(UPDATE_COUNT); + assertThat(attempt.get()).isEqualTo(2); + } finally { + mockSpanner.putStatementResult(StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT)); + } + } + + @Test + public void asyncRunnerCommitAborted() throws Exception { + try { + // Temporarily set the result of the update to 2 rows. + mockSpanner.putStatementResult(StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT + 1L)); + final AtomicInteger attempt = new AtomicInteger(); + AsyncRunner runner = client().runAsync(); + ApiFuture updateCount = + runner.runAsync( + new AsyncWork() { + @Override + public ApiFuture doWorkAsync(final TransactionContext txn) { + if (attempt.get() > 0) { + // Set the result of the update statement back to 1 row. + mockSpanner.putStatementResult( + StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT)); + } + ApiFuture updateCount = txn.executeUpdateAsync(UPDATE_STATEMENT); + if (attempt.incrementAndGet() == 1) { + mockSpanner.abortTransaction(txn); + } + return updateCount; + } + }, + executor); + assertThat(updateCount.get()).isEqualTo(UPDATE_COUNT); + assertThat(attempt.get()).isEqualTo(2); + } finally { + mockSpanner.putStatementResult(StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT)); + } + } + + @Test + public void asyncRunnerUpdateAbortedWithoutGettingResult() throws Exception { + final AtomicInteger attempt = new AtomicInteger(); + AsyncRunner runner = clientWithEmptySessionPool().runAsync(); + ApiFuture result = + runner.runAsync( + new AsyncWork() { + @Override + public ApiFuture doWorkAsync(TransactionContext txn) { + if (attempt.incrementAndGet() == 1) { + mockSpanner.abortTransaction(txn); + } + // This update statement will be aborted, but the error will not propagated to the + // transaction runner and cause the transaction to retry. Instead, the commit call + // will do that. + txn.executeUpdateAsync(UPDATE_STATEMENT); + // Resolving this future will not resolve the result of the entire transaction. The + // transaction result will be resolved when the commit has actually finished + // successfully. + return ApiFutures.immediateFuture(null); + } + }, + executor); + assertThat(result.get()).isNull(); + assertThat(attempt.get()).isEqualTo(2); + assertThat(mockSpanner.getRequestTypes()) + .containsExactly( + BatchCreateSessionsRequest.class, + BeginTransactionRequest.class, + ExecuteSqlRequest.class, + CommitRequest.class, + BeginTransactionRequest.class, + ExecuteSqlRequest.class, + CommitRequest.class); + } + + @Test + public void asyncRunnerCommitFails() throws Exception { + mockSpanner.setCommitExecutionTime( + SimulatedExecutionTime.ofException( + Status.RESOURCE_EXHAUSTED + .withDescription("mutation limit exceeded") + .asRuntimeException())); + AsyncRunner runner = client().runAsync(); + ApiFuture updateCount = + runner.runAsync( + new AsyncWork() { + @Override + public ApiFuture doWorkAsync(TransactionContext txn) { + // This statement will succeed, but the commit will fail. The error from the commit + // will bubble up to the future that is returned by the transaction, and the update + // count returned here will never reach the user application. + return txn.executeUpdateAsync(UPDATE_STATEMENT); + } + }, + executor); + try { + updateCount.get(); + fail("missing expected exception"); + } catch (ExecutionException e) { + assertThat(e.getCause()).isInstanceOf(SpannerException.class); + SpannerException se = (SpannerException) e.getCause(); + assertThat(se.getErrorCode()).isEqualTo(ErrorCode.RESOURCE_EXHAUSTED); + assertThat(se.getMessage()).contains("mutation limit exceeded"); + } + } + + @Test + public void asyncRunnerWaitsUntilAsyncUpdateHasFinished() throws Exception { + AsyncRunner runner = clientWithEmptySessionPool().runAsync(); + ApiFuture res = + runner.runAsync( + new AsyncWork() { + @Override + public ApiFuture doWorkAsync(TransactionContext txn) { + txn.executeUpdateAsync(UPDATE_STATEMENT); + return ApiFutures.immediateFuture(null); + } + }, + executor); + res.get(); + assertThat(mockSpanner.getRequestTypes()) + .containsExactly( + BatchCreateSessionsRequest.class, + BeginTransactionRequest.class, + ExecuteSqlRequest.class, + CommitRequest.class); + } + + @Test + public void asyncRunnerBatchUpdate() throws Exception { + AsyncRunner runner = client().runAsync(); + ApiFuture updateCount = + runner.runAsync( + new AsyncWork() { + @Override + public ApiFuture doWorkAsync(TransactionContext txn) { + return txn.batchUpdateAsync(ImmutableList.of(UPDATE_STATEMENT, UPDATE_STATEMENT)); + } + }, + executor); + assertThat(updateCount.get()).asList().containsExactly(UPDATE_COUNT, UPDATE_COUNT); + } + + @Test + public void asyncRunnerIsNonBlockingWithBatchUpdate() throws Exception { + mockSpanner.freeze(); + AsyncRunner runner = clientWithEmptySessionPool().runAsync(); + ApiFuture res = + runner.runAsync( + new AsyncWork() { + @Override + public ApiFuture doWorkAsync(TransactionContext txn) { + txn.batchUpdateAsync(ImmutableList.of(UPDATE_STATEMENT)); + return ApiFutures.immediateFuture(null); + } + }, + executor); + ApiFuture ts = runner.getCommitTimestamp(); + mockSpanner.unfreeze(); + assertThat(res.get()).isNull(); + assertThat(ts.get()).isNotNull(); + } + + @Test + public void asyncRunnerInvalidBatchUpdate() throws Exception { + AsyncRunner runner = client().runAsync(); + ApiFuture updateCount = + runner.runAsync( + new AsyncWork() { + @Override + public ApiFuture doWorkAsync(TransactionContext txn) { + return txn.batchUpdateAsync( + ImmutableList.of(UPDATE_STATEMENT, INVALID_UPDATE_STATEMENT)); + } + }, + executor); + try { + updateCount.get(); + fail("missing expected exception"); + } catch (ExecutionException e) { + assertThat(e.getCause()).isInstanceOf(SpannerException.class); + SpannerException se = (SpannerException) e.getCause(); + assertThat(se.getErrorCode()).isEqualTo(ErrorCode.INVALID_ARGUMENT); + assertThat(se.getMessage()).contains("invalid statement"); + } + } + + @Test + public void asyncRunnerFireAndForgetInvalidBatchUpdate() throws Exception { + AsyncRunner runner = client().runAsync(); + ApiFuture res = + runner.runAsync( + new AsyncWork() { + @Override + public ApiFuture doWorkAsync(TransactionContext txn) { + txn.batchUpdateAsync(ImmutableList.of(UPDATE_STATEMENT, INVALID_UPDATE_STATEMENT)); + return txn.batchUpdateAsync(ImmutableList.of(UPDATE_STATEMENT, UPDATE_STATEMENT)); + } + }, + executor); + assertThat(res.get()).asList().containsExactly(UPDATE_COUNT, UPDATE_COUNT); + } + + @Test + public void asyncRunnerBatchUpdateAborted() throws Exception { + final AtomicInteger attempt = new AtomicInteger(); + AsyncRunner runner = client().runAsync(); + ApiFuture updateCount = + runner.runAsync( + new AsyncWork() { + @Override + public ApiFuture doWorkAsync(TransactionContext txn) { + if (attempt.incrementAndGet() == 1) { + return txn.batchUpdateAsync( + ImmutableList.of(UPDATE_STATEMENT, UPDATE_ABORTED_STATEMENT)); + } else { + return txn.batchUpdateAsync(ImmutableList.of(UPDATE_STATEMENT, UPDATE_STATEMENT)); + } + } + }, + executor); + assertThat(updateCount.get()).asList().containsExactly(UPDATE_COUNT, UPDATE_COUNT); + assertThat(attempt.get()).isEqualTo(2); + } + + @Test + public void asyncRunnerWithBatchUpdateCommitAborted() throws Exception { + try { + // Temporarily set the result of the update to 2 rows. + mockSpanner.putStatementResult(StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT + 1L)); + final AtomicInteger attempt = new AtomicInteger(); + AsyncRunner runner = client().runAsync(); + ApiFuture updateCount = + runner.runAsync( + new AsyncWork() { + @Override + public ApiFuture doWorkAsync(final TransactionContext txn) { + if (attempt.get() > 0) { + // Set the result of the update statement back to 1 row. + mockSpanner.putStatementResult( + StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT)); + } + ApiFuture updateCount = + txn.batchUpdateAsync(ImmutableList.of(UPDATE_STATEMENT, UPDATE_STATEMENT)); + if (attempt.incrementAndGet() == 1) { + mockSpanner.abortTransaction(txn); + } + return updateCount; + } + }, + executor); + assertThat(updateCount.get()).asList().containsExactly(UPDATE_COUNT, UPDATE_COUNT); + assertThat(attempt.get()).isEqualTo(2); + } finally { + mockSpanner.putStatementResult(StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT)); + } + } + + @Test + public void asyncRunnerBatchUpdateAbortedWithoutGettingResult() throws Exception { + final AtomicInteger attempt = new AtomicInteger(); + AsyncRunner runner = clientWithEmptySessionPool().runAsync(); + ApiFuture result = + runner.runAsync( + new AsyncWork() { + @Override + public ApiFuture doWorkAsync(TransactionContext txn) { + if (attempt.incrementAndGet() == 1) { + mockSpanner.abortTransaction(txn); + } + // This update statement will be aborted, but the error will not propagated to the + // transaction runner and cause the transaction to retry. Instead, the commit call + // will do that. + txn.batchUpdateAsync(ImmutableList.of(UPDATE_STATEMENT, UPDATE_STATEMENT)); + // Resolving this future will not resolve the result of the entire transaction. The + // transaction result will be resolved when the commit has actually finished + // successfully. + return ApiFutures.immediateFuture(null); + } + }, + executor); + assertThat(result.get()).isNull(); + assertThat(attempt.get()).isEqualTo(2); + assertThat(mockSpanner.getRequestTypes()) + .containsExactly( + BatchCreateSessionsRequest.class, + BeginTransactionRequest.class, + ExecuteBatchDmlRequest.class, + CommitRequest.class, + BeginTransactionRequest.class, + ExecuteBatchDmlRequest.class, + CommitRequest.class); + } + + @Test + public void asyncRunnerWithBatchUpdateCommitFails() throws Exception { + mockSpanner.setCommitExecutionTime( + SimulatedExecutionTime.ofException( + Status.RESOURCE_EXHAUSTED + .withDescription("mutation limit exceeded") + .asRuntimeException())); + AsyncRunner runner = client().runAsync(); + ApiFuture updateCount = + runner.runAsync( + new AsyncWork() { + @Override + public ApiFuture doWorkAsync(TransactionContext txn) { + // This statement will succeed, but the commit will fail. The error from the commit + // will bubble up to the future that is returned by the transaction, and the update + // count returned here will never reach the user application. + return txn.batchUpdateAsync(ImmutableList.of(UPDATE_STATEMENT, UPDATE_STATEMENT)); + } + }, + executor); + try { + updateCount.get(); + fail("missing expected exception"); + } catch (ExecutionException e) { + assertThat(e.getCause()).isInstanceOf(SpannerException.class); + SpannerException se = (SpannerException) e.getCause(); + assertThat(se.getErrorCode()).isEqualTo(ErrorCode.RESOURCE_EXHAUSTED); + assertThat(se.getMessage()).contains("mutation limit exceeded"); + } + } + + @Test + public void asyncRunnerWaitsUntilAsyncBatchUpdateHasFinished() throws Exception { + AsyncRunner runner = clientWithEmptySessionPool().runAsync(); + ApiFuture res = + runner.runAsync( + new AsyncWork() { + @Override + public ApiFuture doWorkAsync(TransactionContext txn) { + txn.batchUpdateAsync(ImmutableList.of(UPDATE_STATEMENT)); + return ApiFutures.immediateFuture(null); + } + }, + executor); + res.get(); + assertThat(mockSpanner.getRequestTypes()) + .containsExactly( + BatchCreateSessionsRequest.class, + BeginTransactionRequest.class, + ExecuteBatchDmlRequest.class, + CommitRequest.class); + } + + @Test + public void closeTransactionBeforeEndOfAsyncQuery() throws Exception { + final BlockingQueue results = new SynchronousQueue<>(); + final SettableApiFuture finished = SettableApiFuture.create(); + DatabaseClientImpl clientImpl = (DatabaseClientImpl) client(); + + // There should currently not be any sessions checked out of the pool. + assertThat(clientImpl.pool.getNumberOfSessionsInUse()).isEqualTo(0); + + AsyncRunner runner = clientImpl.runAsync(); + final CountDownLatch dataReceived = new CountDownLatch(1); + final CountDownLatch dataChecked = new CountDownLatch(1); + ApiFuture res = + runner.runAsync( + new AsyncWork() { + @Override + public ApiFuture doWorkAsync(TransactionContext txn) { + try (AsyncResultSet rs = + txn.readAsync( + READ_TABLE_NAME, KeySet.all(), READ_COLUMN_NAMES, Options.bufferRows(1))) { + rs.setCallback( + Executors.newSingleThreadExecutor(), + new ReadyCallback() { + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + dataReceived.countDown(); + try { + while (true) { + switch (resultSet.tryNext()) { + case DONE: + finished.set(true); + return CallbackResponse.DONE; + case NOT_READY: + return CallbackResponse.CONTINUE; + case OK: + dataChecked.await(); + results.put(resultSet.getString(0)); + } + } + } catch (Throwable t) { + finished.setException(t); + return CallbackResponse.DONE; + } + } + }); + } + try { + dataReceived.await(); + return ApiFutures.immediateFuture(null); + } catch (InterruptedException e) { + return ApiFutures.immediateFailedFuture( + SpannerExceptionFactory.propagateInterrupt(e)); + } + } + }, + executor); + // Wait until at least one row has been fetched. At that moment there should be one session + // checked out. + dataReceived.await(); + assertThat(clientImpl.pool.getNumberOfSessionsInUse()).isEqualTo(1); + assertThat(res.isDone()).isFalse(); + dataChecked.countDown(); + // Get the data from the transaction. + List resultList = new ArrayList<>(); + do { + results.drainTo(resultList); + } while (!finished.isDone() || results.size() > 0); + assertThat(finished.get()).isTrue(); + assertThat(resultList).containsExactly("k1", "k2", "k3"); + assertThat(res.get()).isNull(); + assertThat(clientImpl.pool.getNumberOfSessionsInUse()).isEqualTo(0); + } + + @Test + public void asyncRunnerReadRow() throws Exception { + AsyncRunner runner = client().runAsync(); + ApiFuture val = + runner.runAsync( + new AsyncWork() { + @Override + public ApiFuture doWorkAsync(TransactionContext txn) { + return ApiFutures.transform( + txn.readRowAsync(READ_TABLE_NAME, Key.of(1L), READ_COLUMN_NAMES), + new ApiFunction() { + @Override + public String apply(Struct input) { + return input.getString("Value"); + } + }, + MoreExecutors.directExecutor()); + } + }, + executor); + assertThat(val.get()).isEqualTo("v1"); + } + + @Test + public void asyncRunnerRead() throws Exception { + AsyncRunner runner = client().runAsync(); + ApiFuture> val = + runner.runAsync( + new AsyncWork>() { + @Override + public ApiFuture> doWorkAsync(TransactionContext txn) { + return txn.readAsync(READ_TABLE_NAME, KeySet.all(), READ_COLUMN_NAMES) + .toListAsync( + new Function() { + @Override + public String apply(StructReader input) { + return input.getString("Value"); + } + }, + MoreExecutors.directExecutor()); + } + }, + executor); + assertThat(val.get()).containsExactly("v1", "v2", "v3"); + } +} diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadAsyncTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadAsyncTest.java index 50350fd2a14..83d060596c7 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadAsyncTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadAsyncTest.java @@ -123,34 +123,35 @@ public void after() { @Test public void emptyReadAsync() throws Exception { final SettableFuture result = SettableFuture.create(); - AsyncResultSet resultSet = + try (AsyncResultSet resultSet = client .singleUse(TimestampBound.strong()) - .readAsync(EMPTY_READ_TABLE_NAME, KeySet.singleKey(Key.of("k99")), READ_COLUMN_NAMES); - resultSet.setCallback( - executor, - new ReadyCallback() { - @Override - public CallbackResponse cursorReady(AsyncResultSet resultSet) { - try { - while (true) { - switch (resultSet.tryNext()) { - case OK: - fail("received unexpected data"); - case NOT_READY: - return CallbackResponse.CONTINUE; - case DONE: - assertThat(resultSet.getType()).isEqualTo(READ_TABLE_TYPE); - result.set(true); - return CallbackResponse.DONE; + .readAsync(EMPTY_READ_TABLE_NAME, KeySet.singleKey(Key.of("k99")), READ_COLUMN_NAMES)) { + resultSet.setCallback( + executor, + new ReadyCallback() { + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + try { + while (true) { + switch (resultSet.tryNext()) { + case OK: + fail("received unexpected data"); + case NOT_READY: + return CallbackResponse.CONTINUE; + case DONE: + assertThat(resultSet.getType()).isEqualTo(READ_TABLE_TYPE); + result.set(true); + return CallbackResponse.DONE; + } } + } catch (Throwable t) { + result.setException(t); + return CallbackResponse.DONE; } - } catch (Throwable t) { - result.setException(t); - return CallbackResponse.DONE; } - } - }); + }); + } assertThat(result.get()).isTrue(); } From 888edd819565eff78c3a48fbdca58de57d9c8869 Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Sat, 13 Jun 2020 15:43:49 +0200 Subject: [PATCH 36/49] fix: session leaks + code format --- .../cloud/spanner/AsyncResultSetImpl.java | 17 ++++--- .../cloud/spanner/DatabaseClientImplTest.java | 49 ++++++++++--------- .../cloud/spanner/SpannerGaxRetryTest.java | 13 ++--- 3 files changed, 43 insertions(+), 36 deletions(-) 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 a86af434342..3e15dc6e635 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 @@ -223,6 +223,14 @@ public CursorState tryNext() throws SpannerException { return CursorState.NOT_READY; } + private void closeDelegateResultSet() { + try { + delegateResultSet.close(); + } catch (Throwable t) { + log.log(Level.FINE, "Ignoring error from closing delegate result set", t); + } + } + /** * {@link CallbackRunnable} calls the {@link ReadyCallback} registered for this {@link * AsyncResultSet}. @@ -264,6 +272,7 @@ public void run() { switch (response) { case DONE: state = State.DONE; + closeDelegateResultSet(); return; case PAUSE: state = State.PAUSED; @@ -347,8 +356,8 @@ public Void call() throws Exception { if (!stop) { buffer.put(delegateResultSet.getCurrentRowAsStruct()); startCallbackIfNecessary(); + hasNext = delegateResultSet.next(); } - hasNext = delegateResultSet.next(); } catch (Throwable e) { synchronized (monitor) { executionException = SpannerExceptionFactory.newSpannerException(e); @@ -358,11 +367,7 @@ public Void call() throws Exception { } // We don't need any more data from the underlying result set, so we close it as soon as // possible. Any error that might occur during this will be ignored. - try { - delegateResultSet.close(); - } catch (Throwable t) { - log.log(Level.FINE, "Ignoring error from closing delegate result set", t); - } + closeDelegateResultSet(); // Ensure that the callback has been called at least once, even if the result set was // cancelled. diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java index 4a9e3efcf4e..7d2d8cad8fa 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java @@ -1389,32 +1389,33 @@ public void testAsyncQuery() throws Exception { final List receivedResults = new ArrayList<>(); try (AsyncResultSet rs = client.singleUse().executeQueryAsync(Statement.of("SELECT * FROM RANDOM"))) { - resultSetClosed = rs.setCallback( - executor, - new ReadyCallback() { - @Override - public CallbackResponse cursorReady(AsyncResultSet resultSet) { - try { - while (true) { - switch (rs.tryNext()) { - case DONE: - finished.set(true); - return CallbackResponse.DONE; - case NOT_READY: - return CallbackResponse.CONTINUE; - case OK: - receivedResults.add(resultSet.getCurrentRowAsStruct()); - break; - default: - throw new IllegalStateException("Unknown cursor state"); + resultSetClosed = + rs.setCallback( + executor, + new ReadyCallback() { + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + try { + while (true) { + switch (rs.tryNext()) { + case DONE: + finished.set(true); + return CallbackResponse.DONE; + case NOT_READY: + return CallbackResponse.CONTINUE; + case OK: + receivedResults.add(resultSet.getCurrentRowAsStruct()); + break; + default: + throw new IllegalStateException("Unknown cursor state"); + } + } + } catch (Throwable t) { + finished.setException(t); + return CallbackResponse.DONE; } } - } catch (Throwable t) { - finished.setException(t); - return CallbackResponse.DONE; - } - } - }); + }); } assertThat(finished.get()).isTrue(); assertThat(receivedResults.size()).isEqualTo(EXPECTED_ROW_COUNT); diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerGaxRetryTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerGaxRetryTest.java index 1ed8165cd02..f9c3c2040fa 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerGaxRetryTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerGaxRetryTest.java @@ -321,12 +321,13 @@ public void readWriteTransactionTimeout() { mockSpanner.setBeginTransactionExecutionTime(ONE_SECOND); try { TransactionRunner runner = clientWithTimeout.readWriteTransaction(); - runner.run(new TransactionCallable(){ - @Override - public Void run(TransactionContext transaction) throws Exception { - return null; - } - }); + runner.run( + new TransactionCallable() { + @Override + public Void run(TransactionContext transaction) throws Exception { + return null; + } + }); fail("Expected exception"); } catch (SpannerException ex) { assertEquals(ErrorCode.DEADLINE_EXCEEDED, ex.getErrorCode()); From b2a7176d994aacc0b413d6cf19b2da42fae0f2bd Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Sat, 13 Jun 2020 16:48:15 +0200 Subject: [PATCH 37/49] fix: more session leak fixes --- .../cloud/spanner/AbstractReadContext.java | 39 +-- .../cloud/spanner/AsyncResultSetImpl.java | 16 +- .../google/cloud/spanner/ReadAsyncTest.java | 224 ++++++++++-------- 3 files changed, 155 insertions(+), 124 deletions(-) diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java index f4e6596fdf5..bc4a8685648 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java @@ -22,6 +22,8 @@ import static com.google.common.base.Preconditions.checkState; 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.api.gax.core.ExecutorProvider; import com.google.cloud.Timestamp; @@ -51,7 +53,6 @@ import io.opencensus.trace.Span; import io.opencensus.trace.Tracing; import java.util.Map; -import java.util.concurrent.ExecutionException; import java.util.concurrent.atomic.AtomicLong; import javax.annotation.Nullable; import javax.annotation.concurrent.GuardedBy; @@ -759,24 +760,24 @@ static ApiFuture consumeSingleRowAsync(AsyncResultSet resultSet) { // We can safely use a directExecutor here, as we will only be consuming one row, and we will // not be doing any blocking stuff in the handler. final SettableApiFuture row = SettableApiFuture.create(); - resultSet - .setCallback(MoreExecutors.directExecutor(), ConsumeSingleRowCallback.create(row)) - .addListener( - new Runnable() { - @Override - public void run() { - try { - result.set(row.get()); - } catch (ExecutionException e) { - result.setException( - SpannerExceptionFactory.newSpannerException( - e.getCause() == null ? e : e.getCause())); - } catch (InterruptedException e) { - result.setException(SpannerExceptionFactory.propagateInterrupt(e)); - } - } - }, - MoreExecutors.directExecutor()); + ApiFutures.addCallback( + resultSet.setCallback(MoreExecutors.directExecutor(), ConsumeSingleRowCallback.create(row)), + new ApiFutureCallback() { + @Override + public void onFailure(Throwable t) { + result.setException(t); + } + + @Override + public void onSuccess(Void input) { + try { + result.set(row.get()); + } catch (Throwable t) { + result.setException(t); + } + } + }, + MoreExecutors.directExecutor()); return result; } 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 3e15dc6e635..f277388b0b2 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 @@ -16,7 +16,9 @@ package com.google.cloud.spanner; +import com.google.api.core.ApiAsyncFunction; import com.google.api.core.ApiFuture; +import com.google.api.core.ApiFutures; import com.google.api.core.ListenableFutureToApiFuture; import com.google.api.core.SettableApiFuture; import com.google.api.gax.core.ExecutorProvider; @@ -520,10 +522,18 @@ public ApiFuture> toListAsync( Preconditions.checkState(!closed, "This AsyncResultSet has been closed"); Preconditions.checkState( this.state == State.INITIALIZED, "This AsyncResultSet has already been used."); - SettableApiFuture> res = SettableApiFuture.>create(); + final SettableApiFuture> res = SettableApiFuture.>create(); CreateListCallback callback = new CreateListCallback(res, transformer); - setCallback(executor, callback); - return res; + ApiFuture finished = setCallback(executor, callback); + return ApiFutures.transformAsync( + finished, + new ApiAsyncFunction>() { + @Override + public ApiFuture> apply(Void input) throws Exception { + return res; + } + }, + MoreExecutors.directExecutor()); } } diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadAsyncTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadAsyncTest.java index 83d060596c7..13e4c47d082 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadAsyncTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadAsyncTest.java @@ -35,7 +35,6 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableSet; import com.google.common.collect.Iterables; -import com.google.common.util.concurrent.SettableFuture; import io.grpc.Server; import io.grpc.Status; import io.grpc.inprocess.InProcessServerBuilder; @@ -120,39 +119,63 @@ public void after() { mockSpanner.removeAllExecutionTimes(); } + @Test + public void readAsyncPropagatesError() throws Exception { + ApiFuture result; + try (AsyncResultSet resultSet = + client + .singleUse(TimestampBound.strong()) + .readAsync(EMPTY_READ_TABLE_NAME, KeySet.singleKey(Key.of("k99")), READ_COLUMN_NAMES)) { + result = + resultSet.setCallback( + executor, + new ReadyCallback() { + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + throw SpannerExceptionFactory.newSpannerException( + ErrorCode.CANCELLED, "Don't want the data"); + } + }); + } + try { + result.get(); + fail("missing expected exception"); + } catch (ExecutionException e) { + assertThat(e.getCause()).isInstanceOf(SpannerException.class); + SpannerException se = (SpannerException) e.getCause(); + assertThat(se.getErrorCode()).isEqualTo(ErrorCode.CANCELLED); + assertThat(se.getMessage()).contains("Don't want the data"); + } + } + @Test public void emptyReadAsync() throws Exception { - final SettableFuture result = SettableFuture.create(); + ApiFuture result; try (AsyncResultSet resultSet = client .singleUse(TimestampBound.strong()) .readAsync(EMPTY_READ_TABLE_NAME, KeySet.singleKey(Key.of("k99")), READ_COLUMN_NAMES)) { - resultSet.setCallback( - executor, - new ReadyCallback() { - @Override - public CallbackResponse cursorReady(AsyncResultSet resultSet) { - try { - while (true) { - switch (resultSet.tryNext()) { - case OK: - fail("received unexpected data"); - case NOT_READY: - return CallbackResponse.CONTINUE; - case DONE: - assertThat(resultSet.getType()).isEqualTo(READ_TABLE_TYPE); - result.set(true); - return CallbackResponse.DONE; + result = + resultSet.setCallback( + executor, + new ReadyCallback() { + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + while (true) { + switch (resultSet.tryNext()) { + case OK: + fail("received unexpected data"); + case NOT_READY: + return CallbackResponse.CONTINUE; + case DONE: + assertThat(resultSet.getType()).isEqualTo(READ_TABLE_TYPE); + return CallbackResponse.DONE; + } } } - } catch (Throwable t) { - result.setException(t); - return CallbackResponse.DONE; - } - } - }); + }); } - assertThat(result.get()).isTrue(); + assertThat(result.get()).isNull(); } @Test @@ -225,6 +248,7 @@ public void tableNotFound() throws Exception { public void closeTransactionBeforeEndOfAsyncQuery() throws Exception { final BlockingQueue results = new SynchronousQueue<>(); final SettableApiFuture finished = SettableApiFuture.create(); + ApiFuture closed; DatabaseClientImpl clientImpl = (DatabaseClientImpl) client; // There should currently not be any sessions checked out of the pool. @@ -234,30 +258,31 @@ public void closeTransactionBeforeEndOfAsyncQuery() throws Exception { try (ReadOnlyTransaction tx = client.readOnlyTransaction()) { try (AsyncResultSet rs = tx.readAsync(READ_TABLE_NAME, KeySet.all(), READ_COLUMN_NAMES, Options.bufferRows(1))) { - rs.setCallback( - executor, - new ReadyCallback() { - @Override - public CallbackResponse cursorReady(AsyncResultSet resultSet) { - try { - while (true) { - switch (resultSet.tryNext()) { - case DONE: - finished.set(true); - return CallbackResponse.DONE; - case NOT_READY: - return CallbackResponse.CONTINUE; - case OK: - dataReceived.countDown(); - results.put(resultSet.getString(0)); + closed = + rs.setCallback( + executor, + new ReadyCallback() { + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + try { + while (true) { + switch (resultSet.tryNext()) { + case DONE: + finished.set(true); + return CallbackResponse.DONE; + case NOT_READY: + return CallbackResponse.CONTINUE; + case OK: + dataReceived.countDown(); + results.put(resultSet.getString(0)); + } + } + } catch (Throwable t) { + finished.setException(t); + return CallbackResponse.DONE; } } - } catch (Throwable t) { - finished.setException(t); - return CallbackResponse.DONE; - } - } - }); + }); } // Wait until at least one row has been fetched. At that moment there should be one session // checked out. @@ -278,7 +303,7 @@ public CallbackResponse cursorReady(AsyncResultSet resultSet) { assertThat(resultList).containsExactly("k1", "k2", "k3"); // The session will be released back into the pool by the asynchronous result set when it has // returned all rows. As this is done in the background, it could take a couple of milliseconds. - Thread.sleep(10L); + closed.get(); assertThat(clientImpl.pool.getNumberOfSessionsInUse()).isEqualTo(0); } @@ -354,69 +379,64 @@ public void pauseResume() throws Exception { evenStatement, generateKeyValueResultSet(ImmutableSet.of(2, 4, 6, 8, 10)))); final Object lock = new Object(); - final SettableApiFuture evenFinished = SettableApiFuture.create(); - final SettableApiFuture unevenFinished = SettableApiFuture.create(); + ApiFuture evenFinished; + ApiFuture unevenFinished; final CountDownLatch unevenReturnedFirstRow = new CountDownLatch(1); final Deque allValues = new ConcurrentLinkedDeque<>(); try (ReadOnlyTransaction tx = client.readOnlyTransaction()) { try (AsyncResultSet evenRs = tx.executeQueryAsync(evenStatement); AsyncResultSet unevenRs = tx.executeQueryAsync(unevenStatement)) { - unevenRs.setCallback( - executor, - new ReadyCallback() { - @Override - public CallbackResponse cursorReady(AsyncResultSet resultSet) { - try { - while (true) { - switch (resultSet.tryNext()) { - case DONE: - unevenFinished.set(true); - return CallbackResponse.DONE; - case NOT_READY: - return CallbackResponse.CONTINUE; - case OK: - synchronized (lock) { - allValues.add(resultSet.getString("Value")); - } - unevenReturnedFirstRow.countDown(); - return CallbackResponse.PAUSE; + unevenFinished = + unevenRs.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: + synchronized (lock) { + allValues.add(resultSet.getString("Value")); + } + unevenReturnedFirstRow.countDown(); + return CallbackResponse.PAUSE; + } } } - } catch (Throwable t) { - unevenFinished.setException(t); - return CallbackResponse.DONE; - } - } - }); - evenRs.setCallback( - executor, - new ReadyCallback() { - @Override - public CallbackResponse cursorReady(AsyncResultSet resultSet) { - try { - // Make sure the uneven result set has returned the first before we start the even - // results. - unevenReturnedFirstRow.await(); - while (true) { - switch (resultSet.tryNext()) { - case DONE: - evenFinished.set(true); - return CallbackResponse.DONE; - case NOT_READY: - return CallbackResponse.CONTINUE; - case OK: - synchronized (lock) { - allValues.add(resultSet.getString("Value")); + }); + evenFinished = + evenRs.setCallback( + executor, + new ReadyCallback() { + @Override + public CallbackResponse cursorReady(AsyncResultSet resultSet) { + try { + // Make sure the uneven result set has returned the first before we start the + // even + // results. + unevenReturnedFirstRow.await(); + while (true) { + switch (resultSet.tryNext()) { + case DONE: + return CallbackResponse.DONE; + case NOT_READY: + return CallbackResponse.CONTINUE; + case OK: + synchronized (lock) { + allValues.add(resultSet.getString("Value")); + } + return CallbackResponse.PAUSE; } - return CallbackResponse.PAUSE; + } + } catch (InterruptedException e) { + throw SpannerExceptionFactory.propagateInterrupt(e); } } - } catch (Throwable t) { - evenFinished.setException(t); - return CallbackResponse.DONE; - } - } - }); + }); while (!(evenFinished.isDone() && unevenFinished.isDone())) { synchronized (lock) { if (allValues.peekLast() != null) { @@ -435,7 +455,7 @@ public CallbackResponse cursorReady(AsyncResultSet resultSet) { } } assertThat(ApiFutures.allAsList(Arrays.asList(evenFinished, unevenFinished)).get()) - .containsExactly(Boolean.TRUE, Boolean.TRUE); + .containsExactly(null, null); assertThat(allValues) .containsExactly("v1", "v2", "v3", "v4", "v5", "v6", "v7", "v8", "v9", "v10"); } From fba270fbf62f73bc0489a2ce8626e5b011ff1416 Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Sun, 14 Jun 2020 22:10:40 +0200 Subject: [PATCH 38/49] feat: further work on AsyncTransactionManager --- .../spanner/AsyncTransactionManager.java | 37 +- .../spanner/AsyncTransactionManagerImpl.java | 62 +- .../com/google/cloud/spanner/SessionPool.java | 51 +- .../spanner/TransactionContextFutureImpl.java | 207 +++ .../cloud/spanner/TransactionRunnerImpl.java | 26 +- .../spanner/AbstractAsyncTransactionTest.java | 2 +- .../spanner/AsyncTransactionManagerTest.java | 1234 ++++++++++------- 7 files changed, 1021 insertions(+), 598 deletions(-) create mode 100644 google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionContextFutureImpl.java diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManager.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManager.java index 8e8f808d0b8..7cce2296ce4 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManager.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManager.java @@ -18,7 +18,11 @@ import com.google.api.core.ApiFuture; import com.google.cloud.Timestamp; +import com.google.cloud.spanner.AsyncTransactionManager.AsyncTransactionFunction; import com.google.cloud.spanner.TransactionManager.TransactionState; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; /** * An interface for managing the life cycle of a read write transaction including all its retries. @@ -37,13 +41,36 @@ * * @see DatabaseClient#transactionManager() */ -public interface AsyncTransactionManager { +public interface AsyncTransactionManager extends AutoCloseable { + interface TransactionContextFuture extends ApiFuture { + AsyncTransactionStep then(AsyncTransactionFunction function); + } + + interface CommitTimestampFuture extends ApiFuture { + @Override + Timestamp get() throws AbortedException, InterruptedException, ExecutionException; + + @Override + Timestamp get(long timeout, TimeUnit unit) + throws AbortedException, InterruptedException, ExecutionException, TimeoutException; + } + + interface AsyncTransactionStep extends ApiFuture { + AsyncTransactionStep then(AsyncTransactionFunction next); + + CommitTimestampFuture commitAsync(); + } + + interface AsyncTransactionFunction { + ApiFuture apply(TransactionContext txn, I input) throws Exception; + } + /** * Creates a new read write transaction. This must be called before doing any other operation and * can only be called once. To create a new transaction for subsequent retries, see {@link * #resetForRetry()}. */ - ApiFuture beginAsync(); + TransactionContextFuture beginAsync(); /** * Commits the currently active transaction. If the transaction was already aborted, then this @@ -64,7 +91,7 @@ public interface AsyncTransactionManager { * specified by {@link SpannerException#getRetryDelayInMillis()} on the {@code SpannerException} * throw by the previous commit call. */ - ApiFuture resetForRetryAsync(); + TransactionContextFuture resetForRetryAsync(); /** Returns the state of the transaction. */ TransactionState getState(); @@ -73,6 +100,6 @@ public interface AsyncTransactionManager { * Closes the manager. If there is an active transaction, it will be rolled back. Underlying * session will be released back to the session pool. */ - // @Override - // void close(); + @Override + void close(); } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManagerImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManagerImpl.java index b830a8a61b3..31d809dad69 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManagerImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManagerImpl.java @@ -29,7 +29,6 @@ import io.opencensus.trace.Span; import io.opencensus.trace.Tracer; import io.opencensus.trace.Tracing; -import java.util.concurrent.ExecutionException; /** Implementation of {@link AsyncTransactionManager}. */ final class AsyncTransactionManagerImpl implements AsyncTransactionManager, SessionTransaction { @@ -53,29 +52,41 @@ public void setSpan(Span span) { } @Override - public ApiFuture beginAsync() { + public void close() { + txn.close(); + } + + @Override + public TransactionContextFuture beginAsync() { Preconditions.checkState(txn == null, "begin can only be called once"); - txnState = TransactionState.STARTED; - txn = session.newTransaction(); + TransactionContextFuture begin = new TransactionContextFutureImpl(this, internalBeginAsync()); session.setActive(this); + return begin; + } + + private ApiFuture internalBeginAsync() { + final Scope s = tracer.withSpan(span); + txn = session.newTransaction(); + txnState = TransactionState.STARTED; final SettableApiFuture res = SettableApiFuture.create(); final ApiFuture fut = txn.ensureTxnAsync(); - fut.addListener( - tracer.withSpan( - span, - new Runnable() { - @Override - public void run() { - try { - fut.get(); - res.set(txn); - } catch (ExecutionException e) { - res.setException(e.getCause() == null ? e : e.getCause()); - } catch (InterruptedException e) { - res.setException(SpannerExceptionFactory.propagateInterrupt(e)); - } - } - }), + ApiFutures.addCallback( + fut, + new ApiFutureCallback() { + @Override + public void onFailure(Throwable t) { + s.close(); + TraceUtil.endSpanWithFailure(span, t); + res.setException(SpannerExceptionFactory.newSpannerException(t)); + } + + @Override + public void onSuccess(Void result) { + s.close(); + span.end(); + res.set(txn); + } + }, MoreExecutors.directExecutor()); return res; } @@ -128,17 +139,12 @@ public ApiFuture rollbackAsync() { } @Override - public ApiFuture resetForRetryAsync() { + public TransactionContextFuture resetForRetryAsync() { if (txn == null || !txn.isAborted() && txnState != TransactionState.ABORTED) { throw new IllegalStateException( - "resetForRetry can only be called if the previous attempt" + " aborted"); - } - try (Scope s = tracer.withSpan(span)) { - txn = session.newTransaction(); - txn.ensureTxn(); - txnState = TransactionState.STARTED; - return ApiFutures.immediateFuture((TransactionContext) txn); + "resetForRetry can only be called if the previous attempt aborted"); } + return new TransactionContextFutureImpl(this, internalBeginAsync()); } @Override diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java index 6126a077bfc..711a905ac98 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java @@ -974,19 +974,32 @@ public void run() { } @Override - public ApiFuture beginAsync() { - return ApiFutures.transformAsync( - delegate, - new ApiAsyncFunction() { + public void close() { + delegate.addListener( + new Runnable() { @Override - public ApiFuture apply(AsyncTransactionManager input) - throws Exception { - return input.beginAsync(); + public void run() { + session.close(); } }, MoreExecutors.directExecutor()); } + @Override + public TransactionContextFuture beginAsync() { + return new TransactionContextFutureImpl( + this, + ApiFutures.transformAsync( + delegate, + new ApiAsyncFunction() { + @Override + public ApiFuture apply(AsyncTransactionManager input) { + return input.beginAsync(); + } + }, + MoreExecutors.directExecutor())); + } + @Override public ApiFuture commitAsync() { return ApiFutures.transformAsync( @@ -1032,17 +1045,19 @@ public void run() { } @Override - public ApiFuture resetForRetryAsync() { - return ApiFutures.transformAsync( - delegate, - new ApiAsyncFunction() { - @Override - public ApiFuture apply(AsyncTransactionManager input) - throws Exception { - return input.resetForRetryAsync(); - } - }, - MoreExecutors.directExecutor()); + public TransactionContextFuture resetForRetryAsync() { + return new TransactionContextFutureImpl( + this, + ApiFutures.transformAsync( + delegate, + new ApiAsyncFunction() { + @Override + public ApiFuture apply(AsyncTransactionManager input) + throws Exception { + return input.resetForRetryAsync(); + } + }, + MoreExecutors.directExecutor())); } @Override diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionContextFutureImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionContextFutureImpl.java new file mode 100644 index 00000000000..487b9616518 --- /dev/null +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionContextFutureImpl.java @@ -0,0 +1,207 @@ +/* + * 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.ApiFunction; +import com.google.api.core.ApiFuture; +import com.google.api.core.ApiFutureCallback; +import com.google.api.core.ApiFutures; +import com.google.api.core.ForwardingApiFuture; +import com.google.api.core.SettableApiFuture; +import com.google.cloud.Timestamp; +import com.google.cloud.spanner.AsyncTransactionManager.AsyncTransactionFunction; +import com.google.cloud.spanner.AsyncTransactionManager.AsyncTransactionStep; +import com.google.cloud.spanner.AsyncTransactionManager.CommitTimestampFuture; +import com.google.cloud.spanner.AsyncTransactionManager.TransactionContextFuture; +import com.google.common.base.Preconditions; +import com.google.common.util.concurrent.MoreExecutors; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; + +class TransactionContextFutureImpl extends ForwardingApiFuture + implements TransactionContextFuture { + static class CommitTimestampFutureImpl extends ForwardingApiFuture + implements CommitTimestampFuture { + CommitTimestampFutureImpl(ApiFuture delegate) { + super(Preconditions.checkNotNull(delegate)); + } + + @Override + public Timestamp get() throws AbortedException, ExecutionException, InterruptedException { + try { + return super.get(); + } catch (ExecutionException e) { + if (e.getCause() != null && e.getCause() instanceof AbortedException) { + throw (AbortedException) e.getCause(); + } + throw e; + } + } + + @Override + public Timestamp get(long timeout, TimeUnit unit) + throws AbortedException, ExecutionException, InterruptedException, TimeoutException { + try { + return super.get(timeout, unit); + } catch (ExecutionException e) { + if (e.getCause() != null && e.getCause() instanceof AbortedException) { + throw (AbortedException) e.getCause(); + } + throw e; + } + } + } + + class AsyncTransactionStatementImpl extends ForwardingApiFuture + implements AsyncTransactionStep { + final ApiFuture txnFuture; + final SettableApiFuture statementResult; + + AsyncTransactionStatementImpl( + final ApiFuture txnFuture, + ApiFuture input, + final AsyncTransactionFunction function) { + this(SettableApiFuture.create(), txnFuture, input, function); + } + + AsyncTransactionStatementImpl( + SettableApiFuture delegate, + final ApiFuture txnFuture, + ApiFuture input, + final AsyncTransactionFunction function) { + super(delegate); + this.statementResult = delegate; + this.txnFuture = txnFuture; + ApiFutures.addCallback( + input, + new ApiFutureCallback() { + @Override + public void onFailure(Throwable t) { + txnResult.setException(t); + } + + @Override + public void onSuccess(I result) { + try { + ApiFutures.addCallback( + Preconditions.checkNotNull( + function.apply(txnFuture.get(), result), + "AsyncTransactionFunction returned . Did you mean to return ApiFutures.immediateFuture(null)?"), + new ApiFutureCallback() { + @Override + public void onFailure(Throwable t) { + txnResult.setException(t); + } + + @Override + public void onSuccess(O result) { + statementResult.set(result); + } + }, + MoreExecutors.directExecutor()); + } catch (Throwable t) { + txnResult.setException(t); + } + } + }, + MoreExecutors.directExecutor()); + } + + @Override + public AsyncTransactionStatementImpl then(AsyncTransactionFunction next) { + return new AsyncTransactionStatementImpl<>(txnFuture, statementResult, next); + } + + @Override + public CommitTimestampFuture commitAsync() { + ApiFutures.addCallback( + statementResult, + new ApiFutureCallback() { + @Override + public void onFailure(Throwable t) { + txnResult.setException(t); + } + + @Override + public void onSuccess(O result) { + ApiFutures.addCallback( + mgr.commitAsync(), + new ApiFutureCallback() { + @Override + public void onFailure(Throwable t) { + txnResult.setException(t); + } + + @Override + public void onSuccess(Timestamp result) { + txnResult.set(result); + } + }, + MoreExecutors.directExecutor()); + } + }, + MoreExecutors.directExecutor()); + return new CommitTimestampFutureImpl(txnResult); + } + } + + final AsyncTransactionManager mgr; + final SettableApiFuture txnResult = SettableApiFuture.create(); + + TransactionContextFutureImpl( + AsyncTransactionManager mgr, ApiFuture txnFuture) { + super(txnFuture); + this.mgr = mgr; + } + + @Override + public AsyncTransactionStatementImpl then( + AsyncTransactionFunction function) { + return new AsyncTransactionStatementImpl<>( + this, + ApiFutures.transform( + this, + new ApiFunction() { + @Override + public Void apply(TransactionContext input) { + return null; + } + }, + MoreExecutors.directExecutor()), + function); + } + + ApiFuture commitAsync() { + ApiFuture res = mgr.commitAsync(); + ApiFutures.addCallback( + res, + new ApiFutureCallback() { + @Override + public void onFailure(Throwable t) { + txnResult.setException(t); + } + + @Override + public void onSuccess(Timestamp result) { + txnResult.set(result); + } + }, + MoreExecutors.directExecutor()); + return txnResult; + } +} 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 36565fb0f8d..c10b7132852 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 @@ -246,18 +246,6 @@ void commit() { ApiFuture commitAsync() { final SettableApiFuture res = SettableApiFuture.create(); - CommitRequest.Builder builder = - CommitRequest.newBuilder().setSession(session.getName()).setTransactionId(transactionId); - synchronized (lock) { - if (!mutations.isEmpty()) { - List mutationsProto = new ArrayList<>(); - Mutation.toProto(mutations, mutationsProto); - builder.addAllMutations(mutationsProto); - } - // Ensure that no call to buffer mutations that would be lost can succeed. - mutations = null; - } - final CommitRequest commitRequest = builder.build(); final SettableApiFuture latch; synchronized (lock) { latch = finishedAsyncOperations; @@ -268,6 +256,20 @@ ApiFuture commitAsync() { public void run() { try { latch.get(); + CommitRequest.Builder builder = + CommitRequest.newBuilder() + .setSession(session.getName()) + .setTransactionId(transactionId); + synchronized (lock) { + if (!mutations.isEmpty()) { + List mutationsProto = new ArrayList<>(); + Mutation.toProto(mutations, mutationsProto); + builder.addAllMutations(mutationsProto); + } + // Ensure that no call to buffer mutations that would be lost can succeed. + mutations = null; + } + final CommitRequest commitRequest = builder.build(); span.addAnnotation("Starting Commit"); final Span opSpan = tracer.spanBuilderWithExplicitParent(SpannerImpl.COMMIT, span).startSpan(); diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AbstractAsyncTransactionTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AbstractAsyncTransactionTest.java index fbd8e44ffef..1926860162a 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AbstractAsyncTransactionTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AbstractAsyncTransactionTest.java @@ -87,9 +87,9 @@ public static void setup() throws Exception { @AfterClass public static void teardown() throws Exception { - executor.shutdown(); server.shutdown(); server.awaitTermination(); + executor.shutdown(); } @Before diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncTransactionManagerTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncTransactionManagerTest.java index 8b6772e43b9..65681eda092 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncTransactionManagerTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncTransactionManagerTest.java @@ -16,20 +16,24 @@ package com.google.cloud.spanner; -import static com.google.cloud.spanner.MockSpannerTestUtil.*; +import static com.google.cloud.spanner.MockSpannerTestUtil.INVALID_UPDATE_STATEMENT; +import static com.google.cloud.spanner.MockSpannerTestUtil.READ_COLUMN_NAMES; +import static com.google.cloud.spanner.MockSpannerTestUtil.READ_TABLE_NAME; +import static com.google.cloud.spanner.MockSpannerTestUtil.UPDATE_ABORTED_STATEMENT; +import static com.google.cloud.spanner.MockSpannerTestUtil.UPDATE_COUNT; +import static com.google.cloud.spanner.MockSpannerTestUtil.UPDATE_STATEMENT; import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.fail; -import com.google.api.core.ApiAsyncFunction; -import com.google.api.core.ApiFunction; 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.AsyncResultSet.CallbackResponse; -import com.google.cloud.spanner.AsyncResultSet.ReadyCallback; import com.google.cloud.spanner.AsyncRunner.AsyncWork; +import com.google.cloud.spanner.AsyncTransactionManager.AsyncTransactionFunction; +import com.google.cloud.spanner.AsyncTransactionManager.AsyncTransactionStep; +import com.google.cloud.spanner.AsyncTransactionManager.CommitTimestampFuture; +import com.google.cloud.spanner.AsyncTransactionManager.TransactionContextFuture; import com.google.cloud.spanner.MockSpannerServiceImpl.SimulatedExecutionTime; import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult; import com.google.common.base.Function; @@ -41,13 +45,8 @@ import com.google.spanner.v1.ExecuteBatchDmlRequest; import com.google.spanner.v1.ExecuteSqlRequest; import io.grpc.Status; -import java.util.ArrayList; -import java.util.List; -import java.util.concurrent.BlockingQueue; -import java.util.concurrent.CountDownLatch; +import java.util.Arrays; import java.util.concurrent.ExecutionException; -import java.util.concurrent.Executors; -import java.util.concurrent.SynchronousQueue; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; import org.junit.Test; @@ -57,238 +56,343 @@ @RunWith(JUnit4.class) public class AsyncTransactionManagerTest extends AbstractAsyncTransactionTest { + /** + * Static helper methods that simplifies creating {@link AsyncTransactionFunction}s for Java7. + * Java8 and higher can use lambda expressions. + */ + public static class AsyncTransactionManagerHelper { + public static AsyncTransactionFunction readRowAsync( + final String table, final Key key, final Iterable columns) { + return new AsyncTransactionFunction() { + @Override + public ApiFuture apply(TransactionContext txn, I input) throws Exception { + return txn.readRowAsync(table, key, columns); + } + }; + } + + public static AsyncTransactionFunction executeUpdateAsync(Statement statement) { + return executeUpdateAsync(SettableApiFuture.create(), statement); + } + + public static AsyncTransactionFunction executeUpdateAsync( + final SettableApiFuture result, final Statement statement) { + return new AsyncTransactionFunction() { + @Override + public ApiFuture apply(TransactionContext txn, I input) throws Exception { + ApiFuture updateCount = txn.executeUpdateAsync(statement); + ApiFutures.addCallback( + updateCount, + new ApiFutureCallback() { + @Override + public void onFailure(Throwable t) { + result.setException(t); + } + + @Override + public void onSuccess(Long input) { + result.set(input); + } + }, + MoreExecutors.directExecutor()); + return updateCount; + } + }; + } + + public static AsyncTransactionFunction batchUpdateAsync( + final Statement... statements) { + return batchUpdateAsync(SettableApiFuture.create(), statements); + } + + public static AsyncTransactionFunction batchUpdateAsync( + final SettableApiFuture result, final Statement... statements) { + return new AsyncTransactionFunction() { + @Override + public ApiFuture apply(TransactionContext txn, I input) throws Exception { + ApiFuture updateCounts = txn.batchUpdateAsync(Arrays.asList(statements)); + ApiFutures.addCallback( + updateCounts, + new ApiFutureCallback() { + @Override + public void onFailure(Throwable t) { + result.setException(t); + } + + @Override + public void onSuccess(long[] input) { + result.set(input); + } + }, + MoreExecutors.directExecutor()); + return updateCounts; + } + }; + } + } + @Test public void asyncTransactionManagerUpdate() throws Exception { - final SettableApiFuture> commitTimestamp = SettableApiFuture.create(); - final AsyncTransactionManager manager = client().transactionManagerAsync(); - ApiFuture txn = manager.beginAsync(); - ApiFuture updateCount = - ApiFutures.transformAsync( - txn, - new ApiAsyncFunction() { - @Override - public ApiFuture apply(TransactionContext input) throws Exception { - ApiFuture res = input.executeUpdateAsync(UPDATE_STATEMENT); - commitTimestamp.set( - ApiFutures.transformAsync( - res, - new ApiAsyncFunction() { - @Override - public ApiFuture apply(Long input) throws Exception { - return manager.commitAsync(); - } - }, - executor)); - return res; - } - }, - executor); - assertThat(updateCount.get()).isEqualTo(UPDATE_COUNT); - assertThat(commitTimestamp.get().get()).isNotNull(); + final SettableApiFuture updateCount = SettableApiFuture.create(); + + try (AsyncTransactionManager manager = client().transactionManagerAsync()) { + TransactionContextFuture txn = manager.beginAsync(); + while (true) { + try { + CommitTimestampFuture commitTimestamp = + txn.then( + AsyncTransactionManagerHelper.executeUpdateAsync( + updateCount, UPDATE_STATEMENT)) + .commitAsync(); + assertThat(updateCount.get()).isEqualTo(UPDATE_COUNT); + assertThat(commitTimestamp.get()).isNotNull(); + break; + } catch (AbortedException e) { + txn = manager.resetForRetryAsync(); + } + } + } } @Test public void asyncTransactionManagerIsNonBlocking() throws Exception { + SettableApiFuture updateCount = SettableApiFuture.create(); + mockSpanner.freeze(); - final SettableApiFuture> commitTimestamp = SettableApiFuture.create(); - final AsyncTransactionManager manager = client().transactionManagerAsync(); - ApiFuture txn = manager.beginAsync(); - ApiFuture updateCount = - ApiFutures.transformAsync( - txn, - new ApiAsyncFunction() { - @Override - public ApiFuture apply(TransactionContext input) throws Exception { - ApiFuture res = input.executeUpdateAsync(UPDATE_STATEMENT); - commitTimestamp.set( - ApiFutures.transformAsync( - res, - new ApiAsyncFunction() { - @Override - public ApiFuture apply(Long input) throws Exception { - return manager.commitAsync(); - } - }, - executor)); - return res; - } - }, - executor); - mockSpanner.unfreeze(); - assertThat(updateCount.get(10L, TimeUnit.SECONDS)).isEqualTo(UPDATE_COUNT); - assertThat(commitTimestamp.get().get(10L, TimeUnit.SECONDS)).isNotNull(); + try (AsyncTransactionManager manager = clientWithEmptySessionPool().transactionManagerAsync()) { + TransactionContextFuture txn = manager.beginAsync(); + while (true) { + try { + CommitTimestampFuture commitTimestamp = + txn.then( + AsyncTransactionManagerHelper.executeUpdateAsync( + updateCount, UPDATE_STATEMENT)) + .commitAsync(); + mockSpanner.unfreeze(); + assertThat(updateCount.get(10L, TimeUnit.SECONDS)).isEqualTo(UPDATE_COUNT); + assertThat(commitTimestamp.get(10L, TimeUnit.SECONDS)).isNotNull(); + break; + } catch (AbortedException e) { + txn = manager.resetForRetryAsync(); + } + } + } } @Test public void asyncTransactionManagerInvalidUpdate() throws Exception { - final AsyncTransactionManager manager = client().transactionManagerAsync(); - final SettableApiFuture rollbacked = SettableApiFuture.create(); - ApiFuture txnFut = manager.beginAsync(); - ApiFuture updateCount = - ApiFutures.transformAsync( - txnFut, - new ApiAsyncFunction() { - @Override - public ApiFuture apply(TransactionContext txn) throws Exception { - ApiFuture res = txn.executeUpdateAsync(INVALID_UPDATE_STATEMENT); - ApiFutures.addCallback( - res, - new ApiFutureCallback() { - @Override - public void onFailure(Throwable t) { - manager - .rollbackAsync() - .addListener( - new Runnable() { - @Override - public void run() { - rollbacked.set(null); - } - }, - executor); - ; - } - - @Override - public void onSuccess(Long result) { - fail("update should not succeed"); - } - }, - MoreExecutors.directExecutor()); - return res; - } - }, - executor); - try { - updateCount.get(); - fail("missing expected exception"); - } catch (ExecutionException e) { - assertThat(e.getCause()).isInstanceOf(SpannerException.class); - SpannerException se = (SpannerException) e.getCause(); - assertThat(se.getErrorCode()).isEqualTo(ErrorCode.INVALID_ARGUMENT); - assertThat(se.getMessage()).contains("invalid statement"); + try (AsyncTransactionManager manager = client().transactionManagerAsync()) { + TransactionContextFuture txn = manager.beginAsync(); + while (true) { + try { + CommitTimestampFuture commitTimestamp = + txn.then( + AsyncTransactionManagerHelper.executeUpdateAsync( + INVALID_UPDATE_STATEMENT)) + .commitAsync(); + commitTimestamp.get(); + fail("missing expected exception"); + } catch (AbortedException e) { + txn = manager.resetForRetryAsync(); + } catch (ExecutionException e) { + manager.rollbackAsync(); + assertThat(e.getCause()).isInstanceOf(SpannerException.class); + SpannerException se = (SpannerException) e.getCause(); + assertThat(se.getErrorCode()).isEqualTo(ErrorCode.INVALID_ARGUMENT); + assertThat(se.getMessage()).contains("invalid statement"); + break; + } + } } - // This is to ensure that the test case does not end before the session has been returned to the - // pool. - rollbacked.get(); } @Test public void asyncTransactionManagerCommitAborted() throws Exception { + SettableApiFuture updateCount = SettableApiFuture.create(); final AtomicInteger attempt = new AtomicInteger(); - final AsyncTransactionManager manager = clientWithEmptySessionPool().transactionManagerAsync(); - ApiFuture txn = manager.beginAsync(); - while (true) { - try { - final SettableApiFuture> commitTimestamp = SettableApiFuture.create(); - ApiFuture updateCount = - ApiFutures.transformAsync( - txn, - new ApiAsyncFunction() { - @Override - public ApiFuture apply(TransactionContext input) throws Exception { - ApiFuture res = input.executeUpdateAsync(UPDATE_STATEMENT); - commitTimestamp.set( - ApiFutures.transformAsync( - res, - new ApiAsyncFunction() { - @Override - public ApiFuture apply(Long input) throws Exception { - if (attempt.incrementAndGet() == 1) { - mockSpanner.abortAllTransactions(); + try (AsyncTransactionManager manager = clientWithEmptySessionPool().transactionManagerAsync()) { + TransactionContextFuture txn = manager.beginAsync(); + while (true) { + try { + attempt.incrementAndGet(); + CommitTimestampFuture commitTimestamp = + txn.then( + AsyncTransactionManagerHelper.executeUpdateAsync( + updateCount, UPDATE_STATEMENT)) + .then( + new AsyncTransactionFunction() { + @Override + public ApiFuture apply(TransactionContext txn, Long input) + throws Exception { + if (attempt.get() == 1) { + mockSpanner.abortTransaction(txn); + } + return ApiFutures.immediateFuture(null); + } + }) + .commitAsync(); + assertThat(updateCount.get()).isEqualTo(UPDATE_COUNT); + assertThat(commitTimestamp.get()).isNotNull(); + assertThat(attempt.get()).isEqualTo(2); + break; + } catch (AbortedException e) { + txn = manager.resetForRetryAsync(); + } + } + } + } + + @Test + public void asyncTransactionManagerFireAndForgetInvalidUpdate() throws Exception { + final SettableApiFuture updateCount = SettableApiFuture.create(); + + try (AsyncTransactionManager mgr = client().transactionManagerAsync()) { + TransactionContextFuture txn = mgr.beginAsync(); + while (true) { + try { + CommitTimestampFuture ts = + txn.then( + new AsyncTransactionFunction() { + @Override + public ApiFuture apply(TransactionContext txn, Void input) + throws Exception { + // This fire-and-forget update statement should not fail the transaction. + txn.executeUpdateAsync(INVALID_UPDATE_STATEMENT); + ApiFutures.addCallback( + txn.executeUpdateAsync(UPDATE_STATEMENT), + new ApiFutureCallback() { + @Override + public void onFailure(Throwable t) { + updateCount.setException(t); } - return manager.commitAsync(); - } - }, - executor)); - return res; - } - }, - executor); - assertThat(updateCount.get()).isEqualTo(UPDATE_COUNT); - assertThat(commitTimestamp.get().get()).isNotNull(); - assertThat(attempt.get()).isEqualTo(2); - break; - } catch (ExecutionException e) { - assertThat(e.getCause()).isInstanceOf(AbortedException.class); - assertThat(attempt.get()).isEqualTo(1); - txn = manager.resetForRetryAsync(); + + @Override + public void onSuccess(Long result) { + updateCount.set(result); + } + }, + MoreExecutors.directExecutor()); + return updateCount; + } + }) + .commitAsync(); + assertThat(updateCount.get()).isEqualTo(UPDATE_COUNT); + assertThat(ts.get()).isNotNull(); + break; + } catch (AbortedException e) { + txn = mgr.resetForRetryAsync(); + } } } } @Test - public void asyncRunnerFireAndForgetInvalidUpdate() throws Exception { - AsyncRunner runner = client().runAsync(); - ApiFuture res = - runner.runAsync( - new AsyncWork() { - @Override - public ApiFuture doWorkAsync(TransactionContext txn) { - txn.executeUpdateAsync(INVALID_UPDATE_STATEMENT); - return txn.executeUpdateAsync(UPDATE_STATEMENT); - } - }, - executor); - assertThat(res.get()).isEqualTo(UPDATE_COUNT); + public void asyncTransactionManagerChain() throws Exception { + try (AsyncTransactionManager mgr = client().transactionManagerAsync()) { + TransactionContextFuture txn = mgr.beginAsync(); + while (true) { + try { + CommitTimestampFuture ts = + txn.then(AsyncTransactionManagerHelper.executeUpdateAsync(UPDATE_STATEMENT)) + .then( + AsyncTransactionManagerHelper.readRowAsync( + READ_TABLE_NAME, Key.of(1L), READ_COLUMN_NAMES)) + .then( + new AsyncTransactionFunction() { + @Override + public ApiFuture apply(TransactionContext txn, Struct input) + throws Exception { + return ApiFutures.immediateFuture(input.getString("Value")); + } + }) + .then( + new AsyncTransactionFunction() { + @Override + public ApiFuture apply(TransactionContext txn, String input) + throws Exception { + assertThat(input).isEqualTo("v1"); + return ApiFutures.immediateFuture(null); + } + }) + .commitAsync(); + assertThat(ts.get()).isNotNull(); + break; + } catch (AbortedException e) { + txn = mgr.resetForRetryAsync(); + } + } + } } @Test - public void asyncRunnerUpdateAborted() throws Exception { - try { - // Temporarily set the result of the update to 2 rows. - mockSpanner.putStatementResult(StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT + 1L)); - final AtomicInteger attempt = new AtomicInteger(); - AsyncRunner runner = client().runAsync(); - ApiFuture updateCount = - runner.runAsync( - new AsyncWork() { - @Override - public ApiFuture doWorkAsync(TransactionContext txn) { - if (attempt.incrementAndGet() == 1) { - mockSpanner.abortTransaction(txn); - } else { - // Set the result of the update statement back to 1 row. - mockSpanner.putStatementResult( - StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT)); - } - return txn.executeUpdateAsync(UPDATE_STATEMENT); - } - }, - executor); - assertThat(updateCount.get()).isEqualTo(UPDATE_COUNT); - assertThat(attempt.get()).isEqualTo(2); - } finally { - mockSpanner.putStatementResult(StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT)); + public void asyncTransactionManagerChainWithErrorInTheMiddle() throws Exception { + try (AsyncTransactionManager mgr = client().transactionManagerAsync()) { + TransactionContextFuture txn = mgr.beginAsync(); + while (true) { + try { + CommitTimestampFuture ts = + txn.then( + AsyncTransactionManagerHelper.executeUpdateAsync( + INVALID_UPDATE_STATEMENT)) + .then( + new AsyncTransactionFunction() { + @Override + public ApiFuture apply(TransactionContext txn, Long input) + throws Exception { + throw new IllegalStateException("this should not be executed"); + } + }) + .commitAsync(); + ts.get(); + break; + } catch (AbortedException e) { + txn = mgr.resetForRetryAsync(); + } catch (ExecutionException e) { + mgr.rollbackAsync(); + assertThat(e.getCause()).isInstanceOf(SpannerException.class); + SpannerException se = (SpannerException) e.getCause(); + assertThat(se.getErrorCode()).isEqualTo(ErrorCode.INVALID_ARGUMENT); + break; + } + } } } @Test - public void asyncRunnerCommitAborted() throws Exception { - try { + public void asyncTransactionManagerUpdateAborted() throws Exception { + try (AsyncTransactionManager mgr = client().transactionManagerAsync()) { // Temporarily set the result of the update to 2 rows. mockSpanner.putStatementResult(StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT + 1L)); final AtomicInteger attempt = new AtomicInteger(); - AsyncRunner runner = client().runAsync(); - ApiFuture updateCount = - runner.runAsync( - new AsyncWork() { - @Override - public ApiFuture doWorkAsync(final TransactionContext txn) { - if (attempt.get() > 0) { - // Set the result of the update statement back to 1 row. - mockSpanner.putStatementResult( - StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT)); - } - ApiFuture updateCount = txn.executeUpdateAsync(UPDATE_STATEMENT); - if (attempt.incrementAndGet() == 1) { - mockSpanner.abortTransaction(txn); - } - return updateCount; - } - }, - executor); - assertThat(updateCount.get()).isEqualTo(UPDATE_COUNT); + + TransactionContextFuture txn = mgr.beginAsync(); + while (true) { + try { + CommitTimestampFuture ts = + txn.then( + new AsyncTransactionFunction() { + @Override + public ApiFuture apply(TransactionContext txn, Void input) + throws Exception { + if (attempt.incrementAndGet() == 1) { + // Abort the first attempt. + mockSpanner.abortTransaction(txn); + } else { + // Set the result of the update statement back to 1 row. + mockSpanner.putStatementResult( + StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT)); + } + return ApiFutures.immediateFuture(null); + } + }) + .then(AsyncTransactionManagerHelper.executeUpdateAsync(UPDATE_STATEMENT)) + .commitAsync(); + assertThat(ts.get()).isNotNull(); + break; + } catch (AbortedException e) { + txn = mgr.resetForRetryAsync(); + } + } assertThat(attempt.get()).isEqualTo(2); } finally { mockSpanner.putStatementResult(StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT)); @@ -296,248 +400,361 @@ public ApiFuture doWorkAsync(final TransactionContext txn) { } @Test - public void asyncRunnerUpdateAbortedWithoutGettingResult() throws Exception { + public void asyncTransactionManagerUpdateAbortedWithoutGettingResult() throws Exception { final AtomicInteger attempt = new AtomicInteger(); - AsyncRunner runner = clientWithEmptySessionPool().runAsync(); - ApiFuture result = - runner.runAsync( - new AsyncWork() { - @Override - public ApiFuture doWorkAsync(TransactionContext txn) { - if (attempt.incrementAndGet() == 1) { - mockSpanner.abortTransaction(txn); - } - // This update statement will be aborted, but the error will not propagated to the - // transaction runner and cause the transaction to retry. Instead, the commit call - // will do that. - txn.executeUpdateAsync(UPDATE_STATEMENT); - // Resolving this future will not resolve the result of the entire transaction. The - // transaction result will be resolved when the commit has actually finished - // successfully. - return ApiFutures.immediateFuture(null); - } - }, - executor); - assertThat(result.get()).isNull(); - assertThat(attempt.get()).isEqualTo(2); - assertThat(mockSpanner.getRequestTypes()) - .containsExactly( - BatchCreateSessionsRequest.class, - BeginTransactionRequest.class, - ExecuteSqlRequest.class, - CommitRequest.class, - BeginTransactionRequest.class, - ExecuteSqlRequest.class, - CommitRequest.class); + try (AsyncTransactionManager mgr = clientWithEmptySessionPool().transactionManagerAsync()) { + TransactionContextFuture txn = mgr.beginAsync(); + while (true) { + try { + CommitTimestampFuture ts = + txn.then( + new AsyncTransactionFunction() { + @Override + public ApiFuture apply(TransactionContext txn, Void input) + throws Exception { + if (attempt.incrementAndGet() == 1) { + mockSpanner.abortTransaction(txn); + } + // This update statement will be aborted, but the error will not + // propagated to the + // transaction runner and cause the transaction to retry. Instead, the + // commit call + // will do that. + txn.executeUpdateAsync(UPDATE_STATEMENT); + // Resolving this future will not resolve the result of the entire + // transaction. The + // transaction result will be resolved when the commit has actually + // finished + // successfully. + return ApiFutures.immediateFuture(null); + } + }) + .commitAsync(); + assertThat(ts.get()).isNotNull(); + assertThat(attempt.get()).isEqualTo(2); + assertThat(mockSpanner.getRequestTypes()) + .containsExactly( + BatchCreateSessionsRequest.class, + BeginTransactionRequest.class, + ExecuteSqlRequest.class, + CommitRequest.class, + BeginTransactionRequest.class, + ExecuteSqlRequest.class, + CommitRequest.class); + break; + } catch (AbortedException e) { + txn = mgr.resetForRetryAsync(); + } + } + } } @Test - public void asyncRunnerCommitFails() throws Exception { + public void asyncTransactionManagerCommitFails() throws Exception { mockSpanner.setCommitExecutionTime( SimulatedExecutionTime.ofException( Status.RESOURCE_EXHAUSTED .withDescription("mutation limit exceeded") .asRuntimeException())); - AsyncRunner runner = client().runAsync(); - ApiFuture updateCount = - runner.runAsync( - new AsyncWork() { - @Override - public ApiFuture doWorkAsync(TransactionContext txn) { - // This statement will succeed, but the commit will fail. The error from the commit - // will bubble up to the future that is returned by the transaction, and the update - // count returned here will never reach the user application. - return txn.executeUpdateAsync(UPDATE_STATEMENT); - } - }, - executor); - try { - updateCount.get(); - fail("missing expected exception"); - } catch (ExecutionException e) { - assertThat(e.getCause()).isInstanceOf(SpannerException.class); - SpannerException se = (SpannerException) e.getCause(); - assertThat(se.getErrorCode()).isEqualTo(ErrorCode.RESOURCE_EXHAUSTED); - assertThat(se.getMessage()).contains("mutation limit exceeded"); + try (AsyncTransactionManager mgr = client().transactionManagerAsync()) { + TransactionContextFuture txn = mgr.beginAsync(); + while (true) { + try { + txn.then(AsyncTransactionManagerHelper.executeUpdateAsync(UPDATE_STATEMENT)) + .commitAsync() + .get(); + fail("missing expected exception"); + } catch (AbortedException e) { + txn = mgr.resetForRetryAsync(); + } catch (ExecutionException e) { + assertThat(e.getCause()).isInstanceOf(SpannerException.class); + SpannerException se = (SpannerException) e.getCause(); + assertThat(se.getErrorCode()).isEqualTo(ErrorCode.RESOURCE_EXHAUSTED); + assertThat(se.getMessage()).contains("mutation limit exceeded"); + break; + } + } } } @Test - public void asyncRunnerWaitsUntilAsyncUpdateHasFinished() throws Exception { - AsyncRunner runner = clientWithEmptySessionPool().runAsync(); - ApiFuture res = - runner.runAsync( - new AsyncWork() { - @Override - public ApiFuture doWorkAsync(TransactionContext txn) { - txn.executeUpdateAsync(UPDATE_STATEMENT); - return ApiFutures.immediateFuture(null); - } - }, - executor); - res.get(); - assertThat(mockSpanner.getRequestTypes()) - .containsExactly( - BatchCreateSessionsRequest.class, - BeginTransactionRequest.class, - ExecuteSqlRequest.class, - CommitRequest.class); + public void asyncTransactionManagerWaitsUntilAsyncUpdateHasFinished() throws Exception { + try (AsyncTransactionManager mgr = clientWithEmptySessionPool().transactionManagerAsync()) { + TransactionContextFuture txn = mgr.beginAsync(); + while (true) { + try { + txn.then( + new AsyncTransactionFunction() { + @Override + public ApiFuture apply(TransactionContext txn, Void input) + throws Exception { + // Shoot-and-forget update. The commit will still wait for this request to + // finish. + txn.executeUpdateAsync(UPDATE_STATEMENT); + return ApiFutures.immediateFuture(null); + } + }) + .commitAsync() + .get(); + assertThat(mockSpanner.getRequestTypes()) + .containsExactly( + BatchCreateSessionsRequest.class, + BeginTransactionRequest.class, + ExecuteSqlRequest.class, + CommitRequest.class); + break; + } catch (AbortedException e) { + txn = mgr.resetForRetryAsync(); + } + } + } } @Test - public void asyncRunnerBatchUpdate() throws Exception { - AsyncRunner runner = client().runAsync(); - ApiFuture updateCount = - runner.runAsync( - new AsyncWork() { - @Override - public ApiFuture doWorkAsync(TransactionContext txn) { - return txn.batchUpdateAsync(ImmutableList.of(UPDATE_STATEMENT, UPDATE_STATEMENT)); - } - }, - executor); - assertThat(updateCount.get()).asList().containsExactly(UPDATE_COUNT, UPDATE_COUNT); + public void asyncTransactionManagerBatchUpdate() throws Exception { + final SettableApiFuture result = SettableApiFuture.create(); + try (AsyncTransactionManager mgr = client().transactionManagerAsync()) { + TransactionContextFuture txn = mgr.beginAsync(); + while (true) { + try { + txn.then( + AsyncTransactionManagerHelper.batchUpdateAsync( + result, UPDATE_STATEMENT, UPDATE_STATEMENT)) + .commitAsync() + .get(); + break; + } catch (AbortedException e) { + txn = mgr.resetForRetryAsync(); + } + } + } + assertThat(result.get()).asList().containsExactly(UPDATE_COUNT, UPDATE_COUNT); } @Test - public void asyncRunnerIsNonBlockingWithBatchUpdate() throws Exception { + public void asyncTransactionManagerIsNonBlockingWithBatchUpdate() throws Exception { + SettableApiFuture res = SettableApiFuture.create(); mockSpanner.freeze(); - AsyncRunner runner = clientWithEmptySessionPool().runAsync(); - ApiFuture res = - runner.runAsync( - new AsyncWork() { - @Override - public ApiFuture doWorkAsync(TransactionContext txn) { - txn.batchUpdateAsync(ImmutableList.of(UPDATE_STATEMENT)); - return ApiFutures.immediateFuture(null); - } - }, - executor); - ApiFuture ts = runner.getCommitTimestamp(); - mockSpanner.unfreeze(); - assertThat(res.get()).isNull(); - assertThat(ts.get()).isNotNull(); + try (AsyncTransactionManager mgr = clientWithEmptySessionPool().transactionManagerAsync()) { + TransactionContextFuture txn = mgr.beginAsync(); + while (true) { + try { + CommitTimestampFuture ts = + txn.then(AsyncTransactionManagerHelper.batchUpdateAsync(res, UPDATE_STATEMENT)) + .commitAsync(); + mockSpanner.unfreeze(); + assertThat(ts.get()).isNotNull(); + assertThat(res.get()).asList().containsExactly(UPDATE_COUNT); + break; + } catch (AbortedException e) { + txn = mgr.resetForRetryAsync(); + } + } + } } @Test - public void asyncRunnerInvalidBatchUpdate() throws Exception { - AsyncRunner runner = client().runAsync(); - ApiFuture updateCount = - runner.runAsync( - new AsyncWork() { - @Override - public ApiFuture doWorkAsync(TransactionContext txn) { - return txn.batchUpdateAsync( - ImmutableList.of(UPDATE_STATEMENT, INVALID_UPDATE_STATEMENT)); - } - }, - executor); - try { - updateCount.get(); - fail("missing expected exception"); - } catch (ExecutionException e) { - assertThat(e.getCause()).isInstanceOf(SpannerException.class); - SpannerException se = (SpannerException) e.getCause(); - assertThat(se.getErrorCode()).isEqualTo(ErrorCode.INVALID_ARGUMENT); - assertThat(se.getMessage()).contains("invalid statement"); + public void asyncTransactionManagerInvalidBatchUpdate() throws Exception { + SettableApiFuture result = SettableApiFuture.create(); + try (AsyncTransactionManager mgr = client().transactionManagerAsync()) { + TransactionContextFuture txn = mgr.beginAsync(); + while (true) { + try { + txn.then( + AsyncTransactionManagerHelper.batchUpdateAsync( + result, UPDATE_STATEMENT, INVALID_UPDATE_STATEMENT)) + .commitAsync() + .get(); + fail("missing expected exception"); + } catch (AbortedException e) { + txn = mgr.resetForRetryAsync(); + } catch (ExecutionException e) { + assertThat(e.getCause()).isInstanceOf(SpannerException.class); + SpannerException se = (SpannerException) e.getCause(); + assertThat(se.getErrorCode()).isEqualTo(ErrorCode.INVALID_ARGUMENT); + assertThat(se.getMessage()).contains("invalid statement"); + break; + } + } } } @Test - public void asyncRunnerFireAndForgetInvalidBatchUpdate() throws Exception { - AsyncRunner runner = client().runAsync(); - ApiFuture res = - runner.runAsync( - new AsyncWork() { - @Override - public ApiFuture doWorkAsync(TransactionContext txn) { - txn.batchUpdateAsync(ImmutableList.of(UPDATE_STATEMENT, INVALID_UPDATE_STATEMENT)); - return txn.batchUpdateAsync(ImmutableList.of(UPDATE_STATEMENT, UPDATE_STATEMENT)); - } - }, - executor); - assertThat(res.get()).asList().containsExactly(UPDATE_COUNT, UPDATE_COUNT); + public void asyncTransactionManagerFireAndForgetInvalidBatchUpdate() throws Exception { + SettableApiFuture result = SettableApiFuture.create(); + try (AsyncTransactionManager mgr = clientWithEmptySessionPool().transactionManagerAsync()) { + TransactionContextFuture txn = mgr.beginAsync(); + while (true) { + try { + txn.then( + new AsyncTransactionFunction() { + @Override + public ApiFuture apply(TransactionContext txn, Void input) + throws Exception { + txn.batchUpdateAsync( + ImmutableList.of(UPDATE_STATEMENT, INVALID_UPDATE_STATEMENT)); + return ApiFutures.immediateFuture(null); + } + }) + .then( + AsyncTransactionManagerHelper.batchUpdateAsync( + result, UPDATE_STATEMENT, UPDATE_STATEMENT)) + .commitAsync() + .get(); + break; + } catch (AbortedException e) { + txn = mgr.resetForRetryAsync(); + } + } + } + assertThat(result.get()).asList().containsExactly(UPDATE_COUNT, UPDATE_COUNT); + assertThat(mockSpanner.getRequestTypes()) + .containsExactly( + BatchCreateSessionsRequest.class, + BeginTransactionRequest.class, + ExecuteBatchDmlRequest.class, + ExecuteBatchDmlRequest.class, + CommitRequest.class); } @Test - public void asyncRunnerBatchUpdateAborted() throws Exception { + public void asyncTransactionManagerBatchUpdateAborted() throws Exception { final AtomicInteger attempt = new AtomicInteger(); - AsyncRunner runner = client().runAsync(); - ApiFuture updateCount = - runner.runAsync( - new AsyncWork() { - @Override - public ApiFuture doWorkAsync(TransactionContext txn) { - if (attempt.incrementAndGet() == 1) { - return txn.batchUpdateAsync( - ImmutableList.of(UPDATE_STATEMENT, UPDATE_ABORTED_STATEMENT)); - } else { - return txn.batchUpdateAsync(ImmutableList.of(UPDATE_STATEMENT, UPDATE_STATEMENT)); - } - } - }, - executor); - assertThat(updateCount.get()).asList().containsExactly(UPDATE_COUNT, UPDATE_COUNT); + try (AsyncTransactionManager mgr = clientWithEmptySessionPool().transactionManagerAsync()) { + TransactionContextFuture txn = mgr.beginAsync(); + while (true) { + try { + txn.then( + new AsyncTransactionFunction() { + @Override + public ApiFuture apply(TransactionContext txn, Void input) + throws Exception { + if (attempt.incrementAndGet() == 1) { + return txn.batchUpdateAsync( + ImmutableList.of(UPDATE_STATEMENT, UPDATE_ABORTED_STATEMENT)); + } else { + return txn.batchUpdateAsync( + ImmutableList.of(UPDATE_STATEMENT, UPDATE_STATEMENT)); + } + } + }) + .commitAsync() + .get(); + break; + } catch (AbortedException e) { + txn = mgr.resetForRetryAsync(); + } + } + } assertThat(attempt.get()).isEqualTo(2); + // There should only be 1 CommitRequest, as the first attempt should abort already after the + // ExecuteBatchDmlRequest. + assertThat(mockSpanner.getRequestTypes()) + .containsExactly( + BatchCreateSessionsRequest.class, + BeginTransactionRequest.class, + ExecuteBatchDmlRequest.class, + BeginTransactionRequest.class, + ExecuteBatchDmlRequest.class, + CommitRequest.class); } @Test - public void asyncRunnerWithBatchUpdateCommitAborted() throws Exception { - try { + public void asyncTransactionManagerWithBatchUpdateCommitAborted() throws Exception { + try (AsyncTransactionManager mgr = clientWithEmptySessionPool().transactionManagerAsync()) { // Temporarily set the result of the update to 2 rows. mockSpanner.putStatementResult(StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT + 1L)); final AtomicInteger attempt = new AtomicInteger(); - AsyncRunner runner = client().runAsync(); - ApiFuture updateCount = - runner.runAsync( - new AsyncWork() { - @Override - public ApiFuture doWorkAsync(final TransactionContext txn) { - if (attempt.get() > 0) { - // Set the result of the update statement back to 1 row. - mockSpanner.putStatementResult( - StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT)); - } - ApiFuture updateCount = - txn.batchUpdateAsync(ImmutableList.of(UPDATE_STATEMENT, UPDATE_STATEMENT)); - if (attempt.incrementAndGet() == 1) { - mockSpanner.abortTransaction(txn); - } - return updateCount; - } - }, - executor); - assertThat(updateCount.get()).asList().containsExactly(UPDATE_COUNT, UPDATE_COUNT); - assertThat(attempt.get()).isEqualTo(2); + TransactionContextFuture txn = mgr.beginAsync(); + while (true) { + final SettableApiFuture result = SettableApiFuture.create(); + try { + txn.then( + new AsyncTransactionFunction() { + @Override + public ApiFuture apply(TransactionContext txn, Void input) + throws Exception { + if (attempt.get() > 0) { + // Set the result of the update statement back to 1 row. + mockSpanner.putStatementResult( + StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT)); + } + return ApiFutures.immediateFuture(null); + } + }) + .then( + AsyncTransactionManagerHelper.batchUpdateAsync( + result, UPDATE_STATEMENT, UPDATE_STATEMENT)) + .then( + new AsyncTransactionFunction() { + @Override + public ApiFuture apply(TransactionContext txn, long[] input) + throws Exception { + if (attempt.incrementAndGet() == 1) { + mockSpanner.abortTransaction(txn); + } + return ApiFutures.immediateFuture(null); + } + }) + .commitAsync() + .get(); + assertThat(result.get()).asList().containsExactly(UPDATE_COUNT, UPDATE_COUNT); + assertThat(attempt.get()).isEqualTo(2); + break; + } catch (AbortedException e) { + txn = mgr.resetForRetryAsync(); + } + } } finally { mockSpanner.putStatementResult(StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT)); } + assertThat(mockSpanner.getRequestTypes()) + .containsExactly( + BatchCreateSessionsRequest.class, + BeginTransactionRequest.class, + ExecuteBatchDmlRequest.class, + CommitRequest.class, + BeginTransactionRequest.class, + ExecuteBatchDmlRequest.class, + CommitRequest.class); } @Test - public void asyncRunnerBatchUpdateAbortedWithoutGettingResult() throws Exception { + public void asyncTransactionManagerBatchUpdateAbortedWithoutGettingResult() throws Exception { final AtomicInteger attempt = new AtomicInteger(); - AsyncRunner runner = clientWithEmptySessionPool().runAsync(); - ApiFuture result = - runner.runAsync( - new AsyncWork() { - @Override - public ApiFuture doWorkAsync(TransactionContext txn) { - if (attempt.incrementAndGet() == 1) { - mockSpanner.abortTransaction(txn); - } - // This update statement will be aborted, but the error will not propagated to the - // transaction runner and cause the transaction to retry. Instead, the commit call - // will do that. - txn.batchUpdateAsync(ImmutableList.of(UPDATE_STATEMENT, UPDATE_STATEMENT)); - // Resolving this future will not resolve the result of the entire transaction. The - // transaction result will be resolved when the commit has actually finished - // successfully. - return ApiFutures.immediateFuture(null); - } - }, - executor); - assertThat(result.get()).isNull(); + try (AsyncTransactionManager mgr = clientWithEmptySessionPool().transactionManagerAsync()) { + TransactionContextFuture txn = mgr.beginAsync(); + while (true) { + try { + txn.then( + new AsyncTransactionFunction() { + @Override + public ApiFuture apply(TransactionContext txn, Void input) + throws Exception { + if (attempt.incrementAndGet() == 1) { + mockSpanner.abortTransaction(txn); + } + // This update statement will be aborted, but the error will not propagated to + // the + // transaction runner and cause the transaction to retry. Instead, the commit + // call + // will do that. + txn.batchUpdateAsync(ImmutableList.of(UPDATE_STATEMENT, UPDATE_STATEMENT)); + // Resolving this future will not resolve the result of the entire + // transaction. The + // transaction result will be resolved when the commit has actually finished + // successfully. + return ApiFutures.immediateFuture(null); + } + }) + .commitAsync() + .get(); + break; + } catch (AbortedException e) { + txn = mgr.resetForRetryAsync(); + } + } + } assertThat(attempt.get()).isEqualTo(2); assertThat(mockSpanner.getRequestTypes()) .containsExactly( @@ -551,50 +768,64 @@ public ApiFuture doWorkAsync(TransactionContext txn) { } @Test - public void asyncRunnerWithBatchUpdateCommitFails() throws Exception { + public void asyncTransactionManagerWithBatchUpdateCommitFails() throws Exception { mockSpanner.setCommitExecutionTime( SimulatedExecutionTime.ofException( Status.RESOURCE_EXHAUSTED .withDescription("mutation limit exceeded") .asRuntimeException())); - AsyncRunner runner = client().runAsync(); - ApiFuture updateCount = - runner.runAsync( - new AsyncWork() { - @Override - public ApiFuture doWorkAsync(TransactionContext txn) { - // This statement will succeed, but the commit will fail. The error from the commit - // will bubble up to the future that is returned by the transaction, and the update - // count returned here will never reach the user application. - return txn.batchUpdateAsync(ImmutableList.of(UPDATE_STATEMENT, UPDATE_STATEMENT)); - } - }, - executor); - try { - updateCount.get(); - fail("missing expected exception"); - } catch (ExecutionException e) { - assertThat(e.getCause()).isInstanceOf(SpannerException.class); - SpannerException se = (SpannerException) e.getCause(); - assertThat(se.getErrorCode()).isEqualTo(ErrorCode.RESOURCE_EXHAUSTED); - assertThat(se.getMessage()).contains("mutation limit exceeded"); + try (AsyncTransactionManager mgr = clientWithEmptySessionPool().transactionManagerAsync()) { + TransactionContextFuture txn = mgr.beginAsync(); + while (true) { + try { + txn.then( + AsyncTransactionManagerHelper.batchUpdateAsync( + UPDATE_STATEMENT, UPDATE_STATEMENT)) + .commitAsync() + .get(); + fail("missing expected exception"); + } catch (AbortedException e) { + txn = mgr.resetForRetryAsync(); + } catch (ExecutionException e) { + assertThat(e.getCause()).isInstanceOf(SpannerException.class); + SpannerException se = (SpannerException) e.getCause(); + assertThat(se.getErrorCode()).isEqualTo(ErrorCode.RESOURCE_EXHAUSTED); + assertThat(se.getMessage()).contains("mutation limit exceeded"); + break; + } + } } + assertThat(mockSpanner.getRequestTypes()) + .containsExactly( + BatchCreateSessionsRequest.class, + BeginTransactionRequest.class, + ExecuteBatchDmlRequest.class, + CommitRequest.class); } @Test - public void asyncRunnerWaitsUntilAsyncBatchUpdateHasFinished() throws Exception { - AsyncRunner runner = clientWithEmptySessionPool().runAsync(); - ApiFuture res = - runner.runAsync( - new AsyncWork() { - @Override - public ApiFuture doWorkAsync(TransactionContext txn) { - txn.batchUpdateAsync(ImmutableList.of(UPDATE_STATEMENT)); - return ApiFutures.immediateFuture(null); - } - }, - executor); - res.get(); + public void asyncTransactionManagerWaitsUntilAsyncBatchUpdateHasFinished() throws Exception { + try (AsyncTransactionManager mgr = clientWithEmptySessionPool().transactionManagerAsync()) { + TransactionContextFuture txn = mgr.beginAsync(); + while (true) { + try { + txn.then( + new AsyncTransactionFunction() { + @Override + public ApiFuture apply(TransactionContext txn, Void input) + throws Exception { + txn.batchUpdateAsync(ImmutableList.of(UPDATE_STATEMENT)); + return ApiFutures.immediateFuture(null); + } + }) + .commitAsync() + .get(); + break; + } catch (AbortedException e) { + txn = mgr.resetForRetryAsync(); + } + } + } assertThat(mockSpanner.getRequestTypes()) .containsExactly( BatchCreateSessionsRequest.class, @@ -604,98 +835,33 @@ public ApiFuture doWorkAsync(TransactionContext txn) { } @Test - public void closeTransactionBeforeEndOfAsyncQuery() throws Exception { - final BlockingQueue results = new SynchronousQueue<>(); - final SettableApiFuture finished = SettableApiFuture.create(); - DatabaseClientImpl clientImpl = (DatabaseClientImpl) client(); - - // There should currently not be any sessions checked out of the pool. - assertThat(clientImpl.pool.getNumberOfSessionsInUse()).isEqualTo(0); - - AsyncRunner runner = clientImpl.runAsync(); - final CountDownLatch dataReceived = new CountDownLatch(1); - final CountDownLatch dataChecked = new CountDownLatch(1); - ApiFuture res = - runner.runAsync( - new AsyncWork() { - @Override - public ApiFuture doWorkAsync(TransactionContext txn) { - try (AsyncResultSet rs = - txn.readAsync( - READ_TABLE_NAME, KeySet.all(), READ_COLUMN_NAMES, Options.bufferRows(1))) { - rs.setCallback( - Executors.newSingleThreadExecutor(), - new ReadyCallback() { - @Override - public CallbackResponse cursorReady(AsyncResultSet resultSet) { - dataReceived.countDown(); - try { - while (true) { - switch (resultSet.tryNext()) { - case DONE: - finished.set(true); - return CallbackResponse.DONE; - case NOT_READY: - return CallbackResponse.CONTINUE; - case OK: - dataChecked.await(); - results.put(resultSet.getString(0)); - } + public void asyncTransactionManagerReadRow() throws Exception { + ApiFuture val; + try (AsyncTransactionManager mgr = client().transactionManagerAsync()) { + TransactionContextFuture txn = mgr.beginAsync(); + while (true) { + try { + AsyncTransactionStep step; + val = + step = + txn.then( + AsyncTransactionManagerHelper.readRowAsync( + READ_TABLE_NAME, Key.of(1L), READ_COLUMN_NAMES)) + .then( + new AsyncTransactionFunction() { + @Override + public ApiFuture apply(TransactionContext txn, Struct input) + throws Exception { + return ApiFutures.immediateFuture(input.getString("Value")); } - } catch (Throwable t) { - finished.setException(t); - return CallbackResponse.DONE; - } - } - }); - } - try { - dataReceived.await(); - return ApiFutures.immediateFuture(null); - } catch (InterruptedException e) { - return ApiFutures.immediateFailedFuture( - SpannerExceptionFactory.propagateInterrupt(e)); - } - } - }, - executor); - // Wait until at least one row has been fetched. At that moment there should be one session - // checked out. - dataReceived.await(); - assertThat(clientImpl.pool.getNumberOfSessionsInUse()).isEqualTo(1); - assertThat(res.isDone()).isFalse(); - dataChecked.countDown(); - // Get the data from the transaction. - List resultList = new ArrayList<>(); - do { - results.drainTo(resultList); - } while (!finished.isDone() || results.size() > 0); - assertThat(finished.get()).isTrue(); - assertThat(resultList).containsExactly("k1", "k2", "k3"); - assertThat(res.get()).isNull(); - assertThat(clientImpl.pool.getNumberOfSessionsInUse()).isEqualTo(0); - } - - @Test - public void asyncRunnerReadRow() throws Exception { - AsyncRunner runner = client().runAsync(); - ApiFuture val = - runner.runAsync( - new AsyncWork() { - @Override - public ApiFuture doWorkAsync(TransactionContext txn) { - return ApiFutures.transform( - txn.readRowAsync(READ_TABLE_NAME, Key.of(1L), READ_COLUMN_NAMES), - new ApiFunction() { - @Override - public String apply(Struct input) { - return input.getString("Value"); - } - }, - MoreExecutors.directExecutor()); - } - }, - executor); + }); + step.commitAsync().get(); + break; + } catch (AbortedException e) { + txn = mgr.resetForRetryAsync(); + } + } + } assertThat(val.get()).isEqualTo("v1"); } From 8a0ad3ff5db6c35d7a36a277d6b9154d47dd9c42 Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Mon, 15 Jun 2020 08:22:39 +0200 Subject: [PATCH 39/49] fix: fix test failures --- .../spanner/AsyncResultSetImplStressTest.java | 5 +++++ .../spanner/AsyncTransactionManagerTest.java | 13 ++++++------- .../cloud/spanner/MockSpannerServiceImpl.java | 18 +++++++++++++++++- 3 files changed, 28 insertions(+), 8 deletions(-) diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncResultSetImplStressTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncResultSetImplStressTest.java index 0bf7a1a2c96..8c195bb6da8 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncResultSetImplStressTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncResultSetImplStressTest.java @@ -45,7 +45,9 @@ import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; +import org.junit.Rule; import org.junit.Test; +import org.junit.rules.Timeout; import org.junit.runner.RunWith; import org.junit.runners.Parameterized; import org.junit.runners.Parameterized.Parameter; @@ -55,6 +57,9 @@ public class AsyncResultSetImplStressTest { private static final int TEST_RUNS = 1000; + /** Timeout is applied to each test case individually. */ + @Rule public Timeout timeout = new Timeout(30, TimeUnit.SECONDS); + @Parameter(0) public int resultSetSize; diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncTransactionManagerTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncTransactionManagerTest.java index 65681eda092..9d4b5dac2a7 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncTransactionManagerTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncTransactionManagerTest.java @@ -735,15 +735,14 @@ public ApiFuture apply(TransactionContext txn, Void input) mockSpanner.abortTransaction(txn); } // This update statement will be aborted, but the error will not propagated to - // the - // transaction runner and cause the transaction to retry. Instead, the commit - // call - // will do that. + // the transaction manager and cause the transaction to retry. Instead, the + // commit call will do that. txn.batchUpdateAsync(ImmutableList.of(UPDATE_STATEMENT, UPDATE_STATEMENT)); + // Wait for the request to arrive at the server. + mockSpanner.waitForLastRequestToBe(ExecuteBatchDmlRequest.class, 1000L); // Resolving this future will not resolve the result of the entire - // transaction. The - // transaction result will be resolved when the commit has actually finished - // successfully. + // transaction. The transaction result will be resolved when the commit has + // actually finished successfully. return ApiFutures.immediateFuture(null); } }) 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 73d903de5b9..25c3428d386 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 @@ -23,6 +23,7 @@ import com.google.cloud.spanner.TransactionRunnerImpl.TransactionContextImpl; import com.google.common.base.Optional; import com.google.common.base.Preconditions; +import com.google.common.base.Stopwatch; import com.google.common.base.Throwables; import com.google.common.util.concurrent.Uninterruptibles; import com.google.protobuf.AbstractMessage; @@ -90,10 +91,12 @@ import java.util.Random; import java.util.UUID; import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentLinkedDeque; import java.util.concurrent.ConcurrentLinkedQueue; import java.util.concurrent.ConcurrentMap; import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicLong; @@ -493,7 +496,7 @@ private static void checkException(Queue exceptions, boolean keepExce private double abortProbability = 0.0010D; private final Object lock = new Object(); - private final Queue requests = new ConcurrentLinkedQueue<>(); + private final Deque requests = new ConcurrentLinkedDeque<>(); private volatile CountDownLatch freezeLock = new CountDownLatch(0); private final Queue exceptions = new ConcurrentLinkedQueue<>(); private boolean stickyGlobalExceptions = false; @@ -1820,6 +1823,19 @@ public int countRequestsOfType(Class type) { return c; } + public void waitForLastRequestToBe(Class type, long timeoutMillis) + throws InterruptedException, TimeoutException { + Stopwatch watch = Stopwatch.createStarted(); + while (!(this.requests.peekLast() != null + && this.requests.peekLast().getClass().equals(type))) { + Thread.sleep(10L); + if (watch.elapsed(TimeUnit.MILLISECONDS) > timeoutMillis) { + throw new TimeoutException( + "Timeout while waiting for last request to become " + type.getName()); + } + } + } + @Override public void addResponse(AbstractMessage response) { throw new UnsupportedOperationException(); From 80224575e37ed9d0d906ae2de5d1bfbc1d35e370 Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Mon, 15 Jun 2020 22:29:03 +0200 Subject: [PATCH 40/49] fix: fix several race conditions --- .../spanner/AsyncTransactionManagerImpl.java | 21 +- .../com/google/cloud/spanner/SessionPool.java | 129 +----------- .../SessionPoolAsyncTransactionManager.java | 194 ++++++++++++++++++ .../spanner/TransactionContextFutureImpl.java | 37 ++-- .../spanner/AbstractAsyncTransactionTest.java | 49 +++-- .../spanner/AsyncTransactionManagerTest.java | 46 +++-- .../cloud/spanner/MockSpannerServiceImpl.java | 48 +++-- 7 files changed, 303 insertions(+), 221 deletions(-) create mode 100644 google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPoolAsyncTransactionManager.java diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManagerImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManagerImpl.java index 31d809dad69..aca7f6e8cea 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManagerImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManagerImpl.java @@ -25,7 +25,6 @@ import com.google.cloud.spanner.TransactionManager.TransactionState; import com.google.common.base.Preconditions; import com.google.common.util.concurrent.MoreExecutors; -import io.opencensus.common.Scope; import io.opencensus.trace.Span; import io.opencensus.trace.Tracer; import io.opencensus.trace.Tracing; @@ -59,15 +58,17 @@ public void close() { @Override public TransactionContextFuture beginAsync() { Preconditions.checkState(txn == null, "begin can only be called once"); - TransactionContextFuture begin = new TransactionContextFutureImpl(this, internalBeginAsync()); - session.setActive(this); + TransactionContextFuture begin = + new TransactionContextFutureImpl(this, internalBeginAsync(true)); return begin; } - private ApiFuture internalBeginAsync() { - final Scope s = tracer.withSpan(span); - txn = session.newTransaction(); + private ApiFuture internalBeginAsync(boolean setActive) { txnState = TransactionState.STARTED; + txn = session.newTransaction(); + if (setActive) { + session.setActive(this); + } final SettableApiFuture res = SettableApiFuture.create(); final ApiFuture fut = txn.ensureTxnAsync(); ApiFutures.addCallback( @@ -75,15 +76,11 @@ private ApiFuture internalBeginAsync() { new ApiFutureCallback() { @Override public void onFailure(Throwable t) { - s.close(); - TraceUtil.endSpanWithFailure(span, t); res.setException(SpannerExceptionFactory.newSpannerException(t)); } @Override public void onSuccess(Void result) { - s.close(); - span.end(); res.set(txn); } }, @@ -95,7 +92,7 @@ public void onSuccess(Void result) { public ApiFuture commitAsync() { Preconditions.checkState( txnState == TransactionState.STARTED, - "commit can only be invoked if the transaction is in progress"); + "commit can only be invoked if the transaction is in progress. Current state: " + txnState); if (txn.isAborted()) { txnState = TransactionState.ABORTED; return ApiFutures.immediateFailedFuture( @@ -144,7 +141,7 @@ public TransactionContextFuture resetForRetryAsync() { throw new IllegalStateException( "resetForRetry can only be called if the previous attempt aborted"); } - return new TransactionContextFutureImpl(this, internalBeginAsync()); + return new TransactionContextFutureImpl(this, internalBeginAsync(false)); } @Override diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java index 711a905ac98..b6ffa4da8ef 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java @@ -38,7 +38,6 @@ import static com.google.cloud.spanner.MetricRegistryConstants.SPANNER_LABEL_KEYS_WITH_TYPE; import static com.google.cloud.spanner.SpannerExceptionFactory.newSpannerException; -import com.google.api.core.ApiAsyncFunction; import com.google.api.core.ApiFuture; import com.google.api.core.ApiFutures; import com.google.api.core.SettableApiFuture; @@ -809,7 +808,9 @@ public void close() { } closed = true; try { - delegate.close(); + if (delegate != null) { + delegate.close(); + } } finally { session.close(); } @@ -948,130 +949,6 @@ public ApiFuture getCommitTimestamp() { } } - private static class SessionPoolAsyncTransactionManager implements AsyncTransactionManager { - private volatile PooledSessionFuture session; - private final SettableApiFuture delegate = SettableApiFuture.create(); - private volatile ApiFuture commitTimestamp; - - private SessionPoolAsyncTransactionManager(PooledSessionFuture session) { - this.session = session; - this.session.addListener( - new Runnable() { - @Override - public void run() { - try { - delegate.set( - SessionPoolAsyncTransactionManager.this - .session - .get() - .transactionManagerAsync()); - } catch (Throwable t) { - delegate.setException(t); - } - } - }, - MoreExecutors.directExecutor()); - } - - @Override - public void close() { - delegate.addListener( - new Runnable() { - @Override - public void run() { - session.close(); - } - }, - MoreExecutors.directExecutor()); - } - - @Override - public TransactionContextFuture beginAsync() { - return new TransactionContextFutureImpl( - this, - ApiFutures.transformAsync( - delegate, - new ApiAsyncFunction() { - @Override - public ApiFuture apply(AsyncTransactionManager input) { - return input.beginAsync(); - } - }, - MoreExecutors.directExecutor())); - } - - @Override - public ApiFuture commitAsync() { - return ApiFutures.transformAsync( - delegate, - new ApiAsyncFunction() { - @Override - public ApiFuture apply(AsyncTransactionManager input) throws Exception { - ApiFuture res = input.commitAsync(); - res.addListener( - new Runnable() { - @Override - public void run() { - session.close(); - } - }, - MoreExecutors.directExecutor()); - return res; - } - }, - MoreExecutors.directExecutor()); - } - - @Override - public ApiFuture rollbackAsync() { - return ApiFutures.transformAsync( - delegate, - new ApiAsyncFunction() { - @Override - public ApiFuture apply(AsyncTransactionManager input) throws Exception { - ApiFuture res = input.rollbackAsync(); - res.addListener( - new Runnable() { - @Override - public void run() { - session.close(); - } - }, - MoreExecutors.directExecutor()); - return res; - } - }, - MoreExecutors.directExecutor()); - } - - @Override - public TransactionContextFuture resetForRetryAsync() { - return new TransactionContextFutureImpl( - this, - ApiFutures.transformAsync( - delegate, - new ApiAsyncFunction() { - @Override - public ApiFuture apply(AsyncTransactionManager input) - throws Exception { - return input.resetForRetryAsync(); - } - }, - MoreExecutors.directExecutor())); - } - - @Override - public TransactionState getState() { - try { - return delegate.get().getState(); - } catch (InterruptedException e) { - throw SpannerExceptionFactory.propagateInterrupt(e); - } catch (ExecutionException e) { - throw SpannerExceptionFactory.newSpannerException(e.getCause() == null ? e : e.getCause()); - } - } - } - // Exception class used just to track the stack trace at the point when a session was handed out // from the pool. final class LeakedSessionException extends RuntimeException { diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPoolAsyncTransactionManager.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPoolAsyncTransactionManager.java new file mode 100644 index 00000000000..005dfe7b600 --- /dev/null +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPoolAsyncTransactionManager.java @@ -0,0 +1,194 @@ +/* + * 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.ApiAsyncFunction; +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.AsyncTransactionManager.TransactionContextFuture; +import com.google.cloud.spanner.SessionPool.PooledSessionFuture; +import com.google.cloud.spanner.TransactionManager.TransactionState; +import com.google.common.base.Preconditions; +import com.google.common.util.concurrent.MoreExecutors; +import java.util.concurrent.ExecutionException; + +class SessionPoolAsyncTransactionManager implements AsyncTransactionManager { + private TransactionState txnState; + private volatile PooledSessionFuture session; + private final SettableApiFuture delegate = SettableApiFuture.create(); + + SessionPoolAsyncTransactionManager(PooledSessionFuture session) { + this.session = session; + this.session.addListener( + new Runnable() { + @Override + public void run() { + try { + delegate.set( + SessionPoolAsyncTransactionManager.this.session.get().transactionManagerAsync()); + } catch (Throwable t) { + delegate.setException(t); + } + } + }, + MoreExecutors.directExecutor()); + } + + @Override + public void close() { + delegate.addListener( + new Runnable() { + @Override + public void run() { + session.close(); + } + }, + MoreExecutors.directExecutor()); + } + + @Override + public TransactionContextFuture beginAsync() { + Preconditions.checkState(txnState == null, "begin can only be called once"); + txnState = TransactionState.STARTED; + final SettableApiFuture delegateTxnFuture = SettableApiFuture.create(); + ApiFutures.addCallback( + delegate, + new ApiFutureCallback() { + @Override + public void onFailure(Throwable t) { + delegateTxnFuture.setException(t); + } + + @Override + public void onSuccess(AsyncTransactionManager result) { + ApiFutures.addCallback( + result.beginAsync(), + new ApiFutureCallback() { + @Override + public void onFailure(Throwable t) { + delegateTxnFuture.setException(t); + } + + @Override + public void onSuccess(TransactionContext result) { + delegateTxnFuture.set(result); + } + }, + MoreExecutors.directExecutor()); + } + }, + MoreExecutors.directExecutor()); + return new TransactionContextFutureImpl(this, delegateTxnFuture); + + // return new TransactionContextFutureImpl( + // this, + // ApiFutures.transformAsync( + // delegate, + // new ApiAsyncFunction() { + // @Override + // public ApiFuture apply(AsyncTransactionManager input) { + // return input.beginAsync(); + // } + // }, + // MoreExecutors.directExecutor())); + } + + @Override + public ApiFuture commitAsync() { + Preconditions.checkState( + txnState == TransactionState.STARTED, + "commit can only be invoked if the transaction is in progress. Current state: " + txnState); + txnState = TransactionState.COMMITTED; + return ApiFutures.transformAsync( + delegate, + new ApiAsyncFunction() { + @Override + public ApiFuture apply(AsyncTransactionManager input) throws Exception { + ApiFuture res = input.commitAsync(); + // res.addListener( + // new Runnable() { + // @Override + // public void run() { + // session.close(); + // } + // }, + // MoreExecutors.directExecutor()); + return res; + } + }, + MoreExecutors.directExecutor()); + } + + @Override + public ApiFuture rollbackAsync() { + Preconditions.checkState( + txnState == TransactionState.STARTED, + "rollback can only be called if the transaction is in progress"); + txnState = TransactionState.ROLLED_BACK; + return ApiFutures.transformAsync( + delegate, + new ApiAsyncFunction() { + @Override + public ApiFuture apply(AsyncTransactionManager input) throws Exception { + ApiFuture res = input.rollbackAsync(); + res.addListener( + new Runnable() { + @Override + public void run() { + session.close(); + } + }, + MoreExecutors.directExecutor()); + return res; + } + }, + MoreExecutors.directExecutor()); + } + + @Override + public TransactionContextFuture resetForRetryAsync() { + Preconditions.checkState( + txnState != null, "resetForRetry can only be called after the transaction has started."); + txnState = TransactionState.STARTED; + return new TransactionContextFutureImpl( + this, + ApiFutures.transformAsync( + delegate, + new ApiAsyncFunction() { + @Override + public ApiFuture apply(AsyncTransactionManager input) + throws Exception { + return input.resetForRetryAsync(); + } + }, + MoreExecutors.directExecutor())); + } + + @Override + public TransactionState getState() { + try { + return delegate.get().getState(); + } catch (InterruptedException e) { + throw SpannerExceptionFactory.propagateInterrupt(e); + } catch (ExecutionException e) { + throw SpannerExceptionFactory.newSpannerException(e.getCause() == null ? e : e.getCause()); + } + } +} diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionContextFutureImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionContextFutureImpl.java index 487b9616518..3cdd6430424 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionContextFutureImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionContextFutureImpl.java @@ -16,7 +16,6 @@ package com.google.cloud.spanner; -import com.google.api.core.ApiFunction; import com.google.api.core.ApiFuture; import com.google.api.core.ApiFutureCallback; import com.google.api.core.ApiFutures; @@ -35,6 +34,13 @@ class TransactionContextFutureImpl extends ForwardingApiFuture implements TransactionContextFuture { + + /** + * {@link ApiFuture} that returns a commit timestamp. Any {@link AbortedException} that is thrown + * by either the commit call or any other rpc during the transaction will be thrown by the {@link + * #get()} method of this future as an {@link AbortedException} and not as an {@link + * ExecutionException} with an {@link AbortedException} as its cause. + */ static class CommitTimestampFutureImpl extends ForwardingApiFuture implements CommitTimestampFuture { CommitTimestampFutureImpl(ApiFuture delegate) { @@ -172,36 +178,21 @@ public void onSuccess(Timestamp result) { @Override public AsyncTransactionStatementImpl then( AsyncTransactionFunction function) { - return new AsyncTransactionStatementImpl<>( - this, - ApiFutures.transform( - this, - new ApiFunction() { - @Override - public Void apply(TransactionContext input) { - return null; - } - }, - MoreExecutors.directExecutor()), - function); - } - - ApiFuture commitAsync() { - ApiFuture res = mgr.commitAsync(); + final SettableApiFuture input = SettableApiFuture.create(); ApiFutures.addCallback( - res, - new ApiFutureCallback() { + this, + new ApiFutureCallback() { @Override public void onFailure(Throwable t) { - txnResult.setException(t); + input.setException(t); } @Override - public void onSuccess(Timestamp result) { - txnResult.set(result); + public void onSuccess(TransactionContext result) { + input.set(null); } }, MoreExecutors.directExecutor()); - return txnResult; + return new AsyncTransactionStatementImpl<>(this, input, function); } } diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AbstractAsyncTransactionTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AbstractAsyncTransactionTest.java index 1926860162a..bf76ea4f392 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AbstractAsyncTransactionTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AbstractAsyncTransactionTest.java @@ -30,15 +30,16 @@ import static com.google.cloud.spanner.MockSpannerTestUtil.UPDATE_COUNT; import static com.google.cloud.spanner.MockSpannerTestUtil.UPDATE_STATEMENT; -import com.google.api.gax.grpc.testing.LocalChannelProvider; +import com.google.api.core.ApiFunction; import com.google.cloud.NoCredentials; import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult; +import io.grpc.ManagedChannelBuilder; import io.grpc.Server; import io.grpc.Status; -import io.grpc.inprocess.InProcessServerBuilder; +import io.grpc.netty.shaded.io.grpc.netty.NettyServerBuilder; +import java.net.InetSocketAddress; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; -import java.util.concurrent.ScheduledThreadPoolExecutor; import org.junit.After; import org.junit.AfterClass; import org.junit.Before; @@ -48,7 +49,7 @@ public abstract class AbstractAsyncTransactionTest { static MockSpannerServiceImpl mockSpanner; private static Server server; - private static LocalChannelProvider channelProvider; + private static InetSocketAddress address; static ExecutorService executor; Spanner spanner; @@ -74,14 +75,9 @@ public static void setup() throws Exception { StatementResult.exception( UPDATE_ABORTED_STATEMENT, Status.ABORTED.withDescription("Transaction was aborted").asRuntimeException())); - String uniqueName = InProcessServerBuilder.generateName(); - server = - InProcessServerBuilder.forName(uniqueName) - .scheduledExecutorService(new ScheduledThreadPoolExecutor(1)) - .addService(mockSpanner) - .build() - .start(); - channelProvider = LocalChannelProvider.create(uniqueName); + + address = new InetSocketAddress("localhost", 0); + server = NettyServerBuilder.forAddress(address).addService(mockSpanner).build().start(); executor = Executors.newSingleThreadExecutor(); } @@ -93,11 +89,20 @@ public static void teardown() throws Exception { } @Before - public void before() { + public void before() throws Exception { + String endpoint = address.getHostString() + ":" + server.getPort(); spanner = SpannerOptions.newBuilder() .setProjectId(TEST_PROJECT) - .setChannelProvider(channelProvider) + .setChannelConfigurator( + new ApiFunction() { + @Override + public ManagedChannelBuilder apply(ManagedChannelBuilder input) { + input.usePlaintext(); + return input; + } + }) + .setHost("http://" + endpoint) .setCredentials(NoCredentials.getInstance()) .setSessionPoolOption(SessionPoolOptions.newBuilder().setFailOnSessionLeak().build()) .build() @@ -116,6 +121,14 @@ public void before() { .getService(); } + @After + public void after() throws Exception { + spanner.close(); + spannerWithEmptySessionPool.close(); + mockSpanner.removeAllExecutionTimes(); + mockSpanner.reset(); + } + DatabaseClient client() { return spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); } @@ -124,12 +137,4 @@ DatabaseClient clientWithEmptySessionPool() { return spannerWithEmptySessionPool.getDatabaseClient( DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); } - - @After - public void after() { - spanner.close(); - spannerWithEmptySessionPool.close(); - mockSpanner.removeAllExecutionTimes(); - mockSpanner.reset(); - } } diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncTransactionManagerTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncTransactionManagerTest.java index 9d4b5dac2a7..76eec9f1173 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncTransactionManagerTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncTransactionManagerTest.java @@ -38,7 +38,10 @@ import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult; import com.google.common.base.Function; import com.google.common.collect.ImmutableList; +import com.google.common.collect.Iterables; +import com.google.common.collect.Range; import com.google.common.util.concurrent.MoreExecutors; +import com.google.protobuf.AbstractMessage; import com.google.spanner.v1.BatchCreateSessionsRequest; import com.google.spanner.v1.BeginTransactionRequest; import com.google.spanner.v1.CommitRequest; @@ -55,7 +58,6 @@ @RunWith(JUnit4.class) public class AsyncTransactionManagerTest extends AbstractAsyncTransactionTest { - /** * Static helper methods that simplifies creating {@link AsyncTransactionFunction}s for Java7. * Java8 and higher can use lambda expressions. @@ -736,13 +738,11 @@ public ApiFuture apply(TransactionContext txn, Void input) } // This update statement will be aborted, but the error will not propagated to // the transaction manager and cause the transaction to retry. Instead, the - // commit call will do that. + // commit call will do that. Depending on the timing, that will happen + // directly in the transaction manager if the ABORTED error has already been + // returned by the batch update call before the commit call starts. Otherwise, + // the backend will return an ABORTED error for the commit call. txn.batchUpdateAsync(ImmutableList.of(UPDATE_STATEMENT, UPDATE_STATEMENT)); - // Wait for the request to arrive at the server. - mockSpanner.waitForLastRequestToBe(ExecuteBatchDmlRequest.class, 1000L); - // Resolving this future will not resolve the result of the entire - // transaction. The transaction result will be resolved when the commit has - // actually finished successfully. return ApiFutures.immediateFuture(null); } }) @@ -755,15 +755,29 @@ public ApiFuture apply(TransactionContext txn, Void input) } } assertThat(attempt.get()).isEqualTo(2); - assertThat(mockSpanner.getRequestTypes()) - .containsExactly( - BatchCreateSessionsRequest.class, - BeginTransactionRequest.class, - ExecuteBatchDmlRequest.class, - CommitRequest.class, - BeginTransactionRequest.class, - ExecuteBatchDmlRequest.class, - CommitRequest.class); + Iterable> requests = mockSpanner.getRequestTypes(); + int size = Iterables.size(requests); + assertThat(size).isIn(Range.closed(6, 7)); + if (size == 6) { + assertThat(requests) + .containsExactly( + BatchCreateSessionsRequest.class, + BeginTransactionRequest.class, + ExecuteBatchDmlRequest.class, + BeginTransactionRequest.class, + ExecuteBatchDmlRequest.class, + CommitRequest.class); + } else { + assertThat(requests) + .containsExactly( + BatchCreateSessionsRequest.class, + BeginTransactionRequest.class, + ExecuteBatchDmlRequest.class, + CommitRequest.class, + BeginTransactionRequest.class, + ExecuteBatchDmlRequest.class, + CommitRequest.class); + } } @Test 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 25c3428d386..8ba13ddfae3 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 @@ -496,23 +496,22 @@ private static void checkException(Queue exceptions, boolean keepExce private double abortProbability = 0.0010D; private final Object lock = new Object(); - private final Deque requests = new ConcurrentLinkedDeque<>(); + private Deque requests = new ConcurrentLinkedDeque<>(); private volatile CountDownLatch freezeLock = new CountDownLatch(0); - private final Queue exceptions = new ConcurrentLinkedQueue<>(); + private Queue exceptions = new ConcurrentLinkedQueue<>(); private boolean stickyGlobalExceptions = false; - private final ConcurrentMap statementResults = - new ConcurrentHashMap<>(); - private final ConcurrentMap statementGetCounts = new ConcurrentHashMap<>(); - private final ConcurrentMap sessions = new ConcurrentHashMap<>(); + private ConcurrentMap statementResults = new ConcurrentHashMap<>(); + private ConcurrentMap statementGetCounts = new ConcurrentHashMap<>(); + private ConcurrentMap sessions = new ConcurrentHashMap<>(); private ConcurrentMap sessionLastUsed = new ConcurrentHashMap<>(); - private final ConcurrentMap transactions = new ConcurrentHashMap<>(); - private final ConcurrentMap isPartitionedDmlTransaction = + private ConcurrentMap transactions = new ConcurrentHashMap<>(); + private ConcurrentMap isPartitionedDmlTransaction = new ConcurrentHashMap<>(); - private final ConcurrentMap abortedTransactions = new ConcurrentHashMap<>(); + private ConcurrentMap abortedTransactions = new ConcurrentHashMap<>(); private final AtomicBoolean abortNextTransaction = new AtomicBoolean(); private final AtomicBoolean abortNextStatement = new AtomicBoolean(); - private final ConcurrentMap transactionCounters = new ConcurrentHashMap<>(); - private final ConcurrentMap> partitionTokens = new ConcurrentHashMap<>(); + private ConcurrentMap transactionCounters = new ConcurrentHashMap<>(); + private ConcurrentMap> partitionTokens = new ConcurrentHashMap<>(); private ConcurrentMap transactionLastUsed = new ConcurrentHashMap<>(); private int maxNumSessionsInOneBatch = 100; private int maxTotalSessions = Integer.MAX_VALUE; @@ -1639,7 +1638,10 @@ private void ensureMostRecentTransaction(Session session, ByteString transaction throw Status.FAILED_PRECONDITION .withDescription( String.format( - "This transaction has been invalidated by a later transaction in the same session.", + "This transaction has been invalidated by a later transaction in the same session.\nTransaction id: " + + id + + "\nExpected: " + + counter.get(), session.getName())) .asRuntimeException(); } @@ -1862,17 +1864,19 @@ public ServerServiceDefinition getServiceDefinition() { /** Removes all sessions and transactions. Mocked results are not removed. */ @Override public void reset() { - requests.clear(); - sessions.clear(); + requests = new ConcurrentLinkedDeque<>(); + exceptions = new ConcurrentLinkedQueue<>(); + statementGetCounts = new ConcurrentHashMap<>(); + sessions = new ConcurrentHashMap<>(); + sessionLastUsed = new ConcurrentHashMap<>(); + transactions = new ConcurrentHashMap<>(); + isPartitionedDmlTransaction = new ConcurrentHashMap<>(); + abortedTransactions = new ConcurrentHashMap<>(); + transactionCounters = new ConcurrentHashMap<>(); + partitionTokens = new ConcurrentHashMap<>(); + transactionLastUsed = new ConcurrentHashMap<>(); + numSessionsCreated.set(0); - sessionLastUsed.clear(); - transactions.clear(); - isPartitionedDmlTransaction.clear(); - abortedTransactions.clear(); - transactionCounters.clear(); - partitionTokens.clear(); - transactionLastUsed.clear(); - exceptions.clear(); stickyGlobalExceptions = false; freezeLock.countDown(); } From 5e84d344896c53f2f18a4ad2e7cb9b58bc5b4e46 Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Mon, 15 Jun 2020 23:40:59 +0200 Subject: [PATCH 41/49] tests: increase test timeout --- .../com/google/cloud/spanner/AsyncResultSetImplStressTest.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncResultSetImplStressTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncResultSetImplStressTest.java index 8c195bb6da8..c3ad9f45dc4 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncResultSetImplStressTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncResultSetImplStressTest.java @@ -58,7 +58,7 @@ public class AsyncResultSetImplStressTest { private static final int TEST_RUNS = 1000; /** Timeout is applied to each test case individually. */ - @Rule public Timeout timeout = new Timeout(30, TimeUnit.SECONDS); + @Rule public Timeout timeout = new Timeout(120, TimeUnit.SECONDS); @Parameter(0) public int resultSetSize; From fcf37fda8ca063231d13055225f63f48d1508a4d Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Sun, 21 Jun 2020 12:48:19 +0200 Subject: [PATCH 42/49] feat: further towards AsyncTransactionManager --- .../spanner/AsyncTransactionManager.java | 116 ++++++- .../spanner/AsyncTransactionManagerImpl.java | 15 +- .../google/cloud/spanner/DatabaseClient.java | 86 ++++++ .../com/google/cloud/spanner/SessionImpl.java | 2 +- .../com/google/cloud/spanner/SessionPool.java | 2 +- .../SessionPoolAsyncTransactionManager.java | 122 +++++--- .../spanner/TransactionContextFutureImpl.java | 17 +- .../spanner/AsyncTransactionManagerTest.java | 130 ++++++-- .../cloud/spanner/MockSpannerTestUtil.java | 24 ++ .../it/ITTransactionManagerAsyncTest.java | 285 ++++++++++++++++++ 10 files changed, 705 insertions(+), 94 deletions(-) create mode 100644 google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITTransactionManagerAsyncTest.java diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManager.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManager.java index 7cce2296ce4..5c1ad5cc1d4 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManager.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManager.java @@ -1,5 +1,5 @@ /* - * Copyright 2017 Google LLC + * 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. @@ -19,6 +19,8 @@ import com.google.api.core.ApiFuture; import com.google.cloud.Timestamp; import com.google.cloud.spanner.AsyncTransactionManager.AsyncTransactionFunction; +import com.google.cloud.spanner.AsyncTransactionManager.CommitTimestampFuture; +import com.google.cloud.spanner.AsyncTransactionManager.TransactionContextFuture; import com.google.cloud.spanner.TransactionManager.TransactionState; import java.util.concurrent.ExecutionException; import java.util.concurrent.TimeUnit; @@ -30,38 +32,126 @@ * *

    At any point in time there can be at most one active transaction in this manager. When that * transaction is committed, if it fails with an {@code ABORTED} error, calling {@link - * #resetForRetry()} would create a new {@link TransactionContext}. The newly created transaction - * would use the same session thus increasing its lock priority. If the transaction is committed - * successfully, or is rolled back or commit fails with any error other than {@code ABORTED}, the - * manager is considered complete and no further transactions are allowed to be created in it. + * #resetForRetryAsync()} would create a new {@link TransactionContextFuture}. The newly created + * transaction would use the same session thus increasing its lock priority. If the transaction is + * committed successfully, or is rolled back or commit fails with any error other than {@code + * ABORTED}, the manager is considered complete and no further transactions are allowed to be + * created in it. * *

    Every {@code AsyncTransactionManager} should either be committed or rolled back. Failure to do * so can cause resources to be leaked and deadlocks. Easiest way to guarantee this is by calling * {@link #close()} in a finally block. * - * @see DatabaseClient#transactionManager() + * @see DatabaseClient#transactionManagerAsync() */ public interface AsyncTransactionManager extends AutoCloseable { - interface TransactionContextFuture extends ApiFuture { + /** + * {@link ApiFuture} that returns a {@link TransactionContext} and that supports chaining of + * multiple {@link TransactionContextFuture}s to form a transaction. + */ + public interface TransactionContextFuture extends ApiFuture { AsyncTransactionStep then(AsyncTransactionFunction function); } - interface CommitTimestampFuture extends ApiFuture { + /** + * {@link ApiFuture} that returns the commit {@link Timestamp} of a Cloud Spanner transaction that + * is executed using an {@link AsyncTransactionManager}. This future is returned by the call to + * {@link AsyncTransactionStep#commitAsync()} of the last step in the transaction. + */ + public interface CommitTimestampFuture extends ApiFuture { + /** + * Returns the commit timestamp of the transaction. Getting this value should always be done in + * order to ensure that the transaction succeeded. If any of the steps in the transaction fails + * with an uncaught exception, this method will automatically stop the transaction at that point + * and the exception will be returned as the cause of the {@link ExecutionException} that is + * thrown by this method. + * + * @throws AbortedException if the transaction was aborted by Cloud Spanner and needs to be + * retried. + */ @Override Timestamp get() throws AbortedException, InterruptedException, ExecutionException; + /** + * Same as {@link #get()}, but will throw a {@link TimeoutException} if the transaction does not + * finish within the timeout. + */ @Override Timestamp get(long timeout, TimeUnit unit) throws AbortedException, InterruptedException, ExecutionException, TimeoutException; } - interface AsyncTransactionStep extends ApiFuture { + /** + * {@link AsyncTransactionStep} is returned by {@link + * TransactionContextFuture#then(AsyncTransactionFunction)} and {@link + * AsyncTransactionStep#then(AsyncTransactionFunction)} and allows transaction steps that should + * be executed serially to be chained together. Each step can contain one or more statements that + * may execute in parallel. + * + *

    Example usage: + * + * 

    {@code
+   * TransactionContextFuture txnFuture = manager.beginAsync();
+   * final String column = "FirstName";
+   * txnFuture.then(
+   *         new AsyncTransactionFunction() {
+   *           @Override
+   *           public ApiFuture apply(TransactionContext txn, Void input)
+   *               throws Exception {
+   *             return txn.readRowAsync(
+   *                 "Singers", Key.of(singerId), Collections.singleton(column));
+   *           }
+   *         })
+   *     .then(
+   *         new AsyncTransactionFunction() {
+   *           @Override
+   *           public ApiFuture apply(TransactionContext txn, Struct input)
+   *               throws Exception {
+   *             String name = input.getString(column);
+   *             txn.buffer(
+   *                 Mutation.newUpdateBuilder("Singers")
+   *                     .set(column)
+   *                     .to(name.toUpperCase())
+   *                     .build());
+   *             return ApiFutures.immediateFuture(null);
+   *           }
+   *         })
+   * }
    + */ + public interface AsyncTransactionStep extends ApiFuture { + /** + * Adds a step to the transaction chain that should be executed. This step is guaranteed to be + * executed only after the previous step executed successfully. + */ AsyncTransactionStep then(AsyncTransactionFunction next); + /** + * Commits the transaction and returns a {@link CommitTimestampFuture} that will return the + * commit timestamp of the transaction, or throw the first uncaught exception in the transaction + * chain as an {@link ExecutionException}. + */ CommitTimestampFuture commitAsync(); } - interface AsyncTransactionFunction { + /** + * Each step in a transaction chain is defined by an {@link AsyncTransactionFunction}. It receives + * a {@link TransactionContext} and the output value of the previous transaction step as its input + * parameters. The method should return an {@link ApiFuture} that will return the result of this + * step. + */ + public interface AsyncTransactionFunction { + /** + * {@link #apply(TransactionContext, Object)} is called when this transaction step is executed. + * The input value is the result of the previous step, and this method will only be called if + * the previous step executed successfully. + * + * @param txn the {@link TransactionContext} that can be used to execute statements. + * @param input the result of the previous transaction step. + * @return an {@link ApiFuture} that will return the result of this step, and that will be the + * input of the next transaction step. This method should never return null. + * Instead, if the method does not have a return value, the method should return {@link + * ApiFutures#immediateFuture(null)}. + */ ApiFuture apply(TransactionContext txn, I input) throws Exception; } @@ -72,12 +162,6 @@ interface AsyncTransactionFunction { */ TransactionContextFuture beginAsync(); - /** - * Commits the currently active transaction. If the transaction was already aborted, then this - * would throw an {@link AbortedException}. - */ - ApiFuture commitAsync(); - /** * Rolls back the currently active transaction. In most cases there should be no need to call this * explicitly since {@link #close()} would automatically roll back any active transaction. diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManagerImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManagerImpl.java index aca7f6e8cea..082fa827e73 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManagerImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManagerImpl.java @@ -22,6 +22,7 @@ import com.google.api.core.SettableApiFuture; import com.google.cloud.Timestamp; import com.google.cloud.spanner.SessionImpl.SessionTransaction; +import com.google.cloud.spanner.TransactionContextFutureImpl.CommittableAsyncTransactionManager; import com.google.cloud.spanner.TransactionManager.TransactionState; import com.google.common.base.Preconditions; import com.google.common.util.concurrent.MoreExecutors; @@ -30,7 +31,8 @@ import io.opencensus.trace.Tracing; /** Implementation of {@link AsyncTransactionManager}. */ -final class AsyncTransactionManagerImpl implements AsyncTransactionManager, SessionTransaction { +final class AsyncTransactionManagerImpl + implements CommittableAsyncTransactionManager, SessionTransaction { private static final Tracer tracer = Tracing.getTracer(); private final SessionImpl session; @@ -56,9 +58,9 @@ public void close() { } @Override - public TransactionContextFuture beginAsync() { + public TransactionContextFutureImpl beginAsync() { Preconditions.checkState(txn == null, "begin can only be called once"); - TransactionContextFuture begin = + TransactionContextFutureImpl begin = new TransactionContextFutureImpl(this, internalBeginAsync(true)); return begin; } @@ -88,6 +90,13 @@ public void onSuccess(Void result) { return res; } + @Override + public void onError(Throwable t) { + if (t instanceof AbortedException) { + txnState = TransactionState.ABORTED; + } + } + @Override public ApiFuture commitAsync() { Preconditions.checkState( diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClient.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClient.java index 2792fd0c866..d52d1d892e5 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClient.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseClient.java @@ -311,6 +311,92 @@ public interface DatabaseClient { */ AsyncRunner runAsync(); + /** + * Returns an asynchronous transaction manager which allows manual management of transaction + * lifecycle. This API is meant for advanced users. Most users should instead use the {@link + * #runAsync()} API instead. + * + *

    Example of using {@link AsyncTransactionManager} with lambda expressions (Java 8 and + * higher). + * + * 

    {@code
+   * long singerId = 1L;
+   * try (AsyncTransactionManager manager = client.transactionManagerAsync()) {
+   *   TransactionContextFuture txnFut = manager.beginAsync();
+   *   while (true) {
+   *     String column = "FirstName";
+   *     CommitTimestampFuture commitTimestamp =
+   *         txnFut
+   *             .then(
+   *                 (txn, __) ->
+   *                     txn.readRowAsync(
+   *                         "Singers", Key.of(singerId), Collections.singleton(column)))
+   *             .then(
+   *                 (txn, row) -> {
+   *                   String name = row.getString(column);
+   *                   txn.buffer(
+   *                       Mutation.newUpdateBuilder("Singers")
+   *                           .set(column)
+   *                           .to(name.toUpperCase())
+   *                           .build());
+   *                   return ApiFutures.immediateFuture(null);
+   *                 })
+   *             .commitAsync();
+   *     try {
+   *       commitTimestamp.get();
+   *       break;
+   *     } catch (AbortedException e) {
+   *       Thread.sleep(e.getRetryDelayInMillis() / 1000);
+   *       txnFut = manager.resetForRetryAsync();
+   *     }
+   *   }
+   * }
+   * }
    + * + *

    Example of using {@link AsyncTransactionManager} (Java 7). + * + * 

    {@code
+   * final long singerId = 1L;
+   * try (AsyncTransactionManager manager = client().transactionManagerAsync()) {
+   *   TransactionContextFuture txn = manager.beginAsync();
+   *   while (true) {
+   *     final String column = "FirstName";
+   *     CommitTimestampFuture commitTimestamp =
+   *         txn.then(
+   *                 new AsyncTransactionFunction() {
+   *                   @Override
+   *                   public ApiFuture apply(TransactionContext txn, Void input)
+   *                       throws Exception {
+   *                     return txn.readRowAsync(
+   *                         "Singers", Key.of(singerId), Collections.singleton(column));
+   *                   }
+   *                 })
+   *             .then(
+   *                 new AsyncTransactionFunction() {
+   *                   @Override
+   *                   public ApiFuture apply(TransactionContext txn, Struct input)
+   *                       throws Exception {
+   *                     String name = input.getString(column);
+   *                     txn.buffer(
+   *                         Mutation.newUpdateBuilder("Singers")
+   *                             .set(column)
+   *                             .to(name.toUpperCase())
+   *                             .build());
+   *                     return ApiFutures.immediateFuture(null);
+   *                   }
+   *                 })
+   *             .commitAsync();
+   *     try {
+   *       commitTimestamp.get();
+   *       break;
+   *     } catch (AbortedException e) {
+   *       Thread.sleep(e.getRetryDelayInMillis() / 1000);
+   *       txn = manager.resetForRetryAsync();
+   *     }
+   *   }
+   * }
+   * }
    + */ AsyncTransactionManager transactionManagerAsync(); /** diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java index e1cbfca8c53..5798e35202e 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java @@ -241,7 +241,7 @@ public TransactionManager transactionManager() { } @Override - public AsyncTransactionManager transactionManagerAsync() { + public AsyncTransactionManagerImpl transactionManagerAsync() { return new AsyncTransactionManagerImpl(this, currentSpan); } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java index b6ffa4da8ef..90e399fad69 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java @@ -1393,7 +1393,7 @@ public AsyncRunner runAsync() { } @Override - public AsyncTransactionManager transactionManagerAsync() { + public AsyncTransactionManagerImpl transactionManagerAsync() { return delegate.transactionManagerAsync(); } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPoolAsyncTransactionManager.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPoolAsyncTransactionManager.java index 005dfe7b600..55b6102a270 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPoolAsyncTransactionManager.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPoolAsyncTransactionManager.java @@ -24,15 +24,21 @@ import com.google.cloud.Timestamp; import com.google.cloud.spanner.AsyncTransactionManager.TransactionContextFuture; import com.google.cloud.spanner.SessionPool.PooledSessionFuture; +import com.google.cloud.spanner.TransactionContextFutureImpl.CommittableAsyncTransactionManager; import com.google.cloud.spanner.TransactionManager.TransactionState; import com.google.common.base.Preconditions; import com.google.common.util.concurrent.MoreExecutors; -import java.util.concurrent.ExecutionException; +import javax.annotation.concurrent.GuardedBy; -class SessionPoolAsyncTransactionManager implements AsyncTransactionManager { +class SessionPoolAsyncTransactionManager implements CommittableAsyncTransactionManager { + private final Object lock = new Object(); + + @GuardedBy("lock") private TransactionState txnState; + private volatile PooledSessionFuture session; - private final SettableApiFuture delegate = SettableApiFuture.create(); + private final SettableApiFuture delegate = + SettableApiFuture.create(); SessionPoolAsyncTransactionManager(PooledSessionFuture session) { this.session = session; @@ -65,19 +71,21 @@ public void run() { @Override public TransactionContextFuture beginAsync() { - Preconditions.checkState(txnState == null, "begin can only be called once"); - txnState = TransactionState.STARTED; + synchronized (lock) { + Preconditions.checkState(txnState == null, "begin can only be called once"); + txnState = TransactionState.STARTED; + } final SettableApiFuture delegateTxnFuture = SettableApiFuture.create(); ApiFutures.addCallback( delegate, - new ApiFutureCallback() { + new ApiFutureCallback() { @Override public void onFailure(Throwable t) { delegateTxnFuture.setException(t); } @Override - public void onSuccess(AsyncTransactionManager result) { + public void onSuccess(AsyncTransactionManagerImpl result) { ApiFutures.addCallback( result.beginAsync(), new ApiFutureCallback() { @@ -96,40 +104,53 @@ public void onSuccess(TransactionContext result) { }, MoreExecutors.directExecutor()); return new TransactionContextFutureImpl(this, delegateTxnFuture); + } - // return new TransactionContextFutureImpl( - // this, - // ApiFutures.transformAsync( - // delegate, - // new ApiAsyncFunction() { - // @Override - // public ApiFuture apply(AsyncTransactionManager input) { - // return input.beginAsync(); - // } - // }, - // MoreExecutors.directExecutor())); + @Override + public void onError(Throwable t) { + if (t instanceof AbortedException) { + synchronized (lock) { + txnState = TransactionState.ABORTED; + } + } } @Override public ApiFuture commitAsync() { - Preconditions.checkState( - txnState == TransactionState.STARTED, - "commit can only be invoked if the transaction is in progress. Current state: " + txnState); - txnState = TransactionState.COMMITTED; + synchronized (lock) { + Preconditions.checkState( + txnState == TransactionState.STARTED, + "commit can only be invoked if the transaction is in progress. Current state: " + + txnState); + txnState = TransactionState.COMMITTED; + } return ApiFutures.transformAsync( delegate, - new ApiAsyncFunction() { + new ApiAsyncFunction() { @Override - public ApiFuture apply(AsyncTransactionManager input) throws Exception { - ApiFuture res = input.commitAsync(); - // res.addListener( - // new Runnable() { - // @Override - // public void run() { - // session.close(); - // } - // }, - // MoreExecutors.directExecutor()); + public ApiFuture apply(AsyncTransactionManagerImpl input) throws Exception { + final SettableApiFuture res = SettableApiFuture.create(); + ApiFutures.addCallback( + input.commitAsync(), + new ApiFutureCallback() { + @Override + public void onFailure(Throwable t) { + synchronized (lock) { + if (t instanceof AbortedException) { + txnState = TransactionState.ABORTED; + } else { + txnState = TransactionState.COMMIT_FAILED; + } + } + res.setException(t); + } + + @Override + public void onSuccess(Timestamp result) { + res.set(result); + } + }, + MoreExecutors.directExecutor()); return res; } }, @@ -138,15 +159,17 @@ public ApiFuture apply(AsyncTransactionManager input) throws Exceptio @Override public ApiFuture rollbackAsync() { - Preconditions.checkState( - txnState == TransactionState.STARTED, - "rollback can only be called if the transaction is in progress"); - txnState = TransactionState.ROLLED_BACK; + synchronized (lock) { + Preconditions.checkState( + txnState == TransactionState.STARTED, + "rollback can only be called if the transaction is in progress"); + txnState = TransactionState.ROLLED_BACK; + } return ApiFutures.transformAsync( delegate, - new ApiAsyncFunction() { + new ApiAsyncFunction() { @Override - public ApiFuture apply(AsyncTransactionManager input) throws Exception { + public ApiFuture apply(AsyncTransactionManagerImpl input) throws Exception { ApiFuture res = input.rollbackAsync(); res.addListener( new Runnable() { @@ -164,16 +187,19 @@ public void run() { @Override public TransactionContextFuture resetForRetryAsync() { - Preconditions.checkState( - txnState != null, "resetForRetry can only be called after the transaction has started."); - txnState = TransactionState.STARTED; + synchronized (lock) { + Preconditions.checkState( + txnState == TransactionState.ABORTED, + "resetForRetry can only be called after the transaction aborted."); + txnState = TransactionState.STARTED; + } return new TransactionContextFutureImpl( this, ApiFutures.transformAsync( delegate, - new ApiAsyncFunction() { + new ApiAsyncFunction() { @Override - public ApiFuture apply(AsyncTransactionManager input) + public ApiFuture apply(AsyncTransactionManagerImpl input) throws Exception { return input.resetForRetryAsync(); } @@ -183,12 +209,8 @@ public ApiFuture apply(AsyncTransactionManager input) @Override public TransactionState getState() { - try { - return delegate.get().getState(); - } catch (InterruptedException e) { - throw SpannerExceptionFactory.propagateInterrupt(e); - } catch (ExecutionException e) { - throw SpannerExceptionFactory.newSpannerException(e.getCause() == null ? e : e.getCause()); + synchronized (lock) { + return txnState; } } } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionContextFutureImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionContextFutureImpl.java index 3cdd6430424..5922726120e 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionContextFutureImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionContextFutureImpl.java @@ -20,6 +20,7 @@ import com.google.api.core.ApiFutureCallback; import com.google.api.core.ApiFutures; import com.google.api.core.ForwardingApiFuture; +import com.google.api.core.InternalApi; import com.google.api.core.SettableApiFuture; import com.google.cloud.Timestamp; import com.google.cloud.spanner.AsyncTransactionManager.AsyncTransactionFunction; @@ -35,6 +36,12 @@ class TransactionContextFutureImpl extends ForwardingApiFuture implements TransactionContextFuture { + @InternalApi + interface CommittableAsyncTransactionManager extends AsyncTransactionManager { + void onError(Throwable t); + + ApiFuture commitAsync(); + } /** * {@link ApiFuture} that returns a commit timestamp. Any {@link AbortedException} that is thrown * by either the commit call or any other rpc during the transaction will be thrown by the {@link @@ -98,6 +105,7 @@ class AsyncTransactionStatementImpl extends ForwardingApiFuture new ApiFutureCallback() { @Override public void onFailure(Throwable t) { + mgr.onError(t); txnResult.setException(t); } @@ -111,6 +119,7 @@ public void onSuccess(I result) { new ApiFutureCallback() { @Override public void onFailure(Throwable t) { + mgr.onError(t); txnResult.setException(t); } @@ -121,6 +130,7 @@ public void onSuccess(O result) { }, MoreExecutors.directExecutor()); } catch (Throwable t) { + mgr.onError(t); txnResult.setException(t); } } @@ -140,6 +150,7 @@ public CommitTimestampFuture commitAsync() { new ApiFutureCallback() { @Override public void onFailure(Throwable t) { + mgr.onError(t); txnResult.setException(t); } @@ -150,6 +161,7 @@ public void onSuccess(O result) { new ApiFutureCallback() { @Override public void onFailure(Throwable t) { + mgr.onError(t); txnResult.setException(t); } @@ -166,11 +178,11 @@ public void onSuccess(Timestamp result) { } } - final AsyncTransactionManager mgr; + final CommittableAsyncTransactionManager mgr; final SettableApiFuture txnResult = SettableApiFuture.create(); TransactionContextFutureImpl( - AsyncTransactionManager mgr, ApiFuture txnFuture) { + CommittableAsyncTransactionManager mgr, ApiFuture txnFuture) { super(txnFuture); this.mgr = mgr; } @@ -184,6 +196,7 @@ public AsyncTransactionStatementImpl then( new ApiFutureCallback() { @Override public void onFailure(Throwable t) { + mgr.onError(t); input.setException(t); } diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncTransactionManagerTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncTransactionManagerTest.java index 76eec9f1173..e9f33546dcf 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncTransactionManagerTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncTransactionManagerTest.java @@ -29,13 +29,13 @@ import com.google.api.core.ApiFutureCallback; import com.google.api.core.ApiFutures; import com.google.api.core.SettableApiFuture; -import com.google.cloud.spanner.AsyncRunner.AsyncWork; import com.google.cloud.spanner.AsyncTransactionManager.AsyncTransactionFunction; import com.google.cloud.spanner.AsyncTransactionManager.AsyncTransactionStep; import com.google.cloud.spanner.AsyncTransactionManager.CommitTimestampFuture; import com.google.cloud.spanner.AsyncTransactionManager.TransactionContextFuture; import com.google.cloud.spanner.MockSpannerServiceImpl.SimulatedExecutionTime; import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult; +import com.google.cloud.spanner.Options.ReadOption; import com.google.common.base.Function; import com.google.common.collect.ImmutableList; import com.google.common.collect.Iterables; @@ -49,6 +49,7 @@ import com.google.spanner.v1.ExecuteSqlRequest; import io.grpc.Status; import java.util.Arrays; +import java.util.Collections; import java.util.concurrent.ExecutionException; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; @@ -63,6 +64,20 @@ public class AsyncTransactionManagerTest extends AbstractAsyncTransactionTest { * Java8 and higher can use lambda expressions. */ public static class AsyncTransactionManagerHelper { + + public static AsyncTransactionFunction readAsync( + final String table, + final KeySet keys, + final Iterable columns, + final ReadOption... options) { + return new AsyncTransactionFunction() { + @Override + public ApiFuture apply(TransactionContext txn, I input) throws Exception { + return ApiFutures.immediateFuture(txn.readAsync(table, keys, columns, options)); + } + }; + } + public static AsyncTransactionFunction readRowAsync( final String table, final Key key, final Iterable columns) { return new AsyncTransactionFunction() { @@ -73,6 +88,20 @@ public ApiFuture apply(TransactionContext txn, I input) throws Exception }; } + public static AsyncTransactionFunction buffer(Mutation mutation) { + return buffer(ImmutableList.of(mutation)); + } + + public static AsyncTransactionFunction buffer(final Iterable mutations) { + return new AsyncTransactionFunction() { + @Override + public ApiFuture apply(TransactionContext txn, I input) throws Exception { + txn.buffer(mutations); + return ApiFutures.immediateFuture(null); + } + }; + } + public static AsyncTransactionFunction executeUpdateAsync(Statement statement) { return executeUpdateAsync(SettableApiFuture.create(), statement); } @@ -879,25 +908,84 @@ public ApiFuture apply(TransactionContext txn, Struct input) } @Test - public void asyncRunnerRead() throws Exception { - AsyncRunner runner = client().runAsync(); - ApiFuture> val = - runner.runAsync( - new AsyncWork>() { - @Override - public ApiFuture> doWorkAsync(TransactionContext txn) { - return txn.readAsync(READ_TABLE_NAME, KeySet.all(), READ_COLUMN_NAMES) - .toListAsync( - new Function() { - @Override - public String apply(StructReader input) { - return input.getString("Value"); - } - }, - MoreExecutors.directExecutor()); - } - }, - executor); - assertThat(val.get()).containsExactly("v1", "v2", "v3"); + public void asyncTransactionManagerRead() throws Exception { + AsyncTransactionStep> res; + try (AsyncTransactionManager mgr = client().transactionManagerAsync()) { + TransactionContextFuture txn = mgr.beginAsync(); + while (true) { + try { + res = + txn.then( + new AsyncTransactionFunction>() { + @Override + public ApiFuture> apply( + TransactionContext txn, Void input) throws Exception { + return txn.readAsync(READ_TABLE_NAME, KeySet.all(), READ_COLUMN_NAMES) + .toListAsync( + new Function() { + @Override + public String apply(StructReader input) { + return input.getString("Value"); + } + }, + MoreExecutors.directExecutor()); + } + }); + // Commit the transaction. + res.commitAsync().get(); + break; + } catch (AbortedException e) { + txn = mgr.resetForRetryAsync(); + } + } + } + assertThat(res.get()).containsExactly("v1", "v2", "v3"); + } + + @Test + public void asyncTransactionManagerQuery() throws Exception { + mockSpanner.putStatementResult( + StatementResult.query( + Statement.of("SELECT FirstName FROM Singers WHERE ID=1"), + MockSpannerTestUtil.READ_FIRST_NAME_SINGERS_RESULTSET)); + final long singerId = 1L; + try (AsyncTransactionManager manager = client().transactionManagerAsync()) { + TransactionContextFuture txn = manager.beginAsync(); + while (true) { + final String column = "FirstName"; + CommitTimestampFuture commitTimestamp = + txn.then( + new AsyncTransactionFunction() { + @Override + public ApiFuture apply(TransactionContext txn, Void input) + throws Exception { + return txn.readRowAsync( + "Singers", Key.of(singerId), Collections.singleton(column)); + } + }) + .then( + new AsyncTransactionFunction() { + @Override + public ApiFuture apply(TransactionContext txn, Struct input) + throws Exception { + String name = input.getString(column); + txn.buffer( + Mutation.newUpdateBuilder("Singers") + .set(column) + .to(name.toUpperCase()) + .build()); + return ApiFutures.immediateFuture(null); + } + }) + .commitAsync(); + try { + commitTimestamp.get(); + break; + } catch (AbortedException e) { + Thread.sleep(e.getRetryDelayInMillis() / 1000); + txn = manager.resetForRetryAsync(); + } + } + } } } diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MockSpannerTestUtil.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MockSpannerTestUtil.java index c0d3bdc7d1c..cc6784b679a 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MockSpannerTestUtil.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MockSpannerTestUtil.java @@ -124,4 +124,28 @@ static com.google.spanner.v1.ResultSet generateKeyValueResultSet(Iterable() { + @Override + public ApiFuture apply(TransactionContext txn, Void input) + throws Exception { + txn.buffer( + Mutation.newInsertBuilder("T") + .set("K") + .to("Key1") + .set("BoolValue") + .to(true) + .build()); + return ApiFutures.immediateFuture(null); + } + }) + .commitAsync() + .get(); + assertThat(manager.getState()).isEqualTo(TransactionState.COMMITTED); + Struct row = + client.singleUse().readRow("T", Key.of("Key1"), Arrays.asList("K", "BoolValue")); + assertThat(row.getString(0)).isEqualTo("Key1"); + assertThat(row.getBoolean(1)).isTrue(); + break; + } catch (AbortedException e) { + Thread.sleep(e.getRetryDelayInMillis() / 1000); + txn = manager.resetForRetryAsync(); + } + } + } + } + + @Test + public void testInvalidInsert() throws InterruptedException { + try (AsyncTransactionManager manager = client.transactionManagerAsync()) { + TransactionContextFuture txn = manager.beginAsync(); + while (true) { + try { + txn.then( + new AsyncTransactionFunction() { + @Override + public ApiFuture apply(TransactionContext txn, Void input) + throws Exception { + txn.buffer( + Mutation.newInsertBuilder("InvalidTable") + .set("K") + .to("Key1") + .set("BoolValue") + .to(true) + .build()); + return ApiFutures.immediateFuture(null); + } + }) + .commitAsync() + .get(); + fail("Expected exception"); + } catch (AbortedException e) { + Thread.sleep(e.getRetryDelayInMillis() / 1000); + txn = manager.resetForRetryAsync(); + } catch (ExecutionException e) { + assertThat(e.getCause()).isInstanceOf(SpannerException.class); + SpannerException se = (SpannerException) e.getCause(); + assertThat(se.getErrorCode()).isEqualTo(ErrorCode.NOT_FOUND); + // expected + break; + } + } + assertThat(manager.getState()).isEqualTo(TransactionState.COMMIT_FAILED); + // We cannot retry for non aborted errors. + try { + manager.resetForRetryAsync(); + fail("Expected exception"); + } catch (IllegalStateException ex) { + assertNotNull(ex.getMessage()); + } + } + } + + @Test + public void testRollback() throws InterruptedException { + try (AsyncTransactionManager manager = client.transactionManagerAsync()) { + TransactionContextFuture txn = manager.beginAsync(); + while (true) { + txn.then( + new AsyncTransactionFunction() { + @Override + public ApiFuture apply(TransactionContext txn, Void input) throws Exception { + txn.buffer( + Mutation.newInsertBuilder("T") + .set("K") + .to("Key2") + .set("BoolValue") + .to(true) + .build()); + return ApiFutures.immediateFuture(null); + } + }); + try { + manager.rollbackAsync(); + break; + } catch (AbortedException e) { + Thread.sleep(e.getRetryDelayInMillis() / 1000); + txn = manager.resetForRetryAsync(); + } + } + assertThat(manager.getState()).isEqualTo(TransactionState.ROLLED_BACK); + // Row should not have been inserted. + assertThat(client.singleUse().readRow("T", Key.of("Key2"), Arrays.asList("K", "BoolValue"))) + .isNull(); + } + } + + @Test + public void testAbortAndRetry() throws InterruptedException, ExecutionException { + assumeFalse( + "Emulator does not support more than 1 simultanous transaction. " + + "This test would therefore loop indefinetly on the emulator.", + env.getTestHelper().isEmulator()); + + client.write( + Arrays.asList( + Mutation.newInsertBuilder("T").set("K").to("Key3").set("BoolValue").to(true).build())); + try (AsyncTransactionManager manager1 = client.transactionManagerAsync()) { + TransactionContextFuture txn1 = manager1.beginAsync(); + AsyncTransactionManager manager2; + TransactionContextFuture txn2; + AsyncTransactionStep txn2Step1; + while (true) { + try { + AsyncTransactionStep txn1Step1 = + txn1.then( + new AsyncTransactionFunction() { + @Override + public ApiFuture apply(TransactionContext txn, Void input) + throws Exception { + return txn.readRowAsync("T", Key.of("Key3"), Arrays.asList("K", "BoolValue")); + } + }); + manager2 = client.transactionManagerAsync(); + txn2 = manager2.beginAsync(); + txn2Step1 = + txn2.then( + new AsyncTransactionFunction() { + @Override + public ApiFuture apply(TransactionContext txn, Void input) + throws Exception { + return txn.readRowAsync("T", Key.of("Key3"), Arrays.asList("K", "BoolValue")); + } + }); + + AsyncTransactionStep txn1Step2 = + txn1Step1.then( + new AsyncTransactionFunction() { + @Override + public ApiFuture apply(TransactionContext txn, Struct input) + throws Exception { + txn.buffer( + Mutation.newUpdateBuilder("T") + .set("K") + .to("Key3") + .set("BoolValue") + .to(false) + .build()); + return ApiFutures.immediateFuture(null); + } + }); + + txn2Step1.get(); + txn1Step2.commitAsync().get(); + break; + } catch (AbortedException e) { + Thread.sleep(e.getRetryDelayInMillis() / 1000); + // It is possible that it was txn2 that aborted. + // In that case we should just retry without resetting anything. + if (manager1.getState() == TransactionState.ABORTED) { + txn1 = manager1.resetForRetryAsync(); + } + } + } + + // txn2 should have been aborted. + try { + txn2Step1.commitAsync().get(); + fail("Expected to abort"); + } catch (AbortedException e) { + assertThat(manager2.getState()).isEqualTo(TransactionState.ABORTED); + txn2 = manager2.resetForRetryAsync(); + } + AsyncTransactionStep txn2Step2 = + txn2.then( + new AsyncTransactionFunction() { + @Override + public ApiFuture apply(TransactionContext txn, Void input) throws Exception { + txn.buffer( + Mutation.newUpdateBuilder("T") + .set("K") + .to("Key3") + .set("BoolValue") + .to(true) + .build()); + return ApiFutures.immediateFuture(null); + } + }); + txn2Step2.commitAsync().get(); + Struct row = client.singleUse().readRow("T", Key.of("Key3"), Arrays.asList("K", "BoolValue")); + assertThat(row.getString(0)).isEqualTo("Key3"); + assertThat(row.getBoolean(1)).isTrue(); + manager2.close(); + } + } +} From 910b6c7d0eb62666061efaf25e43c2de42b0bc83 Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Sun, 21 Jun 2020 18:15:49 +0200 Subject: [PATCH 43/49] feat: require executor for transaction functions --- .../spanner/AsyncTransactionManager.java | 22 ++- .../spanner/TransactionContextFutureImpl.java | 67 +++++++-- .../spanner/AsyncTransactionManagerTest.java | 128 +++++++++++++----- .../it/ITTransactionManagerAsyncTest.java | 43 ++++-- 4 files changed, 201 insertions(+), 59 deletions(-) diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManager.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManager.java index 5c1ad5cc1d4..d519c68013f 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManager.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncTransactionManager.java @@ -22,7 +22,10 @@ import com.google.cloud.spanner.AsyncTransactionManager.CommitTimestampFuture; import com.google.cloud.spanner.AsyncTransactionManager.TransactionContextFuture; import com.google.cloud.spanner.TransactionManager.TransactionState; +import com.google.common.util.concurrent.ListenableFuture; +import com.google.common.util.concurrent.MoreExecutors; import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executor; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; @@ -50,7 +53,14 @@ public interface AsyncTransactionManager extends AutoCloseable { * multiple {@link TransactionContextFuture}s to form a transaction. */ public interface TransactionContextFuture extends ApiFuture { - AsyncTransactionStep then(AsyncTransactionFunction function); + /** + * Sets the first step to execute as part of this transaction after the transaction has started + * using the specified executor. {@link MoreExecutors#directExecutor()} can be be used for + * lightweight functions, but should be avoided for heavy or blocking operations. See also + * {@link ListenableFuture#addListener(Runnable, Executor)} for further information. + */ + AsyncTransactionStep then( + AsyncTransactionFunction function, Executor executor); } /** @@ -120,10 +130,14 @@ Timestamp get(long timeout, TimeUnit unit) */ public interface AsyncTransactionStep extends ApiFuture { /** - * Adds a step to the transaction chain that should be executed. This step is guaranteed to be - * executed only after the previous step executed successfully. + * Adds a step to the transaction chain that should be executed using the specified executor. + * This step is guaranteed to be executed only after the previous step executed successfully. + * {@link MoreExecutors#directExecutor()} can be be used for lightweight functions, but should + * be avoided for heavy or blocking operations. See also {@link + * ListenableFuture#addListener(Runnable, Executor)} for further information. */ - AsyncTransactionStep then(AsyncTransactionFunction next); + AsyncTransactionStep then( + AsyncTransactionFunction next, Executor executor); /** * Commits the transaction and returns a {@link CommitTimestampFuture} that will return the diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionContextFutureImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionContextFutureImpl.java index 5922726120e..bc8262a5358 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionContextFutureImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionContextFutureImpl.java @@ -30,6 +30,7 @@ import com.google.common.base.Preconditions; import com.google.common.util.concurrent.MoreExecutors; import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executor; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; @@ -88,15 +89,17 @@ class AsyncTransactionStatementImpl extends ForwardingApiFuture AsyncTransactionStatementImpl( final ApiFuture txnFuture, ApiFuture input, - final AsyncTransactionFunction function) { - this(SettableApiFuture.create(), txnFuture, input, function); + final AsyncTransactionFunction function, + Executor executor) { + this(SettableApiFuture.create(), txnFuture, input, function, executor); } AsyncTransactionStatementImpl( SettableApiFuture delegate, final ApiFuture txnFuture, ApiFuture input, - final AsyncTransactionFunction function) { + final AsyncTransactionFunction function, + final Executor executor) { super(delegate); this.statementResult = delegate; this.txnFuture = txnFuture; @@ -113,9 +116,7 @@ public void onFailure(Throwable t) { public void onSuccess(I result) { try { ApiFutures.addCallback( - Preconditions.checkNotNull( - function.apply(txnFuture.get(), result), - "AsyncTransactionFunction returned . Did you mean to return ApiFutures.immediateFuture(null)?"), + runAsyncTransactionFunction(function, txnFuture.get(), result, executor), new ApiFutureCallback() { @Override public void onFailure(Throwable t) { @@ -139,8 +140,9 @@ public void onSuccess(O result) { } @Override - public AsyncTransactionStatementImpl then(AsyncTransactionFunction next) { - return new AsyncTransactionStatementImpl<>(txnFuture, statementResult, next); + public AsyncTransactionStatementImpl then( + AsyncTransactionFunction next, Executor executor) { + return new AsyncTransactionStatementImpl<>(txnFuture, statementResult, next, executor); } @Override @@ -178,6 +180,51 @@ public void onSuccess(Timestamp result) { } } + static ApiFuture runAsyncTransactionFunction( + final AsyncTransactionFunction function, + final TransactionContext txn, + final I input, + Executor executor) + throws Exception { + // Shortcut for common path. + if (executor == MoreExecutors.directExecutor()) { + return Preconditions.checkNotNull( + function.apply(txn, input), + "AsyncTransactionFunction returned . Did you mean to return ApiFutures.immediateFuture(null)?"); + } else { + final SettableApiFuture res = SettableApiFuture.create(); + executor.execute( + new Runnable() { + @Override + public void run() { + try { + ApiFuture functionResult = + Preconditions.checkNotNull( + function.apply(txn, input), + "AsyncTransactionFunction returned . Did you mean to return ApiFutures.immediateFuture(null)?"); + ApiFutures.addCallback( + functionResult, + new ApiFutureCallback() { + @Override + public void onFailure(Throwable t) { + res.setException(t); + } + + @Override + public void onSuccess(O result) { + res.set(result); + } + }, + MoreExecutors.directExecutor()); + } catch (Throwable t) { + res.setException(t); + } + } + }); + return res; + } + } + final CommittableAsyncTransactionManager mgr; final SettableApiFuture txnResult = SettableApiFuture.create(); @@ -189,7 +236,7 @@ public void onSuccess(Timestamp result) { @Override public AsyncTransactionStatementImpl then( - AsyncTransactionFunction function) { + AsyncTransactionFunction function, Executor executor) { final SettableApiFuture input = SettableApiFuture.create(); ApiFutures.addCallback( this, @@ -206,6 +253,6 @@ public void onSuccess(TransactionContext result) { } }, MoreExecutors.directExecutor()); - return new AsyncTransactionStatementImpl<>(this, input, function); + return new AsyncTransactionStatementImpl<>(this, input, function, executor); } } diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncTransactionManagerTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncTransactionManagerTest.java index e9f33546dcf..c5de6542192 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncTransactionManagerTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncTransactionManagerTest.java @@ -49,16 +49,34 @@ import com.google.spanner.v1.ExecuteSqlRequest; import io.grpc.Status; import java.util.Arrays; +import java.util.Collection; import java.util.Collections; import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executor; +import java.util.concurrent.Executors; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; import org.junit.Test; import org.junit.runner.RunWith; -import org.junit.runners.JUnit4; +import org.junit.runners.Parameterized; +import org.junit.runners.Parameterized.Parameter; +import org.junit.runners.Parameterized.Parameters; -@RunWith(JUnit4.class) +@RunWith(Parameterized.class) public class AsyncTransactionManagerTest extends AbstractAsyncTransactionTest { + + @Parameter public Executor executor; + + @Parameters(name = "executor = {0}") + public static Collection data() { + return Arrays.asList( + new Object[][] { + {MoreExecutors.directExecutor()}, + {Executors.newSingleThreadExecutor()}, + {Executors.newFixedThreadPool(4)} + }); + } + /** * Static helper methods that simplifies creating {@link AsyncTransactionFunction}s for Java7. * Java8 and higher can use lambda expressions. @@ -173,7 +191,8 @@ public void asyncTransactionManagerUpdate() throws Exception { CommitTimestampFuture commitTimestamp = txn.then( AsyncTransactionManagerHelper.executeUpdateAsync( - updateCount, UPDATE_STATEMENT)) + updateCount, UPDATE_STATEMENT), + executor) .commitAsync(); assertThat(updateCount.get()).isEqualTo(UPDATE_COUNT); assertThat(commitTimestamp.get()).isNotNull(); @@ -197,7 +216,8 @@ public void asyncTransactionManagerIsNonBlocking() throws Exception { CommitTimestampFuture commitTimestamp = txn.then( AsyncTransactionManagerHelper.executeUpdateAsync( - updateCount, UPDATE_STATEMENT)) + updateCount, UPDATE_STATEMENT), + executor) .commitAsync(); mockSpanner.unfreeze(); assertThat(updateCount.get(10L, TimeUnit.SECONDS)).isEqualTo(UPDATE_COUNT); @@ -219,7 +239,8 @@ public void asyncTransactionManagerInvalidUpdate() throws Exception { CommitTimestampFuture commitTimestamp = txn.then( AsyncTransactionManagerHelper.executeUpdateAsync( - INVALID_UPDATE_STATEMENT)) + INVALID_UPDATE_STATEMENT), + executor) .commitAsync(); commitTimestamp.get(); fail("missing expected exception"); @@ -249,7 +270,8 @@ public void asyncTransactionManagerCommitAborted() throws Exception { CommitTimestampFuture commitTimestamp = txn.then( AsyncTransactionManagerHelper.executeUpdateAsync( - updateCount, UPDATE_STATEMENT)) + updateCount, UPDATE_STATEMENT), + executor) .then( new AsyncTransactionFunction() { @Override @@ -260,7 +282,8 @@ public ApiFuture apply(TransactionContext txn, Long input) } return ApiFutures.immediateFuture(null); } - }) + }, + executor) .commitAsync(); assertThat(updateCount.get()).isEqualTo(UPDATE_COUNT); assertThat(commitTimestamp.get()).isNotNull(); @@ -305,7 +328,8 @@ public void onSuccess(Long result) { MoreExecutors.directExecutor()); return updateCount; } - }) + }, + executor) .commitAsync(); assertThat(updateCount.get()).isEqualTo(UPDATE_COUNT); assertThat(ts.get()).isNotNull(); @@ -324,10 +348,13 @@ public void asyncTransactionManagerChain() throws Exception { while (true) { try { CommitTimestampFuture ts = - txn.then(AsyncTransactionManagerHelper.executeUpdateAsync(UPDATE_STATEMENT)) + txn.then( + AsyncTransactionManagerHelper.executeUpdateAsync(UPDATE_STATEMENT), + executor) .then( AsyncTransactionManagerHelper.readRowAsync( - READ_TABLE_NAME, Key.of(1L), READ_COLUMN_NAMES)) + READ_TABLE_NAME, Key.of(1L), READ_COLUMN_NAMES), + executor) .then( new AsyncTransactionFunction() { @Override @@ -335,7 +362,8 @@ public ApiFuture apply(TransactionContext txn, Struct input) throws Exception { return ApiFutures.immediateFuture(input.getString("Value")); } - }) + }, + executor) .then( new AsyncTransactionFunction() { @Override @@ -344,7 +372,8 @@ public ApiFuture apply(TransactionContext txn, String input) assertThat(input).isEqualTo("v1"); return ApiFutures.immediateFuture(null); } - }) + }, + executor) .commitAsync(); assertThat(ts.get()).isNotNull(); break; @@ -364,7 +393,8 @@ public void asyncTransactionManagerChainWithErrorInTheMiddle() throws Exception CommitTimestampFuture ts = txn.then( AsyncTransactionManagerHelper.executeUpdateAsync( - INVALID_UPDATE_STATEMENT)) + INVALID_UPDATE_STATEMENT), + executor) .then( new AsyncTransactionFunction() { @Override @@ -372,7 +402,8 @@ public ApiFuture apply(TransactionContext txn, Long input) throws Exception { throw new IllegalStateException("this should not be executed"); } - }) + }, + executor) .commitAsync(); ts.get(); break; @@ -415,8 +446,11 @@ public ApiFuture apply(TransactionContext txn, Void input) } return ApiFutures.immediateFuture(null); } - }) - .then(AsyncTransactionManagerHelper.executeUpdateAsync(UPDATE_STATEMENT)) + }, + executor) + .then( + AsyncTransactionManagerHelper.executeUpdateAsync(UPDATE_STATEMENT), + executor) .commitAsync(); assertThat(ts.get()).isNotNull(); break; @@ -459,7 +493,8 @@ public ApiFuture apply(TransactionContext txn, Void input) // successfully. return ApiFutures.immediateFuture(null); } - }) + }, + executor) .commitAsync(); assertThat(ts.get()).isNotNull(); assertThat(attempt.get()).isEqualTo(2); @@ -491,7 +526,9 @@ public void asyncTransactionManagerCommitFails() throws Exception { TransactionContextFuture txn = mgr.beginAsync(); while (true) { try { - txn.then(AsyncTransactionManagerHelper.executeUpdateAsync(UPDATE_STATEMENT)) + txn.then( + AsyncTransactionManagerHelper.executeUpdateAsync(UPDATE_STATEMENT), + executor) .commitAsync() .get(); fail("missing expected exception"); @@ -524,7 +561,8 @@ public ApiFuture apply(TransactionContext txn, Void input) txn.executeUpdateAsync(UPDATE_STATEMENT); return ApiFutures.immediateFuture(null); } - }) + }, + executor) .commitAsync() .get(); assertThat(mockSpanner.getRequestTypes()) @@ -550,7 +588,8 @@ public void asyncTransactionManagerBatchUpdate() throws Exception { try { txn.then( AsyncTransactionManagerHelper.batchUpdateAsync( - result, UPDATE_STATEMENT, UPDATE_STATEMENT)) + result, UPDATE_STATEMENT, UPDATE_STATEMENT), + executor) .commitAsync() .get(); break; @@ -571,7 +610,9 @@ public void asyncTransactionManagerIsNonBlockingWithBatchUpdate() throws Excepti while (true) { try { CommitTimestampFuture ts = - txn.then(AsyncTransactionManagerHelper.batchUpdateAsync(res, UPDATE_STATEMENT)) + txn.then( + AsyncTransactionManagerHelper.batchUpdateAsync(res, UPDATE_STATEMENT), + executor) .commitAsync(); mockSpanner.unfreeze(); assertThat(ts.get()).isNotNull(); @@ -593,7 +634,8 @@ public void asyncTransactionManagerInvalidBatchUpdate() throws Exception { try { txn.then( AsyncTransactionManagerHelper.batchUpdateAsync( - result, UPDATE_STATEMENT, INVALID_UPDATE_STATEMENT)) + result, UPDATE_STATEMENT, INVALID_UPDATE_STATEMENT), + executor) .commitAsync() .get(); fail("missing expected exception"); @@ -626,10 +668,12 @@ public ApiFuture apply(TransactionContext txn, Void input) ImmutableList.of(UPDATE_STATEMENT, INVALID_UPDATE_STATEMENT)); return ApiFutures.immediateFuture(null); } - }) + }, + executor) .then( AsyncTransactionManagerHelper.batchUpdateAsync( - result, UPDATE_STATEMENT, UPDATE_STATEMENT)) + result, UPDATE_STATEMENT, UPDATE_STATEMENT), + executor) .commitAsync() .get(); break; @@ -668,7 +712,8 @@ public ApiFuture apply(TransactionContext txn, Void input) ImmutableList.of(UPDATE_STATEMENT, UPDATE_STATEMENT)); } } - }) + }, + executor) .commitAsync() .get(); break; @@ -712,10 +757,12 @@ public ApiFuture apply(TransactionContext txn, Void input) } return ApiFutures.immediateFuture(null); } - }) + }, + executor) .then( AsyncTransactionManagerHelper.batchUpdateAsync( - result, UPDATE_STATEMENT, UPDATE_STATEMENT)) + result, UPDATE_STATEMENT, UPDATE_STATEMENT), + executor) .then( new AsyncTransactionFunction() { @Override @@ -726,7 +773,8 @@ public ApiFuture apply(TransactionContext txn, long[] input) } return ApiFutures.immediateFuture(null); } - }) + }, + executor) .commitAsync() .get(); assertThat(result.get()).asList().containsExactly(UPDATE_COUNT, UPDATE_COUNT); @@ -774,7 +822,8 @@ public ApiFuture apply(TransactionContext txn, Void input) txn.batchUpdateAsync(ImmutableList.of(UPDATE_STATEMENT, UPDATE_STATEMENT)); return ApiFutures.immediateFuture(null); } - }) + }, + executor) .commitAsync() .get(); break; @@ -822,7 +871,8 @@ public void asyncTransactionManagerWithBatchUpdateCommitFails() throws Exception try { txn.then( AsyncTransactionManagerHelper.batchUpdateAsync( - UPDATE_STATEMENT, UPDATE_STATEMENT)) + UPDATE_STATEMENT, UPDATE_STATEMENT), + executor) .commitAsync() .get(); fail("missing expected exception"); @@ -859,7 +909,8 @@ public ApiFuture apply(TransactionContext txn, Void input) txn.batchUpdateAsync(ImmutableList.of(UPDATE_STATEMENT)); return ApiFutures.immediateFuture(null); } - }) + }, + executor) .commitAsync() .get(); break; @@ -888,7 +939,8 @@ public void asyncTransactionManagerReadRow() throws Exception { step = txn.then( AsyncTransactionManagerHelper.readRowAsync( - READ_TABLE_NAME, Key.of(1L), READ_COLUMN_NAMES)) + READ_TABLE_NAME, Key.of(1L), READ_COLUMN_NAMES), + executor) .then( new AsyncTransactionFunction() { @Override @@ -896,7 +948,8 @@ public ApiFuture apply(TransactionContext txn, Struct input) throws Exception { return ApiFutures.immediateFuture(input.getString("Value")); } - }); + }, + executor); step.commitAsync().get(); break; } catch (AbortedException e) { @@ -930,7 +983,8 @@ public String apply(StructReader input) { }, MoreExecutors.directExecutor()); } - }); + }, + executor); // Commit the transaction. res.commitAsync().get(); break; @@ -962,7 +1016,8 @@ public ApiFuture apply(TransactionContext txn, Void input) return txn.readRowAsync( "Singers", Key.of(singerId), Collections.singleton(column)); } - }) + }, + executor) .then( new AsyncTransactionFunction() { @Override @@ -976,7 +1031,8 @@ public ApiFuture apply(TransactionContext txn, Struct input) .build()); return ApiFutures.immediateFuture(null); } - }) + }, + executor) .commitAsync(); try { commitTimestamp.get(); diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITTransactionManagerAsyncTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITTransactionManagerAsyncTest.java index 2d812d5bbc7..4bdb8804e71 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITTransactionManagerAsyncTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITTransactionManagerAsyncTest.java @@ -38,17 +38,35 @@ import com.google.cloud.spanner.Struct; import com.google.cloud.spanner.TransactionContext; import com.google.cloud.spanner.TransactionManager.TransactionState; +import com.google.common.util.concurrent.MoreExecutors; import java.util.Arrays; +import java.util.Collection; import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executor; +import java.util.concurrent.Executors; import org.junit.BeforeClass; import org.junit.ClassRule; import org.junit.Test; import org.junit.runner.RunWith; -import org.junit.runners.JUnit4; +import org.junit.runners.Parameterized; +import org.junit.runners.Parameterized.Parameter; +import org.junit.runners.Parameterized.Parameters; -@RunWith(JUnit4.class) +@RunWith(Parameterized.class) public class ITTransactionManagerAsyncTest { + @Parameter public Executor executor; + + @Parameters(name = "executor = {0}") + public static Collection data() { + return Arrays.asList( + new Object[][] { + {MoreExecutors.directExecutor()}, + {Executors.newSingleThreadExecutor()}, + {Executors.newFixedThreadPool(4)} + }); + } + @ClassRule public static IntegrationTestEnv env = new IntegrationTestEnv(); private static Database db; private static DatabaseClient client; @@ -87,7 +105,8 @@ public ApiFuture apply(TransactionContext txn, Void input) .build()); return ApiFutures.immediateFuture(null); } - }) + }, + executor) .commitAsync() .get(); assertThat(manager.getState()).isEqualTo(TransactionState.COMMITTED); @@ -124,7 +143,8 @@ public ApiFuture apply(TransactionContext txn, Void input) .build()); return ApiFutures.immediateFuture(null); } - }) + }, + executor) .commitAsync() .get(); fail("Expected exception"); @@ -168,7 +188,8 @@ public ApiFuture apply(TransactionContext txn, Void input) throws Exceptio .build()); return ApiFutures.immediateFuture(null); } - }); + }, + executor); try { manager.rollbackAsync(); break; @@ -209,7 +230,8 @@ public ApiFuture apply(TransactionContext txn, Void input) throws Exception { return txn.readRowAsync("T", Key.of("Key3"), Arrays.asList("K", "BoolValue")); } - }); + }, + executor); manager2 = client.transactionManagerAsync(); txn2 = manager2.beginAsync(); txn2Step1 = @@ -220,7 +242,8 @@ public ApiFuture apply(TransactionContext txn, Void input) throws Exception { return txn.readRowAsync("T", Key.of("Key3"), Arrays.asList("K", "BoolValue")); } - }); + }, + executor); AsyncTransactionStep txn1Step2 = txn1Step1.then( @@ -237,7 +260,8 @@ public ApiFuture apply(TransactionContext txn, Struct input) .build()); return ApiFutures.immediateFuture(null); } - }); + }, + executor); txn2Step1.get(); txn1Step2.commitAsync().get(); @@ -274,7 +298,8 @@ public ApiFuture apply(TransactionContext txn, Void input) throws Exceptio .build()); return ApiFutures.immediateFuture(null); } - }); + }, + executor); txn2Step2.commitAsync().get(); Struct row = client.singleUse().readRow("T", Key.of("Key3"), Arrays.asList("K", "BoolValue")); assertThat(row.getString(0)).isEqualTo("Key3"); From 86f85c0545368160b7d5f93fb0492792f3337014 Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Sun, 21 Jun 2020 19:56:58 +0200 Subject: [PATCH 44/49] revert: remove async connection api from branch --- .../AbstractMultiUseTransaction.java | 11 -- .../connection/AsyncChecksumResultSet.java | 73 ----------- .../cloud/spanner/connection/Connection.java | 3 - .../spanner/connection/ConnectionImpl.java | 49 -------- .../cloud/spanner/connection/DdlBatch.java | 33 ++--- .../cloud/spanner/connection/DmlBatch.java | 8 -- .../connection/ReadWriteTransaction.java | 44 ------- .../connection/SingleUseTransaction.java | 74 ++++-------- .../connection/StatementResultImpl.java | 1 + .../cloud/spanner/connection/UnitOfWork.java | 4 - .../connection/AbstractMockServerTest.java | 7 -- .../connection/AsyncConnectionApiTest.java | 113 ------------------ .../connection/ReadOnlyTransactionTest.java | 11 +- .../connection/SingleUseTransactionTest.java | 12 +- 14 files changed, 43 insertions(+), 400 deletions(-) delete mode 100644 google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AsyncChecksumResultSet.java delete mode 100644 google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AsyncConnectionApiTest.java 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 9f278fb11db..cb8cf3bc557 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,7 +16,6 @@ package com.google.cloud.spanner.connection; -import com.google.cloud.spanner.AsyncResultSet; import com.google.cloud.spanner.ErrorCode; import com.google.cloud.spanner.Options.QueryOption; import com.google.cloud.spanner.ReadContext; @@ -74,16 +73,6 @@ public ResultSet call() throws Exception { }); } - @Override - public AsyncResultSet executeQueryAsync( - final ParsedStatement statement, - final AnalyzeMode analyzeMode, - final QueryOption... options) { - Preconditions.checkArgument(statement.isQuery(), "Statement is not a query"); - checkValidTransaction(); - return getReadContext().executeQueryAsync(statement.getStatement(), options); - } - ResultSet internalExecuteQuery( final ParsedStatement statement, AnalyzeMode analyzeMode, QueryOption... options) { if (analyzeMode == AnalyzeMode.NONE) { diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AsyncChecksumResultSet.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AsyncChecksumResultSet.java deleted file mode 100644 index 95c1077fa60..00000000000 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AsyncChecksumResultSet.java +++ /dev/null @@ -1,73 +0,0 @@ -/* - * Copyright 2020 Google LLC - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package com.google.cloud.spanner.connection; - -import com.google.api.core.ApiFuture; -import com.google.cloud.spanner.AsyncResultSet; -import com.google.cloud.spanner.Options.QueryOption; -import com.google.cloud.spanner.SpannerException; -import com.google.cloud.spanner.StructReader; -import com.google.cloud.spanner.connection.StatementParser.ParsedStatement; -import com.google.common.base.Function; -import com.google.common.collect.ImmutableList; -import java.util.concurrent.Executor; - -class AsyncChecksumResultSet extends ChecksumResultSet implements AsyncResultSet { - private AsyncResultSet delegate; - - AsyncChecksumResultSet( - ReadWriteTransaction transaction, - AsyncResultSet delegate, - ParsedStatement statement, - AnalyzeMode analyzeMode, - QueryOption... options) { - super(transaction, delegate, statement, analyzeMode, options); - this.delegate = delegate; - } - - @Override - public CursorState tryNext() throws SpannerException { - return delegate.tryNext(); - } - - @Override - public ApiFuture setCallback(Executor exec, ReadyCallback cb) { - return delegate.setCallback(exec, cb); - } - - @Override - public void cancel() { - delegate.cancel(); - } - - @Override - public void resume() { - delegate.resume(); - } - - @Override - public ApiFuture> toListAsync( - Function transformer, Executor executor) { - return delegate.toListAsync(transformer, executor); - } - - @Override - public ImmutableList toList(Function transformer) - throws SpannerException { - return delegate.toList(transformer); - } -} 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 9a1dc69a0cd..5247ce2c130 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 @@ -20,7 +20,6 @@ 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; @@ -620,8 +619,6 @@ public interface Connection extends AutoCloseable { */ ResultSet executeQuery(Statement query, QueryOption... options); - AsyncResultSet executeQueryAsync(Statement query, QueryOption... options); - /** * Analyzes a query and returns query plan and/or query execution statistics information. * 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 1e37f3927e9..ce24791859e 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 @@ -17,14 +17,12 @@ package com.google.cloud.spanner.connection; 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; @@ -683,11 +681,6 @@ 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); @@ -724,38 +717,6 @@ private ResultSet parseAndExecuteQuery( "Statement is not a query: " + parsedStatement.getSqlWithoutComments()); } - /** - * Parses the given statement as a query and executes it asynchronously. Throws a {@link - * SpannerException} if the statement is not a query. - */ - private AsyncResultSet parseAndExecuteQueryAsync( - Statement query, AnalyzeMode analyzeMode, QueryOption... options) { - Preconditions.checkNotNull(query); - Preconditions.checkNotNull(analyzeMode); - 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()); - 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); @@ -826,16 +787,6 @@ private ResultSet internalExecuteQuery( } } - 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 transaction.executeQueryAsync(statement, analyzeMode, options); - } - private long internalExecuteUpdate(final ParsedStatement update) { Preconditions.checkArgument( update.getType() == StatementType.UPDATE, "Statement must be an update"); 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 b49f443227b..b18f3fa891c 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 @@ -18,7 +18,6 @@ import com.google.api.gax.longrunning.OperationFuture; 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; @@ -115,27 +114,6 @@ public boolean isReadOnly() { @Override public ResultSet executeQuery( final ParsedStatement statement, AnalyzeMode analyzeMode, QueryOption... options) { - final QueryOption[] internalOptions = verifyQueryForDdlBatch(statement, analyzeMode, options); - Callable callable = - new Callable() { - @Override - public ResultSet call() throws Exception { - return DirectExecuteResultSet.ofResultSet( - dbClient.singleUse().executeQuery(statement.getStatement(), internalOptions)); - } - }; - return asyncExecuteStatement(statement, callable); - } - - @Override - public AsyncResultSet executeQueryAsync( - final ParsedStatement statement, AnalyzeMode analyzeMode, QueryOption... options) { - final QueryOption[] internalOptions = verifyQueryForDdlBatch(statement, analyzeMode, options); - return dbClient.singleUse().executeQueryAsync(statement.getStatement(), internalOptions); - } - - private QueryOption[] verifyQueryForDdlBatch( - ParsedStatement statement, AnalyzeMode analyzeMode, QueryOption... options) { if (options != null) { for (int i = 0; i < options.length; i++) { if (options[i] instanceof InternalMetadataQuery) { @@ -146,7 +124,16 @@ private QueryOption[] verifyQueryForDdlBatch( // Queries marked with internal metadata queries are allowed during a DDL batch. // These can only be generated by library internal methods and may be used to check // whether a database object such as table or an index exists. - return ArrayUtils.remove(options, i); + final QueryOption[] internalOptions = ArrayUtils.remove(options, i); + Callable callable = + new Callable() { + @Override + public ResultSet call() throws Exception { + return DirectExecuteResultSet.ofResultSet( + dbClient.singleUse().executeQuery(statement.getStatement(), internalOptions)); + } + }; + return asyncExecuteStatement(statement, callable); } } } 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 250e7a1cc72..ff38338d623 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 @@ -17,7 +17,6 @@ package com.google.cloud.spanner.connection; import com.google.cloud.Timestamp; -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; @@ -94,13 +93,6 @@ public ResultSet executeQuery( ErrorCode.FAILED_PRECONDITION, "Executing queries is not allowed for DML batches."); } - @Override - public AsyncResultSet executeQueryAsync( - ParsedStatement statement, AnalyzeMode analyzeMode, QueryOption... options) { - throw SpannerExceptionFactory.newSpannerException( - ErrorCode.FAILED_PRECONDITION, "Executing queries is not allowed for DML batches."); - } - @Override public Timestamp getReadTimestamp() { throw SpannerExceptionFactory.newSpannerException( 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 3689d4b8d95..7a0155cbfb8 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 @@ -21,7 +21,6 @@ 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.DatabaseClient; import com.google.cloud.spanner.ErrorCode; import com.google.cloud.spanner.Mutation; @@ -280,21 +279,6 @@ public ResultSet call() throws Exception { } } - @Override - public AsyncResultSet executeQueryAsync( - final ParsedStatement statement, - final AnalyzeMode analyzeMode, - final QueryOption... options) { - Preconditions.checkArgument(statement.isQuery(), "Statement is not a query"); - checkValidTransaction(); - if (retryAbortsInternally) { - AsyncResultSet delegate = super.executeQueryAsync(statement, analyzeMode, options); - return createAndAddAsyncRetryResultSet(delegate, statement, analyzeMode, options); - } else { - return super.executeQueryAsync(statement, analyzeMode, options); - } - } - @Override public long executeUpdate(final ParsedStatement update) { Preconditions.checkNotNull(update); @@ -558,24 +542,6 @@ private ResultSet createAndAddRetryResultSet( return resultSet; } - /** - * Registers a {@link AsyncResultSet} on this transaction that must be checked during a retry, and - * returns a retryable {@link AsyncResultSet}. - */ - private AsyncResultSet createAndAddAsyncRetryResultSet( - AsyncResultSet resultSet, - ParsedStatement statement, - AnalyzeMode analyzeMode, - QueryOption... options) { - if (retryAbortsInternally) { - AsyncChecksumResultSet checksumResultSet = - createAsyncChecksumResultSet(resultSet, statement, analyzeMode, options); - addRetryStatement(checksumResultSet); - return checksumResultSet; - } - return resultSet; - } - /** Registers the statement as a query that should return an error during a retry. */ private void createAndAddFailedQuery( SpannerException e, @@ -793,14 +759,4 @@ ChecksumResultSet createChecksumResultSet( QueryOption... options) { return new ChecksumResultSet(this, delegate, statement, analyzeMode, options); } - - /** Creates a {@link AsyncChecksumResultSet} for this {@link ReadWriteTransaction}. */ - @VisibleForTesting - AsyncChecksumResultSet createAsyncChecksumResultSet( - AsyncResultSet delegate, - ParsedStatement statement, - AnalyzeMode analyzeMode, - QueryOption... options) { - return new AsyncChecksumResultSet(this, delegate, statement, analyzeMode, options); - } } 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 3da17cf26d4..614d0c61e52 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 @@ -19,7 +19,6 @@ import com.google.api.gax.longrunning.OperationFuture; import com.google.cloud.Timestamp; import com.google.cloud.spanner.AbortedException; -import com.google.cloud.spanner.AsyncResultSet; import com.google.cloud.spanner.DatabaseClient; import com.google.cloud.spanner.ErrorCode; import com.google.cloud.spanner.Mutation; @@ -67,7 +66,7 @@ class SingleUseTransaction extends AbstractBaseUnitOfWork { private final DatabaseClient dbClient; private final TimestampBound readOnlyStaleness; private final AutocommitDmlMode autocommitDmlMode; - private ReadOnlyTransaction readOnlyTransaction; + private Timestamp readTimestamp = null; private volatile TransactionManager txManager; private TransactionRunner writeTransaction; private boolean used = false; @@ -169,81 +168,52 @@ public ResultSet executeQuery( Preconditions.checkArgument(statement.isQuery(), "Statement is not a query"); checkAndMarkUsed(); - readOnlyTransaction = dbClient.singleUseReadOnlyTransaction(readOnlyStaleness); + final ReadOnlyTransaction currentTransaction = + dbClient.singleUseReadOnlyTransaction(readOnlyStaleness); Callable callable = new Callable() { @Override public ResultSet call() throws Exception { - // try { - ResultSet rs; - if (analyzeMode == AnalyzeMode.NONE) { - rs = readOnlyTransaction.executeQuery(statement.getStatement(), options); - } else { - rs = - readOnlyTransaction.analyzeQuery( - statement.getStatement(), analyzeMode.getQueryAnalyzeMode()); + try { + ResultSet rs; + if (analyzeMode == AnalyzeMode.NONE) { + rs = currentTransaction.executeQuery(statement.getStatement(), options); + } else { + rs = + currentTransaction.analyzeQuery( + statement.getStatement(), analyzeMode.getQueryAnalyzeMode()); + } + // 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 { + currentTransaction.close(); } - // 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); - // } catch (Exception e) { - // readOnlyTransaction.close(); - // throw e; - // } finally { - // readOnlyTransaction.close(); - // currentTransaction.close(); - // } } }; try { ResultSet res = asyncExecuteStatement(statement, callable); - // readTimestamp = currentTransaction.getReadTimestamp(); + readTimestamp = currentTransaction.getReadTimestamp(); state = UnitOfWorkState.COMMITTED; return res; } catch (Throwable e) { state = UnitOfWorkState.COMMIT_FAILED; throw e; } finally { - readOnlyTransaction.close(); - } - } - - @Override - public AsyncResultSet executeQueryAsync( - final ParsedStatement statement, - final AnalyzeMode analyzeMode, - final QueryOption... options) { - Preconditions.checkNotNull(statement); - Preconditions.checkArgument(statement.isQuery(), "Statement is not a query"); - checkAndMarkUsed(); - - readOnlyTransaction = dbClient.singleUseReadOnlyTransaction(readOnlyStaleness); - try { - AsyncResultSet res = readOnlyTransaction.executeQueryAsync(statement.getStatement(), options); - state = UnitOfWorkState.COMMITTED; - return res; - } catch (Throwable e) { - readOnlyTransaction.close(); - state = UnitOfWorkState.COMMIT_FAILED; - throw e; - // } finally { - // currentTransaction.close(); + currentTransaction.close(); } } @Override public Timestamp getReadTimestamp() { ConnectionPreconditions.checkState( - readOnlyTransaction != null && state == UnitOfWorkState.COMMITTED, - "There is no read timestamp available for this transaction."); - return readOnlyTransaction.getReadTimestamp(); + readTimestamp != null, "There is no read timestamp available for this transaction."); + return readTimestamp; } @Override public Timestamp getReadTimestampOrNull() { - return readOnlyTransaction == null || state != UnitOfWorkState.COMMITTED - ? null - : readOnlyTransaction.getReadTimestamp(); + return readTimestamp; } private boolean hasCommitTimestamp() { 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 ab5610d0723..6221cc447b6 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 @@ -26,6 +26,7 @@ /** Implementation of {@link StatementResult} */ class StatementResultImpl implements StatementResult { + /** {@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 49001cd8d8f..e372229c64c 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 @@ -18,7 +18,6 @@ import com.google.api.core.InternalApi; import com.google.cloud.Timestamp; -import com.google.cloud.spanner.AsyncResultSet; import com.google.cloud.spanner.Mutation; import com.google.cloud.spanner.Options.QueryOption; import com.google.cloud.spanner.ReadContext; @@ -115,9 +114,6 @@ public boolean isActive() { ResultSet executeQuery( ParsedStatement statement, AnalyzeMode analyzeMode, QueryOption... options); - AsyncResultSet executeQueryAsync( - ParsedStatement statement, AnalyzeMode analyzeMode, QueryOption... options); - /** * @return the read timestamp of this transaction. Will throw a {@link SpannerException} if there * is no read timestamp. 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 29532c54830..3497b42bc7d 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 @@ -89,11 +89,6 @@ public abstract class AbstractMockServerTest { Statement.of("INSERT INTO TEST (ID, NAME) VALUES (1, 'test aborted')"); public static final int UPDATE_COUNT = 1; - 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; @@ -117,8 +112,6 @@ public static void startStaticServer() throws IOException { 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 diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AsyncConnectionApiTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AsyncConnectionApiTest.java deleted file mode 100644 index c1381f1b241..00000000000 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AsyncConnectionApiTest.java +++ /dev/null @@ -1,113 +0,0 @@ -/* - * Copyright 2020 Google LLC - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package com.google.cloud.spanner.connection; - -import static com.google.common.truth.Truth.assertThat; - -import com.google.api.core.ApiFuture; -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.connection.ITAbstractSpannerTest.ITConnection; -import com.google.common.base.Function; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.Executors; -import java.util.concurrent.atomic.AtomicInteger; -import org.junit.AfterClass; -import org.junit.Test; -import org.junit.runner.RunWith; -import org.junit.runners.JUnit4; - -@RunWith(JUnit4.class) -public class AsyncConnectionApiTest extends AbstractMockServerTest { - private static final ExecutorService executor = Executors.newSingleThreadExecutor(); - - @AfterClass - public static void stopExecutor() { - executor.shutdown(); - } - - @Test - public void testSimpleSelectAutocommit() throws Exception { - testSimpleSelect( - new Function() { - @Override - public Void apply(Connection input) { - input.setAutocommit(true); - return null; - } - }); - } - - @Test - public void testSimpleSelectReadOnly() throws Exception { - testSimpleSelect( - new Function() { - @Override - public Void apply(Connection input) { - input.setReadOnly(true); - return null; - } - }); - } - - @Test - public void testSimpleSelectReadWrite() throws Exception { - testSimpleSelect( - new Function() { - @Override - public Void apply(Connection input) { - return null; - } - }); - } - - private void testSimpleSelect(Function connectionConfigurator) - throws Exception { - final AtomicInteger rowCount = new AtomicInteger(); - ApiFuture res; - try (ITConnection connection = createConnection()) { - connectionConfigurator.apply(connection); - // Verify that the call is non-blocking. - // mockSpanner.freeze(); - try (AsyncResultSet rs = connection.executeQueryAsync(SELECT_RANDOM_STATEMENT)) { - // mockSpanner.unfreeze(); - 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; - } - } - } - }); - } - res.get(); - assertThat(rowCount.get()).isEqualTo(RANDOM_RESULT_SET_ROW_COUNT); - } - } -} 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 2266a88a259..6918c9f268d 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 @@ -25,7 +25,6 @@ import static org.junit.Assert.fail; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; - import com.google.api.core.ApiFuture; import com.google.cloud.Timestamp; import com.google.cloud.spanner.AsyncResultSet; @@ -137,29 +136,29 @@ public Timestamp getReadTimestamp() { @Override public AsyncResultSet readAsync( String table, KeySet keys, Iterable columns, ReadOption... options) { - throw new UnsupportedOperationException(); + return null; } @Override public AsyncResultSet readUsingIndexAsync( String table, String index, KeySet keys, Iterable columns, ReadOption... options) { - throw new UnsupportedOperationException(); + return null; } @Override public ApiFuture readRowAsync(String table, Key key, Iterable columns) { - throw new UnsupportedOperationException(); + return null; } @Override public ApiFuture readRowUsingIndexAsync( String table, String index, Key key, Iterable columns) { - throw new UnsupportedOperationException(); + return null; } @Override public AsyncResultSet executeQueryAsync(Statement statement, QueryOption... options) { - throw new UnsupportedOperationException(); + return null; } } 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 e7ac02c49c8..1d4b9a99f1d 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 @@ -23,7 +23,6 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; - import com.google.api.core.ApiFuture; import com.google.api.gax.longrunning.OperationFuture; import com.google.cloud.Timestamp; @@ -236,29 +235,29 @@ public Timestamp getReadTimestamp() { @Override public AsyncResultSet readAsync( String table, KeySet keys, Iterable columns, ReadOption... options) { - throw new UnsupportedOperationException(); + return null; } @Override public AsyncResultSet readUsingIndexAsync( String table, String index, KeySet keys, Iterable columns, ReadOption... options) { - throw new UnsupportedOperationException(); + return null; } @Override public ApiFuture readRowAsync(String table, Key key, Iterable columns) { - throw new UnsupportedOperationException(); + return null; } @Override public ApiFuture readRowUsingIndexAsync( String table, String index, Key key, Iterable columns) { - throw new UnsupportedOperationException(); + return null; } @Override public AsyncResultSet executeQueryAsync(Statement statement, QueryOption... options) { - throw new UnsupportedOperationException(); + return null; } } @@ -751,7 +750,6 @@ public void testExecuteQueryWithTimeout() { SingleUseTransaction subject = createSubjectWithTimeout(1L); try { subject.executeQuery(createParsedQuery(SLOW_QUERY), AnalyzeMode.NONE); - fail("missing expected exception"); } catch (SpannerException e) { if (e.getErrorCode() != ErrorCode.DEADLINE_EXCEEDED) { throw e; From cfa8e79955e92c6b132b980646db5f9f392e9cd4 Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Sun, 21 Jun 2020 19:56:58 +0200 Subject: [PATCH 45/49] revert: remove async connection api from branch --- .../connection/ReadOnlyTransactionTest.java | 14 +++++++------- .../connection/SingleUseTransactionTest.java | 18 +++++++++--------- 2 files changed, 16 insertions(+), 16 deletions(-) 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 6918c9f268d..71be33581b1 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 @@ -25,6 +25,13 @@ import static org.junit.Assert.fail; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; +import java.util.Arrays; +import java.util.Calendar; +import java.util.List; +import java.util.concurrent.TimeUnit; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; import com.google.api.core.ApiFuture; import com.google.cloud.Timestamp; import com.google.cloud.spanner.AsyncResultSet; @@ -45,13 +52,6 @@ import com.google.cloud.spanner.connection.StatementParser.StatementType; import com.google.cloud.spanner.connection.UnitOfWork.UnitOfWorkState; import com.google.spanner.v1.ResultSetStats; -import java.util.Arrays; -import java.util.Calendar; -import java.util.List; -import java.util.concurrent.TimeUnit; -import org.junit.Test; -import org.junit.runner.RunWith; -import org.junit.runners.JUnit4; @RunWith(JUnit4.class) public class ReadOnlyTransactionTest { 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 1d4b9a99f1d..c1fa6045f38 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 @@ -23,6 +23,15 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; +import java.util.Arrays; +import java.util.Calendar; +import java.util.List; +import java.util.concurrent.TimeUnit; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; +import org.mockito.invocation.InvocationOnMock; +import org.mockito.stubbing.Answer; import com.google.api.core.ApiFuture; import com.google.api.gax.longrunning.OperationFuture; import com.google.cloud.Timestamp; @@ -48,15 +57,6 @@ import com.google.cloud.spanner.connection.StatementParser.StatementType; 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.List; -import java.util.concurrent.TimeUnit; -import org.junit.Test; -import org.junit.runner.RunWith; -import org.junit.runners.JUnit4; -import org.mockito.invocation.InvocationOnMock; -import org.mockito.stubbing.Answer; @RunWith(JUnit4.class) public class SingleUseTransactionTest { From 46708c8e6e3f60edad29227c1525d3939b205744 Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Sun, 21 Jun 2020 20:26:11 +0200 Subject: [PATCH 46/49] feat: async connection api --- .../AbstractMultiUseTransaction.java | 11 ++ .../connection/AsyncChecksumResultSet.java | 73 +++++++++++ .../cloud/spanner/connection/Connection.java | 3 + .../spanner/connection/ConnectionImpl.java | 49 ++++++++ .../cloud/spanner/connection/DdlBatch.java | 33 +++-- .../cloud/spanner/connection/DmlBatch.java | 8 ++ .../connection/ReadWriteTransaction.java | 44 +++++++ .../connection/SingleUseTransaction.java | 74 ++++++++---- .../connection/StatementResultImpl.java | 1 - .../cloud/spanner/connection/UnitOfWork.java | 4 + .../connection/AbstractMockServerTest.java | 7 ++ .../connection/AsyncConnectionApiTest.java | 113 ++++++++++++++++++ .../ReadOnlyTransactionTest.java.rej | 46 +++++++ .../connection/SingleUseTransactionTest.java | 1 + .../SingleUseTransactionTest.java.rej | 55 +++++++++ 15 files changed, 489 insertions(+), 33 deletions(-) create mode 100644 google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AsyncChecksumResultSet.java create mode 100644 google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AsyncConnectionApiTest.java create mode 100644 google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ReadOnlyTransactionTest.java.rej create mode 100644 google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/SingleUseTransactionTest.java.rej 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 cb8cf3bc557..9f278fb11db 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.cloud.spanner.AsyncResultSet; import com.google.cloud.spanner.ErrorCode; import com.google.cloud.spanner.Options.QueryOption; import com.google.cloud.spanner.ReadContext; @@ -73,6 +74,16 @@ public ResultSet call() throws Exception { }); } + @Override + public AsyncResultSet executeQueryAsync( + final ParsedStatement statement, + final AnalyzeMode analyzeMode, + final QueryOption... options) { + Preconditions.checkArgument(statement.isQuery(), "Statement is not a query"); + checkValidTransaction(); + return getReadContext().executeQueryAsync(statement.getStatement(), options); + } + ResultSet internalExecuteQuery( final ParsedStatement statement, AnalyzeMode analyzeMode, QueryOption... options) { if (analyzeMode == AnalyzeMode.NONE) { diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AsyncChecksumResultSet.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AsyncChecksumResultSet.java new file mode 100644 index 00000000000..95c1077fa60 --- /dev/null +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AsyncChecksumResultSet.java @@ -0,0 +1,73 @@ +/* + * 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.cloud.spanner.AsyncResultSet; +import com.google.cloud.spanner.Options.QueryOption; +import com.google.cloud.spanner.SpannerException; +import com.google.cloud.spanner.StructReader; +import com.google.cloud.spanner.connection.StatementParser.ParsedStatement; +import com.google.common.base.Function; +import com.google.common.collect.ImmutableList; +import java.util.concurrent.Executor; + +class AsyncChecksumResultSet extends ChecksumResultSet implements AsyncResultSet { + private AsyncResultSet delegate; + + AsyncChecksumResultSet( + ReadWriteTransaction transaction, + AsyncResultSet delegate, + ParsedStatement statement, + AnalyzeMode analyzeMode, + QueryOption... options) { + super(transaction, delegate, statement, analyzeMode, options); + this.delegate = delegate; + } + + @Override + public CursorState tryNext() throws SpannerException { + return delegate.tryNext(); + } + + @Override + public ApiFuture setCallback(Executor exec, ReadyCallback cb) { + return delegate.setCallback(exec, cb); + } + + @Override + public void cancel() { + delegate.cancel(); + } + + @Override + public void resume() { + delegate.resume(); + } + + @Override + public ApiFuture> toListAsync( + Function transformer, Executor executor) { + return delegate.toListAsync(transformer, executor); + } + + @Override + public ImmutableList toList(Function transformer) + throws SpannerException { + return delegate.toList(transformer); + } +} 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 5247ce2c130..9a1dc69a0cd 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 @@ -20,6 +20,7 @@ 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; @@ -619,6 +620,8 @@ public interface Connection extends AutoCloseable { */ ResultSet executeQuery(Statement query, QueryOption... options); + AsyncResultSet executeQueryAsync(Statement query, QueryOption... options); + /** * Analyzes a query and returns query plan and/or query execution statistics information. * 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 ce24791859e..1e37f3927e9 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 @@ -17,12 +17,14 @@ package com.google.cloud.spanner.connection; 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; @@ -681,6 +683,11 @@ 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 +724,38 @@ private ResultSet parseAndExecuteQuery( "Statement is not a query: " + parsedStatement.getSqlWithoutComments()); } + /** + * Parses the given statement as a query and executes it asynchronously. Throws a {@link + * SpannerException} if the statement is not a query. + */ + private AsyncResultSet parseAndExecuteQueryAsync( + Statement query, AnalyzeMode analyzeMode, QueryOption... options) { + Preconditions.checkNotNull(query); + Preconditions.checkNotNull(analyzeMode); + 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()); + 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); @@ -787,6 +826,16 @@ private ResultSet internalExecuteQuery( } } + 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 transaction.executeQueryAsync(statement, analyzeMode, options); + } + private long internalExecuteUpdate(final ParsedStatement update) { Preconditions.checkArgument( update.getType() == StatementType.UPDATE, "Statement must be an update"); 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 b18f3fa891c..b49f443227b 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 @@ -18,6 +18,7 @@ import com.google.api.gax.longrunning.OperationFuture; 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; @@ -114,6 +115,27 @@ public boolean isReadOnly() { @Override public ResultSet executeQuery( final ParsedStatement statement, AnalyzeMode analyzeMode, QueryOption... options) { + final QueryOption[] internalOptions = verifyQueryForDdlBatch(statement, analyzeMode, options); + Callable callable = + new Callable() { + @Override + public ResultSet call() throws Exception { + return DirectExecuteResultSet.ofResultSet( + dbClient.singleUse().executeQuery(statement.getStatement(), internalOptions)); + } + }; + return asyncExecuteStatement(statement, callable); + } + + @Override + public AsyncResultSet executeQueryAsync( + final ParsedStatement statement, AnalyzeMode analyzeMode, QueryOption... options) { + final QueryOption[] internalOptions = verifyQueryForDdlBatch(statement, analyzeMode, options); + return dbClient.singleUse().executeQueryAsync(statement.getStatement(), internalOptions); + } + + private QueryOption[] verifyQueryForDdlBatch( + ParsedStatement statement, AnalyzeMode analyzeMode, QueryOption... options) { if (options != null) { for (int i = 0; i < options.length; i++) { if (options[i] instanceof InternalMetadataQuery) { @@ -124,16 +146,7 @@ public ResultSet executeQuery( // Queries marked with internal metadata queries are allowed during a DDL batch. // These can only be generated by library internal methods and may be used to check // whether a database object such as table or an index exists. - final QueryOption[] internalOptions = ArrayUtils.remove(options, i); - Callable callable = - new Callable() { - @Override - public ResultSet call() throws Exception { - return DirectExecuteResultSet.ofResultSet( - dbClient.singleUse().executeQuery(statement.getStatement(), internalOptions)); - } - }; - return asyncExecuteStatement(statement, callable); + return ArrayUtils.remove(options, i); } } } 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 ff38338d623..250e7a1cc72 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 @@ -17,6 +17,7 @@ package com.google.cloud.spanner.connection; import com.google.cloud.Timestamp; +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; @@ -93,6 +94,13 @@ public ResultSet executeQuery( ErrorCode.FAILED_PRECONDITION, "Executing queries is not allowed for DML batches."); } + @Override + public AsyncResultSet executeQueryAsync( + ParsedStatement statement, AnalyzeMode analyzeMode, QueryOption... options) { + throw SpannerExceptionFactory.newSpannerException( + ErrorCode.FAILED_PRECONDITION, "Executing queries is not allowed for DML batches."); + } + @Override public Timestamp getReadTimestamp() { throw SpannerExceptionFactory.newSpannerException( 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 7a0155cbfb8..3689d4b8d95 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 @@ -21,6 +21,7 @@ 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.DatabaseClient; import com.google.cloud.spanner.ErrorCode; import com.google.cloud.spanner.Mutation; @@ -279,6 +280,21 @@ public ResultSet call() throws Exception { } } + @Override + public AsyncResultSet executeQueryAsync( + final ParsedStatement statement, + final AnalyzeMode analyzeMode, + final QueryOption... options) { + Preconditions.checkArgument(statement.isQuery(), "Statement is not a query"); + checkValidTransaction(); + if (retryAbortsInternally) { + AsyncResultSet delegate = super.executeQueryAsync(statement, analyzeMode, options); + return createAndAddAsyncRetryResultSet(delegate, statement, analyzeMode, options); + } else { + return super.executeQueryAsync(statement, analyzeMode, options); + } + } + @Override public long executeUpdate(final ParsedStatement update) { Preconditions.checkNotNull(update); @@ -542,6 +558,24 @@ private ResultSet createAndAddRetryResultSet( return resultSet; } + /** + * Registers a {@link AsyncResultSet} on this transaction that must be checked during a retry, and + * returns a retryable {@link AsyncResultSet}. + */ + private AsyncResultSet createAndAddAsyncRetryResultSet( + AsyncResultSet resultSet, + ParsedStatement statement, + AnalyzeMode analyzeMode, + QueryOption... options) { + if (retryAbortsInternally) { + AsyncChecksumResultSet checksumResultSet = + createAsyncChecksumResultSet(resultSet, statement, analyzeMode, options); + addRetryStatement(checksumResultSet); + return checksumResultSet; + } + return resultSet; + } + /** Registers the statement as a query that should return an error during a retry. */ private void createAndAddFailedQuery( SpannerException e, @@ -759,4 +793,14 @@ ChecksumResultSet createChecksumResultSet( QueryOption... options) { return new ChecksumResultSet(this, delegate, statement, analyzeMode, options); } + + /** Creates a {@link AsyncChecksumResultSet} for this {@link ReadWriteTransaction}. */ + @VisibleForTesting + AsyncChecksumResultSet createAsyncChecksumResultSet( + AsyncResultSet delegate, + ParsedStatement statement, + AnalyzeMode analyzeMode, + QueryOption... options) { + return new AsyncChecksumResultSet(this, delegate, statement, analyzeMode, options); + } } 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 614d0c61e52..3da17cf26d4 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 @@ -19,6 +19,7 @@ import com.google.api.gax.longrunning.OperationFuture; import com.google.cloud.Timestamp; import com.google.cloud.spanner.AbortedException; +import com.google.cloud.spanner.AsyncResultSet; import com.google.cloud.spanner.DatabaseClient; import com.google.cloud.spanner.ErrorCode; import com.google.cloud.spanner.Mutation; @@ -66,7 +67,7 @@ class SingleUseTransaction extends AbstractBaseUnitOfWork { private final DatabaseClient dbClient; private final TimestampBound readOnlyStaleness; private final AutocommitDmlMode autocommitDmlMode; - private Timestamp readTimestamp = null; + private ReadOnlyTransaction readOnlyTransaction; private volatile TransactionManager txManager; private TransactionRunner writeTransaction; private boolean used = false; @@ -168,52 +169,81 @@ public ResultSet executeQuery( Preconditions.checkArgument(statement.isQuery(), "Statement is not a query"); checkAndMarkUsed(); - final ReadOnlyTransaction currentTransaction = - dbClient.singleUseReadOnlyTransaction(readOnlyStaleness); + readOnlyTransaction = dbClient.singleUseReadOnlyTransaction(readOnlyStaleness); Callable callable = new Callable() { @Override public ResultSet call() throws Exception { - try { - ResultSet rs; - if (analyzeMode == AnalyzeMode.NONE) { - rs = currentTransaction.executeQuery(statement.getStatement(), options); - } else { - rs = - currentTransaction.analyzeQuery( - statement.getStatement(), analyzeMode.getQueryAnalyzeMode()); - } - // 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 { - currentTransaction.close(); + // try { + ResultSet rs; + if (analyzeMode == AnalyzeMode.NONE) { + rs = readOnlyTransaction.executeQuery(statement.getStatement(), options); + } else { + rs = + readOnlyTransaction.analyzeQuery( + statement.getStatement(), analyzeMode.getQueryAnalyzeMode()); } + // 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); + // } catch (Exception e) { + // readOnlyTransaction.close(); + // throw e; + // } finally { + // readOnlyTransaction.close(); + // currentTransaction.close(); + // } } }; try { ResultSet res = asyncExecuteStatement(statement, callable); - readTimestamp = currentTransaction.getReadTimestamp(); + // readTimestamp = currentTransaction.getReadTimestamp(); state = UnitOfWorkState.COMMITTED; return res; } catch (Throwable e) { state = UnitOfWorkState.COMMIT_FAILED; throw e; } finally { - currentTransaction.close(); + readOnlyTransaction.close(); + } + } + + @Override + public AsyncResultSet executeQueryAsync( + final ParsedStatement statement, + final AnalyzeMode analyzeMode, + final QueryOption... options) { + Preconditions.checkNotNull(statement); + Preconditions.checkArgument(statement.isQuery(), "Statement is not a query"); + checkAndMarkUsed(); + + readOnlyTransaction = dbClient.singleUseReadOnlyTransaction(readOnlyStaleness); + try { + AsyncResultSet res = readOnlyTransaction.executeQueryAsync(statement.getStatement(), options); + state = UnitOfWorkState.COMMITTED; + return res; + } catch (Throwable e) { + readOnlyTransaction.close(); + state = UnitOfWorkState.COMMIT_FAILED; + throw e; + // } finally { + // currentTransaction.close(); } } @Override public Timestamp getReadTimestamp() { ConnectionPreconditions.checkState( - readTimestamp != null, "There is no read timestamp available for this transaction."); - return readTimestamp; + readOnlyTransaction != null && state == UnitOfWorkState.COMMITTED, + "There is no read timestamp available for this transaction."); + return readOnlyTransaction.getReadTimestamp(); } @Override public Timestamp getReadTimestampOrNull() { - return readTimestamp; + return readOnlyTransaction == null || state != UnitOfWorkState.COMMITTED + ? null + : readOnlyTransaction.getReadTimestamp(); } private boolean hasCommitTimestamp() { 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 6221cc447b6..ab5610d0723 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 @@ -26,7 +26,6 @@ /** Implementation of {@link StatementResult} */ class StatementResultImpl implements StatementResult { - /** {@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 e372229c64c..49001cd8d8f 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 @@ -18,6 +18,7 @@ import com.google.api.core.InternalApi; import com.google.cloud.Timestamp; +import com.google.cloud.spanner.AsyncResultSet; import com.google.cloud.spanner.Mutation; import com.google.cloud.spanner.Options.QueryOption; import com.google.cloud.spanner.ReadContext; @@ -114,6 +115,9 @@ public boolean isActive() { ResultSet executeQuery( ParsedStatement statement, AnalyzeMode analyzeMode, QueryOption... options); + AsyncResultSet executeQueryAsync( + ParsedStatement statement, AnalyzeMode analyzeMode, QueryOption... options); + /** * @return the read timestamp of this transaction. Will throw a {@link SpannerException} if there * is no read timestamp. 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 3497b42bc7d..29532c54830 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 @@ -89,6 +89,11 @@ public abstract class AbstractMockServerTest { Statement.of("INSERT INTO TEST (ID, NAME) VALUES (1, 'test aborted')"); public static final int UPDATE_COUNT = 1; + 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; @@ -112,6 +117,8 @@ public static void startStaticServer() throws IOException { 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 diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AsyncConnectionApiTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AsyncConnectionApiTest.java new file mode 100644 index 00000000000..c1381f1b241 --- /dev/null +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AsyncConnectionApiTest.java @@ -0,0 +1,113 @@ +/* + * 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.common.truth.Truth.assertThat; + +import com.google.api.core.ApiFuture; +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.connection.ITAbstractSpannerTest.ITConnection; +import com.google.common.base.Function; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.AfterClass; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class AsyncConnectionApiTest extends AbstractMockServerTest { + private static final ExecutorService executor = Executors.newSingleThreadExecutor(); + + @AfterClass + public static void stopExecutor() { + executor.shutdown(); + } + + @Test + public void testSimpleSelectAutocommit() throws Exception { + testSimpleSelect( + new Function() { + @Override + public Void apply(Connection input) { + input.setAutocommit(true); + return null; + } + }); + } + + @Test + public void testSimpleSelectReadOnly() throws Exception { + testSimpleSelect( + new Function() { + @Override + public Void apply(Connection input) { + input.setReadOnly(true); + return null; + } + }); + } + + @Test + public void testSimpleSelectReadWrite() throws Exception { + testSimpleSelect( + new Function() { + @Override + public Void apply(Connection input) { + return null; + } + }); + } + + private void testSimpleSelect(Function connectionConfigurator) + throws Exception { + final AtomicInteger rowCount = new AtomicInteger(); + ApiFuture res; + try (ITConnection connection = createConnection()) { + connectionConfigurator.apply(connection); + // Verify that the call is non-blocking. + // mockSpanner.freeze(); + try (AsyncResultSet rs = connection.executeQueryAsync(SELECT_RANDOM_STATEMENT)) { + // mockSpanner.unfreeze(); + 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; + } + } + } + }); + } + res.get(); + assertThat(rowCount.get()).isEqualTo(RANDOM_RESULT_SET_ROW_COUNT); + } + } +} diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ReadOnlyTransactionTest.java.rej b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ReadOnlyTransactionTest.java.rej new file mode 100644 index 00000000000..5544db0c0ef --- /dev/null +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ReadOnlyTransactionTest.java.rej @@ -0,0 +1,46 @@ +diff 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 (rejected hunks) +@@ -26,7 +26,9 @@ import static org.junit.Assert.fail; + import static org.mockito.Mockito.mock; + import static org.mockito.Mockito.when; + ++import com.google.api.core.ApiFuture; + 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.Key; +@@ -131,6 +133,34 @@ public class ReadOnlyTransactionTest { + public Timestamp getReadTimestamp() { + return readTimestamp; + } ++ ++ @Override ++ public AsyncResultSet readAsync( ++ String table, KeySet keys, Iterable columns, ReadOption... options) { ++ throw new UnsupportedOperationException(); ++ } ++ ++ @Override ++ public AsyncResultSet readUsingIndexAsync( ++ String table, String index, KeySet keys, Iterable columns, ReadOption... options) { ++ throw new UnsupportedOperationException(); ++ } ++ ++ @Override ++ public ApiFuture readRowAsync(String table, Key key, Iterable columns) { ++ throw new UnsupportedOperationException(); ++ } ++ ++ @Override ++ public ApiFuture readRowUsingIndexAsync( ++ String table, String index, Key key, Iterable columns) { ++ throw new UnsupportedOperationException(); ++ } ++ ++ @Override ++ public AsyncResultSet executeQueryAsync(Statement statement, QueryOption... options) { ++ throw new UnsupportedOperationException(); ++ } + } + + private ReadOnlyTransaction createSubject() { 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 c1fa6045f38..67fc95f21c7 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 @@ -750,6 +750,7 @@ public void testExecuteQueryWithTimeout() { SingleUseTransaction subject = createSubjectWithTimeout(1L); try { subject.executeQuery(createParsedQuery(SLOW_QUERY), AnalyzeMode.NONE); + fail("missing expected exception"); } catch (SpannerException e) { if (e.getErrorCode() != ErrorCode.DEADLINE_EXCEEDED) { throw e; diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/SingleUseTransactionTest.java.rej b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/SingleUseTransactionTest.java.rej new file mode 100644 index 00000000000..9aafc978573 --- /dev/null +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/SingleUseTransactionTest.java.rej @@ -0,0 +1,55 @@ +diff 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 (rejected hunks) +@@ -24,8 +24,10 @@ import static org.mockito.Mockito.mock; + import static org.mockito.Mockito.verify; + import static org.mockito.Mockito.when; + ++import com.google.api.core.ApiFuture; + import com.google.api.gax.longrunning.OperationFuture; + 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.Key; +@@ -230,6 +232,34 @@ public class SingleUseTransactionTest { + public Timestamp getReadTimestamp() { + return readTimestamp; + } ++ ++ @Override ++ public AsyncResultSet readAsync( ++ String table, KeySet keys, Iterable columns, ReadOption... options) { ++ throw new UnsupportedOperationException(); ++ } ++ ++ @Override ++ public AsyncResultSet readUsingIndexAsync( ++ String table, String index, KeySet keys, Iterable columns, ReadOption... options) { ++ throw new UnsupportedOperationException(); ++ } ++ ++ @Override ++ public ApiFuture readRowAsync(String table, Key key, Iterable columns) { ++ throw new UnsupportedOperationException(); ++ } ++ ++ @Override ++ public ApiFuture readRowUsingIndexAsync( ++ String table, String index, Key key, Iterable columns) { ++ throw new UnsupportedOperationException(); ++ } ++ ++ @Override ++ public AsyncResultSet executeQueryAsync(Statement statement, QueryOption... options) { ++ throw new UnsupportedOperationException(); ++ } + } + + private DdlClient createDefaultMockDdlClient() { +@@ -721,6 +751,7 @@ public class SingleUseTransactionTest { + SingleUseTransaction subject = createSubjectWithTimeout(1L); + try { + subject.executeQuery(createParsedQuery(SLOW_QUERY), AnalyzeMode.NONE); ++ fail("missing expected exception"); + } catch (SpannerException e) { + if (e.getErrorCode() != ErrorCode.DEADLINE_EXCEEDED) { + throw e; From 6dfeac558d7e7b1b0296037ad40e217b50cf06b5 Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Sun, 21 Jun 2020 20:32:02 +0200 Subject: [PATCH 47/49] chore: run code formatter --- .../connection/ReadOnlyTransactionTest.java | 15 ++++++++------- .../connection/SingleUseTransactionTest.java | 19 ++++++++++--------- 2 files changed, 18 insertions(+), 16 deletions(-) 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 71be33581b1..118f596c868 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 @@ -25,13 +25,7 @@ import static org.junit.Assert.fail; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; -import java.util.Arrays; -import java.util.Calendar; -import java.util.List; -import java.util.concurrent.TimeUnit; -import org.junit.Test; -import org.junit.runner.RunWith; -import org.junit.runners.JUnit4; + import com.google.api.core.ApiFuture; import com.google.cloud.Timestamp; import com.google.cloud.spanner.AsyncResultSet; @@ -52,6 +46,13 @@ import com.google.cloud.spanner.connection.StatementParser.StatementType; import com.google.cloud.spanner.connection.UnitOfWork.UnitOfWorkState; import com.google.spanner.v1.ResultSetStats; +import java.util.Arrays; +import java.util.Calendar; +import java.util.List; +import java.util.concurrent.TimeUnit; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; @RunWith(JUnit4.class) public class ReadOnlyTransactionTest { 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 67fc95f21c7..3d94021c294 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 @@ -23,15 +23,7 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; -import java.util.Arrays; -import java.util.Calendar; -import java.util.List; -import java.util.concurrent.TimeUnit; -import org.junit.Test; -import org.junit.runner.RunWith; -import org.junit.runners.JUnit4; -import org.mockito.invocation.InvocationOnMock; -import org.mockito.stubbing.Answer; + import com.google.api.core.ApiFuture; import com.google.api.gax.longrunning.OperationFuture; import com.google.cloud.Timestamp; @@ -57,6 +49,15 @@ import com.google.cloud.spanner.connection.StatementParser.StatementType; 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.List; +import java.util.concurrent.TimeUnit; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; +import org.mockito.invocation.InvocationOnMock; +import org.mockito.stubbing.Answer; @RunWith(JUnit4.class) public class SingleUseTransactionTest { From 6763a59e2374eb4e0bbb4d2bbb8bc3dca195279d Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Sun, 21 Jun 2020 20:46:14 +0200 Subject: [PATCH 48/49] chore: fix flaky test case --- .../spanner/AsyncTransactionManagerTest.java | 16 ++++------------ 1 file changed, 4 insertions(+), 12 deletions(-) diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncTransactionManagerTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncTransactionManagerTest.java index c5de6542192..380fd4fd8c3 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncTransactionManagerTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncTransactionManagerTest.java @@ -480,17 +480,9 @@ public ApiFuture apply(TransactionContext txn, Void input) if (attempt.incrementAndGet() == 1) { mockSpanner.abortTransaction(txn); } - // This update statement will be aborted, but the error will not - // propagated to the - // transaction runner and cause the transaction to retry. Instead, the - // commit call - // will do that. + // This update statement will be aborted, but the error will not propagated to the transaction runner and cause the transaction to retry. Instead, the commit call will do that. txn.executeUpdateAsync(UPDATE_STATEMENT); - // Resolving this future will not resolve the result of the entire - // transaction. The - // transaction result will be resolved when the commit has actually - // finished - // successfully. + // Resolving this future will not resolve the result of the entire transaction. The transaction result will be resolved when the commit has actually finished successfully. return ApiFutures.immediateFuture(null); } }, @@ -498,12 +490,12 @@ public ApiFuture apply(TransactionContext txn, Void input) .commitAsync(); assertThat(ts.get()).isNotNull(); assertThat(attempt.get()).isEqualTo(2); + // The server may receive 1 or 2 commit requests, depending on whether the commitAsync() call already knows that the transaction has aborted or not. In case it is known that the transaction has aborted, it will not attempt to call the Commit RPC, and instead propagate the Aborted error directly. assertThat(mockSpanner.getRequestTypes()) - .containsExactly( + .containsAtLeast( BatchCreateSessionsRequest.class, BeginTransactionRequest.class, ExecuteSqlRequest.class, - CommitRequest.class, BeginTransactionRequest.class, ExecuteSqlRequest.class, CommitRequest.class); From 6843515629214f607feba983d99ec14d74cbf0a9 Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Thu, 6 Aug 2020 09:01:50 +0200 Subject: [PATCH 49/49] feat: async connection api --- .../spanner/PartitionedDMLTransaction.java | 109 --- .../google/cloud/spanner/SpannerOptions.java | 3 - .../connection/AbstractBaseConnection.java | 37 + .../spanner/connection/AsyncConnection.java | 636 ++++++++++++++++++ .../connection/AsyncConnectionImpl.java | 247 +++++++ .../connection/AsyncStatementResult.java | 47 ++ .../spanner/connection/BaseConnection.java | 343 ++++++++++ .../cloud/spanner/connection/Connection.java | 337 +--------- .../spanner/connection/ConnectionImpl.java | 5 - .../ConnectionStatementExecutorImpl.java | 6 +- .../cloud/spanner/connection/DdlBatch.java | 2 +- .../connection/AbstractMockServerTest.java | 4 + .../connection/AsyncConnectionApiTest.java | 21 +- .../spanner/connection/DdlBatchTest.java | 2 +- 14 files changed, 1330 insertions(+), 469 deletions(-) delete mode 100644 google-cloud-spanner/src/main/java/com/google/cloud/spanner/PartitionedDMLTransaction.java create mode 100644 google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AbstractBaseConnection.java create mode 100644 google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AsyncConnection.java create mode 100644 google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AsyncConnectionImpl.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/BaseConnection.java diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/PartitionedDMLTransaction.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/PartitionedDMLTransaction.java deleted file mode 100644 index 0b24e13ada5..00000000000 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/PartitionedDMLTransaction.java +++ /dev/null @@ -1,109 +0,0 @@ -/* - * Copyright 2019 Google LLC - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package com.google.cloud.spanner; - -import static com.google.common.base.Preconditions.checkState; - -import com.google.cloud.spanner.SessionImpl.SessionTransaction; -import com.google.cloud.spanner.spi.v1.SpannerRpc; -import com.google.protobuf.ByteString; -import com.google.spanner.v1.BeginTransactionRequest; -import com.google.spanner.v1.ExecuteSqlRequest; -import com.google.spanner.v1.ExecuteSqlRequest.QueryMode; -import com.google.spanner.v1.Transaction; -import com.google.spanner.v1.TransactionOptions; -import com.google.spanner.v1.TransactionSelector; -import io.opencensus.trace.Span; -import java.util.Map; -import java.util.concurrent.Callable; - -/** Partitioned DML transaction for bulk updates and deletes. */ -class PartitionedDMLTransaction implements SessionTransaction { - private final SessionImpl session; - private final SpannerRpc rpc; - private volatile boolean isValid = true; - - PartitionedDMLTransaction(SessionImpl session, SpannerRpc rpc) { - this.session = session; - this.rpc = rpc; - } - - private ByteString initTransaction() { - final BeginTransactionRequest request = - BeginTransactionRequest.newBuilder() - .setSession(session.getName()) - .setOptions( - TransactionOptions.newBuilder() - .setPartitionedDml(TransactionOptions.PartitionedDml.getDefaultInstance())) - .build(); - Transaction txn = rpc.beginTransaction(request, session.getOptions()); - if (txn.getId().isEmpty()) { - throw SpannerExceptionFactory.newSpannerException( - ErrorCode.INTERNAL, - "Failed to init transaction, missing transaction id\n" + session.getName()); - } - return txn.getId(); - } - - /** - * Executes the {@link Statement} using a partitioned dml transaction with automatic retry if the - * transaction was aborted. - */ - long executePartitionedUpdate(final Statement statement) { - checkState(isValid, "Partitioned DML has been invalidated by a new operation on the session"); - Callable callable = - new Callable() { - @Override - public com.google.spanner.v1.ResultSet call() throws Exception { - ByteString transactionId = initTransaction(); - final ExecuteSqlRequest.Builder builder = - ExecuteSqlRequest.newBuilder() - .setSql(statement.getSql()) - .setQueryMode(QueryMode.NORMAL) - .setSession(session.getName()) - .setTransaction(TransactionSelector.newBuilder().setId(transactionId).build()); - Map stmtParameters = statement.getParameters(); - if (!stmtParameters.isEmpty()) { - com.google.protobuf.Struct.Builder paramsBuilder = builder.getParamsBuilder(); - for (Map.Entry param : stmtParameters.entrySet()) { - paramsBuilder.putFields(param.getKey(), param.getValue().toProto()); - builder.putParamTypes(param.getKey(), param.getValue().getType().toProto()); - } - } - return rpc.executePartitionedDml(builder.build(), session.getOptions()); - } - }; - com.google.spanner.v1.ResultSet resultSet = - SpannerRetryHelper.runTxWithRetriesOnAborted( - callable, rpc.getPartitionedDmlRetrySettings()); - if (!resultSet.hasStats()) { - throw new IllegalArgumentException( - "Partitioned DML response missing stats possibly due to non-DML statement as input"); - } - // For partitioned DML, using the row count lower bound. - return resultSet.getStats().getRowCountLowerBound(); - } - - @Override - public void invalidate() { - isValid = false; - } - - // No-op method needed to implement SessionTransaction interface. - @Override - public void setSpan(Span span) {} -} diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java index 33af42f6f9f..bc3f513ce0d 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java @@ -44,9 +44,6 @@ import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import com.google.common.util.concurrent.ThreadFactoryBuilder; -import com.google.spanner.admin.database.v1.CreateBackupRequest; -import com.google.spanner.admin.database.v1.CreateDatabaseRequest; -import com.google.spanner.admin.database.v1.RestoreDatabaseRequest; import com.google.spanner.v1.ExecuteSqlRequest.QueryOptions; import io.grpc.CallCredentials; import io.grpc.CompressorRegistry; diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AbstractBaseConnection.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AbstractBaseConnection.java new file mode 100644 index 00000000000..7fa9cffa03d --- /dev/null +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AbstractBaseConnection.java @@ -0,0 +1,37 @@ +package com.google.cloud.spanner.connection; + +import com.google.cloud.spanner.connection.ConnectionImpl.LeakedConnectionException; +import org.threeten.bp.Instant; + +abstract class AbstractBaseConnection implements BaseConnection { + private static final String CLOSED_ERROR_MSG = "This connection is closed"; + private static final String ONLY_ALLOWED_IN_AUTOCOMMIT = + "This method may only be called while in autocommit mode"; + private static final String NOT_ALLOWED_IN_AUTOCOMMIT = + "This method may not be called while in autocommit mode"; + + /** + * Exception that is used to register the stacktrace of the code that opened a {@link Connection}. + * This exception is logged if the application closes without first closing the connection. + */ + static class LeakedConnectionException extends RuntimeException { + private static final long serialVersionUID = 7119433786832158700L; + + private LeakedConnectionException() { + super("Connection was opened at " + Instant.now()); + } + } + + private volatile LeakedConnectionException leakedException = new LeakedConnectionException(); + private final SpannerPool spannerPool; + private final StatementParser parser = StatementParser.INSTANCE; + /** + * The {@link ConnectionStatementExecutor} is responsible for translating parsed {@link + * ClientSideStatement}s into actual method calls on this {@link ConnectionImpl}. I.e. the {@link + * ClientSideStatement} 'SET AUTOCOMMIT ON' will be translated into the method call {@link + * ConnectionImpl#setAutocommit(boolean)} with value true. + */ + private final ConnectionStatementExecutor connectionStatementExecutor = + new ConnectionStatementExecutorImpl(this); + +} diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AsyncConnection.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AsyncConnection.java new file mode 100644 index 00000000000..d9d97fc6e88 --- /dev/null +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AsyncConnection.java @@ -0,0 +1,636 @@ +/* + * Copyright 2019 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.spanner.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; +import com.google.cloud.spanner.ReadContext.QueryAnalyzeMode; +import com.google.cloud.spanner.ResultSet; +import com.google.cloud.spanner.SpannerBatchUpdateException; +import com.google.cloud.spanner.SpannerException; +import com.google.cloud.spanner.Statement; +import com.google.cloud.spanner.TimestampBound; +import com.google.cloud.spanner.connection.StatementResult.ResultType; +import java.util.Iterator; +import java.util.concurrent.TimeUnit; + +/** + * Internal connection API for Google Cloud Spanner. This interface may introduce breaking changes + * without prior notice. + * + *

    An asynchronous connection to a Cloud Spanner database. Connections are not designed to be thread-safe. + * + *

    Connections accept a number of additional SQL statements for setting or changing the state of + * a {@link AsyncConnection}. These statements can only be executed using the {@link + * AsyncConnection#execute(Statement)} method: + * + *

      + *
    • SHOW AUTOCOMMIT: Returns the current value of AUTOCOMMIT of this + * connection as a {@link ResultSet} + *
    • SET AUTOCOMMIT=TRUE|FALSE: Sets the value of AUTOCOMMIT for this + * connection + *
    • SHOW READONLY: Returns the current value of READONLY of this + * connection as a {@link ResultSet} + *
    • SET READONLY=TRUE|FALSE: Sets the value of READONLY for this + * connection + *
    • SHOW RETRY_ABORTS_INTERNALLY: Returns the current value of + * RETRY_ABORTS_INTERNALLY of this connection as a {@link ResultSet} + *
    • SET RETRY_ABORTS_INTERNALLY=TRUE|FALSE: Sets the value of + * RETRY_ABORTS_INTERNALLY for this connection + *
    • SHOW AUTOCOMMIT_DML_MODE: Returns the current value of + * AUTOCOMMIT_DML_MODE of this connection as a {@link ResultSet} + *
    • SET AUTOCOMMIT_DML_MODE='TRANSACTIONAL' | 'PARTITIONED_NON_ATOMIC': Sets the + * value of AUTOCOMMIT_DML_MODE for this connection + *
    • SHOW STATEMENT_TIMEOUT: Returns the current value of STATEMENT_TIMEOUT + * of this connection as a {@link ResultSet} + *
    • SET STATEMENT_TIMEOUT='<int64>s|ms|us|ns' | NULL: Sets the value of + * STATEMENT_TIMEOUT for this connection. The supported {@link TimeUnit}s are: + *
        + *
      • s - Seconds + *
      • ms - Milliseconds + *
      • us - Microseconds + *
      • ns - Nanoseconds + *
      + * Setting the STATEMENT_TIMEOUT to NULL will clear the value for the STATEMENT_TIMEOUT on the + * connection. + *
    • SHOW READ_TIMESTAMP: Returns the last READ_TIMESTAMP of this + * connection as a {@link ResultSet} + *
    • SHOW COMMIT_TIMESTAMP: Returns the last COMMIT_TIMESTAMP of this + * connection as a {@link ResultSet} + *
    • SHOW READ_ONLY_STALENESS: Returns the current value of + * READ_ONLY_STALENESS of this connection as a {@link ResultSet} + *
    • + * SET READ_ONLY_STALENESS='STRONG' | 'MIN_READ_TIMESTAMP <timestamp>' | 'READ_TIMESTAMP <timestamp>' | 'MAX_STALENESS <int64>s|ms|mus|ns' | 'EXACT_STALENESS (<int64>s|ms|mus|ns)' + * : Sets the value of READ_ONLY_STALENESS for this connection. + *
    • SHOW OPTIMIZER_VERSION: Returns the current value of + * OPTIMIZER_VERSION of this connection as a {@link ResultSet} + *
    • + * SET OPTIMIZER_VERSION='<version>' | 'LATEST' + * : Sets the value of OPTIMIZER_VERSION for this connection. + *
    • BEGIN [TRANSACTION]: Begins a new transaction. This statement is optional when + * the connection is not in autocommit mode, as a new transaction will automatically be + * started when a query or update statement is issued. In autocommit mode, this statement will + * temporarily put the connection in transactional mode, and return the connection to + * autocommit mode when COMMIT [TRANSACTION] or ROLLBACK [TRANSACTION] + * is executed + *
    • COMMIT [TRANSACTION]: Commits the current transaction + *
    • ROLLBACK [TRANSACTION]: Rollbacks the current transaction + *
    • SET TRANSACTION READ ONLY|READ WRITE: Sets the type for the current + * transaction. May only be executed before a transaction is actually running (i.e. before any + * statements have been executed in the transaction) + *
    • START BATCH DDL: Starts a batch of DDL statements. May only be executed when + * no transaction has been started and the connection is in read/write mode. The connection + * will only accept DDL statements while a DDL batch is active. + *
    • START BATCH DML: Starts a batch of DML statements. May only be executed when + * the connection is in read/write mode. The connection will only accept DML statements while + * a DML batch is active. + *
    • RUN BATCH: Ends the current batch, sends the batched DML or DDL statements to + * Spanner and blocks until all statements have been executed or an error occurs. May only be + * executed when a (possibly empty) batch is active. The statement will return the update + * counts of the batched statements as {@link ResultSet} with an ARRAY<INT64> column. In + * case of a DDL batch, this array will always be empty. + *
    • ABORT BATCH: Ends the current batch and removes any DML or DDL statements from + * the buffer without sending any statements to Spanner. May only be executed when a (possibly + * empty) batch is active. + *
    + * + * Note that Cloud Spanner could abort read/write transactions in the background, and that + * any database call during a read/write transaction could fail with an {@link + * AbortedException}. This also includes calls to {@link ResultSet#next()}. + * + *

    If {@link AsyncConnection#isRetryAbortsInternally()} is true, then the connection will + * silently handle any {@link AbortedException}s by internally re-acquiring all transactional locks + * and verifying (via the use of cryptographic checksums) that no underlying data has changed. If a + * change to the underlying data is detected, then an {@link + * AbortedDueToConcurrentModificationException} error will be thrown. If your application already + * uses retry loops to handle these Aborted errors, then it will be most efficient to set {@link + * AsyncConnection#isRetryAbortsInternally()} to false. + * + *

    Use {@link ConnectionOptions} to create a {@link AsyncConnection}. + */ +@InternalApi +public interface AsyncConnection extends AutoCloseable { + /** Closes this connection. This is a no-op if the {@link AsyncConnection} has alread been closed. */ + @Override + void close(); + + /** @return true if this connection has been closed. */ + boolean isClosed(); + + /** + * Sets autocommit on/off for this {@link AsyncConnection}. Connections in autocommit mode will apply + * any changes to the database directly without waiting for an explicit commit. DDL- and DML + * statements as well as {@link Mutation}s are sent directly to Spanner, and committed + * automatically unless the statement caused an error. The statement is retried in case of an + * {@link AbortedException}. All other errors will cause the underlying transaction to be rolled + * back. + * + *

    A {@link AsyncConnection} that is in autocommit and read/write mode will allow all types of + * statements: Queries, DML, DDL, and Mutations (writes). If the connection is in read-only mode, + * only queries will be allowed. + * + *

    {@link AsyncConnection}s in autocommit mode may also accept partitioned DML statements. See + * {@link AsyncConnection#setAutocommitDmlMode(AutocommitDmlMode)} for more information. + * + * @param autocommit true/false to turn autocommit on/off + */ + void setAutocommit(boolean autocommit); + + /** @return true if this connection is in autocommit mode */ + boolean isAutocommit(); + + /** + * Sets this connection to read-only or read-write. This method may only be called when no + * transaction is active. A connection that is in read-only mode, will never allow any kind of + * changes to the database to be submitted. + * + * @param readOnly true/false to turn read-only mode on/off + */ + void setReadOnly(boolean readOnly); + + /** @return true if this connection is in read-only mode */ + boolean isReadOnly(); + + /** + * Begins a new transaction for this connection. + * + *

      + *
    • 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 + * have been executed in the transaction. + * + * @param transactionMode The transaction mode to use for the current transaction. + *
      + *
    • {@link TransactionMode#READ_ONLY_TRANSACTION} will create a read-only transaction and + * prevent any changes to written to the database through this transaction. The read + * timestamp to be used will be determined based on the current readOnlyStaleness + * setting of this connection. It is recommended to use {@link + * TransactionMode#READ_ONLY_TRANSACTION} instead of {@link + * TransactionMode#READ_WRITE_TRANSACTION} when possible, as read-only transactions do + * not acquire locks on Cloud Spanner, and read-only transactions never abort. + *
    • {@link TransactionMode#READ_WRITE_TRANSACTION} this value is only allowed when the + * connection is not in read-only mode and will create a read-write transaction. If + * {@link AsyncConnection#isRetryAbortsInternally()} is true, each read/write + * transaction will keep track of a running SHA256 checksum for each {@link ResultSet} + * that is returned in order to be able to retry the transaction in case the transaction + * is aborted by Spanner. + *
    + */ + void setTransactionMode(TransactionMode transactionMode); + + /** + * @return the transaction mode of the current transaction. This method may only be called when + * the connection is in a transaction. + */ + TransactionMode getTransactionMode(); + + /** + * @return true if this connection will automatically retry read/write transactions + * that abort. This method may only be called when the connection is in read/write + * transactional mode and no transaction has been started yet. + */ + boolean isRetryAbortsInternally(); + + /** + * Sets whether this connection will internally retry read/write transactions that abort. The + * default is true. When internal retry is enabled, the {@link AsyncConnection} will keep + * track of a running SHA256 checksum of all {@link ResultSet}s that have been returned from Cloud + * Spanner. If the checksum that is calculated during an internal retry differs from the original + * checksum, the transaction will abort with an {@link + * AbortedDueToConcurrentModificationException}. + * + *

    Note that retries of a read/write transaction that calls a non-deterministic function on + * Cloud Spanner, such as CURRENT_TIMESTAMP(), will never be successful, as the data returned + * during the retry will always be different from the original transaction. + * + *

    It is also highly recommended that all queries in a read/write transaction have an ORDER BY + * clause that guarantees that the data is returned in the same order as in the original + * transaction if the transaction is internally retried. The most efficient way to achieve this is + * to always include the primary key columns at the end of the ORDER BY clause. + * + *

    This method may only be called when the connection is in read/write transactional mode and + * no transaction has been started yet. + * + * @param retryAbortsInternally Set to true to internally retry transactions that are + * aborted by Spanner. When set to false, any database call on a transaction that + * has been aborted by Cloud Spanner will throw an {@link AbortedException} instead of being + * retried. Set this to false if your application already uses retry loops to handle {@link + * AbortedException}s. + */ + void setRetryAbortsInternally(boolean retryAbortsInternally); + + /** + * Add a {@link TransactionRetryListener} to this {@link AsyncConnection} for testing and logging + * purposes. The method {@link TransactionRetryListener#retryStarting(Timestamp, long, int)} will + * be called before an automatic retry is started for a read/write transaction on this connection. + * The method {@link TransactionRetryListener#retryFinished(Timestamp, long, int, + * TransactionRetryListener.RetryResult)} will be called after the retry has finished. + * + * @param listener The listener to add to this connection. + */ + void addTransactionRetryListener(TransactionRetryListener listener); + + /** + * Removes one existing {@link TransactionRetryListener} from this {@link AsyncConnection}, if it is + * present (optional operation). + * + * @param listener The listener to remove from the connection. + * @return true if a listener was removed from the connection. + */ + boolean removeTransactionRetryListener(TransactionRetryListener listener); + + /** + * @return an unmodifiable iterator of the {@link TransactionRetryListener}s registered for this + * connection. + */ + Iterator getTransactionRetryListeners(); + + /** + * Sets the mode for executing DML statements in autocommit mode for this connection. This setting + * is only used when the connection is in autocommit mode, and may only be set while the + * transaction is in autocommit mode and not in a temporary transaction. The autocommit + * transaction mode is reset to its default value of {@link AutocommitDmlMode#TRANSACTIONAL} when + * autocommit mode is changed on the connection. + * + * @param mode The DML autocommit mode to use + *

      + *
    • {@link AutocommitDmlMode#TRANSACTIONAL} DML statements are executed as single + * read-write transaction. After successful execution, the DML statement is guaranteed + * to have been applied exactly once to the database + *
    • {@link AutocommitDmlMode#PARTITIONED_NON_ATOMIC} DML statements are executed as + * partitioned DML transactions. If an error occurs during the execution of the DML + * statement, it is possible that the statement has been applied to some but not all of + * the rows specified in the statement. + *
    + */ + void setAutocommitDmlMode(AutocommitDmlMode mode); + + /** + * @return the current {@link AutocommitDmlMode} setting for this connection. This method may only + * be called on a connection that is in autocommit mode and not while in a temporary + * transaction. + */ + AutocommitDmlMode getAutocommitDmlMode(); + + /** + * Sets the staleness to use for the current read-only transaction. This method may only be called + * when the transaction mode of the current transaction is {@link + * TransactionMode#READ_ONLY_TRANSACTION} and there is no transaction that has started, or when + * the connection is in read-only and autocommit mode. + * + * @param staleness The staleness to use for the current but not yet started read-only transaction + */ + void setReadOnlyStaleness(TimestampBound staleness); + + /** + * @return the read-only staleness setting for the current read-only transaction. This method may + * only be called when the current transaction is a read-only transaction, or when the + * connection is in read-only and autocommit mode. + */ + TimestampBound getReadOnlyStaleness(); + + /** + * Sets the query optimizer version to use for this connection. + * + * @param optimizerVersion The query optimizer version to use. Must be a valid optimizer version + * number, the string LATEST or an empty string. The empty string will instruct + * the connection to use the optimizer version that is defined in the environment variable + * SPANNER_OPTIMIZER_VERSION. If no value is specified in the environment + * variable, the default query optimizer of Cloud Spanner is used. + */ + void setOptimizerVersion(String optimizerVersion); + + /** + * Gets the current query optimizer version of this connection. + * + * @return The query optimizer version that is currently used by this connection. + */ + String getOptimizerVersion(); + + /** + * Commits the current transaction of this connection. All mutations that have been buffered + * during the current transaction will be written to the database. + * + *

    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 AsyncConnection#getReadTimestamp()} or {@link AsyncConnection#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. + * + *

    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 AsyncConnection} is in autocommit + * mode and no explicit transaction has been started by calling {@link + * AsyncConnection#beginTransaction()}. If the {@link AsyncConnection} is not in autocommit mode, there + * will always be a transaction. + */ + boolean isInTransaction(); + + /** + * @return true if this connection has a transaction that has started. A transaction + * is automatically started by the first statement that is executed in the transaction. + */ + boolean isTransactionStarted(); + + /** + * Returns the read timestamp of the current/last {@link TransactionMode#READ_ONLY_TRANSACTION} + * transaction, or the read timestamp of the last query in autocommit mode. + * + *
      + *
    • When in autocommit mode: The method will return the read timestamp of the last statement + * if the last statement was a query. + *
    • When in a {@link TransactionMode#READ_ONLY_TRANSACTION} transaction that has started (a + * query has been executed), or that has just committed: The read timestamp of the + * transaction. If the read-only transaction was committed without ever executing a query, + * calling this method after the commit will also throw a {@link SpannerException} + *
    • In all other cases the method will throw a {@link SpannerException}. + *
    + * + * @return the read timestamp of the current/last read-only transaction. + */ + Timestamp getReadTimestamp(); + + /** + * @return the commit timestamp of the last {@link TransactionMode#READ_WRITE_TRANSACTION} + * transaction. This method will throw a {@link SpannerException} if there is no last {@link + * TransactionMode#READ_WRITE_TRANSACTION} transaction (i.e. the last transaction was a {@link + * TransactionMode#READ_ONLY_TRANSACTION}), or if the last {@link + * TransactionMode#READ_WRITE_TRANSACTION} transaction rolled back. It will also throw a + * {@link SpannerException} if the last {@link TransactionMode#READ_WRITE_TRANSACTION} + * transaction was empty when committed. + */ + Timestamp getCommitTimestamp(); + + /** + * Starts a new DDL batch on this connection. A DDL batch allows several DDL statements to be + * grouped into a batch that can be executed as a group. DDL statements that are issued during the + * batch are buffered locally and will return immediately with an OK. It is not guaranteed that a + * DDL statement that has been issued during a batch will eventually succeed when running the + * batch. Aborting a DDL batch will clear the DDL buffer and will have made no changes to the + * database. Running a DDL batch will send all buffered DDL statements to Spanner, and Spanner + * will try to execute these. The result will be OK if all the statements executed successfully. + * If a statement cannot be executed, Spanner will stop execution at that point and return an + * error message for the statement that could not be executed. Preceding statements of the batch + * may have been executed. + * + *

    This method may only be called when the connection is in read/write mode, autocommit mode is + * enabled or no read/write transaction has been started, and there is not already another batch + * active. The connection will only accept DDL statements while a DDL batch is active. + */ + void startBatchDdl(); + + /** + * Starts a new DML batch on this connection. A DML batch allows several DML statements to be + * grouped into a batch that can be executed as a group. DML statements that are issued during the + * batch are buffered locally and will return immediately with an OK. It is not guaranteed that a + * DML statement that has been issued during a batch will eventually succeed when running the + * batch. Aborting a DML batch will clear the DML buffer and will have made no changes to the + * database. Running a DML batch will send all buffered DML statements to Spanner, and Spanner + * will try to execute these. The result will be OK if all the statements executed successfully. + * If a statement cannot be executed, Spanner will stop execution at that point and return {@link + * SpannerBatchUpdateException} for the statement that could not be executed. Preceding statements + * of the batch will have been executed, and the update counts of those statements can be + * retrieved through {@link SpannerBatchUpdateException#getUpdateCounts()}. + * + *

    This method may only be called when the connection is in read/write mode, autocommit mode is + * enabled or no read/write transaction has been started, and there is not already another batch + * active. The connection will only accept DML statements while a DML batch is active. + */ + void startBatchDml(); + + /** + * 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 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. + */ + long[] runBatch(); + + /** + * Clears all buffered statements in the current batch and ends the batch. + * + *

    This method may only be called when a (possibly empty) batch is active. + */ + void abortBatch(); + + /** @return true if a DDL batch is active on this connection. */ + boolean isDdlBatchActive(); + + /** @return true if a DML batch is active on this connection. */ + boolean isDmlBatchActive(); + + /** + * Executes the given statement if allowed in the current {@link TransactionMode} and connection + * state. The returned value depends on the type of statement: + * + *

      + *
    • Queries will return a {@link ResultSet} + *
    • DML statements will return an update count + *
    • DDL statements will return a {@link ResultType#NO_RESULT} + *
    • 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) + *
    + * + * @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 + * query, the method will throw a {@link SpannerException}. + * + * @param query The query statement to execute + * @param options the options to configure the query + * @return a {@link ResultSet} 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. + * + *

    The query plan and query statistics information is contained in {@link + * com.google.spanner.v1.ResultSetStats} that can be accessed by calling {@link + * ResultSet#getStats()} on the returned {@code ResultSet}. + * + * 

    +   * 
+   * {@code
+   * ResultSet resultSet =
+   *     connection.analyzeQuery(
+   *         Statement.of("SELECT SingerId, AlbumId, MarketingBudget FROM Albums"),
+   *         ReadContext.QueryAnalyzeMode.PROFILE);
+   * while (resultSet.next()) {
+   *   // Discard the results. We're only processing because getStats() below requires it.
+   * }
+   * ResultSetStats stats = resultSet.getStats();
+   * }
+   * 
+   *
    + * + * @param query the query statement to execute + * @param queryMode the mode in which to execute the query + */ + AsyncResultSet analyzeQueryAsync(Statement query, QueryAnalyzeMode queryMode); + + /** + * Executes the given statement as a DML statement. If the statement does not contain a valid DML + * statement, the method will throw a {@link SpannerException}. + * + * @param update The update statement to execute + * @return 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 + * AsyncConnection#executeUpdate(Statement)} in a loop. This method returns an array of long integers, + * each representing the number of rows modified by each statement. + * + *

    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 array containing the update counts per statement. + */ + ApiFuture executeBatchUpdateAsync(Iterable updates); + + /** + * Buffers the given mutation locally on the current transaction of this {@link AsyncConnection}. The + * mutation will be written to the database at the next call to {@link AsyncConnection#commit()}. The + * value will not be readable on this {@link AsyncConnection} before the transaction is committed. + * + *

    Calling this method is only allowed when not in autocommit mode. See {@link + * AsyncConnection#write(Mutation)} for writing mutations in autocommit mode. + * + * @param mutation the {@link Mutation} to buffer for writing to the database on the next commit + * @throws SpannerException if the {@link AsyncConnection} is in autocommit mode + */ + void bufferedWrite(Mutation mutation); + + /** + * Buffers the given mutations locally on the current transaction of this {@link AsyncConnection}. The + * mutations will be written to the database at the next call to {@link AsyncConnection#commit()}. The + * values will not be readable on this {@link AsyncConnection} before the transaction is committed. + * + *

    Calling this method is only allowed when not in autocommit mode. See {@link + * AsyncConnection#write(Iterable)} for writing mutations in autocommit mode. + * + * @param mutations the {@link Mutation}s to buffer for writing to the database on the next commit + * @throws SpannerException if the {@link AsyncConnection} is in autocommit mode + */ + void bufferedWrite(Iterable mutations); +} diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AsyncConnectionImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AsyncConnectionImpl.java new file mode 100644 index 00000000000..3f94079f4d6 --- /dev/null +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AsyncConnectionImpl.java @@ -0,0 +1,247 @@ +package com.google.cloud.spanner.connection; + +import com.google.api.core.ApiFuture; +import com.google.cloud.Timestamp; +import com.google.cloud.spanner.AsyncResultSet; +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.Statement; +import com.google.cloud.spanner.TimestampBound; +import java.util.Iterator; + +public class AsyncConnectionImpl implements AsyncConnection { + + @Override + public void close() { + // TODO Auto-generated method stub + + } + + @Override + public boolean isClosed() { + // TODO Auto-generated method stub + return false; + } + + @Override + public void setAutocommit(boolean autocommit) { + // TODO Auto-generated method stub + + } + + @Override + public boolean isAutocommit() { + // TODO Auto-generated method stub + return false; + } + + @Override + public void setReadOnly(boolean readOnly) { + // TODO Auto-generated method stub + + } + + @Override + public boolean isReadOnly() { + // TODO Auto-generated method stub + return false; + } + + @Override + public ApiFuture beginTransactionAsync() { + // TODO Auto-generated method stub + return null; + } + + @Override + public void setTransactionMode(TransactionMode transactionMode) { + // TODO Auto-generated method stub + + } + + @Override + public TransactionMode getTransactionMode() { + // TODO Auto-generated method stub + return null; + } + + @Override + public boolean isRetryAbortsInternally() { + // TODO Auto-generated method stub + return false; + } + + @Override + public void setRetryAbortsInternally(boolean retryAbortsInternally) { + // TODO Auto-generated method stub + + } + + @Override + public void addTransactionRetryListener(TransactionRetryListener listener) { + // TODO Auto-generated method stub + + } + + @Override + public boolean removeTransactionRetryListener(TransactionRetryListener listener) { + // TODO Auto-generated method stub + return false; + } + + @Override + public Iterator getTransactionRetryListeners() { + // TODO Auto-generated method stub + return null; + } + + @Override + public void setAutocommitDmlMode(AutocommitDmlMode mode) { + // TODO Auto-generated method stub + + } + + @Override + public AutocommitDmlMode getAutocommitDmlMode() { + // TODO Auto-generated method stub + return null; + } + + @Override + public void setReadOnlyStaleness(TimestampBound staleness) { + // TODO Auto-generated method stub + + } + + @Override + public TimestampBound getReadOnlyStaleness() { + // TODO Auto-generated method stub + return null; + } + + @Override + public void setOptimizerVersion(String optimizerVersion) { + // TODO Auto-generated method stub + + } + + @Override + public String getOptimizerVersion() { + // TODO Auto-generated method stub + return null; + } + + @Override + public ApiFuture commitAsync() { + // TODO Auto-generated method stub + return null; + } + + @Override + public ApiFuture rollbackAsync() { + // TODO Auto-generated method stub + return null; + } + + @Override + public boolean isInTransaction() { + // TODO Auto-generated method stub + return false; + } + + @Override + public boolean isTransactionStarted() { + // TODO Auto-generated method stub + return false; + } + + @Override + public Timestamp getReadTimestamp() { + // TODO Auto-generated method stub + return null; + } + + @Override + public Timestamp getCommitTimestamp() { + // TODO Auto-generated method stub + return null; + } + + @Override + public void startBatchDdl() { + // TODO Auto-generated method stub + + } + + @Override + public void startBatchDml() { + // TODO Auto-generated method stub + + } + + @Override + public long[] runBatch() { + // TODO Auto-generated method stub + return null; + } + + @Override + public void abortBatch() { + // TODO Auto-generated method stub + + } + + @Override + public boolean isDdlBatchActive() { + // TODO Auto-generated method stub + return false; + } + + @Override + public boolean isDmlBatchActive() { + // TODO Auto-generated method stub + return false; + } + + @Override + public AsyncStatementResult executeAsync(Statement statement) { + // TODO Auto-generated method stub + return null; + } + + @Override + public AsyncResultSet executeQueryAsync(Statement query, QueryOption... options) { + // TODO Auto-generated method stub + return null; + } + + @Override + public AsyncResultSet analyzeQueryAsync(Statement query, QueryAnalyzeMode queryMode) { + // TODO Auto-generated method stub + return null; + } + + @Override + public ApiFuture executeUpdateAsync(Statement update) { + // TODO Auto-generated method stub + return null; + } + + @Override + public ApiFuture executeBatchUpdateAsync(Iterable updates) { + // TODO Auto-generated method stub + return null; + } + + @Override + public void bufferedWrite(Mutation mutation) { + // TODO Auto-generated method stub + + } + + @Override + public void bufferedWrite(Iterable mutations) { + // TODO Auto-generated method stub + + }} 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 00000000000..53de7535632 --- /dev/null +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AsyncStatementResult.java @@ -0,0 +1,47 @@ +/* + * Copyright 2019 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.spanner.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.ClientSideStatementType; +import com.google.cloud.spanner.connection.StatementResult.ResultType; + +/** + * A result of the execution of a statement. Statements that are executed by the {@link + * Connection#execute(com.google.cloud.spanner.Statement)} method could have different types of + * return values. These are wrapped in a {@link AsyncStatementResult}. + */ +@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 AsyncResultType#RESULT_SET}. + * + * @return the {@link AsyncResultSet} held by this result. + */ + AsyncResultSet getResultSet(); + + /** + * 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(); +} diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/BaseConnection.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/BaseConnection.java new file mode 100644 index 00000000000..c740f34e118 --- /dev/null +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/BaseConnection.java @@ -0,0 +1,343 @@ +package com.google.cloud.spanner.connection; + +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.Mutation; +import com.google.cloud.spanner.ResultSet; +import com.google.cloud.spanner.SpannerBatchUpdateException; +import com.google.cloud.spanner.SpannerException; +import com.google.cloud.spanner.TimestampBound; +import com.google.cloud.spanner.Options.QueryOption; +import java.util.Iterator; + +public interface BaseConnection extends AutoCloseable { + /** Closes this connection. This is a no-op if the {@link Connection} has alread been closed. */ + @Override + void close(); + + /** @return true if this connection has been closed. */ + boolean isClosed(); + + /** + * Sets autocommit on/off for this {@link Connection}. Connections in autocommit mode will apply + * any changes to the database directly without waiting for an explicit commit. DDL- and DML + * statements as well as {@link Mutation}s are sent directly to Spanner, and committed + * automatically unless the statement caused an error. The statement is retried in case of an + * {@link AbortedException}. All other errors will cause the underlying transaction to be rolled + * back. + * + *

    A {@link Connection} that is in autocommit and read/write mode will allow all types of + * statements: Queries, DML, DDL, and Mutations (writes). If the connection is in read-only mode, + * only queries will be allowed. + * + *

    {@link Connection}s in autocommit mode may also accept partitioned DML statements. See + * {@link Connection#setAutocommitDmlMode(AutocommitDmlMode)} for more information. + * + * @param autocommit true/false to turn autocommit on/off + */ + void setAutocommit(boolean autocommit); + + /** @return true if this connection is in autocommit mode */ + boolean isAutocommit(); + + /** + * Sets this connection to read-only or read-write. This method may only be called when no + * transaction is active. A connection that is in read-only mode, will never allow any kind of + * changes to the database to be submitted. + * + * @param readOnly true/false to turn read-only mode on/off + */ + void setReadOnly(boolean readOnly); + + /** @return true if this connection is in read-only mode */ + boolean isReadOnly(); + + /** + * 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 + * have been executed in the transaction. + * + * @param transactionMode The transaction mode to use for the current transaction. + *

      + *
    • {@link TransactionMode#READ_ONLY_TRANSACTION} will create a read-only transaction and + * prevent any changes to written to the database through this transaction. The read + * timestamp to be used will be determined based on the current readOnlyStaleness + * setting of this connection. It is recommended to use {@link + * TransactionMode#READ_ONLY_TRANSACTION} instead of {@link + * TransactionMode#READ_WRITE_TRANSACTION} when possible, as read-only transactions do + * not acquire locks on Cloud Spanner, and read-only transactions never abort. + *
    • {@link TransactionMode#READ_WRITE_TRANSACTION} this value is only allowed when the + * connection is not in read-only mode and will create a read-write transaction. If + * {@link Connection#isRetryAbortsInternally()} is true, each read/write + * transaction will keep track of a running SHA256 checksum for each {@link ResultSet} + * that is returned in order to be able to retry the transaction in case the transaction + * is aborted by Spanner. + *
    + */ + void setTransactionMode(TransactionMode transactionMode); + + /** + * @return the transaction mode of the current transaction. This method may only be called when + * the connection is in a transaction. + */ + TransactionMode getTransactionMode(); + + /** + * @return true if this connection will automatically retry read/write transactions + * that abort. This method may only be called when the connection is in read/write + * transactional mode and no transaction has been started yet. + */ + boolean isRetryAbortsInternally(); + + /** + * Sets whether this connection will internally retry read/write transactions that abort. The + * default is true. When internal retry is enabled, the {@link Connection} will keep + * track of a running SHA256 checksum of all {@link ResultSet}s that have been returned from Cloud + * Spanner. If the checksum that is calculated during an internal retry differs from the original + * checksum, the transaction will abort with an {@link + * AbortedDueToConcurrentModificationException}. + * + *

    Note that retries of a read/write transaction that calls a non-deterministic function on + * Cloud Spanner, such as CURRENT_TIMESTAMP(), will never be successful, as the data returned + * during the retry will always be different from the original transaction. + * + *

    It is also highly recommended that all queries in a read/write transaction have an ORDER BY + * clause that guarantees that the data is returned in the same order as in the original + * transaction if the transaction is internally retried. The most efficient way to achieve this is + * to always include the primary key columns at the end of the ORDER BY clause. + * + *

    This method may only be called when the connection is in read/write transactional mode and + * no transaction has been started yet. + * + * @param retryAbortsInternally Set to true to internally retry transactions that are + * aborted by Spanner. When set to false, any database call on a transaction that + * has been aborted by Cloud Spanner will throw an {@link AbortedException} instead of being + * retried. Set this to false if your application already uses retry loops to handle {@link + * AbortedException}s. + */ + void setRetryAbortsInternally(boolean retryAbortsInternally); + + /** + * Add a {@link TransactionRetryListener} to this {@link Connection} for testing and logging + * purposes. The method {@link TransactionRetryListener#retryStarting(Timestamp, long, int)} will + * be called before an automatic retry is started for a read/write transaction on this connection. + * The method {@link TransactionRetryListener#retryFinished(Timestamp, long, int, + * TransactionRetryListener.RetryResult)} will be called after the retry has finished. + * + * @param listener The listener to add to this connection. + */ + void addTransactionRetryListener(TransactionRetryListener listener); + + /** + * Removes one existing {@link TransactionRetryListener} from this {@link Connection}, if it is + * present (optional operation). + * + * @param listener The listener to remove from the connection. + * @return true if a listener was removed from the connection. + */ + boolean removeTransactionRetryListener(TransactionRetryListener listener); + + /** + * @return an unmodifiable iterator of the {@link TransactionRetryListener}s registered for this + * connection. + */ + Iterator getTransactionRetryListeners(); + + /** + * Sets the mode for executing DML statements in autocommit mode for this connection. This setting + * is only used when the connection is in autocommit mode, and may only be set while the + * transaction is in autocommit mode and not in a temporary transaction. The autocommit + * transaction mode is reset to its default value of {@link AutocommitDmlMode#TRANSACTIONAL} when + * autocommit mode is changed on the connection. + * + * @param mode The DML autocommit mode to use + *

      + *
    • {@link AutocommitDmlMode#TRANSACTIONAL} DML statements are executed as single + * read-write transaction. After successful execution, the DML statement is guaranteed + * to have been applied exactly once to the database + *
    • {@link AutocommitDmlMode#PARTITIONED_NON_ATOMIC} DML statements are executed as + * partitioned DML transactions. If an error occurs during the execution of the DML + * statement, it is possible that the statement has been applied to some but not all of + * the rows specified in the statement. + *
    + */ + void setAutocommitDmlMode(AutocommitDmlMode mode); + + /** + * @return the current {@link AutocommitDmlMode} setting for this connection. This method may only + * be called on a connection that is in autocommit mode and not while in a temporary + * transaction. + */ + AutocommitDmlMode getAutocommitDmlMode(); + + /** + * Sets the staleness to use for the current read-only transaction. This method may only be called + * when the transaction mode of the current transaction is {@link + * TransactionMode#READ_ONLY_TRANSACTION} and there is no transaction that has started, or when + * the connection is in read-only and autocommit mode. + * + * @param staleness The staleness to use for the current but not yet started read-only transaction + */ + void setReadOnlyStaleness(TimestampBound staleness); + + /** + * @return the read-only staleness setting for the current read-only transaction. This method may + * only be called when the current transaction is a read-only transaction, or when the + * connection is in read-only and autocommit mode. + */ + TimestampBound getReadOnlyStaleness(); + + /** + * Sets the query optimizer version to use for this connection. + * + * @param optimizerVersion The query optimizer version to use. Must be a valid optimizer version + * number, the string LATEST or an empty string. The empty string will instruct + * the connection to use the optimizer version that is defined in the environment variable + * SPANNER_OPTIMIZER_VERSION. If no value is specified in the environment + * variable, the default query optimizer of Cloud Spanner is used. + */ + void setOptimizerVersion(String optimizerVersion); + + /** + * Gets the current query optimizer version of this connection. + * + * @return The query optimizer version that is currently used by this connection. + */ + String getOptimizerVersion(); + + /** + * @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 + * mode and no explicit transaction has been started by calling {@link + * Connection#beginTransaction()}. If the {@link Connection} is not in autocommit mode, there + * will always be a transaction. + */ + boolean isInTransaction(); + + /** + * @return true if this connection has a transaction that has started. A transaction + * is automatically started by the first statement that is executed in the transaction. + */ + boolean isTransactionStarted(); + + /** + * Returns the read timestamp of the current/last {@link TransactionMode#READ_ONLY_TRANSACTION} + * transaction, or the read timestamp of the last query in autocommit mode. + * + *
      + *
    • When in autocommit mode: The method will return the read timestamp of the last statement + * if the last statement was a query. + *
    • When in a {@link TransactionMode#READ_ONLY_TRANSACTION} transaction that has started (a + * query has been executed), or that has just committed: The read timestamp of the + * transaction. If the read-only transaction was committed without ever executing a query, + * calling this method after the commit will also throw a {@link SpannerException} + *
    • In all other cases the method will throw a {@link SpannerException}. + *
    + * + * @return the read timestamp of the current/last read-only transaction. + */ + Timestamp getReadTimestamp(); + + /** + * @return the commit timestamp of the last {@link TransactionMode#READ_WRITE_TRANSACTION} + * transaction. This method will throw a {@link SpannerException} if there is no last {@link + * TransactionMode#READ_WRITE_TRANSACTION} transaction (i.e. the last transaction was a {@link + * TransactionMode#READ_ONLY_TRANSACTION}), or if the last {@link + * TransactionMode#READ_WRITE_TRANSACTION} transaction rolled back. It will also throw a + * {@link SpannerException} if the last {@link TransactionMode#READ_WRITE_TRANSACTION} + * transaction was empty when committed. + */ + Timestamp getCommitTimestamp(); + + /** + * Starts a new DDL batch on this connection. A DDL batch allows several DDL statements to be + * grouped into a batch that can be executed as a group. DDL statements that are issued during the + * batch are buffered locally and will return immediately with an OK. It is not guaranteed that a + * DDL statement that has been issued during a batch will eventually succeed when running the + * batch. Aborting a DDL batch will clear the DDL buffer and will have made no changes to the + * database. Running a DDL batch will send all buffered DDL statements to Spanner, and Spanner + * will try to execute these. The result will be OK if all the statements executed successfully. + * If a statement cannot be executed, Spanner will stop execution at that point and return an + * error message for the statement that could not be executed. Preceding statements of the batch + * may have been executed. + * + *

    This method may only be called when the connection is in read/write mode, autocommit mode is + * enabled or no read/write transaction has been started, and there is not already another batch + * active. The connection will only accept DDL statements while a DDL batch is active. + */ + void startBatchDdl(); + + /** + * Starts a new DML batch on this connection. A DML batch allows several DML statements to be + * grouped into a batch that can be executed as a group. DML statements that are issued during the + * batch are buffered locally and will return immediately with an OK. It is not guaranteed that a + * DML statement that has been issued during a batch will eventually succeed when running the + * batch. Aborting a DML batch will clear the DML buffer and will have made no changes to the + * database. Running a DML batch will send all buffered DML statements to Spanner, and Spanner + * will try to execute these. The result will be OK if all the statements executed successfully. + * If a statement cannot be executed, Spanner will stop execution at that point and return {@link + * SpannerBatchUpdateException} for the statement that could not be executed. Preceding statements + * of the batch will have been executed, and the update counts of those statements can be + * retrieved through {@link SpannerBatchUpdateException#getUpdateCounts()}. + * + *

    This method may only be called when the connection is in read/write mode, autocommit mode is + * enabled or no read/write transaction has been started, and there is not already another batch + * active. The connection will only accept DML statements while a DML batch is active. + */ + void startBatchDml(); + + /** + * Clears all buffered statements in the current batch and ends the batch. + * + *

    This method may only be called when a (possibly empty) batch is active. + */ + void abortBatch(); + + /** @return true if a DDL batch is active on this connection. */ + boolean isDdlBatchActive(); + + /** @return true if a DML batch is active on this connection. */ + boolean isDmlBatchActive(); + + /** + * 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 + * value will not be readable on this {@link Connection} before the transaction is committed. + * + *

    Calling this method is only allowed when not in autocommit mode. See {@link + * Connection#write(Mutation)} for writing mutations in autocommit mode. + * + * @param mutation the {@link Mutation} to buffer for writing to the database on the next commit + * @throws SpannerException if the {@link Connection} is in autocommit mode + */ + void bufferedWrite(Mutation mutation); + + /** + * Buffers the given mutations locally on the current transaction of this {@link Connection}. The + * mutations will be written to the database at the next call to {@link Connection#commit()}. The + * values will not be readable on this {@link Connection} before the transaction is committed. + * + *

    Calling this method is only allowed when not in autocommit mode. See {@link + * Connection#write(Iterable)} for writing mutations in autocommit mode. + * + * @param mutations the {@link Mutation}s to buffer for writing to the database on the next commit + * @throws SpannerException if the {@link Connection} is in autocommit mode + */ + void bufferedWrite(Iterable mutations); + + /** + * This query option is used internally to indicate that a query is executed by the library itself + * to fetch metadata. These queries are specifically allowed to be executed even when a DDL batch + * is active. + * + *

    NOT INTENDED FOR EXTERNAL USE! + */ + @InternalApi + public static final class InternalMetadataQuery implements QueryOption { + @InternalApi public static final InternalMetadataQuery INSTANCE = new InternalMetadataQuery(); + + private InternalMetadataQuery() {} + } +} 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 9a1dc69a0cd..1bd425c2ecc 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 @@ -17,21 +17,16 @@ package com.google.cloud.spanner.connection; 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; import com.google.cloud.spanner.ReadContext.QueryAnalyzeMode; import com.google.cloud.spanner.ResultSet; -import com.google.cloud.spanner.SpannerBatchUpdateException; import com.google.cloud.spanner.SpannerException; import com.google.cloud.spanner.Statement; -import com.google.cloud.spanner.TimestampBound; import com.google.cloud.spanner.connection.StatementResult.ResultType; -import java.util.Iterator; import java.util.concurrent.TimeUnit; /** @@ -131,48 +126,7 @@ *

    Use {@link ConnectionOptions} to create a {@link Connection}. */ @InternalApi -public interface Connection extends AutoCloseable { - /** Closes this connection. This is a no-op if the {@link Connection} has alread been closed. */ - @Override - void close(); - - /** @return true if this connection has been closed. */ - boolean isClosed(); - - /** - * Sets autocommit on/off for this {@link Connection}. Connections in autocommit mode will apply - * any changes to the database directly without waiting for an explicit commit. DDL- and DML - * statements as well as {@link Mutation}s are sent directly to Spanner, and committed - * automatically unless the statement caused an error. The statement is retried in case of an - * {@link AbortedException}. All other errors will cause the underlying transaction to be rolled - * back. - * - *

    A {@link Connection} that is in autocommit and read/write mode will allow all types of - * statements: Queries, DML, DDL, and Mutations (writes). If the connection is in read-only mode, - * only queries will be allowed. - * - *

    {@link Connection}s in autocommit mode may also accept partitioned DML statements. See - * {@link Connection#setAutocommitDmlMode(AutocommitDmlMode)} for more information. - * - * @param autocommit true/false to turn autocommit on/off - */ - void setAutocommit(boolean autocommit); - - /** @return true if this connection is in autocommit mode */ - boolean isAutocommit(); - - /** - * Sets this connection to read-only or read-write. This method may only be called when no - * transaction is active. A connection that is in read-only mode, will never allow any kind of - * changes to the database to be submitted. - * - * @param readOnly true/false to turn read-only mode on/off - */ - void setReadOnly(boolean readOnly); - - /** @return true if this connection is in read-only mode */ - boolean isReadOnly(); - +public interface Connection extends BaseConnection { /** * Sets the duration the connection should wait before automatically aborting the execution of a * statement. The default is no timeout. Statement timeouts are applied all types of statements, @@ -260,159 +214,6 @@ public interface Connection extends AutoCloseable { */ void beginTransaction(); - /** - * 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 - * have been executed in the transaction. - * - * @param transactionMode The transaction mode to use for the current transaction. - *

      - *
    • {@link TransactionMode#READ_ONLY_TRANSACTION} will create a read-only transaction and - * prevent any changes to written to the database through this transaction. The read - * timestamp to be used will be determined based on the current readOnlyStaleness - * setting of this connection. It is recommended to use {@link - * TransactionMode#READ_ONLY_TRANSACTION} instead of {@link - * TransactionMode#READ_WRITE_TRANSACTION} when possible, as read-only transactions do - * not acquire locks on Cloud Spanner, and read-only transactions never abort. - *
    • {@link TransactionMode#READ_WRITE_TRANSACTION} this value is only allowed when the - * connection is not in read-only mode and will create a read-write transaction. If - * {@link Connection#isRetryAbortsInternally()} is true, each read/write - * transaction will keep track of a running SHA256 checksum for each {@link ResultSet} - * that is returned in order to be able to retry the transaction in case the transaction - * is aborted by Spanner. - *
    - */ - void setTransactionMode(TransactionMode transactionMode); - - /** - * @return the transaction mode of the current transaction. This method may only be called when - * the connection is in a transaction. - */ - TransactionMode getTransactionMode(); - - /** - * @return true if this connection will automatically retry read/write transactions - * that abort. This method may only be called when the connection is in read/write - * transactional mode and no transaction has been started yet. - */ - boolean isRetryAbortsInternally(); - - /** - * Sets whether this connection will internally retry read/write transactions that abort. The - * default is true. When internal retry is enabled, the {@link Connection} will keep - * track of a running SHA256 checksum of all {@link ResultSet}s that have been returned from Cloud - * Spanner. If the checksum that is calculated during an internal retry differs from the original - * checksum, the transaction will abort with an {@link - * AbortedDueToConcurrentModificationException}. - * - *

    Note that retries of a read/write transaction that calls a non-deterministic function on - * Cloud Spanner, such as CURRENT_TIMESTAMP(), will never be successful, as the data returned - * during the retry will always be different from the original transaction. - * - *

    It is also highly recommended that all queries in a read/write transaction have an ORDER BY - * clause that guarantees that the data is returned in the same order as in the original - * transaction if the transaction is internally retried. The most efficient way to achieve this is - * to always include the primary key columns at the end of the ORDER BY clause. - * - *

    This method may only be called when the connection is in read/write transactional mode and - * no transaction has been started yet. - * - * @param retryAbortsInternally Set to true to internally retry transactions that are - * aborted by Spanner. When set to false, any database call on a transaction that - * has been aborted by Cloud Spanner will throw an {@link AbortedException} instead of being - * retried. Set this to false if your application already uses retry loops to handle {@link - * AbortedException}s. - */ - void setRetryAbortsInternally(boolean retryAbortsInternally); - - /** - * Add a {@link TransactionRetryListener} to this {@link Connection} for testing and logging - * purposes. The method {@link TransactionRetryListener#retryStarting(Timestamp, long, int)} will - * be called before an automatic retry is started for a read/write transaction on this connection. - * The method {@link TransactionRetryListener#retryFinished(Timestamp, long, int, - * TransactionRetryListener.RetryResult)} will be called after the retry has finished. - * - * @param listener The listener to add to this connection. - */ - void addTransactionRetryListener(TransactionRetryListener listener); - - /** - * Removes one existing {@link TransactionRetryListener} from this {@link Connection}, if it is - * present (optional operation). - * - * @param listener The listener to remove from the connection. - * @return true if a listener was removed from the connection. - */ - boolean removeTransactionRetryListener(TransactionRetryListener listener); - - /** - * @return an unmodifiable iterator of the {@link TransactionRetryListener}s registered for this - * connection. - */ - Iterator getTransactionRetryListeners(); - - /** - * Sets the mode for executing DML statements in autocommit mode for this connection. This setting - * is only used when the connection is in autocommit mode, and may only be set while the - * transaction is in autocommit mode and not in a temporary transaction. The autocommit - * transaction mode is reset to its default value of {@link AutocommitDmlMode#TRANSACTIONAL} when - * autocommit mode is changed on the connection. - * - * @param mode The DML autocommit mode to use - *

      - *
    • {@link AutocommitDmlMode#TRANSACTIONAL} DML statements are executed as single - * read-write transaction. After successful execution, the DML statement is guaranteed - * to have been applied exactly once to the database - *
    • {@link AutocommitDmlMode#PARTITIONED_NON_ATOMIC} DML statements are executed as - * partitioned DML transactions. If an error occurs during the execution of the DML - * statement, it is possible that the statement has been applied to some but not all of - * the rows specified in the statement. - *
    - */ - void setAutocommitDmlMode(AutocommitDmlMode mode); - - /** - * @return the current {@link AutocommitDmlMode} setting for this connection. This method may only - * be called on a connection that is in autocommit mode and not while in a temporary - * transaction. - */ - AutocommitDmlMode getAutocommitDmlMode(); - - /** - * Sets the staleness to use for the current read-only transaction. This method may only be called - * when the transaction mode of the current transaction is {@link - * TransactionMode#READ_ONLY_TRANSACTION} and there is no transaction that has started, or when - * the connection is in read-only and autocommit mode. - * - * @param staleness The staleness to use for the current but not yet started read-only transaction - */ - void setReadOnlyStaleness(TimestampBound staleness); - - /** - * @return the read-only staleness setting for the current read-only transaction. This method may - * only be called when the current transaction is a read-only transaction, or when the - * connection is in read-only and autocommit mode. - */ - TimestampBound getReadOnlyStaleness(); - - /** - * Sets the query optimizer version to use for this connection. - * - * @param optimizerVersion The query optimizer version to use. Must be a valid optimizer version - * number, the string LATEST or an empty string. The empty string will instruct - * the connection to use the optimizer version that is defined in the environment variable - * SPANNER_OPTIMIZER_VERSION. If no value is specified in the environment - * variable, the default query optimizer of Cloud Spanner is used. - */ - void setOptimizerVersion(String optimizerVersion); - - /** - * Gets the current query optimizer version of this connection. - * - * @return The query optimizer version that is currently used by this connection. - */ - String getOptimizerVersion(); - /** * Commits the current transaction of this connection. All mutations that have been buffered * during the current transaction will be written to the database. @@ -482,87 +283,6 @@ public interface Connection extends AutoCloseable { */ void rollback(); - /** - * @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 - * mode and no explicit transaction has been started by calling {@link - * Connection#beginTransaction()}. If the {@link Connection} is not in autocommit mode, there - * will always be a transaction. - */ - boolean isInTransaction(); - - /** - * @return true if this connection has a transaction that has started. A transaction - * is automatically started by the first statement that is executed in the transaction. - */ - boolean isTransactionStarted(); - - /** - * Returns the read timestamp of the current/last {@link TransactionMode#READ_ONLY_TRANSACTION} - * transaction, or the read timestamp of the last query in autocommit mode. - * - *
      - *
    • When in autocommit mode: The method will return the read timestamp of the last statement - * if the last statement was a query. - *
    • When in a {@link TransactionMode#READ_ONLY_TRANSACTION} transaction that has started (a - * query has been executed), or that has just committed: The read timestamp of the - * transaction. If the read-only transaction was committed without ever executing a query, - * calling this method after the commit will also throw a {@link SpannerException} - *
    • In all other cases the method will throw a {@link SpannerException}. - *
    - * - * @return the read timestamp of the current/last read-only transaction. - */ - Timestamp getReadTimestamp(); - - /** - * @return the commit timestamp of the last {@link TransactionMode#READ_WRITE_TRANSACTION} - * transaction. This method will throw a {@link SpannerException} if there is no last {@link - * TransactionMode#READ_WRITE_TRANSACTION} transaction (i.e. the last transaction was a {@link - * TransactionMode#READ_ONLY_TRANSACTION}), or if the last {@link - * TransactionMode#READ_WRITE_TRANSACTION} transaction rolled back. It will also throw a - * {@link SpannerException} if the last {@link TransactionMode#READ_WRITE_TRANSACTION} - * transaction was empty when committed. - */ - Timestamp getCommitTimestamp(); - - /** - * Starts a new DDL batch on this connection. A DDL batch allows several DDL statements to be - * grouped into a batch that can be executed as a group. DDL statements that are issued during the - * batch are buffered locally and will return immediately with an OK. It is not guaranteed that a - * DDL statement that has been issued during a batch will eventually succeed when running the - * batch. Aborting a DDL batch will clear the DDL buffer and will have made no changes to the - * database. Running a DDL batch will send all buffered DDL statements to Spanner, and Spanner - * will try to execute these. The result will be OK if all the statements executed successfully. - * If a statement cannot be executed, Spanner will stop execution at that point and return an - * error message for the statement that could not be executed. Preceding statements of the batch - * may have been executed. - * - *

    This method may only be called when the connection is in read/write mode, autocommit mode is - * enabled or no read/write transaction has been started, and there is not already another batch - * active. The connection will only accept DDL statements while a DDL batch is active. - */ - void startBatchDdl(); - - /** - * Starts a new DML batch on this connection. A DML batch allows several DML statements to be - * grouped into a batch that can be executed as a group. DML statements that are issued during the - * batch are buffered locally and will return immediately with an OK. It is not guaranteed that a - * DML statement that has been issued during a batch will eventually succeed when running the - * batch. Aborting a DML batch will clear the DML buffer and will have made no changes to the - * database. Running a DML batch will send all buffered DML statements to Spanner, and Spanner - * will try to execute these. The result will be OK if all the statements executed successfully. - * If a statement cannot be executed, Spanner will stop execution at that point and return {@link - * SpannerBatchUpdateException} for the statement that could not be executed. Preceding statements - * of the batch will have been executed, and the update counts of those statements can be - * retrieved through {@link SpannerBatchUpdateException#getUpdateCounts()}. - * - *

    This method may only be called when the connection is in read/write mode, autocommit mode is - * enabled or no read/write transaction has been started, and there is not already another batch - * active. The connection will only accept DML statements while a DML batch is active. - */ - void startBatchDml(); - /** * 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 @@ -578,19 +298,6 @@ public interface Connection extends AutoCloseable { */ long[] runBatch(); - /** - * Clears all buffered statements in the current batch and ends the batch. - * - *

    This method may only be called when a (possibly empty) batch is active. - */ - void abortBatch(); - - /** @return true if a DDL batch is active on this connection. */ - boolean isDdlBatchActive(); - - /** @return true if a DML batch is active on this connection. */ - boolean isDmlBatchActive(); - /** * Executes the given statement if allowed in the current {@link TransactionMode} and connection * state. The returned value depends on the type of statement: @@ -620,8 +327,6 @@ public interface Connection extends AutoCloseable { */ ResultSet executeQuery(Statement query, QueryOption... options); - AsyncResultSet executeQueryAsync(Statement query, QueryOption... options); - /** * Analyzes a query and returns query plan and/or query execution statistics information. * @@ -706,44 +411,4 @@ public interface Connection extends AutoCloseable { * @throws SpannerException if the {@link Connection} is not in autocommit mode */ void write(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 - * value will not be readable on this {@link Connection} before the transaction is committed. - * - *

    Calling this method is only allowed when not in autocommit mode. See {@link - * Connection#write(Mutation)} for writing mutations in autocommit mode. - * - * @param mutation the {@link Mutation} to buffer for writing to the database on the next commit - * @throws SpannerException if the {@link Connection} is in autocommit mode - */ - void bufferedWrite(Mutation mutation); - - /** - * Buffers the given mutations locally on the current transaction of this {@link Connection}. The - * mutations will be written to the database at the next call to {@link Connection#commit()}. The - * values will not be readable on this {@link Connection} before the transaction is committed. - * - *

    Calling this method is only allowed when not in autocommit mode. See {@link - * Connection#write(Iterable)} for writing mutations in autocommit mode. - * - * @param mutations the {@link Mutation}s to buffer for writing to the database on the next commit - * @throws SpannerException if the {@link Connection} is in autocommit mode - */ - void bufferedWrite(Iterable mutations); - - /** - * This query option is used internally to indicate that a query is executed by the library itself - * to fetch metadata. These queries are specifically allowed to be executed even when a DDL batch - * is active. - * - *

    NOT INTENDED FOR EXTERNAL USE! - */ - @InternalApi - public static final class InternalMetadataQuery implements QueryOption { - @InternalApi public static final InternalMetadataQuery INSTANCE = new InternalMetadataQuery(); - - private InternalMetadataQuery() {} - } } 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 1e37f3927e9..d0e3bc02629 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 @@ -683,11 +683,6 @@ 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); diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionStatementExecutorImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionStatementExecutorImpl.java index 27b975b21e5..e7297e05a1a 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionStatementExecutorImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionStatementExecutorImpl.java @@ -73,13 +73,13 @@ public boolean hasDuration() { } /** The connection to execute the statements on. */ - private final ConnectionImpl connection; + private final AbstractBaseConnection connection; - ConnectionStatementExecutorImpl(ConnectionImpl connection) { + ConnectionStatementExecutorImpl(AbstractBaseConnection connection) { this.connection = connection; } - ConnectionImpl getConnection() { + AbstractBaseConnection getConnection() { return connection; } 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 b49f443227b..7c144f3ad4a 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 @@ -27,7 +27,7 @@ import com.google.cloud.spanner.SpannerException; import com.google.cloud.spanner.SpannerExceptionFactory; import com.google.cloud.spanner.Statement; -import com.google.cloud.spanner.connection.Connection.InternalMetadataQuery; +import com.google.cloud.spanner.connection.BaseConnection.InternalMetadataQuery; import com.google.cloud.spanner.connection.StatementParser.ParsedStatement; import com.google.cloud.spanner.connection.StatementParser.StatementType; import com.google.common.annotations.VisibleForTesting; 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 29f924661c0..292b163eefc 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 @@ -141,6 +141,10 @@ public void closeSpannerPool() { protected java.sql.Connection createJdbcConnection() throws SQLException { return DriverManager.getConnection("jdbc:" + getBaseUrl()); } + + AsyncConnection createAsyncConnection() { + return null; + } ITConnection createConnection() { return createConnection( diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AsyncConnectionApiTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AsyncConnectionApiTest.java index c1381f1b241..c20cd4de734 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AsyncConnectionApiTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AsyncConnectionApiTest.java @@ -22,7 +22,6 @@ 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.connection.ITAbstractSpannerTest.ITConnection; import com.google.common.base.Function; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; @@ -44,9 +43,9 @@ public static void stopExecutor() { @Test public void testSimpleSelectAutocommit() throws Exception { testSimpleSelect( - new Function() { + new Function() { @Override - public Void apply(Connection input) { + public Void apply(AsyncConnection input) { input.setAutocommit(true); return null; } @@ -56,9 +55,9 @@ public Void apply(Connection input) { @Test public void testSimpleSelectReadOnly() throws Exception { testSimpleSelect( - new Function() { + new Function() { @Override - public Void apply(Connection input) { + public Void apply(AsyncConnection input) { input.setReadOnly(true); return null; } @@ -68,24 +67,24 @@ public Void apply(Connection input) { @Test public void testSimpleSelectReadWrite() throws Exception { testSimpleSelect( - new Function() { + new Function() { @Override - public Void apply(Connection input) { + public Void apply(AsyncConnection input) { return null; } }); } - private void testSimpleSelect(Function connectionConfigurator) + private void testSimpleSelect(Function connectionConfigurator) throws Exception { final AtomicInteger rowCount = new AtomicInteger(); ApiFuture res; - try (ITConnection connection = createConnection()) { + try (AsyncConnection connection = createAsyncConnection()) { connectionConfigurator.apply(connection); // Verify that the call is non-blocking. - // mockSpanner.freeze(); + mockSpanner.freeze(); try (AsyncResultSet rs = connection.executeQueryAsync(SELECT_RANDOM_STATEMENT)) { - // mockSpanner.unfreeze(); + mockSpanner.unfreeze(); res = rs.setCallback( executor, 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 4f02fb9a367..438028ffda8 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 @@ -42,7 +42,7 @@ import com.google.cloud.spanner.SpannerException; import com.google.cloud.spanner.SpannerExceptionFactory; import com.google.cloud.spanner.Statement; -import com.google.cloud.spanner.connection.Connection.InternalMetadataQuery; +import com.google.cloud.spanner.connection.BaseConnection.InternalMetadataQuery; import com.google.cloud.spanner.connection.StatementParser.ParsedStatement; import com.google.cloud.spanner.connection.StatementParser.StatementType; import com.google.cloud.spanner.connection.UnitOfWork.UnitOfWorkState;