feat: add logical termination for RunQueryResponse (#956)

* feat: add logical termination for RunQueryResponse

* Add one test and comments

* Change code based on feedback

* 🦉 Updates from OwlBot post-processor

See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md

* add comments

* 🦉 Updates from OwlBot post-processor

See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md

Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>

google-cloud-firestore/src/main/java/com/google/cloud/firestore/Query.java

@@ -1330,6 +1330,8 @@ public void stream(@Nonnull final ApiStreamObserver<DocumentSnapshot> responseOb
 
     internalStream(
         new QuerySnapshotObserver() {
+          boolean hasCompleted = false;
+
           @Override
           public void onNext(QueryDocumentSnapshot documentSnapshot) {
             responseObserver.onNext(documentSnapshot);
@@ -1342,6 +1344,8 @@ public void onError(Throwable throwable) {
 
           @Override
           public void onCompleted() {
+            if (hasCompleted) return;
+            hasCompleted = true;
             responseObserver.onCompleted();
           }
         },
@@ -1537,6 +1541,16 @@ public void onResponse(RunQueryResponse response) {
             if (readTime == null) {
               readTime = Timestamp.fromProto(response.getReadTime());
             }
+
+            if (response.hasDone() && response.getDone()) {
+              Tracing.getTracer()
+                  .getCurrentSpan()
+                  .addAnnotation(
+                      "Firestore.Query: Completed",
+                      ImmutableMap.of(
+                          "numDocuments", AttributeValue.longAttributeValue(numDocuments)));
+              documentObserver.onCompleted(readTime);
+            }
           }
 
           @Override
@@ -1640,6 +1654,9 @@ ApiFuture<QuerySnapshot> get(@Nullable ByteString transactionId) {
     internalStream(
         new QuerySnapshotObserver() {
           final List<QueryDocumentSnapshot> documentSnapshots = new ArrayList<>();
+          // The stream's onCompleted could be called more than once,
+          // this flag makes sure only the first one is actually processed.
+          boolean hasCompleted = false;
 
           @Override
           public void onNext(QueryDocumentSnapshot documentSnapshot) {
@@ -1653,6 +1670,9 @@ public void onError(Throwable throwable) {
 
           @Override
           public void onCompleted() {
+            if (hasCompleted) return;
+            hasCompleted = true;
+
             // The results for limitToLast queries need to be flipped since we reversed the
             // ordering constraints before sending the query to the backend.
             List<QueryDocumentSnapshot> resultView =

google-cloud-firestore/src/test/java/com/google/cloud/firestore/LocalFirestoreHelper.java

@@ -311,6 +311,29 @@ public static Answer<RunQueryResponse> queryResponse(
     return streamingResponse(responses, throwable);
   }
 
+  /** Returns a stream of responses when RunQueryResponse.done set to true */
+  public static Answer<RunQueryResponse> queryResponseWithDone(
+      boolean callWithoutOnComplete, String... documentNames) {
+    RunQueryResponse[] responses = new RunQueryResponse[documentNames.length];
+
+    for (int i = 0; i < documentNames.length; ++i) {
+      final RunQueryResponse.Builder runQueryResponse = RunQueryResponse.newBuilder();
+      runQueryResponse.setDocument(
+          Document.newBuilder().setName(documentNames[i]).putAllFields(SINGLE_FIELD_PROTO));
+      runQueryResponse.setReadTime(
+          com.google.protobuf.Timestamp.newBuilder().setSeconds(1).setNanos(2));
+      if (i == (documentNames.length - 1)) {
+        runQueryResponse.setDone(true);
+      }
+      responses[i] = runQueryResponse.build();
+    }
+    if (callWithoutOnComplete) {
+      return streamingResponseWithoutOnComplete(responses);
+    } else {
+      return streamingResponse(responses, null);
+    }
+  }
+
   /** Returns a stream of responses followed by an optional exception. */
   public static <T> Answer<T> streamingResponse(
       final T[] response, @Nullable final Throwable throwable) {
@@ -328,6 +351,18 @@ public static <T> Answer<T> streamingResponse(
     };
   }
 
+  /** Returns a stream of responses even though onComplete() wasn't triggered */
+  public static <T> Answer<T> streamingResponseWithoutOnComplete(final T[] response) {
+    return invocation -> {
+      Object[] args = invocation.getArguments();
+      ResponseObserver<T> observer = (ResponseObserver<T>) args[1];
+      for (T resp : response) {
+        observer.onResponse(resp);
+      }
+      return null;
+    };
+  }
+
   public static ApiFuture<CommitResponse> commitResponse(int adds, int deletes) {
     CommitResponse.Builder commitResponse = CommitResponse.newBuilder();
     commitResponse.getCommitTimeBuilder().setSeconds(0).setNanos(0);

google-cloud-firestore/src/test/java/com/google/cloud/firestore/QueryTest.java

@@ -27,6 +27,7 @@
 import static com.google.cloud.firestore.LocalFirestoreHelper.order;
 import static com.google.cloud.firestore.LocalFirestoreHelper.query;
 import static com.google.cloud.firestore.LocalFirestoreHelper.queryResponse;
+import static com.google.cloud.firestore.LocalFirestoreHelper.queryResponseWithDone;
 import static com.google.cloud.firestore.LocalFirestoreHelper.reference;
 import static com.google.cloud.firestore.LocalFirestoreHelper.select;
 import static com.google.cloud.firestore.LocalFirestoreHelper.startAt;
@@ -952,6 +953,86 @@ public void onCompleted() {
     semaphore.acquire();
   }
 
+  @Test
+  public void successfulReturnWithoutOnComplete() throws Exception {
+    doAnswer(
+            queryResponseWithDone(
+                /* callWithoutOnComplete */ true, DOCUMENT_NAME + "1", DOCUMENT_NAME + "2"))
+        .when(firestoreMock)
+        .streamRequest(
+            runQuery.capture(),
+            streamObserverCapture.capture(),
+            Matchers.<ServerStreamingCallable>any());
+
+    final Semaphore semaphore = new Semaphore(0);
+    final Iterator<String> iterator = Arrays.asList("doc1", "doc2").iterator();
+
+    query.stream(
+        new ApiStreamObserver<DocumentSnapshot>() {
+          @Override
+          public void onNext(DocumentSnapshot documentSnapshot) {
+            assertEquals(iterator.next(), documentSnapshot.getId());
+          }
+
+          @Override
+          public void onError(Throwable throwable) {
+            fail();
+          }
+
+          @Override
+          public void onCompleted() {
+            semaphore.release();
+          }
+        });
+
+    semaphore.acquire();
+  }
+
+  @Test
+  /**
+   * onComplete() will be called twice. The first time is when it detects RunQueryResponse.done set
+   * to true. The second time is when it receives half close
+   */
+  public void successfulReturnCallsOnCompleteTwice() throws Exception {
+    doAnswer(
+            queryResponseWithDone(
+                /* callWithoutOnComplete */ false, DOCUMENT_NAME + "1", DOCUMENT_NAME + "2"))
+        .when(firestoreMock)
+        .streamRequest(
+            runQuery.capture(),
+            streamObserverCapture.capture(),
+            Matchers.<ServerStreamingCallable>any());
+
+    final Semaphore semaphore = new Semaphore(0);
+    final Iterator<String> iterator = Arrays.asList("doc1", "doc2").iterator();
+    final int[] counter = {0};
+
+    query.stream(
+        new ApiStreamObserver<DocumentSnapshot>() {
+          @Override
+          public void onNext(DocumentSnapshot documentSnapshot) {
+            assertEquals(iterator.next(), documentSnapshot.getId());
+          }
+
+          @Override
+          public void onError(Throwable throwable) {
+            fail();
+          }
+
+          @Override
+          public void onCompleted() {
+            counter[0]++;
+            semaphore.release();
+          }
+        });
+
+    semaphore.acquire();
+
+    // Wait for some time to see whether onCompleted() has been called more than once
+    Thread.sleep(200);
+    assertEquals(1, counter[0]);
+  }
+
   @Test
   public void retriesAfterRetryableError() throws Exception {
     final boolean[] returnError = new boolean[] {true};