perf: Skip gRPC trailers for StreamingRead & ExecuteStreamingSql

Author: sakthivelmanii

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

@@ -52,6 +52,7 @@ class GrpcStreamIterator extends AbstractIterator<PartialResultSet>
   private TimeUnit streamWaitTimeoutUnit;
   private long streamWaitTimeoutValue;
   private SpannerException error;
+  private volatile boolean done;
 
   @VisibleForTesting
   GrpcStreamIterator(int prefetchChunks, boolean cancelQueryWhenClientIsClosed) {
@@ -166,11 +167,17 @@ private class ConsumerImpl implements SpannerRpc.ResultStreamConsumer {
     @Override
     public void onPartialResultSet(PartialResultSet results) {
       addToStream(results);
+      if (results.getLast()) {
+        done = true;
+        addToStream(END_OF_STREAM);
+      }
     }
 
     @Override
     public void onCompleted() {
-      addToStream(END_OF_STREAM);
+      if (!done) {
+        addToStream(END_OF_STREAM);
+      }
     }
 
     @Override

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

@@ -1115,4 +1115,58 @@ public void getProtoEnumList() {
           resultSet.getProtoEnum(0, Genre::forNumber);
         });
   }
+
+  @Test
+  public void verifyResultSetWithLastTrue() {
+    long[] longArray = {111, 333, 444, 0, -1, -2234, Long.MAX_VALUE, Long.MIN_VALUE};
+
+    consumer.onPartialResultSet(
+        PartialResultSet.newBuilder()
+            .setMetadata(
+                makeMetadata(Type.struct(Type.StructField.of("f", Type.array(Type.int64())))))
+            .addValues(Value.int64Array(longArray).toProto())
+            .setLast(false)
+            .build());
+    assertThat(resultSet.next()).isTrue();
+    consumer.onPartialResultSet(
+        PartialResultSet.newBuilder()
+            .setMetadata(
+                makeMetadata(Type.struct(Type.StructField.of("f", Type.array(Type.int64())))))
+            .addValues(Value.int64Array(longArray).toProto())
+            .setLast(true)
+            .build());
+    assertThat(resultSet.next()).isTrue();
+    assertThat(resultSet.next()).isFalse();
+    consumer.onCompleted();
+  }
+
+  @Test
+  public void shouldThrowDeadlineExceededIfLastTrueIsNotReceived() {
+    long[] longArray = {111, 333, 444, 0, -1, -2234, Long.MAX_VALUE, Long.MIN_VALUE};
+
+    consumer.onPartialResultSet(
+        PartialResultSet.newBuilder()
+            .setMetadata(
+                makeMetadata(Type.struct(Type.StructField.of("f", Type.array(Type.int64())))))
+            .addValues(Value.int64Array(longArray).toProto())
+            .setLast(false)
+            .build());
+    assertThat(resultSet.next()).isTrue();
+    consumer.onPartialResultSet(
+        PartialResultSet.newBuilder()
+            .setMetadata(
+                makeMetadata(Type.struct(Type.StructField.of("f", Type.array(Type.int64())))))
+            .addValues(Value.int64Array(longArray).toProto())
+            .setLast(false)
+            .build());
+    assertThat(resultSet.next()).isTrue();
+    SpannerException spannerException =
+        assertThrows(
+            SpannerException.class,
+            () -> {
+              assertThat(resultSet.next()).isFalse();
+            });
+    assertEquals("DEADLINE_EXCEEDED: stream wait timeout", spannerException.getMessage());
+    consumer.onCompleted();
+  }
 }