From 7b8bf00a1c4deb2641df969b10cdc16f3acfa1ee Mon Sep 17 00:00:00 2001 From: Sakthivel Subramanian Date: Wed, 14 Oct 2026 16:00:35 +0000 Subject: [PATCH] test(spanner): unflake AsyncResultSetImplTest 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. --- .../cloud/spanner/AsyncResultSetImpl.java | 13 ++++- .../cloud/spanner/AsyncResultSetImplTest.java | 49 +++++++++++++++++-- 2 files changed, 57 insertions(+), 5 deletions(-) diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncResultSetImpl.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncResultSetImpl.java index 3dd5724532b3..f3a2e7d71a48 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncResultSetImpl.java +++ b/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) { diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncResultSetImplTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncResultSetImplTest.java index 23180356cbfa..2e7f17485263 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncResultSetImplTest.java +++ b/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 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 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 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();