perf(spanner): make internal AsyncResultSet row production non-blocking

java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java

@@ -62,7 +62,6 @@
 import java.util.concurrent.ThreadLocalRandom;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
-import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.locks.ReentrantLock;
@@ -572,38 +571,11 @@ private ListenableAsyncResultSet createAsyncResultSet(
       } finally {
         txnLock.unlock();
       }
-      // Make sure that we decrement the counter exactly once, either
-      // when the query is actually executed, or when the result set is closed,
-      // or if something goes wrong when creating the result set.
-      final AtomicBoolean decremented = new AtomicBoolean(false);
       try {
         return new AsyncResultSetImpl(
-            executorProvider,
-            () -> {
-              try {
-                return resultSetSupplier.get();
-              } finally {
-                if (decremented.compareAndSet(false, true)) {
-                  decrementPendingStartsAndSignal();
-                }
-              }
-            },
-            bufferRows) {
-          @Override
-          public void close() {
-            try {
-              super.close();
-            } finally {
-              if (!isUsed() && decremented.compareAndSet(false, true)) {
-                decrementPendingStartsAndSignal();
-              }
-            }
-          }
-        };
+            executorProvider, resultSetSupplier, bufferRows, this::decrementPendingStartsAndSignal);
       } catch (Throwable throwable) {
-        if (decremented.compareAndSet(false, true)) {
-          decrementPendingStartsAndSignal();
-        }
+        decrementPendingStartsAndSignal();
         throw throwable;
       }
     }
@@ -1295,15 +1267,6 @@ CloseableIterator<PartialResultSet> startStream(
               @Nullable ByteString resumeToken,
               AsyncResultSet.StreamMessageListener streamListener,
               XGoogSpannerRequestId requestId) {
-            GrpcStreamIterator stream =
-                new GrpcStreamIterator(
-                    statement,
-                    request.getLastStatement(),
-                    prefetchChunks,
-                    cancelQueryWhenClientIsClosed);
-            if (streamListener != null) {
-              stream.registerListener(streamListener);
-            }
             if (partitionToken != null) {
               request.setPartitionToken(partitionToken);
             }
@@ -1318,6 +1281,24 @@ CloseableIterator<PartialResultSet> startStream(
               request.setTransaction(selector);
             }
             boolean withBeginTransaction = request.getTransaction().hasBegin();
+            GrpcStreamIterator stream =
+                new GrpcStreamIterator(
+                    statement,
+                    request.getLastStatement(),
+                    prefetchChunks,
+                    cancelQueryWhenClientIsClosed,
+                    withBeginTransaction);
+            if (streamListener != null) {
+              stream.registerListener(streamListener);
+            }
+            if (!setStream(stream)) {
+              onStartFailed(
+                  withBeginTransaction,
+                  SpannerExceptionFactory.newSpannerException(
+                      ErrorCode.FAILED_PRECONDITION,
+                      "ResultSet was closed before a transaction id was returned"));
+              return stream;
+            }
             SpannerRpc.StreamingCall call;
             try {
               call =
@@ -1533,12 +1514,6 @@ CloseableIterator<PartialResultSet> startStream(
               @Nullable ByteString resumeToken,
               AsyncResultSet.StreamMessageListener streamListener,
               XGoogSpannerRequestId requestId) {
-            GrpcStreamIterator stream =
-                new GrpcStreamIterator(
-                    lastStatement, prefetchChunks, cancelQueryWhenClientIsClosed);
-            if (streamListener != null) {
-              stream.registerListener(streamListener);
-            }
             TransactionSelector selector = null;
             if (resumeToken != null) {
               builder.setResumeToken(resumeToken);
@@ -1551,6 +1526,24 @@ CloseableIterator<PartialResultSet> startStream(
             }
             builder.setRequestOptions(buildRequestOptions(readOptions));
             boolean withBeginTransaction = builder.getTransaction().hasBegin();
+            GrpcStreamIterator stream =
+                new GrpcStreamIterator(
+                    null,
+                    lastStatement,
+                    prefetchChunks,
+                    cancelQueryWhenClientIsClosed,
+                    withBeginTransaction);
+            if (streamListener != null) {
+              stream.registerListener(streamListener);
+            }
+            if (!setStream(stream)) {
+              onStartFailed(
+                  withBeginTransaction,
+                  SpannerExceptionFactory.newSpannerException(
+                      ErrorCode.FAILED_PRECONDITION,
+                      "ResultSet was closed before a transaction id was returned"));
+              return stream;
+            }
             SpannerRpc.StreamingCall call;
             try {
               call =

java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractResultSet.java

@@ -166,6 +166,22 @@ 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 {@link #hasNext()} and {@link #next()} can be called immediately to yield a
+     * result or EOF without blocking on network I/O and without throwing a stream error.
+     */
+    default boolean canPollWithoutBlocking() {
+      return false;
+    }
   }
 
   static double valueProtoToFloat64(com.google.protobuf.Value proto) {