From 393f3b2b77e1f36e3a13edbe03899641c9295226 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= Date: Mon, 26 Oct 2026 10:35:01 +0100 Subject: [PATCH] refactor(spanner): unify streaming read and query startup Streaming SQL queries and streaming key-value reads previously duplicated the same request setup and stream startup logic in two separate code paths. This change unifies both operations behind a single internal request abstraction so that option configuration, retry settings, and stream lifecycle handling are maintained in one place. As part of unifying stream startup, this also hardens how inlined BeginTransaction state is tracked across the stream lifecycle: - Records whether a request carries an inlined BeginTransaction on the request itself and passes it to the stream iterator at construction time rather than when the gRPC call is attached, keeping the state available across retry backoffs while ensuring it is cleared once the iterator is closed. - Ensures that if starting an inlined-begin stream fails before the stream can be consumed, the pending transaction ID future is failed immediately with a consistent exception so concurrent callers on the same transaction do not wait for a timeout. --- .../cloud/spanner/AbstractReadContext.java | 283 ++---- .../cloud/spanner/GrpcStreamIterator.java | 35 +- .../spanner/MultiUseReadOnlyTransaction.java | 2 +- .../cloud/spanner/ReadOrQueryRequest.java | 300 ++++++ .../cloud/spanner/TransactionRunnerImpl.java | 83 +- .../spanner/AsyncReadOnlyTransactionTest.java | 353 +++++++ .../cloud/spanner/DatabaseClientImplTest.java | 14 +- .../cloud/spanner/GrpcResultSetTest.java | 3 +- .../cloud/spanner/GrpcStreamIteratorTest.java | 82 +- .../cloud/spanner/MockSpannerServiceImpl.java | 5 +- .../cloud/spanner/ReadFormatTestRunner.java | 3 +- ...adWriteTransactionWithInlineBeginTest.java | 912 ++++++++++++++++++ 12 files changed, 1814 insertions(+), 261 deletions(-) create mode 100644 java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ReadOrQueryRequest.java diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java index a769758861f4..8345d78388ec 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java @@ -26,19 +26,17 @@ import com.google.api.core.ApiFutures; import com.google.api.core.SettableApiFuture; import com.google.api.gax.core.ExecutorProvider; -import com.google.api.gax.retrying.RetrySettings; -import com.google.api.gax.rpc.StatusCode.Code; import com.google.cloud.grpc.GcpManagedChannel.ChannelAffinityRef; import com.google.cloud.spanner.AbstractResultSet.CloseableIterator; import com.google.cloud.spanner.AsyncResultSet.CallbackResponse; import com.google.cloud.spanner.AsyncResultSet.ReadyCallback; +import com.google.cloud.spanner.AsyncResultSet.StreamMessageListener; import com.google.cloud.spanner.Options.QueryOption; import com.google.cloud.spanner.Options.ReadOption; import com.google.cloud.spanner.SessionClient.SessionOption; 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.base.Supplier; import com.google.common.util.concurrent.MoreExecutors; import com.google.protobuf.ByteString; import com.google.spanner.v1.DirectedReadOptions; @@ -54,11 +52,10 @@ import com.google.spanner.v1.TransactionSelector; import io.opentelemetry.api.common.Attributes; import java.util.Map; -import java.util.Set; import java.util.concurrent.atomic.AtomicLong; import java.util.logging.Logger; -import javax.annotation.Nullable; import javax.annotation.concurrent.GuardedBy; +import org.jspecify.annotations.Nullable; /** * Abstract base class for all {@link ReadContext}s + concrete implementations of read-only {@link @@ -170,7 +167,7 @@ static void assertTimestampAvailable(boolean available) { TraceWrapper tracer; private final int defaultPrefetchChunks; private final QueryOptions defaultQueryOptions; - private final DirectedReadOptions defaultDirectedReadOptions; + final DirectedReadOptions defaultDirectedReadOptions; private final DecodeMode defaultDecodeMode; private final Clock clock; private volatile String cachedRequestTag; @@ -393,11 +390,19 @@ RequestOptions buildRequestOptions(Options options) { ExecuteSqlRequest.Builder getExecuteSqlRequestBuilder( Statement statement, QueryMode queryMode, Options options, boolean withTransactionSelector) { + return getExecuteSqlRequest( + statement, queryMode, options, /* partitionToken= */ null, withTransactionSelector) + .getBuilder(); + } + + private ReadOrQueryRequest.Query getExecuteSqlRequest( + Statement statement, + QueryMode queryMode, + Options options, + @Nullable ByteString partitionToken, + boolean withTransactionSelector) { ExecuteSqlRequest.Builder builder = - ExecuteSqlRequest.newBuilder() - .setSql(statement.getSql()) - .setQueryMode(queryMode) - .setSession(session.getName()); + ExecuteSqlRequest.newBuilder().setSql(statement.getSql()).setQueryMode(queryMode); addParameters(builder, statement.getParameters()); if (withTransactionSelector) { TransactionSelector selector = getTransactionSelector(); @@ -405,14 +410,6 @@ ExecuteSqlRequest.Builder getExecuteSqlRequestBuilder( builder.setTransaction(selector); } } - if (options.hasDataBoostEnabled()) { - builder.setDataBoostEnabled(options.dataBoostEnabled()); - } - if (options.hasDirectedReadOptions()) { - builder.setDirectedReadOptions(options.directedReadOptions()); - } else if (defaultDirectedReadOptions != null) { - builder.setDirectedReadOptions(defaultDirectedReadOptions); - } if (!isReadOnly()) { builder.setSeqno(getSeqNo()); } @@ -420,8 +417,7 @@ ExecuteSqlRequest.Builder getExecuteSqlRequestBuilder( builder.setLastStatement(options.isLastStatement()); } builder.setQueryOptions(buildQueryOptions(statement.getQueryOptions())); - builder.setRequestOptions(buildRequestOptions(options)); - return builder; + return ReadOrQueryRequest.forQuery(this, builder, statement, options, partitionToken); } static void addParameters(ExecuteSqlRequest.Builder builder, Map stmtParameters) { @@ -472,84 +468,79 @@ ExecuteBatchDmlRequest.Builder getExecuteBatchDmlRequestBuilder( return builder; } - @FunctionalInterface - private interface StreamStarter { - StreamCallResult start( - @Nullable ByteString resumeToken, - XGoogSpannerRequestId requestId, - SpannerRpc.ResultStreamConsumer consumer); - } - - private static final class StreamCallResult { - final SpannerRpc.StreamingCall call; - final boolean withBeginTransaction; - - StreamCallResult(SpannerRpc.StreamingCall call, boolean withBeginTransaction) { - this.call = call; - this.withBeginTransaction = withBeginTransaction; - } - } - - @FunctionalInterface - private interface StreamCall { - SpannerRpc.StreamingCall execute( - SpannerRpc.ResultStreamConsumer consumer, XGoogSpannerRequestId requestId); - } - - private StreamCallResult executeCall( - boolean withBeginTransaction, - SpannerRpc.ResultStreamConsumer consumer, - XGoogSpannerRequestId requestId, - StreamCall call) { - try { - return new StreamCallResult(call.execute(consumer, requestId), withBeginTransaction); - } catch (RuntimeException | Error t) { - onStartFailed(withBeginTransaction, t); - throw t; - } - } - - @Nullable - private TransactionSelector prepareTransactionSelector( - @Nullable ByteString resumeToken, boolean hasTransaction) { - if (resumeToken != null || !hasTransaction) { - return getTransactionSelector(); - } - return null; - } - private ResultSet createStreamingResultSet( - String operationName, - Attributes attributes, - RetrySettings retrySettings, - Set retryableCodes, - Options options, - Supplier streamSupplier, - StreamStarter streamStarter) { + ReadOrQueryRequest request, Attributes attributes, Options options) { + final int prefetchChunks = + options.hasPrefetchChunks() ? options.prefetchChunks() : this.defaultPrefetchChunks; + final boolean lastStatement = options.hasLastStatement() && options.isLastStatement(); ResumableStreamIterator stream = new ResumableStreamIterator( MAX_BUFFERED_CHUNKS, - operationName, - span, - tracer, + request.getOperationName(), + this.span, + this.tracer, attributes, - session.getErrorHandler(), - retrySettings, - retryableCodes, - session.getRequestIdCreator()) { + this.session.getErrorHandler(), + request.getRetrySettings(this.rpc), + request.getRetryableCodes(this.rpc), + this.session.getRequestIdCreator()) { + private volatile boolean closed; + @Override CloseableIterator startStream( @Nullable ByteString resumeToken, - AsyncResultSet.StreamMessageListener streamListener, + @Nullable StreamMessageListener streamListener, XGoogSpannerRequestId requestId) { - GrpcStreamIterator stream = streamSupplier.get(); + if (resumeToken != null) { + request.setResumeToken(resumeToken); + } + // 1. When starting the stream for the first time (resumeToken == null and + // !request.hasTransaction()), look up and set the TransactionSelector. + // If the stream is retried before receiving a resumeToken (resumeToken == null) + // and already has a TransactionSelector, we reuse the existing selector. In + // particular, if the initial request included an inlined BeginTransaction, the + // retry must still send BeginTransaction, and calling getTransactionSelector() + // again before the transaction ID has been received would block waiting on its + // own transactionIdFuture. + // 2. When resuming a stream after receiving a resumeToken (resumeToken != null), + // the TransactionSelector only changes if the initial request included an + // inlined BeginTransaction (request.isWithBeginTransaction()): receiving a + // resumeToken means the first response chunk (and its Transaction metadata) has + // already been processed, so we replace the BeginTransaction selector with the + // transaction ID selector for the resumed stream. + if ((resumeToken == null && !request.hasTransaction()) + || (resumeToken != null && request.isWithBeginTransaction())) { + TransactionSelector selector = getTransactionSelector(); + if (selector != null) { + request.setTransaction(selector); + } + } + boolean withBeginTransaction = request.isWithBeginTransaction(); + GrpcStreamIterator stream = + new GrpcStreamIterator( + request.getStatement(), + lastStatement, + prefetchChunks, + cancelQueryWhenClientIsClosed, + withBeginTransaction); if (streamListener != null) { stream.registerListener(streamListener); } - StreamCallResult result = - streamStarter.start(resumeToken, requestId, stream.consumer()); + SpannerRpc.StreamingCall call; + try { + call = + request.start( + rpc, + stream.consumer(), + getTransactionChannelHint(), + requestId, + isRouteToLeader()); + } catch (Throwable throwable) { + onStartFailed(withBeginTransaction, lastStatement, throwable); + throw throwable; + } session.markUsed(clock.instant()); - stream.setCall(result.call, result.withBeginTransaction); + stream.setCall(call); return stream; } @@ -557,9 +548,28 @@ CloseableIterator startStream( boolean prepareIteratorForRetryOnDifferentGrpcChannel() { return AbstractReadContext.this.prepareRetryOnDifferentGrpcChannel(); } + + @Override + public void close(@Nullable String message) { + // Mark closed before super.close() cancels the underlying gRPC call so that a consumer + // thread woken by the cancellation never reports an inlined-begin failure for a closed + // result set. + this.closed = true; + super.close(message); + } + + @Override + public boolean isWithBeginTransaction() { + return !this.closed && request.isWithBeginTransaction(); + } + + @Override + public boolean isLastStatement() { + return lastStatement; + } }; return new GrpcResultSet( - stream, this, options.hasDecodeMode() ? options.decodeMode() : defaultDecodeMode); + stream, this, options.hasDecodeMode() ? options.decodeMode() : this.defaultDecodeMode); } ResultSet executeQueryInternalWithOptions( @@ -568,47 +578,11 @@ ResultSet executeQueryInternalWithOptions( final Options options, final ByteString partitionToken) { beforeReadOrQuery(); - final int prefetchChunks = - options.hasPrefetchChunks() ? options.prefetchChunks() : defaultPrefetchChunks; - final ExecuteSqlRequest.Builder request = - getExecuteSqlRequestBuilder( - statement, queryMode, options, /* withTransactionSelector= */ false); - if (partitionToken != null) { - request.setPartitionToken(partitionToken); - } return createStreamingResultSet( - SpannerImpl.QUERY, - tracer.createStatementAttributes(statement, options), - rpc.getExecuteQueryRetrySettings(), - rpc.getExecuteQueryRetryableCodes(), - options, - () -> - new GrpcStreamIterator( - statement, - request.getLastStatement(), - prefetchChunks, - cancelQueryWhenClientIsClosed), - (resumeToken, requestId, consumer) -> { - if (resumeToken != null) { - request.setResumeToken(resumeToken); - } - TransactionSelector selector = - prepareTransactionSelector(resumeToken, request.hasTransaction()); - if (selector != null) { - request.setTransaction(selector); - } - return executeCall( - request.getTransaction().hasBegin(), - consumer, - requestId, - (streamConsumer, reqId) -> - rpc.executeQuery( - request.build(), - streamConsumer, - getTransactionChannelHint(), - reqId, - isRouteToLeader())); - }); + getExecuteSqlRequest( + statement, queryMode, options, partitionToken, /* withTransactionSelector= */ false), + this.tracer.createStatementAttributes(statement, options), + options); } static Map getChannelHintOptions( @@ -694,8 +668,7 @@ boolean prepareRetryOnDifferentGrpcChannel() { * Returns the transaction tag for this {@link AbstractReadContext} or null if this * {@link AbstractReadContext} does not have a transaction tag. */ - @Nullable - String getTransactionTag() { + @Nullable String getTransactionTag() { return null; } @@ -715,7 +688,7 @@ public void onDone(boolean withBeginTransaction) { this.session.onReadDone(); } - void onStartFailed(boolean withBeginTransaction, Throwable t) {} + void onStartFailed(boolean withBeginTransaction, boolean lastStatement, Throwable throwable) {} /** * For transactions other than read-write, the MultiplexedSessionPrecommitToken will not be @@ -744,32 +717,17 @@ ResultSet readInternalWithOptions( ByteString partitionToken) { beforeReadOrQuery(); final ReadRequest.Builder builder = - ReadRequest.newBuilder() - .setSession(session.getName()) - .setTable(checkNotNull(table)) - .addAllColumns(columns); + ReadRequest.newBuilder().setTable(checkNotNull(table)).addAllColumns(columns); if (readOptions.hasLimit()) { builder.setLimit(readOptions.limit()); } - keys.appendToProto(builder.getKeySetBuilder()); if (index != null) { builder.setIndex(index); } - if (partitionToken != null) { - builder.setPartitionToken(partitionToken); - } - if (readOptions.hasDataBoostEnabled()) { - builder.setDataBoostEnabled(readOptions.dataBoostEnabled()); - } if (readOptions.hasOrderBy()) { builder.setOrderBy(readOptions.orderBy()); } - if (readOptions.hasDirectedReadOptions()) { - builder.setDirectedReadOptions(readOptions.directedReadOptions()); - } else if (defaultDirectedReadOptions != null) { - builder.setDirectedReadOptions(defaultDirectedReadOptions); - } if (readOptions.hasLockHint()) { if (isReadOnly()) { logger.warning( @@ -779,39 +737,10 @@ ResultSet readInternalWithOptions( builder.setLockHint(readOptions.lockHint()); } } - builder.setRequestOptions(buildRequestOptions(readOptions)); - final int prefetchChunks = - readOptions.hasPrefetchChunks() ? readOptions.prefetchChunks() : defaultPrefetchChunks; - final boolean lastStatement = - readOptions.hasLastStatement() ? readOptions.isLastStatement() : false; return createStreamingResultSet( - SpannerImpl.READ, - tracer.createTableAttributes(table, readOptions), - rpc.getReadRetrySettings(), - rpc.getReadRetryableCodes(), - readOptions, - () -> new GrpcStreamIterator(lastStatement, prefetchChunks, cancelQueryWhenClientIsClosed), - (resumeToken, requestId, consumer) -> { - if (resumeToken != null) { - builder.setResumeToken(resumeToken); - } - TransactionSelector selector = - prepareTransactionSelector(resumeToken, builder.hasTransaction()); - if (selector != null) { - builder.setTransaction(selector); - } - return executeCall( - builder.getTransaction().hasBegin(), - consumer, - requestId, - (streamConsumer, reqId) -> - rpc.read( - builder.build(), - streamConsumer, - getTransactionChannelHint(), - reqId, - isRouteToLeader())); - }); + ReadOrQueryRequest.forRead(this, builder, readOptions, partitionToken), + this.tracer.createTableAttributes(table, readOptions), + readOptions); } private Struct consumeSingleRow(ResultSet resultSet) { diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/GrpcStreamIterator.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/GrpcStreamIterator.java index e49860fbfd08..7b0cb08b5c84 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/GrpcStreamIterator.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/GrpcStreamIterator.java @@ -86,7 +86,7 @@ public String toString() { private boolean callCleared; private boolean closed; @Nullable private String closeMessage; - private volatile boolean withBeginTransaction; + private final boolean withBeginTransaction; private final boolean lastStatement; @Nullable private volatile Timeout streamWaitTimeout; @Nullable private volatile SpannerException error; @@ -97,7 +97,12 @@ public String toString() { @VisibleForTesting GrpcStreamIterator( boolean lastStatement, int prefetchChunks, boolean cancelQueryWhenClientIsClosed) { - this(null, lastStatement, prefetchChunks, cancelQueryWhenClientIsClosed); + this( + /* statement= */ null, + lastStatement, + prefetchChunks, + cancelQueryWhenClientIsClosed, + /* withBeginTransaction= */ false); } @VisibleForTesting @@ -106,9 +111,24 @@ public String toString() { boolean lastStatement, int prefetchChunks, boolean cancelQueryWhenClientIsClosed) { + this( + statement, + lastStatement, + prefetchChunks, + cancelQueryWhenClientIsClosed, + /* withBeginTransaction= */ false); + } + + GrpcStreamIterator( + Statement statement, + boolean lastStatement, + int prefetchChunks, + boolean cancelQueryWhenClientIsClosed, + boolean withBeginTransaction) { this.statement = statement; this.lastStatement = lastStatement; this.prefetchChunks = prefetchChunks; + this.withBeginTransaction = withBeginTransaction; this.consumer = new ConsumerImpl(cancelQueryWhenClientIsClosed); // One extra to allow for END_OF_STREAM message. this.stream = new LinkedBlockingQueue<>(prefetchChunks + 1); @@ -136,12 +156,11 @@ void registerListener(AsyncResultSet.StreamMessageListener streamMessageListener } } - public void setCall(SpannerRpc.StreamingCall call, boolean withBeginTransaction) { + void setCall(SpannerRpc.StreamingCall call) { Timeout timeout = extractTimeout(call); boolean shouldCancel; String messageToCancel; synchronized (this) { - this.withBeginTransaction = withBeginTransaction; shouldCancel = this.closed && !this.callCleared; if (!this.callCleared && !this.closed) { this.call = call; @@ -246,10 +265,10 @@ public boolean isLastStatement() { * *

When an asynchronous {@link AsyncResultSet.StreamMessageListener} is registered, consumption * is gated until the underlying {@link SpannerRpc.StreamingCall} is initialized via {@link - * #setCall(SpannerRpc.StreamingCall, boolean)} so that call context, cancellation hooks, and - * listener notifications are properly established. In synchronous mode (where no listener is - * registered, e.g. {@code streamMessageListener == null}), this gating is not needed because - * callers invoke {@link #next()} directly on the iterator. + * #setCall(SpannerRpc.StreamingCall)} so that call context, cancellation hooks, and listener + * notifications are properly established. In synchronous mode (where no listener is registered, + * e.g. {@code streamMessageListener == null}), this gating is not needed because callers invoke + * {@link #next()} directly on the iterator. */ private boolean isReadyForConsumption() { return callInitialized || streamMessageListener == null; diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiUseReadOnlyTransaction.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiUseReadOnlyTransaction.java index 3a35a0be07fa..05323fb12144 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiUseReadOnlyTransaction.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiUseReadOnlyTransaction.java @@ -483,7 +483,7 @@ public void onDone(boolean withBeginTransaction) { } @Override - void onStartFailed(boolean withBeginTransaction, Throwable throwable) { + void onStartFailed(boolean withBeginTransaction, boolean lastStatement, Throwable throwable) { if (withBeginTransaction) { failTransactionIdFuture(throwable); } diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ReadOrQueryRequest.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ReadOrQueryRequest.java new file mode 100644 index 000000000000..08808e312a94 --- /dev/null +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ReadOrQueryRequest.java @@ -0,0 +1,300 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.spanner; + +import static com.google.common.base.Preconditions.checkNotNull; + +import com.google.api.gax.retrying.RetrySettings; +import com.google.api.gax.rpc.StatusCode.Code; +import com.google.cloud.spanner.spi.v1.SpannerRpc; +import com.google.protobuf.ByteString; +import com.google.spanner.v1.DirectedReadOptions; +import com.google.spanner.v1.ExecuteSqlRequest; +import com.google.spanner.v1.ReadRequest; +import com.google.spanner.v1.RequestOptions; +import com.google.spanner.v1.TransactionSelector; +import java.util.Map; +import java.util.Set; +import org.jspecify.annotations.NullMarked; +import org.jspecify.annotations.Nullable; + +/** + * Package-private abstraction around either an {@link ExecuteSqlRequest.Builder} or a {@link + * ReadRequest.Builder} to unify common request options, streaming request preparation, and + * execution in {@link AbstractReadContext}. + */ +@NullMarked +abstract class ReadOrQueryRequest { + + static Query forQuery( + AbstractReadContext context, + ExecuteSqlRequest.Builder builder, + Statement statement, + Options options, + @Nullable ByteString partitionToken) { + return applyOptions(new Query(builder, statement), context, options, partitionToken); + } + + static Read forRead( + AbstractReadContext context, + ReadRequest.Builder builder, + Options options, + @Nullable ByteString partitionToken) { + return applyOptions(new Read(builder), context, options, partitionToken); + } + + private static T applyOptions( + T request, + AbstractReadContext context, + Options options, + @Nullable ByteString partitionToken) { + request.setSession(context.session.getName()); + if (partitionToken != null) { + request.setPartitionToken(partitionToken); + } + if (options.hasDataBoostEnabled()) { + request.setDataBoostEnabled(options.dataBoostEnabled()); + } + if (options.hasDirectedReadOptions()) { + request.setDirectedReadOptions(options.directedReadOptions()); + } else if (context.defaultDirectedReadOptions != null) { + request.setDirectedReadOptions(context.defaultDirectedReadOptions); + } + request.setRequestOptions(context.buildRequestOptions(options)); + return request; + } + + /** + * Indicates whether the most recently configured {@link TransactionSelector} carries an inlined + * {@code begin} option. The underlying protobuf {@code Builder} is only accessed from {@code + * startStream()}, whereas this field is {@code volatile} so {@link #isWithBeginTransaction()} can + * be read safely from a concurrent {@link GrpcResultSet#close()} call. + */ + private volatile boolean withBeginTransaction; + + private ReadOrQueryRequest(boolean withBeginTransaction) { + this.withBeginTransaction = withBeginTransaction; + } + + abstract String getOperationName(); + + abstract RetrySettings getRetrySettings(SpannerRpc rpc); + + abstract Set getRetryableCodes(SpannerRpc rpc); + + abstract @Nullable Statement getStatement(); + + abstract void setSession(String session); + + abstract void setPartitionToken(ByteString partitionToken); + + abstract void setDataBoostEnabled(boolean dataBoostEnabled); + + abstract void setDirectedReadOptions(DirectedReadOptions directedReadOptions); + + abstract void setRequestOptions(RequestOptions requestOptions); + + abstract void setResumeToken(ByteString resumeToken); + + abstract boolean hasTransaction(); + + final void setTransaction(TransactionSelector selector) { + setTransactionInternal(selector); + this.withBeginTransaction = selector.hasBegin(); + } + + /** + * Returns {@code true} if the current {@link TransactionSelector} on this request includes an + * inlined {@code BeginTransaction} option. Safe to call concurrently with {@code startStream()}. + */ + final boolean isWithBeginTransaction() { + return withBeginTransaction; + } + + abstract void setTransactionInternal(TransactionSelector selector); + + abstract SpannerRpc.StreamingCall start( + SpannerRpc rpc, + SpannerRpc.ResultStreamConsumer consumer, + @Nullable Map channelHint, + XGoogSpannerRequestId requestId, + boolean routeToLeader); + + static final class Query extends ReadOrQueryRequest { + private final ExecuteSqlRequest.Builder builder; + private final Statement statement; + + private Query(ExecuteSqlRequest.Builder builder, Statement statement) { + super(checkNotNull(builder).getTransaction().hasBegin()); + this.builder = builder; + this.statement = checkNotNull(statement); + } + + ExecuteSqlRequest.Builder getBuilder() { + return builder; + } + + @Override + String getOperationName() { + return SpannerImpl.QUERY; + } + + @Override + RetrySettings getRetrySettings(SpannerRpc rpc) { + return rpc.getExecuteQueryRetrySettings(); + } + + @Override + Set getRetryableCodes(SpannerRpc rpc) { + return rpc.getExecuteQueryRetryableCodes(); + } + + @Override + Statement getStatement() { + return statement; + } + + @Override + void setSession(String session) { + builder.setSession(session); + } + + @Override + void setPartitionToken(ByteString partitionToken) { + builder.setPartitionToken(partitionToken); + } + + @Override + void setDataBoostEnabled(boolean dataBoostEnabled) { + builder.setDataBoostEnabled(dataBoostEnabled); + } + + @Override + void setDirectedReadOptions(DirectedReadOptions directedReadOptions) { + builder.setDirectedReadOptions(directedReadOptions); + } + + @Override + void setRequestOptions(RequestOptions requestOptions) { + builder.setRequestOptions(requestOptions); + } + + @Override + void setResumeToken(ByteString resumeToken) { + builder.setResumeToken(resumeToken); + } + + @Override + boolean hasTransaction() { + return builder.hasTransaction(); + } + + @Override + void setTransactionInternal(TransactionSelector selector) { + builder.setTransaction(selector); + } + + @Override + SpannerRpc.StreamingCall start( + SpannerRpc rpc, + SpannerRpc.ResultStreamConsumer consumer, + @Nullable Map channelHint, + XGoogSpannerRequestId requestId, + boolean routeToLeader) { + return rpc.executeQuery(builder.build(), consumer, channelHint, requestId, routeToLeader); + } + } + + static final class Read extends ReadOrQueryRequest { + private final ReadRequest.Builder builder; + + private Read(ReadRequest.Builder builder) { + super(checkNotNull(builder).getTransaction().hasBegin()); + this.builder = builder; + } + + @Override + String getOperationName() { + return SpannerImpl.READ; + } + + @Override + RetrySettings getRetrySettings(SpannerRpc rpc) { + return rpc.getReadRetrySettings(); + } + + @Override + Set getRetryableCodes(SpannerRpc rpc) { + return rpc.getReadRetryableCodes(); + } + + @Override + @Nullable Statement getStatement() { + return null; + } + + @Override + void setSession(String session) { + builder.setSession(session); + } + + @Override + void setPartitionToken(ByteString partitionToken) { + builder.setPartitionToken(partitionToken); + } + + @Override + void setDataBoostEnabled(boolean dataBoostEnabled) { + builder.setDataBoostEnabled(dataBoostEnabled); + } + + @Override + void setDirectedReadOptions(DirectedReadOptions directedReadOptions) { + builder.setDirectedReadOptions(directedReadOptions); + } + + @Override + void setRequestOptions(RequestOptions requestOptions) { + builder.setRequestOptions(requestOptions); + } + + @Override + void setResumeToken(ByteString resumeToken) { + builder.setResumeToken(resumeToken); + } + + @Override + boolean hasTransaction() { + return builder.hasTransaction(); + } + + @Override + void setTransactionInternal(TransactionSelector selector) { + builder.setTransaction(selector); + } + + @Override + SpannerRpc.StreamingCall start( + SpannerRpc rpc, + SpannerRpc.ResultStreamConsumer consumer, + @Nullable Map channelHint, + XGoogSpannerRequestId requestId, + boolean routeToLeader) { + return rpc.read(builder.build(), consumer, channelHint, requestId, routeToLeader); + } + } +} diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java index d6353bfcdb86..ddef5b841b6e 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java @@ -825,36 +825,10 @@ public SpannerException onError( // must be included with the transaction to ensure that any locks that are taken by the // statement are included in the transaction, even if the statement again causes an error // during the retry. + SpannerException exceptionToThrow = e; if (withBeginTransaction) { - if (lastStatement) { - this.transactionIdFuture.setException(e); - } else { - // Simulate an aborted transaction to force a retry with a new transaction. - this.transactionIdFuture.setException( - SpannerExceptionFactory.newSpannerException( - ErrorCode.ABORTED, - "Aborted due to failed initial statement", - SpannerExceptionFactory.createAbortedExceptionWithRetryDelay( - "Aborted due to failed initial statement", e, 0, 1))); - } - } - SpannerException exceptionToThrow; - if (withBeginTransaction - && e.getErrorCode() == ErrorCode.CANCELLED - && e.getMessage().contains(TRANSACTION_CANCELLED_MESSAGE)) { - // If the first statement of a transaction fails because it was invalidated by a later - // transaction, then the transaction should be retried with an explicit BeginTransaction - // RPC. It could be that this occurred because of a previous transaction that timed out or - // was cancelled by the client, but that was sent to Cloud Spanner and that was still active - // on the backend. - exceptionToThrow = - SpannerExceptionFactory.newSpannerException( - ErrorCode.ABORTED, - e.getMessage(), - SpannerExceptionFactory.createAbortedExceptionWithRetryDelay( - "Aborted due to failed initial statement", e, 0, 1)); - } else { - exceptionToThrow = e; + exceptionToThrow = translateInitialStatementError(e); + failTransactionIdFutureForInitialStatement(exceptionToThrow, lastStatement); } if (exceptionToThrow.getErrorCode() == ErrorCode.ABORTED) { @@ -880,6 +854,48 @@ public SpannerException onError( return exceptionToThrow; } + private SpannerException translateInitialStatementError(SpannerException exception) { + if (exception.getErrorCode() == ErrorCode.CANCELLED + && exception.getMessage().contains(TRANSACTION_CANCELLED_MESSAGE)) { + // If the first statement of a transaction fails because it was invalidated by a later + // transaction, then the transaction should be retried with an explicit BeginTransaction + // RPC. It could be that this occurred because of a previous transaction that timed out or + // was cancelled by the client, but that was sent to Cloud Spanner and that was still active + // on the backend. + return SpannerExceptionFactory.newSpannerException( + ErrorCode.ABORTED, + exception.getMessage(), + SpannerExceptionFactory.createAbortedExceptionWithRetryDelay( + "Aborted due to failed initial statement", exception, 0, 1)); + } + return exception; + } + + /** + * Completes {@link #transactionIdFuture} with an exception when the initial statement that + * carried an inlined {@code BeginTransaction} fails. This method is idempotent because a + * synchronous stream-start failure invokes {@link #onStartFailed} immediately and then {@link + * #onError} when {@link GrpcResultSet} catches the thrown exception. + */ + private void failTransactionIdFutureForInitialStatement( + SpannerException exception, boolean lastStatement) { + final SettableApiFuture future = this.transactionIdFuture; + if (future == null || future.isDone()) { + return; + } + if (lastStatement || exception instanceof AbortedException) { + future.setException(exception); + } else { + // Simulate an aborted transaction to force a retry with a new transaction. + future.setException( + SpannerExceptionFactory.newSpannerException( + ErrorCode.ABORTED, + "Aborted due to failed initial statement", + SpannerExceptionFactory.createAbortedExceptionWithRetryDelay( + "Aborted due to failed initial statement", exception, 0, 1))); + } + } + @Override public void onDone(boolean withBeginTransaction) { if (withBeginTransaction @@ -893,6 +909,15 @@ public void onDone(boolean withBeginTransaction) { } } + @Override + void onStartFailed(boolean withBeginTransaction, boolean lastStatement, Throwable throwable) { + if (withBeginTransaction) { + SpannerException exception = + translateInitialStatementError(SpannerExceptionFactory.asSpannerException(throwable)); + failTransactionIdFutureForInitialStatement(exception, lastStatement); + } + } + @Override public void buffer(Mutation mutation) { synchronized (committingLock) { diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncReadOnlyTransactionTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncReadOnlyTransactionTest.java index 7a1791386a17..1af639a7dddf 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncReadOnlyTransactionTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncReadOnlyTransactionTest.java @@ -17,32 +17,50 @@ package com.google.cloud.spanner; import static com.google.cloud.spanner.MockSpannerServiceImpl.NO_EXECUTION_TIME; +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_MULTIPLE_KEY_VALUE_STATEMENT; import static com.google.cloud.spanner.MockSpannerTestUtil.READ_ONE_KEY_VALUE_STATEMENT; +import static com.google.cloud.spanner.MockSpannerTestUtil.READ_TABLE_NAME; 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 org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNull; import static org.junit.Assert.assertSame; import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyBoolean; +import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; import com.google.api.core.ApiFuture; import com.google.cloud.NoCredentials; import com.google.cloud.spanner.MockSpannerServiceImpl.SimulatedExecutionTime; +import com.google.cloud.spanner.spi.v1.GapicSpannerRpc; +import com.google.cloud.spanner.spi.v1.SpannerRpc; import com.google.spanner.v1.BeginTransactionRequest; import com.google.spanner.v1.ExecuteSqlRequest; +import com.google.spanner.v1.ReadRequest; import io.grpc.ManagedChannelBuilder; import io.grpc.Status; import java.util.Arrays; import java.util.Collections; +import java.util.List; +import java.util.Map; import java.util.concurrent.CountDownLatch; +import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; +import org.mockito.AdditionalAnswers; @RunWith(JUnit4.class) public class AsyncReadOnlyTransactionTest extends AbstractAsyncTransactionTest { @@ -711,6 +729,7 @@ public void closeAsyncWithInlinedBeginTransaction() throws Exception { AsyncResultSet resultSet = transaction.executeQueryAsync(READ_ONE_KEY_VALUE_STATEMENT); final CountDownLatch callbackLatch = new CountDownLatch(1); + final AtomicReference callbackError = new AtomicReference<>(); resultSet.setCallback( executor, ignored -> { @@ -722,6 +741,7 @@ public void closeAsyncWithInlinedBeginTransaction() throws Exception { } return AsyncResultSet.CallbackResponse.CONTINUE; } catch (Throwable throwable) { + callbackError.set(throwable); callbackLatch.countDown(); return AsyncResultSet.CallbackResponse.DONE; } @@ -733,6 +753,339 @@ public void closeAsyncWithInlinedBeginTransaction() throws Exception { mockSpanner.unfreeze(); assertTrue("Timeout waiting for callback", callbackLatch.await(10, TimeUnit.SECONDS)); + assertNull(callbackError.get()); + assertNotNull(transaction.getReadTimestamp()); + closeFuture.get(10, TimeUnit.SECONDS); + assertTrue(closeFuture.isDone()); + } + + private SpannerOptions createSpannerOptionsWithFailingStart( + AtomicBoolean failNextStreamingQuery, + AtomicBoolean failNextStreamingRead, + Runnable beforeReadFailure) { + String endpoint = address.getHostString() + ":" + server.getPort(); + return SpannerOptions.newBuilder() + .setProjectId(TEST_PROJECT) + .setChannelConfigurator(ManagedChannelBuilder::usePlaintext) + .setHost("http://" + endpoint) + .setCredentials(NoCredentials.getInstance()) + .setServiceRpcFactory( + spannerOptions -> { + SpannerRpc delegate = new GapicSpannerRpc(spannerOptions); + SpannerRpc spannerRpc = + mock(SpannerRpc.class, AdditionalAnswers.delegatesTo(delegate)); + doAnswer( + invocation -> { + if (failNextStreamingQuery.compareAndSet(true, false)) { + throw new IllegalStateException("unexpected runtime start failure"); + } + Map rpcOptions = invocation.getArgument(2); + return delegate.executeQuery( + invocation.getArgument(0), + invocation.getArgument(1), + rpcOptions, + invocation.getArgument(3), + invocation.getArgument(4)); + }) + .when(spannerRpc) + .executeQuery( + any(ExecuteSqlRequest.class), + any(SpannerRpc.ResultStreamConsumer.class), + any(), + any(XGoogSpannerRequestId.class), + anyBoolean()); + doAnswer( + invocation -> { + if (failNextStreamingRead.compareAndSet(true, false)) { + beforeReadFailure.run(); + throw new IllegalStateException("unexpected runtime read start failure"); + } + Map rpcOptions = invocation.getArgument(2); + return delegate.read( + invocation.getArgument(0), + invocation.getArgument(1), + rpcOptions, + invocation.getArgument(3), + invocation.getArgument(4)); + }) + .when(spannerRpc) + .read( + any(ReadRequest.class), + any(SpannerRpc.ResultStreamConsumer.class), + any(), + any(XGoogSpannerRequestId.class), + anyBoolean()); + return spannerRpc; + }) + .build(); + } + + @Test(timeout = 10000) + public void inlinedBeginStreamStartFailureCompletesTransactionIdFutureWithError() + throws Exception { + AtomicBoolean failNextStreamingQuery = new AtomicBoolean(false); + AtomicBoolean failNextStreamingRead = new AtomicBoolean(false); + CountDownLatch readStartEnteredLatch = new CountDownLatch(1); + CountDownLatch followerStartedLatch = new CountDownLatch(1); + AtomicReference followerThread = new AtomicReference<>(); + AtomicBoolean followerFinished = new AtomicBoolean(false); + SpannerOptions options = + createSpannerOptionsWithFailingStart( + failNextStreamingQuery, + failNextStreamingRead, + () -> { + readStartEnteredLatch.countDown(); + try { + assertTrue(followerStartedLatch.await(5, TimeUnit.SECONDS)); + } catch (InterruptedException interruptedException) { + Thread.currentThread().interrupt(); + throw new RuntimeException(interruptedException); + } + while (followerThread.get().getState() != Thread.State.TIMED_WAITING + && !followerFinished.get()) { + Thread.yield(); + } + }); + + try (Spanner customSpanner = options.getService()) { + DatabaseClient customClient = + customSpanner.getDatabaseClient( + DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); + try (ResultSet warmupResultSet = + customClient.singleUse().executeQuery(READ_ONE_KEY_VALUE_STATEMENT)) { + while (warmupResultSet.next()) {} + } + failNextStreamingQuery.set(true); + ReadOnlyTransaction transaction = + customClient.readOnlyTransaction( + TimestampBound.strong(), + Options.beginTransactionOption(Options.BeginTransactionOption.INLINE)); + + try (ResultSet resultSet = transaction.executeQuery(READ_ONE_KEY_VALUE_STATEMENT)) { + SpannerException firstException = assertThrows(SpannerException.class, resultSet::next); + assertEquals(ErrorCode.UNKNOWN, firstException.getErrorCode()); + + // Even before the first ResultSet is closed, a subsequent query on the same transaction + // must fail immediately with a SpannerException wrapping the start error rather than + // hanging waiting on transactionIdFuture. + try (ResultSet secondResultSet = transaction.executeQuery(READ_ONE_KEY_VALUE_STATEMENT)) { + SpannerException exception = assertThrows(SpannerException.class, secondResultSet::next); + assertEquals(ErrorCode.UNKNOWN, exception.getErrorCode()); + assertTrue(exception.getMessage().contains("unexpected runtime start failure")); + } + + // Calling initFallbackTransaction() (used for multiplexed session partitionRead / + // partitionQuery fallback, which is the other consumer of transactionIdFuture in + // MultiUseReadOnlyTransaction) must also convert the raw Throwable in transactionIdFuture + // to a SpannerException via asSpannerException. + SpannerException initException = + assertThrows( + SpannerException.class, + () -> ((MultiUseReadOnlyTransaction) transaction).initFallbackTransaction()); + assertEquals(ErrorCode.UNKNOWN, initException.getErrorCode()); + assertTrue(initException.getMessage().contains("unexpected runtime start failure")); + } + + ApiFuture closeFuture = transaction.closeAsync(); + closeFuture.get(10, TimeUnit.SECONDS); + assertTrue(closeFuture.isDone()); + + // Also verify an inlined-begin Read failing synchronously while a concurrent follower query + // is actively blocked waiting on transactionIdFuture. + failNextStreamingRead.set(true); + ReadOnlyTransaction readTransaction = + customClient.readOnlyTransaction( + TimestampBound.strong(), + Options.beginTransactionOption(Options.BeginTransactionOption.INLINE)); + Future followerFailureFuture = + executor.submit( + () -> { + try { + assertTrue(readStartEnteredLatch.await(5, TimeUnit.SECONDS)); + followerThread.set(Thread.currentThread()); + try (ResultSet followerResultSet = + readTransaction.executeQuery(READ_ONE_KEY_VALUE_STATEMENT)) { + followerStartedLatch.countDown(); + return assertThrows(SpannerException.class, followerResultSet::next); + } + } finally { + followerFinished.set(true); + } + }); + try (ResultSet readResultSet = + readTransaction.read(READ_TABLE_NAME, KeySet.all(), READ_COLUMN_NAMES)) { + SpannerException leaderException = + assertThrows(SpannerException.class, readResultSet::next); + assertEquals(ErrorCode.UNKNOWN, leaderException.getErrorCode()); + assertTrue(leaderException.getMessage().contains("unexpected runtime read start failure")); + } + + SpannerException followerException = followerFailureFuture.get(5, TimeUnit.SECONDS); + assertEquals(ErrorCode.UNKNOWN, followerException.getErrorCode()); + assertTrue(followerException.getMessage().contains("unexpected runtime read start failure")); + readTransaction.closeAsync().get(10, TimeUnit.SECONDS); + + // When a MultiUseReadOnlyTransaction does not use inlined begin (withBeginTransaction == + // false), or when a subsequent statement fails synchronously after the inlined-begin + // statement has already succeeded, the transaction remains usable for subsequent queries. + try (ReadOnlyTransaction secondTransaction = + customClient.readOnlyTransaction( + TimestampBound.strong(), + Options.beginTransactionOption(Options.BeginTransactionOption.INLINE))) { + try (ResultSet firstResultSet = + secondTransaction.executeQuery(READ_ONE_KEY_VALUE_STATEMENT)) { + while (firstResultSet.next()) {} + } + assertNotNull(secondTransaction.getReadTimestamp()); + + failNextStreamingQuery.set(true); + try (ResultSet failingResultSet = + secondTransaction.executeQuery(READ_ONE_KEY_VALUE_STATEMENT)) { + SpannerException exception = assertThrows(SpannerException.class, failingResultSet::next); + assertEquals(ErrorCode.UNKNOWN, exception.getErrorCode()); + } + try (ResultSet subsequentResultSet = + secondTransaction.executeQuery(READ_ONE_KEY_VALUE_STATEMENT)) { + while (subsequentResultSet.next()) {} + } + } + } + } + + @Test + public void inlinedBeginReadOnlyTransactionStreamRetriesWithMockSpanner() { + try (ResultSet warmupResultSet = + client().singleUse().executeQuery(READ_ONE_KEY_VALUE_STATEMENT)) { + while (warmupResultSet.next()) {} + } + + // 1. Stream retry before receiving a resume token (errorIndex == 0): both requests must carry + // BeginTransaction(ReadOnly). + mockSpanner.clearRequests(); + mockSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.ofStreamException(Status.UNAVAILABLE.asRuntimeException(), 0)); + try (ReadOnlyTransaction transaction = + client() + .readOnlyTransaction( + TimestampBound.strong(), + Options.beginTransactionOption(Options.BeginTransactionOption.INLINE))) { + int rowCount = 0; + try (ResultSet resultSet = transaction.executeQuery(READ_MULTIPLE_KEY_VALUE_STATEMENT)) { + while (resultSet.next()) { + rowCount++; + } + } + assertEquals(3, rowCount); + assertNotNull(transaction.getReadTimestamp()); + + try (ResultSet readResultSet = + transaction.read(READ_TABLE_NAME, KeySet.all(), READ_COLUMN_NAMES)) { + int readRowCount = 0; + while (readResultSet.next()) { + readRowCount++; + } + assertEquals(3, readRowCount); + } + } + assertEquals(0, mockSpanner.countRequestsOfType(BeginTransactionRequest.class)); + List firstRunSqlRequests = + mockSpanner.getRequestsOfType(ExecuteSqlRequest.class); + assertEquals(2, firstRunSqlRequests.size()); + assertTrue(firstRunSqlRequests.get(0).getTransaction().getBegin().hasReadOnly()); + assertTrue(firstRunSqlRequests.get(0).getResumeToken().isEmpty()); + assertTrue(firstRunSqlRequests.get(1).getTransaction().getBegin().hasReadOnly()); + assertTrue(firstRunSqlRequests.get(1).getResumeToken().isEmpty()); + List firstRunReadRequests = mockSpanner.getRequestsOfType(ReadRequest.class); + assertEquals(1, firstRunReadRequests.size()); + assertFalse(firstRunReadRequests.get(0).getTransaction().hasBegin()); + assertFalse(firstRunReadRequests.get(0).getTransaction().getId().isEmpty()); + + // 2. Stream retry on executeQuery after receiving a resume token (errorIndex == 1): the first + // ExecuteSqlRequest carries BeginTransaction(ReadOnly), while the resumed request switches to + // the allocated transaction ID so it does not start a second read-only transaction. + mockSpanner.clearRequests(); + mockSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.ofStreamException(Status.UNAVAILABLE.asRuntimeException(), 1)); + try (ReadOnlyTransaction transaction = + client() + .readOnlyTransaction( + TimestampBound.strong(), + Options.beginTransactionOption(Options.BeginTransactionOption.INLINE))) { + int rowCount = 0; + try (ResultSet resultSet = transaction.executeQuery(READ_MULTIPLE_KEY_VALUE_STATEMENT)) { + while (resultSet.next()) { + rowCount++; + } + } + assertEquals(3, rowCount); + assertNotNull(transaction.getReadTimestamp()); + } + assertEquals(0, mockSpanner.countRequestsOfType(BeginTransactionRequest.class)); + List secondRunSqlRequests = + mockSpanner.getRequestsOfType(ExecuteSqlRequest.class); + assertEquals(2, secondRunSqlRequests.size()); + assertTrue(secondRunSqlRequests.get(0).getTransaction().getBegin().hasReadOnly()); + assertTrue(secondRunSqlRequests.get(0).getResumeToken().isEmpty()); + assertFalse(secondRunSqlRequests.get(1).getTransaction().hasBegin()); + assertFalse(secondRunSqlRequests.get(1).getTransaction().getId().isEmpty()); + assertFalse(secondRunSqlRequests.get(1).getResumeToken().isEmpty()); + + // 3. Stream retry on read after receiving a resume token (errorIndex == 1): the first + // ReadRequest carries BeginTransaction(ReadOnly), while the resumed ReadRequest switches to + // the allocated transaction ID. + mockSpanner.clearRequests(); + mockSpanner.setStreamingReadExecutionTime( + SimulatedExecutionTime.ofStreamException(Status.UNAVAILABLE.asRuntimeException(), 1)); + try (ReadOnlyTransaction transaction = + client() + .readOnlyTransaction( + TimestampBound.strong(), + Options.beginTransactionOption(Options.BeginTransactionOption.INLINE))) { + int rowCount = 0; + try (ResultSet resultSet = + transaction.read(READ_TABLE_NAME, KeySet.all(), READ_COLUMN_NAMES)) { + while (resultSet.next()) { + rowCount++; + } + } + assertEquals(3, rowCount); + assertNotNull(transaction.getReadTimestamp()); + } + assertEquals(0, mockSpanner.countRequestsOfType(BeginTransactionRequest.class)); + List thirdRunReadRequests = mockSpanner.getRequestsOfType(ReadRequest.class); + assertEquals(2, thirdRunReadRequests.size()); + assertTrue(thirdRunReadRequests.get(0).getTransaction().getBegin().hasReadOnly()); + assertTrue(thirdRunReadRequests.get(0).getResumeToken().isEmpty()); + assertFalse(thirdRunReadRequests.get(1).getTransaction().hasBegin()); + assertFalse(thirdRunReadRequests.get(1).getTransaction().getId().isEmpty()); + assertFalse(thirdRunReadRequests.get(1).getResumeToken().isEmpty()); + } + + @Test(timeout = 10000) + public void inlinedBeginReadOnlyTransactionServerErrorCompletesTransactionIdFuture() + throws Exception { + try (ResultSet warmupResultSet = + client().singleUse().executeQuery(READ_ONE_KEY_VALUE_STATEMENT)) { + while (warmupResultSet.next()) {} + } + mockSpanner.clearRequests(); + + ReadOnlyTransaction transaction = + client() + .readOnlyTransaction( + TimestampBound.strong(), + Options.beginTransactionOption(Options.BeginTransactionOption.INLINE)); + try (ResultSet resultSet = transaction.executeQuery(INVALID_UPDATE_STATEMENT)) { + SpannerException firstException = assertThrows(SpannerException.class, resultSet::next); + assertEquals(ErrorCode.INVALID_ARGUMENT, firstException.getErrorCode()); + + try (ResultSet secondResultSet = transaction.executeQuery(READ_ONE_KEY_VALUE_STATEMENT)) { + SpannerException secondException = + assertThrows(SpannerException.class, secondResultSet::next); + assertEquals(ErrorCode.INVALID_ARGUMENT, secondException.getErrorCode()); + } + } + ApiFuture closeFuture = transaction.closeAsync(); closeFuture.get(10, TimeUnit.SECONDS); assertTrue(closeFuture.isDone()); } diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java index 8acb0a7b725b..0187f94218d2 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java @@ -2376,7 +2376,8 @@ public void testBackendPartitionQueryOptions() { .setOptimizerStatisticsPackage("custom-package") .build()) .build(), - Options.directedRead(DIRECTED_READ_OPTIONS1)); + Options.directedRead(DIRECTED_READ_OPTIONS1), + Options.dataBoostEnabled(true)); try (ResultSet rs = transaction.execute(partitions.get(0))) { // Just iterate over the results to execute the query. consumeResults(rs); @@ -2395,6 +2396,8 @@ public void testBackendPartitionQueryOptions() { assertThat(executeSqlRequest.getQueryOptions().getOptimizerStatisticsPackage()) .isEqualTo("custom-package"); assertThat(executeSqlRequest.getDirectedReadOptions()).isEqualTo(DIRECTED_READ_OPTIONS1); + assertTrue(executeSqlRequest.getDataBoostEnabled()); + assertEquals(partitions.get(0).getPartitionToken(), executeSqlRequest.getPartitionToken()); } } @@ -2445,6 +2448,8 @@ public void testBackendPartitionQueryOptions() { assertThat(executeSqlRequest.getQueryOptions().getOptimizerStatisticsPackage()) .isEqualTo("custom-package"); assertThat(executeSqlRequest.getDirectedReadOptions()).isEqualTo(DIRECTED_READ_OPTIONS2); + assertFalse(executeSqlRequest.getDataBoostEnabled()); + assertEquals(partitions.get(0).getPartitionToken(), executeSqlRequest.getPartitionToken()); } } @@ -2472,7 +2477,8 @@ public void testBackendPartitionReadOptions() { "FOO", KeySet.all(), Lists.newArrayList("1"), - Options.directedRead(DIRECTED_READ_OPTIONS1)); + Options.directedRead(DIRECTED_READ_OPTIONS1), + Options.dataBoostEnabled(true)); try (ResultSet rs = transaction.execute(partitions.get(0))) { // Just iterate over the results to execute the query. consumeResults(rs); @@ -2487,6 +2493,8 @@ public void testBackendPartitionReadOptions() { assertThat(requests.get(requests.size() - 1)).isInstanceOf(ReadRequest.class); ReadRequest readRequest = (ReadRequest) requests.get(requests.size() - 1); assertThat(readRequest.getDirectedReadOptions()).isEqualTo(DIRECTED_READ_OPTIONS1); + assertTrue(readRequest.getDataBoostEnabled()); + assertEquals(partitions.get(0).getPartitionToken(), readRequest.getPartitionToken()); } } @@ -2530,6 +2538,8 @@ public void testBackendPartitionReadOptions() { assertThat(requests.get(requests.size() - 1)).isInstanceOf(ReadRequest.class); ReadRequest readRequest = (ReadRequest) requests.get(requests.size() - 1); assertThat(readRequest.getDirectedReadOptions()).isEqualTo(DIRECTED_READ_OPTIONS2); + assertFalse(readRequest.getDataBoostEnabled()); + assertEquals(partitions.get(0).getPartitionToken(), readRequest.getPartitionToken()); } } diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/GrpcResultSetTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/GrpcResultSetTest.java index 85f94ebc6d74..b1972beed96d 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/GrpcResultSetTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/GrpcResultSetTest.java @@ -105,8 +105,7 @@ public void cancel(@Nullable String message) {} @Override public void request(int numMessages) {} - }, - false); + }); consumer = stream.consumer(); resultSet = new GrpcResultSet(stream, new NoOpListener()); } diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/GrpcStreamIteratorTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/GrpcStreamIteratorTest.java index 72d03d533152..b872132aee97 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/GrpcStreamIteratorTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/GrpcStreamIteratorTest.java @@ -69,7 +69,7 @@ public void closeBeforeSetCallCancelsUnderlyingCallWithMessage() { /* cancelQueryWhenClientIsClosed= */ false); streamIterator.close("Cancelled before setCall"); - streamIterator.setCall(call, /* withBeginTransaction= */ false); + streamIterator.setCall(call); verify(call).cancel("Cancelled before setCall"); } @@ -86,7 +86,7 @@ public void closeAfterSetCallCancelsUnderlyingCallWithMessage() { /* prefetchChunks= */ 10, /* cancelQueryWhenClientIsClosed= */ false); - streamIterator.setCall(call, /* withBeginTransaction= */ false); + streamIterator.setCall(call); streamIterator.close("Cancelled after setCall"); verify(call).cancel("Cancelled after setCall"); @@ -104,7 +104,7 @@ public void closeAfterLastTrueChunkConsumedDoesNotCancelCall() { /* prefetchChunks= */ 10, /* cancelQueryWhenClientIsClosed= */ false); - streamIterator.setCall(call, /* withBeginTransaction= */ false); + streamIterator.setCall(call); PartialResultSet lastChunk = PartialResultSet.newBuilder() @@ -134,7 +134,7 @@ public void closeBeforeLastTrueChunkConsumedCancelsCall() { /* prefetchChunks= */ 10, /* cancelQueryWhenClientIsClosed= */ false); - streamIterator.setCall(call, /* withBeginTransaction= */ false); + streamIterator.setCall(call); PartialResultSet lastChunk = PartialResultSet.newBuilder() @@ -176,7 +176,7 @@ public void lastTrueConsumedBeforeSetCallDoesNotCancelWhenClosed() { if (closeBeforeSetCall) { streamIterator.close("Closed before setCall"); } - streamIterator.setCall(call, /* withBeginTransaction= */ false); + streamIterator.setCall(call); if (!closeBeforeSetCall) { streamIterator.close("Closed after setCall"); } @@ -198,7 +198,7 @@ public void consumerPrefetchChunksOneDoesNotDeadlockOnCompletion() { SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); - streamIterator.setCall(call, /* withBeginTransaction= */ false); + streamIterator.setCall(call); PartialResultSet lastChunk = PartialResultSet.newBuilder() @@ -228,7 +228,7 @@ public void consumerPrefetchChunksOneDoesNotDeadlockOnError() { SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); - streamIterator.setCall(call, /* withBeginTransaction= */ false); + streamIterator.setCall(call); PartialResultSet lastChunk = PartialResultSet.newBuilder() @@ -261,7 +261,7 @@ public void requestPrefetchChunksValidatesInitialization() { SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); - streamIterator.setCall(call, /* withBeginTransaction= */ false); + streamIterator.setCall(call); streamIterator.requestPrefetchChunks(); verify(call).request(5); @@ -274,10 +274,14 @@ public void statePreservedAcrossClose() { /* statement= */ null, /* lastStatement= */ true, /* prefetchChunks= */ 2, - /* cancelQueryWhenClientIsClosed= */ false); + /* cancelQueryWhenClientIsClosed= */ false, + /* withBeginTransaction= */ true); + assertTrue(streamIterator.isLastStatement()); + assertTrue(streamIterator.isWithBeginTransaction()); + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); - streamIterator.setCall(call, /* withBeginTransaction= */ true); + streamIterator.setCall(call); assertTrue(streamIterator.isLastStatement()); assertTrue(streamIterator.isWithBeginTransaction()); @@ -320,7 +324,7 @@ public void streamWaitTimeoutThrowsDeadlineExceeded() { /* prefetchChunks= */ 2, /* cancelQueryWhenClientIsClosed= */ false); - streamIterator.setCall(call, /* withBeginTransaction= */ false); + streamIterator.setCall(call); SpannerException thrown = assertThrows(SpannerException.class, streamIterator::hasNext); assertEquals(ErrorCode.DEADLINE_EXCEEDED, thrown.getErrorCode()); @@ -339,7 +343,7 @@ public void closeBeforeSetCallDoesNotRetainCall() { /* cancelQueryWhenClientIsClosed= */ false); streamIterator.close("Closed before setCall"); - streamIterator.setCall(call, /* withBeginTransaction= */ false); + streamIterator.setCall(call); verify(call).cancel("Closed before setCall"); // Since streamIterator was already closed, requestPrefetchChunks should safely no-op @@ -360,7 +364,7 @@ public void closeAfterSetCallClearsCallReference() { /* prefetchChunks= */ 5, /* cancelQueryWhenClientIsClosed= */ false); - streamIterator.setCall(call, /* withBeginTransaction= */ false); + streamIterator.setCall(call); streamIterator.close("Closed after setCall"); verify(call).cancel("Closed after setCall"); @@ -391,7 +395,7 @@ public void concurrentCloseAndSetCallCancelsCallDeterministically() throws Excep () -> { try { startLatch.await(); - streamIterator.setCall(call, /* withBeginTransaction= */ false); + streamIterator.setCall(call); } catch (InterruptedException ignored) { } finally { doneLatch.countDown(); @@ -432,7 +436,7 @@ public void concurrentCloseUnblocksConsumerWaitingOnStream() throws Exception { /* lastStatement= */ false, /* prefetchChunks= */ 5, /* cancelQueryWhenClientIsClosed= */ false); - streamIterator.setCall(call, /* withBeginTransaction= */ false); + streamIterator.setCall(call); AtomicReference consumerError = new AtomicReference<>(); Thread consumerThread = @@ -479,7 +483,7 @@ public void setCallWithNullCallContextDoesNotThrow() { /* lastStatement= */ false, /* prefetchChunks= */ 5, /* cancelQueryWhenClientIsClosed= */ false); - streamIterator.setCall(call, /* withBeginTransaction= */ false); + streamIterator.setCall(call); streamIterator.requestPrefetchChunks(); verify(call).request(5); @@ -498,7 +502,7 @@ public void streamWaitTimeoutWithSecondsUnitThrowsDeadlineExceeded() { /* lastStatement= */ false, /* prefetchChunks= */ 5, /* cancelQueryWhenClientIsClosed= */ false); - streamIterator.setCall(call, /* withBeginTransaction= */ false); + streamIterator.setCall(call); SpannerException spannerException = assertThrows(SpannerException.class, streamIterator::hasNext); @@ -516,7 +520,7 @@ public void duplicateCloseIsIdempotent() { /* lastStatement= */ false, /* prefetchChunks= */ 5, /* cancelQueryWhenClientIsClosed= */ false); - streamIterator.setCall(call, /* withBeginTransaction= */ false); + streamIterator.setCall(call); streamIterator.close("First close"); streamIterator.close("Second close"); @@ -536,7 +540,7 @@ public void requestPrefetchChunksSafelyNoOpsWhenCallCleared() { /* lastStatement= */ false, /* prefetchChunks= */ 5, /* cancelQueryWhenClientIsClosed= */ false); - streamIterator.setCall(call, /* withBeginTransaction= */ false); + streamIterator.setCall(call); streamIterator .consumer() @@ -575,7 +579,7 @@ public void concurrentSetCallAndRequestPrefetchChunks() throws Exception { () -> { try { startLatch.await(); - streamIterator.setCall(call, /* withBeginTransaction= */ false); + streamIterator.setCall(call); } catch (Throwable t) { error.set(t); } finally { @@ -701,7 +705,7 @@ public void closeOnAlreadyCompletedStreamDoesNotDeliverSpuriousError() { /* lastStatement= */ false, /* prefetchChunks= */ 1, /* cancelQueryWhenClientIsClosed= */ false); - streamIterator.setCall(call, /* withBeginTransaction= */ false); + streamIterator.setCall(call); streamIterator .consumer() @@ -776,7 +780,7 @@ public void interruptedTakeNextFromStreamThrowsPropagateInterrupt() throws Excep /* lastStatement= */ false, /* prefetchChunks= */ 5, /* cancelQueryWhenClientIsClosed= */ false); - streamIterator.setCall(call, /* withBeginTransaction= */ false); + streamIterator.setCall(call); AtomicReference thrown = new AtomicReference<>(); Thread consumerThread = @@ -818,7 +822,7 @@ public void dataAvailableAndCanPollWithoutErrorOrBlockingEmptyStream() { SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); - streamIterator.setCall(call, /* withBeginTransaction= */ false); + streamIterator.setCall(call); assertFalse(streamIterator.canPollWithoutErrorOrBlocking()); assertFalse(streamIterator.isDataAvailable()); @@ -835,7 +839,7 @@ public void dataAvailableAndCanPollWithoutErrorOrBlockingWithBufferedChunksAndCo SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); - streamIterator.setCall(call, /* withBeginTransaction= */ false); + streamIterator.setCall(call); PartialResultSet chunk1 = PartialResultSet.newBuilder() @@ -876,7 +880,7 @@ public void dataAvailableAndCanPollWithoutErrorOrBlockingWhenEndOfStreamReached( SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); - streamIterator.setCall(call, /* withBeginTransaction= */ false); + streamIterator.setCall(call); PartialResultSet lastChunk = PartialResultSet.newBuilder() @@ -912,7 +916,7 @@ public void dataAvailableAndCanPollWithoutErrorOrBlockingWhenStreamErrorOccurs() SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); - streamIterator.setCall(call, /* withBeginTransaction= */ false); + streamIterator.setCall(call); SpannerException spannerException = SpannerExceptionFactory.newSpannerException(ErrorCode.UNAVAILABLE, "network unavailable"); @@ -962,7 +966,7 @@ public void listenerRegisteredBeforeSetCallDelaysEventDeliveryUntilCallInitializ // Now initialize call via setCall. SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); - streamIterator.setCall(call, /* withBeginTransaction= */ false); + streamIterator.setCall(call); // Flushing pending message in setCall notifies the listener. assertEquals(chunk, receivedMessage.get()); @@ -982,7 +986,7 @@ public void listenerRegisteredAfterSetCallReceivesMessagesImmediately() { SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); - streamIterator.setCall(call, /* withBeginTransaction= */ false); + streamIterator.setCall(call); AtomicReference receivedMessage = new AtomicReference<>(); AtomicBoolean receivedBlockedState = new AtomicBoolean(); @@ -1033,7 +1037,7 @@ public void dataAvailableAndCanPollWithoutErrorOrBlockingWhenStreamCompletesWith SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); - streamIterator.setCall(call, /* withBeginTransaction= */ false); + streamIterator.setCall(call); // Initial state: No messages received yet. Polling would block on network I/O. assertFalse(streamIterator.canPollWithoutErrorOrBlocking()); @@ -1066,7 +1070,7 @@ public void dataAvailableAndCanPollWithoutErrorOrBlockingWhenAllDataConsumedWith SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); - streamIterator.setCall(call, /* withBeginTransaction= */ false); + streamIterator.setCall(call); PartialResultSet chunk = PartialResultSet.newBuilder() @@ -1109,7 +1113,7 @@ public void dataAvailableWhenServerReturnsErrorAfterDrainingBuffer() { SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); - streamIterator.setCall(call, /* withBeginTransaction= */ false); + streamIterator.setCall(call); PartialResultSet chunk = PartialResultSet.newBuilder() @@ -1174,7 +1178,7 @@ public void concurrentSetCallAndOnPartialResultSet() throws Exception { () -> { try { startLatch.await(); - streamIterator.setCall(call, /* withBeginTransaction= */ false); + streamIterator.setCall(call); } catch (Throwable throwable) { threadError.set(throwable); } finally { @@ -1239,7 +1243,7 @@ public void concurrentSetCallAndOnError() throws Exception { () -> { try { startLatch.await(); - streamIterator.setCall(call, /* withBeginTransaction= */ false); + streamIterator.setCall(call); } catch (Throwable throwable) { threadError.set(throwable); } finally { @@ -1298,7 +1302,7 @@ public void concurrentSetCallAndOnCompleted() throws Exception { () -> { try { startLatch.await(); - streamIterator.setCall(call, /* withBeginTransaction= */ false); + streamIterator.setCall(call); } catch (Throwable throwable) { threadError.set(throwable); } finally { @@ -1394,7 +1398,7 @@ public void listenerRegisteredBeforeSetCallFlushesAllQueuedMessages() { SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); - streamIterator.setCall(call, /* withBeginTransaction= */ false); + streamIterator.setCall(call); verify(listener, times(1)).onStreamMessage(chunk1, false); verify(listener, times(1)).onStreamMessage(chunk2, false); @@ -1413,7 +1417,7 @@ public void listenerRegisteredAfterSetCallFlushesAllQueuedMessages() { SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); - streamIterator.setCall(call, /* withBeginTransaction= */ false); + streamIterator.setCall(call); PartialResultSet chunk1 = PartialResultSet.newBuilder().setResumeToken(ByteString.copyFromUtf8("token-1")).build(); @@ -1468,7 +1472,7 @@ public void listenerRegisteredBeforeSetCallDelaysErrorUntilSetCall() { SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); - streamIterator.setCall(call, /* withBeginTransaction= */ false); + streamIterator.setCall(call); verify(listener, times(1)).onStreamMessage(GrpcStreamIterator.END_OF_STREAM, false); assertFalse(streamIterator.canPollWithoutErrorOrBlocking()); @@ -1499,7 +1503,7 @@ public void listenerRegisteredBeforeSetCallDelaysCompletionUntilSetCall() { SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); - streamIterator.setCall(call, /* withBeginTransaction= */ false); + streamIterator.setCall(call); verify(listener, times(1)).onStreamMessage(GrpcStreamIterator.END_OF_STREAM, false); assertTrue(streamIterator.canPollWithoutErrorOrBlocking()); @@ -1518,7 +1522,7 @@ public void dataAvailableAndCanPollWithoutErrorOrBlockingWithBufferedChunksBefor SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); - streamIterator.setCall(call, /* withBeginTransaction= */ false); + streamIterator.setCall(call); PartialResultSet chunk = PartialResultSet.newBuilder() diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MockSpannerServiceImpl.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MockSpannerServiceImpl.java index a598e64045fa..0134c8f0f892 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MockSpannerServiceImpl.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MockSpannerServiceImpl.java @@ -1840,7 +1840,10 @@ private void returnPartialResultSet( .setTransaction( ignoreInlineBeginRequest.get() ? Transaction.getDefaultInstance() - : Transaction.newBuilder().setId(transactionId).build()) + : (transactionSelector.hasBegin() + ? transactions.getOrDefault( + transactionId, Transaction.newBuilder().setId(transactionId).build()) + : Transaction.newBuilder().setId(transactionId).build())) .build(); } resultSet = resultSet.toBuilder().setMetadata(metadata).build(); diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadFormatTestRunner.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadFormatTestRunner.java index ff26f774b4b0..3081eea2f00a 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadFormatTestRunner.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadFormatTestRunner.java @@ -134,8 +134,7 @@ public void cancel(@Nullable String message) {} @Override public void request(int numMessages) {} - }, - false); + }); consumer = stream.consumer(); resultSet = new GrpcResultSet(stream, new NoOpListener()); diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadWriteTransactionWithInlineBeginTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadWriteTransactionWithInlineBeginTest.java index 492252d486cb..b5591d3072a8 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadWriteTransactionWithInlineBeginTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadWriteTransactionWithInlineBeginTest.java @@ -18,15 +18,29 @@ import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotEquals; import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyBoolean; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mock; +import com.google.api.core.SettableApiFuture; import com.google.api.gax.grpc.testing.LocalChannelProvider; +import com.google.api.gax.rpc.ApiCallContext; import com.google.cloud.NoCredentials; +import com.google.cloud.spanner.MockSpannerServiceImpl.SimulatedExecutionTime; import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult; +import com.google.cloud.spanner.spi.v1.GapicSpannerRpc; +import com.google.cloud.spanner.spi.v1.SpannerRpc; import com.google.common.collect.Iterables; import com.google.protobuf.AbstractMessage; +import com.google.protobuf.ByteString; import com.google.protobuf.ListValue; import com.google.spanner.v1.BeginTransactionRequest; +import com.google.spanner.v1.CommitRequest; import com.google.spanner.v1.ExecuteSqlRequest; import com.google.spanner.v1.ReadRequest; import com.google.spanner.v1.ResultSetMetadata; @@ -35,6 +49,7 @@ import com.google.spanner.v1.TransactionOptions; import com.google.spanner.v1.TransactionOptions.ReadWrite; import com.google.spanner.v1.TypeCode; +import io.grpc.Context; import io.grpc.Server; import io.grpc.Status; import io.grpc.inprocess.InProcessServerBuilder; @@ -43,11 +58,20 @@ import java.util.Collection; import java.util.Collections; import java.util.List; +import java.util.Map; +import java.util.concurrent.CancellationException; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.Future; import java.util.concurrent.ScheduledThreadPoolExecutor; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.BiFunction; +import java.util.function.Supplier; import java.util.stream.Collectors; import org.junit.After; import org.junit.AfterClass; @@ -56,6 +80,8 @@ import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; +import org.mockito.AdditionalAnswers; +import org.mockito.invocation.InvocationOnMock; @RunWith(JUnit4.class) public class ReadWriteTransactionWithInlineBeginTest { @@ -70,6 +96,8 @@ public class ReadWriteTransactionWithInlineBeginTest { Statement.of("SELECT * FROM NON_EXISTENT_TABLE"); private static final long UPDATE_COUNT = 1L; private static final Statement SELECT1 = Statement.of("SELECT 1 AS COL1"); + private static final Statement SELECT1_UNION_ALL_SELECT2 = + Statement.of("SELECT 1 AS COL1 UNION ALL SELECT 2 AS COL1"); private static final ResultSetMetadata SELECT1_METADATA = ResultSetMetadata.newBuilder() .setRowType( @@ -92,6 +120,18 @@ public class ReadWriteTransactionWithInlineBeginTest { .build()) .setMetadata(SELECT1_METADATA) .build(); + private static final com.google.spanner.v1.ResultSet SELECT1_UNION_ALL_SELECT2_RESULTSET = + com.google.spanner.v1.ResultSet.newBuilder() + .addRows( + ListValue.newBuilder() + .addValues(com.google.protobuf.Value.newBuilder().setStringValue("1").build()) + .build()) + .addRows( + ListValue.newBuilder() + .addValues(com.google.protobuf.Value.newBuilder().setStringValue("2").build()) + .build()) + .setMetadata(SELECT1_METADATA) + .build(); private static final TransactionOptions OPTIMISTIC_LOCK_OPTIONS = TransactionOptions.newBuilder() .setReadWrite(ReadWrite.newBuilder().setReadLockMode(ReadWrite.ReadLockMode.OPTIMISTIC)) @@ -105,6 +145,8 @@ public static void startStaticServer() throws Exception { 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_UNION_ALL_SELECT2, SELECT1_UNION_ALL_SELECT2_RESULTSET)); mockSpanner.putStatementResult( StatementResult.exception( INVALID_UPDATE_STATEMENT, @@ -116,6 +158,12 @@ public static void startStaticServer() throws Exception { mockSpanner.putStatementResult( StatementResult.read( "FOO", KeySet.all(), Collections.singletonList("ID"), SELECT1_RESULTSET)); + mockSpanner.putStatementResult( + StatementResult.read( + "FOO_TWO_ROWS", + KeySet.all(), + Collections.singletonList("ID"), + SELECT1_UNION_ALL_SELECT2_RESULTSET)); String uniqueName = InProcessServerBuilder.generateName(); server = @@ -479,6 +527,870 @@ public void abortedBatchUpdate() { assertThat(countTransactionsStarted()).isEqualTo(2); } + private SpannerOptions createSpannerOptionsWithStreamInterceptors( + BiFunction executeQueryInterceptor, + BiFunction readInterceptor) { + return SpannerOptions.newBuilder() + .setProjectId("[PROJECT]") + .setChannelProvider(channelProvider) + .setCredentials(NoCredentials.getInstance()) + .setServiceRpcFactory( + spannerOptions -> { + SpannerRpc delegate = new GapicSpannerRpc(spannerOptions); + SpannerRpc spannerRpc = + mock(SpannerRpc.class, AdditionalAnswers.delegatesTo(delegate)); + doAnswer(invocation -> executeQueryInterceptor.apply(invocation, delegate)) + .when(spannerRpc) + .executeQuery( + any(ExecuteSqlRequest.class), + any(SpannerRpc.ResultStreamConsumer.class), + any(), + any(XGoogSpannerRequestId.class), + anyBoolean()); + doAnswer(invocation -> readInterceptor.apply(invocation, delegate)) + .when(spannerRpc) + .read( + any(ReadRequest.class), + any(SpannerRpc.ResultStreamConsumer.class), + any(), + any(XGoogSpannerRequestId.class), + anyBoolean()); + return spannerRpc; + }) + .build(); + } + + private static SpannerRpc.StreamingCall delegateExecuteQuery( + InvocationOnMock invocation, SpannerRpc delegate) { + Map rpcOptions = invocation.getArgument(2); + return delegate.executeQuery( + invocation.getArgument(0), + invocation.getArgument(1), + rpcOptions, + invocation.getArgument(3), + invocation.getArgument(4)); + } + + private static SpannerRpc.StreamingCall delegateRead( + InvocationOnMock invocation, SpannerRpc delegate) { + Map rpcOptions = invocation.getArgument(2); + return delegate.read( + invocation.getArgument(0), + invocation.getArgument(1), + rpcOptions, + invocation.getArgument(3), + invocation.getArgument(4)); + } + + private SpannerOptions createSpannerOptionsWithFailingStreamStart( + AtomicBoolean failNextStreamStart) { + return createSpannerOptionsWithFailingStreamStart( + failNextStreamStart, + () -> new IllegalStateException("simulated synchronous start failure")); + } + + private SpannerOptions createSpannerOptionsWithFailingStreamStart( + AtomicBoolean failNextStreamStart, Supplier failureSupplier) { + return createSpannerOptionsWithStreamInterceptors( + (invocation, delegate) -> { + if (failNextStreamStart.compareAndSet(true, false)) { + throw failureSupplier.get(); + } + return delegateExecuteQuery(invocation, delegate); + }, + (invocation, delegate) -> { + if (failNextStreamStart.compareAndSet(true, false)) { + throw failureSupplier.get(); + } + return delegateRead(invocation, delegate); + }); + } + + @Test(timeout = 10_000) + public void synchronousStartFailureOfInlinedBeginStatement() { + AtomicBoolean failNextStreamStart = new AtomicBoolean(); + SpannerOptions options = createSpannerOptionsWithFailingStreamStart(failNextStreamStart); + AtomicInteger attempts = new AtomicInteger(); + try (Spanner customSpanner = options.getService()) { + DatabaseClient customClient = + customSpanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]")); + + // Single-use read context exercises the base AbstractReadContext.onStartFailed no-op hook. + failNextStreamStart.set(true); + try (ResultSet singleUseResultSet = customClient.singleUse().executeQuery(SELECT1)) { + SpannerException singleUseError = + assertThrows(SpannerException.class, singleUseResultSet::next); + assertEquals(ErrorCode.UNKNOWN, singleUseError.getErrorCode()); + } + + mockSpanner.reset(); + customClient + .readWriteTransaction() + .run( + transaction -> { + int attempt = attempts.incrementAndGet(); + if (attempt == 1) { + failNextStreamStart.set(true); + try (ResultSet resultSet = transaction.executeQuery(SELECT1)) { + SpannerException firstStatementError = + assertThrows(SpannerException.class, resultSet::next); + assertEquals(ErrorCode.UNKNOWN, firstStatementError.getErrorCode()); + } + } else { + // On the retry attempt, the transaction has already been started, so a + // synchronous start failure has withBeginTransaction == false and does not + // abort the transaction. + failNextStreamStart.set(true); + try (ResultSet failingResultSet = + transaction.read("FOO", KeySet.all(), Collections.singletonList("ID"))) { + SpannerException subsequentError = + assertThrows(SpannerException.class, failingResultSet::next); + assertEquals(ErrorCode.UNKNOWN, subsequentError.getErrorCode()); + } + } + try (ResultSet resultSet = transaction.executeQuery(SELECT1)) { + while (resultSet.next()) {} + } + return null; + }); + assertEquals(2, attempts.get()); + assertEquals(1, countRequests(BeginTransactionRequest.class)); + assertEquals(1, countRequests(CommitRequest.class)); + + // Also verify an inlined-begin Read failing synchronously where the user catches the error + // and only buffers a mutation; CommitRunnable awaits transactionIdFuture and retries with an + // explicit BeginTransaction RPC. + mockSpanner.reset(); + AtomicInteger readAndBufferAttempts = new AtomicInteger(); + customClient + .readWriteTransaction() + .run( + transaction -> { + if (readAndBufferAttempts.incrementAndGet() == 1) { + failNextStreamStart.set(true); + try (ResultSet resultSet = + transaction.read("FOO", KeySet.all(), Collections.singletonList("ID"))) { + SpannerException readError = + assertThrows(SpannerException.class, resultSet::next); + assertEquals(ErrorCode.UNKNOWN, readError.getErrorCode()); + } + } + transaction.buffer(Mutation.newInsertBuilder("FOO").set("ID").to(1L).build()); + return null; + }); + assertEquals(2, readAndBufferAttempts.get()); + assertEquals(1, countRequests(BeginTransactionRequest.class)); + assertEquals(1, countRequests(CommitRequest.class)); + + // Also verify initiateStreaming() failing synchronously before next() is called: only + // TransactionContextImpl.onStartFailed runs during initiateStreaming(), and it must + // immediately complete transactionIdFuture with a simulated AbortedException so a subsequent + // statement on the transaction fails fast with AbortedException and triggers a retry. + mockSpanner.reset(); + AtomicInteger initiateStreamingAttempts = new AtomicInteger(); + customClient + .readWriteTransaction() + .run( + transaction -> { + if (initiateStreamingAttempts.incrementAndGet() == 1) { + failNextStreamStart.set(true); + try (ResultSet resultSet = transaction.executeQuery(SELECT1)) { + assertThrows( + IllegalStateException.class, + () -> + ((StreamingResultSet) resultSet) + .initiateStreaming((message, prefetchChunks) -> {})); + } + } + assertEquals(1L, transaction.executeUpdate(UPDATE_STATEMENT)); + return null; + }); + assertEquals(2, initiateStreamingAttempts.get()); + assertEquals(1, countRequests(BeginTransactionRequest.class)); + assertEquals(1, countRequests(CommitRequest.class)); + } + } + + @Test(timeout = 10_000) + public void synchronousStartFailureWithLastStatementDoesNotRetry() { + AtomicBoolean failNextStreamStart = new AtomicBoolean(); + SpannerOptions options = createSpannerOptionsWithFailingStreamStart(failNextStreamStart); + AtomicInteger attempts = new AtomicInteger(); + try (Spanner customSpanner = options.getService()) { + DatabaseClient customClient = + customSpanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]")); + SpannerException exception = + assertThrows( + SpannerException.class, + () -> + customClient + .readWriteTransaction() + .run( + transaction -> { + attempts.incrementAndGet(); + failNextStreamStart.set(true); + try (ResultSet resultSet = + transaction.executeQuery(SELECT1, Options.lastStatement())) { + SpannerException firstError = + assertThrows(SpannerException.class, resultSet::next); + assertEquals(ErrorCode.UNKNOWN, firstError.getErrorCode()); + } + // Executing a second statement awaits transactionIdFuture in + // getTransactionSelector(). Because the initial statement had + // lastStatement=true, transactionIdFuture was completed with the + // original exception rather than a simulated AbortedException, so + // this rethrows UNKNOWN without retrying the transaction. + try (ResultSet secondResultSet = transaction.executeQuery(SELECT1)) { + secondResultSet.next(); + } + return null; + })); + assertEquals(ErrorCode.UNKNOWN, exception.getErrorCode()); + assertEquals(1, attempts.get()); + + // Also verify the uncaught single-statement path where rollbackAsync() awaits + // transactionIdFuture and completes immediately without hanging or retrying. + AtomicInteger uncaughtAttempts = new AtomicInteger(); + SpannerException uncaughtException = + assertThrows( + SpannerException.class, + () -> + customClient + .readWriteTransaction() + .run( + transaction -> { + uncaughtAttempts.incrementAndGet(); + failNextStreamStart.set(true); + try (ResultSet resultSet = + transaction.executeQuery(SELECT1, Options.lastStatement())) { + resultSet.next(); + } + return null; + })); + assertEquals(ErrorCode.UNKNOWN, uncaughtException.getErrorCode()); + assertEquals(1, uncaughtAttempts.get()); + + // Also verify CommitRunnable.run() awaiting transactionIdFuture when lastStatement=true + // stored a raw IllegalStateException in transactionIdFuture and the caller caught the query + // failure and buffered a mutation. + AtomicInteger commitAttempts = new AtomicInteger(); + SpannerException commitException = + assertThrows( + SpannerException.class, + () -> + customClient + .readWriteTransaction() + .run( + transaction -> { + commitAttempts.incrementAndGet(); + failNextStreamStart.set(true); + try (ResultSet resultSet = + transaction.executeQuery(SELECT1, Options.lastStatement())) { + SpannerException firstError = + assertThrows(SpannerException.class, resultSet::next); + assertEquals(ErrorCode.UNKNOWN, firstError.getErrorCode()); + } + transaction.buffer( + Mutation.newInsertBuilder("FOO").set("ID").to(1L).build()); + return null; + })); + assertEquals(ErrorCode.UNKNOWN, commitException.getErrorCode()); + assertEquals(1, commitAttempts.get()); + + // Also verify initiateStreaming() failing synchronously when lastStatement=true: only + // TransactionContextImpl.onStartFailed runs during initiateStreaming(), and it must + // immediately complete transactionIdFuture with the converted SpannerException so commit() + // fails immediately with UNKNOWN without retrying. + AtomicInteger initiateStreamingLastStatementAttempts = new AtomicInteger(); + SpannerException initiateStreamingCommitException = + assertThrows( + SpannerException.class, + () -> + customClient + .readWriteTransaction() + .run( + transaction -> { + initiateStreamingLastStatementAttempts.incrementAndGet(); + failNextStreamStart.set(true); + try (ResultSet resultSet = + transaction.executeQuery(SELECT1, Options.lastStatement())) { + assertThrows( + IllegalStateException.class, + () -> + ((StreamingResultSet) resultSet) + .initiateStreaming((message, prefetchChunks) -> {})); + } + transaction.buffer( + Mutation.newInsertBuilder("FOO").set("ID").to(1L).build()); + return null; + })); + assertEquals(ErrorCode.UNKNOWN, initiateStreamingCommitException.getErrorCode()); + assertEquals(1, initiateStreamingLastStatementAttempts.get()); + } + } + + @Test(timeout = 10_000) + public void synchronousStartFailureWithAbortedOrInvalidatedTransactionRetries() throws Exception { + AtomicBoolean failNextStreamStart = new AtomicBoolean(); + AtomicReference> failureSupplierRef = new AtomicReference<>(); + SpannerOptions options = + createSpannerOptionsWithFailingStreamStart( + failNextStreamStart, () -> failureSupplierRef.get().get()); + try (Spanner customSpanner = options.getService()) { + DatabaseClient customClient = + customSpanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]")); + + // 1. When an inlined-begin statement with lastStatement=true fails synchronously with a raw + // StatusRuntimeException(ABORTED) while a concurrent follower thread is blocked on + // transactionIdFuture.get() in getTransactionSelector(), onStartFailed converts the raw + // throwable to AbortedException so the follower thread sees e.getCause() instanceof + // AbortedException, marks the transaction aborted, and the transaction retries with an + // explicit BeginTransaction RPC. + mockSpanner.reset(); + AtomicInteger abortedAttempts = new AtomicInteger(); + customClient + .readWriteTransaction() + .run( + transaction -> { + if (abortedAttempts.incrementAndGet() == 1) { + CountDownLatch leaderEnteredStartLatch = new CountDownLatch(1); + CountDownLatch followerStartedLatch = new CountDownLatch(1); + AtomicReference followerThread = new AtomicReference<>(); + SettableApiFuture followerFailure = SettableApiFuture.create(); + failureSupplierRef.set( + () -> { + leaderEnteredStartLatch.countDown(); + try { + assertTrue(followerStartedLatch.await(5, TimeUnit.SECONDS)); + } catch (InterruptedException interruptedException) { + Thread.currentThread().interrupt(); + throw new RuntimeException(interruptedException); + } + while (followerThread.get().getState() != Thread.State.TIMED_WAITING + && !followerFailure.isDone()) { + Thread.yield(); + } + return Status.ABORTED.asRuntimeException(); + }); + failNextStreamStart.set(true); + Thread thread = + new Thread( + () -> { + try { + assertTrue(leaderEnteredStartLatch.await(5, TimeUnit.SECONDS)); + followerStartedLatch.countDown(); + try (ResultSet followerResultSet = + transaction.executeQuery(SELECT1)) { + followerResultSet.next(); + } + } catch (SpannerException spannerException) { + followerFailure.set(spannerException); + } catch (Throwable throwable) { + followerFailure.setException(throwable); + } + }); + followerThread.set(thread); + thread.start(); + + try (ResultSet firstResultSet = + transaction.executeQuery(SELECT1, Options.lastStatement())) { + SpannerException firstException = + assertThrows(SpannerException.class, firstResultSet::next); + assertEquals(ErrorCode.ABORTED, firstException.getErrorCode()); + } + SpannerException followerException = followerFailure.get(5, TimeUnit.SECONDS); + assertTrue(followerException instanceof AbortedException); + thread.join(5_000); + throw followerException; + } + try (ResultSet resultSet = + transaction.executeQuery(SELECT1, Options.lastStatement())) { + while (resultSet.next()) {} + } + return null; + }); + assertEquals(2, abortedAttempts.get()); + assertEquals(1, countRequests(BeginTransactionRequest.class)); + assertEquals(1, countRequests(CommitRequest.class)); + + // 2. When an inlined-begin statement fails synchronously with CANCELLED ("Transaction was + // invalidated by a later transaction in the same session"), ResumableStreamIterator preserves + // withBeginTransaction=true so GrpcResultSet.next() passes withBeginTransaction=true to + // TransactionRunnerImpl.onError, which marks the transaction aborted and retries with an + // explicit BeginTransaction RPC even when lastStatement=true. + mockSpanner.reset(); + failureSupplierRef.set( + () -> + Status.CANCELLED + .withDescription( + "Transaction was invalidated by a later transaction in the same session") + .asRuntimeException()); + AtomicInteger invalidatedAttempts = new AtomicInteger(); + customClient + .readWriteTransaction() + .run( + transaction -> { + if (invalidatedAttempts.incrementAndGet() == 1) { + failNextStreamStart.set(true); + } + try (ResultSet resultSet = + transaction.executeQuery(SELECT1, Options.lastStatement())) { + while (resultSet.next()) {} + } + return null; + }); + assertEquals(2, invalidatedAttempts.get()); + assertEquals(1, countRequests(BeginTransactionRequest.class)); + assertEquals(1, countRequests(CommitRequest.class)); + + // 3. When an inlined-begin statement with lastStatement=true fails synchronously with + // CANCELLED ("Transaction was invalidated by a later transaction in the same session") and + // the caller catches the exception and buffers a mutation, CommitRunnable awaits + // transactionIdFuture, which must hold the translated AbortedException so the transaction + // retries with an explicit BeginTransaction RPC instead of failing commit() with CANCELLED. + mockSpanner.reset(); + AtomicInteger invalidatedAndBufferedAttempts = new AtomicInteger(); + customClient + .readWriteTransaction() + .run( + transaction -> { + if (invalidatedAndBufferedAttempts.incrementAndGet() == 1) { + failNextStreamStart.set(true); + try (ResultSet resultSet = + transaction.executeQuery(SELECT1, Options.lastStatement())) { + SpannerException firstException = + assertThrows(SpannerException.class, resultSet::next); + assertEquals(ErrorCode.ABORTED, firstException.getErrorCode()); + } + } + transaction.buffer(Mutation.newInsertBuilder("FOO").set("ID").to(1L).build()); + return null; + }); + assertEquals(2, invalidatedAndBufferedAttempts.get()); + assertEquals(1, countRequests(BeginTransactionRequest.class)); + assertEquals(1, countRequests(CommitRequest.class)); + + // 4. Same scenario when the CANCELLED ("Transaction was invalidated...") error is returned + // asynchronously on the gRPC stream (exercising TransactionRunnerImpl.onError translating + // the error before completing transactionIdFuture when lastStatement=true). + mockSpanner.reset(); + mockSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.ofException( + Status.CANCELLED + .withDescription( + "Transaction was invalidated by a later transaction in the same session") + .asRuntimeException())); + AtomicInteger streamInvalidatedAndBufferedAttempts = new AtomicInteger(); + customClient + .readWriteTransaction() + .run( + transaction -> { + if (streamInvalidatedAndBufferedAttempts.incrementAndGet() == 1) { + try (ResultSet resultSet = + transaction.executeQuery(SELECT1, Options.lastStatement())) { + SpannerException firstException = + assertThrows(SpannerException.class, resultSet::next); + assertEquals(ErrorCode.ABORTED, firstException.getErrorCode()); + } + } + transaction.buffer(Mutation.newInsertBuilder("FOO").set("ID").to(1L).build()); + return null; + }); + assertEquals(2, streamInvalidatedAndBufferedAttempts.get()); + assertEquals(1, countRequests(BeginTransactionRequest.class)); + assertEquals(1, countRequests(CommitRequest.class)); + } + } + + @Test(timeout = 10_000) + public void cancelDuringRetryBackoffOfInlinedBeginStatementFailsTransactionIdFuture() { + // When an inlined-begin stream fails with a retryable UNAVAILABLE before receiving a resume + // token, ResumableStreamIterator sets stream = null and enters backoffSleep. Note that + // startStream() succeeded, so onStartFailed was not called and transactionIdFuture is still + // pending. If the gRPC Context is cancelled during backoffSleep, GrpcResultSet.next() catches + // the cancellation exception while ResumableStreamIterator.stream is null. The + // isWithBeginTransaction() and isLastStatement() overrides on the anonymous + // ResumableStreamIterator must still report true and lastStatement so onError completes + // transactionIdFuture appropriately (simulated ABORTED when lastStatement=false, or CANCELLED + // without retry when lastStatement=true). + AtomicReference firstAttemptDeliveredErrorLatch = new AtomicReference<>(); + AtomicBoolean deliverTransientUnavailable = new AtomicBoolean(true); + SpannerOptions options = + createSpannerOptionsWithStreamInterceptors( + (invocation, delegate) -> { + if (deliverTransientUnavailable.get()) { + SpannerRpc.ResultStreamConsumer consumer = invocation.getArgument(1); + // Deliver UNAVAILABLE with a 5s retry delay on the first attempt so + // ResumableStreamIterator sets stream = null and parks in + // backoffSleep until the Context is cancelled. + consumer.onError( + SpannerExceptionFactory.newSpannerException( + ErrorCode.UNAVAILABLE, + "simulated transient stream failure", + SpannerExceptionFactory.createAbortedExceptionWithRetryDelay( + "simulated transient stream failure", null, 5, 0))); + CountDownLatch latch = firstAttemptDeliveredErrorLatch.get(); + if (latch != null) { + latch.countDown(); + } + return new SpannerRpc.StreamingCall() { + @Override + public ApiCallContext getCallContext() { + return null; + } + + @Override + public void cancel(String message) {} + + @Override + public void request(int numMessages) {} + }; + } + return delegateExecuteQuery(invocation, delegate); + }, + ReadWriteTransactionWithInlineBeginTest::delegateRead); + + try (Spanner customSpanner = options.getService()) { + DatabaseClient customClient = + customSpanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]")); + + // 1. With Options.lastStatement(): cancelling the Context during backoffSleep must pass + // withBeginTransaction=true and lastStatement=true to onError so transactionIdFuture is + // completed with CANCELLED (not simulated ABORTED), and commit() fails with CANCELLED + // after 1 attempt rather than retrying. + CountDownLatch latch1 = new CountDownLatch(1); + firstAttemptDeliveredErrorLatch.set(latch1); + AtomicInteger lastStatementAttempts = new AtomicInteger(); + SpannerException lastStatementException = + assertThrows( + SpannerException.class, + () -> + customClient + .readWriteTransaction() + .run( + transaction -> { + lastStatementAttempts.incrementAndGet(); + Thread consumerThread = Thread.currentThread(); + try (Context.CancellableContext cancellableContext = + Context.current().withCancellation()) { + Thread canceller = + new Thread( + () -> { + try { + assertTrue(latch1.await(5, TimeUnit.SECONDS)); + } catch (InterruptedException interruptedException) { + Thread.currentThread().interrupt(); + return; + } + while (consumerThread.getState() + != Thread.State.TIMED_WAITING + && consumerThread.isAlive()) { + Thread.yield(); + } + cancellableContext.cancel( + new CancellationException("cancelled during backoff")); + }); + cancellableContext.run( + () -> { + canceller.start(); + try (ResultSet resultSet = + transaction.executeQuery( + SELECT1, Options.lastStatement())) { + SpannerException exception = + assertThrows(SpannerException.class, resultSet::next); + assertEquals(ErrorCode.CANCELLED, exception.getErrorCode()); + } + }); + canceller.join(5_000); + } + transaction.buffer( + Mutation.newInsertBuilder("FOO").set("ID").to(1L).build()); + return null; + })); + assertEquals(ErrorCode.CANCELLED, lastStatementException.getErrorCode()); + assertEquals(1, lastStatementAttempts.get()); + + // 2. Without Options.lastStatement() (lastStatement=false): cancelling the Context during + // backoffSleep must pass withBeginTransaction=true and lastStatement=false to onError so + // transactionIdFuture is completed with a simulated AbortedException, and commit() retries + // the transaction with an explicit BeginTransaction RPC. + mockSpanner.reset(); + CountDownLatch latch2 = new CountDownLatch(1); + firstAttemptDeliveredErrorLatch.set(latch2); + AtomicInteger nonLastStatementAttempts = new AtomicInteger(); + customClient + .readWriteTransaction() + .run( + transaction -> { + if (nonLastStatementAttempts.incrementAndGet() == 1) { + Thread consumerThread = Thread.currentThread(); + try (Context.CancellableContext cancellableContext = + Context.current().withCancellation()) { + Thread canceller = + new Thread( + () -> { + try { + assertTrue(latch2.await(5, TimeUnit.SECONDS)); + } catch (InterruptedException interruptedException) { + Thread.currentThread().interrupt(); + return; + } + while (consumerThread.getState() != Thread.State.TIMED_WAITING + && consumerThread.isAlive()) { + Thread.yield(); + } + deliverTransientUnavailable.set(false); + cancellableContext.cancel( + new CancellationException("cancelled during backoff")); + }); + cancellableContext.run( + () -> { + canceller.start(); + try (ResultSet resultSet = transaction.executeQuery(SELECT1)) { + SpannerException exception = + assertThrows(SpannerException.class, resultSet::next); + assertEquals(ErrorCode.CANCELLED, exception.getErrorCode()); + } + }); + canceller.join(5_000); + } + } + transaction.buffer(Mutation.newInsertBuilder("FOO").set("ID").to(1L).build()); + return null; + }); + assertEquals(2, nonLastStatementAttempts.get()); + assertEquals(1, countRequests(BeginTransactionRequest.class)); + assertEquals(1, countRequests(CommitRequest.class)); + } + } + + @Test(timeout = 10_000) + public void closeBeforeTransactionIdReturnedDoesNotRaceWithCancelledConsumer() throws Exception { + // Deterministic reproduction of the race between GrpcResultSet.close() and a consumer thread + // blocked in GrpcResultSet.next() when the inlined-begin stream has not yet returned a + // transaction ID: + // 1. GrpcResultSet.close() calls iterator.close("ResultSet closed"), which cancels the gRPC + // call and wakes the blocked consumer thread before GrpcResultSet.close() reaches + // listener.onDone(true). + // 2. When the woken consumer thread catches the stream termination error in + // GrpcResultSet.next() + // and calls yieldError(e, iterator.isWithBeginTransaction() && currRow == null, ...), + // iterator.isWithBeginTransaction() must return false because the iterator is already + // closed, so onError does not overwrite transactionIdFuture with a simulated ABORTED before + // onDone(true) sets FAILED_PRECONDITION. + CountDownLatch consumerFinishedNextLatch = new CountDownLatch(1); + SpannerOptions options = + createSpannerOptionsWithStreamInterceptors( + (invocation, delegate) -> { + SpannerRpc.StreamingCall delegateCall = delegateExecuteQuery(invocation, delegate); + return new SpannerRpc.StreamingCall() { + @Override + public ApiCallContext getCallContext() { + return delegateCall.getCallContext(); + } + + @Override + public void cancel(String message) { + delegateCall.cancel(message); + try { + assertTrue(consumerFinishedNextLatch.await(5, TimeUnit.SECONDS)); + } catch (InterruptedException interruptedException) { + Thread.currentThread().interrupt(); + } + } + + @Override + public void request(int numMessages) { + delegateCall.request(numMessages); + } + }; + }, + ReadWriteTransactionWithInlineBeginTest::delegateRead); + + ExecutorService executor = Executors.newFixedThreadPool(2); + try (Spanner customSpanner = options.getService()) { + DatabaseClient customClient = + customSpanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]")); + SpannerException exception = + assertThrows( + SpannerException.class, + () -> + customClient + .readWriteTransaction() + .run( + transaction -> { + mockSpanner.freeze(); + ResultSet resultSet = transaction.executeQuery(SELECT1); + Future nextFuture = + executor.submit( + () -> { + try { + resultSet.next(); + } finally { + consumerFinishedNextLatch.countDown(); + } + }); + try { + mockSpanner.waitForRequestsToContain(ExecuteSqlRequest.class, 5000); + } catch (Exception exceptionFromWait) { + throw new RuntimeException(exceptionFromWait); + } + resultSet.close(); + mockSpanner.unfreeze(); + ExecutionException consumerExecutionException = + assertThrows( + ExecutionException.class, + () -> nextFuture.get(5, TimeUnit.SECONDS)); + assertTrue( + consumerExecutionException.getCause() instanceof SpannerException); + SpannerException consumerException = + (SpannerException) consumerExecutionException.getCause(); + assertNotEquals(ErrorCode.ABORTED, consumerException.getErrorCode()); + return transaction.executeUpdate(UPDATE_STATEMENT); + })); + assertEquals(ErrorCode.FAILED_PRECONDITION, exception.getErrorCode()); + assertTrue( + exception + .getMessage() + .contains("ResultSet was closed before a transaction id was returned")); + } finally { + mockSpanner.unfreeze(); + executor.shutdownNow(); + } + } + + @Test + public void inlinedBeginReadWithStreamRetryBeforeAndAfterResumeToken() { + // 1. Stream retry before receiving a resumeToken (errorIndex == 0): both ReadRequests must + // keep the inlined BeginTransaction selector and an empty resumeToken. + mockSpanner.setStreamingReadExecutionTime( + SimulatedExecutionTime.ofStreamException(Status.UNAVAILABLE.asRuntimeException(), 0)); + List valuesBeforeResumeToken = + client + .readWriteTransaction() + .run( + transaction -> { + List values = new ArrayList<>(); + try (ResultSet resultSet = + transaction.read( + "FOO_TWO_ROWS", KeySet.all(), Collections.singletonList("ID"))) { + while (resultSet.next()) { + values.add(resultSet.getLong(0)); + } + } + return values; + }); + assertEquals(Arrays.asList(1L, 2L), valuesBeforeResumeToken); + assertEquals(0, countRequests(BeginTransactionRequest.class)); + assertEquals(2, countRequests(ReadRequest.class)); + assertEquals(1, countRequests(CommitRequest.class)); + List readRequestsBeforeResumeToken = + mockSpanner.getRequestsOfType(ReadRequest.class); + assertTrue(readRequestsBeforeResumeToken.get(0).getTransaction().getBegin().hasReadWrite()); + assertEquals(ByteString.EMPTY, readRequestsBeforeResumeToken.get(0).getResumeToken()); + assertTrue(readRequestsBeforeResumeToken.get(1).getTransaction().getBegin().hasReadWrite()); + assertEquals(ByteString.EMPTY, readRequestsBeforeResumeToken.get(1).getResumeToken()); + + mockSpanner.reset(); + mockSpanner.removeAllExecutionTimes(); + + // 2. Stream retry after receiving a resumeToken (errorIndex == 1): the first ReadRequest has + // the inlined BeginTransaction selector and an empty resumeToken, while the resumed + // ReadRequest replaces BeginTransaction with the transaction ID and sets resumeToken. + mockSpanner.setStreamingReadExecutionTime( + SimulatedExecutionTime.ofStreamException(Status.UNAVAILABLE.asRuntimeException(), 1)); + List valuesAfterResumeToken = + client + .readWriteTransaction() + .run( + transaction -> { + List values = new ArrayList<>(); + try (ResultSet resultSet = + transaction.read( + "FOO_TWO_ROWS", KeySet.all(), Collections.singletonList("ID"))) { + while (resultSet.next()) { + values.add(resultSet.getLong(0)); + } + } + return values; + }); + assertEquals(Arrays.asList(1L, 2L), valuesAfterResumeToken); + assertEquals(0, countRequests(BeginTransactionRequest.class)); + assertEquals(2, countRequests(ReadRequest.class)); + assertEquals(1, countRequests(CommitRequest.class)); + assertEquals(1, countTransactionsStarted()); + List readRequestsAfterResumeToken = + mockSpanner.getRequestsOfType(ReadRequest.class); + assertTrue(readRequestsAfterResumeToken.get(0).getTransaction().getBegin().hasReadWrite()); + assertEquals(ByteString.EMPTY, readRequestsAfterResumeToken.get(0).getTransaction().getId()); + assertEquals(ByteString.EMPTY, readRequestsAfterResumeToken.get(0).getResumeToken()); + assertFalse(readRequestsAfterResumeToken.get(1).getTransaction().hasBegin()); + assertNotEquals(ByteString.EMPTY, readRequestsAfterResumeToken.get(1).getTransaction().getId()); + assertNotEquals(ByteString.EMPTY, readRequestsAfterResumeToken.get(1).getResumeToken()); + } + + @Test + public void subsequentQueryAndReadWithStreamRetryAfterResumeToken() { + mockSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.ofStreamException(Status.UNAVAILABLE.asRuntimeException(), 1)); + mockSpanner.setStreamingReadExecutionTime( + SimulatedExecutionTime.ofStreamException(Status.UNAVAILABLE.asRuntimeException(), 1)); + + client + .readWriteTransaction() + .run( + transaction -> { + // First statement starts the transaction inline. + assertEquals(1L, transaction.executeUpdate(UPDATE_STATEMENT)); + + // Second statement (query) starts with transaction.id (hasBegin() == false) and + // retries mid-stream with a resumeToken. + List queryValues = new ArrayList<>(); + try (ResultSet queryResultSet = transaction.executeQuery(SELECT1_UNION_ALL_SELECT2)) { + while (queryResultSet.next()) { + queryValues.add(queryResultSet.getLong(0)); + } + } + assertEquals(Arrays.asList(1L, 2L), queryValues); + + // Third statement (read) also starts with transaction.id (hasBegin() == false) and + // retries mid-stream with a resumeToken. + List readValues = new ArrayList<>(); + try (ResultSet readResultSet = + transaction.read("FOO_TWO_ROWS", KeySet.all(), Collections.singletonList("ID"))) { + while (readResultSet.next()) { + readValues.add(readResultSet.getLong(0)); + } + } + assertEquals(Arrays.asList(1L, 2L), readValues); + return null; + }); + + assertEquals(0, countRequests(BeginTransactionRequest.class)); + assertEquals(1, countTransactionsStarted()); + ByteString expectedTransactionId = mockSpanner.getTransactionsStarted().get(0); + + List executeSqlRequests = + mockSpanner.getRequestsOfType(ExecuteSqlRequest.class); + // 1 DML + 2 streaming query attempts. + assertEquals(3, executeSqlRequests.size()); + assertFalse(executeSqlRequests.get(1).getTransaction().hasBegin()); + assertEquals(expectedTransactionId, executeSqlRequests.get(1).getTransaction().getId()); + assertEquals(ByteString.EMPTY, executeSqlRequests.get(1).getResumeToken()); + assertFalse(executeSqlRequests.get(2).getTransaction().hasBegin()); + assertEquals(expectedTransactionId, executeSqlRequests.get(2).getTransaction().getId()); + assertNotEquals(ByteString.EMPTY, executeSqlRequests.get(2).getResumeToken()); + + List readRequests = mockSpanner.getRequestsOfType(ReadRequest.class); + assertEquals(2, readRequests.size()); + assertFalse(readRequests.get(0).getTransaction().hasBegin()); + assertEquals(expectedTransactionId, readRequests.get(0).getTransaction().getId()); + assertEquals(ByteString.EMPTY, readRequests.get(0).getResumeToken()); + assertFalse(readRequests.get(1).getTransaction().hasBegin()); + assertEquals(expectedTransactionId, readRequests.get(1).getTransaction().getId()); + assertNotEquals(ByteString.EMPTY, readRequests.get(1).getResumeToken()); + } + private int countRequests(Class requestType) { int count = 0; for (AbstractMessage msg : mockSpanner.getRequests()) {