verify fix with adding latency

Author: sakthivelmanii

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

@@ -516,6 +516,11 @@ private void initiateProduceRows() {
       }
       produceRowsInitiated = true;
     }
+    try {
+      Thread.sleep(5L);
+    } catch (InterruptedException e) {
+      throw new RuntimeException(e);
+    }
     this.service.execute(new ProduceRowsRunnable());
   }
 

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

@@ -44,7 +44,6 @@
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.Executor;
 import java.util.concurrent.TimeUnit;
-import java.util.concurrent.atomic.AtomicReference;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 import javax.annotation.Nullable;
@@ -71,8 +70,7 @@ abstract class ResumableStreamIterator extends AbstractIterator<PartialResultSet
   private final int maxBufferSize;
   private final ISpan span;
   private final TraceWrapper tracer;
-  private AtomicReference<CloseableIterator<PartialResultSet>> streamCache =
-      new AtomicReference<>();
+  private volatile CloseableIterator<PartialResultSet> stream;
   private ByteString resumeToken;
   private boolean finished;
   /**
@@ -215,16 +213,16 @@ boolean prepareIteratorForRetryOnDifferentGrpcChannel() {
 
   @Override
   public void close(@Nullable String message) {
-    if (streamCache.get() != null) {
-      streamCache.get().close(message);
+    if (stream != null) {
+      stream.close(message);
       span.end();
-      streamCache.set(null);
+      stream = null;
     }
   }
 
   @Override
   public boolean isWithBeginTransaction() {
-    return streamCache.get() != null && streamCache.get().isWithBeginTransaction();
+    return stream != null && stream.isWithBeginTransaction();
   }
 
   @Override
@@ -248,8 +246,8 @@ protected PartialResultSet computeNext() {
         return buffer.pop();
       }
       try {
-        if (streamCache.get().hasNext()) {
-          PartialResultSet next = streamCache.get().next();
+        if (stream.hasNext()) {
+          PartialResultSet next = stream.next();
           boolean hasResumeToken = !next.getResumeToken().isEmpty();
           if (hasResumeToken) {
             resumeToken = next.getResumeToken();
@@ -283,7 +281,7 @@ protected PartialResultSet computeNext() {
             buffer.removeLast();
           }
           assert buffer.isEmpty() || buffer.getLast().getResumeToken().equals(resumeToken);
-          streamCache.set(null);
+          stream = null;
           try (IScope s = tracer.withSpan(span)) {
             long delay = spannerException.getRetryDelayInMillis();
             if (delay != -1) {
@@ -304,7 +302,7 @@ protected PartialResultSet computeNext() {
           if (translated instanceof RetryOnDifferentGrpcChannelException) {
             if (++numAttemptsOnOtherChannel < errorHandler.getMaxAttempts()
                 && prepareIteratorForRetryOnDifferentGrpcChannel()) {
-              streamCache.set(null);
+              stream = null;
               continue;
             }
           }
@@ -325,7 +323,7 @@ private void startGrpcStreaming() {
     System.out.printf(
         "[%s][%s] Inside startGrpcStreaming\n",
         OffsetDateTime.now(), Thread.currentThread().getName());
-    if (streamCache.get() == null) {
+    if (stream == null) {
       span.addAnnotation(
           "Starting/Resuming stream",
           "ResumeToken",
@@ -336,7 +334,7 @@ private void startGrpcStreaming() {
         System.out.printf(
             "[%s][%s] Inside creating stream\n",
             OffsetDateTime.now(), Thread.currentThread().getName());
-        streamCache.set(checkNotNull(startStream(resumeToken, streamMessageListener)));
+        stream = checkNotNull(startStream(resumeToken, streamMessageListener));
       }
     }
     //    }