fix(spanner): cancel queries promptly on close and synchronize error handling

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

@@ -189,6 +189,9 @@ static void assertTimestampAvailable(boolean available) {
   protected static final String NO_TRANSACTION_RETURNED_MSG =
       "The statement did not return a transaction even though one was requested";
 
+  protected static final String RESULT_SET_CLOSED_BEFORE_ID_RETURNED_MESSAGE =
+      "ResultSet was closed before a transaction id was returned";
+
   AbstractReadContext(Builder<?, ?> builder) {
     this.session = builder.session;
     this.cancelQueryWhenClientIsClosed = builder.cancelQueryWhenClientIsClosed;
@@ -526,6 +529,18 @@ CloseableIterator<PartialResultSet> startStream(
             if (streamListener != null) {
               stream.registerListener(streamListener);
             }
+            // setStream() returns false if this ResumableStreamIterator was already closed
+            // concurrently (e.g., via resultSet.close() or transaction.close()). In that case,
+            // setStream() has already invoked stream.close() out-of-lock to prevent leaks. We skip
+            // issuing the RPC call (avoiding unnecessary network I/O) and fail any inlined-begin
+            // future via onClosedBeforeStreamStarted(). We still return the (now closed) stream
+            // instance because startStream() requires a non-null return value to satisfy
+            // checkNotNull() in ResumableStreamIterator.startGrpcStreaming(), which then cleanly
+            // completes stream teardown and wakes any threads waiting on streamLock.
+            if (!setStream(stream)) {
+              onClosedBeforeStreamStarted(withBeginTransaction);
+              return stream;
+            }
             SpannerRpc.StreamingCall call;
             try {
               call =
@@ -690,6 +705,8 @@ public void onDone(boolean withBeginTransaction) {
 
   void onStartFailed(boolean withBeginTransaction, boolean lastStatement, Throwable throwable) {}
 
+  void onClosedBeforeStreamStarted(boolean withBeginTransaction) {}
+
   /**
    * For transactions other than read-write, the MultiplexedSessionPrecommitToken will not be
    * present in the RPC response. In such cases, this method will be a no-op.

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

@@ -39,8 +39,10 @@ class GrpcResultSet extends AbstractResultSet<List<Object>>
   private ResultSetMetadata metadata;
   private GrpcStruct currRow;
   private List<Object> rowData;
+  private volatile SpannerException error;
   private ResultSetStats statistics;
   private volatile boolean closed;
+  private boolean yieldingError;
 
   GrpcResultSet(CloseableIterator<PartialResultSet> iterator, Listener listener) {
     this(iterator, listener, DecodeMode.DIRECT);
@@ -74,6 +76,9 @@ protected GrpcStruct currRow() {
 
   @Override
   public boolean next() throws SpannerException {
+    if (this.error != null) {
+      throw this.error;
+    }
     try {
       if (currRow == null) {
         metadata = iterator.getMetadata();
@@ -102,9 +107,9 @@ public boolean next() throws SpannerException {
         close();
       }
       return hasNext;
-    } catch (Throwable t) {
+    } catch (Throwable throwable) {
       throw yieldError(
-          asSpannerException(t),
+          asSpannerException(throwable),
           iterator.isWithBeginTransaction() && currRow == null,
           iterator.isLastStatement());
     }
@@ -139,15 +144,25 @@ public boolean isDataAvailable() {
 
   @Override
   public void close() {
+    boolean shouldNotifyDone;
     synchronized (this) {
       if (closed) {
         return;
       }
       closed = true;
+      // If a thread is currently in yieldError() dispatching listener.onError(), defer calling
+      // listener.onDone() to yieldError() finally block so that onError() always completes before
+      // onDone() is invoked.
+      shouldNotifyDone = !yieldingError;
     }
     boolean withBeginTransaction = iterator.isWithBeginTransaction();
-    iterator.close("ResultSet closed");
-    listener.onDone(withBeginTransaction);
+    try {
+      iterator.close("ResultSet closed");
+    } finally {
+      if (shouldNotifyDone) {
+        listener.onDone(withBeginTransaction || iterator.isWithBeginTransaction());
+      }
+    }
   }
 
   @Override
@@ -157,9 +172,38 @@ public Type getType() {
   }
 
   private SpannerException yieldError(
-      SpannerException e, boolean beginTransaction, boolean lastStatement) {
-    SpannerException toThrow = listener.onError(e, beginTransaction, lastStatement);
-    close();
-    throw toThrow;
+      SpannerException exception, boolean beginTransaction, boolean lastStatement) {
+    synchronized (this) {
+      if (this.error == null) {
+        this.error = exception;
+      }
+      if (this.closed || this.yieldingError) {
+        return this.error;
+      }
+      this.yieldingError = true;
+    }
+    try {
+      this.error = listener.onError(exception, beginTransaction, lastStatement);
+      throw this.error;
+    } finally {
+      boolean alreadyClosed;
+      synchronized (this) {
+        this.yieldingError = false;
+        alreadyClosed = this.closed;
+      }
+      // If close() was called concurrently while we were dispatching listener.onError(), close()
+      // intentionally skipped calling listener.onDone() to avoid racing with or completing before
+      // onError(). We ensure the underlying iterator is closed before invoking listener.onDone().
+      // Otherwise, if close() was not called yet, proceed with normal closure.
+      if (alreadyClosed) {
+        try {
+          iterator.close("ResultSet closed");
+        } finally {
+          listener.onDone(beginTransaction || iterator.isWithBeginTransaction());
+        }
+      } else {
+        close();
+      }
+    }
   }
 }

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

@@ -119,6 +119,20 @@ public String toString() {
         /* withBeginTransaction= */ false);
   }
 
+  @VisibleForTesting
+  GrpcStreamIterator(
+      boolean lastStatement,
+      int prefetchChunks,
+      boolean cancelQueryWhenClientIsClosed,
+      boolean withBeginTransaction) {
+    this(
+        /* statement= */ null,
+        lastStatement,
+        prefetchChunks,
+        cancelQueryWhenClientIsClosed,
+        withBeginTransaction);
+  }
+
   GrpcStreamIterator(
       Statement statement,
       boolean lastStatement,
@@ -158,9 +172,7 @@ void registerListener(AsyncResultSet.StreamMessageListener streamMessageListener
 
   @VisibleForTesting
   void setCall(SpannerRpc.StreamingCall call, boolean withBeginTransaction) {
-    synchronized (this) {
-      this.withBeginTransaction = withBeginTransaction;
-    }
+    this.withBeginTransaction = withBeginTransaction;
     setCall(call);
   }
 

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

@@ -474,21 +474,32 @@ public SpannerException onError(
   @Override
   public void onDone(boolean withBeginTransaction) {
     if (withBeginTransaction) {
-      failTransactionIdFuture(
-          SpannerExceptionFactory.newSpannerException(
-              ErrorCode.FAILED_PRECONDITION,
-              "ResultSet was closed before a read-only transaction id was returned"));
+      failTransactionIdFutureIfClosedBeforeIdReturned();
     }
     super.onDone(withBeginTransaction);
   }
 
   @Override
   void onStartFailed(boolean withBeginTransaction, boolean lastStatement, Throwable throwable) {
     if (withBeginTransaction) {
-      failTransactionIdFuture(throwable);
+      failTransactionIdFuture(SpannerExceptionFactory.asSpannerException(throwable));
     }
   }
 
+  @Override
+  void onClosedBeforeStreamStarted(boolean withBeginTransaction) {
+    if (withBeginTransaction) {
+      failTransactionIdFutureIfClosedBeforeIdReturned();
+    }
+  }
+
+  private void failTransactionIdFutureIfClosedBeforeIdReturned() {
+    failTransactionIdFuture(
+        SpannerExceptionFactory.newSpannerException(
+            ErrorCode.FAILED_PRECONDITION,
+            "ResultSet was closed before a read-only transaction id was returned"));
+  }
+
   private void failTransactionIdFuture(Throwable throwable) {
     txnLock.lock();
     try {

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

@@ -344,10 +344,9 @@ public boolean isWithBeginTransaction() {
       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().
+    // A subclass that registers its stream early calls setStream(consumer) during stream
+    // initialization. Querying currentInFlightStream live ensures callers observe
+    // withBeginTransaction as soon as the stream is created during an in-flight startStream().
     return currentInFlightStream != null
         ? currentInFlightStream.isWithBeginTransaction()
         : fallbackWithBeginTransaction;
@@ -634,10 +633,8 @@ private void resetStreamOnFailure(Throwable throwable) {
     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().
+        // clearing inFlightStream and calling notifyAll(). 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;

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

@@ -904,11 +904,15 @@ public void onDone(boolean withBeginTransaction) {
         // Context was done (closed) before a transaction id was returned.
         this.transactionIdFuture.setException(
             SpannerExceptionFactory.newSpannerException(
-                ErrorCode.FAILED_PRECONDITION,
-                "ResultSet was closed before a transaction id was returned"));
+                ErrorCode.FAILED_PRECONDITION, RESULT_SET_CLOSED_BEFORE_ID_RETURNED_MESSAGE));
       }
     }
 
+    @Override
+    void onClosedBeforeStreamStarted(boolean withBeginTransaction) {
+      onDone(withBeginTransaction);
+    }
+
     @Override
     void onStartFailed(boolean withBeginTransaction, boolean lastStatement, Throwable throwable) {
       if (withBeginTransaction) {