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 c473d0719f49..e66855049c70 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 @@ -62,7 +62,6 @@ import java.util.concurrent.ThreadLocalRandom; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; -import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicLong; import java.util.concurrent.locks.ReentrantLock; @@ -572,38 +571,11 @@ private ListenableAsyncResultSet createAsyncResultSet( } finally { txnLock.unlock(); } - // Make sure that we decrement the counter exactly once, either - // when the query is actually executed, or when the result set is closed, - // or if something goes wrong when creating the result set. - final AtomicBoolean decremented = new AtomicBoolean(false); try { return new AsyncResultSetImpl( - executorProvider, - () -> { - try { - return resultSetSupplier.get(); - } finally { - if (decremented.compareAndSet(false, true)) { - decrementPendingStartsAndSignal(); - } - } - }, - bufferRows) { - @Override - public void close() { - try { - super.close(); - } finally { - if (!isUsed() && decremented.compareAndSet(false, true)) { - decrementPendingStartsAndSignal(); - } - } - } - }; + executorProvider, resultSetSupplier, bufferRows, this::decrementPendingStartsAndSignal); } catch (Throwable throwable) { - if (decremented.compareAndSet(false, true)) { - decrementPendingStartsAndSignal(); - } + decrementPendingStartsAndSignal(); throw throwable; } } @@ -1295,15 +1267,6 @@ CloseableIterator startStream( @Nullable ByteString resumeToken, AsyncResultSet.StreamMessageListener streamListener, XGoogSpannerRequestId requestId) { - GrpcStreamIterator stream = - new GrpcStreamIterator( - statement, - request.getLastStatement(), - prefetchChunks, - cancelQueryWhenClientIsClosed); - if (streamListener != null) { - stream.registerListener(streamListener); - } if (partitionToken != null) { request.setPartitionToken(partitionToken); } @@ -1318,6 +1281,24 @@ CloseableIterator startStream( request.setTransaction(selector); } boolean withBeginTransaction = request.getTransaction().hasBegin(); + GrpcStreamIterator stream = + new GrpcStreamIterator( + statement, + request.getLastStatement(), + prefetchChunks, + cancelQueryWhenClientIsClosed, + withBeginTransaction); + if (streamListener != null) { + stream.registerListener(streamListener); + } + if (!setStream(stream)) { + onStartFailed( + withBeginTransaction, + SpannerExceptionFactory.newSpannerException( + ErrorCode.FAILED_PRECONDITION, + "ResultSet was closed before a transaction id was returned")); + return stream; + } SpannerRpc.StreamingCall call; try { call = @@ -1533,12 +1514,6 @@ CloseableIterator startStream( @Nullable ByteString resumeToken, AsyncResultSet.StreamMessageListener streamListener, XGoogSpannerRequestId requestId) { - GrpcStreamIterator stream = - new GrpcStreamIterator( - lastStatement, prefetchChunks, cancelQueryWhenClientIsClosed); - if (streamListener != null) { - stream.registerListener(streamListener); - } TransactionSelector selector = null; if (resumeToken != null) { builder.setResumeToken(resumeToken); @@ -1551,6 +1526,24 @@ CloseableIterator startStream( } builder.setRequestOptions(buildRequestOptions(readOptions)); boolean withBeginTransaction = builder.getTransaction().hasBegin(); + GrpcStreamIterator stream = + new GrpcStreamIterator( + null, + lastStatement, + prefetchChunks, + cancelQueryWhenClientIsClosed, + withBeginTransaction); + if (streamListener != null) { + stream.registerListener(streamListener); + } + if (!setStream(stream)) { + onStartFailed( + withBeginTransaction, + SpannerExceptionFactory.newSpannerException( + ErrorCode.FAILED_PRECONDITION, + "ResultSet was closed before a transaction id was returned")); + return stream; + } SpannerRpc.StreamingCall call; try { call = diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractResultSet.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractResultSet.java index 0717cae74f25..1e2b1897419c 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractResultSet.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractResultSet.java @@ -166,6 +166,22 @@ default boolean initiateStreaming(AsyncResultSet.StreamMessageListener streamMes /** it requests the initial prefetch chunks from gRPC stream */ default void requestPrefetchChunks() {} + + /** + * Returns true if data (a chunk, row, EOF, or error) is available to read immediately without + * blocking the calling thread on network I/O. + */ + default boolean isDataAvailable() { + return true; + } + + /** + * Returns true if {@link #hasNext()} and {@link #next()} can be called immediately to yield a + * result or EOF without blocking on network I/O and without throwing a stream error. + */ + default boolean canPollWithoutBlocking() { + return false; + } } static double valueProtoToFloat64(com.google.protobuf.Value proto) { 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 f3a2e7d71a48..a3bd8656f4df 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 @@ -20,7 +20,7 @@ import com.google.api.core.ApiFutures; import com.google.api.core.SettableApiFuture; import com.google.api.gax.core.ExecutorProvider; -import com.google.cloud.spanner.AbstractReadContext.ListenableAsyncResultSet; +import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Function; import com.google.common.base.Preconditions; import com.google.common.base.Supplier; @@ -35,23 +35,24 @@ import java.util.LinkedList; import java.util.List; import java.util.concurrent.BlockingDeque; -import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; import java.util.concurrent.Executor; import java.util.concurrent.Future; import java.util.concurrent.LinkedBlockingDeque; import java.util.logging.Level; import java.util.logging.Logger; +import org.jspecify.annotations.NullMarked; +import org.jspecify.annotations.Nullable; /** Default implementation for {@link AsyncResultSet}. */ +@NullMarked class AsyncResultSetImpl extends ForwardingStructReader - implements ListenableAsyncResultSet, AsyncResultSet.StreamMessageListener { + implements AbstractReadContext.ListenableAsyncResultSet, AsyncResultSet.StreamMessageListener { private static final Logger log = Logger.getLogger(AsyncResultSetImpl.class.getName()); /** State of an {@link AsyncResultSetImpl}. */ - private enum State { + enum State { INITIALIZED, - STREAMING_INITIALIZED, /** SYNC indicates that the {@link ResultSet} is used in sync pattern. */ SYNC, CONSUMING, @@ -73,7 +74,6 @@ private enum State { } static final int DEFAULT_BUFFER_SIZE = 10; - private static final int MAX_WAIT_FOR_BUFFER_CONSUMPTION = 10; private static final SpannerException CANCELLED_EXCEPTION = SpannerExceptionFactory.newSpannerException( ErrorCode.CANCELLED, "This AsyncResultSet has been cancelled"); @@ -90,7 +90,7 @@ private enum State { private final ListeningScheduledExecutorService service; private final BlockingDeque buffer; - private Struct currentRow; + @Nullable private Struct currentRow; /** Supplies the underlying synchronous {@link ResultSet} that will be producing the rows. */ private final Supplier delegateResultSet; @@ -99,15 +99,15 @@ private enum State { * Any exception that occurs while executing the query and iterating over the result set will be * stored in this variable and propagated to the user through {@link #tryNext()}. */ - private volatile SpannerException executionException; + @Nullable private volatile SpannerException executionException; /** * Executor for callbacks. Regardless of the type of executor that is provided, the {@link * AsyncResultSetImpl} will ensure that at most 1 callback call will be active at any one time. */ - private Executor executor; + @Nullable private Executor executor; - private ReadyCallback callback; + @Nullable private ReadyCallback callback; /** * Listeners that will be called when the {@link AsyncResultSetImpl} has finished fetching all @@ -117,15 +117,18 @@ private enum State { private volatile State state = State.INITIALIZED; - /** This variable indicates that produce rows thread is initiated */ - private volatile boolean produceRowsInitiated; + /** Indicates whether a task is currently scheduled or executing ProduceRowsRunnable. */ + private boolean producerRunning; + + /** Indicates whether the result future and cleanup have been completed. */ + private boolean completed; /** * This variable indicates whether all the results from the underlying result set have been read. */ private volatile boolean finished; - private volatile SettableApiFuture result; + @Nullable private volatile SettableApiFuture result; /** * This variable indicates whether {@link #tryNext()} has returned {@link CursorState#DONE} or a @@ -133,52 +136,73 @@ private enum State { */ private volatile boolean cursorReturnedDoneOrException; - /** - * This variable is used to pause the producer when the {@link AsyncResultSet} is paused. The - * production of rows that are put into the buffer is only paused once the buffer is full. - */ - private volatile CountDownLatch pausedLatch = new CountDownLatch(1); - - /** - * This variable is used to pause the producer when the buffer is full and the consumer needs some - * time to catch up. - */ - private volatile CountDownLatch bufferConsumptionLatch = new CountDownLatch(0); + /** Indicates whether a task is currently scheduled or executing CallbackRunnable. */ + private boolean callbackRunning; - /** - * This variable is used to pause the producer when all rows have been put into the buffer, but - * the consumer (the callback) has not yet received and processed all rows. - */ - private volatile CountDownLatch consumingLatch = new CountDownLatch(0); + private final Runnable onStartOrClose; + private boolean startOrCloseNotified; + private boolean delegateInitialized; + private boolean delegateClosed; AsyncResultSetImpl(ExecutorProvider executorProvider, ResultSet delegate, int bufferSize) { - this(executorProvider, Suppliers.ofInstance(Preconditions.checkNotNull(delegate)), bufferSize); + this( + executorProvider, + Suppliers.ofInstance(Preconditions.checkNotNull(delegate)), + bufferSize, + () -> {}, + /* delegateInitialized= */ true); } AsyncResultSetImpl( ExecutorProvider executorProvider, Supplier delegate, int bufferSize) { + this(executorProvider, delegate, bufferSize, () -> {}); + } + + AsyncResultSetImpl( + ExecutorProvider executorProvider, + Supplier delegate, + int bufferSize, + Runnable onStartOrClose) { this( executorProvider, Suppliers.memoize(Preconditions.checkNotNull(delegate)), bufferSize, - true); + Preconditions.checkNotNull(onStartOrClose), + /* delegateInitialized= */ false); } private AsyncResultSetImpl( ExecutorProvider executorProvider, Supplier memoizedDelegate, int bufferSize, - boolean dummy) { + Runnable onStartOrClose, + boolean delegateInitialized) { super(memoizedDelegate); this.executorProvider = Preconditions.checkNotNull(executorProvider); this.delegateResultSet = memoizedDelegate; + this.onStartOrClose = onStartOrClose; + this.delegateInitialized = delegateInitialized; this.service = MoreExecutors.listeningDecorator(executorProvider.getExecutor()); this.buffer = new LinkedBlockingDeque<>(bufferSize); } - boolean isUsed() { + private ResultSet getDelegate() { + synchronized (monitor) { + delegateInitialized = true; + } + return delegateResultSet.get(); + } + + private void notifyStartOrClose() { + boolean shouldNotify = false; synchronized (monitor) { - return state != State.INITIALIZED; + if (!startOrCloseNotified) { + startOrCloseNotified = true; + shouldNotify = true; + } + } + if (shouldNotify) { + onStartOrClose.run(); } } @@ -191,15 +215,27 @@ boolean isUsed() { */ @Override public void close() { + boolean shouldCloseDelegate = false; + boolean shouldShutdownService = false; synchronized (monitor) { if (this.closed) { return; } if (state == State.INITIALIZED || state == State.SYNC) { - delegateResultSet.get().close(); + shouldCloseDelegate = true; + if (executorProvider.shouldAutoClose()) { + shouldShutdownService = true; + } } this.closed = true; } + if (shouldCloseDelegate) { + closeDelegateResultSet(); + notifyStartOrClose(); + } + if (shouldShutdownService) { + service.shutdown(); + } } /** @@ -234,8 +270,6 @@ public CursorState tryNext() throws SpannerException { cursorReturnedDoneOrException = true; throw executionException; } - Preconditions.checkState( - this.callback != null, "tryNext may only be called after a callback has been set."); Preconditions.checkState( this.state == State.CONSUMING, "tryNext may only be called from a DataReady callback. Current state: " @@ -246,22 +280,30 @@ public CursorState tryNext() throws SpannerException { return CursorState.DONE; } } - if (!buffer.isEmpty()) { + Struct nextRow = buffer.poll(); + if (nextRow != null) { // Set the next row from the buffer as the current row of the StructReader. - replaceDelegate(currentRow = buffer.pop()); - synchronized (monitor) { - bufferConsumptionLatch.countDown(); - } + replaceDelegate(currentRow = nextRow); + scheduleProducerIfNecessary(); return CursorState.OK; } return CursorState.NOT_READY; } private void closeDelegateResultSet() { - try { - delegateResultSet.get().close(); - } catch (Throwable t) { - log.log(Level.FINE, "Ignoring error from closing delegate result set", t); + boolean shouldClose = false; + synchronized (monitor) { + if (!delegateClosed) { + delegateClosed = true; + shouldClose = delegateInitialized; + } + } + if (shouldClose) { + try { + delegateResultSet.get().close(); + } catch (Throwable t) { + log.log(Level.FINE, "Ignoring error from closing delegate result set", t); + } } } @@ -274,9 +316,10 @@ private class CallbackRunnable implements Runnable { public void run() { try { while (true) { + ReadyCallback callback; synchronized (monitor) { if (cursorReturnedDoneOrException) { - break; + return; } if (state == State.CANCELLED) { // The callback should always get at least one chance to catch the CANCELLED @@ -285,27 +328,34 @@ public void run() { // also stop, even though the callback has not seen the CANCELLED state. cursorReturnedDoneOrException = true; } + callback = Preconditions.checkNotNull(AsyncResultSetImpl.this.callback); } CallbackResponse response; try { - response = callback.cursorReady(AsyncResultSetImpl.this); - } catch (Throwable e) { + response = + Preconditions.checkNotNull( + callback.cursorReady(AsyncResultSetImpl.this), + "CallbackResponse may not be null"); + } catch (Throwable throwable) { synchronized (monitor) { - if (cursorReturnedDoneOrException - && state == State.CANCELLED - && e instanceof SpannerException - && ((SpannerException) e).getErrorCode() == ErrorCode.CANCELLED) { + cursorReturnedDoneOrException = true; + if (state == State.CANCELLED + && throwable instanceof SpannerException + && ((SpannerException) throwable).getErrorCode() == ErrorCode.CANCELLED) { // The callback did not catch the cancelled exception (which it should have), but // we'll keep the cancelled state. return; } - executionException = SpannerExceptionFactory.asSpannerException(e); - cursorReturnedDoneOrException = true; + executionException = SpannerExceptionFactory.asSpannerException(throwable); } + closeDelegateResultSet(); return; } synchronized (monitor) { if (state == State.CANCELLED) { + if (response == CallbackResponse.DONE) { + cursorReturnedDoneOrException = true; + } if (cursorReturnedDoneOrException) { return; } @@ -317,41 +367,36 @@ public void run() { return; case PAUSE: state = State.PAUSED; - // Make sure no-one else is waiting on the current pause latch and create a new - // one. - pausedLatch.countDown(); - pausedLatch = new CountDownLatch(1); return; - case CONTINUE: + default: if (buffer.isEmpty()) { - // Call the callback once more if the entire result set has been processed but - // the callback has not yet received a CursorState.DONE or a CANCELLED error. - if (finished && !cursorReturnedDoneOrException) { + // Call the callback once more if the entire result set has been processed or an + // exception was encountered, but the callback has not yet received a + // CursorState.DONE or error. + if ((finished || executionException != null) + && !cursorReturnedDoneOrException) { break; } state = State.RUNNING; return; } break; - default: - throw new IllegalStateException("Unknown response: " + response); } } } } } finally { synchronized (monitor) { - // Count down all latches that the producer might be waiting on. - consumingLatch.countDown(); - while (bufferConsumptionLatch.getCount() > 0L) { - bufferConsumptionLatch.countDown(); - } + callbackRunning = false; } + scheduleCallbackIfNecessary(); + checkCompletion(); } } } private final CallbackRunnable callbackRunnable = new CallbackRunnable(); + private final ProduceRowsRunnable produceRowsRunnable = new ProduceRowsRunnable(); /** * {@link ProduceRowsRunnable} reads data from the underlying {@link ResultSet}, places these in @@ -360,132 +405,199 @@ public void run() { private class ProduceRowsRunnable implements Runnable { @Override public void run() { - boolean stop = false; - boolean hasNext = false; - try { - hasNext = delegateResultSet.get().next(); - } catch (Throwable e) { - synchronized (monitor) { - executionException = SpannerExceptionFactory.asSpannerException(e); - } - } try { - while (!stop && hasNext) { - try { - synchronized (monitor) { - stop = shouldStopProducing(); - } - if (!stop) { - while (buffer.remainingCapacity() == 0 && !stop) { - waitIfPaused(); - // The buffer is full and we should let the callback consume a number of rows before - // we proceed with producing any more rows to prevent us from potentially waiting on - // a full buffer repeatedly. - // Wait until at least half of the buffer is available, or if it's a bigger buffer, - // wait until at least 10 rows can be placed in it. - // TODO: Make this more dynamic / configurable? - startCallbackWithBufferLatchIfNecessary( - Math.min( - Math.min(buffer.size() / 2 + 1, buffer.size()), - MAX_WAIT_FOR_BUFFER_CONSUMPTION)); - bufferConsumptionLatch.await(); - synchronized (monitor) { - stop = shouldStopProducing(); - } - } - } - if (!stop) { - buffer.put(delegateResultSet.get().getCurrentRowAsStruct()); - startCallbackIfNecessary(); - hasNext = delegateResultSet.get().next(); - } - } catch (Throwable e) { - synchronized (monitor) { - executionException = SpannerExceptionFactory.asSpannerException(e); - stop = true; + while (true) { + synchronized (monitor) { + if (shouldStopProducer() || buffer.remainingCapacity() == 0 || !isDataAvailable()) { + return; } } - } - // We don't need any more data from the underlying result set, so we close it as soon as - // possible. Any error that might occur during this will be ignored. - closeDelegateResultSet(); - // Ensure that the callback has been called at least once, even if the result set was - // cancelled. - synchronized (monitor) { - finished = true; - stop = cursorReturnedDoneOrException; - } - // Call the callback if there are still rows in the buffer that need to be processed. - while (!stop) { - try { - waitIfPaused(); - startCallbackIfNecessary(); - // Make sure we wait until the callback runner has actually finished. - consumingLatch.await(); + ResultSet delegate = getDelegate(); + boolean hasNext = delegate.next(); + if (hasNext) { + buffer.add(delegate.getCurrentRowAsStruct()); + scheduleCallbackIfNecessary(); + } else { synchronized (monitor) { - stop = cursorReturnedDoneOrException; + finished = true; } - } catch (Throwable e) { - result.setException(e); - return; + closeDelegateResultSet(); + break; } } + scheduleCallbackIfNecessary(); + } catch (Throwable throwable) { + setExecutionException(throwable); + closeDelegateResultSet(); + scheduleCallbackIfNecessary(); } finally { - if (executorProvider.shouldAutoClose()) { - service.shutdown(); - } - for (Runnable listener : listeners) { - listener.run(); - } synchronized (monitor) { - if (executionException != null) { - result.setException(executionException); - } else if (state == State.CANCELLED) { - result.setException(CANCELLED_EXCEPTION); - } else { - result.set(null); - } + producerRunning = false; } + scheduleProducerIfNecessary(); + checkCompletion(); } } + } - /** - * 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 boolean isDataAvailable() { + try { + return StreamingUtil.isDataAvailable(getDelegate()); + } catch (Throwable t) { + // Returning true ensures that any exception encountered while checking for available data or + // initializing the delegate is not swallowed. The producer will proceed to invoke next(), + // where the exception is caught, recorded in executionException, and propagated to the + // caller. + return true; } + } - private void waitIfPaused() throws InterruptedException { - CountDownLatch pause; + private void setExecutionException(Throwable throwable) { + synchronized (monitor) { + if (executionException == null && !state.shouldStop) { + executionException = SpannerExceptionFactory.asSpannerException(throwable); + } + } + } + + private boolean shouldStopProducer() { + return finished || state.shouldStop || state == State.PAUSED || executionException != null; + } + + private boolean canScheduleProducer() { + return !producerRunning + && !shouldStopProducer() + && buffer.remainingCapacity() > 0 + && isDataAvailable(); + } + + private void scheduleProducerIfNecessary() { + synchronized (monitor) { + if (!canScheduleProducer()) { + return; + } + producerRunning = true; + } + try { + service.execute(produceRowsRunnable); + } catch (Throwable throwable) { synchronized (monitor) { - pause = pausedLatch; + producerRunning = false; } - pause.await(); + setExecutionException(throwable); + closeDelegateResultSet(); + scheduleCallbackIfNecessary(); + checkCompletion(); } + } + + private boolean canScheduleCallback() { + return (state == State.RUNNING || state == State.CANCELLED) + && !cursorReturnedDoneOrException + && !callbackRunning + && (!buffer.isEmpty() + || finished + || executionException != null + || state == State.CANCELLED); + } - private void startCallbackIfNecessary() { - startCallbackWithBufferLatchIfNecessary(0); + private void scheduleCallbackIfNecessary() { + Executor callbackExecutor = null; + synchronized (monitor) { + if (canScheduleCallback()) { + if (state == State.RUNNING) { + state = State.CONSUMING; + } + callbackRunning = true; + callbackExecutor = Preconditions.checkNotNull(this.executor); + } + } + if (callbackExecutor != null) { + try { + callbackExecutor.execute(callbackRunnable); + } catch (Throwable throwable) { + setExecutionException(throwable); + synchronized (monitor) { + callbackRunning = false; + cursorReturnedDoneOrException = true; + } + closeDelegateResultSet(); + checkCompletion(); + } } + } - private void startCallbackWithBufferLatchIfNecessary(int bufferLatch) { - synchronized (monitor) { - if ((state == State.RUNNING || state == State.CANCELLED) - && !cursorReturnedDoneOrException) { - consumingLatch = new CountDownLatch(1); - if (bufferLatch > 0) { - bufferConsumptionLatch = new CountDownLatch(bufferLatch); - } - if (state == State.RUNNING) { - state = State.CONSUMING; - } - executor.execute(callbackRunnable); + private void checkCompletion() { + boolean shouldComplete = false; + synchronized (monitor) { + if (!completed && cursorReturnedDoneOrException && !producerRunning && !callbackRunning) { + if (state != State.CANCELLED) { + state = State.DONE; } + completed = true; + shouldComplete = true; + } + } + if (shouldComplete) { + cleanupAndCompleteResult(); + } + } + + private void cleanupAndCompleteResult() { + closeDelegateResultSet(); + notifyStartOrClose(); + buffer.clear(); + currentRow = null; + if (executorProvider.shouldAutoClose()) { + service.shutdown(); + } + synchronized (monitor) { + callback = null; + executor = null; + } + for (Runnable listener : listeners) { + try { + listener.run(); + } catch (Throwable t) { + log.log(Level.WARNING, "Listener threw exception", t); } } + listeners.clear(); + SettableApiFuture resultFuture = Preconditions.checkNotNull(this.result); + if (executionException != null) { + resultFuture.setException(executionException); + } else if (state == State.CANCELLED) { + resultFuture.setException(CANCELLED_EXCEPTION); + } else { + resultFuture.set(null); + } + } + + @VisibleForTesting + int getBufferSize() { + return buffer.size(); + } + + @VisibleForTesting + State getState() { + synchronized (monitor) { + return state; + } + } + + @VisibleForTesting + boolean isClosed() { + synchronized (monitor) { + return closed; + } + } + + @VisibleForTesting + boolean isDelegateClosed() { + synchronized (monitor) { + return delegateClosed; + } } private class InitiateStreamingRunnable implements Runnable { @@ -493,51 +605,45 @@ private class InitiateStreamingRunnable implements Runnable { @Override public void run() { try { - // This method returns true if the underlying result set is a streaming result set (e.g. a - // GrpcResultSet). - // Those result sets will trigger initiateProduceRows() when the first results are received. - // Non-streaming result sets do not trigger this callback, and for those result sets, we - // need to eagerly start the ProduceRowsRunnable. - if (!initiateStreaming(AsyncResultSetImpl.this)) { - initiateProduceRows(); - } + initiateStreaming(AsyncResultSetImpl.this); + scheduleProducerIfNecessary(); } catch (Throwable exception) { - executionException = SpannerExceptionFactory.asSpannerException(exception); - initiateProduceRows(); + setExecutionException(exception); + closeDelegateResultSet(); + scheduleCallbackIfNecessary(); + checkCompletion(); } } } /** Sets the callback for this {@link AsyncResultSet}. */ @Override - public ApiFuture setCallback(Executor exec, ReadyCallback cb) { + public ApiFuture setCallback(Executor executor, ReadyCallback callback) { + SettableApiFuture resultFuture; synchronized (monitor) { Preconditions.checkState(!closed, "This AsyncResultSet has been closed"); Preconditions.checkState( this.state == State.INITIALIZED, "callback may not be set multiple times"); // Start to fetch data and buffer these. - this.result = SettableApiFuture.create(); - this.state = State.STREAMING_INITIALIZED; - this.service.execute(new InitiateStreamingRunnable()); - this.executor = MoreExecutors.newSequentialExecutor(Preconditions.checkNotNull(exec)); - this.callback = Preconditions.checkNotNull(cb); - pausedLatch.countDown(); - return result; + this.result = resultFuture = SettableApiFuture.create(); + this.state = State.RUNNING; + this.executor = MoreExecutors.newSequentialExecutor(Preconditions.checkNotNull(executor)); + this.callback = Preconditions.checkNotNull(callback); } - } - - private void initiateProduceRows() { - synchronized (monitor) { - if (this.state == State.STREAMING_INITIALIZED) { - this.state = State.RUNNING; + try { + this.service.execute(new InitiateStreamingRunnable()); + } catch (Throwable throwable) { + setExecutionException(throwable); + synchronized (monitor) { + cursorReturnedDoneOrException = true; } - produceRowsInitiated = true; + checkCompletion(); } - this.service.execute(new ProduceRowsRunnable()); + return resultFuture; } - Future getResult() { + @Nullable Future getResult() { return result; } @@ -547,22 +653,39 @@ public void cancel() { Preconditions.checkState( state != State.INITIALIZED && state != State.SYNC, "cannot cancel a result set without a callback"); + if (completed) { + return; + } state = State.CANCELLED; - pausedLatch.countDown(); } + closeDelegateResultSet(); + scheduleCallbackIfNecessary(); + checkCompletion(); } @Override public void resume() { + boolean shouldStartCallback = false; + boolean shouldScheduleProducer = false; synchronized (monitor) { Preconditions.checkState( state != State.INITIALIZED && state != State.SYNC, "cannot resume a result set without a callback"); + if (completed) { + return; + } if (state == State.PAUSED) { state = State.RUNNING; - pausedLatch.countDown(); + shouldStartCallback = true; + shouldScheduleProducer = true; } } + if (shouldStartCallback) { + scheduleCallbackIfNecessary(); + } + if (shouldScheduleProducer) { + scheduleProducerIfNecessary(); + } } private static class CreateListCallback implements ReadyCallback { @@ -584,11 +707,11 @@ public CallbackResponse cursorReady(AsyncResultSet resultSet) { case DONE: future.set(builder.build()); return CallbackResponse.DONE; - case NOT_READY: - return CallbackResponse.CONTINUE; case OK: builder.add(transformer.apply(resultSet)); break; + default: + return CallbackResponse.CONTINUE; } } } catch (Throwable t) { @@ -605,11 +728,12 @@ public ApiFuture> toListAsync( Preconditions.checkState(!closed, "This AsyncResultSet has been closed"); Preconditions.checkState( this.state == State.INITIALIZED, "This AsyncResultSet has already been used."); - final SettableApiFuture> res = SettableApiFuture.create(); - CreateListCallback callback = new CreateListCallback<>(res, transformer); - ApiFuture finished = setCallback(executor, callback); - return ApiFutures.transformAsync(finished, ignored -> res, MoreExecutors.directExecutor()); } + final SettableApiFuture> resultFuture = SettableApiFuture.create(); + CreateListCallback callback = new CreateListCallback<>(resultFuture, transformer); + ApiFuture finished = setCallback(executor, callback); + return ApiFutures.transformAsync( + finished, ignored -> resultFuture, MoreExecutors.directExecutor()); } @Override @@ -617,10 +741,10 @@ public List toList(Function transformer) throws SpannerE ApiFuture> future = toListAsync(transformer, MoreExecutors.directExecutor()); try { return future.get(); - } catch (ExecutionException e) { - throw SpannerExceptionFactory.asSpannerException(e.getCause()); - } catch (Throwable e) { - throw SpannerExceptionFactory.asSpannerException(e); + } catch (ExecutionException executionException) { + throw SpannerExceptionFactory.asSpannerException(executionException.getCause()); + } catch (Throwable throwable) { + throw SpannerExceptionFactory.asSpannerException(throwable); } } @@ -632,23 +756,32 @@ public boolean next() throws SpannerException { "Cannot call next() on a result set with a callback."); this.state = State.SYNC; } - boolean res = delegateResultSet.get().next(); - currentRow = res ? delegateResultSet.get().getCurrentRowAsStruct() : null; - return res; + try { + ResultSet delegate = getDelegate(); + boolean hasNext = delegate.next(); + currentRow = hasNext ? delegate.getCurrentRowAsStruct() : null; + return hasNext; + } finally { + notifyStartOrClose(); + } } @Override public ResultSetStats getStats() { - return delegateResultSet.get().getStats(); + return getDelegate().getStats(); } @Override public ResultSetMetadata getMetadata() { - return delegateResultSet.get().getMetadata(); + return getDelegate().getMetadata(); } boolean initiateStreaming(StreamMessageListener streamMessageListener) { - return StreamingUtil.initiateStreaming(delegateResultSet.get(), streamMessageListener); + try { + return StreamingUtil.initiateStreaming(getDelegate(), streamMessageListener); + } finally { + notifyStartOrClose(); + } } @Override @@ -669,19 +802,6 @@ public Struct getCurrentRowAsStruct() { @Override public void onStreamMessage(PartialResultSet partialResultSet, boolean bufferIsFull) { - synchronized (monitor) { - if (produceRowsInitiated) { - return; - } - // if PartialResultSet contains a resume token or buffer size is full, or - // we have reached the end of the stream, we can start the thread. - boolean startJobThread = - !partialResultSet.getResumeToken().isEmpty() - || bufferIsFull - || partialResultSet == GrpcStreamIterator.END_OF_STREAM; - if (startJobThread || state != State.STREAMING_INITIALIZED) { - initiateProduceRows(); - } - } + scheduleProducerIfNecessary(); } } diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ForwardingResultSet.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ForwardingResultSet.java index 3c4883e65862..33f721c90a04 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ForwardingResultSet.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ForwardingResultSet.java @@ -110,4 +110,9 @@ public ResultSetMetadata getMetadata() { public boolean initiateStreaming(AsyncResultSet.StreamMessageListener streamMessageListener) { return StreamingUtil.initiateStreaming(delegate.get(), streamMessageListener); } + + @Override + public boolean isDataAvailable() { + return StreamingUtil.isDataAvailable(delegate.get()); + } } 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 80a9dfcf5331..de314b642e19 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 @@ -39,9 +39,8 @@ class GrpcResultSet extends AbstractResultSet> private ResultSetMetadata metadata; private GrpcStruct currRow; private List rowData; - private SpannerException error; private ResultSetStats statistics; - private boolean closed; + private volatile boolean closed; GrpcResultSet(CloseableIterator iterator, Listener listener) { this(iterator, listener, DecodeMode.DIRECT); @@ -75,9 +74,6 @@ protected GrpcStruct currRow() { @Override public boolean next() throws SpannerException { - if (error != null) { - throw asSpannerException(error); - } try { if (currRow == null) { metadata = iterator.getMetadata(); @@ -132,6 +128,11 @@ public boolean initiateStreaming(AsyncResultSet.StreamMessageListener streamMess return iterator.initiateStreaming(streamMessageListener); } + @Override + public boolean isDataAvailable() { + return closed || iterator.isDataAvailable(); + } + @Override public void close() { synchronized (this) { @@ -140,8 +141,8 @@ public void close() { } closed = true; } - listener.onDone(iterator.isWithBeginTransaction()); iterator.close("ResultSet closed"); + listener.onDone(iterator.isWithBeginTransaction()); } @Override diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/GrpcStreamIterator.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/GrpcStreamIterator.java index e0df4c422e3c..7f4cd37396ec 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/GrpcStreamIterator.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/GrpcStreamIterator.java @@ -41,24 +41,29 @@ class GrpcStreamIterator extends AbstractIterator private static final Logger logger = Logger.getLogger(GrpcStreamIterator.class.getName()); static final PartialResultSet END_OF_STREAM = PartialResultSet.newBuilder().build(); private final int prefetchChunks; - private AsyncResultSet.StreamMessageListener streamMessageListener; + private volatile AsyncResultSet.StreamMessageListener streamMessageListener; private final ConsumerImpl consumer; private final BlockingQueue stream; private final Statement statement; - private SpannerRpc.StreamingCall call; + private volatile SpannerRpc.StreamingCall call; + private volatile boolean callInitialized; private volatile boolean withBeginTransaction; private final boolean lastStatement; + private boolean closed; + private boolean callCleared; + @Nullable private String closeMessage; private TimeUnit streamWaitTimeoutUnit; private long streamWaitTimeoutValue; - private SpannerException error; - private boolean done; + private volatile SpannerException error; + private volatile boolean done; + private volatile boolean endOfStreamReached; @VisibleForTesting GrpcStreamIterator( boolean lastStatement, int prefetchChunks, boolean cancelQueryWhenClientIsClosed) { - this(null, lastStatement, prefetchChunks, cancelQueryWhenClientIsClosed); + this(null, lastStatement, prefetchChunks, cancelQueryWhenClientIsClosed, false); } @VisibleForTesting @@ -67,9 +72,20 @@ class GrpcStreamIterator extends AbstractIterator boolean lastStatement, int prefetchChunks, boolean cancelQueryWhenClientIsClosed) { + this(statement, lastStatement, prefetchChunks, cancelQueryWhenClientIsClosed, false); + } + + @VisibleForTesting + GrpcStreamIterator( + Statement statement, + boolean lastStatement, + int prefetchChunks, + boolean cancelQueryWhenClientIsClosed, + boolean withBeginTransaction) { this.statement = statement; this.lastStatement = lastStatement; this.prefetchChunks = prefetchChunks; + this.withBeginTransaction = withBeginTransaction; this.consumer = new ConsumerImpl(cancelQueryWhenClientIsClosed); // One extra to allow for END_OF_STREAM message. this.stream = new LinkedBlockingQueue<>(prefetchChunks + 1); @@ -84,39 +100,67 @@ void registerListener(AsyncResultSet.StreamMessageListener streamMessageListener } public void setCall(SpannerRpc.StreamingCall call, boolean withBeginTransaction) { - this.call = call; - this.withBeginTransaction = withBeginTransaction; - ApiCallContext callContext = call.getCallContext(); - Duration streamWaitTimeout = - callContext == null ? null : callContext.getStreamWaitTimeoutDuration(); - if (streamWaitTimeout != null) { - // Determine the timeout unit to use. This reduces the precision to seconds if the timeout - // value is more than 1 second, which is lower than the precision that would normally be - // used by the stream watchdog (which uses a precision of 10 seconds by default). - if (streamWaitTimeout.getSeconds() > 0L) { - streamWaitTimeoutValue = streamWaitTimeout.getSeconds(); - streamWaitTimeoutUnit = TimeUnit.SECONDS; - } else if (streamWaitTimeout.getNano() > 0) { - streamWaitTimeoutValue = streamWaitTimeout.getNano(); - streamWaitTimeoutUnit = TimeUnit.NANOSECONDS; + boolean shouldCancel; + String messageToCancel; + synchronized (this) { + this.withBeginTransaction = withBeginTransaction; + if (!this.callCleared) { + this.call = call; } - // Note that if the stream-wait-timeout is zero, we won't set a timeout at all. - // That is consistent with ApiCallContext#withStreamWaitTimeout(Duration.ZERO). + this.callInitialized = true; + ApiCallContext callContext = call.getCallContext(); + Duration streamWaitTimeout = + callContext == null ? null : callContext.getStreamWaitTimeoutDuration(); + if (streamWaitTimeout != null) { + // Determine the timeout unit to use. This reduces the precision to seconds if the timeout + // value is more than 1 second, which is lower than the precision that would normally be + // used by the stream watchdog (which uses a precision of 10 seconds by default). + if (streamWaitTimeout.getSeconds() > 0L) { + streamWaitTimeoutValue = streamWaitTimeout.getSeconds(); + streamWaitTimeoutUnit = TimeUnit.SECONDS; + } else if (streamWaitTimeout.getNano() > 0) { + streamWaitTimeoutValue = streamWaitTimeout.getNano(); + streamWaitTimeoutUnit = TimeUnit.NANOSECONDS; + } + // Note that if the stream-wait-timeout is zero, we won't set a timeout at all. + // That is consistent with ApiCallContext#withStreamWaitTimeout(Duration.ZERO). + } + shouldCancel = this.closed && !this.callCleared; + messageToCancel = this.closeMessage; + } + if (shouldCancel) { + call.cancel(messageToCancel); + return; + } + PartialResultSet pending = stream.peek(); + if (pending != null) { + onStreamMessage(pending); } } @Override public void close(@Nullable String message) { - if (call != null) { - call.cancel(message); + SpannerRpc.StreamingCall callToCancel = null; + synchronized (this) { + if (!closed) { + closed = true; + closeMessage = message; + callToCancel = this.call; + } + } + if (callToCancel != null) { + callToCancel.cancel(message); } } @Override @InternalApi public void requestPrefetchChunks() { - Preconditions.checkState(call != null, "The StreamingCall object is not initialized"); - call.request(prefetchChunks); + Preconditions.checkState(callInitialized, "The StreamingCall object is not initialized"); + SpannerRpc.StreamingCall currentCall = this.call; + if (currentCall != null) { + currentCall.request(prefetchChunks); + } } @Override @@ -129,6 +173,18 @@ public boolean isLastStatement() { return lastStatement; } + @Override + public boolean isDataAvailable() { + return canPollWithoutBlocking() + || ((callInitialized || streamMessageListener == null) && error != null); + } + + @Override + public boolean canPollWithoutBlocking() { + return (callInitialized || streamMessageListener == null) + && (endOfStreamReached || (!stream.isEmpty() && error == null)); + } + @Override protected final PartialResultSet computeNext() { PartialResultSet next; @@ -147,17 +203,31 @@ protected final PartialResultSet computeNext() { throw SpannerExceptionFactory.propagateInterrupt(e); } if (next != END_OF_STREAM) { - call.request(1); + SpannerRpc.StreamingCall currentCall = this.call; + if (currentCall != null) { + currentCall.request(1); + } + if (next.getLast()) { + // All done - close() no longer needs to cancel the call. + synchronized (this) { + callCleared = true; + call = null; + } + } return next; } // All done - close() no longer needs to cancel the call. - call = null; + synchronized (this) { + callCleared = true; + call = null; + } if (error != null) { throw SpannerExceptionFactory.asSpannerException(error); } + endOfStreamReached = true; endOfData(); return null; } @@ -165,7 +235,9 @@ protected final PartialResultSet computeNext() { private void addToStream(PartialResultSet results) { // We assume that nothing from the user will interrupt gRPC event threads. Uninterruptibles.putUninterruptibly(stream, results); - onStreamMessage(results); + if (callInitialized) { + onStreamMessage(results); + } } private class ConsumerImpl implements SpannerRpc.ResultStreamConsumer { @@ -177,16 +249,28 @@ private class ConsumerImpl implements SpannerRpc.ResultStreamConsumer { @Override public void onPartialResultSet(PartialResultSet results) { + boolean isLast = results.getLast(); + if (isLast) { + synchronized (GrpcStreamIterator.this) { + done = true; + } + } addToStream(results); - if (results.getLast()) { - done = true; + if (isLast) { addToStream(END_OF_STREAM); } } @Override public void onCompleted() { - if (!done) { + boolean shouldAddEndOfStream = false; + synchronized (GrpcStreamIterator.this) { + if (!done) { + done = true; + shouldAddEndOfStream = true; + } + } + if (shouldAddEndOfStream) { addToStream(END_OF_STREAM); } } @@ -203,7 +287,16 @@ public void onError(SpannerException exception) { } } error = exception; - addToStream(END_OF_STREAM); + boolean shouldAddEndOfStream = false; + synchronized (GrpcStreamIterator.this) { + if (!done) { + done = true; + shouldAddEndOfStream = true; + } + } + if (shouldAddEndOfStream) { + addToStream(END_OF_STREAM); + } } @Override diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/GrpcValueIterator.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/GrpcValueIterator.java index 09b850c93f3c..25bb463d4246 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/GrpcValueIterator.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/GrpcValueIterator.java @@ -29,6 +29,7 @@ import com.google.spanner.v1.ResultSetStats; import com.google.spanner.v1.TypeCode; import java.util.ArrayList; +import java.util.LinkedList; import java.util.List; import javax.annotation.Nullable; @@ -40,10 +41,12 @@ private enum StreamValue { } private final CloseableIterator stream; + private final LinkedList bufferedChunks = new LinkedList<>(); + private boolean streamFinished; private ResultSetMetadata metadata; private Type type; - private PartialResultSet current; - private int pos; + private volatile PartialResultSet current; + private volatile int pos; private ResultSetStats statistics; private final Listener listener; @@ -131,43 +134,130 @@ boolean initiateStreaming(AsyncResultSet.StreamMessageListener streamMessageList return stream.initiateStreaming(streamMessageListener); } + /** + * Returns the number of complete {@link com.google.protobuf.Value}s required to assemble one row + * without blocking. + * + *

Returns at least {@code 1} even when the result set has 0 columns (such as DML statements + * without a {@code THEN RETURN} clause or {@code QueryMode.PLAN} queries). For 0-column + * statements, {@link GrpcResultSet#next()} calls {@link #ensureReady(StreamValue)} with {@link + * StreamValue#RESULT}, which loops until the entire stream has finished so that the final {@link + * ResultSetStats} is collected. Requiring at least 1 value ensures that {@link + * #isDataAvailable()} waits for {@link #streamFinished} rather than returning {@code true} as + * soon as an initial metadata-only chunk arrives. + */ + private int requiredCompleteValues() { + if (type != null) { + return Math.max(1, type.getStructFields().size()); + } + PartialResultSet firstChunk = bufferedChunks.peekFirst(); + return Math.max(1, firstChunk.getMetadata().getRowType().getFieldsCount()); + } + + /** + * Returns {@code true} if the first chunk of the stream has arrived in {@link #bufferedChunks} + * but is missing the required row-type metadata. In that case, {@link #isDataAvailable()} should + * return {@code true} immediately so {@link #ensureReady(StreamValue)} can surface the metadata + * error without waiting for more chunks. + */ + private boolean hasInvalidFirstChunkMetadata() { + if (type != null) { + return false; + } + PartialResultSet firstChunk = bufferedChunks.peekFirst(); + return firstChunk != null + && (!firstChunk.hasMetadata() || !firstChunk.getMetadata().hasRowType()); + } + + /** + * Counts the total number of complete (fully merged) {@link com.google.protobuf.Value}s currently + * available across the unconsumed tail of {@link #current} and all chunks in {@link + * #bufferedChunks}. + */ + private int countCompleteValues() { + int completeValues = countCompleteValuesInChunk(this.current, this.pos); + for (PartialResultSet chunk : bufferedChunks) { + completeValues += countCompleteValuesInChunk(chunk, 0); + } + return completeValues; + } + + /** + * Counts how many complete values are contributed by {@code chunk} starting at {@code + * startIndex}. + * + *

Whenever a chunk has {@code chunked_value = true} (with a mergeable trailing value), its + * last value is incomplete and will be merged with the first value of the following non-empty + * chunk. Thus each chunk with an incomplete trailing value reduces the number of complete values + * by 1 (either because its last value merges with the next chunk's first value, or, if it is the + * final buffered chunk, because its last value is still awaiting continuation). + */ + private int countCompleteValuesInChunk(@Nullable PartialResultSet chunk, int startIndex) { + if (chunk == null) { + return 0; + } + int count = chunk.getValuesCount(); + int remaining = count - startIndex; + if (remaining <= 0) { + return 0; + } + boolean hasIncompleteTrailingValue = + chunk.getChunkedValue() && isMergeable(chunk.getValues(count - 1).getKindCase()); + return hasIncompleteTrailingValue ? remaining - 1 : remaining; + } + + /** + * Returns {@code true} if metadata is known (either already parsed into {@link #type} or present + * on the first buffered chunk) and enough complete values are buffered to assemble a full row. + */ + private boolean hasCompleteRowAvailable() { + return (type != null || !bufferedChunks.isEmpty()) + && countCompleteValues() >= requiredCompleteValues(); + } + + /** + * Non-blockingly pulls available chunks from {@link #stream} into {@link #bufferedChunks} until a + * complete row is available, EOF is reached, or no further chunks can be polled without blocking. + */ + private void bufferAvailableChunks() { + while (!streamFinished + && !hasInvalidFirstChunkMetadata() + && !hasCompleteRowAvailable() + && stream.canPollWithoutBlocking()) { + if (stream.hasNext()) { + bufferedChunks.addLast(stream.next()); + } else { + streamFinished = true; + } + } + } + + /** + * Returns {@code true} if the next call to {@link GrpcResultSet#next()} is guaranteed not to + * block on network I/O (because a complete row, EOF, or an error/closure is already available). + */ + synchronized boolean isDataAvailable() { + bufferAvailableChunks(); + if (streamFinished || hasInvalidFirstChunkMetadata() || hasCompleteRowAvailable()) { + return true; + } + return stream.isDataAvailable(); + } + Type type() { checkState(type != null, "metadata has not been received"); return type; } private boolean ensureReady(StreamValue requiredValue) throws SpannerException { - while (current == null || pos >= current.getValuesCount()) { - if (!stream.hasNext()) { + while (!hasUnreadValuesInCurrentChunk()) { + if (!advanceToNextChunk()) { return false; } - current = stream.next(); - pos = 0; if (type == null) { - // This is the first message on the stream. - if (!current.hasMetadata() || !current.getMetadata().hasRowType()) { - throw newSpannerException(ErrorCode.INTERNAL, "Missing type metadata in first message"); - } - metadata = current.getMetadata(); - com.google.spanner.v1.Type typeProto = - com.google.spanner.v1.Type.newBuilder() - .setCode(TypeCode.STRUCT) - .setStructType(metadata.getRowType()) - .build(); - try { - type = Type.fromProto(typeProto); - } catch (IllegalArgumentException e) { - throw newSpannerException( - ErrorCode.INTERNAL, "Invalid type metadata: " + e.getMessage(), e); - } - } - // collect the precommit token from each PartialResultSet - if (current.hasPrecommitToken()) { - listener.onPrecommitToken(current.getPrecommitToken()); - } - if (current.hasStats()) { - statistics = current.getStats(); + initializeMetadata(current); } + updatePrecommitTokenAndStats(current); if (requiredValue == StreamValue.METADATA) { return true; } @@ -175,7 +265,65 @@ private boolean ensureReady(StreamValue requiredValue) throws SpannerException { return true; } + private boolean hasUnreadValuesInCurrentChunk() { + return current != null && pos < current.getValuesCount(); + } + + private boolean advanceToNextChunk() { + PartialResultSet nextChunk; + synchronized (this) { + nextChunk = bufferedChunks.pollFirst(); + if (nextChunk == null && streamFinished) { + return false; + } + } + if (nextChunk == null) { + if (!stream.hasNext()) { + synchronized (this) { + streamFinished = true; + } + return false; + } + nextChunk = stream.next(); + } + synchronized (this) { + pos = 0; + current = nextChunk; + } + return true; + } + + private void initializeMetadata(PartialResultSet firstChunk) { + if (!firstChunk.hasMetadata() || !firstChunk.getMetadata().hasRowType()) { + throw newSpannerException(ErrorCode.INTERNAL, "Missing type metadata in first message"); + } + metadata = firstChunk.getMetadata(); + com.google.spanner.v1.Type typeProto = + com.google.spanner.v1.Type.newBuilder() + .setCode(TypeCode.STRUCT) + .setStructType(metadata.getRowType()) + .build(); + try { + type = Type.fromProto(typeProto); + } catch (IllegalArgumentException e) { + throw newSpannerException(ErrorCode.INTERNAL, "Invalid type metadata: " + e.getMessage(), e); + } + } + + private void updatePrecommitTokenAndStats(PartialResultSet chunk) { + if (chunk.hasPrecommitToken()) { + listener.onPrecommitToken(chunk.getPrecommitToken()); + } + if (chunk.hasStats()) { + statistics = chunk.getStats(); + } + } + void close(@Nullable String message) { + synchronized (this) { + bufferedChunks.clear(); + streamFinished = true; + } stream.close(message); } 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..0d6b6e150511 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 @@ -41,9 +41,10 @@ import java.util.concurrent.CountDownLatch; import java.util.concurrent.Executor; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.logging.Level; import java.util.logging.Logger; -import javax.annotation.Nullable; +import org.jspecify.annotations.Nullable; /** * Wraps an iterator over partial result sets, supporting resuming RPCs on error. This class keeps @@ -73,16 +74,31 @@ abstract class ResumableStreamIterator extends AbstractIterator implements CloseableIterator { private final ErrorHandler errorHandler; - private AsyncResultSet.StreamMessageListener streamMessageListener; + private volatile AsyncResultSet.StreamMessageListener streamMessageListener; private final RetrySettings streamingRetrySettings; private final Set retryableCodes; private static final Logger logger = Logger.getLogger(ResumableStreamIterator.class.getName()); private BackOff backOff; - private final LinkedList buffer = new LinkedList<>(); + @VisibleForTesting final LinkedList buffer = new LinkedList<>(); private final int maxBufferSize; private final ISpan span; private final TraceWrapper tracer; - private CloseableIterator stream; + + /** Guards state transitions between {@link #stream} and {@link #closed}. */ + private final Object streamLock = new Object(); + + /** Indicates whether this iterator has been closed. */ + private volatile boolean closed; + + @Nullable private String closeMessage; + private boolean streamSetInStartStream; + private volatile boolean withBeginTransaction; + private volatile boolean lastStatement; + + /** Ensures the tracing span is finalized at most once. */ + private final AtomicBoolean spanEnded = new AtomicBoolean(false); + + private volatile CloseableIterator stream; /** * The number of consecutive failed attempts without any progress on the stream. Reset to zero @@ -98,7 +114,7 @@ abstract class ResumableStreamIterator extends AbstractIterator= totalTimeoutMillis - elapsedMillis; } - private ExponentialBackOff newBackOff() { + @VisibleForTesting + ExponentialBackOff newBackOff() { return new ExponentialBackOff.Builder() .setMultiplier(streamingRetrySettings.getRetryDelayMultiplier()) // All of these values must be > 0. @@ -210,7 +227,8 @@ private ExponentialBackOff newBackOff() { .build(); } - private static long nextBackOffMillis(BackOff backoff) throws SpannerException { + @VisibleForTesting + static long nextBackOffMillis(BackOff backoff) throws SpannerException { try { return backoff.nextBackOffMillis(); } catch (IOException e) { @@ -266,23 +284,118 @@ boolean prepareIteratorForRetryOnDifferentGrpcChannel() { return false; } + /** Ensures the tracing span opened in the constructor is ended exactly once. */ + private void endSpanOnce() { + if (spanEnded.compareAndSet(false, true)) { + span.end(); + } + } + + /** + * Safely attaches a newly initialized gRPC stream. If {@link #close(String)} was called while + * stream acquisition was in flight, the newly acquired stream is immediately closed and discarded + * rather than being resurrected or leaked. + * + * @param streamToAssign the newly opened stream iterator + * @param requestPrefetch whether to eagerly request initial prefetch chunks if attached + */ + private boolean assignStream( + CloseableIterator streamToAssign, boolean requestPrefetch) { + boolean isClosed; + String messageToClose; + boolean shouldRequestPrefetch = false; + synchronized (streamLock) { + isClosed = this.closed; + messageToClose = this.closeMessage; + if (streamToAssign != null) { + this.withBeginTransaction = streamToAssign.isWithBeginTransaction(); + this.lastStatement = streamToAssign.isLastStatement(); + } + if (!isClosed) { + if (!requestPrefetch) { + this.streamSetInStartStream = streamToAssign != null; + this.stream = streamToAssign; + } else if (!this.streamSetInStartStream) { + this.stream = streamToAssign; + shouldRequestPrefetch = true; + } else if (this.stream == streamToAssign) { + shouldRequestPrefetch = true; + } + } + } + if (isClosed) { + if (streamToAssign != null) { + streamToAssign.close(messageToClose); + } + endSpanOnce(); + } else if (shouldRequestPrefetch) { + streamToAssign.requestPrefetchChunks(); + } + return !isClosed; + } + + /** Resets the active stream reference during retries or channel failovers. */ + private void resetStream() { + synchronized (streamLock) { + if (this.stream != null) { + this.withBeginTransaction = this.stream.isWithBeginTransaction(); + this.lastStatement = this.stream.isLastStatement(); + } + this.stream = null; + } + } + + /** + * Registers the active stream early during {@code startStream()} setup so that in-flight RPCs can + * be cancelled immediately if {@link #close(String)} is invoked before {@code startStream()} + * returns. + * + * @return {@code true} if the stream was attached, or {@code false} if this iterator was already + * closed + */ + boolean setStream(CloseableIterator stream) { + return assignStream(stream, /* requestPrefetch= */ false); + } + @Override public void close(@Nullable String message) { - if (stream != null) { - stream.close(message); - span.end(); - stream = null; + CloseableIterator streamToClose; + synchronized (streamLock) { + if (!closed) { + closed = true; + closeMessage = message; + } + streamToClose = this.stream; + if (streamToClose != null) { + this.withBeginTransaction = streamToClose.isWithBeginTransaction(); + this.lastStatement = streamToClose.isLastStatement(); + } + this.stream = null; + } + // Perform stream teardown outside the lock to prevent deadlocks with external callbacks. + if (streamToClose != null) { + streamToClose.close(message); + } + endSpanOnce(); + synchronized (buffer) { + buffer.clear(); } } + boolean isClosed() { + return closed; + } + @Override public boolean isWithBeginTransaction() { - return stream != null && stream.isWithBeginTransaction(); + CloseableIterator currentStream = this.stream; + return currentStream != null ? currentStream.isWithBeginTransaction() : withBeginTransaction; } @Override public boolean isLastStatement() { - return stream != null && stream.isLastStatement(); + CloseableIterator currentStream = this.stream; + return currentStream != null ? currentStream.isLastStatement() : lastStatement; } @Override @@ -293,51 +406,118 @@ public boolean initiateStreaming(AsyncResultSet.StreamMessageListener streamMess return true; } + private boolean isBufferFlushableLocked() { + return !buffer.isEmpty() + && (finished || !safeToRetry || !buffer.getLast().getResumeToken().isEmpty()); + } + + private void bufferChunkLocked(PartialResultSet next) { + boolean hasResumeToken = !next.getResumeToken().isEmpty(); + if (hasResumeToken) { + // Only a resume token that differs from the last seen token is progress: a stream + // that repeatedly returns the token that was used to resume it has not moved past it. + if (!next.getResumeToken().equals(resumeToken)) { + // The stream made progress, so reset the budget for consecutive failed attempts. + attempts = 0; + backOff = null; + retrySequenceStartNanos = -1L; + } + resumeToken = next.getResumeToken(); + safeToRetry = true; + } + buffer.add(next); + if (buffer.size() > maxBufferSize && buffer.getLast().getResumeToken().isEmpty()) { + // We need to flush without a restart token. Errors encountered until we see + // such a token will fail the read. + safeToRetry = false; + } + } + + /** + * Non-blockingly drains available non-error chunks from {@code currentStream} into {@link + * #buffer} until {@link #buffer} becomes flushable (for example, a resume token is seen or {@link + * #maxBufferSize} is exceeded) or the stream reaches EOF. + */ + private void drainAvailableFromStreamLocked( + @Nullable CloseableIterator currentStream) { + if (currentStream == null) { + return; + } + while (!closed + && !finished + && !isBufferFlushableLocked() + && currentStream.canPollWithoutBlocking()) { + if (currentStream.hasNext()) { + bufferChunkLocked(currentStream.next()); + } else { + finished = true; + } + } + } + + @Override + public boolean canPollWithoutBlocking() { + if (closed) { + return false; + } + synchronized (buffer) { + drainAvailableFromStreamLocked(this.stream); + return !closed && (isBufferFlushableLocked() || finished); + } + } + + @Override + public boolean isDataAvailable() { + // If closed, return true to unblock any waiting async consumer so it can observe EOF/closure + // instead of hanging indefinitely. + if (closed || canPollWithoutBlocking()) { + return true; + } + CloseableIterator currentStream = this.stream; + return currentStream != null && currentStream.isDataAvailable(); + } + @Override protected PartialResultSet computeNext() { int numAttemptsOnOtherChannel = 0; Context context = Context.current(); - while (true) { + while (!closed) { // Eagerly start stream before consuming any buffered items. startGrpcStreaming(); + if (closed) { + break; + } // Buffer contains items up to a resume token or has reached capacity: flush. - if (!buffer.isEmpty() - && (finished || !safeToRetry || !buffer.getLast().getResumeToken().isEmpty())) { - return buffer.pop(); + PartialResultSet buffered = null; + synchronized (buffer) { + if (isBufferFlushableLocked()) { + buffered = buffer.pop(); + } + } + if (buffered != null) { + return buffered; + } + // Snapshot the volatile stream reference to guard against concurrent close() nulling the + // field between check and dereference. + CloseableIterator currentStream = this.stream; + if (currentStream == null) { + break; } try { - if (stream.hasNext()) { - PartialResultSet next = stream.next(); - boolean hasResumeToken = !next.getResumeToken().isEmpty(); - if (hasResumeToken) { - // Only a resume token that differs from the last seen token is progress: a stream - // that repeatedly returns the token that was used to resume it has not moved past it. - if (!next.getResumeToken().equals(resumeToken)) { - // The stream made progress, so reset the budget for consecutive failed attempts. - attempts = 0; - backOff = null; - retrySequenceStartNanos = -1L; + if (currentStream.hasNext()) { + PartialResultSet next = currentStream.next(); + synchronized (buffer) { + bufferChunkLocked(next); + if (isBufferFlushableLocked()) { + return buffer.pop(); } - resumeToken = next.getResumeToken(); - safeToRetry = true; - } - // If the buffer is empty and this chunk has a resume token or we cannot resume safely - // anyway, we can yield it immediately rather than placing it in the buffer to be - // returned on the next iteration. - if ((hasResumeToken || !safeToRetry) && buffer.isEmpty()) { - return next; - } - buffer.add(next); - if (buffer.size() > maxBufferSize && buffer.getLast().getResumeToken().isEmpty()) { - // We need to flush without a restart token. Errors encountered until we see - // such a token will fail the read. - safeToRetry = false; } } else { finished = true; - if (buffer.isEmpty()) { - endOfData(); - return null; + synchronized (buffer) { + if (buffer.isEmpty()) { + break; + } } } } catch (SpannerException spannerException) { @@ -347,12 +527,16 @@ protected PartialResultSet computeNext() { continue; } // Check if we should retry the request on a different gRPC channel. - if (resumeToken == null && buffer.isEmpty()) { + boolean bufferIsEmpty; + synchronized (buffer) { + bufferIsEmpty = buffer.isEmpty(); + } + if (resumeToken == null && bufferIsEmpty) { Throwable translated = errorHandler.translateException(spannerException); if (translated instanceof RetryOnDifferentGrpcChannelException) { if (++numAttemptsOnOtherChannel < errorHandler.getMaxAttempts() && prepareIteratorForRetryOnDifferentGrpcChannel()) { - stream = null; + resetStream(); continue; } } @@ -366,6 +550,8 @@ && prepareIteratorForRetryOnDifferentGrpcChannel()) { throw e; } } + endOfData(); + return null; } /** Monotonic time source, overridable for deterministic retry-budget tests. */ @@ -414,28 +600,44 @@ private void handleRetryableException(Context context, SpannerException spannerE span.addAnnotation("Stream broken. Safe to retry", spannerException); logger.log(Level.FINE, "Retryable exception, will sleep and retry", spannerException); // Truncate any items in the buffer before the last retry token. - while (!buffer.isEmpty() && buffer.getLast().getResumeToken().isEmpty()) { - buffer.removeLast(); + synchronized (buffer) { + while (!buffer.isEmpty() && buffer.getLast().getResumeToken().isEmpty()) { + buffer.removeLast(); + } + assert buffer.isEmpty() || buffer.getLast().getResumeToken().equals(resumeToken); } - assert buffer.isEmpty() || buffer.getLast().getResumeToken().equals(resumeToken); - stream = null; - try (IScope s = tracer.withSpan(span)) { + resetStream(); + try (IScope scope = tracer.withSpan(span)) { backoffSleep(context, delayMillis); } } + /** + * Lazily starts the underlying gRPC stream under {@link #streamLock} if streaming has not already + * been initiated and the iterator has not been closed. + */ private void startGrpcStreaming() { - if (stream == null) { - span.addAnnotation( - "Starting/Resuming stream", - "ResumeToken", - resumeToken == null ? "null" : resumeToken.toStringUtf8()); - 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(); + synchronized (streamLock) { + if (stream != null || closed) { + return; + } + streamSetInStartStream = false; + } + span.addAnnotation( + "Starting/Resuming stream", + "ResumeToken", + resumeToken == null ? "null" : resumeToken.toStringUtf8()); + 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. + CloseableIterator streamIterator; + try { + streamIterator = checkNotNull(startStream(resumeToken, streamMessageListener, requestId)); + } catch (Throwable t) { + resetStream(); + throw t; } + assignStream(streamIterator, /* requestPrefetch= */ true); } } diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/StreamingResultSet.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/StreamingResultSet.java index 47b10d852c64..1d23d4a7e2ec 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/StreamingResultSet.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/StreamingResultSet.java @@ -28,4 +28,13 @@ interface StreamingResultSet extends ResultSet { */ @InternalApi boolean initiateStreaming(AsyncResultSet.StreamMessageListener streamMessageListener); + + /** + * Returns true if data (a chunk, row, EOF, or error) is available to read immediately without + * blocking the calling thread on network I/O. + */ + @InternalApi + default boolean isDataAvailable() { + return true; + } } diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/StreamingUtil.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/StreamingUtil.java index 54496d39f965..3798b6b1e023 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/StreamingUtil.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/StreamingUtil.java @@ -27,4 +27,11 @@ static boolean initiateStreaming( } return false; } + + static boolean isDataAvailable(ResultSet resultSet) { + if (resultSet instanceof StreamingResultSet) { + return ((StreamingResultSet) resultSet).isDataAvailable(); + } + return true; + } } diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java index d6353bfcdb86..4a027efd40a5 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java @@ -893,6 +893,15 @@ public void onDone(boolean withBeginTransaction) { } } + @Override + void onStartFailed(boolean withBeginTransaction, Throwable throwable) { + if (withBeginTransaction + && transactionIdFuture != null + && !this.transactionIdFuture.isDone()) { + this.transactionIdFuture.setException(throwable); + } + } + @Override public void buffer(Mutation mutation) { synchronized (committingLock) { diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncReadOnlyTransactionTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncReadOnlyTransactionTest.java index 034af24f8cdd..3e769413b155 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncReadOnlyTransactionTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncReadOnlyTransactionTest.java @@ -36,10 +36,15 @@ import com.google.spanner.v1.ExecuteSqlRequest; import io.grpc.ManagedChannelBuilder; import io.grpc.Status; +import io.opencensus.trace.Tracing; +import io.opentelemetry.api.OpenTelemetry; +import io.opentelemetry.api.common.Attributes; +import io.opentelemetry.api.trace.Span; import java.util.Arrays; import java.util.Collections; import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -686,13 +691,20 @@ public void closeAsyncThenCloseUnusedAsyncResultSetCompletesCleanly() throws Exc try (ResultSet resultSet = client().singleUse().executeQuery(READ_ONE_KEY_VALUE_STATEMENT)) { while (resultSet.next()) {} } + mockSpanner.clearRequests(); ReadOnlyTransaction transaction = client().readOnlyTransaction(); - AsyncResultSet resultSet = transaction.executeQueryAsync(READ_ONE_KEY_VALUE_STATEMENT); - ApiFuture closeFuture = transaction.closeAsync(); - assertFalse(closeFuture.isDone()); - resultSet.close(); - closeFuture.get(10, TimeUnit.SECONDS); - assertTrue(closeFuture.isDone()); + mockSpanner.freeze(); + try { + AsyncResultSet resultSet = transaction.executeQueryAsync(READ_ONE_KEY_VALUE_STATEMENT); + ApiFuture closeFuture = transaction.closeAsync(); + assertFalse(closeFuture.isDone()); + resultSet.close(); + closeFuture.get(10, TimeUnit.SECONDS); + assertTrue(closeFuture.isDone()); + assertEquals(0, mockSpanner.countRequestsOfType(BeginTransactionRequest.class)); + } finally { + mockSpanner.unfreeze(); + } } @Test @@ -736,4 +748,138 @@ public void closeAsyncWithInlinedBeginTransaction() throws Exception { closeFuture.get(10, TimeUnit.SECONDS); assertTrue(closeFuture.isDone()); } + + private static ISpan createCloseOnStartStreamSpan(AtomicReference closeableRef) { + return new OpenTelemetrySpan(Span.getInvalid()) { + @Override + public void addAnnotation(String message, String key, String value) { + if ("Starting/Resuming stream".equals(message)) { + AutoCloseable closeable = closeableRef.getAndSet(null); + if (closeable != null) { + try { + closeable.close(); + } catch (Exception e) { + throw new RuntimeException(e); + } + } + } + } + }; + } + + private static TraceWrapper createCloseOnStartStreamTracer( + AtomicReference closeableRef) { + return new TraceWrapper(Tracing.getTracer(), OpenTelemetry.noop().getTracer("test"), false) { + @Override + ISpan spanBuilderWithExplicitParent( + String spanName, ISpan parentSpan, Attributes attributes) { + return createCloseOnStartStreamSpan(closeableRef); + } + }; + } + + @Test(timeout = 10000) + public void cancelInlinedBeginQueryBeforeSetStreamFailsWaitingQueryAndCompletesCloseAsync() + throws Exception { + try (ResultSet resultSet = client().singleUse().executeQuery(READ_ONE_KEY_VALUE_STATEMENT)) { + while (resultSet.next()) {} + } + mockSpanner.clearRequests(); + + ReadOnlyTransaction transaction = + client() + .readOnlyTransaction( + TimestampBound.strong(), + Options.beginTransactionOption(Options.BeginTransactionOption.INLINE)); + + AtomicReference resultSetRef = new AtomicReference<>(); + ((AbstractReadContext) transaction).tracer = createCloseOnStartStreamTracer(resultSetRef); + + // Close the result set after startGrpcStreaming passes its initial closed check but before + // startStream calls setStream(stream), exercising !setStream(stream) and onStartFailed(true). + AsyncResultSetImpl resultSet = + (AsyncResultSetImpl) transaction.executeQueryAsync(READ_ONE_KEY_VALUE_STATEMENT); + resultSetRef.set(resultSet); + resultSet.initiateStreaming(resultSet); + + // A subsequent query on the same transaction must fail immediately with FAILED_PRECONDITION + // rather than hanging for 60s waiting on transactionIdFuture. + try (ResultSet secondResultSet = transaction.executeQuery(READ_ONE_KEY_VALUE_STATEMENT)) { + SpannerException exception = assertThrows(SpannerException.class, secondResultSet::next); + assertEquals(ErrorCode.FAILED_PRECONDITION, exception.getErrorCode()); + } + + ApiFuture closeFuture = transaction.closeAsync(); + closeFuture.get(10, TimeUnit.SECONDS); + assertTrue(closeFuture.isDone()); + } + + @Test(timeout = 10000) + public void cancelInlinedBeginReadBeforeSetStreamFailsWaitingReadAndCompletesCloseAsync() + throws Exception { + try (ResultSet resultSet = client().singleUse().executeQuery(READ_ONE_KEY_VALUE_STATEMENT)) { + while (resultSet.next()) {} + } + mockSpanner.clearRequests(); + + ReadOnlyTransaction transaction = + client() + .readOnlyTransaction( + TimestampBound.strong(), + Options.beginTransactionOption(Options.BeginTransactionOption.INLINE)); + + AtomicReference resultSetRef = new AtomicReference<>(); + ((AbstractReadContext) transaction).tracer = createCloseOnStartStreamTracer(resultSetRef); + + AsyncResultSetImpl resultSet = + (AsyncResultSetImpl) + transaction.readAsync("TestTable", KeySet.all(), Collections.singletonList("Value")); + resultSetRef.set(resultSet); + resultSet.initiateStreaming(resultSet); + + try (ResultSet secondResultSet = + transaction.read("TestTable", KeySet.all(), Collections.singletonList("Value"))) { + SpannerException exception = assertThrows(SpannerException.class, secondResultSet::next); + assertEquals(ErrorCode.FAILED_PRECONDITION, exception.getErrorCode()); + } + + ApiFuture closeFuture = transaction.closeAsync(); + closeFuture.get(10, TimeUnit.SECONDS); + assertTrue(closeFuture.isDone()); + } + + @Test(timeout = 10000) + public void readWriteTransactionInlinedBeginClosedBeforeSetStreamFailsWaitingQuery() { + try (ResultSet resultSet = client().singleUse().executeQuery(READ_ONE_KEY_VALUE_STATEMENT)) { + while (resultSet.next()) {} + } + mockSpanner.clearRequests(); + + SpannerException commitException = + assertThrows( + SpannerException.class, + () -> + client() + .readWriteTransaction() + .run( + transaction -> { + AtomicReference resultSetRef = new AtomicReference<>(); + ((AbstractReadContext) transaction).tracer = + createCloseOnStartStreamTracer(resultSetRef); + + ResultSet resultSet = + transaction.executeQuery(READ_ONE_KEY_VALUE_STATEMENT); + resultSetRef.set(resultSet); + StreamingUtil.initiateStreaming(resultSet, (chunk, full) -> {}); + + try (ResultSet secondResultSet = + transaction.executeQuery(READ_ONE_KEY_VALUE_STATEMENT)) { + SpannerException exception = + assertThrows(SpannerException.class, secondResultSet::next); + assertEquals(ErrorCode.FAILED_PRECONDITION, exception.getErrorCode()); + } + return null; + })); + assertEquals(ErrorCode.FAILED_PRECONDITION, commitException.getErrorCode()); + } } 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 2e7f17485263..f107b5dc7b52 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 @@ -20,36 +20,61 @@ import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertNull; import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import static org.mockito.Mockito.withSettings; import com.google.api.core.ApiFuture; +import com.google.api.core.SettableApiFuture; import com.google.api.gax.core.ExecutorProvider; import com.google.cloud.spanner.AsyncResultSet.CallbackResponse; import com.google.cloud.spanner.AsyncResultSet.CursorState; import com.google.cloud.spanner.AsyncResultSet.ReadyCallback; +import com.google.cloud.spanner.ErrorHandler.DefaultErrorHandler; +import com.google.cloud.spanner.XGoogSpannerRequestId.NoopRequestIdCreator; +import com.google.cloud.spanner.spi.v1.SpannerRpc.StreamingCall; +import com.google.cloud.spanner.v1.stub.SpannerStubSettings; import com.google.common.base.Function; import com.google.common.collect.Range; import com.google.protobuf.ByteString; import com.google.protobuf.Value; +import com.google.spanner.v1.MultiplexedSessionPrecommitToken; import com.google.spanner.v1.PartialResultSet; +import com.google.spanner.v1.ResultSetMetadata; +import com.google.spanner.v1.ResultSetStats; +import com.google.spanner.v1.Transaction; +import io.opencensus.trace.Tracing; +import io.opentelemetry.api.OpenTelemetry; +import io.opentelemetry.api.trace.Span; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; import java.util.List; +import java.util.Random; import java.util.concurrent.BlockingDeque; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; import java.util.concurrent.Executor; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; +import java.util.concurrent.Future; import java.util.concurrent.LinkedBlockingDeque; +import java.util.concurrent.RejectedExecutionException; import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.Semaphore; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import org.junit.Assume; import org.junit.Before; import org.junit.Test; import org.junit.runner.RunWith; @@ -563,6 +588,7 @@ public void testOnStreamMessageWhenResumeTokenIsPresent() { Mockito.when( delegate.initiateStreaming(Mockito.any(AsyncResultSet.StreamMessageListener.class))) .thenReturn(true); + Mockito.when(delegate.isDataAvailable()).thenReturn(true); rs.setCallback(Executors.newSingleThreadExecutor(), ignored -> CallbackResponse.DONE); rs.onStreamMessage( @@ -584,6 +610,7 @@ public void testOnStreamMessageWhenCurrentBufferSizeReachedPrefetchChunkSize() { Mockito.when( delegate.initiateStreaming(Mockito.any(AsyncResultSet.StreamMessageListener.class))) .thenReturn(true); + Mockito.when(delegate.isDataAvailable()).thenReturn(true); rs.setCallback(Executors.newSingleThreadExecutor(), ignored -> CallbackResponse.DONE); rs.onStreamMessage( @@ -606,7 +633,2817 @@ public void testOnStreamMessageWhenAsyncResultIsCancelled() { rs.cancel(); rs.onStreamMessage( PartialResultSet.newBuilder().addValues(Value.newBuilder().build()).build(), false); - Mockito.verify(mockedProvider.getExecutor(), times(2)).execute(Mockito.any()); + Mockito.verify(mockedProvider.getExecutor(), times(1)).execute(Mockito.any()); + } + } + + @Test + public void testSequentialRowDeliveryUnderConcurrentEvents() throws Exception { + ExecutorService callbackExecutor = Executors.newFixedThreadPool(2); + ExecutorService resumeExecutor = Executors.newFixedThreadPool(2); + final int rowCount = 100; + ResultSet delegate = mock(ResultSet.class); + when(delegate.next()) + .thenAnswer( + new Answer() { + int currentRowIndex = 0; + + @Override + public Boolean answer(InvocationOnMock invocation) { + currentRowIndex++; + return currentRowIndex <= rowCount; + } + }); + when(delegate.getCurrentRowAsStruct()) + .thenAnswer( + new Answer() { + int currentRowIndex = 0; + + @Override + public Struct answer(InvocationOnMock invocation) { + currentRowIndex++; + return Struct.newBuilder().set("ID").to((long) currentRowIndex).build(); + } + }); + + final List receivedRowIds = Collections.synchronizedList(new ArrayList<>()); + final CountDownLatch finishedLatch = new CountDownLatch(1); + + try (AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { + ApiFuture callbackFuture = + asyncResultSet.setCallback( + callbackExecutor, + resultSet -> { + CursorState cursorState; + while ((cursorState = resultSet.tryNext()) == CursorState.OK) { + receivedRowIds.add(resultSet.getLong("ID")); + } + if (cursorState == CursorState.DONE) { + finishedLatch.countDown(); + return CallbackResponse.DONE; + } + return CallbackResponse.CONTINUE; + }); + + final AtomicBoolean testRunning = new AtomicBoolean(true); + resumeExecutor.execute( + () -> { + while (testRunning.get()) { + asyncResultSet.resume(); + Thread.yield(); + } + }); + + assertTrue(finishedLatch.await(10, TimeUnit.SECONDS)); + testRunning.set(false); + assertNull(callbackFuture.get(5, TimeUnit.SECONDS)); + + assertEquals(rowCount, receivedRowIds.size()); + for (int i = 0; i < rowCount; i++) { + assertEquals((long) (i + 1), (long) receivedRowIds.get(i)); + } + } finally { + callbackExecutor.shutdown(); + resumeExecutor.shutdown(); + } + } + + @Test + public void testNonStreamingResultSetWithSmallBufferCapacity() throws Exception { + ExecutorService callbackExecutor = Executors.newSingleThreadExecutor(); + final int rowCount = 50; + ResultSet delegate = mock(ResultSet.class); + when(delegate.next()) + .thenAnswer( + new Answer() { + int currentRowIndex = 0; + + @Override + public Boolean answer(InvocationOnMock invocation) { + currentRowIndex++; + return currentRowIndex <= rowCount; + } + }); + when(delegate.getCurrentRowAsStruct()) + .thenAnswer( + new Answer() { + int currentRowIndex = 0; + + @Override + public Struct answer(InvocationOnMock invocation) { + currentRowIndex++; + return Struct.newBuilder().set("ID").to((long) currentRowIndex).build(); + } + }); + + final List receivedRowIds = new ArrayList<>(); + // Use buffer size = 1 to force maximum contention and buffer exhaustion + try (AsyncResultSetImpl asyncResultSet = new AsyncResultSetImpl(simpleProvider, delegate, 1)) { + ApiFuture callbackFuture = + asyncResultSet.setCallback( + callbackExecutor, + resultSet -> { + CursorState cursorState; + while ((cursorState = resultSet.tryNext()) == CursorState.OK) { + receivedRowIds.add(resultSet.getLong("ID")); + } + if (cursorState == CursorState.DONE) { + return CallbackResponse.DONE; + } + return CallbackResponse.CONTINUE; + }); + + assertNull(callbackFuture.get(10, TimeUnit.SECONDS)); + assertEquals(rowCount, receivedRowIds.size()); + for (int i = 0; i < rowCount; i++) { + assertEquals((long) (i + 1), (long) receivedRowIds.get(i)); + } + } finally { + callbackExecutor.shutdown(); + } + } + + @Test + public void testPostStreamDrainingWithSlowConsumer() throws Exception { + ExecutorService callbackExecutor = Executors.newSingleThreadExecutor(); + final int rowCount = 5; + ResultSet delegate = mock(ResultSet.class); + when(delegate.next()) + .thenAnswer( + new Answer() { + int currentRowIndex = 0; + + @Override + public Boolean answer(InvocationOnMock invocation) { + currentRowIndex++; + return currentRowIndex <= rowCount; + } + }); + when(delegate.getCurrentRowAsStruct()) + .thenAnswer( + new Answer() { + int currentRowIndex = 0; + + @Override + public Struct answer(InvocationOnMock invocation) { + currentRowIndex++; + return Struct.newBuilder().set("ID").to((long) currentRowIndex).build(); + } + }); + + final List receivedRowIds = new ArrayList<>(); + final AtomicBoolean listenerInvoked = new AtomicBoolean(false); + + try (AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { + asyncResultSet.addListener(() -> listenerInvoked.set(true)); + + ApiFuture callbackFuture = + asyncResultSet.setCallback( + callbackExecutor, + resultSet -> { + CursorState cursorState = resultSet.tryNext(); + if (cursorState == CursorState.OK) { + receivedRowIds.add(resultSet.getLong("ID")); + return CallbackResponse.PAUSE; + } else if (cursorState == CursorState.DONE) { + return CallbackResponse.DONE; + } + return CallbackResponse.CONTINUE; + }); + + while (!callbackFuture.isDone()) { + Thread.sleep(10); + asyncResultSet.resume(); + } + + assertNull(callbackFuture.get(5, TimeUnit.SECONDS)); + assertTrue(listenerInvoked.get()); + assertEquals(rowCount, receivedRowIds.size()); + } finally { + callbackExecutor.shutdown(); + } + } + + @Test + public void testStreamingResultSetWithMultipleRowsPerChunk() throws Exception { + ExecutorService callbackExecutor = Executors.newSingleThreadExecutor(); + StreamingResultSet delegate = mock(StreamingResultSet.class); + when(delegate.isDataAvailable()).thenReturn(true); + final int rowCount = 10; + when(delegate.next()) + .thenAnswer( + new Answer() { + int currentRowIndex = 0; + + @Override + public Boolean answer(InvocationOnMock invocation) { + currentRowIndex++; + return currentRowIndex <= rowCount; + } + }); + when(delegate.getCurrentRowAsStruct()) + .thenAnswer( + new Answer() { + int currentRowIndex = 0; + + @Override + public Struct answer(InvocationOnMock invocation) { + currentRowIndex++; + return Struct.newBuilder().set("ID").to((long) currentRowIndex).build(); + } + }); + + try (AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { + when(delegate.initiateStreaming(any(AsyncResultSet.StreamMessageListener.class))) + .thenAnswer( + answer -> { + AsyncResultSet.StreamMessageListener listener = answer.getArgument(0); + // Deliver one chunk containing a resume token representing all rows + listener.onStreamMessage( + PartialResultSet.newBuilder() + .setResumeToken(ByteString.copyFromUtf8("resume-token")) + .build(), + false); + return true; + }); + + final List receivedRowIds = new ArrayList<>(); + ApiFuture callbackFuture = + asyncResultSet.setCallback( + callbackExecutor, + resultSet -> { + CursorState cursorState; + while ((cursorState = resultSet.tryNext()) == CursorState.OK) { + receivedRowIds.add(resultSet.getLong("ID")); + } + if (cursorState == CursorState.DONE) { + return CallbackResponse.DONE; + } + return CallbackResponse.CONTINUE; + }); + + assertNull(callbackFuture.get(5, TimeUnit.SECONDS)); + assertEquals(rowCount, receivedRowIds.size()); + } finally { + callbackExecutor.shutdown(); + } + } + + @Test + public void testRapidConcurrentPauseResume() throws Exception { + ExecutorService callbackExecutor = Executors.newSingleThreadExecutor(); + ExecutorService resumeExecutor = Executors.newFixedThreadPool(4); + final int rowCount = 100; + ResultSet delegate = mock(ResultSet.class); + when(delegate.next()) + .thenAnswer( + new Answer() { + int currentRowIndex = 0; + + @Override + public Boolean answer(InvocationOnMock invocation) { + currentRowIndex++; + return currentRowIndex <= rowCount; + } + }); + when(delegate.getCurrentRowAsStruct()) + .thenAnswer( + new Answer() { + int currentRowIndex = 0; + + @Override + public Struct answer(InvocationOnMock invocation) { + currentRowIndex++; + return Struct.newBuilder().set("ID").to((long) currentRowIndex).build(); + } + }); + + final List receivedRowIds = Collections.synchronizedList(new ArrayList<>()); + final Random random = new Random(); + + try (AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { + ApiFuture callbackFuture = + asyncResultSet.setCallback( + callbackExecutor, + resultSet -> { + CursorState cursorState = resultSet.tryNext(); + if (cursorState == CursorState.OK) { + receivedRowIds.add(resultSet.getLong("ID")); + return random.nextBoolean() ? CallbackResponse.PAUSE : CallbackResponse.CONTINUE; + } else if (cursorState == CursorState.DONE) { + return CallbackResponse.DONE; + } + return CallbackResponse.CONTINUE; + }); + + final AtomicBoolean testRunning = new AtomicBoolean(true); + for (int i = 0; i < 4; i++) { + resumeExecutor.execute( + () -> { + while (testRunning.get()) { + asyncResultSet.resume(); + Thread.yield(); + } + }); + } + + assertNull(callbackFuture.get(10, TimeUnit.SECONDS)); + testRunning.set(false); + + assertEquals(rowCount, receivedRowIds.size()); + for (int i = 0; i < rowCount; i++) { + assertEquals((long) (i + 1), (long) receivedRowIds.get(i)); + } + } finally { + callbackExecutor.shutdown(); + resumeExecutor.shutdown(); + } + } + + @Test + public void testErrorAfterBufferedRowsDelivered() throws Exception { + ExecutorService callbackExecutor = Executors.newSingleThreadExecutor(); + ResultSet delegate = mock(ResultSet.class); + when(delegate.next()) + .thenReturn(true) + .thenReturn(true) + .thenReturn(true) + .thenThrow( + SpannerExceptionFactory.newSpannerException( + ErrorCode.UNAVAILABLE, "temporary network glitch")); + when(delegate.getCurrentRowAsStruct()).thenReturn(mock(Struct.class)); + + final List receivedStates = new ArrayList<>(); + final AtomicInteger rowCount = new AtomicInteger(); + final SettableApiFuture exceptionFuture = SettableApiFuture.create(); + + try (AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { + ApiFuture callbackFuture = + asyncResultSet.setCallback( + callbackExecutor, + resultSet -> { + try { + CursorState cursorState = resultSet.tryNext(); + receivedStates.add(cursorState); + if (cursorState == CursorState.OK) { + rowCount.incrementAndGet(); + return CallbackResponse.CONTINUE; + } + return CallbackResponse.DONE; + } catch (SpannerException e) { + exceptionFuture.set(e); + return CallbackResponse.DONE; + } + }); + + SpannerException callbackException = exceptionFuture.get(5, TimeUnit.SECONDS); + assertEquals(ErrorCode.UNAVAILABLE, callbackException.getErrorCode()); + assertTrue(callbackException.getMessage().contains("temporary network glitch")); + assertEquals(3, rowCount.get()); + + ExecutionException executionException = + assertThrows(ExecutionException.class, () -> callbackFuture.get(5, TimeUnit.SECONDS)); + assertTrue(executionException.getCause() instanceof SpannerException); + assertEquals( + ErrorCode.UNAVAILABLE, ((SpannerException) executionException.getCause()).getErrorCode()); + } finally { + callbackExecutor.shutdown(); + } + } + + @Test + public void testProducerYieldsThreadWhenWaitingForStreamData() throws Exception { + ExecutorService callbackExecutor = Executors.newSingleThreadExecutor(); + GrpcResultSet delegate = mock(GrpcResultSet.class); + AtomicReference listenerRef = new AtomicReference<>(); + CountDownLatch listenerRegistered = new CountDownLatch(1); + when(delegate.initiateStreaming(any(AsyncResultSet.StreamMessageListener.class))) + .thenAnswer( + answer -> { + listenerRef.set(answer.getArgument(0)); + listenerRegistered.countDown(); + return true; + }); + final AtomicInteger nextCallCount = new AtomicInteger(); + when(delegate.next()) + .thenAnswer( + invocation -> { + int count = nextCallCount.incrementAndGet(); + return count <= 2; + }); + when(delegate.getCurrentRowAsStruct()) + .thenAnswer( + invocation -> Struct.newBuilder().set("ID").to((long) nextCallCount.get()).build()); + + // Initially isDataAvailable is true for the first row, then false until chunk 2 arrives + AtomicBoolean secondChunkAvailable = new AtomicBoolean(false); + when(delegate.isDataAvailable()) + .thenAnswer(invocation -> nextCallCount.get() == 0 || secondChunkAvailable.get()); + + try (AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { + List receivedRows = new ArrayList<>(); + CountDownLatch firstRowReceived = new CountDownLatch(1); + + ApiFuture callbackFuture = + asyncResultSet.setCallback( + callbackExecutor, + resultSet -> { + CursorState state; + while ((state = resultSet.tryNext()) == CursorState.OK) { + receivedRows.add(resultSet.getLong("ID")); + firstRowReceived.countDown(); + } + if (state == CursorState.DONE) { + return CallbackResponse.DONE; + } + return CallbackResponse.CONTINUE; + }); + + // Wait for the streaming listener to be registered asynchronously + assertTrue(listenerRegistered.await(5, TimeUnit.SECONDS)); + + // Deliver first message to trigger the producer + listenerRef + .get() + .onStreamMessage( + PartialResultSet.newBuilder() + .setResumeToken(ByteString.copyFromUtf8("chunk-1")) + .build(), + false); + assertTrue(firstRowReceived.await(5, TimeUnit.SECONDS)); + assertEquals(1, receivedRows.size()); + + // While stream is waiting for chunk 2, delegate.isDataAvailable() is false. + // Submit another task to simpleProvider (which is a 1-thread pool!). + // If the producer did not yield, the 1-thread pool would be blocked and this task could NOT + // execute! + ScheduledExecutorService producerExecutor = simpleProvider.getExecutor(); + Future pingTask = producerExecutor.submit(() -> "pong"); + assertEquals("pong", pingTask.get(2, TimeUnit.SECONDS)); + + // Now make chunk 2 available and notify listener + secondChunkAvailable.set(true); + listenerRef + .get() + .onStreamMessage( + PartialResultSet.newBuilder() + .setResumeToken(ByteString.copyFromUtf8("chunk-2")) + .build(), + false); + + assertNull(callbackFuture.get(5, TimeUnit.SECONDS)); + assertEquals(2, receivedRows.size()); + assertEquals(1L, (long) receivedRows.get(0)); + assertEquals(2L, (long) receivedRows.get(1)); + } finally { + callbackExecutor.shutdown(); + } + } + + @Test + public void testStreamInitializationFailurePropagatesToCallbackAndFuture() throws Exception { + ExecutorService callbackExecutor = Executors.newSingleThreadExecutor(); + StreamingResultSet delegate = mock(StreamingResultSet.class); + when(delegate.initiateStreaming(any(AsyncResultSet.StreamMessageListener.class))) + .thenThrow( + SpannerExceptionFactory.newSpannerException( + ErrorCode.INVALID_ARGUMENT, "Invalid query syntax error")); + + try (AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { + final SettableApiFuture exceptionFromCallback = SettableApiFuture.create(); + + ApiFuture callbackFuture = + asyncResultSet.setCallback( + callbackExecutor, + resultSet -> { + try { + resultSet.tryNext(); + } catch (SpannerException exception) { + exceptionFromCallback.set(exception); + } + return CallbackResponse.DONE; + }); + + SpannerException callbackException = exceptionFromCallback.get(5, TimeUnit.SECONDS); + assertEquals(ErrorCode.INVALID_ARGUMENT, callbackException.getErrorCode()); + assertTrue(callbackException.getMessage().contains("Invalid query syntax error")); + + ExecutionException executionException = + assertThrows(ExecutionException.class, () -> callbackFuture.get(5, TimeUnit.SECONDS)); + assertTrue(executionException.getCause() instanceof SpannerException); + assertEquals( + ErrorCode.INVALID_ARGUMENT, + ((SpannerException) executionException.getCause()).getErrorCode()); + } finally { + callbackExecutor.shutdown(); + } + } + + @Test + public void testCancelDoesNotRunListenersUntilCallbackFinishes() throws Exception { + ExecutorService callbackExecutor = Executors.newSingleThreadExecutor(); + ResultSet delegate = mock(ResultSet.class); + when(delegate.next()).thenReturn(true, false); + when(delegate.getCurrentRowAsStruct()).thenReturn(Struct.newBuilder().set("ID").to(1L).build()); + + final CountDownLatch callbackStartedLatch = new CountDownLatch(1); + final CountDownLatch allowCallbackToFinishLatch = new CountDownLatch(1); + final AtomicBoolean listenerRanWhileCallbackExecuting = new AtomicBoolean(false); + final AtomicBoolean listenerRan = new AtomicBoolean(false); + final AtomicBoolean callbackFinished = new AtomicBoolean(false); + + try (AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { + asyncResultSet.addListener( + () -> { + listenerRan.set(true); + if (!callbackFinished.get()) { + listenerRanWhileCallbackExecuting.set(true); + } + }); + + ApiFuture callbackFuture = + asyncResultSet.setCallback( + callbackExecutor, + resultSet -> { + CursorState state = resultSet.tryNext(); + if (state == CursorState.OK) { + callbackStartedLatch.countDown(); + try { + assertTrue(allowCallbackToFinishLatch.await(5, TimeUnit.SECONDS)); + } catch (InterruptedException interruptedException) { + Thread.currentThread().interrupt(); + } + callbackFinished.set(true); + return CallbackResponse.CONTINUE; + } + return CallbackResponse.DONE; + }); + + // Wait until the callback is actively executing inside cursorReady + assertTrue(callbackStartedLatch.await(5, TimeUnit.SECONDS)); + + // Cancel while the callback is still blocked inside cursorReady + asyncResultSet.cancel(); + + // Ensure that listeners did NOT run while the callback was still executing + assertFalse(listenerRan.get()); + assertFalse(listenerRanWhileCallbackExecuting.get()); + + // Release callback so it can complete + allowCallbackToFinishLatch.countDown(); + + // Now verify that the callback future completes with CANCELLED + ExecutionException executionException = + assertThrows(ExecutionException.class, () -> callbackFuture.get(5, TimeUnit.SECONDS)); + assertTrue(executionException.getCause() instanceof SpannerException); + assertEquals( + ErrorCode.CANCELLED, ((SpannerException) executionException.getCause()).getErrorCode()); + + // And the listener has run AFTER callback completed + assertTrue(listenerRan.get()); + assertFalse(listenerRanWhileCallbackExecuting.get()); + } finally { + callbackExecutor.shutdown(); + } + } + + @Test + public void testGrpcStreamIteratorDoneSetOnCompleted() { + GrpcStreamIterator streamIterator = new GrpcStreamIterator(false, 4, false); + assertFalse(streamIterator.isDataAvailable()); + + streamIterator.consumer().onCompleted(); + assertTrue(streamIterator.isDataAvailable()); + assertFalse(streamIterator.hasNext()); + assertTrue(streamIterator.isDataAvailable()); + } + + @Test + public void testBufferClearedOnCancellation() throws Exception { + ExecutorService callbackExecutor = Executors.newSingleThreadExecutor(); + ResultSet delegate = mock(ResultSet.class); + when(delegate.next()).thenReturn(true, true, true, false); + when(delegate.getCurrentRowAsStruct()) + .thenReturn( + Struct.newBuilder().set("ID").to(1L).build(), + Struct.newBuilder().set("ID").to(2L).build(), + Struct.newBuilder().set("ID").to(3L).build()); + + final CountDownLatch callbackStartedLatch = new CountDownLatch(1); + final CountDownLatch allowCancellationLatch = new CountDownLatch(1); + + try (AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { + ApiFuture callbackFuture = + asyncResultSet.setCallback( + callbackExecutor, + resultSet -> { + CursorState state = resultSet.tryNext(); + if (state == CursorState.OK) { + callbackStartedLatch.countDown(); + try { + assertTrue(allowCancellationLatch.await(5, TimeUnit.SECONDS)); + } catch (InterruptedException interruptedException) { + Thread.currentThread().interrupt(); + } + return CallbackResponse.CONTINUE; + } + return CallbackResponse.DONE; + }); + + assertTrue(callbackStartedLatch.await(5, TimeUnit.SECONDS)); + asyncResultSet.cancel(); + allowCancellationLatch.countDown(); + + ExecutionException executionException = + assertThrows(ExecutionException.class, () -> callbackFuture.get(5, TimeUnit.SECONDS)); + assertTrue(executionException.getCause() instanceof SpannerException); + assertEquals( + ErrorCode.CANCELLED, ((SpannerException) executionException.getCause()).getErrorCode()); + assertEquals(0, asyncResultSet.getBufferSize()); + } finally { + callbackExecutor.shutdown(); + } + } + + @Test + public void testForwardingResultSetIsDataAvailable() { + GrpcResultSet grpcDelegate = mock(GrpcResultSet.class); + when(grpcDelegate.isDataAvailable()).thenReturn(false, true); + + ForwardingResultSet innerForwarding = new ForwardingResultSet(grpcDelegate); + ForwardingResultSet outerForwarding = new ForwardingResultSet(innerForwarding); + + assertFalse(outerForwarding.isDataAvailable()); + assertTrue(outerForwarding.isDataAvailable()); + } + + @Test + public void testProducerStopsImmediatelyWhenCallbackThrowsUncheckedException() throws Exception { + ExecutorService callbackExecutor = Executors.newSingleThreadExecutor(); + ResultSet delegate = mock(ResultSet.class); + when(delegate.next()).thenReturn(true); + when(delegate.getCurrentRowAsStruct()).thenReturn(mock(Struct.class)); + try (AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { + ApiFuture future = + asyncResultSet.setCallback( + callbackExecutor, + resultSet -> { + CursorState state = resultSet.tryNext(); + if (state == CursorState.OK) { + throw new RuntimeException("User callback error"); + } + return CallbackResponse.CONTINUE; + }); + + ExecutionException executionException = + assertThrows(ExecutionException.class, () -> future.get(5, TimeUnit.SECONDS)); + assertTrue(executionException.getCause() instanceof SpannerException); + assertTrue(executionException.getCause().getMessage().contains("User callback error")); + assertEquals(0, asyncResultSet.getBufferSize()); + } finally { + callbackExecutor.shutdown(); + } + } + + @Test + public void testRejectedExecutionExceptionOnProducerService() throws Exception { + ScheduledExecutorService failingExecutor = mock(ScheduledExecutorService.class); + Mockito.doThrow(new RejectedExecutionException("Producer rejected")) + .when(failingExecutor) + .execute(any(Runnable.class)); + ExecutorProvider failingProvider = + new ExecutorProvider() { + @Override + public boolean shouldAutoClose() { + return false; + } + + @Override + public ScheduledExecutorService getExecutor() { + return failingExecutor; + } + }; + + ExecutorService callbackExecutor = Executors.newSingleThreadExecutor(); + try (AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl( + failingProvider, mock(ResultSet.class), AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { + ApiFuture future = + asyncResultSet.setCallback(callbackExecutor, resultSet -> CallbackResponse.DONE); + + ExecutionException executionException = + assertThrows(ExecutionException.class, () -> future.get(5, TimeUnit.SECONDS)); + assertTrue(executionException.getCause() instanceof SpannerException); + assertTrue(executionException.getCause().getMessage().contains("Producer rejected")); + } finally { + callbackExecutor.shutdown(); + } + } + + @Test + public void testRejectedExecutionExceptionOnCallbackExecutor() { + Executor failingCallbackExecutor = + command -> { + throw new RejectedExecutionException("Callback rejected"); + }; + ResultSet delegate = mock(ResultSet.class); + when(delegate.next()).thenReturn(true, false); + when(delegate.getCurrentRowAsStruct()).thenReturn(mock(Struct.class)); + + try (AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { + ApiFuture future = + asyncResultSet.setCallback(failingCallbackExecutor, resultSet -> CallbackResponse.DONE); + + ExecutionException executionException = + assertThrows(ExecutionException.class, () -> future.get(5, TimeUnit.SECONDS)); + assertTrue(executionException.getCause() instanceof SpannerException); + assertTrue(executionException.getCause().getMessage().contains("Callback rejected")); + } + } + + @Test + public void testOnStartOrClose() throws Exception { + ResultSet delegate = mock(ResultSet.class); + when(delegate.next()).thenReturn(true, false); + when(delegate.getCurrentRowAsStruct()).thenReturn(mock(Struct.class)); + AtomicInteger asyncNotifications = new AtomicInteger(); + + try (AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl( + simpleProvider, + () -> delegate, + AsyncResultSetImpl.DEFAULT_BUFFER_SIZE, + asyncNotifications::incrementAndGet)) { + assertEquals(0, asyncNotifications.get()); + ExecutorService executor = Executors.newSingleThreadExecutor(); + try { + ApiFuture future = + asyncResultSet.setCallback( + executor, + resultSet -> { + resultSet.tryNext(); + return CallbackResponse.DONE; + }); + future.get(5, TimeUnit.SECONDS); + assertEquals(1, asyncNotifications.get()); + } finally { + executor.shutdown(); + } + } + assertEquals(1, asyncNotifications.get()); + + ResultSet delegateSync = mock(ResultSet.class); + when(delegateSync.next()).thenReturn(true, false); + when(delegateSync.getCurrentRowAsStruct()).thenReturn(mock(Struct.class)); + AtomicInteger syncNotifications = new AtomicInteger(); + try (AsyncResultSetImpl asyncResultSetSync = + new AsyncResultSetImpl( + simpleProvider, + () -> delegateSync, + AsyncResultSetImpl.DEFAULT_BUFFER_SIZE, + syncNotifications::incrementAndGet)) { + assertEquals(0, syncNotifications.get()); + assertTrue(asyncResultSetSync.next()); + assertEquals(1, syncNotifications.get()); + } + assertEquals(1, syncNotifications.get()); + } + + @Test + public void testListenerPreconditionsAndRemoval() throws Exception { + ResultSet delegate = mock(ResultSet.class); + when(delegate.next()).thenReturn(false); + + AtomicBoolean removedListenerCalled = new AtomicBoolean(false); + AtomicBoolean activeListenerCalled = new AtomicBoolean(false); + Runnable removedListener = () -> removedListenerCalled.set(true); + Runnable activeListener = () -> activeListenerCalled.set(true); + + ExecutorService executor = Executors.newSingleThreadExecutor(); + try (AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { + asyncResultSet.addListener(removedListener); + asyncResultSet.addListener(activeListener); + asyncResultSet.removeListener(removedListener); + + ApiFuture future = + asyncResultSet.setCallback( + executor, + resultSet -> { + resultSet.tryNext(); + return CallbackResponse.DONE; + }); + + assertThrows(IllegalStateException.class, () -> asyncResultSet.addListener(() -> {})); + assertThrows( + IllegalStateException.class, () -> asyncResultSet.removeListener(activeListener)); + + future.get(5, TimeUnit.SECONDS); + assertFalse(removedListenerCalled.get()); + assertTrue(activeListenerCalled.get()); + } finally { + executor.shutdown(); + } + } + + @Test + public void testListenerExceptionDoesNotPreventSubsequentListenersOrCompletion() + throws Exception { + ResultSet delegate = mock(ResultSet.class); + when(delegate.next()).thenReturn(false); + + AtomicBoolean secondListenerCalled = new AtomicBoolean(false); + ExecutorService executor = Executors.newSingleThreadExecutor(); + try (AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { + asyncResultSet.addListener( + () -> { + throw new RuntimeException("Listener failure"); + }); + asyncResultSet.addListener(() -> secondListenerCalled.set(true)); + + ApiFuture future = + asyncResultSet.setCallback( + executor, + resultSet -> { + resultSet.tryNext(); + return CallbackResponse.DONE; + }); + + assertNull(future.get(5, TimeUnit.SECONDS)); + assertTrue(secondListenerCalled.get()); + } finally { + executor.shutdown(); + } + } + + @Test + public void testGetStatsAndGetMetadata() { + ResultSet delegate = mock(ResultSet.class); + ResultSetStats expectedStats = ResultSetStats.getDefaultInstance(); + ResultSetMetadata expectedMetadata = ResultSetMetadata.getDefaultInstance(); + when(delegate.getStats()).thenReturn(expectedStats); + when(delegate.getMetadata()).thenReturn(expectedMetadata); + + try (AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { + assertEquals(expectedStats, asyncResultSet.getStats()); + assertEquals(expectedMetadata, asyncResultSet.getMetadata()); + } + } + + @Test + public void testSyncModeOperations() { + ResultSet delegate = mock(ResultSet.class); + Struct mockRow = mock(Struct.class); + when(delegate.next()).thenReturn(true, false); + when(delegate.getCurrentRowAsStruct()).thenReturn(mockRow); + + try (AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { + assertTrue(asyncResultSet.next()); + assertEquals(mockRow, asyncResultSet.getCurrentRowAsStruct()); + + assertThrows( + IllegalStateException.class, + () -> + asyncResultSet.setCallback( + Executors.newSingleThreadExecutor(), resultSet -> CallbackResponse.DONE)); + + assertFalse(asyncResultSet.next()); + assertNull(asyncResultSet.getCurrentRowAsStruct()); + + ResultSet anotherDelegate = mock(ResultSet.class); + try (AsyncResultSetImpl asyncResultSet2 = + new AsyncResultSetImpl( + simpleProvider, anotherDelegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { + asyncResultSet2.setCallback( + Executors.newSingleThreadExecutor(), resultSet -> CallbackResponse.DONE); + assertThrows(IllegalStateException.class, asyncResultSet2::next); + } + + asyncResultSet.close(); + Mockito.verify(delegate, times(1)).close(); + assertThrows(IllegalStateException.class, asyncResultSet::getCurrentRowAsStruct); + } + } + + @Test + public void testExecutorProviderShouldAutoClose() throws Exception { + ScheduledExecutorService realService = Executors.newSingleThreadScheduledExecutor(); + ExecutorProvider autoCloseProvider = + new ExecutorProvider() { + @Override + public boolean shouldAutoClose() { + return true; + } + + @Override + public ScheduledExecutorService getExecutor() { + return realService; + } + }; + + ResultSet delegate = mock(ResultSet.class); + when(delegate.next()).thenReturn(false); + + ExecutorService callbackExecutor = Executors.newSingleThreadExecutor(); + try (AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl( + autoCloseProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { + ApiFuture future = + asyncResultSet.setCallback( + callbackExecutor, + resultSet -> { + resultSet.tryNext(); + return CallbackResponse.DONE; + }); + + future.get(5, TimeUnit.SECONDS); + assertTrue(realService.isShutdown()); + } finally { + callbackExecutor.shutdown(); + realService.shutdown(); + } + } + + @Test + public void testDefaultIsDataAvailable() { + StreamingResultSet defaultStreamingResultSet = + mock(StreamingResultSet.class, Mockito.CALLS_REAL_METHODS); + assertTrue(defaultStreamingResultSet.isDataAvailable()); + + AbstractResultSet.CloseableIterator defaultIterator = + mock(AbstractResultSet.CloseableIterator.class, Mockito.CALLS_REAL_METHODS); + assertTrue(defaultIterator.isDataAvailable()); + } + + @Test + public void testDuplicateCloseIsIdempotent() { + ResultSet delegate = mock(ResultSet.class); + AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE); + asyncResultSet.close(); + verify(delegate, times(1)).close(); + // Second close is a no-op + asyncResultSet.close(); + verify(delegate, times(1)).close(); + } + + @Test + public void testCallbackThrowsCancelledException() throws Exception { + ExecutorService callbackExecutor = Executors.newSingleThreadExecutor(); + ResultSet delegate = mock(ResultSet.class); + when(delegate.next()).thenReturn(true); + when(delegate.getCurrentRowAsStruct()).thenReturn(mock(Struct.class)); + try (AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { + ApiFuture future = + asyncResultSet.setCallback( + callbackExecutor, + resultSet -> { + asyncResultSet.cancel(); + // Calling tryNext() will throw CANCELLED_EXCEPTION out of cursorReady + resultSet.tryNext(); + return CallbackResponse.CONTINUE; + }); + SpannerException exception = assertThrows(SpannerException.class, () -> get(future)); + assertEquals(ErrorCode.CANCELLED, exception.getErrorCode()); + } finally { + callbackExecutor.shutdown(); + } + } + + @Test + public void testSetCallbackInitiateStreamingRejection() { + ScheduledExecutorService rejectingExecutor = mock(ScheduledExecutorService.class); + Mockito.doThrow(new RejectedExecutionException("rejected")) + .when(rejectingExecutor) + .execute(any(Runnable.class)); + ExecutorProvider rejectingProvider = + new ExecutorProvider() { + @Override + public boolean shouldAutoClose() { + return false; + } + + @Override + public ScheduledExecutorService getExecutor() { + return rejectingExecutor; + } + }; + + ResultSet delegate = mock(ResultSet.class); + try (AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl( + rejectingProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { + ApiFuture future = + asyncResultSet.setCallback( + Executors.newSingleThreadExecutor(), resultSet -> CallbackResponse.DONE); + SpannerException exception = assertThrows(SpannerException.class, () -> get(future)); + assertTrue(exception.getMessage().contains("rejected")); + } + } + + @Test + public void testIsDataAvailableExceptionHandling() throws Exception { + ExecutorService callbackExecutor = Executors.newSingleThreadExecutor(); + StreamingResultSet delegate = mock(StreamingResultSet.class); + when(delegate.isDataAvailable()).thenThrow(new RuntimeException("isDataAvailable error")); + when(delegate.next()).thenReturn(false); + + try (AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { + ApiFuture future = + asyncResultSet.setCallback( + callbackExecutor, + resultSet -> { + assertEquals(CursorState.DONE, resultSet.tryNext()); + return CallbackResponse.DONE; + }); + assertNull(future.get(5, TimeUnit.SECONDS)); + } finally { + callbackExecutor.shutdown(); + } + } + + @Test + public void testCallbackProcessingOneRowAtATimeReceivesDone() throws Exception { + ExecutorService callbackExecutor = Executors.newSingleThreadExecutor(); + ResultSet delegate = mock(ResultSet.class); + when(delegate.next()).thenReturn(true, true, false); + when(delegate.getCurrentRowAsStruct()) + .thenReturn( + Struct.newBuilder().set("v").to("a").build(), + Struct.newBuilder().set("v").to("b").build()); + + AtomicInteger rowsSeen = new AtomicInteger(); + AtomicBoolean doneReceived = new AtomicBoolean(); + try (AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { + ApiFuture future = + asyncResultSet.setCallback( + callbackExecutor, + resultSet -> { + CursorState state = resultSet.tryNext(); + if (state == CursorState.OK) { + rowsSeen.incrementAndGet(); + return CallbackResponse.CONTINUE; + } + if (state == CursorState.DONE) { + doneReceived.set(true); + return CallbackResponse.DONE; + } + return CallbackResponse.CONTINUE; + }); + assertNull(future.get(5, TimeUnit.SECONDS)); + assertTrue(doneReceived.get()); + assertEquals(2, rowsSeen.get()); + } finally { + callbackExecutor.shutdown(); + } + } + + @Test + public void testHighVolumeStreamingWithoutBufferStarvation() throws Exception { + ExecutorService callbackExecutor = Executors.newSingleThreadExecutor(); + int totalRowCount = 50_000; + int bufferCapacity = 50; + + StreamingResultSet delegate = + new ForwardingResultSet(mock(ResultSet.class)) { + private int currentRow = 0; + private Struct currentStruct; + + @Override + public boolean next() { + if (currentRow < totalRowCount) { + currentRow++; + currentStruct = Struct.newBuilder().set("ID").to((long) currentRow).build(); + return true; + } + return false; + } + + @Override + public Struct getCurrentRowAsStruct() { + return currentStruct; + } + + @Override + public boolean isDataAvailable() { + return true; + } + + @Override + public boolean initiateStreaming(AsyncResultSet.StreamMessageListener listener) { + return true; + } + }; + + AtomicInteger rowsReceived = new AtomicInteger(); + AtomicBoolean doneReceived = new AtomicBoolean(); + + try (AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl(simpleProvider, delegate, bufferCapacity)) { + ApiFuture future = + asyncResultSet.setCallback( + callbackExecutor, + resultSet -> { + while (true) { + switch (resultSet.tryNext()) { + case OK: + int expectedId = rowsReceived.incrementAndGet(); + assertEquals(expectedId, resultSet.getLong("ID")); + break; + case NOT_READY: + return CallbackResponse.CONTINUE; + case DONE: + doneReceived.set(true); + return CallbackResponse.DONE; + } + } + }); + + assertNull(future.get(30, TimeUnit.SECONDS)); + assertTrue(doneReceived.get()); + assertEquals(totalRowCount, rowsReceived.get()); + } finally { + callbackExecutor.shutdown(); + assertTrue(callbackExecutor.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + @Test + public void testHighVolumeStreamingWithAsyncChunkArrivals() throws Exception { + ExecutorService callbackExecutor = Executors.newSingleThreadExecutor(); + ExecutorService feederExecutor = Executors.newSingleThreadExecutor(); + int totalRowCount = 10_000; + int chunkSize = 100; + int bufferCapacity = 50; + + AtomicInteger availableRows = new AtomicInteger(); + AtomicReference registeredListener = + new AtomicReference<>(); + CountDownLatch streamingInitiated = new CountDownLatch(1); + + StreamingResultSet delegate = + new ForwardingResultSet(mock(ResultSet.class)) { + private int currentRow = 0; + private Struct currentStruct; + + @Override + public boolean next() { + if (currentRow < totalRowCount) { + assertTrue( + "next() should not be called when data is not available", + currentRow < availableRows.get()); + currentRow++; + currentStruct = Struct.newBuilder().set("ID").to((long) currentRow).build(); + return true; + } + return false; + } + + @Override + public Struct getCurrentRowAsStruct() { + return currentStruct; + } + + @Override + public boolean isDataAvailable() { + return currentRow < availableRows.get() || currentRow >= totalRowCount; + } + + @Override + public boolean initiateStreaming(AsyncResultSet.StreamMessageListener listener) { + registeredListener.set(listener); + streamingInitiated.countDown(); + return true; + } + }; + + Semaphore chunkRequested = new Semaphore(1); + AtomicInteger rowsReceived = new AtomicInteger(); + AtomicBoolean doneReceived = new AtomicBoolean(); + + try (AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl(simpleProvider, delegate, bufferCapacity)) { + ApiFuture future = + asyncResultSet.setCallback( + callbackExecutor, + resultSet -> { + while (true) { + switch (resultSet.tryNext()) { + case OK: + int expectedId = rowsReceived.incrementAndGet(); + assertEquals(expectedId, resultSet.getLong("ID")); + break; + case NOT_READY: + if (availableRows.get() < totalRowCount + && chunkRequested.availablePermits() == 0) { + chunkRequested.release(); + } + return CallbackResponse.CONTINUE; + case DONE: + doneReceived.set(true); + return CallbackResponse.DONE; + } + } + }); + + assertTrue(streamingInitiated.await(5, TimeUnit.SECONDS)); + feederExecutor.execute( + () -> { + try { + while (availableRows.get() < totalRowCount) { + chunkRequested.acquire(); + availableRows.addAndGet(chunkSize); + AsyncResultSet.StreamMessageListener listener = registeredListener.get(); + if (listener != null) { + listener.onStreamMessage(PartialResultSet.getDefaultInstance(), false); + } + } + } catch (InterruptedException ignored) { + } + }); + + assertNull(future.get(30, TimeUnit.SECONDS)); + assertTrue(doneReceived.get()); + assertEquals(totalRowCount, rowsReceived.get()); + assertEquals(totalRowCount, availableRows.get()); + } finally { + feederExecutor.shutdownNow(); + callbackExecutor.shutdown(); + assertTrue(feederExecutor.awaitTermination(5, TimeUnit.SECONDS)); + assertTrue(callbackExecutor.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + @Test + public void testSetCallbackTwiceThrowsException() { + AsyncResultSetImpl resultSet = + new AsyncResultSetImpl( + mockedProvider, mock(ResultSet.class), AsyncResultSetImpl.DEFAULT_BUFFER_SIZE); + resultSet.setCallback(mock(Executor.class), mock(ReadyCallback.class)); + IllegalStateException exception = + assertThrows( + IllegalStateException.class, + () -> resultSet.setCallback(mock(Executor.class), mock(ReadyCallback.class))); + assertTrue( + "Expected message to contain 'callback may not be set multiple times', got: " + + exception.getMessage(), + exception.getMessage().contains("callback may not be set multiple times")); + } + + @Test + public void testCancelWithoutCallbackThrowsException() { + AsyncResultSetImpl resultSet = + new AsyncResultSetImpl( + mockedProvider, mock(ResultSet.class), AsyncResultSetImpl.DEFAULT_BUFFER_SIZE); + IllegalStateException exception = assertThrows(IllegalStateException.class, resultSet::cancel); + assertTrue( + "Expected message to contain 'cannot cancel a result set without a callback', got: " + + exception.getMessage(), + exception.getMessage().contains("cannot cancel a result set without a callback")); + } + + @Test + public void testResumeWithoutCallbackThrowsException() { + AsyncResultSetImpl resultSet = + new AsyncResultSetImpl( + mockedProvider, mock(ResultSet.class), AsyncResultSetImpl.DEFAULT_BUFFER_SIZE); + IllegalStateException exception = assertThrows(IllegalStateException.class, resultSet::resume); + assertTrue( + "Expected message to contain 'cannot resume a result set without a callback', got: " + + exception.getMessage(), + exception.getMessage().contains("cannot resume a result set without a callback")); + } + + @Test + public void testCancelAfterCompletionIsNoOp() throws Exception { + ResultSet delegate = mock(ResultSet.class); + when(delegate.next()).thenReturn(false); + ExecutorService callbackExecutor = Executors.newSingleThreadExecutor(); + try (AsyncResultSetImpl resultSet = + new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { + ApiFuture future = + resultSet.setCallback( + callbackExecutor, + readyResultSet -> { + if (readyResultSet.tryNext() == CursorState.DONE) { + return CallbackResponse.DONE; + } + return CallbackResponse.CONTINUE; + }); + future.get(5, TimeUnit.SECONDS); + assertEquals(AsyncResultSetImpl.State.DONE, resultSet.getState()); + resultSet.cancel(); + assertEquals(AsyncResultSetImpl.State.DONE, resultSet.getState()); + } finally { + callbackExecutor.shutdown(); + assertTrue(callbackExecutor.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + @Test + public void testResumeAfterCompletionIsNoOp() throws Exception { + ResultSet delegate = mock(ResultSet.class); + when(delegate.next()).thenReturn(false); + ExecutorService callbackExecutor = Executors.newSingleThreadExecutor(); + try (AsyncResultSetImpl resultSet = + new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { + ApiFuture future = + resultSet.setCallback( + callbackExecutor, + readyResultSet -> { + if (readyResultSet.tryNext() == CursorState.DONE) { + return CallbackResponse.DONE; + } + return CallbackResponse.CONTINUE; + }); + future.get(5, TimeUnit.SECONDS); + assertEquals(AsyncResultSetImpl.State.DONE, resultSet.getState()); + resultSet.resume(); + assertEquals(AsyncResultSetImpl.State.DONE, resultSet.getState()); + } finally { + callbackExecutor.shutdown(); + assertTrue(callbackExecutor.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + @Test + public void testCloseDelegateResultSetIgnoresException() { + ResultSet delegate = mock(ResultSet.class); + doThrow(new RuntimeException("close failed")).when(delegate).close(); + AsyncResultSetImpl resultSet = + new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE); + resultSet.close(); + assertTrue(resultSet.isClosed()); + verify(delegate).close(); + } + + @Test + public void testResumeWhileConsumingDoesNotPreemptivelyUnpause() throws Exception { + int totalRowCount = 5; + ResultSet delegate = mock(ResultSet.class); + when(delegate.next()).thenReturn(true, true, true, true, true, false); + when(delegate.getCurrentRowAsStruct()).thenReturn(mock(Struct.class)); + + CountDownLatch callbackEnteredLatch = new CountDownLatch(1); + CountDownLatch resumeCalledLatch = new CountDownLatch(1); + AtomicInteger rowsReceived = new AtomicInteger(); + AtomicBoolean pausedOnce = new AtomicBoolean(); + + ExecutorService callbackExecutor = Executors.newSingleThreadExecutor(); + try (AsyncResultSetImpl resultSet = + new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { + ApiFuture future = + resultSet.setCallback( + callbackExecutor, + readyResultSet -> { + while (true) { + switch (readyResultSet.tryNext()) { + case OK: + int count = rowsReceived.incrementAndGet(); + if (count == 1 && pausedOnce.compareAndSet(false, true)) { + callbackEnteredLatch.countDown(); + try { + assertTrue(resumeCalledLatch.await(5, TimeUnit.SECONDS)); + } catch (InterruptedException interruptedException) { + Thread.currentThread().interrupt(); + throw new RuntimeException(interruptedException); + } + return CallbackResponse.PAUSE; + } + break; + case NOT_READY: + return CallbackResponse.CONTINUE; + case DONE: + return CallbackResponse.DONE; + } + } + }); + + assertTrue(callbackEnteredLatch.await(5, TimeUnit.SECONDS)); + assertEquals(AsyncResultSetImpl.State.CONSUMING, resultSet.getState()); + // Calling resume() while currently consuming should be a no-op and not pre-emptively cancel + // the upcoming PAUSE. + resultSet.resume(); + resumeCalledLatch.countDown(); + + // Drain the single-threaded callback executor so that CallbackRunnable completes returning + // PAUSE. + callbackExecutor.submit(() -> {}).get(5, TimeUnit.SECONDS); + + // Verify the result set actually transitioned to PAUSED and did not automatically resume. + assertEquals(AsyncResultSetImpl.State.PAUSED, resultSet.getState()); + assertEquals(1, rowsReceived.get()); + assertFalse(future.isDone()); + + // Now resume from the PAUSED state. + resultSet.resume(); + + assertNull(future.get(5, TimeUnit.SECONDS)); + assertEquals(totalRowCount, rowsReceived.get()); + assertEquals(AsyncResultSetImpl.State.DONE, resultSet.getState()); + } finally { + callbackExecutor.shutdown(); + assertTrue(callbackExecutor.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + @Test + public void testResumeWhileRunningDoesNotPreemptivelyUnpause() throws Exception { + int totalRowCount = 3; + ResultSet delegate = mock(ResultSet.class); + when(delegate.next()).thenReturn(true, true, true, false); + when(delegate.getCurrentRowAsStruct()).thenReturn(mock(Struct.class)); + + CountDownLatch pausedLatch = new CountDownLatch(1); + AtomicInteger rowsReceived = new AtomicInteger(); + + ExecutorService callbackExecutor = Executors.newSingleThreadExecutor(); + try (AsyncResultSetImpl resultSet = + new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { + ApiFuture future = + resultSet.setCallback( + callbackExecutor, + readyResultSet -> { + while (true) { + switch (readyResultSet.tryNext()) { + case OK: + int count = rowsReceived.incrementAndGet(); + if (count == 1) { + pausedLatch.countDown(); + return CallbackResponse.PAUSE; + } + break; + case NOT_READY: + return CallbackResponse.CONTINUE; + case DONE: + return CallbackResponse.DONE; + } + } + }); + + // Resume called while in RUNNING state (before first pause) + resultSet.resume(); + + // Wait until callback processed first row and returned PAUSE + assertTrue(pausedLatch.await(5, TimeUnit.SECONDS)); + + // Drain callback executor + callbackExecutor.submit(() -> {}).get(5, TimeUnit.SECONDS); + + // Verify the result set transitioned to PAUSED and stayed paused + assertEquals(AsyncResultSetImpl.State.PAUSED, resultSet.getState()); + assertEquals(1, rowsReceived.get()); + assertFalse(future.isDone()); + + // Resume from PAUSED state + resultSet.resume(); + + assertNull(future.get(5, TimeUnit.SECONDS)); + assertEquals(totalRowCount, rowsReceived.get()); + assertEquals(AsyncResultSetImpl.State.DONE, resultSet.getState()); + } finally { + callbackExecutor.shutdown(); + assertTrue(callbackExecutor.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + @Test + public void testVirtualThreadPerTaskExecutor() throws Exception { + Assume.assumeTrue(isJava21OrHigher()); + + ExecutorService virtualThreadExecutor = + ThreadFactoryUtil.tryCreateVirtualThreadPerTaskExecutor("async-result-set-virtual-thread"); + assertNotNull(virtualThreadExecutor); + try { + ResultSet delegate = mock(ResultSet.class); + when(delegate.next()).thenReturn(true, true, true, false); + when(delegate.getCurrentRowAsStruct()).thenReturn(mock(Struct.class)); + + try (AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl( + simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { + AtomicInteger rowCounter = new AtomicInteger(0); + AtomicBoolean ranOnVirtualThread = new AtomicBoolean(false); + ApiFuture future = + asyncResultSet.setCallback( + virtualThreadExecutor, + resultSet -> { + if (isVirtual(Thread.currentThread())) { + ranOnVirtualThread.set(true); + } + while (true) { + switch (resultSet.tryNext()) { + case OK: + rowCounter.incrementAndGet(); + break; + case NOT_READY: + return CallbackResponse.CONTINUE; + case DONE: + return CallbackResponse.DONE; + } + } + }); + + future.get(5, TimeUnit.SECONDS); + assertEquals(3, rowCounter.get()); + assertTrue(ranOnVirtualThread.get()); + } + } finally { + virtualThreadExecutor.shutdown(); + assertTrue(virtualThreadExecutor.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + private static boolean isJava21OrHigher() { + String[] versionElements = System.getProperty("java.version").split("\\."); + int majorVersion = Integer.parseInt(versionElements[0]); + // Java 1.8 (Java 8) and lower used the format 1.8 etc. + // Java 9 and higher use the format 9.x + if (majorVersion == 1) { + majorVersion = Integer.parseInt(versionElements[1]); + } + return majorVersion >= 21; + } + + private static boolean isVirtual(Thread thread) { + if (thread == null) { + return false; + } + try { + java.lang.reflect.Method isVirtualMethod = Thread.class.getMethod("isVirtual"); + return (Boolean) isVirtualMethod.invoke(thread); + } catch (Exception ignore) { + return false; + } + } + + @Test + public void testCancelWhileProducerIsBlockedInNext() throws Exception { + CountDownLatch producerEnteredNext = new CountDownLatch(1); + CountDownLatch allowNextToComplete = new CountDownLatch(1); + ExecutorService callbackExecutor = Executors.newSingleThreadExecutor(); + + ResultSet delegate = mock(ResultSet.class); + when(delegate.next()) + .thenAnswer( + invocation -> { + producerEnteredNext.countDown(); + assertTrue(allowNextToComplete.await(5, TimeUnit.SECONDS)); + return true; + }); + when(delegate.getCurrentRowAsStruct()).thenReturn(mock(Struct.class)); + + try (AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { + ApiFuture callbackFuture = + asyncResultSet.setCallback( + callbackExecutor, + resultSet -> { + resultSet.tryNext(); + return CallbackResponse.CONTINUE; + }); + + // Wait until producer is actively executing delegate.next() + assertTrue(producerEnteredNext.await(5, TimeUnit.SECONDS)); + + // Cancel from another thread while producer is blocked + asyncResultSet.cancel(); + allowNextToComplete.countDown(); + + ExecutionException executionException = + assertThrows(ExecutionException.class, () -> callbackFuture.get(5, TimeUnit.SECONDS)); + assertTrue(executionException.getCause() instanceof SpannerException); + assertEquals( + ErrorCode.CANCELLED, ((SpannerException) executionException.getCause()).getErrorCode()); + verify(delegate, times(1)).close(); + } finally { + callbackExecutor.shutdown(); + assertTrue(callbackExecutor.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + @Test + public void testSingleRowSplitAcrossChunksDeliveredCleanly() throws Exception { + ExecutorService callbackExecutor = Executors.newSingleThreadExecutor(); + StreamingResultSet delegate = mock(StreamingResultSet.class); + + AtomicBoolean dataAvailable = new AtomicBoolean(false); + when(delegate.isDataAvailable()).thenAnswer(invocation -> dataAvailable.get()); + + when(delegate.next()).thenReturn(true, false); + Struct expectedRow = Struct.newBuilder().set("Value").to("chunked-payload").build(); + when(delegate.getCurrentRowAsStruct()).thenReturn(expectedRow); + + CountDownLatch initiateStreamingCalled = new CountDownLatch(1); + AtomicReference listenerRef = new AtomicReference<>(); + when(delegate.initiateStreaming(any(AsyncResultSet.StreamMessageListener.class))) + .thenAnswer( + invocation -> { + listenerRef.set(invocation.getArgument(0)); + initiateStreamingCalled.countDown(); + return true; + }); + + try (AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { + List receivedRows = new ArrayList<>(); + ApiFuture future = + asyncResultSet.setCallback( + callbackExecutor, + resultSet -> { + while (true) { + switch (resultSet.tryNext()) { + case OK: + receivedRows.add(resultSet.getCurrentRowAsStruct()); + break; + case NOT_READY: + return CallbackResponse.CONTINUE; + case DONE: + return CallbackResponse.DONE; + } + } + }); + + assertTrue(initiateStreamingCalled.await(5, TimeUnit.SECONDS)); + assertNotNull(listenerRef.get()); + + // First chunk arrives with partial data: isDataAvailable is false. + // The producer yields without producing any row. + dataAvailable.set(false); + listenerRef + .get() + .onStreamMessage(PartialResultSet.newBuilder().setChunkedValue(true).build(), false); + + // Second chunk arrives completing the row: isDataAvailable is true. + dataAvailable.set(true); + listenerRef + .get() + .onStreamMessage(PartialResultSet.newBuilder().setChunkedValue(false).build(), false); + + future.get(5, TimeUnit.SECONDS); + assertEquals(1, receivedRows.size()); + assertEquals("chunked-payload", receivedRows.get(0).getString("Value")); + } finally { + callbackExecutor.shutdown(); + assertTrue(callbackExecutor.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + @Test + public void testInitiateStreamingWithZeroRowsDeliversDoneImmediately() throws Exception { + ExecutorService callbackExecutor = Executors.newSingleThreadExecutor(); + StreamingResultSet delegate = mock(StreamingResultSet.class); + when(delegate.next()).thenReturn(false); + when(delegate.isDataAvailable()).thenReturn(true); + + CountDownLatch initiateStreamingCalled = new CountDownLatch(1); + AtomicReference listenerRef = new AtomicReference<>(); + when(delegate.initiateStreaming(any(AsyncResultSet.StreamMessageListener.class))) + .thenAnswer( + invocation -> { + listenerRef.set(invocation.getArgument(0)); + initiateStreamingCalled.countDown(); + return true; + }); + + try (AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { + AtomicBoolean receivedDone = new AtomicBoolean(false); + ApiFuture future = + asyncResultSet.setCallback( + callbackExecutor, + resultSet -> { + CursorState state = resultSet.tryNext(); + if (state == CursorState.DONE) { + receivedDone.set(true); + return CallbackResponse.DONE; + } + return CallbackResponse.CONTINUE; + }); + + assertTrue(initiateStreamingCalled.await(5, TimeUnit.SECONDS)); + assertNotNull(listenerRef.get()); + + // Signal end of stream + listenerRef.get().onStreamMessage(GrpcStreamIterator.END_OF_STREAM, false); + + future.get(5, TimeUnit.SECONDS); + assertTrue(receivedDone.get()); + assertEquals(0, asyncResultSet.getBufferSize()); + } finally { + callbackExecutor.shutdown(); + assertTrue(callbackExecutor.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + @Test + public void testDelegateClosedExactlyOnceOnNormalCompletionAndSubsequentClose() throws Exception { + ExecutorService callbackExecutor = Executors.newSingleThreadExecutor(); + ResultSet delegate = mock(ResultSet.class); + when(delegate.next()).thenReturn(true, false); + when(delegate.getCurrentRowAsStruct()).thenReturn(mock(Struct.class)); + + try (AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { + ApiFuture future = + asyncResultSet.setCallback( + callbackExecutor, + resultSet -> { + while (true) { + switch (resultSet.tryNext()) { + case OK: + break; + case NOT_READY: + return CallbackResponse.CONTINUE; + case DONE: + return CallbackResponse.DONE; + } + } + }); + + future.get(5, TimeUnit.SECONDS); + assertTrue(asyncResultSet.isDelegateClosed()); + verify(delegate, times(1)).close(); + + // Calling close() or cancel() subsequently should be a no-op and never close the delegate + // again + asyncResultSet.close(); + asyncResultSet.close(); + asyncResultSet.cancel(); + verify(delegate, times(1)).close(); + } finally { + callbackExecutor.shutdown(); + assertTrue(callbackExecutor.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + @Test + public void testDelegateClosedExactlyOnceOnCallbackException() throws Exception { + ExecutorService callbackExecutor = Executors.newSingleThreadExecutor(); + ResultSet delegate = mock(ResultSet.class); + when(delegate.next()).thenReturn(true, false); + when(delegate.getCurrentRowAsStruct()).thenReturn(mock(Struct.class)); + + try (AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { + ApiFuture future = + asyncResultSet.setCallback( + callbackExecutor, + resultSet -> { + throw new RuntimeException("test error"); + }); + + ExecutionException executionException = + assertThrows(ExecutionException.class, () -> future.get(5, TimeUnit.SECONDS)); + assertTrue(executionException.getCause() instanceof SpannerException); + assertTrue(asyncResultSet.isDelegateClosed()); + verify(delegate, times(1)).close(); + + asyncResultSet.close(); + verify(delegate, times(1)).close(); + } finally { + callbackExecutor.shutdown(); + assertTrue(callbackExecutor.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + @Test + public void testDelegateClosedExactlyOnceOnConcurrentCancelAndClose() throws Exception { + ExecutorService callbackExecutor = Executors.newFixedThreadPool(4); + ResultSet delegate = mock(ResultSet.class); + when(delegate.next()).thenReturn(true, true, false); + when(delegate.getCurrentRowAsStruct()).thenReturn(mock(Struct.class)); + + try (AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { + CountDownLatch ready = new CountDownLatch(1); + ApiFuture future = + asyncResultSet.setCallback( + callbackExecutor, + resultSet -> { + ready.countDown(); + return CallbackResponse.PAUSE; + }); + + assertTrue(ready.await(5, TimeUnit.SECONDS)); + + int threadCount = 8; + CountDownLatch startLatch = new CountDownLatch(1); + CountDownLatch doneLatch = new CountDownLatch(threadCount); + for (int i = 0; i < threadCount; i++) { + final int index = i; + callbackExecutor.execute( + () -> { + try { + startLatch.await(); + if (index % 2 == 0) { + asyncResultSet.cancel(); + } else { + asyncResultSet.close(); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } finally { + doneLatch.countDown(); + } + }); + } + startLatch.countDown(); + assertTrue(doneLatch.await(5, TimeUnit.SECONDS)); + + assertThrows(ExecutionException.class, () -> future.get(5, TimeUnit.SECONDS)); + assertTrue(asyncResultSet.isDelegateClosed()); + verify(delegate, times(1)).close(); + } finally { + callbackExecutor.shutdown(); + assertTrue(callbackExecutor.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + private static ResultSetMetadata createTestMetadata(Type.StructField... fields) { + return ResultSetMetadata.newBuilder() + .setRowType(Type.struct(fields).toProto().getStructType()) + .build(); + } + + private static Value stringProtoValue(String value) { + return Value.newBuilder().setStringValue(value).build(); + } + + private static TraceWrapper createTestTracer() { + SpannerOptions.resetActiveTracingFramework(); + SpannerOptions.enableOpenTelemetryTraces(); + return new TraceWrapper(Tracing.getTracer(), OpenTelemetry.noop().getTracer(""), false); + } + + private static AbstractResultSet.Listener createNoOpListener() { + return new AbstractResultSet.Listener() { + @Override + public void onTransactionMetadata(Transaction transaction, boolean shouldIncludeId) {} + + @Override + public SpannerException onError( + SpannerException e, boolean withBeginTransaction, boolean lastStatement) { + return e; + } + + @Override + public void onDone(boolean withBeginTransaction) {} + + @Override + public void onPrecommitToken(MultiplexedSessionPrecommitToken token) {} + }; + } + + @Test + public void testSingleThreadExecutorDoesNotDeadlockOnChunkWithoutResumeToken() throws Exception { + ExecutorService callbackExecutor = Executors.newSingleThreadExecutor(); + ExecutorProvider singleThreadProvider = + SpannerOptions.createAsyncExecutorProvider(1, 1L, TimeUnit.SECONDS); + ScheduledExecutorService producerExecutor = singleThreadProvider.getExecutor(); + ResumableStreamIterator resumableStreamIterator = + new ResumableStreamIterator( + 10, + "", + new OpenTelemetrySpan(mock(Span.class)), + createTestTracer(), + DefaultErrorHandler.INSTANCE, + SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetrySettings(), + SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetryableCodes(), + NoopRequestIdCreator.INSTANCE) { + @Override + AbstractResultSet.CloseableIterator startStream( + ByteString resumeToken, + AsyncResultSet.StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + GrpcStreamIterator stream = new GrpcStreamIterator(false, 10, false); + setStream(stream); + if (streamMessageListener != null) { + stream.registerListener(streamMessageListener); + } + stream.setCall(mock(StreamingCall.class), false); + // Deliver first chunk WITHOUT a resume token. + stream + .consumer() + .onPartialResultSet( + PartialResultSet.newBuilder() + .setMetadata(createTestMetadata(Type.StructField.of("col1", Type.string()))) + .addValues(stringProtoValue("row1")) + .build()); + // Schedule second chunk WITH a resume token on the same 1-thread producer executor. + // If ProduceRowsRunnable blocks on stream.take() when consuming chunk 1, this task will + // never execute. + producerExecutor.execute( + () -> { + stream + .consumer() + .onPartialResultSet( + PartialResultSet.newBuilder() + .setResumeToken(ByteString.copyFromUtf8("token1")) + .addValues(stringProtoValue("row2")) + .build()); + stream.consumer().onCompleted(); + }); + return stream; + } + }; + + GrpcResultSet grpcResultSet = new GrpcResultSet(resumableStreamIterator, createNoOpListener()); + try (AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl(singleThreadProvider, grpcResultSet, 10)) { + List rows = new ArrayList<>(); + ApiFuture future = + asyncResultSet.setCallback( + callbackExecutor, + resultSet -> { + while (true) { + switch (resultSet.tryNext()) { + case OK: + rows.add(resultSet.getString(0)); + break; + case NOT_READY: + return CallbackResponse.CONTINUE; + case DONE: + return CallbackResponse.DONE; + } + } + }); + + assertNull(future.get(1, TimeUnit.SECONDS)); + assertEquals(Arrays.asList("row1", "row2"), rows); + } finally { + producerExecutor.shutdownNow(); + callbackExecutor.shutdownNow(); + } + } + + @Test + public void testSingleThreadExecutorDoesNotDeadlockOnSplitRowAcrossChunks() throws Exception { + ExecutorService callbackExecutor = Executors.newSingleThreadExecutor(); + ExecutorProvider singleThreadProvider = + SpannerOptions.createAsyncExecutorProvider(1, 1L, TimeUnit.SECONDS); + ScheduledExecutorService producerExecutor = singleThreadProvider.getExecutor(); + ResumableStreamIterator resumableStreamIterator = + new ResumableStreamIterator( + 10, + "", + new OpenTelemetrySpan(mock(Span.class)), + createTestTracer(), + DefaultErrorHandler.INSTANCE, + SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetrySettings(), + SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetryableCodes(), + NoopRequestIdCreator.INSTANCE) { + @Override + AbstractResultSet.CloseableIterator startStream( + ByteString resumeToken, + AsyncResultSet.StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + GrpcStreamIterator stream = new GrpcStreamIterator(false, 10, false); + setStream(stream); + if (streamMessageListener != null) { + stream.registerListener(streamMessageListener); + } + stream.setCall(mock(StreamingCall.class), false); + // 2-column table: first chunk has Row 1 (2 values) + first column of Row 2 (1 value). + stream + .consumer() + .onPartialResultSet( + PartialResultSet.newBuilder() + .setMetadata( + createTestMetadata( + Type.StructField.of("col1", Type.string()), + Type.StructField.of("col2", Type.string()))) + .setResumeToken(ByteString.copyFromUtf8("token1")) + .addValues(stringProtoValue("r1c1")) + .addValues(stringProtoValue("r1c2")) + .addValues(stringProtoValue("r2c1")) + .build()); + // Schedule second chunk (with second column of Row 2) on the same 1-thread producer + // executor. + producerExecutor.execute( + () -> { + stream + .consumer() + .onPartialResultSet( + PartialResultSet.newBuilder() + .setResumeToken(ByteString.copyFromUtf8("token2")) + .addValues(stringProtoValue("r2c2")) + .build()); + stream.consumer().onCompleted(); + }); + return stream; + } + }; + + GrpcResultSet grpcResultSet = new GrpcResultSet(resumableStreamIterator, createNoOpListener()); + try (AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl(singleThreadProvider, grpcResultSet, 10)) { + List rows = new ArrayList<>(); + ApiFuture future = + asyncResultSet.setCallback( + callbackExecutor, + resultSet -> { + while (true) { + switch (resultSet.tryNext()) { + case OK: + rows.add(resultSet.getString(0) + ":" + resultSet.getString(1)); + break; + case NOT_READY: + return CallbackResponse.CONTINUE; + case DONE: + return CallbackResponse.DONE; + } + } + }); + + assertNull(future.get(1, TimeUnit.SECONDS)); + assertEquals(Arrays.asList("r1c1:r1c2", "r2c1:r2c2"), rows); + } finally { + producerExecutor.shutdownNow(); + callbackExecutor.shutdownNow(); + } + } + + @Test + public void testSingleThreadExecutorDoesNotDeadlockOnChunkedValue() throws Exception { + ExecutorService callbackExecutor = Executors.newSingleThreadExecutor(); + ExecutorProvider singleThreadProvider = + SpannerOptions.createAsyncExecutorProvider(1, 1L, TimeUnit.SECONDS); + ScheduledExecutorService producerExecutor = singleThreadProvider.getExecutor(); + ResumableStreamIterator resumableStreamIterator = + new ResumableStreamIterator( + 10, + "", + new OpenTelemetrySpan(mock(Span.class)), + createTestTracer(), + DefaultErrorHandler.INSTANCE, + SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetrySettings(), + SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetryableCodes(), + NoopRequestIdCreator.INSTANCE) { + @Override + AbstractResultSet.CloseableIterator startStream( + ByteString resumeToken, + AsyncResultSet.StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + GrpcStreamIterator stream = new GrpcStreamIterator(false, 10, false); + setStream(stream); + if (streamMessageListener != null) { + stream.registerListener(streamMessageListener); + } + stream.setCall(mock(StreamingCall.class), false); + // First chunk has a chunked value requiring a continuation chunk. + stream + .consumer() + .onPartialResultSet( + PartialResultSet.newBuilder() + .setMetadata(createTestMetadata(Type.StructField.of("col1", Type.string()))) + .setResumeToken(ByteString.copyFromUtf8("token1")) + .addValues(stringProtoValue("hello_")) + .setChunkedValue(true) + .build()); + // Schedule continuation chunk on the same 1-thread producer executor. + producerExecutor.execute( + () -> { + stream + .consumer() + .onPartialResultSet( + PartialResultSet.newBuilder() + .setResumeToken(ByteString.copyFromUtf8("token2")) + .addValues(stringProtoValue("world")) + .build()); + stream.consumer().onCompleted(); + }); + return stream; + } + }; + + GrpcResultSet grpcResultSet = new GrpcResultSet(resumableStreamIterator, createNoOpListener()); + try (AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl(singleThreadProvider, grpcResultSet, 10)) { + List rows = new ArrayList<>(); + ApiFuture future = + asyncResultSet.setCallback( + callbackExecutor, + resultSet -> { + while (true) { + switch (resultSet.tryNext()) { + case OK: + rows.add(resultSet.getString(0)); + break; + case NOT_READY: + return CallbackResponse.CONTINUE; + case DONE: + return CallbackResponse.DONE; + } + } + }); + + assertNull(future.get(1, TimeUnit.SECONDS)); + assertEquals(Collections.singletonList("hello_world"), rows); + } finally { + producerExecutor.shutdownNow(); + callbackExecutor.shutdownNow(); + } + } + + @Test + public void testMessageArrivingBeforeSetCallDoesNotThrowNpeOrLoseBeginTransaction() + throws Exception { + ExecutorService multiThreadExecutor = Executors.newFixedThreadPool(4); + ExecutorProvider multiThreadProvider = + new ExecutorProvider() { + @Override + public boolean shouldAutoClose() { + return false; + } + + @Override + public ScheduledExecutorService getExecutor() { + return Executors.newScheduledThreadPool(4); + } + }; + AtomicBoolean transactionMetadataShouldIncludeId = new AtomicBoolean(false); + CountDownLatch messageProcessedLatch = new CountDownLatch(1); + + ResumableStreamIterator resumableStreamIterator = + new ResumableStreamIterator( + 10, + "", + new OpenTelemetrySpan(mock(Span.class)), + createTestTracer(), + DefaultErrorHandler.INSTANCE, + SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetrySettings(), + SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetryableCodes(), + NoopRequestIdCreator.INSTANCE) { + @Override + AbstractResultSet.CloseableIterator startStream( + ByteString resumeToken, + AsyncResultSet.StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + GrpcStreamIterator stream = new GrpcStreamIterator(false, 10, false); + setStream(stream); + if (streamMessageListener != null) { + stream.registerListener(streamMessageListener); + } + // Simulate gRPC thread delivering PartialResultSet inside rpc.executeQuery() BEFORE + // stream.setCall(call, true) executes on the caller thread. + Future grpcCallbackFuture = + multiThreadExecutor.submit( + () -> { + stream + .consumer() + .onPartialResultSet( + PartialResultSet.newBuilder() + .setMetadata( + ResultSetMetadata.newBuilder() + .setRowType( + Type.struct( + Type.StructField.of("col1", Type.string())) + .toProto() + .getStructType()) + .setTransaction( + Transaction.newBuilder() + .setId(ByteString.copyFromUtf8("tx1")))) + .setResumeToken(ByteString.copyFromUtf8("token1")) + .addValues(stringProtoValue("v1")) + .build()); + stream.consumer().onCompleted(); + messageProcessedLatch.countDown(); + }); + try { + assertTrue(messageProcessedLatch.await(2, TimeUnit.SECONDS)); + grpcCallbackFuture.get(2, TimeUnit.SECONDS); + // Give ProduceRowsRunnable a brief moment to react if it was prematurely triggered + // before setCall. + Thread.sleep(50); + } catch (Exception e) { + throw new RuntimeException(e); + } + stream.setCall(mock(StreamingCall.class), /* withBeginTransaction= */ true); + return stream; + } + }; + + AbstractResultSet.Listener listener = + new AbstractResultSet.Listener() { + @Override + public void onTransactionMetadata(Transaction transaction, boolean shouldIncludeId) { + transactionMetadataShouldIncludeId.set(shouldIncludeId); + } + + @Override + public SpannerException onError( + SpannerException e, boolean withBeginTransaction, boolean lastStatement) { + return e; + } + + @Override + public void onDone(boolean withBeginTransaction) {} + + @Override + public void onPrecommitToken(MultiplexedSessionPrecommitToken token) {} + }; + + GrpcResultSet grpcResultSet = new GrpcResultSet(resumableStreamIterator, listener); + try (AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl(multiThreadProvider, grpcResultSet, 10)) { + List rows = new ArrayList<>(); + ApiFuture future = + asyncResultSet.setCallback( + multiThreadExecutor, + resultSet -> { + while (true) { + switch (resultSet.tryNext()) { + case OK: + rows.add(resultSet.getString(0)); + break; + case NOT_READY: + return CallbackResponse.CONTINUE; + case DONE: + return CallbackResponse.DONE; + } + } + }); + + assertNull(future.get(5, TimeUnit.SECONDS)); + assertEquals(Collections.singletonList("v1"), rows); + assertTrue( + "Expected onTransactionMetadata to receive shouldIncludeId=true", + transactionMetadataShouldIncludeId.get()); + } finally { + multiThreadExecutor.shutdown(); + assertTrue(multiThreadExecutor.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + @Test + public void testSinglePartialResultSetWithLastTrueAndMoreThanBufferRowsDoesNotCancelStream() + throws Exception { + StreamingCall call = mock(StreamingCall.class); + ResumableStreamIterator resumableStreamIterator = + new ResumableStreamIterator( + 10, + "", + new OpenTelemetrySpan(mock(Span.class)), + createTestTracer(), + DefaultErrorHandler.INSTANCE, + SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetrySettings(), + SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetryableCodes(), + NoopRequestIdCreator.INSTANCE) { + @Override + AbstractResultSet.CloseableIterator startStream( + ByteString resumeToken, + AsyncResultSet.StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + GrpcStreamIterator stream = new GrpcStreamIterator(false, 10, false); + setStream(stream); + if (streamMessageListener != null) { + stream.registerListener(streamMessageListener); + } + stream.setCall(call, false); + PartialResultSet.Builder builder = + PartialResultSet.newBuilder() + .setMetadata( + ResultSetMetadata.newBuilder() + .setRowType( + Type.struct(Type.StructField.of("col1", Type.string())) + .toProto() + .getStructType())) + .setResumeToken(ByteString.copyFromUtf8("r1")) + .setLast(true); + for (int i = 0; i < 20; i++) { + builder.addValues(stringProtoValue("row-" + i)); + } + stream.consumer().onPartialResultSet(builder.build()); + return stream; + } + }; + + GrpcResultSet grpcResultSet = + new GrpcResultSet(resumableStreamIterator, mock(AbstractResultSet.Listener.class)); + ExecutorService callbackExecutor = Executors.newSingleThreadExecutor(); + try (AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl( + simpleProvider, grpcResultSet, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { + List rows = new ArrayList<>(); + ApiFuture future = + asyncResultSet.setCallback( + callbackExecutor, + resultSet -> { + if (resultSet.tryNext() == CursorState.OK) { + rows.add(resultSet.getString(0)); + return CallbackResponse.DONE; + } + return CallbackResponse.CONTINUE; + }); + + assertNull(future.get(5, TimeUnit.SECONDS)); + assertEquals(Collections.singletonList("row-0"), rows); + verify(call, Mockito.never()).cancel(any()); + } finally { + callbackExecutor.shutdownNow(); + } + } + + @Test + public void testMultiplePartialResultSetsWithLastTrueCancelBehavior() throws Exception { + for (boolean stopAfterFirstRow : new boolean[] {true, false}) { + StreamingCall call = mock(StreamingCall.class); + ResumableStreamIterator resumableStreamIterator = + new ResumableStreamIterator( + 10, + "", + new OpenTelemetrySpan(mock(Span.class)), + createTestTracer(), + DefaultErrorHandler.INSTANCE, + SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetrySettings(), + SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetryableCodes(), + NoopRequestIdCreator.INSTANCE) { + @Override + AbstractResultSet.CloseableIterator startStream( + ByteString resumeToken, + AsyncResultSet.StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + GrpcStreamIterator stream = new GrpcStreamIterator(false, 10, false); + setStream(stream); + if (streamMessageListener != null) { + stream.registerListener(streamMessageListener); + } + stream.setCall(call, false); + PartialResultSet.Builder firstChunk = + PartialResultSet.newBuilder() + .setMetadata( + ResultSetMetadata.newBuilder() + .setRowType( + Type.struct(Type.StructField.of("col1", Type.string())) + .toProto() + .getStructType())) + .setResumeToken(ByteString.copyFromUtf8("r1")) + .setLast(false); + for (int i = 0; i < 15; i++) { + firstChunk.addValues(stringProtoValue("row-" + i)); + } + PartialResultSet.Builder secondChunk = + PartialResultSet.newBuilder() + .setResumeToken(ByteString.copyFromUtf8("r2")) + .setLast(true); + for (int i = 15; i < 20; i++) { + secondChunk.addValues(stringProtoValue("row-" + i)); + } + stream.consumer().onPartialResultSet(firstChunk.build()); + stream.consumer().onPartialResultSet(secondChunk.build()); + return stream; + } + }; + + GrpcResultSet grpcResultSet = + new GrpcResultSet(resumableStreamIterator, mock(AbstractResultSet.Listener.class)); + ExecutorService callbackExecutor = Executors.newSingleThreadExecutor(); + try (AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl( + simpleProvider, grpcResultSet, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { + List rows = new ArrayList<>(); + ApiFuture future = + asyncResultSet.setCallback( + callbackExecutor, + resultSet -> { + while (true) { + switch (resultSet.tryNext()) { + case OK: + rows.add(resultSet.getString(0)); + if (stopAfterFirstRow || rows.size() == 20) { + return CallbackResponse.DONE; + } + break; + case NOT_READY: + return CallbackResponse.CONTINUE; + case DONE: + return CallbackResponse.DONE; + } + } + }); + + assertNull(future.get(5, TimeUnit.SECONDS)); + if (stopAfterFirstRow) { + assertEquals(Collections.singletonList("row-0"), rows); + verify(call, times(1)).cancel(any()); + } else { + assertEquals(20, rows.size()); + verify(call, Mockito.never()).cancel(any()); + } + } finally { + callbackExecutor.shutdownNow(); + } + } + } + + @Test + public void testCloseWithAutoCloseExecutorAndThrowingDelegate() { + ScheduledExecutorService service = mock(ScheduledExecutorService.class); + when(mockedProvider.getExecutor()).thenReturn(service); + when(mockedProvider.shouldAutoClose()).thenReturn(true); + ResultSet delegate = mock(ResultSet.class); + doThrow(new RuntimeException("close error")).when(delegate).close(); + AsyncResultSetImpl resultSet = + new AsyncResultSetImpl(mockedProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE); + resultSet.close(); + verify(delegate).close(); + verify(service).shutdown(); + } + + @Test + public void testProducerScheduleRejected() { + ScheduledExecutorService service = mock(ScheduledExecutorService.class); + when(mockedProvider.getExecutor()).thenReturn(service); + AtomicInteger executeCount = new AtomicInteger(); + Mockito.doAnswer( + invocation -> { + Runnable runnable = invocation.getArgument(0); + if (executeCount.incrementAndGet() == 1) { + runnable.run(); + return null; + } + throw new RejectedExecutionException("producer rejected"); + }) + .when(service) + .execute(any(Runnable.class)); + + ResultSet delegate = mock(ResultSet.class); + AsyncResultSetImpl resultSet = + new AsyncResultSetImpl(mockedProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE); + ApiFuture future = + resultSet.setCallback( + Runnable::run, + cursor -> { + cursor.tryNext(); + return CallbackResponse.DONE; + }); + ExecutionException exception = assertThrows(ExecutionException.class, future::get); + assertTrue(exception.getCause() instanceof SpannerException); + verify(delegate).close(); + } + + @Test + public void testProduceRowsStoppedBeforeRun() throws Exception { + List producerTasks = new ArrayList<>(); + ScheduledExecutorService service = mock(ScheduledExecutorService.class); + when(mockedProvider.getExecutor()).thenReturn(service); + Mockito.doAnswer( + invocation -> { + Runnable runnable = invocation.getArgument(0); + if (producerTasks.isEmpty()) { + producerTasks.add(runnable); + runnable.run(); + } else { + producerTasks.add(runnable); + } + return null; + }) + .when(service) + .execute(any(Runnable.class)); + + ResultSet delegate = mock(ResultSet.class); + AsyncResultSetImpl resultSet = + new AsyncResultSetImpl(mockedProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE); + ApiFuture future = + resultSet.setCallback( + Runnable::run, + cursor -> { + cursor.tryNext(); + return CallbackResponse.DONE; + }); + assertEquals(2, producerTasks.size()); + resultSet.cancel(); + // Now run the queued ProduceRowsRunnable after cancellation; it should exit immediately. + producerTasks.get(1).run(); + ExecutionException exception = + assertThrows(ExecutionException.class, () -> future.get(5, TimeUnit.SECONDS)); + assertEquals(ErrorCode.CANCELLED, ((SpannerException) exception.getCause()).getErrorCode()); + } + + @Test + public void testCallbackReturnsNull() { + ResultSet delegate = mock(ResultSet.class); + when(delegate.next()).thenReturn(false); + try (AsyncResultSetImpl resultSet = + new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { + ApiFuture future = resultSet.setCallback(Runnable::run, cursor -> null); + ExecutionException exception = + assertThrows(ExecutionException.class, () -> future.get(5, TimeUnit.SECONDS)); + assertTrue(exception.getCause() instanceof SpannerException); + } + } + + @Test + public void testCallbackRunnableEdgeCases() throws Exception { + // 1. CallbackRunnable starts when cursorReturnedDoneOrException is already true. + List callbackTasks = new ArrayList<>(); + ResultSet delegate1 = mock(ResultSet.class); + StreamingResultSet streamingDelegate1 = mock(StreamingResultSet.class); + when(streamingDelegate1.initiateStreaming(any())).thenReturn(true); + when(streamingDelegate1.isDataAvailable()).thenReturn(false); + try (AsyncResultSetImpl resultSet = + new AsyncResultSetImpl( + simpleProvider, streamingDelegate1, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { + ApiFuture future = + resultSet.setCallback(callbackTasks::add, cursor -> CallbackResponse.DONE); + resultSet.cancel(); + assertEquals(1, callbackTasks.size()); + // Calling tryNext() when cancelled sets cursorReturnedDoneOrException = true before callback + // runs. + assertThrows(SpannerException.class, resultSet::tryNext); + // Calling cancel() again calls checkCompletion() while callbackScheduled is still true. + resultSet.cancel(); + callbackTasks.get(0).run(); + ExecutionException exception = + assertThrows(ExecutionException.class, () -> future.get(5, TimeUnit.SECONDS)); + assertEquals(ErrorCode.CANCELLED, ((SpannerException) exception.getCause()).getErrorCode()); + } + + // 2. Callback throws non-SpannerException and non-CANCELLED SpannerException when cancelled. + for (Throwable toThrow : + new Throwable[] { + new RuntimeException("custom runtime error"), + SpannerExceptionFactory.newSpannerException(ErrorCode.INTERNAL, "custom internal error") + }) { + StreamingResultSet streamingDelegate = mock(StreamingResultSet.class); + when(streamingDelegate.initiateStreaming(any())).thenReturn(true); + when(streamingDelegate.isDataAvailable()).thenReturn(false); + try (AsyncResultSetImpl resultSet = + new AsyncResultSetImpl( + simpleProvider, streamingDelegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { + ApiFuture future = + resultSet.setCallback( + Runnable::run, + cursor -> { + throw (RuntimeException) toThrow; + }); + resultSet.cancel(); + ExecutionException exception = + assertThrows(ExecutionException.class, () -> future.get(5, TimeUnit.SECONDS)); + assertTrue(exception.getCause() instanceof SpannerException); + assertNotNull(exception.getCause().getMessage()); + } + } + + // 3. cancel() called inside callback.cursorReady(), returning DONE vs CONTINUE, and + // getCurrentRowAsStruct() called while state == State.CANCELLED. + for (CallbackResponse responseOnCancel : + new CallbackResponse[] {CallbackResponse.DONE, CallbackResponse.CONTINUE}) { + ResultSet rowDelegate = mock(ResultSet.class); + Struct row = Struct.newBuilder().set("col").to("val").build(); + when(rowDelegate.next()).thenReturn(true, false); + when(rowDelegate.getCurrentRowAsStruct()).thenReturn(row); + try (AsyncResultSetImpl resultSet = + new AsyncResultSetImpl( + simpleProvider, rowDelegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { + AtomicInteger calls = new AtomicInteger(); + ApiFuture future = + resultSet.setCallback( + Runnable::run, + cursor -> { + if (calls.incrementAndGet() == 1) { + assertEquals(CursorState.OK, cursor.tryNext()); + resultSet.cancel(); + // Verify getCurrentRowAsStruct() is allowed when state == State.CANCELLED. + assertEquals(row, cursor.getCurrentRowAsStruct()); + if (responseOnCancel == CallbackResponse.DONE) { + assertThrows(SpannerException.class, cursor::tryNext); + // Calling cancel() while cursorReturnedDoneOrException && callbackRunning. + resultSet.cancel(); + } + return responseOnCancel; + } + cursor.tryNext(); + return CallbackResponse.DONE; + }); + ExecutionException exception = + assertThrows(ExecutionException.class, () -> future.get(5, TimeUnit.SECONDS)); + assertEquals(ErrorCode.CANCELLED, ((SpannerException) exception.getCause()).getErrorCode()); + } + } + + // 4. tryNext() returns DONE, but callback returns CONTINUE instead of DONE. + when(delegate1.next()).thenReturn(false); + try (AsyncResultSetImpl resultSet = + new AsyncResultSetImpl(simpleProvider, delegate1, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { + ApiFuture future = + resultSet.setCallback( + Runnable::run, + cursor -> { + assertEquals(CursorState.DONE, cursor.tryNext()); + return CallbackResponse.CONTINUE; + }); + assertNull(future.get(5, TimeUnit.SECONDS)); + } + + // 5. Callback consumes the only buffered row and returns CONTINUE without calling tryNext() a + // second time, when finished == true vs when executionException != null. + Struct singleRow = Struct.newBuilder().set("col").to("val").build(); + ResultSet finishedAfterOneRowDelegate = mock(ResultSet.class); + when(finishedAfterOneRowDelegate.next()).thenReturn(true, false); + when(finishedAfterOneRowDelegate.getCurrentRowAsStruct()).thenReturn(singleRow); + try (AsyncResultSetImpl resultSet = + new AsyncResultSetImpl( + simpleProvider, finishedAfterOneRowDelegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { + ApiFuture future = + resultSet.setCallback( + Runnable::run, + cursor -> { + if (cursor.tryNext() == CursorState.OK) { + return CallbackResponse.CONTINUE; + } + return CallbackResponse.DONE; + }); + assertNull(future.get(5, TimeUnit.SECONDS)); + } + + CountDownLatch delegateClosed = new CountDownLatch(1); + ResultSet errorAfterOneRowDelegate = mock(ResultSet.class); + when(errorAfterOneRowDelegate.next()) + .thenReturn(true) + .thenThrow(new RuntimeException("error on second row")); + when(errorAfterOneRowDelegate.getCurrentRowAsStruct()).thenReturn(singleRow); + Mockito.doAnswer( + invocation -> { + delegateClosed.countDown(); + return null; + }) + .when(errorAfterOneRowDelegate) + .close(); + ExecutorService singleThreadCallbackExecutor = Executors.newSingleThreadExecutor(); + try (AsyncResultSetImpl resultSet = + new AsyncResultSetImpl( + simpleProvider, errorAfterOneRowDelegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { + ApiFuture future = + resultSet.setCallback( + singleThreadCallbackExecutor, + cursor -> { + try { + assertTrue(delegateClosed.await(5, TimeUnit.SECONDS)); + } catch (InterruptedException e) { + throw SpannerExceptionFactory.propagateInterrupt(e); + } + if (cursor.tryNext() == CursorState.OK) { + return CallbackResponse.CONTINUE; + } + return CallbackResponse.DONE; + }); + assertThrows(ExecutionException.class, () -> future.get(5, TimeUnit.SECONDS)); + } finally { + singleThreadCallbackExecutor.shutdownNow(); + } + } + + @Test + public void testSetExecutionExceptionWhenAlreadyFailedOrStopped() throws Exception { + // 1. Producer fails with exception AND callback executor rejects execution -> + // setExecutionException called twice (second time when executionException != null). + ResultSet failingDelegate = mock(ResultSet.class); + when(failingDelegate.next()).thenThrow(new RuntimeException("delegate error")); + try (AsyncResultSetImpl resultSet = + new AsyncResultSetImpl( + simpleProvider, failingDelegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { + ApiFuture future = + resultSet.setCallback( + command -> { + throw new RejectedExecutionException("callback rejected"); + }, + cursor -> CallbackResponse.DONE); + ExecutionException exception = + assertThrows(ExecutionException.class, () -> future.get(5, TimeUnit.SECONDS)); + assertTrue(exception.getCause() instanceof SpannerException); + } + + // 2. Producer fails with exception AFTER resultSet is cancelled (state.shouldStop == true). + CountDownLatch inNext = new CountDownLatch(1); + CountDownLatch proceedNext = new CountDownLatch(1); + ResultSet cancelWhileNextDelegate = mock(ResultSet.class); + when(cancelWhileNextDelegate.next()) + .thenAnswer( + invocation -> { + inNext.countDown(); + assertTrue(proceedNext.await(5, TimeUnit.SECONDS)); + throw new RuntimeException("error after cancel"); + }); + try (AsyncResultSetImpl resultSet = + new AsyncResultSetImpl( + simpleProvider, cancelWhileNextDelegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { + ApiFuture future = + resultSet.setCallback( + Runnable::run, + cursor -> { + cursor.tryNext(); + return CallbackResponse.DONE; + }); + assertTrue(inNext.await(5, TimeUnit.SECONDS)); + resultSet.cancel(); + proceedNext.countDown(); + ExecutionException exception = + assertThrows(ExecutionException.class, () -> future.get(5, TimeUnit.SECONDS)); + assertEquals(ErrorCode.CANCELLED, ((SpannerException) exception.getCause()).getErrorCode()); + } + } + + @Test + public void testToListAsyncClosedAndNotReadyAndInterrupted() throws Exception { + ResultSet delegate = mock(ResultSet.class); + AsyncResultSetImpl closedResultSet = + new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE); + // Calling getCurrentRowAsStruct() before next() or callback throws IllegalStateException. + assertThrows(IllegalStateException.class, closedResultSet::getCurrentRowAsStruct); + closedResultSet.close(); + assertThrows( + IllegalStateException.class, + () -> closedResultSet.toListAsync(reader -> reader.getString(0), Runnable::run)); + + // Calling cancel(), resume(), or toListAsync() in State.SYNC throws IllegalStateException. + ResultSet syncDelegate = mock(ResultSet.class); + when(syncDelegate.next()).thenReturn(false); + try (AsyncResultSetImpl syncResultSet = + new AsyncResultSetImpl( + simpleProvider, syncDelegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { + assertFalse(syncResultSet.next()); + assertThrows(IllegalStateException.class, syncResultSet::cancel); + assertThrows(IllegalStateException.class, syncResultSet::resume); + assertThrows( + IllegalStateException.class, + () -> syncResultSet.toListAsync(reader -> reader.getString(0), Runnable::run)); + } + + // Test CreateListCallback receiving NOT_READY when buffer size is 1 and there are 3 rows. + ResultSet multiRowDelegate = mock(ResultSet.class); + Struct row = Struct.newBuilder().set("col").to("val").build(); + when(multiRowDelegate.next()).thenReturn(true, true, true, false); + when(multiRowDelegate.getCurrentRowAsStruct()).thenReturn(row); + ExecutorService callbackExecutor = Executors.newSingleThreadExecutor(); + try (AsyncResultSetImpl resultSet = + new AsyncResultSetImpl(simpleProvider, multiRowDelegate, 1)) { + List list = + resultSet + .toListAsync(reader -> reader.getString(0), callbackExecutor) + .get(5, TimeUnit.SECONDS); + assertEquals(3, list.size()); + } finally { + callbackExecutor.shutdownNow(); + } + + // Test toList() when calling thread is interrupted. + StreamingResultSet blockingDelegate = mock(StreamingResultSet.class); + when(blockingDelegate.initiateStreaming(any())).thenReturn(true); + when(blockingDelegate.isDataAvailable()).thenReturn(false); + try (AsyncResultSetImpl resultSet = + new AsyncResultSetImpl( + simpleProvider, blockingDelegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) { + Thread.currentThread().interrupt(); + try { + SpannerException exception = + assertThrows( + SpannerException.class, () -> resultSet.toList(reader -> reader.getString(0))); + assertEquals(ErrorCode.UNKNOWN, exception.getErrorCode()); + } finally { + Thread.interrupted(); + resultSet.cancel(); + } } } + + @Test + public void testGrpcStreamIteratorOnErrorAfterLastDoesNotDeadlockWhenQueueFull() { + GrpcStreamIterator stream = new GrpcStreamIterator(false, 1, false); + StreamingCall call = mock(StreamingCall.class); + stream.setCall(call, false); + PartialResultSet lastChunk = + PartialResultSet.newBuilder() + .setResumeToken(ByteString.copyFromUtf8("r1")) + .setLast(true) + .build(); + // With prefetchChunks = 1, queue capacity is 2. + // onPartialResultSet with last=true adds lastChunk AND END_OF_STREAM, filling the queue. + stream.consumer().onPartialResultSet(lastChunk); + // Calling onError after last=true must not block trying to add another END_OF_STREAM. + stream + .consumer() + .onError(SpannerExceptionFactory.newSpannerException(ErrorCode.CANCELLED, "cancelled")); + assertTrue(stream.hasNext()); + assertEquals(lastChunk, stream.next()); + assertThrows(SpannerException.class, stream::hasNext); + } } diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ForwardingResultSetTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ForwardingResultSetTest.java new file mode 100644 index 000000000000..995b746cee50 --- /dev/null +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ForwardingResultSetTest.java @@ -0,0 +1,94 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.spanner; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import com.google.spanner.v1.ResultSetMetadata; +import com.google.spanner.v1.ResultSetStats; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class ForwardingResultSetTest { + + private interface MockStreamingResultSet extends ResultSet, StreamingResultSet {} + + @Test + public void testDelegation() { + MockStreamingResultSet delegate = mock(MockStreamingResultSet.class); + when(delegate.next()).thenReturn(true, false); + ResultSetStats stats = ResultSetStats.getDefaultInstance(); + ResultSetMetadata metadata = ResultSetMetadata.getDefaultInstance(); + when(delegate.getStats()).thenReturn(stats); + when(delegate.getMetadata()).thenReturn(metadata); + + ForwardingResultSet forwardingResultSet = new ForwardingResultSet(delegate); + + assertTrue(forwardingResultSet.next()); + assertFalse(forwardingResultSet.next()); + assertEquals(stats, forwardingResultSet.getStats()); + assertEquals(metadata, forwardingResultSet.getMetadata()); + assertEquals(delegate, forwardingResultSet.getDelegate()); + + forwardingResultSet.close(); + verify(delegate).close(); + } + + @Test + public void testInitiateStreamingDelegation() { + MockStreamingResultSet delegate = mock(MockStreamingResultSet.class); + AsyncResultSet.StreamMessageListener listener = + mock(AsyncResultSet.StreamMessageListener.class); + when(delegate.initiateStreaming(listener)).thenReturn(true); + + ForwardingResultSet forwardingResultSet = new ForwardingResultSet(delegate); + assertTrue(forwardingResultSet.initiateStreaming(listener)); + verify(delegate).initiateStreaming(listener); + } + + @Test + public void testIsDataAvailableDelegation() { + MockStreamingResultSet delegate = mock(MockStreamingResultSet.class); + when(delegate.isDataAvailable()).thenReturn(false, true); + + ForwardingResultSet forwardingResultSet = new ForwardingResultSet(delegate); + assertFalse(forwardingResultSet.isDataAvailable()); + assertTrue(forwardingResultSet.isDataAvailable()); + } + + @Test + public void testReplaceDelegate() { + ResultSet firstDelegate = mock(ResultSet.class); + ResultSet secondDelegate = mock(ResultSet.class); + when(firstDelegate.next()).thenReturn(true); + when(secondDelegate.next()).thenReturn(false); + + ForwardingResultSet forwardingResultSet = new ForwardingResultSet(firstDelegate); + assertTrue(forwardingResultSet.next()); + + forwardingResultSet.replaceDelegate(secondDelegate); + assertFalse(forwardingResultSet.next()); + assertEquals(secondDelegate, forwardingResultSet.getDelegate()); + } +} 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 4007c972c24e..26ea13da4dda 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 @@ -22,6 +22,9 @@ import static org.junit.Assert.assertFalse; 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; import com.google.api.gax.grpc.GrpcCallContext; import com.google.api.gax.rpc.ApiCallContext; @@ -1254,4 +1257,45 @@ public void shouldThrowDeadlineExceededIfLastTrueIsNotReceived() { assertEquals("DEADLINE_EXCEEDED: stream wait timeout", spannerException.getMessage()); consumer.onCompleted(); } + + @Test + public void testIsDataAvailable() { + assertFalse(resultSet.isDataAvailable()); + + consumer.onPartialResultSet( + PartialResultSet.newBuilder() + .setMetadata(makeMetadata(Type.struct(Type.StructField.of("f", Type.string())))) + .addValues(Value.string("val1").toProto()) + .addValues(Value.string("val2").toProto()) + .build()); + + assertTrue(resultSet.isDataAvailable()); + assertTrue(resultSet.next()); + assertEquals("val1", resultSet.getString(0)); + + assertTrue(resultSet.isDataAvailable()); + assertTrue(resultSet.next()); + assertEquals("val2", resultSet.getString(0)); + + assertFalse(resultSet.isDataAvailable()); + consumer.onCompleted(); + assertTrue(resultSet.isDataAvailable()); + assertFalse(resultSet.next()); + } + + @Test + public void testInitiateStreamingDelegation() { + @SuppressWarnings("unchecked") + AbstractResultSet.CloseableIterator mockStream = + mock(AbstractResultSet.CloseableIterator.class); + AsyncResultSet.StreamMessageListener listener = + mock(AsyncResultSet.StreamMessageListener.class); + when(mockStream.initiateStreaming(listener)).thenReturn(true); + + GrpcResultSet streamingResultSet = new GrpcResultSet(mockStream, new NoOpListener()); + assertTrue(streamingResultSet.initiateStreaming(listener)); + verify(mockStream).initiateStreaming(listener); + + assertFalse(resultSet.initiateStreaming(listener)); + } } diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/GrpcValueIteratorTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/GrpcValueIteratorTest.java new file mode 100644 index 000000000000..15e7423c33cd --- /dev/null +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/GrpcValueIteratorTest.java @@ -0,0 +1,138 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.spanner; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import com.google.cloud.spanner.AbstractResultSet.CloseableIterator; +import com.google.protobuf.Value; +import com.google.spanner.v1.PartialResultSet; +import com.google.spanner.v1.ResultSetMetadata; +import com.google.spanner.v1.StructType; +import com.google.spanner.v1.StructType.Field; +import com.google.spanner.v1.Type; +import com.google.spanner.v1.TypeCode; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class GrpcValueIteratorTest { + @SuppressWarnings("unchecked") + private final CloseableIterator stream = mock(CloseableIterator.class); + + private final AbstractResultSet.Listener listener = mock(AbstractResultSet.Listener.class); + + @Test + public void testIsDataAvailableWhenCurrentIsNull() { + GrpcValueIterator iterator = new GrpcValueIterator(stream, listener); + when(stream.isDataAvailable()).thenReturn(false); + assertFalse(iterator.isDataAvailable()); + + when(stream.isDataAvailable()).thenReturn(true); + assertTrue(iterator.isDataAvailable()); + } + + @Test + public void testIsDataAvailableWhenValuesRemainInCurrent() { + PartialResultSet prs = + PartialResultSet.newBuilder() + .setMetadata( + ResultSetMetadata.newBuilder() + .setRowType( + StructType.newBuilder() + .addFields( + Field.newBuilder() + .setName("col") + .setType(Type.newBuilder().setCode(TypeCode.STRING))))) + .addValues(Value.newBuilder().setStringValue("val1")) + .addValues(Value.newBuilder().setStringValue("val2")) + .build(); + when(stream.hasNext()).thenReturn(true, false); + when(stream.next()).thenReturn(prs); + + GrpcValueIterator iterator = new GrpcValueIterator(stream, listener); + when(stream.isDataAvailable()).thenReturn(false); + + // Initial state before reading: current is null, stream has no data available + assertFalse(iterator.isDataAvailable()); + + // Advance iterator: reads first value "val1", current is set, pos is now 1 (1 < 2 values) + assertTrue(iterator.hasNext()); + assertEquals("val1", iterator.next().getStringValue()); + + // Even though stream.isDataAvailable() is false, iterator has 1 value remaining in current + // chunk + assertFalse(stream.isDataAvailable()); + assertTrue(iterator.isDataAvailable()); + + // Read second value "val2", pos is now 2 (all 2 values consumed from current chunk) + assertTrue(iterator.hasNext()); + assertEquals("val2", iterator.next().getStringValue()); + + // Now all values in current chunk are consumed, so isDataAvailable matches stream + assertFalse(iterator.isDataAvailable()); + when(stream.isDataAvailable()).thenReturn(true); + assertTrue(iterator.isDataAvailable()); + } + + @Test + public void testIsDataAvailableWhenCurrentEndsWithChunkedValueAndStreamHasNoData() { + PartialResultSet prs1 = + PartialResultSet.newBuilder() + .setMetadata( + ResultSetMetadata.newBuilder() + .setRowType( + StructType.newBuilder() + .addFields( + Field.newBuilder() + .setName("col") + .setType(Type.newBuilder().setCode(TypeCode.STRING))))) + .addValues(Value.newBuilder().setStringValue("complete_val")) + .addValues(Value.newBuilder().setStringValue("chunked_part_1_")) + .setChunkedValue(true) + .build(); + PartialResultSet prs2 = + PartialResultSet.newBuilder() + .addValues(Value.newBuilder().setStringValue("chunked_part_2")) + .build(); + when(stream.hasNext()).thenReturn(true, true, false); + when(stream.next()).thenReturn(prs1, prs2); + + GrpcValueIterator iterator = new GrpcValueIterator(stream, listener); + when(stream.isDataAvailable()).thenReturn(false); + + // Read the first complete value; pos is now 1 (pointing to the last value in prs1, which has + // chunked_value = true). + assertTrue(iterator.hasNext()); + assertEquals("complete_val", iterator.next().getStringValue()); + + // Because the remaining value in prs1 is chunked and stream.isDataAvailable() is false, + // advancing the iterator would block waiting for prs2, so isDataAvailable() must be false. + assertFalse(iterator.isDataAvailable()); + + // Once prs2 is available on the stream, isDataAvailable() becomes true. + when(stream.isDataAvailable()).thenReturn(true); + assertTrue(iterator.isDataAvailable()); + assertTrue(iterator.hasNext()); + assertEquals("chunked_part_1_chunked_part_2", iterator.next().getStringValue()); + } +} diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadAsyncTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadAsyncTest.java index 1251ee270faf..652d55ae85eb 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadAsyncTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadAsyncTest.java @@ -19,7 +19,10 @@ import static com.google.cloud.spanner.MockSpannerTestUtil.*; import static com.google.cloud.spanner.SpannerApiFutures.get; import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNull; import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; import static org.junit.Assert.fail; import com.google.api.core.ApiFuture; @@ -33,6 +36,9 @@ import com.google.common.collect.ContiguousSet; import com.google.common.collect.ImmutableSet; import com.google.common.collect.Iterables; +import com.google.protobuf.ByteString; +import com.google.spanner.v1.ExecuteSqlRequest; +import com.google.spanner.v1.ReadRequest; import io.grpc.Server; import io.grpc.Status; import io.grpc.inprocess.InProcessServerBuilder; @@ -49,16 +55,21 @@ import java.util.concurrent.Executors; import java.util.concurrent.ScheduledThreadPoolExecutor; import java.util.concurrent.SynchronousQueue; +import java.util.concurrent.TimeUnit; import org.junit.After; import org.junit.AfterClass; import org.junit.Before; import org.junit.BeforeClass; +import org.junit.Rule; import org.junit.Test; +import org.junit.rules.Timeout; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @RunWith(JUnit4.class) public class ReadAsyncTest { + @Rule public Timeout globalTimeout = Timeout.seconds(60); + private static MockSpannerServiceImpl mockSpanner; private static Server server; private static LocalChannelProvider channelProvider; @@ -114,6 +125,7 @@ public void before() { public void after() { spanner.close(); mockSpanner.removeAllExecutionTimes(); + mockSpanner.clearRequests(); } @Test @@ -261,9 +273,15 @@ public void closeTransactionBeforeEndOfAsyncQuery() throws Exception { dataReceived.await(); } List resultList = new ArrayList<>(); - do { + while (!finished.isDone() || !results.isEmpty()) { results.drainTo(resultList); - } while (!finished.isDone() || results.size() > 0); + if (!finished.isDone() && resultList.size() < 3) { + String next = results.poll(100, TimeUnit.MILLISECONDS); + if (next != null) { + resultList.add(next); + } + } + } assertThat(finished.get()).isTrue(); assertThat(resultList).containsExactly("k1", "k2", "k3"); closed.get(); @@ -332,12 +350,16 @@ public void pauseResume() throws Exception { while (true) { switch (resultSet.tryNext()) { case DONE: + synchronized (lock) { + lock.notifyAll(); + } return CallbackResponse.DONE; case NOT_READY: return CallbackResponse.CONTINUE; case OK: synchronized (lock) { allValues.add(resultSet.getString("Value")); + lock.notifyAll(); } unevenReturnedFirstRow.countDown(); return CallbackResponse.PAUSE; @@ -355,12 +377,16 @@ public void pauseResume() throws Exception { while (true) { switch (resultSet.tryNext()) { case DONE: + synchronized (lock) { + lock.notifyAll(); + } return CallbackResponse.DONE; case NOT_READY: return CallbackResponse.CONTINUE; case OK: synchronized (lock) { allValues.add(resultSet.getString("Value")); + lock.notifyAll(); } return CallbackResponse.PAUSE; } @@ -371,7 +397,11 @@ public void pauseResume() throws Exception { }); while (!(evenFinished.isDone() && unevenFinished.isDone())) { synchronized (lock) { - if (allValues.peekLast() != null) { + if (unevenFinished.isDone()) { + evenRs.resume(); + } else if (evenFinished.isDone()) { + unevenRs.resume(); + } else if (allValues.peekLast() != null) { if (Integer.parseInt(allValues.peekLast().substring(1)) % 2 == 1) { evenRs.resume(); } else { @@ -382,6 +412,9 @@ public void pauseResume() throws Exception { unevenRs.resume(); evenRs.resume(); } + if (!(evenFinished.isDone() && unevenFinished.isDone())) { + lock.wait(100); + } } } } @@ -392,6 +425,115 @@ public void pauseResume() throws Exception { .containsExactly("v1", "v2", "v3", "v4", "v5", "v6", "v7", "v8", "v9", "v10"); } + @Test + public void pauseResumeUnequalStreams() throws Exception { + Statement unevenStatement = + Statement.of( + "SELECT * FROM TestTable WHERE MOD(CAST(SUBSTR(Key, 2) AS INT64), 2) = 1 ORDER BY" + + " CAST(SUBSTR(Key, 2) AS INT64)"); + Statement evenStatement = + Statement.of( + "SELECT * FROM TestTable WHERE MOD(CAST(SUBSTR(Key, 2) AS INT64), 2) = 0 ORDER BY" + + " CAST(SUBSTR(Key, 2) AS INT64)"); + mockSpanner.putStatementResult( + StatementResult.query( + unevenStatement, generateKeyValueResultSet(ImmutableSet.of(1, 3, 5, 7, 9, 11, 13)))); + mockSpanner.putStatementResult( + StatementResult.query( + evenStatement, generateKeyValueResultSet(ImmutableSet.of(0, 2, 4, 6, 8, 10, 12, 14)))); + + final Object lock = new Object(); + ApiFuture evenFinished; + ApiFuture unevenFinished; + final CountDownLatch evenReturnedFirstRow = new CountDownLatch(1); + final Deque allValues = new ConcurrentLinkedDeque<>(); + try (ReadOnlyTransaction tx = client.readOnlyTransaction()) { + try (AsyncResultSet evenRs = tx.executeQueryAsync(evenStatement); + AsyncResultSet unevenRs = tx.executeQueryAsync(unevenStatement)) { + evenFinished = + evenRs.setCallback( + executor, + resultSet -> { + while (true) { + switch (resultSet.tryNext()) { + case DONE: + synchronized (lock) { + lock.notifyAll(); + } + return CallbackResponse.DONE; + case NOT_READY: + return CallbackResponse.CONTINUE; + case OK: + synchronized (lock) { + allValues.add(resultSet.getString("Value")); + lock.notifyAll(); + } + evenReturnedFirstRow.countDown(); + return CallbackResponse.PAUSE; + } + } + }); + unevenFinished = + unevenRs.setCallback( + executor, + resultSet -> { + try { + // Make sure the even result set has returned the first before we start the + // uneven results. + evenReturnedFirstRow.await(); + while (true) { + switch (resultSet.tryNext()) { + case DONE: + synchronized (lock) { + lock.notifyAll(); + } + return CallbackResponse.DONE; + case NOT_READY: + return CallbackResponse.CONTINUE; + case OK: + synchronized (lock) { + allValues.add(resultSet.getString("Value")); + lock.notifyAll(); + } + return CallbackResponse.PAUSE; + } + } + } catch (InterruptedException e) { + throw SpannerExceptionFactory.propagateInterrupt(e); + } + }); + while (!(evenFinished.isDone() && unevenFinished.isDone())) { + synchronized (lock) { + if (unevenFinished.isDone()) { + evenRs.resume(); + } else if (evenFinished.isDone()) { + unevenRs.resume(); + } else if (allValues.peekLast() != null) { + if (Integer.parseInt(allValues.peekLast().substring(1)) % 2 == 1) { + evenRs.resume(); + } else { + unevenRs.resume(); + } + } + if (allValues.size() == 15) { + unevenRs.resume(); + evenRs.resume(); + } + if (!(evenFinished.isDone() && unevenFinished.isDone())) { + lock.wait(100); + } + } + } + } + } + assertThat(ApiFutures.allAsList(Arrays.asList(evenFinished, unevenFinished)).get()) + .containsExactly(null, null); + assertThat(allValues) + .containsExactly( + "v0", "v1", "v2", "v3", "v4", "v5", "v6", "v7", "v8", "v9", "v10", "v11", "v12", "v13", + "v14"); + } + @Test public void cancel() throws Exception { final List values = new LinkedList<>(); @@ -431,6 +573,115 @@ public void cancel() throws Exception { assertThat(values).containsExactly("v1"); } + @Test + public void readAsyncRetriesOnUnavailableHalfway() throws Exception { + int totalRowCount = 50; + int errorIndex = 20; + String retryTableName = "RetryTable"; + mockSpanner.putStatementResult( + StatementResult.read( + retryTableName, + KeySet.all(), + READ_COLUMN_NAMES, + generateKeyValueResultSet(ContiguousSet.closed(1, totalRowCount)))); + mockSpanner.setStreamingReadExecutionTime( + SimulatedExecutionTime.ofStreamException( + Status.UNAVAILABLE.asRuntimeException(), errorIndex)); + mockSpanner.clearRequests(); + + List receivedKeys = new ArrayList<>(); + List receivedValues = new ArrayList<>(); + try (AsyncResultSet resultSet = + client.singleUse().readAsync(retryTableName, KeySet.all(), READ_COLUMN_NAMES)) { + ApiFuture future = + resultSet.setCallback( + executor, + ready -> { + while (true) { + switch (ready.tryNext()) { + case OK: + receivedKeys.add(ready.getString("Key")); + receivedValues.add(ready.getString("Value")); + break; + case NOT_READY: + return CallbackResponse.CONTINUE; + case DONE: + return CallbackResponse.DONE; + } + } + }); + assertNull(future.get(10, TimeUnit.SECONDS)); + } + + assertEquals(totalRowCount, receivedKeys.size()); + assertEquals(totalRowCount, receivedValues.size()); + for (int i = 0; i < totalRowCount; i++) { + assertEquals("k" + (i + 1), receivedKeys.get(i)); + assertEquals("v" + (i + 1), receivedValues.get(i)); + } + + assertEquals(2, mockSpanner.countRequestsOfType(ReadRequest.class)); + ReadRequest initialRequest = mockSpanner.getRequestsOfType(ReadRequest.class).get(0); + assertTrue(initialRequest.getResumeToken().isEmpty()); + + ReadRequest resumeRequest = mockSpanner.getRequestsOfType(ReadRequest.class).get(1); + assertEquals( + ByteString.copyFromUtf8(String.format("%09d", errorIndex)), resumeRequest.getResumeToken()); + } + + @Test + public void executeQueryAsyncRetriesOnUnavailableHalfway() throws Exception { + int totalRowCount = 50; + int errorIndex = 20; + Statement statement = Statement.of("SELECT Key, Value FROM RetryTable"); + mockSpanner.putStatementResult( + StatementResult.query( + statement, generateKeyValueResultSet(ContiguousSet.closed(1, totalRowCount)))); + mockSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.ofStreamException( + Status.UNAVAILABLE.asRuntimeException(), errorIndex)); + mockSpanner.clearRequests(); + + List receivedKeys = new ArrayList<>(); + List receivedValues = new ArrayList<>(); + try (AsyncResultSet resultSet = client.singleUse().executeQueryAsync(statement)) { + ApiFuture future = + resultSet.setCallback( + executor, + ready -> { + while (true) { + switch (ready.tryNext()) { + case OK: + receivedKeys.add(ready.getString("Key")); + receivedValues.add(ready.getString("Value")); + break; + case NOT_READY: + return CallbackResponse.CONTINUE; + case DONE: + return CallbackResponse.DONE; + } + } + }); + assertNull(future.get(10, TimeUnit.SECONDS)); + } + + assertEquals(totalRowCount, receivedKeys.size()); + assertEquals(totalRowCount, receivedValues.size()); + for (int i = 0; i < totalRowCount; i++) { + assertEquals("k" + (i + 1), receivedKeys.get(i)); + assertEquals("v" + (i + 1), receivedValues.get(i)); + } + + assertEquals(2, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)); + ExecuteSqlRequest initialRequest = + mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(0); + assertTrue(initialRequest.getResumeToken().isEmpty()); + + ExecuteSqlRequest resumeRequest = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(1); + assertEquals( + ByteString.copyFromUtf8(String.format("%09d", errorIndex)), resumeRequest.getResumeToken()); + } + private boolean isMultiplexedSessionsEnabled() { if (spanner.getOptions() == null || spanner.getOptions().getSessionPoolOptions() == null) { return false; 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..067ebfeb6a35 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 @@ -18,24 +18,30 @@ import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; 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; +import com.google.api.client.util.BackOff; +import com.google.api.client.util.ExponentialBackOff; import com.google.api.gax.retrying.RetrySettings; +import com.google.api.gax.rpc.StatusCode.Code; import com.google.cloud.spanner.ErrorHandler.DefaultErrorHandler; import com.google.cloud.spanner.XGoogSpannerRequestId.NoopRequestIdCreator; +import com.google.cloud.spanner.spi.v1.SpannerRpc; import com.google.cloud.spanner.v1.stub.SpannerStubSettings; import com.google.common.base.Stopwatch; import com.google.common.collect.AbstractIterator; import com.google.common.collect.ImmutableList; import com.google.common.collect.Lists; import com.google.protobuf.ByteString; -import com.google.protobuf.Duration; import com.google.protobuf.Value; import com.google.rpc.RetryInfo; import com.google.spanner.v1.PartialResultSet; +import io.grpc.Context; import io.grpc.Metadata; import io.grpc.Status; import io.grpc.StatusRuntimeException; @@ -44,11 +50,18 @@ import io.opencensus.trace.Tracing; import io.opentelemetry.api.OpenTelemetry; import io.opentelemetry.api.common.Attributes; +import java.io.IOException; import java.util.ArrayList; +import java.util.Collections; import java.util.Iterator; import java.util.LinkedList; import java.util.List; +import java.util.NoSuchElementException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; import java.util.function.LongSupplier; import javax.annotation.Nullable; import org.junit.Before; @@ -58,6 +71,7 @@ import org.junit.runners.Parameterized.Parameter; import org.junit.runners.Parameterized.Parameters; import org.mockito.Mockito; +import org.threeten.bp.Duration; /** Unit tests for {@link ResumableStreamIterator}. */ @RunWith(Parameterized.class) @@ -88,7 +102,7 @@ private static StatusRuntimeException statusWithRetryInfo(ErrorCode code) { RetryInfo retryInfo = RetryInfo.newBuilder() .setRetryDelay( - Duration.newBuilder() + com.google.protobuf.Duration.newBuilder() .setNanos((int) TimeUnit.MILLISECONDS.toNanos(1L)) .setSeconds(0L)) .build(); @@ -196,7 +210,7 @@ AbstractResultSet.CloseableIterator startStream( @Nullable ByteString resumeToken, AsyncResultSet.StreamMessageListener streamMessageListener, XGoogSpannerRequestId requestId) { - return starter.startStream(resumeToken, null); + return starter.startStream(resumeToken, streamMessageListener); } }; } @@ -750,6 +764,943 @@ public void bufferLimitMissingTokensSafeToRetry() { assertThat(consume(resumableStreamIterator)).containsExactly("a", "b", "c", "d").inOrder(); } + @Test + public void isDataAvailableWithGrpcStreamIterator() { + initWithLimit(512); + GrpcStreamIterator grpcStream = new GrpcStreamIterator(false, 4, false); + Mockito.when(starter.startStream(null, null)).thenReturn(grpcStream); + resumableStreamIterator.setStream(grpcStream); + + // Empty stream -> false + assertFalse(resumableStreamIterator.isDataAvailable()); + + // Chunk without resume token arrives -> drained into buffer, but not flushable yet until a + // resume token or EOF arrives (or maxBufferSize is exceeded). + grpcStream.consumer().onPartialResultSet(resultSet(null, "a")); + assertFalse(resumableStreamIterator.isDataAvailable()); + assertEquals(1, resumableStreamIterator.buffer.size()); + + grpcStream.consumer().onCompleted(); + assertTrue(resumableStreamIterator.isDataAvailable()); + } + + @Test + public void isDataAvailableWithoutResumeTokensWithProductionBufferSize() throws Exception { + initWithLimit(512); + GrpcStreamIterator grpcStream = new GrpcStreamIterator(false, 4, false); + SpannerRpc.StreamingCall call = Mockito.mock(SpannerRpc.StreamingCall.class); + grpcStream.setCall(call, false); + Mockito.when(starter.startStream(null, null)).thenReturn(grpcStream); + resumableStreamIterator.setStream(grpcStream); + + assertFalse(resumableStreamIterator.isDataAvailable()); + + ExecutorService feeder = Executors.newSingleThreadExecutor(); + try { + feeder.submit( + () -> { + for (int i = 0; i < 10; i++) { + grpcStream.consumer().onPartialResultSet(resultSet(null, "val" + i)); + } + grpcStream.consumer().onCompleted(); + }); + + List results = new ArrayList<>(); + while (resumableStreamIterator.hasNext()) { + PartialResultSet prs = resumableStreamIterator.next(); + results.add(prs.getValues(0).getStringValue()); + } + assertEquals(10, results.size()); + for (int i = 0; i < 10; i++) { + assertEquals("val" + i, results.get(i)); + } + } finally { + feeder.shutdown(); + } + } + + @Test + public void isDataAvailableWithResumeTokenInGrpcStream() { + initWithLimit(10); + GrpcStreamIterator grpcStream = new GrpcStreamIterator(false, 4, false); + Mockito.when(starter.startStream(null, null)).thenReturn(grpcStream); + resumableStreamIterator.setStream(grpcStream); + + assertFalse(resumableStreamIterator.isDataAvailable()); + + // Chunk with resume token arrives -> immediately ready to emit! + grpcStream.consumer().onPartialResultSet(resultSet(ByteString.copyFromUtf8("r1"), "a")); + assertTrue(resumableStreamIterator.isDataAvailable()); + } + + @Test + public void isDataAvailableWhenStreamIsNull() { + initWithLimit(10); + resumableStreamIterator.setStream(null); + assertFalse(resumableStreamIterator.isDataAvailable()); + } + + @Test + public void isDataAvailableWhenFinished() { + initWithLimit(10); + ResultSetStream stream = mock(ResultSetStream.class); + when(starter.startStream(null, null)).thenReturn(new ResultSetIterator(stream)); + when(stream.next()).thenReturn(null); + assertFalse(resumableStreamIterator.hasNext()); + assertTrue(resumableStreamIterator.isDataAvailable()); + } + + @Test + public void isDataAvailableWhenClosedReturnsTrue() { + initWithLimit(10); + resumableStreamIterator.close("closed"); + assertTrue(resumableStreamIterator.isClosed()); + assertTrue(resumableStreamIterator.isDataAvailable()); + } + + @Test + public void closeDuringStartStreamClosesStreamAndPreventsResurrection() { + initWithLimit(10); + @SuppressWarnings("unchecked") + AbstractResultSet.CloseableIterator streamIterator = + mock(AbstractResultSet.CloseableIterator.class); + when(starter.startStream(null, null)) + .thenAnswer( + invocation -> { + resumableStreamIterator.close("cancelled while starting stream"); + return streamIterator; + }); + + assertFalse(resumableStreamIterator.hasNext()); + assertTrue(resumableStreamIterator.isClosed()); + verify(streamIterator).close("cancelled while starting stream"); + } + + @Test + public void setStreamWhenAlreadyClosedImmediatelyClosesStream() { + initWithLimit(10); + resumableStreamIterator.close("already closed"); + assertTrue(resumableStreamIterator.isClosed()); + + @SuppressWarnings("unchecked") + AbstractResultSet.CloseableIterator streamIterator = + mock(AbstractResultSet.CloseableIterator.class); + assertFalse(resumableStreamIterator.setStream(streamIterator)); + + verify(streamIterator).close("already closed"); + } + + @Test + public void multipleCloseCallsAreIdempotent() { + initWithLimit(10); + @SuppressWarnings("unchecked") + AbstractResultSet.CloseableIterator streamIterator = + mock(AbstractResultSet.CloseableIterator.class); + resumableStreamIterator.setStream(streamIterator); + + resumableStreamIterator.close("first close"); + resumableStreamIterator.close("second close"); + + assertTrue(resumableStreamIterator.isClosed()); + verify(streamIterator, Mockito.times(1)).close("first close"); + } + + @Test + public void closeWhenStreamIsNullEndsSpan() { + SpannerOptions.resetActiveTracingFramework(); + SpannerOptions.enableOpenTelemetryTraces(); + + io.opentelemetry.api.trace.Span openTelemetrySpan = mock(io.opentelemetry.api.trace.Span.class); + ISpan span = new OpenTelemetrySpan(openTelemetrySpan); + TraceWrapper tracer = mock(TraceWrapper.class); + when(tracer.spanBuilderWithExplicitParent(Mockito.anyString(), Mockito.any(), Mockito.any())) + .thenReturn(span); + + ResumableStreamIterator iterator = + new ResumableStreamIterator( + 10, + "testStream", + span, + tracer, + DefaultErrorHandler.INSTANCE, + SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetrySettings(), + SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetryableCodes(), + NoopRequestIdCreator.INSTANCE) { + @Override + AbstractResultSet.CloseableIterator startStream( + @Nullable ByteString resumeToken, + AsyncResultSet.StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + return starter.startStream(resumeToken, streamMessageListener); + } + }; + + iterator.close("cancelled before streaming"); + assertTrue(iterator.isClosed()); + verify(openTelemetrySpan).end(); + } + + @Test + public void closeWithNullMessage() { + initWithLimit(10); + @SuppressWarnings("unchecked") + AbstractResultSet.CloseableIterator streamIterator = + mock(AbstractResultSet.CloseableIterator.class); + resumableStreamIterator.setStream(streamIterator); + + resumableStreamIterator.close(null); + assertTrue(resumableStreamIterator.isClosed()); + verify(streamIterator).close(null); + } + + @Test + public void isDataAvailableWhenBufferHasItems() { + initWithLimit(10); + ResultSetStream s1 = mock(ResultSetStream.class); + when(starter.startStream(null, null)).thenReturn(new ResultSetIterator(s1)); + when(s1.next()) + .thenReturn(resultSet(null, "chunkWithoutToken")) + .thenReturn(resultSet(ByteString.copyFromUtf8("r1"), "chunkWithToken")) + .thenReturn(null); + + // Consuming the first element causes computeNext to buffer both chunks (up to the resume token) + // and pop the first one, leaving chunkWithToken in the buffer. + assertEquals("chunkWithoutToken", resumableStreamIterator.next().getValues(0).getStringValue()); + + // With chunkWithToken still buffered, isDataAvailable() returns true via buffer.isEmpty() == + // false + assertTrue(resumableStreamIterator.isDataAvailable()); + assertEquals("chunkWithToken", resumableStreamIterator.next().getValues(0).getStringValue()); + } + + @Test + public void isWithBeginTransactionDelegatesToStream() { + initWithLimit(10); + assertFalse(resumableStreamIterator.isWithBeginTransaction()); + + @SuppressWarnings("unchecked") + AbstractResultSet.CloseableIterator streamIterator = + mock(AbstractResultSet.CloseableIterator.class); + when(streamIterator.isWithBeginTransaction()).thenReturn(true); + resumableStreamIterator.setStream(streamIterator); + assertTrue(resumableStreamIterator.isWithBeginTransaction()); + + when(streamIterator.isWithBeginTransaction()).thenReturn(false); + assertFalse(resumableStreamIterator.isWithBeginTransaction()); + } + + @Test + public void isLastStatementDelegatesToStream() { + initWithLimit(10); + assertFalse(resumableStreamIterator.isLastStatement()); + + @SuppressWarnings("unchecked") + AbstractResultSet.CloseableIterator streamIterator = + mock(AbstractResultSet.CloseableIterator.class); + when(streamIterator.isLastStatement()).thenReturn(true); + resumableStreamIterator.setStream(streamIterator); + assertTrue(resumableStreamIterator.isLastStatement()); + + when(streamIterator.isLastStatement()).thenReturn(false); + assertFalse(resumableStreamIterator.isLastStatement()); + } + + @Test + public void initiateStreamingStartsStreamAndRegistersListener() { + initWithLimit(10); + AsyncResultSet.StreamMessageListener streamMessageListener = + mock(AsyncResultSet.StreamMessageListener.class); + @SuppressWarnings("unchecked") + AbstractResultSet.CloseableIterator streamIterator = + mock(AbstractResultSet.CloseableIterator.class); + when(starter.startStream(null, streamMessageListener)).thenReturn(streamIterator); + + assertTrue(resumableStreamIterator.initiateStreaming(streamMessageListener)); + verify(starter).startStream(null, streamMessageListener); + verify(streamIterator).requestPrefetchChunks(); + } + + @Test + public void setStreamDoesNotRequestPrefetchChunks() { + initWithLimit(10); + @SuppressWarnings("unchecked") + AbstractResultSet.CloseableIterator streamIterator = + mock(AbstractResultSet.CloseableIterator.class); + resumableStreamIterator.setStream(streamIterator); + verify(streamIterator, Mockito.never()).requestPrefetchChunks(); + } + + @Test + public void computeNextWhenAlreadyClosedReturnsEndOfData() { + initWithLimit(10); + resumableStreamIterator.close("pre-closed"); + assertFalse(resumableStreamIterator.hasNext()); + assertThrows(NoSuchElementException.class, () -> resumableStreamIterator.next()); + } + + @Test + public void retryOnDifferentGrpcChannelWhenSupported() { + ErrorHandler errorHandler = mock(ErrorHandler.class); + when(errorHandler.getMaxAttempts()).thenReturn(2); + when(errorHandler.translateException(Mockito.any())) + .thenReturn(new RetryOnDifferentGrpcChannelException("channel failed", 1, null)); + + final AtomicInteger channelRetries = new AtomicInteger(0); + ResumableStreamIterator channelRetryIterator = + new ResumableStreamIterator( + 10, + "", + new OpenTelemetrySpan(mock(io.opentelemetry.api.trace.Span.class)), + new TraceWrapper(Tracing.getTracer(), OpenTelemetry.noop().getTracer(""), false), + errorHandler, + SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetrySettings(), + SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetryableCodes(), + NoopRequestIdCreator.INSTANCE) { + @Override + AbstractResultSet.CloseableIterator startStream( + @Nullable ByteString resumeToken, + AsyncResultSet.StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + return starter.startStream(resumeToken, streamMessageListener); + } + + @Override + boolean prepareIteratorForRetryOnDifferentGrpcChannel() { + channelRetries.incrementAndGet(); + return true; + } + }; + + ResultSetStream firstStream = mock(ResultSetStream.class); + ResultSetStream secondStream = mock(ResultSetStream.class); + when(starter.startStream(null, null)) + .thenReturn(new ResultSetIterator(firstStream)) + .thenReturn(new ResultSetIterator(secondStream)); + + when(firstStream.next()) + .thenThrow(new NonRetryableException(ErrorCode.UNAVAILABLE, "transient channel error")); + when(secondStream.next()) + .thenReturn(resultSet(ByteString.copyFromUtf8("r1"), "recoveredValue")) + .thenReturn(null); + + List results = consume(channelRetryIterator); + assertEquals(1, results.size()); + assertEquals("recoveredValue", results.get(0)); + assertEquals(1, channelRetries.get()); + } + + @Test + public void retryOnDifferentGrpcChannelNotAttemptedWhenPrepareReturnsFalse() { + ErrorHandler errorHandler = mock(ErrorHandler.class); + when(errorHandler.getMaxAttempts()).thenReturn(2); + when(errorHandler.translateException(Mockito.any())) + .thenReturn(new RetryOnDifferentGrpcChannelException("channel failed", 1, null)); + + ResumableStreamIterator channelRetryIterator = + new ResumableStreamIterator( + 10, + "", + new OpenTelemetrySpan(mock(io.opentelemetry.api.trace.Span.class)), + new TraceWrapper(Tracing.getTracer(), OpenTelemetry.noop().getTracer(""), false), + errorHandler, + SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetrySettings(), + SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetryableCodes(), + NoopRequestIdCreator.INSTANCE) { + @Override + AbstractResultSet.CloseableIterator startStream( + @Nullable ByteString resumeToken, + AsyncResultSet.StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + return starter.startStream(resumeToken, streamMessageListener); + } + }; + + ResultSetStream firstStream = mock(ResultSetStream.class); + when(starter.startStream(null, null)).thenReturn(new ResultSetIterator(firstStream)); + when(firstStream.next()) + .thenThrow(new NonRetryableException(ErrorCode.UNAVAILABLE, "transient error")); + + SpannerException spannerException = + assertThrows(SpannerException.class, channelRetryIterator::next); + assertEquals(ErrorCode.UNAVAILABLE, spannerException.getErrorCode()); + } + + @Test + public void retryOnDifferentGrpcChannelNotAttemptedWhenMaxAttemptsExceeded() { + ErrorHandler errorHandler = mock(ErrorHandler.class); + when(errorHandler.getMaxAttempts()).thenReturn(1); + when(errorHandler.translateException(Mockito.any())) + .thenReturn(new RetryOnDifferentGrpcChannelException("channel failed", 1, null)); + + ResumableStreamIterator channelRetryIterator = + new ResumableStreamIterator( + 10, + "", + new OpenTelemetrySpan(mock(io.opentelemetry.api.trace.Span.class)), + new TraceWrapper(Tracing.getTracer(), OpenTelemetry.noop().getTracer(""), false), + errorHandler, + SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetrySettings(), + SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetryableCodes(), + NoopRequestIdCreator.INSTANCE) { + @Override + AbstractResultSet.CloseableIterator startStream( + @Nullable ByteString resumeToken, + AsyncResultSet.StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + return starter.startStream(resumeToken, streamMessageListener); + } + + @Override + boolean prepareIteratorForRetryOnDifferentGrpcChannel() { + return true; + } + }; + + ResultSetStream firstStream = mock(ResultSetStream.class); + when(starter.startStream(null, null)).thenReturn(new ResultSetIterator(firstStream)); + when(firstStream.next()) + .thenThrow(new NonRetryableException(ErrorCode.UNAVAILABLE, "transient error")); + + SpannerException spannerException = + assertThrows(SpannerException.class, channelRetryIterator::next); + assertEquals(ErrorCode.UNAVAILABLE, spannerException.getErrorCode()); + } + + @Test + public void retryOnDifferentGrpcChannelNotAttemptedWhenResumeTokenPresent() { + ErrorHandler errorHandler = mock(ErrorHandler.class); + when(errorHandler.getMaxAttempts()).thenReturn(5); + when(errorHandler.translateException(Mockito.any())) + .thenReturn(new RetryOnDifferentGrpcChannelException("channel failed", 1, null)); + + final AtomicBoolean prepareCalled = new AtomicBoolean(false); + ResumableStreamIterator channelRetryIterator = + new ResumableStreamIterator( + 10, + "", + new OpenTelemetrySpan(mock(io.opentelemetry.api.trace.Span.class)), + new TraceWrapper(Tracing.getTracer(), OpenTelemetry.noop().getTracer(""), false), + errorHandler, + SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetrySettings(), + SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetryableCodes(), + NoopRequestIdCreator.INSTANCE) { + @Override + AbstractResultSet.CloseableIterator startStream( + @Nullable ByteString resumeToken, + AsyncResultSet.StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + return starter.startStream(resumeToken, streamMessageListener); + } + + @Override + boolean prepareIteratorForRetryOnDifferentGrpcChannel() { + prepareCalled.set(true); + return true; + } + }; + + ResultSetStream firstStream = mock(ResultSetStream.class); + when(starter.startStream(null, null)).thenReturn(new ResultSetIterator(firstStream)); + when(firstStream.next()) + .thenReturn(resultSet(ByteString.copyFromUtf8("r1"), "firstChunk")) + .thenThrow(new NonRetryableException(ErrorCode.UNAVAILABLE, "error after token")); + + assertEquals("firstChunk", channelRetryIterator.next().getValues(0).getStringValue()); + SpannerException spannerException = + assertThrows(SpannerException.class, channelRetryIterator::next); + assertEquals(ErrorCode.UNAVAILABLE, spannerException.getErrorCode()); + assertFalse(prepareCalled.get()); + } + + @Test + public void newBackOffWithDefaultSettings() { + initWithLimit(10); + ExponentialBackOff backOff = resumableStreamIterator.newBackOff(); + assertEquals(1.0, backOff.getMultiplier(), 0.001); + assertEquals(10, backOff.getInitialIntervalMillis()); + assertEquals(1000, backOff.getMaxIntervalMillis()); + assertEquals(Integer.MAX_VALUE, backOff.getMaxElapsedTimeMillis()); + } + + @Test + public void newBackOffWithCustomSettings() { + RetrySettings customRetrySettings = + RetrySettings.newBuilder() + .setInitialRetryDelay(Duration.ofMillis(50)) + .setMaxRetryDelay(Duration.ofMillis(500)) + .setRetryDelayMultiplier(2.0) + .setTotalTimeout(Duration.ofMillis(5000)) + .build(); + ResumableStreamIterator customIterator = + new ResumableStreamIterator( + 10, + "", + new OpenTelemetrySpan(mock(io.opentelemetry.api.trace.Span.class)), + new TraceWrapper(Tracing.getTracer(), OpenTelemetry.noop().getTracer(""), false), + DefaultErrorHandler.INSTANCE, + customRetrySettings, + SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetryableCodes(), + NoopRequestIdCreator.INSTANCE) { + @Override + AbstractResultSet.CloseableIterator startStream( + @Nullable ByteString resumeToken, + AsyncResultSet.StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + return starter.startStream(resumeToken, streamMessageListener); + } + }; + + ExponentialBackOff backOff = customIterator.newBackOff(); + assertEquals(2.0, backOff.getMultiplier(), 0.001); + assertEquals(50, backOff.getInitialIntervalMillis()); + assertEquals(500, backOff.getMaxIntervalMillis()); + assertEquals(Integer.MAX_VALUE, backOff.getMaxElapsedTimeMillis()); + } + + @Test + public void nextBackOffMillisWhenIOExceptionThrowsSpannerException() throws Exception { + BackOff mockBackOff = mock(BackOff.class); + when(mockBackOff.nextBackOffMillis()).thenThrow(new IOException("Simulated I/O error")); + + SpannerException spannerException = + assertThrows( + SpannerException.class, () -> ResumableStreamIterator.nextBackOffMillis(mockBackOff)); + assertEquals(ErrorCode.INTERNAL, spannerException.getErrorCode()); + } + + @Test + public void backoffSleepWhenContextCancelledThrowsCancellationException() { + initWithLimit(10); + ResultSetStream firstStream = mock(ResultSetStream.class); + when(starter.startStream(null, null)).thenReturn(new ResultSetIterator(firstStream)); + when(firstStream.next()) + .thenThrow(new RetryableException(errorCodeParameter, "transient error")); + + Context.CancellableContext cancellableContext = Context.current().withCancellation(); + cancellableContext.cancel(new RuntimeException("context cancelled by test")); + + cancellableContext.run( + () -> { + SpannerException spannerException = + assertThrows(SpannerException.class, resumableStreamIterator::next); + assertEquals(ErrorCode.CANCELLED, spannerException.getErrorCode()); + }); + } + + @Test + public void unhandledRuntimeExceptionSetsSpanStatusAndThrows() { + initWithLimit(10); + ResultSetStream firstStream = mock(ResultSetStream.class); + when(starter.startStream(null, null)).thenReturn(new ResultSetIterator(firstStream)); + when(firstStream.next()).thenThrow(new IllegalStateException("unexpected runtime exception")); + + IllegalStateException thrown = + assertThrows(IllegalStateException.class, resumableStreamIterator::next); + assertEquals("unexpected runtime exception", thrown.getMessage()); + } + + @Test + public void isRetryableMatchesRetryableCodesEvenIfExceptionIsNotRetryable() { + ResumableStreamIterator iteratorWithRetryableCodes = + new ResumableStreamIterator( + 10, + "", + new OpenTelemetrySpan(mock(io.opentelemetry.api.trace.Span.class)), + new TraceWrapper(Tracing.getTracer(), OpenTelemetry.noop().getTracer(""), false), + DefaultErrorHandler.INSTANCE, + SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetrySettings(), + Collections.singleton(Code.UNAVAILABLE), + NoopRequestIdCreator.INSTANCE) { + @Override + AbstractResultSet.CloseableIterator startStream( + @Nullable ByteString resumeToken, + AsyncResultSet.StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + return null; + } + }; + + NonRetryableException exceptionWithUnavailableCode = + new NonRetryableException(ErrorCode.UNAVAILABLE, "non retryable flag but UNAVAILABLE code"); + assertTrue(iteratorWithRetryableCodes.isRetryable(exceptionWithUnavailableCode)); + + NonRetryableException exceptionWithInvalidArgument = + new NonRetryableException(ErrorCode.INVALID_ARGUMENT, "non retryable code"); + assertFalse(iteratorWithRetryableCodes.isRetryable(exceptionWithInvalidArgument)); + } + + @Test + public void constructorNegativeMaxBufferSizeThrowsIllegalArgumentException() { + assertThrows(IllegalArgumentException.class, () -> initWithLimit(-1)); + } + + @Test + public void constructorNullRetrySettingsThrowsNullPointerException() { + assertThrows( + NullPointerException.class, + () -> + new ResumableStreamIterator( + 10, + "", + new OpenTelemetrySpan(mock(io.opentelemetry.api.trace.Span.class)), + new TraceWrapper(Tracing.getTracer(), OpenTelemetry.noop().getTracer(""), false), + DefaultErrorHandler.INSTANCE, + null, + SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetryableCodes(), + NoopRequestIdCreator.INSTANCE) { + @Override + AbstractResultSet.CloseableIterator startStream( + @Nullable ByteString resumeToken, + AsyncResultSet.StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + return null; + } + }); + } + + @Test + public void constructorNullRetryableCodesThrowsNullPointerException() { + assertThrows( + NullPointerException.class, + () -> + new ResumableStreamIterator( + 10, + "", + new OpenTelemetrySpan(mock(io.opentelemetry.api.trace.Span.class)), + new TraceWrapper(Tracing.getTracer(), OpenTelemetry.noop().getTracer(""), false), + DefaultErrorHandler.INSTANCE, + SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetrySettings(), + null, + NoopRequestIdCreator.INSTANCE) { + @Override + AbstractResultSet.CloseableIterator startStream( + @Nullable ByteString resumeToken, + AsyncResultSet.StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + return null; + } + }); + } + + @Test + public void isDataAvailableWhenBufferHasItemsWithoutResumeTokenDelegatesToUnderlyingStream() { + initWithLimit(10); + @SuppressWarnings("unchecked") + AbstractResultSet.CloseableIterator streamIterator = + mock(AbstractResultSet.CloseableIterator.class); + resumableStreamIterator.setStream(streamIterator); + + PartialResultSet chunkWithoutToken = resultSet(null, "chunk1"); + synchronized (resumableStreamIterator.buffer) { + resumableStreamIterator.buffer.add(chunkWithoutToken); + } + + when(streamIterator.isDataAvailable()).thenReturn(false); + assertFalse(resumableStreamIterator.isDataAvailable()); + + when(streamIterator.isDataAvailable()).thenReturn(true); + assertTrue(resumableStreamIterator.isDataAvailable()); + + PartialResultSet chunkWithToken = resultSet(ByteString.copyFromUtf8("token1"), "chunk2"); + synchronized (resumableStreamIterator.buffer) { + resumableStreamIterator.buffer.add(chunkWithToken); + } + when(streamIterator.isDataAvailable()).thenReturn(false); + assertTrue(resumableStreamIterator.isDataAvailable()); + } + + @Test + public void startStreamFailureAfterSetStreamDoesNotLeaveStaleStream() { + AtomicInteger startCount = new AtomicInteger(); + ResumableStreamIterator iterator = + new ResumableStreamIterator( + 10, + "", + new OpenTelemetrySpan(mock(io.opentelemetry.api.trace.Span.class)), + new TraceWrapper(Tracing.getTracer(), OpenTelemetry.noop().getTracer(""), false), + DefaultErrorHandler.INSTANCE, + SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetrySettings(), + SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetryableCodes(), + NoopRequestIdCreator.INSTANCE) { + @Override + AbstractResultSet.CloseableIterator startStream( + @Nullable ByteString resumeToken, + AsyncResultSet.StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + GrpcStreamIterator stream = + new GrpcStreamIterator( + /* lastStatement= */ false, 10, /* cancelQueryWhenClientIsClosed= */ false); + setStream(stream); + if (startCount.incrementAndGet() == 1) { + throw SpannerExceptionFactory.newSpannerException( + ErrorCode.DEADLINE_EXCEEDED, + "Simulated failure in getTransactionSelector() or rpc.executeQuery()"); + } + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + stream.setCall(call, false); + stream.consumer().onPartialResultSet(resultSet(ByteString.copyFromUtf8("r1"), "val1")); + stream.consumer().onCompleted(); + return stream; + } + }; + + assertThrows(SpannerException.class, () -> iterator.initiateStreaming(null)); + // A subsequent attempt to consume or check availability must not be stuck on the unstarted + // GrpcStreamIterator from the failed first attempt. + assertFalse(iterator.isDataAvailable()); + ExecutorService testExecutor = Executors.newSingleThreadExecutor(); + try { + assertTrue(testExecutor.submit(iterator::hasNext).get(1, TimeUnit.SECONDS)); + assertEquals("val1", iterator.next().getValues(0).getStringValue()); + } catch (Exception e) { + throw new AssertionError("iterator.hasNext() got stuck after failed startStream()", e); + } finally { + testExecutor.shutdownNow(); + } + } + + @Test + public void closeBeforeSetCallCancelsUnderlyingCallWithMessage() { + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + ResumableStreamIterator iterator = + new ResumableStreamIterator( + 10, + "", + new OpenTelemetrySpan(mock(io.opentelemetry.api.trace.Span.class)), + new TraceWrapper(Tracing.getTracer(), OpenTelemetry.noop().getTracer(""), false), + DefaultErrorHandler.INSTANCE, + SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetrySettings(), + SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetryableCodes(), + NoopRequestIdCreator.INSTANCE) { + @Override + AbstractResultSet.CloseableIterator startStream( + @Nullable ByteString resumeToken, + AsyncResultSet.StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + GrpcStreamIterator stream = + new GrpcStreamIterator( + /* lastStatement= */ false, 10, /* cancelQueryWhenClientIsClosed= */ false); + setStream(stream); + // Simulate concurrent close(message) while rpc.executeQuery() is in flight before + // stream.setCall(call, ...) is invoked. + close("Cancelled by test"); + stream.setCall(call, false); + return stream; + } + }; + + iterator.initiateStreaming(null); + verify(call).cancel("Cancelled by test"); + } + + @Test + public void isWithBeginTransactionAndIsLastStatementPreservedAfterClose() { + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + ResumableStreamIterator iterator = + new ResumableStreamIterator( + 10, + "", + new OpenTelemetrySpan(mock(io.opentelemetry.api.trace.Span.class)), + new TraceWrapper(Tracing.getTracer(), OpenTelemetry.noop().getTracer(""), false), + DefaultErrorHandler.INSTANCE, + SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetrySettings(), + SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetryableCodes(), + NoopRequestIdCreator.INSTANCE) { + @Override + AbstractResultSet.CloseableIterator startStream( + @Nullable ByteString resumeToken, + AsyncResultSet.StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + GrpcStreamIterator stream = + new GrpcStreamIterator( + /* lastStatement= */ true, 10, /* cancelQueryWhenClientIsClosed= */ false); + setStream(stream); + stream.setCall(call, /* withBeginTransaction= */ true); + return stream; + } + }; + + iterator.initiateStreaming(null); + assertTrue(iterator.isWithBeginTransaction()); + assertTrue(iterator.isLastStatement()); + + iterator.close("closed"); + assertTrue(iterator.isWithBeginTransaction()); + assertTrue(iterator.isLastStatement()); + } + + @Test + public void closeAfterLastTrueWithResumeTokenDoesNotCancelUnderlyingCall() { + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + ResumableStreamIterator iterator = + new ResumableStreamIterator( + 10, + "", + new OpenTelemetrySpan(mock(io.opentelemetry.api.trace.Span.class)), + new TraceWrapper(Tracing.getTracer(), OpenTelemetry.noop().getTracer(""), false), + DefaultErrorHandler.INSTANCE, + SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetrySettings(), + SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetryableCodes(), + NoopRequestIdCreator.INSTANCE) { + @Override + AbstractResultSet.CloseableIterator startStream( + @Nullable ByteString resumeToken, + AsyncResultSet.StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + GrpcStreamIterator stream = + new GrpcStreamIterator( + /* lastStatement= */ false, 10, /* cancelQueryWhenClientIsClosed= */ false); + setStream(stream); + stream.setCall(call, /* withBeginTransaction= */ false); + PartialResultSet.Builder builder = + PartialResultSet.newBuilder() + .setResumeToken(ByteString.copyFromUtf8("r1")) + .setLast(true); + for (int i = 0; i < 20; i++) { + builder.addValuesBuilder().setStringValue("v" + i); + } + stream.consumer().onPartialResultSet(builder.build()); + return stream; + } + }; + + assertTrue(iterator.hasNext()); + PartialResultSet prs = iterator.next(); + assertEquals(20, prs.getValuesCount()); + // Closing before calling hasNext() a second time should not cancel the stream because the + // PartialResultSet had last=true. + iterator.close("closed early"); + Mockito.verify(call, Mockito.never()).cancel(Mockito.any()); + } + + @Test + public void closeWithMultipleChunksCancelsUnlessLastTrueChunkWasConsumed() { + for (boolean consumeSecondChunk : new boolean[] {false, true}) { + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + ResumableStreamIterator iterator = + new ResumableStreamIterator( + 10, + "", + new OpenTelemetrySpan(mock(io.opentelemetry.api.trace.Span.class)), + new TraceWrapper(Tracing.getTracer(), OpenTelemetry.noop().getTracer(""), false), + DefaultErrorHandler.INSTANCE, + SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetrySettings(), + SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetryableCodes(), + NoopRequestIdCreator.INSTANCE) { + @Override + AbstractResultSet.CloseableIterator startStream( + @Nullable ByteString resumeToken, + AsyncResultSet.StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + GrpcStreamIterator stream = + new GrpcStreamIterator( + /* lastStatement= */ false, 10, /* cancelQueryWhenClientIsClosed= */ false); + setStream(stream); + stream.setCall(call, /* withBeginTransaction= */ false); + PartialResultSet firstChunk = + PartialResultSet.newBuilder() + .setResumeToken(ByteString.copyFromUtf8("r1")) + .setLast(false) + .addValues(Value.newBuilder().setStringValue("v1")) + .build(); + PartialResultSet secondChunk = + PartialResultSet.newBuilder() + .setResumeToken(ByteString.copyFromUtf8("r2")) + .setLast(true) + .addValues(Value.newBuilder().setStringValue("v2")) + .build(); + stream.consumer().onPartialResultSet(firstChunk); + stream.consumer().onPartialResultSet(secondChunk); + return stream; + } + }; + + assertTrue(iterator.hasNext()); + assertEquals("v1", iterator.next().getValues(0).getStringValue()); + if (consumeSecondChunk) { + assertTrue(iterator.hasNext()); + assertEquals("v2", iterator.next().getValues(0).getStringValue()); + } + iterator.close("closed"); + if (consumeSecondChunk) { + Mockito.verify(call, Mockito.never()).cancel(Mockito.any()); + } else { + verify(call).cancel("closed"); + } + } + } + + @Test + public void lastTrueConsumedBeforeSetCallDoesNotCancelWhenClosedBeforeOrAfterSetCall() { + for (boolean closeBeforeSetCall : new boolean[] {true, false}) { + SpannerRpc.StreamingCall call = mock(SpannerRpc.StreamingCall.class); + GrpcStreamIterator stream = + new GrpcStreamIterator( + /* lastStatement= */ false, 10, /* cancelQueryWhenClientIsClosed= */ false); + stream + .consumer() + .onPartialResultSet( + PartialResultSet.newBuilder() + .setResumeToken(ByteString.copyFromUtf8("r1")) + .setLast(true) + .addValues(Value.newBuilder().setStringValue("v1")) + .build()); + assertTrue(stream.hasNext()); + assertEquals("v1", stream.next().getValues(0).getStringValue()); + if (closeBeforeSetCall) { + stream.close("closed before setCall"); + } + stream.setCall(call, /* withBeginTransaction= */ false); + if (!closeBeforeSetCall) { + stream.close("closed after setCall"); + } + Mockito.verify(call, Mockito.never()).cancel(Mockito.any()); + } + } + + @Test + public void concurrentResetDuringStartStreamDoesNotOverwriteStreamOrRequestPrefetch() { + @SuppressWarnings("unchecked") + AbstractResultSet.CloseableIterator initialStream = + mock(AbstractResultSet.CloseableIterator.class); + @SuppressWarnings("unchecked") + AbstractResultSet.CloseableIterator retryStream = + mock(AbstractResultSet.CloseableIterator.class); + AtomicInteger startCount = new AtomicInteger(); + + ResumableStreamIterator iterator = + new ResumableStreamIterator( + 10, + "", + new OpenTelemetrySpan(mock(io.opentelemetry.api.trace.Span.class)), + new TraceWrapper(Tracing.getTracer(), OpenTelemetry.noop().getTracer(""), false), + DefaultErrorHandler.INSTANCE, + SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetrySettings(), + SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetryableCodes(), + NoopRequestIdCreator.INSTANCE) { + @Override + AbstractResultSet.CloseableIterator startStream( + @Nullable ByteString resumeToken, + AsyncResultSet.StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { + if (startCount.incrementAndGet() == 1) { + assertTrue(setStream(initialStream)); + // Simulate a concurrent consumer thread resetting and replacing the stream before + // the first startStream() returns to startGrpcStreaming(). + setStream(retryStream); + return initialStream; + } + return retryStream; + } + }; + + iterator.initiateStreaming(null); + Mockito.verify(initialStream, Mockito.never()).requestPrefetchChunks(); + iterator.close("closing"); + verify(retryStream).close("closing"); + Mockito.verify(initialStream, Mockito.never()).close(Mockito.any()); + } + static PartialResultSet resultSet(@Nullable ByteString resumeToken, String... data) { PartialResultSet.Builder builder = PartialResultSet.newBuilder(); if (resumeToken != null) { diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/StreamingUtilTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/StreamingUtilTest.java new file mode 100644 index 000000000000..f49559dc11d8 --- /dev/null +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/StreamingUtilTest.java @@ -0,0 +1,107 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.spanner; + +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class StreamingUtilTest { + + private interface MockStreamingResultSet extends ResultSet, StreamingResultSet {} + + @Test + public void testInitiateStreamingOnStreamingResultSet() { + MockStreamingResultSet streamingResultSet = mock(MockStreamingResultSet.class); + AsyncResultSet.StreamMessageListener listener = + mock(AsyncResultSet.StreamMessageListener.class); + when(streamingResultSet.initiateStreaming(listener)).thenReturn(true); + + assertTrue(StreamingUtil.initiateStreaming(streamingResultSet, listener)); + verify(streamingResultSet).initiateStreaming(listener); + } + + @Test + public void testInitiateStreamingOnNonStreamingResultSetReturnsFalse() { + ResultSet nonStreamingResultSet = mock(ResultSet.class); + AsyncResultSet.StreamMessageListener listener = + mock(AsyncResultSet.StreamMessageListener.class); + + assertFalse(StreamingUtil.initiateStreaming(nonStreamingResultSet, listener)); + } + + @Test + public void testIsDataAvailableOnStreamingResultSet() { + MockStreamingResultSet streamingResultSet = mock(MockStreamingResultSet.class); + when(streamingResultSet.isDataAvailable()).thenReturn(false, true); + + assertFalse(StreamingUtil.isDataAvailable(streamingResultSet)); + assertTrue(StreamingUtil.isDataAvailable(streamingResultSet)); + } + + @Test + public void testIsDataAvailableOnNonStreamingResultSetReturnsTrue() { + ResultSet nonStreamingResultSet = mock(ResultSet.class); + assertTrue(StreamingUtil.isDataAvailable(nonStreamingResultSet)); + } + + @Test + public void testIsDataAvailableOnNestedForwardingResultSet() { + MockStreamingResultSet underlyingStreamingResultSet = mock(MockStreamingResultSet.class); + ForwardingResultSet innerForwarder = new ForwardingResultSet(underlyingStreamingResultSet); + ForwardingResultSet outerForwarder = new ForwardingResultSet(innerForwarder); + + when(underlyingStreamingResultSet.isDataAvailable()).thenReturn(false); + assertFalse(StreamingUtil.isDataAvailable(outerForwarder)); + + when(underlyingStreamingResultSet.isDataAvailable()).thenReturn(true); + assertTrue(StreamingUtil.isDataAvailable(outerForwarder)); + } + + @Test + public void testInitiateStreamingOnForwardingResultSet() { + MockStreamingResultSet underlyingStreamingResultSet = mock(MockStreamingResultSet.class); + ForwardingResultSet forwarder = new ForwardingResultSet(underlyingStreamingResultSet); + AsyncResultSet.StreamMessageListener listener = + mock(AsyncResultSet.StreamMessageListener.class); + + when(underlyingStreamingResultSet.initiateStreaming(listener)).thenReturn(true); + assertTrue(StreamingUtil.initiateStreaming(forwarder, listener)); + verify(underlyingStreamingResultSet).initiateStreaming(listener); + } + + @Test + public void testIsDataAvailableOnForwardingResultSetSubclassOverridingIsDataAvailable() { + ResultSet nonStreamingDelegate = mock(ResultSet.class); + ForwardingResultSet customForwarder = + new ForwardingResultSet(nonStreamingDelegate) { + @Override + public boolean isDataAvailable() { + return false; + } + }; + + assertFalse(StreamingUtil.isDataAvailable(customForwarder)); + } +} diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITAsyncExamplesTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITAsyncExamplesTest.java index 82493bfdbe2b..f70798b84429 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITAsyncExamplesTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITAsyncExamplesTest.java @@ -53,8 +53,10 @@ import org.junit.AfterClass; import org.junit.BeforeClass; import org.junit.ClassRule; +import org.junit.Rule; import org.junit.Test; import org.junit.experimental.categories.Category; +import org.junit.rules.Timeout; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -63,6 +65,7 @@ @RunWith(JUnit4.class) public class ITAsyncExamplesTest { @ClassRule public static IntegrationTestEnv env = new IntegrationTestEnv(); + @Rule public Timeout globalTimeout = Timeout.seconds(120); private static final String TABLE_NAME = "TestTable"; private static final String INDEX_NAME = "TestTableByValue"; private static final List ALL_COLUMNS = Arrays.asList("Key", "StringValue"); @@ -396,12 +399,16 @@ public void pauseResume() throws Exception { switch (resultSet.tryNext()) { case DONE: evenFinished.set(true); + synchronized (lock) { + lock.notifyAll(); + } return CallbackResponse.DONE; case NOT_READY: return CallbackResponse.CONTINUE; case OK: synchronized (lock) { allValues.add(resultSet.getString("StringValue")); + lock.notifyAll(); } evenReturnedFirstRow.countDown(); return CallbackResponse.PAUSE; @@ -409,6 +416,9 @@ public void pauseResume() throws Exception { } } catch (Throwable t) { evenFinished.setException(t); + synchronized (lock) { + lock.notifyAll(); + } return CallbackResponse.DONE; } }); @@ -424,24 +434,35 @@ public void pauseResume() throws Exception { switch (resultSet.tryNext()) { case DONE: unevenFinished.set(true); + synchronized (lock) { + lock.notifyAll(); + } return CallbackResponse.DONE; case NOT_READY: return CallbackResponse.CONTINUE; case OK: synchronized (lock) { allValues.add(resultSet.getString("StringValue")); + lock.notifyAll(); } return CallbackResponse.PAUSE; } } } catch (Throwable t) { unevenFinished.setException(t); + synchronized (lock) { + lock.notifyAll(); + } return CallbackResponse.DONE; } }); while (!(evenFinished.isDone() && unevenFinished.isDone())) { synchronized (lock) { - if (allValues.peekLast() != null) { + if (unevenFinished.isDone()) { + evenRs.resume(); + } else if (evenFinished.isDone()) { + unevenRs.resume(); + } else if (allValues.peekLast() != null) { if (Integer.parseInt(allValues.peekLast().substring(1)) % 2 == 1) { evenRs.resume(); } else { @@ -452,6 +473,9 @@ public void pauseResume() throws Exception { unevenRs.resume(); evenRs.resume(); } + if (!(evenFinished.isDone() && unevenFinished.isDone())) { + lock.wait(100); + } } } } diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITLargeReadTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITLargeReadTest.java index 83d505e21241..fd594567bcd3 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITLargeReadTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITLargeReadTest.java @@ -17,8 +17,13 @@ package com.google.cloud.spanner.it; import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; +import com.google.api.core.ApiFuture; import com.google.cloud.ByteArray; +import com.google.cloud.spanner.AsyncResultSet; +import com.google.cloud.spanner.AsyncResultSet.CallbackResponse; import com.google.cloud.spanner.Database; import com.google.cloud.spanner.DatabaseClient; import com.google.cloud.spanner.Dialect; @@ -38,6 +43,10 @@ import java.util.Collections; import java.util.List; import java.util.Random; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; import org.junit.AfterClass; import org.junit.BeforeClass; import org.junit.ClassRule; @@ -217,4 +226,87 @@ private void validate(ResultSet resultSet) { } assertThat(i).isEqualTo(numRows); } + + @Test + public void readAsync() throws Exception { + try (AsyncResultSet resultSet = + getClient(dialect.dialect) + .singleUse() + .readAsync( + TABLE_NAME, KeySet.all(), Arrays.asList("Key", "Data", "Fingerprint", "Size"))) { + validateAsync(resultSet); + } + } + + @Test + public void readAsyncWithSmallPrefetchChunks() throws Exception { + try (AsyncResultSet resultSet = + getClient(dialect.dialect) + .singleUse() + .readAsync( + TABLE_NAME, + KeySet.all(), + Arrays.asList("Key", "Data", "Fingerprint", "Size"), + Options.prefetchChunks(1))) { + validateAsync(resultSet); + } + } + + @Test + public void queryAsync() throws Exception { + try (AsyncResultSet resultSet = + getClient(dialect.dialect) + .singleUse() + .executeQueryAsync( + Statement.of( + "SELECT Key, Data, Fingerprint, Size FROM " + TABLE_NAME + " ORDER BY Key"))) { + validateAsync(resultSet); + } + } + + @Test + public void queryAsyncWithSmallPrefetchChunks() throws Exception { + try (AsyncResultSet resultSet = + getClient(dialect.dialect) + .singleUse() + .executeQueryAsync( + Statement.of( + "SELECT Key, Data, Fingerprint, Size FROM " + TABLE_NAME + " ORDER BY Key"), + Options.prefetchChunks(1))) { + validateAsync(resultSet); + } + } + + private void validateAsync(AsyncResultSet resultSet) throws Exception { + ExecutorService executor = Executors.newSingleThreadExecutor(); + try { + AtomicInteger rowCount = new AtomicInteger(); + ApiFuture future = + resultSet.setCallback( + executor, + ready -> { + while (true) { + switch (ready.tryNext()) { + case OK: + int index = rowCount.getAndIncrement(); + assertEquals(index, ready.getLong(0)); + ByteArray data = ready.getBytes(1); + assertEquals(ready.getLong(3), data.length()); + assertEquals(ready.getLong(2), hasher.hashBytes(data.toByteArray()).asLong()); + assertTrue(rowCount.get() <= numRows); + break; + case NOT_READY: + return CallbackResponse.CONTINUE; + case DONE: + return CallbackResponse.DONE; + } + } + }); + future.get(120, TimeUnit.SECONDS); + assertEquals(numRows, rowCount.get()); + } finally { + executor.shutdown(); + assertTrue(executor.awaitTermination(60, TimeUnit.SECONDS)); + } + } }