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,5 +790,5 @@
     <className>com/google/cloud/spanner/connection/Connection</className>
     <method>boolean isAutoBatchDmlUpdateCountVerification()</method>
   </difference>
-
-</differences>
+  
+</differences>

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

@@ -151,6 +151,10 @@ interface CloseableIterator<T> extends Iterator<T> {
 
     boolean isWithBeginTransaction();
 
+    /**
+     * @param streamMessageListener A class object which implements StreamMessageListener
+     * @return true if streaming is supported by the iterator, otherwise false
+     */
     default boolean initiateStreaming(AsyncResultSet.StreamMessageListener streamMessageListener) {
       return false;
     }

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

@@ -467,13 +467,12 @@ public void run() {
         // 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.
+        // need to eagerly start the ProduceRowsRunnable.
         if (!initiateStreaming(AsyncResultSetImpl.this)) {
           initiateProduceRows();
         }
-      } catch (SpannerException e) {
-        executionException = e;
+      } catch (Throwable exception) {
+        executionException = SpannerExceptionFactory.asSpannerException(exception);
         initiateProduceRows();
       }
     }
@@ -499,11 +498,11 @@ public ApiFuture<Void> setCallback(Executor exec, ReadyCallback cb) {
   }
 
   private void initiateProduceRows() {
-    this.service.execute(new ProduceRowsRunnable());
     if (this.state == State.STREAMING_INITIALIZED) {
       this.state = State.RUNNING;
     }
     produceRowsInitiated = true;
+    this.service.execute(new ProduceRowsRunnable());
   }
 
   Future<Void> getResult() {
@@ -642,8 +641,8 @@ public void onStreamMessage(PartialResultSet partialResultSet, boolean bufferIsF
       if (produceRowsInitiated) {
         return;
       }
-      // if PartialResultSet contains resume token or buffer size is ful or
-      // we have reached end of stream, we can start the thread
+      // if PartialResultSet contains a resume token or buffer size is full, or
+      // we have reached the end of the stream, we can start the thread.
       boolean startJobThread =
           !partialResultSet.getResumeToken().isEmpty()
               || bufferIsFull

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

@@ -37,13 +37,12 @@
 class GrpcStreamIterator extends AbstractIterator<PartialResultSet>
     implements CloseableIterator<PartialResultSet> {
   private static final Logger logger = Logger.getLogger(GrpcStreamIterator.class.getName());
-  public static final PartialResultSet END_OF_STREAM = PartialResultSet.newBuilder().build();
+  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;
@@ -60,10 +59,9 @@ 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 * 2) + 1);
+    this.stream = new LinkedBlockingQueue<>(prefetchChunks + 1);
   }
 
   protected final SpannerRpc.ResultStreamConsumer consumer() {
@@ -194,6 +192,6 @@ public boolean cancelQueryWhenClientIsClosed() {
 
   private void onStreamMessage(PartialResultSet partialResultSet) {
     Optional.ofNullable(streamMessageListener)
-        .ifPresent(sl -> sl.onStreamMessage(partialResultSet, prefetchChunks >= stream.size()));
+        .ifPresent(sl -> sl.onStreamMessage(partialResultSet, stream.remainingCapacity() <= 1));
   }
 }

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

@@ -288,34 +288,6 @@ public boolean next() throws SpannerException {
           }
         }
 
-        @Override
-        public boolean initiateStreaming(
-            AsyncResultSet.StreamMessageListener streamMessageListener) {
-          try {
-            boolean streamInitiated = super.initiateStreaming(streamMessageListener);
-            if (!streamInitiated) {
-              synchronized (lock) {
-                session.get().markUsed();
-                sessionUsedForQuery = true;
-              }
-            }
-            if (!streamInitiated && isSingleUse) {
-              close();
-            }
-            return streamInitiated;
-          } catch (SessionNotFoundException e) {
-            throw e;
-          } catch (SpannerException e) {
-            synchronized (lock) {
-              if (!closed && isSingleUse) {
-                session.get().setLastException(e);
-                AutoClosingReadContext.this.close();
-              }
-            }
-            throw e;
-          }
-        }
-
         private boolean internalNext() {
           try {
             boolean ret = super.next();

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

@@ -17,7 +17,7 @@
 package com.google.cloud.spanner;
 
 /** Streaming implementation of ResultSet that supports streaming of chunks */
-public interface StreamingResultSet extends ResultSet {
+interface StreamingResultSet extends ResultSet {
 
   /**
    * Returns the {@link boolean} for this {@link ResultSet}. This method will be used by

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

@@ -16,8 +16,8 @@
 
 package com.google.cloud.spanner;
 
-public class StreamingUtil {
-  public static boolean initiateStreaming(
+class StreamingUtil {
+  static boolean initiateStreaming(
       ResultSet resultSet, AsyncResultSet.StreamMessageListener streamMessageListener) {
     if (resultSet instanceof StreamingResultSet) {
       return ((StreamingResultSet) resultSet).initiateStreaming(streamMessageListener);