sakthivelmanii · Oct 19, 2026 · Oct 19, 2026
@@ -557,6 +557,13 @@ CloseableIterator<PartialResultSet> startStream(
           boolean prepareIteratorForRetryOnDifferentGrpcChannel() {
             return AbstractReadContext.this.prepareRetryOnDifferentGrpcChannel();
           }
+
+          @Override
+          void onStreamClosedWhileStarting(boolean withBeginTransaction) {
+            // The ResultSet was closed while this stream was being started. GrpcResultSet#close()
+            // could therefore not see whether the stream included a BeginTransaction option.
+            AbstractReadContext.this.onDone(withBeginTransaction);
+          }
         };
     return new GrpcResultSet(
         stream, this, options.hasDecodeMode() ? options.decodeMode() : defaultDecodeMode);

@@ -142,7 +142,11 @@ public void close() {
     }
     boolean withBeginTransaction = iterator.isWithBeginTransaction();
     iterator.close("ResultSet closed");
-    listener.onDone(withBeginTransaction);
+    // Also check after closing the iterator, as a stream could have been started concurrently on
+    // another thread after the first check. ResumableStreamIterator remembers whether the stream
+    // that it closed included a BeginTransaction option, and reports a stream that it sees was
+    // closed while it was being started itself.
+    listener.onDone(withBeginTransaction || iterator.isWithBeginTransaction());
   }
 
   @Override

@@ -44,6 +44,7 @@
 import java.util.logging.Level;
 import java.util.logging.Logger;
 import javax.annotation.Nullable;
+import javax.annotation.concurrent.GuardedBy;
 
 /**
  * Wraps an iterator over partial result sets, supporting resuming RPCs on error. This class keeps
@@ -84,6 +85,23 @@ abstract class ResumableStreamIterator extends AbstractIterator<PartialResultSet
   private final TraceWrapper tracer;
   private CloseableIterator<PartialResultSet> stream;
 
+  /**
+   * Set when {@link #close(String)} is called. {@link #close(String)} can be called from a
+   * different thread than the thread that is consuming (and starting) the stream. This flag is used
+   * to ensure that a stream that is started concurrently with a call to {@link #close(String)} is
+   * also closed.
+   */
+  @GuardedBy("this")
+  private boolean closed;
+
+  @GuardedBy("this")
+  @Nullable
+  private String closeMessage;
+
+  /** Whether the stream that was closed by {@link #close(String)} included a BeginTransaction. */
+  @GuardedBy("this")
+  private boolean closedStreamWithBeginTransaction;
+
   /**
    * The number of consecutive failed attempts without any progress on the stream. Reset to zero
    * every time the stream returns a new resume token.
@@ -266,18 +284,45 @@ boolean prepareIteratorForRetryOnDifferentGrpcChannel() {
     return false;
   }
 
+  /**
+   * Called when a stream was started concurrently with a call to {@link #close(String)}. The stream
+   * has already been closed when this method is called. Subclasses can override this method to
+   * handle the fact that the stream was closed before it returned any results, for example to
+   * notify statements that are waiting for a transaction id that will never be returned.
+   */
+  void onStreamClosedWhileStarting(boolean withBeginTransaction) {}
+
   @Override
   public void close(@Nullable String message) {
-    if (stream != null) {
-      stream.close(message);
-      span.end();
+    CloseableIterator<PartialResultSet> streamToClose;
+    synchronized (this) {
+      if (closed) {
+        return;
+      }
+      closed = true;
+      closeMessage = message;
+      streamToClose = stream;
       stream = null;
+      closedStreamWithBeginTransaction =
+          streamToClose != null && streamToClose.isWithBeginTransaction();
+    }
+    if (streamToClose != null) {
+      streamToClose.close(message);
+      span.end();
     }
   }
 
   @Override
-  public boolean isWithBeginTransaction() {
-    return stream != null && stream.isWithBeginTransaction();
+  public synchronized boolean isWithBeginTransaction() {
+    if (stream != null) {
+      return stream.isWithBeginTransaction();
+    }
+    return closedStreamWithBeginTransaction;
+  }
+
+  private synchronized SpannerException newClosedException() {
+    return SpannerExceptionFactory.newSpannerException(
+        ErrorCode.CANCELLED, closeMessage == null ? "ResultSet closed" : closeMessage);
   }
 
   @Override
@@ -426,15 +471,38 @@ private void handleRetryableException(Context context, SpannerException spannerE
 
   private void startGrpcStreaming() {
     if (stream == null) {
+      synchronized (this) {
+        if (closed) {
+          throw newClosedException();
+        }
+      }
       span.addAnnotation(
           "Starting/Resuming stream",
           "ResumeToken",
           resumeToken == null ? "null" : resumeToken.toStringUtf8());
       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();
+        CloseableIterator<PartialResultSet> newStream =
+            checkNotNull(startStream(resumeToken, streamMessageListener, requestId));
+        boolean closedWhileStarting;
+        String message;
+        synchronized (this) {
+          closedWhileStarting = closed;
+          message = closeMessage;
+          if (!closedWhileStarting) {
+            stream = newStream;
+          }
+        }
+        if (closedWhileStarting) {
+          // close() was called while the stream was being started. Close (cancel) the stream that
+          // was just started, as close() could not see it.
+          newStream.close(message);
+          span.end();
+          onStreamClosedWhileStarting(newStream.isWithBeginTransaction());
+          throw newClosedException();
+        }
+        newStream.requestPrefetchChunks();
       }
     }
   }

@@ -19,6 +19,7 @@
 import static com.google.common.truth.Truth.assertThat;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
@@ -48,7 +49,13 @@
 import java.util.Iterator;
 import java.util.LinkedList;
 import java.util.List;
+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.AtomicReference;
 import java.util.function.LongSupplier;
 import javax.annotation.Nullable;
 import org.junit.Before;
@@ -261,6 +268,85 @@ public void closedOCSpan() {
     verify(mockSpan).end(OpenCensusSpan.END_SPAN_OPTIONS);
   }
 
+  @Test
+  public void closeWhileStreamIsBeingStartedClosesTheStartedStream() throws Exception {
+    ResultSetStream s1 = Mockito.mock(ResultSetStream.class);
+    CountDownLatch startStreamEntered = new CountDownLatch(1);
+    CountDownLatch closeCalled = new CountDownLatch(1);
+    Mockito.when(starter.startStream(null, null))
+        .thenAnswer(
+            invocation -> {
+              startStreamEntered.countDown();
+              // Block the stream start until close() has been called on another thread. This
+              // simulates a ResultSet that is closed while the RPC has been sent, but before the
+              // stream has been handed back to the ResumableStreamIterator.
+              assertTrue(closeCalled.await(10, TimeUnit.SECONDS));
+              return new ResultSetIterator(s1) {
+                @Override
+                public boolean isWithBeginTransaction() {
+                  return true;
+                }
+              };
+            });
+    Mockito.when(s1.next()).thenReturn(resultSet(null, "a")).thenReturn(null);
+    AtomicReference<Boolean> reportedWithBeginTransaction = new AtomicReference<>();
+    resumableStreamIterator =
+        new ResumableStreamIterator(
+            Integer.MAX_VALUE,
+            "",
+            new OpenTelemetrySpan(mock(io.opentelemetry.api.trace.Span.class)),
+            new TraceWrapper(Tracing.getTracer(), OpenTelemetry.noop().getTracer(""), false),
+            DefaultErrorHandler.INSTANCE,
+            RetrySettings.newBuilder().build(),
+            SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetryableCodes(),
+            NoopRequestIdCreator.INSTANCE) {
+          @Override
+          AbstractResultSet.CloseableIterator<PartialResultSet> startStream(
+              @Nullable ByteString resumeToken,
+              AsyncResultSet.StreamMessageListener streamMessageListener,
+              XGoogSpannerRequestId requestId) {
+            return starter.startStream(resumeToken, null);
+          }
+
+          @Override
+          void onStreamClosedWhileStarting(boolean withBeginTransaction) {
+            reportedWithBeginTransaction.set(withBeginTransaction);
+          }
+        };
+    ExecutorService executor = Executors.newSingleThreadExecutor();
+    try {
+      Future<Boolean> hasNext = executor.submit(resumableStreamIterator::hasNext);
+      assertTrue(startStreamEntered.await(10, TimeUnit.SECONDS));
+      resumableStreamIterator.close("ResultSet closed");
+      closeCalled.countDown();
+
+      ExecutionException executionException =
+          assertThrows(ExecutionException.class, () -> hasNext.get(10, TimeUnit.SECONDS));
+      assertThat(executionException.getCause()).isInstanceOf(SpannerException.class);
+      SpannerException spannerException = (SpannerException) executionException.getCause();
+      assertEquals(ErrorCode.CANCELLED, spannerException.getErrorCode());
+      // The stream that was started concurrently with close() must be closed (cancelled).
+      verify(s1).close();
+      // The iterator must report that the stream that it closed included a BeginTransaction, as
+      // the ResultSet could not see this when it was closed.
+      assertEquals(Boolean.TRUE, reportedWithBeginTransaction.get());
+      // No new stream should be started for a closed iterator.
+      verify(starter, Mockito.times(1)).startStream(null, null);
+    } finally {
+      executor.shutdownNow();
+    }
+  }
+
+  @Test
+  public void closeBeforeStreamIsStartedDoesNotStartStream() {
+    resumableStreamIterator.close("ResultSet closed");
+
+    SpannerException spannerException =
+        assertThrows(SpannerException.class, () -> resumableStreamIterator.hasNext());
+    assertEquals(ErrorCode.CANCELLED, spannerException.getErrorCode());
+    verify(starter, Mockito.never()).startStream(Mockito.any(), Mockito.any());
+  }
+
   @Test
   public void restart() {
     ResultSetStream s1 = Mockito.mock(ResultSetStream.class);