Addressed comments

Author: sakthivelmanii

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

@@ -230,17 +230,6 @@ interface ReadyCallback {
    * chunk is received from gRPC streaming call.
    */
   interface StreamMessageListener {
-    void onStreamMessage(
-        PartialResultSet partialResultSet,
-        boolean bufferIsFull,
-        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);
+    void onStreamMessage(PartialResultSet partialResultSet, boolean bufferIsFull);
   }
 }

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

@@ -637,10 +637,7 @@ public Struct getCurrentRowAsStruct() {
   }
 
   @Override
-  public void onStreamMessage(
-      PartialResultSet partialResultSet,
-      boolean bufferIsFull,
-      StreamMessageRequestor streamMessageRequestor) {
+  public void onStreamMessage(PartialResultSet partialResultSet, boolean bufferIsFull) {
     synchronized (monitor) {
       if (produceRowsInitiated) {
         return;
@@ -653,8 +650,6 @@ public void onStreamMessage(
               || partialResultSet == GrpcStreamIterator.END_OF_STREAM;
       if (startJobThread || state != State.STREAMING_INITIALIZED) {
         initiateProduceRows();
-      } else {
-        streamMessageRequestor.requestMessages(1);
       }
     }
   }

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

@@ -35,7 +35,7 @@
 /** Adapts a streaming read/query call into an iterator over partial result sets. */
 @VisibleForTesting
 class GrpcStreamIterator extends AbstractIterator<PartialResultSet>
-    implements CloseableIterator<PartialResultSet>, AsyncResultSet.StreamMessageRequestor {
+    implements CloseableIterator<PartialResultSet> {
   private static final Logger logger = Logger.getLogger(GrpcStreamIterator.class.getName());
   public static final PartialResultSet END_OF_STREAM = PartialResultSet.newBuilder().build();
   private AsyncResultSet.StreamMessageListener streamMessageListener;
@@ -146,11 +146,6 @@ private void addToStream(PartialResultSet results) {
     onStreamMessage(results);
   }
 
-  @Override
-  public void requestMessages(int numOfMessages) {
-    call.request(numOfMessages);
-  }
-
   private class ConsumerImpl implements SpannerRpc.ResultStreamConsumer {
     private final boolean cancelQueryWhenClientIsClosed;
 
@@ -199,7 +194,6 @@ 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()));
   }
 }

google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncResultSetImplTest.java

@@ -398,7 +398,7 @@ public void testCallbackIsNotCalledWhilePausedAndCanceled()
       when(delegate.initiateStreaming(any(AsyncResultSet.StreamMessageListener.class)))
           .thenAnswer(
               answer -> {
-                rs.onStreamMessage(PartialResultSet.newBuilder().build(), false, null);
+                rs.onStreamMessage(PartialResultSet.newBuilder().build(), false);
                 return null;
               });
       callbackResult =
@@ -517,26 +517,19 @@ public void testOnStreamMessageWhenResumeTokenIsPresent() {
     StreamingResultSet delegate = mock(StreamingResultSet.class);
     try (AsyncResultSetImpl rs =
         new AsyncResultSetImpl(mockedProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) {
-      AsyncResultSet.StreamMessageRequestor streamMessageRequestor =
-          Mockito.mock(AsyncResultSet.StreamMessageRequestor.class);
       // 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,
-          streamMessageRequestor);
-      Mockito.verify(streamMessageRequestor, times(1)).requestMessages(Mockito.eq(1));
+          PartialResultSet.newBuilder().addValues(Value.newBuilder().build()).build(), false);
 
       rs.onStreamMessage(
           PartialResultSet.newBuilder().setResumeToken(ByteString.copyFromUtf8("test")).build(),
-          false,
-          streamMessageRequestor);
+          false);
       Mockito.verify(mockedProvider.getExecutor(), times(2)).execute(Mockito.any());
-      Mockito.verify(streamMessageRequestor, times(1)).requestMessages(Mockito.eq(1));
     }
   }
 
@@ -545,20 +538,15 @@ public void testOnStreamMessageWhenCurrentBufferSizeReachedPrefetchChunkSize() {
     StreamingResultSet delegate = mock(StreamingResultSet.class);
     try (AsyncResultSetImpl rs =
         new AsyncResultSetImpl(mockedProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) {
-      AsyncResultSet.StreamMessageRequestor streamMessageRequestor =
-          Mockito.mock(AsyncResultSet.StreamMessageRequestor.class);
       // 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,
-          streamMessageRequestor);
+          PartialResultSet.newBuilder().addValues(Value.newBuilder().build()).build(), true);
       Mockito.verify(mockedProvider.getExecutor(), times(2)).execute(Mockito.any());
-      Mockito.verify(streamMessageRequestor, times(0)).requestMessages(Mockito.eq(1));
     }
   }
 
@@ -567,8 +555,6 @@ public void testOnStreamMessageWhenAsyncResultIsCancelled() {
     StreamingResultSet delegate = mock(StreamingResultSet.class);
     try (AsyncResultSetImpl rs =
         new AsyncResultSetImpl(mockedProvider, delegate, AsyncResultSetImpl.DEFAULT_BUFFER_SIZE)) {
-      AsyncResultSet.StreamMessageRequestor streamMessageRequestor =
-          Mockito.mock(AsyncResultSet.StreamMessageRequestor.class);
       // Marking Streaming as supported
       Mockito.when(
               delegate.initiateStreaming(Mockito.any(AsyncResultSet.StreamMessageListener.class)))
@@ -577,11 +563,8 @@ public void testOnStreamMessageWhenAsyncResultIsCancelled() {
       rs.setCallback(Executors.newSingleThreadExecutor(), ignored -> CallbackResponse.DONE);
       rs.cancel();
       rs.onStreamMessage(
-          PartialResultSet.newBuilder().addValues(Value.newBuilder().build()).build(),
-          false,
-          streamMessageRequestor);
+          PartialResultSet.newBuilder().addValues(Value.newBuilder().build()).build(), false);
       Mockito.verify(mockedProvider.getExecutor(), times(2)).execute(Mockito.any());
-      Mockito.verify(streamMessageRequestor, times(0)).requestMessages(Mockito.eq(1));
     }
   }
 }