diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/GrpcResultSet.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/GrpcResultSet.java index 2f3daef47419..5aa14c550ae6 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/GrpcResultSet.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/GrpcResultSet.java @@ -146,8 +146,15 @@ public void close() { closed = true; } boolean withBeginTransaction = iterator.isWithBeginTransaction(); - iterator.close("ResultSet closed"); - listener.onDone(withBeginTransaction); + // Notify the listener before cancelling the stream. Cancelling the stream causes any thread + // that is concurrently blocked in next() to fail with CANCELLED, which would otherwise race + // with this notification and could fail a pending inline-begin transaction with a different + // error than the one that is caused by closing this ResultSet. + try { + listener.onDone(withBeginTransaction); + } finally { + iterator.close("ResultSet closed"); + } } @Override diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/GrpcResultSetTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/GrpcResultSetTest.java index 614ad37fa041..ec087f58bfd2 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/GrpcResultSetTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/GrpcResultSetTest.java @@ -57,7 +57,14 @@ import java.util.List; import java.util.Map; import java.util.UUID; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; import javax.annotation.Nullable; import org.junit.Before; import org.junit.Test; @@ -1306,6 +1313,94 @@ public PartialResultSet next() { assertTrue(listenerReceivedWithBeginTransaction.get()); } + @Test + public void testCloseNotifiesListenerBeforeCancelErrorFromConcurrentReader() throws Exception { + // The first transaction-id outcome that the listener receives wins, just like the + // transactionIdFuture of an inline-begin read/write transaction. + AtomicReference firstOutcome = new AtomicReference<>(); + CountDownLatch readerReportedError = new CountDownLatch(1); + AbstractResultSet.Listener listener = + new NoOpListener() { + @Override + public SpannerException onError( + SpannerException e, boolean withBeginTransaction, boolean lastStatement) { + if (withBeginTransaction) { + firstOutcome.compareAndSet(null, "onError:" + e.getErrorCode()); + } + readerReportedError.countDown(); + return e; + } + + @Override + public void onDone(boolean withBeginTransaction) { + if (withBeginTransaction) { + firstOutcome.compareAndSet(null, "onDone"); + } + } + }; + GrpcStreamIterator iterator = + new GrpcStreamIterator( + /* lastStatement= */ false, 10, /* cancelQueryWhenClientIsClosed= */ false); + SpannerRpc.ResultStreamConsumer streamConsumer = iterator.consumer(); + iterator.setCall( + new SpannerRpc.StreamingCall() { + @Override + public ApiCallContext getCallContext() { + return GrpcCallContext.createDefault(); + } + + @Override + public void cancel(@Nullable String message) { + // Emulate gRPC delivering the CANCELLED status to the stream consumer, and the + // concurrent reader thread processing that error before close() returns. + streamConsumer.onError( + SpannerExceptionFactory.newSpannerException(ErrorCode.CANCELLED, message)); + try { + readerReportedError.await(10, TimeUnit.SECONDS); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + } + + @Override + public void request(int numMessages) {} + }, + /* withBeginTransaction= */ true); + GrpcResultSet resultSet = new GrpcResultSet(iterator, listener); + + AtomicReference readerThread = new AtomicReference<>(); + ExecutorService executor = + Executors.newSingleThreadExecutor( + runnable -> { + Thread thread = new Thread(runnable); + readerThread.set(thread); + return thread; + }); + try { + Future nextFuture = executor.submit(resultSet::next); + // Wait until the reader is blocked waiting for the first chunk of the stream. + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(10); + while (readerThread.get() == null || readerThread.get().getState() != Thread.State.WAITING) { + if (nextFuture.isDone()) { + nextFuture.get(); + throw new AssertionError("Reader thread terminated unexpectedly without blocking"); + } + assertTrue("reader did not block on the stream", System.nanoTime() < deadline); + Thread.sleep(10); + } + resultSet.close(); + + ExecutionException executionException = + assertThrows(ExecutionException.class, () -> nextFuture.get(10, TimeUnit.SECONDS)); + assertTrue(executionException.getCause() instanceof SpannerException); + assertEquals( + ErrorCode.CANCELLED, ((SpannerException) executionException.getCause()).getErrorCode()); + assertEquals("onDone", firstOutcome.get()); + } finally { + executor.shutdownNow(); + } + } + @Test public void testIsDataAvailable() { assertFalse(resultSet.isDataAvailable());