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 685e9a1e31..7f7d02f95d 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 0000000000..cb05204225 --- /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 0000000000..9b7609d4dc --- /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 129924fbbc..607684611c 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 16f40769fa..542c3da477 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 5023f0fafb..4880cfb370 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 bf0a47222b..344c46f82b 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 cfa8b73c4a..fc72793e86 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 0000000000..ea8396b7ed --- /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 0000000000..f9df08563d --- /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 f665d66ada..066cf5123a 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 6b22ba77c3..edbc7976c0 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 0000000000..63bc234a41 --- /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 60eb151f6e..9f03833e72 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 6daa36f1d6..c1acecdbc2 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 077b660576..d319d13de2 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 7f7d02f95d..1b10f3ba34 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 cb05204225..79c4b3b768 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 9b7609d4dc..fcd0bc770a 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 ac29ba2b37..dd9e976906 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 607684611c..f1bd75b3cc 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 753c3f6f39..4cc0ab9b9e 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 9b30b89985..67e546ad5a 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 015e1862d6..0a8bfd2316 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 4880cfb370..facab374cd 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 45289fb3cd..27257bc65e 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 344c46f82b..ef8e08ef2a 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 32dc3b7157..efe9b381ad 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 fc72793e86..ecd04f26e0 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 ea8396b7ed..0bf7a1a2c9 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 f9df08563d..cd5588187e 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 066cf5123a..aa51449970 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 118b2c57fe..54ae39bbc7 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 c1acecdbc2..8735a2eece 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 d319d13de2..061187696a 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 3dcd523c13..1f2df00e05 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 3a7125312b..b83be137a1 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 e6e473779d..f38809615b 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 4e95f8efe8..5a163d4f6d 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 1b10f3ba34..d395345ace 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 542c3da477..904fa4176b 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 facab374cd..0cff2f0687 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 aa51449970..83913f1bf8 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 54ae39bbc7..8e0375bd27 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 0000000000..112da3948b --- /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 bc7dd5498d..49cbfb905d 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 9662047867..4723497a47 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 0000000000..bc46ff8b0f --- /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 b06d3ae152..ac73cc4f1a 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 0cff2f0687..9366488791 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 a529c4c492..7e09da901c 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 ecd04f26e0..6aa1513411 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 de1a09158c..c696459fee 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 497be948cc..a9114bc22b 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 0000000000..c3a14e7bc5 --- /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 83913f1bf8..6c90bba511 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 112da3948b..21442d3f2e 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 fcd0bc770a..87891f3301 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 0000000000..432d6a8645 --- /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 0000000000..6ffb321490 --- /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 dd9e976906..18cbc3ca85 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 f1bd75b3cc..44f386a727 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 0a8bfd2316..672fa51990 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 9366488791..b8e15ad4ab 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 c3a14e7bc5..fbc1f2177d 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 6c90bba511..7d7b919fc1 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 0000000000..00c296391a --- /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 21442d3f2e..06413ae51e 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 72d537f11a..e386a463f2 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 fbc1f2177d..5782ed8e66 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 00c296391a..336ae9d70e 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 06413ae51e..a8410410b4 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 87891f3301..21f9094b24 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 d193ad1c75..879b632d17 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 b8e15ad4ab..c921eea51f 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 a8afa4b642..d1096aaa6a 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 d395345ace..f4a1297035 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 43de2be092..39827647b6 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 ded74ce85a..351b759628 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 672fa51990..1d9c99eeb5 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 c921eea51f..c3b32351ad 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 bdf7ec954f..35184cdf9c 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 6aa1513411..d292ad5228 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 f3f205a529..cc21774dae 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 8735a2eece..5425042083 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 e60522dc1d..0caab4f574 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 061187696a..3446ba9fc8 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 cc522f3f45..ad56c3cc09 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 1f2df00e05..8bcc05aa00 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 432d6a8645..de15d79c7a 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 5782ed8e66..5dbdd1092f 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 f4a1297035..f191a858a1 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 21f9094b24..82ed5aab98 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 0000000000..c5535bc449 --- /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 c3b32351ad..a0942cc19f 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 ef8e08ef2a..4aaa8ae97b 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 d292ad5228..e4fae3e3a6 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 cd5588187e..9359dc6694 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 5dbdd1092f..d8267dcb4e 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 8e0375bd27..7cb7567321 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 0caab4f574..ff211c4e83 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 bc46ff8b0f..721536cb6b 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 82ed5aab98..fcccf05aac 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 18cbc3ca85..3298e6a2ab 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 a0942cc19f..ee09f28341 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 981f924448..0a8f680f99 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 f73549e465..f324905ef5 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 0000000000..7c80632ed2 --- /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 0a8f680f99..fe85ef7c90 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 f324905ef5..53e5041891 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 7c80632ed2..1f9b0dd81d 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 d8267dcb4e..b20cd3b7a8 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 53e5041891..944a7d70f5 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 1f9b0dd81d..c328f16ef9 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 f191a858a1..9fff15ed62 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 39827647b6..c84bef77cf 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 1d9c99eeb5..77f8946bb1 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 efe9b381ad..3b73158bac 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 e4fae3e3a6..0f098f4f6e 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 bfd739d553..f9cee1c488 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 7d7b919fc1..f6814af8dc 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 5425042083..fc85384453 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 3446ba9fc8..077b660576 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 8bcc05aa00..b790ed93b2 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 3b73158bac..34b5e728e1 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 d1096aaa6a..4695ca9270 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 7cb7567321..1fe8ab4170 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 7b248bfb9d..6b0681b588 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 1fe8ab4170..fb48bb8cd0 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 fb48bb8cd0..d55808ece0 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 79c4b3b768..0dcc379e09 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 fcccf05aac..a673968815 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 c5535bc449..78e3505998 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 ee09f28341..a8fad6c4d0 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 4aaa8ae97b..162f749525 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 34b5e728e1..dd91768362 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 0f098f4f6e..259364e704 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 d55808ece0..9caf7e2cd5 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 29c3e52c6a..26c4832697 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 904fa4176b..e87d40fb20 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 382bbb3097..b6fc3fe065 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 7e09da901c..0b4a92f989 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 259364e704..c3e13f2a87 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 a6e589429f..0fc704f5f7 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 a9114bc22b..0334ca4e40 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 5d67de3fc2..cbe78fe0d0 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 f2589ea6d7..73d903de5b 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 fe85ef7c90..c0d3bdc7d1 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 c328f16ef9..c5e2419ba6 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 4695ca9270..2931407b31 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 77f8946bb1..a425ea5611 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 36b61a920a..e89f03178d 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 26bbef4535..1bcb303f72 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 8d1b780432..d71a03a31c 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 9bd989b98f..24edff9e6b 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 acf48c4510..18f45494d8 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 7ad5293259..267f8a363f 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 b370fd9e6a..93f91f3004 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 267f8a363f..451063403e 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 18f45494d8..c57545bad1 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 a673968815..931e87033a 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 944a7d70f5..50350fd2a1 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 d2c732750e..dd13df65e8 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 3c9e80b8a9..b67e3fb943 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 abc4c0a687..ff76c61b3f 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 76e3cdc389..eee25be70c 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 85ae91dd14..7cd460a958 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 7a3ddb3d90..784dbbcd9f 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 9f6bcebc29..3e53f9b484 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 3c69ba5cd7..77eb238bf2 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 bd1f1b4812..738ccd42c5 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 25af3e2c2f..8efb3a4677 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 907ad4e370..79fe4e3607 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 5c8b59d72b..42875a9a8b 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 931e87033a..d8f35c31e2 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 0000000000..a86c5a776c --- /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 0000000000..945498407d --- /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 3298e6a2ab..2792fd0c86 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 44f386a727..5299d47d10 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 d440162213..278b15d967 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 a425ea5611..01fd8d758d 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 451063403e..9b90ff6477 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 0c6bec4ea8..52c35cb713 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 162f749525..753adc1e86 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 c5488ac55d..0d429661ad 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 35184cdf9c..8dbab88314 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 c3e13f2a87..a91c6a90ef 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 cb8cf3bc55..9f278fb11d 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 0000000000..95c1077fa6 --- /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 5247ce2c13..9a1dc69a0c 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 ce24791859..1e37f3927e 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 b18f3fa891..b49f443227 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 ff38338d62..250e7a1cc7 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 7a0155cbfb..3689d4b8d9 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 614d0c61e5..edf4a9a528 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 6221cc447b..ab5610d072 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 e372229c64..49001cd8d8 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 fc027576c1..57ad55480c 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 0334ca4e40..31347ae547 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 cbe78fe0d0..a020997fd7 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 93f91f3004..5100b2723e 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 cc21774dae..cbcb10d25b 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 c57545bad1..538075fe4f 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 b270c227f8..ead4a6bb15 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 1092493331..74d8b1c905 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 3497b42bc7..1e05dd6e12 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 0000000000..f5e2e9ad31 --- /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 28bc456476..e7ac02c49c 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 9fff15ed62..b1d752e6f4 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 0dcc379e09..c44a42994e 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 d8f35c31e2..a92026536b 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 de15d79c7a..3cae49e65b 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 6ffb321490..5b83402919 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 a86c5a776c..e5b5b3bb03 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 945498407d..7c7f8504c7 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 351b759628..3892a018f1 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 01fd8d758d..f0e01cf986 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 9b90ff6477..60282aa74b 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 d0bad0e652..d8145a5a25 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 a91c6a90ef..9373edde2a 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 edf4a9a528..3da17cf26d 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 57ad55480c..fb9dc1da59 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 31347ae547..84d8b1b6be 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 a020997fd7..a3ba51d000 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 5100b2723e..e2655ccc72 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 cbcb10d25b..06eff8ebfb 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 538075fe4f..8e63538079 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 ead4a6bb15..ad32b20586 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 74d8b1c905..a9abea938d 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 1e05dd6e12..29532c5483 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 f5e2e9ad31..c1381f1b24 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 9373edde2a..ed53fba6d2 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 baf7670e50..f0c67e02c1 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 b1d752e6f4..f4e6596fdf 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 a92026536b..a86af43434 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 276915a4b7..ed854b964f 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 cd5e15f5fc..26e86c289a 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 e5b5b3bb03..8e8f808d0b 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 7c7f8504c7..b830a8a61b 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 ed854b964f..6126a077bf 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 ed53fba6d2..36565fb0f8 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 0000000000..fbd8e44ffe --- /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 a3ba51d000..eb00047ca4 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 0000000000..8b6772e43b --- /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 50350fd2a1..83d060596c 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 a86af43434..3e15dc6e63 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 4a9e3efcf4..7d2d8cad8f 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 1ed8165cd0..f9c3c2040f 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 f4e6596fdf..bc4a868564 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 3e15dc6e63..f277388b0b 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 83d060596c..13e4c47d08 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 8e8f808d0b..7cce2296ce 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 b830a8a61b..31d809dad6 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 6126a077bf..711a905ac9 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 0000000000..487b961651 --- /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 36565fb0f8..c10b713285 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 fbd8e44ffe..1926860162 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 8b6772e43b..65681eda09 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 0bf7a1a2c9..8c195bb6da 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 65681eda09..9d4b5dac2a 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 73d903de5b..25c3428d38 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 31d809dad6..aca7f6e8ce 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 711a905ac9..b6ffa4da8e 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 0000000000..005dfe7b60 --- /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 487b961651..3cdd643042 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 1926860162..bf76ea4f39 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 9d4b5dac2a..76eec9f117 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 25c3428d38..8ba13ddfae 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 8c195bb6da..c3ad9f45dc 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 7cce2296ce..5c1ad5cc1d 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 aca7f6e8ce..082fa827e7 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 2792fd0c86..d52d1d892e 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 e1cbfca8c5..5798e35202 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 b6ffa4da8e..90e399fad6 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 005dfe7b60..55b6102a27 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 3cdd643042..5922726120 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 76eec9f117..e9f33546dc 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 c0d3bdc7d1..cc6784b679 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 5c1ad5cc1d..d519c68013 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 5922726120..bc8262a535 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 e9f33546dc..c5de654219 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 2d812d5bbc..4bdb8804e7 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 9f278fb11d..cb8cf3bc55 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 95c1077fa6..0000000000 --- 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 9a1dc69a0c..5247ce2c13 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 1e37f3927e..ce24791859 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 b49f443227..b18f3fa891 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 250e7a1cc7..ff38338d62 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 3689d4b8d9..7a0155cbfb 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 3da17cf26d..614d0c61e5 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 ab5610d072..6221cc447b 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 49001cd8d8..e372229c64 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 29532c5483..3497b42bc7 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 c1381f1b24..0000000000 --- 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 2266a88a25..6918c9f268 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 e7ac02c49c..1d4b9a99f1 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 f5af48f511da94df8dec03a276f42a0908f6d1cd Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Sun, 21 Jun 2020 20:30:38 +0200 Subject: [PATCH 45/49] chore: run code formatter --- .../google/cloud/spanner/connection/ReadOnlyTransactionTest.java | 1 + .../cloud/spanner/connection/SingleUseTransactionTest.java | 1 + 2 files changed, 2 insertions(+) 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 6918c9f268..118f596c86 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,7 @@ 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; 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 1d4b9a99f1..1a8eecf27d 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,7 @@ 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; From b38d16481c121c183aa1b781d50ee1f476954b2f Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Sun, 21 Jun 2020 20:50:01 +0200 Subject: [PATCH 46/49] chore: fix flaky test case --- .../spanner/AsyncTransactionManagerTest.java | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 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 c5de654219..e2299f3615 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 @@ -481,16 +481,12 @@ 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. + // 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. + // transaction. The transaction result will be resolved when the commit + // has actually finished successfully. return ApiFutures.immediateFuture(null); } }, @@ -498,12 +494,14 @@ 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 call to + // commitAsync() already knows that the transaction has aborted. If it does, it will not + // attempt to call the Commit RPC and instead directly propagate the Aborted error. assertThat(mockSpanner.getRequestTypes()) - .containsExactly( + .containsAtLeast( BatchCreateSessionsRequest.class, BeginTransactionRequest.class, ExecuteSqlRequest.class, - CommitRequest.class, BeginTransactionRequest.class, ExecuteSqlRequest.class, CommitRequest.class); From eee881a80e106cd878bd61b1f96419a989937e0a Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Tue, 30 Jun 2020 18:34:44 +0200 Subject: [PATCH 47/49] tests: fix ITs for emulator --- .../java/com/google/cloud/spanner/it/ITAsyncAPITest.java | 4 ++++ .../cloud/spanner/it/ITTransactionManagerAsyncTest.java | 8 ++++++++ 2 files changed, 12 insertions(+) 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 721536cb6b..a2239aa3b9 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 @@ -18,6 +18,7 @@ import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.fail; +import static org.junit.Assume.assumeFalse; import com.google.api.core.ApiFuture; import com.google.cloud.spanner.AsyncResultSet; @@ -278,6 +279,9 @@ public void columnNotFound() throws Exception { @Test public void asyncRunnerFireAndForgetInvalidUpdate() throws Exception { + assumeFalse( + "errors in read/write transactions on emulator are sticky", + env.getTestHelper().isEmulator()); try { assertThat(client.singleUse().readRow("TestTable", Key.of("k999"), ALL_COLUMNS)).isNull(); AsyncRunner runner = client.runAsync(); 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 4bdb8804e7..c802493dec 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 @@ -33,17 +33,20 @@ import com.google.cloud.spanner.ErrorCode; 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.SpannerException; import com.google.cloud.spanner.Struct; import com.google.cloud.spanner.TransactionContext; import com.google.cloud.spanner.TransactionManager.TransactionState; +import com.google.common.collect.ImmutableList; 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.Before; import org.junit.BeforeClass; import org.junit.ClassRule; import org.junit.Test; @@ -84,6 +87,11 @@ public static void setUpDatabase() { client = env.getTestHelper().getDatabaseClient(db); } + @Before + public void clearTable() { + client.write(ImmutableList.of(Mutation.delete("T", KeySet.all()))); + } + @Test public void testSimpleInsert() throws ExecutionException, InterruptedException { try (AsyncTransactionManager manager = client.transactionManagerAsync()) { From a1ef6405d49286055f2f27ebc0861bb3ac7f9296 Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Tue, 30 Jun 2020 19:10:39 +0200 Subject: [PATCH 48/49] fix: SpannerOptions.toBuilder().host should override emulatorHost --- .../main/java/com/google/cloud/spanner/SpannerOptions.java | 5 +++++ 1 file changed, 5 insertions(+) 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 31104f72a0..35a288530f 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 @@ -376,6 +376,11 @@ private Builder() { Builder(SpannerOptions options) { super(options); + if (options.getHost() != null + && this.emulatorHost != null + && !options.getHost().equals(this.emulatorHost)) { + this.emulatorHost = null; + } this.numChannels = options.numChannels; this.sessionPoolOptions = options.sessionPoolOptions; this.prefetchChunks = options.prefetchChunks; From ea45612620d9d5807f4de936aaf90a8e9c26833a Mon Sep 17 00:00:00 2001 From: Olav Loite Date: Tue, 30 Jun 2020 19:57:25 +0200 Subject: [PATCH 49/49] tests: fix potentially hanging test --- .../cloud/spanner/AsyncResultSetImplStressTest.java | 8 ++++++-- 1 file changed, 6 insertions(+), 2 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 c3ad9f45dc..c3383cadda 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 @@ -397,7 +397,7 @@ public CallbackResponse cursorReady(AsyncResultSet resultSet) { } } final AtomicBoolean finished = new AtomicBoolean(false); - // Both resume and cancel result sets randomly. + // Both resume and cancel resultsets randomly. ExecutorService resumeService = createExecService(); resumeService.execute( new Runnable() { @@ -407,6 +407,10 @@ public void run() { // Randomly resume result sets. resultSets.get(random.nextInt(resultSets.size())).resume(); } + // Make sure all result sets finish. + for (AsyncResultSet rs : resultSets) { + rs.resume(); + } } }); ExecutorService cancelService = createExecService(); @@ -438,7 +442,7 @@ public void run() { int index = 0; for (ApiFuture> future : futures) { try { - ImmutableList list = future.get(); + ImmutableList list = future.get(30L, TimeUnit.SECONDS); // 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