refactor(spanner): add non-blocking polling to GrpcStreamIterator

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) {

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 {
   <T> List<T> toList(Function<StructReader, T> 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.
+   *
+   * <p>This listener is purely a notification mechanism to signal that messages or stream events
+   * are buffered in the underlying stream iterator. It is <b>not</b> 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.
+     *
+     * <p>Implementations of this method <b>must be thread-safe, idempotent, and order-agnostic</b>.
+     * Because notifications may be dispatched concurrently by gRPC network threads and iterator
+     * initialization/flushing threads:
+     *
+     * <ul>
+     *   <li>Notifications may arrive concurrently from multiple threads.
+     *   <li>The same {@code partialResultSet} may be delivered more than once.
+     *   <li>Notifications for different chunks may arrive <b>out of order</b>.
+     * </ul>
+     *
+     * <p>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);
   }
 }

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<PartialResultSet> 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,13 +146,29 @@ 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;
       }
       messageToCancel = this.closeMessage;
     }
     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.
+   *
+   * <p>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));
   }
 }