fix(spanner): prevent stream leaks and race conditions on result set close (#14665)

Fixes several race conditions, potential stream leaks, and state loss
when closing streaming result sets:
1. Early Close Race Condition: If ResultSet.close() was called before
asynchronous gRPC stream initialization finished (before setCall was
invoked), the cancellation was silently dropped because the call
reference was still null. When setCall ran moments later, the stream
remained open and was never cancelled. Now, closed state is tracked so
that setCall immediately cancels the stream if closed early.
2. Concurrent Close Idempotency: close() is now synchronized and
extracts the call reference atomically, guaranteeing call.cancel() is
executed at most once, and executing the cancellation outside the lock
to prevent deadlocks.
3. Superfluous Cancellation on Completed Streams: The call handle is now
cleared as soon as a PartialResultSet with last=true is consumed,
avoiding unnecessary RPC cancellations when closing after reading the
final chunk without an extra next() call.
4. END_OF_STREAM Queue Deduplication: Atomically deduplicate
END_OF_STREAM markers to avoid queue overflow or stalls when
prefetchChunks is configured to small values (e.g. prefetchChunks = 1).
5. Inlined Begin Transaction State on Close: In GrpcResultSet.close(),
capture withBeginTransaction before closing the underlying iterator so
listener.onDone() properly unblocks any threads waiting for inlined
transaction IDs.

Author: olavloite

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

@@ -140,8 +140,9 @@ public void close() {
       }
       closed = true;
     }
-    listener.onDone(iterator.isWithBeginTransaction());
+    boolean withBeginTransaction = iterator.isWithBeginTransaction();
     iterator.close("ResultSet closed");
+    listener.onDone(withBeginTransaction);
   }
 
   @Override

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

@@ -26,6 +26,7 @@
 import com.google.common.util.concurrent.Uninterruptibles;
 import com.google.spanner.v1.PartialResultSet;
 import java.time.Duration;
+import java.util.Objects;
 import java.util.Optional;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.LinkedBlockingQueue;
@@ -38,6 +39,40 @@
 @VisibleForTesting
 class GrpcStreamIterator extends AbstractIterator<PartialResultSet>
     implements CloseableIterator<PartialResultSet> {
+
+  @VisibleForTesting
+  static final class Timeout {
+    final long duration;
+    final TimeUnit unit;
+
+    Timeout(long duration, TimeUnit unit) {
+      this.duration = duration;
+      this.unit = Preconditions.checkNotNull(unit);
+    }
+
+    @Override
+    public boolean equals(Object o) {
+      if (this == o) {
+        return true;
+      }
+      if (!(o instanceof Timeout)) {
+        return false;
+      }
+      Timeout that = (Timeout) o;
+      return this.duration == that.duration && this.unit == that.unit;
+    }
+
+    @Override
+    public int hashCode() {
+      return Objects.hash(duration, unit);
+    }
+
+    @Override
+    public String toString() {
+      return duration + " " + unit;
+    }
+  }
+
   private static final Logger logger = Logger.getLogger(GrpcStreamIterator.class.getName());
   static final PartialResultSet END_OF_STREAM = PartialResultSet.newBuilder().build();
   private final int prefetchChunks;
@@ -47,13 +82,15 @@ class GrpcStreamIterator extends AbstractIterator<PartialResultSet>
   private final BlockingQueue<PartialResultSet> stream;
   private final Statement statement;
 
-  private SpannerRpc.StreamingCall call;
+  private volatile SpannerRpc.StreamingCall call;
+  private boolean callCleared;
+  private boolean closed;
+  @Nullable private String closeMessage;
   private volatile boolean withBeginTransaction;
   private final boolean lastStatement;
-  private TimeUnit streamWaitTimeoutUnit;
-  private long streamWaitTimeoutValue;
-  private SpannerException error;
-  private boolean done;
+  @Nullable private volatile Timeout streamWaitTimeout;
+  @Nullable private volatile SpannerException error;
+  private volatile boolean done;
 
   @VisibleForTesting
   GrpcStreamIterator(
@@ -84,8 +121,28 @@ void registerListener(AsyncResultSet.StreamMessageListener streamMessageListener
   }
 
   public void setCall(SpannerRpc.StreamingCall call, boolean withBeginTransaction) {
-    this.call = call;
-    this.withBeginTransaction = withBeginTransaction;
+    Timeout timeout = extractTimeout(call);
+    boolean shouldCancel;
+    String messageToCancel;
+    synchronized (this) {
+      this.withBeginTransaction = withBeginTransaction;
+      shouldCancel = this.closed && !this.callCleared;
+      if (!this.callCleared && !this.closed) {
+        this.call = call;
+      }
+      if (timeout != null) {
+        this.streamWaitTimeout = timeout;
+      }
+      messageToCancel = this.closeMessage;
+    }
+    if (shouldCancel) {
+      call.cancel(messageToCancel);
+    }
+  }
+
+  @VisibleForTesting
+  @Nullable
+  static Timeout extractTimeout(SpannerRpc.StreamingCall call) {
     ApiCallContext callContext = call.getCallContext();
     Duration streamWaitTimeout =
         callContext == null ? null : callContext.getStreamWaitTimeoutDuration();
@@ -94,29 +151,52 @@ public void setCall(SpannerRpc.StreamingCall call, boolean withBeginTransaction)
       // value is more than 1 second, which is lower than the precision that would normally be
       // used by the stream watchdog (which uses a precision of 10 seconds by default).
       if (streamWaitTimeout.getSeconds() > 0L) {
-        streamWaitTimeoutValue = streamWaitTimeout.getSeconds();
-        streamWaitTimeoutUnit = TimeUnit.SECONDS;
+        return new Timeout(streamWaitTimeout.getSeconds(), TimeUnit.SECONDS);
       } else if (streamWaitTimeout.getNano() > 0) {
-        streamWaitTimeoutValue = streamWaitTimeout.getNano();
-        streamWaitTimeoutUnit = TimeUnit.NANOSECONDS;
+        return new Timeout(streamWaitTimeout.getNano(), TimeUnit.NANOSECONDS);
       }
       // Note that if the stream-wait-timeout is zero, we won't set a timeout at all.
       // That is consistent with ApiCallContext#withStreamWaitTimeout(Duration.ZERO).
     }
+    return null;
   }
 
   @Override
   public void close(@Nullable String message) {
-    if (call != null) {
-      call.cancel(message);
+    SpannerRpc.StreamingCall callToCancel = null;
+    synchronized (this) {
+      if (!closed) {
+        closed = true;
+        closeMessage = message;
+        if (!callCleared) {
+          callToCancel = this.call;
+          this.call = null;
+        }
+      }
+    }
+    if (callToCancel != null) {
+      callToCancel.cancel(message);
     }
   }
 
   @Override
   @InternalApi
   public void requestPrefetchChunks() {
-    Preconditions.checkState(call != null, "The StreamingCall object is not initialized");
-    call.request(prefetchChunks);
+    SpannerRpc.StreamingCall callToRequest = null;
+    synchronized (this) {
+      if (this.call != null) {
+        callToRequest = this.call;
+      } else if (!this.closed && !this.callCleared) {
+        // This should never happen during normal operation because startStream(...) always
+        // initializes the call via setCall(...) before returning the GrpcStreamIterator and
+        // calling requestPrefetchChunks(). Reaching this point means requestPrefetchChunks()
+        // was invoked before setCall(...).
+        throw new IllegalStateException("The StreamingCall object is not initialized");
+      }
+    }
+    if (callToRequest != null) {
+      callToRequest.request(prefetchChunks);
+    }
   }
 
   @Override
@@ -131,35 +211,51 @@ public boolean isLastStatement() {
 
   @Override
   protected final PartialResultSet computeNext() {
-    PartialResultSet next;
+    PartialResultSet next = takeNextFromStream();
+    if (next != END_OF_STREAM) {
+      SpannerRpc.StreamingCall currentCall = this.call;
+      if (currentCall != null) {
+        currentCall.request(1);
+      }
+      if (next.getLast()) {
+        clearCall();
+      }
+      return next;
+    }
+
+    clearCall();
+    if (error != null) {
+      throw SpannerExceptionFactory.asSpannerException(error);
+    }
+    endOfData();
+    return null;
+  }
+
+  private PartialResultSet takeNextFromStream() {
     try {
-      if (streamWaitTimeoutUnit != null) {
-        next = stream.poll(streamWaitTimeoutValue, streamWaitTimeoutUnit);
+      Timeout timeout = this.streamWaitTimeout;
+      if (timeout != null) {
+        PartialResultSet next = stream.poll(timeout.duration, timeout.unit);
         if (next == null) {
           throw SpannerExceptionFactory.newSpannerException(
               ErrorCode.DEADLINE_EXCEEDED, "stream wait timeout");
         }
+        return next;
       } else {
-        next = stream.take();
+        return stream.take();
       }
     } catch (InterruptedException e) {
       // Treat interrupt as a request to cancel the read.
       throw SpannerExceptionFactory.propagateInterrupt(e);
     }
-    if (next != END_OF_STREAM) {
-      call.request(1);
-      return next;
-    }
+  }
 
+  private void clearCall() {
     // All done - close() no longer needs to cancel the call.
-    call = null;
-
-    if (error != null) {
-      throw SpannerExceptionFactory.asSpannerException(error);
+    synchronized (this) {
+      callCleared = true;
+      call = null;
     }
-
-    endOfData();
-    return null;
   }
 
   private void addToStream(PartialResultSet results) {
@@ -179,16 +275,13 @@ private class ConsumerImpl implements SpannerRpc.ResultStreamConsumer {
     public void onPartialResultSet(PartialResultSet results) {
       addToStream(results);
       if (results.getLast()) {
-        done = true;
-        addToStream(END_OF_STREAM);
+        addEndOfStream();
       }
     }
 
     @Override
     public void onCompleted() {
-      if (!done) {
-        addToStream(END_OF_STREAM);
-      }
+      addEndOfStream();
     }
 
     @Override
@@ -203,7 +296,20 @@ public void onError(SpannerException exception) {
         }
       }
       error = exception;
-      addToStream(END_OF_STREAM);
+      addEndOfStream();
+    }
+
+    private void addEndOfStream() {
+      boolean shouldAddEndOfStream = false;
+      synchronized (GrpcStreamIterator.this) {
+        if (!done) {
+          done = true;
+          shouldAddEndOfStream = true;
+        }
+      }
+      if (shouldAddEndOfStream) {
+        addToStream(END_OF_STREAM);
+      }
     }
 
     @Override

java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/CloseSpannerWithOpenResultSetTest.java

@@ -46,6 +46,14 @@
 public class CloseSpannerWithOpenResultSetTest extends AbstractMockServerTest {
 
   Spanner createSpanner() {
+    return spannerOptionsBuilder().build().getService();
+  }
+
+  Spanner createSpanner(int prefetchChunks) {
+    return spannerOptionsBuilder().setPrefetchChunks(prefetchChunks).build().getService();
+  }
+
+  private SpannerOptions.Builder spannerOptionsBuilder() {
     return SpannerOptions.newBuilder()
         .setProjectId("p")
         .setHost(String.format("http://localhost:%d", getPort()))
@@ -54,9 +62,7 @@ Spanner createSpanner() {
         .setSessionPoolOption(
             SessionPoolOptions.newBuilder()
                 .setWaitForMinSessionsDuration(Duration.ofSeconds(5L))
-                .build())
-        .build()
-        .getService();
+                .build());
   }
 
   @BeforeClass
@@ -122,4 +128,65 @@ public void testStreamsAreCleanedUp() throws Exception {
           0, ((GapicSpannerRpc) ((SpannerImpl) spanner).getRpc()).getNumActiveResponseObservers());
     }
   }
+
+  @Test(timeout = 30000)
+  public void testQueryWithPrefetchChunksOne() {
+    try (Spanner spanner = createSpanner(1)) {
+      DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
+      for (int i = 0; i < 10; i++) {
+        try (ResultSet resultSet = client.singleUse().executeQuery(SELECT_RANDOM_STATEMENT)) {
+          int count = 0;
+          while (resultSet.next()) {
+            assertNotNull(resultSet.getCurrentRowAsStruct());
+            count++;
+          }
+          assertTrue(count > 0);
+        }
+      }
+      assertEquals(
+          0, ((GapicSpannerRpc) ((SpannerImpl) spanner).getRpc()).getNumActiveResponseObservers());
+    }
+  }
+
+  @Test(timeout = 30000)
+  public void testQueryWithPrefetchChunksOneAndError() {
+    String errorSql = "select * from will_fail";
+    Statement errorStatement = Statement.of(errorSql);
+    mockSpanner.putStatementResult(
+        StatementResult.exception(
+            errorStatement,
+            Status.INVALID_ARGUMENT.withDescription("Invalid syntax").asRuntimeException()));
+    try (Spanner spanner = createSpanner(1)) {
+      DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
+      try (ResultSet resultSet = client.singleUse().executeQuery(errorStatement)) {
+        SpannerException exception = assertThrows(SpannerException.class, resultSet::next);
+        assertEquals(ErrorCode.INVALID_ARGUMENT, exception.getErrorCode());
+      }
+      assertEquals(
+          0, ((GapicSpannerRpc) ((SpannerImpl) spanner).getRpc()).getNumActiveResponseObservers());
+    }
+  }
+
+  @Test(timeout = 30000)
+  public void testCloseResultSetCleansUpResponseObservers() {
+    try (Spanner spanner = createSpanner()) {
+      DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
+      // 1. Partial consumption then close:
+      try (ResultSet resultSet = client.singleUse().executeQuery(SELECT_RANDOM_STATEMENT)) {
+        assertTrue(resultSet.next());
+        assertNotNull(resultSet.getCurrentRowAsStruct());
+      }
+      assertEquals(
+          0, ((GapicSpannerRpc) ((SpannerImpl) spanner).getRpc()).getNumActiveResponseObservers());
+
+      // 2. Full consumption then close:
+      try (ResultSet resultSet = client.singleUse().executeQuery(SELECT_RANDOM_STATEMENT)) {
+        while (resultSet.next()) {
+          assertNotNull(resultSet.getCurrentRowAsStruct());
+        }
+      }
+      assertEquals(
+          0, ((GapicSpannerRpc) ((SpannerImpl) spanner).getRpc()).getNumActiveResponseObservers());
+    }
+  }
 }