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 e013ef6d77b9..57e81c8aed48 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 59e23b54df4e..9285e1c1b087 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; @@ -1399,6 +1403,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; }