diff --git a/google-java-format-1.17.0-all-deps.jar b/google-java-format-1.17.0-all-deps.jar new file mode 100644 index 000000000000..15a42a996046 Binary files /dev/null and b/google-java-format-1.17.0-all-deps.jar differ diff --git a/google-java-format-1.25.2-all-deps.jar b/google-java-format-1.25.2-all-deps.jar new file mode 100644 index 000000000000..026d07607b6e Binary files /dev/null and b/google-java-format-1.25.2-all-deps.jar differ diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java index a769758861f4..14fb34ca803e 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java @@ -557,6 +557,13 @@ CloseableIterator startStream( boolean prepareIteratorForRetryOnDifferentGrpcChannel() { return AbstractReadContext.this.prepareRetryOnDifferentGrpcChannel(); } + + @Override + void onStreamClosedWhileStarting(boolean withBeginTransaction) { + // The ResultSet was closed while this stream was being started. GrpcResultSet#close() + // could therefore not see whether the stream included a BeginTransaction option. + AbstractReadContext.this.onDone(withBeginTransaction); + } }; return new GrpcResultSet( stream, this, options.hasDecodeMode() ? options.decodeMode() : defaultDecodeMode); 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 299e68fe86a2..10b35c0b35a2 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 @@ -142,7 +142,11 @@ public void close() { } boolean withBeginTransaction = iterator.isWithBeginTransaction(); iterator.close("ResultSet closed"); - listener.onDone(withBeginTransaction); + // Also check after closing the iterator, as a stream could have been started concurrently on + // another thread after the first check. ResumableStreamIterator remembers whether the stream + // that it closed included a BeginTransaction option, and reports a stream that it sees was + // closed while it was being started itself. + listener.onDone(withBeginTransaction || iterator.isWithBeginTransaction()); } @Override diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ResumableStreamIterator.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ResumableStreamIterator.java index d6477fe65104..5e4d2052e4bf 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ResumableStreamIterator.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ResumableStreamIterator.java @@ -44,6 +44,7 @@ import java.util.logging.Level; import java.util.logging.Logger; import javax.annotation.Nullable; +import javax.annotation.concurrent.GuardedBy; /** * Wraps an iterator over partial result sets, supporting resuming RPCs on error. This class keeps @@ -84,6 +85,23 @@ abstract class ResumableStreamIterator extends AbstractIterator stream; + /** + * Set when {@link #close(String)} is called. {@link #close(String)} can be called from a + * different thread than the thread that is consuming (and starting) the stream. This flag is used + * to ensure that a stream that is started concurrently with a call to {@link #close(String)} is + * also closed. + */ + @GuardedBy("this") + private boolean closed; + + @GuardedBy("this") + @Nullable + private String closeMessage; + + /** Whether the stream that was closed by {@link #close(String)} included a BeginTransaction. */ + @GuardedBy("this") + private boolean closedStreamWithBeginTransaction; + /** * The number of consecutive failed attempts without any progress on the stream. Reset to zero * every time the stream returns a new resume token. @@ -266,18 +284,45 @@ boolean prepareIteratorForRetryOnDifferentGrpcChannel() { return false; } + /** + * Called when a stream was started concurrently with a call to {@link #close(String)}. The stream + * has already been closed when this method is called. Subclasses can override this method to + * handle the fact that the stream was closed before it returned any results, for example to + * notify statements that are waiting for a transaction id that will never be returned. + */ + void onStreamClosedWhileStarting(boolean withBeginTransaction) {} + @Override public void close(@Nullable String message) { - if (stream != null) { - stream.close(message); - span.end(); + CloseableIterator streamToClose; + synchronized (this) { + if (closed) { + return; + } + closed = true; + closeMessage = message; + streamToClose = stream; stream = null; + closedStreamWithBeginTransaction = + streamToClose != null && streamToClose.isWithBeginTransaction(); + } + if (streamToClose != null) { + streamToClose.close(message); + span.end(); } } @Override - public boolean isWithBeginTransaction() { - return stream != null && stream.isWithBeginTransaction(); + public synchronized boolean isWithBeginTransaction() { + if (stream != null) { + return stream.isWithBeginTransaction(); + } + return closedStreamWithBeginTransaction; + } + + private synchronized SpannerException newClosedException() { + return SpannerExceptionFactory.newSpannerException( + ErrorCode.CANCELLED, closeMessage == null ? "ResultSet closed" : closeMessage); } @Override @@ -426,6 +471,11 @@ private void handleRetryableException(Context context, SpannerException spannerE private void startGrpcStreaming() { if (stream == null) { + synchronized (this) { + if (closed) { + throw newClosedException(); + } + } span.addAnnotation( "Starting/Resuming stream", "ResumeToken", @@ -433,8 +483,26 @@ private void startGrpcStreaming() { try (IScope scope = tracer.withSpan(span)) { // When start a new stream set the Span as current to make the gRPC Span a child of // this Span. - stream = checkNotNull(startStream(resumeToken, streamMessageListener, requestId)); - stream.requestPrefetchChunks(); + CloseableIterator newStream = + checkNotNull(startStream(resumeToken, streamMessageListener, requestId)); + boolean closedWhileStarting; + String message; + synchronized (this) { + closedWhileStarting = closed; + message = closeMessage; + if (!closedWhileStarting) { + stream = newStream; + } + } + if (closedWhileStarting) { + // close() was called while the stream was being started. Close (cancel) the stream that + // was just started, as close() could not see it. + newStream.close(message); + span.end(); + onStreamClosedWhileStarting(newStream.isWithBeginTransaction()); + throw newClosedException(); + } + newStream.requestPrefetchChunks(); } } } diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ResumableStreamIteratorTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ResumableStreamIteratorTest.java index a0f18431315f..ccf84bae81bf 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ResumableStreamIteratorTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ResumableStreamIteratorTest.java @@ -19,6 +19,7 @@ import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -48,7 +49,13 @@ import java.util.Iterator; import java.util.LinkedList; import java.util.List; +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.AtomicReference; import java.util.function.LongSupplier; import javax.annotation.Nullable; import org.junit.Before; @@ -261,6 +268,85 @@ public void closedOCSpan() { verify(mockSpan).end(OpenCensusSpan.END_SPAN_OPTIONS); } + @Test + public void closeWhileStreamIsBeingStartedClosesTheStartedStream() throws Exception { + ResultSetStream s1 = Mockito.mock(ResultSetStream.class); + CountDownLatch startStreamEntered = new CountDownLatch(1); + CountDownLatch closeCalled = new CountDownLatch(1); + Mockito.when(starter.startStream(null, null)) + .thenAnswer( + invocation -> { + startStreamEntered.countDown(); + // Block the stream start until close() has been called on another thread. This + // simulates a ResultSet that is closed while the RPC has been sent, but before the + // stream has been handed back to the ResumableStreamIterator. + assertTrue(closeCalled.await(10, TimeUnit.SECONDS)); + return new ResultSetIterator(s1) { + @Override + public boolean isWithBeginTransaction() { + return true; + } + }; + }); + Mockito.when(s1.next()).thenReturn(resultSet(null, "a")).thenReturn(null); + AtomicReference reportedWithBeginTransaction = new AtomicReference<>(); + resumableStreamIterator = + new ResumableStreamIterator( + Integer.MAX_VALUE, + "", + new OpenTelemetrySpan(mock(io.opentelemetry.api.trace.Span.class)), + new TraceWrapper(Tracing.getTracer(), OpenTelemetry.noop().getTracer(""), false), + DefaultErrorHandler.INSTANCE, + RetrySettings.newBuilder().build(), + SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetryableCodes(), + NoopRequestIdCreator.INSTANCE) { + @Override + AbstractResultSet.CloseableIterator startStream( + @Nullable ByteString resumeToken, + AsyncResultSet.StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + return starter.startStream(resumeToken, null); + } + + @Override + void onStreamClosedWhileStarting(boolean withBeginTransaction) { + reportedWithBeginTransaction.set(withBeginTransaction); + } + }; + ExecutorService executor = Executors.newSingleThreadExecutor(); + try { + Future hasNext = executor.submit(resumableStreamIterator::hasNext); + assertTrue(startStreamEntered.await(10, TimeUnit.SECONDS)); + resumableStreamIterator.close("ResultSet closed"); + closeCalled.countDown(); + + ExecutionException executionException = + assertThrows(ExecutionException.class, () -> hasNext.get(10, TimeUnit.SECONDS)); + assertThat(executionException.getCause()).isInstanceOf(SpannerException.class); + SpannerException spannerException = (SpannerException) executionException.getCause(); + assertEquals(ErrorCode.CANCELLED, spannerException.getErrorCode()); + // The stream that was started concurrently with close() must be closed (cancelled). + verify(s1).close(); + // The iterator must report that the stream that it closed included a BeginTransaction, as + // the ResultSet could not see this when it was closed. + assertEquals(Boolean.TRUE, reportedWithBeginTransaction.get()); + // No new stream should be started for a closed iterator. + verify(starter, Mockito.times(1)).startStream(null, null); + } finally { + executor.shutdownNow(); + } + } + + @Test + public void closeBeforeStreamIsStartedDoesNotStartStream() { + resumableStreamIterator.close("ResultSet closed"); + + SpannerException spannerException = + assertThrows(SpannerException.class, () -> resumableStreamIterator.hasNext()); + assertEquals(ErrorCode.CANCELLED, spannerException.getErrorCode()); + verify(starter, Mockito.never()).startStream(Mockito.any(), Mockito.any()); + } + @Test public void restart() { ResultSetStream s1 = Mockito.mock(ResultSetStream.class);