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 614ad37fa041..36d55cc69c81 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 @@ -108,8 +108,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()) {