test(spanner): unflake AsyncResultSetImplTest (#14656)

Production busy-spin hang in AsyncResultSetImpl when a ReadyCallback
throws an exception, causing ProduceRowsRunnable to loop infinitely when
buffer is full. Triggered by test race on paused flag in
testCallbackIsNotCalledWhilePaused. Fixed by including
cursorReturnedDoneOrException in producer stop condition and adding
lock-based test synchronization.

Author: sakthivelmanii

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

@@ -373,7 +373,7 @@ public void run() {
         while (!stop && hasNext) {
           try {
             synchronized (monitor) {
-              stop = state.shouldStop;
+              stop = shouldStopProducing();
             }
             if (!stop) {
               while (buffer.remainingCapacity() == 0 && !stop) {
@@ -390,7 +390,7 @@ public void run() {
                         MAX_WAIT_FOR_BUFFER_CONSUMPTION));
                 bufferConsumptionLatch.await();
                 synchronized (monitor) {
-                  stop = state.shouldStop;
+                  stop = shouldStopProducing();
                 }
               }
             }
@@ -450,6 +450,15 @@ public void run() {
       }
     }
 
+    /**
+     * Returns true if no more rows should be produced. This is the case if the result set has been
+     * cancelled or is done, or if the callback will not be invoked anymore (e.g. because it threw
+     * an exception). Must be called while holding the monitor.
+     */
+    private boolean shouldStopProducing() {
+      return state.shouldStop || cursorReturnedDoneOrException;
+    }
+
     private void waitIfPaused() throws InterruptedException {
       CountDownLatch pause;
       synchronized (monitor) {

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

@@ -334,6 +334,11 @@ public Boolean answer(InvocationOnMock invocation) throws Throwable {
     final AtomicInteger callbackCounter = new AtomicInteger();
     final BlockingDeque<Object> queue = new LinkedBlockingDeque<>(1);
     final AtomicBoolean paused = new AtomicBoolean();
+    // Clearing the paused flag and calling resume() must be atomic with respect to the callback
+    // setting the paused flag. Otherwise, the test thread could clear the flag before the callback
+    // sets it, and then call resume() after the callback returned PAUSE. The next callback would
+    // then see a stale paused flag, even though resume() was correctly called.
+    final Object pauseLock = new Object();
     try (AsyncResultSetImpl rs =
         new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) {
       ApiFuture<Void> callbackResult =
@@ -345,7 +350,9 @@ public Boolean answer(InvocationOnMock invocation) throws Throwable {
                 try {
                   switch (resultSet.tryNext()) {
                     case OK:
-                      paused.set(true);
+                      synchronized (pauseLock) {
+                        paused.set(true);
+                      }
                       queue.put(new Object());
                       return CallbackResponse.PAUSE;
                     case DONE:
@@ -366,8 +373,10 @@ public Boolean answer(InvocationOnMock invocation) throws Throwable {
           rowCounter++;
         }
         Thread.yield();
-        paused.set(false);
-        rs.resume();
+        synchronized (pauseLock) {
+          paused.set(false);
+          rs.resume();
+        }
       }
       // Empty the queue to ensure we count all elements.
       while (queue.poll() != null) {
@@ -494,6 +503,40 @@ public void callbackReturnsError() throws InterruptedException {
     }
   }
 
+  @Test
+  public void callbackReturnsErrorWhileBufferIsFull_shouldStopProducer() throws Exception {
+    ExecutorService executor = Executors.newSingleThreadExecutor();
+    // Return more rows than fit in the buffer, so the producer will wait for the buffer to be
+    // consumed by the callback. The callback fails instead of consuming any rows.
+    final int simulatedRows = AsyncResultSetImpl.DEFAULT_BUFFER_SIZE * 10;
+    final AtomicInteger producedRows = new AtomicInteger();
+    ResultSet delegate = mock(ResultSet.class);
+    when(delegate.next()).thenAnswer(ignored -> producedRows.incrementAndGet() <= simulatedRows);
+    when(delegate.getCurrentRowAsStruct()).thenReturn(mock(Struct.class));
+    try (AsyncResultSetImpl rs =
+        new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) {
+      ApiFuture<Void> result =
+          rs.setCallback(
+              executor,
+              ignored -> {
+                throw new RuntimeException("async test");
+              });
+      // The result must complete with the error from the callback. Before the fix, the producer
+      // would keep waiting for the (full) buffer to be consumed, and the result never completed.
+      ExecutionException e =
+          assertThrows(ExecutionException.class, () -> result.get(30L, TimeUnit.SECONDS));
+      assertThat(e.getCause()).isInstanceOf(SpannerException.class);
+      SpannerException se = (SpannerException) e.getCause();
+      assertThat(se.getErrorCode()).isEqualTo(ErrorCode.UNKNOWN);
+      assertThat(se.getMessage()).contains("async test");
+      // The producer should stop and release the underlying result set.
+      Mockito.verify(delegate).close();
+      assertThat(producedRows.get()).isLessThan(simulatedRows);
+    } finally {
+      executor.shutdown();
+    }
+  }
+
   @Test
   public void callbackReturnsDoneBeforeEnd_shouldStopIteration() throws Exception {
     Executor executor = Executors.newSingleThreadExecutor();