Addressed comments

Author: sakthivelmanii

google-cloud-spanner/clirr-ignored-differences.xml

@@ -758,7 +758,7 @@
     <className>com/google/cloud/spanner/connection/Connection</className>
     <method>boolean isKeepTransactionAlive()</method>
   </difference>
-  
+
   <!-- Automatic DML batching -->
   <difference>
     <differenceType>7012</differenceType>
@@ -790,9 +790,5 @@
     <className>com/google/cloud/spanner/connection/Connection</className>
     <method>boolean isAutoBatchDmlUpdateCountVerification()</method>
   </difference>
-  <difference>
-    <differenceType>7012</differenceType>
-    <className>com/google/cloud/spanner/ResultSet</className>
-    <method>boolean initiateStreaming(com.google.cloud.spanner.AsyncResultSet$StreamMessageListener)</method>
-  </difference>
-</differences>
+
+</differences>

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

@@ -232,8 +232,7 @@ interface ReadyCallback {
   interface StreamMessageListener {
     void onStreamMessage(
         PartialResultSet partialResultSet,
-        int prefetchChunks,
-        int currentBufferSize,
+        boolean bufferIsFull,
         StreamMessageRequestor streamMessageRequestor);
   }
 

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

@@ -34,7 +34,12 @@
 import java.util.Collection;
 import java.util.LinkedList;
 import java.util.List;
-import java.util.concurrent.*;
+import java.util.concurrent.BlockingDeque;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Executor;
+import java.util.concurrent.Future;
+import java.util.concurrent.LinkedBlockingDeque;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
@@ -46,7 +51,7 @@ class AsyncResultSetImpl extends ForwardingStructReader
   /** State of an {@link AsyncResultSetImpl}. */
   private enum State {
     INITIALIZED,
-    IN_PROGRESS,
+    STREAMING_INITIALIZED,
     /** SYNC indicates that the {@link ResultSet} is used in sync pattern. */
     SYNC,
     CONSUMING,
@@ -400,7 +405,7 @@ public void run() {
             synchronized (monitor) {
               stop = cursorReturnedDoneOrException;
             }
-          } catch (InterruptedException e) {
+          } catch (Throwable e) {
             result.setException(e);
             return;
           }
@@ -415,15 +420,13 @@ public void run() {
         synchronized (monitor) {
           if (executionException != null) {
             result.setException(executionException);
-            throw executionException;
-          }
-          if (state == State.CANCELLED) {
+          } else if (state == State.CANCELLED) {
             result.setException(CANCELLED_EXCEPTION);
-            throw CANCELLED_EXCEPTION;
+          } else {
+            result.set(null);
           }
         }
       }
-      result.set(null);
     }
 
     private void waitIfPaused() throws InterruptedException {
@@ -460,6 +463,12 @@ private class InitiateStreamingRunnable implements Runnable {
     @Override
     public void run() {
       try {
+        // This method returns true if the underlying result set is a streaming result set (e.g. a
+        // GrpcResultSet).
+        // Those result sets will trigger initiateProduceRows() when the first results are received.
+        // Non-streaming result sets do not trigger this callback, and for those result sets, we
+        // need to eagerly
+        // start the ProduceRowsRunnable.
         if (!initiateStreaming(AsyncResultSetImpl.this)) {
           initiateProduceRows();
         }
@@ -480,7 +489,7 @@ public ApiFuture<Void> setCallback(Executor exec, ReadyCallback cb) {
 
       // Start to fetch data and buffer these.
       this.result = SettableApiFuture.create();
-      this.state = State.IN_PROGRESS;
+      this.state = State.STREAMING_INITIALIZED;
       this.service.execute(new InitiateStreamingRunnable());
       this.executor = MoreExecutors.newSequentialExecutor(Preconditions.checkNotNull(exec));
       this.callback = Preconditions.checkNotNull(cb);
@@ -491,7 +500,7 @@ public ApiFuture<Void> setCallback(Executor exec, ReadyCallback cb) {
 
   private void initiateProduceRows() {
     this.service.execute(new ProduceRowsRunnable());
-    if (this.state == State.IN_PROGRESS) {
+    if (this.state == State.STREAMING_INITIALIZED) {
       this.state = State.RUNNING;
     }
     produceRowsInitiated = true;
@@ -607,9 +616,8 @@ public ResultSetMetadata getMetadata() {
     return delegateResultSet.get().getMetadata();
   }
 
-  @Override
   public boolean initiateStreaming(StreamMessageListener streamMessageListener) {
-    return delegateResultSet.get().initiateStreaming(streamMessageListener);
+    return StreamingUtil.initiateStreaming(delegateResultSet.get(), streamMessageListener);
   }
 
   @Override
@@ -631,20 +639,19 @@ public Struct getCurrentRowAsStruct() {
   @Override
   public void onStreamMessage(
       PartialResultSet partialResultSet,
-      int prefetchChunks,
-      int currentBufferSize,
+      boolean bufferIsFull,
       StreamMessageRequestor streamMessageRequestor) {
     synchronized (monitor) {
       if (produceRowsInitiated) {
         return;
       }
-      // if PartialResultSet contains resume token or buffer size is more than configured size or
+      // if PartialResultSet contains resume token or buffer size is ful or
       // we have reached end of stream, we can start the thread
       boolean startJobThread =
           !partialResultSet.getResumeToken().isEmpty()
-              || currentBufferSize >= prefetchChunks
+              || bufferIsFull
               || partialResultSet == GrpcStreamIterator.END_OF_STREAM;
-      if (startJobThread || state != State.IN_PROGRESS) {
+      if (startJobThread || state != State.STREAMING_INITIALIZED) {
         initiateProduceRows();
       } else {
         streamMessageRequestor.requestMessages(1);

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

@@ -23,7 +23,8 @@
 import com.google.spanner.v1.ResultSetStats;
 
 /** Forwarding implementation of ResultSet that forwards all calls to a delegate. */
-public class ForwardingResultSet extends ForwardingStructReader implements ProtobufResultSet {
+public class ForwardingResultSet extends ForwardingStructReader
+    implements ProtobufResultSet, StreamingResultSet {
 
   private Supplier<? extends ResultSet> delegate;
 
@@ -105,6 +106,6 @@ public ResultSetMetadata getMetadata() {
 
   @Override
   public boolean initiateStreaming(AsyncResultSet.StreamMessageListener streamMessageListener) {
-    return delegate.get().initiateStreaming(streamMessageListener);
+    return StreamingUtil.initiateStreaming(delegate.get(), streamMessageListener);
   }
 }

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

@@ -30,7 +30,8 @@
 import javax.annotation.Nullable;
 
 @VisibleForTesting
-class GrpcResultSet extends AbstractResultSet<List<Object>> implements ProtobufResultSet {
+class GrpcResultSet extends AbstractResultSet<List<Object>>
+    implements ProtobufResultSet, StreamingResultSet {
   private final GrpcValueIterator iterator;
   private final Listener listener;
   private final DecodeMode decodeMode;

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

@@ -199,6 +199,7 @@ public boolean cancelQueryWhenClientIsClosed() {
 
   private void onStreamMessage(PartialResultSet partialResultSet) {
     Optional.ofNullable(streamMessageListener)
-        .ifPresent(sl -> sl.onStreamMessage(partialResultSet, prefetchChunks, stream.size(), this));
+        .ifPresent(
+            sl -> sl.onStreamMessage(partialResultSet, prefetchChunks >= stream.size(), this));
   }
 }

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

@@ -82,12 +82,4 @@ public interface ResultSet extends AutoCloseable, StructReader {
   default ResultSetMetadata getMetadata() {
     throw new UnsupportedOperationException("Method should be overridden");
   }
-
-  /**
-   * Returns the {@link boolean} for this {@link ResultSet}. This method will be used by
-   * AsyncResultSet to initiate gRPC streaming
-   */
-  default boolean initiateStreaming(AsyncResultSet.StreamMessageListener streamMessageListener) {
-    return false;
-  }
 }

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

@@ -292,18 +292,17 @@ public boolean next() throws SpannerException {
         public boolean initiateStreaming(
             AsyncResultSet.StreamMessageListener streamMessageListener) {
           try {
-            boolean ret = super.initiateStreaming(streamMessageListener);
-            if (beforeFirst) {
+            boolean streamInitiated = super.initiateStreaming(streamMessageListener);
+            if (!streamInitiated) {
               synchronized (lock) {
                 session.get().markUsed();
-                beforeFirst = false;
                 sessionUsedForQuery = true;
               }
             }
-            if (!ret && isSingleUse) {
+            if (!streamInitiated && isSingleUse) {
               close();
             }
-            return ret;
+            return streamInitiated;
           } catch (SessionNotFoundException e) {
             throw e;
           } catch (SpannerException e) {

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

@@ -0,0 +1,27 @@
+/*
+ * Copyright 2024 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;
+
+/** Streaming implementation of ResultSet that supports streaming of chunks */
+public interface StreamingResultSet extends ResultSet {
+
+  /**
+   * Returns the {@link boolean} for this {@link ResultSet}. This method will be used by
+   * AsyncResultSet to initiate gRPC streaming
+   */
+  boolean initiateStreaming(AsyncResultSet.StreamMessageListener streamMessageListener);
+}

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

@@ -0,0 +1,27 @@
+/*
+ * Copyright 2024 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;
+
+public class StreamingUtil {
+  public static boolean initiateStreaming(
+      ResultSet resultSet, AsyncResultSet.StreamMessageListener streamMessageListener) {
+    if (resultSet instanceof StreamingResultSet) {
+      return ((StreamingResultSet) resultSet).initiateStreaming(streamMessageListener);
+    }
+    return false;
+  }
+}