test(spanner): unflake InlineBeginTransactionTest

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

@@ -146,8 +146,15 @@ public void close() {
       closed = true;
     }
     boolean withBeginTransaction = iterator.isWithBeginTransaction();
-    iterator.close("ResultSet closed");
-    listener.onDone(withBeginTransaction);
+    // Notify the listener before cancelling the stream. Cancelling the stream causes any thread
+    // that is concurrently blocked in next() to fail with CANCELLED, which would otherwise race
+    // with this notification and could fail a pending inline-begin transaction with a different
+    // error than the one that is caused by closing this ResultSet.
+    try {
+      listener.onDone(withBeginTransaction);
+    } finally {
+      iterator.close("ResultSet closed");
+    }
   }
 
   @Override

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

@@ -57,7 +57,14 @@
 import java.util.List;
 import java.util.Map;
 import java.util.UUID;
+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.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
 import javax.annotation.Nullable;
 import org.junit.Before;
 import org.junit.Test;
@@ -1306,6 +1313,94 @@ public PartialResultSet next() {
     assertTrue(listenerReceivedWithBeginTransaction.get());
   }
 
+  @Test
+  public void testCloseNotifiesListenerBeforeCancelErrorFromConcurrentReader() throws Exception {
+    // The first transaction-id outcome that the listener receives wins, just like the
+    // transactionIdFuture of an inline-begin read/write transaction.
+    AtomicReference<String> firstOutcome = new AtomicReference<>();
+    CountDownLatch readerReportedError = new CountDownLatch(1);
+    AbstractResultSet.Listener listener =
+        new NoOpListener() {
+          @Override
+          public SpannerException onError(
+              SpannerException e, boolean withBeginTransaction, boolean lastStatement) {
+            if (withBeginTransaction) {
+              firstOutcome.compareAndSet(null, "onError:" + e.getErrorCode());
+            }
+            readerReportedError.countDown();
+            return e;
+          }
+
+          @Override
+          public void onDone(boolean withBeginTransaction) {
+            if (withBeginTransaction) {
+              firstOutcome.compareAndSet(null, "onDone");
+            }
+          }
+        };
+    GrpcStreamIterator iterator =
+        new GrpcStreamIterator(
+            /* lastStatement= */ false, 10, /* cancelQueryWhenClientIsClosed= */ false);
+    SpannerRpc.ResultStreamConsumer streamConsumer = iterator.consumer();
+    iterator.setCall(
+        new SpannerRpc.StreamingCall() {
+          @Override
+          public ApiCallContext getCallContext() {
+            return GrpcCallContext.createDefault();
+          }
+
+          @Override
+          public void cancel(@Nullable String message) {
+            // Emulate gRPC delivering the CANCELLED status to the stream consumer, and the
+            // concurrent reader thread processing that error before close() returns.
+            streamConsumer.onError(
+                SpannerExceptionFactory.newSpannerException(ErrorCode.CANCELLED, message));
+            try {
+              readerReportedError.await(10, TimeUnit.SECONDS);
+            } catch (InterruptedException e) {
+              Thread.currentThread().interrupt();
+            }
+          }
+
+          @Override
+          public void request(int numMessages) {}
+        },
+        /* withBeginTransaction= */ true);
+    GrpcResultSet resultSet = new GrpcResultSet(iterator, listener);
+
+    AtomicReference<Thread> readerThread = new AtomicReference<>();
+    ExecutorService executor =
+        Executors.newSingleThreadExecutor(
+            runnable -> {
+              Thread thread = new Thread(runnable);
+              readerThread.set(thread);
+              return thread;
+            });
+    try {
+      Future<Boolean> nextFuture = executor.submit(resultSet::next);
+      // Wait until the reader is blocked waiting for the first chunk of the stream.
+      long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(10);
+      while (readerThread.get() == null || readerThread.get().getState() != Thread.State.WAITING) {
+        if (nextFuture.isDone()) {
+          nextFuture.get();
+          throw new AssertionError("Reader thread terminated unexpectedly without blocking");
+        }
+        assertTrue("reader did not block on the stream", System.nanoTime() < deadline);
+        Thread.sleep(10);
+      }
+      resultSet.close();
+
+      ExecutionException executionException =
+          assertThrows(ExecutionException.class, () -> nextFuture.get(10, TimeUnit.SECONDS));
+      assertTrue(executionException.getCause() instanceof SpannerException);
+      assertEquals(
+          ErrorCode.CANCELLED, ((SpannerException) executionException.getCause()).getErrorCode());
+      assertEquals("onDone", firstOutcome.get());
+    } finally {
+      executor.shutdownNow();
+    }
+  }
+
   @Test
   public void testIsDataAvailable() {
     assertFalse(resultSet.isDataAvailable());