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 80a9dfcf5331..299e68fe86a2 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 @@ -140,8 +140,9 @@ public void close() { } closed = true; } - listener.onDone(iterator.isWithBeginTransaction()); + boolean withBeginTransaction = iterator.isWithBeginTransaction(); iterator.close("ResultSet closed"); + listener.onDone(withBeginTransaction); } @Override 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 e0df4c422e3c..78ab074b891e 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 @@ -26,6 +26,7 @@ import com.google.common.util.concurrent.Uninterruptibles; import com.google.spanner.v1.PartialResultSet; import java.time.Duration; +import java.util.Objects; import java.util.Optional; import java.util.concurrent.BlockingQueue; import java.util.concurrent.LinkedBlockingQueue; @@ -38,6 +39,40 @@ @VisibleForTesting class GrpcStreamIterator extends AbstractIterator implements CloseableIterator { + + @VisibleForTesting + static final class Timeout { + final long duration; + final TimeUnit unit; + + Timeout(long duration, TimeUnit unit) { + this.duration = duration; + this.unit = Preconditions.checkNotNull(unit); + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (!(o instanceof Timeout)) { + return false; + } + Timeout that = (Timeout) o; + return this.duration == that.duration && this.unit == that.unit; + } + + @Override + public int hashCode() { + return Objects.hash(duration, unit); + } + + @Override + public String toString() { + return duration + " " + unit; + } + } + private static final Logger logger = Logger.getLogger(GrpcStreamIterator.class.getName()); static final PartialResultSet END_OF_STREAM = PartialResultSet.newBuilder().build(); private final int prefetchChunks; @@ -47,13 +82,15 @@ class GrpcStreamIterator extends AbstractIterator private final BlockingQueue stream; private final Statement statement; - private SpannerRpc.StreamingCall call; + private volatile SpannerRpc.StreamingCall call; + private boolean callCleared; + private boolean closed; + @Nullable private String closeMessage; private volatile boolean withBeginTransaction; private final boolean lastStatement; - private TimeUnit streamWaitTimeoutUnit; - private long streamWaitTimeoutValue; - private SpannerException error; - private boolean done; + @Nullable private volatile Timeout streamWaitTimeout; + @Nullable private volatile SpannerException error; + private volatile boolean done; @VisibleForTesting GrpcStreamIterator( @@ -84,8 +121,28 @@ void registerListener(AsyncResultSet.StreamMessageListener streamMessageListener } public void setCall(SpannerRpc.StreamingCall call, boolean withBeginTransaction) { - this.call = call; - this.withBeginTransaction = withBeginTransaction; + Timeout timeout = extractTimeout(call); + boolean shouldCancel; + String messageToCancel; + synchronized (this) { + this.withBeginTransaction = withBeginTransaction; + shouldCancel = this.closed && !this.callCleared; + if (!this.callCleared && !this.closed) { + this.call = call; + } + if (timeout != null) { + this.streamWaitTimeout = timeout; + } + messageToCancel = this.closeMessage; + } + if (shouldCancel) { + call.cancel(messageToCancel); + } + } + + @VisibleForTesting + @Nullable + static Timeout extractTimeout(SpannerRpc.StreamingCall call) { ApiCallContext callContext = call.getCallContext(); Duration streamWaitTimeout = callContext == null ? null : callContext.getStreamWaitTimeoutDuration(); @@ -94,29 +151,52 @@ public void setCall(SpannerRpc.StreamingCall call, boolean withBeginTransaction) // value is more than 1 second, which is lower than the precision that would normally be // used by the stream watchdog (which uses a precision of 10 seconds by default). if (streamWaitTimeout.getSeconds() > 0L) { - streamWaitTimeoutValue = streamWaitTimeout.getSeconds(); - streamWaitTimeoutUnit = TimeUnit.SECONDS; + return new Timeout(streamWaitTimeout.getSeconds(), TimeUnit.SECONDS); } else if (streamWaitTimeout.getNano() > 0) { - streamWaitTimeoutValue = streamWaitTimeout.getNano(); - streamWaitTimeoutUnit = TimeUnit.NANOSECONDS; + return new Timeout(streamWaitTimeout.getNano(), TimeUnit.NANOSECONDS); } // Note that if the stream-wait-timeout is zero, we won't set a timeout at all. // That is consistent with ApiCallContext#withStreamWaitTimeout(Duration.ZERO). } + return null; } @Override public void close(@Nullable String message) { - if (call != null) { - call.cancel(message); + SpannerRpc.StreamingCall callToCancel = null; + synchronized (this) { + if (!closed) { + closed = true; + closeMessage = message; + if (!callCleared) { + callToCancel = this.call; + this.call = null; + } + } + } + if (callToCancel != null) { + callToCancel.cancel(message); } } @Override @InternalApi public void requestPrefetchChunks() { - Preconditions.checkState(call != null, "The StreamingCall object is not initialized"); - call.request(prefetchChunks); + SpannerRpc.StreamingCall callToRequest = null; + synchronized (this) { + if (this.call != null) { + callToRequest = this.call; + } else if (!this.closed && !this.callCleared) { + // This should never happen during normal operation because startStream(...) always + // initializes the call via setCall(...) before returning the GrpcStreamIterator and + // calling requestPrefetchChunks(). Reaching this point means requestPrefetchChunks() + // was invoked before setCall(...). + throw new IllegalStateException("The StreamingCall object is not initialized"); + } + } + if (callToRequest != null) { + callToRequest.request(prefetchChunks); + } } @Override @@ -131,35 +211,51 @@ public boolean isLastStatement() { @Override protected final PartialResultSet computeNext() { - PartialResultSet next; + PartialResultSet next = takeNextFromStream(); + if (next != END_OF_STREAM) { + SpannerRpc.StreamingCall currentCall = this.call; + if (currentCall != null) { + currentCall.request(1); + } + if (next.getLast()) { + clearCall(); + } + return next; + } + + clearCall(); + if (error != null) { + throw SpannerExceptionFactory.asSpannerException(error); + } + endOfData(); + return null; + } + + private PartialResultSet takeNextFromStream() { try { - if (streamWaitTimeoutUnit != null) { - next = stream.poll(streamWaitTimeoutValue, streamWaitTimeoutUnit); + Timeout timeout = this.streamWaitTimeout; + if (timeout != null) { + PartialResultSet next = stream.poll(timeout.duration, timeout.unit); if (next == null) { throw SpannerExceptionFactory.newSpannerException( ErrorCode.DEADLINE_EXCEEDED, "stream wait timeout"); } + return next; } else { - next = stream.take(); + return stream.take(); } } catch (InterruptedException e) { // Treat interrupt as a request to cancel the read. throw SpannerExceptionFactory.propagateInterrupt(e); } - if (next != END_OF_STREAM) { - call.request(1); - return next; - } + } + private void clearCall() { // All done - close() no longer needs to cancel the call. - call = null; - - if (error != null) { - throw SpannerExceptionFactory.asSpannerException(error); + synchronized (this) { + callCleared = true; + call = null; } - - endOfData(); - return null; } private void addToStream(PartialResultSet results) { @@ -179,16 +275,13 @@ private class ConsumerImpl implements SpannerRpc.ResultStreamConsumer { public void onPartialResultSet(PartialResultSet results) { addToStream(results); if (results.getLast()) { - done = true; - addToStream(END_OF_STREAM); + addEndOfStream(); } } @Override public void onCompleted() { - if (!done) { - addToStream(END_OF_STREAM); - } + addEndOfStream(); } @Override @@ -203,7 +296,20 @@ public void onError(SpannerException exception) { } } error = exception; - addToStream(END_OF_STREAM); + addEndOfStream(); + } + + private void addEndOfStream() { + boolean shouldAddEndOfStream = false; + synchronized (GrpcStreamIterator.this) { + if (!done) { + done = true; + shouldAddEndOfStream = true; + } + } + if (shouldAddEndOfStream) { + addToStream(END_OF_STREAM); + } } @Override diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/CloseSpannerWithOpenResultSetTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/CloseSpannerWithOpenResultSetTest.java index d3aec724c657..e17758d4a0ed 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/CloseSpannerWithOpenResultSetTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/CloseSpannerWithOpenResultSetTest.java @@ -46,6 +46,14 @@ public class CloseSpannerWithOpenResultSetTest extends AbstractMockServerTest { Spanner createSpanner() { + return spannerOptionsBuilder().build().getService(); + } + + Spanner createSpanner(int prefetchChunks) { + return spannerOptionsBuilder().setPrefetchChunks(prefetchChunks).build().getService(); + } + + private SpannerOptions.Builder spannerOptionsBuilder() { return SpannerOptions.newBuilder() .setProjectId("p") .setHost(String.format("http://localhost:%d", getPort())) @@ -54,9 +62,7 @@ Spanner createSpanner() { .setSessionPoolOption( SessionPoolOptions.newBuilder() .setWaitForMinSessionsDuration(Duration.ofSeconds(5L)) - .build()) - .build() - .getService(); + .build()); } @BeforeClass @@ -122,4 +128,65 @@ public void testStreamsAreCleanedUp() throws Exception { 0, ((GapicSpannerRpc) ((SpannerImpl) spanner).getRpc()).getNumActiveResponseObservers()); } } + + @Test(timeout = 30000) + public void testQueryWithPrefetchChunksOne() { + try (Spanner spanner = createSpanner(1)) { + DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of("p", "i", "d")); + for (int i = 0; i < 10; i++) { + try (ResultSet resultSet = client.singleUse().executeQuery(SELECT_RANDOM_STATEMENT)) { + int count = 0; + while (resultSet.next()) { + assertNotNull(resultSet.getCurrentRowAsStruct()); + count++; + } + assertTrue(count > 0); + } + } + assertEquals( + 0, ((GapicSpannerRpc) ((SpannerImpl) spanner).getRpc()).getNumActiveResponseObservers()); + } + } + + @Test(timeout = 30000) + public void testQueryWithPrefetchChunksOneAndError() { + String errorSql = "select * from will_fail"; + Statement errorStatement = Statement.of(errorSql); + mockSpanner.putStatementResult( + StatementResult.exception( + errorStatement, + Status.INVALID_ARGUMENT.withDescription("Invalid syntax").asRuntimeException())); + try (Spanner spanner = createSpanner(1)) { + DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of("p", "i", "d")); + try (ResultSet resultSet = client.singleUse().executeQuery(errorStatement)) { + SpannerException exception = assertThrows(SpannerException.class, resultSet::next); + assertEquals(ErrorCode.INVALID_ARGUMENT, exception.getErrorCode()); + } + assertEquals( + 0, ((GapicSpannerRpc) ((SpannerImpl) spanner).getRpc()).getNumActiveResponseObservers()); + } + } + + @Test(timeout = 30000) + public void testCloseResultSetCleansUpResponseObservers() { + try (Spanner spanner = createSpanner()) { + DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of("p", "i", "d")); + // 1. Partial consumption then close: + try (ResultSet resultSet = client.singleUse().executeQuery(SELECT_RANDOM_STATEMENT)) { + assertTrue(resultSet.next()); + assertNotNull(resultSet.getCurrentRowAsStruct()); + } + assertEquals( + 0, ((GapicSpannerRpc) ((SpannerImpl) spanner).getRpc()).getNumActiveResponseObservers()); + + // 2. Full consumption then close: + try (ResultSet resultSet = client.singleUse().executeQuery(SELECT_RANDOM_STATEMENT)) { + while (resultSet.next()) { + assertNotNull(resultSet.getCurrentRowAsStruct()); + } + } + assertEquals( + 0, ((GapicSpannerRpc) ((SpannerImpl) spanner).getRpc()).getNumActiveResponseObservers()); + } + } } 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 4007c972c24e..85f94ebc6d74 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 @@ -54,6 +54,7 @@ import java.util.List; import java.util.Map; import java.util.UUID; +import java.util.concurrent.atomic.AtomicBoolean; import javax.annotation.Nullable; import org.junit.Before; import org.junit.Test; @@ -1254,4 +1255,51 @@ public void shouldThrowDeadlineExceededIfLastTrueIsNotReceived() { assertEquals("DEADLINE_EXCEEDED: stream wait timeout", spannerException.getMessage()); consumer.onCompleted(); } + + @Test + public void testClosePreservesWithBeginTransactionForListener() { + AtomicBoolean listenerReceivedWithBeginTransaction = new AtomicBoolean(); + AbstractResultSet.Listener listener = + new NoOpListener() { + @Override + public void onDone(boolean withBeginTransaction) { + listenerReceivedWithBeginTransaction.set(withBeginTransaction); + } + }; + AbstractResultSet.CloseableIterator iterator = + new AbstractResultSet.CloseableIterator() { + private boolean closed; + + @Override + public void close(@Nullable String message) { + this.closed = true; + } + + @Override + public boolean isWithBeginTransaction() { + // Emulate ResumableStreamIterator: stream is nulled upon close, + // so isWithBeginTransaction() returns false once closed. + return !closed; + } + + @Override + public boolean isLastStatement() { + return false; + } + + @Override + public boolean hasNext() { + return false; + } + + @Override + public PartialResultSet next() { + throw new java.util.NoSuchElementException(); + } + }; + + GrpcResultSet resultSet = new GrpcResultSet(iterator, listener); + resultSet.close(); + assertTrue(listenerReceivedWithBeginTransaction.get()); + } } diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/GrpcStreamIteratorTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/GrpcStreamIteratorTest.java new file mode 100644 index 000000000000..c451fe9828e6 --- /dev/null +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/GrpcStreamIteratorTest.java @@ -0,0 +1,801 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.spanner; + +import static org.awaitility.Awaitility.await; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyInt; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import com.google.api.gax.grpc.GrpcCallContext; +import com.google.api.gax.rpc.ApiCallContext; +import com.google.cloud.spanner.spi.v1.SpannerRpc; +import com.google.protobuf.ByteString; +import com.google.protobuf.Value; +import com.google.spanner.v1.PartialResultSet; +import java.time.Duration; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class GrpcStreamIteratorTest { + + @Test + public void closeBeforeSetCallCancelsUnderlyingCallWithMessage() { + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); + + GrpcStreamIterator streamIterator = + new GrpcStreamIterator( + /* statement= */ null, + /* lastStatement= */ false, + /* prefetchChunks= */ 10, + /* cancelQueryWhenClientIsClosed= */ false); + + streamIterator.close("Cancelled before setCall"); + streamIterator.setCall(call, /* withBeginTransaction= */ false); + + verify(call).cancel("Cancelled before setCall"); + } + + @Test + public void closeAfterSetCallCancelsUnderlyingCallWithMessage() { + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); + + GrpcStreamIterator streamIterator = + new GrpcStreamIterator( + /* statement= */ null, + /* lastStatement= */ false, + /* prefetchChunks= */ 10, + /* cancelQueryWhenClientIsClosed= */ false); + + streamIterator.setCall(call, /* withBeginTransaction= */ false); + streamIterator.close("Cancelled after setCall"); + + verify(call).cancel("Cancelled after setCall"); + } + + @Test + public void closeAfterLastTrueChunkConsumedDoesNotCancelCall() { + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); + + GrpcStreamIterator streamIterator = + new GrpcStreamIterator( + /* statement= */ null, + /* lastStatement= */ false, + /* prefetchChunks= */ 10, + /* cancelQueryWhenClientIsClosed= */ false); + + streamIterator.setCall(call, /* withBeginTransaction= */ false); + + PartialResultSet lastChunk = + PartialResultSet.newBuilder() + .setResumeToken(ByteString.copyFromUtf8("token-1")) + .setLast(true) + .addValues(Value.newBuilder().setStringValue("value-1")) + .build(); + + streamIterator.consumer().onPartialResultSet(lastChunk); + assertTrue(streamIterator.hasNext()); + assertEquals(lastChunk, streamIterator.next()); + + streamIterator.close("Closing after finished"); + + verify(call, never()).cancel(any()); + } + + @Test + public void closeBeforeLastTrueChunkConsumedCancelsCall() { + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); + + GrpcStreamIterator streamIterator = + new GrpcStreamIterator( + /* statement= */ null, + /* lastStatement= */ false, + /* prefetchChunks= */ 10, + /* cancelQueryWhenClientIsClosed= */ false); + + streamIterator.setCall(call, /* withBeginTransaction= */ false); + + PartialResultSet lastChunk = + PartialResultSet.newBuilder() + .setResumeToken(ByteString.copyFromUtf8("token-1")) + .setLast(true) + .addValues(Value.newBuilder().setStringValue("value-1")) + .build(); + + streamIterator.consumer().onPartialResultSet(lastChunk); + streamIterator.close("Closing before consumption"); + + verify(call).cancel("Closing before consumption"); + } + + @Test + public void lastTrueConsumedBeforeSetCallDoesNotCancelWhenClosed() { + for (boolean closeBeforeSetCall : new boolean[] {true, false}) { + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); + + GrpcStreamIterator streamIterator = + new GrpcStreamIterator( + /* statement= */ null, + /* lastStatement= */ false, + /* prefetchChunks= */ 10, + /* cancelQueryWhenClientIsClosed= */ false); + + PartialResultSet chunk = + PartialResultSet.newBuilder() + .setResumeToken(ByteString.copyFromUtf8("token-1")) + .setLast(true) + .addValues(Value.newBuilder().setStringValue("value-1")) + .build(); + + streamIterator.consumer().onPartialResultSet(chunk); + assertTrue(streamIterator.hasNext()); + assertEquals(chunk, streamIterator.next()); + + if (closeBeforeSetCall) { + streamIterator.close("Closed before setCall"); + } + streamIterator.setCall(call, /* withBeginTransaction= */ false); + if (!closeBeforeSetCall) { + streamIterator.close("Closed after setCall"); + } + + verify(call, never()).cancel(any()); + } + } + + @Test(timeout = 5000) + public void consumerPrefetchChunksOneDoesNotDeadlockOnCompletion() { + // When prefetchChunks is 1, stream queue capacity is 2 (prefetchChunks + 1). + // Delivering last=true followed by onCompleted() must not deadlock. + GrpcStreamIterator streamIterator = + new GrpcStreamIterator( + /* statement= */ null, + /* lastStatement= */ false, + /* prefetchChunks= */ 1, + /* cancelQueryWhenClientIsClosed= */ false); + + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); + streamIterator.setCall(call, /* withBeginTransaction= */ false); + + PartialResultSet lastChunk = + PartialResultSet.newBuilder() + .setResumeToken(ByteString.copyFromUtf8("token-1")) + .setLast(true) + .addValues(Value.newBuilder().setStringValue("value-1")) + .build(); + + streamIterator.consumer().onPartialResultSet(lastChunk); + streamIterator.consumer().onCompleted(); + + assertTrue(streamIterator.hasNext()); + assertEquals(lastChunk, streamIterator.next()); + assertFalse(streamIterator.hasNext()); + } + + @Test(timeout = 5000) + public void consumerPrefetchChunksOneDoesNotDeadlockOnError() { + // When prefetchChunks is 1, stream queue capacity is 2. + // Delivering last=true followed by onError() must not deadlock. + GrpcStreamIterator streamIterator = + new GrpcStreamIterator( + /* statement= */ null, + /* lastStatement= */ false, + /* prefetchChunks= */ 1, + /* cancelQueryWhenClientIsClosed= */ false); + + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); + streamIterator.setCall(call, /* withBeginTransaction= */ false); + + PartialResultSet lastChunk = + PartialResultSet.newBuilder() + .setResumeToken(ByteString.copyFromUtf8("token-1")) + .setLast(true) + .addValues(Value.newBuilder().setStringValue("value-1")) + .build(); + + streamIterator.consumer().onPartialResultSet(lastChunk); + SpannerException spannerException = + SpannerExceptionFactory.newSpannerException(ErrorCode.INTERNAL, "stream error"); + streamIterator.consumer().onError(spannerException); + + assertTrue(streamIterator.hasNext()); + assertEquals(lastChunk, streamIterator.next()); + SpannerException thrown = assertThrows(SpannerException.class, streamIterator::hasNext); + assertEquals(ErrorCode.INTERNAL, thrown.getErrorCode()); + } + + @Test + public void requestPrefetchChunksValidatesInitialization() { + GrpcStreamIterator streamIterator = + new GrpcStreamIterator( + /* statement= */ null, + /* lastStatement= */ false, + /* prefetchChunks= */ 5, + /* cancelQueryWhenClientIsClosed= */ false); + + assertThrows(IllegalStateException.class, streamIterator::requestPrefetchChunks); + + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); + streamIterator.setCall(call, /* withBeginTransaction= */ false); + + streamIterator.requestPrefetchChunks(); + verify(call).request(5); + } + + @Test + public void statePreservedAcrossClose() { + GrpcStreamIterator streamIterator = + new GrpcStreamIterator( + /* statement= */ null, + /* lastStatement= */ true, + /* prefetchChunks= */ 2, + /* cancelQueryWhenClientIsClosed= */ false); + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); + streamIterator.setCall(call, /* withBeginTransaction= */ true); + + assertTrue(streamIterator.isLastStatement()); + assertTrue(streamIterator.isWithBeginTransaction()); + + streamIterator.close("Closed"); + + assertTrue(streamIterator.isLastStatement()); + assertTrue(streamIterator.isWithBeginTransaction()); + } + + @Test + public void onErrorSetsStatementOnException() { + Statement statement = Statement.of("SELECT 1"); + GrpcStreamIterator streamIterator = + new GrpcStreamIterator( + statement, + /* lastStatement= */ false, + /* prefetchChunks= */ 2, + /* cancelQueryWhenClientIsClosed= */ false); + + SpannerException originalException = + SpannerExceptionFactory.newSpannerException(ErrorCode.UNAVAILABLE, "Unavailable"); + streamIterator.consumer().onError(originalException); + + SpannerException thrown = assertThrows(SpannerException.class, streamIterator::hasNext); + assertTrue(thrown.getMessage().contains(statement.getSql())); + } + + @Test + public void streamWaitTimeoutThrowsDeadlineExceeded() { + ApiCallContext callContext = + GrpcCallContext.createDefault().withStreamWaitTimeoutDuration(Duration.ofMillis(50)); + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + when(call.getCallContext()).thenReturn(callContext); + + GrpcStreamIterator streamIterator = + new GrpcStreamIterator( + /* statement= */ null, + /* lastStatement= */ false, + /* prefetchChunks= */ 2, + /* cancelQueryWhenClientIsClosed= */ false); + + streamIterator.setCall(call, /* withBeginTransaction= */ false); + + SpannerException thrown = assertThrows(SpannerException.class, streamIterator::hasNext); + assertEquals(ErrorCode.DEADLINE_EXCEEDED, thrown.getErrorCode()); + } + + @Test + public void closeBeforeSetCallDoesNotRetainCall() { + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); + + GrpcStreamIterator streamIterator = + new GrpcStreamIterator( + /* statement= */ null, + /* lastStatement= */ false, + /* prefetchChunks= */ 5, + /* cancelQueryWhenClientIsClosed= */ false); + + streamIterator.close("Closed before setCall"); + streamIterator.setCall(call, /* withBeginTransaction= */ false); + + verify(call).cancel("Closed before setCall"); + // Since streamIterator was already closed, requestPrefetchChunks should safely no-op + // without requesting chunks from the cancelled call. + streamIterator.requestPrefetchChunks(); + verify(call, never()).request(anyInt()); + } + + @Test + public void closeAfterSetCallClearsCallReference() { + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); + + GrpcStreamIterator streamIterator = + new GrpcStreamIterator( + /* statement= */ null, + /* lastStatement= */ false, + /* prefetchChunks= */ 5, + /* cancelQueryWhenClientIsClosed= */ false); + + streamIterator.setCall(call, /* withBeginTransaction= */ false); + streamIterator.close("Closed after setCall"); + + verify(call).cancel("Closed after setCall"); + // Since streamIterator is closed, requestPrefetchChunks should safely no-op. + streamIterator.requestPrefetchChunks(); + verify(call, never()).request(anyInt()); + } + + @Test(timeout = 10000) + public void concurrentCloseAndSetCallCancelsCallDeterministically() throws Exception { + ExecutorService executor = Executors.newFixedThreadPool(2); + try { + for (int i = 0; i < 50; i++) { + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); + + GrpcStreamIterator streamIterator = + new GrpcStreamIterator( + /* statement= */ null, + /* lastStatement= */ false, + /* prefetchChunks= */ 5, + /* cancelQueryWhenClientIsClosed= */ false); + + CountDownLatch startLatch = new CountDownLatch(1); + CountDownLatch doneLatch = new CountDownLatch(2); + + executor.submit( + () -> { + try { + startLatch.await(); + streamIterator.setCall(call, /* withBeginTransaction= */ false); + } catch (InterruptedException ignored) { + } finally { + doneLatch.countDown(); + } + }); + + executor.submit( + () -> { + try { + startLatch.await(); + streamIterator.close("Concurrent cancellation"); + } catch (InterruptedException ignored) { + } finally { + doneLatch.countDown(); + } + }); + + startLatch.countDown(); + assertTrue(doneLatch.await(5, TimeUnit.SECONDS)); + + verify(call).cancel("Concurrent cancellation"); + streamIterator.requestPrefetchChunks(); + verify(call, never()).request(anyInt()); + } + } finally { + executor.shutdownNow(); + } + } + + @Test(timeout = 5000) + public void concurrentCloseUnblocksConsumerWaitingOnStream() throws Exception { + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); + + GrpcStreamIterator streamIterator = + new GrpcStreamIterator( + /* statement= */ null, + /* lastStatement= */ false, + /* prefetchChunks= */ 5, + /* cancelQueryWhenClientIsClosed= */ false); + streamIterator.setCall(call, /* withBeginTransaction= */ false); + + AtomicReference consumerError = new AtomicReference<>(); + Thread consumerThread = + new Thread( + () -> { + try { + SpannerException spannerException = + assertThrows(SpannerException.class, streamIterator::hasNext); + assertEquals(ErrorCode.CANCELLED, spannerException.getErrorCode()); + } catch (Throwable t) { + consumerError.set(t); + } + }); + consumerThread.start(); + + // Deterministically wait for the consumer thread to block on stream.take() without sleeping. + await() + .atMost(Duration.ofSeconds(5)) + .until(() -> consumerThread.getState() == Thread.State.WAITING); + + streamIterator.close("User cancelled"); + verify(call).cancel("User cancelled"); + + // gRPC layer delivers error on cancellation + streamIterator + .consumer() + .onError(SpannerExceptionFactory.newSpannerException(ErrorCode.CANCELLED, "Cancelled")); + + consumerThread.join(5000); + assertFalse(consumerThread.isAlive()); + if (consumerError.get() != null) { + throw new AssertionError(consumerError.get()); + } + } + + @Test + public void setCallWithNullCallContextDoesNotThrow() { + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + when(call.getCallContext()).thenReturn(null); + + GrpcStreamIterator streamIterator = + new GrpcStreamIterator( + /* statement= */ null, + /* lastStatement= */ false, + /* prefetchChunks= */ 5, + /* cancelQueryWhenClientIsClosed= */ false); + streamIterator.setCall(call, /* withBeginTransaction= */ false); + + streamIterator.requestPrefetchChunks(); + verify(call).request(5); + } + + @Test + public void streamWaitTimeoutWithSecondsUnitThrowsDeadlineExceeded() { + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + when(call.getCallContext()) + .thenReturn( + GrpcCallContext.createDefault().withStreamWaitTimeoutDuration(Duration.ofSeconds(1))); + + GrpcStreamIterator streamIterator = + new GrpcStreamIterator( + /* statement= */ null, + /* lastStatement= */ false, + /* prefetchChunks= */ 5, + /* cancelQueryWhenClientIsClosed= */ false); + streamIterator.setCall(call, /* withBeginTransaction= */ false); + + SpannerException spannerException = + assertThrows(SpannerException.class, streamIterator::hasNext); + assertEquals(ErrorCode.DEADLINE_EXCEEDED, spannerException.getErrorCode()); + } + + @Test + public void duplicateCloseIsIdempotent() { + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); + + GrpcStreamIterator streamIterator = + new GrpcStreamIterator( + /* statement= */ null, + /* lastStatement= */ false, + /* prefetchChunks= */ 5, + /* cancelQueryWhenClientIsClosed= */ false); + streamIterator.setCall(call, /* withBeginTransaction= */ false); + + streamIterator.close("First close"); + streamIterator.close("Second close"); + + verify(call, times(1)).cancel("First close"); + verify(call, never()).cancel("Second close"); + } + + @Test + public void requestPrefetchChunksSafelyNoOpsWhenCallCleared() { + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); + + GrpcStreamIterator streamIterator = + new GrpcStreamIterator( + /* statement= */ null, + /* lastStatement= */ false, + /* prefetchChunks= */ 5, + /* cancelQueryWhenClientIsClosed= */ false); + streamIterator.setCall(call, /* withBeginTransaction= */ false); + + streamIterator + .consumer() + .onPartialResultSet(PartialResultSet.newBuilder().setLast(true).build()); + assertTrue(streamIterator.hasNext()); + assertNotNull(streamIterator.next()); + assertFalse(streamIterator.hasNext()); + + // Call is now cleared; requesting prefetch chunks should safely no-op + streamIterator.requestPrefetchChunks(); + // Only the request during computeNext should have occurred + verify(call, times(1)).request(1); + } + + @Test(timeout = 10000) + public void concurrentSetCallAndRequestPrefetchChunks() throws Exception { + ExecutorService executor = Executors.newFixedThreadPool(2); + try { + for (int i = 0; i < 50; i++) { + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); + + GrpcStreamIterator streamIterator = + new GrpcStreamIterator( + /* statement= */ null, + /* lastStatement= */ false, + /* prefetchChunks= */ 5, + /* cancelQueryWhenClientIsClosed= */ false); + + CountDownLatch startLatch = new CountDownLatch(1); + CountDownLatch doneLatch = new CountDownLatch(2); + AtomicReference error = new AtomicReference<>(); + AtomicBoolean requestPrefetchSucceeded = new AtomicBoolean(false); + + executor.submit( + () -> { + try { + startLatch.await(); + streamIterator.setCall(call, /* withBeginTransaction= */ false); + } catch (Throwable t) { + error.set(t); + } finally { + doneLatch.countDown(); + } + }); + + executor.submit( + () -> { + try { + startLatch.await(); + streamIterator.requestPrefetchChunks(); + requestPrefetchSucceeded.set(true); + } catch (IllegalStateException ignored) { + // If requestPrefetchChunks runs strictly before setCall begins, throwing + // IllegalStateException is expected. + } catch (Throwable t) { + error.set(t); + } finally { + doneLatch.countDown(); + } + }); + + startLatch.countDown(); + assertTrue(doneLatch.await(5, TimeUnit.SECONDS)); + if (error.get() != null) { + throw new AssertionError(error.get()); + } + if (requestPrefetchSucceeded.get()) { + verify(call).request(5); + } else { + verify(call, never()).request(anyInt()); + } + } + } finally { + executor.shutdownNow(); + } + } + + @Test(timeout = 10000) + public void concurrentAddEndOfStreamFromMultipleCallbacks() throws Exception { + ExecutorService executor = Executors.newFixedThreadPool(3); + try { + for (int i = 0; i < 50; i++) { + GrpcStreamIterator streamIterator = + new GrpcStreamIterator( + /* statement= */ null, + /* lastStatement= */ false, + /* prefetchChunks= */ 1, + /* cancelQueryWhenClientIsClosed= */ false); + + SpannerRpc.ResultStreamConsumer consumer = streamIterator.consumer(); + CountDownLatch startLatch = new CountDownLatch(1); + CountDownLatch doneLatch = new CountDownLatch(3); + AtomicReference error = new AtomicReference<>(); + + executor.submit( + () -> { + try { + startLatch.await(); + consumer.onPartialResultSet(PartialResultSet.newBuilder().setLast(true).build()); + } catch (Throwable t) { + error.set(t); + } finally { + doneLatch.countDown(); + } + }); + + executor.submit( + () -> { + try { + startLatch.await(); + consumer.onCompleted(); + } catch (Throwable t) { + error.set(t); + } finally { + doneLatch.countDown(); + } + }); + + executor.submit( + () -> { + try { + startLatch.await(); + consumer.onError( + SpannerExceptionFactory.newSpannerException(ErrorCode.INTERNAL, "test-error")); + } catch (Throwable t) { + error.set(t); + } finally { + doneLatch.countDown(); + } + }); + + startLatch.countDown(); + assertTrue(doneLatch.await(5, TimeUnit.SECONDS)); + if (error.get() != null) { + throw new AssertionError(error.get()); + } + + // Verify that the iterator can be safely drained without blocking or deadlocks. + try { + while (streamIterator.hasNext()) { + streamIterator.next(); + } + } catch (SpannerException e) { + // If onError won the race, an exception is expected. + assertEquals(ErrorCode.INTERNAL, e.getErrorCode()); + } + } + } finally { + executor.shutdownNow(); + } + } + + @Test + public void closeOnAlreadyCompletedStreamDoesNotDeliverSpuriousError() { + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); + + GrpcStreamIterator streamIterator = + new GrpcStreamIterator( + /* statement= */ null, + /* lastStatement= */ false, + /* prefetchChunks= */ 1, + /* cancelQueryWhenClientIsClosed= */ false); + streamIterator.setCall(call, /* withBeginTransaction= */ false); + + streamIterator + .consumer() + .onPartialResultSet(PartialResultSet.newBuilder().setLast(true).build()); + assertTrue(streamIterator.hasNext()); + assertNotNull(streamIterator.next()); + assertFalse(streamIterator.hasNext()); + + // Closing already completed stream should safely no-op without cancelling or setting errors + streamIterator.close("Post-completion close"); + verify(call, never()).cancel(any()); + assertFalse(streamIterator.hasNext()); + } + + @Test + public void testExtractTimeout() { + SpannerRpc.StreamingCall callWithNullContext = mock(SpannerRpc.StreamingCall.class); + when(callWithNullContext.getCallContext()).thenReturn(null); + assertNull(GrpcStreamIterator.extractTimeout(callWithNullContext)); + + SpannerRpc.StreamingCall callWithNoTimeout = mock(SpannerRpc.StreamingCall.class); + when(callWithNoTimeout.getCallContext()).thenReturn(GrpcCallContext.createDefault()); + assertNull(GrpcStreamIterator.extractTimeout(callWithNoTimeout)); + + SpannerRpc.StreamingCall callWithZeroTimeout = mock(SpannerRpc.StreamingCall.class); + when(callWithZeroTimeout.getCallContext()) + .thenReturn(GrpcCallContext.createDefault().withStreamWaitTimeoutDuration(Duration.ZERO)); + assertNull(GrpcStreamIterator.extractTimeout(callWithZeroTimeout)); + + SpannerRpc.StreamingCall callWithSecondsTimeout = mock(SpannerRpc.StreamingCall.class); + when(callWithSecondsTimeout.getCallContext()) + .thenReturn( + GrpcCallContext.createDefault().withStreamWaitTimeoutDuration(Duration.ofSeconds(10))); + assertEquals( + new GrpcStreamIterator.Timeout(10, TimeUnit.SECONDS), + GrpcStreamIterator.extractTimeout(callWithSecondsTimeout)); + + SpannerRpc.StreamingCall callWithNanosTimeout = mock(SpannerRpc.StreamingCall.class); + when(callWithNanosTimeout.getCallContext()) + .thenReturn( + GrpcCallContext.createDefault() + .withStreamWaitTimeoutDuration(Duration.ofNanos(500_000))); + assertEquals( + new GrpcStreamIterator.Timeout(500_000, TimeUnit.NANOSECONDS), + GrpcStreamIterator.extractTimeout(callWithNanosTimeout)); + } + + @Test + public void testTimeoutEqualsHashCodeAndToString() { + GrpcStreamIterator.Timeout timeout1 = new GrpcStreamIterator.Timeout(10, TimeUnit.SECONDS); + GrpcStreamIterator.Timeout timeout2 = new GrpcStreamIterator.Timeout(10, TimeUnit.SECONDS); + GrpcStreamIterator.Timeout timeout3 = new GrpcStreamIterator.Timeout(20, TimeUnit.SECONDS); + GrpcStreamIterator.Timeout timeout4 = new GrpcStreamIterator.Timeout(10, TimeUnit.MILLISECONDS); + + assertEquals(timeout1, timeout2); + assertEquals(timeout1.hashCode(), timeout2.hashCode()); + assertFalse(timeout1.equals(null)); + assertFalse(timeout1.equals("other")); + assertFalse(timeout1.equals(timeout3)); + assertFalse(timeout1.equals(timeout4)); + assertEquals("10 SECONDS", timeout1.toString()); + } + + @Test(timeout = 5000) + public void interruptedTakeNextFromStreamThrowsPropagateInterrupt() throws Exception { + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); + + GrpcStreamIterator streamIterator = + new GrpcStreamIterator( + /* statement= */ null, + /* lastStatement= */ false, + /* prefetchChunks= */ 5, + /* cancelQueryWhenClientIsClosed= */ false); + streamIterator.setCall(call, /* withBeginTransaction= */ false); + + AtomicReference thrown = new AtomicReference<>(); + Thread consumerThread = + new Thread( + () -> { + try { + streamIterator.hasNext(); + } catch (Throwable t) { + thrown.set(t); + } + }); + consumerThread.start(); + + await() + .atMost(Duration.ofSeconds(5)) + .until(() -> consumerThread.getState() == Thread.State.WAITING); + + consumerThread.interrupt(); + consumerThread.join(); + + assertNotNull(thrown.get()); + assertTrue(thrown.get() instanceof SpannerException); + SpannerException spannerException = (SpannerException) thrown.get(); + assertEquals(ErrorCode.CANCELLED, spannerException.getErrorCode()); + assertTrue(spannerException.getCause() instanceof InterruptedException); + } +} diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/InlineBeginTransactionTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/InlineBeginTransactionTest.java index 5161821c96ac..c7d07898f32e 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/InlineBeginTransactionTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/InlineBeginTransactionTest.java @@ -68,6 +68,7 @@ import java.util.concurrent.Future; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.ScheduledThreadPoolExecutor; +import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import org.junit.After; @@ -1708,6 +1709,173 @@ public void testQueryWithInlineBeginDidNotReturnTransaction() { }); } + @Test + public void testQueryWithInlineBeginClosedBeforeTransactionIdReturned() { + DatabaseClient client = + spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]")); + ExecutorService executor = Executors.newFixedThreadPool(2); + try { + SpannerException e = + assertThrows( + SpannerException.class, + () -> + client + .readWriteTransaction() + .run( + transaction -> { + mockSpanner.freeze(); + ResultSet resultSet = transaction.executeQuery(SELECT1); + Future nextFuture = executor.submit(resultSet::next); + try { + mockSpanner.waitForRequestsToContain(ExecuteSqlRequest.class, 5000); + } catch (Exception ex) { + throw new RuntimeException(ex); + } + resultSet.close(); + mockSpanner.unfreeze(); + try { + nextFuture.get(5, TimeUnit.SECONDS); + } catch (Exception ignored) { + } + return transaction.executeUpdate(UPDATE_STATEMENT); + })); + assertEquals(ErrorCode.FAILED_PRECONDITION, e.getErrorCode()); + assertTrue( + e.getMessage().contains("ResultSet was closed before a transaction id was returned")); + } finally { + mockSpanner.unfreeze(); + executor.shutdownNow(); + } + } + + @Test + public void testQueryWithInlineBeginClosedBeforeTransactionIdReturnedUnblocksWaitingThread() + throws Exception { + DatabaseClient client = + spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]")); + ExecutorService executor = Executors.newFixedThreadPool(2); + try { + client + .readWriteTransaction() + .run( + transaction -> { + mockSpanner.freeze(); + ResultSet resultSet = transaction.executeQuery(SELECT1); + Future nextFuture = executor.submit(resultSet::next); + try { + mockSpanner.waitForRequestsToContain(ExecuteSqlRequest.class, 5000); + } catch (Exception ex) { + throw new RuntimeException(ex); + } + // Start a concurrent statement that waits for the transaction id to be returned. + Future updateFuture = + executor.submit(() -> transaction.executeUpdate(UPDATE_STATEMENT)); + + // Give the update statement thread a moment to reach getTransactionSelector(). + // Even if resultSet.close() happens first, updateFuture will immediately see the + // exception. + resultSet.close(); + mockSpanner.unfreeze(); + + // The waiting thread must be unblocked immediately (within 5 seconds, not 60s + // timeout). + ExecutionException updateException = + assertThrows( + ExecutionException.class, () -> updateFuture.get(5, TimeUnit.SECONDS)); + assertTrue(updateException.getCause() instanceof SpannerException); + SpannerException spannerException = (SpannerException) updateException.getCause(); + assertEquals(ErrorCode.FAILED_PRECONDITION, spannerException.getErrorCode()); + assertTrue( + spannerException + .getMessage() + .contains("ResultSet was closed before a transaction id was returned")); + + try { + nextFuture.get(5, TimeUnit.SECONDS); + } catch (Exception ignored) { + } + throw spannerException; + }); + } catch (SpannerException e) { + assertEquals(ErrorCode.FAILED_PRECONDITION, e.getErrorCode()); + assertTrue( + e.getMessage().contains("ResultSet was closed before a transaction id was returned")); + } finally { + mockSpanner.unfreeze(); + executor.shutdownNow(); + } + } + + @Test + public void testQueryWithInlineBeginClosedImmediatelyWithoutNextFollowedByUpdate() { + DatabaseClient client = + spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]")); + Long updateCount = + client + .readWriteTransaction() + .run( + transaction -> { + ResultSet resultSet = transaction.executeQuery(SELECT1); + resultSet.close(); + return transaction.executeUpdate(UPDATE_STATEMENT); + }); + assertEquals(Long.valueOf(UPDATE_COUNT), updateCount); + assertEquals(0, countRequests(BeginTransactionRequest.class)); + assertEquals(1, countRequests(ExecuteSqlRequest.class)); + ExecuteSqlRequest executeSqlRequest = + (ExecuteSqlRequest) mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(0); + assertTrue(executeSqlRequest.getTransaction().hasBegin()); + assertEquals(1, countRequests(CommitRequest.class)); + } + + @Test + public void testQueryWithInlineBeginClosedImmediatelyWithoutNextFollowedByCommit() { + DatabaseClient client = + spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]")); + client + .readWriteTransaction() + .run( + transaction -> { + ResultSet resultSet = transaction.executeQuery(SELECT1); + resultSet.close(); + return null; + }); + assertEquals(1, countRequests(BeginTransactionRequest.class)); + assertEquals(0, countRequests(ExecuteSqlRequest.class)); + assertEquals(1, countRequests(CommitRequest.class)); + } + + @Test + public void testQueryWithInlineBeginClosedImmediatelyWithoutNextConcurrentUpdate() + throws Exception { + DatabaseClient client = + spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]")); + ExecutorService executor = Executors.newFixedThreadPool(2); + try { + Long updateCount = + client + .readWriteTransaction() + .run( + transaction -> { + ResultSet resultSet = transaction.executeQuery(SELECT1); + Future updateFuture = + executor.submit(() -> transaction.executeUpdate(UPDATE_STATEMENT)); + resultSet.close(); + try { + return updateFuture.get(5, TimeUnit.SECONDS); + } catch (Exception e) { + throw SpannerExceptionFactory.asSpannerException(e); + } + }); + assertEquals(Long.valueOf(UPDATE_COUNT), updateCount); + assertEquals(0, countRequests(BeginTransactionRequest.class)); + assertEquals(1, countRequests(ExecuteSqlRequest.class)); + assertEquals(1, countRequests(CommitRequest.class)); + } finally { + executor.shutdownNow(); + } + } + @Test public void testReadWithInlineBeginDidNotReturnTransaction() { runWithIgnoreInlineBegin(