test(spanner): unflake multiUseReadOnlyTransactionCanUseInlineBeginForReadAsync

[sakthivelmanii]

In SessionImplTest#multiUseReadOnlyTransactionCanUseInlineBeginForReadAsync, mock rpc.read invoked stream callbacks synchronously before returning StreamingCall. This caused a background thread to process rows and invoke startStream while this.stream was still unassigned, triggering a duplicate request (expected:<1> but was:<2>).

Fix this deterministically by delaying consumer.onPartialResultSet and consumer.onCompleted until StreamingCall.request(int) is called. In gRPC streaming, this.stream is guaranteed to be assigned before request() is executed, eliminating the race condition without relying on background threads or executors.

[gemini-code-assist]

Code Review

This pull request updates a unit test in SessionImplTest.java to execute the mocked RPC result stream consumer asynchronously using CompletableFuture.runAsync. The review feedback recommends adding exception handling to the asynchronous block to prevent unhandled exceptions from being silently swallowed, which would make debugging test failures difficult.

[gemini-code-assist]

Using CompletableFuture.runAsync without an explicit executor defaults to the system-wide ForkJoinPool.commonPool(). In resource-constrained CI environments (e.g., single-core containers), this can lead to execution delays or starvation, potentially causing test flakiness.

Additionally, any unhandled exceptions thrown inside the async block (such as assertion failures or unexpected null pointers) will be silently swallowed by the CompletableFuture unless explicitly handled, making debugging difficult.

Consider adding exception handling to ensure any failures during the async execution are logged or propagated.

              CompletableFuture.runAsync(
                      () -> {
                        c.onPartialResultSet(resultSet);
                        c.onCompleted();
                      })
                  .exceptionally(
                      t -> {
                        t.printStackTrace();
                        return null;
                      });

[sakthivelmanii]

Let the test fail if there's an exception

[olavloite]

I don't think this really fixes the possibility that this test could flake. It does reduce the probability, but it does not rule it out. I think that this would be a better way to implement this test:

@Test
public void multiUseReadOnlyTransactionCanUseInlineBeginForReadAsync() throws Exception {
  PartialResultSet resultSet = inlineBeginResultSet("async-inline-tx");
  final ArgumentCaptor<SpannerRpc.ResultStreamConsumer> consumerCaptor =
      ArgumentCaptor.forClass(SpannerRpc.ResultStreamConsumer.class);
  final ArgumentCaptor<ReadRequest> requestCaptor = ArgumentCaptor.forClass(ReadRequest.class);

  Mockito.when(rpc.read(requestCaptor.capture(), consumerCaptor.capture(), anyMap(), any(), eq(false)))
      .then(
          invocation -> {
            SpannerRpc.ResultStreamConsumer consumer = invocation.getArgument(1);
            return new SpannerRpc.StreamingCall() {
              private final AtomicBoolean requested = new AtomicBoolean();

              @Override
              public void request(int numMessages) {
                if (requested.compareAndSet(false, true)) {
                  consumer.onPartialResultSet(resultSet);
                  consumer.onCompleted();
                }
              }

              @Override
              public void cancel(@Nullable String message) {}
            };
          });

  try (ReadOnlyTransaction transaction = inlineReadOnlyTransaction()) {
    try (AsyncResultSet resultSetAsync =
        transaction.readAsync("Dummy", KeySet.all(), Collections.singletonList("C"))) {
      resultSetAsync
          .setCallback(
              Runnable::run,
              asyncResultSet -> {
                while (asyncResultSet.tryNext() == AsyncResultSet.CursorState.OK) {}
                return AsyncResultSet.CallbackResponse.CONTINUE;
              })
          .get(5, TimeUnit.SECONDS);
    }
    assertEquals(
        Timestamp.fromProto(Timestamps.parse("2015-10-01T10:54:20.021Z")),
        transaction.getReadTimestamp());
  }

  Mockito.verify(rpc, Mockito.never()).beginTransaction(Mockito.any(), anyMap(), eq(false));
  assertEquals(1, requestCaptor.getAllValues().size());
  assertTrue(requestCaptor.getValue().getTransaction().hasBegin());
}

[sakthivelmanii]

Done. Updated to trigger consumer callbacks inside StreamingCall.request(int).