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..250a2e071452 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 @@ -31,6 +31,7 @@ import com.google.spanner.v1.PartialResultSet; import com.google.spanner.v1.ResultSetMetadata; import com.google.spanner.v1.ResultSetStats; +import java.util.ArrayList; import java.util.Collection; import java.util.LinkedList; import java.util.List; @@ -42,6 +43,7 @@ import java.util.concurrent.LinkedBlockingDeque; import java.util.logging.Level; import java.util.logging.Logger; +import javax.annotation.concurrent.GuardedBy; /** Default implementation for {@link AsyncResultSet}. */ class AsyncResultSetImpl extends ForwardingStructReader @@ -95,6 +97,24 @@ private enum State { /** Supplies the underlying synchronous {@link ResultSet} that will be producing the rows. */ private final Supplier delegateResultSet; + /** + * Callback invoked once when query execution begins or when the result set is closed without + * starting, allowing callers (like {@link MultiUseReadOnlyTransaction}) to track query lifecycle. + */ + private final Runnable onStartOrClose; + + @GuardedBy("monitor") + private boolean startOrCloseNotified; + + @GuardedBy("monitor") + private boolean delegateInitialized; + + @GuardedBy("monitor") + private boolean delegateClosed; + + @GuardedBy("monitor") + private Throwable supplierException; + /** * 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()}. @@ -152,30 +172,107 @@ private enum State { private volatile CountDownLatch consumingLatch = new CountDownLatch(0); AsyncResultSetImpl(ExecutorProvider executorProvider, ResultSet delegate, int bufferSize) { - this(executorProvider, Suppliers.ofInstance(Preconditions.checkNotNull(delegate)), bufferSize); + this( + executorProvider, + Suppliers.ofInstance(Preconditions.checkNotNull(delegate)), + bufferSize, + /* onStartOrClose= */ () -> {}, + /* delegateInitialized= */ true); } AsyncResultSetImpl( ExecutorProvider executorProvider, Supplier delegate, int bufferSize) { + this(executorProvider, delegate, bufferSize, /* onStartOrClose= */ () -> {}); + } + + 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); } + private ResultSet getDelegate() { + try { + synchronized (monitor) { + if (supplierException != null) { + if (supplierException instanceof RuntimeException) { + throw (RuntimeException) supplierException; + } else if (supplierException instanceof Error) { + throw (Error) supplierException; + } + throw SpannerExceptionFactory.asSpannerException(supplierException); + } + } + ResultSet delegate; + try { + delegate = delegateResultSet.get(); + } catch (Throwable t) { + synchronized (monitor) { + supplierException = t; + } + throw t; + } + boolean shouldCloseImmediately = false; + synchronized (monitor) { + if (!delegateInitialized) { + delegateInitialized = true; + if (delegateClosed) { + shouldCloseImmediately = true; + } + } + } + if (shouldCloseImmediately) { + try { + delegate.close(); + } catch (Throwable throwable) { + log.log(Level.FINE, "Ignoring error from closing delegate result set", throwable); + } + } + return delegate; + } finally { + notifyStartOrClose(); + } + } + + private void notifyStartOrClose() { + boolean shouldNotify = false; + synchronized (monitor) { + if (!startOrCloseNotified) { + startOrCloseNotified = true; + shouldNotify = true; + } + } + if (shouldNotify) { + try { + onStartOrClose.run(); + } catch (Throwable throwable) { + log.log(Level.WARNING, "Failed to execute onStartOrClose callback", throwable); + } + } + } + boolean isUsed() { synchronized (monitor) { return state != State.INITIALIZED; @@ -191,15 +288,33 @@ boolean isUsed() { */ @Override public void close() { + boolean shouldCloseDelegate = false; + boolean shouldShutdownService = false; + boolean shouldNotifyClose = false; synchronized (monitor) { if (this.closed) { return; } if (state == State.INITIALIZED || state == State.SYNC) { - delegateResultSet.get().close(); + shouldCloseDelegate = true; + if (executorProvider.shouldAutoClose()) { + shouldShutdownService = true; + } + } + if (state == State.INITIALIZED) { + shouldNotifyClose = true; } this.closed = true; } + if (shouldCloseDelegate) { + closeDelegateResultSet(); + } + if (shouldNotifyClose) { + notifyStartOrClose(); + } + if (shouldShutdownService) { + service.shutdown(); + } } /** @@ -209,14 +324,18 @@ public void close() { */ @Override public void addListener(Runnable listener) { - Preconditions.checkState(state == State.INITIALIZED); - listeners.add(listener); + synchronized (monitor) { + Preconditions.checkState(state == State.INITIALIZED); + listeners.add(listener); + } } @Override public void removeListener(Runnable listener) { - Preconditions.checkState(state == State.INITIALIZED); - listeners.remove(listener); + synchronized (monitor) { + Preconditions.checkState(state == State.INITIALIZED || closed); + listeners.remove(listener); + } } /** @@ -258,10 +377,19 @@ public CursorState tryNext() throws SpannerException { } 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); + } } } @@ -362,15 +490,23 @@ private class ProduceRowsRunnable implements Runnable { public void run() { boolean stop = false; boolean hasNext = false; + ResultSet delegate = null; try { - hasNext = delegateResultSet.get().next(); + synchronized (monitor) { + stop = shouldStopProducing(); + } + if (!stop) { + delegate = getDelegate(); + hasNext = delegate.next(); + } } catch (Throwable e) { synchronized (monitor) { executionException = SpannerExceptionFactory.asSpannerException(e); + stop = true; } } try { - while (!stop && hasNext) { + while (!stop && hasNext && delegate != null) { try { synchronized (monitor) { stop = shouldStopProducing(); @@ -395,9 +531,9 @@ public void run() { } } if (!stop) { - buffer.put(delegateResultSet.get().getCurrentRowAsStruct()); + buffer.put(delegate.getCurrentRowAsStruct()); startCallbackIfNecessary(); - hasNext = delegateResultSet.get().next(); + hasNext = delegate.next(); } } catch (Throwable e) { synchronized (monitor) { @@ -435,7 +571,11 @@ public void run() { if (executorProvider.shouldAutoClose()) { service.shutdown(); } - for (Runnable listener : listeners) { + List listenersToNotify; + synchronized (monitor) { + listenersToNotify = new ArrayList<>(listeners); + } + for (Runnable listener : listenersToNotify) { listener.run(); } synchronized (monitor) { @@ -472,6 +612,7 @@ private void startCallbackIfNecessary() { } private void startCallbackWithBufferLatchIfNecessary(int bufferLatch) { + boolean shouldExecute = false; synchronized (monitor) { if ((state == State.RUNNING || state == State.CANCELLED) && !cursorReturnedDoneOrException) { @@ -482,7 +623,23 @@ private void startCallbackWithBufferLatchIfNecessary(int bufferLatch) { if (state == State.RUNNING) { state = State.CONSUMING; } + shouldExecute = true; + } + } + if (shouldExecute) { + try { executor.execute(callbackRunnable); + } catch (Throwable t) { + synchronized (monitor) { + consumingLatch.countDown(); + while (bufferConsumptionLatch.getCount() > 0L) { + bufferConsumptionLatch.countDown(); + } + if (state == State.CONSUMING) { + state = State.RUNNING; + } + } + throw t; } } } @@ -502,7 +659,9 @@ public void run() { initiateProduceRows(); } } catch (Throwable exception) { - executionException = SpannerExceptionFactory.asSpannerException(exception); + synchronized (monitor) { + executionException = SpannerExceptionFactory.asSpannerException(exception); + } initiateProduceRows(); } } @@ -511,30 +670,63 @@ public void run() { /** Sets the callback for this {@link AsyncResultSet}. */ @Override public ApiFuture setCallback(Executor exec, ReadyCallback cb) { - 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; + Preconditions.checkNotNull(exec, "exec cannot be null"); + Preconditions.checkNotNull(cb, "cb cannot be null"); + boolean shouldNotify = false; + boolean shouldShutdownService = false; + boolean shouldCloseDelegate = false; + try { + 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.executor = MoreExecutors.newSequentialExecutor(exec); + this.callback = cb; + this.state = State.STREAMING_INITIALIZED; + try { + this.service.execute(new InitiateStreamingRunnable()); + } catch (Throwable throwable) { + this.closed = true; + shouldNotify = true; + shouldCloseDelegate = true; + if (executorProvider.shouldAutoClose()) { + shouldShutdownService = true; + } + throw throwable; + } + pausedLatch.countDown(); + return result; + } + } finally { + if (shouldCloseDelegate) { + closeDelegateResultSet(); + } + if (shouldNotify) { + notifyStartOrClose(); + } + if (shouldShutdownService) { + service.shutdown(); + } } } private void initiateProduceRows() { + boolean shouldExecute = false; synchronized (monitor) { - if (this.state == State.STREAMING_INITIALIZED) { - this.state = State.RUNNING; + if (!produceRowsInitiated) { + produceRowsInitiated = true; + if (this.state == State.STREAMING_INITIALIZED) { + this.state = State.RUNNING; + } + shouldExecute = true; } - produceRowsInitiated = true; } - this.service.execute(new ProduceRowsRunnable()); + if (shouldExecute) { + this.service.execute(new ProduceRowsRunnable()); + } } Future getResult() { @@ -601,15 +793,12 @@ public CallbackResponse cursorReady(AsyncResultSet resultSet) { @Override public ApiFuture> toListAsync( Function transformer, Executor executor) { - synchronized (monitor) { - 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()); - } + Preconditions.checkNotNull(transformer, "transformer cannot be null"); + Preconditions.checkNotNull(executor, "executor cannot be null"); + final SettableApiFuture> res = SettableApiFuture.create(); + CreateListCallback callback = new CreateListCallback<>(res, transformer); + ApiFuture finished = setCallback(executor, callback); + return ApiFutures.transformAsync(finished, ignored -> res, MoreExecutors.directExecutor()); } @Override @@ -627,28 +816,41 @@ public List toList(Function transformer) throws SpannerE @Override public boolean next() throws SpannerException { synchronized (monitor) { + Preconditions.checkState(!closed, "This AsyncResultSet has been closed"); Preconditions.checkState( this.state == State.INITIALIZED || this.state == State.SYNC, "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; + ResultSet delegate = getDelegate(); + boolean res = delegate.next(); + currentRow = res ? delegate.getCurrentRowAsStruct() : null; return res; } + /** + * Returns query execution statistics. Calling this method after the result set has closed is + * allowed because callers frequently inspect statistics or query plans after consuming all rows + * and finishing the query. + */ @Override public ResultSetStats getStats() { - return delegateResultSet.get().getStats(); + return getDelegate().getStats(); } + /** + * Returns query metadata. Calling this method after the result set has closed is allowed so + * callers can inspect column and type information after consuming all rows. If the result set was + * closed before execution started, {@link #getDelegate()} ensures the underlying delegate is + * closed immediately to prevent resource leaks. + */ @Override public ResultSetMetadata getMetadata() { - return delegateResultSet.get().getMetadata(); + return getDelegate().getMetadata(); } boolean initiateStreaming(StreamMessageListener streamMessageListener) { - return StreamingUtil.initiateStreaming(delegateResultSet.get(), streamMessageListener); + return StreamingUtil.initiateStreaming(getDelegate(), streamMessageListener); } @Override diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiUseReadOnlyTransaction.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiUseReadOnlyTransaction.java index 838106ed56a6..2bab26666fc2 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiUseReadOnlyTransaction.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiUseReadOnlyTransaction.java @@ -38,7 +38,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.locks.ReentrantLock; import javax.annotation.Nullable; @@ -313,38 +312,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; } } 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..13e4d784a589 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 @@ -22,9 +22,13 @@ import static org.junit.Assert.assertFalse; 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.never; import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import static org.mockito.Mockito.withSettings; @@ -34,10 +38,14 @@ import com.google.cloud.spanner.AsyncResultSet.CursorState; import com.google.cloud.spanner.AsyncResultSet.ReadyCallback; import com.google.common.base.Function; +import com.google.common.base.Supplier; import com.google.common.collect.Range; +import com.google.common.util.concurrent.MoreExecutors; import com.google.protobuf.ByteString; import com.google.protobuf.Value; import com.google.spanner.v1.PartialResultSet; +import com.google.spanner.v1.ResultSetMetadata; +import com.google.spanner.v1.ResultSetStats; import java.util.List; import java.util.concurrent.BlockingDeque; import java.util.concurrent.CountDownLatch; @@ -46,6 +54,7 @@ import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.LinkedBlockingDeque; +import java.util.concurrent.RejectedExecutionException; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; @@ -609,4 +618,356 @@ public void testOnStreamMessageWhenAsyncResultIsCancelled() { Mockito.verify(mockedProvider.getExecutor(), times(2)).execute(Mockito.any()); } } + + @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 testGetStatsAndGetMetadataAfterClose() { + ResultSet delegate = mock(ResultSet.class); + ResultSetStats expectedStats = ResultSetStats.getDefaultInstance(); + ResultSetMetadata expectedMetadata = ResultSetMetadata.getDefaultInstance(); + when(delegate.getStats()).thenReturn(expectedStats); + when(delegate.getMetadata()).thenReturn(expectedMetadata); + + AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE); + asyncResultSet.close(); + assertEquals(expectedStats, asyncResultSet.getStats()); + assertEquals(expectedMetadata, asyncResultSet.getMetadata()); + } + + @Test + public void testGetStatsAndGetMetadataAfterCloseWithoutPriorExecutionClosesDelegate() { + ResultSet delegate = mock(ResultSet.class); + ResultSetStats expectedStats = ResultSetStats.getDefaultInstance(); + ResultSetMetadata expectedMetadata = ResultSetMetadata.getDefaultInstance(); + when(delegate.getStats()).thenReturn(expectedStats); + when(delegate.getMetadata()).thenReturn(expectedMetadata); + + AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl( + simpleProvider, () -> delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE); + asyncResultSet.close(); + verify(delegate, never()).close(); + + assertEquals(expectedStats, asyncResultSet.getStats()); + verify(delegate, times(1)).close(); + + assertEquals(expectedMetadata, asyncResultSet.getMetadata()); + // Closing once is sufficient + verify(delegate, times(1)).close(); + } + + @Test + public void testOnStartOrCloseCalledOnNext() { + ResultSet delegate = mock(ResultSet.class); + when(delegate.next()).thenReturn(true, false); + AtomicInteger count = new AtomicInteger(); + try (AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl(simpleProvider, () -> delegate, 10, count::incrementAndGet)) { + assertEquals(0, count.get()); + assertTrue(asyncResultSet.next()); + assertEquals(1, count.get()); + assertFalse(asyncResultSet.next()); + assertEquals(1, count.get()); + } + assertEquals(1, count.get()); + } + + @Test + public void testOnStartOrCloseCalledOnCloseBeforeStartWithoutInitializingDelegate() { + @SuppressWarnings("unchecked") + Supplier delegateSupplier = mock(Supplier.class); + AtomicInteger count = new AtomicInteger(); + AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl(simpleProvider, delegateSupplier, 10, count::incrementAndGet); + assertEquals(0, count.get()); + asyncResultSet.close(); + assertEquals(1, count.get()); + verify(delegateSupplier, never()).get(); + + // Calling close a second time is idempotent. + asyncResultSet.close(); + assertEquals(1, count.get()); + } + + @Test + public void testGetStatsOrMetadataOnUnstartedResultSetNotifiesStartOrClose() { + ResultSet delegate = mock(ResultSet.class); + ResultSetMetadata expectedMetadata = ResultSetMetadata.getDefaultInstance(); + when(delegate.getMetadata()).thenReturn(expectedMetadata); + AtomicInteger count = new AtomicInteger(); + + try (AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl(simpleProvider, () -> delegate, 10, count::incrementAndGet)) { + assertEquals(0, count.get()); + assertEquals(expectedMetadata, asyncResultSet.getMetadata()); + assertEquals(1, count.get()); + } + assertEquals(1, count.get()); + } + + @Test + public void testConstructorWithoutOnStartOrCloseCallback() { + ResultSet delegate = mock(ResultSet.class); + when(delegate.next()).thenReturn(true, false); + try (AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl(simpleProvider, () -> delegate, 10)) { + assertTrue(asyncResultSet.next()); + assertFalse(asyncResultSet.next()); + } + } + + @Test + public void testOnStartOrCloseCalledWhenDelegateSupplierFailsOnGetMetadata() { + AtomicInteger count = new AtomicInteger(); + try (AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl( + simpleProvider, + () -> { + throw new RuntimeException("delegate creation failed"); + }, + 10, + count::incrementAndGet)) { + assertEquals(0, count.get()); + assertThrows(RuntimeException.class, asyncResultSet::getMetadata); + assertEquals(1, count.get()); + } + } + + @Test + public void testOnStartOrCloseCalledOnSetCallback() throws Exception { + ResultSet delegate = mock(ResultSet.class); + when(delegate.next()).thenReturn(false); + AtomicInteger count = new AtomicInteger(); + try (AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl(simpleProvider, () -> delegate, 10, count::incrementAndGet)) { + assertEquals(0, count.get()); + ApiFuture future = + asyncResultSet.setCallback( + MoreExecutors.directExecutor(), ignored -> CallbackResponse.DONE); + future.get(10, TimeUnit.SECONDS); + assertEquals(1, count.get()); + } + } + + @Test + public void testOnStartOrCloseCalledOnInitiateStreamingFailure() { + AtomicInteger count = new AtomicInteger(); + AtomicInteger supplierInvocations = new AtomicInteger(); + try (AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl( + simpleProvider, + () -> { + supplierInvocations.incrementAndGet(); + throw new RuntimeException("stream start failure"); + }, + 10, + count::incrementAndGet)) { + assertEquals(0, count.get()); + ApiFuture future = + asyncResultSet.setCallback( + MoreExecutors.directExecutor(), ignored -> CallbackResponse.DONE); + assertThrows(ExecutionException.class, () -> future.get(10, TimeUnit.SECONDS)); + assertEquals(1, count.get()); + assertEquals(1, supplierInvocations.get()); + } + } + + @Test + public void testNextAfterCloseThrowsIllegalStateException() { + ResultSet delegate = mock(ResultSet.class); + AsyncResultSetImpl asyncResultSet = new AsyncResultSetImpl(simpleProvider, () -> delegate, 10); + asyncResultSet.close(); + IllegalStateException exception = + assertThrows(IllegalStateException.class, asyncResultSet::next); + assertEquals("This AsyncResultSet has been closed", exception.getMessage()); + } + + @Test + public void testCloseAfterSetCallbackDoesNotNotifyBeforeDelegateStarts() throws Exception { + CountDownLatch startSupplierLatch = new CountDownLatch(1); + CountDownLatch inSupplierLatch = new CountDownLatch(1); + ResultSet delegate = mock(ResultSet.class); + when(delegate.next()).thenReturn(false); + AtomicInteger count = new AtomicInteger(); + Supplier supplier = + () -> { + inSupplierLatch.countDown(); + try { + startSupplierLatch.await(); + } catch (InterruptedException e) { + throw new RuntimeException(e); + } + return delegate; + }; + + AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl(simpleProvider, supplier, 10, count::incrementAndGet); + ApiFuture future = + asyncResultSet.setCallback( + MoreExecutors.directExecutor(), ignored -> CallbackResponse.DONE); + + assertTrue(inSupplierLatch.await(5, TimeUnit.SECONDS)); + // At this point, the delegate supplier is executing, but has not completed. + // Closing the AsyncResultSet while streaming is starting should not immediately trigger + // onStartOrClose because streaming has already been initiated. + asyncResultSet.close(); + assertEquals(0, count.get()); + + // Allow the supplier to complete. + startSupplierLatch.countDown(); + future.get(10, TimeUnit.SECONDS); + + // Now count should be 1. + assertEquals(1, count.get()); + } + + @Test + public void testCloseImmediatelyAfterSetCallback() throws Exception { + ResultSet delegate = mock(ResultSet.class); + when(delegate.next()).thenReturn(false); + AtomicInteger count = new AtomicInteger(); + AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl(simpleProvider, () -> delegate, 10, count::incrementAndGet); + ApiFuture future = + asyncResultSet.setCallback( + MoreExecutors.directExecutor(), ignored -> CallbackResponse.DONE); + asyncResultSet.close(); + future.get(10, TimeUnit.SECONDS); + assertEquals(1, count.get()); + verify(delegate).close(); + } + + @Test + public void testOnStartOrCloseExceptionDoesNotMaskPrimaryException() { + Supplier delegateSupplier = + () -> { + throw SpannerExceptionFactory.newSpannerException( + ErrorCode.NOT_FOUND, "database not found"); + }; + Runnable throwingCallback = + () -> { + throw new RuntimeException("boom in onStartOrClose"); + }; + + try (AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl(simpleProvider, delegateSupplier, 10, throwingCallback)) { + SpannerException exception = assertThrows(SpannerException.class, asyncResultSet::next); + assertEquals(ErrorCode.NOT_FOUND, exception.getErrorCode()); + assertTrue(exception.getMessage().contains("database not found")); + } + } + + @Test + public void testSetCallbackRejectsNullArguments() { + ResultSet delegate = mock(ResultSet.class); + AtomicInteger count = new AtomicInteger(); + AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl(simpleProvider, () -> delegate, 10, count::incrementAndGet); + + assertThrows( + NullPointerException.class, + () -> asyncResultSet.setCallback(null, ignored -> CallbackResponse.DONE)); + assertThrows( + NullPointerException.class, + () -> asyncResultSet.setCallback(MoreExecutors.directExecutor(), null)); + + // Because arguments were null, setCallback failed before mutating state. + // The result set should still be closeable cleanly, notifying start or close once. + assertEquals(0, count.get()); + asyncResultSet.close(); + assertEquals(1, count.get()); + } + + @Test + public void testSetCallbackRejectedExecutionNotifiesStartOrClose() { + ScheduledExecutorService mockService = mock(ScheduledExecutorService.class); + doThrow(new RejectedExecutionException()).when(mockService).execute(any(Runnable.class)); + ExecutorProvider provider = mock(ExecutorProvider.class, withSettings().withoutAnnotations()); + when(provider.getExecutor()).thenReturn(mockService); + when(provider.shouldAutoClose()).thenReturn(true); + + AtomicInteger count = new AtomicInteger(); + AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl(provider, () -> mock(ResultSet.class), 10, count::incrementAndGet); + + assertThrows( + RejectedExecutionException.class, + () -> + asyncResultSet.setCallback( + MoreExecutors.directExecutor(), ignored -> CallbackResponse.DONE)); + + // Rejection should immediately trigger start or close notification so that callers + // waiting on query initialization are unblocked. + assertEquals(1, count.get()); + verify(mockService).shutdown(); + + // Closing afterwards is a no-op and does not notify a second time. + asyncResultSet.close(); + assertEquals(1, count.get()); + } + + @Test + public void testSetCallbackRejectedExecutionClosesInitializedDelegate() { + ScheduledExecutorService mockService = mock(ScheduledExecutorService.class); + doThrow(new RejectedExecutionException()).when(mockService).execute(any(Runnable.class)); + ExecutorProvider provider = mock(ExecutorProvider.class, withSettings().withoutAnnotations()); + when(provider.getExecutor()).thenReturn(mockService); + when(provider.shouldAutoClose()).thenReturn(true); + + ResultSet delegate = mock(ResultSet.class); + try (AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl(provider, () -> delegate, 10, () -> {})) { + // Access metadata to initialize the delegate eagerly + assertNull(asyncResultSet.getMetadata()); + + assertThrows( + RejectedExecutionException.class, + () -> + asyncResultSet.setCallback( + MoreExecutors.directExecutor(), ignored -> CallbackResponse.DONE)); + } + + verify(delegate).close(); + verify(mockService).shutdown(); + } + + @Test + public void testRemoveListenerAfterCloseDoesNotThrow() { + ResultSet delegate = mock(ResultSet.class); + AsyncResultSetImpl asyncResultSet = new AsyncResultSetImpl(simpleProvider, () -> delegate, 10); + Runnable listener = () -> {}; + asyncResultSet.addListener(listener); + asyncResultSet.close(); + + // Calling removeListener after close should not throw IllegalStateException + asyncResultSet.removeListener(listener); + } + + @Test + public void testToListAsyncRejectsNullArguments() { + ResultSet delegate = mock(ResultSet.class); + try (AsyncResultSetImpl asyncResultSet = + new AsyncResultSetImpl(simpleProvider, () -> delegate, 10)) { + assertThrows( + NullPointerException.class, + () -> asyncResultSet.toListAsync(null, MoreExecutors.directExecutor())); + assertThrows( + NullPointerException.class, () -> asyncResultSet.toListAsync(reader -> null, null)); + } + } }