test(spanner): unflake ReadWriteTransactionWithInlineBeginTest

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

@@ -1179,6 +1179,12 @@ public void closeBeforeTransactionIdReturnedDoesNotRaceWithCancelledConsumer() t
     //    closed, so onError does not overwrite transactionIdFuture with a simulated ABORTED before
     //    onDone(true) sets FAILED_PRECONDITION.
     CountDownLatch consumerFinishedNextLatch = new CountDownLatch(1);
+    // Counted down when the initial prefetch request is issued on the call. This only happens after
+    // startStream() has returned and the stream has been registered with the result set iterator,
+    // which guarantees that a subsequent close() cancels the call from the closing thread rather
+    // than running before stream registration.
+    CountDownLatch streamRegisteredLatch = new CountDownLatch(1);
+    AtomicReference<Thread> consumerThread = new AtomicReference<>();
     SpannerOptions options =
         createSpannerOptionsWithStreamInterceptors(
             (invocation, delegate) -> {
@@ -1192,16 +1198,21 @@ public ApiCallContext getCallContext() {
                 @Override
                 public void cancel(String message) {
                   delegateCall.cancel(message);
-                  try {
-                    assertTrue(consumerFinishedNextLatch.await(5, TimeUnit.SECONDS));
-                  } catch (InterruptedException interruptedException) {
-                    Thread.currentThread().interrupt();
+                  // Only wait if called from an external thread (e.g. resultSet.close()),
+                  // not if cancel() is invoked by the consumer thread itself.
+                  if (Thread.currentThread() != consumerThread.get()) {
+                    try {
+                      assertTrue(consumerFinishedNextLatch.await(5, TimeUnit.SECONDS));
+                    } catch (InterruptedException interruptedException) {
+                      Thread.currentThread().interrupt();
+                    }
                   }
                 }
 
                 @Override
                 public void request(int numMessages) {
                   delegateCall.request(numMessages);
+                  streamRegisteredLatch.countDown();
                 }
               };
             },
@@ -1224,17 +1235,14 @@ public void request(int numMessages) {
                             Future<?> nextFuture =
                                 executor.submit(
                                     () -> {
+                                      consumerThread.set(Thread.currentThread());
                                       try {
                                         resultSet.next();
                                       } finally {
                                         consumerFinishedNextLatch.countDown();
                                       }
                                     });
-                            try {
-                              mockSpanner.waitForRequestsToContain(ExecuteSqlRequest.class, 5000);
-                            } catch (Exception exceptionFromWait) {
-                              throw new RuntimeException(exceptionFromWait);
-                            }
+                            assertTrue(streamRegisteredLatch.await(5, TimeUnit.SECONDS));
                             resultSet.close();
                             mockSpanner.unfreeze();
                             ExecutionException consumerExecutionException =