From c8c9b3c531252cd5e78902f608b3172d0cfd6499 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= Date: Mon, 2 Nov 2026 12:55:10 +0100 Subject: [PATCH] fix(spanner): cancel queries promptly on close and synchronize error handling When a query or transaction is cancelled or closed while the network request is still being established, the underlying gRPC stream was previously left unregistered, allowing the query to continue running in the background and wasting database and network resources. Additionally, closing a query concurrently with an error could cause completion callbacks to race ahead of error handlers, corrupting transaction state and masking the underlying failure. This change fixes several concurrency edge cases in query execution: - Prompt query cancellation: Registers streams immediately so that queries cancelled or closed during startup are aborted right away, or skipped entirely if already closed, preventing redundant network RPCs. - Clean failure for early closes: When a query that initiates a transaction is closed before the request finishes, waiting callers fail fast with a clear error rather than hanging indefinitely for a transaction ID. - Coordinated error and close callbacks: Guarantees that query errors are fully dispatched before completion notifications run, ensuring transactions observe the actual database failure and can retry appropriately instead of being masked by a generic "result set closed" error. - Fast failure on repeated reads: Once a query stream encounters an error, subsequent read calls immediately surface the original failure rather than attempting further operations. --- .../cloud/spanner/AbstractReadContext.java | 17 + .../google/cloud/spanner/GrpcResultSet.java | 60 ++- .../cloud/spanner/GrpcStreamIterator.java | 18 +- .../spanner/MultiUseReadOnlyTransaction.java | 21 +- .../spanner/ResumableStreamIterator.java | 13 +- .../cloud/spanner/TransactionRunnerImpl.java | 8 +- .../spanner/AbstractReadContextTest.java | 433 ++++++++++++++- .../spanner/AsyncReadOnlyTransactionTest.java | 249 +++++++++ .../cloud/spanner/GrpcResultSetTest.java | 495 ++++++++++++++++++ ...adWriteTransactionWithInlineBeginTest.java | 264 +++++++++- .../spanner/ResumableStreamIteratorTest.java | 121 +++-- 11 files changed, 1616 insertions(+), 83 deletions(-) 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 8345d78388ec..39d2653ce196 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 @@ -189,6 +189,9 @@ static void assertTimestampAvailable(boolean available) { protected static final String NO_TRANSACTION_RETURNED_MSG = "The statement did not return a transaction even though one was requested"; + protected static final String RESULT_SET_CLOSED_BEFORE_ID_RETURNED_MESSAGE = + "ResultSet was closed before a transaction id was returned"; + AbstractReadContext(Builder builder) { this.session = builder.session; this.cancelQueryWhenClientIsClosed = builder.cancelQueryWhenClientIsClosed; @@ -526,6 +529,18 @@ CloseableIterator startStream( if (streamListener != null) { stream.registerListener(streamListener); } + // setStream() returns false if this ResumableStreamIterator was already closed + // concurrently (e.g., via resultSet.close() or transaction.close()). In that case, + // setStream() has already invoked stream.close() out-of-lock to prevent leaks. We skip + // issuing the RPC call (avoiding unnecessary network I/O) and fail any inlined-begin + // future via onClosedBeforeStreamStarted(). We still return the (now closed) stream + // instance because startStream() requires a non-null return value to satisfy + // checkNotNull() in ResumableStreamIterator.startGrpcStreaming(), which then cleanly + // completes stream teardown and wakes any threads waiting on streamLock. + if (!setStream(stream)) { + onClosedBeforeStreamStarted(withBeginTransaction); + return stream; + } SpannerRpc.StreamingCall call; try { call = @@ -690,6 +705,8 @@ public void onDone(boolean withBeginTransaction) { void onStartFailed(boolean withBeginTransaction, boolean lastStatement, Throwable throwable) {} + void onClosedBeforeStreamStarted(boolean withBeginTransaction) {} + /** * For transactions other than read-write, the MultiplexedSessionPrecommitToken will not be * present in the RPC response. In such cases, this method will be a no-op. diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/GrpcResultSet.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/GrpcResultSet.java index 2f3daef47419..92aa99df000e 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/GrpcResultSet.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/GrpcResultSet.java @@ -39,8 +39,10 @@ class GrpcResultSet extends AbstractResultSet> private ResultSetMetadata metadata; private GrpcStruct currRow; private List rowData; + private volatile SpannerException error; private ResultSetStats statistics; private volatile boolean closed; + private boolean yieldingError; GrpcResultSet(CloseableIterator iterator, Listener listener) { this(iterator, listener, DecodeMode.DIRECT); @@ -74,6 +76,9 @@ protected GrpcStruct currRow() { @Override public boolean next() throws SpannerException { + if (this.error != null) { + throw this.error; + } try { if (currRow == null) { metadata = iterator.getMetadata(); @@ -102,9 +107,9 @@ public boolean next() throws SpannerException { close(); } return hasNext; - } catch (Throwable t) { + } catch (Throwable throwable) { throw yieldError( - asSpannerException(t), + asSpannerException(throwable), iterator.isWithBeginTransaction() && currRow == null, iterator.isLastStatement()); } @@ -139,15 +144,25 @@ public boolean isDataAvailable() { @Override public void close() { + boolean shouldNotifyDone; synchronized (this) { if (closed) { return; } closed = true; + // If a thread is currently in yieldError() dispatching listener.onError(), defer calling + // listener.onDone() to yieldError() finally block so that onError() always completes before + // onDone() is invoked. + shouldNotifyDone = !yieldingError; } boolean withBeginTransaction = iterator.isWithBeginTransaction(); - iterator.close("ResultSet closed"); - listener.onDone(withBeginTransaction); + try { + iterator.close("ResultSet closed"); + } finally { + if (shouldNotifyDone) { + listener.onDone(withBeginTransaction || iterator.isWithBeginTransaction()); + } + } } @Override @@ -157,9 +172,38 @@ public Type getType() { } private SpannerException yieldError( - SpannerException e, boolean beginTransaction, boolean lastStatement) { - SpannerException toThrow = listener.onError(e, beginTransaction, lastStatement); - close(); - throw toThrow; + SpannerException exception, boolean beginTransaction, boolean lastStatement) { + synchronized (this) { + if (this.error == null) { + this.error = exception; + } + if (this.closed || this.yieldingError) { + return this.error; + } + this.yieldingError = true; + } + try { + this.error = listener.onError(exception, beginTransaction, lastStatement); + throw this.error; + } finally { + boolean alreadyClosed; + synchronized (this) { + this.yieldingError = false; + alreadyClosed = this.closed; + } + // If close() was called concurrently while we were dispatching listener.onError(), close() + // intentionally skipped calling listener.onDone() to avoid racing with or completing before + // onError(). We ensure the underlying iterator is closed before invoking listener.onDone(). + // Otherwise, if close() was not called yet, proceed with normal closure. + if (alreadyClosed) { + try { + iterator.close("ResultSet closed"); + } finally { + listener.onDone(beginTransaction || iterator.isWithBeginTransaction()); + } + } else { + close(); + } + } } } 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 312e841a188c..e4856f4a20b8 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 @@ -119,6 +119,20 @@ public String toString() { /* withBeginTransaction= */ false); } + @VisibleForTesting + GrpcStreamIterator( + boolean lastStatement, + int prefetchChunks, + boolean cancelQueryWhenClientIsClosed, + boolean withBeginTransaction) { + this( + /* statement= */ null, + lastStatement, + prefetchChunks, + cancelQueryWhenClientIsClosed, + withBeginTransaction); + } + GrpcStreamIterator( Statement statement, boolean lastStatement, @@ -158,9 +172,7 @@ void registerListener(AsyncResultSet.StreamMessageListener streamMessageListener @VisibleForTesting void setCall(SpannerRpc.StreamingCall call, boolean withBeginTransaction) { - synchronized (this) { - this.withBeginTransaction = withBeginTransaction; - } + this.withBeginTransaction = withBeginTransaction; setCall(call); } 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 05323fb12144..838106ed56a6 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 @@ -474,10 +474,7 @@ public SpannerException onError( @Override public void onDone(boolean withBeginTransaction) { if (withBeginTransaction) { - failTransactionIdFuture( - SpannerExceptionFactory.newSpannerException( - ErrorCode.FAILED_PRECONDITION, - "ResultSet was closed before a read-only transaction id was returned")); + failTransactionIdFutureIfClosedBeforeIdReturned(); } super.onDone(withBeginTransaction); } @@ -485,10 +482,24 @@ public void onDone(boolean withBeginTransaction) { @Override void onStartFailed(boolean withBeginTransaction, boolean lastStatement, Throwable throwable) { if (withBeginTransaction) { - failTransactionIdFuture(throwable); + failTransactionIdFuture(SpannerExceptionFactory.asSpannerException(throwable)); } } + @Override + void onClosedBeforeStreamStarted(boolean withBeginTransaction) { + if (withBeginTransaction) { + failTransactionIdFutureIfClosedBeforeIdReturned(); + } + } + + private void failTransactionIdFutureIfClosedBeforeIdReturned() { + failTransactionIdFuture( + SpannerExceptionFactory.newSpannerException( + ErrorCode.FAILED_PRECONDITION, + "ResultSet was closed before a read-only transaction id was returned")); + } + private void failTransactionIdFuture(Throwable throwable) { txnLock.lock(); try { diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ResumableStreamIterator.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ResumableStreamIterator.java index 4ea252f90e0b..e13245b1ef1d 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ResumableStreamIterator.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ResumableStreamIterator.java @@ -344,10 +344,9 @@ public boolean isWithBeginTransaction() { currentInFlightStream = this.inFlightStream; fallbackWithBeginTransaction = this.withBeginTransaction; } - // A subclass that registers its stream early calls setStream(consumer) before - // consumer.setCall(call, true) sets withBeginTransaction on the underlying stream. Querying - // currentInFlightStream live ensures callers observe true as soon as setCall() runs during an - // in-flight startStream(). + // A subclass that registers its stream early calls setStream(consumer) during stream + // initialization. Querying currentInFlightStream live ensures callers observe + // withBeginTransaction as soon as the stream is created during an in-flight startStream(). return currentInFlightStream != null ? currentInFlightStream.isWithBeginTransaction() : fallbackWithBeginTransaction; @@ -634,10 +633,8 @@ private void resetStreamOnFailure(Throwable throwable) { synchronized (streamLock) { if (this.inFlightStream != null) { // Refresh withBeginTransaction (a non-blocking field read on the stream iterator) before - // clearing inFlightStream and calling notifyAll(). A subclass that registers its stream - // early calls setStream(consumer) before consumer.setCall(call, true), and any follower - // thread woken by notifyAll() must immediately observe the updated flag via - // isWithBeginTransaction(). + // clearing inFlightStream and calling notifyAll(). Any follower thread woken by notifyAll() + // must immediately observe the updated flag via isWithBeginTransaction(). this.withBeginTransaction = this.inFlightStream.isWithBeginTransaction(); if (!this.closed) { failedStream = this.inFlightStream; 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 ddef5b841b6e..f058de043d40 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 @@ -904,11 +904,15 @@ public void onDone(boolean withBeginTransaction) { // Context was done (closed) before a transaction id was returned. this.transactionIdFuture.setException( SpannerExceptionFactory.newSpannerException( - ErrorCode.FAILED_PRECONDITION, - "ResultSet was closed before a transaction id was returned")); + ErrorCode.FAILED_PRECONDITION, RESULT_SET_CLOSED_BEFORE_ID_RETURNED_MESSAGE)); } } + @Override + void onClosedBeforeStreamStarted(boolean withBeginTransaction) { + onDone(withBeginTransaction); + } + @Override void onStartFailed(boolean withBeginTransaction, boolean lastStatement, Throwable throwable) { if (withBeginTransaction) { diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AbstractReadContextTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AbstractReadContextTest.java index 37827aff5c75..1c193087ecdd 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AbstractReadContextTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AbstractReadContextTest.java @@ -19,15 +19,26 @@ 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.assertThrows; import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyBoolean; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import static org.mockito.Mockito.withSettings; +import com.google.api.core.SettableApiFuture; import com.google.api.gax.core.ExecutorProvider; +import com.google.api.gax.retrying.RetrySettings; import com.google.cloud.spanner.Options.RpcPriority; +import com.google.cloud.spanner.XGoogSpannerRequestId.NoopRequestIdCreator; import com.google.cloud.spanner.spi.v1.SpannerRpc; import com.google.cloud.spanner.spi.v1.SpannerRpc.Option; +import com.google.protobuf.ByteString; import com.google.spanner.v1.DirectedReadOptions; import com.google.spanner.v1.DirectedReadOptions.IncludeReplicas; import com.google.spanner.v1.DirectedReadOptions.ReplicaSelection; @@ -41,12 +52,20 @@ import com.google.spanner.v1.RequestOptions; import com.google.spanner.v1.RequestOptions.Priority; import com.google.spanner.v1.SessionName; +import com.google.spanner.v1.TransactionOptions; import com.google.spanner.v1.TransactionSelector; +import io.opencensus.trace.Tracing; +import io.opentelemetry.api.OpenTelemetry; import java.util.ArrayList; import java.util.Collection; import java.util.Collections; import java.util.List; import java.util.Map; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; import org.junit.Before; import org.junit.Test; import org.junit.runner.RunWith; @@ -97,7 +116,7 @@ TestReadContextWithTag build() { } } - private final class TestReadContext extends AbstractReadContext { + private class TestReadContext extends AbstractReadContext { TestReadContext(TestReadContextBuilder builder) { super(builder); } @@ -401,4 +420,416 @@ public void testGetExecuteSqlRequestBuilderWithDirectedReadOptions() { false); assertEquals(DIRECTED_READ_OPTIONS, request.getDirectedReadOptions()); } + + private static ISpan createTestSpan() { + if (SpannerOptions.getActiveTracingFramework() + .equals(SpannerOptions.TracingFramework.OPEN_TELEMETRY)) { + return new OpenTelemetrySpan(mock(io.opentelemetry.api.trace.Span.class)); + } else { + return new OpenCensusSpan(io.opencensus.trace.BlankSpan.INSTANCE); + } + } + + @Test + public void testSynchronousStartFailureDoesNotInvokeOnErrorTwice() { + SessionImpl session = mock(SessionImpl.class); + when(session.getName()).thenReturn("session-1"); + SpannerRpc rpc = mock(SpannerRpc.class); + when(rpc.getExecuteQueryRetrySettings()).thenReturn(RetrySettings.newBuilder().build()); + when(rpc.getExecuteQueryRetryableCodes()).thenReturn(Collections.emptySet()); + when(rpc.executeQuery(any(), any(), any(), any(), anyBoolean())) + .thenThrow( + SpannerExceptionFactory.newSpannerException( + ErrorCode.FAILED_PRECONDITION, "start failure")); + SpannerImpl spanner = mock(SpannerImpl.class); + SpannerOptions spannerOptions = mock(SpannerOptions.class); + when(spanner.getOptions()).thenReturn(spannerOptions); + when(session.getSpanner()).thenReturn(spanner); + when(session.getRequestIdCreator()).thenReturn(NoopRequestIdCreator.INSTANCE); + + AtomicInteger onErrorCalls = new AtomicInteger(); + TestReadContext testContext = + new TestReadContext( + new TestReadContextBuilder() + .setSession(session) + .setRpc(rpc) + .setDefaultQueryOptions(defaultQueryOptions) + .setExecutorProvider( + mock(ExecutorProvider.class, withSettings().withoutAnnotations())) + .setTracer( + new TraceWrapper( + Tracing.getTracer(), OpenTelemetry.noop().getTracer(""), false)) + .setSpan(createTestSpan())) { + @Override + public SpannerException onError( + SpannerException exception, boolean withBeginTransaction, boolean lastStatement) { + onErrorCalls.incrementAndGet(); + return super.onError(exception, withBeginTransaction, lastStatement); + } + }; + + try (ResultSet resultSet = testContext.executeQuery(Statement.of("SELECT 1"))) { + assertThrows(SpannerException.class, resultSet::next); + } + assertEquals(1, onErrorCalls.get()); + } + + @Test + public void testCloseDuringStartStreamDoesNotInvokeOnDoneTwice() { + SessionImpl session = mock(SessionImpl.class); + when(session.getName()).thenReturn("session-1"); + SpannerRpc rpc = mock(SpannerRpc.class); + when(rpc.getExecuteQueryRetrySettings()).thenReturn(RetrySettings.newBuilder().build()); + when(rpc.getExecuteQueryRetryableCodes()).thenReturn(Collections.emptySet()); + SpannerImpl spanner = mock(SpannerImpl.class); + SpannerOptions spannerOptions = mock(SpannerOptions.class); + when(spanner.getOptions()).thenReturn(spannerOptions); + when(session.getSpanner()).thenReturn(spanner); + when(session.getRequestIdCreator()).thenReturn(NoopRequestIdCreator.INSTANCE); + + AtomicInteger onDoneCalls = new AtomicInteger(); + TestReadContext testContext = + new TestReadContext( + new TestReadContextBuilder() + .setSession(session) + .setRpc(rpc) + .setDefaultQueryOptions(defaultQueryOptions) + .setExecutorProvider( + mock(ExecutorProvider.class, withSettings().withoutAnnotations())) + .setTracer( + new TraceWrapper( + Tracing.getTracer(), OpenTelemetry.noop().getTracer(""), false)) + .setSpan(createTestSpan())) { + @Override + public void onDone(boolean withBeginTransaction) { + onDoneCalls.incrementAndGet(); + super.onDone(withBeginTransaction); + } + }; + + try (ResultSet resultSet = testContext.executeQuery(Statement.of("SELECT 1"))) { + when(rpc.executeQuery(any(), any(), any(), any(), anyBoolean())) + .thenAnswer( + invocation -> { + resultSet.close(); + return mock(SpannerRpc.StreamingCall.class); + }); + + assertThrows(SpannerException.class, resultSet::next); + } + assertEquals(1, onDoneCalls.get()); + } + + @Test + public void testCloseDuringStartStreamInvokesOnDoneOnlyOnceWithBeginTransaction() { + SessionImpl session = mock(SessionImpl.class); + when(session.getName()).thenReturn("session-1"); + SpannerRpc rpc = mock(SpannerRpc.class); + when(rpc.getExecuteQueryRetrySettings()).thenReturn(RetrySettings.newBuilder().build()); + when(rpc.getExecuteQueryRetryableCodes()).thenReturn(Collections.emptySet()); + when(rpc.executeQuery(any(), any(), any(), any(), anyBoolean())) + .thenReturn(mock(SpannerRpc.StreamingCall.class)); + SpannerImpl spanner = mock(SpannerImpl.class); + SpannerOptions spannerOptions = mock(SpannerOptions.class); + when(spanner.getOptions()).thenReturn(spannerOptions); + when(session.getSpanner()).thenReturn(spanner); + when(session.getRequestIdCreator()).thenReturn(NoopRequestIdCreator.INSTANCE); + + AtomicInteger onDoneCalls = new AtomicInteger(); + AtomicBoolean closedBeforeStreamStarted = new AtomicBoolean(); + AtomicReference resultSetReference = new AtomicReference<>(); + TestReadContext testContext = + new TestReadContext( + new TestReadContextBuilder() + .setSession(session) + .setRpc(rpc) + .setDefaultQueryOptions(defaultQueryOptions) + .setExecutorProvider( + mock(ExecutorProvider.class, withSettings().withoutAnnotations())) + .setTracer( + new TraceWrapper( + Tracing.getTracer(), OpenTelemetry.noop().getTracer(""), false)) + .setSpan(createTestSpan())) { + @Override + public void onDone(boolean withBeginTransaction) { + onDoneCalls.incrementAndGet(); + super.onDone(withBeginTransaction); + } + + @Override + void onClosedBeforeStreamStarted(boolean withBeginTransaction) { + closedBeforeStreamStarted.set(withBeginTransaction); + } + + @Override + TransactionSelector getTransactionSelector() { + ResultSet currentResultSet = resultSetReference.get(); + if (currentResultSet != null) { + currentResultSet.close(); + } + return TransactionSelector.newBuilder() + .setBegin( + TransactionOptions.newBuilder() + .setReadOnly(TransactionOptions.ReadOnly.getDefaultInstance())) + .build(); + } + }; + + try (ResultSet resultSet = testContext.executeQuery(Statement.of("SELECT 1"))) { + resultSetReference.set(resultSet); + assertThrows(SpannerException.class, resultSet::next); + } + assertEquals(1, onDoneCalls.get()); + assertTrue(closedBeforeStreamStarted.get()); + } + + @Test(timeout = 10_000) + public void testCloseDuringPrepareRequestWhichThrowsCallsOnDoneExactlyOnce() { + SessionImpl session = mock(SessionImpl.class); + when(session.getName()).thenReturn("session-1"); + SpannerRpc rpc = mock(SpannerRpc.class); + when(rpc.getExecuteQueryRetrySettings()).thenReturn(RetrySettings.newBuilder().build()); + when(rpc.getExecuteQueryRetryableCodes()).thenReturn(Collections.emptySet()); + SpannerImpl spanner = mock(SpannerImpl.class); + SpannerOptions spannerOptions = mock(SpannerOptions.class); + when(spanner.getOptions()).thenReturn(spannerOptions); + when(session.getSpanner()).thenReturn(spanner); + when(session.getRequestIdCreator()).thenReturn(NoopRequestIdCreator.INSTANCE); + + AtomicInteger onDoneCalls = new AtomicInteger(); + AtomicReference resultSetReference = new AtomicReference<>(); + TestReadContext testContext = + new TestReadContext( + new TestReadContextBuilder() + .setSession(session) + .setRpc(rpc) + .setDefaultQueryOptions(defaultQueryOptions) + .setExecutorProvider( + mock(ExecutorProvider.class, withSettings().withoutAnnotations())) + .setTracer( + new TraceWrapper( + Tracing.getTracer(), OpenTelemetry.noop().getTracer(""), false)) + .setSpan(createTestSpan())) { + @Override + public void onDone(boolean withBeginTransaction) { + onDoneCalls.incrementAndGet(); + super.onDone(withBeginTransaction); + } + + @Override + TransactionSelector getTransactionSelector() { + ResultSet currentResultSet = resultSetReference.get(); + if (currentResultSet != null) { + currentResultSet.close(); + } + throw new RuntimeException("simulated prepare failure"); + } + }; + + try (ResultSet resultSet = testContext.executeQuery(Statement.of("SELECT 1"))) { + resultSetReference.set(resultSet); + assertThrows(SpannerException.class, resultSet::next); + } + assertEquals(1, onDoneCalls.get()); + verify(rpc, never()).executeQuery(any(), any(), any(), any(), anyBoolean()); + } + + @Test(timeout = 10_000) + public void testCloseDuringResetStreamAfterStartThrowsCallsOnDoneExactlyOnce() { + SessionImpl session = mock(SessionImpl.class); + when(session.getName()).thenReturn("session-1"); + SpannerRpc rpc = mock(SpannerRpc.class); + when(rpc.getExecuteQueryRetrySettings()).thenReturn(RetrySettings.newBuilder().build()); + when(rpc.getExecuteQueryRetryableCodes()).thenReturn(Collections.emptySet()); + SpannerImpl spanner = mock(SpannerImpl.class); + SpannerOptions spannerOptions = mock(SpannerOptions.class); + when(spanner.getOptions()).thenReturn(spannerOptions); + when(session.getSpanner()).thenReturn(spanner); + when(session.getRequestIdCreator()).thenReturn(NoopRequestIdCreator.INSTANCE); + + AtomicInteger onDoneCalls = new AtomicInteger(); + AtomicReference resultSetReference = new AtomicReference<>(); + ISpan mockSpan = mock(ISpan.class); + doAnswer( + invocation -> { + ResultSet currentResultSet = resultSetReference.get(); + if (currentResultSet != null) { + currentResultSet.close(); + } + return null; + }) + .when(mockSpan) + .addAnnotation(eq("Stream broken. Not safe to retry"), any(Throwable.class)); + + TraceWrapper tracer = mock(TraceWrapper.class); + when(tracer.spanBuilderWithExplicitParent(any(), any(), any())).thenReturn(mockSpan); + + TestReadContext testContext = + new TestReadContext( + new TestReadContextBuilder() + .setSession(session) + .setRpc(rpc) + .setDefaultQueryOptions(defaultQueryOptions) + .setExecutorProvider( + mock(ExecutorProvider.class, withSettings().withoutAnnotations())) + .setTracer(tracer) + .setSpan(mockSpan)) { + @Override + public void onDone(boolean withBeginTransaction) { + onDoneCalls.incrementAndGet(); + super.onDone(withBeginTransaction); + } + }; + + when(rpc.executeQuery(any(), any(), any(), any(), anyBoolean())) + .thenThrow( + SpannerExceptionFactory.newSpannerException( + ErrorCode.INTERNAL, "simulated start failure")); + + try (ResultSet resultSet = testContext.executeQuery(Statement.of("SELECT 1"))) { + resultSetReference.set(resultSet); + assertThrows(SpannerException.class, resultSet::next); + } + assertEquals(1, onDoneCalls.get()); + } + + @Test(timeout = 10_000) + public void testStartStreamThrowsConcurrentlyWithCloseFailsTransactionWithStartError() + throws Exception { + SessionImpl session = mock(SessionImpl.class); + when(session.getName()).thenReturn("session-1"); + SpannerRpc rpc = mock(SpannerRpc.class); + when(rpc.getExecuteQueryRetrySettings()).thenReturn(RetrySettings.newBuilder().build()); + when(rpc.getExecuteQueryRetryableCodes()).thenReturn(Collections.emptySet()); + SpannerImpl spanner = mock(SpannerImpl.class); + SpannerOptions spannerOptions = mock(SpannerOptions.class); + when(spanner.getOptions()).thenReturn(spannerOptions); + when(session.getSpanner()).thenReturn(spanner); + when(session.getRequestIdCreator()).thenReturn(NoopRequestIdCreator.INSTANCE); + + SettableApiFuture transactionIdFuture = SettableApiFuture.create(); + AtomicReference resultSetReference = new AtomicReference<>(); + ISpan mockSpan = mock(ISpan.class); + doAnswer( + invocation -> { + ResultSet currentResultSet = resultSetReference.get(); + if (currentResultSet != null) { + currentResultSet.close(); + } + return null; + }) + .when(mockSpan) + .addAnnotation(eq("Stream broken. Not safe to retry"), any(Throwable.class)); + + TraceWrapper tracer = mock(TraceWrapper.class); + when(tracer.spanBuilderWithExplicitParent(any(), any(), any())).thenReturn(mockSpan); + + TestReadContext testContext = + new TestReadContext( + new TestReadContextBuilder() + .setSession(session) + .setRpc(rpc) + .setDefaultQueryOptions(defaultQueryOptions) + .setExecutorProvider( + mock(ExecutorProvider.class, withSettings().withoutAnnotations())) + .setTracer(tracer) + .setSpan(mockSpan)) { + @Override + TransactionSelector getTransactionSelector() { + return TransactionSelector.newBuilder() + .setBegin( + TransactionOptions.newBuilder() + .setReadOnly(TransactionOptions.ReadOnly.getDefaultInstance())) + .build(); + } + + @Override + void onStartFailed( + boolean withBeginTransaction, boolean lastStatement, Throwable throwable) { + if (withBeginTransaction && !transactionIdFuture.isDone()) { + transactionIdFuture.setException(throwable); + } + } + + @Override + public void onDone(boolean withBeginTransaction) { + if (withBeginTransaction && !transactionIdFuture.isDone()) { + transactionIdFuture.setException( + SpannerExceptionFactory.newSpannerException( + ErrorCode.FAILED_PRECONDITION, + "Read-only transaction was closed before the transaction ID was returned by Cloud Spanner")); + } + super.onDone(withBeginTransaction); + } + }; + + when(rpc.executeQuery(any(), any(), any(), any(), anyBoolean())) + .thenThrow( + SpannerExceptionFactory.newSpannerException( + ErrorCode.INTERNAL, "simulated start failure")); + + try (ResultSet resultSet = testContext.executeQuery(Statement.of("SELECT 1"))) { + resultSetReference.set(resultSet); + assertThrows(SpannerException.class, resultSet::next); + } + + ExecutionException executionException = + assertThrows(ExecutionException.class, () -> transactionIdFuture.get(1, TimeUnit.SECONDS)); + assertTrue(executionException.getCause() instanceof SpannerException); + SpannerException spannerException = (SpannerException) executionException.getCause(); + assertEquals(ErrorCode.INTERNAL, spannerException.getErrorCode()); + assertTrue(spannerException.getMessage().contains("simulated start failure")); + verify(session).onReadDone(); + } + + @Test + public void testCloseFreshResultSetBeforeNextPreventsRpcAndThrowsOnNext() { + SessionImpl session = mock(SessionImpl.class); + when(session.getName()).thenReturn("session-1"); + SpannerRpc rpc = mock(SpannerRpc.class); + when(rpc.getExecuteQueryRetrySettings()).thenReturn(RetrySettings.newBuilder().build()); + when(rpc.getExecuteQueryRetryableCodes()).thenReturn(Collections.emptySet()); + SpannerImpl spanner = mock(SpannerImpl.class); + SpannerOptions spannerOptions = mock(SpannerOptions.class); + when(spanner.getOptions()).thenReturn(spannerOptions); + when(session.getSpanner()).thenReturn(spanner); + when(session.getRequestIdCreator()).thenReturn(NoopRequestIdCreator.INSTANCE); + + AtomicInteger onDoneCalls = new AtomicInteger(); + AtomicInteger onErrorCalls = new AtomicInteger(); + TestReadContext testContext = + new TestReadContext( + new TestReadContextBuilder() + .setSession(session) + .setRpc(rpc) + .setDefaultQueryOptions(defaultQueryOptions) + .setExecutorProvider( + mock(ExecutorProvider.class, withSettings().withoutAnnotations())) + .setTracer( + new TraceWrapper( + Tracing.getTracer(), OpenTelemetry.noop().getTracer(""), false)) + .setSpan(createTestSpan())) { + @Override + public void onDone(boolean withBeginTransaction) { + onDoneCalls.incrementAndGet(); + super.onDone(withBeginTransaction); + } + + @Override + public SpannerException onError( + SpannerException error, boolean withBeginTransaction, boolean lastStatement) { + onErrorCalls.incrementAndGet(); + return super.onError(error, withBeginTransaction, lastStatement); + } + }; + + ResultSet resultSet = testContext.executeQuery(Statement.of("SELECT 1")); + resultSet.close(); + + SpannerException exception = assertThrows(SpannerException.class, resultSet::next); + assertEquals(ErrorCode.INTERNAL, exception.getErrorCode()); + verify(rpc, never()).executeQuery(any(), any(), any(), any(), anyBoolean()); + assertEquals(1, onDoneCalls.get()); + assertEquals(0, onErrorCalls.get()); + } } 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 1af639a7dddf..04f797ceb64b 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 @@ -48,6 +48,10 @@ import com.google.spanner.v1.ReadRequest; import io.grpc.ManagedChannelBuilder; import io.grpc.Status; +import io.opencensus.trace.Tracing; +import io.opentelemetry.api.OpenTelemetry; +import io.opentelemetry.api.common.Attributes; +import io.opentelemetry.api.trace.Span; import java.util.Arrays; import java.util.Collections; import java.util.List; @@ -1089,4 +1093,249 @@ public void inlinedBeginReadOnlyTransactionServerErrorCompletesTransactionIdFutu closeFuture.get(10, TimeUnit.SECONDS); assertTrue(closeFuture.isDone()); } + + private static ISpan createCloseOnStartStreamSpan( + AtomicReference closeableReference) { + return new OpenTelemetrySpan(Span.getInvalid()) { + @Override + public void addAnnotation(String message, String key, String value) { + if ("Starting/Resuming stream".equals(message)) { + AutoCloseable closeable = closeableReference.getAndSet(null); + if (closeable != null) { + try { + closeable.close(); + } catch (Exception exception) { + throw new RuntimeException(exception); + } + } + } + } + }; + } + + private static TraceWrapper createCloseOnStartStreamTracer( + AtomicReference closeableReference) { + return new TraceWrapper(Tracing.getTracer(), OpenTelemetry.noop().getTracer("test"), false) { + @Override + ISpan spanBuilderWithExplicitParent( + String spanName, ISpan parentSpan, Attributes attributes) { + return createCloseOnStartStreamSpan(closeableReference); + } + + @Override + IScope withSpan(ISpan span) { + return () -> {}; + } + }; + } + + @Test(timeout = 10000) + public void cancelInlinedBeginQueryBeforeSetStreamFailsWaitingQueryAndCompletesCloseAsync() + throws Exception { + try (ResultSet resultSet = client().singleUse().executeQuery(READ_ONE_KEY_VALUE_STATEMENT)) { + while (resultSet.next()) {} + } + mockSpanner.clearRequests(); + + ReadOnlyTransaction transaction = + client() + .readOnlyTransaction( + TimestampBound.strong(), + Options.beginTransactionOption(Options.BeginTransactionOption.INLINE)); + + AtomicReference resultSetReference = new AtomicReference<>(); + ((AbstractReadContext) transaction).tracer = createCloseOnStartStreamTracer(resultSetReference); + + // Close the result set after startGrpcStreaming passes its initial closed check but before + // startStream calls setStream(stream), exercising !setStream(stream) and onDone(true). + try (AsyncResultSetImpl resultSet = + (AsyncResultSetImpl) transaction.executeQueryAsync(READ_ONE_KEY_VALUE_STATEMENT)) { + resultSetReference.set(resultSet); + resultSet.initiateStreaming(resultSet); + } + + // A subsequent query on the same transaction must fail immediately with FAILED_PRECONDITION + // rather than hanging for 60s waiting on transactionIdFuture. + try (ResultSet secondResultSet = transaction.executeQuery(READ_ONE_KEY_VALUE_STATEMENT)) { + SpannerException exception = assertThrows(SpannerException.class, secondResultSet::next); + assertEquals(ErrorCode.FAILED_PRECONDITION, exception.getErrorCode()); + } + + ApiFuture closeFuture = transaction.closeAsync(); + closeFuture.get(10, TimeUnit.SECONDS); + assertTrue(closeFuture.isDone()); + } + + @Test(timeout = 10000) + public void cancelInlinedBeginReadBeforeSetStreamFailsWaitingReadAndCompletesCloseAsync() + throws Exception { + try (ResultSet resultSet = client().singleUse().executeQuery(READ_ONE_KEY_VALUE_STATEMENT)) { + while (resultSet.next()) {} + } + mockSpanner.clearRequests(); + + ReadOnlyTransaction transaction = + client() + .readOnlyTransaction( + TimestampBound.strong(), + Options.beginTransactionOption(Options.BeginTransactionOption.INLINE)); + + AtomicReference resultSetReference = new AtomicReference<>(); + ((AbstractReadContext) transaction).tracer = createCloseOnStartStreamTracer(resultSetReference); + + try (AsyncResultSetImpl resultSet = + (AsyncResultSetImpl) + transaction.readAsync("TestTable", KeySet.all(), Collections.singletonList("Value"))) { + resultSetReference.set(resultSet); + resultSet.initiateStreaming(resultSet); + } + + try (ResultSet secondResultSet = + transaction.read("TestTable", KeySet.all(), Collections.singletonList("Value"))) { + SpannerException exception = assertThrows(SpannerException.class, secondResultSet::next); + assertEquals(ErrorCode.FAILED_PRECONDITION, exception.getErrorCode()); + } + + ApiFuture closeFuture = transaction.closeAsync(); + closeFuture.get(10, TimeUnit.SECONDS); + assertTrue(closeFuture.isDone()); + } + + @Test(timeout = 10000) + public void readWriteTransactionInlinedBeginClosedBeforeSetStreamFailsWaitingQuery() { + try (ResultSet resultSet = client().singleUse().executeQuery(READ_ONE_KEY_VALUE_STATEMENT)) { + while (resultSet.next()) {} + } + mockSpanner.clearRequests(); + + SpannerException commitException = + assertThrows( + SpannerException.class, + () -> + client() + .readWriteTransaction() + .run( + transaction -> { + AtomicReference resultSetReference = + new AtomicReference<>(); + ((AbstractReadContext) transaction).tracer = + createCloseOnStartStreamTracer(resultSetReference); + + try (ResultSet resultSet = + transaction.executeQuery(READ_ONE_KEY_VALUE_STATEMENT)) { + resultSetReference.set(resultSet); + StreamingUtil.initiateStreaming(resultSet, (chunk, full) -> {}); + } + + try (ResultSet secondResultSet = + transaction.executeQuery(READ_ONE_KEY_VALUE_STATEMENT)) { + SpannerException exception = + assertThrows(SpannerException.class, secondResultSet::next); + assertEquals(ErrorCode.FAILED_PRECONDITION, exception.getErrorCode()); + } + return null; + })); + assertEquals(ErrorCode.FAILED_PRECONDITION, commitException.getErrorCode()); + } + + @Test(timeout = 10000) + public void + inlinedBeginStartFailureWithRuntimeExceptionRacingWithCloseFailsWaitingQueryWithSpannerException() + throws Exception { + AtomicBoolean failNextStreamingQuery = new AtomicBoolean(false); + String endpoint = address.getHostString() + ":" + server.getPort(); + SpannerOptions options = + 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()); + return spannerRpc; + }) + .build(); + + 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)); + + AtomicReference resultSetReference = new AtomicReference<>(); + ((AbstractReadContext) transaction).tracer = + new TraceWrapper(Tracing.getTracer(), OpenTelemetry.noop().getTracer("test"), false) { + @Override + ISpan spanBuilderWithExplicitParent( + String spanName, ISpan parentSpan, Attributes attributes) { + return new OpenTelemetrySpan(Span.getInvalid()) { + @Override + public void addAnnotation(String message, Throwable throwable) { + if ("Stream broken. Not safe to retry".equals(message)) { + AutoCloseable closeable = resultSetReference.getAndSet(null); + if (closeable != null) { + try { + closeable.close(); + } catch (Exception exception) { + throw new RuntimeException(exception); + } + } + } + } + }; + } + + @Override + IScope withSpan(ISpan span) { + return () -> {}; + } + }; + + try (ResultSet firstResultSet = transaction.executeQuery(READ_ONE_KEY_VALUE_STATEMENT)) { + resultSetReference.set(firstResultSet); + assertThrows(SpannerException.class, firstResultSet::next); + } + + 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")); + } + + 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/GrpcResultSetTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/GrpcResultSetTest.java index 36d55cc69c81..be54e4b7bcff 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 @@ -20,24 +20,30 @@ 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.assertSame; import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; +import com.google.api.core.SettableApiFuture; import com.google.api.gax.grpc.GrpcCallContext; +import com.google.api.gax.retrying.RetrySettings; import com.google.api.gax.rpc.ApiCallContext; import com.google.cloud.ByteArray; import com.google.cloud.Date; import com.google.cloud.Timestamp; +import com.google.cloud.spanner.ErrorHandler.DefaultErrorHandler; import com.google.cloud.spanner.SingerProto.Genre; import com.google.cloud.spanner.SingerProto.SingerInfo; +import com.google.cloud.spanner.XGoogSpannerRequestId.NoopRequestIdCreator; import com.google.cloud.spanner.spi.v1.SpannerRpc; import com.google.common.base.Function; import com.google.common.base.Strings; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; +import com.google.common.util.concurrent.Uninterruptibles; import com.google.protobuf.ByteString; import com.google.spanner.v1.ExecuteSqlRequest.QueryMode; import com.google.spanner.v1.MultiplexedSessionPrecommitToken; @@ -46,6 +52,9 @@ import com.google.spanner.v1.ResultSetMetadata; import com.google.spanner.v1.ResultSetStats; import com.google.spanner.v1.Transaction; +import io.opencensus.trace.BlankSpan; +import io.opencensus.trace.Tracing; +import io.opentelemetry.api.OpenTelemetry; import java.math.BigDecimal; import java.math.BigInteger; import java.nio.charset.StandardCharsets; @@ -56,8 +65,17 @@ import java.util.Collections; import java.util.List; import java.util.Map; +import java.util.NoSuchElementException; import java.util.UUID; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; import javax.annotation.Nullable; import org.junit.Before; import org.junit.Test; @@ -1376,4 +1394,481 @@ public void testInitiateStreamingDelegation() { assertFalse(resultSet.initiateStreaming(listener)); } + + @Test(timeout = 10_000) + public void closeRaceCompletesTransactionIdFuture() throws Exception { + SettableApiFuture transactionIdFuture = SettableApiFuture.create(); + CountDownLatch preCloseCheckReturnedFalse = new CountDownLatch(1); + CountDownLatch attachAttempted = new CountDownLatch(1); + AtomicReference closerThread = new AtomicReference<>(); + + AbstractResultSet.Listener listener = + new NoOpListener() { + @Override + public void onDone(boolean withBeginTransaction) { + if (withBeginTransaction) { + transactionIdFuture.setException( + SpannerExceptionFactory.newSpannerException( + ErrorCode.FAILED_PRECONDITION, + AbstractReadContext.RESULT_SET_CLOSED_BEFORE_ID_RETURNED_MESSAGE)); + } + } + }; + + AtomicBoolean closed = new AtomicBoolean(); + AtomicBoolean attached = new AtomicBoolean(); + + AbstractResultSet.CloseableIterator iterator = + new AbstractResultSet.CloseableIterator() { + @Override + public void close(@Nullable String message) { + if (Thread.currentThread() == closerThread.get()) { + preCloseCheckReturnedFalse.countDown(); + Uninterruptibles.awaitUninterruptibly(attachAttempted, 5, TimeUnit.SECONDS); + } + closed.set(true); + } + + @Override + public boolean isWithBeginTransaction() { + return attached.get(); + } + + @Override + public boolean isLastStatement() { + return false; + } + + @Override + public boolean hasNext() { + return false; + } + + @Override + public PartialResultSet next() { + throw new NoSuchElementException(); + } + }; + + try (GrpcResultSet resultSet = new GrpcResultSet(iterator, listener)) { + Thread closer = new Thread(resultSet::close, "closer"); + closerThread.set(closer); + closer.start(); + assertTrue(preCloseCheckReturnedFalse.await(5, TimeUnit.SECONDS)); + + // After the initial pre-close isWithBeginTransaction() returned false and before + // iterator.close() finishes, simulate a concurrent thread attaching a stream with + // withBeginTransaction=true. + attached.set(true); + attachAttempted.countDown(); + closer.join(5_000); + } + + assertTrue(closed.get()); + assertTrue( + "transactionIdFuture was never completed when closed raced with stream attachment", + transactionIdFuture.isDone()); + } + + @Test(timeout = 10_000) + public void closeWhileStreamIsStartingInvokesOnDoneSynchronously() throws Exception { + AtomicInteger onDoneCalls = new AtomicInteger(); + AtomicBoolean observedWithBeginTransaction = new AtomicBoolean(); + AbstractResultSet.Listener listener = + new NoOpListener() { + @Override + public void onDone(boolean withBeginTransaction) { + observedWithBeginTransaction.set(withBeginTransaction); + onDoneCalls.incrementAndGet(); + } + }; + + CountDownLatch startStreamEntered = new CountDownLatch(1); + CountDownLatch allowStartStreamToComplete = new CountDownLatch(1); + + ResumableStreamIterator resumableStreamIterator = + new ResumableStreamIterator( + 10, + "", + new OpenCensusSpan(BlankSpan.INSTANCE), + new TraceWrapper(Tracing.getTracer(), OpenTelemetry.noop().getTracer(""), false), + DefaultErrorHandler.INSTANCE, + RetrySettings.newBuilder().build(), + Collections.emptySet(), + NoopRequestIdCreator.INSTANCE) { + @Override + AbstractResultSet.CloseableIterator startStream( + @Nullable ByteString resumeToken, + @Nullable AsyncResultSet.StreamMessageListener streamListener, + XGoogSpannerRequestId requestId) { + AbstractResultSet.CloseableIterator stream = + mock(AbstractResultSet.CloseableIterator.class); + when(stream.isWithBeginTransaction()).thenReturn(true); + setStream(stream); + startStreamEntered.countDown(); + Uninterruptibles.awaitUninterruptibly(allowStartStreamToComplete, 5, TimeUnit.SECONDS); + return stream; + } + }; + + ExecutorService executor = Executors.newSingleThreadExecutor(); + try (GrpcResultSet resultSet = new GrpcResultSet(resumableStreamIterator, listener)) { + Future streamInitFuture = executor.submit(resultSet::next); + assertTrue(startStreamEntered.await(5, TimeUnit.SECONDS)); + + // Close while startStream is in progress. + // onDone() must be invoked synchronously during close() with the correct withBeginTransaction + // value. + resultSet.close(); + assertEquals(1, onDoneCalls.get()); + assertTrue(observedWithBeginTransaction.get()); + + // Allow startStream to finish. + allowStartStreamToComplete.countDown(); + assertThrows(Exception.class, streamInitFuture::get); + + // onDone() was not invoked a second time. + assertEquals(1, onDoneCalls.get()); + } finally { + executor.shutdownNow(); + } + } + + @Test(timeout = 10_000) + public void concurrentCloseDuringYieldErrorDoesNotInvokeOnDoneBeforeOnErrorCompletes() + throws Exception { + CountDownLatch onErrorEntered = new CountDownLatch(1); + CountDownLatch allowOnErrorToFinish = new CountDownLatch(1); + SettableApiFuture transactionIdFuture = SettableApiFuture.create(); + List callbackOrder = Collections.synchronizedList(new ArrayList<>()); + AtomicInteger onDoneCalls = new AtomicInteger(); + + AbstractResultSet.Listener listener = + new NoOpListener() { + @Override + public SpannerException onError( + SpannerException exception, boolean withBeginTransaction, boolean lastStatement) { + onErrorEntered.countDown(); + Uninterruptibles.awaitUninterruptibly(allowOnErrorToFinish, 5, TimeUnit.SECONDS); + callbackOrder.add("onError"); + if (withBeginTransaction && !transactionIdFuture.isDone()) { + transactionIdFuture.setException(exception); + } + return exception; + } + + @Override + public void onDone(boolean withBeginTransaction) { + callbackOrder.add("onDone"); + onDoneCalls.incrementAndGet(); + if (withBeginTransaction && !transactionIdFuture.isDone()) { + transactionIdFuture.setException( + SpannerExceptionFactory.newSpannerException( + ErrorCode.FAILED_PRECONDITION, + AbstractReadContext.RESULT_SET_CLOSED_BEFORE_ID_RETURNED_MESSAGE)); + } + } + }; + + AbstractResultSet.CloseableIterator failingIterator = + new AbstractResultSet.CloseableIterator() { + @Override + public void close(@Nullable String message) {} + + @Override + public boolean isWithBeginTransaction() { + return true; + } + + @Override + public boolean isLastStatement() { + return false; + } + + @Override + public boolean hasNext() { + throw SpannerExceptionFactory.newSpannerException( + ErrorCode.INVALID_ARGUMENT, "invalid sql"); + } + + @Override + public PartialResultSet next() { + throw new NoSuchElementException(); + } + }; + + ExecutorService executor = Executors.newSingleThreadExecutor(); + try (GrpcResultSet resultSet = new GrpcResultSet(failingIterator, listener)) { + Future nextFuture = executor.submit(resultSet::next); + assertTrue(onErrorEntered.await(5, TimeUnit.SECONDS)); + + // Concurrent close() while yieldError() is inside listener.onError(). + // close() must not invoke listener.onDone() before listener.onError() completes. + resultSet.close(); + assertEquals(0, onDoneCalls.get()); + + allowOnErrorToFinish.countDown(); + ExecutionException nextException = + assertThrows(ExecutionException.class, () -> nextFuture.get(5, TimeUnit.SECONDS)); + assertTrue(nextException.getCause() instanceof SpannerException); + assertEquals( + ErrorCode.INVALID_ARGUMENT, ((SpannerException) nextException.getCause()).getErrorCode()); + + assertEquals(Arrays.asList("onError", "onDone"), callbackOrder); + assertEquals(1, onDoneCalls.get()); + + ExecutionException futureException = + assertThrows( + ExecutionException.class, () -> transactionIdFuture.get(5, TimeUnit.SECONDS)); + assertTrue(futureException.getCause() instanceof SpannerException); + assertEquals( + ErrorCode.INVALID_ARGUMENT, + ((SpannerException) futureException.getCause()).getErrorCode()); + + // Subsequent next() on another thread observes the volatile error. + SpannerException repeatedException = assertThrows(SpannerException.class, resultSet::next); + assertEquals(ErrorCode.INVALID_ARGUMENT, repeatedException.getErrorCode()); + } finally { + executor.shutdownNow(); + } + } + + @Test(timeout = 10_000) + public void initiateStreamingFailureFailsInlinedBeginTransactionAndPreventsSecondRpcOnNext() + throws Exception { + SettableApiFuture transactionIdFuture = SettableApiFuture.create(); + AtomicInteger onErrorCalls = new AtomicInteger(); + AtomicInteger onDoneCalls = new AtomicInteger(); + AtomicBoolean onErrorObservedWithBeginTransaction = new AtomicBoolean(); + + AbstractResultSet.Listener listener = + new NoOpListener() { + @Override + public SpannerException onError( + SpannerException exception, boolean withBeginTransaction, boolean lastStatement) { + onErrorCalls.incrementAndGet(); + onErrorObservedWithBeginTransaction.set(withBeginTransaction); + if (withBeginTransaction && !transactionIdFuture.isDone()) { + transactionIdFuture.setException(exception); + } + return exception; + } + + @Override + public void onDone(boolean withBeginTransaction) { + onDoneCalls.incrementAndGet(); + } + }; + + AtomicInteger initiateStreamingCalls = new AtomicInteger(); + AtomicInteger hasNextCalls = new AtomicInteger(); + AbstractResultSet.CloseableIterator failingIterator = + new AbstractResultSet.CloseableIterator() { + @Override + public void close(@Nullable String message) {} + + @Override + public boolean isWithBeginTransaction() { + return true; + } + + @Override + public boolean isLastStatement() { + return false; + } + + @Override + public boolean initiateStreaming( + AsyncResultSet.StreamMessageListener streamMessageListener) { + initiateStreamingCalls.incrementAndGet(); + throw SpannerExceptionFactory.newSpannerException( + ErrorCode.UNAVAILABLE, "stream startup failed"); + } + + @Override + public boolean hasNext() { + hasNextCalls.incrementAndGet(); + throw SpannerExceptionFactory.newSpannerException( + ErrorCode.UNAVAILABLE, "stream startup failed"); + } + + @Override + public PartialResultSet next() { + throw new NoSuchElementException(); + } + }; + + try (GrpcResultSet resultSet = new GrpcResultSet(failingIterator, listener)) { + SpannerException initiateException = + assertThrows( + SpannerException.class, + () -> resultSet.initiateStreaming(mock(AsyncResultSet.StreamMessageListener.class))); + assertEquals(ErrorCode.UNAVAILABLE, initiateException.getErrorCode()); + assertEquals(1, initiateStreamingCalls.get()); + assertEquals(0, onErrorCalls.get()); + assertEquals(0, onDoneCalls.get()); + + // Calling next() on the failed result set calls yieldError, notifying listener + SpannerException nextException = assertThrows(SpannerException.class, resultSet::next); + assertEquals(ErrorCode.UNAVAILABLE, nextException.getErrorCode()); + assertEquals(1, hasNextCalls.get()); + assertEquals(1, onErrorCalls.get()); + assertTrue(onErrorObservedWithBeginTransaction.get()); + assertEquals(1, onDoneCalls.get()); + + // Subsequent next() immediately rethrows the cached error without invoking hasNext() again + SpannerException cachedException = assertThrows(SpannerException.class, resultSet::next); + assertSame(nextException, cachedException); + assertEquals(1, hasNextCalls.get()); + + // onDone() was not called a second time + assertEquals(1, onDoneCalls.get()); + } + } + + @Test(timeout = 10_000) + public void yieldErrorDoesNotInvokeOnErrorOrOnDoneTwiceWhenAlreadyClosed() { + AtomicInteger onErrorCalls = new AtomicInteger(); + AtomicInteger onDoneCalls = new AtomicInteger(); + AbstractResultSet.Listener listener = + new NoOpListener() { + @Override + public SpannerException onError( + SpannerException exception, boolean withBeginTransaction, boolean lastStatement) { + onErrorCalls.incrementAndGet(); + return exception; + } + + @Override + public void onDone(boolean withBeginTransaction) { + onDoneCalls.incrementAndGet(); + } + }; + + AbstractResultSet.CloseableIterator failingIterator = + new AbstractResultSet.CloseableIterator() { + @Override + public void close(@Nullable String message) {} + + @Override + public boolean isWithBeginTransaction() { + return false; + } + + @Override + public boolean isLastStatement() { + return false; + } + + @Override + public boolean hasNext() { + throw SpannerExceptionFactory.newSpannerException(ErrorCode.INTERNAL, "first error"); + } + + @Override + public PartialResultSet next() { + throw new NoSuchElementException(); + } + }; + + try (GrpcResultSet resultSet = new GrpcResultSet(failingIterator, listener)) { + assertThrows(SpannerException.class, resultSet::next); + assertEquals(1, onErrorCalls.get()); + assertEquals(1, onDoneCalls.get()); + + // Subsequent next() returns the cached error and does NOT call onError or onDone again + assertThrows(SpannerException.class, resultSet::next); + assertEquals(1, onErrorCalls.get()); + assertEquals(1, onDoneCalls.get()); + + // Explicit close() does not invoke onDone() again + resultSet.close(); + assertEquals(1, onDoneCalls.get()); + } + } + + @Test + public void errorOccurringWhileClosedCachesErrorAndSubsequentNextCallsThrow() { + AtomicInteger onErrorCalls = new AtomicInteger(); + AtomicInteger onDoneCalls = new AtomicInteger(); + AbstractResultSet.Listener listener = + new NoOpListener() { + @Override + public SpannerException onError( + SpannerException e, boolean withBeginTransaction, boolean lastStatement) { + onErrorCalls.incrementAndGet(); + return e; + } + + @Override + public void onDone(boolean withBeginTransaction) { + onDoneCalls.incrementAndGet(); + } + }; + + List results = + ImmutableList.of( + PartialResultSet.newBuilder() + .setMetadata(makeMetadata(Type.struct(Type.StructField.of("col", Type.string())))) + .addValues(com.google.protobuf.Value.newBuilder().setStringValue("val1").build()) + .build()); + + AtomicReference resultSetRef = new AtomicReference<>(); + AbstractResultSet.CloseableIterator iterator = + new AbstractResultSet.CloseableIterator() { + private int index = 0; + + @Override + public void close(@Nullable String message) {} + + @Override + public boolean isWithBeginTransaction() { + return false; + } + + @Override + public boolean isLastStatement() { + return false; + } + + @Override + public boolean hasNext() { + if (index == 0) { + return true; + } + // Simulate a concurrent close() happening just before yieldError handles the failure + resultSetRef.get().close(); + throw SpannerExceptionFactory.newSpannerException( + ErrorCode.UNAVAILABLE, "Stream unavailable"); + } + + @Override + public PartialResultSet next() { + if (index < results.size()) { + return results.get(index++); + } + throw new NoSuchElementException(); + } + }; + + GrpcResultSet testResultSet = new GrpcResultSet(iterator, listener); + resultSetRef.set(testResultSet); + + // First next() consumes the first row successfully + assertTrue(testResultSet.next()); + assertEquals("val1", testResultSet.getString("col")); + + // Next next() call encounters error while close() was triggered + SpannerException firstException = assertThrows(SpannerException.class, testResultSet::next); + assertEquals(ErrorCode.UNAVAILABLE, firstException.getErrorCode()); + // onError was skipped because onDone was already dispatched during close() + assertEquals(0, onErrorCalls.get()); + assertEquals(1, onDoneCalls.get()); + + // Subsequent next() calls continue to throw the cached error rather than returning false + SpannerException secondException = assertThrows(SpannerException.class, testResultSet::next); + assertSame(firstException, secondException); + assertEquals(0, onErrorCalls.get()); + assertEquals(1, onDoneCalls.get()); + } } 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 b5591d3072a8..df117d4691c2 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 @@ -53,6 +53,10 @@ import io.grpc.Server; import io.grpc.Status; import io.grpc.inprocess.InProcessServerBuilder; +import io.opencensus.trace.Tracing; +import io.opentelemetry.api.OpenTelemetry; +import io.opentelemetry.api.common.Attributes; +import io.opentelemetry.api.trace.Span; import java.util.ArrayList; import java.util.Arrays; import java.util.Collection; @@ -1179,6 +1183,7 @@ public void closeBeforeTransactionIdReturnedDoesNotRaceWithCancelledConsumer() t // closed, so onError does not overwrite transactionIdFuture with a simulated ABORTED before // onDone(true) sets FAILED_PRECONDITION. CountDownLatch consumerFinishedNextLatch = new CountDownLatch(1); + AtomicReference consumerThread = new AtomicReference<>(); SpannerOptions options = createSpannerOptionsWithStreamInterceptors( (invocation, delegate) -> { @@ -1192,10 +1197,12 @@ public ApiCallContext getCallContext() { @Override public void cancel(String message) { delegateCall.cancel(message); - try { - assertTrue(consumerFinishedNextLatch.await(5, TimeUnit.SECONDS)); - } catch (InterruptedException interruptedException) { - Thread.currentThread().interrupt(); + if (Thread.currentThread() != consumerThread.get()) { + try { + assertTrue(consumerFinishedNextLatch.await(5, TimeUnit.SECONDS)); + } catch (InterruptedException interruptedException) { + Thread.currentThread().interrupt(); + } } } @@ -1224,6 +1231,7 @@ public void request(int numMessages) { Future nextFuture = executor.submit( () -> { + consumerThread.set(Thread.currentThread()); try { resultSet.next(); } finally { @@ -1391,6 +1399,254 @@ public void subsequentQueryAndReadWithStreamRetryAfterResumeToken() { assertNotEquals(ByteString.EMPTY, readRequests.get(1).getResumeToken()); } + @Test(timeout = 10_000) + public void asynchronousStartFailureOfInlinedBeginStatementRetriesTransaction() { + mockSpanner.putStatementResult(StatementResult.query(SELECT1, SELECT1_RESULTSET)); + AtomicBoolean failNextStreamingQuery = new AtomicBoolean(); + AtomicInteger executeQueryCalls = new AtomicInteger(); + SpannerOptions options = + 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 -> { + executeQueryCalls.incrementAndGet(); + if (failNextStreamingQuery.compareAndSet(true, false)) { + throw SpannerExceptionFactory.newSpannerException( + ErrorCode.FAILED_PRECONDITION, + "simulated asynchronous 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()); + return spannerRpc; + }) + .build(); + AtomicInteger attempts = new AtomicInteger(); + ExecutorService asyncExecutor = Executors.newSingleThreadExecutor(); + try (Spanner customSpanner = options.getService()) { + DatabaseClient customClient = + customSpanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]")); + customClient + .readWriteTransaction() + .run( + transaction -> { + int attempt = attempts.incrementAndGet(); + if (attempt == 1) { + failNextStreamingQuery.set(true); + try (AsyncResultSet asyncResultSet = transaction.executeQueryAsync(SELECT1)) { + final SettableApiFuture errorSeen = SettableApiFuture.create(); + asyncResultSet.setCallback( + asyncExecutor, + resultSet -> { + try { + while (resultSet.tryNext() == AsyncResultSet.CursorState.OK) {} + return AsyncResultSet.CallbackResponse.CONTINUE; + } catch (Throwable throwable) { + errorSeen.set(true); + return AsyncResultSet.CallbackResponse.DONE; + } + }); + try { + errorSeen.get(5, TimeUnit.SECONDS); + } catch (Exception exception) { + throw new RuntimeException(exception); + } + } + } + try (ResultSet resultSet = transaction.executeQuery(SELECT1)) { + while (resultSet.next()) {} + } + return null; + }); + } finally { + asyncExecutor.shutdown(); + try { + assertTrue(asyncExecutor.awaitTermination(5, TimeUnit.SECONDS)); + } catch (InterruptedException interruptedException) { + throw new RuntimeException(interruptedException); + } + } + assertEquals(2, attempts.get()); + assertEquals(2, executeQueryCalls.get()); + assertEquals(1, countRequests(ExecuteSqlRequest.class)); + } + + @Test(timeout = 10_000) + public void closeAsyncResultSetBeforeMetadataCompletesTransactionIdFuture() { + mockSpanner.putStatementResult(StatementResult.query(SELECT1, SELECT1_RESULTSET)); + SpannerOptions options = + SpannerOptions.newBuilder() + .setProjectId("[PROJECT]") + .setChannelProvider(channelProvider) + .setCredentials(NoCredentials.getInstance()) + .build(); + ExecutorService executor = Executors.newFixedThreadPool(2); + try (Spanner customSpanner = options.getService()) { + DatabaseClient customClient = + customSpanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]")); + try { + SpannerException exception = + assertThrows( + SpannerException.class, + () -> + customClient + .readWriteTransaction() + .run( + transaction -> { + mockSpanner.freeze(); + Future nextFuture; + try (AsyncResultSet resultSet = + transaction.executeQueryAsync(SELECT1)) { + nextFuture = executor.submit(resultSet::next); + try { + mockSpanner.waitForRequestsToContain( + ExecuteSqlRequest.class, 5000); + } catch (InterruptedException interruptedException) { + throw new RuntimeException(interruptedException); + } + } + mockSpanner.unfreeze(); + try { + nextFuture.get(5, TimeUnit.SECONDS); + } catch (Exception ignored) { + } + try (ResultSet secondResultSet = transaction.executeQuery(SELECT1)) { + secondResultSet.next(); + } + return null; + })); + assertEquals(ErrorCode.FAILED_PRECONDITION, exception.getErrorCode()); + assertTrue( + exception + .getMessage() + .contains("ResultSet was closed before a transaction id was returned")); + } finally { + mockSpanner.unfreeze(); + executor.shutdownNow(); + try { + assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS)); + } catch (InterruptedException interruptedException) { + throw new RuntimeException(interruptedException); + } + } + } + } + + @Test(timeout = 10_000) + public void startFailureRacingWithCloseSimulatesAbortedAndRetriesTransaction() { + mockSpanner.putStatementResult(StatementResult.query(SELECT1, SELECT1_RESULTSET)); + AtomicBoolean failFirstStreamingQuery = new AtomicBoolean(true); + SpannerOptions options = + 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 -> { + if (failFirstStreamingQuery.compareAndSet(true, false)) { + throw SpannerExceptionFactory.newSpannerException( + ErrorCode.UNAVAILABLE, "simulated 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()); + return spannerRpc; + }) + .build(); + AtomicInteger attempts = new AtomicInteger(); + try (Spanner customSpanner = options.getService()) { + DatabaseClient customClient = + customSpanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]")); + customClient + .readWriteTransaction() + .run( + transaction -> { + int attempt = attempts.incrementAndGet(); + if (attempt == 1) { + AtomicReference resultSetReference = new AtomicReference<>(); + ((AbstractReadContext) transaction).tracer = + new TraceWrapper( + Tracing.getTracer(), OpenTelemetry.noop().getTracer("test"), false) { + @Override + ISpan spanBuilderWithExplicitParent( + String spanName, ISpan parentSpan, Attributes attributes) { + return new OpenTelemetrySpan(Span.getInvalid()) { + @Override + public void addAnnotation(String message, Throwable throwable) { + if ("Stream broken. Not safe to retry".equals(message)) { + AutoCloseable closeable = resultSetReference.getAndSet(null); + if (closeable != null) { + try { + closeable.close(); + } catch (Exception exception) { + throw new RuntimeException(exception); + } + } + } + } + }; + } + + @Override + IScope withSpan(ISpan span) { + return () -> {}; + } + }; + try (ResultSet firstResultSet = transaction.executeQuery(SELECT1)) { + resultSetReference.set(firstResultSet); + assertThrows(SpannerException.class, firstResultSet::next); + } + } + // On attempt 1, the second statement observes the simulated ABORTED from + // transactionIdFuture (set by onStartFailed before the racing close()'s onDone() + // could set FAILED_PRECONDITION), triggering a retry on attempt 2 with an explicit + // BeginTransaction RPC. + try (ResultSet secondResultSet = transaction.executeQuery(SELECT1)) { + while (secondResultSet.next()) {} + } + return null; + }); + } + assertEquals(2, attempts.get()); + assertEquals(1, countRequests(BeginTransactionRequest.class)); + } + private int countRequests(Class requestType) { int count = 0; for (AbstractMessage msg : mockSpanner.getRequests()) { diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ResumableStreamIteratorTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ResumableStreamIteratorTest.java index e1ecb1e528f5..12a00a990c05 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ResumableStreamIteratorTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ResumableStreamIteratorTest.java @@ -791,7 +791,7 @@ public void isDataAvailableWithGrpcStreamIterator() { initWithLimit(512); GrpcStreamIterator grpcStream = new GrpcStreamIterator(false, 4, false); SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); - grpcStream.setCall(call, false); + grpcStream.setCall(call); Mockito.when(starter.startStream(null, null)).thenReturn(grpcStream); resumableStreamIterator.initiateStreaming(null); @@ -817,7 +817,7 @@ public void isDataAvailableWithGrpcStreamIteratorWhenErrorOccurs() { initWithLimit(512); GrpcStreamIterator grpcStream = new GrpcStreamIterator(false, 4, false); SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); - grpcStream.setCall(call, false); + grpcStream.setCall(call); when(starter.startStream(null, null)).thenReturn(grpcStream); resumableStreamIterator.initiateStreaming(null); @@ -837,7 +837,7 @@ public void isDataAvailableWithoutResumeTokensWithProductionBufferSize() { initWithLimit(512); GrpcStreamIterator grpcStream = new GrpcStreamIterator(false, 4, false); SpannerRpc.StreamingCall call = Mockito.mock(SpannerRpc.StreamingCall.class); - grpcStream.setCall(call, false); + grpcStream.setCall(call); Mockito.when(starter.startStream(null, null)).thenReturn(grpcStream); resumableStreamIterator.initiateStreaming(null); @@ -912,7 +912,7 @@ public void canPollWithoutErrorOrBlockingWithGrpcStreamIterator() { /* lastStatement= */ false, /* prefetchChunks= */ 4, /* cancelQueryWhenClientIsClosed= */ false); - grpcStream.setCall(mock(SpannerRpc.StreamingCall.class), false); + grpcStream.setCall(mock(SpannerRpc.StreamingCall.class)); when(starter.startStream(null, null)).thenReturn(grpcStream); resumableStreamIterator.initiateStreaming(null); @@ -941,8 +941,8 @@ public void canPollWithoutErrorOrBlockingWithGrpcStreamIterator() { public void closeDuringStartStreamClosesStreamAndPreventsResurrection() { initWithLimit(10); SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); - GrpcStreamIterator grpcStream = new GrpcStreamIterator(true, 4, false); - grpcStream.setCall(call, true); + GrpcStreamIterator grpcStream = new GrpcStreamIterator(true, 4, false, true); + grpcStream.setCall(call); when(starter.startStream(null, null)) .thenAnswer( invocation -> { @@ -963,7 +963,7 @@ public void multipleCloseCallsAreIdempotent() { initWithLimit(10); SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); GrpcStreamIterator grpcStream = new GrpcStreamIterator(false, 4, false); - grpcStream.setCall(call, false); + grpcStream.setCall(call); when(starter.startStream(null, null)).thenReturn(grpcStream); resumableStreamIterator.initiateStreaming(null); grpcStream.consumer().onPartialResultSet(resultSet(ByteString.copyFromUtf8("r1"), "a")); @@ -1022,8 +1022,8 @@ public void isDataAvailableWhenBufferHasItems() { assertFalse(resumableStreamIterator.isLastStatement()); SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); - GrpcStreamIterator grpcStream = new GrpcStreamIterator(true, 4, false); - grpcStream.setCall(call, true); + GrpcStreamIterator grpcStream = new GrpcStreamIterator(true, 4, false, true); + grpcStream.setCall(call); when(starter.startStream(null, null)).thenReturn(grpcStream); Transaction expectedTransaction = @@ -1062,7 +1062,7 @@ public void isWithBeginTransactionAndLastStatementFalseWhenNotSetOnGrpcStreamIte initWithLimit(10); SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); GrpcStreamIterator grpcStream = new GrpcStreamIterator(false, 4, false); - grpcStream.setCall(call, false); + grpcStream.setCall(call); when(starter.startStream(null, null)).thenReturn(grpcStream); resumableStreamIterator.initiateStreaming(null); @@ -1082,7 +1082,7 @@ public void initiateStreamingStartsStreamAndDeliversRowsToAsyncResultSet() throw if (listener != null) { grpcStream.registerListener(listener); } - grpcStream.setCall(call, false); + grpcStream.setCall(call); grpcStream .consumer() .onPartialResultSet( @@ -1144,11 +1144,11 @@ public void isDataAvailableHoldingBackNonResumeTokenChunksPreventsDataDuplicatio initWithLimit(10); GrpcStreamIterator grpcStream1 = new GrpcStreamIterator(false, 4, false); SpannerRpc.StreamingCall call1 = mock(SpannerRpc.StreamingCall.class); - grpcStream1.setCall(call1, false); + grpcStream1.setCall(call1); GrpcStreamIterator grpcStream2 = new GrpcStreamIterator(false, 4, false); SpannerRpc.StreamingCall call2 = mock(SpannerRpc.StreamingCall.class); - grpcStream2.setCall(call2, false); + grpcStream2.setCall(call2); when(starter.startStream(null, null)).thenReturn(grpcStream1); when(starter.startStream(ByteString.copyFromUtf8("r1"), null)).thenReturn(grpcStream2); @@ -1188,7 +1188,7 @@ public void asynchronousChunkProductionWithInterleavedIsDataAvailableAndNext() t initWithLimit(512); GrpcStreamIterator grpcStream = new GrpcStreamIterator(false, 4, false); SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); - grpcStream.setCall(call, false); + grpcStream.setCall(call); when(starter.startStream(null, null)).thenReturn(grpcStream); resumableStreamIterator.initiateStreaming(null); @@ -1414,12 +1414,15 @@ CloseableIterator startStream( startCount.incrementAndGet(); GrpcStreamIterator grpcStream = new GrpcStreamIterator( - /* lastStatement= */ true, 10, /* cancelQueryWhenClientIsClosed= */ false); + /* lastStatement= */ true, + 10, + /* cancelQueryWhenClientIsClosed= */ false, + /* withBeginTransaction= */ true); if (streamMessageListener != null) { grpcStream.registerListener(streamMessageListener); } setStream(grpcStream); - grpcStream.setCall(call, /* withBeginTransaction= */ true); + grpcStream.setCall(call); throw new RuntimeException("start failure"); } }; @@ -1467,7 +1470,7 @@ public void drainPathBufferLimitMissingTokensBecomesUnsafeToRetry() { initWithLimit(1); GrpcStreamIterator grpcStream = new GrpcStreamIterator(false, 4, false); SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); - grpcStream.setCall(call, false); + grpcStream.setCall(call); when(starter.startStream(null, null)).thenReturn(grpcStream); resumableStreamIterator.initiateStreaming(null); @@ -1529,7 +1532,7 @@ CloseableIterator startStream( /* lastStatement= */ false, 10, /* cancelQueryWhenClientIsClosed= */ false); setStream(stream); close("closing before setCall"); - stream.setCall(call, false); + stream.setCall(call); return stream; } }; @@ -1550,9 +1553,12 @@ CloseableIterator startStream( XGoogSpannerRequestId requestId) { GrpcStreamIterator stream = new GrpcStreamIterator( - /* lastStatement= */ true, 10, /* cancelQueryWhenClientIsClosed= */ false); + /* lastStatement= */ true, + 10, + /* cancelQueryWhenClientIsClosed= */ false, + /* withBeginTransaction= */ true); setStream(stream); - stream.setCall(call, /* withBeginTransaction= */ true); + stream.setCall(call); return stream; } }; @@ -1581,7 +1587,7 @@ CloseableIterator startStream( new GrpcStreamIterator( /* lastStatement= */ false, 10, /* cancelQueryWhenClientIsClosed= */ false); setStream(stream); - stream.setCall(call, false); + stream.setCall(call); PartialResultSet.Builder builder = PartialResultSet.newBuilder() .setResumeToken(ByteString.copyFromUtf8("r1")) @@ -1618,7 +1624,7 @@ CloseableIterator startStream( new GrpcStreamIterator( /* lastStatement= */ false, 10, /* cancelQueryWhenClientIsClosed= */ false); setStream(stream); - stream.setCall(call, false); + stream.setCall(call); PartialResultSet firstChunk = PartialResultSet.newBuilder() .setResumeToken(ByteString.copyFromUtf8("r1")) @@ -1657,7 +1663,7 @@ public void closeAfterDrainingChunkPreventsSubsequentReads() { initWithLimit(10); GrpcStreamIterator grpcStream = new GrpcStreamIterator(false, 4, false); SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); - grpcStream.setCall(call, false); + grpcStream.setCall(call); when(starter.startStream(null, null)).thenReturn(grpcStream); resumableStreamIterator.initiateStreaming(null); grpcStream.consumer().onPartialResultSet(resultSet(null, "bufferedBeforeClose")); @@ -1695,7 +1701,7 @@ CloseableIterator startStream( XGoogSpannerRequestId requestId) { GrpcStreamIterator stream = new GrpcStreamIterator(false, 10, false); setStream(stream); - stream.setCall(call, false); + stream.setCall(call); stream.consumer().onPartialResultSet(resultSet(null, "chunkAfterClose")); return stream; } @@ -1838,7 +1844,7 @@ CloseableIterator startStream( }) .when(call) .cancel(Mockito.any()); - stream.setCall(call, false); + stream.setCall(call); return stream; } }; @@ -1883,7 +1889,7 @@ CloseableIterator startStream( XGoogSpannerRequestId requestId) { close("closing"); setStream(stream); - stream.setCall(call, false); + stream.setCall(call); return stream; } }; @@ -1922,11 +1928,11 @@ CloseableIterator startStream( .consumer() .onError(new RetryableException(errorCodeParameter, "transient failure")); // setCall() flushes pending messages and invokes streamMessageListener synchronously. - stream.setCall(call, false); + stream.setCall(call); Uninterruptibles.awaitUninterruptibly( allowFirstStartStreamToReturn, 5, TimeUnit.SECONDS); } else { - stream.setCall(call, false); + stream.setCall(call); stream .consumer() .onPartialResultSet( @@ -2041,7 +2047,7 @@ CloseableIterator startStream( startStreamEntered.countDown(); Uninterruptibles.awaitUninterruptibly(allowStartStreamToFinish, 5, TimeUnit.SECONDS); GrpcStreamIterator stream = new GrpcStreamIterator(false, 10, false); - stream.setCall(mock(SpannerRpc.StreamingCall.class), false); + stream.setCall(mock(SpannerRpc.StreamingCall.class)); stream.consumer().onCompleted(); return stream; } @@ -2151,8 +2157,11 @@ CloseableIterator startStream( startCount.incrementAndGet(); GrpcStreamIterator stream = new GrpcStreamIterator( - /* lastStatement= */ true, 10, /* cancelQueryWhenClientIsClosed= */ false); - stream.setCall(call, /* withBeginTransaction= */ true); + /* lastStatement= */ true, + 10, + /* cancelQueryWhenClientIsClosed= */ false, + /* withBeginTransaction= */ true); + stream.setCall(call); return stream; } }; @@ -2187,7 +2196,7 @@ CloseableIterator startStream( new GrpcStreamIterator( /* lastStatement= */ false, 10, /* cancelQueryWhenClientIsClosed= */ false); setStream(stream); - stream.setCall(call, /* withBeginTransaction= */ false); + stream.setCall(call); close("closed before startStream threw"); throw new RuntimeException("startStream threw after close"); } @@ -2265,7 +2274,7 @@ public boolean isDataAvailable() { return super.isDataAvailable(); } }; - racingStream.setCall(call, false); + racingStream.setCall(call); when(starter.startStream(null, null)).thenReturn(racingStream); resumableStreamIterator.initiateStreaming(null); @@ -2300,7 +2309,7 @@ public boolean canPollWithoutErrorOrBlocking() { return super.canPollWithoutErrorOrBlocking(); } }; - racingStream.setCall(call, false); + racingStream.setCall(call); when(starter.startStream(null, null)).thenReturn(racingStream); resumableStreamIterator.initiateStreaming(null); @@ -2453,12 +2462,15 @@ CloseableIterator startStream( startCount.incrementAndGet(); GrpcStreamIterator stream = new GrpcStreamIterator( - /* lastStatement= */ true, 10, /* cancelQueryWhenClientIsClosed= */ false); + /* lastStatement= */ true, + 10, + /* cancelQueryWhenClientIsClosed= */ false, + /* withBeginTransaction= */ true); if (streamMessageListener != null) { stream.registerListener(streamMessageListener); } setStream(stream); - stream.setCall(call, /* withBeginTransaction= */ true); + stream.setCall(call); consumerWaitingLatch.countDown(); Uninterruptibles.awaitUninterruptibly( allowStartStreamToThrowError, 5, TimeUnit.SECONDS); @@ -2531,16 +2543,18 @@ CloseableIterator startStream( XGoogSpannerRequestId requestId) { GrpcStreamIterator stream = new GrpcStreamIterator( - /* lastStatement= */ true, 10, /* cancelQueryWhenClientIsClosed= */ false) { + /* lastStatement= */ true, + 10, + /* cancelQueryWhenClientIsClosed= */ false, + /* withBeginTransaction= */ true) { @Override public void close(@Nullable String message) { streamCloseCount.incrementAndGet(); super.close(message); } }; - // Register the stream early before stream.setCall(...) sets withBeginTransaction. setStream(stream); - stream.setCall(call, /* withBeginTransaction= */ true); + stream.setCall(call); setCallCompleted.countDown(); Uninterruptibles.awaitUninterruptibly(allowStartStreamToReturn, 5, TimeUnit.SECONDS); return stream; @@ -2585,7 +2599,7 @@ public void canPollWithoutErrorOrBlockingStopsDrainingAndReturnsFalseWhenClosedC initWithLimit(10); SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); GrpcStreamIterator grpcStream = new GrpcStreamIterator(false, 10, false); - grpcStream.setCall(call, false); + grpcStream.setCall(call); when(starter.startStream(null, null)).thenReturn(grpcStream); resumableStreamIterator.initiateStreaming(null); @@ -2631,16 +2645,16 @@ CloseableIterator startStream( startCount.incrementAndGet(); GrpcStreamIterator stream = new GrpcStreamIterator( - /* lastStatement= */ true, 10, /* cancelQueryWhenClientIsClosed= */ false); + /* lastStatement= */ true, + 10, + /* cancelQueryWhenClientIsClosed= */ false, + /* withBeginTransaction= */ true); if (streamMessageListener != null) { stream.registerListener(streamMessageListener); } - // Register the stream early before stream.setCall(call, withBeginTransaction) sets - // withBeginTransaction = true. setStream(stream); - assertFalse(stream.isWithBeginTransaction()); - stream.setCall(call, /* withBeginTransaction= */ true); assertTrue(stream.isWithBeginTransaction()); + stream.setCall(call); stream.requestPrefetchChunks(); return stream; } @@ -2702,8 +2716,11 @@ CloseableIterator startStream( XGoogSpannerRequestId requestId) { GrpcStreamIterator stream = new GrpcStreamIterator( - /* lastStatement= */ true, 10, /* cancelQueryWhenClientIsClosed= */ false); - stream.setCall(call, /* withBeginTransaction= */ true); + /* lastStatement= */ true, + 10, + /* cancelQueryWhenClientIsClosed= */ false, + /* withBeginTransaction= */ true); + stream.setCall(call); stream .consumer() .onError(new RetryableException(errorCodeParameter, "transient error")); @@ -2768,8 +2785,8 @@ CloseableIterator startStream( @Nullable StreamMessageListener streamMessageListener, XGoogSpannerRequestId requestId) { startCount.incrementAndGet(); - GrpcStreamIterator stream = new GrpcStreamIterator(true, 10, false); - stream.setCall(call, true); + GrpcStreamIterator stream = new GrpcStreamIterator(true, 10, false, true); + stream.setCall(call); return stream; } }; @@ -2800,12 +2817,12 @@ CloseableIterator startStream( @Nullable StreamMessageListener streamMessageListener, XGoogSpannerRequestId requestId) { int attempt = streamStartCount.incrementAndGet(); - GrpcStreamIterator stream = new GrpcStreamIterator(true, 10, false); + GrpcStreamIterator stream = new GrpcStreamIterator(true, 10, false, true); if (streamMessageListener != null) { stream.registerListener(streamMessageListener); } setStream(stream); - stream.setCall(call, true); + stream.setCall(call); if (attempt == 1) { stream .consumer() @@ -2902,7 +2919,7 @@ CloseableIterator startStream( XGoogSpannerRequestId requestId) { GrpcStreamIterator stream = new GrpcStreamIterator(false, 10, false); setStream(stream); - stream.setCall(call, false); + stream.setCall(call); stream.consumer().onPartialResultSet(resultSet(null, "chunk1")); return stream; }