AsyncResultSetImplTest.java

/*
 * Copyright 2020 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 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.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;

import com.google.api.core.ApiFuture;
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.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;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Executor;
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;
import java.util.concurrent.atomic.AtomicInteger;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
import org.mockito.Mockito;
import org.mockito.invocation.InvocationOnMock;
import org.mockito.stubbing.Answer;

@RunWith(JUnit4.class)
public class AsyncResultSetImplTest {
  private ExecutorProvider mockedProvider;
  private ExecutorProvider simpleProvider;

  @Before
  public void setup() {
    mockedProvider = mock(ExecutorProvider.class, withSettings().withoutAnnotations());
    when(mockedProvider.getExecutor()).thenReturn(mock(ScheduledExecutorService.class));
    simpleProvider = SpannerOptions.createAsyncExecutorProvider(1, 1L, TimeUnit.SECONDS);
  }

  @SuppressWarnings("unchecked")
  @Test
  public void close() {
    AsyncResultSetImpl rs =
        new AsyncResultSetImpl(
            mockedProvider, mock(ResultSet.class), AsyncResultSetImpl.DEFAULT_BUFFER_SIZE);
    rs.close();
    // Closing a second time should be a no-op.
    rs.close();

    // The following methods are not allowed to call after closing the result set.
    assertThrows(
        IllegalStateException.class,
        () -> rs.setCallback(mock(Executor.class), mock(ReadyCallback.class)));
    assertThrows(IllegalStateException.class, () -> rs.toList(mock(Function.class)));
    assertThrows(
        IllegalStateException.class,
        () -> rs.toListAsync(mock(Function.class), mock(Executor.class)));

    // The following methods are allowed on a closed result set.
    AsyncResultSetImpl rs2 =
        new AsyncResultSetImpl(
            mockedProvider, mock(ResultSet.class), AsyncResultSetImpl.DEFAULT_BUFFER_SIZE);
    rs2.setCallback(mock(Executor.class), mock(ReadyCallback.class));
    rs2.close();
    rs2.cancel();
    rs2.resume();
  }

  @Test
  public void tryNextNotAllowed() {
    try (AsyncResultSetImpl rs =
        new AsyncResultSetImpl(
            mockedProvider, mock(ResultSet.class), AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) {
      rs.setCallback(mock(Executor.class), mock(ReadyCallback.class));
      IllegalStateException e = assertThrows(IllegalStateException.class, rs::tryNext);
      assertThat(e.getMessage()).contains("tryNext may only be called from a DataReady callback.");
    }
  }

  @Test
  public void toList() {
    ResultSet delegate = mock(ResultSet.class);
    when(delegate.next()).thenReturn(true, true, true, false);
    when(delegate.getCurrentRowAsStruct()).thenReturn(mock(Struct.class));
    try (AsyncResultSetImpl rs =
        new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) {
      List<Object> list = rs.toList(ignored -> new Object());
      assertThat(list).hasSize(3);
    }
  }

  @Test
  public void toListPropagatesError() {
    ResultSet delegate = mock(ResultSet.class);
    when(delegate.next())
        .thenThrow(
            SpannerExceptionFactory.newSpannerException(
                ErrorCode.INVALID_ARGUMENT, "invalid query"));
    try (AsyncResultSetImpl rs =
        new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) {
      SpannerException e =
          assertThrows(SpannerException.class, () -> rs.toList(ignored -> new Object()));
      assertThat(e.getErrorCode()).isEqualTo(ErrorCode.INVALID_ARGUMENT);
      assertThat(e.getMessage()).contains("invalid query");
    }
  }

  @Test
  public void toListAsync() throws InterruptedException, ExecutionException {
    ExecutorService executor = Executors.newFixedThreadPool(1);
    ResultSet delegate = mock(ResultSet.class);
    when(delegate.next()).thenReturn(true, true, true, false);
    when(delegate.getCurrentRowAsStruct()).thenReturn(mock(Struct.class));
    try (AsyncResultSetImpl rs =
        new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) {
      ApiFuture<List<Object>> future = rs.toListAsync(ignored -> new Object(), executor);
      assertThat(future.get()).hasSize(3);
    }
    executor.shutdown();
  }

  @Test
  public void toListAsyncPropagatesError() {
    ExecutorService executor = Executors.newFixedThreadPool(1);
    ResultSet delegate = mock(ResultSet.class);
    when(delegate.next())
        .thenThrow(
            SpannerExceptionFactory.newSpannerException(
                ErrorCode.INVALID_ARGUMENT, "invalid query"));
    try (AsyncResultSetImpl rs =
        new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) {
      ExecutionException e =
          assertThrows(
              ExecutionException.class,
              () -> rs.toListAsync(ignored -> new Object(), executor).get());
      assertThat(e.getCause()).isInstanceOf(SpannerException.class);
      SpannerException se = SpannerExceptionFactory.asSpannerException(e.getCause());
      assertThat(se.getErrorCode()).isEqualTo(ErrorCode.INVALID_ARGUMENT);
      assertThat(se.getMessage()).contains("invalid query");
    }
    executor.shutdown();
  }

  @Test
  public void withCallback() throws InterruptedException {
    Executor executor = Executors.newSingleThreadExecutor();
    ResultSet delegate = mock(ResultSet.class);
    when(delegate.next()).thenReturn(true, true, true, false);
    when(delegate.getCurrentRowAsStruct()).thenReturn(mock(Struct.class));
    final AtomicInteger callbackCounter = new AtomicInteger();
    final AtomicInteger rowCounter = new AtomicInteger();
    final CountDownLatch finishedLatch = new CountDownLatch(1);
    try (AsyncResultSetImpl rs =
        new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) {
      rs.setCallback(
          executor,
          resultSet -> {
            callbackCounter.incrementAndGet();
            CursorState state;
            while ((state = resultSet.tryNext()) == CursorState.OK) {
              rowCounter.incrementAndGet();
            }
            if (state == CursorState.DONE) {
              finishedLatch.countDown();
            }
            return CallbackResponse.CONTINUE;
          });
    }
    finishedLatch.await();
    // There should be between 1 and 5 callbacks, depending on the timing of the threads.
    // Normally, there should be just 1 callback.
    assertThat(callbackCounter.get()).isIn(Range.closed(1, 5));
    assertThat(rowCounter.get()).isEqualTo(3);
  }

  @Test
  public void callbackReceivesError() throws InterruptedException {
    Executor executor = Executors.newSingleThreadExecutor();
    ResultSet delegate = mock(ResultSet.class);
    when(delegate.next())
        .thenThrow(
            SpannerExceptionFactory.newSpannerException(
                ErrorCode.INVALID_ARGUMENT, "invalid query"));
    final BlockingDeque<Exception> receivedErr = new LinkedBlockingDeque<>(1);
    try (AsyncResultSetImpl rs =
        new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) {
      rs.setCallback(
          executor,
          resultSet -> {
            try {
              resultSet.tryNext();
              receivedErr.push(new Exception("missing expected exception"));
            } catch (SpannerException e) {
              receivedErr.push(e);
            }
            return CallbackResponse.DONE;
          });
    }
    Exception e = receivedErr.take();
    assertThat(e).isInstanceOf(SpannerException.class);
    SpannerException se = (SpannerException) e;
    assertThat(se.getErrorCode()).isEqualTo(ErrorCode.INVALID_ARGUMENT);
    assertThat(se.getMessage()).contains("invalid query");
  }

  @Test
  public void callbackReceivesErrorHalfwayThrough() throws InterruptedException {
    Executor executor = Executors.newSingleThreadExecutor();
    ResultSet delegate = mock(ResultSet.class);
    when(delegate.next())
        .thenReturn(true)
        .thenThrow(
            SpannerExceptionFactory.newSpannerException(
                ErrorCode.INVALID_ARGUMENT, "invalid query"));
    when(delegate.getCurrentRowAsStruct()).thenReturn(mock(Struct.class));
    final AtomicInteger rowCount = new AtomicInteger();
    final BlockingDeque<Exception> receivedErr = new LinkedBlockingDeque<>(1);
    try (AsyncResultSetImpl rs =
        new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) {
      rs.setCallback(
          executor,
          resultSet -> {
            try {
              if (resultSet.tryNext() != CursorState.DONE) {
                rowCount.incrementAndGet();
                return CallbackResponse.CONTINUE;
              }
            } catch (SpannerException e) {
              receivedErr.push(e);
            }
            return CallbackResponse.DONE;
          });
    }
    Exception e = receivedErr.take();
    assertThat(e).isInstanceOf(SpannerException.class);
    SpannerException se = (SpannerException) e;
    assertThat(se.getErrorCode()).isEqualTo(ErrorCode.INVALID_ARGUMENT);
    assertThat(se.getMessage()).contains("invalid query");
    assertThat(rowCount.get()).isEqualTo(1);
  }

  @Test
  public void pauseResume() throws InterruptedException {
    Executor executor = Executors.newSingleThreadExecutor();
    ResultSet delegate = mock(ResultSet.class);
    when(delegate.next()).thenReturn(true, true, true, false);
    when(delegate.getCurrentRowAsStruct()).thenReturn(mock(Struct.class));
    final AtomicInteger callbackCounter = new AtomicInteger();
    final BlockingDeque<Object> queue = new LinkedBlockingDeque<>(1);
    final AtomicBoolean finished = new AtomicBoolean(false);
    try (AsyncResultSetImpl rs =
        new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) {
      rs.setCallback(
          executor,
          resultSet -> {
            callbackCounter.incrementAndGet();
            CursorState state = resultSet.tryNext();
            if (state == CursorState.OK) {
              try {
                queue.put(new Object());
              } catch (InterruptedException e) {
                // Finish early if an error occurs.
                return CallbackResponse.DONE;
              }
              return CallbackResponse.PAUSE;
            }
            finished.set(true);
            return CallbackResponse.DONE;
          });
      int rowCounter = 0;
      while (!finished.get()) {
        Object o = queue.poll(1L, TimeUnit.MILLISECONDS);
        if (o != null) {
          rowCounter++;
        }
        rs.resume();
      }
      // There should be exactly 4 callbacks as we only consume one row per callback.
      assertThat(callbackCounter.get()).isEqualTo(4);
      assertThat(rowCounter).isEqualTo(3);
    }
  }

  @Test
  public void testCallbackIsNotCalledWhilePaused() throws InterruptedException, ExecutionException {
    Executor executor = Executors.newSingleThreadExecutor();
    final int simulatedRows = 100;
    ResultSet delegate = mock(ResultSet.class);
    when(delegate.next())
        .thenAnswer(
            new Answer<Boolean>() {
              int row = 0;

              @Override
              public Boolean answer(InvocationOnMock invocation) throws Throwable {
                row++;
                return row <= simulatedRows;
              }
            });
    when(delegate.getCurrentRowAsStruct()).thenReturn(mock(Struct.class));
    final AtomicInteger callbackCounter = new AtomicInteger();
    final BlockingDeque<Object> queue = new LinkedBlockingDeque<>(1);
    final AtomicBoolean paused = new AtomicBoolean();
    // Clearing the paused flag and calling resume() must be atomic with respect to the callback
    // setting the paused flag. Otherwise, the test thread could clear the flag before the callback
    // sets it, and then call resume() after the callback returned PAUSE. The next callback would
    // then see a stale paused flag, even though resume() was correctly called.
    final Object pauseLock = new Object();
    try (AsyncResultSetImpl rs =
        new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) {
      ApiFuture<Void> callbackResult =
          rs.setCallback(
              executor,
              resultSet -> {
                assertFalse(paused.get());
                callbackCounter.incrementAndGet();
                try {
                  switch (resultSet.tryNext()) {
                    case OK:
                      synchronized (pauseLock) {
                        paused.set(true);
                      }
                      queue.put(new Object());
                      return CallbackResponse.PAUSE;
                    case DONE:
                      return CallbackResponse.DONE;
                    case NOT_READY:
                      return CallbackResponse.CONTINUE;
                    default:
                      throw new IllegalStateException();
                  }
                } catch (InterruptedException e) {
                  throw SpannerExceptionFactory.propagateInterrupt(e);
                }
              });
      int rowCounter = 0;
      while (!callbackResult.isDone()) {
        Object o = queue.poll(1L, TimeUnit.MILLISECONDS);
        if (o != null) {
          rowCounter++;
        }
        Thread.yield();
        synchronized (pauseLock) {
          paused.set(false);
          rs.resume();
        }
      }
      // Empty the queue to ensure we count all elements.
      while (queue.poll() != null) {
        rowCounter++;
      }
      // Assert that we can get the result from the callback future without any exceptions. That
      // indicates that the callback function never failed with an unexpected exception.
      assertNull(callbackResult.get());
      assertThat(callbackCounter.get()).isEqualTo(simulatedRows + 1);
      assertThat(rowCounter).isEqualTo(simulatedRows);
    }
  }

  @Test
  public void testCallbackIsNotCalledWhilePausedAndCanceled() {
    ExecutorService executor = Executors.newSingleThreadExecutor();
    StreamingResultSet delegate = mock(StreamingResultSet.class);

    final AtomicInteger callbackCounter = new AtomicInteger();
    ApiFuture<Void> callbackResult;

    try (AsyncResultSetImpl rs =
        new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) {

      when(delegate.initiateStreaming(any(AsyncResultSet.StreamMessageListener.class)))
          .thenAnswer(
              answer -> {
                rs.onStreamMessage(PartialResultSet.newBuilder().build(), false);
                return null;
              });
      callbackResult =
          rs.setCallback(
              executor,
              resultSet -> {
                callbackCounter.getAndIncrement();
                return CallbackResponse.PAUSE;
              });

      rs.cancel();

      SpannerException exception = assertThrows(SpannerException.class, () -> get(callbackResult));
      assertEquals(ErrorCode.CANCELLED, exception.getErrorCode());
      assertEquals(1, callbackCounter.get());
    } finally {
      executor.shutdown();
    }
  }

  @Test
  public void cancel() throws InterruptedException {
    Executor executor = Executors.newSingleThreadExecutor();
    ResultSet delegate = mock(ResultSet.class);
    when(delegate.next()).thenReturn(true, true, true, false);
    when(delegate.getCurrentRowAsStruct()).thenReturn(mock(Struct.class));
    final AtomicInteger callbackCounter = new AtomicInteger();
    final BlockingDeque<Object> queue = new LinkedBlockingDeque<>(1);
    final AtomicBoolean finished = new AtomicBoolean(false);
    try (AsyncResultSetImpl rs =
        new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) {
      rs.setCallback(
          executor,
          resultSet -> {
            callbackCounter.incrementAndGet();
            try {
              CursorState state = resultSet.tryNext();
              if (state == CursorState.OK) {
                try {
                  queue.put(new Object());
                } catch (InterruptedException e) {
                  // Finish early if an error occurs.
                  return CallbackResponse.DONE;
                }
              }
              // Pause after 2 rows to make sure that no more data is consumed until the cancel
              // call has been received.
              return callbackCounter.get() == 2
                  ? CallbackResponse.PAUSE
                  : CallbackResponse.CONTINUE;
            } catch (SpannerException e) {
              if (e.getErrorCode() == ErrorCode.CANCELLED) {
                finished.set(true);
              }
            }
            return CallbackResponse.DONE;
          });
      int rowCounter = 0;
      while (!finished.get()) {
        Object o = queue.poll(1L, TimeUnit.MILLISECONDS);
        if (o != null) {
          rowCounter++;
        }
        if (rowCounter == 2) {
          // Cancel the result set and then resume it to get the cancelled error.
          rs.cancel();
          rs.resume();
        }
      }
      assertThat(callbackCounter.get()).isIn(Range.closed(2, 4));
      assertThat(rowCounter).isIn(Range.closed(2, 3));
    }
  }

  @Test
  public void callbackReturnsError() throws InterruptedException {
    Executor executor = Executors.newSingleThreadExecutor();
    ResultSet delegate = mock(ResultSet.class);
    when(delegate.next()).thenReturn(true, true, true, false);
    when(delegate.getCurrentRowAsStruct()).thenReturn(mock(Struct.class));
    final AtomicInteger callbackCounter = new AtomicInteger();
    try (AsyncResultSetImpl rs =
        new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) {
      rs.setCallback(
          executor,
          resultSet -> {
            callbackCounter.incrementAndGet();
            throw new RuntimeException("async test");
          });
      ExecutionException e = assertThrows(ExecutionException.class, () -> rs.getResult().get());
      assertThat(e.getCause()).isInstanceOf(SpannerException.class);
      SpannerException se = SpannerExceptionFactory.asSpannerException(e.getCause());
      assertThat(se.getErrorCode()).isEqualTo(ErrorCode.UNKNOWN);
      assertThat(se.getMessage()).contains("async test");
      assertThat(callbackCounter.get()).isEqualTo(1);
    }
  }

  @Test
  public void callbackReturnsErrorWhileBufferIsFull_shouldStopProducer() throws Exception {
    ExecutorService executor = Executors.newSingleThreadExecutor();
    // Return more rows than fit in the buffer, so the producer will wait for the buffer to be
    // consumed by the callback. The callback fails instead of consuming any rows.
    final int simulatedRows = AsyncResultSetImpl.DEFAULT_BUFFER_SIZE * 10;
    final AtomicInteger producedRows = new AtomicInteger();
    ResultSet delegate = mock(ResultSet.class);
    when(delegate.next()).thenAnswer(ignored -> producedRows.incrementAndGet() <= simulatedRows);
    when(delegate.getCurrentRowAsStruct()).thenReturn(mock(Struct.class));
    try (AsyncResultSetImpl rs =
        new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) {
      ApiFuture<Void> result =
          rs.setCallback(
              executor,
              ignored -> {
                throw new RuntimeException("async test");
              });
      // The result must complete with the error from the callback. Before the fix, the producer
      // would keep waiting for the (full) buffer to be consumed, and the result never completed.
      ExecutionException e =
          assertThrows(ExecutionException.class, () -> result.get(30L, TimeUnit.SECONDS));
      assertThat(e.getCause()).isInstanceOf(SpannerException.class);
      SpannerException se = SpannerExceptionFactory.asSpannerException(e.getCause());
      assertThat(se.getErrorCode()).isEqualTo(ErrorCode.UNKNOWN);
      assertThat(se.getMessage()).contains("async test");
      // The producer should stop and release the underlying result set.
      Mockito.verify(delegate).close();
      assertThat(producedRows.get()).isLessThan(simulatedRows);
    } finally {
      executor.shutdown();
    }
  }

  @Test
  public void callbackReturnsDoneBeforeEnd_shouldStopIteration() throws Exception {
    Executor executor = Executors.newSingleThreadExecutor();
    ResultSet delegate = mock(ResultSet.class);
    when(delegate.next()).thenReturn(true, true, true, false);
    when(delegate.getCurrentRowAsStruct()).thenReturn(mock(Struct.class));
    try (AsyncResultSetImpl rs =
        new AsyncResultSetImpl(simpleProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) {
      rs.setCallback(
          executor,
          // Not calling resultSet.tryNext() means that it will also never return DONE.
          // Instead the callback indicates that it does not want any more rows.
          ignored -> CallbackResponse.DONE);
      rs.getResult().get(10L, TimeUnit.SECONDS);
    }
  }

  @Test
  public void testOnStreamMessageWhenResumeTokenIsPresent() {
    StreamingResultSet delegate = mock(StreamingResultSet.class);
    try (AsyncResultSetImpl rs =
        new AsyncResultSetImpl(mockedProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) {
      // Marking Streaming as supported
      Mockito.when(
              delegate.initiateStreaming(Mockito.any(AsyncResultSet.StreamMessageListener.class)))
          .thenReturn(true);

      rs.setCallback(Executors.newSingleThreadExecutor(), ignored -> CallbackResponse.DONE);
      rs.onStreamMessage(
          PartialResultSet.newBuilder().addValues(Value.newBuilder().build()).build(), false);

      rs.onStreamMessage(
          PartialResultSet.newBuilder().setResumeToken(ByteString.copyFromUtf8("test")).build(),
          false);
      Mockito.verify(mockedProvider.getExecutor(), times(2)).execute(Mockito.any());
    }
  }

  @Test
  public void testOnStreamMessageWhenCurrentBufferSizeReachedPrefetchChunkSize() {
    StreamingResultSet delegate = mock(StreamingResultSet.class);
    try (AsyncResultSetImpl rs =
        new AsyncResultSetImpl(mockedProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) {
      // Marking Streaming as supported
      Mockito.when(
              delegate.initiateStreaming(Mockito.any(AsyncResultSet.StreamMessageListener.class)))
          .thenReturn(true);

      rs.setCallback(Executors.newSingleThreadExecutor(), ignored -> CallbackResponse.DONE);
      rs.onStreamMessage(
          PartialResultSet.newBuilder().addValues(Value.newBuilder().build()).build(), true);
      Mockito.verify(mockedProvider.getExecutor(), times(2)).execute(Mockito.any());
    }
  }

  @Test
  public void testOnStreamMessageWhenAsyncResultIsCancelled() {
    StreamingResultSet delegate = mock(StreamingResultSet.class);
    try (AsyncResultSetImpl rs =
        new AsyncResultSetImpl(mockedProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) {
      // Marking Streaming as supported
      Mockito.when(
              delegate.initiateStreaming(Mockito.any(AsyncResultSet.StreamMessageListener.class)))
          .thenReturn(true);

      rs.setCallback(Executors.newSingleThreadExecutor(), ignored -> CallbackResponse.DONE);
      rs.cancel();
      rs.onStreamMessage(
          PartialResultSet.newBuilder().addValues(Value.newBuilder().build()).build(), false);
      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<ResultSet> 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<Void> 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<Void> 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<ResultSet> 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<Void> 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<Void> 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<ResultSet> 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));
    }
  }
}