From 8f73ee36e445b99cf633f4f7aeda7caae47cdf16 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= Date: Fri, 16 Oct 2026 15:15:56 +0200 Subject: [PATCH] fix(spanner): prevent stream leaks and race conditions on result set close Fixes several race conditions, potential stream leaks, and state loss when closing streaming result sets: 1. Early Close Race Condition: If ResultSet.close() was called before asynchronous gRPC stream initialization finished (before setCall was invoked), the cancellation was silently dropped because the call reference was still null. When setCall ran moments later, the stream remained open and was never cancelled. Now, closed state is tracked so that setCall immediately cancels the stream if closed early. 2. Concurrent Close Idempotency: close() is now synchronized and extracts the call reference atomically, guaranteeing call.cancel() is executed at most once, and executing the cancellation outside the lock to prevent deadlocks. 3. Superfluous Cancellation on Completed Streams: The call handle is now cleared as soon as a PartialResultSet with last=true is consumed, avoiding unnecessary RPC cancellations when closing after reading the final chunk without an extra next() call. 4. END_OF_STREAM Queue Deduplication: Atomically deduplicate END_OF_STREAM markers to avoid queue overflow or stalls when prefetchChunks is configured to small values (e.g. prefetchChunks = 1). 5. Inlined Begin Transaction State on Close: In GrpcResultSet.close(), capture withBeginTransaction before closing the underlying iterator so listener.onDone() properly unblocks any threads waiting for inlined transaction IDs. --- .../google/cloud/spanner/GrpcResultSet.java | 3 +- .../cloud/spanner/GrpcStreamIterator.java | 178 +++- .../CloseSpannerWithOpenResultSetTest.java | 73 +- .../cloud/spanner/GrpcResultSetTest.java | 48 ++ .../cloud/spanner/GrpcStreamIteratorTest.java | 801 ++++++++++++++++++ .../spanner/InlineBeginTransactionTest.java | 168 ++++ 6 files changed, 1231 insertions(+), 40 deletions(-) create mode 100644 java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/GrpcStreamIteratorTest.java 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(