fix

Author: sakthivelmanii

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

@@ -55,6 +55,7 @@
 import com.google.cloud.Tuple;
 import com.google.cloud.grpc.GrpcTransportOptions;
 import com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory;
+import com.google.cloud.spanner.AsyncResultSet.StreamMessageListener;
 import com.google.cloud.spanner.Options.QueryOption;
 import com.google.cloud.spanner.Options.ReadOption;
 import com.google.cloud.spanner.Options.TransactionOption;
@@ -266,6 +267,55 @@ private ResultSet wrap(final CachedResultSetSupplier resultSetSupplier) {
       return new ForwardingResultSet(resultSetSupplier) {
         private boolean beforeFirst = true;
 
+        @Override
+        public boolean initiateStreaming(StreamMessageListener streamMessageListener) {
+          while (true) {
+            try {
+              return internalInitiateStreaming(streamMessageListener);
+            } catch (SessionNotFoundException e) {
+              while (true) {
+                // Keep the replace-if-possible outside the try-block to let the exception bubble up
+                // if it's too late to replace the session.
+                replaceSessionIfPossible(e);
+                try {
+                  replaceDelegate(resultSetSupplier.reload());
+                  break;
+                } catch (SessionNotFoundException snfe) {
+                  e = snfe;
+                  // retry on yet another session.
+                }
+              }
+            }
+          }
+        }
+
+        private boolean internalInitiateStreaming(final StreamMessageListener streamMessageListener) {
+          try {
+            boolean ret = super.initiateStreaming(streamMessageListener);
+            if (beforeFirst) {
+              synchronized (lock) {
+                session.get().markUsed();
+                beforeFirst = false;
+                sessionUsedForQuery = true;
+              }
+            }
+            if (!ret && isSingleUse) {
+              close();
+            }
+            return ret;
+          } catch (SessionNotFoundException e) {
+            throw e;
+          } catch (SpannerException e) {
+            synchronized (lock) {
+              if (!closed && isSingleUse) {
+                session.get().setLastException(e);
+                AutoClosingReadContext.this.close();
+              }
+            }
+            throw e;
+          }
+        }
+
         @Override
         public boolean next() throws SpannerException {
           while (true) {