From d6a447461f4be9f4d17c65b1a5dc4c2452053408 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= Date: Sun, 25 Oct 2026 12:24:17 +0100 Subject: [PATCH] fix(spanner): synchronize stream lifecycle and retry buffering in ResumableStreamIterator Asynchronous result sets interact with ResumableStreamIterator from multiple threads: the caller thread initiating or closing the stream, gRPC callback threads notifying that messages or errors have arrived, and consumer threads polling for readiness and reading rows. Previously, stream lifecycle state and the retry buffer were not synchronized across these interactions, which caused several issues: 1. Duplicate or prematurely consumed streams during startup: When a gRPC callback fired while the initial stream was still being started and prefetched, the consumer thread could observe a partially initialized stream or start a second stream concurrently. Conversely, if a retryable error occurred immediately and the consumer thread reset the stream to retry, the finishing startup call could overwrite the reset state with the failed stream. 2. Uncancelled streams and deadlocks on concurrent close: Closing a result set while a stream was being created or retried could leave the newly created stream open, close it multiple times, or hold an internal lock while invoking external close callbacks. In addition, transaction metadata flags (whether the stream included an inlined BeginTransaction or was the last statement) were lost when the active stream reference was cleared on close or during retry backoff, preventing callers from properly cleaning up transactions and sessions. 3. Premature readiness signals for unresumable chunks: ResumableStreamIterator buffers chunks that do not carry a resume token so that a transient failure can be retried transparently without yielding duplicate data. However, non-blocking readiness checks delegated directly to the underlying gRPC stream and reported data as available as soon as a single tokenless chunk arrived. This caused asynchronous consumers to wake up and block on network I/O waiting for a resume token or buffer limit. How this change fixes these problems: - Coordinates stream startup so only one thread starts a stream at a time while concurrent callers wait for startup to complete, and tracks an in-flight stream separately from the published stream. A stream is only exposed to consumers and readiness checks after startup and initial prefetching have completed, while still allowing concurrent close calls to cancel an in-flight stream immediately and tearing down streams outside the lock. - Preserves stream transaction flags across stream resets, startup failures, and concurrent closes so upper layers always know whether an inlined transaction or final statement was in progress, and ensures the tracing span is always ended exactly once. - Drains available non-error chunks into the internal retry buffer during non-blocking readiness checks and only reports data as ready to poll when the buffer can actually be flushed (a resume token has been seen, the buffer is full, or the stream has reached the end), or when an error or closure needs to be surfaced. --- .../cloud/spanner/AbstractResultSet.java | 16 +- .../spanner/ResumableStreamIterator.java | 489 +++- .../spanner/ResumableStreamIteratorTest.java | 2222 ++++++++++++++++- 3 files changed, 2632 insertions(+), 95 deletions(-) diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractResultSet.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractResultSet.java index 2a1c781c23d7..4a179414e4c1 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractResultSet.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractResultSet.java @@ -158,7 +158,8 @@ interface CloseableIterator extends Iterator { /** * @param streamMessageListener A class object which implements StreamMessageListener - * @return true if streaming is supported by the iterator, otherwise false + * @return true if streaming is supported by the iterator, otherwise false. Also returns false + * if the iterator has already been closed, in which case no stream is started. */ default boolean initiateStreaming(AsyncResultSet.StreamMessageListener streamMessageListener) { return false; @@ -170,6 +171,15 @@ default void requestPrefetchChunks() {} /** * Returns true if data (a chunk, row, EOF, or error) is available to read immediately without * blocking the calling thread on network I/O. + * + *

Single-Consumer Threading Contract: This method drains underlying stream + * state into internal buffers and MUST only be called by the designated consuming thread. It + * must NEVER be invoked concurrently with {@link #hasNext()} or {@link #next()}. + * + *

Note on Pending Retryable Errors: If a retryable stream error is pending, + * this method returns {@code true} so the consumer is prompted to call {@link #next()} to + * handle backoff and reconnect. In that specific scenario, the subsequent call to {@link + * #next()} may block for backoff sleep and stream re-establishment. */ default boolean isDataAvailable() { return true; @@ -179,6 +189,10 @@ default boolean isDataAvailable() { * Returns true if the iterator can be polled immediately (via {@link #hasNext()} or {@link * #next()}) to yield a chunk or EOF without blocking on network I/O and without throwing a * stream error. + * + *

Single-Consumer Threading Contract: This method drains underlying stream + * state and MUST only be called by the designated consuming thread. It must NEVER be invoked + * concurrently with {@link #hasNext()} or {@link #next()}. */ default boolean canPollWithoutErrorOrBlocking() { return false; diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ResumableStreamIterator.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ResumableStreamIterator.java index d6477fe65104..4ea252f90e0b 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ResumableStreamIterator.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ResumableStreamIterator.java @@ -28,6 +28,7 @@ import com.google.api.gax.retrying.RetrySettings; import com.google.api.gax.rpc.StatusCode.Code; import com.google.cloud.spanner.AbstractResultSet.CloseableIterator; +import com.google.cloud.spanner.AsyncResultSet.StreamMessageListener; import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Preconditions; import com.google.common.collect.AbstractIterator; @@ -36,14 +37,16 @@ import io.grpc.Context; import io.opentelemetry.api.common.Attributes; import java.io.IOException; -import java.util.LinkedList; +import java.util.ArrayDeque; import java.util.Set; import java.util.concurrent.CountDownLatch; import java.util.concurrent.Executor; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.logging.Level; import java.util.logging.Logger; -import javax.annotation.Nullable; +import org.jspecify.annotations.NullMarked; +import org.jspecify.annotations.Nullable; /** * Wraps an iterator over partial result sets, supporting resuming RPCs on error. This class keeps @@ -68,21 +71,48 @@ * so a configured maxAttempts of N bounds the total number of RPC attempts without progress by N*N, * not by N. The total timeout is measured in wall-clock time from the first failure of the sequence * and therefore spans both layers. + * + *

Threading Contract: This iterator follows a single-consumer model. + * Invocations of {@link #hasNext()}, {@link #next()}, {@link #computeNext()}, {@link + * #isDataAvailable()}, and {@link #canPollWithoutErrorOrBlocking()} must be serialized by the + * caller on a single thread. Concurrency is supported for {@link #close(String)} and listener + * registration, but simultaneous consumption or polling from multiple threads is not supported. + * Invoking {@link #close(String)} while the consumer thread is blocked in {@link #next()} cancels + * the underlying stream and causes {@link #computeNext()} to return end-of-data rather than + * propagating the resulting {@code CANCELLED} stream error. */ +@NullMarked @VisibleForTesting abstract class ResumableStreamIterator extends AbstractIterator implements CloseableIterator { private final ErrorHandler errorHandler; - private AsyncResultSet.StreamMessageListener streamMessageListener; + @Nullable private volatile StreamMessageListener streamMessageListener; private final RetrySettings streamingRetrySettings; private final Set retryableCodes; private static final Logger logger = Logger.getLogger(ResumableStreamIterator.class.getName()); - private BackOff backOff; - private final LinkedList buffer = new LinkedList<>(); + @Nullable private BackOff backOff; + private final ArrayDeque buffer = new ArrayDeque<>(); private final int maxBufferSize; private final ISpan span; private final TraceWrapper tracer; - private CloseableIterator stream; + + /** Guards state transitions between {@link #stream} and {@link #closed}. */ + private final Object streamLock = new Object(); + + /** Indicates whether this iterator has been closed. */ + private volatile boolean closed; + + @Nullable private String closeMessage; + private volatile boolean withBeginTransaction; + private volatile boolean lastStatement; + + /** Ensures the tracing span is finalized at most once. */ + private final AtomicBoolean spanEnded = new AtomicBoolean(false); + + @Nullable private volatile CloseableIterator stream; + @Nullable private CloseableIterator inFlightStream; + @Nullable private volatile Throwable startFailure; + private boolean startingStream; /** * The number of consecutive failed attempts without any progress on the stream. Reset to zero @@ -97,8 +127,8 @@ abstract class ResumableStreamIterator extends AbstractIterator startStream( @Nullable ByteString resumeToken, - AsyncResultSet.StreamMessageListener streamMessageListener, + @Nullable StreamMessageListener streamMessageListener, XGoogSpannerRequestId requestId); /** @@ -268,104 +305,159 @@ boolean prepareIteratorForRetryOnDifferentGrpcChannel() { @Override public void close(@Nullable String message) { - if (stream != null) { - stream.close(message); - span.end(); - stream = null; + // Closing while another thread is blocked in next() makes computeNext() return end-of-data + // (via its closed checks) instead of surfacing the CANCELLED error reported by the cancelled + // underlying stream. + CloseableIterator streamToClose = null; + synchronized (streamLock) { + if (!closed) { + closed = true; + closeMessage = message; + streamToClose = this.stream != null ? this.stream : this.inFlightStream; + // Leave inFlightStream set until startGrpcStreaming() finishes so the post-startStream() + // setStream(streamIterator) check sees inFlightStream == streamToRegister and does not + // close streamIterator a second time. + this.stream = null; + streamLock.notifyAll(); + } + } + synchronized (buffer) { + buffer.clear(); } + // Perform stream teardown outside the lock to prevent deadlocks with external callbacks. + if (streamToClose != null) { + streamToClose.close(message); + } + endSpanOnce(); + } + + @VisibleForTesting + boolean isClosed() { + return closed; } @Override public boolean isWithBeginTransaction() { - return stream != null && stream.isWithBeginTransaction(); + CloseableIterator currentInFlightStream; + boolean fallbackWithBeginTransaction; + synchronized (streamLock) { + currentInFlightStream = this.inFlightStream; + fallbackWithBeginTransaction = this.withBeginTransaction; + } + // A subclass that registers its stream early calls setStream(consumer) before + // consumer.setCall(call, true) sets withBeginTransaction on the underlying stream. Querying + // currentInFlightStream live ensures callers observe true as soon as setCall() runs during an + // in-flight startStream(). + return currentInFlightStream != null + ? currentInFlightStream.isWithBeginTransaction() + : fallbackWithBeginTransaction; } @Override public boolean isLastStatement() { - return stream != null && stream.isLastStatement(); + return this.lastStatement; } @Override @InternalApi - public boolean initiateStreaming(AsyncResultSet.StreamMessageListener streamMessageListener) { - this.streamMessageListener = streamMessageListener; + public boolean initiateStreaming(StreamMessageListener streamMessageListener) { + synchronized (streamLock) { + if (closed) { + return false; + } + this.streamMessageListener = streamMessageListener; + } startGrpcStreaming(); return true; } @Override + public boolean canPollWithoutErrorOrBlocking() { + synchronized (buffer) { + if (closed) { + return false; + } + if (isBufferFlushableLocked() || finished) { + return true; + } + } + drainAvailableFromStream(this.stream); + synchronized (buffer) { + return !closed && (isBufferFlushableLocked() || finished); + } + } + + @Override + public boolean isDataAvailable() { + // If closed, return true to unblock any waiting async consumer so it can observe EOF/closure + // instead of hanging indefinitely. + // This loop is bounded: an iteration only loops when currentStream has a non-error chunk or EOF + // ready to poll, and the next iteration's canPollWithoutErrorOrBlocking() call drains that into + // buffer until either buffer becomes flushable (at most maxBufferSize + 1 chunks) or + // currentStream has nothing more that is immediately available. + while (true) { + if (closed || canPollWithoutErrorOrBlocking()) { + return true; + } + CloseableIterator currentStream = this.stream; + if (currentStream == null || !currentStream.isDataAvailable()) { + return closed || startFailure != null; + } + if (!currentStream.canPollWithoutErrorOrBlocking()) { + return true; + } + } + } + + @Override + @Nullable protected PartialResultSet computeNext() { int numAttemptsOnOtherChannel = 0; Context context = Context.current(); - while (true) { + while (!closed) { // Eagerly start stream before consuming any buffered items. startGrpcStreaming(); // Buffer contains items up to a resume token or has reached capacity: flush. - if (!buffer.isEmpty() - && (finished || !safeToRetry || !buffer.getLast().getResumeToken().isEmpty())) { - return buffer.pop(); + PartialResultSet buffered = pollBufferedChunk(); + if (buffered != null) { + return buffered; + } + if (finished) { + break; + } + // Snapshot the volatile stream reference to guard against concurrent close() nulling the + // field between check and dereference. + CloseableIterator currentStream = this.stream; + if (currentStream == null) { + break; } try { - if (stream.hasNext()) { - PartialResultSet next = stream.next(); - boolean hasResumeToken = !next.getResumeToken().isEmpty(); - if (hasResumeToken) { - // Only a resume token that differs from the last seen token is progress: a stream - // that repeatedly returns the token that was used to resume it has not moved past it. - if (!next.getResumeToken().equals(resumeToken)) { - // The stream made progress, so reset the budget for consecutive failed attempts. - attempts = 0; - backOff = null; - retrySequenceStartNanos = -1L; - } - resumeToken = next.getResumeToken(); - safeToRetry = true; - } - // If the buffer is empty and this chunk has a resume token or we cannot resume safely - // anyway, we can yield it immediately rather than placing it in the buffer to be - // returned on the next iteration. - if ((hasResumeToken || !safeToRetry) && buffer.isEmpty()) { - return next; - } - buffer.add(next); - if (buffer.size() > maxBufferSize && buffer.getLast().getResumeToken().isEmpty()) { - // We need to flush without a restart token. Errors encountered until we see - // such a token will fail the read. - safeToRetry = false; - } - } else { - finished = true; - if (buffer.isEmpty()) { - endOfData(); - return null; - } - } + readNextFromStreamIntoBuffer(currentStream); } catch (SpannerException spannerException) { if (safeToRetry && isRetryable(spannerException)) { - handleRetryableException(context, spannerException); - - continue; + try { + handleRetryableException(context, spannerException); + continue; + } catch (Throwable throwable) { + span.setStatus(throwable); + endSpanOnce(); + throw throwable; + } } // Check if we should retry the request on a different gRPC channel. - if (resumeToken == null && buffer.isEmpty()) { - Throwable translated = errorHandler.translateException(spannerException); - if (translated instanceof RetryOnDifferentGrpcChannelException) { - if (++numAttemptsOnOtherChannel < errorHandler.getMaxAttempts() - && prepareIteratorForRetryOnDifferentGrpcChannel()) { - stream = null; - continue; - } - } + if (shouldRetryOnDifferentChannel(spannerException, ++numAttemptsOnOtherChannel)) { + resetStream(); + continue; } - span.addAnnotation("Stream broken. Not safe to retry", spannerException); - span.setStatus(spannerException); + recordFailure(spannerException); throw spannerException; - } catch (RuntimeException e) { - span.addAnnotation("Stream broken. Not safe to retry", e); - span.setStatus(e); - throw e; + } catch (RuntimeException | Error throwable) { + recordFailure(throwable); + throw throwable; } } + endOfData(); + return null; } /** Monotonic time source, overridable for deterministic retry-budget tests. */ @@ -385,7 +477,6 @@ long checkRetryBudgetAndGetDelay(SpannerException spannerException) { "Stream broken. Not retrying because the maximum number of attempts has been" + " exhausted", spannerException); - span.setStatus(spannerException); throw spannerException; } // Determine the retry delay: either the delay that the server included in the error, or @@ -403,7 +494,6 @@ long checkRetryBudgetAndGetDelay(SpannerException spannerException) { span.addAnnotation( "Stream broken. Not retrying because the total timeout has been exhausted", spannerException); - span.setStatus(spannerException); throw spannerException; } return delayMillis; @@ -413,19 +503,50 @@ private void handleRetryableException(Context context, SpannerException spannerE long delayMillis = checkRetryBudgetAndGetDelay(spannerException); span.addAnnotation("Stream broken. Safe to retry", spannerException); logger.log(Level.FINE, "Retryable exception, will sleep and retry", spannerException); - // Truncate any items in the buffer before the last retry token. - while (!buffer.isEmpty() && buffer.getLast().getResumeToken().isEmpty()) { - buffer.removeLast(); + // Discard any uncommitted chunks in the buffer (they will be replayed from resumeToken). + synchronized (buffer) { + buffer.clear(); } - assert buffer.isEmpty() || buffer.getLast().getResumeToken().equals(resumeToken); - stream = null; - try (IScope s = tracer.withSpan(span)) { + resetStream(); + try (IScope scope = tracer.withSpan(span)) { backoffSleep(context, delayMillis); } } + /** + * Lazily starts the underlying gRPC stream if streaming has not already been initiated and the + * iterator has not been closed. Concurrent callers (for example, when {@link + * #initiateStreaming(StreamMessageListener)} triggers an asynchronous {@code onStreamMessage} + * callback before {@code startGrpcStreaming()} returns) wait for the in-flight startup to finish + * rather than starting a second stream or consuming an uninitialized stream. + */ private void startGrpcStreaming() { - if (stream == null) { + // Fast path: lock-free check when stream is already running. + if (stream != null) { + return; + } + synchronized (streamLock) { + while (startingStream && !closed) { + try { + streamLock.wait(); + } catch (InterruptedException interruptedException) { + Thread.currentThread().interrupt(); + throw newSpannerExceptionForCancellation(Context.current(), interruptedException); + } + } + if (stream != null || closed) { + return; + } + if (startFailure instanceof Error) { + throw (Error) startFailure; + } + if (startFailure != null) { + throw (RuntimeException) startFailure; + } + startingStream = true; + } + CloseableIterator streamIterator; + try { span.addAnnotation( "Starting/Resuming stream", "ResumeToken", @@ -433,10 +554,17 @@ private void startGrpcStreaming() { try (IScope scope = tracer.withSpan(span)) { // When start a new stream set the Span as current to make the gRPC Span a child of // this Span. - stream = checkNotNull(startStream(resumeToken, streamMessageListener, requestId)); - stream.requestPrefetchChunks(); + streamIterator = checkNotNull(startStream(resumeToken, streamMessageListener, requestId)); + if (setStream(streamIterator)) { + streamIterator.requestPrefetchChunks(); + } } + } catch (RuntimeException | Error throwable) { + resetStreamOnFailure(throwable); + recordFailure(throwable); + throw throwable; } + publishStartedStream(streamIterator); } boolean isRetryable(SpannerException spannerException) { @@ -444,4 +572,195 @@ boolean isRetryable(SpannerException spannerException) { || retryableCodes.contains( GrpcStatusCode.of(spannerException.getErrorCode().getGrpcStatusCode()).getCode()); } + + /** + * Registers the in-flight stream early during {@code startStream()} setup so that in-flight RPCs + * can be cancelled immediately if {@link #close(String)} is invoked before {@code startStream()} + * returns, without publishing the stream to consumers before it is fully initialized. + * + * @return {@code true} if the stream was attached, or {@code false} if this iterator was already + * closed + */ + boolean setStream(CloseableIterator streamToRegister) { + boolean streamWithBeginTransaction = streamToRegister.isWithBeginTransaction(); + boolean streamLastStatement = streamToRegister.isLastStatement(); + boolean isClosed; + String messageToClose; + synchronized (streamLock) { + this.withBeginTransaction = streamWithBeginTransaction; + this.lastStatement = streamLastStatement; + // Only the single thread that set startingStream = true in startGrpcStreaming() calls + // setStream(), at most once early inside startStream() and once immediately after + // startStream() returns with the same instance. inFlightStream is always cleared before + // startingStream is reset to false, so inFlightStream is either null or streamToRegister. + if (this.inFlightStream == streamToRegister) { + return !this.closed; + } + isClosed = this.closed; + messageToClose = this.closeMessage; + this.inFlightStream = streamToRegister; + } + // If closed concurrently during acquisition, tear down the stream out-of-lock to avoid leaks. + if (isClosed) { + streamToRegister.close(messageToClose); + } + return !isClosed; + } + + private void publishStartedStream(CloseableIterator streamIterator) { + synchronized (streamLock) { + if (!this.closed) { + this.stream = streamIterator; + } + this.inFlightStream = null; + this.startingStream = false; + streamLock.notifyAll(); + } + } + + /** + * Resets the active stream reference during retries or channel failovers. Only called after + * {@link #startGrpcStreaming()} has already succeeded ({@code startFailure == null}) and a stream + * error is later encountered while reading from the published {@link #stream}. + */ + private void resetStream() { + synchronized (streamLock) { + this.stream = null; + } + } + + private void resetStreamOnFailure(Throwable throwable) { + CloseableIterator failedStream = null; + synchronized (streamLock) { + if (this.inFlightStream != null) { + // Refresh withBeginTransaction (a non-blocking field read on the stream iterator) before + // clearing inFlightStream and calling notifyAll(). A subclass that registers its stream + // early calls setStream(consumer) before consumer.setCall(call, true), and any follower + // thread woken by notifyAll() must immediately observe the updated flag via + // isWithBeginTransaction(). + this.withBeginTransaction = this.inFlightStream.isWithBeginTransaction(); + if (!this.closed) { + failedStream = this.inFlightStream; + } + } + this.inFlightStream = null; + this.startingStream = false; + this.startFailure = throwable; + streamLock.notifyAll(); + } + if (failedStream != null) { + failedStream.close(throwable.getMessage()); + } + } + + /** Ensures the tracing span opened in the constructor is ended exactly once. */ + private void endSpanOnce() { + if (spanEnded.compareAndSet(false, true)) { + span.end(); + } + } + + private boolean isBufferFlushableLocked() { + return !buffer.isEmpty() + && (finished || !safeToRetry || !buffer.getLast().getResumeToken().isEmpty()); + } + + private void bufferChunkLocked(PartialResultSet next) { + boolean hasResumeToken = !next.getResumeToken().isEmpty(); + if (hasResumeToken) { + // Only a resume token that differs from the last seen token is progress: a stream + // that repeatedly returns the token that was used to resume it has not moved past it. + if (!next.getResumeToken().equals(resumeToken)) { + // The stream made progress, so reset the budget for consecutive failed attempts. + attempts = 0; + backOff = null; + retrySequenceStartNanos = -1L; + } + resumeToken = next.getResumeToken(); + safeToRetry = true; + } + buffer.add(next); + if (buffer.size() > maxBufferSize && !hasResumeToken) { + // We need to flush without a restart token. Errors encountered until we see + // such a token will fail the read. + safeToRetry = false; + } + } + + /** + * Non-blockingly drains available non-error chunks from {@code currentStream} into {@link + * #buffer} until {@link #buffer} becomes flushable (for example, a resume token is seen or {@link + * #maxBufferSize} is exceeded) or the stream reaches EOF. + */ + private void drainAvailableFromStream( + @Nullable CloseableIterator currentStream) { + if (currentStream == null) { + return; + } + // Under the single-consumer contract, canPollWithoutErrorOrBlocking() returning true guarantees + // that the next item in currentStream is a non-error PartialResultSet or EOF, so + // readNextFromStreamIntoBuffer(currentStream) will not block or throw a stream error unless + // close() is invoked concurrently while reading (which readNextFromStreamIntoBuffer handles). + while (currentStream.canPollWithoutErrorOrBlocking()) { + if (readNextFromStreamIntoBuffer(currentStream)) { + return; + } + } + } + + private boolean readNextFromStreamIntoBuffer(CloseableIterator currentStream) { + try { + if (currentStream.hasNext()) { + PartialResultSet next = currentStream.next(); + synchronized (buffer) { + if (closed) { + return true; + } + bufferChunkLocked(next); + return isBufferFlushableLocked(); + } + } else { + synchronized (buffer) { + finished = true; + } + return true; + } + } catch (SpannerException spannerException) { + if (closed) { + return true; + } + throw spannerException; + } + } + + @Nullable + private PartialResultSet pollBufferedChunk() { + synchronized (buffer) { + if (isBufferFlushableLocked()) { + return buffer.pop(); + } + } + return null; + } + + private boolean shouldRetryOnDifferentChannel( + SpannerException spannerException, int numAttemptsOnOtherChannel) { + boolean canRetry; + synchronized (buffer) { + canRetry = resumeToken == null && buffer.isEmpty(); + } + if (canRetry) { + Throwable translated = errorHandler.translateException(spannerException); + return translated instanceof RetryOnDifferentGrpcChannelException + && numAttemptsOnOtherChannel < errorHandler.getMaxAttempts() + && prepareIteratorForRetryOnDifferentGrpcChannel(); + } + return false; + } + + private void recordFailure(Throwable throwable) { + span.addAnnotation("Stream broken. Not safe to retry", throwable); + span.setStatus(throwable); + endSpanOnce(); + } } diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ResumableStreamIteratorTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ResumableStreamIteratorTest.java index a0f18431315f..e1ecb1e528f5 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ResumableStreamIteratorTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ResumableStreamIteratorTest.java @@ -18,24 +18,37 @@ import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; +import com.google.api.core.ApiFuture; +import com.google.api.gax.core.ExecutorProvider; import com.google.api.gax.retrying.RetrySettings; +import com.google.cloud.spanner.AbstractResultSet.CloseableIterator; +import com.google.cloud.spanner.AsyncResultSet.CallbackResponse; +import com.google.cloud.spanner.AsyncResultSet.StreamMessageListener; import com.google.cloud.spanner.ErrorHandler.DefaultErrorHandler; import com.google.cloud.spanner.XGoogSpannerRequestId.NoopRequestIdCreator; +import com.google.cloud.spanner.spi.v1.SpannerRpc; import com.google.cloud.spanner.v1.stub.SpannerStubSettings; import com.google.common.base.Stopwatch; import com.google.common.collect.AbstractIterator; import com.google.common.collect.ImmutableList; import com.google.common.collect.Lists; +import com.google.common.util.concurrent.MoreExecutors; +import com.google.common.util.concurrent.Uninterruptibles; import com.google.protobuf.ByteString; import com.google.protobuf.Duration; import com.google.protobuf.Value; import com.google.rpc.RetryInfo; import com.google.spanner.v1.PartialResultSet; +import com.google.spanner.v1.Transaction; +import io.grpc.Context; import io.grpc.Metadata; import io.grpc.Status; import io.grpc.StatusRuntimeException; @@ -48,9 +61,19 @@ import java.util.Iterator; import java.util.LinkedList; import java.util.List; +import java.util.NoSuchElementException; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; import java.util.function.LongSupplier; -import javax.annotation.Nullable; +import org.jspecify.annotations.Nullable; +import org.junit.After; import org.junit.Before; import org.junit.Test; import org.junit.runner.RunWith; @@ -63,9 +86,8 @@ @RunWith(Parameterized.class) public class ResumableStreamIteratorTest { interface Starter { - AbstractResultSet.CloseableIterator startStream( - @Nullable ByteString resumeToken, - AsyncResultSet.StreamMessageListener streamMessageListener); + CloseableIterator startStream( + @Nullable ByteString resumeToken, @Nullable StreamMessageListener streamMessageListener); } interface ResultSetStream { @@ -115,13 +137,18 @@ static class NonRetryableException extends SpannerException { } static class ResultSetIterator extends AbstractIterator - implements AbstractResultSet.CloseableIterator { + implements CloseableIterator { final ResultSetStream stream; ResultSetIterator(ResultSetStream stream) { this.stream = stream; } + @Override + public boolean isDataAvailable() { + return false; + } + @Override protected PartialResultSet computeNext() { PartialResultSet next = stream.next(); @@ -158,6 +185,11 @@ public void setUp() { initWithLimit(Integer.MAX_VALUE); } + @After + public void tearDown() { + SpannerOptions.resetActiveTracingFramework(); + } + private void initWithLimit(int maxBufferSize) { initWithLimitAndRetrySettings( maxBufferSize, @@ -174,13 +206,17 @@ private void initWithLimitAndRetrySettings(int maxBufferSize, RetrySettings retr new TraceWrapper(Tracing.getTracer(), OpenTelemetry.noop().getTracer(""), false)); } + private static OpenTelemetrySpan mockOpenTelemetrySpan() { + return new OpenTelemetrySpan(mock(io.opentelemetry.api.trace.Span.class)); + } + private void initWithLimitAndRetrySettings( int maxBufferSize, RetrySettings retrySettings, TraceWrapper tracer) { resumableStreamIterator = new ResumableStreamIterator( maxBufferSize, "", - new OpenTelemetrySpan(mock(io.opentelemetry.api.trace.Span.class)), + mockOpenTelemetrySpan(), tracer, DefaultErrorHandler.INSTANCE, retrySettings, @@ -192,11 +228,11 @@ long nanoTime() { } @Override - AbstractResultSet.CloseableIterator startStream( + CloseableIterator startStream( @Nullable ByteString resumeToken, - AsyncResultSet.StreamMessageListener streamMessageListener, + @Nullable StreamMessageListener streamMessageListener, XGoogSpannerRequestId requestId) { - return starter.startStream(resumeToken, null); + return starter.startStream(resumeToken, streamMessageListener); } }; } @@ -750,6 +786,927 @@ public void bufferLimitMissingTokensSafeToRetry() { assertThat(consume(resumableStreamIterator)).containsExactly("a", "b", "c", "d").inOrder(); } + @Test + public void isDataAvailableWithGrpcStreamIterator() { + initWithLimit(512); + GrpcStreamIterator grpcStream = new GrpcStreamIterator(false, 4, false); + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + grpcStream.setCall(call, false); + Mockito.when(starter.startStream(null, null)).thenReturn(grpcStream); + resumableStreamIterator.initiateStreaming(null); + + // Empty stream -> false + assertFalse(resumableStreamIterator.isDataAvailable()); + + // Chunk without resume token arrives -> drained into buffer (replenishing gRPC flow-control + // credit), but not flushable yet until a resume token or EOF arrives (or maxBufferSize is + // exceeded). + grpcStream.consumer().onPartialResultSet(resultSet(null, "a")); + assertFalse(resumableStreamIterator.isDataAvailable()); + verify(call, Mockito.times(1)).request(1); + assertFalse(grpcStream.canPollWithoutErrorOrBlocking()); + + grpcStream.consumer().onCompleted(); + assertTrue(resumableStreamIterator.isDataAvailable()); + assertEquals("a", resumableStreamIterator.next().getValues(0).getStringValue()); + assertFalse(resumableStreamIterator.hasNext()); + } + + @Test + public void isDataAvailableWithGrpcStreamIteratorWhenErrorOccurs() { + initWithLimit(512); + GrpcStreamIterator grpcStream = new GrpcStreamIterator(false, 4, false); + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + grpcStream.setCall(call, false); + when(starter.startStream(null, null)).thenReturn(grpcStream); + resumableStreamIterator.initiateStreaming(null); + + // Initial chunk without resume token arrives -> drained into buffer, but not flushable yet. + grpcStream.consumer().onPartialResultSet(resultSet(null, "a")); + assertFalse(resumableStreamIterator.isDataAvailable()); + verify(call, Mockito.times(1)).request(1); + assertFalse(grpcStream.canPollWithoutErrorOrBlocking()); + + // Stream encounters an error -> isDataAvailable() must return true to unblock the consumer. + grpcStream.consumer().onError(new RetryableException(errorCodeParameter, "transient error")); + assertTrue(resumableStreamIterator.isDataAvailable()); + } + + @Test(timeout = 10_000) + public void isDataAvailableWithoutResumeTokensWithProductionBufferSize() { + initWithLimit(512); + GrpcStreamIterator grpcStream = new GrpcStreamIterator(false, 4, false); + SpannerRpc.StreamingCall call = Mockito.mock(SpannerRpc.StreamingCall.class); + grpcStream.setCall(call, false); + Mockito.when(starter.startStream(null, null)).thenReturn(grpcStream); + resumableStreamIterator.initiateStreaming(null); + + assertFalse(resumableStreamIterator.isDataAvailable()); + + for (int i = 0; i < 512; i++) { + grpcStream.consumer().onPartialResultSet(resultSet(null, "val" + i)); + // Draining each chunk from grpcStream's bounded queue into ResumableStreamIterator.buffer. + // Since buffer size <= 512, safeToRetry remains true and isDataAvailable() returns false. + assertFalse(resumableStreamIterator.isDataAvailable()); + } + + // Pushing chunk 513 exceeds maxBufferSize (512). This flips safeToRetry to false, + // making the buffer flushable and isDataAvailable() true BEFORE the stream completes. + grpcStream.consumer().onPartialResultSet(resultSet(null, "val512")); + assertTrue(resumableStreamIterator.isDataAvailable()); + + // Once the stream completes, buffered chunks remain flushable. + grpcStream.consumer().onCompleted(); + assertTrue(resumableStreamIterator.isDataAvailable()); + + List results = new ArrayList<>(); + while (resumableStreamIterator.hasNext()) { + PartialResultSet partialResultSet = resumableStreamIterator.next(); + results.add(partialResultSet.getValues(0).getStringValue()); + } + assertEquals(513, results.size()); + for (int i = 0; i < 513; i++) { + assertEquals("val" + i, results.get(i)); + } + } + + @Test + public void isDataAvailableWhenStreamIsNull() { + initWithLimit(10); + assertFalse(resumableStreamIterator.isDataAvailable()); + } + + @Test + public void isDataAvailableWhenFinished() { + initWithLimit(10); + ResultSetStream stream = mock(ResultSetStream.class); + when(starter.startStream(null, null)).thenReturn(new ResultSetIterator(stream)); + when(stream.next()).thenReturn(null); + assertFalse(resumableStreamIterator.hasNext()); + assertTrue(resumableStreamIterator.isDataAvailable()); + } + + @Test + public void isDataAvailableWhenClosedReturnsTrue() { + initWithLimit(10); + resumableStreamIterator.close("closed"); + assertTrue(resumableStreamIterator.isClosed()); + assertTrue(resumableStreamIterator.isDataAvailable()); + } + + @Test + public void canPollWithoutErrorOrBlockingWhenClosedReturnsFalse() { + initWithLimit(10); + assertFalse(resumableStreamIterator.canPollWithoutErrorOrBlocking()); + + resumableStreamIterator.close("closed"); + assertTrue(resumableStreamIterator.isClosed()); + assertFalse(resumableStreamIterator.canPollWithoutErrorOrBlocking()); + } + + @Test + public void canPollWithoutErrorOrBlockingWithGrpcStreamIterator() { + initWithLimit(512); + GrpcStreamIterator grpcStream = + new GrpcStreamIterator( + /* lastStatement= */ false, + /* prefetchChunks= */ 4, + /* cancelQueryWhenClientIsClosed= */ false); + grpcStream.setCall(mock(SpannerRpc.StreamingCall.class), false); + when(starter.startStream(null, null)).thenReturn(grpcStream); + resumableStreamIterator.initiateStreaming(null); + + assertFalse(resumableStreamIterator.canPollWithoutErrorOrBlocking()); + + // Non-resume-token chunk arrives -> buffered, but not flushable yet + grpcStream.consumer().onPartialResultSet(resultSet(null, "a")); + assertFalse(resumableStreamIterator.canPollWithoutErrorOrBlocking()); + + // Resume token chunk arrives -> flushable! + grpcStream.consumer().onPartialResultSet(resultSet(ByteString.copyFromUtf8("r1"), "b")); + assertTrue(resumableStreamIterator.canPollWithoutErrorOrBlocking()); + + // Consuming both items empties the buffer + assertEquals("a", resumableStreamIterator.next().getValues(0).getStringValue()); + assertEquals("b", resumableStreamIterator.next().getValues(0).getStringValue()); + assertFalse(resumableStreamIterator.canPollWithoutErrorOrBlocking()); + + // EOF arrives -> flushable again (yields EOF) + grpcStream.consumer().onCompleted(); + assertTrue(resumableStreamIterator.canPollWithoutErrorOrBlocking()); + assertFalse(resumableStreamIterator.hasNext()); + } + + @Test(timeout = 10_000) + public void closeDuringStartStreamClosesStreamAndPreventsResurrection() { + initWithLimit(10); + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + GrpcStreamIterator grpcStream = new GrpcStreamIterator(true, 4, false); + grpcStream.setCall(call, true); + when(starter.startStream(null, null)) + .thenAnswer( + invocation -> { + resumableStreamIterator.close("cancelled while starting stream"); + return grpcStream; + }); + + assertFalse(resumableStreamIterator.hasNext()); + assertTrue(resumableStreamIterator.isClosed()); + assertTrue(resumableStreamIterator.isWithBeginTransaction()); + assertTrue(resumableStreamIterator.isLastStatement()); + verify(call, Mockito.times(1)).cancel("cancelled while starting stream"); + verify(call, Mockito.never()).request(Mockito.anyInt()); + } + + @Test(timeout = 10_000) + public void multipleCloseCallsAreIdempotent() { + initWithLimit(10); + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + GrpcStreamIterator grpcStream = new GrpcStreamIterator(false, 4, false); + grpcStream.setCall(call, false); + when(starter.startStream(null, null)).thenReturn(grpcStream); + resumableStreamIterator.initiateStreaming(null); + grpcStream.consumer().onPartialResultSet(resultSet(ByteString.copyFromUtf8("r1"), "a")); + + resumableStreamIterator.close("first close"); + resumableStreamIterator.close("second close"); + + assertTrue(resumableStreamIterator.isClosed()); + assertFalse(resumableStreamIterator.hasNext()); + verify(call, Mockito.times(1)).cancel("first close"); + verify(call, Mockito.never()).cancel("second close"); + } + + @Test + public void closeWhenStreamIsNullEndsSpan() { + SpannerOptions.resetActiveTracingFramework(); + SpannerOptions.enableOpenTelemetryTraces(); + + io.opentelemetry.api.trace.Span openTelemetrySpan = mock(io.opentelemetry.api.trace.Span.class); + ISpan span = new OpenTelemetrySpan(openTelemetrySpan); + TraceWrapper tracer = mock(TraceWrapper.class); + when(tracer.spanBuilderWithExplicitParent(Mockito.anyString(), Mockito.any(), Mockito.any())) + .thenReturn(span); + initWithLimitAndRetrySettings(10, RetrySettings.newBuilder().build(), tracer); + + resumableStreamIterator.close("cancelled before streaming"); + assertTrue(resumableStreamIterator.isClosed()); + verify(openTelemetrySpan).end(); + } + + @Test + public void isDataAvailableWhenBufferHasItems() { + initWithLimit(10); + ResultSetStream stream = mock(ResultSetStream.class); + when(starter.startStream(null, null)).thenReturn(new ResultSetIterator(stream)); + when(stream.next()) + .thenReturn(resultSet(null, "chunkWithoutToken")) + .thenReturn(resultSet(ByteString.copyFromUtf8("r1"), "chunkWithToken")) + .thenReturn(null); + + // Consuming the first element causes computeNext to buffer both chunks (up to the resume token) + // and pop the first one, leaving chunkWithToken in the buffer. + assertEquals("chunkWithoutToken", resumableStreamIterator.next().getValues(0).getStringValue()); + + // With chunkWithToken still buffered, isDataAvailable() returns true via buffer.peekFirst() + // != null + assertTrue(resumableStreamIterator.isDataAvailable()); + assertEquals("chunkWithToken", resumableStreamIterator.next().getValues(0).getStringValue()); + } + + @Test + public void + isWithBeginTransactionAndLastStatementPropagateFromGrpcStreamIteratorToGrpcResultSet() { + initWithLimit(10); + assertFalse(resumableStreamIterator.isWithBeginTransaction()); + assertFalse(resumableStreamIterator.isLastStatement()); + + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + GrpcStreamIterator grpcStream = new GrpcStreamIterator(true, 4, false); + grpcStream.setCall(call, true); + when(starter.startStream(null, null)).thenReturn(grpcStream); + + Transaction expectedTransaction = + Transaction.newBuilder().setId(ByteString.copyFromUtf8("tx-1")).build(); + grpcStream + .consumer() + .onPartialResultSet( + PartialResultSet.newBuilder() + .setMetadata( + MockSpannerTestUtil.READ_KEY_VALUE_METADATA.toBuilder() + .setTransaction(expectedTransaction) + .build()) + .setResumeToken(ByteString.copyFromUtf8("r1")) + .setLast(true) + .addValues(Value.newBuilder().setStringValue("k1")) + .addValues(Value.newBuilder().setStringValue("v1")) + .build()); + grpcStream.consumer().onCompleted(); + + AbstractResultSet.Listener listener = mockErrorAndDoneListener(); + try (GrpcResultSet grpcResultSet = new GrpcResultSet(resumableStreamIterator, listener)) { + assertTrue(grpcResultSet.next()); + assertEquals("k1", grpcResultSet.getString(0)); + assertEquals("v1", grpcResultSet.getString(1)); + assertTrue(resumableStreamIterator.isWithBeginTransaction()); + assertTrue(resumableStreamIterator.isLastStatement()); + assertFalse(grpcResultSet.next()); + } + + verify(listener).onTransactionMetadata(expectedTransaction, true); + verify(listener).onDone(true); + } + + @Test + public void isWithBeginTransactionAndLastStatementFalseWhenNotSetOnGrpcStreamIterator() { + initWithLimit(10); + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + GrpcStreamIterator grpcStream = new GrpcStreamIterator(false, 4, false); + grpcStream.setCall(call, false); + when(starter.startStream(null, null)).thenReturn(grpcStream); + + resumableStreamIterator.initiateStreaming(null); + assertFalse(resumableStreamIterator.isWithBeginTransaction()); + assertFalse(resumableStreamIterator.isLastStatement()); + } + + @Test(timeout = 10_000) + public void initiateStreamingStartsStreamAndDeliversRowsToAsyncResultSet() throws Exception { + initWithLimit(10); + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + when(starter.startStream(Mockito.isNull(), Mockito.any())) + .thenAnswer( + invocation -> { + StreamMessageListener listener = invocation.getArgument(1); + GrpcStreamIterator grpcStream = new GrpcStreamIterator(false, 4, false); + if (listener != null) { + grpcStream.registerListener(listener); + } + grpcStream.setCall(call, false); + grpcStream + .consumer() + .onPartialResultSet( + PartialResultSet.newBuilder() + .setMetadata(MockSpannerTestUtil.READ_KEY_VALUE_METADATA) + .setResumeToken(ByteString.copyFromUtf8("r1")) + .setLast(true) + .addValues(Value.newBuilder().setStringValue("k1")) + .addValues(Value.newBuilder().setStringValue("v1")) + .build()); + grpcStream.consumer().onCompleted(); + return grpcStream; + }); + + ExecutorProvider executorProvider = + SpannerOptions.createAsyncExecutorProvider(1, 1L, TimeUnit.SECONDS); + List receivedRows = new ArrayList<>(); + try (GrpcResultSet grpcResultSet = + new GrpcResultSet(resumableStreamIterator, mockErrorAndDoneListener()); + AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl( + executorProvider, grpcResultSet, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { + ApiFuture callbackFuture = + asyncResultSet.setCallback( + MoreExecutors.directExecutor(), + resultSet -> { + while (true) { + switch (resultSet.tryNext()) { + case OK: + receivedRows.add(resultSet.getString(0) + ":" + resultSet.getString(1)); + break; + case DONE: + return CallbackResponse.DONE; + case NOT_READY: + return CallbackResponse.CONTINUE; + } + } + }); + callbackFuture.get(5, TimeUnit.SECONDS); + } finally { + executorProvider.getExecutor().shutdown(); + } + + verify(call).request(4); + assertEquals(ImmutableList.of("k1:v1"), receivedRows); + } + + @Test + public void initiateStreamingWhenClosedReturnsFalse() { + initWithLimit(10); + resumableStreamIterator.close("already closed"); + StreamMessageListener streamMessageListener = mock(StreamMessageListener.class); + assertFalse(resumableStreamIterator.initiateStreaming(streamMessageListener)); + verify(starter, Mockito.never()).startStream(Mockito.any(), Mockito.any()); + } + + @Test + public void isDataAvailableHoldingBackNonResumeTokenChunksPreventsDataDuplicationOnRetry() { + initWithLimit(10); + GrpcStreamIterator grpcStream1 = new GrpcStreamIterator(false, 4, false); + SpannerRpc.StreamingCall call1 = mock(SpannerRpc.StreamingCall.class); + grpcStream1.setCall(call1, false); + + GrpcStreamIterator grpcStream2 = new GrpcStreamIterator(false, 4, false); + SpannerRpc.StreamingCall call2 = mock(SpannerRpc.StreamingCall.class); + grpcStream2.setCall(call2, false); + + when(starter.startStream(null, null)).thenReturn(grpcStream1); + when(starter.startStream(ByteString.copyFromUtf8("r1"), null)).thenReturn(grpcStream2); + + // Initial stream delivers chunk 1 with resume token "r1". + grpcStream1.consumer().onPartialResultSet(resultSet(ByteString.copyFromUtf8("r1"), "val1")); + assertTrue(resumableStreamIterator.hasNext()); + assertEquals("val1", resumableStreamIterator.next().getValues(0).getStringValue()); + + // Chunk 2 arrives WITHOUT a resume token. isDataAvailable() MUST return false to hold it back. + grpcStream1.consumer().onPartialResultSet(resultSet(null, "val2")); + assertFalse(resumableStreamIterator.isDataAvailable()); + + // Stream 1 encounters a transient retryable error before a resume token arrives. + grpcStream1 + .consumer() + .onError(new RetryableException(errorCodeParameter, "transient network disconnect")); + + // Resumed stream restarts at token "r1" and replays "val2", followed by "val3" with token "r2". + grpcStream2.consumer().onPartialResultSet(resultSet(null, "val2")); + grpcStream2.consumer().onPartialResultSet(resultSet(ByteString.copyFromUtf8("r2"), "val3")); + grpcStream2.consumer().onCompleted(); + + // Error on stream 1 cannot be polled without blocking, but isDataAvailable() returns true to + // prompt the consumer to call hasNext() and trigger the retry onto stream 2. + assertFalse(resumableStreamIterator.canPollWithoutErrorOrBlocking()); + assertTrue(resumableStreamIterator.isDataAvailable()); + assertTrue(resumableStreamIterator.hasNext()); + assertEquals("val2", resumableStreamIterator.next().getValues(0).getStringValue()); + assertTrue(resumableStreamIterator.hasNext()); + assertEquals("val3", resumableStreamIterator.next().getValues(0).getStringValue()); + assertFalse(resumableStreamIterator.hasNext()); + } + + @Test + public void asynchronousChunkProductionWithInterleavedIsDataAvailableAndNext() throws Exception { + initWithLimit(512); + GrpcStreamIterator grpcStream = new GrpcStreamIterator(false, 4, false); + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + grpcStream.setCall(call, false); + when(starter.startStream(null, null)).thenReturn(grpcStream); + resumableStreamIterator.initiateStreaming(null); + + int totalChunks = 50; + ExecutorService producerExecutor = Executors.newSingleThreadExecutor(); + try { + Future producerFuture = + producerExecutor.submit( + () -> { + for (int i = 0; i < totalChunks; i++) { + ByteString token = i % 3 == 0 ? ByteString.copyFromUtf8("token" + i) : null; + grpcStream.consumer().onPartialResultSet(resultSet(token, "chunk_" + i)); + Thread.yield(); + } + grpcStream.consumer().onCompleted(); + }); + + long deadlineNanos = System.nanoTime() + TimeUnit.SECONDS.toNanos(10); + List receivedValues = new ArrayList<>(); + while (receivedValues.size() < totalChunks) { + if (System.nanoTime() > deadlineNanos) { + fail("Timed out waiting for chunks to be produced and consumed"); + } + if (resumableStreamIterator.isDataAvailable() && resumableStreamIterator.hasNext()) { + PartialResultSet chunk = resumableStreamIterator.next(); + receivedValues.add(chunk.getValues(0).getStringValue()); + } else { + Thread.yield(); + } + } + producerFuture.get(10, TimeUnit.SECONDS); + + assertEquals(totalChunks, receivedValues.size()); + for (int i = 0; i < totalChunks; i++) { + assertEquals("chunk_" + i, receivedValues.get(i)); + } + assertFalse(resumableStreamIterator.hasNext()); + } finally { + producerExecutor.shutdownNow(); + } + } + + @Test + public void computeNextWhenAlreadyClosedReturnsEndOfData() { + initWithLimit(10); + resumableStreamIterator.close("pre-closed"); + assertFalse(resumableStreamIterator.hasNext()); + assertThrows(NoSuchElementException.class, () -> resumableStreamIterator.next()); + } + + @Test + public void retryOnDifferentGrpcChannelWhenSupported() { + ErrorHandler errorHandler = mock(ErrorHandler.class); + when(errorHandler.getMaxAttempts()).thenReturn(2); + when(errorHandler.translateException(Mockito.any())) + .thenReturn(new RetryOnDifferentGrpcChannelException("channel failed", 1, null)); + + final AtomicInteger channelRetries = new AtomicInteger(0); + ResumableStreamIterator channelRetryIterator = + new TestResumableStreamIterator(errorHandler) { + @Override + CloseableIterator startStream( + @Nullable ByteString resumeToken, + @Nullable StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + return starter.startStream(resumeToken, streamMessageListener); + } + + @Override + boolean prepareIteratorForRetryOnDifferentGrpcChannel() { + channelRetries.incrementAndGet(); + return true; + } + }; + + ResultSetStream firstStream = mock(ResultSetStream.class); + ResultSetStream secondStream = mock(ResultSetStream.class); + when(starter.startStream(null, null)) + .thenReturn(new ResultSetIterator(firstStream)) + .thenReturn(new ResultSetIterator(secondStream)); + + when(firstStream.next()) + .thenThrow(new NonRetryableException(ErrorCode.UNAVAILABLE, "transient channel error")); + when(secondStream.next()).thenReturn(resultSet(null, "success")).thenReturn(null); + + assertEquals("success", channelRetryIterator.next().getValues(0).getStringValue()); + assertEquals(1, channelRetries.get()); + } + + @Test + public void retryOnDifferentGrpcChannelNotAttemptedWhenPrepareReturnsFalse() { + ErrorHandler errorHandler = mock(ErrorHandler.class); + when(errorHandler.getMaxAttempts()).thenReturn(2); + when(errorHandler.translateException(Mockito.any())) + .thenReturn(new RetryOnDifferentGrpcChannelException("channel failed", 1, null)); + + ResumableStreamIterator channelRetryIterator = + new TestResumableStreamIterator(errorHandler) { + @Override + CloseableIterator startStream( + @Nullable ByteString resumeToken, + @Nullable StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + return starter.startStream(resumeToken, streamMessageListener); + } + }; + + ResultSetStream firstStream = mock(ResultSetStream.class); + when(starter.startStream(null, null)).thenReturn(new ResultSetIterator(firstStream)); + when(firstStream.next()) + .thenThrow(new NonRetryableException(ErrorCode.UNAVAILABLE, "transient channel error")); + + SpannerException spannerException = + assertThrows(SpannerException.class, channelRetryIterator::next); + assertEquals(ErrorCode.UNAVAILABLE, spannerException.getErrorCode()); + } + + @Test + public void retryOnDifferentGrpcChannelNotAttemptedWhenMaxAttemptsExceeded() { + ErrorHandler errorHandler = mock(ErrorHandler.class); + when(errorHandler.getMaxAttempts()).thenReturn(1); + when(errorHandler.translateException(Mockito.any())) + .thenReturn(new RetryOnDifferentGrpcChannelException("channel failed", 1, null)); + + ResumableStreamIterator channelRetryIterator = + new TestResumableStreamIterator(errorHandler) { + @Override + CloseableIterator startStream( + @Nullable ByteString resumeToken, + @Nullable StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + return starter.startStream(resumeToken, streamMessageListener); + } + + @Override + boolean prepareIteratorForRetryOnDifferentGrpcChannel() { + return true; + } + }; + + ResultSetStream firstStream = mock(ResultSetStream.class); + when(starter.startStream(null, null)).thenReturn(new ResultSetIterator(firstStream)); + when(firstStream.next()) + .thenThrow(new NonRetryableException(ErrorCode.UNAVAILABLE, "transient error")); + + SpannerException spannerException = + assertThrows(SpannerException.class, channelRetryIterator::next); + assertEquals(ErrorCode.UNAVAILABLE, spannerException.getErrorCode()); + } + + @Test + public void retryOnDifferentGrpcChannelNotAttemptedWhenResumeTokenPresent() { + ErrorHandler errorHandler = mock(ErrorHandler.class); + when(errorHandler.getMaxAttempts()).thenReturn(5); + when(errorHandler.translateException(Mockito.any())) + .thenReturn(new RetryOnDifferentGrpcChannelException("channel failed", 1, null)); + + final AtomicBoolean prepareCalled = new AtomicBoolean(false); + ResumableStreamIterator channelRetryIterator = + new TestResumableStreamIterator(errorHandler) { + @Override + CloseableIterator startStream( + @Nullable ByteString resumeToken, + @Nullable StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + return starter.startStream(resumeToken, streamMessageListener); + } + + @Override + boolean prepareIteratorForRetryOnDifferentGrpcChannel() { + prepareCalled.set(true); + return true; + } + }; + + ResultSetStream firstStream = mock(ResultSetStream.class); + when(starter.startStream(null, null)).thenReturn(new ResultSetIterator(firstStream)); + when(firstStream.next()) + .thenReturn(resultSet(ByteString.copyFromUtf8("r1"), "firstChunk")) + .thenThrow(new NonRetryableException(ErrorCode.UNAVAILABLE, "error after token")); + + assertEquals("firstChunk", channelRetryIterator.next().getValues(0).getStringValue()); + SpannerException spannerException = + assertThrows(SpannerException.class, channelRetryIterator::next); + assertEquals(ErrorCode.UNAVAILABLE, spannerException.getErrorCode()); + assertFalse(prepareCalled.get()); + } + + @Test + public void unhandledRuntimeExceptionSetsSpanStatusAndThrows() { + ISpan mockSpan = mock(ISpan.class); + initWithLimitAndRetrySettings( + 10, RetrySettings.newBuilder().build(), mockTracerWithSpan(mockSpan)); + + ResultSetStream firstStream = mock(ResultSetStream.class); + when(starter.startStream(null, null)).thenReturn(new ResultSetIterator(firstStream)); + when(firstStream.next()).thenThrow(new IllegalStateException("unexpected runtime exception")); + + IllegalStateException thrown = + assertThrows(IllegalStateException.class, resumableStreamIterator::next); + assertEquals("unexpected runtime exception", thrown.getMessage()); + verify(mockSpan) + .addAnnotation( + Mockito.eq("Stream broken. Not safe to retry"), Mockito.any(Throwable.class)); + verify(mockSpan).setStatus(Mockito.any(Throwable.class)); + verify(mockSpan).end(); + } + + @Test(timeout = 10_000) + public void startStreamFailureAfterSetStreamDoesNotDeadlockOrStartSecondStream() + throws Exception { + AtomicInteger startCount = new AtomicInteger(); + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + ISpan span = mock(ISpan.class); + + ResumableStreamIterator iterator = + new TestResumableStreamIterator(mockTracerWithSpan(span)) { + @Override + CloseableIterator startStream( + @Nullable ByteString resumeToken, + @Nullable StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + startCount.incrementAndGet(); + GrpcStreamIterator grpcStream = + new GrpcStreamIterator( + /* lastStatement= */ true, 10, /* cancelQueryWhenClientIsClosed= */ false); + if (streamMessageListener != null) { + grpcStream.registerListener(streamMessageListener); + } + setStream(grpcStream); + grpcStream.setCall(call, /* withBeginTransaction= */ true); + throw new RuntimeException("start failure"); + } + }; + + AbstractResultSet.Listener listener = mockErrorAndDoneListener(); + ExecutorProvider executorProvider = + SpannerOptions.createAsyncExecutorProvider(1, 1L, TimeUnit.SECONDS); + AtomicReference callbackException = new AtomicReference<>(); + try (GrpcResultSet grpcResultSet = new GrpcResultSet(iterator, listener); + AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl( + executorProvider, grpcResultSet, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { + ApiFuture callbackFuture = + asyncResultSet.setCallback( + MoreExecutors.directExecutor(), + resultSet -> { + try { + resultSet.tryNext(); + } catch (SpannerException spannerException) { + callbackException.set(spannerException); + } + return CallbackResponse.DONE; + }); + + ExecutionException executionException = + assertThrows(ExecutionException.class, () -> callbackFuture.get(5, TimeUnit.SECONDS)); + assertTrue(executionException.getCause() instanceof SpannerException); + assertTrue(executionException.getCause().getMessage().contains("start failure")); + } finally { + executorProvider.getExecutor().shutdown(); + } + + assertTrue(callbackException.get().getMessage().contains("start failure")); + assertEquals(1, startCount.get()); + verify(call, Mockito.times(1)).cancel("start failure"); + verify(span, Mockito.times(1)).setStatus(Mockito.any(Throwable.class)); + verify(span, Mockito.times(1)).end(); + verify(listener) + .onError(Mockito.any(SpannerException.class), Mockito.eq(true), Mockito.eq(true)); + verify(listener).onDone(true); + } + + @Test + public void drainPathBufferLimitMissingTokensBecomesUnsafeToRetry() { + initWithLimit(1); + GrpcStreamIterator grpcStream = new GrpcStreamIterator(false, 4, false); + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + grpcStream.setCall(call, false); + when(starter.startStream(null, null)).thenReturn(grpcStream); + resumableStreamIterator.initiateStreaming(null); + + // Initial chunk with resume token. + grpcStream.consumer().onPartialResultSet(resultSet(ByteString.copyFromUtf8("r1"), "a")); + assertTrue(resumableStreamIterator.hasNext()); + assertEquals("a", resumableStreamIterator.next().getValues(0).getStringValue()); + + // Deliver 2 chunks without tokens. Buffer limit is 1, so the drain path exceeds maxBufferSize + // and transitions safeToRetry to false. + grpcStream.consumer().onPartialResultSet(resultSet(null, "b")); + grpcStream.consumer().onPartialResultSet(resultSet(null, "c")); + assertTrue(resumableStreamIterator.canPollWithoutErrorOrBlocking()); + assertEquals("b", resumableStreamIterator.next().getValues(0).getStringValue()); + assertEquals("c", resumableStreamIterator.next().getValues(0).getStringValue()); + + // While safeToRetry is false and buffer is empty, even a single tokenless chunk is immediately + // flushable via canPollWithoutErrorOrBlocking() / isDataAvailable(). + grpcStream.consumer().onPartialResultSet(resultSet(null, "d1")); + assertTrue(resumableStreamIterator.canPollWithoutErrorOrBlocking()); + assertEquals("d1", resumableStreamIterator.next().getValues(0).getStringValue()); + + // Receiving a chunk with a resume token restores safeToRetry to true. + grpcStream.consumer().onPartialResultSet(resultSet(ByteString.copyFromUtf8("r2"), "e")); + assertTrue(resumableStreamIterator.isDataAvailable()); + assertEquals("e", resumableStreamIterator.next().getValues(0).getStringValue()); + + // Now that safeToRetry is true again, a single tokenless chunk is held back again. + grpcStream.consumer().onPartialResultSet(resultSet(null, "f")); + assertFalse(resumableStreamIterator.isDataAvailable()); + + // Exceeding maxBufferSize again transitions safeToRetry back to false; a subsequent tokenless + // chunk "h" is yielded directly via consumeNextFromStream() and the error fails without retry. + grpcStream.consumer().onPartialResultSet(resultSet(null, "g")); + assertTrue(resumableStreamIterator.canPollWithoutErrorOrBlocking()); + grpcStream.consumer().onPartialResultSet(resultSet(null, "h")); + grpcStream.consumer().onError(new RetryableException(errorCodeParameter, "transient error")); + + assertEquals("f", resumableStreamIterator.next().getValues(0).getStringValue()); + assertEquals("g", resumableStreamIterator.next().getValues(0).getStringValue()); + assertEquals("h", resumableStreamIterator.next().getValues(0).getStringValue()); + SpannerException spannerException = + assertThrows(SpannerException.class, resumableStreamIterator::next); + assertEquals(errorCodeParameter, spannerException.getErrorCode()); + } + + @Test + public void closeBeforeSetCallCancelsUnderlyingCallWithMessage() { + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + ResumableStreamIterator iterator = + new TestResumableStreamIterator() { + @Override + CloseableIterator startStream( + @Nullable ByteString resumeToken, + @Nullable StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + GrpcStreamIterator stream = + new GrpcStreamIterator( + /* lastStatement= */ false, 10, /* cancelQueryWhenClientIsClosed= */ false); + setStream(stream); + close("closing before setCall"); + stream.setCall(call, false); + return stream; + } + }; + + iterator.initiateStreaming(null); + verify(call).cancel("closing before setCall"); + } + + @Test + public void isWithBeginTransactionAndIsLastStatementPreservedAfterClose() { + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + ResumableStreamIterator iterator = + new TestResumableStreamIterator() { + @Override + CloseableIterator startStream( + @Nullable ByteString resumeToken, + @Nullable StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + GrpcStreamIterator stream = + new GrpcStreamIterator( + /* lastStatement= */ true, 10, /* cancelQueryWhenClientIsClosed= */ false); + setStream(stream); + stream.setCall(call, /* withBeginTransaction= */ true); + return stream; + } + }; + + iterator.initiateStreaming(null); + assertTrue(iterator.isWithBeginTransaction()); + assertTrue(iterator.isLastStatement()); + + iterator.close("closed"); + verify(call).cancel("closed"); + assertTrue(iterator.isWithBeginTransaction()); + assertTrue(iterator.isLastStatement()); + } + + @Test + public void closeAfterLastTrueWithResumeTokenDoesNotCancelUnderlyingCall() { + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + ResumableStreamIterator iterator = + new TestResumableStreamIterator() { + @Override + CloseableIterator startStream( + @Nullable ByteString resumeToken, + @Nullable StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + GrpcStreamIterator stream = + new GrpcStreamIterator( + /* lastStatement= */ false, 10, /* cancelQueryWhenClientIsClosed= */ false); + setStream(stream); + stream.setCall(call, false); + PartialResultSet.Builder builder = + PartialResultSet.newBuilder() + .setResumeToken(ByteString.copyFromUtf8("r1")) + .setLast(true); + for (int i = 0; i < 20; i++) { + builder.addValuesBuilder().setStringValue("v" + i); + } + stream.consumer().onPartialResultSet(builder.build()); + return stream; + } + }; + + assertTrue(iterator.hasNext()); + PartialResultSet partialResultSet = iterator.next(); + assertEquals(20, partialResultSet.getValuesCount()); + // Closing before calling hasNext() a second time should not cancel the stream because the + // PartialResultSet had last=true. + iterator.close("closed early"); + Mockito.verify(call, Mockito.never()).cancel(Mockito.any()); + } + + @Test + public void closeWithMultipleChunksCancelsUnlessLastTrueChunkWasConsumed() { + for (boolean consumeSecondChunk : new boolean[] {false, true}) { + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + ResumableStreamIterator iterator = + new TestResumableStreamIterator() { + @Override + CloseableIterator startStream( + @Nullable ByteString resumeToken, + @Nullable StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + GrpcStreamIterator stream = + new GrpcStreamIterator( + /* lastStatement= */ false, 10, /* cancelQueryWhenClientIsClosed= */ false); + setStream(stream); + stream.setCall(call, false); + PartialResultSet firstChunk = + PartialResultSet.newBuilder() + .setResumeToken(ByteString.copyFromUtf8("r1")) + .setLast(false) + .addValues(Value.newBuilder().setStringValue("v1")) + .build(); + PartialResultSet secondChunk = + PartialResultSet.newBuilder() + .setResumeToken(ByteString.copyFromUtf8("r2")) + .setLast(true) + .addValues(Value.newBuilder().setStringValue("v2")) + .build(); + stream.consumer().onPartialResultSet(firstChunk); + stream.consumer().onPartialResultSet(secondChunk); + return stream; + } + }; + + assertTrue(iterator.hasNext()); + assertEquals("v1", iterator.next().getValues(0).getStringValue()); + if (consumeSecondChunk) { + assertTrue(iterator.hasNext()); + assertEquals("v2", iterator.next().getValues(0).getStringValue()); + } + iterator.close("closed"); + if (consumeSecondChunk) { + Mockito.verify(call, Mockito.never()).cancel(Mockito.any()); + } else { + verify(call).cancel("closed"); + } + } + } + + @Test(timeout = 10_000) + public void closeAfterDrainingChunkPreventsSubsequentReads() { + initWithLimit(10); + GrpcStreamIterator grpcStream = new GrpcStreamIterator(false, 4, false); + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + grpcStream.setCall(call, false); + when(starter.startStream(null, null)).thenReturn(grpcStream); + resumableStreamIterator.initiateStreaming(null); + grpcStream.consumer().onPartialResultSet(resultSet(null, "bufferedBeforeClose")); + assertFalse(resumableStreamIterator.canPollWithoutErrorOrBlocking()); + verify(call, Mockito.times(1)).request(1); + assertFalse(grpcStream.canPollWithoutErrorOrBlocking()); + + resumableStreamIterator.close("closed with buffered chunk"); + assertFalse(resumableStreamIterator.canPollWithoutErrorOrBlocking()); + assertFalse(resumableStreamIterator.hasNext()); + } + + @Test(timeout = 10_000) + public void closeDuringStreamNextTerminatesCleanlyAndCancelsOnce() { + AtomicReference iteratorReference = new AtomicReference<>(); + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + // GrpcStreamIterator.computeNext() invokes call.request(1) after dequeuing a chunk and before + // returning it to ResumableStreamIterator.readNextFromStreamIntoBuffer(). Closing the + // ResumableStreamIterator inside call.request(1) deterministically exercises the race where + // close() completes while currentStream.next() is in flight. + Mockito.doAnswer( + invocation -> { + iteratorReference.get().close("closed while fetching"); + return null; + }) + .when(call) + .request(1); + + ResumableStreamIterator iterator = + new TestResumableStreamIterator() { + @Override + CloseableIterator startStream( + @Nullable ByteString resumeToken, + @Nullable StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + GrpcStreamIterator stream = new GrpcStreamIterator(false, 10, false); + setStream(stream); + stream.setCall(call, false); + stream.consumer().onPartialResultSet(resultSet(null, "chunkAfterClose")); + return stream; + } + }; + iteratorReference.set(iterator); + + assertFalse(iterator.hasNext()); + assertFalse(iterator.canPollWithoutErrorOrBlocking()); + verify(call).cancel("closed while fetching"); + } + static PartialResultSet resultSet(@Nullable ByteString resumeToken, String... data) { PartialResultSet.Builder builder = PartialResultSet.newBuilder(); if (resumeToken != null) { @@ -783,6 +1740,1253 @@ protected String computeNext() { }; } + @Test + public void retryOnDifferentGrpcChannelNotAttemptedWhenBufferNotEmptyAndResumeTokenNull() { + ErrorHandler errorHandler = mock(ErrorHandler.class); + when(errorHandler.getMaxAttempts()).thenReturn(2); + when(errorHandler.translateException(Mockito.any())) + .thenReturn(new RetryOnDifferentGrpcChannelException("channel failed", 1, null)); + + ResumableStreamIterator channelRetryIterator = + new TestResumableStreamIterator(errorHandler) { + @Override + CloseableIterator startStream( + @Nullable ByteString resumeToken, + @Nullable StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + return starter.startStream(resumeToken, streamMessageListener); + } + + @Override + boolean prepareIteratorForRetryOnDifferentGrpcChannel() { + return true; + } + }; + + ResultSetStream firstStream = mock(ResultSetStream.class); + when(starter.startStream(null, null)).thenAnswer(inv -> new ResultSetIterator(firstStream)); + when(firstStream.next()) + .thenReturn(resultSet(null, "chunkWithoutToken")) + .thenThrow(new NonRetryableException(ErrorCode.INTERNAL, "transient error")); + + SpannerException spannerException = + assertThrows(SpannerException.class, channelRetryIterator::next); + assertEquals(ErrorCode.INTERNAL, spannerException.getErrorCode()); + } + + @Test + public void retryOnDifferentGrpcChannelNotAttemptedWhenTranslatedExceptionIsNotChannelRetry() { + ResultSetStream firstStream = mock(ResultSetStream.class); + when(starter.startStream(null, null)).thenReturn(new ResultSetIterator(firstStream)); + when(firstStream.next()) + .thenThrow(new NonRetryableException(ErrorCode.INTERNAL, "non-channel error")); + + SpannerException spannerException = + assertThrows(SpannerException.class, resumableStreamIterator::next); + assertEquals(ErrorCode.INTERNAL, spannerException.getErrorCode()); + } + + @Test(timeout = 10_000) + public void spanSetStatusCalledExactlyOnceWhenRetryBudgetExhausted() { + ISpan childSpan = mock(ISpan.class); + RetrySettings retrySettings = RetrySettings.newBuilder().setMaxAttempts(1).build(); + + ResumableStreamIterator iterator = + new TestResumableStreamIterator(mockTracerWithSpan(childSpan), retrySettings) { + @Override + CloseableIterator startStream( + @Nullable ByteString resumeToken, + @Nullable StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + ResultSetStream stream = mock(ResultSetStream.class); + when(stream.next()) + .thenThrow(new RetryableException(errorCodeParameter, "retryable error")); + return new ResultSetIterator(stream); + } + }; + + SpannerException spannerException = assertThrows(SpannerException.class, iterator::next); + assertEquals(errorCodeParameter, spannerException.getErrorCode()); + Mockito.verify(childSpan, Mockito.times(1)).setStatus(Mockito.any(Throwable.class)); + Mockito.verify(childSpan).end(); + } + + @Test(timeout = 10_000) + public void closeDuringStreamNextWithCancelledExceptionTerminatesCleanly() throws Exception { + CountDownLatch consumerStarted = new CountDownLatch(1); + AtomicReference consumerThreadRef = new AtomicReference<>(); + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + + ResumableStreamIterator iterator = + new TestResumableStreamIterator() { + @Override + CloseableIterator startStream( + @Nullable ByteString resumeToken, + @Nullable StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + GrpcStreamIterator stream = new GrpcStreamIterator(false, 10, false); + setStream(stream); + Mockito.doAnswer( + invocation -> { + String message = invocation.getArgument(0); + stream + .consumer() + .onError( + SpannerExceptionFactory.newSpannerException( + ErrorCode.CANCELLED, message)); + return null; + }) + .when(call) + .cancel(Mockito.any()); + stream.setCall(call, false); + return stream; + } + }; + + ExecutorService executor = Executors.newSingleThreadExecutor(); + try { + Future hasNextFuture = + executor.submit( + () -> { + consumerThreadRef.set(Thread.currentThread()); + consumerStarted.countDown(); + return iterator.hasNext(); + }); + + assertTrue(consumerStarted.await(5, TimeUnit.SECONDS)); + awaitThreadWaiting(consumerThreadRef.get()); + + // Closing the iterator cancels the active gRPC call, which delivers ErrorCode.CANCELLED to + // the GrpcStreamIterator blocked in stream.take(). ResumableStreamIterator must treat this + // post-close cancellation as clean EOF rather than throwing CANCELLED to the caller. + iterator.close("cancelled by user"); + + assertFalse(hasNextFuture.get(5, TimeUnit.SECONDS)); + assertTrue(iterator.isClosed()); + verify(call).cancel("cancelled by user"); + } finally { + executor.shutdownNow(); + } + } + + @Test(timeout = 10_000) + public void startStreamWhenClosedCancelsStreamCallAndMarksIteratorClosed() { + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + GrpcStreamIterator stream = new GrpcStreamIterator(false, 10, false); + + ResumableStreamIterator iterator = + new TestResumableStreamIterator() { + @Override + CloseableIterator startStream( + @Nullable ByteString resumeToken, + @Nullable StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + close("closing"); + setStream(stream); + stream.setCall(call, false); + return stream; + } + }; + + iterator.initiateStreaming(null); + Mockito.verify(call, Mockito.times(1)).cancel("closing"); + assertTrue(iterator.isClosed()); + assertFalse(iterator.hasNext()); + } + + @Test(timeout = 10_000) + public void initiateStreamingConcurrentWithImmediateRetryableErrorDoesNotResurrectFailedStream() + throws Exception { + AtomicInteger streamAttempts = new AtomicInteger(); + CountDownLatch consumerStarted = new CountDownLatch(1); + CountDownLatch allowFirstStartStreamToReturn = new CountDownLatch(1); + AtomicReference consumerThreadRef = new AtomicReference<>(); + + ResumableStreamIterator iterator = + new TestResumableStreamIterator() { + @Override + CloseableIterator startStream( + @Nullable ByteString resumeToken, + @Nullable StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + int attempt = streamAttempts.incrementAndGet(); + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + GrpcStreamIterator stream = new GrpcStreamIterator(false, 10, false); + if (streamMessageListener != null) { + stream.registerListener(streamMessageListener); + } + setStream(stream); + if (attempt == 1) { + // Simulate an immediate retryable gRPC error arriving before setCall() completes. + stream + .consumer() + .onError(new RetryableException(errorCodeParameter, "transient failure")); + // setCall() flushes pending messages and invokes streamMessageListener synchronously. + stream.setCall(call, false); + Uninterruptibles.awaitUninterruptibly( + allowFirstStartStreamToReturn, 5, TimeUnit.SECONDS); + } else { + stream.setCall(call, false); + stream + .consumer() + .onPartialResultSet( + PartialResultSet.newBuilder() + .setResumeToken(ByteString.copyFromUtf8("r1")) + .setLast(true) + .addValues(Value.newBuilder().setStringValue("recovered-row")) + .build()); + stream.consumer().onCompleted(); + } + return stream; + } + }; + + ExecutorService executor = Executors.newFixedThreadPool(2); + try { + AtomicReference> consumerFutureRef = new AtomicReference<>(); + Future initiateFuture = + executor.submit( + () -> + iterator.initiateStreaming( + (chunk, full) -> { + if (consumerFutureRef.get() == null) { + consumerFutureRef.set( + executor.submit( + () -> { + consumerThreadRef.set(Thread.currentThread()); + consumerStarted.countDown(); + return iterator.next(); + })); + } + })); + + assertTrue(consumerStarted.await(5, TimeUnit.SECONDS)); + awaitThreadWaiting(consumerThreadRef.get()); + allowFirstStartStreamToReturn.countDown(); + + assertTrue(initiateFuture.get(5, TimeUnit.SECONDS)); + PartialResultSet result = consumerFutureRef.get().get(5, TimeUnit.SECONDS); + assertEquals("recovered-row", result.getValues(0).getStringValue()); + assertFalse(iterator.hasNext()); + assertEquals(2, streamAttempts.get()); + } finally { + executor.shutdownNow(); + } + } + + @Test(timeout = 10_000) + public void isDataAvailableAndReadinessPollingDoNotTouchUninitializedStreamDuringStartStream() + throws Exception { + CountDownLatch setStreamCalled = new CountDownLatch(1); + CountDownLatch allowStartStreamToFail = new CountDownLatch(1); + AtomicInteger startStreamCalls = new AtomicInteger(); + + ResumableStreamIterator iterator = + new TestResumableStreamIterator() { + @Override + CloseableIterator startStream( + @Nullable ByteString resumeToken, + @Nullable StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + startStreamCalls.incrementAndGet(); + GrpcStreamIterator uninitializedStream = new GrpcStreamIterator(false, 10, false); + setStream(uninitializedStream); + setStreamCalled.countDown(); + Uninterruptibles.awaitUninterruptibly(allowStartStreamToFail, 5, TimeUnit.SECONDS); + throw SpannerExceptionFactory.newSpannerException( + ErrorCode.PERMISSION_DENIED, "synchronous start failure"); + } + }; + + ExecutorService executor = Executors.newSingleThreadExecutor(); + try { + Future initiateFuture = + executor.submit(() -> iterator.initiateStreaming((chunk, full) -> {})); + assertTrue(setStreamCalled.await(5, TimeUnit.SECONDS)); + + // While startStream() is still in flight and setCall() has not been called on the + // GrpcStreamIterator, polling readiness must not touch or block on the uninitialized stream. + assertFalse(iterator.canPollWithoutErrorOrBlocking()); + assertFalse(iterator.isDataAvailable()); + + allowStartStreamToFail.countDown(); + ExecutionException executionException = + assertThrows(ExecutionException.class, () -> initiateFuture.get(5, TimeUnit.SECONDS)); + assertTrue(executionException.getCause() instanceof SpannerException); + assertEquals( + ErrorCode.PERMISSION_DENIED, + ((SpannerException) executionException.getCause()).getErrorCode()); + assertEquals(1, startStreamCalls.get()); + } finally { + executor.shutdownNow(); + } + } + + @Test(timeout = 10_000) + public void waitingConsumerInterruptedDuringStartStreamThrowsCancelledAndPreservesInterruptFlag() + throws Exception { + CountDownLatch startStreamEntered = new CountDownLatch(1); + CountDownLatch allowStartStreamToFinish = new CountDownLatch(1); + CountDownLatch consumerStarted = new CountDownLatch(1); + AtomicReference consumerThreadRef = new AtomicReference<>(); + AtomicBoolean interruptFlagPreserved = new AtomicBoolean(); + + ResumableStreamIterator iterator = + new TestResumableStreamIterator() { + @Override + CloseableIterator startStream( + @Nullable ByteString resumeToken, + @Nullable StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + startStreamEntered.countDown(); + Uninterruptibles.awaitUninterruptibly(allowStartStreamToFinish, 5, TimeUnit.SECONDS); + GrpcStreamIterator stream = new GrpcStreamIterator(false, 10, false); + stream.setCall(mock(SpannerRpc.StreamingCall.class), false); + stream.consumer().onCompleted(); + return stream; + } + }; + + ExecutorService executor = Executors.newFixedThreadPool(2); + try { + Future initiateFuture = executor.submit(() -> iterator.initiateStreaming(null)); + assertTrue(startStreamEntered.await(5, TimeUnit.SECONDS)); + + Future consumerFuture = + executor.submit( + () -> { + consumerThreadRef.set(Thread.currentThread()); + consumerStarted.countDown(); + try { + iterator.hasNext(); + } finally { + interruptFlagPreserved.set(Thread.interrupted()); + } + }); + + assertTrue(consumerStarted.await(5, TimeUnit.SECONDS)); + awaitThreadWaiting(consumerThreadRef.get()); + consumerThreadRef.get().interrupt(); + + ExecutionException executionException = + assertThrows(ExecutionException.class, () -> consumerFuture.get(5, TimeUnit.SECONDS)); + assertTrue(executionException.getCause() instanceof SpannerException); + assertEquals( + ErrorCode.CANCELLED, ((SpannerException) executionException.getCause()).getErrorCode()); + assertTrue(interruptFlagPreserved.get()); + + allowStartStreamToFinish.countDown(); + assertTrue(initiateFuture.get(5, TimeUnit.SECONDS)); + } finally { + executor.shutdownNow(); + } + } + + @Test(timeout = 10_000) + public void waitingConsumerWokenByConcurrentCloseDuringStartStreamTerminatesCleanly() + throws Exception { + CountDownLatch startStreamEntered = new CountDownLatch(1); + CountDownLatch allowStartStreamToFinish = new CountDownLatch(1); + CountDownLatch consumerStarted = new CountDownLatch(1); + AtomicReference consumerThreadRef = new AtomicReference<>(); + AtomicInteger startStreamCalls = new AtomicInteger(); + + ResumableStreamIterator iterator = + new TestResumableStreamIterator() { + @Override + CloseableIterator startStream( + @Nullable ByteString resumeToken, + @Nullable StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + startStreamCalls.incrementAndGet(); + GrpcStreamIterator stream = new GrpcStreamIterator(false, 10, false); + setStream(stream); + startStreamEntered.countDown(); + Uninterruptibles.awaitUninterruptibly(allowStartStreamToFinish, 5, TimeUnit.SECONDS); + return stream; + } + }; + + ExecutorService executor = Executors.newFixedThreadPool(2); + try { + Future initiateFuture = executor.submit(() -> iterator.initiateStreaming(null)); + assertTrue(startStreamEntered.await(5, TimeUnit.SECONDS)); + + Future consumerFuture = + executor.submit( + () -> { + consumerThreadRef.set(Thread.currentThread()); + consumerStarted.countDown(); + return iterator.hasNext(); + }); + + assertTrue(consumerStarted.await(5, TimeUnit.SECONDS)); + awaitThreadWaiting(consumerThreadRef.get()); + + iterator.close("closed while consumer waiting"); + assertFalse(consumerFuture.get(5, TimeUnit.SECONDS)); + + allowStartStreamToFinish.countDown(); + assertTrue(initiateFuture.get(5, TimeUnit.SECONDS)); + assertEquals(1, startStreamCalls.get()); + } finally { + executor.shutdownNow(); + } + } + + @Test(timeout = 10_000) + public void requestPrefetchChunksFailureClosesInFlightStream() { + AtomicInteger startCount = new AtomicInteger(); + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + Mockito.doThrow(new RuntimeException("prefetch failed")).when(call).request(10); + ISpan span = mock(ISpan.class); + + ResumableStreamIterator iterator = + new TestResumableStreamIterator(mockTracerWithSpan(span)) { + @Override + CloseableIterator startStream( + @Nullable ByteString resumeToken, + @Nullable StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + startCount.incrementAndGet(); + GrpcStreamIterator stream = + new GrpcStreamIterator( + /* lastStatement= */ true, 10, /* cancelQueryWhenClientIsClosed= */ false); + stream.setCall(call, /* withBeginTransaction= */ true); + return stream; + } + }; + + RuntimeException initialFailure = + assertThrows(RuntimeException.class, () -> iterator.initiateStreaming(null)); + assertEquals("prefetch failed", initialFailure.getMessage()); + verify(call, Mockito.times(1)).cancel("prefetch failed"); + verify(span, Mockito.times(1)).setStatus(Mockito.any(Throwable.class)); + verify(span, Mockito.times(1)).end(); + + RuntimeException subsequentFailure = assertThrows(RuntimeException.class, iterator::hasNext); + assertEquals("prefetch failed", subsequentFailure.getMessage()); + assertEquals(1, startCount.get()); + } + + @Test(timeout = 10_000) + public void closeDuringStartStreamFailureDoesNotDoubleCloseStream() { + AtomicInteger startCount = new AtomicInteger(); + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + ISpan span = mock(ISpan.class); + + ResumableStreamIterator iterator = + new TestResumableStreamIterator(mockTracerWithSpan(span)) { + @Override + CloseableIterator startStream( + @Nullable ByteString resumeToken, + @Nullable StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + startCount.incrementAndGet(); + GrpcStreamIterator stream = + new GrpcStreamIterator( + /* lastStatement= */ false, 10, /* cancelQueryWhenClientIsClosed= */ false); + setStream(stream); + stream.setCall(call, /* withBeginTransaction= */ false); + close("closed before startStream threw"); + throw new RuntimeException("startStream threw after close"); + } + }; + + assertThrows(RuntimeException.class, () -> iterator.initiateStreaming(null)); + verify(call, Mockito.times(1)).cancel("closed before startStream threw"); + verify(call, Mockito.never()).cancel("startStream threw after close"); + verify(span, Mockito.times(1)).setStatus(Mockito.any(Throwable.class)); + verify(span, Mockito.times(1)).end(); + assertFalse(iterator.hasNext()); + assertEquals(1, startCount.get()); + } + + @Test(timeout = 10_000) + public void contextCancellationDuringRetryBackoffInComputeNextSetsSpanStatusAndEndsSpan() { + ISpan childSpan = mock(ISpan.class); + ISpan currentSpan = mock(ISpan.class); + TraceWrapper tracer = mockTracerWithSpan(childSpan); + when(tracer.getCurrentSpan()).thenReturn(currentSpan); + + Context.CancellableContext cancellableContext = Context.current().withCancellation(); + Mockito.doAnswer( + invocation -> { + cancellableContext.cancel(new InterruptedException("cancelled during retry backoff")); + return null; + }) + .when(currentSpan) + .addAnnotation(Mockito.eq("Backing off"), Mockito.eq("Delay"), Mockito.anyLong()); + + ResumableStreamIterator iterator = + new TestResumableStreamIterator(tracer) { + @Override + CloseableIterator startStream( + @Nullable ByteString resumeToken, + @Nullable StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + ResultSetStream stream = mock(ResultSetStream.class); + when(stream.next()) + .thenThrow(new RetryableException(errorCodeParameter, "transient error")); + return new ResultSetIterator(stream); + } + }; + + try { + cancellableContext.run( + () -> { + SpannerException cancelledException = + assertThrows(SpannerException.class, iterator::next); + assertEquals(ErrorCode.CANCELLED, cancelledException.getErrorCode()); + }); + } finally { + cancellableContext.close(); + } + Mockito.verify(childSpan, Mockito.times(1)).setStatus(Mockito.any(Throwable.class)); + Mockito.verify(childSpan, Mockito.times(1)).end(); + } + + @Test + public void isDataAvailableLoopsWhenTokenlessChunkArrivesConcurrentlyOnStream() { + initWithLimit(10); + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + AtomicBoolean deliveredDuringIsDataAvailable = new AtomicBoolean(false); + + GrpcStreamIterator racingStream = + new GrpcStreamIterator(false, 4, false) { + @Override + public boolean isDataAvailable() { + // Simulate a tokenless chunk arriving on the GrpcStreamIterator right after the + // initial canPollWithoutErrorOrBlocking() drain check returned false, before + // currentStream.isDataAvailable() is evaluated. + if (deliveredDuringIsDataAvailable.compareAndSet(false, true)) { + consumer().onPartialResultSet(resultSet(null, "tokenlessChunk")); + } + return super.isDataAvailable(); + } + }; + racingStream.setCall(call, false); + + when(starter.startStream(null, null)).thenReturn(racingStream); + resumableStreamIterator.initiateStreaming(null); + + // Without the loop in isDataAvailable(), this would erroneously return true even though only a + // single tokenless chunk arrived (which cannot be flushed without blocking on network I/O). + assertFalse(resumableStreamIterator.isDataAvailable()); + verify(call, Mockito.times(1)).request(1); + assertFalse(racingStream.canPollWithoutErrorOrBlocking()); + + racingStream.consumer().onCompleted(); + assertTrue(resumableStreamIterator.isDataAvailable()); + assertEquals("tokenlessChunk", resumableStreamIterator.next().getValues(0).getStringValue()); + assertFalse(resumableStreamIterator.hasNext()); + } + + @Test(timeout = 10_000) + public void concurrentCloseDuringIsDataAvailableReturnsTrue() { + initWithLimit(10); + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + AtomicBoolean closedDuringDrainCheck = new AtomicBoolean(false); + + GrpcStreamIterator racingStream = + new GrpcStreamIterator(false, 4, false) { + @Override + public boolean canPollWithoutErrorOrBlocking() { + // Simulate a concurrent close() completing right after the initial + // canPollWithoutErrorOrBlocking() check in isDataAvailable() returns false. + if (closedDuringDrainCheck.compareAndSet(false, true)) { + resumableStreamIterator.close("closed concurrently during isDataAvailable"); + } + return super.canPollWithoutErrorOrBlocking(); + } + }; + racingStream.setCall(call, false); + + when(starter.startStream(null, null)).thenReturn(racingStream); + resumableStreamIterator.initiateStreaming(null); + + assertTrue(resumableStreamIterator.isDataAvailable()); + verify(call).cancel("closed concurrently during isDataAvailable"); + } + + @Test(timeout = 10_000) + public void startGrpcStreamingExceptionInTracerDoesNotDeadlockWaitingConsumer() throws Exception { + CountDownLatch withSpanEntered = new CountDownLatch(1); + CountDownLatch allowWithSpanToThrow = new CountDownLatch(1); + CountDownLatch consumerStarted = new CountDownLatch(1); + AtomicReference consumerThreadRef = new AtomicReference<>(); + + ISpan mockSpan = mock(ISpan.class); + TraceWrapper tracer = mock(TraceWrapper.class); + when(tracer.spanBuilderWithExplicitParent(Mockito.anyString(), Mockito.any(), Mockito.any())) + .thenReturn(mockSpan); + when(tracer.withSpan(mockSpan)) + .thenAnswer( + invocation -> { + withSpanEntered.countDown(); + Uninterruptibles.awaitUninterruptibly(allowWithSpanToThrow, 5, TimeUnit.SECONDS); + throw new RuntimeException("tracer withSpan failed"); + }); + + ResumableStreamIterator iterator = + new TestResumableStreamIterator(tracer) { + @Override + CloseableIterator startStream( + @Nullable ByteString resumeToken, + @Nullable StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + throw new AssertionError("startStream should not be reached"); + } + }; + + ExecutorService executor = Executors.newFixedThreadPool(2); + try { + Future initiateFuture = executor.submit(() -> iterator.initiateStreaming(null)); + assertTrue(withSpanEntered.await(5, TimeUnit.SECONDS)); + + Future consumerFuture = + executor.submit( + () -> { + consumerThreadRef.set(Thread.currentThread()); + consumerStarted.countDown(); + return iterator.hasNext(); + }); + + assertTrue(consumerStarted.await(5, TimeUnit.SECONDS)); + awaitThreadWaiting(consumerThreadRef.get()); + allowWithSpanToThrow.countDown(); + + ExecutionException leaderException = + assertThrows(ExecutionException.class, () -> initiateFuture.get(2, TimeUnit.SECONDS)); + assertEquals("tracer withSpan failed", leaderException.getCause().getMessage()); + + // Without covering tracer.withSpan(span) in the try/catch block of startGrpcStreaming(), + // consumerFuture permanently deadlocks in streamLock.wait() and times out here. + ExecutionException waiterException = + assertThrows(ExecutionException.class, () -> consumerFuture.get(2, TimeUnit.SECONDS)); + assertEquals("tracer withSpan failed", waiterException.getCause().getMessage()); + } finally { + executor.shutdownNow(); + } + } + + @Test(timeout = 10_000) + public void concurrentWaiterInStartGrpcStreamingDoesNotStartSecondStreamWhenLeaderFails() + throws Exception { + CountDownLatch startStreamEntered = new CountDownLatch(1); + CountDownLatch allowStartStreamToFail = new CountDownLatch(1); + CountDownLatch consumerStarted = new CountDownLatch(1); + AtomicReference consumerThreadRef = new AtomicReference<>(); + AtomicInteger startStreamCalls = new AtomicInteger(); + + ResumableStreamIterator iterator = + new TestResumableStreamIterator() { + @Override + CloseableIterator startStream( + @Nullable ByteString resumeToken, + @Nullable StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + startStreamCalls.incrementAndGet(); + startStreamEntered.countDown(); + Uninterruptibles.awaitUninterruptibly(allowStartStreamToFail, 5, TimeUnit.SECONDS); + throw SpannerExceptionFactory.newSpannerException( + ErrorCode.PERMISSION_DENIED, "leader startStream failed"); + } + }; + + ExecutorService executor = Executors.newFixedThreadPool(2); + try { + Future initiateFuture = executor.submit(() -> iterator.initiateStreaming(null)); + assertTrue(startStreamEntered.await(5, TimeUnit.SECONDS)); + + Future consumerFuture = + executor.submit( + () -> { + consumerThreadRef.set(Thread.currentThread()); + consumerStarted.countDown(); + return iterator.hasNext(); + }); + + assertTrue(consumerStarted.await(5, TimeUnit.SECONDS)); + awaitThreadWaiting(consumerThreadRef.get()); + allowStartStreamToFail.countDown(); + + ExecutionException leaderException = + assertThrows(ExecutionException.class, () -> initiateFuture.get(5, TimeUnit.SECONDS)); + assertTrue(leaderException.getCause() instanceof SpannerException); + assertEquals( + ErrorCode.PERMISSION_DENIED, + ((SpannerException) leaderException.getCause()).getErrorCode()); + + ExecutionException waiterException = + assertThrows(ExecutionException.class, () -> consumerFuture.get(5, TimeUnit.SECONDS)); + assertTrue(waiterException.getCause() instanceof SpannerException); + assertEquals( + ErrorCode.PERMISSION_DENIED, + ((SpannerException) waiterException.getCause()).getErrorCode()); + + // The waiting consumer thread must fail with the leader's startup failure rather than + // initiating a second startStream() call. + assertEquals(1, startStreamCalls.get()); + } finally { + executor.shutdownNow(); + } + } + + @Test(timeout = 10_000) + public void errorInStartStreamDoesNotDeadlockWaitingFollowerAndClosesInFlightStreamAndSpan() + throws Exception { + AtomicInteger startCount = new AtomicInteger(); + CountDownLatch consumerWaitingLatch = new CountDownLatch(1); + CountDownLatch allowStartStreamToThrowError = new CountDownLatch(1); + AtomicReference consumerThreadRef = new AtomicReference<>(); + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + ISpan span = mock(ISpan.class); + + ResumableStreamIterator iterator = + new TestResumableStreamIterator(mockTracerWithSpan(span)) { + @Override + CloseableIterator startStream( + @Nullable ByteString resumeToken, + @Nullable StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + startCount.incrementAndGet(); + GrpcStreamIterator stream = + new GrpcStreamIterator( + /* lastStatement= */ true, 10, /* cancelQueryWhenClientIsClosed= */ false); + if (streamMessageListener != null) { + stream.registerListener(streamMessageListener); + } + setStream(stream); + stream.setCall(call, /* withBeginTransaction= */ true); + consumerWaitingLatch.countDown(); + Uninterruptibles.awaitUninterruptibly( + allowStartStreamToThrowError, 5, TimeUnit.SECONDS); + throw new AssertionError("fatal error during startStream"); + } + }; + + AbstractResultSet.Listener listener = mockErrorAndDoneListener(); + ExecutorService executor = Executors.newFixedThreadPool(2); + try (GrpcResultSet grpcResultSet = new GrpcResultSet(iterator, listener)) { + Future initiateFuture = executor.submit(() -> iterator.initiateStreaming(null)); + assertTrue(consumerWaitingLatch.await(5, TimeUnit.SECONDS)); + + CountDownLatch followerEntered = new CountDownLatch(1); + Future followerFuture = + executor.submit( + () -> { + consumerThreadRef.set(Thread.currentThread()); + followerEntered.countDown(); + return grpcResultSet.next(); + }); + + assertTrue(followerEntered.await(5, TimeUnit.SECONDS)); + awaitThreadWaiting(consumerThreadRef.get()); + allowStartStreamToThrowError.countDown(); + + ExecutionException leaderException = + assertThrows(ExecutionException.class, () -> initiateFuture.get(2, TimeUnit.SECONDS)); + assertTrue(leaderException.getCause() instanceof AssertionError); + assertEquals("fatal error during startStream", leaderException.getCause().getMessage()); + + ExecutionException followerException = + assertThrows(ExecutionException.class, () -> followerFuture.get(2, TimeUnit.SECONDS)); + assertTrue(followerException.getCause() instanceof SpannerException); + assertTrue(followerException.getCause().getCause() instanceof AssertionError); + assertEquals( + "fatal error during startStream", followerException.getCause().getCause().getMessage()); + + verify(listener, Mockito.times(1)) + .onError(Mockito.any(SpannerException.class), Mockito.eq(true), Mockito.eq(true)); + verify(listener, Mockito.times(1)).onDone(true); + + assertEquals(1, startCount.get()); + verify(call, Mockito.times(1)).cancel("fatal error during startStream"); + verify(span, Mockito.times(1)).setStatus(Mockito.any(Throwable.class)); + verify(span, Mockito.times(1)).end(); + assertTrue(iterator.isWithBeginTransaction()); + assertTrue(iterator.isLastStatement()); + } finally { + executor.shutdownNow(); + } + } + + @Test(timeout = 10_000) + public void + concurrentCloseDuringStartStreamAfterSetCallCapturesWithBeginTransactionAndClosesStreamOnce() + throws Exception { + CountDownLatch setCallCompleted = new CountDownLatch(1); + CountDownLatch allowStartStreamToReturn = new CountDownLatch(1); + AtomicInteger streamCloseCount = new AtomicInteger(); + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + ISpan span = mock(ISpan.class); + + ResumableStreamIterator iterator = + new TestResumableStreamIterator(mockTracerWithSpan(span)) { + @Override + CloseableIterator startStream( + @Nullable ByteString resumeToken, + @Nullable StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + GrpcStreamIterator stream = + new GrpcStreamIterator( + /* lastStatement= */ true, 10, /* cancelQueryWhenClientIsClosed= */ false) { + @Override + public void close(@Nullable String message) { + streamCloseCount.incrementAndGet(); + super.close(message); + } + }; + // Register the stream early before stream.setCall(...) sets withBeginTransaction. + setStream(stream); + stream.setCall(call, /* withBeginTransaction= */ true); + setCallCompleted.countDown(); + Uninterruptibles.awaitUninterruptibly(allowStartStreamToReturn, 5, TimeUnit.SECONDS); + return stream; + } + }; + + AbstractResultSet.Listener listener = mockErrorAndDoneListener(); + ExecutorService executor = Executors.newSingleThreadExecutor(); + try (GrpcResultSet grpcResultSet = new GrpcResultSet(iterator, listener)) { + Future initiateFuture = executor.submit(() -> iterator.initiateStreaming(null)); + assertTrue(setCallCompleted.await(5, TimeUnit.SECONDS)); + + // Close while startStream() is still in flight after setCall(call, true) has completed. + grpcResultSet.close(); + + // Immediately after close() returns (before startStream() returns), withBeginTransaction and + // lastStatement must already reflect the in-flight stream's state, and GrpcResultSet.close() + // must have notified listener.onDone(true). + verify(listener, Mockito.times(1)).onDone(true); + assertTrue(iterator.isWithBeginTransaction()); + assertTrue(iterator.isLastStatement()); + assertEquals(1, streamCloseCount.get()); + verify(call, Mockito.times(1)).cancel("ResultSet closed"); + + allowStartStreamToReturn.countDown(); + assertTrue(initiateFuture.get(5, TimeUnit.SECONDS)); + + // Returning the same in-flight stream from startStream() to startGrpcStreaming() must not + // close the stream a second time. + assertEquals(1, streamCloseCount.get()); + verify(call, Mockito.times(1)).cancel("ResultSet closed"); + verify(span, Mockito.times(1)).end(); + assertTrue(iterator.isWithBeginTransaction()); + assertTrue(iterator.isLastStatement()); + } finally { + executor.shutdownNow(); + } + } + + @Test(timeout = 10_000) + public void canPollWithoutErrorOrBlockingStopsDrainingAndReturnsFalseWhenClosedConcurrently() { + initWithLimit(10); + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + GrpcStreamIterator grpcStream = new GrpcStreamIterator(false, 10, false); + grpcStream.setCall(call, false); + when(starter.startStream(null, null)).thenReturn(grpcStream); + resumableStreamIterator.initiateStreaming(null); + + // Enqueue two tokenless chunks so drainAvailableFromStream() would normally pull both. + grpcStream.consumer().onPartialResultSet(resultSet(null, "chunk1")); + grpcStream.consumer().onPartialResultSet(resultSet(null, "chunk2")); + + // Close the ResumableStreamIterator while the first chunk is being dequeued in + // GrpcStreamIterator.computeNext() -> call.request(1). + Mockito.doAnswer( + invocation -> { + resumableStreamIterator.close("closed while draining"); + return null; + }) + .when(call) + .request(1); + + assertFalse(resumableStreamIterator.canPollWithoutErrorOrBlocking()); + assertFalse(resumableStreamIterator.hasNext()); + // Must stop draining immediately after the first chunk rather than pulling the second chunk: + // chunk2 must still be sitting unconsumed in grpcStream. + assertTrue(grpcStream.canPollWithoutErrorOrBlocking()); + assertEquals("chunk2", grpcStream.next().getValues(0).getStringValue()); + verify(call, Mockito.times(1)).request(1); + verify(call, Mockito.times(1)).cancel("closed while draining"); + } + + @Test(timeout = 10_000) + public void + startStreamFailureAfterSetCallReportsDataAvailableAndPreservesWithBeginTransactionForGrpcResultSet() { + AtomicInteger startCount = new AtomicInteger(); + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + Mockito.doThrow(new RuntimeException("prefetch failed after setCall")).when(call).request(10); + ISpan span = mock(ISpan.class); + + ResumableStreamIterator iterator = + new TestResumableStreamIterator(mockTracerWithSpan(span)) { + @Override + CloseableIterator startStream( + @Nullable ByteString resumeToken, + @Nullable StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + startCount.incrementAndGet(); + GrpcStreamIterator stream = + new GrpcStreamIterator( + /* lastStatement= */ true, 10, /* cancelQueryWhenClientIsClosed= */ false); + if (streamMessageListener != null) { + stream.registerListener(streamMessageListener); + } + // Register the stream early before stream.setCall(call, withBeginTransaction) sets + // withBeginTransaction = true. + setStream(stream); + assertFalse(stream.isWithBeginTransaction()); + stream.setCall(call, /* withBeginTransaction= */ true); + assertTrue(stream.isWithBeginTransaction()); + stream.requestPrefetchChunks(); + return stream; + } + }; + + AbstractResultSet.Listener listener = mockErrorAndDoneListener(); + try (GrpcResultSet grpcResultSet = new GrpcResultSet(iterator, listener)) { + RuntimeException initialFailure = + assertThrows( + RuntimeException.class, + () -> grpcResultSet.initiateStreaming(mock(StreamMessageListener.class))); + assertEquals("prefetch failed after setCall", initialFailure.getMessage()); + + // Even before GrpcResultSet or the caller closes the iterator, the recorded startup failure + // is immediately available to surface without blocking on network I/O, and + // resetStreamOnFailure() must have captured withBeginTransaction = true from inFlightStream. + assertFalse(iterator.canPollWithoutErrorOrBlocking()); + assertTrue(iterator.isDataAvailable()); + assertTrue(iterator.isWithBeginTransaction()); + assertTrue(iterator.isLastStatement()); + + SpannerException surfacedException = + assertThrows(SpannerException.class, grpcResultSet::next); + assertTrue(surfacedException.getMessage().contains("prefetch failed after setCall")); + } + assertEquals(1, startCount.get()); + verify(listener) + .onError(Mockito.any(SpannerException.class), Mockito.eq(true), Mockito.eq(true)); + verify(listener).onDone(true); + verify(call, Mockito.times(1)).cancel("prefetch failed after setCall"); + verify(span, Mockito.times(1)).setStatus(Mockito.any(Throwable.class)); + verify(span, Mockito.times(1)).end(); + } + + @Test(timeout = 10_000) + public void withBeginTransactionAndLastStatementPreservedWhenCancelledDuringRetryBackoff() + throws Exception { + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + ISpan span = mock(ISpan.class); + ISpan currentSpan = mock(ISpan.class); + TraceWrapper tracer = mockTracerWithSpan(span); + when(tracer.getCurrentSpan()).thenReturn(currentSpan); + + Context.CancellableContext cancellableContext = Context.current().withCancellation(); + Mockito.doAnswer( + invocation -> { + cancellableContext.cancel(new InterruptedException("cancelled during retry backoff")); + return null; + }) + .when(currentSpan) + .addAnnotation(Mockito.eq("Backing off"), Mockito.eq("Delay"), Mockito.anyLong()); + + ResumableStreamIterator iterator = + new TestResumableStreamIterator(tracer) { + @Override + CloseableIterator startStream( + @Nullable ByteString resumeToken, + @Nullable StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + GrpcStreamIterator stream = + new GrpcStreamIterator( + /* lastStatement= */ true, 10, /* cancelQueryWhenClientIsClosed= */ false); + stream.setCall(call, /* withBeginTransaction= */ true); + stream + .consumer() + .onError(new RetryableException(errorCodeParameter, "transient error")); + return stream; + } + }; + + AbstractResultSet.Listener listener = mockErrorAndDoneListener(); + try (GrpcResultSet grpcResultSet = new GrpcResultSet(iterator, listener)) { + SpannerException thrown = + cancellableContext.call(() -> assertThrows(SpannerException.class, grpcResultSet::next)); + assertEquals(ErrorCode.CANCELLED, thrown.getErrorCode()); + } finally { + cancellableContext.close(); + } + + verify(listener) + .onError(Mockito.any(SpannerException.class), Mockito.eq(true), Mockito.eq(true)); + verify(listener).onDone(true); + verify(span, Mockito.times(1)).setStatus(Mockito.any(Throwable.class)); + verify(span, Mockito.times(1)).end(); + } + + @Test + public void unhandledErrorInStreamNextSetsSpanStatusAndThrows() { + ISpan mockSpan = mock(ISpan.class); + initWithLimitAndRetrySettings( + 10, RetrySettings.newBuilder().build(), mockTracerWithSpan(mockSpan)); + + ResultSetStream firstStream = mock(ResultSetStream.class); + when(starter.startStream(null, null)).thenReturn(new ResultSetIterator(firstStream)); + when(firstStream.next()).thenThrow(new AssertionError("unexpected error in stream.next")); + + AssertionError thrown = assertThrows(AssertionError.class, resumableStreamIterator::next); + assertEquals("unexpected error in stream.next", thrown.getMessage()); + verify(mockSpan) + .addAnnotation( + Mockito.eq("Stream broken. Not safe to retry"), Mockito.any(Throwable.class)); + verify(mockSpan).setStatus(Mockito.any(Throwable.class)); + verify(mockSpan).end(); + } + + @Test(timeout = 10_000) + public void closeDuringRequestPrefetchChunksCancelsOnceAndTerminatesCleanly() { + AtomicInteger startCount = new AtomicInteger(); + AtomicReference iteratorRef = new AtomicReference<>(); + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + Mockito.doAnswer( + invocation -> { + iteratorRef.get().close("closed during prefetch"); + return null; + }) + .when(call) + .request(10); + + ISpan span = mock(ISpan.class); + ResumableStreamIterator iterator = + new TestResumableStreamIterator(mockTracerWithSpan(span)) { + @Override + CloseableIterator startStream( + @Nullable ByteString resumeToken, + @Nullable StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + startCount.incrementAndGet(); + GrpcStreamIterator stream = new GrpcStreamIterator(true, 10, false); + stream.setCall(call, true); + return stream; + } + }; + iteratorRef.set(iterator); + + assertTrue(iterator.initiateStreaming(null)); + assertFalse(iterator.hasNext()); + assertEquals(1, startCount.get()); + verify(call, Mockito.times(1)).cancel("closed during prefetch"); + verify(span, Mockito.times(1)).end(); + } + + @Test(timeout = 10_000) + public void + retryableStreamErrorDuringInitialStartStreamIsRetriedWhileSynchronousStartStreamExceptionFailsFast() { + // 1. A retryable stream error delivered via the gRPC stream consumer during initial + // startStream() is retried transparently and yields rows from the restarted stream. + AtomicInteger streamStartCount = new AtomicInteger(); + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + ISpan successSpan = mock(ISpan.class); + Transaction expectedTransaction = + Transaction.newBuilder().setId(ByteString.copyFromUtf8("tx-1")).build(); + ResumableStreamIterator retryingIterator = + new TestResumableStreamIterator(mockTracerWithSpan(successSpan)) { + @Override + CloseableIterator startStream( + @Nullable ByteString resumeToken, + @Nullable StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + int attempt = streamStartCount.incrementAndGet(); + GrpcStreamIterator stream = new GrpcStreamIterator(true, 10, false); + if (streamMessageListener != null) { + stream.registerListener(streamMessageListener); + } + setStream(stream); + stream.setCall(call, true); + if (attempt == 1) { + stream + .consumer() + .onError(new RetryableException(errorCodeParameter, "initial stream failed")); + } else { + stream + .consumer() + .onPartialResultSet( + PartialResultSet.newBuilder() + .setMetadata( + MockSpannerTestUtil.READ_KEY_VALUE_METADATA.toBuilder() + .setTransaction(expectedTransaction) + .build()) + .setResumeToken(ByteString.copyFromUtf8("r1")) + .setLast(true) + .addValues(Value.newBuilder().setStringValue("k1")) + .addValues(Value.newBuilder().setStringValue("v1")) + .build()); + stream.consumer().onCompleted(); + } + return stream; + } + }; + + AbstractResultSet.Listener successListener = mockErrorAndDoneListener(); + try (GrpcResultSet grpcResultSet = new GrpcResultSet(retryingIterator, successListener)) { + grpcResultSet.initiateStreaming(mock(StreamMessageListener.class)); + assertTrue(grpcResultSet.next()); + assertEquals("k1", grpcResultSet.getString(0)); + assertEquals("v1", grpcResultSet.getString(1)); + assertFalse(grpcResultSet.next()); + } + assertEquals(2, streamStartCount.get()); + verify(successListener).onTransactionMetadata(expectedTransaction, true); + verify(successListener).onDone(true); + verify(successSpan, Mockito.times(1)).end(); + + // 2. A synchronous exception thrown directly by startStream() itself (even with a retryable + // error code) is a fatal pre-stream failure: it records the failure on the span, ends the span + // once, and fails fast on subsequent calls without re-invoking startStream(). + AtomicInteger syncFailStartCount = new AtomicInteger(); + ISpan failSpan = mock(ISpan.class); + ResumableStreamIterator failingIterator = + new TestResumableStreamIterator(mockTracerWithSpan(failSpan)) { + @Override + CloseableIterator startStream( + @Nullable ByteString resumeToken, + @Nullable StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + syncFailStartCount.incrementAndGet(); + throw new RetryableException(errorCodeParameter, "synchronous startStream failure"); + } + }; + + AbstractResultSet.Listener failListener = mockErrorAndDoneListener(); + try (GrpcResultSet failingResultSet = new GrpcResultSet(failingIterator, failListener)) { + SpannerException initialException = + assertThrows( + SpannerException.class, + () -> failingResultSet.initiateStreaming(mock(StreamMessageListener.class))); + assertEquals(errorCodeParameter, initialException.getErrorCode()); + assertTrue(failingIterator.isDataAvailable()); + + SpannerException subsequentException = + assertThrows(SpannerException.class, failingResultSet::next); + assertEquals(errorCodeParameter, subsequentException.getErrorCode()); + } + assertEquals(1, syncFailStartCount.get()); + verify(failSpan, Mockito.times(1)).setStatus(Mockito.any(Throwable.class)); + verify(failSpan, Mockito.times(1)).end(); + } + + @Test(timeout = 10_000) + public void concurrentCloseDuringReadinessPollingDoesNotThrowCancelledException() { + for (boolean pollViaIsDataAvailable : new boolean[] {false, true}) { + AtomicReference iteratorRef = new AtomicReference<>(); + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + Mockito.doAnswer( + invocation -> { + iteratorRef.get().close("closed while polling"); + throw SpannerExceptionFactory.newSpannerException( + ErrorCode.CANCELLED, "closed while polling"); + }) + .when(call) + .request(1); + + ISpan span = mock(ISpan.class); + ResumableStreamIterator iterator = + new TestResumableStreamIterator(mockTracerWithSpan(span)) { + @Override + CloseableIterator startStream( + @Nullable ByteString resumeToken, + @Nullable StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + GrpcStreamIterator stream = new GrpcStreamIterator(false, 10, false); + setStream(stream); + stream.setCall(call, false); + stream.consumer().onPartialResultSet(resultSet(null, "chunk1")); + return stream; + } + }; + iteratorRef.set(iterator); + assertTrue(iterator.initiateStreaming(null)); + + if (pollViaIsDataAvailable) { + assertTrue(iterator.isDataAvailable()); + } else { + assertFalse(iterator.canPollWithoutErrorOrBlocking()); + } + assertFalse(iterator.canPollWithoutErrorOrBlocking()); + assertTrue(iterator.isDataAvailable()); + assertFalse(iterator.hasNext()); + verify(call, Mockito.times(1)).cancel("closed while polling"); + verify(span, Mockito.times(1)).end(); + } + } + + private abstract static class TestResumableStreamIterator extends ResumableStreamIterator { + TestResumableStreamIterator() { + this(new TraceWrapper(Tracing.getTracer(), OpenTelemetry.noop().getTracer(""), false)); + } + + TestResumableStreamIterator(TraceWrapper tracer) { + this( + tracer, + DefaultErrorHandler.INSTANCE, + SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetrySettings()); + } + + TestResumableStreamIterator(ErrorHandler errorHandler) { + this( + new TraceWrapper(Tracing.getTracer(), OpenTelemetry.noop().getTracer(""), false), + errorHandler, + SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetrySettings()); + } + + TestResumableStreamIterator(TraceWrapper tracer, RetrySettings retrySettings) { + this(tracer, DefaultErrorHandler.INSTANCE, retrySettings); + } + + TestResumableStreamIterator( + TraceWrapper tracer, ErrorHandler errorHandler, RetrySettings retrySettings) { + super( + 10, + "", + mockOpenTelemetrySpan(), + tracer, + errorHandler, + retrySettings, + SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetryableCodes(), + NoopRequestIdCreator.INSTANCE); + } + } + + private static TraceWrapper mockTracerWithSpan(ISpan span) { + TraceWrapper tracer = mock(TraceWrapper.class); + when(tracer.spanBuilderWithExplicitParent(Mockito.any(), Mockito.any(), Mockito.any())) + .thenReturn(span); + when(tracer.withSpan(span)).thenReturn(() -> {}); + when(tracer.getCurrentSpan()).thenReturn(span); + return tracer; + } + + private static AbstractResultSet.Listener mockErrorAndDoneListener() { + AbstractResultSet.Listener listener = mock(AbstractResultSet.Listener.class); + when(listener.onError(Mockito.any(), Mockito.anyBoolean(), Mockito.anyBoolean())) + .thenAnswer(invocation -> invocation.getArgument(0)); + return listener; + } + + private static void awaitThreadWaiting(Thread thread) { + long deadlineNanos = System.nanoTime() + TimeUnit.SECONDS.toNanos(5); + while (thread.getState() != Thread.State.WAITING + && thread.getState() != Thread.State.TIMED_WAITING) { + if (System.nanoTime() > deadlineNanos) { + fail("Timed out waiting for thread to enter WAITING state; current=" + thread.getState()); + } + Thread.yield(); + } + } + static List consume(Iterator iterator) { return Lists.newArrayList(stringIterator(iterator)); }