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 0717cae74f25..2a1c781c23d7 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 @@ -166,6 +166,23 @@ default boolean initiateStreaming(AsyncResultSet.StreamMessageListener streamMes /** it requests the initial prefetch chunks from gRPC stream */ 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. + */ + default boolean isDataAvailable() { + return true; + } + + /** + * 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. + */ + default boolean canPollWithoutErrorOrBlocking() { + return false; + } } static double valueProtoToFloat64(com.google.protobuf.Value proto) { diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncResultSet.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncResultSet.java index 2b3225bfc598..6d9025eb455a 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncResultSet.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncResultSet.java @@ -17,6 +17,7 @@ package com.google.cloud.spanner; import com.google.api.core.ApiFuture; +import com.google.api.core.InternalApi; import com.google.common.base.Function; import com.google.spanner.v1.PartialResultSet; import java.util.List; @@ -226,10 +227,36 @@ interface ReadyCallback { List toList(Function transformer) throws SpannerException; /** - * An interface to register the listener for streaming gRPC request. It will be called when a - * chunk is received from gRPC streaming call. + * An internal notification listener for a streaming gRPC request. + * + *

This listener is purely a notification mechanism to signal that messages or stream events + * are buffered in the underlying stream iterator. It is not a data consumer and does not + * process or reconstruct rows. All query results and data chunks must be consumed sequentially + * via the iterator's FIFO queue (e.g. {@link ResultSet#next()}). */ + @InternalApi interface StreamMessageListener { + /** + * Invoked to notify that a message or stream event has arrived in the stream's buffer. + * + *

Implementations of this method must be thread-safe, idempotent, and order-agnostic. + * Because notifications may be dispatched concurrently by gRPC network threads and iterator + * initialization/flushing threads: + * + *

    + *
  • Notifications may arrive concurrently from multiple threads. + *
  • The same {@code partialResultSet} may be delivered more than once. + *
  • Notifications for different chunks may arrive out of order. + *
+ * + *

Implementations must tolerate out-of-order and duplicate notifications without adverse + * side effects, and must not assume that this method delivers chunks sequentially for data + * processing. + * + * @param partialResultSet the chunk received from the stream, or an empty sentinel {@link + * PartialResultSet} if the stream has completed + * @param bufferIsFull whether the stream iterator's internal buffer capacity is currently full + */ void onStreamMessage(PartialResultSet partialResultSet, boolean bufferIsFull); } } diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/GrpcStreamIterator.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/GrpcStreamIterator.java index 78ab074b891e..e49860fbfd08 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/GrpcStreamIterator.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/GrpcStreamIterator.java @@ -76,7 +76,7 @@ public String toString() { private static final Logger logger = Logger.getLogger(GrpcStreamIterator.class.getName()); static final PartialResultSet END_OF_STREAM = PartialResultSet.newBuilder().build(); private final int prefetchChunks; - private AsyncResultSet.StreamMessageListener streamMessageListener; + private volatile AsyncResultSet.StreamMessageListener streamMessageListener; private final ConsumerImpl consumer; private final BlockingQueue stream; @@ -91,6 +91,8 @@ public String toString() { @Nullable private volatile Timeout streamWaitTimeout; @Nullable private volatile SpannerException error; private volatile boolean done; + private volatile boolean callInitialized; + private volatile boolean endOfStreamReached; @VisibleForTesting GrpcStreamIterator( @@ -118,6 +120,20 @@ protected final SpannerRpc.ResultStreamConsumer consumer() { void registerListener(AsyncResultSet.StreamMessageListener streamMessageListener) { this.streamMessageListener = Preconditions.checkNotNull(streamMessageListener); + // If setCall() has already completed, flush notifications for any messages that were buffered + // before this listener was registered. + // + // This check is safe from races with concurrent setCall() execution: + // 1. streamMessageListener is set before reading callInitialized (both volatile). If this check + // observes false, setCall() is guaranteed to observe the non-null listener and perform the + // flush itself. + // 2. If both registerListener() and setCall() run the flush concurrently, duplicate deliveries + // are safe because onStreamMessage() is idempotent. + if (this.callInitialized) { + for (PartialResultSet pending : stream) { + onStreamMessage(pending); + } + } } public void setCall(SpannerRpc.StreamingCall call, boolean withBeginTransaction) { @@ -130,6 +146,7 @@ public void setCall(SpannerRpc.StreamingCall call, boolean withBeginTransaction) if (!this.callCleared && !this.closed) { this.call = call; } + this.callInitialized = true; if (timeout != null) { this.streamWaitTimeout = timeout; } @@ -137,6 +154,21 @@ public void setCall(SpannerRpc.StreamingCall call, boolean withBeginTransaction) } if (shouldCancel) { call.cancel(messageToCancel); + return; + } + // Flush notifications for any messages that arrived before setCall() completed. + // While callInitialized was false, addToStream() suppressed onStreamMessage() notifications + // to prevent asynchronous consumers from processing rows before the underlying StreamingCall + // was established. + // + // This is safe with concurrent stream activity because: + // 1. LinkedBlockingQueue's iterator is weakly consistent and thread-safe. + // 2. Since callInitialized was set to true above, any chunk added concurrently during or after + // this loop will trigger onStreamMessage() in addToStream(). If this loop also visits a + // concurrently added chunk during iteration, duplicate onStreamMessage() invocations are + // completely idempotent in AsyncResultSetImpl. + for (PartialResultSet pending : stream) { + onStreamMessage(pending); } } @@ -209,6 +241,45 @@ public boolean isLastStatement() { return lastStatement; } + /** + * Returns true if this iterator is ready to allow polling or error delivery. + * + *

When an asynchronous {@link AsyncResultSet.StreamMessageListener} is registered, consumption + * is gated until the underlying {@link SpannerRpc.StreamingCall} is initialized via {@link + * #setCall(SpannerRpc.StreamingCall, boolean)} so that call context, cancellation hooks, and + * listener notifications are properly established. In synchronous mode (where no listener is + * registered, e.g. {@code streamMessageListener == null}), this gating is not needed because + * callers invoke {@link #next()} directly on the iterator. + */ + private boolean isReadyForConsumption() { + return callInitialized || streamMessageListener == null; + } + + @Override + public boolean canPollWithoutErrorOrBlocking() { + // A client can poll without error or blocking on network I/O if: + // 1. The iterator is ready for consumption (call is initialized or iterator is synchronous). + // 2. AND either: + // a) The stream has already reached EOF (computeNext() will return null immediately), OR + // b) The next message at the head of the stream is an unconsumed data chunk (which next() + // will return immediately without error and without I/O), OR + // c) The next message is EOF and no stream error occurred. + PartialResultSet head = stream.peek(); + return isReadyForConsumption() + && (endOfStreamReached || (head != null && (head != END_OF_STREAM || error == null))); + } + + @Override + public boolean isDataAvailable() { + // Data is available without blocking on network I/O if: + // 1. canPollWithoutErrorOrBlocking() is true (a normal chunk or EOF is ready to read + // immediately), OR + // 2. An error has occurred and is buffered in memory, ready to be thrown immediately by + // computeNext() without waiting on network I/O (assuming the iterator is ready for + // consumption). + return canPollWithoutErrorOrBlocking() || (isReadyForConsumption() && error != null); + } + @Override protected final PartialResultSet computeNext() { PartialResultSet next = takeNextFromStream(); @@ -227,6 +298,7 @@ protected final PartialResultSet computeNext() { if (error != null) { throw SpannerExceptionFactory.asSpannerException(error); } + endOfStreamReached = true; endOfData(); return null; } @@ -261,7 +333,9 @@ private void clearCall() { private void addToStream(PartialResultSet results) { // We assume that nothing from the user will interrupt gRPC event threads. Uninterruptibles.putUninterruptibly(stream, results); - onStreamMessage(results); + if (callInitialized) { + onStreamMessage(results); + } } private class ConsumerImpl implements SpannerRpc.ResultStreamConsumer { @@ -320,6 +394,8 @@ public boolean cancelQueryWhenClientIsClosed() { private void onStreamMessage(PartialResultSet partialResultSet) { Optional.ofNullable(streamMessageListener) - .ifPresent(sl -> sl.onStreamMessage(partialResultSet, stream.remainingCapacity() <= 1)); + .ifPresent( + listener -> + listener.onStreamMessage(partialResultSet, stream.remainingCapacity() <= 1)); } } diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/GrpcStreamIteratorTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/GrpcStreamIteratorTest.java index c451fe9828e6..72d03d533152 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/GrpcStreamIteratorTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/GrpcStreamIteratorTest.java @@ -24,6 +24,7 @@ import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyBoolean; import static org.mockito.ArgumentMatchers.anyInt; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; @@ -33,17 +34,21 @@ import com.google.api.gax.grpc.GrpcCallContext; import com.google.api.gax.rpc.ApiCallContext; +import com.google.cloud.spanner.AbstractResultSet.CloseableIterator; import com.google.cloud.spanner.spi.v1.SpannerRpc; import com.google.protobuf.ByteString; import com.google.protobuf.Value; import com.google.spanner.v1.PartialResultSet; import java.time.Duration; +import java.util.NoSuchElementException; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; +import org.jspecify.annotations.Nullable; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -798,4 +803,772 @@ public void interruptedTakeNextFromStreamThrowsPropagateInterrupt() throws Excep assertEquals(ErrorCode.CANCELLED, spannerException.getErrorCode()); assertTrue(spannerException.getCause() instanceof InterruptedException); } + + @Test + public void dataAvailableAndCanPollWithoutErrorOrBlockingEmptyStream() { + GrpcStreamIterator streamIterator = + new GrpcStreamIterator( + /* statement= */ null, + /* lastStatement= */ false, + /* prefetchChunks= */ 10, + /* cancelQueryWhenClientIsClosed= */ false); + + assertFalse(streamIterator.canPollWithoutErrorOrBlocking()); + assertFalse(streamIterator.isDataAvailable()); + + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); + streamIterator.setCall(call, /* withBeginTransaction= */ false); + + assertFalse(streamIterator.canPollWithoutErrorOrBlocking()); + assertFalse(streamIterator.isDataAvailable()); + } + + @Test + public void dataAvailableAndCanPollWithoutErrorOrBlockingWithBufferedChunksAndConsumption() { + GrpcStreamIterator streamIterator = + new GrpcStreamIterator( + /* statement= */ null, + /* lastStatement= */ false, + /* prefetchChunks= */ 10, + /* cancelQueryWhenClientIsClosed= */ false); + + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); + streamIterator.setCall(call, /* withBeginTransaction= */ false); + + PartialResultSet chunk1 = + PartialResultSet.newBuilder() + .setResumeToken(ByteString.copyFromUtf8("token-1")) + .addValues(Value.newBuilder().setStringValue("value-1")) + .build(); + PartialResultSet chunk2 = + PartialResultSet.newBuilder() + .setResumeToken(ByteString.copyFromUtf8("token-2")) + .addValues(Value.newBuilder().setStringValue("value-2")) + .build(); + + streamIterator.consumer().onPartialResultSet(chunk1); + assertTrue(streamIterator.canPollWithoutErrorOrBlocking()); + assertTrue(streamIterator.isDataAvailable()); + + streamIterator.consumer().onPartialResultSet(chunk2); + assertTrue(streamIterator.canPollWithoutErrorOrBlocking()); + assertTrue(streamIterator.isDataAvailable()); + + assertEquals(chunk1, streamIterator.next()); + assertTrue(streamIterator.canPollWithoutErrorOrBlocking()); + assertTrue(streamIterator.isDataAvailable()); + + assertEquals(chunk2, streamIterator.next()); + assertFalse(streamIterator.canPollWithoutErrorOrBlocking()); + assertFalse(streamIterator.isDataAvailable()); + } + + @Test + public void dataAvailableAndCanPollWithoutErrorOrBlockingWhenEndOfStreamReached() { + GrpcStreamIterator streamIterator = + new GrpcStreamIterator( + /* statement= */ null, + /* lastStatement= */ false, + /* prefetchChunks= */ 10, + /* cancelQueryWhenClientIsClosed= */ false); + + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); + streamIterator.setCall(call, /* withBeginTransaction= */ false); + + PartialResultSet lastChunk = + PartialResultSet.newBuilder() + .setResumeToken(ByteString.copyFromUtf8("token-1")) + .setLast(true) + .addValues(Value.newBuilder().setStringValue("value-1")) + .build(); + + streamIterator.consumer().onPartialResultSet(lastChunk); + assertTrue(streamIterator.canPollWithoutErrorOrBlocking()); + assertTrue(streamIterator.isDataAvailable()); + + assertEquals(lastChunk, streamIterator.next()); + // The END_OF_STREAM token is still queued in the stream queue, so data availability remains + // true. + assertTrue(streamIterator.canPollWithoutErrorOrBlocking()); + assertTrue(streamIterator.isDataAvailable()); + + assertFalse(streamIterator.hasNext()); + // After EOF is processed by computeNext(), endOfStreamReached is set to true. + assertTrue(streamIterator.canPollWithoutErrorOrBlocking()); + assertTrue(streamIterator.isDataAvailable()); + } + + @Test + public void dataAvailableAndCanPollWithoutErrorOrBlockingWhenStreamErrorOccurs() { + GrpcStreamIterator streamIterator = + new GrpcStreamIterator( + /* statement= */ null, + /* lastStatement= */ false, + /* prefetchChunks= */ 10, + /* cancelQueryWhenClientIsClosed= */ false); + + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); + streamIterator.setCall(call, /* withBeginTransaction= */ false); + + SpannerException spannerException = + SpannerExceptionFactory.newSpannerException(ErrorCode.UNAVAILABLE, "network unavailable"); + streamIterator.consumer().onError(spannerException); + + // An error is present, so polling cannot succeed without throwing an exception. + assertFalse(streamIterator.canPollWithoutErrorOrBlocking()); + // However, data (the error) is available immediately without blocking on I/O. + assertTrue(streamIterator.isDataAvailable()); + + SpannerException thrown = assertThrows(SpannerException.class, streamIterator::next); + assertEquals(ErrorCode.UNAVAILABLE, thrown.getErrorCode()); + } + + @Test + public void listenerRegisteredBeforeSetCallDelaysEventDeliveryUntilCallInitialized() { + GrpcStreamIterator streamIterator = + new GrpcStreamIterator( + /* statement= */ null, + /* lastStatement= */ false, + /* prefetchChunks= */ 10, + /* cancelQueryWhenClientIsClosed= */ false); + + AtomicReference receivedMessage = new AtomicReference<>(); + AtomicBoolean receivedBlockedState = new AtomicBoolean(); + streamIterator.registerListener( + (partialResultSet, isBlocked) -> { + receivedMessage.set(partialResultSet); + receivedBlockedState.set(isBlocked); + }); + + PartialResultSet chunk = + PartialResultSet.newBuilder() + .setResumeToken(ByteString.copyFromUtf8("token-1")) + .addValues(Value.newBuilder().setStringValue("value-1")) + .build(); + + // Delivering chunk before setCall must not notify listener yet. + streamIterator.consumer().onPartialResultSet(chunk); + assertNull(receivedMessage.get()); + + // Because a listener is registered but callInitialized is false, canPollWithoutErrorOrBlocking + // is false. + assertFalse(streamIterator.canPollWithoutErrorOrBlocking()); + assertFalse(streamIterator.isDataAvailable()); + + // Now initialize call via setCall. + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); + streamIterator.setCall(call, /* withBeginTransaction= */ false); + + // Flushing pending message in setCall notifies the listener. + assertEquals(chunk, receivedMessage.get()); + assertFalse(receivedBlockedState.get()); + assertTrue(streamIterator.canPollWithoutErrorOrBlocking()); + assertTrue(streamIterator.isDataAvailable()); + } + + @Test + public void listenerRegisteredAfterSetCallReceivesMessagesImmediately() { + GrpcStreamIterator streamIterator = + new GrpcStreamIterator( + /* statement= */ null, + /* lastStatement= */ false, + /* prefetchChunks= */ 1, + /* cancelQueryWhenClientIsClosed= */ false); + + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); + streamIterator.setCall(call, /* withBeginTransaction= */ false); + + AtomicReference receivedMessage = new AtomicReference<>(); + AtomicBoolean receivedBlockedState = new AtomicBoolean(); + streamIterator.registerListener( + (partialResultSet, isBlocked) -> { + receivedMessage.set(partialResultSet); + receivedBlockedState.set(isBlocked); + }); + + PartialResultSet chunk = + PartialResultSet.newBuilder() + .setResumeToken(ByteString.copyFromUtf8("token-1")) + .addValues(Value.newBuilder().setStringValue("value-1")) + .build(); + + streamIterator.consumer().onPartialResultSet(chunk); + assertEquals(chunk, receivedMessage.get()); + // Stream queue has capacity 2 (prefetchChunks 1 + 1), and contains 1 element, so remaining + // capacity is 1 <= 1 (isBlocked = true). + assertTrue(receivedBlockedState.get()); + assertTrue(streamIterator.canPollWithoutErrorOrBlocking()); + assertTrue(streamIterator.isDataAvailable()); + } + + @Test + public void closeBeforeSetCallSafelyNoOpsRequestPrefetchChunks() { + GrpcStreamIterator streamIterator = + new GrpcStreamIterator( + /* statement= */ null, + /* lastStatement= */ false, + /* prefetchChunks= */ 10, + /* cancelQueryWhenClientIsClosed= */ false); + + streamIterator.close("Cancelled before setCall"); + // Since streamIterator was already closed, requestPrefetchChunks should safely no-op + // without throwing IllegalStateException, even though setCall was never invoked. + streamIterator.requestPrefetchChunks(); + } + + @Test + public void dataAvailableAndCanPollWithoutErrorOrBlockingWhenStreamCompletesWithZeroRows() { + GrpcStreamIterator streamIterator = + new GrpcStreamIterator( + /* statement= */ null, + /* lastStatement= */ false, + /* prefetchChunks= */ 10, + /* cancelQueryWhenClientIsClosed= */ false); + + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); + streamIterator.setCall(call, /* withBeginTransaction= */ false); + + // Initial state: No messages received yet. Polling would block on network I/O. + assertFalse(streamIterator.canPollWithoutErrorOrBlocking()); + assertFalse(streamIterator.isDataAvailable()); + + // Server finishes query with 0 rows via onCompleted(). + streamIterator.consumer().onCompleted(); + + // The stream has an EOF token queued. Polling will immediately return false/EOF without + // blocking. + assertTrue(streamIterator.canPollWithoutErrorOrBlocking()); + assertTrue(streamIterator.isDataAvailable()); + + // Calling hasNext() consumes the EOF marker. + assertFalse(streamIterator.hasNext()); + + // After EOF has been processed, the stream remains non-blocking and ready for callers. + assertTrue(streamIterator.canPollWithoutErrorOrBlocking()); + assertTrue(streamIterator.isDataAvailable()); + } + + @Test + public void dataAvailableAndCanPollWithoutErrorOrBlockingWhenAllDataConsumedWithPendingEof() { + GrpcStreamIterator streamIterator = + new GrpcStreamIterator( + /* statement= */ null, + /* lastStatement= */ false, + /* prefetchChunks= */ 10, + /* cancelQueryWhenClientIsClosed= */ false); + + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); + streamIterator.setCall(call, /* withBeginTransaction= */ false); + + PartialResultSet chunk = + PartialResultSet.newBuilder() + .setResumeToken(ByteString.copyFromUtf8("token-1")) + .addValues(Value.newBuilder().setStringValue("value-1")) + .build(); + + // Server sends a data chunk followed immediately by completion. + streamIterator.consumer().onPartialResultSet(chunk); + streamIterator.consumer().onCompleted(); + + // Both data chunk and EOF are in the queue. + assertTrue(streamIterator.canPollWithoutErrorOrBlocking()); + assertTrue(streamIterator.isDataAvailable()); + + // Client consumes the data chunk. + assertEquals(chunk, streamIterator.next()); + + // All data chunks are consumed, but there is still a pending EOF in the stream queue. + // Fetching the EOF does not block on network I/O. + assertTrue(streamIterator.canPollWithoutErrorOrBlocking()); + assertTrue(streamIterator.isDataAvailable()); + + // Client verifies stream completion. + assertFalse(streamIterator.hasNext()); + + // Stream has reached EOF. Both methods remain true. + assertTrue(streamIterator.canPollWithoutErrorOrBlocking()); + assertTrue(streamIterator.isDataAvailable()); + } + + @Test + public void dataAvailableWhenServerReturnsErrorAfterDrainingBuffer() { + GrpcStreamIterator streamIterator = + new GrpcStreamIterator( + /* statement= */ null, + /* lastStatement= */ false, + /* prefetchChunks= */ 10, + /* cancelQueryWhenClientIsClosed= */ false); + + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); + streamIterator.setCall(call, /* withBeginTransaction= */ false); + + PartialResultSet chunk = + PartialResultSet.newBuilder() + .setResumeToken(ByteString.copyFromUtf8("token-1")) + .addValues(Value.newBuilder().setStringValue("value-1")) + .build(); + + // Chunk arrives and is consumed by client. + streamIterator.consumer().onPartialResultSet(chunk); + assertEquals(chunk, streamIterator.next()); + + // Buffer is now drained and client is waiting for server. + assertFalse(streamIterator.canPollWithoutErrorOrBlocking()); + assertFalse(streamIterator.isDataAvailable()); + + // Server encounters an error. + SpannerException spannerException = + SpannerExceptionFactory.newSpannerException( + ErrorCode.DEADLINE_EXCEEDED, "deadline exceeded"); + streamIterator.consumer().onError(spannerException); + + // Because an error occurred, normal polling cannot proceed. + assertFalse(streamIterator.canPollWithoutErrorOrBlocking()); + // But data availability is true because the error is ready to be thrown immediately without + // network I/O. + assertTrue(streamIterator.isDataAvailable()); + + SpannerException thrown = assertThrows(SpannerException.class, streamIterator::next); + assertEquals(ErrorCode.DEADLINE_EXCEEDED, thrown.getErrorCode()); + } + + @Test(timeout = 10000) + public void concurrentSetCallAndOnPartialResultSet() throws Exception { + ExecutorService executorService = Executors.newFixedThreadPool(2); + try { + for (int i = 0; i < 50; i++) { + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); + + GrpcStreamIterator streamIterator = + new GrpcStreamIterator( + /* statement= */ null, + /* lastStatement= */ false, + /* prefetchChunks= */ 5, + /* cancelQueryWhenClientIsClosed= */ false); + + AtomicInteger listenerInvocations = new AtomicInteger(); + streamIterator.registerListener( + (message, isBufferFull) -> listenerInvocations.incrementAndGet()); + + PartialResultSet chunk = + PartialResultSet.newBuilder() + .setResumeToken(ByteString.copyFromUtf8("token-" + i)) + .addValues(Value.newBuilder().setStringValue("value-" + i)) + .build(); + + CountDownLatch startLatch = new CountDownLatch(1); + CountDownLatch doneLatch = new CountDownLatch(2); + AtomicReference threadError = new AtomicReference<>(); + + executorService.submit( + () -> { + try { + startLatch.await(); + streamIterator.setCall(call, /* withBeginTransaction= */ false); + } catch (Throwable throwable) { + threadError.set(throwable); + } finally { + doneLatch.countDown(); + } + }); + + executorService.submit( + () -> { + try { + startLatch.await(); + streamIterator.consumer().onPartialResultSet(chunk); + } catch (Throwable throwable) { + threadError.set(throwable); + } finally { + doneLatch.countDown(); + } + }); + + startLatch.countDown(); + assertTrue(doneLatch.await(5, TimeUnit.SECONDS)); + if (threadError.get() != null) { + throw new AssertionError(threadError.get()); + } + + assertTrue(streamIterator.canPollWithoutErrorOrBlocking()); + assertTrue(streamIterator.isDataAvailable()); + assertTrue(listenerInvocations.get() >= 1); + assertEquals(chunk, streamIterator.next()); + assertFalse(streamIterator.canPollWithoutErrorOrBlocking()); + assertFalse(streamIterator.isDataAvailable()); + } + } finally { + executorService.shutdownNow(); + } + } + + @Test(timeout = 10000) + public void concurrentSetCallAndOnError() throws Exception { + ExecutorService executorService = Executors.newFixedThreadPool(2); + try { + for (int i = 0; i < 50; i++) { + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); + + GrpcStreamIterator streamIterator = + new GrpcStreamIterator( + /* statement= */ null, + /* lastStatement= */ false, + /* prefetchChunks= */ 5, + /* cancelQueryWhenClientIsClosed= */ false); + + SpannerException spannerException = + SpannerExceptionFactory.newSpannerException( + ErrorCode.UNAVAILABLE, "network unavailable " + i); + + CountDownLatch startLatch = new CountDownLatch(1); + CountDownLatch doneLatch = new CountDownLatch(2); + AtomicReference threadError = new AtomicReference<>(); + + executorService.submit( + () -> { + try { + startLatch.await(); + streamIterator.setCall(call, /* withBeginTransaction= */ false); + } catch (Throwable throwable) { + threadError.set(throwable); + } finally { + doneLatch.countDown(); + } + }); + + executorService.submit( + () -> { + try { + startLatch.await(); + streamIterator.consumer().onError(spannerException); + } catch (Throwable throwable) { + threadError.set(throwable); + } finally { + doneLatch.countDown(); + } + }); + + startLatch.countDown(); + assertTrue(doneLatch.await(5, TimeUnit.SECONDS)); + if (threadError.get() != null) { + throw new AssertionError(threadError.get()); + } + + assertFalse(streamIterator.canPollWithoutErrorOrBlocking()); + assertTrue(streamIterator.isDataAvailable()); + SpannerException thrown = assertThrows(SpannerException.class, streamIterator::next); + assertEquals(ErrorCode.UNAVAILABLE, thrown.getErrorCode()); + } + } finally { + executorService.shutdownNow(); + } + } + + @Test(timeout = 10000) + public void concurrentSetCallAndOnCompleted() throws Exception { + ExecutorService executorService = Executors.newFixedThreadPool(2); + try { + for (int i = 0; i < 50; i++) { + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); + + GrpcStreamIterator streamIterator = + new GrpcStreamIterator( + /* statement= */ null, + /* lastStatement= */ false, + /* prefetchChunks= */ 5, + /* cancelQueryWhenClientIsClosed= */ false); + + CountDownLatch startLatch = new CountDownLatch(1); + CountDownLatch doneLatch = new CountDownLatch(2); + AtomicReference threadError = new AtomicReference<>(); + + executorService.submit( + () -> { + try { + startLatch.await(); + streamIterator.setCall(call, /* withBeginTransaction= */ false); + } catch (Throwable throwable) { + threadError.set(throwable); + } finally { + doneLatch.countDown(); + } + }); + + executorService.submit( + () -> { + try { + startLatch.await(); + streamIterator.consumer().onCompleted(); + } catch (Throwable throwable) { + threadError.set(throwable); + } finally { + doneLatch.countDown(); + } + }); + + startLatch.countDown(); + assertTrue(doneLatch.await(5, TimeUnit.SECONDS)); + if (threadError.get() != null) { + throw new AssertionError(threadError.get()); + } + + assertTrue(streamIterator.canPollWithoutErrorOrBlocking()); + assertTrue(streamIterator.isDataAvailable()); + assertFalse(streamIterator.hasNext()); + assertTrue(streamIterator.canPollWithoutErrorOrBlocking()); + assertTrue(streamIterator.isDataAvailable()); + } + } finally { + executorService.shutdownNow(); + } + } + + @Test + public void closeableIteratorDefaultMethodsReturnExpectedValues() { + CloseableIterator iterator = + new CloseableIterator() { + @Override + public void close(@Nullable String message) {} + + @Override + public boolean isWithBeginTransaction() { + return false; + } + + @Override + public boolean isLastStatement() { + return false; + } + + @Override + public boolean hasNext() { + return false; + } + + @Override + public String next() { + throw new NoSuchElementException(); + } + }; + + assertTrue(iterator.isDataAvailable()); + assertFalse(iterator.canPollWithoutErrorOrBlocking()); + } + + @Test + public void listenerRegisteredBeforeSetCallFlushesAllQueuedMessages() { + GrpcStreamIterator streamIterator = + new GrpcStreamIterator( + /* statement= */ null, + /* lastStatement= */ false, + /* prefetchChunks= */ 5, + /* cancelQueryWhenClientIsClosed= */ false); + + AsyncResultSet.StreamMessageListener listener = + mock(AsyncResultSet.StreamMessageListener.class); + streamIterator.registerListener(listener); + + PartialResultSet chunk1 = + PartialResultSet.newBuilder().setResumeToken(ByteString.copyFromUtf8("token-1")).build(); + PartialResultSet chunk2 = + PartialResultSet.newBuilder().setResumeToken(ByteString.copyFromUtf8("token-2")).build(); + + streamIterator.consumer().onPartialResultSet(chunk1); + streamIterator.consumer().onPartialResultSet(chunk2); + + verify(listener, never()).onStreamMessage(any(), anyBoolean()); + assertFalse(streamIterator.canPollWithoutErrorOrBlocking()); + assertFalse(streamIterator.isDataAvailable()); + + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); + streamIterator.setCall(call, /* withBeginTransaction= */ false); + + verify(listener, times(1)).onStreamMessage(chunk1, false); + verify(listener, times(1)).onStreamMessage(chunk2, false); + assertTrue(streamIterator.canPollWithoutErrorOrBlocking()); + assertTrue(streamIterator.isDataAvailable()); + } + + @Test + public void listenerRegisteredAfterSetCallFlushesAllQueuedMessages() { + GrpcStreamIterator streamIterator = + new GrpcStreamIterator( + /* statement= */ null, + /* lastStatement= */ false, + /* prefetchChunks= */ 5, + /* cancelQueryWhenClientIsClosed= */ false); + + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); + streamIterator.setCall(call, /* withBeginTransaction= */ false); + + PartialResultSet chunk1 = + PartialResultSet.newBuilder().setResumeToken(ByteString.copyFromUtf8("token-1")).build(); + PartialResultSet chunk2 = + PartialResultSet.newBuilder().setResumeToken(ByteString.copyFromUtf8("token-2")).build(); + + streamIterator.consumer().onPartialResultSet(chunk1); + streamIterator.consumer().onPartialResultSet(chunk2); + + assertTrue(streamIterator.canPollWithoutErrorOrBlocking()); + assertTrue(streamIterator.isDataAvailable()); + + AsyncResultSet.StreamMessageListener listener = + mock(AsyncResultSet.StreamMessageListener.class); + streamIterator.registerListener(listener); + + verify(listener, times(1)).onStreamMessage(chunk1, false); + verify(listener, times(1)).onStreamMessage(chunk2, false); + assertTrue(streamIterator.canPollWithoutErrorOrBlocking()); + assertTrue(streamIterator.isDataAvailable()); + + PartialResultSet chunk3 = + PartialResultSet.newBuilder().setResumeToken(ByteString.copyFromUtf8("token-3")).build(); + streamIterator.consumer().onPartialResultSet(chunk3); + + verify(listener, times(1)).onStreamMessage(chunk3, false); + assertEquals(chunk1, streamIterator.next()); + assertEquals(chunk2, streamIterator.next()); + assertEquals(chunk3, streamIterator.next()); + } + + @Test + public void listenerRegisteredBeforeSetCallDelaysErrorUntilSetCall() { + GrpcStreamIterator streamIterator = + new GrpcStreamIterator( + /* statement= */ null, + /* lastStatement= */ false, + /* prefetchChunks= */ 5, + /* cancelQueryWhenClientIsClosed= */ false); + + AsyncResultSet.StreamMessageListener listener = + mock(AsyncResultSet.StreamMessageListener.class); + streamIterator.registerListener(listener); + + SpannerException spannerException = + SpannerExceptionFactory.newSpannerException(ErrorCode.UNAVAILABLE, "temporary outage"); + streamIterator.consumer().onError(spannerException); + + verify(listener, never()).onStreamMessage(any(), anyBoolean()); + assertFalse(streamIterator.canPollWithoutErrorOrBlocking()); + assertFalse(streamIterator.isDataAvailable()); + + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); + streamIterator.setCall(call, /* withBeginTransaction= */ false); + + verify(listener, times(1)).onStreamMessage(GrpcStreamIterator.END_OF_STREAM, false); + assertFalse(streamIterator.canPollWithoutErrorOrBlocking()); + assertTrue(streamIterator.isDataAvailable()); + + SpannerException thrown = assertThrows(SpannerException.class, streamIterator::next); + assertEquals(ErrorCode.UNAVAILABLE, thrown.getErrorCode()); + } + + @Test + public void listenerRegisteredBeforeSetCallDelaysCompletionUntilSetCall() { + GrpcStreamIterator streamIterator = + new GrpcStreamIterator( + /* statement= */ null, + /* lastStatement= */ false, + /* prefetchChunks= */ 5, + /* cancelQueryWhenClientIsClosed= */ false); + + AsyncResultSet.StreamMessageListener listener = + mock(AsyncResultSet.StreamMessageListener.class); + streamIterator.registerListener(listener); + + streamIterator.consumer().onCompleted(); + + verify(listener, never()).onStreamMessage(any(), anyBoolean()); + assertFalse(streamIterator.canPollWithoutErrorOrBlocking()); + assertFalse(streamIterator.isDataAvailable()); + + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); + streamIterator.setCall(call, /* withBeginTransaction= */ false); + + verify(listener, times(1)).onStreamMessage(GrpcStreamIterator.END_OF_STREAM, false); + assertTrue(streamIterator.canPollWithoutErrorOrBlocking()); + assertTrue(streamIterator.isDataAvailable()); + assertFalse(streamIterator.hasNext()); + } + + @Test + public void dataAvailableAndCanPollWithoutErrorOrBlockingWithBufferedChunksBeforeError() { + GrpcStreamIterator streamIterator = + new GrpcStreamIterator( + /* statement= */ null, + /* lastStatement= */ false, + /* prefetchChunks= */ 10, + /* cancelQueryWhenClientIsClosed= */ false); + + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + when(call.getCallContext()).thenReturn(GrpcCallContext.createDefault()); + streamIterator.setCall(call, /* withBeginTransaction= */ false); + + PartialResultSet chunk = + PartialResultSet.newBuilder() + .setResumeToken(ByteString.copyFromUtf8("token-1")) + .addValues(Value.newBuilder().setStringValue("value-1")) + .build(); + + streamIterator.consumer().onPartialResultSet(chunk); + SpannerException spannerException = + SpannerExceptionFactory.newSpannerException(ErrorCode.UNAVAILABLE, "network unavailable"); + streamIterator.consumer().onError(spannerException); + + // Chunk is buffered at the head of the stream: it can be polled without error or blocking. + assertTrue(streamIterator.canPollWithoutErrorOrBlocking()); + assertTrue(streamIterator.isDataAvailable()); + + // Consuming the buffered chunk succeeds without error. + assertEquals(chunk, streamIterator.next()); + + // Now only the error/END_OF_STREAM remains in the stream: it cannot be polled without error. + assertFalse(streamIterator.canPollWithoutErrorOrBlocking()); + assertTrue(streamIterator.isDataAvailable()); + + SpannerException thrown = assertThrows(SpannerException.class, streamIterator::next); + assertEquals(ErrorCode.UNAVAILABLE, thrown.getErrorCode()); + } + + @Test + public void synchronousModeAllowsImmediatePollingBeforeSetCall() { + GrpcStreamIterator streamIterator = + new GrpcStreamIterator( + /* statement= */ null, + /* lastStatement= */ false, + /* prefetchChunks= */ 5, + /* cancelQueryWhenClientIsClosed= */ false); + + PartialResultSet chunk = + PartialResultSet.newBuilder() + .setResumeToken(ByteString.copyFromUtf8("token-1")) + .addValues(Value.newBuilder().setStringValue("value-1")) + .build(); + + // In synchronous mode (streamMessageListener == null), pre-setCall chunks are immediately + // available. + streamIterator.consumer().onPartialResultSet(chunk); + assertTrue(streamIterator.canPollWithoutErrorOrBlocking()); + assertTrue(streamIterator.isDataAvailable()); + + assertEquals(chunk, streamIterator.next()); + assertFalse(streamIterator.canPollWithoutErrorOrBlocking()); + assertFalse(streamIterator.isDataAvailable()); + } }