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)); }