fix: Avoid blocking thread usage in AsyncResultSet

Author: sakthivelmanii

google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java

@@ -768,9 +768,11 @@ ResultSet executeQueryInternalWithOptions(
             rpc.getExecuteQueryRetrySettings(),
             rpc.getExecuteQueryRetryableCodes()) {
           @Override
-          CloseableIterator<PartialResultSet> startStream(@Nullable ByteString resumeToken) {
+          CloseableIterator<PartialResultSet> startStream(@Nullable ByteString resumeToken,
+                                                          AsyncResultSet.StreamMessageListener streamListener) {
             GrpcStreamIterator stream =
                 new GrpcStreamIterator(statement, prefetchChunks, cancelQueryWhenClientIsClosed);
+            stream.registerListener(streamListener);
             if (partitionToken != null) {
               request.setPartitionToken(partitionToken);
             }
@@ -791,8 +793,8 @@ CloseableIterator<PartialResultSet> startStream(@Nullable ByteString resumeToken
                     getTransactionChannelHint(),
                     isRouteToLeader());
             session.markUsed(clock.instant());
-            call.request(prefetchChunks);
             stream.setCall(call, request.getTransaction().hasBegin());
+            call.request(prefetchChunks);
             return stream;
           }
 
@@ -959,9 +961,11 @@ ResultSet readInternalWithOptions(
             rpc.getReadRetrySettings(),
             rpc.getReadRetryableCodes()) {
           @Override
-          CloseableIterator<PartialResultSet> startStream(@Nullable ByteString resumeToken) {
+          CloseableIterator<PartialResultSet> startStream(@Nullable ByteString resumeToken,
+                                                          AsyncResultSet.StreamMessageListener streamListener) {
             GrpcStreamIterator stream =
                 new GrpcStreamIterator(prefetchChunks, cancelQueryWhenClientIsClosed);
+            stream.registerListener(streamListener);
             TransactionSelector selector = null;
             if (resumeToken != null) {
               builder.setResumeToken(resumeToken);
@@ -980,8 +984,8 @@ CloseableIterator<PartialResultSet> startStream(@Nullable ByteString resumeToken
                     getTransactionChannelHint(),
                     isRouteToLeader());
             session.markUsed(clock.instant());
-            call.request(prefetchChunks);
             stream.setCall(call, /* withBeginTransaction = */ builder.getTransaction().hasBegin());
+            call.request(prefetchChunks);
             return stream;
           }
 

google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractResultSet.java

@@ -150,6 +150,10 @@ interface CloseableIterator<T> extends Iterator<T> {
     void close(@Nullable String message);
 
     boolean isWithBeginTransaction();
+
+    default boolean initiateStreaming(AsyncResultSet.StreamMessageListener streamMessageListener) {
+      return false;
+    }
   }
 
   static double valueProtoToFloat64(com.google.protobuf.Value proto) {

google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncResultSet.java

@@ -18,6 +18,8 @@
 
 import com.google.api.core.ApiFuture;
 import com.google.common.base.Function;
+import com.google.spanner.v1.PartialResultSet;
+
 import java.util.List;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Executor;
@@ -223,4 +225,20 @@ interface ReadyCallback {
    * @param transformer function which will be used to transform the row. It should not return null.
    */
   <T> List<T> toList(Function<StructReader, T> transformer) throws SpannerException;
+
+  /**
+   * An interface to register the listener for streaming gRPC request. It will be called when a chunk is received
+   * from gRPC streaming call.
+   */
+  interface StreamMessageListener {
+    void onStreamMessage(PartialResultSet partialResultSet, int prefetchChunks, int currentBufferSize, StreamMessageRequestor streamMessageRequestor);
+  }
+
+  /**
+   * An interface to request more messages from the gRPC streaming call. It will be implemented by the class which has access
+   * to SpannerRpc.StreamingCall object
+   */
+  interface StreamMessageRequestor {
+    void requestMessages(int numOfMessages);
+  }
 }

google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncResultSetImpl.java

@@ -18,7 +18,6 @@
 
 import com.google.api.core.ApiFuture;
 import com.google.api.core.ApiFutures;
-import com.google.api.core.ListenableFutureToApiFuture;
 import com.google.api.core.SettableApiFuture;
 import com.google.api.gax.core.ExecutorProvider;
 import com.google.cloud.spanner.AbstractReadContext.ListenableAsyncResultSet;
@@ -29,28 +28,26 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.util.concurrent.ListeningScheduledExecutorService;
 import com.google.common.util.concurrent.MoreExecutors;
+import com.google.spanner.v1.PartialResultSet;
 import com.google.spanner.v1.ResultSetMetadata;
 import com.google.spanner.v1.ResultSetStats;
+
 import java.util.Collection;
 import java.util.LinkedList;
 import java.util.List;
-import java.util.concurrent.BlockingDeque;
-import java.util.concurrent.Callable;
-import java.util.concurrent.CountDownLatch;
-import java.util.concurrent.ExecutionException;
-import java.util.concurrent.Executor;
-import java.util.concurrent.Future;
-import java.util.concurrent.LinkedBlockingDeque;
+import java.util.concurrent.*;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
 /** Default implementation for {@link AsyncResultSet}. */
-class AsyncResultSetImpl extends ForwardingStructReader implements ListenableAsyncResultSet {
+class AsyncResultSetImpl extends ForwardingStructReader implements ListenableAsyncResultSet, AsyncResultSet.StreamMessageListener {
   private static final Logger log = Logger.getLogger(AsyncResultSetImpl.class.getName());
 
+
   /** State of an {@link AsyncResultSetImpl}. */
   private enum State {
     INITIALIZED,
+    IN_PROGRESS,
     /** SYNC indicates that the {@link ResultSet} is used in sync pattern. */
     SYNC,
     CONSUMING,
@@ -120,7 +117,7 @@ private enum State {
    */
   private volatile boolean finished;
 
-  private volatile ApiFuture<Void> result;
+  private volatile SettableApiFuture<Void> result;
 
   /**
    * This variable indicates whether {@link #tryNext()} has returned {@link CursorState#DONE} or a
@@ -329,12 +326,12 @@ public void run() {
   private final CallbackRunnable callbackRunnable = new CallbackRunnable();
 
   /**
-   * {@link ProduceRowsCallable} reads data from the underlying {@link ResultSet}, places these in
+   * {@link ProduceRowsRunnable} reads data from the underlying {@link ResultSet}, places these in
    * the buffer and dispatches the {@link CallbackRunnable} when data is ready to be consumed.
    */
-  private class ProduceRowsCallable implements Callable<Void> {
+  private class ProduceRowsRunnable implements Runnable {
     @Override
-    public Void call() throws Exception {
+    public void run() {
       boolean stop = false;
       boolean hasNext = false;
       try {
@@ -393,12 +390,17 @@ public Void call() throws Exception {
         }
         // Call the callback if there are still rows in the buffer that need to be processed.
         while (!stop) {
-          waitIfPaused();
-          startCallbackIfNecessary();
-          // Make sure we wait until the callback runner has actually finished.
-          consumingLatch.await();
-          synchronized (monitor) {
-            stop = cursorReturnedDoneOrException;
+          try {
+            waitIfPaused();
+            startCallbackIfNecessary();
+            // Make sure we wait until the callback runner has actually finished.
+            consumingLatch.await();
+            synchronized (monitor) {
+              stop = cursorReturnedDoneOrException;
+            }
+          } catch (InterruptedException e) {
+            result.setException(e);
+            return;
           }
         }
       } finally {
@@ -410,14 +412,16 @@ public Void call() throws Exception {
         }
         synchronized (monitor) {
           if (executionException != null) {
+            result.setException(executionException);
             throw executionException;
           }
           if (state == State.CANCELLED) {
+            result.setException(CANCELLED_EXCEPTION);
             throw CANCELLED_EXCEPTION;
           }
         }
       }
-      return null;
+      result.set(null);
     }
 
     private void waitIfPaused() throws InterruptedException {
@@ -449,6 +453,21 @@ private void startCallbackWithBufferLatchIfNecessary(int bufferLatch) {
     }
   }
 
+  private class InitiateStreamingRunnable implements Runnable {
+
+    @Override
+    public void run() {
+      try {
+        if(!initiateStreaming(AsyncResultSetImpl.this)) {
+          initiateProduceRows();
+        }
+      } catch (SpannerException e) {
+        executionException = e;
+        initiateProduceRows();
+      }
+    }
+  }
+
   /** Sets the callback for this {@link AsyncResultSet}. */
   @Override
   public ApiFuture<Void> setCallback(Executor exec, ReadyCallback cb) {
@@ -458,16 +477,21 @@ public ApiFuture<Void> setCallback(Executor exec, ReadyCallback cb) {
           this.state == State.INITIALIZED, "callback may not be set multiple times");
 
       // Start to fetch data and buffer these.
-      this.result =
-          new ListenableFutureToApiFuture<>(this.service.submit(new ProduceRowsCallable()));
+      this.result = SettableApiFuture.create();
+      this.state = State.IN_PROGRESS;
+      this.service.execute(new InitiateStreamingRunnable());
       this.executor = MoreExecutors.newSequentialExecutor(Preconditions.checkNotNull(exec));
       this.callback = Preconditions.checkNotNull(cb);
-      this.state = State.RUNNING;
       pausedLatch.countDown();
       return result;
     }
   }
 
+  private void initiateProduceRows() {
+    this.service.execute(new ProduceRowsRunnable());
+    this.state = State.RUNNING;
+  }
+
   Future<Void> getResult() {
     return result;
   }
@@ -480,6 +504,7 @@ public void cancel() {
           "cannot cancel a result set without a callback");
       state = State.CANCELLED;
       pausedLatch.countDown();
+      this.result.setException(CANCELLED_EXCEPTION);
     }
   }
 
@@ -578,6 +603,11 @@ public ResultSetMetadata getMetadata() {
     return delegateResultSet.get().getMetadata();
   }
 
+  @Override
+  public boolean initiateStreaming(StreamMessageListener streamMessageListener) {
+    return delegateResultSet.get().initiateStreaming(streamMessageListener);
+  }
+
   @Override
   protected void checkValidState() {
     synchronized (monitor) {
@@ -593,4 +623,21 @@ public Struct getCurrentRowAsStruct() {
     checkValidState();
     return currentRow;
   }
+
+  @Override
+  public void onStreamMessage(PartialResultSet partialResultSet, int prefetchChunks, int currentBufferSize, StreamMessageRequestor streamMessageRequestor) {
+    synchronized (monitor) {
+      if (state == State.IN_PROGRESS) {
+        // if PartialResultSet contains resume token or buffer size is more than configured size or we have reached
+        // end of stream, we can start the thread
+        boolean startJobThread = !partialResultSet.getResumeToken().isEmpty()
+                || currentBufferSize > prefetchChunks || partialResultSet == GrpcStreamIterator.END_OF_STREAM;
+        if (startJobThread){
+          initiateProduceRows();
+        } else {
+          streamMessageRequestor.requestMessages(1);
+        }
+      }
+    }
+  }
 }

google-cloud-spanner/src/main/java/com/google/cloud/spanner/ForwardingResultSet.java

@@ -102,4 +102,9 @@ public ResultSetStats getStats() {
   public ResultSetMetadata getMetadata() {
     return delegate.get().getMetadata();
   }
+
+  @Override
+  public boolean initiateStreaming(AsyncResultSet.StreamMessageListener streamMessageListener) {
+      return delegate.get().initiateStreaming(streamMessageListener);
+  }
 }

google-cloud-spanner/src/main/java/com/google/cloud/spanner/GrpcResultSet.java

@@ -123,6 +123,11 @@ public ResultSetMetadata getMetadata() {
     return metadata;
   }
 
+  @Override
+  public boolean initiateStreaming(AsyncResultSet.StreamMessageListener streamMessageListener) {
+    return iterator.initiateStreaming(streamMessageListener);
+  }
+
   @Override
   public void close() {
     synchronized (this) {

google-cloud-spanner/src/main/java/com/google/cloud/spanner/GrpcStreamIterator.java

@@ -23,24 +23,28 @@
 import com.google.common.collect.AbstractIterator;
 import com.google.common.util.concurrent.Uninterruptibles;
 import com.google.spanner.v1.PartialResultSet;
+import org.threeten.bp.Duration;
+
+import javax.annotation.Nullable;
+import java.util.Optional;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.TimeUnit;
 import java.util.logging.Level;
 import java.util.logging.Logger;
-import javax.annotation.Nullable;
-import org.threeten.bp.Duration;
 
 /** Adapts a streaming read/query call into an iterator over partial result sets. */
 @VisibleForTesting
 class GrpcStreamIterator extends AbstractIterator<PartialResultSet>
-    implements CloseableIterator<PartialResultSet> {
+    implements CloseableIterator<PartialResultSet>, AsyncResultSet.StreamMessageRequestor {
   private static final Logger logger = Logger.getLogger(GrpcStreamIterator.class.getName());
-  private static final PartialResultSet END_OF_STREAM = PartialResultSet.newBuilder().build();
+  public static final PartialResultSet END_OF_STREAM = PartialResultSet.newBuilder().build();
+  private AsyncResultSet.StreamMessageListener streamMessageListener;
 
   private final ConsumerImpl consumer;
   private final BlockingQueue<PartialResultSet> stream;
   private final Statement statement;
+  private final int prefetchChunks;
 
   private SpannerRpc.StreamingCall call;
   private volatile boolean withBeginTransaction;
@@ -57,15 +61,20 @@ class GrpcStreamIterator extends AbstractIterator<PartialResultSet>
   GrpcStreamIterator(
       Statement statement, int prefetchChunks, boolean cancelQueryWhenClientIsClosed) {
     this.statement = statement;
+    this.prefetchChunks = prefetchChunks;
     this.consumer = new ConsumerImpl(cancelQueryWhenClientIsClosed);
     // One extra to allow for END_OF_STREAM message.
-    this.stream = new LinkedBlockingQueue<>(prefetchChunks + 1);
+    this.stream = new LinkedBlockingQueue<>((prefetchChunks * 2) + 1);
   }
 
   protected final SpannerRpc.ResultStreamConsumer consumer() {
     return consumer;
   }
 
+  public void registerListener(AsyncResultSet.StreamMessageListener streamMessageListener) {
+    this.streamMessageListener = streamMessageListener;
+  }
+
   public void setCall(SpannerRpc.StreamingCall call, boolean withBeginTransaction) {
     this.call = call;
     this.withBeginTransaction = withBeginTransaction;
@@ -135,6 +144,12 @@ protected final PartialResultSet computeNext() {
   private void addToStream(PartialResultSet results) {
     // We assume that nothing from the user will interrupt gRPC event threads.
     Uninterruptibles.putUninterruptibly(stream, results);
+    onStreamMessage(results);
+  }
+
+  @Override
+  public void requestMessages(int numOfMessages) {
+    call.request(numOfMessages);
   }
 
   private class ConsumerImpl implements SpannerRpc.ResultStreamConsumer {
@@ -182,4 +197,9 @@ public boolean cancelQueryWhenClientIsClosed() {
       return this.cancelQueryWhenClientIsClosed;
     }
   }
+
+  private void onStreamMessage(PartialResultSet partialResultSet) {
+    Optional.ofNullable(streamMessageListener).ifPresent(
+            sl -> sl.onStreamMessage(partialResultSet, prefetchChunks, stream.size(), this));
+  }
 }

google-cloud-spanner/src/main/java/com/google/cloud/spanner/GrpcValueIterator.java

@@ -127,6 +127,10 @@ ResultSetStats getStats() {
     return statistics;
   }
 
+  boolean initiateStreaming(AsyncResultSet.StreamMessageListener streamMessageListener) {
+    return stream.initiateStreaming(streamMessageListener);
+  }
+
   Type type() {
     checkState(type != null, "metadata has not been received");
     return type;