From c6448fc6411de01b3453e58bcbfca633e0849166 Mon Sep 17 00:00:00 2001
From: Ehsan <ehsann@google.com>
Date: Wed, 10 Jan 2024 15:01:39 -0800
Subject: [PATCH] feat: Logical termination for firestore.getAll(...). (#1517)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* feat: Logical termination for firestore.getAll(...).

* Using existing unit tests to verify the behaviour

* 🦉 Updates from OwlBot post-processor

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

* transaction unit tests are not mocking response correctly, so it cannot run with logical termination

---------

Co-authored-by: cherylEnkidu <cheryllin@google.com>
Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
---
 .../google/cloud/firestore/FirestoreImpl.java  | 10 ++++++++++
 .../google/cloud/firestore/FirestoreTest.java  |  8 ++++----
 .../cloud/firestore/LocalFirestoreHelper.java  | 18 +++++++++++++++++-
 3 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/FirestoreImpl.java b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/FirestoreImpl.java
index 38e037b33..4b54f7dcb 100644
--- a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/FirestoreImpl.java
+++ b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/FirestoreImpl.java
@@ -222,6 +222,7 @@ void getAll(
     ResponseObserver<BatchGetDocumentsResponse> responseObserver =
         new ResponseObserver<BatchGetDocumentsResponse>() {
           int numResponses;
+          boolean hasCompleted = false;
 
           @Override
           public void onStart(StreamController streamController) {}
@@ -265,6 +266,13 @@ public void onResponse(BatchGetDocumentsResponse response) {
                 return;
             }
             apiStreamObserver.onNext(documentSnapshot);
+
+            // Logical termination: if we have already received as many documents as we had
+            // requested, we can
+            // raise the results without waiting for the termination from the server.
+            if (numResponses == documentReferences.length) {
+              onComplete();
+            }
           }
 
           @Override
@@ -277,6 +285,8 @@ public void onError(Throwable throwable) {
 
           @Override
           public void onComplete() {
+            if (hasCompleted) return;
+            hasCompleted = true;
             tracer
                 .getCurrentSpan()
                 .addAnnotation(TraceUtil.SPAN_NAME_BATCHGETDOCUMENTS + ": Complete");
diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/FirestoreTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/FirestoreTest.java
index d4df6220b..638fa19ba 100644
--- a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/FirestoreTest.java
+++ b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/FirestoreTest.java
@@ -23,7 +23,7 @@
 import static com.google.cloud.firestore.LocalFirestoreHelper.arrayUnion;
 import static com.google.cloud.firestore.LocalFirestoreHelper.commit;
 import static com.google.cloud.firestore.LocalFirestoreHelper.commitResponse;
-import static com.google.cloud.firestore.LocalFirestoreHelper.getAllResponse;
+import static com.google.cloud.firestore.LocalFirestoreHelper.getAllResponseWithoutOnComplete;
 import static com.google.cloud.firestore.LocalFirestoreHelper.transform;
 import static com.google.cloud.firestore.LocalFirestoreHelper.update;
 import static org.junit.Assert.assertEquals;
@@ -80,7 +80,7 @@ public void encodeFieldPath() {
 
   @Test
   public void illegalFieldPath() throws Exception {
-    doAnswer(getAllResponse(SINGLE_FIELD_PROTO))
+    doAnswer(getAllResponseWithoutOnComplete(SINGLE_FIELD_PROTO))
         .when(firestoreMock)
         .streamRequest(
             getAllCapture.capture(),
@@ -110,7 +110,7 @@ public void exposesOptions() {
   @Test
   public void getAll() throws Exception {
     doAnswer(
-            getAllResponse(
+            getAllResponseWithoutOnComplete(
                 SINGLE_FIELD_PROTO, SINGLE_FIELD_PROTO, SINGLE_FIELD_PROTO, SINGLE_FIELD_PROTO))
         .when(firestoreMock)
         .streamRequest(
@@ -132,7 +132,7 @@ public void getAll() throws Exception {
 
   @Test
   public void getAllWithFieldMask() throws Exception {
-    doAnswer(getAllResponse(SINGLE_FIELD_PROTO))
+    doAnswer(getAllResponseWithoutOnComplete(SINGLE_FIELD_PROTO))
         .when(firestoreMock)
         .streamRequest(
             getAllCapture.capture(),
diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/LocalFirestoreHelper.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/LocalFirestoreHelper.java
index 9ee7f6ff5..a87188402 100644
--- a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/LocalFirestoreHelper.java
+++ b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/LocalFirestoreHelper.java
@@ -262,6 +262,16 @@ public static Map<String, Object> map() {
 
   public static Answer<BatchGetDocumentsResponse> getAllResponse(
       final Map<String, Value>... fields) {
+    return getAllResponseImpl(true, fields);
+  }
+
+  public static Answer<BatchGetDocumentsResponse> getAllResponseWithoutOnComplete(
+      final Map<String, Value>... fields) {
+    return getAllResponseImpl(false, fields);
+  }
+
+  public static Answer<BatchGetDocumentsResponse> getAllResponseImpl(
+      boolean withOnComplete, final Map<String, Value>... fields) {
     BatchGetDocumentsResponse[] responses = new BatchGetDocumentsResponse[fields.length];
 
     for (int i = 0; i < fields.length; ++i) {
@@ -281,7 +291,13 @@ public static Answer<BatchGetDocumentsResponse> getAllResponse(
       responses[i] = response.build();
     }
 
-    return streamingResponse(responses, null);
+    if (withOnComplete) {
+      return streamingResponse(responses, null);
+    } else {
+      // Verify with logical termination, the return of results no longer depends on calling
+      // OnComplete.
+      return streamingResponseWithoutOnComplete(responses);
+    }
   }
 
   public static ApiFuture<Empty> rollbackResponse() {