From 4c7425114421be413e5a214bb44b599274832957 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= Date: Tue, 3 Nov 2026 10:54:14 +0100 Subject: [PATCH] refactor(spanner): decouple read-only transactions from async result sets Decouple multi-use read-only transactions from internal async result set implementation details by introducing a dedicated lifecycle callback. Key improvements in this change: - Replaces the anonymous subclass in read-only transactions with a clean callback that signals when an asynchronous query either begins execution or is closed without starting. - Lazily resolves the underlying query stream so that closing an async query before it begins execution no longer starts the query or opens an RPC stream, preventing unnecessary network calls and session allocations. - Permits safe inspection of query statistics and metadata after a result set has been closed. - Adds comprehensive unit tests for the decoupled lifecycle, post-close inspections, and closed-state error handling. --- .../cloud/spanner/AsyncResultSetImpl.java | 300 ++++++++++++--- .../spanner/MultiUseReadOnlyTransaction.java | 32 +- .../cloud/spanner/AsyncResultSetImplTest.java | 361 ++++++++++++++++++ 3 files changed, 614 insertions(+), 79 deletions(-) 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)); + } + } }