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