test(spanner): unflake AsyncResultSetImplTest

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();